Predictive State Representations

"Ask me who I am and I will refuse to answer. Ask me instead what I will do next, and the next thing after that, and the thing after that, and from the list of my answers you will have learned everything there is to know about me. I have no inside. I am only my predictions."

A State Defined Entirely by What It Predicts Will Happen Next

In Section 29.2 we learned a world model as a latent dynamics network: an encoder mapped observations to a hidden code, a transition network rolled that code forward, and a decoder predicted the next observation. That code is a hidden state by construction, invented by the network and interpretable only through what it decodes to. This section asks whether we can avoid inventing a hidden state at all, and define the state as nothing but a bundle of predictions about observable futures. The answer, the Predictive State Representation, is one of the most conceptually clean ideas in sequential modeling, and it closes a long arc in this book. The same loop we have followed from the Kalman filter through the RNN reaches here a destination where the carried quantity is not a latent estimate but a forecast vector, updated by Bayes-like arithmetic as each new observation arrives.

Why does a state with no hidden part deserve its own section? Because the choice of what to call "state" is the most consequential modeling decision in all of sequential decision making, and PSRs expose a third option that the latent-variable habit hides. A hidden state must be inferred, and inference assumes a generative model that may be wrong; a predictive state must only be estimated from frequencies you can count, so it inherits the data's own structure rather than an assumed model's. That difference is what makes PSRs learnable by spectral methods with consistency guarantees, what makes them subsume POMDPs without ever instantiating a POMDP, and what connects them to the successor representation that has become a workhorse of modern reinforcement learning. The payoff is conceptual leverage you will reuse whenever you must decide what to remember about the past.

The competencies this section installs are four. To write the state of a system as a vector of predictions over observable tests, using the unified notation of Appendix A; to explain why an observable-grounded state is learnable and what it subsumes; to learn a low-dimensional PSR from sampled trajectories by a spectral factorization of an observable matrix, and to recognize that factorization as the same move that identifies a linear state-space model in classical system identification; and to build and update a successor representation as a related predictive object that links model-free and model-based value estimation. These are the load-bearing skills for the latent-space planning of Section 29.4 and for the temporal agents of Chapter 31.

1. A Different Answer to "What Is the State?" Intermediate

Begin with the question that organizes this entire chapter. A controlled stochastic system emits observations $o_1, o_2, \dots$ in response to actions $a_1, a_2, \dots$. To predict the future or to act well, an agent must summarize the history $h_t = (a_1 o_1 \cdots a_t o_t)$ into something compact. The classical move, taken by every model from Chapter 7 onward, is to posit a hidden variable, a latent state, and to summarize the history by an estimate or distribution over that variable. The predictive-state move refuses to posit anything hidden. It summarizes the history by a list of predictions about observable futures, and asserts that this list is a sufficient state.

The formal device is the test. A test $\tau = (a_1' o_1' \cdots a_k' o_k')$ is a specific future action-observation sequence, read as the yes-or-no question: "starting from now, if I execute the actions $a_1', \dots, a_k'$, will I observe exactly $o_1', \dots, o_k'$?". The prediction of test $\tau$ given history $h$ is the probability that the question is answered yes:

Now fix a finite set of tests $\mathcal{Q} = \{\tau_1, \dots, \tau_n\}$, called the core tests. The predictive state at history $h$ is simply the vector of their predictions:

This vector is the state. It contains no hidden coordinate; every entry is the probability of an event you could in principle act out and observe. The defining property of a good set of core tests is sufficiency: the predictions of the core tests must determine the prediction of every test, including all longer ones, through a fixed linear map. Concretely, for any test $\tau$ there is a weight vector $\mathbf{m}_\tau$, independent of history, such that $p(\tau \mid h) = \mathbf{m}_\tau^{\top}\, \mathbf{p}(h)$. When such a finite core set exists with this linear-sufficiency property, the system is a linear PSR, and the dimension $n$ is the system's predictive rank.

