You have an SFT model and a reward model. The simplest thing to try is to sample responses, score them, and push up the log-probability of the ones that scored well, REINFORCE from §15.3. Run that for a while and you get a model that scores higher and higher while its text gets stranger. It repeats phrases the reward model happens to like, grows every answer to the maximum length, or drifts into a style no person would write, because the reward model was trained on comparisons of ordinary SFT outputs and has no idea what to say about strange ones. Ziegler et al. (2019) saw a sharper version of this when a sign bug made their model optimize for the text labelers hated most, and it did so fluently.

The fix that every RLHF system since 2019 uses is a leash. Reward the policy for what the reward model likes and charge it for every bit of probability it moves away from the SFT model. This section builds that objective, shows the distribution it aims for, turns it into a per-token reward that PPO can use, runs generalized advantage estimation over tokens, writes the clipped and masked loss, assembles the full training step, and counts what the four models in memory cost.

Four models

PPO-based RLHF keeps four networks, and it helps to name them before any equation.

The policy and reference start identical, so at the first step the KL penalty is zero. The value model is a whole Transformer with a scalar head, the critic of §15.4, and it is usually initialized from the reward model (InstructGPT, Stiennon et al. 2020, and Huang et al. 2024 all do this). The reward model already maps a conversation to one number on the reward's own scale, which is a far better start than a fresh head on the policy, even though it was only ever trained on finished responses. We write its parameters $\psi$ because $\phi$ belongs to the reward model.

The KL-regularized objective

Treat the prompt $x$ as the state and the whole response $y$ as one action, the view of §15.7. The objective is

with $\beta > 0$ the KL coefficient. In plain words, get high reward, and pay $\beta$ for every nat by which your response distribution differs from the reference's.

Section 15.7 gave the penalty's three jobs. It keeps the policy where the reward model was trained and can be trusted (§16.3 watched a reward model go from 96% right to 78% wrong as the policy moved), it keeps the model fluent and knowledgeable, and it stops the policy from collapsing onto one response. One more property is worth spelling out for RLHF. The penalty is the reverse KL, an average over the policy's own samples, which is mode-seeking (§0.5). It charges the policy for putting mass where the reference has little and lets it drop some of the reference's modes freely. That is what we want, since the point is to stop writing some of what the SFT model writes, without inventing text it would never write.

Recall from §15.7, which stated it and proved it in two lines, that the objective has a closed-form maximizer, the reference tilted toward high reward,

which §17.1 derives in full, in two ways, and turns into DPO. We only need what it says. In plain words, the ideal RLHF policy is the reference model with every response's probability multiplied by $e^{r/\beta}$ and renormalized. Three consequences matter for practice.

• RLHF can only reweight what the SFT model already writes. A response the reference gives zero probability stays at zero for every $\beta$. A response it gives tiny probability needs a large KL to become likely. That is one reason the quality and diversity of the SFT model bound what RLHF can reach (§14.1).
• Making one response certain costs $-\log\pi_{\text{ref}}(y\mid x)$ nats. If the reward model wildly overrates a response the reference writes 3% of the time, the policy can put all its mass there for a KL of $\log(1/0.03) \approx 3.5$ nats. Overoptimization in §16.5 is measured in exactly these units.
• $Z$ is intractable. It sums over every possible response, so nobody can compute $\pi^\star$ for a real model. PPO approaches it by sampling, and DPO sidesteps it algebraically.

The per-token form, and what RLHF adds to it

PPO works with per-step rewards. A response is $y = (y_1, \dots, y_T)$ and the state before token $t$ is $s_t = (x, y_{<t})$. Section 15.7 showed that the sequence log-ratio splits into per-token log-ratios, $\log\frac{\pi_\theta(y\mid x)}{\pi_{\text{ref}}(y\mid x)} = \sum_t k_t$ with $k_t = \log\pi_\theta(y_t\mid s_t) - \log\pi_{\text{ref}}(y_t\mid s_t)$, which turns the objective into a token-level MDP with the shaped reward

and that putting this penalty in the reward, rather than differentiating it in the loss, gives the correct gradient. In plain words, every token pays a KL toll for how much more likely the policy made it than the reference did, and the reward model's single score arrives as a bonus on the final token. The tolls and the bonus sum to $r_\phi(x,y) - \beta\sum_t k_t$, whose expectation is exactly $J$ for this prompt. InstructGPT describes it as a per-token KL penalty from the SFT model at each token, and it is the form in Ziegler et al. (2019), Stiennon et al. (2020), and Huang et al. (2024).

