Implementing GRPO from Scratch: DeepSeek-R1's Reinforcement Learning Engine

This is the start of a series where I re-implement ideas from key AI papers from scratch — and I'm beginning with the piece that matters most in DeepSeek-R1: GRPO.

The paper introduces two models. DeepSeek-R1-Zero is trained with pure reinforcement learning — no supervised fine-tuning at all — while DeepSeek-R1 uses a multi-phase approach combining SFT and RL. What's remarkable is that the pure-RL model, DeepSeek-R1-Zero, learned to solve hard problems, catch its own errors and inconsistencies, and self-correct with explanations — behaviour that emerged from RL alone rather than being explicitly taught.

I'm keeping the scope of this implementation tight: just GRPO. And instead of a frontier-scale run, I use the distilled model DeepSeek-R1-Distill-Qwen-1.5B so the whole thing fits on modest hardware while staying faithful to the algorithm.

Title page of the DeepSeek-R1 paper: 'DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning' by DeepSeek-AI, with the abstract introducing DeepSeek-R1-Zero and DeepSeek-R1. — The DeepSeek-R1 paper — the basis for this implementation series.

A Deep Dive into GRPO

Group Relative Policy Optimization (GRPO) has become the core RL algorithm behind reasoning models like DeepSeek-R1. Its headline advantage is eliminating the massive memory and compute overhead of a traditional critic (value) network — while still producing a stable learning signal.

1. Rule-Based Deterministic Rewards

Rather than relying on a separate, subjective neural reward model — which can be biased or gamed — GRPO can learn directly from rigid, objective validation functions:

Format checking: regex or a compiler verifies that the solution is formatted correctly — e.g. reasoning wrapped in <think>...</think> and the final output in <answer>...</answer>.
Correctness validation: for math, logic, or code, a Python interpreter or unit-test environment programmatically checks the final answer against ground truth.
Reasoning alignment: the intermediate steps are expected to progress naturally toward the provided answer.

2. Group Relative Advantage

Instead of judging a single output against a baseline predicted by a critic network, GRPO generates multiple outputs for the same prompt and scores them relative to each other. The advantage for each item in the group is computed from group-normalized rewards:

text

advantage = (reward - mean(group_rewards)) / std(group_rewards)

Memory efficiency: dropping the critic network — typically the same size as the actor model — saves roughly 50% of GPU memory during training.
Self-correction: on hard prompts where every output scores poorly, subtracting the group mean means the least-bad answer still gets a positive advantage, nudging the model toward incremental improvement.

The Core Objective Function

GRPO maximizes a clipped objective adapted for groups, penalizing drift from a reference policy with a KL-divergence term.

The GRPO objective function: expectation over prompts q and a group of G sampled outputs of the mean over the group of min(ratio * advantage, clip(ratio, 1-epsilon, 1+epsilon) * advantage) minus beta times the KL divergence between the current policy and the reference policy. — The GRPO objective. Daunting at first glance, but each component plays a specific role in stabilizing learning.

Probability ratio r(θ): how much more (or less) likely the current policy is to generate a sequence compared to the old/reference policy.
Clipping mechanism (ε): stops the policy from changing too drastically in a single step when an output yields a huge positive advantage — typically ε = 0.2.
KL penalty (β·D_KL): an elastic band. With no critic to stabilize training, this keeps the active model close to the base reference so it doesn't lose its core language and formatting abilities.

Diagram comparing PPO and GRPO. PPO routes a single output o through reference, reward, and value models. GRPO drops the value model and instead samples a group of outputs o_1...o_G, scoring each to get rewards r_1...r_G, with KL computed against the reference model. — PPO vs GRPO. PPO needs a separate value model to estimate the baseline; GRPO removes it entirely and derives the baseline from the group of sampled outputs. Diagram from the DeepSeekMath paper. Resource: https://ghost.oxen.ai/why-grpo-is-important-and-how-it-works/

“For each iteration: sample a group of outputs per prompt, compute rule-based rewards, normalize them into group-relative advantages, then update the policy by maximizing the clipped ratio objective minus a KL penalty against the frozen reference model.”

✶

Building It in Pure PyTorch

