World Models: Learning to Simulate Consequences

Model-based reinforcement learning · technical foundations

Suppose an agent sees a camera frame, presses a button, and needs to decide what to do next. A model-free policy maps the current input directly to an action. A world-model agent can ask: “If I press left, what state and reward are likely next? What if I press right?” That internal counterfactual is the source of planning.

The formal object

In a partially observed environment, the true state st is hidden. The agent receives an observation ot and takes an action at. A learned world model maintains a useful internal state zt:

z_t = encoder(o_t, history)
(z_{t+1}, r_t, done_t) ~ model(z_t, a_t)

The model approximates the environment's transition and reward dynamics. The distribution symbol matters: real environments can be stochastic, so a good model may need to represent several plausible futures rather than one average future.

Why a latent state?

Raw pixels contain enormous detail: texture, lighting, and background motion. Most of it is irrelevant to the next decision. The encoder compresses observations into a latent representation that tries to preserve decision-relevant information.

Learning and using the model

Training usually begins with trajectories of observations, actions, and rewards. The model is optimized so that its predictions match what actually followed. Once it is adequate, the agent can generate imagined rollouts:

z0 = encode(current_observation)
for t in horizon:
    action = policy(z)
    z, reward = world_model(z, action)
    learn_from_imagined_trajectory()

Dreamer propagates value-learning signals through imagined latent trajectories. MuZero takes a different route: its learned model predicts the quantities needed by tree search—reward, value, and policy—without reconstructing every observation. The model is therefore judged by decisions, not visual fidelity.

The catch: imagination can lie

Prediction errors compound. If each step is slightly wrong, a long rollout can enter states the model has never seen. The policy may then exploit a hallucinated reward loophole. This is why model-based agents often use short-horizon replanning, uncertainty estimates, conservative objectives, or continual collection of real data.

A useful mental comparison

Primary sources

• Ha & Schmidhuber, World Models
• Hafner et al., Dream to Control
• Schrittwieser et al., MuZero

Next: Ask me to derive a simple latent-dynamics loss, compare Dreamer with MuZero, or build a tiny world model in Python.