Policy Gradient Methods: REINFORCE Algorithm & Theory

With a trained reward model that captures human preferences, we face a basic challenge: how do we use this reward signal to improve our language model? The reward model tells us whether a complete generated response is good or bad, but our language model makes decisions token by token, one word at a time. We need a way to propagate that final reward signal back through the entire generation process to influence all of the individual token choices that collectively produced the outcome.

This challenge has a name in machine learning: the credit assignment problem. When a student gets a high score on an exam, you want to reinforce the study habits that led to that success, but determining which specific habits mattered most is hard. Language model alignment faces the same difficulty. A response might consist of hundreds of tokens, and only some of those tokens made the response good. The reward model scores the whole response, not individual tokens. We need a principled way to distribute credit across all the decisions that contributed to the outcome.

This is precisely where reinforcement learning enters the picture. Unlike supervised learning, where we have explicit targets for each prediction (we know the correct next token at every position), reinforcement learning can work with sparse feedback that arrives after an entire trajectory rather than at each step. Policy gradient methods provide a mathematically principled framework for optimizing a model's behavior based on exactly this kind of reward signal.

The core idea behind policy gradients is elegant in its simplicity. Think of it as a trial-and-error learning system operating at the level of complete responses. The model generates a response, receives a reward signal, and then adjusts its parameters to make similar responses more (or less) likely in the future depending on whether the reward was positive or negative. Over thousands of such updates, the model learns to generate the kinds of responses that tend to receive high rewards. The mathematical machinery we develop in this chapter makes this intuition precise and computationally tractable.

In this chapter, we develop the theory behind policy gradients from first principles, derive the key mathematical results that make training possible, and implement the REINFORCE algorithm. Along the way, we encounter the variance problem that makes naive REINFORCE impractical at scale, and we develop variance reduction techniques including baselines, advantages, and reward-to-go. This foundation is needed for understanding PPO and the complete RLHF pipeline that follows in subsequent chapters.

In reinforcement learning terminology, a policy is a function that maps states to actions. For language models, this mapping has a natural interpretation that becomes clear once we examine how text generation works. When you ask a language model to complete a sentence, it does not produce the entire response instantaneously. Instead, it generates one token at a time, each choice depending on everything that came before. This sequential decision-making process is precisely what reinforcement learning was designed to handle.

Think of a language model as an author writing a letter one word at a time. At each moment, the author has a specific context in mind (the prompt and everything written so far), and from that context, must choose the next word. Some words are clearly more appropriate given the context; others would seem out of place. The author's internal sense of which words fit a context is exactly what a language model encodes as a probability distribution. That probability distribution is the policy.

To understand why this framing is so powerful, consider what happens during text generation. The model receives a prompt, which establishes the initial context. Based on this context, the model must decide which token to produce first. This decision changes the context, as the newly generated token becomes part of the history. The model then faces a new decision: given the original prompt plus the first generated token, what should the second token be? This process continues, with each token choice altering the state and presenting a fresh decision problem.

The reinforcement learning vocabulary for this process is precise and useful. At each timestep $t$, the system occupies a state $s_t$, the agent takes an action $a_t$, the state transitions to $s_{t+1}$, and eventually the agent receives a reward $R$. For language generation, this maps directly:

• State $s_t$: The current context, consisting of the prompt plus all tokens generated so far
• Action $a_t$: The next token to generate
• Policy $\pi_\theta(a_t|s_t)$: The probability the model assigns to token $a_t$ given context $s_t$

This framing reveals why language model alignment is fundamentally a reinforcement learning problem. At each timestep, the model takes an action (selects a token), the state changes (the context grows), and eventually the complete sequence receives a reward from our reward model. The challenge is learning which tokens led to that reward. Unlike a game where individual moves might receive immediate feedback, language generation only reveals whether the response was good or bad after the entire sequence is complete. This delayed reward structure is exactly the scenario where reinforcement learning excels.

One important property of this MDP (Markov Decision Process) formulation is that the state transition is deterministic given the action. Once we pick token $y_t$, the new state is simply the old state with $y_t$ appended. There is no external randomness in the transition. All the stochasticity lives in the policy: the model samples tokens from its probability distribution. This simplification makes the mathematics more tractable than general RL settings where transitions are random.

