RLHF Pipeline › Worked Example: Tracing a Single RLHF Training Step

To make the pipeline concrete, let's trace a single complete training step from prompt to policy update, using specific numerical values at each stage. This worked example connects all the abstract components into a coherent sequence.

Setup: Suppose we have already completed SFT and trained a reward model. The policy starts from the SFT model checkpoint. We use a KL coefficient $\beta = 0.05$ and a PPO clip epsilon of $0.2$.

Step 1: Sample a prompt from the training distribution. We draw the prompt "What is gradient descent?" from our prompt library.

Step 2: Generate a response from the current policy. The policy samples token by token, creating the response: "Gradient descent is an optimization algorithm that iteratively adjusts parameters in the direction that minimizes a loss function." Alongside the response, we record the log-probabilities assigned by the policy to each token, which we'll call $\log \pi_\theta(a_t | s_t)$ for each token position $t$.

Step 3: Score the response using the reward model. The reward model receives the full prompt-response pair and outputs a raw score. Suppose it outputs $R_{RM} = 1.42$.

Step 4: Compute the KL penalty. We also pass the prompt-response pair through the frozen reference model (the SFT model) and compute token-level log-probabilities. The KL divergence at each token position $t$ is:

$$\text{KL}_t = \pi_\theta(a_t | s_t) \cdot \left(\log \pi_\theta(a_t | s_t) - \log \pi_{ref}(a_t | s_t)\right)$$

Summing across all token positions gives the total sequence-level KL divergence. Suppose this comes out to $D_{KL} = 0.28$.

Step 5: Compute the total reward. Combining the reward model score with the KL penalty:

$$\begin{aligned} R_{total} &= R_{RM} - \beta \cdot D_{KL} \\ &= 1.42 - 0.05 \times 0.28 \\ &= 1.42 - 0.014 \\ &= 1.406 \end{aligned}$$

where:

• $R_{RM} = 1.42$: the raw reward model score for this response
• $\beta = 0.05$: the KL penalty coefficient
• $D_{KL} = 0.28$: the total sequence-level KL divergence

At this early stage of training, the KL penalty is small because the policy has barely moved from its SFT initialization. The total reward is almost entirely the raw reward model score.

Step 6: Estimate advantages. The advantage tells us how much better this response was than the average response to this prompt. Using a value function baseline with estimated value $V(s) = 1.15$, the advantage is approximately:

$$A = R_{total} - V(s) = 1.406 - 1.15 = 0.256$$

A positive advantage means this response was better than expected. The policy update will increase the probability of generating similar responses.

Step 7: Compute the PPO ratio. For each token, we compute the ratio of the current policy's probability to the old policy's probability (from before this update):

$$\text{ratio}_t = \frac{\pi_\theta(a_t | s_t)}{\pi_{\text{old}}(a_t | s_t)} = \exp\left(\log \pi_\theta(a_t | s_t) - \log \pi_{\text{old}}(a_t | s_t)\right)$$

Suppose the mean ratio across tokens is $1.08$, meaning the policy has moved slightly toward creating this response.

Step 8: Apply the clipped objective. Since $1.08 < 1 + 0.2 = 1.2$, the ratio is within the clipping bounds. The surrogate objective for this response is $1.08 \times 0.256 = 0.277$. The gradient update increases the probability of generating this response, pushing the policy in a direction that the reward model rates more highly.

After this single step, the policy is infinitesimally more likely to generate responses similar to "Gradient descent is an optimization algorithm..." when asked about gradient descent. Repeating this process across thousands of prompts and responses gradually sculpts the policy's behavior across the entire prompt distribution.