Sufficient State and State Abstraction

Imagine a robot arm learning to stack blocks while a camera watches from across the room. In the background, an advertising billboard cycles through colors every few seconds. Every frame the robot's camera captures contains the billboard, and its current color predicts the next frame with high accuracy. Yet no policy for stacking blocks should care about it.

This little scene captures one of the central design questions in world modeling. The raw observation stream contains everything, including what matters and what does not. An observation-prediction loss assigns error to the billboard, so under limited capacity it competes with task-relevant features. If the policy only needs to stack blocks, the billboard contributes nothing to its decision. The question this chapter answers is: what exactly must a representation preserve, and what can it safely throw away?

Notice how different this question is from the ones we asked in earlier chapters. Those chapters asked how to build models that fit data or how to plan given a model. This chapter asks a question that comes logically before both: given that a representation must be smaller than the raw experience it summarizes, which parts deserve to survive the compression? Likelihood and reconstruction losses pressure a model to retain predictable or reconstructable detail; control losses assign weight according to decisions. This chapter makes that choice explicit.

The answer is subtle because there are two different standards of "must." A representation can preserve far more predictive detail than control needs. It can also preserve the information needed for an optimal decision while predicting very little of the observation stream. In finite-capacity learning, predictive and control objectives can therefore pull in different directions. This chapter develops the vocabulary for that choice: predictive and control-sufficient state, bisimulation as a formal notion of behavioral equivalence, minimal abstractions and nuisance invariance, and task-aware representation learning that ties what a model keeps to what an agent needs.

Building on Part III Chapter 1, we take observations as given: the agent sees $o_t$, acts with $a_t$, and does not see the environment's true state $x_t$. The job of a state abstraction is to compress the observation history $h_t = (o_1, a_1, \dots, o_t)$ into something a dynamics model and a planner can work with. The previous chapters of Part II told us what the compression has to feed: a transition model, a reward model, and a planner. This chapter tells us what information those consumers require. Keep the consumers in mind as we go: a representation is never sufficient in the abstract. It is sufficient for a purpose, and the first step in any principled design is naming that purpose precisely.

The word "sufficient" hides a fork in the road. Before formalizing anything, let's separate the two standards that the rest of the chapter builds on. Both share the same skeleton: a quantity of interest computed from the full history should be recoverable from the compressed state. They differ in what that quantity is, and under limited data or capacity, optimizing one quantity can hurt the other.

A representation is predictively sufficient if it supports accurate forecasts of future observations. The intuition is straightforward: if you handed a forecaster the compressed state instead of the full history, its predictions of what comes next would be no worse. Everything the history knew about the future, the compression kept. Formally, let $z_t = \phi(h_t)$ be the encoded state. Predictive sufficiency means that for any horizon $T$ and any future action sequence $a_{t:t+T-1}$,

where:

• $h_t = (o_1, a_1, \dots, o_t)$: the full observation-action history available at time $t$
• $z_t = \phi(h_t)$: the encoded (compressed) state produced by the abstraction map $\phi$
• $a_{t:t+T-1}$: the sequence of actions $a_t, a_{t+1}, \dots, a_{t+T-1}$ (one action per transition step across the horizon) used for conditioning; predicting the $T$ future observations requires exactly these $T$ actions
• $o_{t+1:t+T}$: the block of future observations $o_{t+1}, \dots, o_{t+T}$ whose distribution we want to forecast
• $p(\cdot \mid \cdot)$: the conditional probability of the future observations under the environment's true dynamics

Two technical points about this definition deserve a pause. First, the condition holds for all horizons and all action sequences, including sequences the current policy does not visit. A representation that predicts well under the current behavior policy but fails under a different action pattern is not predictively sufficient; it has memorized a slice of the future rather than summarized the dynamics. Second, $a_{t:t+T-1}$ is treated as an externally fixed intervention. The compressed state must support forecasts under actions the agent has not yet taken, rather than merely reproduce correlations under the behavior policy.

