make sure time kv cache can be passed back in during generation

dreamer4/dreamer4.py

@@ -2069,6 +2069,7 @@ class DynamicsWorldModel(Module):
         image_width = None,
         return_decoded_video = None,
         context_signal_noise = 0.1,       # they do a noising of the past, this was from an old diffusion world modeling paper from EPFL iirc
+        time_kv_cache: Tensor | None = None,
         use_time_kv_cache = True,
         return_rewards_per_frame = False,
         return_agent_actions = False,
@@ -2120,13 +2121,10 @@ class DynamicsWorldModel(Module):
 
         # maybe return rewards
 
+        decoded_rewards = None
         if return_rewards_per_frame:
             decoded_rewards = empty((batch_size, 0), device = self.device, dtype = torch.float32)
 
-        # handle maybe time kv cache
-
-        time_kv_cache = None
-
         # while all the frames of the video (per latent) is not generated
 
         while latents.shape[1] < time_steps:

pyproject.toml

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

tests/test_dreamer.py

@@ -473,3 +473,33 @@ def test_tokenizer_trainer():
     )
 
     trainer()
+
+def test_cache_generate():
+    from dreamer4.dreamer4 import DynamicsWorldModel
+
+    dynamics = DynamicsWorldModel(
+        dim = 16,
+        dim_latent = 16,
+        max_steps = 64,
+        num_tasks = 4,
+        num_latent_tokens = 4,
+        depth = 4,
+        num_spatial_tokens = 1,
+        pred_orig_latent = True,
+        num_discrete_actions = 4,
+        attn_dim_head = 16,
+        prob_no_shortcut_train = 0.1,
+        num_residual_streams = 1
+    )
+
+    generated, time_kv_cache = dynamics.generate(1, return_time_kv_cache = True)
+    generated, time_kv_cache = dynamics.generate(1, time_kv_cache = time_kv_cache, return_time_kv_cache = True)
+    generated, time_kv_cache = dynamics.generate(1, time_kv_cache = time_kv_cache, return_time_kv_cache = True)
+
+@param('vectorized', (False, True))
+def test_online_rl(
+    vectorized
+):
+    from dreamer4.mocks import MockEnv
+
+    mock_env = MockEnv((256, 256), vectorized = vectorized, batch_size = 4)