PPO: on-policy policy gradients that climb past DQN's ceiling

Code: rl_ablations/models/ppo, built on the rl-factory engine.

This tutorial explains how our PPO agent actually learns — the RolloutBuffer, the advantage calculation, and the loss function — starting from the very first policy-gradient idea (REINFORCE, 1992) and building up historically through A2C to PPO. It assumes you know basic calculus and a little probability, and nothing else. Every term is defined the first time it appears. Wherever we describe code, it lives in rl_factory/training/rollout.py (the buffer

• advantage math) and rl_ablations/models/ppo/learner.py (the loss).

If you only remember one sentence: we nudge the network to make actions that turned out better-than-expected more likely, and we do it carefully enough that we can squeeze several gradient steps out of each batch of experience without the training blowing up.

Part 0 — The problem we’re solving

An agent interacts with an environment (for us, Super Mario Bros). At each timestep it sees a state s (four stacked game frames), picks an action a (press a button), and receives a reward r (roughly, how far right Mario moved). This repeats until Mario dies or finishes the level — one episode.

The agent’s behaviour is a policy, written π(a | s) (“pi of a given s”): a function that, given a state, outputs a probability for each possible action. Our policy is a neural network. “Learning” means adjusting the network’s weights θ (theta) so that the policy collects more total reward.

We measure “total reward” with the return. The return from timestep t, written G_t, is the sum of all rewards from t onward, with later rewards discounted by a factor γ (gamma, between 0 and 1, e.g. 0.99):

G_t = r_t + γ·r_{t+1} + γ²·r_{t+2} + γ³·r_{t+3} + ...

Discounting says “a reward now is worth a bit more than the same reward later.” It also keeps the sum finite for long episodes. The goal of the whole enterprise:

Adjust θ to maximize the expected return — the average G we’d get if we let the policy play many times.

Everything below is a story about how to compute a good gradient — a direction to push θ — for that goal.

Part 1 — REINFORCE (1992): the original idea

1.1 The intuition

Suppose Mario plays a whole episode and gets a great return. Whatever actions he took, they probably weren’t all brilliant — but on average they were good, so let’s make all of them a little more likely next time. If the episode went badly, make those actions less likely. Scale the nudge by how good the outcome was.

That’s REINFORCE in one sentence. Now the math that makes it precise.

1.2 The policy-gradient formula

We want to increase expected return J(θ). The famous policy-gradient theorem says its gradient can be written as an average over experience the policy itself generated:

∇J(θ)  =  E[  Σ_t  ∇ log π(a_t | s_t) · G_t  ]

Let’s decode every piece:

• E[ … ] — “the average over many episodes played by the current policy.”
• Σ_t — sum over every timestep in the episode.
• π(a_t | s_t) — the probability the policy assigned to the action it actually took.
• log π(a_t | s_t) — the natural log of that probability. (Logs show up because of a calculus identity, ∇π = π · ∇log π; they also turn the tiny products of probabilities into stable sums. You don’t need to derive this — just know ∇log π(a_t|s_t) is “the direction in weight space that makes action a_t more likely in state s_t.”)
• G_t — the return that followed. This is the crucial weight: a number, not something we differentiate through.

Read the formula in plain English: for each action taken, step in the direction that raises its probability, scaled by the return that followed. Good returns → push those actions up. Bad returns → the scaling is small (or, once we add a baseline in Part 2, negative) → push down.

1.3 From gradient to “loss”

Deep-learning libraries minimize a loss by gradient descent. We want to maximize J, so we minimize its negative. The quantity whose gradient equals the formula above is:

Loss = − Σ_t  log π(a_t | s_t) · G_t

We compute this number, call .backward(), and the optimizer does the rest. Note G_t is treated as a fixed constant (we detach it from the graph): we are not trying to change the rewards, only the probabilities.

1.4 Why REINFORCE isn’t enough

Two problems, and the rest of history is about fixing them:

1. It’s extremely noisy (high variance). G_t depends on everything that happened after t — every later action, every bit of randomness in the game. The same good action can be attached to a huge return in one episode and a tiny one in the next, purely by luck. Training on such a jittery signal is slow and unstable.

