streaming-deep-rl 0.3.0__tar.gz → 0.3.2__tar.gz

{streaming_deep_rl-0.3.0 → streaming_deep_rl-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streaming-deep-rl
-Version: 0.3.0
+Version: 0.3.2
 Summary: Streaming Deep Reinforcement Learning
 Project-URL: Homepage, https://pypi.org/project/streaming-deep-rl/
 Project-URL: Repository, https://codeberg.org/lucidrains/streaming-deep-rl

{streaming_deep_rl-0.3.0 → streaming_deep_rl-0.3.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "streaming-deep-rl"
-version = "0.3.0"
+version = "0.3.2"
 description = "Streaming Deep Reinforcement Learning"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{streaming_deep_rl-0.3.0 → streaming_deep_rl-0.3.2}/streaming_deep_rl/streaming_deep_rl.py

@@ -5,8 +5,9 @@ from typing import NamedTuple, Literal
 
 import torch
 import torch.nn.functional as F
-from torch import nn, tensor, is_tensor, cat, stack
+from torch import nn, tensor, is_tensor, cat, Tensor
 from torch.nn import Module, Linear, Sequential
+from torch.nn.utils import parameters_to_vector
 
 import einx
 from einops import rearrange, reduce
@@ -14,7 +15,7 @@ from einops.layers.torch import Rearrange
 
 from discrete_continuous_embed_readout import Readout, Embed
 
-from torch_einops_utils import tree_map_tensor, pack_with_inverse, masked_mean
+from torch_einops_utils import tree_map_tensor, masked_mean
 
 from ema_pytorch import EMA
 
@@ -59,6 +60,10 @@ def l2norm(t):
 
 # return types
 
+class SSLLosses(NamedTuple):
+    total: Tensor
+    recon: Tensor
+
 class UpdateMetrics(NamedTuple):
     td_error: float
     value_pred: float
@@ -68,6 +73,7 @@ class UpdateMetrics(NamedTuple):
     critic_trace_norm: float
     actor_scale: float
     critic_scale: float
+    chunked_td_lambda: float
 
 # initialization
 
@@ -102,23 +108,35 @@ def sparse_init_(
 # orthogonal projection for SPR in streaming setting
 # https://arxiv.org/abs/2602.09396
 
-def orthog_project(x, y):
-    x, inverse_pack = pack_with_inverse(x, '*')
-    y, _ = pack_with_inverse(y, '*')
+def orthog_project_tensors(
+    grads: list[Tensor],
+    moms: list[Tensor]
+):
+    if len(grads) == 0:
+        return []
+
+    flat_grad = parameters_to_vector(grads)
+    flat_mom = parameters_to_vector(moms)
 
-    dtype = x.dtype
+    dtype = flat_grad.dtype
 
-    if x.device.type != 'mps':
-        x, y = x.double(), y.double()
+    if flat_grad.device.type != 'mps':
+        flat_grad, flat_mom = flat_grad.double(), flat_mom.double()
 
-    unit = l2norm(y)
+    unit = l2norm(flat_mom)
+    orthog = flat_grad - (flat_grad * unit).sum(dim = -1, keepdim = True) * unit
 
-    parallel = (x * unit).sum(dim = -1, keepdim = True) * unit
-    orthog = x - parallel
+    orthog = orthog.to(dtype)
 
-    orthog = inverse_pack(orthog, '*')
+    orthog_grads = []
+    offset = 0
 
-    return orthog.to(dtype)
+    for g in grads:
+        numel = g.numel()
+        orthog_grads.append(orthog[offset:offset + numel].view_as(g))
+        offset += numel
+
+    return orthog_grads
 
 # Simplicial Embeddings
 # Lavoie et al - https://arxiv.org/abs/2204.00616
@@ -127,6 +145,8 @@ class SEM(Module):
     def __init__(
         self,
         dim,
+        dim_in = None,
+        project_in = None,
         temperature = 0.1,
         dim_simplex = 8,
         pre_layernorm = False
@@ -134,6 +154,11 @@ class SEM(Module):
         super().__init__()
         assert divisible_by(dim, dim_simplex), f'{dim} must be divisible by {dim_simplex}'
 
+        dim_in = default(dim_in, dim)
+        project_in = default(project_in, dim_in != dim)
+
+        self.embedder = Linear(dim_in, dim, bias = False) if project_in else nn.Identity()
+
         self.dim = dim
         self.dim_simplex = dim_simplex
         self.temperature = temperature
@@ -144,6 +169,7 @@ class SEM(Module):
         self,
         t
     ):
+        t = self.embedder(t)
         t = self.norm(t)
         t = rearrange(t, '... (l v) -> ... l v', v = self.dim_simplex)
         t = (t / self.temperature).softmax(dim = -1)
@@ -161,7 +187,7 @@ class SelfPredictRepr(Module):
         target_embed_from_ema = True, # whether the target embed is from an EMA, or from the model itself (traditional vs sigreg from lejepa)
         sigreg_weight = 0.,
         dim_hidden_expand_factor = 4,
-        use_sem = False,
+        use_sem = True,
         sem_dim_simplex = 8,
         sem_temperature = 0.1
     ):
@@ -171,22 +197,26 @@ class SelfPredictRepr(Module):
         dim_predict = default(dim_predict, dim_embed)
         dim_mlp_hidden = int(dim_embed * dim_hidden_expand_factor)
 
-        self.to_action_embed = Embed(dim_embed, num_discrete = num_discrete_actions, num_continuous = num_continuous_actions)
+        # when SEM is active, projects dim_embed -> dim_predict before the projector
+
+        dim_projector_in = dim_predict if use_sem else dim_embed
+
+        self.to_action_embed = Embed(dim_projector_in, num_discrete = num_discrete_actions, num_continuous = num_continuous_actions)
 
         self.target_embed_from_ema = target_embed_from_ema
 
         self.to_projection = nn.Sequential(
-            nn.RMSNorm(dim_embed),
-            MLP(dim_embed, dim_mlp_hidden, dim_predict)
+            nn.RMSNorm(dim_projector_in),
+            MLP(dim_projector_in, dim_mlp_hidden, dim_predict)
         )
 
-        self.to_prediction = MLP(dim_predict + dim_embed, dim_mlp_hidden, dim_predict)
+        self.to_prediction = MLP(dim_predict + dim_projector_in, dim_mlp_hidden, dim_predict)
 
         self.apply_sigreg = sigreg_weight > 0.
         self.sigreg_weight = sigreg_weight
 
         self.use_sem = use_sem
-        self.sem = SEM(dim_predict, temperature = sem_temperature, dim_simplex = sem_dim_simplex) if use_sem else nn.Identity()
+        self.sem = SEM(dim_predict, dim_in = dim_embed, project_in = True, temperature = sem_temperature, dim_simplex = sem_dim_simplex) if use_sem else nn.Identity()
 
     def forward(
         self,
@@ -195,6 +225,10 @@ class SelfPredictRepr(Module):
         next_state_embed,
         ema_to_projection = None
     ):
+        if self.use_sem:
+            state_embed = self.sem(state_embed)
+            next_state_embed = self.sem(next_state_embed)
+
         projected = self.to_projection(state_embed)
 
         # maybe no
@@ -207,22 +241,20 @@ class SelfPredictRepr(Module):
 
         next_projected = decorator(target_to_projection)(next_state_embed)
 
-        if self.use_sem:
-            projected = self.sem(projected)
-            next_projected = self.sem(next_projected)
-
         action_embed = self.to_action_embed(actions)
 
         predicted = self.to_prediction(cat((projected, action_embed), dim = -1))
 
         # cosine sim loss
 
-        loss = F.mse_loss(l2norm(predicted), l2norm(next_projected))
+        recon_loss = F.mse_loss(l2norm(predicted), l2norm(next_projected))
+
+        total_loss = recon_loss
 
-        if not self.apply_sigreg:
-            return loss
+        if self.apply_sigreg:
+            total_loss = total_loss + self.sigreg_weight * sigreg(projected)
 
-        return loss + self.sigreg_weight * sigreg(projected)
+        return SSLLosses(total = total_loss, recon = recon_loss)
 
 # Dominik Schmidt https://arxiv.org/abs/2312.10812
 
@@ -313,7 +345,7 @@ class LAPO(Module):
 
         # lapo actually does a decoder back to state in pixel space, although probably not crucial given SPR
 
-        return recon_loss + pred_action_loss
+        return SSLLosses(total = recon_loss + pred_action_loss, recon = recon_loss)
 
 # sigreg from lejepa
 
@@ -550,7 +582,8 @@ class StreamingACLambda(Module):
         use_hl_gauss = False,
         hl_gauss_val_min = -3.,
         hl_gauss_val_max = 3.,
-        hl_gauss_num_bins = 64
+        hl_gauss_num_bins = 64,
+        use_chunked_td = False
     ):
         super().__init__()
 
@@ -559,6 +592,11 @@ class StreamingACLambda(Module):
         assert not (use_minto and not use_critic_ema), 'minto can only be used with critic ema'
         self.use_minto = use_minto
 
+        # chunked td
+
+        assert not (use_chunked_td and ssl_type is None), 'chunked td requires a predictive model (ssl_type)'
+        self.use_chunked_td = use_chunked_td
+
         # delightful pg
 
         self.use_delightful_pg = use_delightful_pg
@@ -846,7 +884,8 @@ class StreamingACLambda(Module):
 
             # actor with entropy regularization
 
-            action_logits = self.actor_with_readout(state)
+            actor_embed = self.actor(state)
+            action_logits = self.readout(actor_embed)
             log_prob = self.readout.log_prob(action_logits, action).mean()
             entropy = self.readout.entropy(action_logits).mean()
 
@@ -884,12 +923,12 @@ class StreamingACLambda(Module):
 
                 # loss
 
-                spr_loss = spr_module(embed, action, next_embed, ema_to_projection = ema_to_projection)
+                ssl_losses = spr_module(embed, action, next_embed, ema_to_projection = ema_to_projection)
 
-                # gradients
+                # gradients (use total loss for backprop)
 
                 spr_params = list(spr_module.parameters())
-                spr_grads_tup = torch.autograd.grad(spr_loss, network_params + spr_params, retain_graph = True, allow_unused = True)
+                spr_grads_tup = torch.autograd.grad(ssl_losses.total, network_params + spr_params, retain_graph = True, allow_unused = True)
                 spr_grads_tup = tuple(default(g, torch.zeros_like(p)) for g, p in zip(spr_grads_tup, network_params + spr_params))
 
                 spr_grad = {name: grad for (name, _), grad in zip(network_with_readout.named_parameters(), spr_grads_tup[:len(network_params)])}
@@ -898,35 +937,41 @@ class StreamingACLambda(Module):
                 # orth¹: project SPR encoder grads orthogonal to their own EMA momentum
                 # this decorrelates the highly correlated streaming data
 
-                names = list(spr_grad.keys())
-                all_grads = [spr_grad[n] for n in names] + spr_only_grad
-                all_moms = [spr_momentum[n] for n in names] + list(spr_aux_momentum.values())
+                enc_names = [n for n in spr_grad.keys() if n.startswith('0.')]
+                enc_grads = [spr_grad[n] for n in enc_names]
+                enc_moms = [spr_momentum[n] for n in enc_names]
+
+                enc_grads_proj = orthog_project_tensors(enc_grads, enc_moms)
 
-                all_grads_proj = [orthog_project(g, m) for g, m in zip(all_grads, all_moms)]
+                aux_grads = spr_only_grad
+                aux_moms = list(spr_aux_momentum.values())
+                aux_grads_proj = orthog_project_tensors(aux_grads, aux_moms)
 
                 # momentum update
 
-                for grad, mom in zip(all_grads, all_moms):
+                for grad, mom in zip([*enc_grads, *aux_grads], [*enc_moms, *aux_moms]):
                     mom.mul_(self.spr_orth_beta).add_(grad, alpha = 1. - self.spr_orth_beta)
 
-                for name, proj_grad in zip(names, all_grads_proj[:len(names)]):
+                for name, proj_grad in zip(enc_names, enc_grads_proj):
                     spr_grad[name] = proj_grad
 
-                spr_only_grad = all_grads_proj[len(names):]
+                spr_only_grad = aux_grads_proj
 
                 # update SPR auxiliary networks via SGD with dedicated spr_lr
 
                 for param, grad in zip(spr_params, spr_only_grad):
                     param.data.add_(grad, alpha = -self.spr_lr)
 
-                return spr_grad
+                # return recon-only loss for chunked-td (not contaminated by sigreg / action loss)
+
+                return spr_grad, ssl_losses.recon.detach()
 
             # get the SPR loss
 
             actor_spr_grad = None
+            actor_spr_loss = None
             if exists(self.ssl_type):
-                actor_embed = self.actor(state)
-                actor_spr_grad = get_spr_grad(
+                actor_spr_grad, actor_spr_loss = get_spr_grad(
                     actor_embed,
                     self.actor,
                     self.actor_with_readout,
@@ -941,8 +986,9 @@ class StreamingACLambda(Module):
             # critic SPR loss
 
             critic_spr_grad = None
+            critic_spr_loss = None
             if exists(self.ssl_type):
-                critic_spr_grad = get_spr_grad(
+                critic_spr_grad, critic_spr_loss = get_spr_grad(
                     critic_embed,
                     self.critic,
                     self.critic_full,
@@ -970,6 +1016,12 @@ class StreamingACLambda(Module):
 
         decay = self.eligibility_trace_decay * self.discount_factor
 
+        chunked_td_lambda = 1.
+
+        if self.use_chunked_td and exists(critic_spr_loss):
+            chunked_td_lambda = max(0., 1. - critic_spr_loss.item() / 2.)
+            decay = decay * chunked_td_lambda
+
         for name, trace in self.actor_trace.items():
             trace.mul_(decay).add_(actor_grad[name] * gate)
 
@@ -1018,12 +1070,24 @@ class StreamingACLambda(Module):
 
             global_scale = (kappa * td_error_factor * trace_sum).reciprocal().clamp(max = 1.)
 
+            rl_updates = {name: adapted_grads[name] * lr for name, _ in params.named_parameters()}
+
+            if exists(spr_grads):
+                enc_names = [n for n in spr_grads.keys() if n.startswith('0.')]
+                enc_spr_updates = [-spr_grads[n] * self.spr_lr_param_update for n in enc_names]
+                enc_rl_updates = [rl_updates[n] for n in enc_names]
+                
+                enc_spr_updates_proj = orthog_project_tensors(enc_spr_updates, enc_rl_updates)
+                
+                spr_updates_proj = {n: g for n, g in zip(enc_names, enc_spr_updates_proj)}
+            else:
+                spr_updates_proj = {}
+
             for name, param in params.named_parameters():
-                update = adapted_grads[name] * lr
+                update = rl_updates[name]
 
-                if exists(spr_grads) and name in spr_grads:
-                    spr_update = -spr_grads[name] * self.spr_lr_param_update
-                    update = update + orthog_project(spr_update, update)
+                if name in spr_updates_proj:
+                    update = update + spr_updates_proj[name]
 
                 update = update * global_scale
 
@@ -1109,7 +1173,8 @@ class StreamingACLambda(Module):
             actor_trace_norm = actor_trace_norm.item(),
             critic_trace_norm = critic_trace_norm.item(),
             actor_scale = actor_scale.item(),
-            critic_scale = critic_scale.item()
+            critic_scale = critic_scale.item(),
+            chunked_td_lambda = chunked_td_lambda
         )
 
     @torch.no_grad()