hongshaorou/dreamer4
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
lucidrains 2025-10-21 09:15:32 -07:00
parent 7ba3988fb9
commit b34128d3d0
3 changed files with 33 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
dreamer4/dreamer4.py
View File

@ -2069,6 +2069,7 @@ class DynamicsWorldModel(Module):
image_width = None, image_width = None,
return_decoded_video = 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 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, use_time_kv_cache = True,
return_rewards_per_frame = False, return_rewards_per_frame = False,
return_agent_actions = False, return_agent_actions = False,
@ -2120,13 +2121,10 @@ class DynamicsWorldModel(Module):
# maybe return rewards # maybe return rewards
decoded_rewards = None
if return_rewards_per_frame: if return_rewards_per_frame:
decoded_rewards = empty((batch_size, 0), device = self.device, dtype = torch.float32) 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 all the frames of the video (per latent) is not generated
while latents.shape[1] < time_steps: while latents.shape[1] < time_steps:

2
pyproject.toml
View File

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

30
tests/test_dreamer.py
View File

@ -473,3 +473,33 @@ def test_tokenizer_trainer():
) )
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)