2. It needs complete episodes and fresh data. You can’t compute G_t until the episode ends, and because the average is “over episodes the current policy generated,” the moment you update θ the data you collected is stale and must be thrown away. This is what on-policy means: you may only learn from data your current policy produced. (Contrast DQN, which is off-policy and reuses a giant replay buffer.)

Part 2 — Baselines and the birth of the critic

2.1 Subtracting a baseline

Here’s a small idea with a big payoff. Take the policy gradient and subtract a baseline b(s_t) — any number that depends on the state but not on the action:

∇J(θ)  =  E[  Σ_t  ∇ log π(a_t | s_t) · ( G_t − b(s_t) )  ]

Mathematically this changes nothing about the average (the baseline term averages out to zero, because probabilities sum to 1). But it can dramatically reduce the variance — the jitter — if we choose b well. Intuitively: instead of asking “was the return big?”, we ask “was the return bigger than what we normally get from this state?” That comparison is far more stable than the raw return.

2.2 The best baseline is the value function

The natural choice of baseline is the value function V(s): the expected return if you start in state s and follow the current policy. “From this spot, how well do we usually do?”

We don’t know V, so we learn it too, with a second neural network (or, in our code, a second head on the same network — see ActorCriticNet in rl_ablations/models/ppo/model.py). This second learner is called the critic; the policy is the actor. The critic doesn’t choose actions; it just predicts how good states are, to serve as the baseline.

Now G_t − V(s_t) has a name and a meaning: the advantage.

Advantage A(s_t, a_t)  ≈  G_t − V(s_t)
                       =  (what actually happened)  −  (what we expected)

• A > 0: this action did better than average from this state → make it more likely.
• A < 0: it did worse than average → make it less likely.

The advantage is the honest “was this action a good idea?” signal. From here on, the policy gradient is:

Loss_policy = − Σ_t  log π(a_t | s_t) · A_t

Same shape as REINFORCE, but with the advantage in place of the raw return. This is the heart of every method that follows.

Part 3 — A2C (Advantage Actor-Critic): the practical loop

A2C is the algorithm that turns the above into something you can actually run efficiently. It contributes three practical ideas, and all three are visible in our RolloutBuffer and OnPolicyStrategy.

3.1 Bootstrapping: don’t wait for the episode to end

Waiting for a full episode to compute G_t is slow and noisy. Instead, use the critic to estimate the tail. The simplest advantage estimate uses just one real reward plus the critic’s guess for the rest:

δ_t  =  r_t  +  γ·V(s_{t+1})  −  V(s_t)

This δ_t (delta) is the TD error (“temporal-difference error”). Read it as: “the reward I actually got, plus what the critic thinks the next state is worth, minus what the critic thought this state was worth.” If that’s positive, things went better than the critic expected. Using a prediction of the future to improve an estimate of the present is called bootstrapping. It lets us learn from short slices of experience instead of whole episodes.

3.2 Fixed-length rollouts across many parallel environments

Rather than one episode at a time, A2C runs N environments in parallel and collects a fixed number of steps T from each — a rollout of shape T × N. This is exactly what our RolloutBuffer stores, and what OnPolicyStrategy.step() fills by stepping all N envs in lockstep (rl_factory/training/on_policy.py). More envs = more, less-correlated data per update = steadier gradients. We’ll come back to how episodes of different lengths fit into a fixed T × N grid in Part 5 — it’s the question that trips everyone up.

3.3 An entropy bonus to keep exploring

A policy can prematurely become over-confident — always pressing the same button — and stop discovering better strategies. To fight this, A2C adds an entropy bonus. Entropy is a measure of how spread-out a probability distribution is: high when the policy is unsure (spreads probability across actions), low when it’s certain (all probability on one action). We add a small reward for keeping entropy high, which gently encourages the policy to keep its options open. In the loss it appears as a term we subtract (because we minimize loss but want to increase entropy).

3.4 Training the critic

The critic is trained by plain regression: make V(s_t) predict the actual return G_t (in practice, the bootstrapped target). The loss is the mean squared error between prediction and target — literally average of (V(s_t) − target_t)². As the critic gets better, the advantages get more accurate, which makes the actor’s updates better — a virtuous cycle.

At this point we essentially have A2C. PPO changes only two things: a smarter advantage estimate (GAE, Part 4) and a safety mechanism that lets us reuse each rollout several times (clipping, Part 6).

Part 4 — GAE: dialing the bias/variance knob

