The recipe extends: Search-R1

The cleanest demonstration is Search-R1 (Jin et al. 2025). Take the R1 training recipe, but instead of letting the model think to itself, let it call a search engine in the middle of its reasoning.

The training data is question-answering, specifically the kind that requires looking facts up. “Curious is a women’s fragrance by a singer born in what city and state?” The model can’t reasonably know the answer from memorized training data and needs to chain together facts it doesn’t have.

The trajectory is interleaved: reason for a bit, decide a search would help, issue a query, read the retrieved passages, reason more, maybe search again, eventually answer. In practice:

<think> I need to find the singer behind the fragrance "Curious". </think>
<search> Curious fragrance singer </search>
<information> [Wikipedia: Curious is a fragrance by Britney Spears...] </information>
<think> So I need to find where Britney Spears was born. </think>
<search> Britney Spears birthplace </search>
<information> [...Spears was born in McComb, Mississippi...] </information>
<think> McComb, Mississippi. That's the answer. </think>
<answer> McComb, Mississippi </answer>

The training loop:

for each training question:
    trajectory = []
    while not answered:
        generate next chunk from the model
        append chunk to trajectory
        if chunk contains <search> query </search>:
            results = retriever(query)
            append <information> results </information> to trajectory
        if chunk contains <answer> ... </answer>:
            answered = True
    reward = exact_match(extracted_answer, ground_truth)
    update the policy, but only on tokens the model wrote,
    NOT on tokens that came from retrieved passages

The last line is the wrinkle worth understanding.

Setup-wise, this is the R1 recipe with one change. Same policy-gradient RL (Search-R1 tests both PPO and GRPO, with PPO more stable in their main setting). Same outcome-only reward (exact match between predicted and ground-truth answer). No process rewards, no learned reward model, no human-graded trajectories. The only difference is that when the model produces </search>, the training loop pauses generation, runs the query, and pastes the top results back into the trajectory before the model continues.

And it works. Trained on Natural Questions and HotpotQA, evaluated across seven QA datasets, Qwen2.5-7B with Search-R1 (PPO) gets 43% average accuracy versus 30% for vanilla retrieval-augmented generation and 28% for R1-style RL with no search at all. The gap is consistent across both in-distribution and out-of-distribution test sets.

What’s striking is what the model learned without being told. The reward checks only the final answer. Nothing in the training signal says “search more often” or “search more carefully.” But over training, the average number of searches per trajectory grows. The model learns to search when uncertain, to search again if the first results aren’t good enough, and sometimes to do a verifying search after it thinks it has the answer.

None of this was hand-designed. The policy gradient finds these behaviors because trajectories that include them tend to end with correct answers more often than trajectories that don’t. The same mechanism that gave R1 its self-correcting math behavior gives Search-R1 its search behavior.

When the action space gets richer: ToolRL

Search-R1 works because search is a forgiving tool. The action space is small (queries to one engine). Verification is sharp (did the final answer match?). Search itself is robust, since even mediocre queries usually retrieve something useful.

Real tool use is messier. An agentic model might have access to dozens of tools, each with named parameters, some required and some optional, taking strings or numbers or structured objects. The model can fail by picking the wrong tool, by using wrong parameters, by getting parameter names right but values wrong, or by forgetting to call a tool that was needed.

ToolRL (Qian et al. 2026) looked at training in this richer setting. Two findings: one expected, one surprising.

What this means for the models you use

When you watch Claude search the web, call a calculator, or run code, there’s a version of this kind of training behind it. Not literally Search-R1 or ToolRL (production training pipelines are more complex and not public), but the conceptual approach is the same: RL on rollouts where the model has tool access and there’s a clear definition of what a successful interaction looks like.

This explains some uneven patterns. Models are smoother with tools they saw a lot of during training (search, code execution, basic file operations) and rougher with niche or company-specific tools that didn’t get the same coverage. They’re also better at tools where success is checkable. For instance, a search returning useful results is verifiable.

The harness around the model (tool schemas, validation, retries, fallback rules) does some of this work too. But none of it substitutes for the model itself having the right reflexes: knowing when to call which tool, what arguments to pass, and when to skip the tool entirely. That part has to come from training.

What’s next, and the series wrap

This is the last piece in the series on RL training for LLMs. The arc:

• REINFORCE foundations: the setup these algorithms live in. MDPs, returns, credit assignment, value vs policy, the RL objective.
• REINFORCE gradient: the policy gradient derived. REINFORCE is weighted supervised learning.
• PPO: the workhorse policy-gradient algorithm; REINFORCE plus five fixes for variance and data efficiency.
• GRPO: the variant that drops the value network and powered DeepSeek’s reasoning models.
• DPO: the offline simplification that collapses RLHF into a single supervised loss.
• R1: what GRPO produces when you point it at math and code, with reasoning behavior consolidating into a long-chain-of-thought policy.
• Search-R1 and ToolRL: what happens when you extend the recipe to tool use.

R1-Zero is the result. Everything else is interpretation.

The setup is simple. Start from DeepSeek-V3-Base (Liu et al. 2024), a 671B model pretrained on text but never instruction-tuned. No SFT, no RLHF, no human-written reasoning examples.

Run GRPO on it (PPO with the critic removed and a group-relative advantage in its place; the GRPO piece covers the details). Use math and code prompts where the answer is checkable. Two rewards:

• Accuracy: exact-match against the ground-truth answer for math, test execution for code.
• Format: the output uses the <think>...</think><answer>...</answer> template.

Both rule-based. No neural reward model, no process supervision, no verifier trained on preferences.

What happens during training

AIME 2024 accuracy climbs from 15.6% to 71.0%, and to 86.7% with sixteen samples and majority voting (Guo et al. 2025). The surprising part is response length. Average response length grows monotonically across training.

The reward function never mentions length. The model discovers through trajectory comparison under GRPO that longer traces correlate with getting the answer right.

The paper flags an “aha moment”: mid-derivation on a competition math problem, the model produces Wait, wait. Wait. That’s an aha moment I can flag here and re-evaluates a step it was about to commit to. Self-reflective tokens (wait, check, verify, mistake, wrong) grow alongside response length. None were in the reward function or the prompt template. They appeared because trajectories that included them got reward more reliably than trajectories without them.

R1 is what you build when the policy needs to be a product

R1’s contribution is different from R1-Zero’s. It’s the recipe for taking the raw reasoning behavior and making it a deployable assistant.

Test-time compute is downstream of training

Snell’s result (Snell et al. 2024) predates R1 and shows that a smaller model with adaptive test-time compute can match a 14× larger model on easy and medium-difficulty problems. But Snell got there by building substantial external scaffolding around the base LLM: a separately trained process reward model (PRM; scoring intermediate steps of a reasoning trace), a revision model fine-tuned on self-correction trajectories, beam search over PRM outputs, best-of-N with verifier-weighted aggregation. The scaffolding extracts the gains.

