dreamer4 0.1.4__py3-none-any.whl → 0.1.10__py3-none-any.whl

dreamer4/dreamer4.py

@@ -14,7 +14,7 @@ from torch.nested import nested_tensor
 from torch.distributions import Normal, kl
 from torch.nn import Module, ModuleList, Embedding, Parameter, Sequential, Linear, RMSNorm, Identity
 from torch import nn, cat, stack, arange, tensor, Tensor, is_tensor, full, zeros, ones, randint, rand, randn, randn_like, empty, full, linspace, arange
-from torch.utils._pytree import tree_flatten, tree_unflatten
+from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten
 
 import torchvision
 from torchvision.models import VGG16_Weights
@@ -27,6 +27,8 @@ from x_mlps_pytorch.normed_mlp import create_mlp
 
 from hyper_connections import get_init_and_expand_reduce_stream_functions
 
+from vit_pytorch.vit_with_decorr import DecorrelationLoss
+
 from assoc_scan import AssocScan
 
 # ein related
@@ -68,10 +70,14 @@ except ImportError:
 
 LinearNoBias = partial(Linear, bias = False)
 
-TokenizerLosses = namedtuple('TokenizerLosses', ('recon', 'lpips'))
+TokenizerLosses = namedtuple('TokenizerLosses', ('recon', 'lpips', 'time_decorr', 'space_decorr'))
 
 WorldModelLosses = namedtuple('WorldModelLosses', ('flow', 'rewards', 'discrete_actions', 'continuous_actions'))
 
+AttentionIntermediates = namedtuple('AttentionIntermediates', ('next_kv_cache', 'normed_inputs'))
+
+TransformerIntermediates = namedtuple('TransformerIntermediates', ('next_kv_cache', 'normed_time_inputs', 'normed_space_inputs'))
+
 MaybeTensor = Tensor | None
 
 @dataclass
@@ -91,6 +97,14 @@ class Experience:
     agent_index: int = 0
     is_from_world_model: bool = True
 
+    def cpu(self):
+        return self.to(torch.device('cpu'))
+
+    def to(self, device):
+        experience_dict = asdict(self)
+        experience_dict = tree_map(lambda t: t.to(device) if is_tensor(t) else t, experience_dict)
+        return Experience(**experience_dict)
+
 def combine_experiences(
     exps: list[Experiences]
 ) -> Experience:
@@ -1313,7 +1327,7 @@ class Attention(Module):
         self,
         tokens, # (b n d)
         kv_cache = None,
-        return_kv_cache = False,
+        return_intermediates = False,
         rotary_pos_emb = None,
         attend_fn: Callable | None = None
     ):
@@ -1367,10 +1381,10 @@ class Attention(Module):
 
         out = inverse_packed_batch(out)
 
-        if not return_kv_cache:
+        if not return_intermediates:
             return out
 
-        return out, stack((k, v))
+        return out, AttentionIntermediates(stack((k, v)), tokens)
 
 # feedforward
 
@@ -1484,7 +1498,7 @@ class AxialSpaceTimeTransformer(Module):
         self,
         tokens,                          # (b t s d)
         kv_cache: Tensor | None = None,  # (y 2 b h t d)
-        return_kv_cache = False
+        return_intermediates = False
 
     ): # (b t s d) | (y 2 b h t d)
 
@@ -1525,6 +1539,11 @@ class AxialSpaceTimeTransformer(Module):
 
         rotary_pos_emb = self.time_rotary(rotary_seq_len, offset = rotary_pos_offset)
 
+        # normed attention inputs
+
+        normed_time_attn_inputs = []
+        normed_space_attn_inputs = []
+
         # attention
 
         tokens = self.expand_streams(tokens)
@@ -1545,12 +1564,12 @@ class AxialSpaceTimeTransformer(Module):
 
             # attention layer
 
