Ignoring the state distribution in TRPO is not an approximation

Every derivation of TRPO reaches the same step and handles it the same way. You want to know how much better a new policy is than the old one, the honest expression involves the new policy's state distribution, you cannot sample that without already having deployed the new policy — so you swap in the old one and write something like "assume the state distribution changes negligibly for small updates."

I took that at face value for a long time, and it bothered me, because the assumption is plainly false. Policies do not change state distributions negligibly. A small parameter change can route an agent down an entirely different branch of the environment. If the whole method rested on that being small, the method should not work as well as it does.

It does not rest on that. The swap is exact to first order, and the reason has nothing to do with the state distribution being stable — it is a property of the advantage function that makes the state-distribution term drop out no matter how violently it moves.

Where the expression comes from

The starting point is the performance difference lemma (Kakade & Langford, 2002). For any two policies,

J(\pi') - J(\pi) = \mathbb{E}_{\tau \sim \pi'}\left[\sum_{t \ge 0} \gamma^t A^{\pi}(s_t, a_t)\right]

Most write-ups quote this and move on, which is a shame, because the proof is three lines and it is the good kind of proof — the kind where the result stops being a fact you memorised.

Proof. Expand

A^\pi(s,a) = \mathbb{E}_{s'}[r + \gamma V^\pi(s') - V^\pi(s)]

along a trajectory drawn from

\pi'

