PPO for Language Models: Adapting RL to Text Generation

In the previous chapter, we explored the PPO algorithm as a general-purpose policy gradient method with clipped objectives and value function estimation. Now we face a important question: how do we apply these ideas to language models? The translation is difficult because traditional reinforcement learning and text generation have basic differences. Language models were designed as next-token predictors, not as agents acting in environments. They were trained to model the statistical patterns of human language, not to maximize cumulative rewards over time. Yet with some careful reframing, we can view text generation through the lens of sequential decision-making and apply PPO to steer models toward human-preferred outputs.

To appreciate why this translation matters, think about what we want from a language model assistant. It should answer accurately and help the user without causing harm. Supervised fine-tuning on human-written demonstrations can get us partway there, but it has a basic limitation: it trains the model to imitate what humans wrote rather than directly optimizing how people evaluate its responses. Reinforcement learning from human feedback (RLHF) provides a framework for that optimization, and PPO is the optimization engine that makes it work.

The central challenge is bridging two very different worlds. In classical RL, an agent acts in a well-defined environment with clear states, actions, and transitions. In language generation, the "environment" is the open-ended space of human communication. There are no physics, no game rules, and no terminal conditions beyond the end-of-sequence token. The reward signal comes not from an environment simulator but from a separate neural network that has learned to predict human preferences. Everything must be carefully adapted to handle the peculiarities of text: the enormous vocabulary, the sequential nature of generation, the delayed and sparse feedback, and the ever-present risk of the model drifting into incoherent or harmful behaviors.

Think of the adaptation problem as similar to teaching a chess-playing algorithm to write poetry. The chess algorithm knows how to optimize sequences of decisions to maximize a score, and poetry writing also involves sequences of decisions toward a goal. But the action space, state representation, reward structure, and constraints are completely different. Translating the optimization machinery requires both conceptual insight about the parallels and careful engineering to handle the practical differences.

This chapter connects reinforcement learning to the structure of language generation step by step. We will see how an LLM naturally is a stochastic policy, why the vocabulary forms a massive discrete action space, and how rewards propagate through generated sequences. Most importantly, we will understand why constraining the policy to stay close to its original behavior is needed for stable training. Each of these concepts builds on the last, forming a complete picture of how PPO turns language model behavior.

In reinforcement learning, a policy maps states to action probabilities. Given a state $s$, the policy $\pi(a|s)$ tells us the probability of taking action $a$. This definition describes decision-making: the agent observes its state, checks its policy, and chooses an action. Language models do the same thing: given a context, the model outputs a probability distribution over the next token. The parallel is exact, and recognizing it is the conceptual foundation for everything that follows.

Think of the language model's forward pass as the policy evaluation step in RL. Every time you call the model with a sequence of tokens, it is computing the policy, reading the current state and determining how to act. The transformer layers process the input, build rich contextual representations through self-attention, and project the final hidden state to a distribution over vocabulary tokens. This distribution is the policy output. Sampling from it is the action selection step. Appending the sampled token to the context creates the next state, and the cycle repeats until the model generates a stop token.

To see this clearly, consider what happens when you type a prompt. The model processes your input, builds internal representations through its transformer layers, and produces a probability distribution over its vocabulary. This distribution assigns higher probabilities to tokens that would naturally continue the text and lower probabilities to tokens that would seem out of place. When the model generates a response, it samples from this distribution (or selects greedily), appends the chosen token to the context, and repeats the process. Each step involves observing a state and selecting an action according to a probability distribution, which is exactly what a policy does.

The key insight is that autoregressive generation is policy rollout in disguise. Starting from the initial state consisting of the prompt, we sample actions according to the policy, each action extends the state, and we continue until reaching a terminal condition. This is a precise mathematical correspondence that allows us to apply policy gradient theorems and optimization techniques to language generation.

Let $x = (x_1, x_2, \ldots, x_n)$ denote the input prompt and $y = (y_1, y_2, \ldots, y_T)$ the generated response. At each generation step $t$, the language model computes the conditional probability of the next token given everything that came before:

where:

• $\pi_\theta$: the policy defined by model parameters $\theta$
• $y_t$: the token generated at the current step $t$
• $x$: the input prompt sequence
• $y_{<t}$: the sequence of tokens generated prior to step $t$, that is $(y_1, \ldots, y_{t-1})$

This formula represents the probability of generating token $y_t$ given the full context of the prompt and previous tokens. The notation emphasizes that the policy depends on everything that came before: the original prompt establishes the task, and each previously generated token shapes what should come next. The model's parameters $\theta$ encode the learned patterns that determine how context maps to token probabilities.

