Reinforcement learning

Think of teaching a dog to sit. You can't show it the right answer: you can only wait for it to try something and give it a treat when the something was good. Over many tries the dog does more of what earned treats and less of what didn't. Nobody ever told it what "sit" means; it worked it out from a score.

That is reinforcement learning (RL). An agent (the dog, or a language model) takes an action (sits, or writes an answer), the world hands back a reward (a treat, or a score from a grader), and the agent adjusts itself so that rewarded actions become more likely. The agent's current habits, written as a probability for every action, are called its policy.

Compare this with ordinary supervised learning (primer.ml.losses), where every example comes with the correct answer attached. In RL there is no answer key, only a score after the fact. That is exactly the situation a language model is in once pretraining is over: for "prove this theorem" or "write a helpful reply" there is no single correct text to copy, but a checker or a judge can say how good an attempt was. RL is how the model learns from those judgements (primer.ml.training_stages places it in the training pipeline).

flowchart LR
  P[Policy<br/>a probability for every action] -->|sample| A[Action]
  A --> E[Environment<br/>a slot machine, a grader, a user]
  E --> R[Reward<br/>one number]
  R -->|nudge the policy| P

The simplest RL problem is a row of slot machines, called a multi-armed bandit (a slot machine is a "one-armed bandit"). Machine A pays out 20% of the time, B 50% and C 80%, but the agent isn't told that. Each pull pays 1 or 0. The agent must find C by pulling and seeing what happens.

A policy here is three probabilities, one per arm. A policy that picks each arm a third of the time earns, on average, (0.2 + 0.5 + 0.8) / 3 = 0.5 per pull. A policy that always picks C earns 0.8. Learning means moving from the first policy to the second using nothing but the 1s and 0s.

This one number, the average reward a policy expects, is what RL maximises:

policy = [1/3, 1/3, 1/3]
# average payout of arms A, B, C
R = [0.2, 0.5, 0.8]
# J = Σ_a π(a) R(a)
round(sum(p * r for p, r in zip(policy, R)), 3)  # → 0.5
round(sum(p * r for p, r in zip([0, 0, 1], R)), 3)  # → 0.8

The agent can't compute J, because it doesn't know R. It can only sample: pull an arm, see a 1 or a 0. Every method below turns those samples into an estimate of which way to move θ to make J bigger. Trying an arm you're unsure of is called exploration; sticking with the best arm so far is exploitation. A sampled policy does some of both automatically, as long as no arm's probability has collapsed to zero.

In code: Bandit hides the win chances and pays out one pull at a time; expected_reward is the formula above.

Everyday picture A football coach reviews the tape after a match. For every play that led to a goal, they tell the team "a bit more of that"; for plays that went nowhere, nothing. They don't need to know why the play worked. Repeat over hundreds of matches and the team drifts towards the plays that score.

Tiny worked example Start with logits (0, 0, 0), so each arm has probability ⅓. The agent pulls C and wins: reward 1. The rule, explained next, says: add to each logit learning rate × reward × (1 if it's the chosen arm, else 0, minus that arm's probability).

One lucky pull moved C from 33% to 45%. Had the pull paid 0, nothing would have moved. Had the agent pulled A and won (A wins sometimes too), A would have gone up instead. The rule is noisy, one pull at a time, but on average the arm that wins most gets pushed up most.

flowchart LR
  L[Logits θ] --> S[softmax<br/>probabilities π]
  S -->|sample| A[Action a]
  A --> ENV[Pull the arm] --> R[Reward R]
  S --> G["∇ log π(a)<br/>= one-hot(a) − π"]
  A --> G
  G --> M["× R × learning rate"]
  R --> M
  M -->|add to| L

The log-probability trick, decoded

We want the gradient of J: for each logit, how much J rises if the logit rises a little (see primer.notation for gradients from scratch). The difficulty is that J is an average over actions we can only sample. The trick rewrites the gradient as an average too, so a sample estimates it:

