World models and learned simulators

A world model is a simulator you learned instead of wrote. Its value is data efficiency and planning; its danger is that the policy will find and exploit its errors.

A model-free policy learns a mapping from observation to action. It has no internal notion of what happens next; it has only a compressed memory of what worked. Ask it to handle a situation slightly outside its data and it has nothing to reason with.

A world model learns the dynamics instead: given the current state and a proposed action, what does the next state look like? With that, a policy can search — try a thousand action sequences in imagination, pick the best, execute the first step, repeat. It can also train in imagination, which is where the data efficiency comes from.

Latent dynamics, not pixel prediction#

The naive formulation predicts the next image. This is expensive and wasteful: most pixels are irrelevant to control, and the model spends its capacity rendering wood grain.

The productive formulation, which the Dreamer line of work established, learns dynamics in a compact latent space:

encoder:     o_t          -> z_t              (observation to latent)
transition:  z_t, a_t     -> ẑ_{t+1}          (imagine the next latent)
decoder:     z_t          -> ô_t              (only used for training signal)
reward:      z_t          -> r̂_t             (what the policy optimizes)

Only the transition model runs during planning. A 30-dimensional latent rolled forward 50 steps is trivial; 50 steps of image generation is not. That gap is why the architecture looks the way it does.

The decoder exists to force the latent to retain information. It is a training device, not a component of the deployed system.

Two design choices dominate quality:

Stochastic versus deterministic latents. Real dynamics are partly unpredictable — a cloth settles differently each time. A purely deterministic transition model averages over the possibilities and produces blurred, low-confidence predictions. The standard resolution is a recurrent state with both a deterministic path (carries information forward reliably) and a stochastic one (models genuine uncertainty).

Discrete versus continuous latents. Categorical latents — a set of one-hot vectors rather than a Gaussian — turn out to work notably better for visual domains. The usual explanation is that they suit multimodal futures: a discrete code can represent "the cup tipped over" or "the cup did not" without having to represent the average of the two, which is not a state the world can be in.

Planning in imagination#

With a transition model you can do model-predictive control without writing any physics. The workhorse is a sampling-based planner:

def cem_plan(world_model, z0, horizon=12, iters=6, n=512, elite=64):
    """Cross-entropy method: sample action sequences, keep the best, refit,
    repeat. Embarrassingly parallel and needs no gradients through dynamics."""
    mu = torch.zeros(horizon, ACTION_DIM)
    sigma = torch.ones(horizon, ACTION_DIM) * 0.5

    for _ in range(iters):
        seqs = (mu + sigma * torch.randn(n, horizon, ACTION_DIM)).clamp(-1, 1)
        z = z0.expand(n, -1)
        total = torch.zeros(n)
        for t in range(horizon):
            z = world_model.transition(z, seqs[:, t])
            total += world_model.reward(z) * (GAMMA ** t)

        top = total.topk(elite).indices
        mu = seqs[top].mean(0)
        sigma = seqs[top].std(0) + 1e-3        # floor, or it collapses
    return mu[0]                               # execute one step, then re-plan

The sigma floor is not cosmetic. Without it the distribution collapses to a point after a few iterations and the planner stops exploring, which manifests as a policy that commits early to a mediocre plan and never reconsiders.

The exploitation problem#

Here is the failure that defines the field. The planner is an optimizer, and it is optimizing against a learned reward and a learned dynamics model. Optimizers find errors. If there is a region of latent space where the model wrongly predicts high reward — and there always is, because the model was fit on finite data — the planner will find it and go there with great enthusiasm.

The result is a policy that achieves spectacular imagined returns and does nothing useful in reality. This is not an edge case; it is the default outcome without mitigation.

The mitigations that work:

Video models as world models#

The recent thread is to skip the compact latent entirely and use a large video-generation model as the dynamics model: condition on the current frame and an action, generate the next frames. These models have absorbed an enormous amount of visual physics from web-scale video, and their predictions look remarkably plausible.

Plausible is doing a lot of work in that sentence. The honest assessment as we see it today:

• They are genuinely good at appearance dynamics — how a scene looks as things move — and much weaker on precise contact and force.
• They are expensive enough that planning with them in a control loop is not yet practical. They are used to generate training data or to evaluate candidate plans offline, not to run MPC at 10 Hz.
• Physical consistency is not guaranteed. Objects can drift in mass, occluded objects can fail to reappear, and small errors accumulate over long rollouts in ways that are hard to detect automatically.

The direction is promising and the current practical use is narrower than the demos imply. We would use one to augment a dataset today; we would not put one inside a control loop.

Where they earn their place#

World models are worth the complexity in three situations:

Data is expensive and resets are cheap. The classic model-based advantage: an order of magnitude fewer environment interactions to reach the same performance. For robots, where every episode costs a human's time, this is the argument.

The task requires lookahead. Anything where the right action now depends on a consequence several steps out — pushing an object to a position from which it can be grasped, sequencing operations with ordering constraints. Reactive policies handle these badly.

You need to ask counterfactual questions. "What if the arm had gone left?" is answerable with a dynamics model and not otherwise. This is useful for debugging and for offline evaluation of policy changes without touching hardware.

They are not worth it when the task is short-horizon and reactive, when demonstrations are cheap, or when the contact dynamics are the whole problem and your model will be wrong about exactly the thing that matters. In those cases a well-collected imitation dataset and a good action head beat a learned simulator, at a fraction of the engineering.