dreamer4 0.0.93__py3-none-any.whl → 0.0.95__py3-none-any.whl

dreamer4/dreamer4.py

@@ -72,20 +72,22 @@ TokenizerLosses = namedtuple('TokenizerLosses', ('recon', 'lpips'))
 
 WorldModelLosses = namedtuple('WorldModelLosses', ('flow', 'rewards', 'discrete_actions', 'continuous_actions'))
 
+MaybeTensor = Tensor | None
+
 @dataclass
 class Experience:
     latents: Tensor
-    video: Tensor | None = None
-    proprio: Tensor | None = None
-    agent_embed: Tensor | None = None
+    video: MaybeTensor = None
+    proprio: MaybeTensor = None
+    agent_embed: MaybeTensor = None
     rewards: Tensor | None = None
-    actions: tuple[Tensor, Tensor] | None = None
-    log_probs: tuple[Tensor, Tensor] | None = None
-    old_action_unembeds: tuple[Tensor, Tensor] | None = None
-    values: Tensor | None = None
+    actions: tuple[MaybeTensor, MaybeTensor] | None = None
+    log_probs: tuple[MaybeTensor, MaybeTensor] | None = None
+    old_action_unembeds: tuple[MaybeTensor, MaybeTensor] | None = None
+    values: MaybeTensor = None
     step_size: int | None = None
-    lens: Tensor | None = None
-    is_truncated: Tensor | None = None
+    lens: MaybeTensor = None
+    is_truncated: MaybeTensor = None
     agent_index: int = 0
     is_from_world_model: bool = True
 
@@ -850,9 +852,10 @@ class ActionEmbedder(Module):
 
     def kl_div(
         self,
-        src: tuple[Tensor | None, Tensor | None],
-        tgt: tuple[Tensor | None, Tensor | None]
-    ) -> tuple[Tensor | None, Tensor | None]:
+        src: tuple[MaybeTensor, MaybeTensor],
+        tgt: tuple[MaybeTensor, MaybeTensor],
+        reduce_across_num_actions = True
+    ) -> tuple[MaybeTensor, MaybeTensor]:
 
         src_discrete, src_continuous = src
         tgt_discrete, tgt_continuous = tgt
@@ -894,6 +897,15 @@ class ActionEmbedder(Module):
 
             continuous_kl_div = kl.kl_divergence(src_normal, tgt_normal)
 
+        # maybe reduce
+
+        if reduce_across_num_actions:
+            if exists(discrete_kl_div):
+                discrete_kl_div = discrete_kl_div.sum(dim = -1)
+
+            if exists(continuous_kl_div):
+                continuous_kl_div = continuous_kl_div.sum(dim = -1)
+
         return discrete_kl_div, continuous_kl_div
 
     def forward(
@@ -2225,8 +2237,8 @@ class DynamicsWorldModel(Module):
         max_timesteps = 16,
         env_is_vectorized = False,
         use_time_kv_cache = True,
-        store_agent_embed = False,
-        store_old_action_unembeds = False,
+        store_agent_embed = True,
+        store_old_action_unembeds = True,
     ):
         assert exists(self.video_tokenizer)
 
@@ -2391,7 +2403,7 @@ class DynamicsWorldModel(Module):
             actions = (discrete_actions, continuous_actions),
             log_probs = (discrete_log_probs, continuous_log_probs),
             values = values,
-            old_action_unembeds = self.action_embedder.unembed(acc_policy_embed, pred_head_index = 0) if store_old_action_unembeds else None,
+            old_action_unembeds = self.action_embedder.unembed(acc_policy_embed, pred_head_index = 0) if exists(acc_policy_embed) and store_old_action_unembeds else None,
             agent_embed = acc_agent_embed if store_agent_embed else None,
             step_size = step_size,
             agent_index = agent_index,
@@ -2667,7 +2679,8 @@ class DynamicsWorldModel(Module):
         return_agent_actions = False,
         return_log_probs_and_values = False,
         return_time_kv_cache = False,
-        store_agent_embed = False
+        store_agent_embed = True,
+        store_old_action_unembeds = True
 
     ): # (b t n d) | (b c t h w)
 
@@ -2947,7 +2960,7 @@ class DynamicsWorldModel(Module):
             video = video,
             proprio = proprio if has_proprio else None,
             agent_embed = acc_agent_embed if store_agent_embed else None,
-            old_action_unembeds = self.action_embedder.unembed(acc_policy_embed, pred_head_index = 0) if store_old_action_unembeds else None,
+            old_action_unembeds = self.action_embedder.unembed(acc_policy_embed, pred_head_index = 0) if exists(acc_policy_embed) and store_old_action_unembeds else None,
             step_size = step_size,
             agent_index = agent_index,
             lens = experience_lens,

{dreamer4-0.0.93.dist-info → dreamer4-0.0.95.dist-info}/METADATA

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

dreamer4-0.0.95.dist-info/RECORD

@@ -0,0 +1,8 @@
+dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
+dreamer4/dreamer4.py,sha256=pSwW4DpHMF78ortsHLVHDdWzrsjse6bIVxX3oolA-Ao,118572
+dreamer4/mocks.py,sha256=TfqOB_Gq6N_GggBYwa6ZAJQx38ntlYbXZe23Ne4jshw,2502
+dreamer4/trainers.py,sha256=D2b7WTgTHElLhIWLFgl2Ct2knGJLTk91HHpC5UkNvG0,14028
+dreamer4-0.0.95.dist-info/METADATA,sha256=VWiRy5xotYsn9HL5EFymn9N8j8-_wFKbYeTA5k6E0z4,3065
+dreamer4-0.0.95.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dreamer4-0.0.95.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+dreamer4-0.0.95.dist-info/RECORD,,

dreamer4-0.0.93.dist-info/RECORD

@@ -1,8 +0,0 @@
-dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
-dreamer4/dreamer4.py,sha256=B2Bk6JJO9MVTWwss9hOP1k6SBiEr56ijNOa3PiidPnY,118120
-dreamer4/mocks.py,sha256=TfqOB_Gq6N_GggBYwa6ZAJQx38ntlYbXZe23Ne4jshw,2502
-dreamer4/trainers.py,sha256=D2b7WTgTHElLhIWLFgl2Ct2knGJLTk91HHpC5UkNvG0,14028
-dreamer4-0.0.93.dist-info/METADATA,sha256=FhVnlhfeloUPMiFqqJ5qR6fqdd7YmN1-gXykkOTPF_A,3065
-dreamer4-0.0.93.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dreamer4-0.0.93.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-dreamer4-0.0.93.dist-info/RECORD,,