Reinforcement learning for policies and reasoning

Why RL after imitation

New to RL? Start with Reinforcement learning from scratch, which builds up the vocabulary used here (state, reward, return, value, baseline) on a small grid world.

Behaviour cloning (the Imitation learning module) copies the data, mistakes included, and never sees the consequences of its own actions. Reinforcement learning (RL) optimises an outcome directly: the policy acts, a reward scores the result, and the policy changes to make high-reward behaviour more likely. Recent language and driving models use both in sequence: supervised fine-tuning (SFT) first, so the model is competent, then RL to improve what imitation cannot.

The policy gradient idea

A policy gives a probability to each action. After sampling an action and observing its reward:

• if the reward is better than expected, make that action more likely;
• if worse, make it less likely.

The update direction for one sample is

$\nabla_\theta J \approx A\,\nabla_\theta \log \pi_\theta(a \mid s)$

where $\pi_\theta$ is the policy and $A$ is the advantage. The advantage is the reward minus a baseline (what you expected). Without a baseline the gradient is correct on average but very noisy. The whole family of methods below differs mainly in how they obtain the baseline and how they keep updates safe.

Estimating the advantage: GAE

PPO uses a learned value network (critic) to turn rewards into advantages. The one-step estimate uses the critic's guess of the next state, $\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$; the full Monte Carlo estimate uses every real reward until the end. Generalized Advantage Estimation (1506.02438) blends them with a dial $\lambda$:

$\lambda = 0$ is the one-step estimate: low variance, but only as good as the critic. $\lambda = 1$ is the Monte Carlo return minus $V(s_t)$: unbiased, but it accumulates the noise of every step. The reference code measures both on a chain of 20 steps with noisy rewards and a critic that is wrong in the middle:

def gae(
    rewards: torch.Tensor,
    values: torch.Tensor,
    gamma: float = 0.99,
    lam: float = 0.95,
) -> torch.Tensor:
    """Generalized advantage estimation along one episode.

    delta_t = r_t + gamma V(s_t+1) - V(s_t) is the one-step TD error. The advantage
    is the discounted sum A_t = sum_k (gamma lam)^k delta_t+k. With lam = 0 it is the
    one-step estimate (low variance, biased by the critic's errors); with lam = 1 it is
    the full Monte Carlo return minus V(s_t) (unbiased, high variance).

    Args:
        rewards: Shape (T,).
        values: Critic values V(s_0 ... s_T), shape (T + 1,); V(s_T) is 0 at a terminal.
    """
    deltas = rewards + gamma * values[1:] - values[:-1]
    advantages = torch.zeros_like(rewards)
    running = 0.0
    for t in reversed(range(len(rewards))):
        running = deltas[t] + gamma * lam * running
        advantages[t] = running
    return advantages

Typical practice is $\lambda$ around 0.95: enough bootstrapping to tame the variance, enough real reward to wash out most critic error. The cost of all this is the critic itself: another network that must be trained and held in memory, and whose mistakes feed into the policy's gradient.

PPO

Proximal Policy Optimization keeps each update small by clipping the ratio between the new and old action probabilities, so one lucky batch cannot move the policy far. It estimates the baseline with a learned value network (a critic), which is another model to train and store.

What the clip does

Let $r = \pi_{\text{new}}(a \mid s) / \pi_{\text{old}}(a \mid s)$ and $A$ the advantage. The objective for one sample is $\min\big(rA,\ \mathrm{clip}(r, 1-\epsilon, 1+\epsilon)\,A\big)$:

def ppo_objective(
    ratio: torch.Tensor, advantage: torch.Tensor, clip: float = 0.2
) -> torch.Tensor:
    """PPO's clipped surrogate: min(r A, clip(r, 1 - eps, 1 + eps) A)."""
    return torch.minimum(ratio * advantage, ratio.clamp(1 - clip, 1 + clip) * advantage)


def clip_gradient_is_nonzero(ratio: float, advantage: float, clip: float = 0.2) -> bool:
    """Whether the clipped surrogate still has a gradient in the ratio.

    With a positive advantage the gradient is cut once the ratio exceeds 1 + eps (the
    action is already more likely enough); with a negative one, once it falls below
    1 - eps. In the other direction the gradient stays, so a bad step can be undone.
    """
    r = torch.tensor(ratio, requires_grad=True)
    ppo_objective(r, torch.tensor(advantage), clip).backward()
    return bool(r.grad != 0)

