Policy Gradient Methods › Policy Gradient Methods

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.