Sufficient State and State Abstraction › Code Implementation
World Models HandbookAll chapters Sufficient State and State Abstraction › Code ImplementationOne section · Sections in this chapter7 / 10 ↓
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 will implement the lamp gridworld, compute its exact bisimulation partition by partition refinement, verify that the optimal value function is constant on classes, and then use two small supervised probes to compare forecasting with value estimation. The exact partition supplies the control certificate; the probes measure, on a fixed dataset, how much the lamp feature contributes to each prediction target.
As throughout the book, we keep computation and rendering separate: all fitting and simulation happen first, and the visualization cells only draw from precomputed results.
import numpy as np
## Environment structure: position in 0..4, lamp bit in {0, 1}
n_pos, n_lamp = 5, 2
lamp_flip = 0.3
actions = [-1, 0, +1]
n_actions = len(actions)
states = [(p, l) for p in range(n_pos) for l in range(n_lamp)]
state_index = {s: i for i, s in enumerate(states)}
n_states = len(states)
def reward(s):
return 1.0 if s[0] == n_pos - 1 else 0.0
def transitions(s, a):
"""Outcome distribution for taking action a in state s."""
next_pos = min(max(s[0] + a, 0), n_pos - 1)
return [
(1 - lamp_flip, (next_pos, s[1])),
(lamp_flip, (next_pos, 1 - s[1])),
]The state space has ten elements: five positions times two lamp values. The transition function is exact, so we can run value iteration to get the ground-truth optimal values. This ground truth is what lets us check the preservation theorem numerically rather than taking it on faith.
gamma = 0.9
V = np.zeros(n_states)
for _ in range(500):
Q = np.zeros((n_states, n_actions))
for i, s in enumerate(states):
for k, a in enumerate(actions):
Q[i, k] = reward(s) + gamma * sum(
p * V[state_index[s2]] for p, s2 in transitions(s, a)
)
V_new = Q.max(axis=1)
if np.max(np.abs(V_new - V)) < 1e-12:
V = V_new
break
V = V_new
Q = np.array(
[
[
reward(s)
+ gamma * sum(p * V[state_index[s2]] for p, s2 in transitions(s, a))
for k, a in enumerate(actions)
]
for i, s in enumerate(states)
]
)
greedy = Q.argmax(axis=1)Next, the partition refinement algorithm. We start from blocks keyed by reward and repeatedly split any block whose members differ in the block-distribution of their successors, for any action. The implementation applies to finite tabular MDPs whose reward and aggregated-probability signatures can be compared exactly. With probabilities produced by different floating-point arithmetic paths, a numerical tolerance or canonicalization policy is required.
def coarsest_bisimulation(states, actions, trans_fn, reward_fn):
"""Exact partition refinement for a finite, exactly comparable MDP table."""
by_reward = {}
for s in states:
reward_signature = tuple(reward_fn(s, a) for a in actions)
by_reward.setdefault(reward_signature, []).append(s)
partition = [tuple(block) for block in by_reward.values()]
def block_of(s):
return next(i for i, b in enumerate(partition) if s in b)
while True:
new_partition = []
for block in partition:
subgroups = {}
for s in block:
sig = []
for a in actions:
dist = {}
for p, s2 in trans_fn(s, a):
j = block_of(s2)
dist[j] = dist.get(j, 0.0) + p
sig.append(tuple(sorted(dist.items())))
subgroups.setdefault(tuple(sig), []).append(s)
new_partition.extend(tuple(g) for g in subgroups.values())
if len(new_partition) == len(partition):
return partition
partition = new_partition
partition = coarsest_bisimulation(
states,
actions,
transitions,
lambda s, a: reward(s),
)
class_of = {}
for c, block in enumerate(partition):
for s in block:
class_of[s] = cStates: 10, bisimulation classes: 5 class 0: members [(0, 0), (0, 1)], V in [6.5610, 6.5610], spread 0.00e+00 class 1: members [(1, 0), (1, 1)], V in [7.2900, 7.2900], spread 0.00e+00 class 2: members [(2, 0), (2, 1)], V in [8.1000, 8.1000], spread 0.00e+00 class 3: members [(3, 0), (3, 1)], V in [9.0000, 9.0000], spread 0.00e+00 class 4: members [(4, 0), (4, 1)], V in [10.0000, 10.0000], spread 0.00e+00 Max within-class value spread: 0.00e+00
The partition collapsed ten states into five classes, and the optimal value function is numerically constant within every class. That is the preservation theorem made concrete: the two-element lamp dimension carries no control information, and bisimulation removes it without changing any optimal decision. Note also the structure of the classes: each contains one position and both lamp values, exactly matching the analytic derivation.
Now the quantitative contrast between the two objectives. We sample a seeded dataset of transitions and fit two small probes on two versions of the representation: one with the lamp feature included and one without. A logistic probe for next-lamp prediction tests predictive content, and a linear probe tests whether the feature helps estimate optimal value. The reported scores are in-sample diagnostics on this synthetic dataset, not generalization estimates or standalone certificates of control sufficiency.
from sklearn.linear_model import LinearRegression, LogisticRegression
rng = np.random.default_rng(0)
n_samples = 4000
X_pos = np.zeros((n_samples, n_pos))
X_lamp = np.zeros((n_samples, 1))
y_next_lamp = np.zeros(n_samples)
y_value = np.zeros(n_samples)
for n in range(n_samples):
s = states[rng.integers(n_states)]
a = actions[rng.integers(n_actions)]
next_pos = min(max(s[0] + a, 0), n_pos - 1)
next_lamp = s[1] if rng.random() > lamp_flip else 1 - s[1]
X_pos[n, s[0]] = 1.0
X_lamp[n, 0] = s[1]
y_next_lamp[n] = next_lamp
y_value[n] = V[state_index[s]]
## Noisy value targets: probes must estimate V from finite data, and an
## exactly interpolable target (R^2 = 1.0 for both probes) would make the
## control-side comparison vacuous.
y_value += rng.normal(0.0, 0.15 * y_value.std(), size=y_value.shape)
X_with = np.hstack([X_pos, X_lamp])
X_without = X_pos
acc_with = (
LogisticRegression(max_iter=200)
.fit(X_with, y_next_lamp)
.score(X_with, y_next_lamp)
)
acc_without = (
LogisticRegression(max_iter=200)
.fit(X_without, y_next_lamp)
.score(X_without, y_next_lamp)
)
r2_with = LinearRegression().fit(X_with, y_value).score(X_with, y_value)
r2_without = (
LinearRegression().fit(X_without, y_value).score(X_without, y_value)
)Next-lamp prediction accuracy, with lamp feature: 0.704 Next-lamp prediction accuracy, without lamp feature: 0.519 Value prediction R^2, with lamp feature: 0.9784 Value prediction R^2, without lamp feature: 0.9784
On this seeded sample, knowing the lamp bit raises in-sample next-lamp accuracy from to , close to the theoretical comparison between a balanced coin flip and the Bayes-optimal . The in-sample value-probe score is with either representation because the lamp contributes nothing to value. The same feature and dataset therefore produce different conclusions for the two consumers.
Finally, two figures. First, the value function plotted separately for each lamp value, showing the curves coincide exactly, alongside the bisimulation class structure itself. Then a bar chart of the probe scores.
V_by_pos = np.array(
[[V[state_index[(p, l)]] for p in range(n_pos)] for l in range(n_lamp)]
)
class_grid = np.zeros((n_lamp, n_pos), dtype=int)
for s, c in class_of.items():
class_grid[s[1], s[0]] = cThese two arrays hold the per-lamp value curves and the class index for every (position, lamp) cell, precomputed so the figure cells only contain presentation logic.
Read the two probes together. The left pair of bars shows that the lamp feature carries real predictive information; removing it drops prediction to a coin flip. The right pair shows that value estimation is untouched by its removal. An abstraction guided purely by prediction error would keep the lamp; an abstraction justified by bisimulation discards it with proof that no optimal decision changes. The full pipeline, environment construction, exact quotient, and probes, is small enough to modify: try adding a second, slower nuisance variable, or a lamp that occasionally blocks movement, and watch which classes split. The blocking variant is the more instructive of the two: the moment the lamp acquires any causal influence on reachable states or rewards, the class structure shatters, and a representation that had learned lamp-invariance from the benign version will be confidently, silently wrong on the new one. That is the augmentation wager from the invariance section, played out in miniature.