Trust Region Policy Optimization (TRPO)

Introduction

Part 1 ended with a problem. The policy gradient tells us which direction to move the model, but not how far. If the step is too small, learning crawls. If it is too large, the model starts writing bad answers, learns from those bad answers, and falls into a death spiral from which it does not recover. So before anything else, we need an answer to a simple question: how large can an update be before it becomes dangerous?

The obvious answer is to turn down the learning rate. But the learning rate only controls how much the weights move, and the weights are not what we care about. What matters is how much the model’s output changes: the probabilities it assigns to each next token. In a large network, the two are only loosely connected. Moving one weight by $0.1$ may barely change the output, while moving another by just $0.001$ may change which token the model picks. A learning rate that is safe for the sensitive weights is painfully slow for all the others, and one that is fast enough for the others is dangerous for the sensitive ones. There is no single right value.

In 2015, Trust Region Policy Optimization (TRPO) [1] took a different route: instead of limiting how far the weights move, it limits how far the output moves. Before each update, TRPO fixes the maximum amount by which the model’s next-token probabilities are allowed to change. Every new model within that limit is considered safe, and together they form the trust region: the neighborhood around the current model where we can trust our estimate of whether an update helps. Figure 1 shows the idea, with three steps in the same direction from the current parameters ${\theta }_{\text{old}}$, but of different sizes:

The small step is safe but barely helps. The large step overshoots the peak, and the reward collapses. The step that stays inside the trust region is the one we want, and the rest of this post is about how TRPO finds it.

Throughout this post, ${\theta }_{\text{old}}$ are the current parameters and ${\pi }_{{\theta }_{\text{old}}}$ is the model before the update, while $\theta$ and ${\pi }_{\theta }$ are the parameters and the model after it. To find the right step, TRPO has to answer three questions. First, how do we measure how much the output has changed? This is the job of the KL divergence. Second, how do we evaluate a new model without generating new answers for it? This is where importance sampling comes in. And third, how do we find the best update inside the trust region? This question takes us from the TRPO objective, through the Fisher information matrix and the natural gradient, all the way to conjugate gradient and line search.

Measuring “How Far”: KL Divergence

At every step, an LLM outputs a probability distribution over its vocabulary. So to measure how much the model has changed, we need to compare two distributions: the old model’s and the new model’s. The standard tool for this is the Kullback–Leibler (KL) divergence [2]:

$$D_{\text{KL}}(P\Vert Q)=\sum _{i}P(i)\log \frac{P(i)}{Q(i)}\text{\hspace{1em}(1)}$$

Here $P$ is the old model’s distribution, $Q$ is the new model’s, and $i$ runs over every token in the vocabulary. For each token, the ratio $\frac{P(i)}{Q(i)}$ compares its probability under the two models: if they agree, the ratio is $1$, $\log 1=0$, and the token adds nothing. Each term is weighted by $P(i)$, so tokens the old model picks often count more. The total is $0$ when the two distributions are identical and grows as they become more different. (The KL is not symmetric, $D_{\text{KL}}(P\Vert Q)\ne D_{\text{KL}}(Q\Vert P)$, which is why it is called a divergence rather than a distance.)

That was one prediction. An LLM makes a prediction at every token of every response, so TRPO averages the KL over all the states $s_t$ the old model visits:

$${\bar{D}}_{\text{KL}}({\theta }_{\text{old}},\theta )={\mathbb{E}}_{s_t\sim {\pi }_{{\theta }_{\text{old}}}}\left[D_{\text{KL}}\right({\pi }_{{\theta }_{\text{old}}}(\cdot \mid s_t)\Vert {\pi }_{\theta }(\cdot \mid s_t)\left)\right]\text{\hspace{1em}(2)}$$

TRPO requires this average KL to stay below a small number $\delta$, for example $0.01$. This $\delta$ is the size of the trust region: with $\delta =0.01$, Update A would be allowed and Update B would not.

We can now check whether a new model is inside the trust region. The next question is how to check whether it is actually better than the old one.

Reusing Old Samples: Importance Sampling