4.1 The trade-off

We now have two ways to estimate the advantage, and they sit at opposite extremes:

• One-step TD (δ_t from §3.1): uses one real reward and trusts the critic for everything after. Low variance (only one step of randomness), but biased — if the critic is wrong, the estimate is wrong.
• Full Monte Carlo (G_t − V(s_t), the REINFORCE-style version): uses all the real rewards to the episode’s end. Unbiased, but high variance — it drags in all that future noise.

“Bias” = systematically off in one direction (a bad-but-consistent ruler). “Variance” = jittery, different each time. You’d like a knob to trade one for the other. Generalized Advantage Estimation (GAE) is that knob.

4.2 The GAE formula

GAE blends all the multi-step estimates with an exponentially decaying weight controlled by a new parameter λ (lambda, between 0 and 1):

A_t  =  δ_t  +  (γλ)·δ_{t+1}  +  (γλ)²·δ_{t+2}  +  (γλ)³·δ_{t+3}  +  ...

• λ = 0: only δ_t survives → one-step TD (low variance, more bias).
• λ = 1: everything survives with only γ decay → Monte Carlo (unbiased, high variance).
• λ ≈ 0.95 (our default): mostly trusts a few real steps, leans on the critic for the far future. In practice this is the sweet spot.

4.3 How the code computes it — the backward loop

You would never sum that infinite series directly. There’s a beautiful trick: the same series can be computed with a single backward pass, because each step’s advantage is just its own TD error plus a discounted copy of the next step’s advantage. Here is the actual code from compute_gae in rl_factory/training/rollout.py, annotated:

for t in reversed(range(horizon)):          # walk backward through the T steps
    nonterminal = 1.0 - dones[t]            # 0 if the episode ended at step t, else 1
    delta = rewards[t] + gamma * next_values * nonterminal - values[t]   # δ_t
    last_advantage = delta + gamma * gae_lambda * nonterminal * last_advantage  # A_t
    advantages[t] = last_advantage
    next_values = values[t]                 # for the next (earlier) iteration

Walking backward, last_advantage carries the accumulated (γλ)·δ + (γλ)²·δ + … from all later steps, so each line adds one fresh δ_t and decays the rest — the whole infinite sum, in one cheap loop.

The nonterminal = 1 − dones[t] factor is the important safety catch. On the step where an episode ended, it becomes 0, which does two things at once: it drops the γ·V(s_{t+1}) term (you must not bootstrap the value of a state in a different episode) and it zeroes the carried-over last_advantage (rewards from after a reset must not leak backward into the previous episode). This is what lets many independent episodes share one array safely.

4.4 The value target falls out for free

At the end, compute_gae also returns the returns, computed as:

returns = advantages + values

Rearranged, that’s return_t = A_t + V(s_t) — an improved estimate of G_t. These are the regression targets the critic is trained against (§3.4). So one backward pass produces both things the loss needs: the advantages (for the actor) and the returns (for the critic).

Part 5 — The RolloutBuffer, concretely

Now we can read the buffer (rl_factory/training/rollout.py) end to end. It is the on-policy counterpart of a replay buffer: where DQN keeps a big pool of old transitions to sample randomly, the RolloutBuffer holds exactly one fresh batch of experience, learns from it, and throws it away — because on-policy methods may only learn from current-policy data (Part 1).

It has a strict three-phase life:

Phase 1 — Accumulate (add)

During collection, OnPolicyStrategy steps all N envs and calls add(...) once per timestep, appending that step’s slice to per-field lists. Per step it records, for all N envs:

Phase 2 — Finalize (compute_advantages)

Once T steps are collected, this stacks the lists into (T, N) arrays, runs compute_gae (Part 4) to get advantages and returns, and flattens everything from (T, N, …) to (T·N, …). After this the buffer is just a flat bag of T·N independent training examples, each a tuple (state, action, old_log_prob, advantage, return) — that’s the RolloutBatch NamedTuple.

Phase 3 — Serve (minibatches)

Yields shuffled minibatches (small random subsets) as tensors on the GPU. The learner asks for these repeatedly — several passes over the same rollout — which is where PPO’s clipping becomes essential (Part 6).

5.1 The question everyone asks: different episode lengths

“Each environment’s episode has a different length. How does a fixed T × N buffer cope?”

