Reinforcement learning in one sitting: policy, reward, baseline, advantage

Reasoning training uses reinforcement learning (RL). You do not need prior RL knowledge. Four ideas, then one tiny program.

#The four ideas

The training idea: make actions that led to above-average rewards more likely, and actions that led to below-average rewards less likely. That is all of policy gradients.

#A tiny example you can run

A policy with three actions, where action 2 is the only one that earns a reward. We sample a group of 8 attempts, compute each one's advantage against the group average, and nudge the probabilities:

import random, math
random.seed(0)

logits = [0.0, 0.0, 0.0]            # the policy's parameters: all actions equally likely

def probs(l):
    m = max(l); e = [math.exp(x - m) for x in l]; t = sum(e)
    return [v / t for v in e]

lr = 0.5
for step in range(200):
    p = probs(logits)
    acts = random.choices(range(3), weights=p, k=8)        # a group of 8 attempts
    rewards = [1.0 if a == 2 else 0.0 for a in acts]       # reward only for action 2
    base = sum(rewards) / len(rewards)                     # baseline = group average
    for a, r in zip(acts, rewards):
        adv = r - base                                     # advantage
        for j in range(3):                                 # gradient of log-probability of action a
            grad = (1.0 if j == a else 0.0) - p[j]
            logits[j] += lr * adv * grad / len(acts)

print([round(x, 3) for x in probs(logits)])

Output: [0.003, 0.003, 0.994]. The policy started at one third each and learned to pick action 2 about 99% of the time.

#Reading the loop

1. Sample a group from the current policy.
2. Score each attempt with the reward.
3. Subtract the group average to get advantages.
4. Update: for each attempt, raise the probability of its action in proportion to its advantage (and lower it if the advantage is negative). The expression (1 if j == a else 0) - p[j] is the gradient of the log-probability of action a with respect to the logits.

This is the shape of every RL-for-LLMs method. In a real model the "action" is a whole generated answer, the log-probability is the sum over its tokens, and the update flows into millions of weights through backpropagation instead of three numbers.

#Why the baseline matters

Suppose almost every attempt earns reward 1. Without a baseline, each of them pushes its actions up with the same strength, and the signal is mostly noise. With the group average subtracted, an attempt is only reinforced if it is better than its siblings. That reduces the variance of the updates, so training is steadier.

#Where it can go wrong

• All attempts equal. All right or all wrong gives zero advantage everywhere, so nothing is learned. Problems that are too easy or too hard waste compute.
• Reward hacking. The model finds a way to earn reward without solving the task (an answer format that fools the checker). The verifier has to be honest.
• Collapse. The policy becomes too narrow, too fast, and stops exploring. Real methods limit how far the policy may move per step (clipping and a KL penalty).
• Noisy rewards. A flawed verifier teaches the wrong behaviour confidently.

The training lessons return to each of these.

#Exercise

1. In a group of 8 attempts, 2 are correct. What is the baseline, and what is the advantage of a correct attempt and of an incorrect one?
2. What happens to the update if 8 of 8 are correct?
3. In the program above, change lr to 5.0 and run it a few times. What do you see, and why?

Answers: 1. The baseline is 2/8 = 0.25; a correct attempt has advantage 0.75 and an incorrect one −0.25. 2. The baseline is 1, every advantage is 0, and the policy does not change. 3. Large steps make the policy jump to one action quickly and may lock onto the wrong one early, or oscillate; a smaller learning rate is steadier, which is why training recipes also limit the update size.

Next: setting up and running the experiments.