Why is this true? Because the slope of a probability equals the probability times the slope of its log (, the chain rule applied to log). So $\nabla J = \sum_a \nabla\pi(a) R(a) = \sum_a \pi(a) \nabla\log\pi(a) R(a)\pi(a)$ is an average over samples from π. The update is then plain gradient ascent: , with learning rate (primer.ml.optimizers).

With the numbers: at the uniform policy the true gradient is (−0.1, 0, +0.1): push C up, A down, leave B (which pays exactly the average) alone. One sample, "pulled C, got 1", estimates it as 1 × (−⅓, −⅓, ⅔). The step with α = 0.5 gives logits (−0.167, −0.167, 0.333) and probabilities (0.274, 0.274, 0.452), as in the table.

In Python:

import math
z = [0.0, 0.0, 0.0]
pi = [math.exp(z_k) / sum(math.exp(v) for v in z) for z_k in z]
# the true gradient: Σ_a π(a) R(a) (1[k=a] − π(k)), for each logit k
R = [0.2, 0.5, 0.8]
true_grad = [sum(pi[a] * R[a] * ((k == a) - pi[k]) for a in range(3)) for k in range(3)]
[round(g, 3) for g in true_grad]  # → [-0.1, 0.0, 0.1]
# one sample: pulled C (a = 2), reward 1
a, reward, alpha = 2, 1.0, 0.5
grad_log_pi = [(k == a) - pi[k] for k in range(3)]
[round(g, 3) for g in grad_log_pi]  # → [-0.333, -0.333, 0.667]
z = [z_k + alpha * reward * g for z_k, g in zip(z, grad_log_pi)]
[round(z_k, 3) for z_k in z]  # → [-0.167, -0.167, 0.333]
[round(math.exp(z_k) / sum(math.exp(v) for v in z), 3) for z_k in z]  # → [0.274, 0.274, 0.452]

This algorithm is called REINFORCE (Williams, 1992). Run it for 500 pulls and the policy finds arm C:

Why it matters in practice. A language model is exactly this kind of policy, with a vocabulary of tokens as its arms, and one sampled answer is a string of sampled tokens. REINFORCE applies unchanged: sum the log probabilities of every token in the answer, and scale the gradient by the answer's reward. Every method below (PPO, GRPO) is REINFORCE with repairs.

In code: grad_log_prob is one-hot minus the probabilities, reinforce_step is one update, worked_reinforce_step is the table above, and train_reinforce runs the whole loop against a Bandit.

Everyday picture A teacher whose class all scores between 90 and 100 learns nothing by being told "you got a 92". What matters is whether 92 is above or below the class average. Raw scores that are all large and positive make every attempt look good; only the difference from typical tells you which way to go.

Tiny worked example Two arms, a 50/50 policy, and every pull pays a lot: arm 1 always pays 10, arm 2 always pays 12. Arm 2 is better, so the logit of arm 2 should rise. Look at the REINFORCE estimate for that logit:

Without a baseline the estimate is −5 or +6 depending on the coin flip. It averages to +0.5, the right answer, but any single sample points the wrong way half the time, and violently. Subtract the average reward, 11, first and every sample says +0.5. Same average, no noise at all.

Why is it allowed? Because the baseline's contribution averages to zero: $\mathbb{E}[, b, \nabla \log \pi(a)] = b \sum_a \nabla \pi(a) = b, \nabla \sum_a \pi(a) = b, \nabla 1 = 0$. Probabilities always add to 1, so pushing all of them up is impossible; the baseline only removes noise, never signal.

With the numbers: without a baseline the estimate's variance (the average squared distance from its mean, see primer.notation) is (25 + 36)/2 − 0.5² = 30.25. With b = 11 it is 0.

In Python:

# arm 1 pays 10, arm 2 pays 12; the 1[arm 2] − π(arm 2) factor for each pull
pulls = [(10, -0.5), (12, +0.5)]
def mean_and_variance(b):
    estimates = [(reward - b) * direction for reward, direction in pulls]
    mean = sum(estimates) / 2
    return mean, sum((e - mean) ** 2 for e in estimates) / 2
