Policy-gradient reinforcement learning built up from REINFORCE to GRPO, in plain PyTorch with no RL library, on a toy task with a verifiable reward. The point isn't the task — it's to make the one idea every algorithm in this family shares visible, and then show exactly what each successive algorithm adds on top of it.
Every algorithm here computes the same update:
grad(J) = sum_t grad( log pi(token_t | prompt, tokens_<t) ) * A
a per-token log-probability, summed over the sequence, multiplied by one
scalar advantage A for the whole sequence. That is exactly how RLHF, RLOO
and GRPO compute gradients on real language models — the completion is a
single action (the bandit formulation: one reward, one advantage, per
sequence), and the only thing that changes between algorithms is how A is
estimated.
| Algorithm | Advantage A |
Learned critic? |
|---|---|---|
| REINFORCE | reward - mean(reward) over the batch |
no |
| Actor–Critic | reward - V(prompt) |
yes |
| PPO | reward - V(prompt), standardised, reused over clipped epochs |
yes |
| RLOO | reward_i - mean(reward_{j != i}) within a group |
no |
| GRPO | (reward_i - mean) / std within a group |
no |
Reading top to bottom is the narrative: start with the crudest baseline, add a learned one, make it reuse data safely, then drop the critic again by letting a group of samples play its role — RLOO one way, GRPO one step further. GRPO is the algorithm behind the sibling grpo-math-reasoning repo; here it's implemented from scratch rather than delegated to a trainer, so the mechanism underneath is on show.
A deliberately tiny stand-in for "language model completes a prompt, a
verifier scores it": the prompt is a target integer T, the policy emits L
digits one at a time, and the reward is 1 exactly when the digits sum to
T, else 0. Sparse, exact, and verifiable — the same shape as a
#### <number> check on a maths answer — but it trains in minutes on a CPU,
so the focus stays on the algorithms.
src/pgprog/
env.py # the digit-sum task and its reward
policy.py # shared GRU policy, critic, and rollout mechanics
algorithms/
advantages.py # the five advantage estimators, as pure functions
reinforce.py # ... and one train_step each, differing only in A
actor_critic.py
ppo.py
rloo.py
grpo.py
stats.py # bootstrap confidence intervals
scripts/
train.py # train one algorithm
compare.py # all five across seeds: learning curves + gradient variance
tests/ # pure-logic unit tests, CPU-only
uv venv && source .venv/bin/activate
uv pip install -e ".[dev]"
python scripts/train.py --algo grpo --steps 400
python scripts/compare.py --seeds 5 --steps 1500compare.py reports two things, because "which one wins" is the less
interesting question:
- Final success rate across several seeds, with bootstrap confidence intervals — so a win has to survive the seed-to-seed noise to count.
- Empirical gradient variance at a fixed policy: freeze an untrained policy, recompute each estimator's gradient from fresh rollouts many times, and measure how much the gradient-norm varies. This measures the thing the whole family is designed around — variance reduction — rather than merely asserting it.
Every algorithm sees the same number of distinct prompts per step, so the learning-curve comparison isn't confounded by prompt diversity. The group methods (RLOO, GRPO) then draw several completions per prompt, so they spend more total rollouts per step — that extra sampling is their real cost and is reported alongside the scores rather than hidden.
5 seeds, 1500 steps (400 for RLOO/GRPO's prompt-batches), seq_len=5, CPU. Same
number of distinct prompts per step across all five (BATCH = 64,
GROUP_SIZE = 8 — see What the comparison measures).
| Algorithm | Final success rate | 95% CI |
|---|---|---|
| REINFORCE | 0.499 | [0.399, 0.582] |
| Actor–Critic | 0.486 | [0.296, 0.687] |
| PPO | 0.810 | [0.754, 0.858] |
| GRPO | 0.961 | [0.926, 0.995] |
| RLOO | 0.982 | [0.960, 1.000] |
The two critic-free, group-baseline methods win clearly, with PPO a clear second — matching the ordering the field converged on for LLM RLHF, for the same reason: a per-prompt group baseline is a tighter variance reducer here than a single learned critic value.
Gradient-norm variance at a fixed, untrained policy (50 resamples, matched on total rollouts, not prompts — see above):
| Algorithm | Gradient-norm variance | Mean norm |
|---|---|---|
| REINFORCE | 0.0024 | 0.092 |
| RLOO | 0.0042 | 0.093 |
| GRPO | 0.0256 | 0.237 |
REINFORCE and RLOO produce gradients of similar magnitude and variance at
this untrained checkpoint — RLOO's advantage shows up over training, not in a
single-step snapshot, since a leave-one-out baseline only pays off once
rewards start to differ within a group. GRPO's extra / std term inflates
both the norm and its variance whenever the group's reward std is small,
which is expected: normalising by a noisy, near-zero denominator early in
training amplifies the estimate rather than stabilising it.
Three things here were bugs first and lessons second; each is documented at the point in the code where it bites.
- A categorical policy on a sparse reward collapses without an entropy
bonus. Left alone, it converges to emitting one constant token before it
ever learns to read the prompt — a stable local optimum once exploration
pressure is gone. Every
train_stepadds an annealed entropy term for exactly this reason (pgprog.policy.rollout). - PPO needs its advantages standardised, not just clipped. Early on every
completion fails, so the untrained critic makes the raw advantage
spuriously positive across the whole batch and PPO is told every action was
above average. Without per-batch standardisation PPO scores near zero at
any learning rate (
pgprog.algorithms.ppo). - PPO's importance ratio is per-token, and it needs a smaller step. A sequence-level ratio is the product of the per-token ratios, volatile enough to saturate the clip range; and reusing each batch over several epochs means the same nominal learning rate moves several times as far, so PPO uses a lower one.
MIT.