Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Sufficient State and State Abstraction › Worked Example: The Lamp Gridworld

Michael BrenndoerferJune 14, 20268 min read

World Models HandbookAll chapters

Part 3 of 12Representing Agents and Worlds

Chapter 14Sufficient State and State Abstraction › Worked Example: The Lamp Gridworld

Reading section 6 of 10 · Worked Example: The Lamp Gridworld · 8 minSection 6 of 10 · 8 min

Section 6 of 10 · 8 min read

Sufficient State and State Abstraction › Worked Example: The Lamp GridworldOne section · Sections in this chapter6 / 10 ↓

Sections in this chapter10

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

Let's make all of this concrete with the smallest example that separates the two notions of sufficiency. The example is deliberately constructed so that a single binary feature sits exactly on the fault line between the two standards: indispensable for prediction, worthless for control. Small enough to compute exactly, structured enough that every definition from the chapter has a visible referent.

The environment is a one-dimensional corridor with five positions $p \in \{0, 1, 2, 3, 4\}$ . An unrelated lamp is mounted at the far end, and its color bit $\ell \in \{0, 1\}$ flips with probability $0.3$ each step. The agent's observation is the pair $(p, \ell)$ . Three actions move the position left, hold, or right, with the position clipped at the boundaries. Reward is $1$ whenever the agent stands at position $4$ , and $0$ otherwise. We use discount $\gamma = 0.9$ . Note the two key structural facts before computing anything: the lamp evolves autonomously, unaffected by every action, and the reward depends on position alone. These two facts are precisely what bisimulation will exploit.

What does prediction require? To predict the next observation exactly, you need both components: the position determines where you will be, and the lamp bit determines the distribution of the next lamp color, since $P(\ell_{t+1} = \ell_t) = 0.7$ . Knowing both yields probabilities $0.7/0.3$ . If the current lamp bit is omitted and its conditional marginal is balanced, as it is under the stationary distribution and the uniform sampling used below, the next lamp color is a coin flip. A predictively sufficient state must include $\ell$ .

What does control require? The reward and position dynamics never consult the lamp. Intuitively, $(p, 0)$ and $(p, 1)$ should be interchangeable: same rewards available, same moves available, and the lamp evolves identically from either value. Bisimulation makes this precise, and the partition refinement algorithm confirms it:

Round 0. Group states by immediate reward: all states with $p < 4$ form one block; the two states with $p = 4$ form another. Notice that the lamp bit plays no role here already, because reward ignores it.
Round 1. Check transitions block by block. From $(3, \ell)$ , moving right lands in the $p = 4$ block. From $(2, \ell)$ , moving right lands in the $p \le 3$ block. Different successor blocks, so the non-goal block splits: $\{p = 3\}$ separates from $\{p \le 2\}$ . The states $(3, 0)$ and $(3, 1)$ still produce the same signature because their successors land in the same blocks with the same probabilities. The position dimension splits; the lamp dimension does not.
Rounds 2 and 3. The same argument peels off $p = 2$ , then $p = 1$ , leaving $p = 0$ as its own block. At every round, the signature of a state is computed from the blocks its successors land in, and the lamp bit never changes which block a successor lands in, because $(p', 0)$ and $(p', 1)$ are always in the same block.

The fixed point has five classes, each containing the two lamp variants of one position. Ten observations compress to five behavioral states, and the optimal value function is constant on each class:

V^*(p, 0) = V^*(p, 1) = \gamma^{4-p} \cdot 10, \qquad p \le 4,

where:

$\gamma = 0.9$ : the discount factor
$4 - p$ : the number of steps needed to reach the reward position $p = 4$ by moving right at every step
$\gamma^{4-p}$ : the discount applied to the first reward, collected $4-p$ steps from now
$10$ : the geometric series $\sum_{k=0}^{\infty} \gamma^k = 1/(1-\gamma) = 1/(1-0.9)$ , since the agent can stand at $p = 4$ and collect reward $1$ every step thereafter

To see why this holds, step through the derivation. The optimal policy moves right whenever $p < 4$ , so from position $p$ the first reward of $1$ arrives after $4 - p$ steps, and one arrives at every step after that. The optimal value is therefore the discounted sum of that reward stream:

\begin{aligned} V^*(p, \ell) &= \sum_{k=4-p}^{\infty} \gamma^{k} \cdot 1 && \text{(first reward at step } 4-p\text{)} \\ &= \gamma^{4-p} \sum_{k=0}^{\infty} \gamma^{k} && \text{(reindex the sum)} \\ &= \gamma^{4-p} \cdot \frac{1}{1 - \gamma} && \text{(geometric series)} \\ &= \gamma^{4-p} \cdot 10 && \text{(substitute } \gamma = 0.9\text{)} \end{aligned}