The model's parameters $\theta$ define the policy. Our goal is to adjust these parameters so that the model generates responses that receive high rewards. This means we need to find parameter values that make high-reward sequences more probable and low-reward sequences less probable. The question is: how do we know which direction to adjust the parameters? This is the central technical question that policy gradient theory answers.

To optimize our policy, we need a clear mathematical objective. In RLHF, we want to maximize the expected reward over all possible responses the model might generate. This expectation is important because language generation is inherently stochastic. The same prompt can lead to many different responses depending on which tokens get sampled at each step. Some responses will be excellent, others will be mediocre, and still others will be poor. We want to tune the model so that, on average, it produces responses that receive high rewards.

Think of the objective function as a report card for a policy, but one computed by averaging over all possible exams the policy might face. A policy scores well on this report card if, when we sample random responses from it, those responses tend to be good. The objective does not ask that every response be perfect; it asks that the distribution over responses be skewed toward high-quality outcomes.

Given a prompt $x$, the model generates a response $y = (y_1, y_2, \ldots, y_T)$ by sampling tokens according to its policy. The reward model then scores the complete response, giving us $R(x, y)$. Our objective is to find parameters $\theta$ that maximize the expected reward:

where:

• $J(\theta)$: the objective function we want to maximize
• $\theta$: the parameters of the language model
• $\mathbb{E}$: the expectation over sequences sampled from the policy
• $\pi_\theta(\cdot|x)$: the policy (language model) distribution conditioned on prompt $x$
• $R(x, y)$: the reward scalar for the generated response $y$

This expectation is taken over all possible responses the model could generate. Some responses will be excellent (high reward), others mediocre, others poor. The expectation weights each by its probability under the current policy. This weighting matters: if the model rarely generates a particular response, that response contributes little to the expected value, even if it happens to have high reward. A policy that occasionally produces brilliant responses but usually produces mediocre ones will score lower than a policy that consistently produces good responses.

We can expand this expectation using the probability of generating the complete sequence:

where:

• $\sum_{y}$: summation over all possible output sequences in the vocabulary
• $\pi_\theta(y|x)$: probability of generating sequence $y$ given prompt $x$
• $R(x, y)$: reward for the specific sequence $y$

This summation form makes the objective's structure transparent, but it also reveals a computational challenge. The number of possible sequences is astronomically large: for a vocabulary of 50,000 tokens and sequences of length 100, there are $50000^{100}$ possible sequences. We cannot enumerate them all. We need a way to optimize this objective without explicitly summing over all sequences.

The probability $\pi_\theta(y|x)$ represents the likelihood of generating the entire response $y$. It is computed as the product of individual token probabilities:

where:

• $T$: the length of the generated sequence
• $y_t$: the token generated at timestep $t$
• $y_{<t}$: the history of tokens generated before step $t$, that is $(y_1, y_2, \ldots, y_{t-1})$
• $\pi_\theta(y_t | x, y_{<t})$: the probability the model assigns to token $y_t$ given the full context

This autoregressive factorization, which we have seen throughout our discussion of language models in earlier chapters, connects the policy formulation directly to how transformers and LSTMs generate text. Each factor in the product corresponds to one step of the generation process, one call to the model's forward function, one decision about which token to emit next. The full sequence probability is simply the product of all these individual decisions.

You might wonder why we cannot just directly maximize $J(\theta)$ by computing its gradient. The obstacle is that $J(\theta)$ involves an expectation over sequences sampled from $\pi_\theta$, and sampling is not a differentiable operation. When we sample a token from a probability distribution, we make a discrete choice, and discrete choices are not differentiable with respect to the parameters that shape that distribution.

Think of it this way: if you tweak the parameters of the model slightly so that the probability of token "excellent" increases by 0.001, you cannot directly measure how that change affects the expected reward. The reward depends on which specific tokens get sampled, and sampling is random. We need a mathematical trick to convert the gradient of an expectation into an expectation of gradients that we can estimate from samples.

