streaming-deep-rl 0.2.2__tar.gz → 0.2.5__tar.gz

{streaming_deep_rl-0.2.2 → streaming_deep_rl-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streaming-deep-rl
-Version: 0.2.2
+Version: 0.2.5
 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
@@ -36,6 +36,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: discrete-continuous-embed-readout>=0.2.0
 Requires-Dist: einops>=0.8.1
+Requires-Dist: einx>=0.3.0
 Requires-Dist: ema-pytorch
 Requires-Dist: hl-gauss-pytorch
 Requires-Dist: torch-einops-utils
@@ -228,3 +229,15 @@ $ uv run train_lunar.py
     url     = {https://arxiv.org/abs/2602.09396},
 }
 ```
+
+```bibtex
+@misc{maes2026leworldmodelstableendtoendjointembedding,
+    title   = {LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels}, 
+    author  = {Lucas Maes and Quentin Le Lidec and Damien Scieur and Yann LeCun and Randall Balestriero},
+    year    = {2026},
+    eprint  = {2603.19312},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2603.19312}, 
+}
+```

{streaming_deep_rl-0.2.2 → streaming_deep_rl-0.2.5}/README.md

@@ -179,3 +179,15 @@ $ uv run train_lunar.py
     url     = {https://arxiv.org/abs/2602.09396},
 }
 ```
+
+```bibtex
+@misc{maes2026leworldmodelstableendtoendjointembedding,
+    title   = {LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels}, 
+    author  = {Lucas Maes and Quentin Le Lidec and Damien Scieur and Yann LeCun and Randall Balestriero},
+    year    = {2026},
+    eprint  = {2603.19312},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2603.19312}, 
+}
+```

{streaming_deep_rl-0.2.2 → streaming_deep_rl-0.2.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "streaming-deep-rl"
-version = "0.2.2"
+version = "0.2.5"
 description = "Streaming Deep Reinforcement Learning"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -25,6 +25,7 @@ classifiers=[
 
 dependencies = [
     "discrete-continuous-embed-readout>=0.2.0",
+    "einx>=0.3.0",
     "einops>=0.8.1",
     "ema-pytorch",
     "hl-gauss-pytorch",

{streaming_deep_rl-0.2.2 → streaming_deep_rl-0.2.5}/streaming_deep_rl/streaming_deep_rl.py

@@ -9,16 +9,19 @@ import torch.nn.functional as F
 from torch import nn, tensor, is_tensor, cat, stack
 from torch.nn import Module, Linear, Sequential
 
-from einops import reduce
+import einx
+from einops import rearrange, reduce
 
 from discrete_continuous_embed_readout import Readout, Embed
 
-from torch_einops_utils import tree_map_tensor, pack_with_inverse
+from torch_einops_utils import tree_map_tensor, pack_with_inverse, masked_mean
 
 from ema_pytorch import EMA
 
 from hl_gauss_pytorch import HLGaussLayer
 
+from x_mlps_pytorch import MLP
+
 from streaming_deep_rl.buffer_dict import BufferDict
 
 # helpers
@@ -115,12 +118,15 @@ class SelfPredictRepr(Module):
         num_discrete_actions = 0,
         num_continuous_actions = 0,
         dim_predict = None,
-        target_embed_from_ema = True # whether the target embed is from an EMA, or from the model itself (traditional vs sigreg from lejepa)
+        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
     ):
         super().__init__()
         assert num_discrete_actions > 0 or num_continuous_actions > 0
 
         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)
 
@@ -128,12 +134,13 @@ class SelfPredictRepr(Module):
 
         self.to_projection = nn.Sequential(
             nn.RMSNorm(dim_embed),
-            nn.Linear(dim_embed, dim_predict, bias = False)
+            MLP(dim_embed, dim_mlp_hidden, dim_predict)
         )
 
-        self.to_prediction = nn.Sequential(
-            nn.Linear(dim_predict + dim_embed, dim_predict, bias = False)
-        )
+        self.to_prediction = MLP(dim_predict + dim_embed, dim_mlp_hidden, dim_predict)
+
+        self.apply_sigreg = sigreg_weight > 0.
+        self.sigreg_weight = sigreg_weight
 
     def forward(
         self,
@@ -155,7 +162,54 @@ class SelfPredictRepr(Module):
 
         # cosine sim loss
 
-        return F.mse_loss(l2norm(predicted), l2norm(next_projected))
+        loss = F.mse_loss(l2norm(predicted), l2norm(next_projected))
+
+        if not self.apply_sigreg:
+            return loss
+
+        return loss + self.sigreg_weight * sigreg_loss(projected)
+
+# sigreg from lejepa
+
+def sigreg_loss(
+    x,
+    num_slices = 1024,
+    domain = (-5, 5),
+    num_knots = 17,
+    mask = None
+):
+    # Randall Balestriero - https://arxiv.org/abs/2511.08544
+
+    dim, device = x.shape[-1], x.device
+
+    # slice sampling
+
+    rand_projs = torch.randn((num_slices, dim), device = device)
+    rand_projs = l2norm(rand_projs)
+
+    # integration points
+
+    t = torch.linspace(*domain, num_knots, device = device)
+
+    # theoretical CF for N(0, 1) and Gauss. window
+
+    exp_f = (-0.5 * t.square()).exp()
+
+    # empirical CF
+
+    x_t = einx.dot('... d, m d -> (...) m', x, rand_projs)
+
+    x_t = einx.multiply('n m, t -> n m t', x_t, t)
+    ecf = (1j * x_t).exp()
+
+    mask_1d = rearrange(mask, '... -> (...)') if exists(mask) else None
+    ecf = masked_mean(ecf, mask_1d, dim = 0)
+
+    # weighted L2 distance
+
+    err = ecf.sub(exp_f).abs().square().mul(exp_f)
+
+    return torch.trapz(err, t, dim = -1).mean()
 
 # online normalization from Welford in 1962
 
@@ -332,6 +386,9 @@ class StreamingACLambda(Module):
         critic_ema_beta = 0.7,
         actor_self_predict_repr = False,
         critic_self_predict_repr = False,
+        spr_target_embed_from_ema = True,
+        spr_sigreg_weight = 0.,
+        spr_dim_hidden_expand_factor = 4,
         entropy_weight = 0.01,
         init_sparsity = 0.9,
         dim_critic = 128,
@@ -406,22 +463,30 @@ class StreamingACLambda(Module):
 
         if actor_self_predict_repr:
 
-            actor_use_ema |= True
+            if spr_target_embed_from_ema:
+                actor_use_ema |= True
 
             self.actor_spr = SelfPredictRepr(
                 dim_actor,
                 num_discrete_actions = num_discrete_actions,
                 num_continuous_actions = num_continuous_actions,
+                target_embed_from_ema = spr_target_embed_from_ema,
+                sigreg_weight = spr_sigreg_weight,
+                dim_hidden_expand_factor = spr_dim_hidden_expand_factor
             )
 
         if critic_self_predict_repr:
 
-            use_critic_ema |= True
+            if spr_target_embed_from_ema:
+                use_critic_ema |= True
 
             self.critic_spr = SelfPredictRepr(
                 dim_critic,
                 num_discrete_actions = num_discrete_actions,
                 num_continuous_actions = num_continuous_actions,
+                target_embed_from_ema = spr_target_embed_from_ema,
+                sigreg_weight = spr_sigreg_weight,
+                dim_hidden_expand_factor = spr_dim_hidden_expand_factor
             )
 
         # actor