Handout #19 — Policy Gradients and RLHF

A reward model can score a response without telling us how to differentiate through the act of generating it. Policy gradients use sampled responses and their rewards to update the language model. A reference policy makes departure from its response distribution costly.

1Optimize the distribution that generates responses

1.1 Two models, two different training problems

In lecture 18, comparisons trained $R_\phi(x,y)$ to predict preferences. Now hold $\phi$ fixed and update $\theta$. For a fixed prompt distribution $D$, the unregularized objective is

The scorer tells us how valuable a completed response appears; the policy determines how often that response occurs. Reward-model training does not itself change the policy. Policy training can use a learned scorer, an executable test, or another numerical reward; the score need not be differentiable with respect to the generated text.

1.2 A frozen scorer does not make the objective constant

For a particular fixed string $y$, $R_\phi(x,y)$ has no direct dependence on $\theta$. But $J$ averages that score under a distribution that does depend on $\theta$. Changing the distribution changes the expectation.

In an ordinary sampled rollout, integer token choices are not differentiable. Calling reward.backward() on a detached score cannot update the policy. Instead, retain or recompute the policy’s differentiable log-probabilities for the sampled tokens, and use the reward as a numerical coefficient.

2A gradient from a sampled response

2.1 REINFORCE: score times log-probability gradient

Sample $M$ prompts independently from $D$ and $K$ responses per prompt independently from the current policy. A Monte Carlo estimator of the gradient is

This is the policy-gradient or REINFORCE estimator of Williams (1992), as used in slides 26–28. The reward is a coefficient on a vector gradient. Averaging controls the scale as the number of samples changes. Unequal response counts need explicit prompt weighting if prompts should count equally.

A positive coefficient pushes up the sampled response’s log-probability for a sufficiently small isolated step; a negative one pushes it down. A zero coefficient contributes nothing. Shared parameters and other samples can change the net update. A low positive reward still reinforces that sampled response; “lower than another response” is not the same as “negative.”

2.2 Compute every possible sample in a two-response world

Let a single scalar $\theta$ set $p=P(A)=\sigma(\theta)$ and $P(B)=1-p$. Suppose $R(A)=3$, $R(B)=1$, and initially $p=0.25$. Then

The two log-probability derivatives are $\frac{d\log p}{d\theta}=1-p=0.75$ and $\frac{d\log(1-p)}{d\theta}=-p=-0.25$. The possible one-sample estimates are:

The expected estimate is $0.25(2.25)+0.75(-0.25)=0.375$, exactly the derivative above. An individual sample can point against the exact gradient. “Unbiased” describes the estimator’s expectation; it does not promise an improvement on every step.

2.3 Optional: where the estimator comes from

3From one response to many generated tokens

3.1 One terminal reward weights a sum of token log-probabilities

Autoregressive factorization gives

The same terminal reward weights each generated token’s log-probability gradient. This is coarse credit assignment: a sequence score does not identify which token was responsible. Summing the token log-probability gradients gives the sequence log-probability gradient.

Include EOS if it was generated and scored; exclude padding and prompt prediction targets. The reward model can read the prompt even though prompt targets are excluded from the policy loss. Different examples may have different response lengths. Summing over response tokens and then averaging over responses estimates the sequence objective; dividing each response by its length changes the objective’s weighting.

3.2 Minimize a negative surrogate to perform ascent

For fresh on-policy samples, use a minimization loss whose gradient is the negative sampled policy gradient:

The numerical value of this surrogate is not an estimate of $-J$; its gradient supplies the desired update at the sampling policy. Do not compare its raw value across changing rollouts as if it measured response quality.

The following original teaching function is included in the downloadable CPU lab. logits[:, t] predicts targets[:, t]; shift inputs/targets before calling it. Every target ID must be valid, including masked positions.