mean_and_variance(b=0)  # → (0.5, 30.25)
mean_and_variance(b=11)  # → (0.5, 0.0)

flowchart LR
  R[Reward R] --> MINUS["R − b"]
  B["Baseline b<br/>average reward so far"] --> MINUS
  MINUS --> ADV{Advantage A}
  ADV -->|positive: better than typical| UP[make the action<br/>more likely]
  ADV -->|negative: worse than typical| DOWN[make the action<br/>less likely]

To see it matter, give every arm of our bandit 5 extra points: rewards are now 5 or 6 instead of 0 or 1, and nothing about which arm is best has changed.

Why it matters in practice. Every practical policy-gradient method uses a baseline. PPO learns one with a second network, the value network (or critic), which predicts the expected reward from each state. GRPO, below, gets one for free by comparing several answers to the same prompt.

In code: gradient_estimate_stats computes the exact mean and variance of the estimate, sampled_gradient_variance measures it from real pulls, and train_reinforce subtracts a running average when baseline=True.

Everyday picture A chef tests a new recipe on one evening's diners. It would be wasteful to use their comments for just one small tweak, so the chef makes several rounds of changes from the same comment cards. But the further the recipe drifts from what the diners actually ate, the less their comments apply, so the chef caps each change: never more than 20% more or less of any ingredient per round.

For a language model, sampling answers is the expensive part (every answer is a full generation), so PPO (Proximal Policy Optimization) reuses each batch of answers for several gradient steps. It needs a way to tell how far the policy has moved since the batch was sampled, and a brake.

Tiny worked example The probability ratio compares the policy now with the policy that generated the sample. A ratio of 1.5 means the current policy is 50% more likely to produce that answer than when it was sampled. With a clip range ε = 0.2 the ratio is allowed to count only between 0.8 and 1.2:

The slope of ρ·A is exactly REINFORCE's gradient scaled by ρ (because ), so inside the band PPO is REINFORCE with importance weighting. Outside it, the clipped term is flat: that sample stops pushing.

With the numbers: the four rows of the table, computed:

In Python:

def clip(x, lo, hi):
    return max(lo, min(x, hi))
def L_clip(rho, A, eps=0.2):
    return min(rho * A, clip(rho, 1 - eps, 1 + eps) * A)
[round(L_clip(rho, A), 2) for rho, A in [(1.1, 3), (1.5, 2), (0.5, -1), (1.5, -1)]]  # → [3.3, 2.4, -0.8, -1.5]
# past 1 + ε with a positive advantage, a higher ratio earns nothing more
L_clip(1.3, 2) == L_clip(1.6, 2)  # → True

flowchart TB
  OLD[Policy π_old] -->|generate a batch| B[Answers + rewards]
  B --> V[Value network<br/>predicts expected reward]
  V --> ADV[Advantages A_t]
  B --> ADV
  ADV --> LOOP
  subgraph LOOP["Several passes over the same batch"]
    RATIO["ratio ρ = π_θ / π_old"] --> CLIP["clip to 1 ± ε, take the min"]
    CLIP --> KL["subtract β × KL to the reference model"]
    KL --> STEP[gradient step on θ]
    STEP --> RATIO
  end
  LOOP -->|π_θ becomes the new π_old| OLD

To see the clip work, take one batch of 16 pulls from the bandit (C won all four of its pulls, A lost all four) and make 50 passes over it:

In code: ppo_clipped_objective is the formula, clipped_policy_update makes several passes over one batch with or without the clip, and ppo_drift_experiment is the 50-pass comparison.

The KL leash: stay close to where you started

Everyday picture A dog on a long leash can explore, but it can't run off a cliff. In RL for language models, the leash ties the policy to a frozen copy of the model it started from, the reference model (usually the model after supervised fine-tuning).