\mathbb{E}_{\tau \sim \pi'}\Big[\sum_t \gamma^t\big(r_t + \gamma V^\pi(s_{t+1}) - V^\pi(s_t)\big)\Big]

The reward terms sum to

J(\pi')

. The value terms telescope — each

\gamma^{t+1}V^\pi(s_{t+1})

cancels against the next step's

-\gamma^{t+1}V^\pi(s_{t+1})

— leaving only

-\mathbb{E}[V^\pi(s_0)] = -J(\pi)

\square

Note the mismatch the lemma leaves behind, because it is the source of all the difficulty: the expectation is over $\pi'$ , while the advantage function belongs to $\pi$ . We have data from $\pi$ and need an expectation over a policy we have not run.

The swap

Rewrite the lemma over the discounted state visitation measure $\rho$ , then replace $\rho_{\pi'}$ with $\rho_{\pi_{\text{old}}}$ — the step in question — and use importance sampling on the action:

L_{\theta_{\text{old}}}(\theta) = J(\theta_{\text{old}}) + \mathbb{E}_{s \sim \rho_{\text{old}},\, a \sim \pi_{\text{old}}}\left[\frac{\pi_\theta(a \mid s)}{\pi_{\text{old}}(a \mid s)} A^{\pi_{\text{old}}}(s,a)\right]

This is estimable from old data, which is the point. The question is what it cost.

Why the cost is second order

Claim. $L(\theta_{\text{old}}) = J(\theta_{\text{old}})$ and $\nabla_\theta L \big|_{\theta_{\text{old}}} = \nabla_\theta J \big|_{\theta_{\text{old}}}$ .

Proof. Write

J(\theta) = J(\theta_{\text{old}}) + \sum_s \rho_{\pi_\theta}(s)\, h_\theta(s)

with

h_\theta(s) := \sum_a \pi_\theta(a\mid s) A^{\pi_{\text{old}}}(s,a)

. Differentiate by the product rule:

\nabla_\theta J\Big|_{\text{old}} = \underbrace{\sum_s \big(\nabla_\theta \rho_{\pi_\theta}(s)\big)\Big|_{\text{old}} h_{\text{old}}(s)}_{(\mathrm{I})} + \underbrace{\sum_s \rho_{\text{old}}(s)\, \nabla_\theta h_\theta(s)\Big|_{\text{old}}}_{(\mathrm{II})}

Now

h_{\text{old}}(s) = \mathbb{E}_{a \sim \pi_{\text{old}}}[A^{\pi_{\text{old}}}(s,a)] = 0

for every

s

, because the advantage function is zero-mean under its own policy. So

(\mathrm{I}) = 0

regardless of

\nabla_\theta \rho

, and

(\mathrm{II})

is exactly

\nabla_\theta L

\square

The derivative of the state distribution is not small. It is annihilated — multiplied by a quantity that is identically zero at $\theta_{\text{old}}$ .

Two things I find worth keeping from this. First, the error from the swap is $O(\|\theta - \theta_{\text{old}}\|^2)$ , not $O(\|\theta - \theta_{\text{old}}\|)$ , which is the actual licence to take a small step and trust the surrogate — the entire foundation of trust-region methods, resting on a cancellation rather than on an empirical claim about environments.

Second, the proof requires the advantage function specifically. Run it with $Q^{\pi_{\text{old}}}$ in place of $A^{\pi_{\text{old}}}$ and $h_{\text{old}}(s) = V^{\pi_{\text{old}}}(s) \ne 0$ , so term $(\mathrm{I})$ survives and the swap really is a rough approximation. The usual story for subtracting a baseline is variance reduction. This is a second, structural reason, and I had never seen the two connected.

The bound nobody uses

First-order agreement says the error is second order but not how big the constant is. The monotonic improvement theorem supplies one: with $\varepsilon = \max_{s,a}|A^{\pi_{\text{old}}}(s,a)|$ ,

J(\theta) \ge L_{\theta_{\text{old}}}(\theta) - C \cdot D_{\mathrm{KL}}^{\max}, \qquad C = \frac{2\varepsilon\gamma}{(1-\gamma)^2}

Maximise the right-hand side and $J$ cannot decrease. It is a genuine minorise-maximise guarantee, and it is unusable: at $\gamma = 0.99$ , $(1-\gamma)^{-2} = 10^4$ . The step it sanctions is far too small to train with.

So TRPO keeps the shape and throws away the theorem. The KL penalty becomes a hard constraint $\bar{D}_{\mathrm{KL}} \le \delta$ with $\delta$ a hyperparameter, and the max over states becomes a mean because you cannot estimate a max over a continuous state space. Both moves void the guarantee. This seems worth saying plainly: the theory picks the shape of the algorithm — constrain a divergence between policies, not a distance between parameters — and then practice discards the constants. It is a derivation that earns its keep by ruling out alternatives, not by being followed.

One asymmetry that is not arbitrary

The remaining piece: TRPO expands the objective to first order and the KL constraint to second. That looks inconsistent until you check what the terms are. At $\theta_{\text{old}}$ the KL is zero and, because the score function has zero mean, its gradient is zero too. So a first-order expansion of the constraint gives $0 \le \delta$ — no constraint at all. The second-order term is the first one that exists:

D_{\mathrm{KL}}\big(\pi_{\text{old}} \,\|\, \pi_{\theta_{\text{old}} + \Delta}\big) = \tfrac{1}{2}\Delta^\top F \Delta + O(\|\Delta\|^3)

where $F$ is the Fisher information matrix. The objective, by contrast, has a perfectly good non-zero gradient, and its Hessian would be a mess — it carries the advantage function as a weight, so it is not positive semi-definite and its estimate is far noisier than $F$ , which depends only on the policy. Hence the division of labour: the objective's first-order information picks the direction, the KL's second-order structure decides how far. Solving that QP gives the natural gradient $F^{-1}g$ , and PPO is what you get when you decide the whole apparatus is not worth it and clip the probability ratio instead — keeping the shape, dropping the guarantee, one more time.

What I am still unsure about

How much of PPO's behaviour the trust-region story explains at all. Clipping bounds a per-sample ratio; it does not bound KL, and across several epochs on one batch the divergence can go wherever it likes. That implementations need KL early-stopping on top suggests the clip is doing something other than what the lineage says it does.
Whether first-order agreement is the right thing to want. It licenses a small step. Modern RLHF takes large effective steps on enormous batches, where the second-order term is not obviously negligible.
What the zero-mean cancellation looks like under GRPO, where the advantage is a group-normalised reward rather than $Q - V$ . Group normalisation makes it zero-mean within a group by construction, which looks like it should preserve the argument — but "zero-mean over sampled completions" and "zero-mean under the policy at every state" are not the same statement, and I have not worked out whether the gap matters.

Working notes The full derivation — policy gradient theorem, performance difference lemma, the first-order result above, the three equivalent forms of the Fisher matrix, and the trust-region QP — is in this PDF (14 pages, with the KL expansion and the natural-gradient invariance checked numerically). Everything here is derived rather than cited; corrections welcome.