Read the plot as a rule: you may make a good action more likely, but only up to 20 percent more likely per update; you may make a bad action less likely, but only down to 80 percent. The gradient vanishes only when the update has already moved far enough in the helpful direction. A test checks all four cases: positive advantage with $r = 1.3$ has no gradient, with $r = 0.7$ it does (the policy made a good action less likely, so it is allowed to repair that); the mirror holds for negative advantage. This is a trust region enforced by a clamp, not by an explicit constraint.

PPO usually runs several epochs of mini-batch updates on each batch of rollouts. The clip is what makes that safe: after a few updates most ratios have left the band and stop contributing.

Penalising drift from a reference: the KL term

For language and driving models the policy must also stay near a reference (the SFT model), or it exploits the reward model (RL for alignment). The penalty is $\mathrm{KL}(\pi \,\|\, \pi_{\text{ref}})$, but with a language model you can only sample from $\pi$, so it is estimated from the sampled sequence. Three common single-sample estimators, with $r = \pi_{\text{ref}}/\pi$ at the sample:

def kl_estimators(
    logp: torch.Tensor, ref_logp: torch.Tensor
) -> dict[str, torch.Tensor]:
    """Three single-sample estimators of KL(pi || ref) from samples x ~ pi.

    With r = ref(x) / pi(x):

    * k1 = -log r: unbiased, high variance, can be negative;
    * k2 = (log r)^2 / 2: biased, low variance, always non-negative;
    * k3 = (r - 1) - log r: unbiased, low variance, always non-negative (used in GRPO).
    """
    log_ratio = ref_logp - logp
    return {
        "k1": -log_ratio,
        "k2": 0.5 * log_ratio**2,
        "k3": log_ratio.exp() - 1 - log_ratio,
    }

For two three-outcome distributions the true KL is 0.0877; over 200,000 samples k1 averages 0.0877 with standard deviation 0.405, k2 averages 0.0858 (biased) with 0.055, and k3 averages 0.0877 with 0.074. k3 is the best of both, which is why GRPO uses it (its loss code is below). The estimator matters because the penalty's gradient is what trains the model, and a noisy or negative estimate makes training erratic.

GRPO

Group Relative Policy Optimization, introduced in the DeepSeekMath paper, removes the critic. For each prompt it samples a group of outputs, scores each, and uses the group itself as the baseline: the advantage of an output is its reward compared with the group's average (often divided by the group's standard deviation). Fewer models in memory, which matters for large language and vision-language models.

For each prompt, sample $G$ outputs with rewards $r_1,\dots,r_G$ and compute

$A_i = \frac{r_i - \mathrm{mean}(r)}{\mathrm{std}(r)}$

Then raise the probability of outputs with positive advantage and lower the others, while staying close to a reference model.

A subtle shrinkage

The group mean includes the output being scored. For a group of $G$, the baseline for output $i$ contains $r_i/G$ of its own reward, which removes that share of its signal. The average gradient is therefore shrunk by $(G-1)/G$: half the true size for $G = 2$, three quarters for $G = 4$, 88 percent for $G = 8$. A leave-one-out baseline (the mean of the other $G-1$ outputs) is unbiased. The reference code measures all of them on the four-armed bandit:

MEANS = torch.tensor([0.2, 0.5, 0.8, 0.3])


