From 38cba80068306e2d3d8957bd9f6586b9ce006329 Mon Sep 17 00:00:00 2001 From: lucidrains Date: Tue, 4 Nov 2025 06:05:11 -0800 Subject: [PATCH] readme --- README.md | 27 +++++++++++++++++++++++---- 1 file changed, 23 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 685729a..5b608a9 100644 --- a/README.md +++ b/README.md @@ -28,9 +28,16 @@ tokenizer = VideoTokenizer( image_width = 256 ) +video = torch.randn(2, 3, 10, 256, 256) + +# learn the tokenizer + +loss = tokenizer(video) +loss.backward() # ler + # dynamics world model -dynamics = DynamicsWorldModel( +world_model = DynamicsWorldModel( dim = 512, dim_latent = 32, video_tokenizer = tokenizer, @@ -46,7 +53,7 @@ rewards = torch.randn(2, 10) # learn dynamics / behavior cloned model -loss = dynamics( +loss = world_model( video = video, rewards = rewards, discrete_actions = discrete_actions @@ -58,7 +65,7 @@ loss.backward() # then generate dreams -dreams = dynamics.generate( +dreams = world_model.generate( 10, batch_size = 2, return_decoded_video = True, @@ -67,7 +74,19 @@ dreams = dynamics.generate( # learn from the dreams -actor_loss, critic_loss = dynamics.learn_from_experience(dreams) +actor_loss, critic_loss = world_model.learn_from_experience(dreams) + +(actor_loss + critic_loss).backward() + +# learn from environment + +from dreamer4.mocks import MockEnv + +mock_env = MockEnv((256, 256), vectorized = True, num_envs = 4) + +experience = world_model.interact_with_env(mock_env, max_timesteps = 8, env_is_vectorized = True) + +actor_loss, critic_loss = world_model.learn_from_experience(experience) (actor_loss + critic_loss).backward() ```