To know whether a new model ${\pi }_{\theta }$ is better, we need to estimate how good its answers are, for example, the average reward they would get. The direct way is to let ${\pi }_{\theta }$ generate answers and score them. But generation is the slowest part of RL for LLMs, and while searching for the best update we want to evaluate many candidate models. Instead, we want to reuse the answers that ${\pi }_{{\theta }_{\text{old}}}$ has already generated. Importance sampling lets us do exactly that.

In general, if we sample from a distribution $p$ but want the average under another distribution $q$, we weight each sample by $\frac{q(y)}{p(y)}$:

$${\mathbb{E}}_{y\sim q}[f(y)]={\mathbb{E}}_{y\sim p}\left[\frac{q(y)}{p(y)}f(y)\right]\text{\hspace{1em}(3)}$$

In our example, $p$ is the old model, $q$ is the new model, and $f$ is the reward.

For an LLM, the same weight is applied to every token $y_t$ of every response, given its state $s_t$. It is called the probability ratio:

$${\rho }_t(\theta )=\frac{{\pi }_{\theta }(y_t\mid s_t)}{{\pi }_{{\theta }_{\text{old}}}(y_t\mid s_t)}\text{\hspace{1em}(5)}$$

There is one more change. In the example, we weighted the reward of each answer, but TRPO weights the advantage $A_t$ of each token (Eq. (3) in Part 1) instead. As in the policy gradient, the reward only says how good the whole answer was, while the advantage says whether each token was better or worse than expected, which is what the model needs to learn which tokens to make more or less likely. Weighting each token’s advantage by its probability ratio gives the surrogate objective:

$$L(\theta )={\mathbb{E}}_{s_t,y_t\sim {\pi }_{{\theta }_{\text{old}}}}[{\rho }_t(\theta )A_t]\text{\hspace{1em}(6)}$$

“Surrogate” means substitute. The objective we really care about is the new model’s actual performance, but measuring it would require generating new answers. $L(\theta )$ stands in for it, using only the old model’s answers. To make $L(\theta )$ larger, the new model has to make tokens with a positive advantage more likely (${\rho }_t>1$) and tokens with a negative advantage less likely (${\rho }_t<1$), the same idea as the policy gradient. In fact, at $\theta ={\theta }_{\text{old}}$ every ratio is $1$, and the gradient of $L$ is exactly the policy gradient in Eq. (2) in Part 1.

There is a catch, though. The estimate is only reliable when the new model is close to the old one. If they are very different, some ratios become very large, a few tokens dominate the average, and the estimate becomes noisy. This gives a second reason for the trust region: besides keeping the policy from breaking, it keeps $L(\theta )$ accurate.

This probability ratio ${\rho }_t(\theta )$ is also the core of PPO.

We now have both pieces: $L(\theta )$ tells us how good a new model is, and the KL tells us how far it is from the old one.

The TRPO Objective

TRPO combines the two pieces into one optimization problem:

$$\underset{\theta }{\max }L(\theta )\text{subject to}{\bar{D}}_{\text{KL}}({\theta }_{\text{old}},\theta )\le \delta \text{\hspace{1em}(7)}$$

In plain words: find the new model that improves the surrogate objective the most, while keeping the change in its output within $\delta$.

Why does this make training stable? Let $J(\theta )$ be the model’s true performance (its expected reward), the same objective whose gradient we took in Eq. (2) in Part 1. The TRPO paper [1], building on earlier work by Kakade and Langford [3], proves that:

$$J(\theta )\ge J({\theta }_{\text{old}})+L(\theta )-C{D}_{\text{KL}}^{\max }({\theta }_{\text{old}},\theta )\text{\hspace{1em}(8)}$$

In words: the new model’s true performance is at least the old performance, plus the surrogate objective, minus a penalty that grows with the KL divergence. Here $D_{\text{KL}}^{\max }$ is like ${\bar{D}}_{\text{KL}}$ but takes the largest KL over all states instead of the average, and $C$ is a positive constant that depends on the size of the largest advantage. So if an update improves the surrogate and keeps the KL small, the true performance cannot decrease. This is the monotonic improvement guarantee, exactly what vanilla policy gradient lacked.

