Math Reasoning as an MDP

How can we train language models to solve hard reasoning problems—multi-step mathematics, formal proofs, complex word problems? The conventional approach of next-token prediction on human-written solutions faces a fundamental limitation: high-quality reasoning data is scarce. It is estimated that we will run out of high-quality internet text for training by 2028, and expert-level mathematical solutions are already in short supply. Moreover, simply imitating human solutions produces models that "assert" their way through problems without genuine understanding—when they encounter an error in their reasoning chain, they cannot recognize or recover from it.

Reinforcement learning offers a compelling alternative. Instead of imitating fixed solutions, we can let the model explore different reasoning strategies and learn from the outcomes. The mathematical structure of reasoning problems maps naturally onto the RL framework.

This formulation highlights the key challenge: the reward signal arrives only at the very end of a multi-step reasoning chain. The model must learn which intermediate steps were good and which were harmful—a classic temporal credit assignment problem. The rest of this lecture explores how different RL techniques address this challenge, from classical methods to the modern "thinking" models that have achieved state-of-the-art results.

Data Scaling: SFT, RFT, and RL

Before diving into specific algorithms, it helps to understand the landscape of training approaches and how they scale with data. There are three basic strategies for improving a language model's math reasoning ability, each representing a different point on the imitation-to-RL spectrum.

Supervised Finetuning (SFT)

The simplest approach: collect problems and their oracle solutions (written by humans or stronger models), then maximize the likelihood of the solution conditioned on the problem. The key scaling axes are the number of unique questions and the number of oracle answers per question. Empirically, performance improves as both axes grow, but the gains are slow and eventually plateau.

Rejection Sampling Finetuning (RFT)

RFT goes a step further: sample multiple solutions from the learner itself, filter for those that produce the correct final answer, and train on the correct ones. This introduces on-policy data—the model learns from its own successful reasoning attempts rather than from external demonstrations. The hope is that on-policy solutions are more "natural" for the model and therefore easier to learn from.

Reinforcement Learning

RL uses both correct and incorrect solutions produced by the learner. Rather than discarding failures, it extracts learning signal from them—understanding why a solution failed is just as informative as understanding why another succeeded. This is the key advantage of RL over imitation-based methods.

The Problem of Spurious Steps

A critical observation from Setlur et al. (NeurIPS 2024) is that on-policy imitation (RFT) eventually degrades if you train on too much of it. The model's performance on the training distribution may continue to improve, but generalization to new problems deteriorates. Why does this happen?

The culprit is spurious steps—reasoning steps that happen to appear in correct solutions but are not actually valid mathematical reasoning. When the model generates a solution that arrives at the right answer despite containing a flawed intermediate step, RFT treats the entire solution as correct and trains on it. The model memorizes these spurious patterns and learns to rely on them.

This phenomenon is closely related to causal confusion in imitation learning: when the policy is conditioned on its own past actions (as in autoregressive generation), it can learn to rely on spurious correlations between past actions and future success. On the training distribution, these correlations are reliable enough to maintain performance, but they do not transfer to new problems.

The solution is to identify and penalize spurious steps. But how can we tell which steps are spurious and which are genuinely good? This leads us to the concept of per-step advantages.

Credit Assignment via Advantages

The key insight for addressing spurious steps is that we need to evaluate the quality of individual reasoning steps, not just entire solutions. The natural tool for this is the advantage function from reinforcement learning, which measures how much better or worse a particular action is compared to the average action in that state.

Step-Level Value Functions

Consider a partial solution consisting of steps $c_1, c_2, \ldots, c_i$ generated for problem $x$. The Q-value after step $i$ is the probability that the model will reach the correct answer if it continues from this partial solution:

In practice, this is estimated by performing multiple rollouts—completing the solution from the current partial state using the model's own policy—and computing the fraction that arrive at the correct answer. If 8 out of 12 rollouts from a given partial solution succeed, the estimated Q-value is $0.66$.

The Advantage Function

The advantage of step $c_i$ measures the relative change in value caused by committing to that particular step:

Consider the following example. A model solving a system of equations produces six steps. After step $c_1$ (defining variables), the rollout success rate is $0.66$. After step $c_2$ (writing equations), it stays at $0.66$. After step $c_3$ (proposing an elimination strategy), it drops to $0.0$—a clear signal that this step was harmful. The advantage for $c_3$ is $0.0 - 0.66 = -0.66$. Step $c_4$ (executing the elimination) brings the success rate back to $1.0$, giving it an advantage of $+1.0$. The advantage function precisely identifies which steps helped and which ones hurt.

Advantage-Filtered Training

With per-step advantages in hand, we can use them to improve both imitation learning and offline RL. Setlur et al. (NeurIPS 2024) proposed two approaches.

