dreamer4 0.0.83__py3-none-any.whl → 0.0.84__py3-none-any.whl

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,6 +2132,8 @@ 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)
@@ -2251,6 +2255,8 @@ class DynamicsWorldModel(Module):
             video = cat((video, next_frame), dim = 2)
             rewards = safe_cat((rewards, reward), dim = 1)
 
+            acc_agent_embed = safe_cat((acc_agent_embed, agent_embed), dim = 1)
+
         # package up one experience for learning
 
         batch, device = latents.shape[0], latents.device
@@ -2262,6 +2268,7 @@ 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_truncated,
@@ -2491,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)
 
@@ -2543,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
@@ -2651,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:
@@ -2747,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.83.dist-info → dreamer4-0.0.84.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreamer4
-Version: 0.0.83
+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.84.dist-info/RECORD

@@ -0,0 +1,8 @@
+dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
+dreamer4/dreamer4.py,sha256=jldgU8K1fSx8Eb9v1VYUhtbVSYZnHIEtQgx-WMnzep4,111820
+dreamer4/mocks.py,sha256=TfqOB_Gq6N_GggBYwa6ZAJQx38ntlYbXZe23Ne4jshw,2502
+dreamer4/trainers.py,sha256=D2b7WTgTHElLhIWLFgl2Ct2knGJLTk91HHpC5UkNvG0,14028
+dreamer4-0.0.84.dist-info/METADATA,sha256=4D7elrDaDKYR2VhD94MZ5gsCFohvKIAkD1SFf_atx3w,3065
+dreamer4-0.0.84.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dreamer4-0.0.84.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+dreamer4-0.0.84.dist-info/RECORD,,

dreamer4-0.0.83.dist-info/RECORD

@@ -1,8 +0,0 @@
-dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
-dreamer4/dreamer4.py,sha256=t9g-fJVRPn4nefpNiKwqkIkKZtIeJAn5V4ruNCJ5a9A,111265
-dreamer4/mocks.py,sha256=TfqOB_Gq6N_GggBYwa6ZAJQx38ntlYbXZe23Ne4jshw,2502
-dreamer4/trainers.py,sha256=D2b7WTgTHElLhIWLFgl2Ct2knGJLTk91HHpC5UkNvG0,14028
-dreamer4-0.0.83.dist-info/METADATA,sha256=pYIw5Pj41JhfJ4Hp9AfegXrxnhIZ9Fk98kF7lY3nAZk,3065
-dreamer4-0.0.83.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dreamer4-0.0.83.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-dreamer4-0.0.83.dist-info/RECORD,,