def group_gradient_samples(
    group: int, baseline: str, trials: int = 20000, noise: float = 0.3, seed: int = 0
) -> torch.Tensor:
    """Gradient estimates for the best arm's logit with a group baseline.

    A softmax policy over four arms (uniform logits) samples ``group`` arms per trial.
    The estimate is (1 / G) sum_i (R_i - b_i) d log pi(a_i) / d logit_best, with:

    * ``"none"``: b = 0;
    * ``"mean"``: b = the group's mean reward (includes sample i itself);
    * ``"leave_one_out"``: b = the mean of the other G - 1 rewards;
    * ``"true"``: b = the policy's expected reward (not available in practice).

    Returns:
        One estimate per trial, shape (trials,).
    """
    generator = torch.Generator().manual_seed(seed)
    arms = torch.randint(0, 4, (trials, group), generator=generator)
    rewards = MEANS[arms] + noise * torch.randn(trials, group, generator=generator)
    if baseline == "none":
        base = torch.zeros(trials, 1)
    elif baseline == "mean":
        base = rewards.mean(dim=1, keepdim=True)
    elif baseline == "leave_one_out":
        base = (rewards.sum(dim=1, keepdim=True) - rewards) / (group - 1)
    elif baseline == "true":
        base = MEANS.mean().expand(trials, 1)
    else:
        raise ValueError(f"unknown baseline {baseline!r}")
    score = (arms == 2).float() - 0.25  # d log pi(a) / d logit_best for uniform logits
    return ((rewards - base) * score).mean(dim=1)

(Mean ± standard deviation.) The group mean buys the lowest variance at the price of a slightly shrunk gradient, which Adam's step normalisation largely hides. Leave-one-out keeps the gradient exact and still halves the variance of having no baseline at $G = 8$. Any baseline is far better than none, and larger groups help both; a group of 2 is nearly useless.

Dividing by the group's standard deviation is a separate choice. It makes updates comparable across prompts of different reward scale, but it also inflates prompts where all outputs are nearly equal (a tiny standard deviation makes noise look like signal), which is why the loss code adds a small $\epsilon$ and why some later variants drop the division.

DeepSeek-R1 showed this style of RL on verifiable rewards (answers that can be checked) can produce long reasoning behaviour. Alpamayo-R1 applies GRPO to a driving VLA.

Exercise: GRPO on a four-armed bandit

Four actions have mean rewards 0.2, 0.5, 0.8 and 0.3, with noise. A softmax policy samples a group of 8 per step and updates with group-normalised advantages. We ran it:

import numpy as np

rng = np.random.default_rng(0)
true_reward = np.array([0.2, 0.5, 0.8, 0.3])  # mean reward of each of four actions
logits = np.zeros(4)
GROUP = 8
LR = 0.5

for step in range(201):
    probs = np.exp(logits) / np.exp(logits).sum()
    actions = rng.choice(4, size=GROUP, p=probs)
    rewards = true_reward[actions] + 0.1 * rng.standard_normal(GROUP)
    advantages = (rewards - rewards.mean()) / (rewards.std() + 1e-8)
    grad = np.zeros(4)
    for action, advantage in zip(actions, advantages):
        onehot = np.eye(4)[action]
        grad += advantage * (onehot - probs)
    logits += LR * grad / GROUP
    if step in (0, 20, 50, 200):
        print(f"step {step:3d}: probabilities {np.round(probs, 2)}")

Output:

step   0: probabilities [0.25 0.25 0.25 0.25]
step  20: probabilities [0.02 0.04 0.91 0.02]
step  50: probabilities [0.   0.01 0.99 0.  ]
step 200: probabilities [0. 0. 1. 0.]

The policy concentrates on the best action (index 2) without any critic.

Two experiments to try: make the reward noise larger than the gap between the rewards, and set the group size to 2. The figure suggests learning becomes slower and more variable.

This toy omits the pieces real systems need: a KL penalty to a reference model, clipping, and sampling whole token sequences instead of one action.

Reference implementation (PyTorch)

