Chapter 6

Reinforcement Learning (Policy Gradients)

The main optimization family for pushing a model toward higher reward while trying not to let it drift too far from useful behavior.

Foundations 1-3
SFT 4
Reward Signals 5
Policy Optimization 6-9
Preferences & Data 10-13
Controls & Product 14-17

The RLHF loop is policy optimization over text

A prompt enters the current policy, the model samples a completion, the completion gets scored, and that score drives a gradient update.

The key translation from classical RL is that the trajectory is a token sequence and the policy is an autoregressive language model.

Once you see that mapping, many RL concepts become much less mysterious.

Overview of the RLHF training loop. A prompt from the dataset is passed to the tuned policy, which generates a completion. The reward model scores this completion, while the frozen initial model computes log probabilities on the same text to calculate a KL penalty that prevents excessive drift. The combined reward signal then drives a reinforcement learning update to the policy parameters.
Chapter 6 Overview of the RLHF training loop. A prompt from the dataset is passed to the tuned policy, which generates a completion. The reward model scores this completion, while the frozen initial model computes log probabilities on the same text to calculate a KL penalty that prevents excessive drift. The combined reward signal then drives a reinforcement learning update to the policy parameters.
Basic REINFORCE architecture for language models. The shaped reward combines the reward model score with a KL penalty from the reference model. We build on this structure throughout the chapter.
Chapter 6 Basic REINFORCE architecture for language models. The shaped reward combines the reward model score with a KL penalty from the reference model. We build on this structure throughout the chapter.
REINFORCE Leave-One-Out (RLOO) architecture. Multiple completions per prompt provide a leave-one-out baseline for advantage estimation without learning a value function.
Chapter 6 REINFORCE Leave-One-Out (RLOO) architecture. Multiple completions per prompt provide a leave-one-out baseline for advantage estimation without learning a value function.

REINFORCE, RLOO, PPO, and GRPO are different variance-control strategies

All of these algorithms are trying to answer the same question: how do we turn noisy sampled rewards into stable updates on a huge language model?

REINFORCE is the clean baseline, RLOO uses multiple completions to build a baseline, PPO adds clipped trust-region style control plus a value function, and GRPO normalizes rewards within a group.

The differences matter because RLHF is often bottlenecked by variance and instability more than by missing reward signal.

PPO architecture. A learned value function enables Generalized Advantage Estimation (GAE) for per-token advantages, used with a clipped surrogate objective.
Chapter 6 PPO architecture. A learned value function enables Generalized Advantage Estimation (GAE) for per-token advantages, used with a clipped surrogate objective.
Visualization of the different regions of the PPO objective for a hypothetical advantage. The "trust region" would be described as the region where the log-ratio is within $1\pm\varepsilon$.
Chapter 6 Visualization of the different regions of the PPO objective for a hypothetical advantage. The "trust region" would be described as the region where the log-ratio is within 1±ε1\pm\varepsilon.

KL regularization is the governor

RL can improve reward quickly by moving the policy into strange regions. KL penalties keep the new policy near a reference model so the update stays behaviorally sane.

This is not a cosmetic term. It often determines whether the run improves or collapses.

A good mental model is that reward says where to go and KL says how recklessly you are allowed to get there.

Value function training uses on-policy rollouts to compute targets. The model predicts $V_t$ at each token, which is trained via MSE against the target return $\hat{V}_t$. The advantage $A_t = \hat{V}_t - V_t$ then weights the policy gradient update.
Chapter 6 Value function training uses on-policy rollouts to compute targets. The model predicts VtV_t at each token, which is trained via MSE against the target return V^t\hat{V}_t. The advantage At=V^tVtA_t = \hat{V}_t - V_t then weights the policy gradient update.
GRPO architecture. Advantages are normalized relative to the group mean and standard deviation. The KL penalty is applied directly in the loss rather than shaping the reward.
Chapter 6 GRPO architecture. Advantages are normalized relative to the group mean and standard deviation. The KL penalty is applied directly in the loss rather than shaping the reward.

As systems scale, orchestration becomes part of the algorithm

Large RLHF runs are not only about a loss function. They are also about queues, actor-learner splits, asynchronous rollout generation, and throughput constraints.

That engineering shape changes what algorithms are practical in real training systems.

For teaching purposes, the clean math is only half the story. The deployment system is the other half.

A comparison of the generation-update phases for synchronous or asynchronous RL training following Noukhovitch et al. 2024.
Chapter 6 A comparison of the generation-update phases for synchronous or asynchronous RL training following Noukhovitch et al. 2024.
An example distributed RL system, where two queues are managed to pass data to the learner and actor GPUs, which can both be synchronized with a distributed computing library such as Ray. Olmo Team 2025, license CC-BY.
Chapter 6 An example distributed RL system, where two queues are managed to pass data to the learner and actor GPUs, which can both be synchronized with a distributed computing library such as Ray. Olmo Team 2025, license CC-BY.

Review

5 quick checks

Questions with revealed answers.

1 What is the simplest way to describe RLHF policy optimization?

Sample completions from the current model, score them, and update the model to make high-scoring completions more likely.

2 Why do so many policy gradient variants exist in RLHF?

Because the core problem is turning noisy rewards on sampled text into updates that are stable enough for large language models.

3 What role does KL regularization play?

It restrains the policy from drifting too far from a reference model while reward optimization is happening.

4 How is PPO different from REINFORCE at a high level?

PPO adds extra control through clipping and typically a value function, making updates less brittle than plain REINFORCE.

5 Why is systems engineering part of the policy-gradient story?

Because rollout generation, actor-learner coordination, and asynchronous pipelines determine what is tractable at scale.