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

dreamer4/dreamer4.py

@@ -1179,10 +1179,11 @@ def special_token_mask(q, k, seq_len, num_tokens, special_attend_only_itself = F
 
 def block_mask_special_tokens_right(
     seq_len,
-    num_tokens
+    num_tokens,
+    special_attend_only_itself = False
 ):
     def inner(b, h, q, k):
-        return special_token_mask(q, k, seq_len, num_tokens)
+        return special_token_mask(q, k, seq_len, num_tokens, special_attend_only_itself)
     return inner
 
 def compose_mask(mask1, mask2):
@@ -1331,6 +1332,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 +1345,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)
@@ -1493,7 +1494,8 @@ class AxialSpaceTimeTransformer(Module):
 
         # attend functions for space and time
 
-        use_flex = exists(flex_attention) and tokens.is_cuda
+        has_kv_cache = exists(kv_cache) 
+        use_flex = exists(flex_attention) and tokens.is_cuda and not has_kv_cache # KV cache shape breaks flex attention TODO: Fix
 
         attend_kwargs = dict(use_flex = use_flex, softclamp_value = self.attn_softclamp_value, special_attend_only_itself = self.special_attend_only_itself, device = device)
 
@@ -1505,14 +1507,12 @@ class AxialSpaceTimeTransformer(Module):
 
         time_attn_kv_caches = []
 
-        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 +1687,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
         )
 
@@ -1847,7 +1848,7 @@ class VideoTokenizer(Module):
 
         losses = (recon_loss, lpips_loss)
 
-        return total_loss, TokenizerLosses(losses)
+        return total_loss, TokenizerLosses(*losses)
 
 # dynamics model, axial space-time transformer
 
@@ -2104,7 +2105,7 @@ class DynamicsWorldModel(Module):
 
         self.ppo_eps_clip = ppo_eps_clip
         self.value_clip = value_clip
-        self.policy_entropy_weight = value_clip
+        self.policy_entropy_weight = policy_entropy_weight
 
         # pmpo related
 
@@ -2127,7 +2128,7 @@ class DynamicsWorldModel(Module):
         self.flow_loss_normalizer = LossNormalizer(1)
         self.reward_loss_normalizer = LossNormalizer(multi_token_pred_len)
         self.discrete_actions_loss_normalizer = LossNormalizer(multi_token_pred_len) if num_discrete_actions > 0 else None
-        self.continuous_actions_loss_normalizer = LossNormalizer(multi_token_pred_len) if num_discrete_actions > 0 else None
+        self.continuous_actions_loss_normalizer = LossNormalizer(multi_token_pred_len) if num_continuous_actions > 0 else None
 
         self.latent_flow_loss_weight = latent_flow_loss_weight
 
@@ -2358,6 +2359,9 @@ class DynamicsWorldModel(Module):
             elif len(env_step_out) == 4:
                 next_frame, reward, terminated, truncated = env_step_out
 
+            elif len(env_step_out) == 5:
+                next_frame, reward, terminated, truncated, info = env_step_out
+
             # update episode lens
 
             episode_lens = torch.where(is_terminated, episode_lens, episode_lens + 1)
@@ -2456,8 +2460,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
 
@@ -2492,7 +2496,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
 
@@ -3085,8 +3089,8 @@ class DynamicsWorldModel(Module):
         if latents.ndim == 4:
             latents = rearrange(latents, 'b t v d -> b t v 1 d') # 1 latent edge case
 
-        assert latents.shape[-2:] == self.latent_shape
-        assert latents.shape[2] == self.num_video_views
+        assert latents.shape[-2:] == self.latent_shape, f'latents must have shape {self.latent_shape}, got {latents.shape[-2:]}'
+        assert latents.shape[2] == self.num_video_views, f'latents must have {self.num_video_views} views, got {latents.shape[2]}'
 
         # variables
 
@@ -3510,7 +3514,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)
@@ -3554,7 +3558,7 @@ class DynamicsWorldModel(Module):
                 discrete_mask = rearrange(discrete_mask, 'b t mtp -> mtp b t')
 
             if exists(continuous_actions):