Advantage-Filtered RFT

The simplest approach is to filter training data at the step level using advantages. Rather than training on entire correct solutions, we retain only the steps with high advantage (from correct solutions) and discard steps with low or negative advantage (from incorrect solutions). This produces a cleaner training set that avoids the spurious step problem.

Concretely, from an incorrect solution generated by the current policy $\pi$, extract any step $c_i$ whose advantage $A_i \gg 0$—even though the overall solution is wrong, this step was beneficial. Conversely, from a correct solution, discard any step $c_i$ whose advantage $A_i \ll 0$—this step happened to appear in a correct solution but was actually harmful. The resulting filtered dataset is used for standard supervised finetuning.

Per-Step Offline RL (DPO)

A more powerful approach uses the advantages to construct preference pairs at the step level for offline RL via DPO. The idea is to create training pairs where:

• From an incorrect solution: identify a step with negative advantage and pair the partial solution up to that point (as context) with the good continuation (from a rollout that succeeds) as the preferred response and the bad step as the dispreferred response.
• From a correct solution: identify a step with negative advantage and similarly construct a preference pair that teaches the model to avoid that step.

This approach retains the partial rollouts from the current policy for training, creating DPO preference pairs at the granularity of individual steps rather than complete solutions.

The reason for this dramatic improvement is clear: per-step RL extracts far more learning signal from each problem. Instead of treating each solution as an atomic unit, it identifies exactly which steps are good and bad, and constructs training signal from each individual step. Moreover, it learns from incorrect solutions that imitation methods would simply discard.

Online RL with Outcome and Process Rewards

While offline methods like advantage-filtered RFT and per-step DPO are effective, online RL can do even better by generating fresh data from the evolving policy during training. The basic recipe is straightforward: sample solutions from the current policy, assign rewards based on correctness, and update using policy gradient methods.

Outcome Reward RL

The simplest online approach uses binary outcome rewards: $r(x, y) = 1$ if the final answer is correct, $r(x, y) = 0$ otherwise. The policy gradient update with these sparse rewards takes the form:

Various policy gradient algorithms can be used: REINFORCE, PPO (Proximal Policy Optimization), or GRPO (Group Relative Policy Optimization, used by DeepSeek-R1). The choice of algorithm affects training stability and sample efficiency, but the core idea is the same: reinforce correct solutions, do not reinforce (or penalize) incorrect ones.

Process Advantage Verifiers (PAVs)

Outcome-only rewards are sparse—the model receives no signal about which of its $N$ reasoning steps were good until the very end. This makes credit assignment difficult and exploration slow. Setlur et al. (ICLR 2025) proposed Process Advantage Verifiers (PAVs): parametric models that predict per-step advantages and provide them as dense reward bonuses during online RL.

The key design question is: what rollout policy should be used to compute the training labels for the PAV? The optimal rollout policy is the one that produces advantages most effectively distinguishing good and bad steps across all policies the RL agent will encounter during training. In practice, the PAV is trained from a snapshot of the policy and then held fixed during RL training. Despite this approximation, PAVs provide substantial benefits.

A particularly important finding is that PAVs enable the model to solve hard questions that it could never solve with outcome rewards alone. On problems where the base policy has near-zero success rate even with 256 independent attempts, PAV-augmented RL discovers correct solutions. The dense reward signal guides exploration through the vast space of possible reasoning chains, providing partial credit for steps that make progress even when the final answer is not yet correct.

Takeaways from Classical RL for Reasoning

The classical RL methods covered so far reveal three important lessons:

• Imitation scales, but slowly and with pitfalls. SFT and RFT improve with more data, but they induce spurious steps that are memorized by the model and cause generalization failures. On-policy imitation (RFT) eventually degrades when trained too aggressively on limited initial states.
• Per-step credit assignment is essential. The advantage function identifies which individual reasoning steps help and which ones hurt. Both advantage-filtered RFT and per-step DPO leverage this information to achieve dramatically better data efficiency than whole-solution methods.
• Online RL with dense rewards is most effective. Combining online policy gradient methods with PAV dense rewards gives the best performance, achieving 5-6x sample efficiency over sparse outcome rewards. The dense signal is especially critical for hard problems where exploration is otherwise intractable.

Training "Thinking" Models via RL

The second half of this lecture turns to a more recent development: the emergence of "thinking" models like DeepSeek-R1 and Kimi K1.5 that produce extended reasoning traces before arriving at an answer. These models have achieved remarkable performance on mathematical competitions, coding benchmarks, and scientific reasoning tasks. What changed?

What Changed: Same Objectives, New Action Space