R1 displaced the stack. No PRM. No beam search. No revision model. No best-of-N at inference. The model produces one long CoT in one decoding pass. The DeepSeek paper classifies PRMs and MCTS-style search (Monte Carlo Tree Search) as unsuccessful attempts in their own development.

The scaffolding got internalized into the policy via training. Wait, that’s wrong is the model running its own verifier mid-trace. Backtracking is the model doing its own revision. Exploring an alternative approach is the model running its own search.

When somebody says a reasoning model “spends more compute at inference,” what they usually mean is that the model produces a longer CoT. The longer CoT exists because training made the longer CoT effective.

Verifiable rewards are the source and the boundary

Every successful application of this recipe shares a common feature: outcomes can be checked cheaply and unambiguously. Math problems with deterministic final answers. Coding problems with executable test suites. Formal theorem proving with proof checkers.

Reasoning behavior consolidates in R1-Zero because GRPO can attribute reward credit reliably across thousands of tokens based on a single end-of-trace check. The trace is long, the signal is sparse, the per-step contribution is invisible. None of that matters as long as the final-answer check is trustworthy. You compare trajectories, reward the ones that ended right, and over enough samples the policy gradient finds whatever in the middle of the trace was helping. Self-correction phrases, hypothesis exploration, longer length all win the comparison on average, when verification is reliable.

Run the same recipe with a learned reward model (say, a neural judge trained on human preferences for “good reasoning”) and the signal becomes noisy and gameable under optimization pressure. Rule-based rewards on verifiable domains don’t.

What’s next

R1-Zero is one demonstration in a broader pattern. The same recipe extends to multi-turn tool-using settings if outcomes remain verifiable. Search-R1 (Jin et al. 2025) is the cleanest demonstration: extend the recipe to a model that interleaves reasoning with search-engine calls, and the model learns to use the tool well as part of its reasoning policy.

The next piece covers what changes when you move from single-turn verifiable reasoning to multi-turn tool-using settings and how the boundaries of “verifiable” get tested when you add tools and external state.

In the broader arc:

• PPO is the workhorse algorithm.
• DPO is the offline alternative.
• GRPO is the variant for verifiable reward settings.
• R1 is what GRPO produces when you point it at math and code on a strong base model.

Why PPO + RLHF was painful

PPO + RLHF is a multi-stage pipeline. You start with an SFT model, then:

1. Collect a preference dataset: pairs of (prompt, chosen response, rejected response).
2. Train a separate reward model on this dataset, learning to score (prompt, response) pairs.
3. Run PPO using the reward model as the reward signal, generating fresh rollouts during training.

Training requires four LLM-sized models in memory: the SFT model serves as the reference for KL anchoring, the policy is the LLM model being trained, the reward model scores rollouts, and PPO needs a value network on top.

The KL anchor matters because of what happens without it. The policy will drift toward outputs that score high under the reward model but bear no relationship to coherent text. Reward models, being LLMs themselves, have adversarial inputs: sequences that exploit weaknesses in the reward model and produce arbitrarily high scores while reading as gibberish. The KL term keeps the policy close to the SFT model.

Mathematically, the RLHF objective is:

In words: maximize expected reward, subject to a KL penalty that keeps you close to the reference model. The hyperparameter controls the tradeoff.

The punchline of DPO is that this objective has a closed-form optimal policy, and that closed form contains all the information you need to train against preference data directly: no reward model, no rollouts, no PPO. We’ll work through the derivation in a few paragraphs. But first, the conceptual move.

The DPO surprise

DPO’s pitch is structural. The RLHF pipeline has two coupled training problems: train a reward model that captures human preferences, then use that reward model to train a policy. DPO observes that these two steps are mathematically linked. Specifically, given the KL-regularized RLHF objective above, the optimal policy can be expressed as a function of the reward.

If you can express the reward in terms of the optimal policy, you don’t need to train a reward model separately. You can plug the implicit reward directly into the preference modeling framework, and now your policy is being trained to match preferences directly.

The result: a single supervised loss on preference pairs. No rollouts. No value network. No reward model.

Two models in memory: the policy being trained, and the reference model. That’s it. Compare to PPO + RLHF’s four. The memory and simplicity savings are significant.

But before showing the math, I want to land the right intuition for what DPO is doing, because there’s a misreading of DPO that’s tempting and wrong.

The key intuition: relative to the reference

The wrong reading of DPO: “It’s preference training. Increase chosen, decrease rejected.” Roughly what happens, but it misses the structural point.

DPO doesn’t push chosen up and rejected down in absolute terms. The intuition: if the reference already strongly preferred chosen, DPO barely moves. If the reference was indifferent or wrong, DPO pushes hard.

This is what preserves the KL anchor from RLHF. Naive preference training (just push chosen up, rejected down) would let the policy drift arbitrarily far from the reference. The “above what the reference already was” part is what keeps DPO solving the same KL-regularized problem PPO + RLHF solves. And it’s free: the KL anchor falls out of the derivation. No separate KL term in the loss.

The derivation: where the DPO loss comes from

I’m going to walk through this in four steps.

What the loss is actually doing

Look at the structure of the DPO loss. There are two log-probability ratios: one for the chosen response, one for the rejected. Each compares the current policy’s probability to the reference policy’s probability. The loss wants the chosen log-ratio to be larger than the rejected log-ratio.

This is the formal version of the intuition I described earlier. The training pressure isn’t “increase chosen, decrease rejected.” It’s “increase the chosen-over-rejected log-ratio relative to the reference.” If the reference model already gave higher probability to the chosen response, the loss is already partially satisfied; the gradient is small. If the reference was wrong (gave higher probability to the rejected response), the gradient is large, pushing hard against the reference.

The KL anchor isn’t a separate term. It’s built into the structure of the loss. Both numerators in the log-ratios are the policy being trained; both denominators are the reference. The reference’s behavior is the implicit constraint on what the policy can do.

The original RLHF objective had a KL penalty as an added term. DPO collapses everything: the reward, the KL anchor, the policy update; into a single supervised loss where the structure of the equation enforces all three.

What DPO meant for the field

DPO didn’t replace online RL. It carved out a specific lane. What DPO won is the open-source ecosystem of budget-constrained preference tuning: for teams without lab-scale compute, DPO is the dominant approach. The structural distinction that matters isn’t PPO vs. DPO; it’s online (fresh rollouts during training) vs. offline (a fixed preference dataset). Online has an exploration advantage; offline is simpler. For most use cases, the simplicity wins.

