dreamer4 0.0.102__py3-none-any.whl → 0.1.1__py3-none-any.whl

dreamer4/dreamer4.py

@@ -1331,6 +1331,12 @@ class Attention(Module):
         q = self.q_heads_rmsnorm(q)
         k = self.k_heads_rmsnorm(k)
 
+        # rotary
+
+        if exists(rotary_pos_emb):
+            q = apply_rotations(rotary_pos_emb, q)
+            k = apply_rotations(rotary_pos_emb, k)
+
         # caching
 
         if exists(kv_cache):
@@ -1338,12 +1344,6 @@ class Attention(Module):
             k = cat((ck, k), dim = -2)
             v = cat((cv, v), dim = -2)
 
-        # rotary
-
-        if exists(rotary_pos_emb):
-            q = apply_rotations(rotary_pos_emb, q)
-            k = apply_rotations(rotary_pos_emb, k)
-
         # attention
 
         attend_fn = default(attend_fn, naive_attend)
@@ -1507,12 +1507,11 @@ class AxialSpaceTimeTransformer(Module):
 
         has_kv_cache = exists(kv_cache)
 
-
         if has_kv_cache:
             past_tokens, tokens = tokens[:, :-1], tokens[:, -1:]
 
             rotary_seq_len = 1
-            rotary_pos_offset = past_tokens.shape[-2]
+            rotary_pos_offset = past_tokens.shape[1]
         else:
             rotary_seq_len = time
             rotary_pos_offset = 0
@@ -1687,6 +1686,7 @@ class VideoTokenizer(Module):
             time_block_every = time_block_every,
             num_special_spatial_tokens = num_latent_tokens,
             num_residual_streams = num_residual_streams,
+            special_attend_only_itself = True,
             final_norm = True
         )
 
@@ -2429,6 +2429,7 @@ class DynamicsWorldModel(Module):
         normalize_advantages = None,
         eps = 1e-6
     ):
+        assert isinstance(experience, Experience)
 
         latents = experience.latents
         actions = experience.actions
@@ -2441,7 +2442,7 @@ class DynamicsWorldModel(Module):
         step_size = experience.step_size
         agent_index = experience.agent_index
 
-        assert all([*map(exists, (old_log_probs, actions, old_values, rewards, step_size))]), 'the generations need to contain the log probs, values, and rewards for policy optimization'
+        assert all([*map(exists, (old_log_probs, actions, old_values, rewards, step_size))]), 'the generations need to contain the log probs, values, and rewards for policy optimization - world_model.generate(..., return_log_probs_and_values = True)'
 
         batch, time = latents.shape[0], latents.shape[1]
 
@@ -2455,8 +2456,8 @@ class DynamicsWorldModel(Module):
         if exists(experience.lens):
             mask_for_gae = lens_to_mask(experience.lens, time)
 
-            rewards = rewards.masked_fill(mask_for_gae, 0.)
-            old_values = old_values.masked_fill(mask_for_gae, 0.)
+            rewards = rewards.masked_fill(~mask_for_gae, 0.)
+            old_values = old_values.masked_fill(~mask_for_gae, 0.)
 
         # calculate returns
 
@@ -2491,7 +2492,7 @@ class DynamicsWorldModel(Module):
 
             # mean, var - todo - handle distributed
 
-            returns_mean, returns_var = returns.mean(), returns.var()
+            returns_mean, returns_var = returns_for_stats.mean(), returns_for_stats.var()
 
             # ema
 
@@ -2694,12 +2695,22 @@ class DynamicsWorldModel(Module):
         return_rewards_per_frame = False,
         return_agent_actions = False,
         return_log_probs_and_values = False,
+        return_for_policy_optimization = False,
         return_time_kv_cache = False,
         store_agent_embed = True,
         store_old_action_unembeds = True
 
     ): # (b t n d) | (b c t h w)
 
+        # handy flag for returning generations for rl
+
+        if return_for_policy_optimization:
+            return_agent_actions |= True
+            return_log_probs_and_values |= True
+            return_rewards_per_frame |= True
+
+        # more variables
+
         has_proprio = self.has_proprio
         was_training = self.training
         self.eval()
@@ -2769,6 +2780,19 @@ class DynamicsWorldModel(Module):
 
             curr_time_steps = latents.shape[1]
 
+            # determine whether to take an extra step if
+            # (1) using time kv cache
+            # (2) decoding anything off agent embedding (rewards, actions, etc)
+
+            take_extra_step = (
+                use_time_kv_cache or
+                return_rewards_per_frame or
+                store_agent_embed or
+                return_agent_actions
+            )
+
+            # prepare noised latent / proprio inputs
+
             noised_latent = randn((batch_size, 1, self.num_video_views, *latent_shape), device = self.device)
 
             noised_proprio = None
@@ -2776,7 +2800,10 @@ class DynamicsWorldModel(Module):
             if has_proprio:
                 noised_proprio = randn((batch_size, 1, self.dim_proprio), device = self.device)
 
-            for step in range(num_steps):
+            # denoising steps
+
+            for step in range(num_steps + int(take_extra_step)):
+
                 is_last_step = (step + 1) == num_steps
 
                 signal_levels = full((batch_size, 1), step * step_size, dtype = torch.long, device = self.device)
@@ -2819,6 +2846,11 @@ class DynamicsWorldModel(Module):
                 if use_time_kv_cache and is_last_step:
                     time_kv_cache = next_time_kv_cache
 
