GRPO from scratch: implement the RL algorithm behind DeepSeek-R1 in ~200 lines of PyTorch

Why GRPO, and why now

Supervised fine-tuning has a ceiling built into it: a model trained to imitate demonstrations can never systematically exceed the quality of those demonstrations. For years that was fine — imitation got us chatbots. But reasoning broke the pattern. You can't imitate your way to a correct proof if none of your training data contains one.

The breakthrough, now called RLVR — reinforcement learning with verifiable rewards — was embarrassingly simple in hindsight: for problems where the answer is checkable (math, code against unit tests, games with rules), you don't need demonstrations at all. You need exploration plus a checker. Sample many candidate answers, keep the ones that check out, and the model discovers reasoning strategies nobody demonstrated. DeepSeek-R1-Zero took this to the extreme: pure RL starting from a base model, no SFT, with rule-based accuracy and format rewards — and reasoning emerged.

GRPO is the algorithm that made this cheap enough to run at scale. Classical PPO needs a critic — a second value network as large as the policy — to estimate how good each answer is relative to average. GRPO deletes the critic entirely: for each prompt, sample a group of completions, and let the group be its own baseline. A completion that scores above its group's average gets reinforced; one below gets suppressed. No value network, no extra parameters, no second training loop. That single deletion is why GRPO, not PPO, became the default post-training algorithm for reasoning models.

What you'll need

• Python 3.10+ with PyTorch installed (CPU build is fine — pip install torch --index-url https://download.pytorch.org/whl/cpu). The reference run used torch 2.14.1 on CPU.
• About ten minutes, most of it the training loop. No GPU, no API keys, no downloads beyond PyTorch.
• One file. Everything below concatenates into a single runnable script, top to bottom. I'll flag the few places where I simplified the real algorithm so you know exactly what the toy leaves out.

Step 1 — Build a world with a checkable answer

GRPO needs a reward function, and the whole RLVR bet is that the reward can be computed, not learned. Our toy world is single-digit addition: the prompt is "3+4=", the model writes digits, and the reward is 1 if the answer equals the true sum, else 0. This is the direct ancestor of R1's accuracy reward — a program checking an answer, not a neural network guessing at quality.

"""Minimal GRPO from scratch: teach a tiny char-level transformer single-digit
addition using only a verifiable reward. Runs on CPU in minutes.

Pipeline:
  1. Build the toy world (addition problems, train/held-out split).
  2. Define the policy: a tiny causal transformer over 13 characters.
  3. SFT warm-start so the policy is imperfect but non-random.
  4. GRPO: sample groups, score with a rule-based reward, group-relative
     advantages, clipped surrogate objective + KL penalty to the SFT reference.
  5. Compare against continued-SFT on identical data/steps; plot curves.
"""
import random
import time

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(7)
random.seed(7)
torch.set_num_threads(2)


# --- 1. The world: single-digit addition, with a checkable answer ---
CHARS = "0123456789+=<"          # '<' is the end-of-answer token
stoi = {c: i for i, c in enumerate(CHARS)}
itos = {i: c for c, i in stoi.items()}
V = len(CHARS)
EOS = stoi["<"]

def encode_prompt(a, b):          # "3+4=" -> token ids
    return [stoi[c] for c in f"{a}+{b}="]

def encode_target(a, b):          # "7<" -> answer + end token
    return [stoi[c] for c in f"{a + b}<"]

def build_data():
    pairs = [(a, b) for a in range(10) for b in range(10)]
    random.shuffle(pairs)
    return pairs[:80], pairs[80:]          # train, held-out

def reward(gen_ids, a, b):        # the whole trick: answers are checkable
    s = "".join(itos[i] for i in gen_ids if i != EOS)
    return 1.0 if s == str(a + b) else 0.0

# --------------------------------------------------------------------------

Note what this buys us: the reward function is three lines, deterministic, and unhackable in the ways learned reward models are hackable. (When your reward is learned, models find loopholes in it instead of solving your task — that's reward hacking, and it's worth understanding before you scale any of this up.)

Step 2 — A tiny policy that writes text

The "policy" π_θ is just a language model: it reads the prompt tokens and generates answer tokens. Ours is a 2-layer, 4-head character transformer, about 103k parameters — small enough to train on a laptop CPU in minutes, large enough to actually learn the toy. Two helpers round it out: sample(), which generates completions (with temperature, for exploration), and accuracy(), which measures greedy-decoding performance on a dataset.

# 2. The policy: a tiny causal transformer (char-level language model)
# --------------------------------------------------------------------------
class Block(nn.Module):
    def __init__(self, d, n_head):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, n_head, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, h):
        T = h.size(1)
        causal = torch.triu(torch.ones(T, T, dtype=torch.bool, device=h.device), 1)
        a, _ = self.attn(self.ln1(h), self.ln1(h), self.ln1(h), attn_mask=causal)
        h = h + a
        return h + self.mlp(self.ln2(h))