The value does not depend on $\ell$ , which is exactly the statement that the lamp bit is control-irrelevant: the two lamp variants of each position are bisimilar, so they must share the same optimal value. The preservation theorem predicted this before any computation: equal values on classes are a consequence of the definition, not a coincidence of this environment.

Notice what the refinement never did: it never needed the observation model. Bisimulation is defined purely on reward and transition structure, which is exactly why it can discard the lamp. The lamp's predictiveness, however strong, never enters the calculation. A prediction-driven abstraction, by contrast, is forced to keep the lamp, because the lamp bit is the best predictor of the next lamp bit, and any forecaster that drops it pays a real loss in likelihood. The example is tiny, but the pattern scales: every task-irrelevant yet predictable feature of a high-dimensional observation is a billboard waiting to consume your model's capacity.

The chapter's central distinction can be drawn on a single pair of axes. Keeping the lamp gridworld fixed, we vary what the representation retains and score it against two consumers at once: a next-lamp forecaster and the optimal value function. The forecaster is the Bayes-optimal conditional predictor given the retained features, so the prediction error shown is the best any representation with that retention level could achieve. The value gap is a useful diagnostic for this example, but scalar value equality alone is not a general certificate of the strong action-value notion of control sufficiency; that certificate comes from the exact bisimulation argument above.

In[6]:

Code

import numpy as np

gamma_s = 0.9
lamp_flip_s = 0.3
n_pos_s = 5
states_s = [(p, l) for p in range(n_pos_s) for l in range(2)]

## Ground-truth optimal values via value iteration (tolerance 1e-12)
idx_s = {s: i for i, s in enumerate(states_s)}
V_s = np.zeros(len(states_s))
for _ in range(500):
    Q_s = np.zeros((len(states_s), 3))
    for i, s in enumerate(states_s):
        for k, a in enumerate([-1, 0, 1]):
            npos = min(max(s[0] + a, 0), n_pos_s - 1)
            Q_s[i, k] = (1.0 if s[0] == 4 else 0.0) + gamma_s * (
                (1 - lamp_flip_s) * V_s[idx_s[(npos, s[1])]]
                + lamp_flip_s * V_s[idx_s[(npos, 1 - s[1])]]
            )
    Vn = Q_s.max(axis=1)
    if np.max(np.abs(Vn - V_s)) < 1e-12:
        V_s = Vn
        break
    V_s = Vn


## Retention levels: 0 keeps position only, 1 keeps position + lamp.
## Predictive cost: expected next-lamp cross-entropy under the Bayes-optimal
## forecaster given the retained features.
def lamp_entropy(p_lamp):
    q = max(p_lamp, 1e-12)
    return -(p_lamp * np.log(q) + (1 - p_lamp) * np.log(1 - q))


pred_err = np.array(
    [
        lamp_entropy(0.5),  # position only: next lamp is a coin flip
        lamp_entropy(lamp_flip_s),  # with lamp: residual is the flip entropy
    ]
)

## Optimal-value diagnostic: max absolute value error after averaging within
## each representation class. The position-only classes merge lamp values.
V_position_only = np.array(
    [
        np.mean([V_s[idx_s[(p, lamp)]] for lamp in range(2)])
        for p in range(n_pos_s)
    ]
)
V_position_only_lifted = np.array([V_position_only[p] for p, _ in states_s])
ctrl_err = np.array([np.max(np.abs(V_s - V_position_only_lifted)), 0.0])

retention_labels = ["Position only", "Position + lamp"]

The prediction errors are analytic entropies of the lamp transition, and the optimal-value gaps are computed from exact value iteration, so no sampling is involved.

Out[7]:

Visualization

Prediction error drops when the lamp is kept while value error stays at zero. — Prediction error and optimal-value gap as a function of what the representation retains. Adding the lamp reduces predictive error to its Bayes floor, the residual flip entropy, while the value gap remains zero. Exact bisimulation, not the scalar value diagnostic alone, establishes strong control sufficiency here.

One retention decision, two consumers, two different responses. The prediction curve falls to its Bayes floor once the lamp is retained, while the optimal-value gap stays at zero. Together with the exact bisimulation proof, the picture shows why representation quality must be evaluated per consumer rather than collapsed into one scalar.