The equality says the compressed state $z_t$ carries the same information about future observations as the entire history $h_t$: conditioning on the full history adds nothing once $z_t$ is known. All information about the future observations that was present in the history survives the compression. This is the Markov condition from Part II Chapter 1, restated as a property of a learned map: $z_t$ is a sufficient statistic for the observation-process future. Read it as a statistical statement about information flow: the map $\phi$ may collapse histories, but it collapses only along directions that the future distribution cannot distinguish.

This is the standard that likelihood-trained world models target, although finite data, finite capacity, and imperfect optimization do not guarantee that they attain it. Maximizing the likelihood of future observations pressures the latent state to encode anything that improves prediction, including billboard colors, background textures, and other features with zero bearing on the task. The objective asks whether a feature reduces prediction error, and the billboard does. The belief-state perspective of Part II Chapter 2 makes the same point constructively: for a known POMDP, the exact filtered belief $b_t$ is predictively sufficient, since it summarizes the observation history for forecasting under the model. It is a standard construction of a predictively sufficient state, though symmetries can permit further compression.

A representation is control-sufficient if it preserves everything needed to act optimally. The standard here is not "can you predict what you will see" but "can you decide what to do." Write $Q^*(h,a)$ for the optimal expected discounted return after taking action $a$ at history $h$ and acting optimally thereafter. A strong, directly checkable form of control sufficiency is

where:

• $h$: a full observation-action history, the input to the optimal action-value function on the left-hand side
• $\phi(h)$: the compressed state obtained by applying the abstraction map to $h$
• $Q^*(\cdot,a)$: the expected discounted return after action $a$ and optimal continuation, applied to histories on the left and encoded states on the right

On the right-hand side $Q^*$ is applied in the abstracted space, so the notation ranges over histories on the left and encoded states on the right. Preserving every action value preserves their ordering and therefore the maximizing action set. A weaker policy-only definition asks only for an optimal policy that is a function of $z_t$ and attains the history-optimal return from every history; it need not preserve suboptimal action values. Preserving the scalar $V^*(h)$ alone would not be enough because two histories can have equal optimal values but require different actions. This chapter uses the stronger action-value definition unless stated otherwise.

This requirement is weaker than preserving the full future distribution because it constrains one action value per available action, not every observable outcome. Large amounts of information can be destroyed as long as all optimal action comparisons are preserved. It is still demanding because those action values solve a long-horizon optimization problem. The billboard fails the relevance test: it changes neither reward nor any reward-relevant transition, so no optimal action depends on it.

It is tempting to assume that a low prediction loss automatically yields good decisions, or that control-sufficient state is automatically a good predictor. The second implication fails even for exact sufficiency. The first requires more care: an exactly predictive state can support control when rewards are observable from the predicted trajectories, but a finite-capacity model trained on average prediction loss may still discard a small control-critical signal.

Control-sufficient without predictive sufficiency. Consider an agent moving down a corridor toward a reward at the far end. A state variable holding "steps remaining to the goal" supports the optimal policy perfectly. It predicts essentially nothing about the observations along the way: which posters hang on which walls, whether a door on the left is open. The agent can act optimally while forecasting little of the stream because the reward ignores most of it.

A predictive objective can miss control-critical information. A perfectly predictive state retains the cue in the following example; the failure comes from capacity and loss weighting, not from exact predictive sufficiency. Suppose one pixel in a $128 \times 128$ camera image turns red exactly one step before a reward becomes available. A model trained to minimize uniformly averaged pixel error gives that pixel one sixteen-thousandth of the aggregate weight, since a $128 \times 128$ image has $16{,}384$ pixels. Under limited capacity, the representation can discard it with little change in the aggregate score, crippling the downstream planner. A feature can be critical for control while contributing almost nothing to average prediction error. Flip the numbers around to see the asymmetry clearly: a feature occupying one pixel can be worth an entire reward, while a feature occupying half the screen, like the billboard, can be worth nothing. Average pixel error measures appearance, not consequences.

