dreamer4 0.0.77__tar.gz → 0.0.79__tar.gz

{dreamer4-0.0.77 → dreamer4-0.0.79}/PKG-INFO

@@ -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.77 → dreamer4-0.0.79}/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-0.0.77 → dreamer4-0.0.79}/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 → dreamer4-0.0.79}/pyproject.toml

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