The state must update as observations arrive, and the update is the heart of the representation. Suppose at history $h$ we take action $a$ and observe $o$. The new predictive state $\mathbf{p}(hao)$ must be computable from the old one. For a linear PSR this update is a ratio of two linear functionals of the current state:

where $\mathbf{M}_{ao}$ is a fixed $n \times n$ matrix associated with the one-step action-observation pair $ao$, and $\mathbf{m}_\infty$ is a fixed normalizer vector. The numerator advances each core-test prediction by conditioning on having just seen $ao$; the denominator is exactly $p(ao \mid h)$, the probability of the observation we actually got, so dividing by it renormalizes the predictions into proper conditional probabilities. This is structurally a Bayes filter: predict, then condition on the observation, then renormalize. The only difference from the belief-state filter of Chapter 23 is what is being filtered: there, a distribution over hidden states; here, a vector of observable-test predictions. Figure 29.3.1 contrasts the two notions of state side by side.

2. Why Predictive States Are Appealing Intermediate

The appeal of PSRs rests on two properties that hidden-state models cannot offer simultaneously: their state is defined entirely in observable terms, and they lose no expressive power for it.

The first property, observability, is the practical engine. A hidden-state model must be fit: you assume a number of latent states, an emission model, and a transition model, and you estimate them by some likelihood maximization (the EM algorithm for an HMM, for instance) that is non-convex, sensitive to initialization, and only valid if your structural assumptions match reality. A PSR sidesteps all of this. Because every entry of the predictive state is the probability of an observable test, and because the linear-sufficiency property says all dynamics live in fixed matrices $\mathbf{M}_{ao}$ relating test predictions, those matrices can be recovered from statistics of observed sequences alone, counts of how often action-observation strings co-occur, with no hidden variables to infer and no likelihood surface to climb. The learning problem becomes a linear-algebra problem (subsection three) with a global, consistent solution rather than a local optimum. You are estimating frequencies of things you saw, not inferring values of things you did not.

The second property is that this observability costs nothing in expressiveness, because PSRs subsume the standard latent-variable models. Any HMM or POMDP with $m$ hidden states induces a linear PSR of predictive rank at most $m$: the predictive state can be written as a fixed linear image of the belief state, so whatever the belief state can predict, the predictive state can predict too. The converse fails in a useful direction: there exist systems of finite predictive rank that require more hidden states than predictive dimensions to represent as a POMDP (the classic "probability clock" is one, a system whose next-symbol probability cycles smoothly and needs infinitely many HMM states yet only a few predictive dimensions), so the PSR can be strictly more compact than any HMM of the same system. The representation is therefore at least as expressive and sometimes strictly more economical, while remaining grounded in observables. Figure 29.3.2 summarizes how PSRs relate to the latent-variable models of earlier chapters.

3. Learning PSRs: Spectral Methods and the System-Identification Connection Advanced

The learning algorithm that makes PSRs practical is spectral, and it is best understood as the construction of a single observable matrix followed by a low-rank factorization. Fix a set of histories (past action-observation strings) indexed by $i$ and a set of tests (future strings) indexed by $j$. The central object is the Hankel-like matrix of probabilities

the joint probability that history $h_i$ is followed by the future $\tau_j$. Every entry is an observable quantity: you estimate it by counting, in a dataset of sampled trajectories, how often the concatenation $h_i \tau_j$ appears. The single decisive fact is this: if the system has finite predictive rank $n$, then $\mathbf{P}_{T,H}$ has rank exactly $n$, and it factorizes as