What §15.7 did not need, and every RLHF implementation does, is the handling of real generations.

• Which policy's log-probabilities? The $k_t$ are computed once, at rollout time, with the policy that generated the samples, $\pi_{\text{old}}$, and treated as constants for all the optimization epochs that follow.
• A single $k_t$ can be negative. If the policy sampled a token the reference liked more, that token earns a small positive reward. Only the expectation of the sum is a KL. Some implementations put the always-nonnegative $k_3$ estimate in the reward instead, or the exact per-token KL summed over the vocabulary, because both have lower variance as estimates of the KL's value. Neither is a free swap. Used as a fixed reward, only $k_1 = k_t$ gives the exact gradient of $J$ (§15.7 proves this). The exact per-token KL, for instance, depends on the state alone, so as a detached reward it gives no gradient at all through the token it is charged on, and its direct derivative has to be added to the loss separately.
• Where does the response end? Batched generation usually runs every sequence for the same number of steps, so tokens after a sequence's first end-of-sequence token are garbage that must be masked out. The response mask is built from the first end-of-sequence token, which is included, because choosing to stop was an action.
• What if it never ends? A response cut off at the length limit was never finished, and the reward model's opinion of a truncated answer means little. Huang et al. (2024) replace its score with a fixed penalty, $-1$ in their runs, which stops the policy from rambling to the limit.

import numpy as np

def rlhf_rewards(tokens, logp_old, logp_ref, scores, beta, eos_id, no_eos_score=-1.0):
    # placeholder: treats every token as real and puts the score on the last column
    B, T = tokens.shape
    mask = np.ones((B, T), bool)
    rewards = -beta * (logp_old - logp_ref)
    rewards[:, -1] += scores

Advantages over tokens

With per-token rewards in hand, the rest is the actor-critic machinery of Chapter 15, applied to a sequence of tokens. The value model predicts $V_\psi(s_t)$, the expected sum of future shaped rewards from state $s_t$ (in practice the value head is read at the position of each response token, so it has seen $x$ and $y_{<t}$ and predicts the rest). The TD error of token $t$ and the generalized advantage estimate of §15.5 are

with $V_\psi(s_{T+1}) = 0$ after the last token, because the episode is over, and $\hat A_{T+1} = 0$. The value targets are the returns $\hat R_t = \hat A_t + V_\psi(s_t)$. In plain words, each token's advantage is an exponentially weighted sum of how much better the following steps went than the critic expected.

RLHF uses $\gamma = 1$, no discounting, because the episode is short and the reward that matters arrives at the end. A discount below 1 would shrink the final score's influence on early tokens for no good reason. With $\gamma = 1$ the knob that matters is $\lambda$. At $\lambda = 1$ the advantage of token $t$ is the full return from $t$ minus the baseline $V_\psi(s_t)$, unbiased but noisy. At $\lambda = 0$ it is the one-step TD error, which leans on the critic. Stiennon et al. (2020) and Huang et al. (2024) used $\lambda = 0.95$.

A worked example by hand

Take a 4-token response with $\beta = 0.1$ and a reward model score of 1.2. The rollout and reference log-probabilities of its tokens are $(-1.2, -0.4, -2.0, -0.1)$ and $(-1.5, -0.5, -1.0, -0.1)$, so the log-ratios are $k = (0.3, 0.1, -1.0, 0)$ and the shaped rewards are $\tilde r = (-0.03, -0.01, 0.1, 1.2)$. The third token was less likely under the policy than under the reference, so it earns a small positive reward. Let the critic's values be $V = (0.5, 0.6, 0.9, 1.1)$. Then

• $\delta_4 = 1.2 + 0 - 1.1 = 0.10$, and $\hat A_4 = 0.10$,
• $\delta_3 = 0.1 + 1.1 - 0.9 = 0.30$, and $\hat A_3 = 0.30 + 0.95 \times 0.10 = 0.395$,
• $\delta_2 = -0.01 + 0.9 - 0.6 = 0.29$, and $\hat A_2 = 0.29 + 0.95\times 0.395 = 0.665$,
• $\delta_1 = -0.03 + 0.6 - 0.5 = 0.07$, and $\hat A_1 = 0.07 + 0.95 \times 0.665 = 0.702$.