To make this asymmetry concrete before we formalize it, consider two features in a 128 x 128 camera stream: a billboard that fills half the frame and carries zero control consequence, and a one-pixel reward cue that gates the entire task. The next cell compares each feature's share of pixels, which is also its loss weight under equal per-pixel squared error, with its control relevance, measured as the normalized drop in optimal value if the feature is discarded.

The billboard wins on pixel share and loses on consequences; the cue does the reverse. Under the thought experiment's equal per-pixel residual, uniform averaging ranks them in the wrong order for control. Different residual magnitudes or loss weights can change the prediction-side ranking, which is precisely why the loss definition matters.

This asymmetry is a practical trap. As we saw in Part II Chapter 6, model-based reinforcement learning compounds rollout error, and the errors that matter for planning are not uniformly distributed over prediction targets. An agent whose latent state faithfully models billboards but has drifted on a small, reward-critical signal will plan confidently toward bad outcomes. The planner consumes the model's predictions as if they were ground truth, and no amount of search can compensate for a signal that was never encoded. The formal tool for saying precisely which information must be preserved is bisimulation, and we turn to it now.

When we compress states, we are implicitly declaring some states interchangeable. Two histories, or two states, land in the same bucket of the abstraction, and the model treats every member of a bucket identically. The question is which equivalence relation is the right one: which partition of the state space collapses the right things together and keeps the right things apart. For control, the answer has a classical name, and it is worth understanding it deeply because nearly every modern abstraction method is an approximation of it.

Two features of this definition deserve attention. First, the second condition is self-referential: whether two states are equivalent depends on where their successors land, which depends on the equivalence relation itself. This circularity is resolved by a fixed point, and the coarsest bisimulation partition exists and is unique for a finite MDP. Partition refinement starts from blocks whose reward vectors already agree (or includes rewards in the refinement signature), then repeatedly splits blocks whose members disagree on successor-class probabilities. Each strict split refines a partition of a finite set, so the process terminates at the unique coarsest stable partition. Second, the definition quantifies over all actions. A pair of states that matches under the current policy's actions but diverges under an untried action is not bisimilar.

The reason bisimulation is useful for decision making is a preservation theorem. If $s \sim t$, then $Q^*(s,a) = Q^*(t,a)$ for every action, so $V^*(s) = V^*(t)$ and the optimal action sets agree. You can prove this by induction on value iteration: rewards are equal by condition 1, and each backup averages a class-constant value function over next states, whose class probabilities agree by condition 2. Since value iteration converges to $V^*$, equality propagates to the limit. The converse does not hold for a fixed task: non-bisimilar states can still share an optimal value or optimal action.

The quotient MDP obtained by collapsing each bisimulation class to a single state is the coarsest exact bisimulation quotient. Planning in the quotient returns the same optimal action values, and an optimal quotient policy lifts back to the original states. This guarantee is stronger than preserving only one task's optimal policy: it preserves actionwise reward and transition behavior at the class level. The distinction matters because a task-specific control abstraction can sometimes merge non-bisimilar states that happen to require the same decision. The construction echoes an idea from Part II Chapter 3 on system identification: define equivalence by what an external observer interacting with the system can detect, then retain what those interactions distinguish.

Historically, the concept has several ancestors worth knowing. Bisimulation originated in concurrency theory as an equivalence between labeled transition systems, introduced by Park and developed extensively by Milner. There it answered the question of when two programs are behaviorally interchangeable, that is, when no experiment, no sequence of observations, can tell them apart. Its probabilistic generalization is due to Larsen and Skou, who extended the notion from deterministic observation of labels to the probabilistic setting where transitions are distributions and the "experiments" become tests on probabilities of reaching groups of states. In the Markov chain literature, the same idea appears as lumpability: a Markov chain can be aggregated over a partition exactly when transition probabilities out of each block depend only on the block, which is condition 2 with rewards removed. And in decision theory, Ravindran and Barto's MDP homomorphisms generalize bisimulation to mappings that may also collapse actions, capturing settings where different actions are behaviorally equivalent in some states, as when a robot with five degrees of freedom has two joint configurations that produce identical effects in a given posture. The common thread across all of these is behavioral indistinguishability as the criterion for collapse.