The formulation is clean, but how do we actually solve it? This is where the Fisher matrix comes in.

The Fisher Information Matrix

Eq. (7) is hard to solve directly, because $L$ and the KL depend on billions of parameters in complicated ways. But the trust region gives us a shortcut. Every update is small, so we only need to know how these two functions behave near the current parameters ${\theta }_{\text{old}}$. And near a single point, even complicated functions look like simple shapes.

Figure 2 shows the two shapes TRPO uses. Near ${\theta }_{\text{old}}$, the objective $L$ looks like a straight line, and the KL looks like a quadratic, a curve shaped like a bowl.

Replacing a function with a simple shape near a point is called a Taylor approximation. A straight line uses the function’s slope (its first derivative), so it is called a first-order approximation. A quadratic also uses its curvature (its second derivative), so it is called a second-order approximation. Let’s look at each shape, writing $\Delta \theta =\theta -{\theta }_{\text{old}}$ for the update step.

Approximating the Objective

We have just seen the full Taylor approximation with its second-order term, but for the objective we will not need it. For a small step, the second-order term is much smaller than the first-order term: it grows with $\Delta {\theta }^2$ instead of $\Delta \theta$, so a step of $0.01$ contributes only about $0.0001$. And the objective only has to tell us which direction improves the model. How far to go is decided by the KL, and that is where the second-order term and the Hessian will matter. So for $L$, we keep only the first two terms, the value and the slope:

$$L(\theta )\approx L({\theta }_{\text{old}})+\nabla L^{\top }\Delta \theta$$

Both terms are easy to fill in. At ${\theta }_{\text{old}}$, every probability ratio is $1$, so the value is $L({\theta }_{\text{old}})=\mathbb{E}[A_t]=0$: on average, the old model’s own tokens are neither better nor worse than expected. And the slope $\nabla L$ at ${\theta }_{\text{old}}$ is exactly the policy gradient, which we call $g$. So the objective becomes a straight line through zero:

$$L(\theta )\approx g^{\top }\Delta \theta =\sum _i{g}_i\Delta {\theta }_i\text{\hspace{1em}(9)}$$

For each parameter, this multiplies its slope $g_i$ by how far we move it, $\Delta {\theta }_i$, and adds everything up.

Approximating the KL Divergence

For the KL, we cannot stop at the first-order term. Applying the same Taylor approximation, this time with all three terms, gives:

$${\bar{D}}_{\text{KL}}({\theta }_{\text{old}},\theta )\approx {\bar{D}}_{\text{KL}}({\theta }_{\text{old}},{\theta }_{\text{old}})+\nabla {\bar{D}}_{\text{KL}}^{\top }\Delta \theta +\frac{1}{2}\Delta {\theta }^{\top }F\Delta \theta$$

where $F$ is the Hessian of the KL at ${\theta }_{\text{old}}$. This time, the first two terms both vanish. The value is the KL of the old model with itself, which is $0$. And since the KL is never negative, ${\theta }_{\text{old}}$ sits at the very bottom of a valley, where the ground is flat, so the slope is $0$ as well (right panel of Figure 2).

Since the first two terms are zero, we keep only the last one, the second-order term. It describes the shape of the valley, a quadratic:

$${\bar{D}}_{\text{KL}}({\theta }_{\text{old}},\theta )\approx \frac{1}{2}\Delta {\theta }^{\top }F\Delta \theta \text{\hspace{1em}(10)}$$

With both approximations, the hard problem in Eq. (7) becomes a much simpler one: go as far as possible along a straight line, while keeping the quadratic below $\delta$:

$$\underset{\Delta \theta }{\max }{g}^{\top }\Delta \theta \text{subject to}\frac{1}{2}\Delta {\theta }^{\top }F\Delta \theta \le \delta \text{\hspace{1em}(11)}$$

The matrix $F$ is an $N\times N$ matrix that describes how steep the walls of the valley are in every direction. This matrix has a name, the Fisher Information Matrix, and the whole trust region now depends on it: the steeper the walls in some direction, the smaller the step we can take that way. So before we use it, let’s see what $F$ tells us about the model.