-                continuous_action_targets, continuous_mask = create_multi_token_prediction_targets(discrete_actions, self.multi_token_pred_len)
+                continuous_action_targets, continuous_mask = create_multi_token_prediction_targets(continuous_actions, self.multi_token_pred_len)
                 continuous_action_targets = rearrange(continuous_action_targets, 'b t mtp ... -> mtp b t ...')
                 continuous_mask = rearrange(continuous_mask, 'b t mtp -> mtp b t')
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreamer4
-Version: 0.1.0
+Version: 0.1.4
 Summary: Dreamer 4
 Project-URL: Homepage, https://pypi.org/project/dreamer4/
 Project-URL: Repository, https://github.com/lucidrains/dreamer4
@@ -57,10 +57,16 @@ 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
+
+## Appreciation
+
+- [@dirkmcpherson](https://github.com/dirkmcpherson) for fixes to typo errors and unpassed arguments!
+
 ## Install
 
 ```bash
-$ pip install dreamer4-pytorch
+$ pip install dreamer4
 ```
 
 ## Usage
@@ -79,9 +85,16 @@ tokenizer = VideoTokenizer(
     image_width = 256
 )
 
+video = torch.randn(2, 3, 10, 256, 256)
+
+# learn the tokenizer
+
+loss = tokenizer(video)
+loss.backward() # ler
+
 # dynamics world model
 
-dynamics = DynamicsWorldModel(
+world_model = DynamicsWorldModel(
     dim = 512,
     dim_latent = 32,
     video_tokenizer = tokenizer,
@@ -97,7 +110,7 @@ rewards = torch.randn(2, 10)
 
 # learn dynamics / behavior cloned model
 
-loss = dynamics(
+loss = world_model(
     video = video,
     rewards = rewards,
     discrete_actions = discrete_actions
@@ -109,7 +122,7 @@ loss.backward()
 
 # then generate dreams
 
-dreams = dynamics.generate(
+dreams = world_model.generate(
     10,
     batch_size = 2,
     return_decoded_video = True,
@@ -118,7 +131,19 @@ dreams = dynamics.generate(
 
 # learn from the dreams
 
-actor_loss, critic_loss = dynamics.learn_from_experience(dreams)
+actor_loss, critic_loss = world_model.learn_from_experience(dreams)
+
+(actor_loss + critic_loss).backward()
+
+# learn from environment
+
+from dreamer4.mocks import MockEnv
+
+mock_env = MockEnv((256, 256), vectorized = True, num_envs = 4)
+
+experience = world_model.interact_with_env(mock_env, max_timesteps = 8, env_is_vectorized = True)
+
+actor_loss, critic_loss = world_model.learn_from_experience(experience)
 
 (actor_loss + critic_loss).backward()
 ```
@@ -137,4 +162,4 @@ actor_loss, critic_loss = dynamics.learn_from_experience(dreams)
 }
 ```
 
-*the conquest of nature is to be achieved through number and measure* - angels to Descartes, in a dream, the story goes.
+*the conquest of nature is to be achieved through number and measure - angels to Descartes in a dream*

dreamer4-0.1.4.dist-info/RECORD

@@ -0,0 +1,8 @@
+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,,

dreamer4-0.1.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
-dreamer4/dreamer4.py,sha256=YB724hMjBYDNhApo2x_52oXIeH5GGQo8Q2pB2lkCq_s,120297
-dreamer4/mocks.py,sha256=TfqOB_Gq6N_GggBYwa6ZAJQx38ntlYbXZe23Ne4jshw,2502
-dreamer4/trainers.py,sha256=JsnJwQJcbI_75KBTNddG6b7QVkO6LD1N_HQiVe-VnCM,15087
-dreamer4-0.1.0.dist-info/METADATA,sha256=kDq66Il_WDNKR2NP9wrVu3fMUIVW-pWySq3CP2ANZ2s,4273
-dreamer4-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dreamer4-0.1.0.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-dreamer4-0.1.0.dist-info/RECORD,,