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

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.95.dist-info → dreamer4-0.0.96.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreamer4
-Version: 0.0.95
+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.96.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=JsnJwQJcbI_75KBTNddG6b7QVkO6LD1N_HQiVe-VnCM,15087
+dreamer4-0.0.96.dist-info/METADATA,sha256=yIipeLorTHctLFwpd3bPExL6-nG5LxTfxPImf5BngN4,3065
+dreamer4-0.0.96.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dreamer4-0.0.96.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+dreamer4-0.0.96.dist-info/RECORD,,

dreamer4-0.0.95.dist-info/RECORD

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