Perhaps surprisingly, the RL training objectives used by thinking models are essentially the same as those covered in the first half of this lecture. DeepSeek-R1 uses GRPO (a policy gradient method); Kimi K1.5 uses a variant of advantage-induced policy alignment (APA). The reward functions include minor modifications but remain fundamentally outcome-based.

The crucial difference is in the action space. The base models underlying thinking models have been trained (through careful pretraining and instruction tuning) to perform "macro" actions—high-level cognitive operations that go beyond individual reasoning steps:

• Verification: Re-checking a previous computation or claim ("Let me verify this result by substituting back into the original equation...")
• Backtracking: Recognizing an error and returning to an earlier point in the reasoning ("Wait, that approach won't work because... Let me try a different strategy.")
• Subgoal decomposition: Breaking a complex problem into manageable sub-problems ("To solve this, I need to first establish X, then use it to derive Y.")
• Backward chaining: Working backward from the desired conclusion ("If the answer were Z, then we would need... So let me check whether that holds.")

These meta-cognitive strategies dramatically expand the effective action space available to the model. RL can then discover which strategies to apply and when, using the same outcome-reward framework.

Formulation with Compute Budgets

Training thinking models can be formalized as optimizing performance subject to a total compute constraint per problem. Let $y$ denote the full response (including the reasoning trace), and let $|y|$ denote its length in tokens. The objective is:

where $B$ is the token budget per problem. This can be optimized via outcome-reward RL (as in DeepSeek-R1) or via SFT/RFT (collect data, filter by correctness, maximize likelihood).

RL Amplifies Response Length

An empirical observation documented by the Kimi K1.5 team is that RL training substantially increases the average response length. This is not an accident—it is a natural consequence of the reward structure. Longer responses that include verification, backtracking, and multiple solution attempts are more likely to arrive at the correct answer. RL discovers this strategy and amplifies it.

However, longer responses are not universally better. There is a tension between the benefits of more "thinking" (higher accuracy) and the cost (more computation, slower inference). The compute budget constraint $B$ in the formulation above captures this trade-off: the model must learn to allocate its thinking budget efficiently, spending more time on hard problems and less on easy ones.

Test-Time Compute Scaling

Thinking models introduce a new paradigm for scaling language model performance. Traditionally, improving a model meant training a bigger model—more parameters, more data, more training compute. Test-time scaling offers an alternative: keep the model size fixed, but allow it to use more compute at inference time by generating longer reasoning traces.

Snell et al. (ICLR 2025, Oral) formalized this as a resource allocation problem: given a fixed total inference budget, how should we distribute compute across problems of varying difficulty? Their key finding is that scaling test-time compute optimally can be more effective than scaling model parameters. A smaller model that "thinks longer" can outperform a larger model that answers immediately.

The results speak for themselves: thinking models now top the leaderboards on mathematical competition benchmarks (AIME, AMC), coding competitions, and scientific reasoning tasks. OpenAI's o3 and o4-mini models, which use this paradigm, have achieved state-of-the-art performance across a range of challenging benchmarks.

Open Questions and Summary

Despite the impressive results, many fundamental questions about RL for reasoning remain open. Three areas seem particularly important for future research.

Desiderata and Formulation

How should we formulate the training problem for thinking models? The current approach—outcome-reward RL with a compute budget—is effective but coarse. A richer formulation might treat the problem as an adaptation problem: given a new test problem, how should the model adapt its reasoning strategy in real time? This connects to meta-RL and the literature on learning to learn.

Dense Rewards Beyond Outcome Verification

Outcome rewards alone are sufficient for current thinking models, but they become increasingly sparse as problems grow harder and reasoning chains grow longer. Process reward models (like PAVs) offer one path to denser signals, but designing reward functions that reliably distinguish good reasoning from bad reasoning—as opposed to good answers from bad answers—remains challenging.

Why RL Outperforms SFT

There is growing empirical evidence that RL substantially outperforms SFT for reasoning tasks, even when both methods use the same base model and data. The theoretical reasons for this gap are not fully understood. One hypothesis is that RL's ability to learn from negative examples (failed solutions) provides information that SFT cannot access. Another is that RL's on-policy nature allows it to correct errors in the model's current reasoning patterns, while SFT can only push toward a fixed target distribution. Understanding this gap more precisely could lead to even more effective training methods.

In summary, RL for reasoning has progressed from classical methods (advantage-filtered imitation, per-step offline RL, dense-reward online RL) to the modern thinking model paradigm (extended reasoning traces, macro-level cognitive strategies, test-time compute scaling). The old RL recipes—policy gradients, advantage functions, credit assignment—remain at the core of the newest systems. What has changed is the scale (millions of tokens per solution), the action space (macro-level cognitive strategies), and the ambition (solving mathematical olympiad problems). The field is moving fast, but the foundational ideas from classical RL continue to do the heavy lifting.