Exact bisimulation is binary: two states either are or are not equivalent. Real learned representations need graded similarity. Embeddings live in continuous spaces, optimization needs smooth targets, and most pairs of states in a rich environment are neither identical nor maximally different. The bisimulation metric of Ferns, Panangaden, and Precup provides it, turning the yes-or-no question of bisimulation into a continuous distance with quantitative guarantees.

Define a distance $d$ between states as the fixed point of

where:

• $s, t$: two states of the MDP whose behavioral distance we measure
• $r(s, a)$: the expected immediate reward for taking action $a$ in state $s$, so the first term penalizes reward disagreement
• $P(\cdot \mid s, a)$: the transition distribution over next states when taking action $a$ in state $s$
• $W_d(\mu, \nu)$: the 1-Wasserstein (earth mover's) distance between next-state distributions $\mu$ and $\nu$, using the current state distance $d(x,y)$ as the transport cost
• $\gamma$: the discount factor, which weights how much successor disagreement matters

For bounded rewards, bounded candidate distances, and $0 \leq \gamma < 1$, the map has contraction factor at most $\gamma$. Iteration from any bounded initial distance therefore converges geometrically to the unique fixed point; the reward-difference initialization used below is a convenient starting point, not a theorem requirement. Successor disagreement is discounted by $\gamma$, matching how future rewards enter the return.

The structure of the formula is worth reading closely. The first term says states are close if their immediate rewards agree. The Wasserstein term says states are close if an adversary moving probability mass between their successor distributions pays little under the current metric. The earth mover's picture is the right intuition: to transform the successor distribution of $s$ into that of $t$, you must shovel probability mass from one next state to another, and the cost of each shovel-load is the metric distance between the two next states. If the successors are behaviorally similar, the shoveling is cheap and the states are close; if it requires moving mass between behaviorally distant regions, the distance is large. It is self-referential in exactly the way bisimulation is, and the contraction converts that circularity into a computable fixed point. Two important consequences:

• $d(s, t) = 0$ if and only if $s \sim t$, so the metric recovers exact bisimulation as its zero set.
• Under the displayed normalization, bounded rewards, and $0 \leq \gamma < 1$, $d$ upper-bounds the optimal value difference: $\lvert V^*(s) - V^*(t) \rvert \le d(s, t)$.

That second bullet is the headline. It converts a geometric statement about embeddings into a decision-theoretic statement about performance: if your encoder maps states so that nearby states in the embedding have small bisimulation distance, then collapsing them, averaging over them, or generalizing between them changes the optimal value by a bounded, controllable amount. This is the certificate that generic similarity measures such as pixel distance cannot provide. Two frames differing by a billboard color are far apart in pixel space and at distance zero in the bisimulation metric; two frames differing by a one-pixel reward cue are nearly identical in pixel space and far apart in the metric.

We can see this inversion numerically on a gridworld we can compute exactly. Using the lamp corridor structure (five positions, a lamp bit that flips with probability 0.3, reward only at position 4, discount 0.9), we iterate the bisimulation metric fixed point from its reward-difference initialization. The map is a $\gamma$-contraction; for this instance, the maximum change between iterates first falls below 1e-12 on sweep 263. The value-bound guarantee then lets us compare value differences with the metric.

The lamp-variant curve lies exactly at zero: the metric recovers the bisimulation relation as its zero set, so those pairs are provably interchangeable. The adjacent-position curve is strictly positive and increases as the reward at position 4 gets closer, matching the value-difference bound. Pixel distance would order these pairs in exactly the wrong way.

Exact bisimulation is expensive: the fixed-point iteration over all state pairs is at least quadratic in the state count, and it depends on the full reward and transition model, which is exactly what we do not have when learning from data. In a visual control problem with continuous observations, there is no tabular state list to iterate over at all. Three relaxations matter in practice, and each trades a piece of the guarantee for tractability.

• On-policy bisimulation metrics replace the max over actions with an expectation under a fixed policy $\pi$, giving a distance $d_\pi$ that quantifies behavioral equivalence under that policy alone. This can be estimated from the transitions the policy generates rather than the response to every action everywhere, and it connects to policy iteration: as the policy improves, the metric can be recomputed. The catch is that $d_\pi$ can merge states that a better policy would distinguish. If the current policy never takes the action that separates two states, the metric has no evidence of the difference.
• $\pi$-bisimulation and planning sufficiency weaken the requirement from optimal value under all policies to value under a family of policies, which calibrates the abstraction to the tasks in scope. The wider the family, the closer the requirement comes to exact bisimulation; the narrower, the more task-specific it becomes.
• Sampled and learned metrics replace the exact Wasserstein step with surrogates: distances between learned embeddings, reward prediction differences plus observation distances, or contrastive objectives. The deep methods later in this chapter all live here. The common move is to use the fixed-point equation as a training signal rather than a computation: update the embedding so that the distance between embeddings of two states approximates their reward difference plus a discounted successor-distance term, and let the iteration happen in the parameters rather than in a table.

The trade is between exact, task-general guarantees and scalable estimation. Exact bisimulation covers all policies but needs a known model and tabular access. On-policy metrics are blind to actions the policy has not tried. Learned surrogates scale to pixels but need explicit approximation-error control before the exact bound applies. Any claim of behavioral fidelity should therefore name the policy family and task distribution.

Sufficiency alone does not make a representation good. The identity map is perfectly sufficient and perfectly useless. It keeps everything, predicts everything, decides everything, and compresses nothing. The second half of the design problem is minimality: among all sufficient representations, prefer the one that discards the most. Sufficiency sets a lower bound on what to keep; minimality sets a pressure to keep nothing beyond it. A good abstraction lives exactly at the intersection.

For exact bisimulation, minimality has a crisp answer within the class of bisimulation-respecting partitions. Collapsing each bisimulation class to a single point yields the coarsest exact quotient: any coarser partition violates reward agreement or successor-class agreement. This is not necessarily the smallest representation for one fixed policy or task, because non-bisimilar states can happen to share optimal actions. The number of quotient classes is an abstract-state cardinality, not a geometric latent dimension. In the corridor example, exact bisimulation yields five classes rather than ten because it removes the lamp bit while keeping the five behaviorally distinct positions.

This gives us a precise language for a failure mode we flagged in the previous section. Reconstruction loss rewards an encoder for retaining every feature the decoder can use, including billboards, background texture, and lighting. By itself, it does not target a minimal control-sufficient representation. With enough capacity, any lossless invertible encoding can minimize reconstruction loss; under a bottleneck, compression follows the weighting of reconstruction errors and regularization rather than control consequences.

The complement of sufficiency is invariance. Call a transformation of the observation a nuisance if it changes $o_t$ but not the sufficient state. The billboard color cycle is a nuisance. So are camera pose jitter, lighting changes, texture swaps on irrelevant objects, and the identity of visually distinct but dynamically identical objects. The unifying property is that the transformation moves pixels without moving consequences: after applying it, every reward, every transition probability, and every optimal decision is unchanged. Formally, if a group $G$ acts on observations and each $g \in G$ is a nuisance, we want our encoder, viewed as acting on single observations here (an extension of $\phi$ to histories), to satisfy

where:

• $G$: a group of transformations acting on observations, each element of which is a nuisance (for example, rotations, color shifts, or lamp-color flips)
• $g \cdot o$: the observation $o$ after transformation $g$ has been applied
• $\phi$: the encoder mapping observations to the compressed state (here applied to single observations; in the predictive-sufficiency section it was applied to histories)

The group structure means transformations compose and invert. The equation makes every orbit $\{g \cdot o : g \in G\}$ lie inside one fiber of $\phi$, where a fiber is the set of observations mapped to the same encoded state. Group orbits need not exhaust behavioral equivalence: two observations can be interchangeable without being connected by a chosen transformation. For an ideal task-specific abstraction, each fiber should be a behavioral-equivalence class, while the nuisance orbits supply verified subsets that must remain inside those fibers.

Invariance is not free and not always wise. A transformation is a nuisance only relative to a task distribution. Cropping can be appropriate when camera translation truly leaves every task-relevant relation intact, but it is dangerous when reward depends on location, such as placing an object in a designated zone. Every augmentation is an implicit claim about what the controller may ignore. If that claim is wrong, no downstream function of the fixed invariant representation can recover the discarded evidence. Retraining the encoder from retained raw observations may restore it, but the current representation cannot.

The evidence lower bound used for a variational autoencoder contains a reconstruction term and a prior-matching term. Reconstruction error pressures the latent to carry details the decoder cannot predict from its other context. A conditional or autoregressive decoder can model some uncertainty without encoding every realized detail, so the pressure depends on decoder structure and likelihood choice. Increasing reconstruction weight still favors observation detail according to that likelihood, not control consequence, which creates the tension with task-specific abstraction.

This does not make such systems useless, and we will meet successful ones in the architecture chapters of Part V. It means their representations are abstractions only by side effect. Compression happens because the latent is small relative to the observation, so the model is forced to prioritize, but the prioritization follows reconstruction error, not control consequence. The remedy is to add terms that pull toward control relevance: reward prediction heads, value function auxiliary losses, or explicit bisimulation-shaped distances. Each of these adds a gradient pathway whose pressure is "encode this because it matters for decisions" rather than "encode this because it appears in the pixels." The system in Part III Chapter 3, which treats the latent state as a random variable with a learned prior, will give us the machinery for making this trade-off explicit in the objective.

Once we accept that representations should serve control, a design space opens. This section walks through it. The unifying idea is that the consumers of a representation, the losses that grade it, and the invariances baked into it jointly determine what the encoder keeps and what it discards, so each of them is a design lever to be set deliberately rather than inherited by default.

The general recipe: train an encoder $\phi$ jointly with objectives that reward keeping control-relevant information and penalties or losses that push out the rest. The consumers of the representation determine its shape. Every item on the following list is a different answer to the question "what gradient signal tells the encoder that a feature matters for decisions?," and they differ in how directly they measure decision relevance and how expensive they are to compute. Common choices include:

• Reward prediction heads. Attach a head predicting $r_t$ from $\phi(h_t)$ (and often $r_{t+1}$ from the transition). Reward relevance is the cheapest available proxy for control relevance, and it costs almost nothing to add. Its weakness is myopia: a feature that affects reward only three steps from now, through a chain of intermediate events, exerts little gradient on a one-step reward head, so the encoder may underweight features with long-delay consequences.
• Value function auxiliary tasks. Train a value estimate alongside the representation so that the embedding space organizes itself by long-term return rather than immediate appearance. States with similar futures land close together. This repairs the myopia of reward heads, since the value function integrates reward over the whole future, but it inherits the non-stationarity of the value estimate itself: early in training, the value signal is noise, and the encoder's early organization is shaped by that noise.
• Bisimulation regression. Compute or approximate target distances $d(s_i, s_j)$ and train the encoder so that embedding distances match them. This directly optimizes for the equivalence structure described above, and it inherits the guarantees of the metric to the extent the targets are accurate. In practice the targets come from the sampled or on-policy variants, so the guarantee is correspondingly softened.
• Augmentation invariance. Apply hand-designed nuisance transformations and enforce that $\phi$ maps augmented copies to the same embedding. This encodes "these differences do not matter" without needing reward information, which makes it usable in reward-free pretraining. The cost is the one identified earlier: every augmentation is a wager that the transformation really is a nuisance, and the wager is settled by downstream performance.
• Contrastive future prediction. Pull together observation pairs that lead to similar futures and push apart pairs that do not, which learns a predictively organized space without reconstruction. This is predictive sufficiency implemented with a selective loss: instead of scoring every pixel, it scores whether the representation can recognize which futures match. The latent-variable machinery behind this appears in the next chapter.
• Decision-focused latent model learning. Train the model to predict quantities used by planning rather than reconstruct observations. MuZero, for example, learns latent dynamics that predict reward, policy, and value for tree search without an observation decoder. That objective encourages control-relevant state but does not certify exact control sufficiency. We will return to this lineage in Part VIII.

These options are not mutually exclusive, and the strongest systems combine several: a reconstruction or contrastive term to organize the space, a reward or value term to orient it toward control, and augmentation invariances to carve out known nuisances. The art is in the weighting, because the terms can pull against each other exactly as the two sufficiency standards do.

Task-awareness has a price. A representation tuned to one reward function may discard information that a second task needs. An agent trained for corridor navigation has no pressure to encode the appearance of manipulable objects; a later grasping task may require those discarded features. Formally, using the action-value definition above, $\phi$ is control-sufficient for a task family when $Q^*_{\mathcal{T}}(h,a) = Q^*_{\mathcal{T}}(\phi(h),a)$ for every task $\mathcal{T}$ in the family and every action $a$, where:

• $Q^*_{\mathcal{T}}(h,a)$: the optimal expected discounted return for task $\mathcal{T}$ after taking action $a$ at history $h$
• $\mathcal{T}$: a task drawn from the task family. Widening the family adds preservation constraints, so a minimal exact sufficient state cannot become smaller.

The connection to predictive sufficiency needs one more condition. If the task family contains arbitrary bounded rewards on future observation trajectories, and action values are preserved for every forced future action sequence, then any distinguishable action-conditioned future can be turned into a distinguishing task. Under those conditions, universal task sufficiency recovers predictive sufficiency. With only scalar optimal values, that conclusion would fail because maximization can hide which named action produced which outcome distribution. Task-awareness buys compression by narrowing the family, at the cost of transfer.

This is why purely task-aware representations are brittle foundation-model ingredients, while purely task-agnostic ones waste capacity. Practical systems hedge, for example by training one encoder against an ensemble of reward functions, so the representation must serve every reward in the ensemble and therefore preserves their union of relevant features, or by keeping a general-purpose predictive model and extracting task-specific abstractions on demand, so the expensive general summary is computed once and the cheap task-specific compression is recomputed per task. The right point on this spectrum is an open research question, and it resurfaces when we discuss foundation world models in Part IX.

• Decide the consumer first: a rollout forecaster, a planner, a value function, or a task family. The consumer determines which sufficiency standard applies, and every subsequent choice follows from it.
• Prefer objectives that test long-horizon decision equivalence rather than relying only on one-step reconstruction. A one-step target can miss distinctions whose consequences emerge only after several transitions.
• Make every invariance explicit and justified by the task distribution. Write down, for each augmentation, the argument that the reward family cannot depend on what it destroys. If you cannot write the argument, do not apply the augmentation.
• Keep a reconstruction or prediction path available as one audit of what may have been discarded. A sufficiently capable, well-optimized decoder or probe can reveal retained information, but a failed reconstruction is not proof of information loss: limited decoder capacity or optimization can also be responsible. Treat failures as candidates for targeted follow-up tests against the task family.

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:

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:

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.

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.

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.

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.

The exact constructions in this chapter are directly computable for small, known MDPs. They require access to the reward and transition model, and dense metric procedures update at least a quadratic number of state pairs. Continuous or high-dimensional systems require additional structure, numerical approximation, or learned surrogates rather than tabular enumeration. Those approximations reintroduce error: value bounds depend on reward, transition, and metric estimation accuracy, and a learned representation can merge states whose relevant difference was absent from its training signal.

Task-awareness introduces its own fragility. Control-sufficiency is defined relative to a reward, and a representation distilled for one reward has no obligation to serve another. Worse, on-policy approximations inherit the current policy's blindness: if the policy never opens a particular door, the metric sees no behavioral difference behind it and may collapse those states, exactly erasing the information a future policy improvement would need. This is the representation-learning face of the model-exploitation problem from Part II Chapter 6: the agent's abstractions and its behavior co-evolve, and each can carve away information the other will later want. The dynamics are worth spelling out. A policy shapes the data distribution; the data distribution shapes the metric or value signal; the signal shapes the abstraction; and the abstraction constrains what the next policy iteration can even perceive. A loop with that structure can spiral: each cycle of policy improvement and representation refit removes a little more of the evidence that improvement depends on. Practical systems respond with optimistic exploration around abstraction boundaries, periodic re-fitting of the metric as the policy changes, or deliberately over-complete representations when the task family is uncertain, but each remedy trades away some of the compression that made abstraction attractive.

There is also an evaluation gap worth taking seriously. We can verify exact bisimulation on a known MDP, but for a learned representation in a complex environment we have no ground-truth partition to compare against. The practical criteria are indirect, and each has blind spots. You can check closed-loop performance, but a representation can score well on the training task while being fragile one task over, since training-task performance only audits the one consumer the training loop exercised. You can check invariance under augmentations, but that only audits the nuisances you thought to enumerate, and the dangerous nuisances are by definition the ones you did not. You can check value consistency across held-out states, but a value head can mask gaps in the underlying dynamics model, producing accurate-looking values built on a distorted internal geography. Multi-step rollout fidelity in the abstracted space, planning speed and quality in the quotient, and transfer to new reward functions form a more honest battery, and we will develop evaluation methodology systematically in Part XI.

These ideas give a precise answer to "what may a world model forget." When an exact quotient has far fewer states, each planning backup operates on the smaller quotient while preserving the action values guaranteed by bisimulation. Learned metrics and decision-focused latent models extend the same question to visual control, but without the exact certificate unless approximation error is controlled. The predictive-versus-control distinction also gives a practical diagnostic: when a world model predicts well and controls poorly, inspect whether its loss preserved the signals the planner consumes.

• Sufficiency is relative to a consumer. A state is predictively sufficient if it determines action-conditioned future observation distributions, and control-sufficient if it preserves the action values needed to choose optimally. Control can need far less than prediction, while finite-capacity prediction losses can underweight control-critical information.
• Bisimulation gives an exact behavioral quotient. Two states are bisimilar when they match in reward and in transition probabilities over equivalence classes for every action. Its quotient preserves action values and policies and is the coarsest exact bisimulation partition, though a fixed-task control abstraction can sometimes be coarser. Bisimulation metrics grade the equivalence with a bound on value differences.
• Minimality is separate from sufficiency. Reconstruction loss rewards retained observation detail but does not by itself target a minimal control abstraction. Nuisance invariance is a task-relative choice of what the fixed representation will ignore.
• Task-aware representation learning makes the trade-off explicit, adding control-relevant objectives and invariances at the cost of transfer across task families and coupling to the current policy. Under the universal task-family conditions defined above, preserving every task and forced action sequence recovers predictive sufficiency.

We now know what a world model's latent state must contain and what it may discard. The next question is how to represent what remains, especially when the state should carry uncertainty rather than a point estimate. The next chapter, Latent-Variable Representations, develops exactly that machinery: latent states as random variables with learned posteriors, and the variational machinery that trains them.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about sufficient state and state abstraction.