The returns are $\hat R = \hat A + V = (1.202, 1.265, 1.295, 1.2)$. Every token's advantage is positive, because the response did better than the critic predicted from the start. The cell checks the arithmetic.

import numpy as np

def gae(rewards, values, mask, gamma=1.0, lam=0.95):
    raise NotImplementedError

The clipped loss, masked

The policy update is PPO's clipped surrogate (§15.6), applied at every response token. Rollouts are generated once with $\pi_{\text{old}}$, then reused for a few optimization epochs, so the new policy's tokens are reweighted by the probability ratio

and the policy loss over a batch with response mask $m_{b,t}$ (1 on real response tokens of sequence $b$, 0 on prompt and padding) is

with clip range $\epsilon$, typically 0.2. In plain words, raise the probability of tokens with positive advantage and lower those with negative advantage, but once a token's probability has moved more than $\epsilon$ relative to the rollout policy in the helpful direction, stop pushing. Here $\epsilon$ is the clip range, not the label-flip rate of §16.2.

The value model is trained by regression to the returns, often with its own clipping so it cannot move far from its rollout-time prediction $V_{\text{old}}$ either,

and the two are combined as $\mathcal{L}^{\text{clip}} + c_v\mathcal{L}^V$ when the policy and value share an optimizer (Huang et al. used $\epsilon_v = 0.2$ and $c_v = 0.1$). An entropy bonus, standard in PPO for games, is usually left out of RLHF, because the KL penalty to the reference already keeps the policy from collapsing.

Two practical points about the mask. First, the prompt tokens are never actions, so they are masked out of every loss, exactly as in the SFT loss of §14.3. Second, how you average matters. The formula above averages over all response tokens in the batch, so long responses carry more weight than short ones. Averaging within each sequence first and then over sequences weights every response equally, which changes which responses the update favors and interacts with length. That choice is at the heart of the length-bias analysis of Dr. GRPO in §18.4.

Advantages are usually whitened before the loss, shifted to mean 0 and scaled to standard deviation 1 over the real tokens of the batch, which makes the step size independent of the reward scale (§16.6 has the details and pitfalls).

import numpy as np

def ppo_policy_loss(logp_new, logp_old, adv, mask, eps=0.2):
    ratio = np.exp(logp_new - logp_old)            # placeholder: no clipping and no mask
    return float(-(ratio * adv).mean()), 0.0

Choosing beta, and the adaptive controller

The KL coefficient sets the exchange rate between reward and drift, and the right value depends on the reward model's scale, which is arbitrary (§16.3). Published values span an order of magnitude. Stiennon et al. (2020) used $\beta = 0.05$, InstructGPT $\beta = 0.02$, Llama 2 used 0.01 for its 7B and 13B models and 0.005 for 34B and 70B, and Anthropic used 0.001 and noted it might be nearly unnecessary (Bai et al. 2022). These numbers are not comparable across papers, because each reward model has its own scale.

Ziegler et al. (2019) avoided choosing by targeting a KL instead. After each batch, with measured KL $\hat\kappa$ and target $\kappa_{\text{targ}}$, the adaptive KL controller sets

with $K_\beta = 0.1$. In plain words, if the policy has drifted more than the target, raise the price of drift by up to 2% per step, and lower it if it has drifted less. Their targets were 6 nats for descriptiveness and 10 for sentiment. The clip keeps one noisy batch from swinging $\beta$, and the multiplicative update keeps $\beta$ positive.

import numpy as np

def update_beta(beta, kl, target, k_beta=0.1):
    return beta + k_beta * (kl - target)      # placeholder: additive and unclipped

def run_controller(beta0, kls, target):
    return []

The pretraining mix

The KL penalty keeps the policy near the SFT model on the prompts it is trained on. It says nothing about everything else the model used to be good at. InstructGPT found that PPO made its models worse on several public NLP benchmarks, SQuADv2 and DROP among them, an alignment tax paid for following instructions better. Their fix, called PPO-ptx, keeps a little pretraining running inside RL. With $z$ a sequence of text drawn from the pretraining distribution $\mathcal{D}_{\text{pretrain}}$ and $\pi_\theta(z)$ the probability the policy assigns to the whole sequence, InstructGPT maximizes