class TinyGPT(nn.Module):
    def __init__(self, n_layer=2, n_head=4, d=64):
        super().__init__()
        self.tok, self.pos = nn.Embedding(V, d), nn.Embedding(16, d)
        self.blocks = nn.ModuleList([Block(d, n_head) for _ in range(n_layer)])
        self.ln = nn.LayerNorm(d)
        self.head = nn.Linear(d, V, bias=False)

    def forward(self, x):
        t = torch.arange(x.size(1), device=x.device)
        h = self.tok(x) + self.pos(t)
        for b in self.blocks:
            h = b(h)
        return self.head(self.ln(h))

def n_params(m):
    return sum(p.numel() for p in m.parameters())

@torch.no_grad()
def sample(model, prompt, max_new=4, temp=1.0, greedy=False, return_logp_old=False):
    # generate a completion; optionally also return per-token log-probs
    # under the sampling policy (the GRPO importance ratio needs them)
    x = torch.tensor([prompt])
    gen, logp_old = [], []
    for _ in range(max_new):
        logits = model(x)[:, -1, :]
        if greedy:
            nxt = logits.argmax(-1, keepdim=True)
        else:
            lp = F.log_softmax(logits / temp, dim=-1)
            nxt = torch.multinomial(lp.exp(), 1)
            logp_old.append(lp.gather(1, nxt).item())
        x = torch.cat([x, nxt], dim=1)
        gen.append(nxt.item())
        if nxt.item() == EOS:
            break
    if return_logp_old:
        return gen, logp_old
    return gen

@torch.no_grad()
def accuracy(model, pairs):
    return sum(reward(sample(model, encode_prompt(a, b), greedy=True), a, b)
               for a, b in pairs) / len(pairs)

# --------------------------------------------------------------------------

Step 3 — Warm-start with SFT, and find its ceiling

Almost nobody runs RL from a randomly initialized policy — exploration from pure noise is hopeless, because the model must stumble onto a correct answer before the reward can teach it anything. The standard recipe (followed by every reasoning lab) is: supervised fine-tuning first, to teach the format of an answer, then RL to teach correctness. Our SFT phase trains on the 80 labeled pairs for eighteen short epochs (batched by answer length, with the prompt tokens masked out of the loss):

# 3. SFT warm-start (imperfect: good enough to sample some correct answers)
# --------------------------------------------------------------------------
def sft_train(model, pairs, epochs, lr=3e-3):
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    # bucket by answer length so no padding is needed (answers are 1-2 digits)
    buckets = {}
    for a, b in pairs:
        ids = encode_prompt(a, b) + encode_target(a, b)
        buckets.setdefault(len(ids), []).append(ids)
    for _ in range(epochs):
        for bucket in buckets.values():
            random.shuffle(bucket)
            x = torch.tensor(bucket)
            inp, tgt = x[:, :-1], x[:, 1:].clone()
            tgt[:, :3] = -100                      # don't train on the prompt "d+d="
            loss = F.cross_entropy(model(inp).reshape(-1, V), tgt.reshape(-1),
                                   ignore_index=-100)
            opt.zero_grad(); loss.backward(); opt.step()

# --------------------------------------------------------------------------

train_pairs, held_pairs = build_data()
model = TinyGPT()
print(f"policy: {n_params(model):,} params, vocab={V}")
sft_train(model, train_pairs, epochs=18)
a_tr, a_ho = accuracy(model, train_pairs), accuracy(model, held_pairs)
print(f"after SFT warm-start: train={a_tr:.2f} heldout={a_ho:.2f}")

