dreamer4 0.0.81__tar.gz → 0.0.83__tar.gz

{dreamer4-0.0.81 → dreamer4-0.0.83}/PKG-INFO

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

{dreamer4-0.0.81 → dreamer4-0.0.83}/dreamer4/dreamer4.py

@@ -1819,6 +1819,7 @@ class DynamicsWorldModel(Module):
         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,
@@ -2029,6 +2030,8 @@ class DynamicsWorldModel(Module):
         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.))
 
@@ -2115,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
@@ -2125,11 +2130,21 @@ 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)
+        is_truncated = full((batch,), False, device = device)
+
+        episode_lens = full((batch,), 0, device = device)
+
         # maybe time kv cache
 
         time_kv_cache = None
 
-        for i in range(max_timesteps + 1):
+        step_index = 0
+
+        while not is_terminated.all():
+            step_index += 1
 
             latents = self.video_tokenizer(video, return_latents = True)
 
@@ -2187,7 +2202,40 @@ 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
+                terminated = full((batch,), False)
+                truncated = full((batch,), False)
+
+            elif len(env_step_out) == 3:
+                next_frame, reward, terminated = env_step_out
+                truncated = full((batch,), False)
+
+            elif len(env_step_out) == 4:
+                next_frame, reward, terminated, truncated = env_step_out
+
+            # update episode lens
+
+            episode_lens = torch.where(is_terminated, episode_lens, episode_lens + 1)
+
+            # update `is_terminated`
+
+            # (1) - environment says it is terminated
+            # (2) - previous step is truncated (this step is for bootstrap value)
+
+            is_terminated |= (terminated | is_truncated)
+
+            # update `is_truncated`
+
+            if step_index <= max_timesteps:
+                is_truncated |= truncated
+
+            if step_index == max_timesteps:
+                # if the step index is at the max time step allowed, set the truncated flag, if not already terminated
+
+                is_truncated |= ~is_terminated
 
             # batch and time dimension
 
@@ -2216,7 +2264,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_truncated,
+            lens = episode_lens,
             is_from_world_model = False
         )
 
@@ -2284,13 +2333,22 @@ class DynamicsWorldModel(Module):
 
             decay = 1. - self.reward_ema_decay
 
-            # todo - handle distributed
+            # 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

dreamer4-0.0.83/dreamer4/mocks.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+from random import choice
+
+import torch
+from torch import tensor, empty, randn, randint
+from torch.nn import Module
+
+from einops import repeat
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+# mock env
+
+class MockEnv(Module):
+    def __init__(
+        self,
+        image_shape,
+        reward_range = (-100, 100),
+        num_envs = 1,
+        vectorized = False,
+        terminate_after_step = None,
+        rand_terminate_prob = 0.05,
+        can_truncate = False,
+        rand_truncate_prob = 0.05,
+    ):
+        super().__init__()
+        self.image_shape = image_shape
+        self.reward_range = reward_range
+
+        self.num_envs = num_envs
+        self.vectorized = vectorized
+        assert not (vectorized and num_envs == 1)
+
+        # mocking termination and truncation
+
+        self.can_terminate = exists(terminate_after_step)
+        self.terminate_after_step = terminate_after_step
+        self.rand_terminate_prob = rand_terminate_prob
+
+        self.can_truncate = can_truncate
+        self.rand_truncate_prob = rand_truncate_prob
+
+        self.register_buffer('_step', tensor(0))
+
+    def get_random_state(self):
+        return randn(3, *self.image_shape)
+
+    def reset(
+        self,
+        seed = None
+    ):
+        self._step.zero_()
+        state = self.get_random_state()
+
+        if self.vectorized:
+            state = repeat(state, '... -> b ...', b = self.num_envs)
+
+        return state
+
+    def step(
+        self,
+        actions,
+    ):
+        state = self.get_random_state()
+
+        reward = empty(()).uniform_(*self.reward_range)
+
+        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:
+            shape = (self.num_envs,) if self.vectorized else (1,)
+            valid_step = self._step > self.terminate_after_step
+
+            terminate = (torch.rand(shape) < self.rand_terminate_prob) & valid_step
+
+            out = (*out, terminate)
+
+            # maybe truncation
+
+            if self.can_truncate:
+                truncate = (torch.rand(shape) < self.rand_truncate_prob) & valid_step & ~terminate
+                out = (*out, truncate)
+
+        self._step.add_(1)
+
+        return out

{dreamer4-0.0.81 → dreamer4-0.0.83}/pyproject.toml

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

{dreamer4-0.0.81 → dreamer4-0.0.83}/tests/test_dreamer.py

@@ -612,9 +612,13 @@ def test_cache_generate():
 
 @param('vectorized', (False, True))
 @param('use_signed_advantage', (False, True))
+@param('env_can_terminate', (False, True))
+@param('env_can_truncate', (False, True))
 def test_online_rl(
     vectorized,
-    use_signed_advantage
+    use_signed_advantage,
+    env_can_terminate,
+    env_can_truncate
 ):
     from dreamer4.dreamer4 import DynamicsWorldModel, VideoTokenizer
 
@@ -649,7 +653,14 @@ def test_online_rl(
     from dreamer4.mocks import MockEnv
     from dreamer4.dreamer4 import combine_experiences
 
-    mock_env = MockEnv((256, 256), vectorized = vectorized, num_envs = 4)
+    mock_env = MockEnv(
+        (256, 256),
+        vectorized = vectorized,
+        num_envs = 4,
+        terminate_after_step = 2 if env_can_terminate else None,
+        can_truncate = env_can_truncate,
+        rand_terminate_prob = 0.1
+    )
 
     # manually
 

dreamer4-0.0.81/dreamer4/mocks.py

@@ -1,62 +0,0 @@
-from __future__ import annotations
-from random import choice
-
-import torch
-from torch import tensor, empty, randn, randint
-from torch.nn import Module
-
-from einops import repeat
-
-# mock env
-
-class MockEnv(Module):
-    def __init__(
-        self,
-        image_shape,
-        reward_range = (-100, 100),
-        num_envs = 1,
-        vectorized = False
-    ):
-        super().__init__()
-        self.image_shape = image_shape
-        self.reward_range = reward_range
-
-        self.num_envs = num_envs
-        self.vectorized = vectorized
-        assert not (vectorized and num_envs == 1)
-
-        self.register_buffer('_step', tensor(0))
-
-    def get_random_state(self):
-        return randn(3, *self.image_shape)
-
-    def reset(
-        self,
-        seed = None
-    ):
-        self._step.zero_()
-        state = self.get_random_state()
-
-        if self.vectorized:
-            state = repeat(state, '... -> b ...', b = self.num_envs)
-
-        return state
-
-    def step(
-        self,
-        actions,
-    ):
-        state = self.get_random_state()
-
-        reward = empty(()).uniform_(*self.reward_range)
-
-        if not self.vectorized:
-            return state, reward
-
-        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)
-
-        return state, reward