In mid-2023, doing RLHF was something a small number of well-resourced labs could do. PPO was complex and engineering-intensive. The infrastructure for online preference tuning was confined to teams that had built it. Most people who wanted to fine-tune a model on preference data couldn’t realistically do it.

By late 2023, that had changed. DPO + qLoRA + 4-bit quantization meant the entire preference-tuning recipe could be run on a single GPU. The Zephyr models, the Tülu series, the wave of open instruction-tuned models that followed: all of these were downstream of DPO making the pipeline accessible. The simplification of preference tuning meant orders of magnitude more people could do it.

Appendix: Deriving the optimal policy

This is the gory derivation I deferred in Step 2: the manipulation that takes the KL-regularized RLHF objective and produces the closed-form optimal policy. I worked through this in late 2023, a few months after the DPO paper came out, because I wanted to convince myself the algorithm was real rather than take the result on faith. What follows is essentially that derivation, cleaned up. Skip this section if you’re satisfied with the form of the result. Read on if you want to see how the partition function shows up.

We start with the RLHF objective:

The KL divergence can be written as an expectation:

So the inner expectation in our objective becomes:

Divide by to convert the maximization into a minimization, and the sign of the reward term flips:

Use the identity to rewrite the reward term:

Combine the two logarithms:

The denominator inside the log isn’t a probability distribution: it doesn’t sum to one over . We can fix this by introducing a normalizer :

and defining:

By construction, is a valid probability distribution: it’s nonnegative everywhere and sums to one over . Now we substitute it back into our objective. Add and subtract inside the bracket (no change, since these cancel):

Use the identity on the first term:

Combine the first two log terms using :

Recognize the denominator inside the log as :

The term doesn’t depend on , so it’s a constant with respect to the optimization. We can drop it. What’s left is exactly a KL divergence:

KL divergence is minimized when the two distributions are equal. So the minimizer is :

That’s the closed-form optimal policy. The rest of the DPO derivation (Steps 3 and 4 in the main text) follows from rearranging this expression to extract the implicit reward and plugging it into Bradley-Terry.

The partition function is the technical price of admission for this derivation. It would be intractable to compute directly as it requires summing over all possible completions, but it cancels out in the final DPO loss because Bradley-Terry only cares about reward differences. That cancellation is what makes DPO work as a practical algorithm.

Why drop the critic?

In PPO, four model roles are active during training:

• Policy: the LLM being trained
• Value network: the critic, also trained, typically same architecture as the policy
• Reference model: frozen copy of the SFT model, for KL anchoring
• Reward model: frozen, scores complete responses

That’s four LLM-sized things in the picture. The actual memory cost depends on a lot of implementation details (whether the policy and value share a body, how parameters are sharded across GPUs, optimizer state for the trained models) but the structural fact stands: PPO involves four model roles, and reducing that count is a real win.

Two of those roles you can’t avoid. The reference model is structurally necessary; without it you have no KL anchor and the policy drifts into reward-hacking gibberish. The reward model (or verifier) provides the training signal in the first place.

But the value network? It’s only there to compute advantages. If you can compute advantages another way, you can drop it.

A natural question is whether the reward model could just take over. It can’t, and the reason matters.

The reward model and the value network do fundamentally different jobs. The reward model answers: “How good is this complete response?” It takes a (prompt, response) pair and outputs a scalar. It was trained on human preferences over finished outputs, so it has no signal for half-finished responses. The value network answers a different question: “Given this partial state, what total reward do I expect by the end?” It takes a state (prompt plus tokens-so-far) and predicts where things are heading.

These are different questions. One scores finished things. The other predicts where things are heading. The reward model can’t do the second job because no human ever rated half-responses during preference data collection.

So the value network isn’t redundant. PPO needs it because GAE (the per-token advantage estimator) needs to know the expected outcome at every state along the trajectory, and only the value network can provide that.

GRPO’s move isn’t to substitute the reward model for the value network. It’s to give up on per-token advantage estimation entirely.

Group-relative advantages

Here’s the new advantage computation. For each prompt, generate rollouts (typically 4 to 16). Score each one with the verifier. The advantage of the -th rollout is:

That’s it. Each rollout gets one scalar advantage (“how much better than my sibling rollouts did this one do?”) applied identically to every token in that rollout.

This is a different kind of advantage than what PPO uses. PPO had a different advantage per token, derived from GAE; GRPO has the same advantage for every token in a rollout, derived from how the rollout’s reward compared to its siblings’. The credit is assigned at the trajectory level, not the token level. That’s the cost of dropping the value network: less granular credit assignment. But it’s also why GRPO works without per-token value estimates.

Let me make this concrete with the example I’ll use throughout: training a model to solve grade-school math word problems.

“Maya is buying tickets for a concert. Adult tickets cost $12 each and child tickets cost $8 each. She buys 9 tickets total and spends $92. How many adult tickets did she buy?”

The answer is 5. The verifier extracts the model’s final boxed answer and checks: 1 if correct, 0 if not.

Sample 4 rollouts ():

• Rollout A: sets up the equation , solves cleanly, ends with \boxed{5}. Reward: 1.
• Rollout B: tries and but makes an arithmetic slip, ends with \boxed{6}. Reward: 0.
• Rollout C: solves it correctly via testing values, ends with \boxed{5}. Reward: 1.
• Rollout D: confuses adult and child prices, ends with \boxed{4}. Reward: 0.

Compute the group statistics. The mean is straightforward:

For the standard deviation, take squared deviations from the mean and sum them:

Then divide by (only 3 of the 4 deviations are independent once the mean is fixed) and take the square root:

(Some implementations divide by instead, giving . The choice rescales all advantages by the same factor and doesn’t change their signs, but it can matter when comparing hyperparameters across codebases.)

Advantages are then :

• (rollout A)
• (rollout B)
• (rollout C)
• (rollout D)

In the gradient update, every token in rollouts A and C gets multiplied by , and every token in rollouts B and D gets multiplied by . The policy gets pushed toward the trajectories that solved the problem, away from the ones that didn’t.

The intuition is clean: of these attempts, which ones got it right? Reinforce those.

Everything else looks PPO-shaped

With the advantage computed, the rest of the algorithm is structurally identical to PPO.

The clipped policy loss has the same form as PPO. Same importance ratio:

Same min-of-clipped-and-unclipped construction for the per-token policy loss:

The only difference from PPO: is now the trajectory-level group-relative advantage rather than the token-level GAE advantage. Every token in a given rollout shares the same . Clipping mechanism, importance ratio, asymmetric min; all unchanged.

The multi-epoch inner loop is the same: collect rollouts, do epochs of mini-batched gradient steps, discard, repeat.

