Proximal Policy Optimization Algorithms, annotated

2.1 Policy gradient methods · original

“While it is appealing to perform multiple steps of optimization on this loss L^PG using the same trajectory, doing so is not well-justified, and empirically it often leads to destructively large policy updates.”Schulman et al. (2017), §2.1

Everyday picture

A basketball player practising free throws leans their form a little towards whatever they did before a shot that went in, and a little away from whatever they did before a miss. They never need to know the physics of why. That is a policy gradient: make actions that did better than typical more likely, and actions that did worse less likely.

Tiny example

For a softmax policy, the direction in logit space that makes action a more likely is “1 for a, minus every action's probability”. At the uniform start that is (−⅓, −⅓, ⅔) for C, (⅔, −⅓, −⅓) for A and (−⅓, ⅔, −⅓) for B. Scale each by its sample's advantage and average over the batch:

(1 × (−⅓, −⅓, ⅔) + (−1) × (⅔, −⅓, −⅓) + 0.5 × (−⅓, ⅔, −⅓)) / 3 = (−7/6, 1/3, 5/6) / 3 = (−0.389, 0.111, 0.278).

Push A down, C up, and B up a little. That is the step the batch asks for.

In words: “the estimated policy gradient is the batch average of the direction that makes each sampled action more likely, weighted by how much better than typical that action did. is simply a number whose gradient is exactly that, so a library that computes gradients automatically can produce ĝ from it.”

With the numbers: ĝ = (−0.389, 0.111, 0.278), as above. itself at the start is log(⅓) × (1 − 1 + 0.5) / 3 = −1.0986 × 0.1667 = −0.183. Its value means little; its slope is what gets used.

In Python:

import math
pi = [1 / 3, 1 / 3, 1 / 3]
# the batch: (action, advantage) for C, A, B; actions 0, 1, 2 are A, B, C
batch = [(2, 1.0), (0, -1.0), (1, 0.5)]
# ∇_θ log π_θ(a) for a softmax policy: 1[k = a] − π(k), for each logit k
def grad_log_pi(a):
    return [(k == a) - pi[k] for k in range(3)]
# ĝ = Ê_t[∇ log π(a_t) Â_t]: the batch average
g_hat = [sum(grad_log_pi(a)[k] * A for a, A in batch) / len(batch) for k in range(3)]
[round(g, 3) for g in g_hat]  # → [-0.389, 0.111, 0.278]
# L^PG = Ê_t[log π(a_t) Â_t]
round(sum(math.log(pi[a]) * A for a, A in batch) / len(batch), 3)  # → -0.183

The trouble, which the quote above names, is running many steps on with the same batch. Its gradient never switches off: as long as A's advantage is −1, every step pushes A's log-probability further down, even when A's probability is already 0.001. The figure in section 3 shows the result on our batch.

Why it matters today

This is REINFORCE with a baseline, the starting point of every method on this page. For a language model, the actions are tokens and a sample is a whole answer, but the formula is unchanged. The lesson builds it as grad_log_prob and reinforce_step.

2.2 Trust region methods · original

“This problem can efficiently be approximately solved using the conjugate gradient algorithm, after making a linear approximation to the objective and a quadratic approximation to the constraint.”Schulman et al. (2017), §2.2

Everyday picture

A hiker in fog can see the slope only a few metres around them. The map they draw from what they see is trustworthy inside that circle and fiction outside it. So they take the best step inside the circle, then look again. The circle is the trust region. TRPO measures the circle's size with the KL divergence between the old and new policy: how differently they behave, on average.

Tiny example

Suppose a step would move the policy from (⅓, ⅓, ⅓) to (0.2, 0.3, 0.5). The probability ratios, new over old, are 0.6 for A, 0.9 for B and 1.5 for C. Weight each sample's advantage by its ratio and average: (1.5 × 1 + 0.6 × (−1) + 0.9 × 0.5) / 3 = 1.35 / 3 = 0.45. That is the surrogate objective: an estimate, from the old batch, of how much better the new policy would do. The KL divergence between the two policies is 0.070. With a trust region of δ = 0.01 this step is too big, and TRPO would not take it.

In words: “TRPO: make the ratio-weighted advantage as large as possible, but keep the average KL divergence between the old and new policy below δ. The penalty version (5) instead subtracts β times that KL divergence and drops the hard limit.”

With the numbers: surrogate 0.45; KL 0.070 > δ = 0.01, so the constraint fails. With a penalty of β = 1 the step scores 0.45 − 0.070 = 0.380.

In Python:

