dreamer4 0.0.80__py3-none-any.whl → 0.0.82__py3-none-any.whl

dreamer4/dreamer4.py

@@ -1817,6 +1817,9 @@ class DynamicsWorldModel(Module):
         continuous_action_loss_weight: float | list[float] = 1.,
         num_latent_genes = 0,                       # for carrying out evolution within the dreams https://web3.arxiv.org/abs/2503.19037
         num_residual_streams = 1,
+        keep_reward_ema_stats = False,
+        reward_ema_decay = 0.998,
+        reward_quantile_filter = (0.05, 0.95),
         gae_discount_factor = 0.997,
         gae_lambda = 0.95,
         ppo_eps_clip = 0.2,
@@ -2022,6 +2025,16 @@ class DynamicsWorldModel(Module):
         self.value_clip = value_clip
         self.policy_entropy_weight = value_clip
 
+        # rewards related
+
+        self.keep_reward_ema_stats = keep_reward_ema_stats
+        self.reward_ema_decay = reward_ema_decay
+
+        self.register_buffer('reward_quantile_filter', tensor(reward_quantile_filter), persistent = False)
+
+        self.register_buffer('ema_returns_mean', tensor(0.))
+        self.register_buffer('ema_returns_var', tensor(1.))
+
         # loss related
 
         self.flow_loss_normalizer = LossNormalizer(1)
@@ -2105,6 +2118,8 @@ class DynamicsWorldModel(Module):
         else:
             video = rearrange(init_frame, 'c vh vw -> 1 c 1 vh vw')
 
+        batch, device = video.shape[0], video.device
+
         # accumulate
 
         rewards = None
@@ -2115,6 +2130,11 @@ class DynamicsWorldModel(Module):
         values = None
         latents = None
 
+        # keep track of termination, for setting the `is_truncated` field on Experience and for early stopping interaction with env
+
+        is_terminated = full((batch,), False, device = device)
+        episode_lens = full((batch,), 0, device = device)
+
         # maybe time kv cache
 
         time_kv_cache = None
@@ -2177,7 +2197,22 @@ class DynamicsWorldModel(Module):
 
             # 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))
+            env_step_out = env.step((sampled_discrete_actions, sampled_continuous_actions))
+
+            if len(env_step_out) == 2:
+                next_frame, reward = env_step_out
+                terminate = full((batch,), False)
+
+            elif len(env_step_out) == 3:
+                next_frame, reward, terminate = env_step_out
+
+            # update episode lens
+
+            episode_lens = torch.where(is_terminated, episode_lens, episode_lens + 1)
+
+            # update `is_terminated`
+
+            is_terminated |= terminate
 
             # batch and time dimension
 
@@ -2193,6 +2228,11 @@ class DynamicsWorldModel(Module):
             video = cat((video, next_frame), dim = 2)
             rewards = safe_cat((rewards, reward), dim = 1)
 
+            # early break out if all terminated
+
+            if is_terminated.all():
+                break
+
         # package up one experience for learning
 
         batch, device = latents.shape[0], latents.device
@@ -2206,7 +2246,8 @@ class DynamicsWorldModel(Module):
             values = values,
             step_size = step_size,
             agent_index = agent_index,
-            lens = full((batch,), max_timesteps + 1, device = device),
+            is_truncated = ~is_terminated,
+            lens = episode_lens,
             is_from_world_model = False
         )
 
@@ -2267,11 +2308,41 @@ class DynamicsWorldModel(Module):
 
         world_model_forward_context = torch.no_grad if only_learn_policy_value_heads else nullcontext
 
+        # maybe keep track returns statistics and normalize returns and values before calculating advantage, as done in dreamer v3
+
+        if self.keep_reward_ema_stats:
+            ema_returns_mean, ema_returns_var = self.ema_returns_mean, self.ema_returns_var
+
+            decay = 1. - self.reward_ema_decay
+
+            # quantile filter
+
+            lo, hi = torch.quantile(returns, self.reward_quantile_filter).tolist()
+            returns_for_stats = returns.clamp(lo, hi)
+
+            # mean, var - todo - handle distributed
+
+            returns_mean, returns_var = returns.mean(), returns.var()
+
+            # ema
+
+            ema_returns_mean.lerp_(returns_mean, decay)
+            ema_returns_var.lerp_(returns_var, decay)
+
+            # normalize
+
+            ema_returns_std = ema_returns_var.clamp(min = 1e-5).sqrt()
+
+            normed_returns = (returns - ema_returns_mean) / ema_returns_std
+            normed_old_values = (old_values - ema_returns_mean) / ema_returns_std
+
+            advantage = normed_returns - normed_old_values
+        else:
+            advantage = returns - old_values
+
         # apparently they just use the sign of the advantage
         # https://arxiv.org/abs/2410.04166v1
 