The answer is that the buffer never stores episodes — it stores a fixed-size time window. Every env is stepped exactly T times per rollout, in lockstep, so the buffer is always a full T × N rectangle, never ragged. Episodes simply don’t line up with the window or with each other, and that’s fine:

• When an env’s episode ends partway through the window, the environment worker auto-resets and immediately returns the first frame of the next episode (see VectorEnv in rl_factory/vector_env.py). That env just keeps producing steps — now belonging to a new episode. One env’s column can contain several episode boundaries; another’s may contain none.
• The only record that a boundary happened is the done flag at that step. And GAE’s nonterminal = 1 − done factor (§4.3) walls off the advantage math at each boundary, so reward never leaks from one episode into another, and each env’s column is handled independently.

Picture T = 5, N = 3 (· = ordinary step, D = episode ended and reset here):

        env0   env1   env2
 t=0     ·      ·      ·
 t=1     ·      D      ·
 t=2     D      ·      ·
 t=3     ·      ·      D
 t=4     ·      ·      ·     ← at the window's edge, these episodes are still unfinished

This whole grid is one rollout. It contains several partial episodes of different lengths, yet it’s a perfect rectangle. Two edges get special handling:

• Episodes still running at t = 4 (the window’s end): we can’t compute their tail, so we ask the critic for one last estimate, V(s_T), and seed the backward pass with it (the last_values argument to compute_gae, obtained via agent.bootstrap_values). That’s why we can learn from a slice of an unfinished episode.
• Episodes that ended inside the window: the done mask zeroes the bootstrap there, as described.

After Phase 2 flattens and Phase 3 shuffles, episode identity is gone entirely — every sample is an independent (state, action, old_log_prob, advantage, return). That’s safe because everything temporal was already baked into advantage and return during GAE. The gradient step doesn’t need to know which env or episode a sample came from.

Part 6 — PPO: reusing each rollout without blowing up

6.1 The motivation

Plain A2C is wasteful: collect a rollout, take one gradient step, throw it away. Collection (stepping the game) is expensive; we’d love to take several gradient steps per rollout. But there’s a catch rooted in Part 1: the policy-gradient average is only valid for the policy that collected the data. After one update, the policy has moved, and taking another step on the same data is technically using the wrong distribution. Push too far and the update becomes nonsense — the policy can collapse. PPO (“Proximal Policy Optimization,” 2017) is the mechanism that makes reusing a rollout safe.

6.2 The probability ratio

To reuse old data with an updated policy, we reweight each sample by how much more (or less) likely the new policy is to take the action than the old one did:

ratio  =  π_new(a | s)  /  π_old(a | s)  =  exp( new_log_prob − old_log_prob )

(The exp of a difference of logs is just the ratio of the probabilities — cheaper and more stable to compute this way.) This is why the buffer stored old_log_probs back at collection time: it’s the “before” snapshot the ratio needs. A ratio of 1 means “the policy hasn’t changed its mind about this action”; above 1 means “now more likely”; below 1, “now less likely.”

The reweighted objective for one action is ratio · advantage. If we optimized just this, we could still push the ratio arbitrarily far and break things — which is exactly the danger.

6.3 The clip: PPO’s one real idea

PPO caps how far the ratio is allowed to help. It compares the plain objective against a clipped version where the ratio is confined to [1 − ε, 1 + ε] (ε ≈ 0.2), and keeps the more pessimistic (smaller) of the two. Here’s the actual code, from _clipped_policy_loss in rl_ablations/models/ppo/learner.py:

ratio     = torch.exp(new_log_probs - old_log_probs)
unclipped = ratio * advantages
clipped   = torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps) * advantages
return -torch.min(unclipped, clipped).mean()

Why min of the clipped and unclipped versions? It removes any incentive to move the policy too far in the helpful direction:

• Good action (advantage > 0): we want to raise its probability, so the ratio wants to grow. But once the ratio passes 1 + ε, the clipped term stops increasing, and min picks it — so there’s no extra reward for pushing further. The update is content to nudge, not lunge.
• Bad action (advantage < 0): we want to lower its probability. The clip floors the ratio at 1 − ε, and min (of two negative numbers) keeps the penalty from being softened by a runaway ratio — again capping how far one update moves.