Tiny worked example An answer earns reward 1.0. The policy now gives it probability 0.6; the reference gave it 0.3. With leash strength β = 0.1, the reward actually used for training is 1.0 − 0.1 × ln(0.6 / 0.3) = 1.0 − 0.1 × 0.693 = 0.931. The policy pays a small fee for having doubled that answer's probability.

import math
R, beta = 1.0, 0.1
pi, pi_ref = 0.6, 0.3
# R' = R − β log(π / π_ref)
round(R - beta * math.log(pi / pi_ref), 3)  # → 0.931
round(R - beta * math.log(0.15 / pi_ref), 3)  # → 1.069

Why it matters in practice. This is the "penalty for drifting" in the RLHF loop of primer.ml.training_stages, and the same β appears in DPO. It stops the policy from forgetting fluent language while it chases reward, and it is the first line of defence against reward hacking, below.

In code: kl_penalised_reward is R′; clipped_policy_update adds the leash's gradient when given a reference policy and a β, and primer.ml.training_stages.kl_divergence computes the KL itself.

Everyday picture Instead of hiring an examiner to predict how hard each exam question is, a teacher gives the same question to eight students and marks each answer relative to the others on that question. On an easy question, getting it right is expected and earns little credit; on a hard one, the only right answer stands out.

PPO's baseline comes from the value network, a second model as large as the policy that has to be trained alongside it. GRPO (Group Relative Policy Optimization) throws the value network away. For each prompt it samples a group of answers and uses the group's own average as the baseline.

Tiny worked example The prompt is "3 + 4 =". The model samples four answers and a checker scores them 1 if the answer is 7, else 0.

A lone right answer in a mostly wrong group earns a big advantage: it's rare, so it's strong evidence. A group that is all right (or all wrong) earns nothing: there is no contrast, so there is nothing to learn from.

import statistics
def advantages(r):
    mu, sd = statistics.mean(r), statistics.pstdev(r)
    return [round((r_i - mu) / sd, 2) if sd else 0.0 for r_i in r]
advantages([1, 0, 0, 1])  # → [1.0, -1.0, -1.0, 1.0]
advantages([1, 0, 0, 0])  # → [1.73, -0.58, -0.58, -0.58]
advantages([1, 1, 1, 1])  # → [0.0, 0.0, 0.0, 0.0]

The rest of GRPO is PPO: the same ratio, the same clip, the same KL leash to a reference model, averaged over the group. (Some implementations divide by the sample standard deviation, with G − 1, rather than the population one; the idea is identical.)

flowchart LR
  subgraph PPO["PPO"]
    P1[Prompt] --> A1[one answer]
    A1 --> RM1[reward]
    A1 --> VN[value network<br/>a second big model]
    RM1 --> AD1[advantage = reward − value]
    VN --> AD1
  end
  subgraph GRPO["GRPO"]
    P2[Prompt] --> G1[answer 1] & G2[answer 2] & G3[answer ...] & G4[answer G]
    G1 & G2 & G3 & G4 --> VER[verifier<br/>checks each answer]
    VER --> NORM[normalise within the group<br/>mean and std]
    NORM --> AD2[advantage per answer]
  end

Verifiable rewards, and why GRPO trains reasoning

A verifiable reward comes from a check that can't be argued with: does the arithmetic equal 7, do the unit tests pass, does the proof check. No learned reward model, so no learned blind spots (see reward hacking, next). Here is GRPO on a toy "language model" that answers eight addition prompts with a single digit token. It starts out about 23% accurate and leans towards off-by-one mistakes.

Why it matters in practice. This recipe, a verifier on the final answer plus GRPO, is how DeepSeek-R1-Zero learned to reason: rewarded only for correct final answers (and a required format), the model learned by itself to write longer chains of thought, to check its work and to back up from mistakes, because those behaviours raised the chance of a correct final answer. Each token in a long chain of thought shares its answer's advantage, so the whole chain is reinforced or discouraged together. See primer.ml.reasoning for what that training produces.