The probability of a complete response $y$ given prompt $x$ is the product of these conditional probabilities across all timesteps. This product structure means that the probability of a response decomposes multiplicatively into the probability of each individual token choice. Understanding this decomposition is needed for understanding how policy gradient methods can work here: updating the parameters $\theta$ changes the probability of individual tokens at each step, and those changes compound across the response to shift the overall probability of different response types.

Comparing language models to standard RL shows a direct correspondence between the two frameworks. The state at each step is the concatenation of the prompt and all tokens generated so far: $s_t = (x, y_1, \ldots, y_{t-1})$. The action is the next token to generate: $a_t = y_t$. The policy is the LLM's softmax output distribution over the vocabulary: $\pi_\theta(a_t|s_t)$. The trajectory is the complete prompt-response pair: $\tau = (x, y)$.

Each of these mappings changes the interpretation for how we set up the optimization. The state includes the entire generation history, so the policy is not Markovian in the traditional sense. The model cannot simply look at one step and decide what to do; it must consider the full context. This is why transformer architectures, which can attend to arbitrarily distant tokens, are well-suited for this role. Older architectures like RNNs, which compress history into a fixed-size vector, are at a disadvantage here.

The autoregressive generation process that produces a response is exactly a policy rollout. Starting from the initial state $s_1 = x$, we sample actions according to the policy, each action extends the state, and we continue until generating a stop token. This precise mathematical correspondence allows us to apply policy gradient methods to language generation. Policy optimization theorems and techniques apply to steering language models toward desired behaviors.

The action space in language generation is the model's vocabulary. This discrete set of tokens defines every possible action the policy can take at each step. Unlike continuous control problems where actions might be forces or velocities, or even discrete games where actions represent button presses, the language model's action space consists of linguistic units. These might be complete words, word fragments, punctuation marks, or special tokens that signal the end of generation.

Think of the vocabulary as a massive library of building blocks, and generating a response as repeatedly choosing which block to place next. Each choice is constrained by what came before: you cannot place a closing parenthesis before an opening one without creating grammatical nonsense, and you cannot begin a response mid-sentence. The model has learned an enormous amount about which sequences of choices produce coherent, meaningful text, but at each step it must select from the entire library.

For modern language models, this vocabulary typically contains between 30,000 and 100,000 tokens. The exact size depends on the tokenization algorithm used during pretraining. Byte-pair encoding, the most common approach, creates vocabularies that balance coverage of common words with the ability to represent rare words through subword decomposition. The vocabulary must be large enough to efficiently represent common text patterns while remaining small enough for the softmax computation to be tractable. This tradeoff explains why vocabulary sizes cluster around 32,000 to 65,000 tokens for most production models.

This scale is much larger than traditional RL domains. In classic control, action spaces are small. For example, a robot arm might have six joints. In Atari games, agents choose among roughly 18 discrete actions representing joystick directions and button combinations. Language models must select from tens of thousands of possible tokens at every single step. This represents an increase of three to four orders of magnitude in the number of discrete choices.

The enormous action space has several important consequences for how PPO behaves when applied to language models. Each of these implications changes how we design and tune the algorithm.

Exploration is implicit in language generation. With such a large action space, the model cannot systematically try all options. There is no way to enumerate the consequences of every possible token choice. Instead, exploration emerges from the inherent stochasticity of sampling from high-entropy distributions during generation. When the model is uncertain, it assigns probability mass to many tokens, and sampling naturally explores these alternatives. This is quite different from classical RL, where explicit exploration strategies like epsilon-greedy or optimistic initialization are necessary. In language generation, sampling temperature takes the place of explicit exploration: higher temperature spreads probability mass more evenly, encouraging exploration; lower temperature concentrates it on the most probable tokens, encouraging exploitation.

Credit assignment is diffuse across the token sequence. When a response receives a reward, determining which specific token choices contributed to that reward is challenging. Did the response succeed because of a precise word choice in the third sentence, or because of the overall argument structure established from the first token? The signal must somehow propagate back through dozens or hundreds of individual decisions. The advantage function addresses this by using a value function to assess the quality of each intermediate state, but the value function itself must learn from the same sparse reward signal. This creates a bootstrapping challenge: we need good value estimates to produce good advantages, but we need good advantages to produce good policy updates that lead to rewards that improve value estimates.

Despite these challenges, the discrete nature of the action space simplifies some aspects of PPO. We can compute exact action probabilities rather than approximating them, as we would need to do in continuous action spaces. We can directly enumerate the KL divergence between policies by summing over vocabulary positions. The softmax function gives us a proper probability distribution that sums to one, avoiding the density estimation challenges that arise with continuous distributions.