What the Fisher Matrix Tells Us

Each entry of $F$ tells us how steep the KL’s walls are in some direction, which is the same as asking: if we move the parameters in this direction, how much does the model’s output change? A large value means the output is sensitive in that direction, and even a small step changes it a lot. A small value means the output barely reacts.

This is exactly what we were missing at the start of this post: a way to know which weights are sensitive and which are not, something a single learning rate cannot capture.

So the trust region, the set of steps whose KL stays below $\delta$, is not a circle in parameter space. It is an ellipse: narrow along sensitive directions, where only tiny steps are safe, and wide along insensitive ones, where large steps are fine.

Computing the Fisher Matrix

As a Hessian, $F$ would require second derivatives of the whole network, which are expensive to compute. Fortunately, it can be rewritten using first derivatives only (see the derivation below):

$$F={\mathbb{E}}_{s_t,y_t\sim {\pi }_{{\theta }_{\text{old}}}}[{\nabla }_{\theta }\log {\pi }_{\theta }(y_t\mid s_t){\nabla }_{\theta }\log {\pi }_{\theta }(y_t\mid s_t{)}^{\top }]\text{\hspace{1em}(12)}$$

The gradient inside, ${\nabla }_{\theta }\log {\pi }_{\theta }(y_t\mid s_t)$, is called the score function. It is the same gradient as in the policy gradient: the direction that makes token $y_t$ more likely. So $F$ is the average outer product of the score function with itself. In practice, it is estimated from $M$ sampled tokens:

$$F\approx \frac{1}{M}\sum _{m=1}^M{\nabla }_{\theta }\log {\pi }_{\theta }(y_m\mid s_m){\nabla }_{\theta }\log {\pi }_{\theta }(y_m\mid s_m{)}^{\top }\text{\hspace{1em}(13)}$$

Solving for the Update Step

We now have everything we need to update the model. The policy gradient $g$ tells us which direction improves the objective, and the Fisher matrix $F$ tells us how sensitive the model is in each direction. What is still missing is the actual update step $\Delta \theta$: how far to move each parameter. It comes from solving the simplified problem in Eq. (11): go as far as possible along the line $g^{\top }\Delta \theta$, while keeping the quadratic $\frac{1}{2}\Delta {\theta }^{\top }F\Delta \theta$ below $\delta$:

$$\underset{\Delta \theta }{\max }{g}^{\top }\Delta \theta \text{subject to}\frac{1}{2}\Delta {\theta }^{\top }F\Delta \theta \le \delta$$

Unlike Eq. (7), this problem can be solved exactly. We will find the answer in three parts: where the best step ends, which direction it points, and how long it is.

Where the best step ends. The objective $g^{\top }\Delta \theta$ is a straight line, so it keeps growing as long as we keep moving in a good direction. There is no reason to stop early, so the best step goes all the way to the edge of the trust region, where the constraint holds with equality:

$$\frac{1}{2}\Delta {\theta }^{\top }F\Delta \theta =\delta$$

In parameter space, this edge is an ellipse: the set of all steps whose KL is exactly $\delta$.

Which direction it points. The best step lies somewhere on this ellipse, but where? The standard tool for maximizing a function under an equality constraint like this one is Lagrange multipliers [4], and the idea behind them is geometric, as shown in Figure 3.

Pick any point on the edge, such as $A$, and split $g$ into two parts: one that points along the edge, and one that points straight out of it. If the part along the edge is not zero, we can slide a little along the edge in that direction. The KL stays exactly $\delta$, but $g^{\top }\Delta \theta$ increases, so $A$ cannot be the best point. Sliding stops helping only where the part of $g$ along the edge is zero, that is, where $g$ points straight out of the ellipse, like at $B$. This is the best step.