Here is the core challenge stated precisely: we need to compute $\nabla_\theta J(\theta)$ so we can update our parameters via gradient ascent. But the expectation involves sampling from the policy itself, which depends on $\theta$. This creates a circular dependency that seems difficult to resolve. The gradient depends on how changing $\theta$ affects which sequences get sampled, but sampling is a discrete, non-differentiable operation. We cannot simply backpropagate through the sampling process the way we would through a continuous neural network layer.

The policy gradient theorem provides an elegant solution to this problem. Rather than trying to differentiate through sampling directly, it rewrites the gradient in a form that we can estimate using samples from the current policy. The key insight is that we can express the gradient of an expectation in terms of an expectation of gradients that we can compute.

Let's derive the policy gradient step by step. Starting with our objective written as an explicit sum:

We take the gradient with respect to $\theta$. Since $R(x,y)$ is produced by a frozen reward model and does not depend on $\theta$, only the policy probability is differentiated:

This is still a sum over all sequences, which we cannot compute directly. We need to transform it into an expectation. The trick is to multiply and divide by $\pi_\theta(y|x)$, introducing a factor of 1 without changing the value:

Now we apply the log-derivative trick (also called the likelihood ratio trick or REINFORCE trick). This technique appears throughout machine learning and statistics whenever we need to compute gradients of expectations. The core identity is:

This follows directly from the chain rule of differentiation. Recall that $\frac{d}{d\theta} \log f(\theta) = \frac{1}{f(\theta)} \frac{df}{d\theta}$, which rearranges to $\frac{df/d\theta}{f(\theta)} = \frac{d}{d\theta} \log f(\theta)$. The logarithm converts a multiplicative relationship into an additive one, which proves extremely useful when we deal with sequence probabilities that are products of many terms.

Substituting back:

A sum over all possible sequences weighted by their probabilities is exactly an expectation under the policy:

where:

• $\nabla_\theta J(\theta)$: gradient of the objective function with respect to parameters
• $\nabla_\theta \log \pi_\theta(y|x)$: gradient of the log-probability of the sequence, called the "score function"
• $R(x, y)$: the reward, acting as a scalar weight on the gradient

This is the policy gradient theorem for our setting. It tells us something remarkable: to compute the gradient of expected reward, we can sample sequences from our policy, compute the gradient of their log-probabilities, and multiply by their rewards. The score function $\nabla_\theta \log \pi_\theta(y|x)$ points in the direction in parameter space that would increase the probability of sequence $y$. Multiplying by the reward scales this direction: high rewards give large positive contributions, low rewards give small or negative contributions.

The transformation from an intractable sum to a Monte Carlo estimate is the key practical insight. We cannot enumerate all possible sequences; for even modest vocabulary sizes and sequence lengths, this would involve more sequences than atoms in the universe. But we can sample sequences from our policy, and the policy gradient theorem tells us that averaging over these samples gives us an unbiased estimate of the true gradient.

The unbiasedness guarantee is important. It means that if we sample enough sequences, our gradient estimate will converge to the true gradient. We might not get the exact gradient from a single sample, but in expectation across many samples, we get the right answer. This is what allows reinforcement learning to work despite the discrete, non-differentiable nature of sampling.

For language models, we need to express the log-probability of an entire sequence in terms of individual token log-probabilities. This connection is needed because our models compute token-level probabilities, not sequence-level probabilities directly. Fortunately, the autoregressive structure of language models makes this conversion straightforward.

Using the autoregressive factorization from the previous section, we can apply the logarithm to convert the product of probabilities into a sum of log-probabilities:

where:

• $\sum_{t=1}^{T}$: sum over all timesteps in the sequence
• $\log \pi_\theta(y_t | x, y_{<t})$: log-probability of the specific token $y_t$ given context

The logarithm converts the product of probabilities into a sum of log-probabilities. This is where the log-derivative trick pays off: instead of dealing with products of many small numbers (which can underflow numerically when sequence lengths are in the hundreds), we work with sums of log-probabilities, which remain in a numerically stable range.

