Policy gradients: from REINFORCE to PPO

The previous section learned what actions are worth and then acted greedily on that. This section learns the behaviour itself: make the actions that preceded a good outcome more likely and the others less. The idea works and is extremely noisy, so most of the section is about what gets added to make it usable, namely a baseline, a second network that predicts what to expect, and a limit on how far a single update may move the policy. That last one is PPO, which is what language model post-training is built on.

Skip the values, learn the behaviour#

A policy network outputs probabilities over actions. After an episode, make the actions that led to high return more likely and the others less likely. That is the whole idea. The maths makes it precise: the gradient of expected return is the average of the gradient of log probability of each action, weighted by the return that followed it.

It is worth seeing why that weighting is the right one rather than a plausible one. You cannot differentiate the reward with respect to the policy's parameters, because the environment is not differentiable and nobody gave you its equations. What you can differentiate is the probability of the action you took. Raising the log probability of an action by one unit multiplies how often that action happens; weighting that push by the return that followed means good episodes push their own actions up hard and bad episodes push theirs down. Average over enough episodes and the pushes that were luck cancel, while the ones that were the policy survive. Enough episodes is the entire problem, and the rest of this section is about reducing how many are needed.

An environment you can write down#

Cart-pole is the standard first problem: a pole hinged on a cart, two actions, push left or push right, one point per time step the pole stays within twelve degrees of upright and the cart within 2.4 units of the centre. It is four lines of physics and worth writing out, because an environment you wrote is an environment you can instrument.

import math, torch, torch.nn as nn

GRAV, M_CART, M_POLE, HALF_LEN, FORCE, DT = 9.8, 1.0, 0.1, 0.5, 10.0, 0.02
TOTAL_MASS, POLE_ML = M_CART + M_POLE, M_POLE * HALF_LEN

def cart_step(state, action):
    """The standard cart-pole equations. State is (x, x_dot, angle, angle_dot)."""
    x, x_dot, th, th_dot = state.tolist()
    force = FORCE if action == 1 else -FORCE
    cos, sin = math.cos(th), math.sin(th)
    temp = (force + POLE_ML * th_dot ** 2 * sin) / TOTAL_MASS
    th_acc = (GRAV * sin - cos * temp) / (HALF_LEN * (4 / 3 - M_POLE * cos ** 2 / TOTAL_MASS))
    x_acc = temp - POLE_ML * th_acc * cos / TOTAL_MASS
    nxt = torch.tensor([x + DT * x_dot, x_dot + DT * x_acc,
                        th + DT * th_dot, th_dot + DT * th_acc], dtype=torch.float32)
    fallen = abs(nxt[0].item()) > 2.4 or abs(nxt[2].item()) > 12 * math.pi / 180
    return nxt, 1.0, fallen                       # one point per step survived

def run_episode(policy, rng, cap=500):
    state = torch.rand(4, generator=rng) * 0.1 - 0.05
    log_probs, rewards = [], []
    for t in range(cap):
        dist = torch.distributions.Categorical(logits=policy(state))
        a = dist.sample()
        log_probs.append(dist.log_prob(a))
        state, r, fallen = cart_step(state, a.item())
        rewards.append(r)
        if fallen:
            break
    return log_probs, rewards

def rewards_to_go(rewards, gamma=0.99):
    out, G = [], 0.0
    for r in reversed(rewards):
        G = r + gamma * G
        out.insert(0, G)
    return torch.tensor(out)

Note what the reward-to-go function does. Each action is credited with the rewards that came after it, not with the whole episode's total, because an action cannot have caused a reward that arrived before it. That is free variance reduction and it is the only change from the naive formula that requires no justification beyond causality.

REINFORCE#

def train(seed, baseline, episodes=400, lr=1e-2):
    torch.manual_seed(seed)
    rng = torch.Generator().manual_seed(seed)
    policy = nn.Sequential(nn.Linear(4, 32), nn.Tanh(), nn.Linear(32, 2))
    opt = torch.optim.Adam(policy.parameters(), lr=lr)
    lengths = []
    for ep in range(episodes):
        log_probs, rewards = run_episode(policy, rng)
        lengths.append(len(rewards))
        G = rewards_to_go(rewards)
        weights = (G - G.mean()) / (G.std() + 1e-8) if baseline else G
        loss = -(torch.stack(log_probs) * weights).sum()
        opt.zero_grad(); loss.backward(); opt.step()
    return lengths

for baseline in (False, True):
    for seed in (0, 1, 2):
        lengths = train(seed, baseline)
        print(f"baseline={baseline!s:5s} seed {seed}: "
              f"first 50 episodes {sum(lengths[:50]) / 50:6.1f} steps, "
              f"last 50 {sum(lengths[-50:]) / 50:6.1f}, best {max(lengths)}")