+                # early break if taking an extra step for agent embedding off cleaned latents for decoding
+
+                if take_extra_step and is_last_step:
+                    break
+
                 # maybe proprio
 
                 if has_proprio:
@@ -3021,7 +3053,7 @@ class DynamicsWorldModel(Module):
         latent_is_noised = False,
         return_all_losses = False,
         return_intermediates = False,
-        add_autoregressive_action_loss = False,
+        add_autoregressive_action_loss = True,
         update_loss_ema = None,
         latent_has_view_dim = False
     ):
@@ -3478,7 +3510,7 @@ class DynamicsWorldModel(Module):
 
             reward_losses = F.cross_entropy(reward_pred, reward_targets, reduction = 'none')
 
-            reward_losses = reward_losses.masked_fill(reward_loss_mask, 0.)
+            reward_losses = reward_losses.masked_fill(~reward_loss_mask, 0.)
 
             if is_var_len:
                 reward_loss = reward_losses[loss_mask_without_last].mean(dim = 0)

{dreamer4-0.0.102.dist-info → dreamer4-0.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreamer4
-Version: 0.0.102
+Version: 0.1.1
 Summary: Dreamer 4
 Project-URL: Homepage, https://pypi.org/project/dreamer4/
 Project-URL: Repository, https://github.com/lucidrains/dreamer4
@@ -53,11 +53,75 @@ Description-Content-Type: text/markdown
 
 <img src="./dreamer4-fig2.png" width="400px"></img>
 
-## Dreamer 4 (wip)
+## Dreamer 4
 
 Implementation of Danijar's [latest iteration](https://arxiv.org/abs/2509.24527v1) for his [Dreamer](https://danijar.com/project/dreamer4/) line of work
 
-[Temporary Discord](https://discord.gg/MkACrrkrYR)
+## Install
+
+```bash
+$ pip install dreamer4
+```
+
+## Usage
+
+```python
+import torch
+from dreamer4 import VideoTokenizer, DynamicsWorldModel
+
+# video tokenizer, learned through MAE + lpips
+
+tokenizer = VideoTokenizer(
+    dim = 512,
+    dim_latent = 32,
+    patch_size = 32,
+    image_height = 256,
+    image_width = 256
+)
+
+# dynamics world model
+
+dynamics = DynamicsWorldModel(
+    dim = 512,
+    dim_latent = 32,
+    video_tokenizer = tokenizer,
+    num_discrete_actions = 4,
+    num_residual_streams = 1
+)
+
+# state, action, rewards
+
+video = torch.randn(2, 3, 10, 256, 256)
+discrete_actions = torch.randint(0, 4, (2, 10, 1))
+rewards = torch.randn(2, 10)
+
+# learn dynamics / behavior cloned model
+
+loss = dynamics(
+    video = video,
+    rewards = rewards,
+    discrete_actions = discrete_actions
+)
+
+loss.backward()
+
+# do the above with much data
+
+# then generate dreams
+
+dreams = dynamics.generate(
+    10,
+    batch_size = 2,
+    return_decoded_video = True,
+    return_for_policy_optimization = True
+)
+
+# learn from the dreams
+
+actor_loss, critic_loss = dynamics.learn_from_experience(dreams)
+
+(actor_loss + critic_loss).backward()
+```
 
 ## Citation
 
@@ -72,3 +136,5 @@ Implementation of Danijar's [latest iteration](https://arxiv.org/abs/2509.24527v
     url     = {https://arxiv.org/abs/2509.24527}, 
 }
 ```
+
+*the conquest of nature is to be achieved through number and measure - angels to Descartes in a dream*

dreamer4-0.1.1.dist-info/RECORD

@@ -0,0 +1,8 @@
+dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
+dreamer4/dreamer4.py,sha256=_gWp08k7tf2VCUv7uqkXKZQugnqJqXPb1-o7_34SA9c,120365
+dreamer4/mocks.py,sha256=TfqOB_Gq6N_GggBYwa6ZAJQx38ntlYbXZe23Ne4jshw,2502
+dreamer4/trainers.py,sha256=JsnJwQJcbI_75KBTNddG6b7QVkO6LD1N_HQiVe-VnCM,15087
+dreamer4-0.1.1.dist-info/METADATA,sha256=2zkBv1BHvGpb6onAnEFsKnPK2KD-0vH8K1nFDBVlpyU,4247
+dreamer4-0.1.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dreamer4-0.1.1.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+dreamer4-0.1.1.dist-info/RECORD,,

dreamer4-0.0.102.dist-info/RECORD

@@ -1,8 +0,0 @@
-dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
-dreamer4/dreamer4.py,sha256=3qeVN3qdvx7iPxA0OBXw_yy5Re6rX6FIKITH9bp6RBs,119202
-dreamer4/mocks.py,sha256=TfqOB_Gq6N_GggBYwa6ZAJQx38ntlYbXZe23Ne4jshw,2502
-dreamer4/trainers.py,sha256=JsnJwQJcbI_75KBTNddG6b7QVkO6LD1N_HQiVe-VnCM,15087
-dreamer4-0.0.102.dist-info/METADATA,sha256=xxVL1sFimb0azSD5sDOEzugY7rBT6oDek4YdiIS8m18,3066
-dreamer4-0.0.102.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dreamer4-0.0.102.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-dreamer4-0.0.102.dist-info/RECORD,,