World Models for Planning

A world model is a learned internal representation of an environment's transition dynamics, reward structure, and observation distribution, enabling an agent to perform mental simulation — planning action sequences by rolling out imagined trajectories without interacting with the real environment. World models matter because they offer sample efficiency gains that purely model-free methods cannot achieve, decoupling policy improvement from costly real-world experience. Remarkably, the same latent-space dynamics models that enable planning in pixel-space environments are structurally equivalent to the predictive coding hierarchies proposed by neuroscientists decades before deep learning existed.

The core insight

The central insight is that planning is just optimization over an imagined future: if you possess a sufficiently accurate model $\hat{p}(s_{t+1} \mid s_t, a_t)$ and reward signal $\hat{r}(s_t, a_t)$ , you can find high-return action sequences by backpropagating through the model's rollouts — effectively turning a sequential decision problem into a differentiable program. The power lies in where you represent states: modern world models operate in a compressed latent space, not in raw observation space, so the dynamics are smoother, lower-dimensional, and far cheaper to simulate. The aha moment is recognizing that a world model is not a side module that assists planning — it is the substrate in which planning becomes tractable. This reframing unifies model predictive control, Dyna-style architectures, and latent imagination under a single principle.

Common trapCommon Misconception

Practitioners routinely conflate model accuracy with planning utility — assuming that a more faithful reconstruction of observations (lower pixel MSE, higher ELBO) implies better plans. This is wrong because planning quality depends on the geometry of the latent dynamics near high-value trajectories, not on global reconstruction fidelity. A model can reconstruct backgrounds perfectly while representing action-relevant state transitions poorly; conversely, a heavily abstracted latent model with poor pixel reconstruction can support near-optimal planning if it preserves the topology of the MDP's value landscape. The correct mental model: evaluate world models by the regret of the policies they produce, not by reconstruction metrics. This distinction drives much of the design philosophy behind DreamerV3 and TD-MPC2, where representation learning objectives are explicitly shaped by the planning task rather than reconstruction loss alone.

How It Works

Formal Setup

