GRPO from Scratch — Training a Reasoning Policy End to End

The Problem

Lesson 09/13 built GRPO around a tabular token policy. Every state (target, position, numbers used so far) carried its own logits. That design is perfect for pedagogy: every symbol in the objective maps to a named dictionary entry, and the gradient update is one line. But it is not how engineers ship reasoning models.

A table cannot generalize. If the prompt target or the allowed digits change, the tabular policy has no parameters to share; it must allocate a new logits vector for every unseen prefix. A real LLM is a single function approximator: the same weights produce sensible next-token distributions for prompts the model has never seen, because representations for "prefix contains a +" and "target is large" are reused across examples.

A monolithic script does not survive production. In the reference GRPO codebase, the verifier is a separate process that emits rewards, the trainer manages rollout buffers and optimization, and the evaluator runs greedy decode plus metric logging. The capstone therefore splits the same math into modules: verifier_reward, train, and evaluate. Each has a narrow contract, its own tests, and can be replaced independently.

This lesson upgrades the 09/13 recipe in two ways: a two-layer MLP policy replaces the table, and the code is organized so the pieces can move to different machines later.

The Concept

From table to function approximator

The policy is a tiny neural network with one hidden layer:

The vocabulary is {"2", "3", "5", "+", "*"}. Autoregressive decoding is identical to an LLM: at each position the model emits a distribution over the vocabulary, a grammar mask zeros out illegal tokens (a number slot cannot pick an operator, and no number may be reused), and a token is sampled or taken greedily. The input dimension is 5 + |TARGETS|, the hidden layer has 48 units, and the output dimension is 5. The whole network is trained with plain SGD.

Because the same weights score every prefix, the model can transfer structure: learning that 2*3 equals 6 for target 6 helps it avoid the same prefix when target 7 appears.

GRPO: the same objective on a functional policy

The verifier is unchanged from 09/13. It pays a format reward, an accuracy reward for hitting the target under real operator precedence, and the Kimi k1.5 length-control term:

For each prompt the trainer samples a group of G = 8 responses and normalizes their rewards into advantages:

The optimization objective is the PPO clipped surrogate plus the k_3 KL penalty against a frozen reference copy of the initialization:

The only difference from 09/13 is that \pi_\theta is now an MLP. The gradient is computed by back-propagating coeff * (onehot(token) - probs) through the two layers, where coeff = adv * rho - beta * (1 - exp(log_ref - log_new)).

Why split into modules

The code is organized into three separable concerns:

flowchart LR
    Prompt[(target)] --> Trainer[train: sample group + update policy]
    Trainer --> Verifier[verifier_reward: check answer]
    Verifier --> Adv[advantages]
    Adv --> Trainer
    Trainer --> Eval[evaluate: greedy success + curves]

• Verifier is a pure function from (target, tokens) to a scalar. It has no parameters, no gradients, and can be replaced by a unit-test runner or a JSON-schema validator without touching the trainer.
• Trainer owns the rollout buffer, the group baseline, the clipped surrogate, and the parameter update. It knows nothing about how expressions are evaluated; it only consumes rewards.
• Evaluator decodes greedily and reports success rate, mean group reward, and mean response length. It is the harness you would run every N steps in a real training job.

This split makes the ablations possible: the trainer can turn the KL penalty or the group normalization on and off without the verifier or evaluator noticing.

Build It

The implementation lives in code/main.py. It is intentionally small — the MLP, verifier, trainer, evaluator, and ablation harness fit in one file — but every GRPO symbol from the objective has a named constant at the top of the file.

Run the standard configuration:

consoleCopypython3 main.py

You will see per-iteration stats and a final greedy evaluation. The script then runs two ablations with the same seed and iteration budget and prints a comparison.

Run an ablation:

consoleCopypython3 main.py --ablation=no-kl
python3 main.py --ablation=no-group-norm

• no-kl sets beta = 0. The policy can chase rewards without any anchor to the initialization. Watch the KL to the reference grow and the format/success rate degrade compared with the standard run.
• no-group-norm replaces the per-group baseline with a global running mean of all rewards. The policy still learns, but the baseline is noisier and slower; after the same number of iterations the greedy success rate lags behind standard GRPO.

Verify It

code/tests/test_main.py locks down the pieces that make the capstone a capstone:

• the verifier respects operator precedence and rejects malformed expressions;
• group advantages are zero-mean and exactly zero when all rewards tie;
• the k_3 estimator is non-negative and zero for identical policies;
• the MLP policy is deterministic after a copy, produces well-formed greedy output, and its token log-probabilities sum to one;
• seeded end-to-end training raises greedy success rate from the cold start;
• removing the KL penalty produces measurably larger KL drift from the reference;
• removing group normalization produces lower greedy success after the same iteration budget;
• the command-line ablation flag is parsed correctly.

Why This Matters for Engineers

• Generalization is the point of a neural policy. The tabular lesson proved the math; the MLP version proves the math can share parameters across states, which is the only reason LLM-scale GRPO is tractable.
• Verifiable rewards are still the right default. Math, code, schemas, and retrievable facts should be checked by programs. Reserve learned reward models for genuinely open-ended quality.
• The KL coefficient is a safety dial, not a nuisance. no-kl shows that the optimizer will drift into higher-KL distributions even on a toy task. In production that drift is the difference between "reasoning" and "gibberish that happens to pass the checker."
• Group size is the resolution of your baseline. no-group-norm shows that a baseline computed far from the current prompt dilutes the signal. Per-group normalization is not a cosmetic normalization; it is the value-function replacement.
• Modular boundaries pay off in testing and maintenance. A verifier you can swap, a trainer you can run offline, and an evaluator you can run in CI are the shape of every production RLVR system.

Hints

• If the success rate stalls below 0.5 after 200 iterations, raise GROUP or reduce LR. Sparse binary rewards need either enough samples per prompt or a smaller step size.
• A negative mean group reward with near-zero format violations usually means the policy found a valid expression family that never hits targets. Increase BETA to pull it back toward the reference.
• Length inflation appears in the mean-response-length curve before it shows up in accuracy. The k1.5 term exists to keep that curve flat.
• In no-group-norm mode, the first few prompts seed the global baseline; if the first prompt is all-fail, every later high-reward rollout gets an oversized advantage and the policy may overshoot.

Key Terms