The same ideas on a toy sequence task: three "prompts", each with a target token sequence, and a reward equal to the share of matching tokens. The policy is a table of per-position logits so the code stays about GRPO. Tests are in libs/tutorial/tests/test_grpo.py and include a check that the clipping stops the gradient for a sample that has already moved far enough; run uv run --group tutorial python -m libs.tutorial.grpo. The building blocks (GAE, the clip's gradient regions, the three KL estimators and the group baselines) are in libs/tutorial/policy_gradient.py, with tests in libs/tutorial/tests/test_policy_gradient.py.

def group_advantages(rewards: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    """Scores each output against the other outputs for the same prompt.

    Args:
        rewards: Shape (prompts, group_size).
        eps: Added to the standard deviation so a constant group gives zeros.

    Returns:
        (reward - group mean) / group std, same shape as ``rewards``.
    """
    mean = rewards.mean(dim=1, keepdim=True)
    std = rewards.std(dim=1, keepdim=True)
    return (rewards - mean) / (std + eps)

def grpo_loss(
    logp: torch.Tensor,
    old_logp: torch.Tensor,
    ref_logp: torch.Tensor,
    advantages: torch.Tensor,
    clip: float = 0.2,
    kl_coef: float = 0.04,
) -> torch.Tensor:
    """Clipped policy-gradient loss with a KL penalty to a reference policy.

    All log-probabilities are of whole sampled sequences, shape (prompts, group).

    Args:
        logp: Log-probability under the policy being trained.
        old_logp: Log-probability under the policy that produced the samples.
        ref_logp: Log-probability under the frozen reference policy.
        advantages: Output of ``group_advantages``.
        clip: PPO-style clip range for the probability ratio.
        kl_coef: Weight of the KL penalty.

    Returns:
        A scalar loss to minimise.
    """
    ratio = (logp - old_logp).exp()
    clipped = ratio.clamp(1 - clip, 1 + clip)
    surrogate = torch.minimum(ratio * advantages, clipped * advantages)
    # Non-negative estimator of KL(policy || reference) from samples.
    delta = ref_logp - logp
    kl = delta.exp() - delta - 1
    return -(surrogate - kl_coef * kl).mean()

def train(
    steps: int = 150, group: int = 8, epochs: int = 2, seed: int = 0
) -> list[float]:
    """Runs GRPO and returns the mean reward after each step."""
    torch.manual_seed(seed)
    policy = TablePolicy(len(TARGETS))
    reference = TablePolicy(len(TARGETS))
    optimizer = torch.optim.Adam(policy.parameters(), lr=0.1)
    history = []
    for _ in range(steps):
        with torch.no_grad():
            sequences = policy.sample(group)
            rewards = reward(sequences)
            old_logp = policy.log_prob(sequences)
            ref_logp = reference.log_prob(sequences)
        advantages = group_advantages(rewards)
        for _ in range(epochs):
            loss = grpo_loss(policy.log_prob(sequences), old_logp, ref_logp, advantages)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        history.append(rewards.mean().item())
    return history

Reward design in Alpamayo-R1

The paper (Sec. 5.3) combines three kinds of reward, checked here against its PDF:

The paper also adds a safety reward that penalizes unsafe or physically implausible trajectories and reports that it further lowers the close encounter rate. The weights of the terms are not stated. It also describes selecting RL training data where the model's output distribution disagrees most with a reward-based target, mixed with random samples.

Reward hacking

A policy maximises the reward you wrote, not the one you meant. Examples to look for:

• If consistency is rewarded alone, the model can write trivial reasoning that always matches.
• If only collision is penalised, standing still scores well.
• If an LLM critic grades reasoning, the policy may learn to flatter the critic.

That is why the recipe combines several rewards and why evaluation should use a separate closed-loop test, not the training reward.

PPO versus GRPO

For long language outputs with one end-of-sequence reward, a critic is hard to train and the group is cheap, which is why GRPO took over in reasoning RL. For control problems with dense rewards, PPO's per-step credit assignment remains the better fit.

Check yourself

1. Why is a baseline needed in a policy gradient, and what does GRPO use instead of a critic?
2. What would you expect if the group size is 2 instead of 8?
3. Why do SFT and RL come in that order?
4. Give a way a driving policy could exploit a reward based only on collisions.
5. What does the PPO clip do when the advantage is positive and the ratio is already 1.3? And when it is 0.7?
6. What does lambda control in GAE, and what goes wrong at each extreme?
7. Why is the group mean a slightly biased baseline, and what fixes it?

Quiz

Answer each question; feedback appears as soon as you choose.

Sources

• Schulman et al., Proximal Policy Optimization Algorithms, 2017.
• Schulman et al., High-Dimensional Continuous Control Using Generalized Advantage Estimation, 2015.
• Shao et al., DeepSeekMath, 2024 (GRPO).
• DeepSeek-AI, DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning, 2025.
• Ouyang et al., Training language models to follow instructions with human feedback, 2022 (RLHF).
• Alpamayo-R1, arXiv 2511.00088, Sec. 5.3.