6 Training reasoning models with reinforcement learning

This chapter explains how reasoning performance in large language models can be improved through training-time scaling, with reinforcement learning serving as the main method. It frames reinforcement learning for LLMs as a post-training stage that shapes how a pre-trained model uses its knowledge, rather than simply learning next-token prediction. The chapter also distinguishes this approach from inference-time scaling, while noting that both can be combined to further boost accuracy.

It then contrasts reinforcement learning with human feedback and reinforcement learning with verifiable rewards. Human feedback relies on annotated preferences and an extra reward model, which makes it flexible but costly and complex. Verifiable rewards remove that extra model by using deterministic checks, such as exact-answer verification for math problems, making the process simpler, reproducible, and easier to scale where reliable automatic validation exists. The chapter uses this setup to train a reasoning model on math tasks, focusing on correctness rather than intermediate reasoning supervision.

Finally, the chapter walks through a practical implementation of group relative policy optimization. Multiple answers are sampled for each prompt, scored by a verifier, converted into group-relative advantages, and combined with sequence-level log probabilities to produce a policy gradient loss. These components are assembled into a full training loop with optimization, gradient clipping, logging, and checkpointing, and the resulting small-scale run shows improved reasoning behavior and stronger MATH-500 accuracy, demonstrating that verifiable-reward reinforcement learning can unlock useful reasoning in a base model even with a simplified training recipe.

A mental model of the topics covered in this book. This chapter focuses on techniques that improve reasoning with additional training (stage 4). Specifically, this chapter covers reinforcement learning.

Conceptual comparison of inference-time scaling and training-time scaling. Increasing compute at inference improves accuracy by spending more resources per answer generation, while increasing compute during training improves accuracy by investing more resources upfront.

Common training stages for LLMs. The ordering of the reasoning training and preference tuning stages can vary, and some pipelines interleave reasoning and preference tuning rather than strictly sequencing them.

Roadmap of this chapter. After a brief introduction to reinforcement learning (RL) for LLMs in this section, we discuss the difference between two RL stages, RLHF and RLVR, in the next section. Then, we focus on implementing RLVR using the GRPO algorithm in the remainder of this chapter.

Two-stage overview of reinforcement learning with human feedback (RLHF). First, a reward model is trained on human-ranked responses to assign a preference score to each. Second, the LLM is updated using these reward scores within an RL objective to encourage preferred responses and discourage less desirable ones.

Overview of reinforcement learning with verifiable rewards (RLVR). The LLM generates a response that is evaluated by a deterministic verifier, for example the math verifier from chapter 3, which assigns a correctness label used as a reward signal within an RL objective to update the model.

After introducing the two main reinforcement learning approaches for LLMs, RLHF and RLVR, the remaining sections focus on implementing RLVR using the GRPO algorithm, from dataset loading to implementing the full training loop.

High-level overview of the GRPO algorithm for RLVR using a chef analogy. Multiple rollouts are generated and scored, relative advantages are computed, and a policy gradient objective with a KL-based regularization (loss) term is used to update the model parameters.

Step-by-step GRPO update for RLVR. (1) A prompt is sampled and multiple rollouts are generated. (2) Each rollout is scored using a verifiable reward. (3) Group-relative advantages are computed from these rewards. (4) The log-probability of each rollout under the current model is calculated. (5) Advantages and log-probabilities are combined to form the policy gradient loss. (6) A KL regularization term against a reference model is added, and the resulting total loss is used to update the model parameters.

In stage 5, we load the pre-trained model (this section) and dataset (next section) that we will use for the model training.

Structure and split of the MATH dataset. The full dataset contains about 12,500 problems that are divided into a 500-problem test set (MATH-500), which we used for model evaluation in chapter 3. A non-overlapping set of 12,000 problems is used for training in this chapter.

After outlining the RLVR method and GRPO algorithm, the following sections implement the individual GRPO stages that we need to train the LLM via verifiable rewards on the MATH dataset.

Step-by-step GRPO update for RLVR (without KL loss term). We begin by prompting the LLM to generate the different rollouts.

The second stage in the GRPO pipeline computes the rewards for each rollout the LLM generated in the previous section.

The third GRPO stage computes the advantage values from the answer (rollout) correctness rewards.

The fourth stage in the GRPO pipeline computes log-probabilities for each rollout, which is related to the logprob scorer we developed in the previous chapter.

The fifth stage in the GRPO pipeline computes the policy gradient loss that we use to update the model. Stage number 6, the model weight update, will be implemented as part of the training loop later.

The complete GRPO workflow where (1) multiple rollouts are generated for a prompt, (2) assigned correctness rewards, (3) converted into group-relative advantages, and (4) combined with log probabilities to (5) compute the policy gradient loss. The loss gradients (6) will be computed and used to update the model in the next section.

After implementing the individual GRPO stages, we now implement the surrounding training loop to update the model weights.

Outline of the training loop. The overall structure follows a standard deep learning training loop. The key difference lies in how the loss is computed: instead of a standard supervised objective, the loss is obtained via the GRPO stages (stage 4).

The final step of this chapter discusses how we can load the saved model checkpoints and evaluate them.

• Reinforcement learning (RL) can be used to train LLMs on human preference labels and verifiable rewards.
• RL is typically applied as post-training on top of a pre-trained base model, and it can be inserted at different stages of an LLM pipeline, including reasoning training and preference tuning.
• RL with human feedback (RLHF) optimizes for human preferences via a two-stage setup: train a reward model from ranked responses, then use reward scores to update the LLM.
• RL with verifiable rewards (RLVR) simplifies RLHF by replacing learned reward models with deterministic, automatically computed verifiers (for example, math answer checking)
• We focussed on RLVR for math reasoning.
• We used GRPO as the policy optimization algorithm that turns verifier rewards into parameter updates; because GRPO directly optimizes the model using sequence-level rewards without requiring a separate value model, it is particularly convenient.
• GRPO is a more resource-friendly alternative to other RL algorithms for LLMs because it avoids training a separate value model and instead derives learning signals from comparisons within a group of sampled rollouts.
• A "rollout" refers to a full model answer (completion) for a prompt; rewards, advantages, and log-probabilities are computed from the rollout in later steps.
• Rewards are computed from a verifier that only grants a reward if the final answer is both correct and extractable in a required format like "\boxed{}".
• Raw rewards are transformed into advantages by normalizing each rollout reward relative to the group mean and standard deviation.
• GRPO also relies on sequence-level log-probabilities, which are computed by summing token log-probabilities over the generated answer tokens.
• Sequence log-probabilities, together with the advantages, form the core policy-gradient objective in GRPO.
• The full GRPO loss computation is combined into a single function that performs rollout sampling, reward computation, advantage calculation, log-prob computation, and policy-gradient loss calculation.
• The surrounding training loop is a standard deep learning loop, with the key difference being that the loss comes from GRPO rather than conventional classification losses.
• Training is resource-intensive because each step requires generating multiple, potentially long rollouts, but even short GRPO runs can increase MATH-500 accuracy from 15% to 47%.

FAQ