Measured result: 17.5% train accuracy, 0% held-out. Eighteen epochs teach the model that something goes after the equals sign, but mostly it collapses to memorized guesses — at this point it answers nearly everything with "12". This is the ceiling imitation gives us on 80 pairs, and it's the floor RL has to beat. Freeze a copy of this model now: it becomes π_ref, the reference policy that will keep GRPO honest.

Step 4 — The GRPO update, in three moving parts

Here is the whole algorithm. For a batch of prompts, sample G completions per prompt from the current policy, score each with the reward function, and compute the group-relative advantage — the z-score of each reward within its group:

Âi = (ri − μ) / σ  — computed over the G completions for one prompt.

A completion that beat its group's average gets a positive advantage; one that lost gets negative. Then the policy update is PPO's clipped surrogate objective on those advantages, minus a KL penalty against the frozen reference model:

# 4. GRPO: groups, verifiable rewards, clipped objective, KL penalty
# --------------------------------------------------------------------------
def grpo_step(model, ref, opt, batch_pairs, G=8, eps=0.2, beta=0.02, temp=1.0):
    # --- rollout: G completions per prompt, scored by the rule-based reward ---
    seqs, masks, advs, old_lp, all_r = [], [], [], [], []
    for a, b in batch_pairs:
        prompt = encode_prompt(a, b)
        gens, lps, rs = [], [], []
        for _ in range(G):
            g, lp = sample(model, prompt, temp=temp, return_logp_old=True)
            gens.append(g); lps.append(lp); rs.append(reward(g, a, b))
        all_r.extend(rs)
        r = torch.tensor(rs)
        adv = (r - r.mean()) / (r.std() + 1e-8)   # group-relative advantage
        for g, lp, av in zip(gens, lps, adv.tolist()):
            full = prompt + g
            seqs.append(full)
            m = [0.0] * len(prompt) + [1.0] * len(g)   # loss only on completion
            masks.append(m); old_lp.append([0.0] * len(prompt) + lp)
            advs.append(av)
    L = max(len(s) for s in seqs)
    x = torch.tensor([s + [EOS] * (L - len(s)) for s in seqs])
    mask = torch.tensor([m + [0.0] * (L - len(m)) for m in masks])
    old = torch.tensor([o + [0.0] * (L - len(o)) for o in old_lp])
    adv_t = torch.tensor(advs).unsqueeze(1)

    # --- objective: clipped surrogate + KL(ref) ---
    logits = model(x[:, :-1])
    logp = F.log_softmax(logits, dim=-1)
    # log-prob of the actually generated token at each completion position
    gen_tok = x[:, 1:]
    new_lp = logp.gather(2, gen_tok.unsqueeze(-1)).squeeze(-1)
    ratio = torch.exp(new_lp - old[:, 1:])
    surr = torch.min(ratio * adv_t, torch.clamp(ratio, 1 - eps, 1 + eps) * adv_t)
    with torch.no_grad():
        ref_lp = F.log_softmax(ref(x[:, :-1]), dim=-1)
    kl = (ref_lp.exp() * (ref_lp - logp)).sum(-1)   # exact KL, vocab is tiny
    loss = (-(surr * mask[:, 1:]).sum() + beta * (kl * mask[:, 1:]).sum()) \
        / mask[:, 1:].sum().clamp_min(1)
    opt.zero_grad(); loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    return sum(all_r) / max(1, len(all_r))

# --------------------------------------------------------------------------

