Trust Region Policy Optimization

In policy gradient methods, we parameterize a policy $\pi_\theta$ and update parameters $\theta$ to maximize expected cumulative reward. The update rule looks simple: $\theta \leftarrow \theta + \alpha \nabla_\theta J(\theta)$. But underneath that simplicity lies a dangerous asymmetry.

A slightly-too-large step can push the policy into a region where it takes terrible actions. Terrible actions produce terrible trajectories. Terrible trajectories produce misleading gradients. Misleading gradients push the policy even further into bad territory. This vicious cycle — the performance collapse spiral — is the central failure mode of vanilla policy gradients.

The core issue is that parameter space doesn't reflect policy space. A tiny change in $\theta$ might barely change the policy, or it might completely flip which actions get chosen. Standard gradient descent has no way to tell the difference — it only sees parameter distance, not behavioral distance.

TRPO's key insight starts with the performance difference lemma (Kakade & Langford, 2002): the difference in total reward between any two policies $\pi'$ and $\pi$ can be written exactly as the expected advantage of $\pi'$ over $\pi$, evaluated under $\pi'$'s own state distribution:

This identity is exact but circular — computing $d^{\pi'}$ requires running the new policy $\pi'$, which is what we're trying to design. Schulman et al. break the circularity by constructing a surrogate objective $L_\pi(\pi')$ that uses the old policy's state distribution instead. They then prove a lower bound: the true improvement $J(\pi') - J(\pi)$ is at least $L_\pi(\pi') - C \cdot D_{KL}^{max}(\pi, \pi')$, where $C$ is a penalty constant and $D_{KL}^{max}$ is the maximum KL divergence across states.

In words: if the surrogate goes up and the policies stay close in KL divergence, the true performance is guaranteed to go up too. This transforms a hard RL problem into a tractable optimization with constraints.

Since we can't sample from the new policy's state distribution (we haven't deployed it yet), TRPO replaces it with the old policy's distribution. The surrogate objective uses importance sampling to re-weight actions:

Think of it as asking: "If I replay the trajectories I already collected but pretend the new policy chose those actions, how much better would the advantage be?" The probability ratio $\frac{\pi_\theta(a|s)}{\pi_{\theta_{old}}(a|s)}$ corrects for the fact that the new policy would choose actions with different probabilities.

But this surrogate alone is dangerous. Without a constraint, the optimizer would push the probability ratios to extreme values — making the new policy wildly different from the old one. The importance sampling correction becomes unreliable, and the surrogate stops predicting true performance. This is exactly the performance-collapse problem dressed in new math.

TRPO's solution: constrain the KL divergence between old and new policies to stay within a small trust region.

The TRPO optimization problem is a constrained optimization over thousands or millions of policy parameters. How do we actually solve it?

Step 1: Approximate the KL constraint. Near the current parameters, the KL divergence is well-approximated by a quadratic form using the Fisher information matrix $F$: $\bar{D}_{KL}(\theta_{old}, \theta) \approx \frac{1}{2} (\theta - \theta_{old})^T F (\theta - \theta_{old})$. This turns the trust region from an abstract probability-distribution constraint into a concrete ellipsoid in parameter space.

Step 2: Solve for the update direction. With a linear approximation to the objective and the quadratic KL approximation, the optimal update direction is $F^{-1} g$, where $g$ is the policy gradient — this is exactly the natural gradient. It accounts for the geometry of the policy distribution, stepping more cautiously in directions where the policy is sensitive and more boldly where it's robust.

Step 3: Use conjugate gradient. Computing $F^{-1} g$ directly requires inverting a matrix with millions of entries — impractical. Instead, TRPO uses the conjugate gradient algorithm to solve $F x = g$ iteratively, needing only matrix-vector products $F v$ which are computed cheaply via automatic differentiation.

Step 4: Backtracking line search. The quadratic approximation to KL is only accurate locally, so TRPO performs a line search: it tries the full natural-gradient step, checks whether the KL constraint is truly satisfied and the surrogate actually improved, and if not, shrinks the step by a factor (typically 0.8) and tries again.

The full TRPO update cycle brings together all the pieces. At each iteration the algorithm collects trajectories, estimates advantages, solves a constrained optimization, and updates the policy — with a guarantee that performance never degrades.

Standard gradient descent constrains the step in parameter space: $\|\theta_{new} - \theta_{old}\|_2 \le \epsilon$. But parameter distance is a poor proxy for behavioral distance. Consider a softmax policy over two actions with logits $[0, 0]$: changing them to $[0.01, 0.01]$ barely changes anything, but changing to $[0, 100]$ makes the policy deterministic — both might have the same $L_2$ distance from $[0, 0]$.

KL divergence measures change in the distribution of actions, which is what actually matters. A step of size $\delta$ in KL space guarantees that the probability of every action changes by at most a bounded amount. The Fisher information matrix, which defines the KL geometry locally, automatically scales: it makes the step conservative in high-sensitivity directions and permissive in low-sensitivity ones.

TRPO describes two sampling schemes for estimating the surrogate objective:

Single path — the simpler approach. Run the current policy to collect full trajectories, then use all state-action pairs to estimate the objective and constraint. This is what most implementations use, and what you'd choose for environments that are expensive to reset.

Vine — a more sample-efficient scheme for cheap-to-reset environments. From each sampled state, branch out with multiple rollouts using different actions. This gives lower-variance estimates because you compare actions from the same state, reducing the confounding effect of different state visitation.

Schulman et al. tested TRPO on two challenging domains:

Continuous control (MuJoCo). Simulated robots learning to swim, hop, and walk. TRPO consistently outperformed vanilla policy gradients and the cross-entropy method, learning stable gaits with monotonic reward improvement. The key result: performance never collapsed during training, even with minimal hyperparameter tuning.

Atari games. TRPO also learned to play Atari games directly from pixel inputs, using a convolutional neural network as the policy. While not matching DQN's peak scores on all games, TRPO demonstrated that a single on-policy algorithm could handle both continuous and discrete action spaces — a significant generality advantage.

TRPO's theoretical elegance came at a computational cost: the conjugate gradient solve and line search make each update significantly more expensive than a standard gradient step. This motivated Schulman's follow-up work, Proximal Policy Optimization (PPO), which approximates TRPO's trust-region constraint with a simple clipped surrogate objective that can be optimized with first-order methods.

PPO became the dominant RL algorithm in practice — powering RLHF for ChatGPT, Claude, and other language models — but its success is built entirely on TRPO's conceptual framework. TRPO also directly inspired Generalized Advantage Estimation (GAE), which provides variance-reduced advantage estimates that are now standard in all policy-gradient methods.

TRPO's lasting contribution is not just an algorithm — it is the idea that policy optimization should be safe by construction. Every time a language model is fine-tuned with RLHF, the underlying machinery is enforcing trust regions in one form or another, ensuring that the model improves without forgetting how to be helpful.