Autoregressive Token World Models › Worked Example: A Tiny Multimodal Token World Model
World Models HandbookAll chapters Autoregressive Token World Models › Worked Example: A Tiny Multimodal Token World ModelOne section · Sections in this chapter7 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
We now build a deterministic one-dimensional world and compare prediction under recorded and generated contexts. It has three entities:
- An agent at a fine position that moves left, right, or holds, and is clipped at the walls.
- Ball 1, which moves one cell to the right each step, wrapping around.
- Ball 2 moves three cells per step with wraparound. Both balls' discrete position sequences repeat after 16 steps because 1 and 3 are coprime to 16; ball 2 has three times the displacement per step, not a shorter sampled repetition period.
The miniature uses the interleaved-token pattern of some categorical world models, without a learned visual tokenizer. An initial observation has three tokens; each transition adds an action followed by agent and ball positions. The disjoint ranges serve as exact encodings of these variables.
Action-before-observation order makes the control available at every next-frame prediction, and action targets will be masked from the loss. The two independent balls demonstrate multiple temporal processes sharing a stream; they do not simulate occlusion, visual rendering, or interacting-object physics.
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
torch.manual_seed(0)
torch.set_num_threads(2)
torch.use_deterministic_algorithms(True)
L = 16 # 1-D world size
N_ACTIONS = 3 # 0 = left, 1 = right, 2 = stay
AGENT_BASE, BALL1_BASE, BALL2_BASE = 3, 3 + L, 3 + 2 * L
VOCAB = 3 * L + 3 # 3 action tokens + 48 observation tokens
def agent_token(ap):
return AGENT_BASE + ap
def ball1_token(b):
return BALL1_BASE + b
def ball2_token(b):
return BALL2_BASE + b
def obs_tokens(ap, b1, b2):
return agent_token(ap), ball1_token(b1), ball2_token(b2)
def env_step(ap, b1, b2, a):
"""One environment transition. Works on scalars or numpy arrays."""
delta = np.where(a == 1, 1, np.where(a == 0, -1, 0))
ap2 = np.clip(ap + delta, 0, L - 1)
return ap2, (b1 + 1) % L, (b2 + 3) % L
def make_episode(rng, n_steps):
"""Serialize one episode into an interleaved action/observation stream."""
ap, b1, b2 = rng.integers(0, L, size=3)
toks = [int(t) for t in obs_tokens(ap, b1, b2)]
for _ in range(n_steps):
a = int(rng.integers(0, N_ACTIONS))
toks.append(a)
ap, b1, b2 = env_step(ap, b1, b2, a)
toks.extend(int(t) for t in obs_tokens(ap, b1, b2))
return toksThe vocabulary is 51 symbols: tokens 0–2 are actions, 3–18 are agent positions, 19–34 are ball-1 positions, and 35–50 are ball-2 positions. Next we generate a dataset of episodes and split it into training and validation sets. We also keep one short demo episode whose token stream we will visualize.
rng = np.random.default_rng(7)
## Episodes of 31 steps (127 tokens) cover every sequence position the
## open-loop rollouts below will visit, so rollout degradation cannot be
## blamed on position embeddings the model never trained on.
N_TRAIN, N_VAL, EP_STEPS = 4000, 500, 31
train_eps = [make_episode(rng, EP_STEPS) for _ in range(N_TRAIN)]
val_eps = [make_episode(rng, EP_STEPS) for _ in range(N_VAL)]
demo_tokens = make_episode(np.random.default_rng(3), 4)
def to_xy(episodes):
"""Inputs are all tokens but the last; targets are all but the first."""
x = torch.tensor([e[:-1] for e in episodes], dtype=torch.long)
y = torch.tensor([e[1:] for e in episodes], dtype=torch.long)
return x, y
X_tr, Y_tr = to_xy(train_eps)
X_val, Y_val = to_xy(val_eps)Here is what one episode looks like as a token stream, colored by modality. The stream opens with one initial observation frame (agent position, then the two ball positions), and after that the alternation is strict: action, then the three observation tokens it caused. This layout means each observation token is conditioned on the most recent action, and each observation token within a frame is conditioned on its earlier siblings, exactly the factorization from the dynamics section above.
The small GPT uses learned token and position embeddings, two pre-norm blocks, and a full-vocabulary output head. Upper-triangular masking excludes later inputs; shifting targets ensures the current input does not already reveal the target. Both are needed for valid next-token supervision.
Token ranges identify modalities directly. Position embeddings and the fixed predecessor-type pattern provide additional slot and timing cues. The shared embedding and categorical head illustrate production components, while the toy omits visual tokenization, modality constraints, caching, and large-data concerns.
class Block(nn.Module):
def __init__(self, d, n_heads):
super().__init__()
self.ln1 = nn.LayerNorm(d)
self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True)
self.ln2 = nn.LayerNorm(d)
self.mlp = nn.Sequential(
nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d)
)
def forward(self, x, mask):
h = self.ln1(x)
a, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
x = x + a
x = x + self.mlp(self.ln2(x))
return x
class TinyGPT(nn.Module):
def __init__(self, vocab, d=64, n_layers=2, n_heads=2, max_len=128):
super().__init__()
self.tok = nn.Embedding(vocab, d)
self.pos = nn.Embedding(max_len, d)
self.blocks = nn.ModuleList(
[Block(d, n_heads) for _ in range(n_layers)]
)
self.ln_f = nn.LayerNorm(d)
self.head = nn.Linear(d, vocab)
def forward(self, idx):
if (
idx.ndim != 2
or idx.shape[0] < 1
or not 1 <= idx.shape[1] <= self.pos.num_embeddings
):
raise ValueError(
"Expected a nonempty batch within the position table"
)
T = idx.shape[1]
pos = torch.arange(T, device=idx.device)
x = self.tok(idx) + self.pos(pos)[None]
mask = torch.triu(
torch.full((T, T), float("-inf"), device=idx.device), diagonal=1
)
for blk in self.blocks:
x = blk(x, mask)
return self.head(self.ln_f(x))We train only on observation-token targets. Action prediction is unnecessary for this dynamics-only objective, although an auxiliary policy objective can serve other designs. The two independent initial ball positions are unpredictable from their earlier prefix; the reporting metric excludes these positions to evaluate transitions.
The 40-step and 200-step models begin from identical parameters and use the same seeded minibatch stream for their shared first 40 updates. This compares two points along the same teacher-forced training trajectory rather than conflating budget with initialization. It remains one seed and one task, not an architecture or scaling study.
Residual prediction error and feedback can interact, but two training budgets alone do not identify exposure bias. We therefore compare both models on paired initial states and actions, report the same per-frame exact-match event under teacher forcing and open-loop prediction, and distinguish per-step match from flawless survival.
def train_teacher_forced(model, X, Y, steps, batch=128, lr=1e-3, seed=101):
model.train()
batch_rng = torch.Generator().manual_seed(seed)
opt = torch.optim.Adam(model.parameters(), lr=lr)
n = X.shape[0]
losses = []
for _ in range(steps):
idx = torch.randint(0, n, (batch,), generator=batch_rng)
logits = model(X[idx])
loss_all = F.cross_entropy(
logits.reshape(-1, VOCAB), Y[idx].reshape(-1), reduction="none"
)
world = (Y[idx] >= 3).reshape(-1).float() # mask out action targets
loss = (loss_all * world).sum() / world.sum()
opt.zero_grad()
loss.backward()
opt.step()
losses.append(loss.item())
return losses
import copy
model_tf = TinyGPT(VOCAB)
model_short = copy.deepcopy(model_tf) # identical initialization
losses_tf = train_teacher_forced(model_tf, X_tr, Y_tr, steps=200)
losses_short = train_teacher_forced(model_short, X_tr, Y_tr, steps=40)teacher-forced loss (200 steps): 4.1260 -> 0.5646 teacher-forced loss (40 steps): 4.1260 -> 2.7346
Training losses describe fit to the supervised token objective, not calibrated uncertainty or decision usefulness. The printed first and final losses make any decrease a measured result. Held-out accuracy is broken down by modality so that different prediction difficulties remain visible.
@torch.no_grad()
def teacher_forced_accuracy(model, X, Y):
model.eval()
pred = model(X).argmax(-1)
ranges = {
"agent": (AGENT_BASE, BALL1_BASE),
"ball 1": (BALL1_BASE, BALL2_BASE),
"ball 2": (BALL2_BASE, VOCAB),
}
# The first two observation targets of an episode are the initial ball
# positions, drawn uniformly and unpredictable from any context; exclude
# them so this report measures transitions, not unpredictable initial states.
# This pooled token metric differs from per-frame exact match below.
predictable = torch.arange(Y.shape[1]) >= 2
acc = {}
for name, (lo, hi) in ranges.items():
m = predictable & (Y >= lo) & (Y < hi)
acc[name] = (((pred == Y) & m).sum() / m.sum()).item()
m = predictable & (Y >= 3)
acc["all obs"] = (((pred == Y) & m).sum() / m.sum()).item()
return acc
tf_acc = teacher_forced_accuracy(model_tf, X_val, Y_val)
short_acc = teacher_forced_accuracy(model_short, X_val, Y_val)Teacher-forced next-token accuracy on held-out episodes: modality 200 steps 40 steps agent 97.07% 23.92% ball 1 87.26% 6.75% ball 2 82.70% 6.77% all obs 89.01% 12.48%
For evaluation, each model receives the same sampled initial states and actions. The deterministic environment removes ambiguity from the reference trajectory: a mismatch is a prediction error. It does not identify whether that error arose from ordinary one-step error, altered context, finite capacity, or optimization. Matched teacher-forced predictions on these very trajectories provide a cleaner comparison than pooled token accuracy against frame-level rollout success.
Greedy decoding makes the evaluation deterministic conditional on parameters and inputs. It selects the highest model probability, which maximizes one-step accuracy only when the conditional ranking is correct. Calibration alone is insufficient, and trajectory-level behavior need not improve over sampling. We report greedy results without interpreting them as a bound on sampling damage.
def rollout_reference(rng, n_episodes, horizon, offset=0):
"""Paired true/shifted reference paths under the same exogenous controls."""
if n_episodes < 1 or horizon < 1 or 3 + 4 * horizon > 128:
raise ValueError(
"Positive batch/horizon required within the position table"
)
initial = rng.integers(0, L, (3, n_episodes))
actions = rng.integers(0, N_ACTIONS, (n_episodes, horizon))
shown = initial.copy()
shown[0] = (shown[0] + offset) % L
true_states, shown_states = initial.copy(), shown.copy()
truths = np.empty((horizon, 3, n_episodes), dtype=np.int64)
alternatives = np.empty_like(truths)
bases = np.array([AGENT_BASE, BALL1_BASE, BALL2_BASE])[:, None]
for t in range(horizon):
true_states = np.asarray(env_step(*true_states, actions[:, t]))
shown_states = np.asarray(env_step(*shown_states, actions[:, t]))
truths[t] = true_states + bases
alternatives[t] = shown_states + bases
return shown + bases, actions, truths, alternatives
@torch.no_grad()
def batched_rollout(
model,
rng,
n_episodes,
horizon,
corrupt_agent_offset=None,
mode="model",
score_alternative=False,
):
"""Return Boolean (horizon, modality, episode) correctness on paired cases.
mode='teacher' feeds reference tokens after each prediction.
mode='oracle' follows exact dynamics from the supplied initial state.
"""
if mode not in {"model", "teacher", "oracle"}:
raise ValueError("Unknown rollout mode")
if model is not None:
model.eval()
initial, actions, truths, alternatives = rollout_reference(
rng, n_episodes, horizon, offset=corrupt_agent_offset or 0
)
targets = alternatives if score_alternative else truths
seq = torch.zeros(n_episodes, 3 + 4 * horizon, dtype=torch.long)
seq[:, :3] = torch.from_numpy(initial.T.copy())
ok = np.zeros((horizon, 3, n_episodes), dtype=bool)
cur = 3
for t in range(horizon):
seq[:, cur] = torch.from_numpy(actions[:, t].copy())
cur += 1
for k in range(3):
if mode == "oracle":
pred = torch.from_numpy(alternatives[t, k].copy())
else:
pred = model(seq[:, :cur])[:, -1].argmax(-1)
ok[t, k] = pred.numpy() == targets[t, k]
seq[:, cur] = (
torch.from_numpy(truths[t, k].copy())
if mode == "teacher"
else pred
)
cur += 1
return ok
def summarize(ok):
if ok.ndim != 3 or ok.shape[0] == 0 or ok.shape[1] != 3 or ok.shape[2] == 0:
raise ValueError("Expected nonempty (horizon, 3, episodes) correctness")
exact = ok.all(axis=1)
step_acc = exact.mean(axis=1)
survival = np.logical_and.accumulate(exact, axis=0).mean(axis=1)
return step_acc, survival
HORIZON = 24
N_EVAL = 256
ok_tf = batched_rollout(model_tf, np.random.default_rng(21), N_EVAL, HORIZON)
ok_short = batched_rollout(
model_short, np.random.default_rng(21), N_EVAL, HORIZON
)
step_tf, surv_tf = summarize(ok_tf)
step_short, surv_short = summarize(ok_short)
teacher_tf = batched_rollout(
model_tf, np.random.default_rng(21), N_EVAL, HORIZON, mode="teacher"
)
teacher_short = batched_rollout(
model_short, np.random.default_rng(21), N_EVAL, HORIZON, mode="teacher"
)
matched_tf, _ = summarize(teacher_tf)
matched_short, _ = summarize(teacher_short)
oracle_clean = batched_rollout(
None, np.random.default_rng(21), N_EVAL, HORIZON, mode="oracle"
)
assert oracle_clean.all()
assert np.allclose(losses_tf[:40], losses_short, rtol=0, atol=0)Open-loop rollouts (256 episodes, horizon 24): quantity 200 steps 40 steps per-step exact match, step 1 85.9% 0.0% per-step exact match, final step 5.1% 0.0% flawless rollouts through horizon 0.0% 0.0% matched teacher-forced frame accuracy, mean 75.7% 0.1%
The matched teacher-forced curve requires all three transition tokens to be correct at a given horizon, with each target evaluated from its true preceding prefix. Open-loop exact match requires the same three targets while feeding predictions back. Both use the same initial states, actions, and positions. The gap therefore avoids the old token-versus-frame aggregation mismatch, although it remains a comparison of conditioning regimes rather than a universal measure of exposure bias.
Per-step exact match can recover after earlier failures. Flawless survival cannot: it is the fraction with every step correct so far and is nonincreasing by definition. Survival is at most current-step match at each horizon, but need not have a uniformly steeper slope. We estimate survival by intersecting events within each rollout, not by multiplying marginal accuracies. The printed values and curves determine this run's behavior without imposing a decay shape.
Fine-Tuning on the Model's Own Rollouts
We next compare self-generated-context corrective fine-tuning with continued teacher forcing. Both start from the same 200-step weights and use the same newly sampled states and action sequences, 60 optimizer updates, batch size, learning rate, and 16-step supervision horizon. Both evaluate every prefix separately, so their forward/backward call counts and token lengths are matched. Wall-clock timing can still differ. Evaluation extends to 24 steps, beyond the fine-tuning horizon.
The generated-context branch predicts along its own token history but receives targets from the original reference trajectory. These labels ask it to return to that trajectory, not to obey the environment transition starting from each hallucinated state. The original initial state and known actions make reference labels available; fixed recordings could supply the same labels without an interactive oracle. Some generated contexts can be ambiguous or inconsistent, so useful correction is not guaranteed. No dataset is aggregated, and no DAgger theorem is claimed.
def corrective_finetune(
model, rng, rounds=60, batch=64, horizon=16, lr=5e-4, self_context=True
):
"""Matched update budget; choose generated or true preceding context."""
if rounds < 1 or batch < 1 or horizon < 1 or 3 + 4 * horizon > 128:
raise ValueError("Invalid training budget or context length")
opt = torch.optim.Adam(model.parameters(), lr=lr)
model.train()
for _ in range(rounds):
opt.zero_grad()
ap = rng.integers(0, L, batch)
b1 = rng.integers(0, L, batch)
b2 = rng.integers(0, L, batch)
actions = rng.integers(0, N_ACTIONS, (batch, horizon))
seq = torch.zeros(batch, 3 + 4 * horizon, dtype=torch.long)
seq[:, 0] = torch.from_numpy(np.ascontiguousarray(AGENT_BASE + ap))
seq[:, 1] = torch.from_numpy(np.ascontiguousarray(BALL1_BASE + b1))
seq[:, 2] = torch.from_numpy(np.ascontiguousarray(BALL2_BASE + b2))
cur = 3
for t in range(horizon):
seq[:, cur] = torch.from_numpy(np.ascontiguousarray(actions[:, t]))
cur += 1
ap, b1, b2 = env_step(ap, b1, b2, actions[:, t])
for truth in (AGENT_BASE + ap, BALL1_BASE + b1, BALL2_BASE + b2):
logits = model(seq[:, :cur])
loss = F.cross_entropy(
logits[:, -1], torch.from_numpy(np.ascontiguousarray(truth))
)
(loss / (3 * horizon)).backward()
with torch.no_grad():
seq[:, cur] = (
logits[:, -1].argmax(-1)
if self_context
else torch.from_numpy(np.ascontiguousarray(truth))
)
cur += 1
opt.step()
return model
model_corrective = copy.deepcopy(model_tf)
model_extra = copy.deepcopy(model_tf)
corrective_finetune(model_corrective, np.random.default_rng(13))
corrective_finetune(model_extra, np.random.default_rng(13), self_context=False)
ok_corrective = batched_rollout(
model_corrective, np.random.default_rng(21), N_EVAL, HORIZON
)
ok_extra = batched_rollout(
model_extra, np.random.default_rng(21), N_EVAL, HORIZON
)
step_corrective, surv_corrective = summarize(ok_corrective)
step_extra, surv_extra = summarize(ok_extra)
corrective_acc = teacher_forced_accuracy(model_corrective, X_val, Y_val)
extra_acc = teacher_forced_accuracy(model_extra, X_val, Y_val)Matched 60-update continuation from the same 200-step weights: base: pooled teacher-forced token accuracy=0.8901; frame match mean steps 1-16=0.3032, 17-24=0.0698, final=0.0508; flawless through 24=0.0000 generated context: pooled teacher-forced token accuracy=0.8519; frame match mean steps 1-16=0.7190, 17-24=0.0137, final=0.0039; flawless through 24=0.0000 true context: pooled teacher-forced token accuracy=0.7940; frame match mean steps 1-16=0.9761, 17-24=0.0005, final=0.0000; flawless through 24=0.0000
In the seeded CPU run (PyTorch 2.13.0, NumPy 2.4.2), mean frame match over steps 1–16 was 30.3% for the base model, 71.9% after generated-context correction, and 97.6% after continued teacher forcing. Over steps 17–24, the same figures were 7.0%, 1.4%, and 0.05%. All three had zero flawless survival through step 24. The extra teacher-forced updates explain a stronger in-horizon gain than corrective contexts in this run; neither continuation improves final-step accuracy or 24-step flawless survival.
The comparison is intentionally allowed to fail. Additional teacher forcing can improve ordinary fitting, while corrective training may help within its supervision horizon and harm later predictions. Use the printed results for each branch rather than attributing every gain to generated-context training. A rule-based predictor can exactly simulate this finite deterministic system, but that does not prove this neural architecture will learn it perfectly. In stochastic worlds, imperfect exact-path matching can persist even for the true conditional distribution and is not itself exposure bias.
Corrupted Context: Testing Sensitivity to Initial Evidence
Finally, shift the supplied initial agent position by seven modulo sixteen while keeping the reference trajectory and actions unchanged. Because a modulo shift permutes a uniform initial distribution, the shown starting states remain in distribution. They are wrong relative to the paired reference, not intrinsically invalid inputs.
The oracle control follows the exact dynamics from the shifted initial state. If even this exact predictor disagrees with the original reference, that discrepancy is a consequence of changed evidence rather than exposure bias. A second learned-model score against the shifted oracle trajectory distinguishes fidelity to the supplied alternative state from agreement with the original trajectory. The intervention does not remove ordinary neural prediction errors.
Clean and shifted runs use paired actions and initial states. We measure agent and ball accuracy and count re-synchronization explicitly. Generated and observed tokens share the input interface without provenance labels, but this does not force equal reliance on every token. The test concerns sensitivity to initial evidence, a phenomenon that also occurs in recurrent predictors.
ok_corrupt = batched_rollout(
model_tf, np.random.default_rng(21), 256, HORIZON, corrupt_agent_offset=7
)
agent_clean = ok_tf[:, 0, :].mean(axis=1)
agent_corrupt = ok_corrupt[:, 0, :].mean(axis=1)
balls_clean = ok_tf[:, 1:, :].mean(axis=(1, 2))
balls_corrupt = ok_corrupt[:, 1:, :].mean(axis=(1, 2))
oracle_shifted = batched_rollout(
None,
np.random.default_rng(21),
N_EVAL,
HORIZON,
corrupt_agent_offset=7,
mode="oracle",
)
ok_alternative = batched_rollout(
model_tf,
np.random.default_rng(21),
N_EVAL,
HORIZON,
corrupt_agent_offset=7,
score_alternative=True,
)
agent_oracle = oracle_shifted[:, 0, :].mean(axis=1)
agent_alternative = ok_alternative[:, 0, :].mean(axis=1)
oracle_resynced = oracle_shifted[:, 0, :].any(axis=0).sum()Exact shifted oracle re-synchronized at least once: 5/256 Shifted neural agent accuracy against shifted oracle: 0.9150 Effect of one shifted agent token in the initial context: agent-token accuracy, clean start : 90.4% agent-token accuracy, corrupted start: 0.6% ball-token accuracy, clean start : 45.8% ball-token accuracy, corrupted start: 45.9%
In the same 256 paired episodes, the exact shifted oracle re-synchronized with the original agent trajectory at least once in 5 episodes. The learned agent averaged 0.6% accuracy against the original trajectory after the shift, but 91.5% against the alternative trajectory implied by its supplied initial state. Clean and shifted ball-token means were 45.8% and 45.9%. Those near-equal ball means do not establish attention independence.
Agent clipping under common actions can merge distinct positions; wraparound motion preserves circular displacement instead. The exact shifted-state oracle makes this distinction visible without attributing recovery to learned attention. Clipping itself is nonexpansive in Euclidean position and strictly contracts only when it shortens a difference; a larger surrounding transition can still expand differences. Learned-model accuracy changes do not establish whether attention routes information between independent variables. In video, spatial and temporal paths may spread errors, but raster order does not imply neighboring patches are affected first. These are hypotheses for intervention tests, not diagnoses from appearance alone.