KL anchoring is also the same idea, though implemented differently. PPO often folds the KL term into the per-token rewards before computing advantages. GRPO can’t easily do this, as rewards in GRPO arrive only at the end of each rollout, with no per-token reward to fold into. So GRPO adds the KL as a separate per-token loss term:

This is the simple log-ratio estimator of KL; same as PPO uses. The two per-token losses then combine into a single scalar that gets backpropped:

Two separate per-token loss components, summed with a weight , averaged across all tokens in the mini-batch. Both contribute their own gradients to the policy parameters.

If you understand PPO, you understand most of GRPO. Two things change: how the advantage gets computed, and what models are in memory. Everything else carries over.

Why this works for reasoning specifically

GRPO didn’t have to be a reasoning algorithm. The math is generic. You could use it for any task with a scalar reward at the end. But there’s a structural reason it works especially well for reasoning.

Reasoning tasks have two properties that are awkward for PPO:

Rewards are usually verifiable. Math problems have right answers. Code can be tested. The reward is a deterministic check (“did the code pass the tests?”), not a learned approximation of human preferences. This is the world of Reinforcement Learning with Verifiable Rewards (RLVR) and it’s the world GRPO was made for.

Reward is genuinely sparse. You don’t know if a chain of thought is good until you see whether it produced the right final answer. Intermediate tokens can look identical between a trajectory that ends correctly and one that ends wrong, until the very end.

PPO handles sparse rewards via the value network, which has to learn “how good does this partial reasoning trajectory look?” That’s a hard regression problem. Reasoning trajectories can look identical for many tokens, then diverge wildly at the answer. The value network struggles to fit this signal cleanly, and the resulting advantages are noisy.

GRPO sidesteps the problem entirely. It doesn’t try to predict per-token value. It just compares whole rollouts to each other.

There’s also a practical reason GRPO and verifiable rewards pair well. GRPO needs rollouts per prompt to compute group statistics: meaning the reward function gets called times per prompt. If the reward function is a few lines of Python that runs in microseconds (a math correctness checker, a code test runner), scaling to is free. If the reward function is a 7B-parameter reward model requiring a full forward pass, inference cost gets expensive fast. GRPO’s economics work in the verifier setting.

This is also where the practical scope of “verifiable” matters. Verifiable rewards work cleanly for math (correctness check on the final answer), code (test cases), formal proofs (proof assistants), and constrained generation (regex compliance). They don’t work for tasks like “is this response helpful?” or “is this writing tasteful?”. Those are inherently subjective and need RLHF-style learned rewards. So GRPO + RLVR doesn’t replace RLHF for general assistant training. It’s a complementary tool for a specific (and increasingly important) class of tasks.

What R1-Zero showed

When DeepSeek published the R1 paper in early 2025, they actually published two models, and the distinction is informative.

DeepSeek-R1 is the deployable model. Trained with a multi-stage pipeline: cold-start SFT, then GRPO + RLVR, then more SFT, then more RL. This is the one people actually use.

DeepSeek-R1-Zero is the surprising one. Trained with pure RL (GRPO + verifiable rewards) directly on the DeepSeek-V3 base model, with no SFT stage at all. R1-Zero showed that reasoning behavior could emerge from RL alone, without any supervised demonstrations of how to reason. The DeepSeek paper documents the specific behaviors that emerged during training:

• Generating progressively longer chains of thought as training continued: the model learning, as training progresses, that more inference-time compute helped it solve harder problems.
• Reflecting on its own solutions, revisiting and evaluating earlier reasoning steps.
• Exploring alternative approaches when a first attempt didn’t pan out.
• Critiquing intermediate steps and catching its own errors.

None of this was programmed in. The RL process simply rewarded the model for producing correct answers in the right format, and the rest emerged. The DeepSeek paper calls this a “self-evolution”: the model learning, through exploration and reward, how to reason without ever being shown an example.

R1-Zero’s outputs were rough at the surface. The model would mix languages mid-response (sometimes Chinese in the middle of an English answer), use inconsistent formatting, occasionally produce incoherent passages. Strong reasoning, weak presentation. That’s what made the deployable R1 require additional SFT stages on top, to clean up the surface presentation while keeping the reasoning capability.

The R1-Zero result changed how people thought about what RL could do. The conventional wisdom had been that you needed SFT to give the model a starting distribution before any RL could help. R1-Zero showed that’s not strictly true for capabilities that have verifiable signals. Pure RL on a base model can produce reasoning. You don’t need to demonstrate how to reason; you just need to be able to verify when it worked.

The current consensus pipeline for production reasoning models combines: SFT for assistant behavior, preference tuning for helpfulness/harmlessness/honesty, RLVR for reasoning capability. Each stage does a job the others can’t do well. R1-Zero showed the RL stage could work without the prior stages; R1 showed why production models include them anyway.

Where GRPO fits in the broader landscape

The post-training landscape as of 2026 has split into a few distinct lanes.

Frontier general RLHF (proprietary, OpenAI / Anthropic / Google) still uses PPO-based or proprietary variants of PPO. Frontier labs have the engineering resources to handle PPO’s complexity and care about every fraction of a percent of quality. PPO is online, it generates new rollouts and learns from them, whereas DPO is offline, learning from a fixed preference dataset. That exploratory advantage seems to matter at the frontier. The frontier labs don’t share their pipelines publicly, so it’s hard to know exactly what they do, but PPO and its descendants remain in use.

Open-source preference tuning has largely migrated to DPO and its variants. DPO is much simpler than PPO: no rollouts, no value network, no reward model, just a clever loss on preference pairs. For most open-source teams, the simplicity is worth more than whatever quality edge PPO might provide.

Reasoning models use GRPO + RLVR. This is the lane DeepSeek pioneered with R1, and most open-source reasoning models follow the same recipe.

The practical takeaway: PPO and DPO/variants share the general-RLHF space (frontier vs open-source split); GRPO + RLVR owns reasoning. For practitioners, the choice depends on what you have. Preference data and a tight budget? DPO. Preference data and frontier-quality ambitions? PPO + RLHF. Verifiable rewards (math, code, theorem proving)? GRPO + RLVR.

Summary

In the broader arc that this post sits in, the next piece is DPO, the offline alternative that displaced PPO in much of the open-source ecosystem. Conceptually different from GRPO, DPO collapses the whole RL pipeline into a single supervised loss, but solves a related problem.

There’s a broader point worth naming here. PPO is hard. The four-models setup, the value network’s instability, and the sensitivity to hyperparameters are real engineering challenges. For years, these kept serious RL research mostly confined to a handful of well-resourced labs. GRPO didn’t just simplify an algorithm. It made RL training accessible to teams that couldn’t have run a stable PPO pipeline if they tried. It was a story about an algorithm simple enough that more people could actually use it.