import math
pi_old = [1 / 3, 1 / 3, 1 / 3]
pi_new = [0.2, 0.3, 0.5]
batch = [(2, 1.0), (0, -1.0), (1, 0.5)]
# π_θ(a_t) / π_θold(a_t), for A, B, C
[round(pi_new[k] / pi_old[k], 2) for k in range(3)]  # → [0.6, 0.9, 1.5]
surrogate = sum(pi_new[a] / pi_old[a] * A for a, A in batch) / len(batch)
round(surrogate, 2)  # → 0.45
# KL[π_old, π_new] = Σ_a π_old(a) log(π_old(a) / π_new(a)); one state, so Ê_t is this one number
kl = sum(p * math.log(p / q) for p, q in zip(pi_old, pi_new))
round(kl, 3)  # → 0.07
delta, beta = 0.01, 1.0
kl <= delta  # → False
# the penalty form, equation (5)
round(surrogate - beta * kl, 3)  # → 0.38

Why is a ratio allowed here at all? The batch was sampled from the old policy, but we want to judge the new one. Weighting each old sample by how much more (or less) likely the new policy makes it is importance sampling: it turns an average over the old policy's actions into an estimate of the average over the new policy's. The estimate is good only while the two policies are close, which is exactly what the constraint protects.

The paper explains why TRPO uses a hard limit rather than the penalty: no single β works well across problems, or even within one problem as learning changes its character. It then adds that simply choosing a fixed β and running SGD on (5) is not enough to match TRPO. Something else is needed.

Why it matters today

The trust-region idea (improve, but only where your estimate is trustworthy) is the ancestor of every “leash” in modern fine-tuning. The conjugate-gradient machinery that TRPO needed to enforce it is what PPO throws away.

The full objective (equation 9) · original

Everyday picture

A sports team has players and a commentator. The players (the policy) choose what to do; the commentator (the value network) keeps saying how well things are likely to go from here. The commentator's guess is the “typical” that advantages are measured against, so it has to be trained too. That split is called actor-critic. When players and commentator share one brain (one network), a single loss has to train both, and it helps to add a small bonus for keeping options open, so the team keeps trying new plays.

Tiny example

Take sample C with the candidate policy: its term is 1.2 (from the table in §3). Suppose the value network predicted 0.5 for this situation and the target it should have predicted is 0.81 (where 0.81 comes from is the next subsection): squared error (0.5 − 0.81)² = 0.0961. The candidate policy's entropy is 1.030. With the paper's Atari settings c₁ = 1 and c₂ = 0.01, this sample contributes 1.2 − 1 × 0.0961 + 0.01 × 1.030 = 1.114.

In words: “for each timestep, take PPO's clipped term, subtract a multiple of the value network's squared error, add a small multiple of the policy's entropy, and average; then climb this with gradient ascent.”

With the numbers: 1.2 − 1 × (0.5 − 0.81)² + 0.01 × 1.030 = 1.2 − 0.096 + 0.010 = 1.114.

In Python:

import math
L_clip_t = 1.2
V_theta, V_targ = 0.5, 0.81
c1, c2 = 1.0, 0.01
pi_new = [0.2, 0.3, 0.5]
# L^VF = (V_θ(s_t) − V_targ)²
L_vf = (V_theta - V_targ) ** 2
round(L_vf, 4)  # → 0.0961
# S = −Σ π log π: the entropy
S = -sum(p * math.log(p) for p in pi_new)
round(S, 3)  # → 1.03
round(L_clip_t - c1 * L_vf + c2 * S, 3)  # → 1.114

The minus sign on the value error is because the whole expression is maximized: climbing it means making the value error smaller. The entropy bonus rewards spreading probability out, the same exploration pressure that keeps the bandit from committing too early. In the robot experiments of §6.1 the paper uses separate networks for the policy and the value function (so c₁ does not matter) and no entropy bonus.

Advantage estimates (equations 10 to 12) · original

Everyday picture

How good was a chess move? You could wait for the end of the game, but the result depends on dozens of later moves, so it is a noisy verdict. You could ask a commentator how the position looks one move later, but the commentator may be wrong. Generalized advantage estimation blends the two: look ahead a few real steps, then trust the commentator, with a dial λ that sets how far you look before trusting.

Tiny example

A three-step episode earns rewards 0, 0, 1 (the only reward comes at the end). The value network's guesses for the three situations are 0.5, 0.6 and 0.8, and 0 after the episode ends. With discount γ = 0.9, each step's surprise is δ_t = reward + γ × next guess − this guess:

• δ₀ = 0 + 0.9 × 0.6 − 0.5 = 0.04
• δ₁ = 0 + 0.9 × 0.8 − 0.6 = 0.12
• δ₂ = 1 + 0.9 × 0 − 0.8 = 0.20

