In reinforcement learning, almost every quantity one cares about (expected return, safety costs, exploration bonuses) is a functional of a single object: the state–action occupancy measure. That makes the space of occupancy measures the natural arena for RL, and it comes equipped with a geometry. This post is the beginning of a series that explores the shape of that space and what it can tell us about decision problems and the design of learning algorithms. In this first post, we'll see how the Bellman equations bend an otherwise flat probability space into a curved surface by coupling the policy to the states it visits.
Fix a finite Markov decision process with state space $\mathcal{S}$, action space $\mathcal{A}$, transition kernel $P(s' \mid s,a)$, reward $r : \mathcal{S}\times\mathcal{A}\to\mathbb{R}$, initial-state distribution $\mu \in \Delta(\mathcal{S})$, and discount $\gamma\in(0,1)$.
A transparent way to package a policy's long-run behaviour is through a generative model. Draw a random horizon $T \sim \mathrm{Geom}(1-\gamma)$, so that $\Pr(T = t) = (1-\gamma)\gamma^t$; run $\pi$ from $\mu$ for exactly $T$ steps; record the terminal pair $(S_T, A_T)$. Its law is
$$\omega^\pi(s,a) \;=\; \sum_{t=0}^{\infty}(1-\gamma)\gamma^t\,\Pr_\pi(S_t = s,\, A_t = a \mid S_0 \sim \mu),$$a genuine probability distribution on $\mathcal{S}\times\mathcal{A}$, normalised by construction. This is the discounted state–action occupancy measure of $\pi$. (Equivalently it is the stationary state–action marginal of the chain augmented with an absorbing state entered with probability $1-\gamma$ at each step; the mixture view is simply more transparent about normalisation.)
The occupancy is useful because the return is linear in it,
$$J(\pi) \;=\; \mathbb{E}_\pi\!\left[\sum_{t=0}^\infty \gamma^t\, r(S_t,A_t)\right] \;=\; \frac{1}{1-\gamma}\,\langle r,\, \omega^\pi\rangle, \quad \langle r,\omega\rangle := \sum_{s,a} r(s,a)\,\omega(s,a).$$The bracket $\langle r,\omega\rangle$ pairs two objects of different type: $r$ is a function on $\mathcal{S}\times\mathcal{A}$, while $\omega$ is a measure on it. The pairing requires no inner product, no notion of length or angle, only the duality between functions and measures. Reward and occupancy live in dual spaces, and the return is their contraction. Because $J$ is linear in $\omega$, the occupancy is a sufficient statistic for policy optimisation, and the set of achievable occupancies is where the geometry must be studied.
Which distributions arise as occupancies? Exactly those satisfying the Bellman flow equations,
$$\sum_{a} \omega(s,a) \;=\; (1-\gamma)\,\mu(s) \;+\; \gamma \sum_{s',a'} P(s \mid s', a')\,\omega(s', a'), \qquad \forall s,$$a conservation-of-flow condition: the mass entering $s$ (initial plus discounted transitions) equals the mass leaving it. Any $\omega$ in the simplex meeting these constraints is the occupancy of the policy $\pi(a \mid s) = \omega(s,a)/\sum_{a'}\omega(s,a')$. Write $\mathcal{O}$ for the resulting set: the occupancy polytope, an affine slice of $\Delta(\mathcal{S}\times\mathcal{A})$. We approach its geometry in two stages: first the ambient simplex with no dynamics, then the warp the Bellman equations impose.
Before dynamics enters, there is a purely kinematic fact about $\Delta(\mathcal{S}\times\mathcal{A})$. Any joint distribution factors as
$$\omega(s,a) \;=\; \rho(s)\,\pi(a\mid s), \qquad \rho(s) = \sum_a \omega(s,a),$$into a state marginal $\rho \in \Delta(\mathcal{S})$ and a conditional policy $\pi(\cdot\mid s) \in \Delta(\mathcal{A})$ for each state. The Fisher–Rao metric on the simplex, the inner product $g_\omega(u,v) = \sum_{s,a} u(s,a)\,v(s,a)/\omega(s,a)$ on tangent vectors, then decomposes along this factorisation into two blocks that are orthogonal at every point:
$$g_\omega \;=\; \underbrace{g^{\mathcal{S}}_{\rho}}_{\text{marginal block}} \;\oplus\; \underbrace{\sum_{s} \rho(s)\, g^{\mathcal{A}}_{\pi(\cdot\mid s)}}_{\text{conditional block}}.$$This is the chain rule for Fisher information, and it has a probabilistic reading. The conditional block is exactly Kakade's natural-policy-gradient metric. The weight $\rho(s)$ in front of each fibre is not an arbitrary choice: it is the visitation probability, which by the law of total variance is exactly how much state $s$ contributes to the overall variance of a function on $\mathcal{S}\times\mathcal{A}$. The cost of varying the policy at $s$ is proportional to how often you visit $s$. This structure is called a warped product: a base manifold (the state marginals $\Delta(\mathcal{S})$) fibred over policy space ($\prod_s \Delta(\mathcal{A})$), the fibre metric at each base point scaled by the base coordinate.
For two states and two actions the interior is parameterised by $(\beta, p, q)$, with $\beta = \rho(s_1)$ the free state marginal and $p = \pi(a_1\mid s_0)$, $q = \pi(a_1\mid s_1)$ the policy parameters. The metric reads
$$ g \;=\; \underbrace{\frac{\mathrm{d}\beta^2}{\beta(1-\beta)}}_{\text{marginal block}} \;+\; \underbrace{(1-\beta)\,\frac{\mathrm{d}p^2}{p(1-p)} \;+\; \beta\,\frac{\mathrm{d}q^2}{q(1-q)}}_{\text{conditional block}}, $$so the fibre ellipse has semi-axes $\sqrt{1-\beta}$ and $\sqrt{\beta}$: visiting $s_1$ more often compresses the $s_0$-fibre and expands the $s_1$-fibre. Drag $\beta$ below.
As $\beta \to 0$ the $s_1$-fibre axis collapses: when $s_1$ is almost never visited, varying $\pi(\cdot\mid s_1)$ barely shifts the occupancy, so it carries almost no Fisher cost. The axes stay perpendicular; only their lengths change. That is the warp, and the perpendicularity is the orthogonal (within/between) decomposition of the Fisher information.
So far there is no MDP, only a product simplex with $\beta, p, q$ as free parameters. The dynamics enter not as a new metric but as a constraint. Collect the Bellman flow into a linear operator $B$ acting on measures, so that the achievable set is $\mathcal{O} = \{\omega : B\omega = (1-\gamma)\mu\}$. Given a policy $(p,q)$ the marginal $\beta$ is then uniquely determined: it is no longer free. The occupancy polytope is the graph of the forced marginal $\beta = \beta(p,q)$, a two-dimensional surface inside the three-dimensional cube.
For a two-state recurrent MDP (from $s_0$, action $a_1$ moves to $s_1$ with probability $p$; from $s_1$, action $a_1$ returns to $s_0$ with probability $q$) the flow equations at $\gamma = 0.9$ give
$$ \beta(p,q) \;=\; \rho(s_1) \;=\; \frac{\gamma\,p}{(1-\gamma)+\gamma(p+q)}, $$increasing in $p$ and decreasing in $q$. Below, the floor is policy space, freely parameterised by $(p,q)$; the teal sheet is $\mathcal{O}$, lifted off the floor by $\beta(p,q)$; the dashed line is the lift of the current policy; the purple ellipse is the conditional block, lying in the horizontal plane; the green arrow is the marginal block, perpendicular to it.
Increasing $p$ pushes mass toward $s_1$, so the occupancy point climbs the sheet; increasing $q$ returns it to $s_0$, so it sinks. The lift height is $\beta(p,q)$, set by the Bellman solve: the policy is free, the marginal is not. The purple ellipse is the conditional block, in the horizontal policy plane; the green arrow is the marginal block, perpendicular in the ambient cube. Near the boundary the ellipse collapses in one direction while the arrow diverges (capped with a chevron when it leaves the frame): the conditional block shrinking and the marginal block blowing up are the same boundary singularity, viewed at right angles.
Pulling the ambient Fisher metric back to $\mathcal{O}$ through the embedding $(p,q)\mapsto \omega^\pi$ means substituting $\mathrm{d}\beta = \beta_p\,\mathrm{d}p + \beta_q\,\mathrm{d}q$ into the ambient metric, giving
$$ g_{\mathcal{O}} \;=\; \underbrace{\frac{(\beta_p\,\mathrm{d}p + \beta_q\,\mathrm{d}q)^2}{\beta(1-\beta)}}_{\text{marginal (tilt) contribution}} \;+\; \underbrace{(1-\beta)\,\frac{\mathrm{d}p^2}{p(1-p)} \;+\; \beta\,\frac{\mathrm{d}q^2}{q(1-q)}}_{\text{conditional block}}. $$Orthogonality is ambient, not induced. In the full cube the marginal and conditional blocks are separate summands with no cross term. On the occupancy surface this is broken, because $\beta$ is now slaved to $(p,q)$: the marginal contribution supplies a cross term $\beta_p\beta_q\,\mathrm{d}p\,\mathrm{d}q/(\beta(1-\beta))$. The two glyphs in the figure depict the ambient orthogonality; the actual geodesics on $\mathcal{O}$ mix both blocks.
The warp is the coupling. The slopes $(\beta_p,\beta_q)$, measuring how steeply the Bellman map lifts the sheet, are the coupling coefficients. A policy change drags the visitation by $\mathrm{d}\beta = \beta_p\,\mathrm{d}p + \beta_q\,\mathrm{d}q$; this visitation drag feeds back through the $1/(\beta(1-\beta))$ factor, which diverges at the boundary. Reaching a deterministic policy costs infinite Fisher length precisely because $\beta \to 0$ or $1$ there. The same factor $\rho(s)$ that shrinks a rarely-visited fibre in the conditional block is what down-weights that state in Kakade's gradient; a fibre axis collapsing and the marginal block blowing up are one boundary singularity seen at right angles.
The drag term is exactly what Kakade's metric omits: the natural policy gradient uses the conditional block alone and drops the marginal contribution, treating the visitation as if it were frozen. The gap between the two is not opaque. The visitation sensitivity admits a Neumann expansion in the transition kernel,
$$ \mathrm{d}\rho \;=\; \gamma \sum_{k\ge 0}\gamma^k\, T_\pi^{k}\,(\delta T)\,\rho, $$a discounted sum of a one-step kick propagated $k$ steps forward under the current dynamics. Truncating the sum at order $K$ interpolates between Kakade's metric ($K=-1$, drop the drag entirely) and the exact pullback ($K=\infty$), with error decaying like $\gamma^{K+1}$; each retained term is one rollout of the model. The Kakade metric is thus the zeroth member of a hierarchy, and the drag is the tail one chooses how much of to keep.
The Fisher–Rao metric degenerates at the boundary of the simplex, where some $\omega(s,a) \to 0$: the denominator blows up, and the closed simplex is not a Riemannian manifold. The correct structure is a stratified Riemannian space, a union of manifolds indexed by support sets, each carrying its own nondegenerate Fisher metric, glued along the face lattice. The simplest example is $\Delta^2$, with its $2^3-1=7$ strata. Pull them apart below.
The seven pieces are the strata, indexed by support set. Dashed lines show the closure order (the face lattice). On each stratum the restricted Fisher metric is nondegenerate; the blow-up of the ambient metric at the boundary is the logarithm's singularity, not a property of the geometry itself.
The same stratification applies to $\mathcal{O}$, now with an MDP reading for each stratum. With full-support $\mu$ the two-state two-action polytope is a quadrilateral, an affine slice of the tetrahedron $\Delta^3(\mathcal{S}\times\mathcal{A})$ cut out by the flow equations. Pull it apart below.
Same grammar, MDP meaning. The four vertices are the deterministic policies $(a_i,a_j)$, each with two-coordinate support rather than a single Dirac. Each edge pins one state’s action and frees the other: that is the occupancy polytope of a smaller MDP, so every face is itself an occupancy polytope carrying its own Fisher metric one dimension down. With $\mu$ concentrated rather than full-support the map from policies to vertices stops being injective: deterministic policies differing only at an unreachable state collapse to one vertex.
Concretely $\mathcal{O}$ is a planar cross-section of $\Delta^3(\mathcal{S}\times\mathcal{A})$. Each vertex (a deterministic policy) lands on an edge of the tetrahedron, the two-coordinate-support stratum for that policy; the interior of $\mathcal{O}$ lies in the interior of $\Delta^3$. The figure below pulls the slice apart, tethering each vertex to the simplex edge it came from.
$\mathcal{O}$ is the affine slice of the tetrahedron cut out by the flow equations. Its four vertices rest on four of the six edges of $\Delta^3$; the amber tethers mark which. The two untouched edges are the “same-state” supports $\{s_0a_0,s_0a_1\}$ and $\{s_1a_0,s_1a_1\}$: with full-support $\mu$ no occupancy can put all its mass on one state, so those strata stay empty. The slice’s own stratification is the trace of the simplex’s stratification on the cutting plane.
What this essay has not done is ask which metric on $\mathcal{O}$ is the right one. On a bare simplex the answer is essentially unique: Chentsov's theorem singles out the Fisher–Rao metric by invariance. On the occupancy polytope, because the Bellman equations couple states to actions, the symmetry group splits and uniqueness fails; the visitation drag term and the Kakade metric are independently invariant, leaving a genuine choice. That story, together with the cleaner algebraic treatment via the Fisher cometric and its relationship to reward shaping and value-function baselines, is the subject of the next post.