Every policy gradient method so far uses each batch of experience for exactly one gradient step and then throws it away. For a robot simulator that is wasteful. For a language model it is painful, because the experience is generated text, and generating a batch of long responses from a large model can take longer than the gradient step itself. The obvious fix is to take many gradient steps on each batch. The obvious fix fails, and understanding why it fails, and what to do instead, leads to the most widely used RL algorithm of the last decade, proximal policy optimization (PPO), the algorithm behind InstructGPT and the RLHF stage of §16.4.

This section is the payoff for the chapter. We will see why large policy steps break things in a way they never do in supervised learning, reuse old data correctly with importance sampling, derive the surrogate objective and the trust-region idea behind TRPO, and then derive PPO's clipped objective case by case. Then we write a complete PPO in NumPy, with GAE, a critic, an entropy bonus, minibatch epochs, and the diagnostics practitioners watch, and run the ablation that shows what the clip is for.

Why big policy steps break

In supervised learning, a bad gradient step is cheap. The next batch comes from the same fixed dataset, its gradient points back toward good parameters, and the damage is undone. In RL the next batch comes from the new policy. If a step makes the policy bad, the data it collects is bad (a robot that falls over immediately learns only about falling over), and the policy may never see the states it would need to recover. Errors compound through the data distribution.

Three facts make large steps dangerous.

• The gradient is local. The policy gradient tells you the slope of $J(\theta)$ at the current parameters and nothing about what happens a long way off.
• Distance in parameters is not distance in behavior. A change of 1 in the logits of a softmax that sits at $(0, 0)$ moves the probabilities from $(0.5, 0.5)$ to $(0.73, 0.27)$. The same change at logits $(6, 0)$ moves the probability of the second action from $0.0025$ to $0.0009$, which barely matters in absolute terms. Measured by the KL divergence from the old distribution to the new one, the first step moves the policy by 0.12 nats and the second by 0.0009, more than a hundred times less, for the same step in parameters. The geometry that matters is the distance between the action distributions, which a fixed learning rate in parameter space ignores.
• Reusing a batch overfits to it. Take many steps on one batch's advantage estimates and the policy chases their noise. Actions that happened to look good once get their probability driven toward 1, the policy becomes nearly deterministic, its softmax gradient vanishes, and the exploration needed to find out it was wrong is gone.

We will watch the third failure happen in the gridworld at the end of the section. First, the tools for doing better.

Reusing data with importance sampling

Suppose the data was collected by an old policy $\pi_{\theta_{\text{old}}}$ and we want to evaluate, or differentiate, an expectation under a new one. Importance sampling corrects for the mismatch. For any distributions $p$ and $q$ with $q(x) > 0$ wherever $p(x) > 0$,

In plain words, to average under $p$ using samples from $q$, reweight each sample by how much more (or less) likely $p$ makes it than $q$ did. The weight $p(x)/q(x)$ is the likelihood ratio of §0.5.

The estimator is unbiased, but its variance can be enormous. Its second moment is $\mathbb{E}_q[(p/q)^2f^2] = \mathbb{E}_p[(p/q)f^2]$, which blows up wherever $p$ puts mass that $q$ rarely visits. For two unit-variance Gaussians whose means differ by $\Delta$, the weights have second moment $\mathbb{E}_q[(p/q)^2] = e^{\Delta^2}$, so a shift of 3 standard deviations makes the weights' variance about $e^{9} - 1 \approx 8{,}100$. A common diagnostic is the effective sample size, $(\sum_iw_i)^2/\sum_iw_i^2$ for weights $w_i$, which says how many samples from $p$ the reweighted batch is worth.

At a shift of 3 the hundred thousand samples are worth about forty, and the estimate is 20% off. Importance sampling works only when the two distributions are close. For policies, this tells us what to reweight. A whole trajectory's probability ratio is a product of $T$ per-step ratios $\pi_\theta(a_t\mid s_t)/\pi_{\theta_{\text{old}}}(a_t\mid s_t)$ (the dynamics cancel, as in §15.3), and a product of many ratios that are each near 1 can still be wildly far from 1, so trajectory-level weights have variance that grows exponentially with the horizon. Policy optimization methods therefore reweight one action at a time and handle the change in the state distribution differently, by staying close.

import numpy as np

def importance_estimate(f, logp_target, logp_behavior):
    raise NotImplementedError

The surrogate objective

How much better is a new policy $\pi'$ than the old one $\pi$? There is an exact answer, the performance difference lemma of Kakade and Langford (2002),

