11.4 — World Models and Predictive Action Planning

Learning objectives

After this lesson, you should be able to:

• Explain what a world model predicts.
• Fit a dynamics model and see how its error grows with horizon.
• Plan with sampled imagined rollouts.
• Explain how planners exploit model errors and how to limit it.

The one-sentence idea

Try it in imagination first.

Imagination is only as good as the model, and planners find its flaws.

Predict first, then act

A policy maps what it sees to what to do. A world model maps what it sees and what it plans to do to what will happen next. With one, a robot can imagine the outcomes of several candidate actions and choose the best before moving anything. Ha and Schmidhuber's “World Models” builds generative neural-network models of RL environments, trained quickly without labels into a compressed representation of space and time; features from the model feed a small controller, and an agent can even be trained entirely inside the model's generated environment, with the policy then transferred to the real one.^[1] DreamerV3 learns a model of the environment and improves its behaviour by imagining future scenarios, and a single configuration is reported to outperform specialised methods across over 150 diverse tasks.^[2]

Why robots want this: real trials are slow and risky (lesson 11.1), imagined trials are fast and safe, if the model is right. The catch is in that last clause, and the experiments below show it.

Learning a model from experience

The simplest world model is a regression: given the current state and an action, predict the next state. A hand moves in the plane with some inertia; we collect random experience and fit a linear model, then check how well it predicts one step and many steps ahead:

import numpy as np
rng = np.random.default_rng(0)

def true_step(s, a):                                    # state: [x, y, vx, vy]; action: acceleration [ax, ay]
    pos, vel = s[:2], s[2:]
    vel = 0.9 * vel + 0.1 * a + rng.normal(0, 0.002, 2)
    pos = pos + 0.1 * vel
    return np.concatenate([pos, vel])

# collect random experience
S, A, S2 = [], [], []
s = np.zeros(4)
for _ in range(2000):
    a = rng.normal(0, 1.0, 2)
    s2 = true_step(s, a)
    S.append(s); A.append(a); S2.append(s2); s = s2
X = np.column_stack([np.array(S), np.array(A)]); Y = np.array(S2)
M, *_ = np.linalg.lstsq(X, Y, rcond=None)                # the world model: next state = [state, action] @ M

def model_step(s, a): return np.concatenate([s, a]) @ M

# compare imagined and real trajectories for the same action sequence
actions = rng.normal(0, 1.0, (10, 2))
real, imagined = np.zeros(4), np.zeros(4)
errors = []
for a in actions:
    real = true_step(real, a); imagined = model_step(imagined, a)
    errors.append(np.linalg.norm(real[:2] - imagined[:2]))
print("position error after 1 step:  ", round(float(errors[0]), 5))
print("position error after 10 steps:", round(float(errors[-1]), 5))
print("error grew by about", round(float(errors[-1] / max(errors[0], 1e-12)), 1), "times")

position error after 1 step:   0.00013
position error after 10 steps: 0.00129
error grew by about 9.8 times

Toy model. A linear system that a linear model can capture almost exactly. Real worlds have contact and clutter, and the errors grow much faster.

Even a good model's errors accumulate with the horizon, exactly like the compounding errors of lesson 1.1. Planning far ahead with an imperfect model means trusting its worst predictions.

Planning with the model

With a model, a robot can try many action sequences in imagination and pick the best (“random shooting” model-predictive control): sample candidate sequences, roll each out in the model, score the imagined end positions and execute only the first action of the best sequence, then re-plan. Here the planner steers the hand to a goal using the model of the previous example (re-created in the cell):

import numpy as np
rng = np.random.default_rng(1)

def true_step(s, a):
    pos, vel = s[:2], s[2:]
    vel = 0.9 * vel + 0.1 * a + rng.normal(0, 0.002, 2)
    return np.concatenate([pos + 0.1 * vel, vel])

# the model: fitted from random experience (as before)
S, A, S2, s = [], [], [], np.zeros(4)
for _ in range(2000):
    a = rng.normal(0, 1.0, 2); s2 = true_step(s, a); S.append(s); A.append(a); S2.append(s2); s = s2
M = np.linalg.lstsq(np.column_stack([np.array(S), np.array(A)]), np.array(S2), rcond=None)[0]
model_step = lambda s, a: np.concatenate([s, a]) @ M

def plan(state, goal, horizon=15, candidates=300):
    best, best_cost = None, np.inf
    for _ in range(candidates):
        seq = rng.normal(0, 1.5, (horizon, 2))
        s = state.copy()
        for a in seq: s = model_step(s, a)
        cost = np.linalg.norm(s[:2] - goal) + 0.3 * np.linalg.norm(s[2:])          # end near the goal, slowly
        if cost < best_cost: best, best_cost = seq[0], cost
    return best

state, goal = np.zeros(4), np.array([1.0, -0.5])
for step in range(60):
    state = true_step(state, plan(state, goal))                                     # act in the real world with the first action
print("distance from the goal after 60 steps:", round(float(np.linalg.norm(state[:2] - goal)), 3))

distance from the goal after 60 steps: 0.205

Starting 1.1 away from the goal, the hand ends 0.2 away after 60 steps, using a model learned only from random motions: no demonstrations and no reward learning. That is far from exact. Random-shooting planners are crude, and more candidates, a longer run or a better cost would get closer; the point is that planning in imagination works at all from such little information.

The danger: a planner exploits model errors

A planner searches for actions that look best according to the model. If the model is wrong somewhere, the planner will find that place and love it. Suppose the model was only trained on small velocities. Outside that range it extrapolates, here wrongly, and an optimistic error gives the planner a fantasy: a huge “free” move.

import numpy as np

def true_gain(a):  return 0.1 * np.clip(a, -1, 1)                       # reality: actuators saturate at an acceleration of 1
def model_gain(a): return 0.1 * a                                       # the model, trained only on |a| < 1, extrapolates linearly

for a in (0.5, 1.0, 3.0, 10.0):
    print(f"acceleration {a:>4}: model predicts a move of {model_gain(a):5.2f}, reality gives {true_gain(a):5.2f}")
best = max((0.5, 1.0, 3.0, 10.0), key=model_gain)
print("the planner picks acceleration", best, "because the model promises", model_gain(best), "but reality delivers", true_gain(best))

acceleration  0.5: model predicts a move of  0.05, reality gives  0.05
acceleration  1.0: model predicts a move of  0.10, reality gives  0.10
acceleration  3.0: model predicts a move of  0.30, reality gives  0.10
acceleration 10.0: model predicts a move of  1.00, reality gives  0.10
the planner picks acceleration 10.0 because the model promises 1.0 but reality delivers 0.1

The remedies in practice: keep planning horizons short and re-plan often, stay near the data the model was trained on, use uncertainty estimates and penalise actions the model is unsure about, and keep checking against reality (lessons 3.5 and 7.7).

World models and VLAs

Current VLAs (phases 5 and 6) mostly predict actions directly; world models predict consequences. The two ideas combine in many ways in current research: a world model can generate extra training experience, evaluate candidate actions, or supply a learned notion of how scenes change. Which combinations work well on real robots is an open question (lesson 11.7). What is not in doubt is the trade-off shown above: imagination is cheap and fast, and only as good as the model.

Common misconceptions

Engineering takeaways

1. Model: (state, action) → next state.
2. Errors accumulate.
3. Plan by imagined rollouts, re-plan often.
4. Stay near the data.

Imagine cheaply, verify in the world.

Knowledge check

Try each question first, then reveal the explanation.

What’s next

Next: splitting long tasks into steps a policy can execute.

Back to the VLA roadmap