dreamer4 0.0.77__py3-none-any.whl → 0.0.79__py3-none-any.whl

dreamer4/dreamer4.py

@@ -13,7 +13,7 @@ import torch.nn.functional as F
 from torch.nested import nested_tensor
 from torch.distributions import Normal
 from torch.nn import Module, ModuleList, Embedding, Parameter, Sequential, Linear, RMSNorm, Identity
-from torch import nn, cat, stack, arange, tensor, Tensor, is_tensor, zeros, ones, randint, rand, randn, randn_like, empty, full, linspace, arange
+from torch import nn, cat, stack, arange, tensor, Tensor, is_tensor, full, zeros, ones, randint, rand, randn, randn_like, empty, full, linspace, arange
 from torch.utils._pytree import tree_flatten, tree_unflatten
 
 import torchvision
@@ -83,6 +83,7 @@ class Experience:
     values: Tensor | None = None
     step_size: int | None = None
     lens: Tensor | None = None
+    is_truncated: Tensor | None = None
     agent_index: int = 0
     is_from_world_model: bool = True
 
@@ -99,7 +100,10 @@ def combine_experiences(
         batch, time, device = *latents.shape[:2], latents.device
 
         if not exists(exp.lens):
-            exp.lens = torch.full((batch,), time, device = device)
+            exp.lens = full((batch,), time, device = device)
+
+        if not exists(exp.is_truncated):
+            exp.is_truncated = full((batch,), True, device = device)
 
     # convert to dictionary
 
@@ -2115,7 +2119,7 @@ class DynamicsWorldModel(Module):
 
         time_kv_cache = None
 
-        for _ in range(max_timesteps):
+        for i in range(max_timesteps + 1):
 
             latents = self.video_tokenizer(video, return_latents = True)
 
@@ -2141,6 +2145,15 @@ class DynamicsWorldModel(Module):
 
             one_agent_embed = agent_embed[..., -1:, agent_index, :]
 
+            # values
+
+            value_bins = self.value_head(one_agent_embed)
+            value = self.reward_encoder.bins_to_scalar_value(value_bins)
+
+            values = safe_cat((values, value), dim = 1)
+
+            # policy embed
+
             policy_embed = self.policy_head(one_agent_embed)
 
             # sample actions
@@ -2162,11 +2175,6 @@ class DynamicsWorldModel(Module):
             discrete_log_probs = safe_cat((discrete_log_probs, one_discrete_log_probs), dim = 1)
             continuous_log_probs = safe_cat((continuous_log_probs, one_continuous_log_probs), dim = 1)
 
-            value_bins = self.value_head(one_agent_embed)
-            value = self.reward_encoder.bins_to_scalar_value(value_bins)
-
-            values = safe_cat((values, value), dim = 1)
-
             # pass the sampled action to the environment and get back next state and reward
 
             next_frame, reward = env.step((sampled_discrete_actions, sampled_continuous_actions))
@@ -2187,6 +2195,8 @@ class DynamicsWorldModel(Module):
 
         # package up one experience for learning
 
+        batch, device = latents.shape[0], latents.device
+
         one_experience = Experience(
             latents = latents,
             video = video[:, :, :-1],
@@ -2196,6 +2206,7 @@ class DynamicsWorldModel(Module):
             values = values,
             step_size = step_size,
             agent_index = agent_index,
+            lens = full((batch,), max_timesteps + 1, device = device),
             is_from_world_model = False
         )
 
@@ -2224,6 +2235,22 @@ class DynamicsWorldModel(Module):
 
         assert all([*map(exists, (old_log_probs, actions, old_values, rewards, step_size))]), 'the generations need to contain the log probs, values, and rewards for policy optimization'
 
+        batch, time = latents.shape[0], latents.shape[1]
+
+        # calculate returns
+
+        # mask out anything after the `lens`, which may include a bootstrapped node at the very end if `is_truncated = True`
+
+        if not exists(experience.is_truncated):
+            experience.is_truncated = full((batch,), True, device = latents.device)
+
+        mask_for_gae = lens_to_mask(experience.lens, time)
+
+        rewards = rewards.masked_fill(mask_for_gae, 0.)
+        old_values = old_values.masked_fill(mask_for_gae, 0.)
+
+        # calculate returns
+
         returns = calc_gae(rewards, old_values, gamma = self.gae_discount_factor, lam = self.gae_lambda, use_accelerated = self.gae_use_accelerated)
 
         # handle variable lengths
@@ -2232,8 +2259,8 @@ class DynamicsWorldModel(Module):
         is_var_len = exists(experience.lens)
 
         if is_var_len:
-            lens = experience.lens
-            mask = lens_to_mask(lens, max_time)
+            learnable_lens = experience.lens - experience.is_truncated.long() # if is truncated, remove the last one, as it is bootstrapped value
+            mask = lens_to_mask(learnable_lens, max_time)
 
         # determine whether to finetune entire transformer or just learn the heads
 
@@ -2387,7 +2414,7 @@ class DynamicsWorldModel(Module):
         assert 0 < num_steps <= self.max_steps, f'number of steps {num_steps} must be between 0 and {self.max_steps}'
 
         if isinstance(tasks, int):
-            tasks = torch.full((batch_size,), tasks, device = self.device)
+            tasks = full((batch_size,), tasks, device = self.device)
 
         assert not exists(tasks) or tasks.shape[0] == batch_size
 
@@ -2624,7 +2651,7 @@ class DynamicsWorldModel(Module):
         # returning agent actions, rewards, and log probs + values for policy optimization
 
         batch, device = latents.shape[0], latents.device
-        experience_lens = torch.full((batch,), time_steps, device = device)
+        experience_lens = full((batch,), time_steps, device = device)
 
         gen = Experience(
             latents = latents,

dreamer4/trainers.py

@@ -287,7 +287,7 @@ class DreamTrainer(Module):
         for _ in range(self.num_train_steps):
 
             dreams = self.unwrapped_model.generate(
-                self.generate_timesteps,
+                self.generate_timesteps + 1, # plus one for bootstrap value
                 batch_size = self.batch_size,
                 return_rewards_per_frame = True,
                 return_agent_actions = True,

{dreamer4-0.0.77.dist-info → dreamer4-0.0.79.dist-info}/METADATA

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

dreamer4-0.0.79.dist-info/RECORD

@@ -0,0 +1,8 @@
+dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
+dreamer4/dreamer4.py,sha256=B-mNoFEVC_fk4ZOAr9CKSniQfFrvtV5aXG9jmlRxoBA,108095
+dreamer4/mocks.py,sha256=Oi91Yv1oK0E-Wz-KDkf79xoyWzIXCvMLCr0WYCpJDLA,1482
+dreamer4/trainers.py,sha256=D2b7WTgTHElLhIWLFgl2Ct2knGJLTk91HHpC5UkNvG0,14028
+dreamer4-0.0.79.dist-info/METADATA,sha256=6QM1bmGCCQeYkYW9HD3OQgNoKYlo-ILQIJ2e03aeZYk,3065
+dreamer4-0.0.79.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dreamer4-0.0.79.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+dreamer4-0.0.79.dist-info/RECORD,,

dreamer4-0.0.77.dist-info/RECORD

@@ -1,8 +0,0 @@
-dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
-dreamer4/dreamer4.py,sha256=v6W9psBXMCGo9K3F6xCxm49TEu3jFUqDvo3pf_bsBQo,107099
-dreamer4/mocks.py,sha256=Oi91Yv1oK0E-Wz-KDkf79xoyWzIXCvMLCr0WYCpJDLA,1482
-dreamer4/trainers.py,sha256=898ye9Y1mqxGZnU_gfQS6pECibZwwyA43sL7wK_JHAU,13993
-dreamer4-0.0.77.dist-info/METADATA,sha256=8XGi2Q8meZfrz2K1cpbVH1P34wODclA0JD7OG0xldaQ,3065
-dreamer4-0.0.77.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dreamer4-0.0.77.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-dreamer4-0.0.77.dist-info/RECORD,,