The − and .mean() turn it into a loss we minimize, averaged over the minibatch. The net effect: each rollout can be reused for several epochs of minibatch updates (our defaults: epochs passes × num_minibatches each) because no single step is allowed to move the policy far from the one that collected the data. That safety is the entire reason PPO is sample- efficient and stable.

6.4 A drift monitor: approximate KL

As a diagnostic, the learner also computes approx_kl = mean(old_log_prob − new_log_prob), a cheap estimate of how far the policy has drifted from the collector (KL is a standard “distance between two probability distributions”). We only log it here, but it’s the standard signal for “we’ve reused this rollout enough; stop before we drift too far.”

Part 7 — Putting the whole loss together

One PPO gradient step (_gradient_step in rl_ablations/models/ppo/learner.py) combines three terms we’ve now met. In plain words:

total_loss  =        clipped_policy_loss          # Part 6 — make good actions likelier, safely
            +  value_coef  · value_loss           # Part 3.4 — critic predicts returns better
            −  entropy_coef · entropy             # Part 3.3 — keep exploring

and in the actual code (_total_loss):

return policy_loss + config.value_coef * value_loss - config.entropy_coef * entropy

Reading the three pieces:

1. policy_loss — the clipped surrogate from Part 6. Drives the actor. Uses the advantages from GAE and the stored old_log_probs.
2. value_loss — F.mse_loss(values, returns): mean squared error between the critic’s current predictions and the GAE returns. Drives the critic. value_coef (≈ 0.5) sets how much we care about critic accuracy relative to the policy.
3. entropy — the policy’s entropy, subtracted so that minimizing the loss increases it. entropy_coef (a small number like 0.01) sets how hard we push for exploration.

The step itself (_descend) is standard deep learning with one extra guardrail:

self._optimizer.zero_grad()          # clear old gradients
loss.backward()                      # compute new gradients
clip_grad_norm_(..., max_grad_norm)  # cap the gradient size (a second safety belt)
self._optimizer.step()               # nudge the weights

clip_grad_norm_ scales the gradient down if it’s too large — a blunt but effective way to stop a single freak batch from taking a giant, destabilizing step. It complements PPO’s probability clipping: one bounds how far the policy moves in probability, the other bounds how big the weight update is.

One more detail you’ll see in _normalize: before use, advantages are often whitened (shifted and scaled to mean 0, standard deviation 1) across the minibatch. This doesn’t change their signs (good stays good), just keeps their scale consistent from batch to batch, which makes the learning rate behave predictably. It’s a practical trick, not part of the theory.

Part 8 — The historical arc, in one table

Each method kept everything before it and added one idea. PPO is A2C-with-GAE plus a clip — which is why, once you understand the RolloutBuffer (collection + advantages) and the three- term loss, you understand the whole lineage.

Part 9 — Trace it yourself

The best way to lock this in is to follow one rollout through the code:

1. Collect — OnPolicyStrategy.step() in rl_factory/training/on_policy.py: the loop that steps all N envs T times and calls buffer.add(...).
2. Bootstrap + advantages — the same method calls agent.bootstrap_values(...) then buffer.compute_advantages(...), which runs compute_gae in rl_factory/training/rollout.py. Print advantages and returns and sanity-check a few by hand.
3. Learn — PPOLearner.learn(rollout) in rl_ablations/models/ppo/learner.py loops over buffer.minibatches(...) calling _gradient_step. Read _clipped_policy_loss, the value_loss, and _total_loss, and match each to Parts 6–7 above.

If you can point at the line where each row of the Part 8 table lives, you’re done. ```

Results on Mario 1-1

Plugged into the on-policy strategy, PPO overtakes DQN. Where DQN plateaued around a mean return of ~2774, PPO keeps climbing to ~4545 over 15M steps — the head-to-head curve is in the DQN post. Same Trainer, same environment; the win comes from optimizing the policy directly instead of routing every improvement through an argmax over Q-values.

Where PPO runs out of road

PPO is a strong baseline on 1-1 — but 1-1 hands out a fairly dense reward (move right, get reward). The moment reward is sparse — long stretches with no feedback until a distant payoff — a pure reward-maximizer has nothing to climb. The next posts add an intrinsic reward: a bonus for surprise, so the agent keeps exploring when the environment goes quiet. First up is ICM, curiosity from predicting the consequences of your own actions → ICM.