Each position in the generated sequence requires selecting from this massive action space. A 100-token response involves 100 sequential decisions, each choosing among 32,000 options. Even a short response spans a large search space of $32000^{100}$ possible sequences. This is why language generation uses learned policies to navigate the space efficiently.

The state in language generation grows with each action. Unlike a game where the agent might return to previously visited states, or a robot navigation task where the agent can revisit locations, language generation always moves forward. Each generated token permanently extends the context. The state at step $t+1$ contains everything from state $t$ plus the newly generated token. There is no going back, no undo, and no way to visit the same state twice unless the exact same sequence of tokens has been generated before.

Think of the growing state as a conversation transcript. Once you have said something, you cannot unsay it. Every subsequent word must be consistent with everything already written. If the model commits to a particular argument in sentence one, it must carry that argument through to the end. If it introduces a character in paragraph one, that character constrains what can happen later. This irreversibility creates a basic asymmetry between early and late tokens: early choices constrain later ones, giving them outsized influence on the final response.

Formally, the state at step $t$ is:

where:

• $s_t$: the state at timestep $t$, consisting of the complete history seen so far
• $x_i$: the $i$-th token of the prompt, for $i = 1, \ldots, n$
• $n$: the number of tokens in the prompt
• $y_j$: the $j$-th generated token, for $j = 1, \ldots, t-1$

This formulation captures the needed nature of autoregressive generation. The state is not a compact summary of the situation; it is the complete history of what has been written. The model must condition on all of this information to decide what should come next. A word that appeared ten sentences ago might be important for maintaining coherence, while a word from three tokens ago might determine grammatical constraints on the current position. Transformers handle this beautifully through self-attention, which allows the model to selectively focus on any previous token regardless of distance.

The key insight is that this state space is effectively infinite and largely unique to each trajectory. Two different prompts lead to entirely different state spaces, and even the same prompt with different generation paths explores different states. Unlike board games where states might repeat (the same chess position can arise from different move orders) or continuous control where the system might return to similar configurations, language generation creates fresh territory with every token. Each state is a unique point in the space of possible text prefixes.

The growing state affects implementation in important ways. The model must process longer sequences as it generates tokens. Each new token requires attending to all previous tokens, which increases computation quadratically with sequence length under standard attention. This is where techniques like KV caching, which we discuss in the chapter on efficient inference, become needed for efficient generation during training rollouts. By caching the key and value projections from previous steps, the model avoids recomputing them for each new token, changing per-step complexity from quadratic to linear in the already-generated sequence.

The generated trajectory illustrates how the context (state) expands with each step. Even a short response creates a sequence of unique states, as the growing history fundamentally changes the input to the policy at every decision point. At step 0, the policy sees only the prompt. At step 1, it sees the prompt plus one generated token. By the final step, it sees the entire conversation history. This expansion means that the policy faces a different decision problem at every timestep, even though the underlying question remains the same.

Reward models and RL algorithms use rewards differently, and this mismatch creates one of the central challenges for PPO in language models. The reward model, as we discussed in the chapter on Reward Modeling, takes a complete prompt-response pair and outputs a single scalar score. It evaluates the response as a whole, considering whether it is helpful and coherent as well as accurate and safe. But PPO operates on trajectories with per-timestep rewards. The algorithm expects to receive a reward signal at each step, allowing it to compute advantages and update the policy accordingly.

Think of it this way: the reward model is like a book critic who reads an entire novel and then gives it one star rating. PPO is like a writing coach who wants feedback on every sentence so it can tell the author what to keep and what to cut. The critic's overall assessment captures something real, but the writing coach needs more granular guidance. We have to translate the critic's single rating into per-sentence feedback, which inevitably involves some imprecision.

The standard approach assigns the reward to the final token:

where:

• $r_t$: the reward assigned at step $t$
• $R(x, y)$: the scalar score from the reward model for the complete response
• $T$: the length of the generated sequence (final step)
• $x$: the input prompt
• $y$: the complete generated response

This approach reflects the reality that we only know how good a response is once it is complete. We cannot reliably evaluate a partial response because quality often depends on the full content, including how the response concludes. A response that starts brilliantly but ends poorly deserves a low score, but you would not know that until you reach the end. So we wait until generation finishes, compute the reward, and assign it to the final step.

Sparse rewards make credit assignment difficult. The policy must learn which of its many token choices contributed to the final reward. PPO addresses this through its advantage estimation, but the challenge remains significant. A response might receive a low reward because of a single poor word choice, but that signal must propagate back through dozens of preceding tokens. The value function must learn to predict, from any intermediate state, what the expected final reward will be. This prediction is what allows advantages to differentiate between tokens: some tokens lead to states with high expected rewards, others to states with lower expectations.

