Tokenized Game and Environment Models › A worked example: the key, the door, and the token
World Models HandbookAll chapters Tokenized Game and Environment Models › A worked example: the key, the door, and the tokenOne 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
Now that the concepts are on the table, let us build a minimal example that exercises them. We will use a tile world with an agent, a key, a locked goal, a wall, and a hazard. It is small enough to reason about completely and rich enough to exhibit the failures we care about.
The story we will demonstrate:
- A discrete tokenizer that preserves the key sprite yields a model that behaves correctly on a plan that goes to the key first.
- A lossy tokenizer that aliases the key sprite into "empty" yields a model with similar one-step accuracy but a broken multi-step plan: the model cannot tell whether the agent has the key, so it cannot open the door.
- A restricted dataset that never includes key pickups yields a model with no in-support entry for that transition. Our rollout applies an explicit stay-put fallback, rather than estimating a confident prediction.
The dynamics model is an empirical count table over observed transitions: for each pair, we store next-observation outcomes with reward and termination sums. It is a pedagogical analogue, not an autoregressive Transformer or a smoothed estimator. The toy isolates representation and data support without large-model training. A more expressive model might use action history to recover some information hidden from the current lossy observation, but changing the model class alone does not guarantee that it will learn the missing pickup rule or handle an unobserved action branch.
Defining the toy world
We start by defining the grid, the tile alphabet, and the transition function.
from collections import defaultdict
import numpy as np
## Grid and tile alphabet
H, W = 5, 5
N = H * W
EMPTY, WALL, KEY, HAZARD, GOAL, AGENT = 0, 1, 2, 3, 4, 5
TILE_NAMES = ["empty", "wall", "key", "hazard", "goal", "agent"]
## Fixed tile positions
KEY_POS = 6 # (row 1, col 1)
HAZARD_POS = 3 # (row 0, col 3)
WALL_POS = 13 # (row 2, col 3)
GOAL_POS = 24 # (row 4, col 4)
## Actions: 0=up, 1=down, 2=left, 3=right
MOVES = np.array([[-1, 0], [1, 0], [0, -1], [0, 1]])
ACTION_NAMES = ["up", "down", "left", "right"]The transition function is deterministic. The agent moves, walls block, hazard ends the episode with a penalty, and the goal is only enterable once the agent holds the key.
def step(state, a):
"""Advance the toy world by one action. Returns (next_state, reward, done)."""
pos, has_key = state
r, c = divmod(int(pos), W)
dr, dc = MOVES[int(a)]
nr, nc = r + int(dr), c + int(dc)
if not (0 <= nr < H and 0 <= nc < W):
nr, nc = r, c
npos = nr * W + nc
if npos == WALL_POS:
npos = pos
new_has_key = has_key or (npos == KEY_POS)
reward, done = 0.0, False
if npos == HAZARD_POS:
reward, done = -1.0, True
elif npos == GOAL_POS:
if new_has_key:
reward, done = 1.0, True
else:
npos = pos # locked door
return (npos, new_has_key), reward, doneRendering writes a tile to each cell and then overlays the agent sprite. The key sprite is drawn only while the agent does not hold the key.
def render(state):
pos, has_key = state
grid = np.full(N, EMPTY, dtype=np.int64)
grid[WALL_POS] = WALL
grid[HAZARD_POS] = HAZARD
grid[GOAL_POS] = GOAL
if not has_key:
grid[KEY_POS] = KEY
grid[pos] = AGENT
return gridA glance at the world makes the layout clear.
The picture shows the layout that the tokenizer must encode. What we will manipulate next is not the picture but the tokenization scheme that reads it, because the scheme decides which of these tiles survive into the model's input.
Tokenization: full and lossy
The full observation is the integer grid. The lossy observation replaces every KEY tile with EMPTY. This mimics a tokenizer that assigns no code to the key sprite, either because the codebook was trained with a small or because the reconstruction loss did not care about a tiny detail.
def lossy(obs):
aliased = obs.copy()
aliased[aliased == KEY] = EMPTY
return aliased
def agent_pos(obs):
return int(np.where(obs == AGENT)[0][0])For reachable states in this toy world, the full frame reveals whether the agent holds the key: the key sprite disappears after pickup. The lossy tokenizer removes that sprite in both cases. At the same apparent position, a key-held and a key-unheld state can therefore produce the same lossy observation but different outcomes for an action toward the goal. The lossy current observation is not Markov-sufficient; an action history that records the pickup could restore the distinction.
Generating training data
We generate trajectories under a uniform exploration policy and record tuples.
EPISODES = 3000
HORIZON = 30
rng = np.random.default_rng(0)
transitions = []
for _ in range(EPISODES):
state = (0, False)
obs = render(state)
for _ in range(HORIZON):
a = int(rng.integers(4))
next_state, r, done = step(state, a)
next_obs = render(next_state)
transitions.append((obs, a, next_obs, float(r), bool(done)))
state, obs = next_state, next_obs
if done:
break
perm = np.random.default_rng(3).permutation(len(transitions))
transitions = [transitions[i] for i in perm]
split = int(0.85 * len(transitions))
train_trans = transitions[:split]
test_trans = transitions[split:]We also generate a restricted dataset. Under the restricted policy, the agent never takes an action that would step onto the key cell, so key pickups never occur in the training data.
def restricted_behavior(state, rng):
options = []
for a in range(4):
(npos, _), _, _ = step(state, a)
if npos != KEY_POS:
options.append(a)
return int(rng.choice(options)) if options else int(rng.integers(4))
rng_restricted = np.random.default_rng(23)
restricted_trans = []
for _ in range(EPISODES):
state = (0, False)
obs = render(state)
for _ in range(HORIZON):
a = restricted_behavior(state, rng_restricted)
next_state, r, done = step(state, a)
next_obs = render(next_state)
restricted_trans.append((obs, a, next_obs, float(r), bool(done)))
state, obs = next_state, next_obs
if done:
breakThe restricted policy still visits most of the grid. What it avoids is the single transition into KEY_POS.
A count-based dynamics model
The model is a lookup table. For each observed pair we store the empirical distribution over next observations and the mean reward and termination under that outcome.
def fit_count_model(transitions, tokenizer=lambda x: x):
table = defaultdict(lambda: defaultdict(lambda: [0, 0.0, 0.0]))
for obs, a, next_obs, r, done in transitions:
key = (tuple(tokenizer(obs).tolist()), int(a))
nkey = tuple(tokenizer(next_obs).tolist())
entry = table[key][nkey]
entry[0] += 1
entry[1] += r
entry[2] += float(done)
return table
def predict_next(table, obs, a, tokenizer=lambda x: x):
key = (tuple(tokenizer(obs).tolist()), int(a))
dist = table.get(key)
if not dist:
return None, 0.0, 0.0
nkey = max(dist, key=lambda k: dist[k][0])
count, r_sum, d_sum = dist[nkey]
total = sum(v[0] for v in dist.values())
return np.array(nkey, dtype=np.int64), r_sum / count, d_sum / countThe predict_next function tabulates complete next-frame outcomes instead of factorizing tokens within a frame. Its empirical frame mode need not match greedy token-by-token decoding, especially when a lossy observation aliases distinct hidden states. We use the table to isolate representation and coverage effects, not as an implementation of the earlier autoregressive equation. The table conditions on the current tokenized observation and action; it has no longer history.
We fit three such models: full tokenization on the rich data, whose tokens distinguish the key sprite from the empty tile; lossy tokenization on the rich data, whose tokens alias the key sprite into the empty tile; and full tokenization on the restricted data, whose transitions never include a key pickup.
model_full = fit_count_model(train_trans)
model_lossy = fit_count_model(train_trans, tokenizer=lossy)
model_restricted = fit_count_model(restricted_trans)One-step accuracy looks fine
Before we test the models on multi-step plans, let us look at the metric that is easy to compute and misleading: held-out next-token accuracy.
def one_step_accuracy(model, test_trans, tokenizer=lambda x: x):
correct = 0
total = 0
for obs, a, next_obs, r, done in test_trans:
pred, _, _ = predict_next(model, obs, a, tokenizer=tokenizer)
if pred is None:
continue
total += 1
if np.array_equal(pred, tokenizer(next_obs)):
correct += 1
return correct / max(total, 1), total
acc_full, n_full = one_step_accuracy(model_full, test_trans)
acc_lossy, n_lossy = one_step_accuracy(model_lossy, test_trans, tokenizer=lossy)Full tokenization, held-out one-step accuracy: 1.000 (10393 transitions) Lossy tokenization, held-out one-step accuracy: 0.997 (10393 transitions) Restricted data, held-out one-step accuracy: 1.000 (4230 transitions)
The full and lossy models score comparably here, but the lossy model is evaluated in its own token space, where the key sprite has been removed. The restricted model's perfect score is more selective: its denominator excludes 6,163 of the 10,393 held-out queries because it has no table entry for them. Accuracy must be read alongside coverage, and neither number alone tells us how the planned key-to-door sequence will unfold.
Multi-step plan: key first, then the door
We hand the model a plan that goes down to the key, picks it up, walks to the goal, and stops. The same plan is executed in the true environment, in the full model, in the lossy model, and in the restricted model. The four resulting distance-to-goal curves reveal which models preserve the pickup rule and which ones stall short of the goal.
start_state = (0, False)
## Actions: down, right, down, down, down, right, right, right
plan = [1, 3, 1, 1, 1, 3, 3, 3]
def rollout_true(state0, actions):
state = state0
obs_hist = [render(state)]
rewards = []
dones = []
for a in actions:
state, r, done = step(state, a)
obs_hist.append(render(state))
rewards.append(r)
dones.append(done)
if done:
break
return obs_hist, rewards, dones
def rollout_model(model, obs0, actions):
obs = obs0.copy()
hist = [obs.copy()]
rewards = []
dones = []
for a in actions:
pred, r_hat, d_hat = predict_next(model, obs, a, tokenizer=lambda x: x)
if pred is None:
pred = obs.copy()
obs = pred
hist.append(obs.copy())
rewards.append(r_hat)
dones.append(d_hat)
return hist, rewards, donesWe roll out all four trajectories and compute the L1 distance from the agent's tile to the goal at each step.
obs0 = render(start_state)
true_hist, true_rewards, true_dones = rollout_true(start_state, plan)
full_hist, full_rewards, full_dones = rollout_model(model_full, obs0, plan)
lossy_hist, lossy_rewards, lossy_dones = rollout_model(
model_lossy, lossy(obs0), plan
)
restricted_hist, restricted_rewards, restricted_dones = rollout_model(
model_restricted, obs0, plan
)
def distance_curve(hist):
goal_r, goal_c = divmod(GOAL_POS, W)
dists = []
for obs in hist:
pos = agent_pos(obs)
r, c = divmod(pos, W)
dists.append(abs(r - goal_r) + abs(c - goal_c))
return np.array(dists, dtype=float)
true_dist = distance_curve(true_hist)
full_dist = distance_curve(full_hist)
lossy_dist = distance_curve(lossy_hist)
restricted_dist = distance_curve(restricted_hist)True reward trajectory: ['+0.0', '+0.0', '+0.0', '+0.0', '+0.0', '+0.0', '+0.0', '+1.0'] Full model reward traj: ['+0.0', '+0.0', '+0.0', '+0.0', '+0.0', '+0.0', '+0.0', '+1.0'] Lossy model reward traj: ['+0.0', '+0.0', '+0.0', '+0.0', '+0.0', '+0.0', '+0.0', '+0.0'] Restricted model reward: ['+0.0', '+0.0', '+0.0', '+0.0', '+0.0', '+0.0', '+0.0', '+0.0'] True termination step: 8 Full model termination: 8 Lossy model termination: never Restricted termination: never
The true environment earns the reward and terminates at the goal. The full model predicts the same state path, reward, and termination step; the lossy and restricted models miss the pickup and do not predict a terminal goal on this branch. To compare all eight prescribed actions on the same chart, the model rollouts record predicted termination without stopping early. This is a fixed-action diagnostic, not a policy run.
The plot below shows the agent's distance to the goal under this fixed action sequence. The true environment reaches zero (goal achieved), and the full model follows it exactly. The lossy model stalls one step away at distance one because the aliased key tile makes the pickup-invisible observation indistinguishable across key-held states, so the model's argmax next-state never enters the goal. The restricted model reaches distance one one step later in the plan because the pickup transition was absent from its training data, so its prediction at the pickup step is wrong and stalls until it recovers. The two impaired models both end at distance one, but for different reasons.
The chart is a fixed-action branch test. It does not score picture similarity or next-token likelihood; it compares the predicted outcome of one specified plan with the environment's outcome. Two models miss the goal despite high conditional one-step accuracy. This is not a closed-loop test of a policy choosing new actions from feedback.
Out-of-support queries and explicit fallback
The restricted model deserves a closer look. It was trained on thousands of trajectories from the same grid, so it knows the world quite well, except that it never saw a key held. We can measure the effect crisply: what fraction of the key-held observations in the held-out test set does the restricted model have an entry for?
def has_key_from_obs(obs):
# Under the full tokenization, the key sprite is absent if and only if the agent holds it.
return not np.any(obs == KEY)
key_held_trans = [t for t in test_trans if has_key_from_obs(t[0])]
def support_fraction(model, transitions):
"""Fraction of test transitions with a matching observation-action entry."""
if not transitions:
return 0.0
hits = 0
for obs, a, next_obs, r, done in transitions:
key = (tuple(obs.tolist()), int(a))
if key in model:
hits += 1
return hits / len(transitions)
support_full = support_fraction(model_full, key_held_trans)
support_lossy = support_fraction(model_lossy, key_held_trans)
support_restricted = support_fraction(model_restricted, key_held_trans)Held-out key-held transitions: 5811 Full-model in-support fraction: 1.000 Lossy-model in-support fraction: 1.000 Restricted-model in-support fraction: 0.000
Both impaired models predict the wrong outcome for the fixed plan, but this code does not measure their confidence. The lossy model chooses a frequent next observation from aliased data. The restricted model has no entry for the pickup action; predict_next returns None, and rollout_model substitutes the unchanged observation. That fallback is deliberately simple, not a general safety guarantee. A neural autoregressive model normally emits a distribution even for a poorly supported query, so its uncertainty and calibration need separate tests.
The toy failure is benign, but the diagnostic carries over: inspect whether a proposed action branch lies within the model's data support, and test its predicted consequences against the real environment before trusting it for control.