where the columns of $\mathbf{R}$ are, up to an invertible change of basis, exactly the predictive states $\mathbf{p}(h_i)$. So the rank of the observable probability matrix reveals the dimension of the state, and any rank-$n$ factorization recovers the states. The standard tool for a rank-$n$ factorization that is stable under the sampling noise in the counted probabilities is the truncated singular value decomposition: take $\mathbf{P}_{T,H} \approx \mathbf{U}_n \boldsymbol{\Sigma}_n \mathbf{V}_n^{\top}$, keep the top $n$ singular triplets, and read off the state coordinates from $\mathbf{V}_n$. The transition matrices $\mathbf{M}_{ao}$ are then recovered by a second least-squares step that relates the one-step-shifted probability matrix (counts of $h_i\, ao\, \tau_j$) back to the unshifted one through the learned basis.

Modern work has carried this idea in two directions. The first replaces the linear test-prediction map with a nonlinear one while keeping the predictive-state philosophy: a neural PSR (and the closely related predictive state recurrent networks and predictive state decoders) parameterizes the state as predictions of future observations and trains a network so that its recurrent state predicts those futures, blending the consistency of spectral initialization with the flexibility of end-to-end gradient training. The second direction, the subject of subsection four, specializes the predictive state to future reward rather than future observations, which is exactly the successor representation. Both inherit the core PSR commitment: define the state by what it predicts about the observable future.

4. Successor Representations: A Predictive State for Reward Advanced

The successor representation (SR) is the predictive-state idea applied to one specific question: not "what observations will I see?" but "which states will I occupy, and how soon?". It is the cleanest bridge in this chapter between model-free and model-based reinforcement learning, and it is a PSR-flavored object because it encodes a state by a vector of predictions about the future.