Three moving parts, each doing one job. (Implementation notes for the careful reader: completions have different lengths, so they're padded into one batch and the loss is masked to completion tokens only — the prompt tokens never receive gradient. The per-token log-probs under the sampling policy are stored at rollout time, which is the PPO-correct pattern for the importance ratio; gradients are norm-clipped at 1.0.)

1. Group-relative advantages — the critic deletion. PPO asks "how much better than expected was this answer?", where "expected" comes from a learned value network. GRPO asks "how much better than its siblings was this answer?" — a z-score inside the sampled group. Same variance-reduction effect, zero extra parameters. This is the entire intellectual content of the "Group Relative" in GRPO.
2. The clipped surrogate — the trust region. The ratio π_θ/π_old measures how far the update wants to move. Clipping it to [1−ε, 1+ε] means a single lucky group can't yank the policy off a cliff. This is inherited straight from PPO, and it's what makes online RL stable enough to run. One honest footnote: because our toy takes a single gradient step per rollout batch, the stored "old" log-probs still match the current policy when the loss is computed — the ratio is exactly 1 and the clipping never binds here. It becomes load-bearing the moment you do multiple epochs per batch, as TRL does. The term is included because it belongs to the objective, not because this toy exercises it.
3. The KL anchor — don't forget what SFT taught you. The β·KL(π_θ‖π_ref) term penalizes drifting too far from the frozen warm-start model. Without it, RL happily destroys the answer format to chase reward quirks; with it, exploration stays in the neighborhood of sensible outputs. (Our vocab is 13 tokens, so we compute the KL exactly; at LLM scale you'd use an unbiased estimator like the k3 one from the PPO literature.)

Two simplifications to be upfront about: we take one gradient step per rollout batch (PPO-style training does several epochs over the same rollouts), and we use outcome supervision — one reward per completion, broadcast to all its tokens — exactly the mode the DeepSeekMath paper used for math. The shape of the real algorithm is all here.

Step 5 — Train it and measure everything

Now the experiment. From the same warm-started model, run two branches: 160 GRPO steps (reward only — no new labels), and 160 more SFT steps (labels) as the honest baseline. Greedy accuracy is logged every 10 steps on both the training pairs and the 20 held-out pairs:

# --- 5. The experiment: GRPO (reward only) vs continued SFT (labels) ---
t0 = time.time()
ref = TinyGPT(); ref.load_state_dict(model.state_dict())  # frozen reference
for p in ref.parameters(): p.requires_grad_(False)

# GRPO branch
gmodel = TinyGPT(); gmodel.load_state_dict(model.state_dict())
gopt = torch.optim.AdamW(gmodel.parameters(), lr=5e-4)
gh_tr, gh_ho, g_rew = [a_tr], [a_ho], []
STEPS, BS, G = 160, 16, 8
for step in range(STEPS):
    batch = random.sample(train_pairs, BS)
    mr = grpo_step(gmodel, ref, gopt, batch, G=G)
    g_rew.append(mr)
    if (step + 1) % 10 == 0:
        gh_tr.append(accuracy(gmodel, train_pairs)); gh_ho.append(accuracy(gmodel, held_pairs))
        print(f"grpo step {step+1:3d}: mean_reward={sum(g_rew[-10:])/10:.2f} "
              f"train={gh_tr[-1]:.2f} heldout={gh_ho[-1]:.2f}", flush=True)

# continued-SFT branch: same data, same number of optimizer steps
smodel = TinyGPT(); smodel.load_state_dict(model.state_dict())
sh_tr, sh_ho = [a_tr], [a_ho]
for step in range(STEPS):
    sft_train(smodel, random.sample(train_pairs, BS), epochs=1, lr=5e-4)
    if (step + 1) % 10 == 0:
        sh_tr.append(accuracy(smodel, train_pairs)); sh_ho.append(accuracy(smodel, held_pairs))
        print(f"sft  step {step+1:3d}: train={sh_tr[-1]:.2f} heldout={sh_ho[-1]:.2f}", flush=True)

print(f"\nFINAL  GRPO: train={gh_tr[-1]:.2f} heldout={gh_ho[-1]:.2f}")
print(f"FINAL    SFT: train={sh_tr[-1]:.2f} heldout={sh_ho[-1]:.2f}")
print(f"elapsed {time.time()-t0:.0f}s")

That's the complete script — concatenate the five blocks and run it. On the reference machine (a 2-core CPU, torch 2.14.1) it finished in 337 seconds.

Reading the curves: what actually happened

The left panel is the money chart: mean verifiable reward per GRPO step climbs from 0.08 to 0.42. The policy produces correct answers more than five times as often at the end as at the start — learned from a reward function alone, with no new labels, no critic, no human feedback of any kind.

The right panel tells the fuller story. Greedy accuracy on the 80 training pairs goes 17.5% → 50% under GRPO (40 of 80 pairs correct, up from 14). Continued SFT — the same warm start plus 160 more steps of labeled data — reaches 71%. And the held-out curves sit at essentially zero for both methods.

