Policy gradient

Overview

Andrej Karpathy’s post “Pong from Pixels” (2016) trains a network on Atari Pong with the same loss as our Catch. It names the ideas behind it; this lesson works through them with small, checkable examples.

The problem

Catch had states, actions and rewards. Written as one structure, that is the standard way RL describes any task:

Markov decision process (MDP)

a set of states, the actions available in each, the rule for which state comes next, and the reward on the way.

Transition

one step from a state, through an action, to the next state. It may be random: the same action can lead to different next states.

The word “Markov” is a promise about the state:

Markov property

the current state holds everything needed to choose well. Knowing the past adds nothing.

Is one screen a Markov state? In Catch, yes: the ball always falls straight down, so its position tells you everything. In Pong it is not:

• In one frame, a ball moving left and a ball moving right look identical, but the right paddle move is opposite.
• With two frames, the motion is visible. Karpathy feeds his network the difference between the current frame and the previous one.

So: deciding what goes into the state is part of setting up the problem. If the state is not Markov, no policy can do well, however long it trains.

The estimator

What training wants is a high expected reward: the average reward the policy would earn over many tries. Call it \(J\), and the policy’s numbers \(\theta\).

Gradient descent needs the slope of \(J\) with respect to \(\theta\). But the reward comes from the game, after a move picked by rolling a die. There is no formula for it to differentiate, and you cannot backpropagate through a die roll.

The trick is one line of calculus. The slope of a probability can be written with the slope of its log, \(\nabla p = p\,\nabla \log p\). Put that inside the average and the slope of the expected reward becomes an average you can measure:

In words:

• Sample actions from the policy and record each one’s reward \(r\).
• For each action, find the direction that makes it more likely: \(\nabla \log p(a)\).
• Weight that direction by the reward, and average.

Score function

\(\nabla \log p(a)\): the direction in which the policy’s numbers must move to make action \(a\) more likely.

Score-function estimator

the average of reward × score function over sampled actions. It estimates the slope of the expected reward without differentiating the reward.

The smallest example: a policy that picks one number \(x\) from a bell curve centred at \(\mu\). The reward is +1 if \(x\) lands between 1 and 2.5, and −1 anywhere else. Training should move \(\mu\) into that window.

def push(x, mu):
    """Which way to move mu to make x more likely: the slope of log p(x) with respect to mu."""
    return (x - mu) / SIGMA ** 2


def estimate(mu, samples, baseline=False, offset=0.0):
    """The score-function estimate of how expected reward changes with mu."""
    rewards = [reward(x, offset) for x in samples]
    b = sum(rewards) / len(rewards) if baseline else 0.0     # subtract the batch's average reward
    return sum((r - b) * push(x, mu) for x, r in zip(samples, rewards)) / len(samples)

"""Why the policy-gradient loss works, on the smallest possible example.

Run: python3 site/ml/foundations/examples/policy_gradient.py

A policy picks a number x from a bell curve centred at mu (spread 1). The reward is +1 when x lands
in the window [1, 2.5] and -1 everywhere else. Training should move mu into the window, using only
samples and their rewards: no formula for the reward, no derivative of it.
"""

import math
import random

LOW, HIGH = 1.0, 2.5        # the rewarded window
SIGMA = 1.0                 # the bell curve's spread (fixed)


def reward(x, offset=0.0):
    """+1 inside the window, -1 outside. `offset` adds the same amount to every reward."""
    return (1.0 if LOW <= x <= HIGH else -1.0) + offset


def push(x, mu):
    """Which way to move mu to make x more likely: the slope of log p(x) with respect to mu."""
    return (x - mu) / SIGMA ** 2


def estimate(mu, samples, baseline=False, offset=0.0):
    """The score-function estimate of how expected reward changes with mu."""
    rewards = [reward(x, offset) for x in samples]
    b = sum(rewards) / len(rewards) if baseline else 0.0     # subtract the batch's average reward
    return sum((r - b) * push(x, mu) for x, r in zip(samples, rewards)) / len(samples)


def exact_slope(mu):
    """The true slope of expected reward, from the bell curve's formula (only for checking)."""
    pdf = lambda z: math.exp(-z * z / 2) / math.sqrt(2 * math.pi)
    return 2 / SIGMA * (pdf((LOW - mu) / SIGMA) - pdf((HIGH - mu) / SIGMA))


def train(steps=25, n=20, lr=0.5, baseline=False, offset=0.0, seed=0):
    rng = random.Random(seed)
    mu, history = 0.0, []
    for _ in range(steps):
        samples = [rng.gauss(mu, SIGMA) for _ in range(n)]
        g = estimate(mu, samples, baseline, offset)
        history.append({"mu": mu, "samples": samples, "rewards": [reward(x, offset) for x in samples], "estimate": g})
        mu += lr * g                                           # step uphill on expected reward
    return mu, history


def spread(mu=0.0, n=20, trials=2000, baseline=False, offset=0.0, seed=1):
    """How much single-batch estimates scatter around the true slope."""
    rng = random.Random(seed)
    values = [estimate(mu, [rng.gauss(mu, SIGMA) for _ in range(n)], baseline, offset) for _ in range(trials)]
    mean = sum(values) / trials
    return mean, math.sqrt(sum((v - mean) ** 2 for v in values) / trials)