What does “straight out of the edge” mean mathematically? Let $c(\Delta \theta )=\frac{1}{2}\Delta {\theta }^{\top }F\Delta \theta$. Along the edge, $c$ stays equal to $\delta$, just like the height stays the same along a contour line on a map. The gradient $\nabla c$ points in the direction where $c$ increases fastest, and since $c$ does not change along the edge, $\nabla c$ must be perpendicular to it, just as uphill on a map is always perpendicular to the contour lines. So $\nabla c$ points straight out of the edge. For our quadratic, it is:

$$\nabla c=\nabla \left(\frac{1}{2}\Delta {\theta }^{\top }F\Delta \theta \right)=F\Delta \theta$$

which is the many-variable version of $\frac{d}{dx}\left(\frac{1}{2}f{x}^2\right)=fx$, using the symmetry of $F$. This is the purple arrow in Figure 3.

So at $B$, both $g$ and $F\Delta \theta$ point straight out of the edge. Two vectors pointing in the same direction differ only in length, so one is a scaled version of the other:

$$g=\lambda F\Delta \theta$$

The scaling factor $\lambda >0$ is called the Lagrange multiplier. Multiplying both sides by $F^{-1}$ gives:

$$\Delta \theta =\frac{1}{\lambda }{F}^{-1}g\text{\hspace{1em}(14)}$$

This fixes the direction of the step: it points along $F^{-1}g$. Its length is still unknown, because we do not know $\lambda$ yet.

How long it is. To find $\lambda$, we use the fact that the best step ends exactly on the edge. Starting from the edge condition and substituting the direction from Eq. (14):

$$\begin{array}{rl}\delta & =\frac{1}{2}\Delta {\theta }^{\top }F\Delta \theta \\ & =\frac{1}{2}\left(\frac{1}{\lambda }{F}^{-1}g{)}^{\top }F\right(\frac{1}{\lambda }{F}^{-1}g)\\ & =\frac{1}{2{\lambda }^2}{g}^{\top }{F}^{-1}F{F}^{-1}g\\ & =\frac{1}{2{\lambda }^2}{g}^{\top }{F}^{-1}g\end{array}$$

Here we used $(F^{-1}g{)}^{\top }=g^{\top }{F}^{-1}$, since $F^{-1}$ is symmetric, and $F^{-1}F=I$. Solving for $\lambda$:

$$\frac{1}{\lambda }=\sqrt{\frac{2\delta }{g^{\top }{F}^{-1}g}}$$

The TRPO update. Putting the direction and the length together gives the full update step:

$$\Delta \theta =\sqrt{\frac{2\delta }{g^{\top }{F}^{-1}g}}{F}^{-1}g\text{\hspace{1em}(15)}$$

The direction, $F^{-1}g$, is called the natural gradient [5], [6]. The ordinary gradient $g$ treats all parameters the same, while multiplying by $F^{-1}$ divides each direction by its sensitivity, so the step becomes smaller along sensitive directions and larger along insensitive ones. The length, $\sqrt{2\delta /g^{\top }{F}^{-1}g}$, stretches or shrinks the step so that it ends exactly on the edge of the trust region, where the KL equals $\delta$. Notice that there is no learning rate in this update. Its role is taken by $\delta$, which is easier to interpret: the maximum change in the model’s output allowed per update.

Computing the Update for Large Models

Eq. (15) gives the ideal update, but it cannot be computed as written for an LLM. The problem is $F^{-1}$. For a 7-billion-parameter model, $F$ has $(7\times 10^9{)}^2=4.9\times 10^{19}$ entries, about 98 exabytes of memory in 16-bit precision, millions of times more than a GPU has. It cannot be stored, let alone inverted. TRPO works around this with two techniques: conjugate gradient to compute the direction without ever building $F$, and line search to make sure the final step is actually safe.

Conjugate gradient. Look again at Eq. (15). We never need the matrix $F^{-1}$ itself, only the vector $x=F^{-1}g$. Finding this vector is the same as solving:

$$Fx=g\text{\hspace{1em}(16)}$$

