World Models

A world model is a learned predictive model of environment dynamics: given the current state (or a history of observations) and an action, it predicts what happens next. Instead of interacting with the real environment (or an expensive simulator) for every learning update, an agent can “imagine” rollouts inside its own model. This makes world models a cornerstone of sample-efficient, model-based robot learning, and connects directly to how we formalize sequential decision making in Markov Decision Processes.

The core intuition: much of the cost in robot learning is the interaction itself (wear on hardware, slow real-time rollouts, safety risk). If we can compress experience into a model of “how the world responds,” we can do most of the heavy computation offline and in imagination.

What a world model predicts

Formally, in an MDP with states $s_t$, actions $a_t$, and rewards $r_t$, a world model approximates the transition and reward functions:

$${\hat{s}}_{t+1}\sim p_{\theta }(s_{t+1}\mid s_t,a_t),{\hat{r}}_t\sim p_{\theta }(r_t\mid s_t,a_t)$$

For robots we rarely observe the true state; we get high-dimensional observations $o_t$ (camera images, proprioception). So the model must also handle partial observability, effectively working with a belief over states.

Latent dynamics models

Predicting raw pixels many steps ahead is wasteful and unstable. The dominant approach instead learns a compact latent state $z_t$ and rolls the dynamics forward in that latent space. A typical recurrent state-space model (RSSM, as used in the Dreamer family) factorizes into:

$$\begin{array}{rl}\text{Encoder (posterior):}& z_t\sim q_{\theta }(z_t\mid z_{t-1},a_{t-1},o_t)\\ \text{Dynamics (prior):}& {\hat{z}}_t\sim p_{\theta }(z_t\mid z_{t-1},a_{t-1})\\ \text{Decoder:}& {\hat{o}}_t\sim p_{\theta }(o_t\mid z_t)\\ \text{Reward head:}& {\hat{r}}_t\sim p_{\theta }(r_t\mid z_t)\end{array}$$

The model is trained to (1) reconstruct observations and rewards and (2) keep the prior (which predicts without seeing $o_t$) close to the posterior (which does see it). The training objective is a variational (ELBO-style) loss:

\mathcal{L} = \underbrace{\mathbb{E}_q\!\left[\sum_t \log p_\theta(o_t \mid z_t) + \log p_\theta(r_t \mid z_t)\right]}_{\text{reconstruction}} - \underbrace{\sum_t \mathrm{KL}\!\left(q_\theta(z_t \mid \cdot)\,\|\,p_\theta(z_t \mid \cdot)\right)}_{\text{dynamics consistency}}}

The KL term is what makes the model usable for imagination: it forces the prior alone (no future observations) to be a good one-step predictor, so we can chain the prior forward to generate long latent rollouts. Latents are often split into a deterministic recurrent part and a stochastic part, which stabilizes long-horizon prediction while still capturing uncertainty.

Why latent, not pixels

Rolling out in a low-dimensional latent space is orders of magnitude cheaper than rendering images, avoids compounding pixel-level errors, and lets the reward/value predictors operate on a semantically meaningful representation.

Using a world model for planning

Once we have a model, one option is to skip learning an explicit policy and instead plan at decision time. Given the current latent state, we search for an action sequence that maximizes predicted return:

$$a_{t:t+H}^{\ast }=\arg \underset{a_{t:t+H}}{\max }\mathbb{E}\left[\sum _{k=0}^H{\gamma }^k{\hat{r}}_{t+k}\right]$$

Sampling-based planners such as the Cross-Entropy Method (CEM) or Model Predictive Path Integral (MPPI) evaluate many candidate action sequences in imagination, keep the best, and re-plan each step (model predictive control). This is the approach in PlaNet and in many MPC-based robot controllers.

CEM planning loop (per control step)

1. Sample $N$ action sequences from a Gaussian over the horizon $H$.
2. Roll each out through the latent dynamics model, sum predicted rewards.
3. Keep the top-$k$ “elite” sequences.
4. Refit the Gaussian’s mean and variance to the elites; repeat a few iterations.
5. Execute the first action of the best sequence, observe, re-plan.

Using a world model for policy learning (Dreamer style)

Planning at every step is expensive. The Dreamer line of work instead learns an actor and a critic entirely from imagined latent rollouts. After fitting the world model on replayed experience, it generates trajectories purely in latent space and trains:

• an actor ${\pi }_{\varphi }(a_t\mid z_t)$ to maximize predicted returns, and
• a critic $v_{\psi }(z_t)$ estimating the value of latent states, using a $\lambda$-return target so gradients flow through the differentiable model.

$$\varphi \leftarrow \arg \underset{\varphi }{\max }{\mathbb{E}}_{{\pi }_{\varphi },p_{\theta }}\left[\sum _t{V}^{\lambda }(z_t)\right]$$

Because imagined rollouts are cheap and differentiable, thousands of them can be generated per real environment step, giving strong sample efficiency. DreamerV2/V3 showed this scales across very different domains with a single set of hyperparameters, and real-robot variants (e.g. DayDreamer) train quadrupeds and manipulators directly from limited real interaction.

Model-based vs model-free tradeoffs

Axis	Model-based (world model)	Model-free
Sample efficiency	High: learns from imagined data	Low: needs many real interactions
Compute per update	High: must train and query the model	Lower per update
Asymptotic performance	Can be capped by model bias	Often higher given unlimited data
Failure mode	Model exploitation, compounding error	High variance, data hunger
Robot fit	Strong when interaction is costly/unsafe	Fine in fast, cheap simulators

The central tension is model bias: a policy optimized against an imperfect model may exploit its errors, choosing actions that look great in imagination but fail in reality. Mitigations include keeping rollouts short, propagating uncertainty (ensembles, stochastic latents), and continually correcting the model with fresh real data.

Common pitfalls

• Model exploitation: the policy finds and abuses regions where the model is wrong. Short imagined horizons and uncertainty penalties help.
• Compounding one-step error: small per-step prediction errors accumulate over long rollouts. This is why latent consistency (the KL term) and $\lambda$-returns matter.
• Reconstruction bias: pixel-reconstruction objectives spend capacity on visually large but task-irrelevant detail (backgrounds) while ignoring small, decision-critical objects.
• Distribution shift: the model is only accurate on states the agent has visited; imagination drifts into unseen regions unless the replay buffer is refreshed.

Connections

World models are one route across the Sim2Real gap: a learned model can serve as an adaptable, data-driven simulator, and “real to sim” model fitting is exactly world-model learning on hardware. The generative rollout idea also overlaps with Diffusion Policy, where a diffusion model captures multimodal action or trajectory distributions rather than one-step dynamics. And the whole framing rests on the Markov Decision Processes formalism: a world model is just a learned approximation of the MDP’s transition and reward functions.