In code: group_advantages is the formula, verify is the checker, pretrained_logits is the weak starting model, and train_grpo runs the loop.

Everyday picture A school pays tutors by the number of pages of homework feedback they write. Feedback gets longer, not better. An RL agent is the most literal-minded employee imaginable: it optimises the number you wrote down, not the thing you meant. When those two differ, it finds the difference. This is reward hacking (also called specification gaming), and it is Goodhart's law in code: when a measure becomes a target, it ceases to be a good measure.

Tiny worked example The goal is a good answer, and good answers here are about 4 sentences long. People rated answers of 1 to 4 sentences, and on those, longer really was better. A reward model fitted to their ratings learns a straight line: "each sentence is worth 0.19 points". Nobody ever rated a 10-sentence answer, so nobody told the reward model it was bad:

The reward model's favourite answer, 10 sentences, is the worst one.

n = [1, 2, 3, 4]
q = [1 - ((n_i - 4) / 4) ** 2 for n_i in n]
q  # → [0.4375, 0.75, 0.9375, 1.0]
n_bar, q_bar = sum(n) / 4, sum(q) / 4
w1 = sum((a - n_bar) * (b - q_bar) for a, b in zip(n, q)) / sum((a - n_bar) ** 2 for a in n)
w0 = q_bar - w1 * n_bar
(w0, w1)  # → (0.3125, 0.1875)
# the reward model's score for a 10-sentence answer, and its true quality
(w0 + w1 * 10, 1 - ((10 - 4) / 4) ** 2)  # → (2.1875, -1.25)

flowchart LR
  GOAL[What we want<br/>helpful answers] -->|people rate a sample| DATA[Ratings]
  DATA -->|fit| RM[Reward model<br/>a proxy for the goal]
  RM -->|reward| OPT[RL optimiser]
  OPT --> POL[Policy]
  POL -->|drifts to where<br/>proxy and goal disagree| GAP[Gap:<br/>high reward, low quality]
  KL[KL leash] -.->|limits drift| POL
  VER[Verifiable reward] -.->|no learned gap| OPT

Now train the starting model (which writes 2 or 3 sentences, a little too short) against each reward:

The right-hand panel uses a closed form: the best policy under a KL leash has an exact formula.

import math
ref, R = [0.5, 0.5], [1.0, 0.0]
def best_policy(beta):
    w = [p * math.exp(r / beta) for p, r in zip(ref, R)]
    return [round(w_a / sum(w), 5) for w_a in w]
best_policy(1.0)  # → [0.73106, 0.26894]
best_policy(0.1)  # → [0.99995, 5e-05]
best_policy(100.0)  # → [0.5025, 0.4975]

This is the same formula DPO starts from (primer.ml.training_stages): it is why β means the same thing in RLHF and DPO.

Why it matters in practice. Reward hacking shows up wherever RL does. A boat-racing game agent that learned to circle forever collecting bonus targets instead of finishing the race. RLHF'd chat models that learned long, confident, flattering answers score well with raters (length bias and sycophancy). Coding models rewarded for passing tests that learned to edit or special-case the tests. The defences, strongest first:

1. Verifiable rewards where the task allows: run the tests, check the answer. A check has no learned blind spot (though a buggy check does).
2. Better reward models: rate the policy's current outputs and retrain, so the reward model sees the regions the policy is exploring; use ensembles; penalise known exploits such as length directly.
3. A KL leash to keep the policy near the data the reward was fit on.
4. Watch a held-out measure of the real goal (human review, a separate evaluation set, see primer.agents.evals) and stop when it turns down, even if the reward is still climbing.

In code: true_quality is the goal, fit_reward_model and proxy_reward are the flawed reward, optimise_lengths trains against either with an optional leash, and kl_regularised_optimum with leash_sweep gives the closed-form best policy for each β.

Answer each one out loud or on paper before you open it. If you can explain it, you know it.