World Models

Intuition

A chess engine does not play moves on a real board and wait for consequences; it imagines future positions and evaluates them. A world model gives a neural agent the same capability: instead of collecting every training example by interacting with the real environment (which is slow, expensive, or dangerous), the agent rolls out imagined trajectories inside a learned model of the environment.

The idea predates deep learning: Schmidhuber proposed self-modeling networks in 1990, and Sutton’s “Dyna” architecture (1991) interleaved real and imagined rollouts in tabular RL [6]. What modern deep world models add is the ability to model rich, high-dimensional observations (pixels, language) using neural representations.

What makes a model a “world model”?

A model $M$ is a world model if it predicts future states $s_{t+1}$ from current state $s_t$ and action $a_t$:

A pure generative model (e.g., an unconditional video diffusion model) generates plausible-looking sequences but does not have an explicit action-conditioning mechanism. A world model is action-conditioned: different actions lead to different predicted futures, and the model’s internal dynamics encode the causal structure of the environment.

A world model supports:

• Planning: roll out multiple action sequences, pick the one leading to the highest estimated return.
• Policy training: generate imagined experience to train a policy without real environment interaction.
• Counterfactual reasoning: ask “what would have happened if I had taken action $a'$?”

Ha and Schmidhuber (2018): Dream to learn

The Ha and Schmidhuber “World Models” paper [1] introduced the V-M-C (Vision-Memory-Controller) architecture:

• V (VAE): a convolutional variational autoencoder that compresses each frame $o_t$ to a low-dimensional latent vector $z_t$.
• M (MDN-RNN): a mixture-density RNN that models the distribution over future latent states $P(z_{t+1} | z_t, a_t, h_t)$, where $h_t$ is the hidden state of the RNN. The mixture-density output allows the model to capture multi-modal next-state distributions.
• C (linear controller): a small linear layer that maps $[z_t, h_t]$ to actions. Because C is linear, it can be optimised cheaply with CMA-ES (a black-box evolution strategy).

The key experiment: the controller C can be trained entirely on rollouts generated inside the M model (the “dream”), then deployed in the real environment. On CarRacing-v0, the policy trained in the dream matched or exceeded policies trained in the real environment, demonstrating that world-model planning can serve as a training environment.

Dreamer and DreamerV3

Dreamer [4] and its successor DreamerV3 [2] replace the V-M-C decomposition with an end-to-end differentiable world model called the Recurrent State Space Model (RSSM), trained jointly with the policy.

DreamerV3 [2] extends the architecture with:

• Symlog predictions (predict the symlog-transformed reward to handle varying reward scales across domains).
• KL balancing (separately scale the KL towards the prior vs the posterior to prevent posterior collapse).
• Free bits (a minimum KL per dimension to prevent degenerate solutions).

With these changes, DreamerV3 trains a single set of hyperparameters that works across Atari, continuous control, Minecraft, and BSuite — the first model-based RL agent to achieve human-level performance in Minecraft (collecting diamonds) from scratch.

MuZero: planning without a reconstruction decoder

MuZero [5] takes a different architectural stance: it learns a world model without an observation decoder. Instead of reconstructing $o_t$, the latent state $s_t$ is learned end-to-end to be useful for value and policy prediction. Three learned functions:

• Representation $h_\theta$: $o_t \to s_t$ (encode observations to latent).
• Dynamics $g_\theta$: $(s_t, a_t) \to (s_{t+1}, r_t)$ (transition + reward).
• Prediction $f_\theta$: $s_t \to (p_t, v_t)$ (policy logits + value estimate).

Planning uses MCTS over the latent state space. Because the dynamics function is learned, MuZero can plan in games (Go, chess, shogi, Atari) without access to a simulator, achieving superhuman performance on all of them.

Video generation as world modeling

A video generation model that accepts action tokens as conditioning implicitly learns a world model: given frame $o_t$ and action $a_t$, generate $o_{t+1}$. Two exemplars:

GAIA-1 (Wayve, 2023): a 9 billion parameter autoregressive video model trained on driving data. Given a driving video and a sequence of actions (steering, throttle), GAIA-1 generates photorealistic future frames consistent with the vehicle dynamics. The generated video can be used to train downstream driving policies, providing a simulation engine that generalises to novel scenarios.

Genie (Google DeepMind, 2024) [3]: a 11 billion parameter foundation world model trained on 2D platformer gameplay video, without any action labels. Genie learns to infer latent actions from video transitions via a spatiotemporal video tokenizer, a latent action model, and a dynamics model. Given a single image, it can generate an interactive environment: pressing left/right/jump produces consistent future frames obeying the inferred game physics. The model generalises to novel hand-drawn scene images never seen in training, demonstrating that world model knowledge can transfer across visual domains.

Open problems

Partial observability: real environments are partially observable — the agent never sees the full state. DreamerV3 handles this via the stochastic latent $z_t$, but long-horizon memory (remembering which doors were locked three rooms ago) remains challenging.

Physical plausibility: video generation models can produce visually coherent but physically implausible sequences (objects passing through walls, incorrect rigid-body dynamics). Training on physics-engine-simulated data or incorporating physics priors is an active area.

Long-horizon consistency: world models tend to drift over many steps. A 10-step imagined rollout is plausible; a 1000-step rollout diverges from the true environment dynamics. Hierarchical world models (planning at multiple time scales) and uncertainty-aware planning (pulling back to real experience when the model is uncertain) are current research directions.

Scaling world models: LeCun (2022) [7] argues that world models are the central missing ingredient for human-level autonomous intelligence, proposing a Joint Embedding Predictive Architecture (JEPA) that learns self-supervised world models in abstract feature space rather than pixel space. Whether scaling video-generation-style world models eventually solves this is an open question.