with $\beta = 0.02$ and $\gamma_{\text{ptx}} = 27.8$ (Ouyang et al. 2022). Setting $\gamma_{\text{ptx}} = 0$ gives plain PPO. In plain words, get high reward on your own samples, pay for drifting from SFT on those samples, and keep assigning high likelihood to ordinary text.

The two regularizers pull in different ways, and the difference is the point. The KL term is the reverse KL to the SFT model, measured on the policy's own samples, so it never looks at text the policy has stopped writing. The new term, up to a constant, is the forward KL from the pretraining data, since $\mathbb{E}_{z\sim\mathcal{D}_{\text{pretrain}}}[\log\pi_\theta(z)] = -H(\mathcal{D}_{\text{pretrain}}) - D_{\mathrm{KL}}\big(\mathcal{D}_{\text{pretrain}}\Vert\pi_\theta\big)$, where $H$ is the entropy of the data and does not depend on $\theta$. Forward KL is mode-covering (§0.5). It charges the policy for every kind of text in the corpus it has become unable to produce, which is exactly the capability loss the reverse KL cannot see. InstructGPT tested the alternative directly. With no pretraining mix, raising $\beta$ a hundredfold, to 2.0, still did not fix the regressions, and it cost a large drop in reward.

Mechanically it is cheap. For each PPO minibatch, InstructGPT computed the PPO gradient and a language-modeling gradient on a pretraining minibatch in consecutive steps, multiplied the second by $\gamma_{\text{ptx}}$, and accumulated both before the optimizer step, drawing 8 times more pretraining examples than RL episodes over the run. A value of 27.8 worked across model sizes from 1.3B to 175B, the labelers' ratings were insensitive to its exact value, and larger values traded away validation reward. It is not free in a subtler sense. The objective is no longer the one whose maximizer is the tilt $\pi_{\text{ref}}e^{r_\phi/\beta}/Z$, so the closed form above, and the DPO derivation built on it, describe PPO without the mix. And it did not remove every regression, which is why later stacks also measure the tax directly (§16.6).

One full training step

Here is the whole loop, in the order the computation happens. Each iteration takes a batch of prompts, InstructGPT and Llama 2 used 512 per batch.

1. Rollout. Generate one response per prompt with the current policy, $\pi_{\text{old}} = \pi_\theta$, at temperature 1 (some recipes use lower, Huang et al. used 0.7). Truncate at the first end-of-sequence token.
2. Score. Run the frozen reward model once per response, read at the last token. Apply the missing-end-of-sequence penalty if you use one. Optionally normalize the scores (§16.6).
3. Log-probabilities. One forward pass of the policy and one of the reference over prompt plus response give $\log\pi_{\text{old}}(y_t\mid s_t)$ and $\log\pi_{\text{ref}}(y_t\mid s_t)$ at every response token.
4. Values. One forward pass of the value model gives $V_{\text{old}}(s_t)$ at every response token.
5. Rewards and advantages. Build the shaped rewards $\tilde r_t$, run GAE to get $\hat A_t$ and $\hat R_t$, and whiten the advantages over the batch's real tokens.
6. Optimize. For a few epochs (InstructGPT 1, Stiennon et al. and Huang et al. 4), split the batch into minibatches, recompute $\log\pi_\theta$ and $V_\psi$ with gradients, and step on $\mathcal{L}^{\text{clip}} + c_v\mathcal{L}^V$, plus $\gamma_{\text{ptx}}$ times a language-modeling loss on pretraining text for PPO-ptx.
7. Bookkeeping. Log the mean score, the mean KL $\sum_t k_t$, response length, the clip fraction, and an approximate KL between $\pi_\theta$ and $\pi_{\text{old}}$. Update $\beta$ if it is adaptive.

Everything but step 6 runs without gradients. Here is the loop in NumPy on a toy that is small enough to see completely. The "language" has 6 tokens, token 0 is end-of-sequence, responses have at most 8 tokens, and the policy and reference are tables of logits, one row per position, so a token's log-probability depends only on where it is. A stand-in reward model likes token 3, dislikes token 5, and charges $-1$ for a response that never ends. The value model is a table of one value per position, fit by averaging returns.

