dreamer4 0.0.82__tar.gz → 0.0.84__tar.gz

{dreamer4-0.0.82 → dreamer4-0.0.84}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreamer4
-Version: 0.0.82
+Version: 0.0.84
 Summary: Dreamer 4
 Project-URL: Homepage, https://pypi.org/project/dreamer4/
 Project-URL: Repository, https://github.com/lucidrains/dreamer4

{dreamer4-0.0.82 → dreamer4-0.0.84}/dreamer4/dreamer4.py

@@ -77,6 +77,7 @@ class Experience:
     latents: Tensor
     video: Tensor | None = None
     proprio: Tensor | None = None
+    agent_embed: Tensor | None = None,
     rewards: Tensor | None = None
     actions: tuple[Tensor, Tensor] | None = None
     log_probs: tuple[Tensor, Tensor] | None = None
@@ -2105,7 +2106,8 @@ class DynamicsWorldModel(Module):
         step_size = 4,
         max_timesteps = 16,
         env_is_vectorized = False,
-        use_time_kv_cache = True
+        use_time_kv_cache = True,
+        store_agent_embed = True
     ):
         assert exists(self.video_tokenizer)
 
@@ -2130,16 +2132,23 @@ class DynamicsWorldModel(Module):
         values = None
         latents = None
 
+        acc_agent_embed = None
+
         # keep track of termination, for setting the `is_truncated` field on Experience and for early stopping interaction with env
 
         is_terminated = full((batch,), False, device = device)
+        is_truncated = full((batch,), False, device = device)
+
         episode_lens = full((batch,), 0, device = device)
 
         # maybe time kv cache
 
         time_kv_cache = None
 
-        for i in range(max_timesteps + 1):
+        step_index = 0
+
+        while not is_terminated.all():
+            step_index += 1
 
             latents = self.video_tokenizer(video, return_latents = True)
 
@@ -2201,10 +2210,15 @@ class DynamicsWorldModel(Module):
 
             if len(env_step_out) == 2:
                 next_frame, reward = env_step_out
-                terminate = full((batch,), False)
+                terminated = full((batch,), False)
+                truncated = full((batch,), False)
 
             elif len(env_step_out) == 3:
-                next_frame, reward, terminate = env_step_out
+                next_frame, reward, terminated = env_step_out
+                truncated = full((batch,), False)
+
+            elif len(env_step_out) == 4:
+                next_frame, reward, terminated, truncated = env_step_out
 
             # update episode lens
 
@@ -2212,7 +2226,20 @@ class DynamicsWorldModel(Module):
 
             # update `is_terminated`
 
-            is_terminated |= terminate
+            # (1) - environment says it is terminated
+            # (2) - previous step is truncated (this step is for bootstrap value)
+
+            is_terminated |= (terminated | is_truncated)
+
+            # update `is_truncated`
+
+            if step_index <= max_timesteps:
+                is_truncated |= truncated
+
+            if step_index == max_timesteps:
+                # if the step index is at the max time step allowed, set the truncated flag, if not already terminated
+
+                is_truncated |= ~is_terminated
 
             # batch and time dimension
 
@@ -2228,10 +2255,7 @@ class DynamicsWorldModel(Module):
             video = cat((video, next_frame), dim = 2)
             rewards = safe_cat((rewards, reward), dim = 1)
 
-            # early break out if all terminated
-
-            if is_terminated.all():
-                break
+            acc_agent_embed = safe_cat((acc_agent_embed, agent_embed), dim = 1)
 
         # package up one experience for learning
 
@@ -2244,9 +2268,10 @@ class DynamicsWorldModel(Module):
             actions = (discrete_actions, continuous_actions),
             log_probs = (discrete_log_probs, continuous_log_probs),
             values = values,
+            agent_embed = acc_agent_embed if store_agent_embed else None,
             step_size = step_size,
             agent_index = agent_index,
-            is_truncated = ~is_terminated,
+            is_truncated = is_truncated,
             lens = episode_lens,
             is_from_world_model = False
         )
@@ -2473,6 +2498,7 @@ class DynamicsWorldModel(Module):
         return_agent_actions = False,
         return_log_probs_and_values = False,
         return_time_kv_cache = False,
+        store_agent_embed = False
 
     ): # (b t n d) | (b c t h w)
 
@@ -2525,6 +2551,10 @@ class DynamicsWorldModel(Module):
         decoded_continuous_log_probs = None
         decoded_values = None
 
+        # maybe store agent embed
+
+        acc_agent_embed = None
+
         # maybe return rewards
 
         decoded_rewards = None
