What Is a World Model?

A world model is a function you can call instead of the world.

Everything else in the field is about which function, learned from what, and how far you can trust it. This post builds the smallest honest version: a 5x5 gridworld, a log of random transitions, two models fitted to that log, a planner that runs entirely inside each model, and a measurement of how fast the imagined trajectory parts company with the real one. Every number below comes from running the code shown.

The smallest world#

The environment is a 5x5 grid with five interior walls. The agent starts at (0, 0), the goal is (4, 4), each step costs -0.04, and reaching the goal pays +1. Moves into a wall or off the edge leave the agent where it is.

import numpy as np
from collections import defaultdict, Counter
 
N = 5
WALLS = {(1, 1), (2, 1), (3, 1), (1, 3), (2, 3)}   # interior walls
GOAL = (4, 4)
MOVES = [(0, 1), (1, 0), (0, -1), (-1, 0)]          # right, down, left, up  (row, col)
 
def step(s, a):
    """True dynamics: deterministic moves, walls and edges block."""
    r, c = s
    dr, dc = MOVES[a]
    nr, nc = r + dr, c + dc
    if not (0 <= nr < N and 0 <= nc < N) or (nr, nc) in WALLS:
        nr, nc = r, c
    s2 = (nr, nc)
    reward = 1.0 if s2 == GOAL else -0.04
    return s2, reward
 
STATES = [(r, c) for r in range(N) for c in range(N) if (r, c) not in WALLS]

step is the territory. It has the signature

A world model is anything with the same signature that was not handed to you. You learn it from experience, and then you are allowed to call it as often as you like without paying for real interaction. That asymmetry is the whole point: real steps are expensive (robots break, patients are real, games are slow), model calls are cheap.

The data: a log of what happened#

The agent does not see step. It sees a log of transitions produced by some behaviour, here a uniformly random walk that resets on reaching the goal.

def log_random(n, rng):
    data, s = [], (0, 0)
    for _ in range(n):
        a = int(rng.integers(4))
        s2, rew = step(s, a)
        data.append((s, a, rew, s2))
        s = (0, 0) if s2 == GOAL else s2
    return data

