Policy gradients and advantages

A policy π(a|s) is the model's distribution over actions in each state. Policy gradient methods change that distribution directly, so that actions which led to high returns become more likely. Language models are policies over tokens, so this is the update that trains them after supervised learning.

REINFORCE

REINFORCE is the simplest policy gradient. Sample a trajectory, compute its return G, and push up the log probability of every action in it, scaled by G:

∇θ J ∝ E[ ∇θ log πθ(a | s) · G ]

In plain words: if a completion scored well, make its tokens more likely. If it scored poorly, make them less likely.

Why the update is noisy

Two samples of the same prompt can score very differently by chance. The gradient then swings from batch to batch, so learning is slow or unstable. This is the variance problem.

A baseline b(s) fixes this. It is a number that depends on the state but not on the action. Subtracting it leaves the expected gradient unchanged and shrinks the noise:

∇θ J ∝ E[ ∇θ log πθ(a | s) · (G − b(s)) ]

When the baseline is the state's expected value, G − V(s) is an advantage. It says how much better this action was than the policy's average in that state.

Walking through a baseline

Take the prompt "2 + 2 = ?" and two completions. Completion A answers 4 and gets reward 1. Completion B answers 5 and gets reward 0.

Use the mean reward, 0.5, as the baseline:

advantage of A = 1 − 0.5 = +0.5    (make "4" more likely)
advantage of B = 0 − 0.5 = −0.5    (make "5" less likely)

Now take three completions with rewards 0, 1, 0. The mean is 1/3:

advantages = 0 − 1/3 = −1/3,   1 − 1/3 = +2/3,   0 − 1/3 = −1/3

The advantages add up to zero. The good completion is pushed up, and the bad ones are pushed down by the same total amount.

Common mistakes

• Using raw REINFORCE with sparse rewards. A completion that scores 0 contributes nothing, so the losers are never pushed down. A baseline fixes that.
• Choosing a baseline that depends on the action. That biases the gradient. The baseline must depend only on the state.
• Reading one batch as the truth. Advantages from a few samples are noisy. Use enough samples per prompt.

Check yourself

1. Why does subtracting a baseline keep the gradient unbiased? Answer: the baseline does not depend on the action, so its expected contribution to the gradient is zero.
2. What does a positive advantage mean? Answer: the action did better than the policy's average in that state, so it should become more likely.

Key takeaways

• A policy gradient raises the probability of actions in proportion to the return that followed.
• REINFORCE is simple, but high variance makes it slow.
• A baseline reduces variance without biasing the update.
• An advantage compares an action with the average. It is the main signal in modern methods.

Further learning

Next, see how PPO adds a limit to how far each update can move the policy, in PPO and the KL leash.