The rest of this implementation walks through a minimal but complete GRPO training loop — no TRL, no black boxes. We install the dependencies and import the essentials.

bash

pip install --upgrade torchao peft transformers accelerate

python

import os
import re
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import get_peft_model, LoraConfig, TaskType

1. Setup & Configuration

We set the GRPO hyperparameters. group_size is kept at 4, and we define XML tags — <think>...</think> for reasoning and <answer>...</answer> for the final answer — so the format reward has something deterministic to check against.

python

device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"

# GRPO Hyperparameters:
GROUP_SIZE = 4       # Num of generations (group) per prompt
EPSILON_CLIP = 0.2   # PPO ratio clipping boundary
BETA_KL = 0.04       # KL divergence penalty
LEARNING_RATE = 2e-6
MAX_PROMPT_LEN = 128
MAX_GEN_LEN = 128

SYSTEM_PROMPT = (
    "Respond exactly in this format:\n"
    "<think>\n...\n</think>\n"
    "<answer>\n...\n</answer>"
)

2. A Verifiable Dataset

Datasets broadly fall into two categories: verifiable and non-verifiable. A verifiable dataset contains tasks whose correctness can be objectively validated by a deterministic function or rule — the expected outcome is unambiguous, so evaluation is fully automatable. A non-verifiable dataset needs semantic judgment, which usually means an LLM-based evaluator (or another capable judge model) has to assess the response.

I focus exclusively on a verifiable dataset here — just 5 examples, kept small because of compute limits — where automated deterministic evaluation gives a reliable performance signal.

python

class MathReasoningDataset(Dataset):
    def __init__(self):
        # Sample deterministic reasoning tasks
        self.data = [
            {"q": "If A = 5 and B = A + 3, what is B?", "a": "8"},
            {"q": "Solve for x: 2x - 4 = 10", "a": "7"},
            {"q": "What is the square root of 81?", "a": "9"},
            {"q": "A box has 3 red and 4 blue balls. Total balls?", "a": "7"},
            {"q": "What is 15 minus 6?", "a": "9"}
        ]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]
        full_prompt = f"{SYSTEM_PROMPT}\n\nProblem: {item['q']}\n"
        return {'prompt': full_prompt, 'target': item['a']}

3. Verifiable Rule-Based Rewards

With the dataset in place, we write deterministic reward functions. The scheme is deliberately simple: +0.5 for correctly following the XML format, and +1 for getting the answer right.

python

def compute_rewards(completions, targets):
    """Computes rule-based rewards for format verification and correctness."""
    rewards = []
    for completion, target in zip(completions, targets):
        score = 0.0
        # Format Reward: Strict enforcement of XML tags
        if completion.startswith('<think>') and "</think>" in completion:
            if '<answer>' in completion and completion.endswith('</answer>'):
                score += 0.5
        # Correctness Reward: Final value lookup extraction
        parts = completion.split("</answer>")
        extracted_answer = parts[-2].split("<answer>")[-1] if len(parts) >= 2 else ""
        if target.strip() in extracted_answer:
            score += 1

        rewards.append(score)
    return torch.tensor(rewards, dtype=torch.float32, device=device)

4. Log-Probability Tracking

Next we need a utility that computes the log-probability the model assigns to every generated token in a sequence. This is what lets us form the probability ratio at the heart of the GRPO objective.

python

def get_action_log_probs(model, input_ids, attention_mask, prompt_len):
    """Computes exact per-token log probabilities for generated tokens only."""
    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    logits = outputs.logits

    # Shift outputs to align each logit with the token it predicts
    shift_logits = logits[:, prompt_len - 1 : -1, :].contiguous()
    shift_labels = input_ids[:, prompt_len:].contiguous()

    log_probs = F.log_softmax(shift_logits, dim=1)
    # Gather the log-prob of each actually-generated token
    selected_log_probs = torch.gather(
        log_probs, dim=-1, index=shift_labels.unsqueeze(-1)
    ).squeeze(-1)

    # Mask out padding so only real generated tokens count
    gen_mask = attention_mask[:, prompt_len:]

    return selected_log_probs * gen_mask

5. Initialization Pipeline

