Trust Region Policy Optimization (TRPO)

TRPO improves a policy only inside a small KL divergence budget around the current behavior, so each update is guaranteed safe instead of risking a collapse.

The plain gradient step jumps outside the KL trust region while the TRPO step stops on its edge, turning a policy-breaking leap into a safe gain of 0.39

Why Does This Exist?

Vanilla policy gradients step in parameter space: $\theta_{new} = \theta + \alpha g$ , where $g$ is the gradient of expected return. The trouble is that a small move in parameters can mean a huge move in behavior. One oversized update can push the action probabilities off a cliff, and because the next batch of data comes from the ruined policy, training never recovers.

Schulman and colleagues (2015) fixed this with a constraint instead of a prayer. TRPO maximizes the usual surrogate advantage objective but only inside a trust region: the mean KL divergence between the old and new policy must stay under a budget $\delta$ , typically 0.01. That single bound is what buys monotonic improvement.

The one idea you need first

A policy gradient (REINFORCE) climbs expected return by weighting log-probability gradients with observed returns. TRPO keeps that direction but rescales and clips the stride using curvature, which is why this page assumes you know the REINFORCE update and focuses on the constraint machinery. The broader family, including the PPO simplification that replaced TRPO in most codebases, lives in trust region and proximal methods.

Think of It Like This

A dog on a fixed-length leash

You are walking a strong dog toward the park gate, which is uphill in return. Your stride toward the gate is the gradient $g$ : the direction that looks best right now. The dog is the policy: it moves where you pull, but it has its own momentum and can bolt.

The leash is the trust region $\delta$ . It lets the dog roam and sniff, but the collar stops it at a fixed radius no matter how hard you stride. A smart walker does not take baby steps just in case; they take the biggest stride the leash allows, then check the collar. When the leash is taut, you have moved as far as safety permits, exactly what the $\beta$ scale in the worked example computes.

Where the analogy stops: a leash measures plain Euclidean distance, the same in every direction. TRPO measures KL divergence between probability distributions, so the safe radius bends with the Fisher curvature: short where the policy is twitchy, long where it is numb.

How It Actually Works

TRPO solves a constrained problem. With $\rho$ the old state distribution, $r(\theta) = \pi_\theta(a \mid s) / \pi_{old}(a \mid s)$ the probability ratio, and $A_{old}$ the old advantage:

\max_\theta \mathbb{E}[r(\theta) A_{old}] \quad \text{subject to} \quad \overline{D}_{KL}(\pi_{old} \parallel \pi_\theta) \le \delta

Here $\overline{D}_{KL}$ is the mean KL divergence over states, and $\delta$ (often 0.01) is the budget. Directly solving this is too costly, so TRPO approximates.

The natural gradient step

Linearize the objective around $\theta_{old}$ with gradient $g$ , and quadratize the KL with the Fisher matrix $F$ (its second derivative at $\theta_{old}$ , where the first derivative is zero):

\max_{\Delta\theta} g^\top \Delta\theta \quad \text{subject to} \quad \tfrac{1}{2}\Delta\theta^\top F \Delta\theta \le \delta

The solution points along the natural gradient $x = F^{-1} g$ with a scale that lands exactly on the boundary:

\beta = \sqrt{\frac{2\delta}{g^\top F^{-1} g}}, \qquad \Delta\theta = \beta F^{-1} g

Forming $F$ for a network with $10^5$ weights is hopeless, so TRPO solves $Fx = g$ with conjugate gradients, which needs only Fisher-vector products $Fv$ from two backward passes. A backtracking line search then halves the step until the true (non-approximated) KL and surrogate both pass.

Worked example

Two parameters, diagonal Fisher $F = \mathrm{diag}(2.0, 4.0)$ , gradient $g = [1.0, 2.0]$ , budget $\delta = 0.05$ . Note $\theta_2$ has double the curvature, so it is double the twitchy.