Let the environment be a partially observable Markov decision process (POMDP) $(\mathcal{S}, \mathcal{A}, \mathcal{O}, T, R, \Omega, \gamma)$ where $T(s' \mid s, a)$ is the true transition kernel, $R(s,a)$ the reward function, and $\Omega(o \mid s)$ the observation emission. A world model learns approximate counterparts $(\hat{T}_\theta, \hat{R}_\theta, \hat{\Omega}_\theta)$ jointly with an encoder $q_\phi(z_t \mid o_{\leq t}, a_{< t})$ mapping histories to latent states. The Recurrent State Space Model (RSSM) factorizes the latent state into deterministic $h_t$ and stochastic $z_t$ components:

\quad h_t = f_\theta(h_{t-1},\, z_{t-1},\, a_{t-1}), \qquad z_t \sim p_\theta(z_t \mid h_t)

Explanation: The deterministic recurrent state $h_t$ carries long-horizon memory through the sequence, while the stochastic component $z_t$ captures aleatoric uncertainty at each step, enabling the model to represent distributions over futures rather than point predictions.

Symbols: $h_t \in \mathbb{R}^{d_h}$ — deterministic hidden state (GRU output); $z_t$ — categorical or Gaussian stochastic latent; $f_\theta$ — recurrent cell; $a_{t-1}$ — previous action; $p_\theta$ — prior over stochastic latent given history.

Variational Training Objective

The world model is trained by maximizing a sequence-level Evidence Lower BOund (ELBO) over observed trajectories:

\quad \mathcal{L}(\theta, \phi) = \mathbb{E}_{q_\phi}\!\left[\sum_{t=1}^{T} \underbrace{\ln \hat{\Omega}_\theta(o_t \mid h_t, z_t)}_{\text{reconstruction}} - \underbrace{\beta\, D_{\mathrm{KL}}\!\bigl(q_\phi(z_t \mid h_t, o_t)\,\|\, p_\theta(z_t \mid h_t)\bigr)}_{\text{representation regularization}}\right]

Explanation: The reconstruction term forces the latent state to encode information sufficient to explain observations, while the KL term aligns the posterior with the model's own prior, ensuring the prior can generate coherent imagined rollouts at plan time without access to future observations.

Symbols: $q_\phi(z_t \mid h_t, o_t)$ — posterior encoder (uses current observation); $p_\theta(z_t \mid h_t)$ — prior (observation-free, used during imagination); $\beta$ — KL weight balancing reconstruction vs. regularization; $D_{\mathrm{KL}}$ — Kullback-Leibler divergence.

Latent Imagination and Policy Optimization

With the world model fixed, imagined trajectories are generated entirely in latent space starting from an encoded real state $z_0$ :

\quad \hat{z}_{t+1} \sim p_\theta(\hat{z}_{t+1} \mid h_{t+1}),\quad h_{t+1} = f_\theta(h_t, \hat{z}_t, a_t),\quad a_t \sim \pi_\psi(a_t \mid h_t, \hat{z}_t)

Explanation: The agent unrolls imagined trajectories of length $H$ by alternating between the learned prior transition and the current policy, collecting a synthetic dataset on which any RL algorithm can operate without touching the real environment. The policy $\pi_\psi$ is optimized using $\lambda$ -returns computed over imagined rollouts. For differentiable world models, gradients flow through time back to $\psi$ :

\quad \mathcal{J}(\psi) = \mathbb{E}_{\pi_\psi, p_\theta}\!\left[\sum_{t=0}^{H-1} \gamma^t \hat{R}_\theta(h_t, \hat{z}_t, a_t) + \gamma^H V_\xi(h_H, \hat{z}_H)\right]

Explanation: The policy objective is the discounted sum of predicted rewards along an imagined horizon plus a bootstrapped value estimate at the terminal imagined state, all of which are differentiable with respect to $\psi$ via the reparameterization trick when $z_t$ is continuous or straight-through estimators when categorical.

Symbols: $H$ — imagination horizon; $\gamma$ — discount factor; $V_\xi$ — learned value function; $\hat{R}_\theta$ — world model's reward head; $\pi_\psi$ — policy with parameters $\psi$ .

Model Predictive Control Variant

For systems where a differentiable policy is undesirable, planning proceeds via sampling-based MPC using the Cross-Entropy Method (CEM) or MPPI:

\quad a^*_{0:H} = \arg\max_{a_{0:H}} \mathbb{E}_{p_\theta}\!\left[\sum_{t=0}^{H} \gamma^t \hat{R}_\theta(h_t, \hat{z}_t, a_t)\right]

Explanation: Rather than learning a parametric policy, this formulation treats planning as an online optimization problem at each decision step — the world model serves as a surrogate objective over action sequences, and CEM iteratively refits a Gaussian proposal toward high-return candidates.

researcher · Richard S. Sutton | b. 1949 | Dyna Architecture and Model-Based RL Foundations

Sutton trained under Andrew Barto at UMass Amherst and developed much of the foundational theory of temporal-difference learning and policy gradient methods. His 1991 Dyna paper introduced the landmark insight that model-generated ("imagined") transitions can be interleaved with real experience during Q-learning, prefiguring all modern latent-imagination approaches by three decades. Sutton later formalized the deadly triad problem — the instability arising from combining function approximation, bootstrapping, and off-policy learning — which remains the central challenge in stabilizing world model training today.

researcher · David Ha | b. 1980 | World Models Paper and Latent-Space Planning

Ha worked as a researcher at Google Brain and later as Chief Scientist at Stability AI, bridging generative modeling and reinforcement learning. His 2018 "World Models" paper with Jürgen Schmidhuber demonstrated that a compressed VAE+MDN-RNN model of a visual environment could support policy training entirely within the dream, achieving competitive performance on Car Racing and VizDoom without any real-environment policy gradient steps. The surprising personal fact: Ha originally trained as a traditional artist and composer before transitioning to ML, and explicitly draws on aesthetic intuitions about compression and representation when designing latent architectures.

researcher · Danijar Hafner | b. 1993 | DreamerV1/V2/V3 and RSSM Design

Hafner completed his PhD at UC Berkeley and conducted postdoctoral work at Google DeepMind, where he developed the Dreamer series of world model agents. His Recurrent State Space Model (RSSM) — separating deterministic and stochastic latent components — became the dominant architecture for latent-imagination planning, and DreamerV3 (2023) achieved human-level performance on the majority of Atari games while simultaneously solving continuous locomotion and Minecraft diamond collection, with a single fixed hyperparameter set across all domains. Remarkably, DreamerV3 was the first algorithm to collect a diamond in Minecraft from scratch without human demonstrations or task-specific reward shaping.

Figure

World Model Planning Architecture

Figure

Comparison: Model-Free vs. World Model Planning

Intuition

The Mental Model

Think of a grandmaster chess player versus a beginner. Both can look at the board — the analogue of observing $o_t$ — but the grandmaster has an internalized mental model of chess dynamics: they can rapidly simulate dozens of future board states in their head, prune implausible lines, and identify the move that leads to a favorable position five steps later. The beginner must "try it and see," relying purely on reactive pattern matching. A learned world model is exactly this internalized chess engine — compressed, approximate, but sufficiently accurate over a planning horizon to guide good decisions. The critical subtlety: the grandmaster doesn't simulate the full board with pixel-perfect fidelity. They think in terms of piece relationships, pawn structure, king safety — an abstract latent space that captures what matters for planning while discarding irrelevant details (exact piece color saturation, board grain texture). This is why world models learned in latent space outperform those that try to predict raw observations step by step: the right level of abstraction makes futures predictable enough to be useful without requiring the impossible precision of full sensory reconstruction.

Historical Development

Origins · 1 minhow this was discovered

History

1989

Sutton introduces the Dyna architecture — the first algorithmic framework explicitly interleaving model-generated experience with real transitions during Q-learning, establishing the theoretical basis for model-based RL.

1990

Schmidhuber proposes "A possibility for implementing curiosity and boredom" — using a learned world model's prediction error as an intrinsic reward signal, anticipating modern curiosity-driven exploration.

2002