Now we wire everything together. We load the tokenizer and set the padding token to the EOS token, since decoder-only LMs typically don't have a dedicated pad token.

Instead of fine-tuning the full model, we use LoRA (Low-Rank Adaptation) to make training far more efficient. Rather than updating the full weight matrix W, LoRA represents the update as W + ΔW where ΔW = BA — two small low-rank trainable matrices A and B of rank r. Only these are learned, drastically cutting the number of trainable parameters while preserving performance. Here the LoRA adapters are applied to the query (Q), key (K), value (V), and output (O) projection matrices of the attention mechanism.

We then load a second, frozen copy of the model as the reference model. It stays fixed throughout training and serves as the baseline for the policy ratio and KL divergence, ensuring the policy doesn't drift too far from its initial behaviour. Finally, we initialize an AdamW optimizer over the trainable LoRA parameters.

python

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# Load the policy network we will optimize, with LoRA adapters
policy_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map={"": device}
)
peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=8, lora_alpha=16, lora_dropout=0.05,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"]
)
policy_model = get_peft_model(policy_model, peft_config)

# Load the frozen reference model (the stable baseline)
ref_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map={"": device}
)
ref_model.eval()
for param in ref_model.parameters():
    param.requires_grad = False

optimiser = torch.optim.AdamW(policy_model.parameters(), lr=LEARNING_RATE)
dataset = MathReasoningDataset()
dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

6. The GRPO Training Loop

For each sample, the policy model generates a group of candidate responses via stochastic sampling. Those responses are scored by our reward function, the rewards are normalized within the group into relative advantages, and we optimize the GRPO objective: compute the policy ratio between the current and reference policies, add the KL-divergence penalty, combine into a single loss, and backpropagate into the LoRA parameters. Throughout, we track total loss, policy loss, KL divergence, mean reward, and mean policy ratio.

python

policy_model.train()
global_step = 0

for epoch in range(15):
    for batch in dataloader:
        prompt_str = batch["prompt"][0]
        target_str = batch["target"][0]

        prompt_inputs = tokenizer(
            prompt_str, return_tensors="pt",
            max_length=MAX_PROMPT_LEN, truncation=True
        ).to(device)
        prompt_len = prompt_inputs["input_ids"].shape[1]

        # STEP 1: Rollout — generate a group of responses from the policy
        completions, rolled_input_ids = [], []
        with torch.no_grad():
            for g in range(GROUP_SIZE):
                gen_out = policy_model.generate(
                    **prompt_inputs,
                    max_new_tokens=MAX_GEN_LEN,
                    do_sample=True, temperature=0.9, top_p=0.95,
                    pad_token_id=tokenizer.pad_token_id,
                    return_dict_in_generate=True, output_scores=False
                )
                rolled_input_ids.append(gen_out.sequences[0])
                gen_text = tokenizer.decode(
                    gen_out.sequences[0][prompt_len:], skip_special_tokens=True
                )
                completions.append(gen_text)

        # STEP 2: Pad all generated sequences to equal length
        max_len_in_group = max(t.size(0) for t in rolled_input_ids)
        padded_input_ids, padded_attn_masks = [], []
        for tensor in rolled_input_ids:
            pad_needed = max_len_in_group - tensor.size(0)
            padded = F.pad(tensor, (0, pad_needed), value=tokenizer.pad_token_id) \
                     if pad_needed > 0 else tensor
            mask = padded.ne(tokenizer.pad_token_id).long()
            padded_input_ids.append(padded.unsqueeze(0))
            padded_attn_masks.append(mask.unsqueeze(0))
        input_ids_tensor = torch.cat(padded_input_ids, dim=0)
        attn_mask_tensor = torch.cat(padded_attn_masks, dim=0)

        # STEP 3: Reward evaluation -> group-normalized advantages
        rewards = compute_rewards(completions, [target_str] * GROUP_SIZE)
        mean_r = rewards.mean()
        std_r = rewards.std() + 1e-8
        advantages = (rewards - mean_r) / std_r

        # STEP 4-5: Reference and current policy log-probs
        with torch.no_grad():
            ref_log_probs = get_action_log_probs(
                ref_model, input_ids_tensor, attn_mask_tensor, prompt_len
            )
        optimiser.zero_grad()
        policy_log_probs = get_action_log_probs(
            policy_model, input_ids_tensor, attn_mask_tensor, prompt_len
        )

        # STEP 6: Ratio, KL penalty, clipped surrogate objective
        ratios = torch.exp(policy_log_probs - ref_log_probs)
        kl_div = (torch.exp(ref_log_probs - policy_log_probs)
                  - (ref_log_probs - policy_log_probs) - 1.0)
        expanded_adv = advantages.unsqueeze(-1).expand_as(ratios)
        surr1 = ratios * expanded_adv
        surr2 = torch.clamp(ratios, 1.0 - EPSILON_CLIP, 1.0 + EPSILON_CLIP) * expanded_adv
        gen_mask = attn_mask_tensor[:, prompt_len:]

        # STEP 7: Final loss = clipped policy loss + KL regularization
        policy_loss = -torch.sum(torch.min(surr1, surr2) * gen_mask) / gen_mask.sum()
        kl_loss = torch.sum(kl_div * gen_mask) / gen_mask.sum()
        total_loss = policy_loss + BETA_KL * kl_loss

        # STEP 8: Backprop and update the LoRA parameters
        total_loss.backward()
        optimiser.step()
        global_step += 1