def policy_loss(logits, targets, response_mask, rewards):
    """Aligned [B,T,V] logits, [B,T] targets/mask, [B] rewards."""
    if logits.ndim != 3 or targets.shape != logits.shape[:2]:
        raise ValueError("expected aligned [B,T,V] logits and [B,T] targets")
    if response_mask.shape != targets.shape or response_mask.dtype != torch.bool:
        raise ValueError("response_mask must be boolean [B,T]")
    if rewards.shape != (logits.shape[0],):
        raise ValueError("expected one reward per response")
    if not response_mask.any(dim=1).all():
        raise ValueError("each row needs at least one response target")
    token_logp = logits.log_softmax(-1).gather(
        -1, targets.unsqueeze(-1)).squeeze(-1)
    sequence_logp = token_logp.masked_fill(~response_mask, 0).sum(-1)
    return -(rewards.detach() * sequence_logp).mean()

Sample IDs and compute rewards without tracking a reward-model graph, then run a policy forward pass on those IDs with autograd enabled. The scoring pass must represent the distribution used for sampling; in this simplified lab, stochastic layers such as dropout are disabled, and sampling and scoring use unchanged parameters. Evaluation mode can disable dropout while still allowing gradients.

After optimizer.zero_grad(), call loss.backward() and optimizer.step(). Then collect fresh samples. Repeated optimization over old samples needs additional machinery beyond the estimator shown here.

3.3 Generating a response versus scoring a given response

Computing $p_\theta(y\mid x)$ for a supplied response uses the same next-token conditionals as generation, but the computation can be scheduled differently. A “lookup” is cheap only after the model has computed the relevant logits.

For example, suppose the prompt is Explain space elevators (three tokens) and the supplied response is A cable idea EOS (four tokens in this toy tokenizer). Feed Explain space elevators A cable idea as the six input tokens. The last prompt position predicts A; the following three positions predict cable, idea, and EOS. All four target IDs are already known, so their conditional probabilities can be computed in one causally masked pass. The mask still prevents access to future response tokens.

Generating those four tokens instead needs one prompt-prefill pass followed by three decode passes. KV caching reuses earlier keys and values; it does not remove the dependence on the token just selected. See Handout #12’s cache discussion and the Transformers cache explanation.

4The feedback loop creates a new evaluation problem

4.1 Which data train which model?

1. Initialize: begin with an SFT policy that already produces useful responses.
2. Collect comparisons: generate candidate responses and have humans compare them under a stated rubric.
3. Fit and validate the scorer: train $R_\phi$ on comparisons, with held-out prompts and comparisons to evaluate generalization.
4. Optimize the policy: freeze the scorer, collect fresh policy responses, score them, and update $p_\theta$. Add the reference penalty from §5.
5. Optional fresh feedback: obtain comparisons on updated-policy responses and retrain the scorer between policy-training phases.

Slides 29–32 separate these steps. Fresh comparisons can address response-distribution shift as the policy changes; evaluate updated outputs independently and update the scorer between policy-training phases.

4.2 Optimizing a proxy can amplify its mistakes

A reward model can be accurate on ordinary held-out comparisons and still be exploitable by a policy that seeks its highest scores. For example, it might confuse verbosity or flattery with usefulness. Reward hacking can produce fluent text; detecting it requires more than spotting gibberish.

Consider this entirely synthetic single-prompt example. The “human quality” column is an independent toy utility, not something available to the optimizer:

Optional empirical illustration · When a higher learned reward stops helping

Synthetic evaluation: a fixed 6B reward model provides “gold” judgments, not direct human ratings. These PPO runs use a zero KL-penalty coefficient; the axis measures how far the policy nevertheless moved. This illustration supports slides 33–34; it does not compare §5’s β values.

4B · Optional figure reading: separate the two scores

Follow the purple 3M curves. After gold score peaks, what happens to proxy score? Would maximizing proxy score select the gold evaluator’s best checkpoint? Does this measure human satisfaction?

Interpret the curves

Proxy score keeps rising while gold score falls, so maximizing the proxy passes the best gold-score region. The evaluator is another model; human satisfaction was not directly measured here.

A useful audit samples outputs from successive policy checkpoints, holds prompts fixed for comparisons, randomizes answer order, and asks evaluators to apply a defined quality rubric without seeing scorer values. Report predicted reward alongside independent quality and the rate of specific errors. Reward-model generalization and policy quality are separate measurements.