With λ = 0.8, each advantage adds up the surprises from that step on, shrinking by γλ = 0.72 per step: Â₂ = 0.20, Â₁ = 0.12 + 0.72 × 0.20 = 0.264, Â₀ = 0.04 + 0.72 × 0.264 = 0.230. With λ = 1 the first advantage is 0.31: the true discounted reward 0.9² × 1 = 0.81, minus the guess 0.5. That 0.81 is the target the value network was chasing in the previous subsection.

In words: “(10): the advantage is the discounted rewards actually received until the end of the segment, plus the value network's guess for what comes after, minus its guess for where we stood. (11, 12): equivalently, add up each step's surprise (reward plus discounted next guess, minus this guess), shrinking by γλ per step; with λ = 1 this is exactly (10).”

With the numbers: δ = (0.04, 0.12, 0.20); with λ = 0.8, Â = (0.230, 0.264, 0.200); with λ = 1, Â₀ = −0.5 + 0 + 0.9 × 0 + 0.81 × 1 = 0.31, and the δ sum gives the same 0.04 + 0.9 × 0.12 + 0.81 × 0.20 = 0.31.

In Python:

r = [0.0, 0.0, 1.0]
# V(s_0), V(s_1), V(s_2), and V(s_T) = 0 after the episode ends
V = [0.5, 0.6, 0.8, 0.0]
gamma, T = 0.9, 3
# δ_t = r_t + γ V(s_(t+1)) − V(s_t)
delta = [r[t] + gamma * V[t + 1] - V[t] for t in range(T)]
[round(d, 2) for d in delta]  # → [0.04, 0.12, 0.2]
def gae(lam):
    # Â_t = Σ_l (γλ)^l δ_(t+l)
    return [sum((gamma * lam) ** l * delta[t + l] for l in range(T - t)) for t in range(T)]
[round(A, 3) for A in gae(0.8)]  # → [0.23, 0.264, 0.2]
# equation (10) at t = 0: −V(s_0) + r_0 + γ r_1 + γ² r_2 + γ³ V(s_3)
round(-V[0] + sum(gamma ** k * r[k] for k in range(T)) + gamma ** T * V[T], 3)  # → 0.31
# λ = 1 gives the same number
round(gae(1.0)[0], 3)  # → 0.31

A note on the printed equations: in (10) and (11) the paper writes the exponent of the last reward and the last δ as T − t + 1. Counting powers (the reward k steps after t carries γ^k, and the last reward, r_T−1, is T − 1 − t steps after t) gives T − t − 1, which is what is shown here and what the Python checks.

Try it: drag λ and γ and watch the three advantages. At λ = 0 each advantage is just its own step's surprise; at λ = 1 each one sees every real reward to the end.

Reading it: each row is one step of the episode. The first number is the surprise δ_t, which depends only on γ; the bar and the second number are the advantage Â_t (a striped bar means a negative advantage, which appears when γ is small enough that the next state's guess no longer covers this one's). The last step's advantage never changes with λ: nothing comes after it. The earliest step changes most, because it has the most future to blend in. Small λ leans on the value network (less noise, but its mistakes pass straight through); λ near 1 leans on real rewards (unbiased, but noisier over long episodes). The paper uses γ = 0.99 and λ = 0.95 in every experiment.

Algorithm 1: PPO, actor-critic style · original

Everyday picture

A coach sends N scouts out at once. Each plays T moves with the current playbook and brings back notes. The coach grades every move against the commentator's expectations, then studies the whole pile of notes K times over, in small handfuls, adjusting the playbook as they go. Then the new playbook goes out with the scouts and the cycle repeats.

Reading it: read from the top down, then follow the wire on the left back up. The stacked boxes at the top are N copies of the frozen old policy, each collecting T timesteps of experience in its own copy of the environment. The value network (the small box on the right) feeds the advantage estimates of the previous subsection. Everything lands in one batch of N·T timesteps, which also records each action's probability under the old policy, needed for the ratio. The wide box below the batch is the part that is new in PPO: K full passes over that batch, in minibatches of M timesteps, climbing the clipped objective. Finally the old policy is overwritten with the new one, and the loop starts again with fresh experience. The only expensive part, running the environment, happens once per iteration; the batch is then reused K times.

Why it matters today

Swap “environment” for “a language model writing answers to prompts” and “reward” for “a reward model's score”, and this diagram is the reinforcement learning stage of RLHF. The costly step there is generation, which is why reusing each batch for several epochs matters even more.

