dreamer4 0.0.94__tar.gz → 0.0.96__tar.gz

{dreamer4-0.0.94 → dreamer4-0.0.96}/PKG-INFO

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

{dreamer4-0.0.94 → dreamer4-0.0.96}/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(

{dreamer4-0.0.94 → dreamer4-0.0.96}/dreamer4/trainers.py

@@ -396,13 +396,20 @@ class SimTrainer(Module):
         old_values = experience.values
         rewards = experience.rewards
 
+        has_agent_embed = exists(experience.agent_embed)
+        agent_embed = experience.agent_embed
+
         discrete_actions, continuous_actions = experience.actions
         discrete_log_probs, continuous_log_probs = experience.log_probs
 
+        discrete_old_action_unembeds, continuous_old_action_unembeds = default(experience.old_action_unembeds, (None, None))
+
         # handle empties
 
         empty_tensor = torch.empty_like(rewards)
 
+        agent_embed = default(agent_embed, empty_tensor)
+
         has_discrete = exists(discrete_actions)
         has_continuous = exists(continuous_actions)
 
@@ -412,6 +419,9 @@ class SimTrainer(Module):
         discrete_log_probs = default(discrete_log_probs, empty_tensor)
         continuous_log_probs = default(continuous_log_probs, empty_tensor)
 
+        discrete_old_action_unembeds = default(discrete_old_action_unembeds, empty_tensor)
+        continuous_old_action_unembeds = default(discrete_old_action_unembeds, empty_tensor)
+
         # create the dataset and dataloader
 
         dataset = TensorDataset(
@@ -420,6 +430,9 @@ class SimTrainer(Module):
             continuous_actions,
             discrete_log_probs,
             continuous_log_probs,
+            agent_embed,
+            discrete_old_action_unembeds,
+            continuous_old_action_unembeds,
             old_values,
             rewards
         )
@@ -434,6 +447,9 @@ class SimTrainer(Module):
                 continuous_actions,
                 discrete_log_probs,
                 continuous_log_probs,
+                agent_embed,
+                discrete_old_action_unembeds,
+                continuous_old_action_unembeds,
                 old_values,
                 rewards
             ) in dataloader:
@@ -448,10 +464,17 @@ class SimTrainer(Module):
                     continuous_log_probs if has_continuous else None
                 )
 
+                old_action_unembeds = (
+                    discrete_old_action_unembeds if has_discrete else None,
+                    continuous_old_action_unembeds if has_continuous else None
+                )
+
                 batch_experience = Experience(
                     latents = latents,
                     actions = actions,
                     log_probs = log_probs,
+                    agent_embed = agent_embed if has_agent_embed else None,
+                    old_action_unembeds = old_action_unembeds,
                     values = old_values,
                     rewards = rewards,
                     step_size = step_size,

{dreamer4-0.0.94 → dreamer4-0.0.96}/pyproject.toml

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

{dreamer4-0.0.94 → dreamer4-0.0.96}/tests/test_dreamer.py

@@ -352,8 +352,8 @@ def test_action_embedder():
 
     discrete_kl_div, continuous_kl_div = embedder.kl_div((discrete_logits, continuous_mean_log_var), (discrete_logits_tgt, continuous_mean_log_var_tgt))
 
-    assert discrete_kl_div.shape == (2, 3, 2)
-    assert continuous_kl_div.shape == (2, 3, 2)
+    assert discrete_kl_div.shape == (2, 3)
+    assert continuous_kl_div.shape == (2, 3)
 
     # return discrete split by number of actions