5Make departure from a reference policy costly

5.1 Three roles: current policy, reference, and scorer

Choose a fixed reference $q(y\mid x)$, commonly a copy of the initial SFT policy. The trainable policy may start with identical weights; only its copy is updated. The reference supplies response probabilities, while the reward model supplies scalar scores. Keep the reference fixed during this training phase.

Slides 35–36 replace pure reward maximization with

The direction matters: this is current-policy-to-reference KL, averaged over responses from the current policy. A larger $\beta$ charges more for the same departure at a fixed reward scale. It is a soft penalty, not a hard guarantee of closeness or quality. Scaling reward scores changes the tradeoff unless $\beta$ is scaled too.

5.2 What is KL divergence?

The Kullback–Leibler (KL) divergence compares two probability distributions over the same outcomes. Here the outcomes are complete responses to one fixed prompt $x$. Temporarily write $p(y)=p_\theta(y\mid x)$ for the current policy and $q(y)=q(y\mid x)$ for the reference. Both distributions sum to one. For the derivations below, use a fixed maximum generation length so the response space is finite. Define

Read this from the inside out. The ratio $p(y)/q(y)$ compares the probability of the same response under the two models. Its log is positive when the current policy makes that response more likely, and negative when it makes it less likely. Multiplying by $p(y)$ and summing averages these log-ratios over responses from the current policy. Thus KL compares the full distributions, not just the most likely answer.

Another interpretation connects KL to the negative log-probabilities used in language-model losses. Define entropy $H(p)=-\sum_y p(y)\log p(y)$ and cross-entropy $H(p,q)=-\sum_y p(y)\log q(y)$. Then

KL is the extra expected negative log-probability from evaluating responses drawn from $p$ with $q$ instead of $p$. Here $p$ is the current policy, not a ground-truth human distribution.

For example, let $p=(0.8,0.2)$ and $q=(0.5,0.5)$ over responses A and B:

Summing the last column gives $D_{\rm KL}(p\|q)\approx0.192745$ nats. Reversing the roles gives $D_{\rm KL}(q\|p)\approx0.223144$ nats, a different value. One B sample has a negative log-ratio; neither that sample nor an equally weighted average over A and B is the forward KL above. A finite sample mean can be negative even though its expectation is nonnegative.

5.3 How the response penalty becomes a KL term

Slide 35 subtracts a log-ratio from each response’s reward. We will show why the expected value of this per-response subtraction is a KL penalty. For a fixed prompt, define

Now average over $y\sim p_\theta(\cdot\mid x)$. Fix $x$ and abbreviate $p(y)=p_\theta(y\mid x)$, $R(y)=R_\phi(x,y)$, and $\widetilde R(y)=\widetilde R_\theta(x,y)$:

Finally average over prompts, restoring the conditioning on $x$:

This is the objective on slide 36. The average subtracted penalty is $\beta D_{\rm KL}$; its contribution to a maximization objective has a minus sign. The sampled log-ratio becomes a KL term when averaged under the current policy. Drawing responses from another distribution changes this calculation.

For $p=(0.8,0.2)$, $q=(0.5,0.5)$, rewards $(3,1)$, and β = 1, the adjusted rewards are approximately 2.529996 for A and 1.916291 for B. Averaging them directly gives

B receives a positive adjustment because its log-ratio is negative. The overall expected adjustment is nevertheless a subtraction of 0.192745. In a real LM, exact enumeration is infeasible; a mean of log-ratios from fresh current-policy samples estimates the KL expectation.

For an autoregressive response, the same penalty can be accumulated over generated tokens:

Both models score the same tokens under the same generated prefixes. Include generated EOS and exclude prompt targets and padding, as in §3. The sum along one response is still a sampled sequence log-ratio; its expectation over current-policy responses is the sequence-distribution KL.

5.4 Use the adjusted reward in a policy-gradient update

For fresh current-policy samples, use the adjusted reward as a detached coefficient in the REINFORCE surrogate. This gives an unbiased estimator of $\nabla F$ under the stated assumptions. Differentiating only the sampled adjusted-reward values would miss the effect of changing the response probabilities.