# baseline=False seed 0: first 50 episodes   36.8 steps, last 50  293.6, best 500
# baseline=False seed 1: first 50 episodes   12.5 steps, last 50  500.0, best 500
# baseline=False seed 2: first 50 episodes   39.2 steps, last 50   14.1, best 500
# baseline=True  seed 0: first 50 episodes   82.5 steps, last 50  376.3, best 500
# baseline=True  seed 1: first 50 episodes   36.8 steps, last 50  500.0, best 500
# baseline=True  seed 2: first 50 episodes   33.7 steps, last 50  324.2, best 500

The third row is the one to sit with. Same code, same hyperparameters, seed 2, and the policy ends at 14 steps per episode having at some point reached the 500-step cap. It did not fail to learn; it learned and then destroyed what it had. A large update after a run of lucky episodes moved the policy somewhere that collects worse episodes, which produce worse updates. Nothing in the loss function notices, because there is no held-out set and no ground truth, only the agent's own increasingly bad data. If you had run one seed and it had been this one, the honest conclusion would have been that the code is broken.

Variance, and why the baseline matters#

Returns vary enormously from episode to episode by luck, so raw REINFORCE gradients are very noisy. Subtracting a baseline, a typical return for that situation, keeps the gradient unbiased but far less noisy: actions are judged by whether they did better than usual, not by the raw total.

That claim is testable, and testing it turns up something the tutorials do not say. Take a half-trained policy, sample 200 episodes, and compute the gradient four ways: weighted by the raw reward-to-go, minus a constant baseline measured from other episodes, minus the mean of the episode's own reward-to-go, and the standardised version the training loop above uses. Then compare not just the variance but the variance against the signal — a weighting that divides everything by 60 reduces variance by 3,600 and has told you nothing.

torch.manual_seed(0)
rng = torch.Generator().manual_seed(0)
policy = nn.Sequential(nn.Linear(4, 32), nn.Tanh(), nn.Linear(32, 2))
opt = torch.optim.Adam(policy.parameters(), lr=1e-2)
for ep in range(150):                                    # a half-trained policy to measure at
    lp, rw = run_episode(policy, rng)
    G = rewards_to_go(rw)
    loss = -(torch.stack(lp) * ((G - G.mean()) / (G.std() + 1e-8))).sum()
    opt.zero_grad(); loss.backward(); opt.step()

probe = [rewards_to_go(run_episode(policy, rng)[1]).mean().item() for _ in range(50)]
b = sum(probe) / len(probe)                              # a state-independent baseline

def gradient(log_probs, weights):
    loss = -(torch.stack(log_probs) * weights).sum()
    opt.zero_grad(); loss.backward(retain_graph=True)
    return torch.cat([p.grad.flatten().clone() for p in policy.parameters()])

rows = {"raw reward-to-go": [], "minus a constant": [], "minus episode mean": [], "standardised": []}
for k in range(200):
    lp, rw = run_episode(policy, rng)
    G = rewards_to_go(rw)
    rows["raw reward-to-go"].append(gradient(lp, G))
    rows["minus a constant"].append(gradient(lp, G - b))
    rows["minus episode mean"].append(gradient(lp, G - G.mean()))
    rows["standardised"].append(gradient(lp, (G - G.mean()) / (G.std() + 1e-8)))

print(f"baseline b = {b:.1f}, measured over 200 episodes")
for name, grads in rows.items():
    stack = torch.stack(grads)
    var = stack.var(0).sum().item()                      # total variance of the estimator
    signal = stack.mean(0).norm().item()                 # the direction it agrees on
    print(f"  {name:20s} variance {var:11.1f}   mean gradient {signal:8.3f}   "
          f"ratio {signal / var ** 0.5:.3f}")
# baseline b = 66.4, measured over 200 episodes
#   raw reward-to-go     variance   4397940.5   mean gradient  299.672   ratio 0.143
#   minus a constant     variance   1080077.6   mean gradient  230.449   ratio 0.222
#   minus episode mean   variance   1043069.0   mean gradient  142.781   ratio 0.140
#   standardised         variance      1666.7   mean gradient    6.059   ratio 0.148

The constant baseline does what the textbook says. Variance falls fourfold and the ratio rises from 0.143 to 0.222, which is a real improvement: over half again as much usable signal per episode, for the cost of tracking one number. Nothing about it is subtle.