REINFORCE in one breath

Reinforcement learning is the framework where an agent takes actions in an environment, the environment responds with reward and a new state, and the agent’s job is to find a strategy, a policy, that collects as much reward as possible. Unlike supervised learning, the agent generates its own training data by acting. Bad policies generate bad data. That’s what makes RL hard in a way classification never is.

The simplest policy gradient algorithm is REINFORCE. Run the policy. Observe what happens. Update the policy so good actions become more likely and bad ones less likely. The objective being maximized:

Reading the parts: is the probability the policy assigns to the action that was actually taken in state . is the return from time : the sum of future rewards from that point to the end of the trajectory.

for episode in range(num_episodes):
    # 1. Roll out a trajectory under the current policy
    states, actions, rewards = rollout(env, policy)

    # 2. Compute returns G_t for each timestep, working backward
    returns = []
    G = 0
    for r in reversed(rewards):
        G = r + gamma * G
        returns.insert(0, G)

    # 3. Form the loss and update
    loss = 0
    for s_t, a_t, G_t in zip(states, actions, returns):
        logits = policy(s_t)
        log_probs = F.log_softmax(logits, dim=-1)
        loss = loss - G_t * log_probs[a_t]

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

The two problems

REINFORCE works. People used it. Two problems become obvious the moment you scale it.

Problem 1: high variance. The “total reward that followed an action” varies wildly across rollouts even when the action itself was fine. Take the same action twice from the same state under a stochastic policy and you might see very different total rewards because what happened afterward was different. The training signal jumps around. Learning is slow and unstable.

Problem 2: data inefficiency. A trajectory is expensive to collect. For an LLM, sampling a 500-token response takes real GPU time. REINFORCE throws each one away after a single gradient update, because after that update the policy has changed and the old data is no longer “from the right distribution”: a sense we’ll make precise when we get to importance sampling.

PPO fixes both. The first three fixes attack variance. The last two attack data inefficiency. Each fix introduces machinery the next one needs, so they’re easiest to read in order.

PPO as five fixes