-            tokens, next_kv_cache = attn(
+            tokens, attn_intermediates = attn(
                 tokens,
                 rotary_pos_emb = layer_rotary_pos_emb,
                 attend_fn = attend_fn,
                 kv_cache = maybe_kv_cache,
-                return_kv_cache = True
+                return_intermediates = True
             )
 
             tokens = post_attn_rearrange(tokens)
@@ -1562,7 +1581,13 @@ class AxialSpaceTimeTransformer(Module):
             # save kv cache if is time layer
 
             if layer_is_time:
-                time_attn_kv_caches.append(next_kv_cache)
+                time_attn_kv_caches.append(attn_intermediates.next_kv_cache)
+
+            # save time attention inputs for decorr
+
+            space_or_time_inputs = normed_time_attn_inputs if layer_is_time else normed_space_attn_inputs
+
+            space_or_time_inputs.append(attn_intermediates.normed_inputs)
 
         tokens = self.reduce_streams(tokens)
 
@@ -1572,10 +1597,16 @@ class AxialSpaceTimeTransformer(Module):
             # just concat the past tokens back on for now, todo - clean up the logic
             out = cat((past_tokens, out), dim = 1)
 
-        if not return_kv_cache:
+        if not return_intermediates:
             return out
 
-        return out, stack(time_attn_kv_caches)
+        intermediates = TransformerIntermediates(
+            stack(time_attn_kv_caches),
+            stack(normed_time_attn_inputs),
+            stack(normed_space_attn_inputs)
+        )
+
+        return out, intermediates
 
 # video tokenizer
 
@@ -1601,12 +1632,15 @@ class VideoTokenizer(Module):
         per_image_patch_mask_prob = (0., 0.9), # probability of patch masking appears to be per image probabilities drawn uniformly between 0. and 0.9 - if you are a phd student and think i'm mistakened, please open an issue
         lpips_loss_network: Module | None = None,
         lpips_loss_weight = 0.2,
+        encoder_add_decor_aux_loss = False,
+        decor_auxx_loss_weight = 0.1,
+        decorr_sample_frac = 0.25,
         nd_rotary_kwargs: dict = dict(
             rope_min_freq = 1.,
             rope_max_freq = 10000.,
             rope_p_zero_freqs = 0.
         ),
-        num_residual_streams = 1
+        num_residual_streams = 1,
     ):
         super().__init__()
 
@@ -1701,6 +1735,14 @@ class VideoTokenizer(Module):
         if self.has_lpips_loss:
             self.lpips = LPIPSLoss(lpips_loss_network)
 
+        # decorr aux loss
+        # https://arxiv.org/abs/2510.14657
+
+        self.encoder_add_decor_aux_loss = encoder_add_decor_aux_loss
+        self.decorr_aux_loss_weight = decor_auxx_loss_weight
+
+        self.decorr_loss = DecorrelationLoss(decorr_sample_frac, soft_validate_num_sampled = True) if encoder_add_decor_aux_loss else None
+
     @property
     def device(self):
         return self.zero.device
@@ -1814,7 +1856,7 @@ class VideoTokenizer(Module):
 
         # encoder attention
 
-        tokens = self.encoder_transformer(tokens)
+        tokens, (_, time_attn_normed_inputs, space_attn_normed_inputs) = self.encoder_transformer(tokens, return_intermediates = True)
 
         # latent bottleneck
 
@@ -1836,17 +1878,25 @@ class VideoTokenizer(Module):
         if self.has_lpips_loss:
             lpips_loss = self.lpips(video, recon_video)
 
+        time_decorr_loss = space_decorr_loss = self.zero
+
+        if self.encoder_add_decor_aux_loss:
+            time_decorr_loss = self.decorr_loss(time_attn_normed_inputs)
+            space_decorr_loss = self.decorr_loss(space_attn_normed_inputs)
+
         # losses
 
         total_loss = (
             recon_loss +
-            lpips_loss * self.lpips_loss_weight
+            lpips_loss * self.lpips_loss_weight +
+            time_decorr_loss * self.decorr_aux_loss_weight +
+            space_decorr_loss * self.decorr_aux_loss_weight
         )
 
         if not return_all_losses:
             return total_loss
 
-        losses = (recon_loss, lpips_loss)
+        losses = (recon_loss, lpips_loss, decorr_loss)
 
         return total_loss, TokenizerLosses(*losses)
 