-        advantage = returns - old_values
-
         if use_signed_advantage:
             advantage = advantage.sign()
         else:

dreamer4/mocks.py

@@ -7,6 +7,11 @@ from torch.nn import Module
 
 from einops import repeat
 
+# helpers
+
+def exists(v):
+    return v is not None
+
 # mock env
 
 class MockEnv(Module):
@@ -15,7 +20,9 @@ class MockEnv(Module):
         image_shape,
         reward_range = (-100, 100),
         num_envs = 1,
-        vectorized = False
+        vectorized = False,
+        terminate_after_step = None,
+        rand_terminate_prob = 0.05
     ):
         super().__init__()
         self.image_shape = image_shape
@@ -25,6 +32,12 @@ class MockEnv(Module):
         self.vectorized = vectorized
         assert not (vectorized and num_envs == 1)
 
+        # mocking termination
+
+        self.can_terminate = exists(terminate_after_step)
+        self.terminate_after_step = terminate_after_step
+        self.rand_terminate_prob = rand_terminate_prob
+
         self.register_buffer('_step', tensor(0))
 
     def get_random_state(self):
@@ -50,13 +63,23 @@ class MockEnv(Module):
 
         reward = empty(()).uniform_(*self.reward_range)
 
-        if not self.vectorized:
-            return state, reward
+        if self.vectorized:
+            discrete, continuous = actions
+            assert discrete.shape[0] == self.num_envs, f'expected batch of actions for {self.num_envs} environments'
+
+            state = repeat(state, '... -> b ...', b = self.num_envs)
+            reward = repeat(reward, ' -> b', b = self.num_envs)
+
+        out = (state, reward)
+
+        if self.can_terminate:
+            terminate = (
+                (torch.rand((self.num_envs)) < self.rand_terminate_prob) &
+                (self._step > self.terminate_after_step)
+            )
 
-        discrete, continuous = actions
-        assert discrete.shape[0] == self.num_envs, f'expected batch of actions for {self.num_envs} environments'
+            out = (*out, terminate)
 
-        state = repeat(state, '... -> b ...', b = self.num_envs)
-        reward = repeat(reward, ' -> b', b = self.num_envs)
+        self._step.add_(1)
 
-        return state, reward
+        return out

{dreamer4-0.0.80.dist-info → dreamer4-0.0.82.dist-info}/METADATA

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

dreamer4-0.0.82.dist-info/RECORD

@@ -0,0 +1,8 @@
+dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
+dreamer4/dreamer4.py,sha256=1WeHCCJO-wa85Ki9uP8FR70pFTro46LF4WyFMxSD90I,110556
+dreamer4/mocks.py,sha256=S1kiENV3kHM1L6pBOLDqLVCZJg7ZydEmagiNb8sFIXc,2077
+dreamer4/trainers.py,sha256=D2b7WTgTHElLhIWLFgl2Ct2knGJLTk91HHpC5UkNvG0,14028
+dreamer4-0.0.82.dist-info/METADATA,sha256=by1KbrcKVy0epdLDLAiSN1MXR_FBgyD40Ol_Mn5iZNM,3065
+dreamer4-0.0.82.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dreamer4-0.0.82.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+dreamer4-0.0.82.dist-info/RECORD,,

dreamer4-0.0.80.dist-info/RECORD

@@ -1,8 +0,0 @@
-dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
-dreamer4/dreamer4.py,sha256=65MN3_ZqB3w5oXiWvpjRIHxxTwsACmnpw1HTd0GVJkU,108143
-dreamer4/mocks.py,sha256=Oi91Yv1oK0E-Wz-KDkf79xoyWzIXCvMLCr0WYCpJDLA,1482
-dreamer4/trainers.py,sha256=D2b7WTgTHElLhIWLFgl2Ct2knGJLTk91HHpC5UkNvG0,14028
-dreamer4-0.0.80.dist-info/METADATA,sha256=kyaYELFKTxlaWrjMYHGFiY1LtuOKl1_fBC1GSd9eo2A,3065
-dreamer4-0.0.80.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dreamer4-0.0.80.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-dreamer4-0.0.80.dist-info/RECORD,,