Sufficient State and State Abstraction › Task-Aware Representation Learning

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.