Each term $\log \pi_\theta(y_t | s_t)$ is simply the log-softmax output of our model at position $t$, evaluated at the selected token. This is exactly what we compute during the forward pass when training language models with cross-entropy loss. When the model processes a sequence, it produces logits for every position, and applying log-softmax gives us log-probabilities. To compute the log-probability of a specific token choice, we simply index into this log-probability vector at the chosen token's index.

The gradient of the full sequence log-probability then follows from the linearity of differentiation. The gradient of a sum is the sum of gradients:

where:

• $s_t$: the state (context) at timestep $t$, equivalent to $(x, y_{<t})$
• $\nabla_\theta \log \pi_\theta(y_t | s_t)$: gradient of the log-probability for a single token decision

Each term in the sum corresponds to one token's contribution to the overall sequence probability. We compute each term by running a standard backward pass through the model at that timestep.

Putting this together, the policy gradient becomes a sum over token-level contributions, all weighted by the same sequence-level reward:

where:

• $\sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(y_t | s_t)$: the accumulated gradient from all tokens in the sequence
• $R(x, y)$: the final reward, which scales the gradient of every token choice equally

The same reward $R(x, y)$ multiplies every token's gradient. This is the basic credit assignment mechanism of REINFORCE: even though we only receive reward at the end, it propagates back to influence all token choices. Every token that contributed to the response receives the same credit or blame for the final outcome. This is both a strength and a weakness of the basic approach. It allows learning from sparse rewards, but it does not distinguish between tokens that were important for the reward and tokens that were irrelevant to it.

The policy gradient theorem gives us a theoretical result: the gradient of expected reward equals the expected product of the score function and the reward. REINFORCE (also called the Monte Carlo policy gradient) is the algorithm that turns this theory into practice through Monte Carlo estimation. The name comes from the idea that we reinforce behaviors that lead to good outcomes: when a sequence receives high reward, we reinforce the probability of generating that exact sequence.

Think of REINFORCE as a simple loop: try something, evaluate whether it worked, update your approach based on the result. A human learning to write better essays follows essentially the same process. They write an essay, receive feedback (reward), and adjust their writing style to emphasize the patterns that received positive feedback and reduce the patterns that received criticism. REINFORCE formalizes this intuition mathematically.

The algorithm proceeds in four steps for each training iteration:

1. Sample a complete sequence from the current policy
2. Compute the reward for that sequence using the reward model
3. For each token in the sequence, compute the gradient of its log-probability scaled by the reward
4. Sum these token-level gradients and update parameters via gradient ascent

In practice, we implement gradient ascent by minimizing the negative objective (gradient descent on the negation):

Notice the use of torch.no_grad() during generation. We generate the sequence without tracking gradients, then separately compute the loss and allow gradients to flow through that computation. This is a memory optimization: we do not need to store the full computation graph during generation, only during the loss computation.

Let's trace through a concrete numerical example to build intuition. Suppose our model generates a three-token response, and at each position the model assigns the following probabilities to the token it chose:

• Token 1: probability 0.30 (for example, "The")
• Token 2: probability 0.50 (for example, "answer")
• Token 3: probability 0.20 (for example, "is")

The reward model gives this short response a reward of $R = +2.0$.

Step 1: Compute per-token log-probabilities.

Step 2: Sum to get the sequence log-probability.

This corresponds to a sequence probability of $e^{-3.506} \approx 0.030$, which matches $0.30 \times 0.50 \times 0.20 = 0.030$ directly.

Step 3: Compute the REINFORCE loss.

Step 4: Interpret the gradient.

When we call loss.backward(), the gradient will be computed with respect to model parameters. This gradient points in the direction that would increase the loss, which means the optimizer (doing gradient descent to minimize the loss) will adjust parameters in the direction that decreases the loss. Decreasing the loss means increasing $\log \pi_\theta(y|x) \cdot R$, which means increasing the log-probability of generating this sequence. Since $R = +2.0$ is positive, this is exactly what we want: after the update, the model is more likely to generate this sequence.

