Dreamer 4: Training Agents Inside of Scalable World Models
11 min read
First learn a fast, action-conditioned video simulator. Then freeze it and let the policy improve through reinforcement learning on imagined experience, without another environment interaction.
training recipe
Learn the simulator, attach an agent, then practice inside it
The environment is never queried during policy improvement. Every imagined trajectory starts from a real dataset context and continues inside the learned world model.
World model pretraining
Train a causal video tokenizer, then learn latent dynamics from tokenized video. Actions are optional, so unlabeled clips still teach appearance, motion, and game mechanics.
video → latent simulator
Agent finetuning
Insert task tokens into the dynamics transformer. Multi-token heads learn the next eight actions and rewards while video prediction continues.
policy + reward heads
Imagination training
Generate rollouts with the policy and dynamics model. Train a value head with lambda-returns and improve the policy with PMPO while keeping the transformer frozen.
stronger offline policy
world model
Compress each frame, then predict the next latent from actions
The tokenizer and dynamics model share a block-causal 2D transformer pattern. Most layers mix space within a frame; every fourth layer reaches back through time.
Causal tokenizer
Masked image patches and learned latent tokens enter an encoder. A low-dimensional tanh bottleneck preserves time-compressed state, and the decoder reconstructs frames one by one.
Interactive dynamics
Clean next latent
Decode only for visualization. The policy can train directly on the latent rollout.
400M tokenizer
Continuous video representations with temporal compression.
1.6B dynamics model
256 spatial tokens and a 192-frame Minecraft context.
One efficient pattern
Space-time factorization, GQA, and sparse temporal layers.
shortcut forcing
Ask the model how large a denoising step to take
Conditioning on both signal level τ and requested step size d lets one network learn fine flow targets and bootstrap larger jumps. Dreamer 4 needs four model evaluations per generated frame.
Four learned jumps from noise to data
X-prediction
Predict the clean latent to reduce accumulating high-frequency error.
Ramp weight
w(τ) = 0.9τ + 0.1 focuses capacity where the target is informative.
Noised context
Past inputs use τctx = 0.1 for robustness to generated imperfections.
21.4 FPS
final 256-token model on one H100
4 steps
versus 64 for diffusion forcing
FVD 57
versus 306 for the naive baseline
policy improvement
The frozen simulator becomes an unlimited practice environment
The policy is constrained by a frozen behavioral prior, but it can discover better sequences than those demonstrated in the offline dataset.
Dataset context
Start from a real recorded state.
Policy action
Sample mouse and keyboard inputs.
Imagined transition
Dynamics predicts z, reward, and value.
PMPO update
Use the sign of each lambda-return advantage.
figures from the paper
Architecture, offline control, and action grounding
These crops preserve the authors' original diagrams. Together they show what the model is, what imagination training changes, and how little paired action data is needed.
The mental model
Dreamer 4 separates visual detail from decision making. A causal tokenizer compresses each video frame into continuous latents, and a large dynamics transformer learns how those latents change under low-level actions. Because missing actions can be replaced by a learned embedding, every video can teach the model about the world even when only a small subset teaches it how controls affect that world.
Interactive generation is the hard systems constraint. A normal diffusion model may need dozens of network evaluations for every frame. Shortcut forcing conditions the dynamics model on the requested denoising step size, so it learns to take four large, accurate steps. Predicting the clean latent directly, lightly noising past context, and factorizing attention keep long rollouts stable and fast enough for real-time interaction on one GPU.
After pretraining, task tokens expose small policy and reward heads inside the world model. Reinforcement learning then begins from recorded contexts and rolls the frozen simulator forward under the policy. The value head learns lambda-returns, while PMPO moves the policy toward actions with positive advantage and away from actions with negative advantage, constrained by the behavioral-cloning prior. That is how an offline dataset becomes an environment for additional practice.
Sources