With a discount factor of $\gamma = 1.0$, which is standard in RLHF, every token in the sequence receives the same return equal to the final reward. This uniform signal might seem uninformative at first glance. If every token gets the same return, how can the algorithm distinguish good tokens from bad ones? The answer lies in the advantage function, which compares the actual return to the expected return under the value function. A token that leads to a higher-than-expected return receives a positive advantage, while a token that leads to a lower-than-expected return receives a negative advantage. This differential signal, created by the value function's predictions, is what enables learning even with sparse rewards.

The most important adaptation for language model PPO is the KL divergence penalty. This constraint measures how far the policy drifts from its starting point and is an important regularizer. Without this constraint, the optimized policy can diverge substantially from the original model, often finding degenerate solutions that maximize reward without producing useful responses.

Think of the KL penalty as a tether connecting the evolving policy to its starting point. As PPO updates the model's parameters to increase reward, the tether pulls back, preventing the model from wandering too far into unknown territory. The reference model, frozen at initialization, represents a baseline of linguistic competence built through pretraining and supervised fine-tuning. The KL penalty says: "you can move to improve reward, but you must not move so far that you forget what you already know."

Without the KL penalty, the policy often finds degenerate solutions. The policy is free to move anywhere in the space of possible token distributions. If the reward model has exploitable shortcuts that yield high scores without corresponding quality, the unconstrained policy will find them. It might learn to generate repetitive phrases that the reward model scores highly. It might produce outputs that superficially resemble good responses while lacking substance. It might drift so far from natural language that it generates text no human would write. One classic failure mode is mode collapse, where the policy learns to always generate one particular type of response that scores well on average, losing the diversity and adaptability that make a language model useful. The KL penalty prevents these failure modes by keeping the policy anchored to the reference distribution.

The KL penalty modifies the reward at each timestep, converting the sparse terminal reward into a dense per-token signal. The shaped reward at step $t$ is:

where:

• $\tilde{r}_t$: the shaped reward at step $t$ used for PPO training
• $r_t$: the original sparse reward (nonzero only at the final step)
• $\beta$: the KL penalty coefficient, controlling the strength of the regularization
• $\text{KL}_t$: the KL divergence contribution at this step

The KL contribution at each step measures the divergence between the current policy and the reference for the specific token generated. Rather than summing over the full vocabulary (which would require querying both policies over all 50,000+ tokens), we use the log ratio for the chosen action as a tractable per-token approximation:

where:

• $\pi_\theta(y_t | s_t)$: the probability of the chosen token under the current policy
• $\pi_{\text{ref}}(y_t | s_t)$: the probability of the chosen token under the reference policy
• $y_t$: the token generated at step $t$
• $s_t$: the context (state) at step $t$

This formula is easy to interpret. When the current policy assigns higher probability to a token than the reference did, the log ratio is positive, and subtracting it from the reward penalizes this divergence. When the current policy assigns lower probability than the reference, the log ratio is negative, and subtracting it provides a bonus (the policy is moving closer to the reference for that action). The net effect is that the policy is discouraged from making dramatic probability changes in either direction. It can shift probabilities to improve reward, but only within bounds set by $\beta$.

The key insight is that this per-token KL penalty serves double duty. First, it constrains policy drift globally, preventing the model from drifting too far from its original behavior. Second, it provides a dense reward signal at every token, not just at the end of generation. This density matters for credit assignment: instead of all tokens receiving the same sparse terminal reward, each token now receives a small bonus or penalty based on how much the policy changed from the reference for that particular choice.

The reference policy $\pi_{\text{ref}}$ is typically the model after supervised fine-tuning (SFT) but before any RL training. This anchor prevents the policy from drifting into regions of token space that the original model considered highly unlikely. The SFT model represents our best current understanding of how to generate helpful, coherent text. By constraining the RL policy to stay near this baseline, we ensure that the optimized model retains the linguistic competence learned during pretraining and SFT.

The coefficient $\beta$ controls the strength of this constraint. Higher values keep the policy closer to the reference but limit learning; the policy cannot deviate much even when doing so would improve reward. Lower values allow more exploration but risk instability and reward hacking; the policy might find degenerate solutions that maximize reward while creating poor text. Success in RLHF depends on finding the right balance. We will explore the mathematical properties and tuning of this penalty in detail in the upcoming chapter on KL Divergence Penalty.

Notice how the KL penalty transforms the reward signal. The original sparse reward only provides signal at the final token; all intermediate tokens receive zero reward. After KL shaping, every token receives a reward component based on how much the policy diverges from the reference. Tokens where the policy assigns higher probability than the reference receive penalties (the KL term is positive, so it subtracts from the reward). Tokens where the policy assigns lower probability than the reference receive bonuses (the KL term is negative, so subtracting it adds to the reward). This transformation converts the sparse terminal reward into a dense per-token signal that guides the policy at every step.