Negative reward case: Suppose instead $R = -1.5$. Then the loss equals $-(-3.506) \times (-1.5) = -5.259$. Minimizing (making more negative) this loss requires making $-\log \pi \cdot R$ decrease. Since $R = -1.5 < 0$, decreasing $-\log \pi \cdot (-1.5)$ means decreasing $\log \pi$, making the sequence less likely. This is again the correct behavior: negative reward makes the sequence less probable.

The key insight is that REINFORCE increases the probability of sequences that receive high rewards and decreases the probability of sequences that receive low or negative rewards. Over many updates, the policy gradually shifts to favor the kinds of responses that the reward model prefers. This is the essence of learning from rewards: try things, evaluate what works, do more of what works.

REINFORCE has a fundamental weakness that prevents it from being directly used for training large language models: extremely high variance in its gradient estimates. This variance explains why every practical RLHF system uses more sophisticated algorithms. It directly motivates PPO, which we cover in the next chapter.

To understand where the variance comes from, recall our Monte Carlo gradient estimate:

where:

• $N$: the batch size (number of sampled sequences)
• $y^{(i)}$: the $i$-th sampled sequence in the batch
• $R(x, y^{(i)})$: the reward for the $i$-th sequence

We are estimating the true gradient by averaging over $N$ samples. If $N = 1$, our estimate equals the product of a random gradient direction and a random reward. Both quantities vary across samples, and their product varies even more. Several sources contribute to the overall variance:

Sampling variance. We estimate an expectation using a small batch of samples. Different samples from the same policy can have wildly different rewards. Consider a prompt where some responses are brilliant (reward +10) and others are terrible (reward -10), but most are mediocre (reward near 0). If we happen to sample one of the rare brilliant responses, we get a huge positive gradient. If we sample a terrible one, we get a huge negative gradient. The expected gradient might be small and stable, but individual estimates can be orders of magnitude larger.

Reward magnitude. Large rewards amplify gradient noise. A reward of 100 versus 0.1 changes gradient magnitudes by a factor of 1000. If the reward scale varies across prompts or across training, the gradient scale varies too. This makes it hard to choose a stable learning rate that works well across all prompts.

Long sequences. With more tokens, there is more opportunity for randomness in the sampling process. Each token choice is stochastic, and these random choices compound. The same prompt can lead to vastly different sequences purely due to sampling randomness early in generation. A small difference in the first few tokens can cascade into completely different responses.

To see why the variance matters practically, imagine training a model with batch size 1 on a task where rewards range from -10 to +10. One batch might produce a sample with reward +8, causing a large positive gradient update. The next batch produces a sample with reward -7, causing a large negative gradient update. The model is being pushed back and forth, making little net progress. It is like trying to navigate by compass in a room where someone randomly rotates the compass between readings.

The variance decreases at rate $O(1/N)$ with batch size $N$, but this means we need exponentially more samples to achieve linear improvements in stability. To reduce variance by a factor of 10, we need 10 times more samples. To reduce it by a factor of 100, we need 100 times more samples. For large language models where each sample requires generating a complete response (potentially hundreds of tokens) through a billion-parameter model, and where each generated response must be scored by a separate reward model, this sample inefficiency becomes a serious practical barrier.

The practical consequence is that naive REINFORCE is essentially unusable for training large language models. The gradient noise dominates the signal, training progress is erratic, and the compute cost is prohibitive. This motivates variance reduction techniques, which we develop in the following sections.

The main point for reducing variance comes from a useful mathematical property of the policy gradient: subtracting any constant from the reward does not change the expected gradient, but it can substantially reduce variance. This property allows us to center the rewards around zero. This provides clearer learning signals, without introducing any bias into our gradient estimates.

To understand why centering helps, consider the extreme case where all rewards are positive, say ranging from +5 to +15. Without centering, every gradient update pushes to increase the probability of all sampled sequences. The model receives mixed signals: "increase this sequence, increase that sequence too, increase everything." The optimizer cannot easily distinguish which responses are better than others because every response receives a positive gradient push, just of different magnitudes.