@@ -2633,6 +2663,10 @@ class DynamicsWorldModel(Module):
 
                 decoded_rewards = cat((decoded_rewards, pred_reward), dim = 1)
 
+            # maybe store agent embed
+
+            acc_agent_embed = safe_cat((acc_agent_embed, agent_embed), dim = 1)
+
             # decode the agent actions if needed
 
             if return_agent_actions:
@@ -2729,6 +2763,7 @@ class DynamicsWorldModel(Module):
             latents = latents,
             video = video,
             proprio = proprio if has_proprio else None,
+            agent_embed = acc_agent_embed if store_agent_embed else None,
             step_size = step_size,
             agent_index = agent_index,
             lens = experience_lens,

{dreamer4-0.0.82 → dreamer4-0.0.84}/dreamer4/mocks.py

@@ -22,7 +22,9 @@ class MockEnv(Module):
         num_envs = 1,
         vectorized = False,
         terminate_after_step = None,
-        rand_terminate_prob = 0.05
+        rand_terminate_prob = 0.05,
+        can_truncate = False,
+        rand_truncate_prob = 0.05,
     ):
         super().__init__()
         self.image_shape = image_shape
@@ -32,12 +34,15 @@ class MockEnv(Module):
         self.vectorized = vectorized
         assert not (vectorized and num_envs == 1)
 
-        # mocking termination
+        # mocking termination and truncation
 
         self.can_terminate = exists(terminate_after_step)
         self.terminate_after_step = terminate_after_step
         self.rand_terminate_prob = rand_terminate_prob
 
+        self.can_truncate = can_truncate
+        self.rand_truncate_prob = rand_truncate_prob
+
         self.register_buffer('_step', tensor(0))
 
     def get_random_state(self):
@@ -72,14 +77,21 @@ class MockEnv(Module):
 
         out = (state, reward)
 
+
         if self.can_terminate:
-            terminate = (
-                (torch.rand((self.num_envs)) < self.rand_terminate_prob) &
-                (self._step > self.terminate_after_step)
-            )
+            shape = (self.num_envs,) if self.vectorized else (1,)
+            valid_step = self._step > self.terminate_after_step
+
+            terminate = (torch.rand(shape) < self.rand_terminate_prob) & valid_step
 
             out = (*out, terminate)
 
+            # maybe truncation
+
+            if self.can_truncate:
+                truncate = (torch.rand(shape) < self.rand_truncate_prob) & valid_step & ~terminate
+                out = (*out, truncate)
+
         self._step.add_(1)
 
         return out

{dreamer4-0.0.82 → dreamer4-0.0.84}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dreamer4"
-version = "0.0.82"
+version = "0.0.84"
 description = "Dreamer 4"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{dreamer4-0.0.82 → dreamer4-0.0.84}/tests/test_dreamer.py

@@ -613,10 +613,14 @@ def test_cache_generate():
 @param('vectorized', (False, True))
 @param('use_signed_advantage', (False, True))
 @param('env_can_terminate', (False, True))
+@param('env_can_truncate', (False, True))
+@param('store_agent_embed', (False, True))
 def test_online_rl(
     vectorized,
     use_signed_advantage,
-    env_can_terminate
+    env_can_terminate,
+    env_can_truncate,
+    store_agent_embed
 ):
     from dreamer4.dreamer4 import DynamicsWorldModel, VideoTokenizer
 
@@ -656,16 +660,20 @@ def test_online_rl(
         vectorized = vectorized,
         num_envs = 4,
         terminate_after_step = 2 if env_can_terminate else None,
+        can_truncate = env_can_truncate,
         rand_terminate_prob = 0.1
     )
 
     # manually
 
-    one_experience = world_model_and_policy.interact_with_env(mock_env, max_timesteps = 8, env_is_vectorized = vectorized)
-    another_experience = world_model_and_policy.interact_with_env(mock_env, max_timesteps = 16, env_is_vectorized = vectorized)
+    one_experience = world_model_and_policy.interact_with_env(mock_env, max_timesteps = 8, env_is_vectorized = vectorized, store_agent_embed = store_agent_embed)
+    another_experience = world_model_and_policy.interact_with_env(mock_env, max_timesteps = 16, env_is_vectorized = vectorized, store_agent_embed = store_agent_embed)
 
     combined_experience = combine_experiences([one_experience, another_experience])
 
+    if store_agent_embed:
+        assert exists(combined_experience.agent_embed)
+
     actor_loss, critic_loss = world_model_and_policy.learn_from_experience(combined_experience, use_signed_advantage = use_signed_advantage)
 
     actor_loss.backward()