Each tuple (s, a, r, s') is one labelled example of the function the model must learn. Learning a world model is supervised learning on this log. The difficulty is not the loss function. It is that the log covers only the states the behaviour policy happened to visit.

Step 1: a lookup table#

The simplest model memorises. For every (s, a) pair seen in the log, store the most common next state and the observed reward. For pairs never seen, guess that nothing happens.

class TableModel:
    def __init__(self, data):
        self.next = defaultdict(Counter); self.rew = {}
        for s, a, r, s2 in data:
            self.next[(s, a)][s2] += 1; self.rew[(s, a)] = r
    def __call__(self, s, a):
        if (s, a) not in self.next:
            return s, -0.04            # unseen: guess "nothing happens"
        s2 = self.next[(s, a)].most_common(1)[0][0]
        return s2, self.rew[(s, a)]
    def coverage(self):
        return len(self.next) / (len(STATES) * 4)
 
def one_step_acc(model):
    hits = [model(s, a)[0] == step(s, a)[0] for s in STATES for a in range(4)]
    return np.mean(hits)

The table is exact wherever it has data and blind everywhere else. Its accuracy is therefore a direct function of coverage:

Accuracy is higher than coverage because the "nothing happens" guess is right for many unseen moves into walls and edges. That is a lucky prior, not knowledge. The bottleneck is plain: a table cannot say anything about a pair it has not seen, and in any interesting environment most pairs are never seen. The tabular model in the companion post hits the same wall from the other side, with pixels.

Step 2: a model that generalises#

To answer for unseen pairs, the model must share structure across states. The smallest such model is linear: predict the next (row, col) from the current (row, col), a bias, and a one-hot action.

def feats(s, a):
    oh = np.zeros(4); oh[a] = 1
    return np.concatenate([[s[0], s[1], 1.0], oh])
 
class LinearModel:
    """s' = W [row, col, 1, onehot(a)]; reward from a separate linear head."""
    def __init__(self, data):
        X = np.array([feats(s, a) for s, a, _, _ in data])
        Y = np.array([[s2[0], s2[1]] for _, _, _, s2 in data], float)
        R = np.array([r for _, _, r, _ in data])
        self.W, *_ = np.linalg.lstsq(X, Y, rcond=None)
        self.v, *_ = np.linalg.lstsq(X, R, rcond=None)
    def __call__(self, s, a):
        y = feats(s, a) @ self.W
        s2 = tuple(int(v) for v in np.clip(np.rint(y), 0, N - 1))
        return s2, float(feats(s, a) @ self.v)

Fitted on 1000 logged steps, the action rows of W come out as roughly +0.49 column for "right", +0.71 row for "down", -0.38 column for "left", and -0.60 row for "up". The model has learned "actions move you", shrunk toward zero because blocked moves pull the average down. After rounding, its one-step accuracy is 0.71 at 1000 samples and 0.74 at 200. More data does not help, because the error is not variance. It is bias: walls are not a linear function of position.

That is the trade every world model makes. The table has no inductive bias and no generalisation. The linear model has a strong inductive bias and generalises confidently, including through walls. Replacing feats with a one-hot of (s, a) turns the linear model back into the table, so the two are endpoints of one dial, and neural networks sit in between.

Step 3: imagine, then plan#

A model earns its keep when something plans inside it. Value iteration needs exactly the function (s, a) -> (s', r), so it can run on either model without touching step:

def plan(model, gamma=0.95, iters=200):
    """Value iteration *inside the model*: never calls step()."""
    V = {s: 0.0 for s in STATES}
    for _ in range(iters):
        for s in STATES:
            if s == GOAL: continue
            V[s] = max(r + gamma * V.get(s2, 0.0) for s2, r in (model(s, a) for a in range(4)))
    def pi(s):
        return max(range(4), key=lambda a: model(s, a)[1] + gamma * V.get(model(s, a)[0], 0.0))
    return pi
 
def run(pi, start=(0, 0), limit=50):
    s, ret = start, 0.0
    for t in range(limit):
        s, r = step(s, pi(s)); ret += r
        if s == GOAL: return t + 1, round(ret, 2)
    return None, round(ret, 2)

The plan is computed in imagination and then executed in the real grid. The result, for models trained on the same logs:

Eight steps is the shortest path. With 200 logged random steps the table already contains enough of the world to plan optimally, even at 78% coverage, because the planner only needs the transitions on and near the good path.

The linear plan fails in an instructive way. From (0, 1) the planner chooses "down", because the model predicts (0, 1) -> (1, 1), and (1, 1) is a wall. The model does not know walls exist; the planner happily routes through one, the real environment refuses, and the agent presses against the same wall for 50 steps. The planner did its job perfectly. It optimised the model, and the model was wrong in exactly the place the optimiser went looking. This is the first law of model-based control: a planner is an adversary of its model's errors.

Step 4: measure where imagination diverges#

One-step accuracy hides how errors interact over a rollout. The direct measurement is to start the real environment and the model in the same state, feed both the same random action sequence, and record how often they still agree after k steps.

def divergence(model, H, trials, rng):
    """P(imagined state == real state) after k steps of the same random actions."""
    match = np.zeros(H)
    for _ in range(trials):
        s_real = s_img = STATES[rng.integers(len(STATES))]
        for k in range(H):
            a = int(rng.integers(4))
            s_real, _ = step(s_real, a)
            s_img, _ = model(s_img, a)
            match[k] += s_real == s_img
    return match / trials

Over 2000 rollouts per model:

This is compounding error, measured. Once the imagined state differs from the real one, the model is answering a question about the wrong state, and every later prediction inherits the mistake.

Why errors compound, and why not exactly geometrically#

If each step were wrong independently with probability $\epsilon$, agreement after $k$ steps would be

For the good table, $\epsilon \approx 0.02$ gives $0.98^{20} \approx 0.67$; the measured 0.72 is close. For the linear model, $0.71^{5} \approx 0.18$, but the measured value at $k = 5$ is 0.39. The independence assumption is too pessimistic here because a grid is small and full of absorbing structure: a wrong prediction against an edge can land both trajectories back in the same cell, and the trajectories re-merge. Continuous state spaces have no such mercy. A small position error changes the next velocity, which changes the next position, and drift grows rather than resets; the planning post measures that case.

The practical consequences follow directly:

• Short imagined rollouts are more trustworthy than long ones. Most successful agents keep imagination horizons in the tens of steps.
• Error matters most where the planner goes, not where the data is. Average one-step accuracy over a fixed test set is the wrong headline number.
• Replanning from real observations resets the divergence clock. A plan executed open-loop for twenty steps inherits all twenty steps of error; a plan recomputed every step inherits one.

What to try next#

1. Replace the random logging policy with the table planner's own behaviour, refit, and replan in a loop. Coverage concentrates on the useful path; this is the seed of Dyna-style learning.
2. Make step stochastic (a 10% slip to a random neighbour) and change TableModel to sample from the stored counts instead of taking the mode. Measure divergence as a distance between distributions, not an equality.
3. Add one-hot row and column features to LinearModel and watch one-step accuracy climb toward the table's. That is the inductive-bias dial in practice.
4. Train a small MLP on the same features and report divergence at k = 1, 5, 20 alongside the table. One-step accuracy and 20-step agreement will not rank models the same way; that gap is the subject of evaluating a world model.
5. Swap the grid for pixels and the table stops existing. That is where latent world models come in, starting with Dreamer and the latent space.