where $A^\pi$ is the advantage function of the old policy and the trajectory is drawn from the new one. In plain words, the improvement is the total advantage, as judged by the old policy, of the actions the new policy takes in the states it visits.

Writing the expectation with the discounted visitation $\rho^{\pi'}$ of §15.3 gives $J(\pi') - J(\pi) = \sum_s\rho^{\pi'}(s)\sum_a\pi'(a\mid s)A^\pi(s,a)$. Everything here is computable from old data except $\rho^{\pi'}$, the states the new policy would visit. The surrogate pretends the state distribution has not changed, and uses importance sampling for the actions,

where $d^\pi = (1-\gamma)\rho^\pi$ is the visitation normalized into a probability distribution, which is what sampling states from the old policy's rollouts approximates. Two facts make it the right object. At $\pi' = \pi$ the ratio is 1 and $\sum_a\pi(a\mid s)A^\pi(s,a) = 0$, so $L_\pi(\pi) = J(\pi)$. And its gradient at $\theta = \theta_{\text{old}}$ is the policy gradient, since $\nabla_\theta\frac{\pi_\theta(a\mid s)}{\pi_{\text{old}}(a\mid s)}\big|_{\theta_{\text{old}}} = \nabla_\theta\log\pi_\theta(a\mid s)\big|_{\theta_{\text{old}}}$, which turns $\nabla L$ into the policy gradient theorem with advantages. So the surrogate matches $J$ to first order. Far from $\pi$, the ignored change in $\rho$ makes it wrong, by an amount that grows with how different the policies are.

Schulman et al. (2015) bounded that error. This is a theorem (TRPO, Theorem 1 and its corollary). For any two policies,

The right side equals $J(\pi)$ at $\pi' = \pi$ and lower-bounds $J$ everywhere, so maximizing it can never decrease $J$. That is a minorize-maximize scheme with a monotone improvement guarantee. Its practical lesson is the shape, not the constant. The surrogate can be trusted only within a region where the policy has not changed much, measured by KL divergence.

Here is the surrogate against the true objective in the gridworld, stepping from the uniform policy along the exact policy gradient by increasing amounts.

Near the old policy, at a maximum KL of a few hundredths, the surrogate and the true objective agree closely. As the step grows they drift apart, and the surrogate stops being a reliable guide (in this problem it happens to be pessimistic, in others it is optimistic, and with estimated advantages it is noisy as well). The TRPO bound is valid at every step and useless as a step-size rule. Its constant, over 200 here, makes the penalty swamp any improvement after the tiniest step, which is why TRPO itself replaces the penalty with a constraint whose size is chosen by hand.

TRPO, a trust region measured in KL

Trust region policy optimization (Schulman et al. 2015) maximizes the surrogate subject to a limit on how far the policy moves, measured by the average KL divergence over the states the old policy visits,

with $\delta_{\text{KL}} = 0.01$ in the paper's experiments. To solve it, approximate both sides around $\theta_{\text{old}}$. The objective is linear to first order, with slope $\mathbf{g}$, the policy gradient. The KL is zero with zero gradient at $\theta_{\text{old}}$ (it is minimized there), so to second order it is a quadratic form,

where $\mathbf{F}$ is the Fisher information matrix of the policy, the Hessian of the KL at its minimum. The approximate problem is to maximize $\mathbf{g}^\top\Delta$ subject to $\frac12\Delta^\top\mathbf{F}\Delta \le \delta_{\text{KL}}$. Setting the gradient of the Lagrangian $\mathbf{g}^\top\Delta - \mu(\frac12\Delta^\top\mathbf{F}\Delta - \delta_{\text{KL}})$ to zero gives $\Delta = \mu^{-1}\mathbf{F}^{-1}\mathbf{g}$, and making the constraint tight fixes the scale,

The direction $\mathbf{F}^{-1}\mathbf{g}$ is the natural policy gradient (Kakade 2001), the steepest ascent direction when distance is measured by KL between the policies rather than Euclidean distance between the parameters, so, for small steps, it does not depend on how the policy happens to be parameterized. TRPO computes it without ever forming $\mathbf{F}$, by running conjugate gradient with Fisher-vector products (each one a double backward pass), and then does a backtracking line search along $\Delta$ until the true sampled KL is within $\delta_{\text{KL}}$ and the surrogate actually improves.

TRPO works, and its trust region is enforced almost exactly, but the machinery is heavy. It is second-order, it needs a line search, and it does not combine easily with architectures that share parameters between the policy and the value function, or with dropout. PPO was designed to get most of the benefit with first-order optimization alone.

import numpy as np

def natural_step(z, adv, delta):
    raise NotImplementedError

The exercise has a striking conclusion. For a softmax policy, the natural gradient step adds a multiple of the advantage itself to the logits, $z_a \leftarrow z_a + \eta A(s, a)$, which is the policy $\pi_{\text{new}}(a\mid s)\propto\pi_{\text{old}}(a\mid s)\,e^{\eta A(s,a)}$. Natural policy gradient on tabular softmax policies is a soft version of policy iteration. We will meet that exponential tilt again as the closed-form solution of the KL-regularized objective in §15.7 and §17.1.

PPO, clipping the incentive to move

Schulman, Wolski, Dhariwal, Radford, and Klimov (2017) asked whether a first-order method could keep the policy inside a trust region without computing one. Write the per-sample probability ratio as

so that the surrogate (without the constant $J(\pi_{\text{old}})$) is $\mathbb{E}_t[\rho_t(\theta)\hat A_t]$, an average over the timesteps of the batch, with $\hat A_t$ from GAE (§15.5). (The PPO paper calls this ratio $r_t(\theta)$ and the objective $L$. We use $\rho_t$ to keep $r_t$ for rewards and $J$ for objectives we maximize.) Maximizing the plain surrogate over many epochs drives each $\rho_t$ as far as it will go in the direction of its advantage's sign, which is exactly the overfitting described above. PPO's clipped objective removes the incentive to do that,

with $\epsilon = 0.2$ as the paper's default, where $\operatorname{clip}(x, l, u) = \max(l, \min(x, u))$.

The formula is compact enough to be confusing, so take it one case at a time. For a single sample, what matters is the sign of the advantage and where the ratio is.

Positive advantage, $\hat A_t > 0$ (the action was better than expected, and we want to make it more likely). The clipped term is $\min(\rho_t, 1+\epsilon)\hat A_t$ for $\rho_t \ge 1 - \epsilon$, and since both terms are increasing in $\rho_t$, the minimum is $\min(\rho_t, 1+\epsilon)\hat A_t$. The objective rises with $\rho_t$ until $\rho_t = 1 + \epsilon$ and is flat beyond it. Once the action's probability has grown by 20%, this sample stops pushing.

Negative advantage, $\hat A_t < 0$ (the action was worse than expected, and we want to make it less likely). Multiplying by a negative number flips the order, and the objective is $\max(\rho_t, 1-\epsilon)\hat A_t$. It rises as $\rho_t$ falls, until $\rho_t = 1 - \epsilon$, and is flat below. Once the action's probability has shrunk by 20%, this sample stops pushing.

The ratio moved the wrong way. If $\hat A_t > 0$ but $\rho_t < 1 - \epsilon$ (other updates made a good action much less likely), the unclipped term $\rho_t\hat A_t$ is the smaller one, so it is selected, and its gradient pulls the ratio back up. Likewise for $\hat A_t < 0$ and $\rho_t > 1+\epsilon$. Clipping only ever removes the incentive to move further in the helpful direction. It never stops the objective from correcting a step that went the harmful way.

Since $\nabla_\theta\rho_t = \rho_t\nabla_\theta\log\pi_\theta(a_t\mid s_t)$, the full gradient of one sample is $\rho_t\hat A_t\nabla_\theta\log\pi_\theta(a_t\mid s_t)$ when the unclipped term is selected and zero otherwise. At $\theta = \theta_{\text{old}}$ every ratio is 1, so the first epoch's first step is exactly the vanilla policy gradient with advantages. The clip only matters on later steps over the same data.

Two more properties are worth stating precisely. The clipped objective is a pessimistic bound on the unclipped one. Each sample's term is the minimum of the unclipped term and something else, so $J^{\text{CLIP}} \le \mathbb{E}_t[\rho_t\hat A_t]$, with equality whenever every ratio is inside $[1-\epsilon, 1+\epsilon]$. And clipping is not a constraint. Nothing stops a ratio from leaving the interval, because other samples' gradients move shared parameters, and a sample in the flat region just stops resisting. Engstrom et al. (2020) measured that ratios in real PPO runs do routinely leave the clip range, and that the trust region PPO actually maintains depends heavily on the rest of the recipe.

import numpy as np

def ppo_clip_loss(logp_new, logp_old, adv, eps):
    raise NotImplementedError

The KL-penalty variant

The PPO paper also tried a softer alternative to the clip, a penalty on the sampled KL divergence, $\mathbb{E}_t[\rho_t\hat A_t] - \beta\,\mathbb{E}_t[D_{\mathrm{KL}}(\pi_{\theta_{\text{old}}}(\cdot\mid s_t)\Vert\pi_\theta(\cdot\mid s_t))]$, with $\beta$ adapted after each update toward a target KL $d_{\text{targ}}$. If the measured KL $d$ is below $d_{\text{targ}}/1.5$, halve $\beta$. If it is above $1.5\,d_{\text{targ}}$, double it. On their continuous-control benchmark the clipped objective with $\epsilon = 0.2$ scored best (an average normalized score of 0.82, against 0.74 for the best adaptive-KL setting and $-0.39$ for the unclipped surrogate with no penalty, which on one task did worse than the random initial policy). The KL penalty survives in a different role in RLHF, as a penalty toward a fixed reference model rather than toward the previous policy (§15.7), and adaptive controllers in the same spirit keep that KL near a target (Ziegler et al. 2019 use a proportional controller on $\beta$).

The rest of the objective, value loss and entropy

PPO trains a critic $V_\psi$ (a value network with parameters $\psi$) alongside the policy, on the value targets $\hat R_t = \hat A_t + V_{\text{old}}(s_t)$ from §15.5, with a squared error $\mathcal{L}^V = \mathbb{E}_t[(V_\psi(s_t) - \hat R_t)^2]$, and adds an entropy bonus $H[\pi_\theta(\cdot\mid s_t)] = -\sum_a\pi_\theta(a\mid s_t)\log\pi_\theta(a\mid s_t)$ that discourages the policy from becoming deterministic too early. The combined objective, maximized each iteration, is

with coefficients $c_v$ (the paper's $c_1$, often 0.5 or 1) and $c_H$ (the paper's $c_2$, 0.01 in its Atari experiments and 0 in MuJoCo). When the policy and the critic are separate networks the value term does not affect $\theta$ and the two can be trained independently. When they share a body, $c_v$ balances the two gradients.

The entropy's gradient with respect to softmax logits is worth deriving once, since you will implement it. From $H = -\sum_i\pi_i\log\pi_i$ and $\partial\pi_i/\partial z_j = \pi_i(\mathbb{1}[i=j] - \pi_j)$,

In plain words, the entropy bonus lowers the logits of actions that are more likely than the "typical" action ($\log\pi_j > -H$) and raises the others, pulling the policy toward uniform.

Many implementations also clip the value function, mirroring the policy clip, and use the larger of the clipped and unclipped squared errors. Huang et al. (2022) document it as one of PPO's standard details. Andrychowicz et al. (2021), in a study of on-policy RL with over 250,000 agents trained, found that value clipping did not help and recommended against it. It is a detail you should be able to implement and to question.

import numpy as np

def ppo_value_loss(v_new, v_old, returns, clip=None):
    raise NotImplementedError

The algorithm

Here is PPO in the actor-critic form of the paper's Algorithm 1, with the details almost every implementation adds.

1. Run the current policy in $N$ parallel environments for $T$ steps each. Store states, actions, the log-probabilities $\log\pi_{\theta_{\text{old}}}(a_t\mid s_t)$, rewards, values, next-state values, and termination and truncation flags.
2. Compute GAE advantages and value targets once, with the critic as it was during collection (§15.5).
3. For $K$ epochs, shuffle the $NT$ samples into minibatches. For each minibatch, normalize its advantages, compute the ratios against the stored old log-probabilities, and take a gradient step on the combined objective (in practice with Adam and a global gradient-norm clip, typically at 0.5).
4. Monitor the approximate KL between the old and new policy on the batch (the $k_3$ estimator $(\rho - 1) - \log\rho$ of §0.5, averaged) and the clip fraction, the share of samples with $|\rho - 1| > \epsilon$. Many implementations stop the epochs early if the KL exceeds a target.
5. Set $\theta_{\text{old}}\leftarrow\theta$ and repeat.

The paper's MuJoCo settings were $T = 2048$, 10 epochs, minibatches of 64, Adam with step size $3\times10^{-4}$, $\gamma = 0.99$, and $\lambda = 0.95$. Its Atari settings used 8 actors, $T = 128$, 3 epochs, $\epsilon = 0.1$ annealed to zero, $c_v = 1$, and $c_H = 0.01$.

A complete PPO in NumPy

Now we build it, for the gridworld, with a tabular softmax policy (one logit per state and action) and a tabular critic, so that every gradient can be written by hand and we can compare the result with $V^\star$ from §15.2. The structure is exactly that of a deep RL implementation. Only the networks are missing. The gridworld is the one defined above. First the reference answer, from value iteration, which we use only to grade the result (PPO itself never sees $P$).

Next, collection. Sixteen copies of the gridworld run in lockstep, episodes are cut off after 50 steps (a truncation, not a termination), and the environments persist from one iteration to the next, so a rollout of 32 steps usually ends in the middle of episodes, which the last-step truncation flag handles. GAE is the function from Exercise 15.5b.

The update is the heart of it. Advantages and targets are computed once. Then, for each epoch and minibatch, the advantages are normalized, the ratios are computed against the stored log-probabilities, the clipped surrogate's gradient flows into the logits through the score $\mathbf{e}_a - \boldsymbol\pi$, the entropy bonus adds $-\boldsymbol\pi\odot(\log\boldsymbol\pi + H)$, and the critic takes a regression step toward its targets. Each line maps to an equation above.

And the outer loop, which alternates collection and update and prints the diagnostics a practitioner would watch.

In 40 iterations, about 20,000 transitions, PPO climbs from $-0.41$ (the random walker) to about 0.34, within 0.01 of the optimal 0.351. Its greedy action is optimal in every cell, and its critic has learned the value of its own policy. The policy is still slightly random, because the entropy bonus keeps it from collapsing completely, which costs a little return. Watch the diagnostics. Early on the approximate KL per update is about 0.01 and over 20% of the samples end up clipped, because the advantages are large and the policy wants to move a lot. As the policy converges, both fall toward zero. A KL that grows instead, or a clip fraction near 1, is the usual sign of a learning rate that is too high.

What the clip is for

Now the ablation. Make the updates aggressive, with a smaller batch (8 environments, 16 steps), 30 epochs over each batch in one big minibatch, and a large learning rate, and train with and without the clip.

Without the clip, three of the five runs end far below the optimum, at $-0.23$, $-0.01$, and $-0.32$, two of them closer to the random walker's $-0.41$ than to $V^\star$. Their policies are nearly deterministic, with entropy close to zero, and they have locked onto wrong actions, for instance walking into the edge of the grid forever from the top-left cell. Thirty epochs on a small batch let a few lucky advantage estimates drive some ratios to enormous values, the softmax saturated, and once the policy stopped exploring, no new data could correct it. With the clip, each sample stops pushing once its ratio leaves $[0.8, 1.2]$, the same aggressive settings are safe, and every run ends between 0.31 and 0.34. This is the whole argument for PPO, in one table. It lets you reuse a batch for many epochs, which is what makes it sample-efficient, without letting the reuse destroy the policy.

The implementation details that matter

PPO's paper is short, and its reference implementation contains many choices the paper barely mentions. Engstrom et al. (2020) showed that these "code-level optimizations" account for most of PPO's advantage over TRPO in their experiments, and Huang et al. (2022) catalogued 37 of them, 13 of which they call core. The ones you should know, and why each exists, are these.

The large study of Andrychowicz et al. (2021) gives complementary advice from their experiments, among it to start with a clipping threshold of 0.25, to initialize the last policy layer with 100 times smaller weights, to always normalize observations, to use GAE with $\lambda = 0.9$, and to treat the discount factor as one of the most important hyperparameters. RLHF adds its own layer of details on top of these (reward whitening, KL control, separate value models, and more), which §16.6 covers.

In PyTorch

In PyTorch the PPO loss is a few lines, and autograd handles the clip's piecewise gradient automatically. Here we compute it for a batch from a small policy network, check its gradient with respect to the log-probabilities against the analytic one, and take one optimizer step with the usual gradient clipping.

Read more

• Schulman, Wolski, Dhariwal, Radford & Klimov (2017). Proximal Policy Optimization Algorithms. The clipped objective, the adaptive KL variant, and the comparison table quoted above.
• Schulman, Levine, Moritz, Jordan & Abbeel (2015). Trust Region Policy Optimization. The surrogate bound, the KL trust region, and the conjugate-gradient implementation.
• Kakade & Langford (2002). Approximately Optimal Approximate Reinforcement Learning. The performance difference lemma and conservative policy iteration.
• Kakade (2001). A Natural Policy Gradient. The natural gradient applied to policies, and its connection to policy iteration.
• Huang et al. (2022). The 37 Implementation Details of Proximal Policy Optimization. Every detail of the reference implementation, explained and reproduced.
• Engstrom et al. (2020). Implementation Matters in Deep Policy Gradients: A Case Study on PPO and TRPO. How much of PPO's performance comes from code-level choices.
• Andrychowicz et al. (2021). What Matters In On-Policy Reinforcement Learning? A Large-Scale Empirical Study. Recommendations from training over 250,000 agents.

Watch next

▶

▶

▶