print("Training finished successfully.")

That's the entire algorithm — no critic network anywhere in the loop. The baseline that PPO would normally learn with a value model is replaced by the mean reward across the group, which is why GRPO is so much lighter to train. Even on five toy problems and a 1.5B distilled model, you can watch the mean reward climb and the KL term keep the policy anchored to the reference.

✶

Training Dynamics

We can monitor the training process using four key metrics: the loss, reward, policy ratio, and KL divergence.

Four training plots over 10 steps. Top-left: total and policy loss, fluctuating between roughly 0.6 and -0.45. Top-right: mean reward with a ±1 std band, rising from about 0.25 to 0.75. Bottom-left: KL divergence, on the order of 1e-5 and hovering around zero. Bottom-right: mean policy ratio, staying essentially at 1.0 with only a tiny deviation late in training. — Figure 4. GRPO training metrics over 10 steps — loss, reward (±1 std), KL divergence, and mean policy ratio.

From these plots, we can observe that the mean reward improves from approximately 0.25 to 0.75, indicating that the policy is learning to generate better responses over time.

The KL divergence remains tightly controlled and very close to zero, suggesting that the updated policy does not drift significantly from the reference policy. This indicates stable optimization without any signs of instability.

Similarly, the policy ratio stays very close to 1 throughout training, which implies that the policy updates are small and remain within the desired clipping range. This is exactly the behavior we expect in GRPO, as it prevents overly aggressive updates.

Although the policy loss fluctuates and briefly becomes negative, this is completely normal for policy gradient methods. Unlike supervised learning, policy loss is not expected to decrease monotonically, and these fluctuations do not indicate failed optimization.

In this demonstration, we intentionally trained the model for only 10 epochs to keep the example simple and easy to understand. If we were to train for more epochs, the learning trends would become even clearer, and we would expect the reward to stabilize further while maintaining controlled policy updates.

✶

References & Further Reading

The DeepSeek-R1 paper details how models can reach state-of-the-art reasoning primarily through reinforcement learning rather than traditional supervised fine-tuning. The core breakthrough is GRPO (Group Relative Policy Optimization), a highly efficient algorithm that reduces the computational overhead of standard RL by eliminating the separate value network — saving large amounts of GPU memory and compute.

Implementing GRPO from Scratch: DeepSeek-R1's Reinforcement Learning Engine

A Deep Dive into GRPO

1. Rule-Based Deterministic Rewards

2. Group Relative Advantage

The Core Objective Function

Building It in Pure PyTorch

1. Setup & Configuration

2. A Verifiable Dataset

3. Verifiable Rule-Based Rewards

4. Log-Probability Tracking

5. Initialization Pipeline

6. The GRPO Training Loop

Training Dynamics

References & Further Reading

Key Papers

Hands-On Resources & Repositories