The PPO objective for language models combines the clipped surrogate loss from standard PPO with the KL-shaped reward structure we just described. Understanding how these pieces fit together reveals why PPO works for alignment: it maximizes expected reward while maintaining two separate safeguards against instability. The KL penalty prevents the policy from drifting too far from the reference in absolute terms, while the clipping mechanism prevents any single update step from moving too far from the data collection policy.

Think of the full PPO objective as a three-way balancing act. The reward signal pulls the policy toward higher-scoring behaviors. The KL penalty pulls it back toward the reference model. The clipping mechanism ensures that each individual gradient step stays within a safe neighborhood. These three forces together create a training dynamic that is far more stable than any one of them alone.

Combining all these elements, the PPO objective for language models takes the following form. For a batch of prompt-response pairs, we compute:

where:

• $\theta$: the parameters of the language model being optimized
• $\mathbb{E}_{(x,y) \sim \pi_{\theta_{\text{old}}}}$: the expectation over trajectories generated by the policy that collected the data
• $\rho_t(\theta)$: the probability ratio between the current and old policies at step $t$
• $\hat{A}_t$: the estimated advantage at step $t$, showing how much better the chosen action was compared to the average
• $T$: the length of the generated sequence
• $\epsilon$: the clipping parameter (typically 0.1 or 0.2) that defines the trust region

The min operator serves a important purpose in this objective. It takes the conservative lower bound between the unclipped and clipped objectives. This keeps updates cannot be too aggressive. If the advantage is positive, showing that the chosen action was better than expected, clipping limits how much we increase the probability of the action. If the advantage is negative, showing the action was worse than expected, clipping limits how much we decrease the probability. This prevents large updates that could collapse the policy.

The probability ratio $\rho_t(\theta)$ is the core mechanism that allows PPO to reuse trajectory data across multiple gradient steps. Without it, policy gradient methods can only take one gradient step per batch of data, making them extremely sample-inefficient. With the ratio, we can check whether the current policy has drifted too far from the data-collection policy and clip the gradient accordingly.

The ratio is defined as:

where:

• $\pi_\theta(y_t | x, y_{<t})$: the probability of the token $y_t$ under the current policy (being updated)
• $\pi_{\theta_{\text{old}}}(y_t | x, y_{<t})$: the probability of the same token under the policy that collected the data (frozen for this update step)
• $y_t$: the token generated at step $t$
• $x$: the input prompt
• $y_{<t}$: the sequence of tokens generated prior to step $t$

This ratio measures how the policy has changed since the data was collected. A ratio of 1.0 means the policy assigns exactly the same probability as before. A ratio greater than 1 means the policy now assigns higher probability to this token. A ratio less than 1 means the probability has decreased. By basing updates on this ratio rather than directly on log probabilities, PPO can perform multiple gradient steps on the same batch of data without the policy drifting too far from the data-collection policy.

The key insight here is why we use the ratio rather than log probabilities directly. In standard REINFORCE, we would multiply the log probability by the advantage and take a gradient step. This works for a single step but fails for multiple steps: after updating the policy, the gradient computed from the old trajectories is no longer valid because the old trajectories were generated by a different policy. The ratio corrects for this by weighting each token's contribution by how much the current policy differs from the data-collection policy.

The advantage $\hat{A}_t$ is estimated using Generalized Advantage Estimation (GAE) applied to the KL-shaped rewards. GAE is a technique that balances the bias-variance tradeoff in advantage estimation by combining multi-step returns through exponential weighting:

where:

• $\gamma$: the discount factor (usually near 1.0 for RLHF)
• $\lambda$: the GAE smoothing parameter (typically 0.95)
• $T$: the length of the generated sequence
• $\delta_{t+l}$: the temporal difference (TD) error at step $t+l$

This summation calculates an exponentially weighted average of future TD errors. The parameter $\lambda$ controls the bias-variance trade-off: higher values rely more on observed returns (reducing bias but increasing variance), while lower values rely more on value estimates (reducing variance but potentially introducing bias from imperfect value predictions). Setting $\lambda = 1$ recovers Monte Carlo returns, while $\lambda = 0$ gives pure one-step TD learning. The typical value of 0.95 balances these considerations.

The TD error $\delta_t$ measures the surprise at each step: the difference between the observed reward plus the estimated value of the next state, and the estimated value of the current state:

where:

• $\tilde{r}_t$: the KL-shaped reward at step $t$
• $\gamma$: the discount factor
• $V(s_t)$: the value function estimate for state $s_t$ (how good is the current context expected to be?)
• $V(s_{t+1})$: the value function estimate for the next state (defined as 0 if $t=T$)