The gradient line is the only thing that needs a word. For a softmax over logits $\mathbf{z}$, $\partial\log\pi(a)/\partial\mathbf{z} = \mathbf{e}_a - \boldsymbol{\pi}$, and the surrogate's derivative with respect to $\log\pi_\theta(y_t\mid s_t)$ is $\rho_t\hat A_t$ where the unclipped term is selected and 0 where the clip is active. Multiplying the two and summing over tokens gives the update. With $\beta = 0$, the score climbs furthest and the KL grows to almost 5 nats. A larger $\beta$ buys a smaller KL at the price of a lower score, the trade-off the tilt formula described, now reached by sampling instead of by enumerating every response.

In PyTorch

The same loop with neural networks. The policy, the reference, and the value model are tiny GRU language models (the value model has a scalar head), the reward model is the same stand-in function, and generation is ordinary autoregressive sampling. Every step of the list above appears in order.

Within a dozen iterations the policy fills all eight slots with 3s and accepts the penalty for never ending, a score of 3.0 (ending on the eighth token would have scored 3.5), and the KL has jumped to about 11 nats. With a large learning rate and $\beta = 0.05$ the leash is loose, and even this tiny run shows the pattern of §16.5, a policy that runs to the highest-scoring corner the reward allows and travels far from the reference to get there.

In a real system every one of these lines is a distributed operation over billions of parameters. Generation runs in a fast inference engine, the four models are sharded across GPUs, and the policy's fresh weights have to be copied into the inference engine after each update. Libraries such as TRL, OpenRLHF, and verl exist mostly to do that plumbing.

What four models cost

Memory. A trained model in mixed precision with Adam needs about 16 bytes per parameter (2 for bf16 weights, 2 for gradients, and 12 for fp32 master weights and the two Adam moments, §10.1). A frozen model needs only its bf16 weights, 2 bytes per parameter. With all four models of size $N$, that is $16N + 16N + 2N + 2N = 36N$ bytes before activations and the KV cache for generation.

That is more than twice the memory of SFT on the same model, which needs 16 bytes per parameter. Common savings are a value model smaller than the policy (as in InstructGPT), offloading the frozen models to CPU between uses (Huang et al. did this for their 6.9B run), sharing a trunk between policy and value (cheaper, though the two losses can interfere), LoRA adapters so the reference is the base weights with the adapter switched off, and dropping the value model altogether, which is what GRPO and RLOO do (§18.4).

Compute. Count FLOPs per response token with the rule of §6.7, $2N$ for a forward pass and $6N$ for forward plus backward. One iteration with all four models at size $N$ and $E$ optimization epochs costs about $2N$ to generate, $2N$ each for the policy, reference, reward, and value forward passes of steps 3 and 4, and $6N$ per epoch each to train the policy and the value model, roughly $(10 + 12E)N$ per token, or $22N$ at $E = 1$ and $58N$ at $E = 4$. (Implementations that keep the log-probabilities computed during generation skip the policy's extra pass and pay $(8 + 12E)N$, but inference engines often produce slightly different numbers than the training code, and a ratio built on mismatched log-probabilities is not 1 at the first step.) Generation is only $2N$ of that, yet it dominates the wall clock, because decoding produces one token per forward pass for the whole batch and is limited by memory bandwidth, not arithmetic (§12.1). The OpenRLHF authors report that the inference phase often takes over 90% of RLHF runtime (Hu et al. 2024), which is why every modern framework hands generation to a dedicated inference engine such as vLLM.

Read more

• Ziegler et al. (2019). Fine-Tuning Language Models from Human Preferences. The KL-penalized reward, the adaptive controller, and the first RLHF on a language model.
• Stiennon et al. (2020). Learning to summarize from human feedback. The setup most later implementations copied, including a value model initialized from the reward model and $\lambda = 0.95$.
• Ouyang et al. (2022). Training language models to follow instructions with human feedback. The per-token KL penalty, PPO-ptx, and the hyperparameters in Appendix C.
• Huang et al. (2024). The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization. More than twenty details, from padding to the end-of-sequence penalty, that decide whether a reproduction works.
• Schulman et al. (2017). Proximal Policy Optimization Algorithms. The clipped objective this section applies to tokens.
• Korbak et al. (2022). RL with KL penalties is better viewed as Bayesian inference. The tilted distribution as a posterior, a useful way to think about what the KL term does.
• Hu et al. (2024). OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework. How production RLHF systems place four models and an inference engine on a cluster.

Watch next

▶

▶

▶