6.1 Comparison of surrogate objectives · original

Everyday picture

Before comparing PPO with other people's algorithms, the paper holds a tryout between its own variants: no clip and no penalty, the clip at three widths, and the KL penalty (fixed or adaptive) at several strengths. Each variant plays seven different sports and is graded on one common scale.

The setup

• 7 simulated robot tasks in OpenAI Gym on the MuJoCo physics engine (HalfCheetah, Hopper, InvertedDoublePendulum, InvertedPendulum, Reacher, Swimmer, Walker2d), one million timesteps of training each, 3 random seeds per task.
• The policy is a small network with two hidden layers of 64 units and tanh activations, outputting the mean of a Gaussian over actions. The policy and value function do not share parameters, and there is no entropy bonus.
• Each run is scored by its average total reward over the last 100 episodes, then shifted and scaled per task so that a random policy scores 0 and the best result scores 1, and averaged over the 21 runs.

Reading it: each bar is one variant's average normalized score; the solid vertical line is 0 (a random policy) and the dashed one is 1 (the best result on each task). The three clipped variants (filled in the accent colour) are all at or near the top, with ε = 0.2 best at 0.82. The KL-penalty variants (grey) do reasonably, 0.62 to 0.74, but depend more on their setting. The dashed outline pointing left of zero is the plain surrogate with neither clip nor penalty: reusing the batch without a brake made one task, HalfCheetah, worse than random, and dragged the average to −0.39. That single bar is the case for the whole paper.

The paper also reports trying clipping in log space, with no improvement. For completeness, the hyperparameters behind these runs (Table 3 of the paper) are: horizon T = 2048, Adam step size 3 × 10⁻⁴, 10 epochs, minibatch size 64, γ = 0.99, λ = 0.95.

Why it matters today

Robustness was the paper's goal, and this table is its evidence. Each setting uses one set of hyperparameters across all seven tasks. In the full table, ε = 0.2 beats every KL setting (adaptive ones score 0.68 to 0.74, fixed ones 0.62 to 0.72), and removing the brake altogether is a disaster. The lesson's code uses the same ε = 0.2.

6.2 Comparison with other algorithms, continuous control · original

Everyday picture

Now PPO with ε = 0.2 plays the same seven robot tasks against tuned versions of the established methods: TRPO, the cross-entropy method, vanilla policy gradient with an adaptive step size, A2C, and A2C with a trust region.

What the paper reports

The learning curves (the paper's Figure 3, one panel per task over one million timesteps) show PPO outperforming the previous methods on almost all of the seven tasks. The paper gives the curves only as plots, so this page does not redraw them: read them in the original.

Why it matters

This was the headline for practitioners: the simple method was also the strongest on the standard benchmark of the day.

6.3 Showcase: humanoid running and steering · original

Everyday picture: a stick-figure robot learns to run towards a flag, turn when the flag moves, and get up again when cubes knock it over. The paper trains three such 3D humanoid tasks in Roboschool (running forward; running to a flag that moves every 200 timesteps or when reached; and the same while being pelted by cubes), with 32 or 128 parallel actors. The learning curves (Figure 4) run for 50 to 100 million timesteps, and the still frames (Figure 5) show the robot turning towards a new target. The paper notes that concurrent work by Heess et al. used the adaptive KL variant of §4 to learn locomotion for 3D robots.

Why it matters: high-dimensional control with many parallel actors is exactly the setting where TRPO's second-order step was expensive; PPO handled it with plain minibatch gradient steps.

6.4 Comparison with other algorithms on Atari · original

Everyday picture

The last tryout is 49 Atari video games, against well-tuned A2C and ACER. A game is “won” by whichever algorithm has the higher score averaged over three runs.

What the paper reports

• Scored by average reward over the whole of training (which favours fast learning): PPO won 30 games, ACER 18, A2C 1, with no ties.
• Scored by average reward over the last 100 episodes (which favours final performance): ACER won 28, PPO 19, A2C 1, with 1 tie.

So on Atari PPO beats A2C clearly and is roughly level with ACER, and the paper points out (§1) that it gets there while being much simpler. The Atari settings (Table 5) use 8 actors, T = 128, 3 epochs, a clip range of 0.1 × α and an Adam step size of 2.5 × 10⁻⁴ × α, where α shrinks linearly from 1 to 0 over training, plus c₁ = 1 and c₂ = 0.01 in the objective of equation (9). Appendix B gives learning curves and scores against A2C for all 49 games.

Why it matters

One algorithm, with modest retuning, worked on both continuous robot control and pixel-based games. That generality, more than any single score, is why PPO became the default.