For the two-response model with $p=\sigma(\theta)$ and reference probability $q_A$, the exact derivative is

This gives the explorer below an exact answer to compare with each possible sampled update. At $p=0.8$, $q_A=0.5$, $(R_A,R_B)=(3,1)$, and $\beta=1$, the derivative is about 0.098193, smaller than the unregularized derivative 0.32.

KL charges for departure from the reference; it cannot correct every error in the reference or scorer. Continue evaluating policy quality independently, even when KL stays small.

6Practice: explain, calculate, repair

6.1 A timed 20-point practice set

# Already aligned log-probabilities, each row = prompt + response + pad.
sequence_logp = token_logp.mean(dim=1)
loss = (rewards * sequence_logp).mean()
loss.backward()
optimizer.step()

optimizer.zero_grad()
sequence_logp = token_logp.masked_fill(~response_mask, 0).sum(dim=1)
loss = -(rewards.detach() * sequence_logp).mean()
loss.backward()
optimizer.step()

6.2 Optional question bank for future quiz preparation

7Optional: implementation checks and a real code reading

7.1 Run the small CPU lab

Download 19_policy_gradient_checks.py and run python3 19_policy_gradient_checks.py in an environment with PyTorch. It needs no GPU, model download, network, or CartPole environment. It enumerates every outcome where possible, so it can compare estimators against exact gradients rather than judging them from one noisy training run.

• Recover exact reward gradients by averaging per-outcome score-function gradients.
• Check positive/negative reward signs, constant-reward cancellation, and the prompt-only baseline example.
• Verify sequence sums, response masks, EOS placement, detached rewards, zero masked-position gradients, and bulk-versus-prefix scoring parity.
• Compare exact KL differentiation against a detached shaped-reward estimator; expose the missing sampling-probability term in a naive implementation.
• Check the reward-hacking table and the timed set’s binary-gradient and KL calculations.

7.2 Read PyTorch’s REINFORCE example

reinforce.py · select_action · lines 54–60 (verbatim)

def select_action(state):
    state = torch.from_numpy(state).float().unsqueeze(0)
    probs = policy(state)
    m = Categorical(probs)
    action = m.sample()
    policy.saved_log_probs.append(m.log_prob(action))
    return action.item()

The action is a sampled integer, while m.log_prob(action) retains the graph from the policy probabilities. Returning action.item() to the environment is fine; detaching the saved log-probability would destroy the learning path. This is §1’s discrete boundary and §3’s retained log-probability in a two-action environment.

reinforce.py · finish_episode · lines 72–77 (verbatim)

    for log_prob, R in zip(policy.saved_log_probs, returns):
        policy_loss.append(-log_prob * R)
    optimizer.zero_grad()
    policy_loss = torch.cat(policy_loss).sum()
    policy_loss.backward()
    optimizer.step()

The minus sign makes a minimizer perform reward-weighted ascent, as in §3. The upstream code sums over time steps; its separate returns values are discounted future rewards computed earlier in the function. It also normalizes those returns within the episode. That normalization is an additional algorithmic choice and does not inherit the simple prompt-only-baseline argument in §2.3. CartPole’s time-step rewards and changing states differ from this handout’s single terminal response reward.

This file has no language-model response mask, reference policy, or KL penalty. Its policy has dropout; the original code saves log-probabilities during the sampling forward pass. Our deterministic teaching loss instead recomputes log-probabilities with dropout disabled.

7.3 Optional sources and a scope check

• Williams (1992), Simple statistical gradient-following algorithms for connectionist reinforcement learning: REINFORCE and the score-function estimator.
• Stiennon et al. (2020), Learning to summarize from human feedback: the summarization application cited on slide 36.
• Ouyang et al. (2022), Training language models to follow instructions with human feedback: a complete applied pipeline, with optimization machinery beyond this lecture.

For the assigned scope, you should be able to trace “sample → score → weight log-probability gradient → update policy,” and explain what the KL penalty adds. The optional derivations justify those operations; implementing a production RLHF trainer is not required.