insideout Get Pro

10 · Inside the RL step: policy gradients and PPO

Pro

Training and Tuning, lesson 10

Inside the RL step: policy gradients and PPO

How a reward becomes a change in the model's weights.

In "Learning from preferences" and "Rewarding checkable answers", a model writes answers, something scores them, and training nudges the model toward the better ones. This lesson opens up that nudge: how a score becomes a change in the weights, and how PPO keeps each change small. You need softmax and its gradient, from "Backprop through softmax and cross-entropy" in Neural Networks.

By the end, you can:

  1. Work out one REINFORCE step for a softmax policy by hand.
  2. Subtract a baseline from rewards, and say why it cuts the noise.
  3. Work out PPO's clipped term for any ratio and advantage.

Part 1. Make what scored well more likely

Work out one REINFORCE step for a softmax policy by hand.

Shrink a language model to one choice. Asked "What is 7 × 8?", it picks one of three answers, with logits 2.0 for 54, 1.0 for 56 and 0.0 for 48. A checker gives reward 1 for the right answer and 0 otherwise. It never sees the right answer, only the reward for what it said.

Reinforcement learning calls whatever picks the actions a policy, and writes the chance it gives an action as π. So π(56) is the chance the model answers 56.

Example 1.1

One REINFORCE step

Treat the three logits as the weights. The model samples 56 and gets reward 1. With a step size of 0.5, how do the logits move, and how likely is 56 afterwards?

  1. Turn logits into chancessoftmax(2.0, 1.0, 0.0) = (0.665, 0.245, 0.090)

    The favourite, 54, is wrong.

  2. Take the gradient of ln π(56)(0 − 0.665, 1 − 0.245, 0 − 0.090) = (−0.665, +0.755, −0.090)

    For a softmax it's 1 for the sampled answer, minus each answer's chance.

  3. Scale by step size and reward0.5 × 1 × (−0.665, +0.755, −0.090) = (−0.333, +0.378, −0.045)

    A reward of 0 would leave the logits where they are.

  4. Move the logits(2.0, 1.0, 0.0) → (1.667, 1.378, −0.045)

    Each rival drops in proportion to its own chance.

  5. Recompute the chancessoftmax → (0.518, 0.388, 0.094)

    From about one in four to almost two in five.

Answer56 rises from 0.245 to 0.388.

The gradient of ln π(56) is the cross-entropy gradient from "Backprop through softmax and cross-entropy" with its sign flipped: the direction that raises 56's chance fastest. The reward sets how far to go. Averaged over many samples, these pushes follow the gradient of the average reward, the number you want to raise.

beforeafter0.6650.518542.0 → 1.6670.2450.388561.0 → 1.3780.0900.094480.0 → −0.045chance of each answer
Figure 1One rewarded sample of 56 moves chance toward it. Under each answer is its logit before and after. 48's logit hardly moves, because the push on each rival is scaled by its own chance.

Policy gradient and REINFORCE

A policy gradient method changes the weights θ to raise the average reward. The simplest, REINFORCE (Williams, 1992), samples an action, scores it with a reward R, and steps along ∇ ln π(action), read "the gradient of the action's log-probability". α is the step size.

θ ← θ + α × R × ∇ ln π(action)

A real answer is many tokens, and ln π(answer) sums their log-probabilities. So every token gets the same push from the one reward at the end, helpful or not, which makes each step noisy.

Your turn

A policy starts with three equal logits: 0, 0 and 0. It samples the first action and gets reward 2. With a step size of 0.3, what are the new logits?

Each chance is 1/3, so the gradient of ln π(first) is (2/3, −1/3, −1/3). Times 0.3 × 2 = 0.6, that's (+0.4, −0.2, −0.2). Moving only the chosen logit gives (0.6, 0, 0), and not subtracting its own chance gives (0.6, −0.2, −0.2). (0.2, −0.1, −0.1) forgets the reward.

Part 2. Better or worse than expected

Subtract a baseline from rewards, and say why it cuts the noise.

Keep going with Pro

You've done the first part. Pro opens the rest of this lesson, and every other one.

Still to come
  1. Better or worse than expected
  2. Small steps on reused answers

Training and Tuning overviewSpotted a mistake? Tell us.