if __name__ == "__main__":
    print(f"true slope at mu = 0: {exact_slope(0.0):.3f}")
    for offset in (0, 20):
        for b in (False, True):
            mean, sd = spread(baseline=b, offset=offset)
            mu, _ = train(baseline=b, offset=offset)
            print(f"offset {offset:2d}, {'baseline' if b else 'raw     '}: estimates average {mean:.3f}, "
                  f"spread {sd:.3f}; mu after 25 steps {mu:.2f}")
    gamma = 0.9
    print("discounted credit for the 7 moves of one drop, reward at the end:",
          ", ".join(f"{gamma ** (6 - t):.2f}" for t in range(7)))

For a bell curve, the direction that makes \(x\) more likely is simply toward \(x\): push is \(x - \mu\) when the spread is 1. Step through training, 20 samples per step:

How to read the figure:

• Dots are the 20 samples: blue for reward +1, orange for −1.
• Each small arrow starts at \(\mu\) and shows which way one sample pulls it: reward × push. A rewarded sample pulls \(\mu\) toward itself; a punished one pushes \(\mu\) away from itself.
• The big arrow is their average: the estimate. \(\mu\) moves by learning rate 0.5 × that estimate.

Is the estimate right? At \(\mu = 0\) the true slope, from the bell curve’s formula, is . Over 2,000 batches of 20 samples, the estimates average : one batch is noisy, but it points the right way on average.

This is exactly Catch. The policy there is the network, \(\log p\) comes from log_softmax, and loss.backward() computes \(\nabla \log p\) for every weight at once. Minimising \(-r \log p\) is stepping along this estimate.

REINFORCE

training a policy with the score-function estimator: the name of the method in Catch and in Karpathy’s Pong.

Baseline

Only relative goodness should matter. Add 20 to every reward, and nothing changes about which outcomes are better. The true slope stays the same.

The raw estimate does not take that well. With every reward near +20, every sample pulls \(\mu\) toward itself, hard. The pulls mostly cancel, but what is left over is large and random.

The fix is to subtract a baseline, here the batch’s average reward, before weighting. Then a sample counts as good only if it beat the others:

Baseline

a reference value subtracted from every reward before it weights the push. It does not change the slope on average, only the noise.

Advantage

reward minus baseline: how much better than usual an action turned out.

Measured over 2,000 batches at \(\mu = 0\):

The baseline’s average sits a little below the true slope, because each sample’s own reward is part of the batch average it is compared with. With 20 samples that costs about 5%.

Training shows the same. In the bell-curve figure above, choose Rewards +20: \(\mu\) lurches. Then choose Rewards +20, baseline: the run looks like the plain one. Over 20 random seeds, after 25 steps:

Karpathy goes one step further and standardises the rewards of a batch: subtract the mean, divide by the spread. Then roughly the better half of the actions is encouraged and the worse half discouraged, whatever the game’s reward scale.

So: in practice the reward is almost never used raw. It is turned into an advantage first.

Discount

In Catch, all seven moves of a drop got the same ±1. In Pong, a point can be decided by a move made dozens of frames earlier, and moves after it had no effect on it.

A discount factor \(\gamma\) gives each move credit for the rewards that follow it, with later rewards counting less: each step of delay multiplies by \(\gamma\).

For one drop of Catch with \(\gamma = 0.9\) and the reward at the end, the seven moves would get:

Return

the discounted sum of the rewards that follow a move: \(r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \dots\)

Karpathy uses \(\gamma = 0.99\) and resets the sum at every point scored, so each move is judged by the point it helped decide.

Catch did not need it: every move in a drop can still change where the paddle ends up. Long tasks, like the robot arm reaching for a cube, do.

The family

Karpathy wrote in 2016 that this was, “modulo some details”, the state of the art. The core ideas on this page are still in every modern method. What changed is how they are combined:

• REINFORCE (this lesson, Catch, Pong): sample, score each action by its return, push.
• Actor-critic: a second network learns to predict the return from each state and serves as the baseline. The “actor” is the policy; the “critic” is the predictor.
• PPO (proximal policy optimization, 2017): actor-critic that limits how far one update can move the policy. It became the default policy-gradient method, in games, robotics and language models.
• Q-learning (the robot-arm workshops): learn how much each action is worth, then pick the best. Karpathy called it “not a great algorithm”; that is too harsh today. Improved versions are strong, and actor-critic methods combine both ideas.
• Model-based RL: learn a model of the world, then plan or practise inside it, needing far fewer real tries. MuZero and DreamerV3 are well-known examples.
• RL for language models: chat models are tuned with rewards from human preferences, first with PPO. Recent reasoning models use simpler variants such as GRPO, which is close to REINFORCE with the rewards of a group of sampled answers standardised: Karpathy’s trick, at a much larger scale.

What has not changed is the main weakness Karpathy pointed out: the method needs a lot of experience. His Pong agent played about 200,000 games; a person needs a few. Much of the progress since has gone into getting more from each try.

Recap

• reward(x, offset) is the game: it scores a sample and offers no derivative.
• push(x, mu) is the score function: the direction that makes a sample more likely.
• estimate(mu, samples) averages reward × push: the score-function estimate of the slope of expected reward.
• baseline=True subtracts the batch’s average reward first, turning rewards into advantages and cutting the noise.
• mu += lr * g steps uphill. In Catch, loss = -(rewards * taken).mean() and optimizer.step() do the same for every weight.

The loss \(-r \log p\) is not a trick that happens to work. It is the slope of the expected reward, estimated from the network’s own tries, and every modern policy-gradient method is built on it.