Toolkit: rewards and policy gradients

Why this is in a driving course

Imitation copies the data, mistakes included, and never sees the result of its own actions. Reinforcement learning (RL) adds a different signal: try something, score the result with a reward, and make high-reward behaviour more likely. Alpamayo-R1 is fine-tuned this way after imitation, with rewards for reasoning quality, for consistency between reasoning and action, and for the trajectory. This page gives the vocabulary and the one algorithm idea you need. The full treatment starts at Reinforcement learning from scratch.

The vocabulary

• A policy $\pi(a \mid s)$ gives a probability for each action in a state. A driving VLA is a policy whose "action" can be a whole reasoning trace and trajectory.
• A reward $r$ scores what happened: progress, a collision, a comfortable ride, a correct explanation.
• The return is the sum of rewards over an episode, often with later rewards discounted.
• The goal is a policy with the highest expected return.

For a model that outputs a whole trajectory and trace at once, an "episode" can be a single output scored once. That is the setting of GRPO below.

The policy gradient

How do you change the weights to raise the expected return when the reward is not differentiable (a collision checker, a judge model)? The policy gradient says: sample an action, then push up its log-probability in proportion to how good it was,

Good outputs become more likely, bad ones less likely, and the reward function never needs a gradient. This estimator is called REINFORCE.

Why the baseline $b$ matters

Without a baseline ($b = 0$), every sampled action is pushed up in proportion to its raw reward. If all rewards are positive, every action is pushed up, and the useful signal is only the small difference between them, buried in noise. Subtracting a baseline (what you expected to get) leaves only "better or worse than expected":

def policy_gradient_samples(
    means: tuple[float, float], baseline: float, samples: int = 5000, seed: int = 0
) -> torch.Tensor:
    """One-sample REINFORCE estimates of d(expected reward)/d(logit of arm 1).

    The policy is a fair coin over two arms (logits 0, 0). Each estimate is
    (reward - baseline) * d log pi(arm) / d logit_1, which is +-0.5 here.

    Returns:
        Shape (samples,): the individual gradient estimates.
    """
    generator = torch.Generator().manual_seed(seed)
    arms = torch.randint(0, 2, (samples,), generator=generator)
    rewards = torch.tensor(means)[arms] + torch.randn(samples, generator=generator)
    score = torch.where(arms == 1, 0.5, -0.5)
    return (rewards - baseline) * score

The baseline does not change what the gradient points at on average; it only removes noise. Same mean, a fifth of the spread.

GRPO: a baseline from a group

Where does the baseline come from? Classic methods (PPO) train a second network, a critic, to predict it. For a large VLA that is a second large model. GRPO (Group Relative Policy Optimization) avoids it: for each input, sample a group of outputs, score them all, and use the group's own average as the baseline.

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)

Each output's advantage is its reward minus the group mean, divided by the group's spread. Outputs better than their siblings are reinforced, worse ones suppressed. A group of 2 gives a poor baseline; 8 or more is typical. This is the method Alpamayo-R1 uses.

Staying close to a reference

An RL-trained model will find whatever the reward actually rewards, including things nobody intended: a reward for "no collisions" alone is maximised by never moving. Two defences appear in every recipe in this course:

• Reward design: several terms (progress, comfort, safety, consistency) so no single loophole wins.
• A penalty for drifting from the reference model (the imitation-trained one), measured by the KL divergence between the two. It keeps RL as a refinement of competent behaviour rather than a search for exploits.

This is also why RL comes after imitation: RL needs outputs that sometimes earn reward, and a competent starting point to stay near.

Where this shows up in the course

• Reinforcement learning: PPO's clip, the GAE dial, the KL estimators and GRPO in depth, with Alpamayo-R1's reward design and reward hacking.
• Data and closed loop: choosing training samples where the model is weakest.
• Safety: why a reward is a preference, not a guarantee.

Check yourself

1. Why does the policy gradient not need the reward to be differentiable?
2. What does a baseline change about the gradient estimate, and what does it leave unchanged?
3. Where does GRPO get its baseline, and what does that save compared with PPO?
4. Give a reward for driving that a policy could exploit, and one way to close the loophole.
5. Why does RL come after imitation in driving and language-model recipes?

Quiz

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

Go deeper

Reinforcement learning from scratch builds the vocabulary on a grid world; Reinforcement learning covers PPO, GAE, GRPO and reward design; RL for alignment covers reward models and the KL penalty in depth.