Advantage = (what happened) − (what we'd expect on average from this state)

advantages = []
gae = 0
for t in reversed(range(len(rewards))):
    delta = rewards[t] + gamma * values[t + 1] - values[t]
    gae = delta + gamma * lam * gae
    advantages.insert(0, gae)

Fix 5: Clip the importance ratio

The four fixes above (advantage, value network, GAE, importance sampling) were already standard prior to PPO.

To address the instability of importance sampling, TRPO (Trust Region Policy Optimization) added an explicit constraint: the new policy must stay within a KL-divergence bound of the old. This works, but it requires solving a constrained optimization with second-order methods and the implementation is a pain.

PPO’s contribution to RL, is a single trick that gets most of the stability of TRPO without TRPO’s complexity: If a sample is already pushing the policy in some direction, stop letting it push further.

Concretely:

In words: take the importance-corrected advantage; also take a version where the ratio is clipped to ; use whichever is smaller. The hyperparameter controls “how far is too far”: typically 0.2, so ratios are constrained to roughly before clipping kicks in.

The asymmetric min matters. The interesting question is when the min picks the unclipped value and when it picks the clipped one.

Advantage	Ratio	Unclipped	Clipped	min picks	What this means
(good)	(moved toward)	(large +)	(smaller +)	clipped	Already pushed right way; cap further reward
(good)	(moved away)	(small +)	(larger +)	unclipped	Wrong direction; let full corrective gradient through
(bad)	(moved toward)	(very −)	(less −)	unclipped	Wrong direction; let full corrective gradient through
(bad)	(moved away)	(less −)	(more −)	clipped	Already pushed right way; cap further punishment

The pattern: when the policy has moved in the right direction, the min picks the clipped value, and we stop reinforcing that direction beyond the band. When the policy has moved in the wrong direction, the min picks the unclipped value, and the full corrective gradient comes through.

Three lines of code:

ratio = torch.exp(log_probs_new - log_probs_old)  # = pi_new / pi_old
surr1 = ratio * advantages
surr2 = torch.clamp(ratio, 1 - epsilon, 1 + epsilon) * advantages
policy_loss = -torch.min(surr1, surr2).mean()

That’s PPO. Five fixes, two problems solved.

The unpacking chain

I want to spend the rest of this piece on how to compute the clipped objective:

In particular, let’s drill through how to calculate the advantage . The advantage is the deepest piece because it is defined recursively, through several layers of intermediate quantities. The plan: start from the GAE definition at the top and unpack downward until we hit primitives we can actually compute.

Level 1: GAE as a weighted average of n-step advantages:

This is what GAE is by definition. The parameter controls a bias-variance trade-off. At , GAE collapses to the 1-step estimate : take one real step of reward, then use the value function to estimate everything that comes after. This is low variance because only one stochastic reward enters, but it inherits whatever errors has. Larger shifts weight toward longer-horizon estimates, which use more actual rewards and rely less on . That lowers the bias but accumulates variance from the rewards themselves. The coefficients form a geometric series that sums to 1, so GAE is a properly normalized weighted average.

The expression is in terms of n-step advantages , which themselves need unpacking. Conceptually, an n-step advantage estimates how much better the trajectory turned out than the value function had predicted, by looking steps ahead before bootstrapping with the value function.

Level 2: n-step advantages as sums of TD residuals:

Each n-step advantage is a sum of TD residuals weighted by . The TD residual is a 1-step “surprise”: the gap between what actually happened in one step and what the value function had predicted. Summing of these gives the -step advantage. The n-step structure is now reduced to TD residuals. We still haven’t said what a TD residual is.

Level 3: TD residuals expanded into rewards and value predictions:

A TD residual is defined as:

Reading left to right: the actual reward at time , plus the discounted value of where we ended up, minus what we had predicted before taking the action. A one-step prediction error of the value function.

Plugging the definition into the 2-step advantage:

Distributing the in the second bracket:

The and cancel, leaving:

This is the standard 2-step formula. The cancellation is what is meant by the interior value terms telescope away: all the intermediate value predictions drop out, leaving only the rewards along the trajectory and the value function evaluated at the endpoints.

At this level, GAE has been reduced to two computable building blocks: per-step rewards and critic predictions . The critic predictions are direct outputs of the value network: a forward pass. They don’t need further unpacking. The rewards do.

Level 4: the per-step reward.

In standard RL, is whatever the environment gives you. Atari emits a score, a board game emits +1 or -1 at the end. There is nothing to unpack: is a primitive.

For LLMs, this is the level where the application-specific machinery shows up. RLHF, RLVR, and other variants all attach at exactly this level: they each define differently. That definition is what turns a generic policy gradient into RLHF, RLVR, or any other LLM-RL variant. Everything above this level (GAE, n-step advantages, TD residuals, the clipped objective) stays the same.

Two common forms:

• RLHF. A separate reward model is trained on human preference data (pairs of “this response is better than that one”). The reward model then scores each completion the policy generates, and that score is .
• RLVR. A verifier checks the answer against ground truth: did the math come out right, did the code pass the tests, did the extracted JSON match the schema? The verifier emits a numeric score, and that is .

Most LLM-RL recipes also add a KL penalty to : a term that penalizes the policy for drifting too far from a frozen reference (usually the base model before RL). This keeps the trained model from collapsing into degenerate high-reward outputs.

Apart from these substitutions, the algorithm is unchanged. The clipped objective, the GAE advantage, the value model, the policy update: all the same. Switching from RLHF to RLVR is switching what is, nothing more.

The full chain, top to bottom: clipped objective → GAE advantage → n-step advantages → TD residuals → rewards and value predictions → for LLMs, log-probs and a learned or verified reward signal. Each level answers “but how do you actually compute the thing one level up?” Eventually you bottom out at things produced by neural-network forward passes, plus whatever reward signal your problem provides.

This is the picture I wish I’d had when I first read the PPO paper. The clipped objective at the top, the unpacking chain underneath, and the application-specific reward sitting at the bottom waiting to be plugged in. Once you have it, the rest is just engineering.

What’s next

PPO was the original workhorse for RLHF, but the field has moved. Two algorithms have displaced it in different settings:

• DPO (Direct Preference Optimization) skips the reward model and the RL machinery entirely. Train directly on preference pairs with a clever loss derived from the Bradley-Terry math that justifies RLHF. Much simpler: no PPO, no rollouts, no value network. It also works surprisingly well. It’s displaced PPO in many open-source pipelines.
• GRPO (Group Relative Policy Optimization) is a PPO variant from DeepSeek that drops the value network and computes advantages by comparing rollouts to each other. Memory-efficient and well-suited for verifiable reward settings like math and code, where you don’t need a learned reward model, and a programmatic verifier provides the signal. GRPO is what powers most recent reasoning models.

Both are easier to understand once PPO is in your head. DPO is “what if we collapsed PPO into a single supervised loss?” GRPO is “what if we kept PPO but dropped the critic?”

1. The objective and its loss form

Restating, with the notation now grounded:

• is the objective we want to maximize.
• is the network’s parameters.
• is the probability the policy assigns to action in state .
• is the return from time : positive when the trajectory was good, negative when bad.

Neural networks are typically trained by minimizing a loss, so define

Minimizing is equivalent to maximizing . The sign flip is purely a notational convenience for using minimize-by-default optimizers.

2. Why the loss form does the right thing

Look at one term in isolation: , where .

Good action (). If increases, increases (becomes less negative), so decreases. Minimizing encourages to grow.

Bad action (). Then is positive, so . If decreases, becomes more negative, so decreases. Minimizing encourages to shrink.

So the form of encodes the right thing: minimizing it pushes good-action probabilities up and bad-action probabilities down. This is the easy part of the argument. It tells us that if probabilities move appropriately, the loss decreases. The harder question is whether gradient descent on actually causes them to move that way. For that, we need to look one layer deeper.

3. A useful reframe: REINFORCE is supervised learning with two tweaks

Before diving into the gradient mechanics, here’s a mental model that makes the REINFORCE loss feel less arbitrary and connects it to something you already know.

In standard supervised classification, the cross-entropy loss is

where is the true label for input . Minimizing this pushes the network’s predicted probability of the correct class upward.

Now compare to REINFORCE:

These are the same form, with two changes:

1. Replace the label with the sampled action. We don’t know what action should have been taken. There’s no oracle telling us “the correct move from was .” So we use the action the policy actually sampled as a kind of pseudo-label.
2. Multiply each term by . Supervised learning treats every example as equally important. Every label is “correct” by definition. REINFORCE doesn’t have that luxury, so it weights each pseudo-labeled example by how good the outcome turned out to be. Positive scales the term up (do more of this); negative flips its sign (do less of this).

Read this way, REINFORCE is just supervised learning where (a) the agent generates its own pseudo-labels by sampling, and (b) each example is weighted by how good the outcome turned out to be.

This means the gradient mechanics we’re about to derive are exactly the gradient mechanics of supervised cross-entropy, scaled by . If you’ve ever computed the gradient of cross-entropy through a softmax, you’ve already done most of the work. The next few sections are just making that work explicit and showing how the scaling rides along.

4. The network outputs logits, not probabilities

The neural network does not directly output probabilities. It outputs logits, which are then converted to probabilities via softmax. With two actions:

For example, with and :

Logits determine probabilities, and changing logits changes probabilities. Gradient descent updates the logits (and ultimately the weights that produce them):

So the question becomes: what is , and does its sign push the logit in a direction that makes the action’s probability move correctly?

5. The chain rule

Focus on one step: . By the chain rule,

We need .

6. Deriving the key softmax gradient

Starting from the softmax and using :

Differentiate term by term. The first term gives . For the second, let . Then (the piece doesn’t depend on ), so

Putting it together:

This is a clean, elegant result: as the action becomes more confident (), the gradient shrinks toward zero. There’s nothing left to push.

7. The gradient on the loss

Substituting back:

This is the core equation. With it, we can check that gradient descent does the right thing.

Use and .

Good action (): increases, which pushes up.

Bad action (): decreases, which pushes down.

8. Closing the loop

Increasing enlarges the numerator of the softmax and pushes up; decreasing it pushes down. Combined with the gradient signs from the previous section:

9. The general update rule

Everything above was for a single timestep and a single logit. For the full sum across the trajectory, the same logic gives

and the gradient-descent step on becomes

This is the REINFORCE update:

Each parameter gets nudged in the direction , the direction that would increase the probability of the action that was actually taken, scaled by , the return that followed. Good actions () get reinforced; bad actions () get suppressed.

1. RL is a different kind of learning problem

In supervised learning, you have a dataset of inputs and labels, and you train a model to map one to the other. Reinforcement learning has no dataset. You have:

• an environment with states,
• an agent that takes actions and moves between states,
• a reward signal that tells the agent (after the fact) how well things are going.

The agent’s job is to figure out a strategy, a policy, that collects as much reward as possible over time. Unlike supervised learning, the data is generated by the agent’s own behavior, which means a bad policy generates bad data. This is what makes RL hard in a way classification never is.

Throughout these notes we’ll use a small grid world. The agent occupies a square and can move up, down, left, or right. A few squares are terminal: in one corner, in another, and a “penalty” square that punishes with . Most squares are empty.

2. The Markov Decision Process

The grid world is an instance of a more general structure called a Markov Decision Process (MDP). An MDP has:

• a set of states ,
• a set of actions ,
• a reward function : how much reward you get for taking action in state ,
• transition dynamics : what state you end up in next.

The basic transition mechanic is straightforward: at each timestep, the agent is in state , picks action , and the environment responds by giving reward and moving the agent to a new state . We write this as .

Concretely in the grid world: if the agent is at and chooses “right,” then , , , and (assuming a per-step fuel cost; more on that below). The next iteration starts from , which becomes the new “current state,” and the cycle repeats until the agent hits a terminal square.

The next state depends on which action was taken. From state :

• Action leads to
• Action leads to
• Action leads to
• Action leads to

In a deterministic environment like this grid world, uniquely determines . In a stochastic environment, taking action in state gives a distribution over possible next states. Think of a robot whose wheels sometimes slip. We’ll come back to this distinction when we write the Bellman equation in its general form.

3. Trajectories and returns

The agent acting in the environment generates a trajectory: the full sequence of states, actions, and rewards from the start of an episode to the end:

The Greek letter (“tau”) is the standard symbol for it. A trajectory is also sometimes called an “episode” or a “rollout.” All three terms mean the same thing: one run of the agent through the environment, from start to finish.

The return

The return from time onward is the sum of future rewards:

This is the most important quantity in REINFORCE. Three things to internalize about it.

is forward-looking. It only counts what happens from onward. Rewards collected before timestep don’t enter into it. Each timestep in a trajectory has its own , and earlier timesteps generally have more future ahead of them.

is empirical, not predicted. It’s a number you measure after running an episode, by summing the rewards you actually got. It’s not a parameter, not an output of any model. Just an observation.

is different from , and the distinction matters. is the single-step reward at timestep : what the agent got right then. is the cumulative return from onward: the sum of all rewards from that point to the end of the episode.

To make this concrete, here’s a trajectory in the grid world:

The rewards are . The returns at each timestep:

• 
• 
• 
• 

Notice: but . They’re not the same number, and they’re not measuring the same thing.

Earlier timesteps had to pay more fuel costs before reaching the , so their returns are smaller. From at the end, the agent is one step from the reward; from at the start, it had to walk all the way there.

4. Values and the Bellman equation

Before talking about how to learn a policy, it helps to ask what makes a state “good.” A natural answer: a state is good if you can collect a lot of reward starting from it. Define the value of a state as the return you’d expect if you played optimally from there.

In the simplest version of the grid world, where only terminal squares give reward, every non-terminal state satisfies:

where is the state reached by taking action from . The value of a state is the best value among its neighbors. Notice the action-dependence: depends on which you take. Writing it as (or sometimes ) makes that explicit and avoids the slightly sloppy notation that hides the dependence.

This is the Bellman equation in its simplest form, and it lets you propagate values from the terminal squares throughout the grid.

There’s a wrinkle: with this rule alone, the agent has no incentive to hurry. Wandering forever before reaching the is as good as walking straight there. To fix this, add a small per-step reward ; a “fuel cost”:

Now the agent prefers shorter paths.

A second refinement is the discount factor , which says future rewards matter less than present ones: the same reason a dollar today is worth more than a dollar next year. With both:

You might wonder why we need when we already have the step cost . They look redundant: both make the agent prefer short paths in the grid world. But they’re solving different problems:

• The step cost makes long paths expensive in absolute terms. Every step subtracts 1 from the total return. A 3-step path to nets ; a 10-step path nets .
• The discount factor makes future rewards worth less than present ones. A reward of received in 3 steps is worth today; received in 10 steps, it’s worth .

The practical reasons does work the step cost can’t:

1. encodes uncertainty about the future. A reward you might get in 100 steps is less trustworthy than one you’ll get in 2 steps. The world might change, the model might be wrong. captures this: distant predictions get less weight because they’re less reliable. A step cost just charges you for moving; it doesn’t model uncertainty.

2. doesn’t require knowing the right magnitude. A step cost only works if you tune it. Set in a grid where rewards are and the agent barely cares about path length; set and the agent refuses to move. is scale-free in this sense: creates a roughly-100-step planning horizon regardless of reward scale.

3. They compose differently with positive rewards along the way. In a “stay alive” task where each surviving step gives , a negative step cost would subtract from a signal that’s supposed to be additive. still works there: it just makes the agent prefer reward sooner rather than later.

So: the step cost is a reward-design choice (you, specifying the task, decide moving is bad), while is a property of the agent’s planning horizon (the agent decides distant futures matter less). They live at different levels.

Once is in the picture, the return is also typically written with discounting:

This is the version that shows up in most modern RL writing. The undiscounted form earlier was a special case with .

5. The general Bellman equation and expectations

In a stochastic environment, taking action in state doesn’t deterministically land you in : it gives a distribution over next states. The fully general Bellman optimality equation handles this with an expectation:

or written out as a sum:

The two forms are the same equation. Going from the first to the second is just unpacking the expectation, so it’s worth being clear about what an expectation is.

The expected value of a random variable is the weighted average of its possible values, where the weights are the probabilities:

That’s it. Roll a fair six-sided die: outcomes each with probability , and . List every possible outcome, multiply by how likely it is, add up.

A conditional expectation is the same weighted average, but using the conditional probabilities instead: “given that I know , what’s the expected value of ?”

Apply this to the Bellman expectation: the random thing is (taking action in state gives a random next state), and the quantity being averaged is . List all possible next states , weight each by , and sum:

Concrete example. Suppose the agent is in state , takes action , but the environment is noisy:

• 80% chance you actually go right, landing in with reward
• 20% chance you slip up, landing in with reward

Then the expected value inside the Bellman equation is

Each possible outcome contributes its value, weighted by how likely it is.

In the deterministic case, for one specific and for everything else. The sum collapses to a single term , and the expectation disappears. That’s why the simpler form is fine for the grid world.

6. vs : two flavors of value

So far we’ve only talked about , the value of a state. There’s a closely related object , the value of a state-action pair. Both come up constantly in RL, and they answer slightly different questions.

: the value of a state. Answers: “How good is it to be in state ?” Specifically, the expected return starting from and following the policy from there:

One number per state.

: the value of a state-action pair. Answers: “How good is it to take action from state ?” The expected return if you take in state and then follow the policy:

One number per state-action pair. With 4 actions, is four numbers, one for each action you might take from .

The two are related. is what you get when you average over the actions the policy would take:

If the policy is greedy (always picks the best action), this collapses to . So is the summary, is the breakdown.

Why does this matter? Because is more useful for acting. Suppose you have a value function and you need to choose an action.

• With : you’d need to look ahead one step for each candidate action, see what state you’d land in, and check the value there. This requires knowing the environment dynamics.
• With : you read off the values directly and pick the action with the highest one. No model of the environment needed.

This is why deep RL methods that learn value functions almost always learn , not . DQN (Deep Q-Network, the famous 2013 Atari paper) learns directly: the network takes a state and outputs one number per action; you act greedily by picking the argmax. You don’t need to know how the Atari game works internally; the -values tell you which button to press.

7. Why neural networks?

Iterating Bellman across a 16-square grid is fine. Iterating across the state space of Go, which has more positions than there are atoms in the universe, is not. You can’t store one number per state, let alone visit each state repeatedly to update it.

The fix is to use a neural network as a function approximator in place of a giant lookup table. The network takes a state as input and outputs either:

• a value estimate ( or ): “from this state, you can expect roughly this much return,” or
• a policy: “from this state, here are the probabilities of each action.”

These two choices give the two main families of deep RL.

A value network approximates or . It’s trained by enforcing the Bellman equation: the network’s prediction at should match times its prediction at the next state . Once trained, you act greedily: pick the action with the highest predicted value. DQN is the canonical example.

A policy network approximates the policy directly. It takes a state and outputs probabilities over actions. There’s no Bellman equation, no value estimation. You train the network so good actions become more likely and bad actions become less likely. This is the world REINFORCE lives in.

Why use one over the other? Policy networks have several advantages worth knowing:

• They naturally output stochastic policies (next section), which matters for exploration.
• They handle continuous action spaces gracefully: a value network would need to argmax over an infinite set.
• They’re often easier to train end-to-end with backpropagation, which is the whole point of writing as something you can take a gradient of.

Value networks tend to be more sample-efficient when they work, but harder to stabilize. The two approaches aren’t mutually exclusive: actor-critic methods use both.

8. Stochastic policies and explore/exploit

A deterministic policy outputs the single best action: “from , always go right.” A stochastic policy outputs a distribution: “from , go right with probability , up with , down with , left with .”

Why prefer the stochastic version? Imagine the agent has found the reward and learned to walk to it. With a deterministic policy, it will never deviate and will never discover the reward sitting in another corner. Randomness is what lets the agent explore. A stochastic policy mostly exploits what it knows but occasionally tries something else, which is how it keeps improving.

This is also why the policy network outputs a softmax over actions rather than a single chosen action. The softmax gives a differentiable, naturally-stochastic parameterization; exactly what the gradient mechanics in Part II need.

9. The policy notation

We’ve referenced “the policy” repeatedly. It’s now time to give it a proper symbol.

is the policy: the neural network that, given a state, outputs a probability distribution over actions. The Greek letter is the standard symbol for a policy. The subscript reminds you that the policy is parameterized by the network’s weights . Different weights produce different policies.

Three notations all show up in practice and mean closely related things:

• on its own refers to the whole policy as an object: the full mapping from states to action distributions.
• is the distribution over actions when the agent is in state : for a 4-action grid world, that’s 4 numbers summing to 1.
• is a single number: “the probability that policy (with weights ) assigns to action when in state .”

Concretely, for the grid world: the network takes a state like as input, runs it through some layers, ends with a softmax, and outputs , , , . Four probabilities summing to 1. Sample one of them to get the action the agent actually takes.

10. The objective from a single trajectory

We can finally write down what REINFORCE is trying to do.

Run the agent. It produces a trajectory . For each timestep in that trajectory, you can compute two things:

• : the return that followed (you get this by summing rewards from onward).
• : the probability the policy assigned to the action that was actually taken.

The objective to maximize is:

Notice the structure: each term in the sum has two factors: the return , and the log-probability of the action that produced it.

11. The objective as an expectation

The single-trajectory form is what you actually compute in code. But to be precise about what really is, you need to think of it as an expectation.

Each time you run the agent, you get a different trajectory, because the policy is stochastic and the environment may be too. The “true” objective averages over all the trajectories the policy could possibly produce:

Reading the new piece of notation:

• reads as “trajectory sampled from ”: meaning is generated by running the policy in the environment.
• reads as “the expectation, where the randomness comes from sampling trajectories from the policy.”

So in plain English: “run the policy infinitely many times, compute the bracketed sum on each trajectory, and average the results.”

12. Putting the training loop together

We now have all the pieces. The REINFORCE training loop:

1. Roll out a trajectory. Starting from some initial state, sample actions from the current policy until the episode ends. Record states, actions, and rewards along the way.

2. Compute the returns. For each timestep in the trajectory, sum the (discounted) rewards from onward to get .

3. Form the objective. Plug into . This is your single-sample estimate of the true expected objective.

4. Update the parameters. Take a gradient step on .

5. Repeat. With updated parameters, roll out a new trajectory and do it again.

A useful intuition: the policy network is being shown its own past behavior and told do more of what worked, less of what didn’t. The return is the supervision signal: it plays the role that the label plays in supervised learning, except the agent generates it for itself by interacting with the environment.

import torch
import torch.nn.functional as F

# Hyperparameters
learning_rate = 1e-3
gamma = 0.99
num_episodes = 10000

# Policy network: state -> action logits
policy = PolicyNetwork()  # outputs logits over actions
optimizer = torch.optim.Adam(policy.parameters(), lr=learning_rate)

for episode in range(num_episodes):
    # 1. Roll out a trajectory
    states, actions, rewards = [], [], []
    state = env.reset()
    done = False
    while not done:
        logits = policy(state)                     # action logits for this state
        probs = F.softmax(logits, dim=-1)          # action probabilities
        action = torch.multinomial(probs, 1).item()  # sample an action
        next_state, reward, done = env.step(action)

        states.append(state)
        actions.append(action)
        rewards.append(reward)
        state = next_state

    # 2. Compute returns G_t for each timestep (working backward)
    returns = []
    G = 0
    for r in reversed(rewards):
        G = r + gamma * G
        returns.insert(0, G)
    returns = torch.tensor(returns)

    # 3. Form the loss
    # For each timestep, we want G_t * log pi(a_t | s_t).
    # We sum over t and negate (since optimizers minimize).
    loss = 0
    for s_t, a_t, G_t in zip(states, actions, returns):
        logits = policy(s_t)
        log_probs = F.log_softmax(logits, dim=-1)
        loss = loss - G_t * log_probs[a_t]

    # 4. Update parameters
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

13. The bridge to the gradient mechanics

This is the conceptual frame in which REINFORCE makes sense. We have an objective , written either as a single-trajectory sum or as an expectation over trajectories, and we want to maximize it. Equivalently, we minimize by gradient descent.

What this frame doesn’t tell you is why gradient descent on actually causes good actions to become more probable. That question lives one layer down. It’s about how the gradient flows through the softmax to the logits, and from the logits back through the network to the parameters. That’s what Part II is about.