Natural direction: $x = F^{-1} g = [1.0/2.0, 2.0/4.0] = [0.5000, 0.5000]$ . The raw gradient pushed twice as hard on $\theta_2$ ; curvature equalized it.

Curvature form: $g^\top x = 1.0 \times 0.5 + 2.0 \times 0.5 = 1.5000$ .

Scale: $\beta = \sqrt{2 \times 0.05 / 1.5000} = \sqrt{1/15} \approx 0.2582$ .

Step: $\Delta\theta = 0.2582 \times [0.5, 0.5] = [0.1291, 0.1291]$ .

Check the KL: $\tfrac{1}{2}(2.0 \times 0.1291^2 + 4.0 \times 0.1291^2) = \tfrac{1}{2} \times 6.0 \times 0.016667 \approx 0.0500 = \delta$ . Right on the boundary. Expected gain: $g^\top \Delta\theta = 0.1291 + 0.2582 = 0.3873 > 0$ , so line search accepts immediately.

Code

This snippet computes the analytic TRPO step for a diagonal Fisher, the same arithmetic as the worked example. Real code uses conjugate gradients instead of dividing, but the numbers are identical here.

1import numpy as np2
3
4def trpo_step(grad: np.ndarray, fisher_diag: np.ndarray, delta: float = 0.05) -> tuple[np.ndarray, float, float]:5    """Return (step, beta, quadratic KL) for a diagonal Fisher matrix."""6    nat = grad / fisher_diag  # natural gradient direction x = F^-1 g7    denom = float(grad @ nat)  # curvature form g^T F^-1 g8    beta = float(np.sqrt(2.0 * delta / denom))9    step = beta * nat10    kl = 0.5 * float(step @ (fisher_diag * step))11    return step, beta, kl12
13
14g = np.array([1.0, 2.0])15step, beta, kl = trpo_step(g, np.array([2.0, 4.0]))16print(np.round(step, 4))  # -> [0.1291 0.1291]17print(round(beta, 4))  # -> 0.258218print(round(kl, 4))  # -> 0.0519print(round(float(g @ step), 4))  # -> 0.3873

The step lands exactly on the KL budget of 0.05 while gaining 0.3873 in surrogate value, which is why no backtracking is needed on this instance.

Watch Out For

Treating the budget like a learning rate

$\delta$ looks like a step size knob, so people crank it to 0.2 or 0.5 to speed up slow runs. That reintroduces the exact collapse TRPO was built to prevent: the quadratic KL approximation only holds near $\theta_{old}$ , and a big $\delta$ accepts steps where the true KL is far larger than the estimate. Symptom: return climbs for a few iterations, then falls off a cliff in one update and flatlines.

Start at $\delta = 0.01$ and log the true mean KL every iteration. If the line search rejects most steps, your curvature estimate (or advantage scale) is off; shrink $\delta$ , do not grow it.

Running conjugate gradients without damping

The Fisher is only guaranteed positive semidefinite, and with sampling noise the product $Fv$ can have zero or negative curvature along some direction. Undamped CG then divides by near zero and returns a garbage direction or NaNs, which the line search silently rejects every time. Symptom: the policy stops updating entirely while losses look fine.

Add damping $(F + \lambda I)x = g$ with $\lambda \approx 10^{-3}$ to $10^{-2}$ . If every line search still rejects, raise $\lambda$ before touching anything else.

The Quick Version

TRPO maximizes surrogate advantage subject to mean KL divergence staying under a budget $\delta$ , usually 0.01.
The update follows the natural gradient $F^{-1} g$ , which rescales each parameter by how much it actually moves behavior.
Conjugate gradients solve for the direction using only Fisher-vector products, never forming the full matrix.
The step scale $\beta$ lands exactly on the trust region boundary, and a line search verifies the true KL before committing.
Damping the Fisher solve and keeping $\delta$ small are what separate a stable run from a silent stall or a collapse.