The TD error captures whether the transition was better or worse than expected. If the reward plus next-state value exceeds the current-state value, the transition was unexpectedly good, and the TD error is positive. If the sum falls short, the transition was unexpectedly bad, and the TD error is negative. These signals, accumulated through GAE, produce advantage estimates that indicate which actions led to better outcomes than the value function predicted.

The positive advantage at the final step reflects the high reward, while negative values indicate steps that yielded lower-than-expected value. The probability ratios stay within the trust region defined by the clipping parameter. This shows how PPO maintains stability. When a ratio exceeds $1 + \epsilon$ or falls below $1 - \epsilon$, the clipped ratio takes over, preventing the gradient from pushing the policy further in that direction.

To solidify the abstract mechanics, let us trace through a complete PPO update for a short response, computing every quantity step by step. We will use a simplified three-token response to keep the numbers manageable.

Setup: Consider a prompt "What is 2+2?" and a response with three tokens: "The", "answer", "is" (we omit the actual number for brevity). The reward model scores this as 0.7 (somewhat helpful but incomplete). The reference model and current policy have slightly different log probabilities for each token.

Step 1: Assign the sparse reward. The terminal reward of 0.7 goes to the final token, while earlier tokens receive zero:

Step 2: Compute the per-token KL divergence. Suppose the log probabilities are:

The KL contribution at each step is the log ratio:

Step 3: Apply KL shaping with $\beta = 0.2$. The shaped rewards are:

Token 1 receives a small penalty because the policy moved toward higher probability than the reference. Token 2 receives a small bonus because the policy moved toward lower probability than the reference. Token 3 receives most of the terminal reward, slightly discounted by the KL penalty.

Step 4: Compute TD errors and GAE advantages. Suppose the value function estimates are:

With $\gamma = 1.0$, the TD errors are:

With $\lambda = 0.95$, the GAE advantages work backward:

All three advantages are positive, meaning all three tokens led to outcomes better than what the value function predicted. This makes sense: the response received a reward of 0.7 while the value function only predicted around 0.3 to 0.6 for the intermediate states.

Step 5: Normalize advantages. We subtract the mean and divide by the standard deviation:

After normalization, tokens 1 and 2 have positive advantages (they were above average), while token 3 has a negative advantage (it was below average even though the reward was positive, because the value function had already predicted a high value for that state).

Step 6: Compute probability ratios. Suppose during the gradient update epoch, the policy has shifted slightly:

The ratio for token 1 is $\exp(-0.85 - (-0.9)) = \exp(0.05) \approx 1.05$. The advantage is positive, so the policy is trying to increase this token's probability. The ratio of 1.05 is within the trust region $[0.8, 1.2]$, so clipping does not activate.

Step 7: Apply the clipped objective. For token 1 with advantage $\hat{A}_1^{\text{norm}} = 1.10$:

Both are equal here because the ratio is within the trust region. The gradient flows through the unclipped objective. For a token where the ratio had exceeded 1.2, clipping would have capped the objective at $1.2 \times \hat{A}$, preventing the gradient from pushing the probability any higher.

This walkthrough shows the complete flow: sparse reward becomes shaped reward through KL penalty, shaped reward becomes advantages through GAE, and advantages become policy gradients through the clipped surrogate objective. Each step has a clear purpose, and together they produce a training signal that is both informative and stable.

Let us now assemble a more complete implementation showing how these pieces fit together in a training step. This simplified version captures the needed structure while omitting some production details like distributed training, gradient accumulation, and advanced memory management. This example shows the flow of data and computation.

The overall structure follows a pattern that will look familiar from supervised learning, but with key differences. We generate responses before computing loss, which makes the training loop longer. We maintain two versions of log probabilities, old and new, to compute the probability ratio. We track the KL divergence as a health metric. These differences reflect the basic distinction between RL and supervised learning: in RL, the data distribution depends on the current policy.

These hyperparameters define the constraints for the optimization. The clipping epsilon and KL coefficient are particularly necessary for preventing the model from collapsing or drifting too far from its original capabilities. The value coefficient balances the policy and value function losses, and the GAE lambda controls the bias-variance trade-off in advantage estimation.

The following example demonstrates how PPO training proceeds at a high level. We use mock components to illustrate the data flow without requiring actual large models. This demonstration shows the needed rhythm of PPO training: generate responses, score them, compute advantages, and update the policy.

The loop has a distinct structure that differs from supervised fine-tuning. We first generate a batch of responses from the current policy (the rollout phase). We then score those responses with the reward model and compute all the ingredients needed for the update: log probabilities, values, KL divergence, advantages. Only then do we update the parameters. Some implementations perform multiple gradient steps on the same batch (the epochs parameter in PPO), with the probability ratio clipping preventing any single epoch from updating too aggressively.

