Sufficient State and State Abstraction › Bisimulation and Behavioral Equivalence

Michael BrenndoerferJune 14, 202612 min read
World Models HandbookAll chapters
Part 3 of 12Representing Agents and Worlds
Chapter 14Sufficient State and State Abstraction › Bisimulation and Behavioral Equivalence
Reading section 3 of 10 · Bisimulation and Behavioral Equivalence · 12 minSection 3 of 10 · 12 min

Book contents

Section 3 of 10 · 12 min read

Sufficient State and State Abstraction › Bisimulation and Behavioral EquivalenceOne section · Sections in this chapter3 / 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

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.

Exact bisimulation

Bisimulation

Two states ss and tt of an MDP are bisimilar, written s∼ts \sim t, if and only if for every action aa:

  1. Reward agreement: r(s,a)=r(t,a)r(s, a) = r(t, a).
  2. Transition agreement: for every equivalence class CC of ∼\sim itself, P(C∣s,a)=P(C∣t,a)P(C \mid s, a) = P(C \mid t, a).

That is, bisimilar states yield the same immediate reward under every action, and their transition distributions cannot be distinguished by any question expressible in terms of the equivalence classes.

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∼ts \sim t, then Q∗(s,a)=Q∗(t,a)Q^*(s,a) = Q^*(t,a) for every action, so V∗(s)=V∗(t)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∗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.

Bisimulation metrics: equivalence with a dial

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 dd between states as the fixed point of

d(s,t)  =  max⁡a(∣r(s,a)−r(t,a)∣  +  γ Wd(P(⋅∣s,a), P(⋅∣t,a))),\begin{aligned} d(s, t) \;&=\; \max_a \Big( \lvert r(s, a) - r(t, a) \rvert \;+\; \gamma \, W_d\big(P(\cdot \mid s, a),\, P(\cdot \mid t, a)\big) \Big), \end{aligned}

where:

  • s,ts, t: two states of the MDP whose behavioral distance we measure
  • r(s,a)r(s, a): the expected immediate reward for taking action aa in state ss, so the first term penalizes reward disagreement
  • P(⋅∣s,a)P(\cdot \mid s, a): the transition distribution over next states when taking action aa in state ss
  • Wd(μ,ν)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)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≤γ<10 \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 ss into that of tt, 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)=0d(s, t) = 0 if and only if s∼ts \sim t, so the metric recovers exact bisimulation as its zero set.
  • Under the displayed normalization, bounded rewards, and 0≤γ<10 \leq \gamma < 1, dd upper-bounds the optimal value difference: ∣V∗(s)−V∗(t)∣≤d(s,t)\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.

In[4]:
Code
import numpy as np
from scipy.optimize import linear_sum_assignment

gamma_m = 0.9
lamp_flip_m = 0.3
n_pos_m, n_lamp_m = 5, 2
actions_m = [-1, 0, +1]
states_m = [(p, l) for p in range(n_pos_m) for l in range(n_lamp_m)]
idx_m = {s: i for i, s in enumerate(states_m)}
n_m = len(states_m)


def reward_m(s):
    return 1.0 if s[0] == n_pos_m - 1 else 0.0


def trans_m(s, a):
    np_ = min(max(s[0] + a, 0), n_pos_m - 1)
    return [(1 - lamp_flip_m, (np_, s[1])), (lamp_flip_m, (np_, 1 - s[1]))]


R = np.array([[reward_m(s) for s in states_m] for _ in actions_m])
## Next-state probability matrices, one per action
P = np.zeros((len(actions_m), n_m, n_m))
for k, a in enumerate(actions_m):
    for i, s in enumerate(states_m):
        for pr, s2 in trans_m(s, a):
            P[k, i, idx_m[s2]] += pr

M = 10  # atoms per unit mass; transition probabilities are multiples of 0.1


def w1(mu, nu, ground):
    """Exact 1-Wasserstein distance between discrete distributions under the
    ground metric `ground`, via optimal assignment of equal-mass atoms."""
    src = np.repeat(np.arange(n_m), (M * mu).astype(int))
    dst = np.repeat(np.arange(n_m), (M * nu).astype(int))
    C = ground[np.ix_(src, dst)]
    ri, ci = linear_sum_assignment(C)
    return C[ri, ci].sum() / M


## Fixed-point iteration, gamma-contraction, tolerance 1e-12
D = np.abs(R[:, :, None] - R[:, None, :]).max(axis=0)
for metric_sweeps in range(1, 501):
    D_new = np.zeros_like(D)
    for k in range(len(actions_m)):
        for i in range(n_m):
            for j in range(n_m):
                # Wasserstein under the current metric D as ground distance
                W = w1(P[k, i], P[k, j], D)
                D_new[i, j] = max(
                    D_new[i, j],
                    np.abs(R[k][i] - R[k][j]) + gamma_m * W,
                )
    metric_delta = np.max(np.abs(D_new - D))
    if metric_delta < 1e-12:
        D = D_new
        break
    D = D_new

assert metric_delta < 1e-12

## Paired distances: same position, different lamp (bisimilar) versus
## adjacent positions with the same lamp (behaviorally distinct)
lamp_pair_dist = np.array(
    [D[idx_m[(p, 0)], idx_m[(p, 1)]] for p in range(n_pos_m)]
)
adjacent_dist = np.array(
    [D[idx_m[(p, 0)], idx_m[(p + 1, 0)]] for p in range(n_pos_m - 1)]
)

The iteration is seeded from a convenient reward-difference initialization; the contraction theorem does not require that particular starting point. The Wasserstein step is solved exactly by optimal assignment over these discrete transport plans, so the result is deterministic.

Out[5]:
Visualization
Zero metric distance across lamp variants, positive and rising distance across positions.
Bisimulation metric distance between the two lamp variants of each position (exactly zero) and between adjacent positions (strictly positive, growing toward the reward). The metric collapses nuisances and separates control-relevant differences.

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.

Approximate and policy-relative variants

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π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π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.