The other two rows are the controls, and they kill a claim people make constantly. Subtracting the episode's own mean drops the variance by the same fourfold and the ratio comes out at 0.140, which is no better than doing nothing at all — it removed as much signal as noise. Standardising drops the variance by a factor of 2,600, which looks spectacular and is entirely a change of scale: the mean gradient shrinks by the same factor and the ratio finishes at 0.148. What standardising buys in the training loop above is a step size that does not depend on how long episodes have become, which matters a great deal to Adam and is not variance reduction. Report a variance number without its signal and you can prove anything.

Actor-critic#

Make the baseline smart: train a second network, the critic, to predict the expected return from each state. The difference between what happened and what the critic expected is the advantage. The actor is pushed towards actions with positive advantage.

A critic is a per-state baseline and it is legal, because it is a function of the state alone and never of the action taken. That is what makes it the right answer where the constant above is a workable one: the constant says whether this episode beat the average episode, while the critic says whether this action beat what was expected from this position, which is a far sharper question when some states are simply better than others. The price is a second network with its own loss, its own learning rate and its own way of being wrong, and a critic that lags the policy injects a bias of its own. Generalised advantage estimation, the lambda you see in every PPO configuration, is a dial between trusting the critic and trusting the observed rewards.

PPO: do not move too far#

ratio = pi_new(a | s) / pi_old(a | s)       how much more likely the action became

L = mean( min( ratio * A,  clip(ratio, 1 - eps, 1 + eps) * A ) )      eps around 0.2

PPO reuses each batch of experience for several updates, which is efficient but dangerous: the policy can drift far from the one that collected the data. The clip removes any incentive to push the probability ratio beyond a small band. Taking the minimum means the objective never rewards going further. It is crude and very robust, which is why it became the default.

The objective is four lines of arithmetic, and printing its slope is more convincing than any description of it. Slope here means how much the objective changes as the probability ratio changes, which is the only thing an optimiser responds to.

EPS = 0.2

def ppo_term(ratio, adv):
    return min(ratio * adv, min(max(ratio, 1 - EPS), 1 + EPS) * adv)

print("ratio    advantage +1            advantage -1")
print("         plain    ppo   slope    plain    ppo   slope")
for ratio in (0.6, 0.9, 1.0, 1.1, 1.4, 2.0):
    cells = ""
    for adv in (1.0, -1.0):
        slope = (ppo_term(ratio + 1e-4, adv) - ppo_term(ratio - 1e-4, adv)) / 2e-4
        cells += f"  {ratio * adv:+6.2f} {ppo_term(ratio, adv):+6.2f}  {slope:+5.2f} "
    print(f"{ratio:5.2f} {cells}")
# ratio    advantage +1            advantage -1
#          plain    ppo   slope    plain    ppo   slope
#  0.60    +0.60  +0.60  +1.00    -0.60  -0.80  +0.00
#  0.90    +0.90  +0.90  +1.00    -0.90  -0.90  -1.00
#  1.00    +1.00  +1.00  +1.00    -1.00  -1.00  -1.00
#  1.10    +1.10  +1.10  +1.00    -1.10  -1.10  -1.00
#  1.40    +1.40  +1.20  +0.00    -1.40  -1.40  -1.00
#  2.00    +2.00  +1.20  +0.00    -2.00  -2.00  -1.00

Read the slope columns and the design is obvious. For a good action, advantage +1, the gradient is a constant +1 until the action has become 1.2 times as likely as it was, and then it is exactly zero: no matter how good the action looked, this batch will not push it further. For a bad action, advantage -1, the gradient keeps pushing the probability down until the action is 0.8 times as likely, and then stops. The plain columns show what would happen without the clip — an objective that grows without limit in both directions, which is precisely the incentive to take one enormous step off the data you have.

Notice that the clip is one-sided in each case. With a positive advantage the objective is still free to move downwards, and with a negative advantage still free to move up: the min never blocks the direction that undoes an overshoot. That asymmetry is the whole trick, and it is why PPO is safe to run for several epochs over one batch of experience while plain policy gradient is not.

Where this lands for an AI engineer: PPO is the algorithm underneath RLHF, and the clip is the reason a preference-training run does not usually destroy the model outright. Two further guards appear there, both covered in the section on language models in this topic — a KL penalty against the starting model, which is a different constraint from the clip and does a different job, and the option of skipping the sampling loop entirely with DPO. Implementing PPO from scratch is not a good use of your time; being able to say what the ratio, the advantage, the clip and the min each do, and why removing any of them breaks something specific, is.

What to skip: trust region derivations, the natural gradient, and the conjugate-gradient machinery in TRPO. They are the reason PPO has the shape it has, and knowing that PPO is a cheap approximation of a trust region is the entire useful content. Read a clean implementation instead — CleanRL's single-file version is a few hundred lines — and check each term against the table above.