The conjugate gradient (CG) algorithm [7] solves this kind of equation step by step (for a gentle introduction, see [8]). It starts from a guess and improves it, and the key point is that it never needs $F$ itself, only products $Fv$ of $F$ with some vector $v$. These Fisher-vector products can be computed with two backward passes through the network, using memory proportional to $N$ instead of $N^2$. In practice, about 10 CG iterations give a good enough approximation of $x\approx F^{-1}g$. And once we have $x$, the step size in Eq. (15) is easy to compute, because $g^{\top }{F}^{-1}g=g^{\top }x$.

Line search. The step from Eq. (15) relies on the linear and quadratic approximations, which are only accurate close to ${\theta }_{\text{old}}$. If the full step goes too far, the actual KL may end up larger than $\delta$, or the surrogate objective may not actually improve. So before accepting the step, TRPO checks it with the exact $L$ and KL, and shrinks it if needed:

$${\theta }_j={\theta }_{\text{old}}+{\beta }^j\Delta \theta ,j=0,1,2,\dots \text{\hspace{1em}(18)}$$

With $\beta =0.5$, this tries the full step, then half, then a quarter, and so on. TRPO accepts the first step where the actual KL is within $\delta$ and $L$ actually improves. If no step passes after about 10 tries, it skips the update. This is called a backtracking line search [4]. Each check only needs forward passes on the batch we already have, so no new text has to be generated.

The TRPO Algorithm, Step by Step

Putting all the pieces together, TRPO alternates between generating responses with the current model and computing one constrained update from them:

Lines 3 and 4 are the only ones that generate text or need rewards. Everything after them works on the same batch: the gradient, the conjugate gradient iterations, and the line search need only forward and backward passes through the network. This is also where the cost of TRPO comes from, as we see next.

Limitations of TRPO

TRPO solved the stability problem of vanilla policy gradient. But its second-order method comes at a high cost, especially for LLMs:

• It’s expensive. Each update needs about 10 conjugate gradient iterations, each with an extra pass through the network, plus more forward passes for the line search. For an LLM, where a single pass is already costly, this multiplies the training cost.
• It doesn’t fit the standard training loop. Deep learning models are trained with first-order optimizers like Adam, which take many small steps on mini-batches. TRPO instead computes one carefully constrained step per large batch, which does not fit this setup.
• It’s hard to implement. Fisher-vector products, conjugate gradient, and line search add a lot of complexity. TRPO is also incompatible with architectures that share parameters between the policy and the value model, or that include noise such as dropout [10].
• The guarantee is only approximate. The monotonic improvement theory assumes the exact objective, but the algorithm uses Taylor approximations and an averaged KL. The line search reduces this gap but does not remove it.

TRPO showed that limiting how much the policy changes in each update makes RL stable. The open question was whether the same stability could be achieved with ordinary first-order gradient descent.

This is the question PPO answered, and it is the topic of Part 3.

References

1. Schulman, J., Levine, S., Moritz, P., Jordan, M. I., & Abbeel, P. (2015). Trust Region Policy Optimization. ICML.
2. Kullback, S., & Leibler, R. A. (1951). On Information and Sufficiency. The Annals of Mathematical Statistics, 22(1), 79–86.
3. Kakade, S., & Langford, J. (2002). Approximately Optimal Approximate Reinforcement Learning. ICML.
4. Nocedal, J., & Wright, S. J. (2006). Numerical Optimization (2nd ed.). Springer.
5. Amari, S. (1998). Natural Gradient Works Efficiently in Learning. Neural Computation, 10(2), 251–276.
6. Kakade, S. M. (2001). A Natural Policy Gradient. NeurIPS.
7. Hestenes, M. R., & Stiefel, E. (1952). Methods of Conjugate Gradients for Solving Linear Systems. Journal of Research of the National Bureau of Standards, 49(6), 409–436.
8. Shewchuk, J. R. (1994). An Introduction to the Conjugate Gradient Method Without the Agonizing Pain. Carnegie Mellon University.
9. Pearlmutter, B. A. (1994). Fast Exact Multiplication by the Hessian. Neural Computation, 6(1), 147–160.
10. Schulman, J., Wolski, F., Dhariwal, P., Radford, A., & Klimov, O. (2017). Proximal Policy Optimization Algorithms. arXiv:1707.06347.