The training metrics reveal several important dynamics that you should monitor during real RLHF runs. The policy loss oscillates as the model learns to balance reward maximization against the clipping constraint. The value loss declines as the critic learns to predict returns. This provides better advantage estimates over time. The KL divergence tracks how far the policy drifts from the reference, a key quantity we want to keep bounded. If KL grows too large, the policy is moving into territory where the reward model may not be reliable. The probability ratio stays near 1.0 because the clipping mechanism prevents extreme updates. This shows that the trust region constraint is functioning as intended.

The success of PPO training depends heavily on correctly setting its hyperparameters. Unlike supervised learning where the learning rate is often the primary concern, PPO for language models has multiple interacting parameters that require careful tuning. Understanding what each parameter controls is the first step toward debugging training failures.

The key parameters for PPO training are:

• clip_epsilon: The clipping threshold (typically 0.1 or 0.2) that constrains the policy update. Smaller values produce more conservative updates and can slow convergence. Larger values allow faster learning but risk instability. A common heuristic is to start with 0.2 and reduce it if training becomes unstable.
• kl_coef ($\beta$): Coefficient for the KL penalty term, controlling how closely the policy must stay to the reference model. This is perhaps the most important hyperparameter for preventing reward hacking. Higher values produce safer but slower learning. In practice, some implementations use an adaptive $\beta$ that increases when KL exceeds a target threshold and decreases when KL falls below it.
• value_coef: Weight for the value function loss in the total optimization objective. Balancing this against the policy loss affects how quickly the value function adapts relative to the policy. If the value function lags too far behind, advantage estimates become unreliable.
• gamma: Discount factor for future rewards. In RLHF, this is typically set to 1.0 since we care equally about all tokens in the response. Using a discount factor less than 1.0 would bias the model toward placing quality in early tokens rather than later ones, which is generally not desired.
• lam: The GAE smoothing parameter that balances bias and variance in advantage estimation. Values around 0.95 are standard. Lower values give more biased but lower-variance advantages; higher values give less biased but higher-variance advantages.

Implementing PPO for large language models involves several practical challenges beyond the algorithmic core. These engineering concerns often dominate the difficulty of real-world deployments and require careful attention. Getting the algorithm correct is necessary but not sufficient; the system must also be computationally efficient, numerically stable, and reliable to the diversity of prompts and responses encountered in practice.

Memory management is a practical constraint. During training, you must store activations for the policy model, reference model, and value model simultaneously. The reference model can be loaded in half precision or quantized to reduce memory footprint. Some implementations share the backbone between policy and value models, adding only a small value head. This sharing reduces memory requirements but couples the representations, which may affect optimization dynamics. The value head needs to represent very different information from the language modeling head, and sharing representations can create competing gradients.

Batch construction requires careful thought. Unlike supervised learning where examples are independent, PPO batches consist of complete generation trajectories. Generation is inherently sequential, making large-batch collection time-consuming. Implementations typically generate multiple responses in parallel across many prompts, using efficient batched inference to maximize throughput. The batch must contain enough diversity to provide a stable estimate of the gradient. Using too small a batch leads to high variance in the gradient estimate, causing oscillations in training. Using too large a batch may require too much memory or too long a generation phase between updates.

Response generation during training uses sampling rather than greedy decoding. This exploration is needed for PPO to discover high-reward responses that might differ from the reference policy's preferred outputs. Temperature and other sampling parameters become training hyperparameters that affect the exploration-exploitation balance. Too low a temperature leads to insufficient exploration; too high a temperature produces incoherent responses that receive low rewards and provide noisy gradient signals.

Advantage normalization stabilizes training significantly. Normalizing advantages to have zero mean and unit variance across each batch prevents any single trajectory from dominating the gradient. Without normalization, a few outlier responses with extreme advantages could destabilize training by creating large gradient updates. This is a simple intervention with a large practical impact.

To make PPO training tractable for large models, practitioners use several optimization strategies:

• Load the reference model in 8-bit quantized format, saving approximately 7 GB versus FP16 for a 7B model.
• Use gradient checkpointing on the policy model to trade computation for activation memory.
• Apply LoRA (Low-Rank Adaptation) to the policy model, which reduces the optimizer state memory dramatically by only maintaining momentum terms for the low-rank adapter parameters.
• Use DeepSpeed ZeRO-stage 3 to shard optimizer states, gradients, and parameters across multiple GPUs.

Applying PPO to language models, while effective, comes with significant challenges that require careful navigation. Understanding these limitations helps explain why the field has explored alternative methods and also sets realistic expectations for what PPO can and cannot achieve.