Be precise about what each of those facts means:

• GRPO works. 27 of the 80 training pairs flipped from wrong to right under pure reward signal (one pair regressed — RL isn't magic, it's stochastic hill-climbing). Before and after, greedy decoding, on training prompts:

  Prompt	SFT warm-start	After GRPO
  2+5	"12" ✗	"7" ✓
  7+8	"12" ✗	"15" ✓
  8+2	"12" ✗	"10" ✓
  2+4	"12" ✗	"6" ✓
  2+6	"11" ✗	"8" ✓
  4+4	"12" ✗	"8" ✓
  3+3	"10" ✗	"5" ✗
  9+9	"13" ✗	"15" ✗
  8+1	"1" ✗	"10" ✗
  2+1	"1" ✗	"6" ✗

• Labels are still the most sample-efficient signal — when you have them. Continued SFT beat GRPO 71% to 50% on the training set. This is exactly why the real pipeline is SFT then GRPO, not GRPO instead of SFT: imitation is the fastest way to absorb what you can demonstrate, and RL is what you reach for when the demonstrations run out — or never existed, as in R1-Zero.
• Neither method generalizes here, and that's honest. Held-out accuracy is ~0% for both. A 103k-parameter model on 80 memorized pairs memorizes; it does not discover the algorithm of addition. The demo proves the mechanism — the loop learns from rewards — not generalization. At LLM scale the same loop generalizes because the base model already knows arithmetic from pretraining; RL reshapes reasoning behavior, not facts.

From this toy to a real reasoning model

The jump from this script to production GRPO is engineering and scale, not a different idea. The reference implementation most people actually run is Hugging Face's GRPOTrainer in TRL, and its knobs map one-to-one onto what you just built:

from trl import GRPOConfig, GRPOTrainer

config = GRPOConfig(
    output_dir="grpo-reasoning",
    loss_type="grpo",            # the DeepSeekMath objective from Step 4
    beta=0.001,                  # KL penalty vs the reference model
    epsilon=10.0,                # DAPO-style clip-higher bound
    num_generations=16,          # G: completions sampled per prompt
    max_completion_length=32768, # long chains of thought need room
    num_train_epochs=1,
    per_device_train_batch_size=2,
)
trainer = GRPOTrainer(
    model=model,                                   # your SFT warm-start
    reward_funcs=[accuracy_reward, format_reward], # verifiable rewards, à la R1
    args=config,
    train_dataset=prompts,                         # prompts only — no answers needed
)
trainer.train()

What changes at scale: G grows (16+ generations per prompt), completions get long (tens of thousands of tokens of chain-of-thought), rewards get richer (accuracy plus format plus sometimes tool-use checks), and the KL uses an unbiased estimator instead of the exact computation. Later refinements you'll see cited — DAPO's dynamic sampling (skip groups with zero reward variance, since they carry no signal) and Dr. GRPO's removal of length normalization — are patches on this same loop, not replacements.

Which approach should you use?

GRPO is one tool in the post-training drawer. Pick by what supervision you actually have:

The frontier recipe, as DeepSeek ran it for R1, stacks them: SFT to teach format, GRPO to teach reasoning, then further rounds of SFT and GRPO on the improved model. Each stage covers what the previous one can't.

The takeaway

Reasoning models weren't taught to reason. They were selected for reasoning: generate widely, check rigorously, reinforce what worked. GRPO is the selection machinery, and its core idea fits in one sentence — sample a group, keep what beat the group's average, stay close to what you already knew. No critic, no value network, no learned reward model. You just implemented it in ~200 lines, watched a reward curve climb 5× from nothing but a three-line checker, and measured exactly where imitation ends and reinforcement begins. The next time someone says "R1 learned to reason with RL," you'll know precisely which loop they mean — because you've run it.

---

The complete script, the plotting code, and the raw run metrics for every number in this article are described exactly as run above — torch 2.14.1 on CPU, seeds fixed at 7, 337 seconds wall-clock. If you extend the toy (larger G, longer training, a harder world like multi-digit arithmetic), the interesting question is the same one the labs ask: at what point does the held-out curve start moving?