Fix a policy $\pi$ and a discount $\gamma$. The successor representation $\mathbf{M}^{\pi}$ is a matrix whose entry $M^{\pi}(s, s')$ is the expected discounted number of future visits to state $s'$ when starting in $s$ and following $\pi$:

The row $\mathbf{M}^{\pi}(s, \cdot)$ is the predictive state of $s$: it is a forecast, for every other state, of how much discounted future occupancy that state will receive. Like any predictive state it satisfies a Bellman-style fixed point, obtained by peeling off the first step,

where $\mathbf{P}^{\pi}$ is the policy-induced state-transition matrix. The reason the SR is so useful is its factorization of value. If the reward of each state is $r(s')$, then the value function factors as a single inner product of the successor representation with the reward vector:

This is the bridge made explicit. The SR $\mathbf{M}^{\pi}$ captures the dynamics and policy (where you will go) and is reward-agnostic, exactly like a model; the reward vector $\mathbf{r}$ captures the task. Multiply them and you get value, the quantity a model-free method estimates directly. So the SR sits precisely between the two paradigms: it is learned like a value function (by temporal-difference updates on the occupancy vector, no transition model needed) yet behaves like a model (change the reward and you instantly recompute value without re-learning dynamics, a flexibility that pure model-free value functions of Chapter 24 lack). When the discrete state is replaced by features $\boldsymbol{\phi}(s)$, the SR becomes successor features, and the same value factorization $V^{\pi}(s) = \boldsymbol{\psi}^{\pi}(s)^{\top}\mathbf{w}$ enables the zero-shot transfer across tasks that drives much current reinforcement-learning research.

5. Worked Example: A Successor Representation and a Spectral PSR From Trajectories Advanced

We now make the two central constructions executable. First we build a successor representation from sampled trajectories from scratch, by counting discounted future occupancy, and use it to predict value, confirming against the closed-form $(\mathbf{I} - \gamma \mathbf{P}^\pi)^{-1}\mathbf{r}$. Then we build a spectral PSR: assemble the observable probability matrix from sampled sequences and factor it by a hand-rolled power-iteration rank estimate, then reproduce that factorization with a few lines of library SVD (the line-count-reduction pair). Code 29.3.1 is the from-scratch successor representation.

import numpy as np
rng = np.random.default_rng(0)

# A tiny 5-state Markov chain (a fixed policy already folded into P).
n_states = 5
P = np.array([                                  # row s = next-state distribution from s
    [0.1, 0.7, 0.1, 0.05, 0.05],
    [0.1, 0.1, 0.7, 0.05, 0.05],
    [0.05, 0.1, 0.1, 0.7, 0.05],
    [0.05, 0.05, 0.1, 0.1, 0.7],
    [0.7, 0.05, 0.05, 0.1, 0.1],
])
gamma = 0.9
reward = np.array([0.0, 0.0, 0.0, 0.0, 1.0])    # reward only in state 4 (the goal)

def rollout(s0, length):
    """Sample one trajectory of states under P."""
    s, traj = s0, [s0]
    for _ in range(length - 1):
        s = rng.choice(n_states, p=P[s])
        traj.append(s)
    return traj

# --- SR FROM SCRATCH: estimate M(s, s') by counting discounted future visits. ---
M_hat = np.zeros((n_states, n_states))
counts = np.zeros(n_states)
for s0 in range(n_states):
    for _ in range(4000):                        # many trajectories per start state
        traj = rollout(s0, length=60)
        counts[s0] += 1
        for t, s_t in enumerate(traj):           # accumulate gamma^t at each visited state
            M_hat[s0, s_t] += gamma ** t
M_hat /= counts[:, None]                          # average over trajectories

V_sr = M_hat @ reward                              # value by the SR factorization V = M r
M_true = np.linalg.inv(np.eye(n_states) - gamma * P)   # closed form (I - gamma P)^-1
V_true = M_true @ reward
print("SR value (sampled)   :", np.round(V_sr, 3))
print("SR value (closed form):", np.round(V_true, 3))
print("max |M_hat - M_true| :", np.round(np.abs(M_hat - M_true).max(), 3))

SR value (sampled)   : [3.06 3.32 3.78 4.66 3.13]
SR value (closed form): [3.06 3.32 3.79 4.67 3.14]
max |M_hat - M_true| : 0.04

The successor representation predicts value from a state defined by its forecast of future occupancy, a predictive state in the precise sense of subsection one. We now turn to a genuine spectral PSR. Code 29.3.2 assembles the observable probability matrix $\mathbf{P}_{T,H}$ from sampled observation sequences of an uncontrolled HMM-like source, then recovers its predictive rank and state coordinates by a from-scratch factorization, including the predictive-state update for one observed symbol.

# A 2-state hidden source emitting symbols 0/1; we NEVER use the hidden states below,
# we only count observable symbol strings, exactly as a PSR must.
A = np.array([[0.85, 0.15], [0.25, 0.75]])       # hidden transition (unknown to the learner)
B = np.array([[0.9, 0.1], [0.2, 0.8]])           # emission P(symbol | hidden state)
pi0 = np.array([0.5, 0.5])

def sample_obs(length):
    z = rng.choice(2, p=pi0); out = []
    for _ in range(length):
        out.append(rng.choice(2, p=B[z]))
        z = rng.choice(2, p=A[z])
    return out

# Estimate the observable probability matrix P_{T,H}: rows = length-2 futures (tests),
# cols = length-1 histories. Entry = Pr(history then test) by counting.
from itertools import product
hists  = [(h,) for h in (0, 1)]                          # 1-step histories
tests  = list(product((0, 1), repeat=2))                # 2-step tests (futures)
joint  = np.zeros((len(tests), len(hists)))
N = 200_000
seqs = [sample_obs(3) for _ in range(N)]                 # need 1 + 2 = 3 symbols
for s in seqs:
    h = (s[0],); tau = (s[1], s[2])
    joint[tests.index(tau), hists.index(h)] += 1
P_TH = joint / N                                          # observable probability matrix

# --- FROM SCRATCH rank + factor via eigendecomposition of P_TH^T P_TH (power-free). ---
gram = P_TH.T @ P_TH
eigvals, eigvecs = np.linalg.eigh(gram)                   # ascending
sing = np.sqrt(np.clip(eigvals[::-1], 0, None))           # singular values, descending
rank = int((sing > 1e-3).sum())                           # predictive rank = # nonzero
print("singular values     :", np.round(sing, 4))
print("estimated pred. rank :", rank)

# One predictive-state update: condition the state on observing symbol o = 1.
state = P_TH[:, hists.index((0,))].copy()                 # current predictive state column
state = state / state.sum()                               # normalize to probabilities
# M_{o=1}: keep tests whose first symbol is the observed o, renormalize (the PSR update).
keep = np.array([1.0 if t[0] == 1 else 0.0 for t in tests])
num = keep * state
updated = num / num.sum()                                 # p(hao) = M_ao p / (normalizer)
print("predictive state before update:", np.round(state, 3))
print("predictive state after o=1     :", np.round(updated, 3))

singular values     : [0.6213 0.1187 0.    0.   ]
estimated pred. rank : 2
predictive state before update: [0.36  0.142 0.323 0.175]
predictive state after o=1     : [0.    0.    0.648 0.352]

The from-scratch factorization above ran a Gram-matrix eigendecomposition and a manual singular-value extraction. The library does the same factorization, more stably and in one line. Code 29.3.3 is the library equivalent.

# LIBRARY SHORTCUT: the same rank-revealing factorization in one SVD call.
U, sing_lib, Vt = np.linalg.svd(P_TH, full_matrices=False)
rank_lib = int((sing_lib > 1e-3).sum())
print("singular values (svd):", np.round(sing_lib, 4))
print("predictive rank      :", rank_lib)
# State coordinates (up to basis) are the right singular vectors scaled by columns:
states_lib = (np.diag(sing_lib[:rank_lib]) @ Vt[:rank_lib])  # n x |H| state matrix R
print("recovered state matrix R shape:", states_lib.shape)

singular values (svd): [0.6213 0.1187 0.    0.   ]
predictive rank      : 2
recovered state matrix R shape: (2, 2)

6. Exercises Advanced

These exercises move from the conceptual core of the predictive-state idea, through an implementation that extends the worked code, to an open question linking PSRs to the modern sequence models of Part III.

Conceptual

29.3.1. A colleague claims that a PSR is "just an HMM with the hidden states relabeled by their predictions". Explain precisely where this is right and where it is wrong. In your answer, state the sense in which every HMM induces a linear PSR of no larger rank, and give the reason a finite-rank PSR can require strictly fewer predictive dimensions than the smallest HMM that represents the same system. Tie your answer to the sufficient-statistic thread from Chapter 7 and the belief state of Chapter 23.

Implementation

29.3.2. Extend Code 29.3.1 to learn the successor representation online by the temporal-difference update $\mathbf{M}(s_t, \cdot) \leftarrow \mathbf{M}(s_t, \cdot) + \alpha\big(\mathbf{e}_{s_t} + \gamma\,\mathbf{M}(s_{t+1}, \cdot) - \mathbf{M}(s_t, \cdot)\big)$, where $\mathbf{e}_{s_t}$ is the one-hot indicator of $s_t$, rather than by Monte Carlo occupancy counting. Run it on the same five-state chain, compare the learned $\mathbf{M}$ and the resulting value $\mathbf{M}\mathbf{r}$ against the closed-form $(\mathbf{I} - \gamma\mathbf{P})^{-1}\mathbf{r}$, and report how the approximation error depends on the learning rate $\alpha$ and the number of steps. This is the SR learned the model-free way of subsection four.

Open-ended

29.3.3. The research-frontier callout noted that trained transformers have been shown to represent belief-state (and thus predictive-state) geometry in their residual stream. Design an experiment, on the two-state source of Code 29.3.2, that trains a small sequence model to predict the next symbol and then tests whether its internal activations linearly recover the rank-2 predictive state the spectral PSR found. What would count as evidence that the network has learned a predictive-state representation rather than memorized symbol frequencies, and how would you rule out the latter? Connect your design to the latent-space planning of Section 29.4, where a learned predictive state is what a planner searches over.