Sparse rewards make credit assignment fundamentally difficult. When a complete response receives a reward, determining which tokens contributed positively and which detracted is imprecise at best. The advantage function provides some differentiation through the value function's predictions, but it operates through imperfect estimates. The value function can only learn what patterns in intermediate states correlate with final rewards; it cannot directly observe the causal relationships between specific word choices and response quality. A model might learn that responses beginning with "Certainly! Here is..." receive high rewards, without learning anything about the actual quality of what follows. This imprecision can lead to slow learning, gaming of surface patterns, and inconsistent behavior across similar prompts.

The high computational cost of PPO creates a significant barrier to adoption and iteration. PPO requires generating complete responses during training, which is much slower than the teacher-forcing paradigm used in supervised learning. During teacher forcing, the model processes all tokens in parallel because it always sees the correct previous token. During generation, tokens must be produced sequentially. Each training step involves running the policy model to generate text, running it again to compute log probabilities, running the reference model for KL computation, running the value model for advantage estimation, and running the reward model for scoring. This pipeline is sequential in several places and requires keeping multiple large models in memory simultaneously. A single PPO training run can cost 10 to 100 times more compute than the equivalent supervised fine-tuning run, making extensive hyperparameter search impractical.

The sensitivity to hyperparameters creates both reproducibility challenges and practical fragility. The clipping parameter, KL coefficient, learning rate, batch size, number of PPO epochs, advantage normalization, and many other choices all interact in complex ways. Settings that work well for one model size, one task, or one reward model may fail on another configuration. Small changes to any hyperparameter can lead to training instability, where the policy diverges or collapses, or to reward hacking, where the policy finds degenerate solutions. This sensitivity is one of the primary motivations for exploring simpler alternatives like Direct Preference Optimization (DPO), which we will discuss in an upcoming chapter, that eliminate many of these hyperparameters by eliminating the RL training loop entirely.

The reliance on a learned reward model introduces all the challenges we discussed in the Reward Hacking chapter, and PPO's optimization pressure exacerbates them. The reward model is a proxy for human preferences, and optimizing it too aggressively can lead to responses that score highly according to the model while being less useful or even harmful. The reward model was trained on a finite dataset of human comparisons and generalizes imperfectly to the full distribution of prompts the policy will encounter. In the regions of prompt space that were underrepresented in the reward model training data, the reward model may give unreliable signals that the policy will happily exploit. The KL penalty mitigates but does not eliminate this risk, because a sufficiently flexible policy can find high-reward behaviors within the KL constraint that still fail human evaluation.

Finally, the training distribution mismatch between the supervised fine-tuning phase and the RL phase can create instabilities. The SFT model was trained to predict specific human-written responses. The RL policy must generate diverse responses that receive high reward, which may include phrasings and styles very different from the SFT training data. This distribution shift means that the value function must extrapolate far from its training distribution when evaluating novel responses, potentially giving noisy advantage estimates that destabilize the policy update. Some practitioners address this by warming up the value function with supervised training before starting RL, or by periodically refreshing the reference model to track the evolving policy.

This chapter translated the PPO algorithm from its general reinforcement learning formulation to the specific setting of language model alignment. We started with the basic observation that autoregressive generation is policy rollout: the language model is a policy, the vocabulary is the action space, the growing context is the state, and a generated response is a trajectory. This conceptual mapping enables a precise application of policy gradient theory to language generation.

The key conceptual mappings are:

• A language model is a stochastic policy, with states being the growing context and actions being vocabulary tokens at each position
• The action space is the vocabulary, typically containing 30,000 to 100,000 discrete options, many orders of magnitude larger than classical RL environments
• Rewards are assigned sparsely, with the reward model's score appearing only at the final token because quality can only be evaluated for complete responses
• The KL divergence penalty turns sparse rewards into dense per-token signal while simultaneously preventing policy collapse and reward hacking
• GAE computes advantages that differentiate between token choices by estimating the value of intermediate states under the current policy

The PPO objective for language models combines the clipped surrogate loss with KL-shaped rewards, creating a training signal that balances reward maximization against staying close to the reference distribution. This balance is important for stable training and for avoiding the reward hacking that would otherwise plague unconstrained optimization against a learned reward proxy.

Practical deployment requires significant engineering beyond the algorithmic core, including memory optimization through quantization and parameter-efficient methods, careful batch construction to ensure gradient quality, and reliable monitoring of KL divergence, probability ratios, and reward trends during training.

In the next chapter on the RLHF Pipeline, we will see how PPO fits into the complete workflow that turns a pretrained language model into an aligned assistant, including the data collection, reward model training, and iterative refinement stages that surround the PPO optimization we have studied here.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about applying PPO to language models.