With a baseline of 10 (the approximate mean), rewards become -5 to +5. Now the gradient clearly distinguishes: "increase sequences above the baseline, decrease sequences below it." Sequences that scored 14 get positive reinforcement; sequences that scored 6 get negative reinforcement. This clearer signal leads to faster learning and lower variance in gradient estimates.

Consider the gradient estimator with a baseline $b$:

where:

• $b$: a baseline value (which must be independent of the action $y$, though it can depend on the state $x$)
• $(R(x, y) - b)$: the "centered" reward, which reduces variance by removing the mean level

To verify this does not change the expected gradient, we need to show that the baseline term contributes zero in expectation. We check this by computing the expected value of the baseline term directly:

We pulled $b$ outside the expectation because it does not depend on $y$. Now we need to evaluate the expected value of the score function $\nabla_\theta \log \pi_\theta(y|x)$. This expectation is always zero:

This result shows that the score function always has zero mean under the policy that generates it. The reason is elegant: the score function measures how changing parameters affects the log-probability of different outcomes. Since increasing the probability of some outcomes necessarily decreases the probability of others (they must sum to one), these effects balance out exactly. Mathematically, the sum of all probabilities is fixed at 1, so its gradient must be zero.

This result implies that we can subtract any baseline without introducing bias. The question becomes: what baseline minimizes variance?

The optimal baseline is the expected reward under the current policy:

where:

• $b^*$: the variance-minimizing baseline value
• $\mathbb{E}[R(x, y)]$: the expected reward under the current policy

This is called a value baseline, and in practice we often estimate it with a learned value function $V(x)$ that takes the prompt $x$ as input and predicts the expected reward. With this baseline, the quantity $(R(x, y) - b^*)$ becomes the advantage: how much better (or worse) this specific sequence is compared to what we expect on average. The advantage is positive when a response exceeds expectations and negative when it falls short.

The advantage has a natural interpretation: it measures how surprising a reward is relative to what we expected. A response with advantage +5 did much better than typical; one with advantage -3 did worse. By training on advantages rather than raw rewards, we give the model clearer information about which responses are good relative to the typical performance level. This is the same intuition a manager uses when giving feedback: rather than saying "you scored 85 on this review," it is more informative to say "you scored 15 points above the team average."

Another variance reduction technique exploits a basic causal property of sequential decision-making: actions can only affect future rewards, not past ones. A decision made at time $t$ cannot retroactively change what happened at times $1$ through $t-1$. Yet in basic REINFORCE, we multiply every token's gradient by the full sequence reward, even though early tokens cannot have influenced any "future" rewards from the perspective of later tokens.

Think of a recipe as an analogy. If a chef makes three decisions (which oil to use, which spice to add, whether to reduce the heat), and the final dish receives a good review, it would be wrong to credit the choice of oil for the flavor benefits that the reduced heat provided. The choice of oil could not have influenced whether the sauce reduced properly, since the oil was chosen before the heat decision. Yet basic REINFORCE does exactly this: it credits the oil choice for the reduced-heat flavor improvement.

Instead of multiplying each token's gradient by the total reward, we can use only the rewards that could have been causally influenced by that action. This is called the reward-to-go formulation. The reward-to-go at timestep $t$ sums all rewards from that step forward:

where:

• $R_t$: reward-to-go from timestep $t$
• $r_{t'}$: immediate reward received at timestep $t'$
• The sum starts at the current timestep $t$ and runs to the end of the sequence $T$

This captures only the rewards that could have been influenced by the action at time $t$, ignoring any rewards that occurred before. The policy gradient with reward-to-go is:

where:

• $\mathbb{E}_{y \sim \pi_\theta}$: expectation over sequences sampled from the policy
• $y_t$: the token generated at timestep $t$
• $R_t$: the cumulative reward from time $t$ onwards (reward-to-go), ignoring past rewards

This does not change the expected gradient (the proof is analogous to the baseline proof: terms involving rewards from before time $t$ vanish in expectation because they are independent of the action at time $t$) but reduces variance by removing terms that add noise without giving useful credit assignment signal.

In the typical RLHF setup where reward is only given at the end, reward-to-go equals the final reward at all timesteps. However, this structure becomes important when we add auxiliary rewards or penalties, such as the KL divergence penalty we will introduce in subsequent chapters. That penalty applies at each token, creating a richer reward structure where reward-to-go provides meaningful variance reduction by correctly attributing per-token contributions to the overall gradient.

The two techniques can be combined. We can use a per-token value function $V(s_t)$ as a baseline for each timestep, subtracting it from the reward-to-go at that position:

where:

• $A_t$: the per-token advantage at timestep $t$
• $R_t$: the reward-to-go from timestep $t$
• $V(s_t)$: the value function predicting expected future reward from state $s_t$

This per-token advantage is what actor-critic algorithms compute. The actor is the policy network being trained; the critic is the value function giving the baseline. By training the critic to accurately predict expected future rewards, we get better baselines, which gives us lower-variance advantage estimates, which gives us more stable policy updates.

Let's put everything together into a complete REINFORCE implementation with baseline and a proper training loop. We will build on the components developed throughout this chapter.

The first component we need is the value network, which learns to predict the expected reward for a given prompt:

The plots reveal the characteristic noisiness of REINFORCE training. While the value network learns to predict expected rewards (value loss decreases), the policy loss remains highly variable. This variance is the primary motivation for more sophisticated algorithms like PPO, which we cover in the next chapter.

Beyond baselines and reward-to-go, several additional techniques help control variance in practice. Together, these techniques form a toolkit that practitioners reach for when raw REINFORCE proves too unstable.

Reward normalization standardizes rewards across a batch to have zero mean and unit variance:

This technique prevents large reward magnitudes from causing gradient explosions while preserving the relative ordering of rewards. It is particularly useful when the reward scale varies across training: perhaps early in training rewards cluster near zero, while later in training they spread across a wider range. Normalization ensures the learning rate remains appropriate throughout.

The key insight is that what matters for policy improvement is not the absolute scale of rewards but their relative ordering. Whether responses score 5, 7, and 9, or 500, 700, and 900, the model should learn the same qualitative lessons. Normalization makes the gradient magnitudes consistent regardless of the reward scale.

Entropy regularization adds a bonus for maintaining diversity in the policy's probability distribution, preventing premature convergence to deterministic behavior:

Gradient clipping limits the magnitude of gradient updates by rescaling the gradient vector when its norm exceeds a threshold:

As we discussed in the chapter on gradient clipping, this technique prevents individual updates from being too large, which is especially important given REINFORCE's high variance. A single high-reward sample can produce an enormous gradient that catastrophically overwrites learned behavior. Clipping ensures no single update can move parameters too far from their current values.

The choice of max_norm involves a tradeoff: too small and learning slows unnecessarily; too large and you lose the protection against large updates. Values between 0.5 and 5.0 are common in practice, with 1.0 being a typical starting point.

In production RLHF systems, these variance reduction techniques are applied together. A typical implementation would:

1. Generate a batch of responses (multiple responses per prompt for better variance reduction)
2. Score each response with the reward model
3. Normalize rewards within the batch to zero mean and unit variance
4. Compute baseline predictions from the value network
5. Compute per-sequence advantages (normalized reward minus baseline)
6. Compute the REINFORCE loss using advantages instead of raw rewards
7. Backpropagate and clip gradients before applying the update
8. Add entropy bonus to prevent policy collapse

Each technique addresses a different source of instability: normalization handles reward scale variation, baselines address absolute reward level, clipping prevents catastrophic updates, and entropy regularization prevents premature convergence.

REINFORCE, while theoretically elegant and mathematically well-grounded, has significant practical limitations that make it challenging to apply directly to large language model alignment. Understanding these limitations is important both for appreciating why more sophisticated algorithms exist and for knowing when simpler approaches might be appropriate.

Sample efficiency. REINFORCE requires many samples to produce reliable gradient estimates. For large language models generating hundreds of tokens, each sample is expensive: it requires a full autoregressive forward pass through a model with billions of parameters, potentially followed by a separate reward model evaluation. Training to convergence can require millions of sampled sequences. At a cost of perhaps several seconds per sample on expensive hardware, the total compute required becomes prohibitive. The $O(1/N)$ variance reduction from batching means you need 100 times more samples just to reduce variance by a factor of 10 compared to having a perfect gradient estimate.

Credit assignment at the token level. The credit assignment problem is only partially solved. REINFORCE attributes equal credit to all tokens in a sequence based on the sequence-level reward. This ignores the reality that different tokens contribute differently to response quality. A response might be mostly excellent with one factually incorrect sentence, yet every token in the response receives the same blame for that error. The model cannot learn to avoid that specific error while keeping the rest of its behavior intact. More sophisticated algorithms like actor-critic methods attempt finer-grained credit assignment by training a value function that estimates the expected future reward from each intermediate state.

Update stability. REINFORCE provides no guarantees about the size of policy updates. A lucky high-reward sample can push the policy dramatically in one direction, only for the next batch to push it back in the opposite direction. This oscillation wastes computation and can destabilize training. More seriously, if the policy changes too much in one update, future samples will be drawn from a different distribution, making old gradient estimates less relevant. This mismatch between the policy at sample time and the policy at update time can compound over multiple steps, leading to divergence. The PPO algorithm, covered in the next chapter, directly addresses this by constraining how much the policy can change per update.

On-policy requirement and sample reuse. Policy gradient methods are on-policy: the gradient estimate is valid only for the policy that generated the samples. Once we update the policy, the samples become stale, and we must generate new ones. This is in contrast to off-policy methods (like Q-learning) that can reuse old samples from a replay buffer. The on-policy requirement means we cannot reuse samples across multiple gradient steps, further reducing sample efficiency. Every parameter update requires generating a fresh batch of sequences.

Despite these limitations, REINFORCE remains important as the foundational algorithm for understanding all more sophisticated policy gradient methods. PPO, TRPO, and other production algorithms all build on the same core gradient estimator derived in this chapter. Understanding REINFORCE's mechanics, failure modes, and the mathematical properties that enable variance reduction is prerequisite knowledge for understanding why those algorithms make the design choices they do. The log-derivative trick, the zero-mean score function, the baseline property, and the causal structure of reward-to-go all appear in more advanced algorithms, just in more sophisticated forms.

Policy gradient methods provide the mathematical foundation for optimizing language models using reward signals. The key insights from this chapter are:

• Language models are policies. The autoregressive generation process naturally maps to the reinforcement learning framework: states are contexts, actions are tokens, and the policy $\pi_\theta(a_t|s_t)$ is the model's next-token probability distribution. This framing connects language model alignment directly to decades of reinforcement learning research.

• The policy gradient theorem enables optimization. By applying the log-derivative trick, we rewrite the gradient of expected reward as $\mathbb{E}[\nabla_\theta \log \pi_\theta(y|x) \cdot R(y)]$. This turns an intractable sum over all possible sequences into a Monte Carlo estimate we can compute from samples.

• REINFORCE is simple but high-variance. The basic algorithm samples sequences, computes their rewards, and updates parameters in proportion to $\log \pi(y) \cdot R$. This works in principle but requires many samples for stable estimates, which makes it impractical at scale for large language models.

• Baselines reduce variance without introducing bias. Subtracting a baseline $b$ from rewards yields the same expected gradient but lower variance. The score function has zero mean in expectation, which is the key mathematical fact that makes baselines work. The optimal baseline is the expected reward, leading to the advantage formulation $A = R - V$.

• Reward-to-go exploits causal structure. Actions cannot affect past rewards, only future ones. Using only future rewards in the gradient estimate removes noise from past outcomes while preserving the unbiasedness of the estimator.

• Additional techniques help in practice. Reward normalization, entropy regularization, and gradient clipping all contribute to more stable training. These techniques address different failure modes: scale variation, policy collapse, and catastrophic updates respectively.

The variance problem with REINFORCE motivates the development of more sophisticated algorithms. In the next chapter, we will see how Proximal Policy Optimization (PPO) addresses these issues by constraining policy updates and using careful advantage estimation with generalized advantage estimation (GAE), making RL-based alignment practical for production language models.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about policy gradient methods and the REINFORCE algorithm.