World models

The idea

A world model predicts how the world changes: given the current state (or observation) and an action, what comes next? A policy chooses actions; a world model answers "what would happen if". With one you can:

• plan by imagining several futures and picking the best;
• train a policy on imagined experience instead of the real robot or road;
• simulate rare situations that are hard or dangerous to collect.

A world model you can run

Before the large systems, the smallest one. A pendulum has state (angle, angular velocity); the action is a torque. A network learns the one-step map $s_{t+1} = s_t + f_\theta(s_t, a_t)$ from trajectories under random actions, and we ask three questions the big systems also face.

class DynamicsModel(nn.Module):
    """Predicts the change in state from (state, action): s' = s + f(s, a)."""

    def __init__(self, hidden: int = 64) -> None:
        """Creates the MLP."""
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(3, hidden),
            nn.Tanh(),
            nn.Linear(hidden, hidden),
            nn.Tanh(),
            nn.Linear(hidden, 2),
        )

    def forward(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
        """The predicted next state."""
        return state + self.net(torch.cat([state, action[..., None]], dim=-1))


def train_model(
    states: torch.Tensor,
    actions: torch.Tensor,
    unroll: int = 1,
    steps: int = 1500,
    seed: int = 0,
) -> DynamicsModel:
    """Fits the model on windows of ``unroll`` steps.

    With ``unroll`` = 1 this is ordinary one-step prediction (teacher forcing: every
    input is a real state). With more, the model is rolled forward on its own
    predictions inside the loss and penalised at every step, which exposes it during
    training to the states it will drift into at test time.
    """
    torch.manual_seed(seed)
    model = DynamicsModel()
    optimizer = torch.optim.Adam(model.parameters(), lr=3e-3)
    n, horizon = actions.shape
    for _ in range(steps):
        i = torch.randint(0, n, (128,))
        t = torch.randint(0, horizon - unroll + 1, (128,))
        state = states[i, t]
        loss = 0.0
        for k in range(unroll):
            state = model(state, actions[i, t + k])
            loss = loss + ((state - states[i, t + k + 1]) ** 2).mean()
        optimizer.zero_grad()
        (loss / unroll).backward()
        optimizer.step()
    return model

1. Does the error compound? Roll the model forward on its own predictions and compare with the true system. The error after one step is tiny (0.004) and grows steadily with the horizon (0.03 after 10 steps, 0.11 after 40, mean of 3 seeds), because each prediction feeds the next one and small mistakes accumulate. This is the compounding error of the Imitation learning module, now in the model instead of the policy.

2. What reduces it? Train on rolled-out sequences, so the model learns to be accurate from its own predicted states, not only from real ones. With clean data this changes nothing. With noisy observations (the realistic case: a camera or lidar pipeline never reports the true state), one-step training learns from corrupted inputs and drifts; training on 10-step rollouts cuts the 40-step error from 0.31 to 0.13:

3. Is it good enough to plan with? Random-shooting model-predictive control samples 256 action sequences, imagines each 12 steps ahead in the model, scores them against a target angle and executes the first action of the best one, re-planning at every step. This is what "plan by imagining several futures" means in code:

@torch.no_grad()
def plan_action(
    step_fn,
    state: torch.Tensor,
    horizon: int = 12,
    candidates: int = 256,
    seed: int = 0,
) -> torch.Tensor:
    """Random-shooting model-predictive control for one state.

    Samples ``candidates`` action sequences, imagines each ``horizon`` steps with
    ``step_fn`` (the true system or a learned model), scores the imagined trajectory by
    its distance from the target angle plus a small action cost, and returns the first
    action of the best sequence. The caller re-plans at every step.
    """
    generator = torch.Generator().manual_seed(seed)
    actions = torch.rand(candidates, horizon, generator=generator) * 2 - 1
    s = state.expand(candidates, 2)
    cost = torch.zeros(candidates)
    for t in range(horizon):
        s = step_fn(s, actions[:, t])
        cost = (
            cost
            + (s[:, 0] - TARGET) ** 2
            + 0.1 * s[:, 1] ** 2
            + 0.01 * actions[:, t] ** 2
        )
    return actions[cost.argmin(), 0]


@torch.no_grad()
def control_error(step_fn, episodes: int = 20, steps: int = 60, seed: int = 0) -> float:
    """Mean |theta - target| over the last 20 steps, planning with ``step_fn`` but
    acting in the real pendulum."""
    generator = torch.Generator().manual_seed(seed)
    errors = []
    for _ in range(episodes):
        state = (torch.rand(2, generator=generator) - 0.5) * 2
        trace = []
        for t in range(steps):
            action = plan_action(step_fn, state, seed=t)
            state = true_step(state, action)
            trace.append((state[0] - TARGET).abs())
        errors.append(torch.stack(trace[-20:]).mean())
    return torch.stack(errors).mean().item()

Controlling the real pendulum, the final angle error is 0.035 when planning in the true dynamics, 0.037 with the clean learned model and 0.044 to 0.053 with noisy-data models: nearly the same, and far better than doing nothing (0.98). The surprise is that the control error barely changes between the models whose rollout errors differ threefold. Re-planning at every step hides model error: the plan is only trusted for its first action, and the next observation corrects the drift. A model that is bad at 40 steps can still be good at 1. The error that matters is the one at the horizon you actually trust, which is also why driving systems re-plan many times per second.

The limits appear when the model must be trusted for longer: imagining many steps to train a policy on dreamed experience (Dreamer, below), or generating minutes of video to evaluate a driving policy. There the drift measured in question 1 is the whole problem.

Two families

1. Compact latent models for control. Learn a small latent state from images, a dynamics model in that latent space, and train behaviour inside it.

• World Models (1803.10122, Ha and Schmidhuber, 2018) learns a compressed spatial and temporal representation (with a VAE, the VAE module) and shows an agent can be trained entirely inside its own generated "dream" and transferred back.
• Dreamer (1912.01603, 2019) learns behaviours purely by latent imagination, back-propagating value gradients through imagined trajectories, and reports better data efficiency on 20 visual control tasks.
• DreamerV3 (2301.04104, 2023) extends the approach to diverse domains with one configuration.

2. Large generative video models. Predict future video (or video tokens), conditioned on actions or text.

• GAIA-1 (2309.17080, 2023) is a generative world model for driving. It maps video, text and action inputs to discrete tokens and predicts the next token, giving control over ego behaviour and scene features.
• Genie (2402.15391, 2024) is an 11B model trained from unlabelled internet video that makes action-controllable worlds, using a video tokenizer, an autoregressive dynamics model and a latent action model.
• Cosmos World Foundation Model Platform (2501.03575, 2025) is NVIDIA's open platform for Physical AI world models, covering video curation, pretrained world foundation models, post-training examples and tokenizers.

Notice that the building blocks are the ones from earlier modules: tokenisers (the Tokens, ViT and VLMs and VAE modules), autoregressive transformers (the Transformers module) or diffusion and flow models (the Diffusion and Flow matching modules).

World models and VLAs

A VLA (the VLA models module) maps observation and instruction straight to actions. It does not, by itself, imagine consequences. A world action model adds a prediction of the future (images, latent state or both) next to the action. The aim is better planning, richer training signal from unlabelled video, and a way to evaluate policies without a robot.

Alpamayo-R1 is not a world model. Its paper lists world-model integration among future directions and says its current approach "predicts actions from observed states" (in its future-work section, which also mentions self-supervised auxiliary tasks and hierarchical policies).

World action models

A world action model (WAM) joins the two halves: one model that forecasts the future and outputs actions. A 2026 survey (2606.20781) defines WAMs as "embodied predictive-action models that make a forecast of the future available to action", separates them from plain video generators and from VLA policies, and groups them by what they must generate: rendered future frames, latent futures, or no video at all.

The open question has been whether the model needs to imagine the future at inference, which is slow (iterative video denoising before every action), or only during training:

• Fast-WAM (2603.16666) keeps video prediction as a co-training objective but skips it at test time. In controlled comparisons it stays competitive with imagine-then-execute variants, while removing the video co-training hurts much more; it runs at 190 ms, which the authors report as over 4 times faster. Their conclusion: video prediction's main value may be the better representation it trains, not the frames it generates.
• SimWAM (2608.07468) applies the same idea to driving: future-video prediction only during training, an attention mask that keeps action prediction independent of future frames so planning runs without video generation, and RL on a compositional driving reward. It reports 91.9 PDMS on NAVSIM and zero-shot transfer to nuScenes.

This is the same trade as keeping perception and prediction outputs as auxiliary tasks (the From modular stacks to end-to-end module): extra prediction targets shape the representation during training and cost nothing on the car.

Reference implementation (PyTorch)

libs/tutorial/world_model.py; tests are in libs/tutorial/tests/test_world_model.py: collected trajectories follow the true dynamics, rollout error grows with the horizon, rollout training beats one-step training on noisy data by a factor of about two, planning in the true model reaches the target while doing nothing does not, and a learned model is good enough to plan with. Run uv run --group tutorial python -m libs.tutorial.world_model.

Limits to keep in mind

• Compounding error. A model rolled forward on its own predictions drifts, the same issue as in the Imitation learning module.
• Plausible is not physical. A generated video can look right and violate dynamics or contact.
• Evaluation. A policy evaluated in a learned simulator inherits the simulator's mistakes; check against real or physics-based closed-loop tests.
• Cost. Large video models are expensive to train and run.

Check yourself

1. What can you do with a world model that you cannot do with a policy alone?
2. How does Dreamer's imagination differ from GAIA-1's generation?
3. Why might a policy trained inside a world model fail on the real robot?
4. Which earlier module's technique does a video tokenizer rely on?
5. Why did the pendulum's control error barely change when the rollout error differed threefold?
6. Why does training on rolled-out sequences help with noisy observations but not with clean ones?

Quiz

Answer each question; feedback appears as soon as you choose.

Sources

Every paper above is linked inline; titles and identifiers were checked on arXiv.