@@ -2435,6 +2485,8 @@ class DynamicsWorldModel(Module):
     ):
         assert isinstance(experience, Experience)
 
+        experience = experience.to(self.device)
+
         latents = experience.latents
         actions = experience.actions
         old_log_probs = experience.log_probs
@@ -3325,7 +3377,7 @@ class DynamicsWorldModel(Module):
 
             # attention
 
-            tokens, next_time_kv_cache = self.transformer(tokens, kv_cache = time_kv_cache, return_kv_cache = True)
+            tokens, (next_time_kv_cache, *_) = self.transformer(tokens, kv_cache = time_kv_cache, return_intermediates = True)
 
             # unpack
 

dreamer4/trainers.py

@@ -528,7 +528,7 @@ class SimTrainer(Module):
 
                 total_experience += num_experience
 
-                experiences.append(experience)
+                experiences.append(experience.cpu())
 
             combined_experiences = combine_experiences(experiences)
 

{dreamer4-0.1.4.dist-info → dreamer4-0.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreamer4
-Version: 0.1.4
+Version: 0.1.10
 Summary: Dreamer 4
 Project-URL: Homepage, https://pypi.org/project/dreamer4/
 Project-URL: Repository, https://github.com/lucidrains/dreamer4
@@ -44,6 +44,7 @@ Requires-Dist: hl-gauss-pytorch
 Requires-Dist: hyper-connections>=0.2.1
 Requires-Dist: torch>=2.4
 Requires-Dist: torchvision
+Requires-Dist: vit-pytorch>=1.15.3
 Requires-Dist: x-mlps-pytorch>=0.0.29
 Provides-Extra: examples
 Provides-Extra: test
@@ -57,7 +58,7 @@ Description-Content-Type: text/markdown
 
 Implementation of Danijar's [latest iteration](https://arxiv.org/abs/2509.24527v1) for his [Dreamer](https://danijar.com/project/dreamer4/) line of work
 
-[Discord channel](https://discord.gg/ab4BEk3W) for collaborating with other researchers interested in this work
+[Discord channel](https://discord.gg/PmGR7KRwxq) for collaborating with other researchers interested in this work
 
 ## Appreciation
 
@@ -90,7 +91,7 @@ video = torch.randn(2, 3, 10, 256, 256)
 # learn the tokenizer
 
 loss = tokenizer(video)
-loss.backward() # ler
+loss.backward()
 
 # dynamics world model
 

dreamer4-0.1.10.dist-info/RECORD

@@ -0,0 +1,8 @@
+dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
+dreamer4/dreamer4.py,sha256=_xr_XJGfqhCabRV0vnue4zypHZ4kXeUDZp1N6RF2AoY,122988
+dreamer4/mocks.py,sha256=TfqOB_Gq6N_GggBYwa6ZAJQx38ntlYbXZe23Ne4jshw,2502
+dreamer4/trainers.py,sha256=h_BMi-P2QMVi-IWQCkejPmyA0UzHgKtE1n7Qn1-IrxE,15093
+dreamer4-0.1.10.dist-info/METADATA,sha256=oTK9b_fWDCQTC89Y30OBY_2BzJJ6ih25BzgO0D-SApg,4973
+dreamer4-0.1.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dreamer4-0.1.10.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+dreamer4-0.1.10.dist-info/RECORD,,

dreamer4-0.1.4.dist-info/RECORD

@@ -1,8 +0,0 @@
-dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
-dreamer4/dreamer4.py,sha256=ghestMgz7B1oEqBRR0XkkdWe0kkh7bshhzmi6-n-XIs,120790
-dreamer4/mocks.py,sha256=TfqOB_Gq6N_GggBYwa6ZAJQx38ntlYbXZe23Ne4jshw,2502
-dreamer4/trainers.py,sha256=JsnJwQJcbI_75KBTNddG6b7QVkO6LD1N_HQiVe-VnCM,15087
-dreamer4-0.1.4.dist-info/METADATA,sha256=GkzuqKtNJJCSh5FycWJOr49253_w926biJkSz9ic4TQ,4941
-dreamer4-0.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dreamer4-0.1.4.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-dreamer4-0.1.4.dist-info/RECORD,,