Reinforcement Learning Principles Refresher

Supervised fine-tuning alone often struggles to capture the full spectrum of desired behaviors for large language models. Reinforcement Learning (RL) provides a framework to optimize LLMs based on a learned signal representing human preferences. A background in Reinforcement Learning (RL) is typically expected. A focused refresher is provided on the principles most pertinent to RLHF, particularly those underlying Proximal Policy Optimization (PPO).

Framing LLM Generation as an RL Problem

At its core, RL deals with agents learning to make sequences of decisions in an environment to maximize a cumulative reward. We can frame the task of generating text with an LLM using the standard RL formalism, specifically the Markov Decision Process (MDP):

State ( $s$ ): The current context, typically comprising the initial prompt and the sequence of tokens generated so far.
Action ( $a$ ): The next token to be generated by the LLM. The action space is the vocabulary of the language model.
Policy ( $\pi_\theta(a|s)$ ): The LLM itself, parameterized by $\theta$ . It defines a probability distribution over the next possible tokens (actions) given the current state (prompt + previous tokens).
Reward ( $R(s, a)$ ): This is the critical component learned from human feedback. After a sequence (response) is generated, a reward signal reflects how well that sequence aligns with human preferences. This is often assigned only at the end of the sequence, based on the output of a Reward Model (RM), which we will detail in Chapter 3.
Transition ( $P(s'|s, a)$ ): In text generation, the transition is typically deterministic. Given the current state $s$ (e.g., "Translate 'hello' to French:") and the chosen action $a$ (e.g., the token "Bonjour"), the next state $s'$ is simply the concatenation (e.g., "Translate 'hello' to French: Bonjour").

The objective in this RL setting is to adjust the LLM's parameters $\theta$ to find a policy $\pi_\theta$ that maximizes the expected cumulative reward (often called the expected return) for the generated sequences:

J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_{t=0}^{T} \gamma^t R(s_t, a_t) \right]

Here, $\tau = (s_0, a_0, s_1, a_1, ...)$ represents a full trajectory (e.g., a complete generated response), $\pi_\theta$ dictates the probability of generating that trajectory, $T$ is the length of the sequence, and $\gamma$ is a discount factor (often set to 1 for finite-horizon text generation).

Policy Gradients and the Need for Stability

Policy Gradient methods directly optimize the policy parameters $\theta$ by performing gradient ascent on the objective $J(\theta)$ . A common form of the policy gradient theorem gives us:

\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t|s_t) \hat{A}_t \right]

Where $\hat{A}_t$ is an estimate of the advantage function at timestep $t$ . The advantage $A(s, a) = Q(s, a) - V(s)$ represents how much better taking action $a$ in state $s$ is compared to the average action according to the current policy, estimated by the value function $V(s)$ . Intuitively, this increases the probability of actions that lead to higher-than-expected rewards and decreases the probability of actions leading to lower-than-expected rewards.

While effective, basic policy gradient methods can suffer from high variance and instability, especially when dealing with high-dimensional parameter spaces like those found in LLMs. A single large gradient update, driven by a particularly high or low reward sample, could drastically alter the policy, potentially leading to a collapse in performance. This instability motivated the development of more sophisticated algorithms like PPO.

Proximal Policy Optimization (PPO)

PPO has become the standard RL algorithm for fine-tuning LLMs in the RLHF pipeline due to its relative simplicity, stability, and strong empirical performance. It addresses the instability issue of vanilla policy gradients by constraining how much the policy can change in each update step.

PPO optimizes a surrogate objective function that incorporates a mechanism to discourage large policy updates. The most common variant uses a clipped objective:

L^{CLIP}(\theta) = \mathbb{E}_t \left[ \min \left( r_t(\theta) \hat{A}_t, \text{clip}(r_t(\theta), 1 - \epsilon, 1 + \epsilon) \hat{A}_t \right) \right]

Let's break this down:

$\mathbb{E}_t [...]$ indicates taking the average over a batch of collected experience (timesteps).
$r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}$ is the probability ratio between the current policy $\pi_\theta$ (being optimized) and the old policy $\pi_{\theta_{old}}$ (used to collect the data). It measures how much the policy has changed.
$\hat{A}_t$ is the estimated advantage for timestep $t$ , often calculated using Generalized Advantage Estimation (GAE), which helps reduce variance. This requires learning a value function $V(s)$ (the "critic") alongside the policy (the "actor").
$\text{clip}(r_t(\theta), 1 - \epsilon, 1 + \epsilon)$ restricts the ratio $r_t(\theta)$ to stay within the interval $[1 - \epsilon, 1 + \epsilon]$ . The hyperparameter $\epsilon$ (e.g., 0.2) defines the clipping range.
$\min(...)$ $min (...)$ takes the minimum of two terms:
1. The unclipped objective $r_t(\theta) \hat{A}_t$ .
2. The clipped objective $\text{clip}(r_t(\theta), 1 - \epsilon, 1 + \epsilon) \hat{A}_t$ .

The effect of this clipping is to limit the influence of actions where the policy ratio $r_t(\theta)$ moves outside the $[1-\epsilon, 1+\epsilon]$ boundary. If the advantage $\hat{A}_t$ is positive (meaning the action was better than average), the objective is clipped from above if the policy change is too large ( $r_t > 1+\epsilon$ ), preventing overly aggressive updates. If the advantage $\hat{A}_t$ is negative (the action was worse than average), the objective is clipped from below if the policy change is too large ( $r_t < 1-\epsilon$ ), preventing overly drastic reductions in probability for that action. This keeps the new policy close to the old policy, ensuring more stable learning.

The KL Divergence Constraint in RLHF

While PPO's clipped objective inherently encourages smaller updates, RLHF implementations often add an explicit penalty term based on the Kullback-Leibler (KL) divergence between the current policy $\pi_\theta$ and a reference policy $\pi_{ref}$ . Typically, $\pi_{ref}$ is the initial SFT model.

This KL penalty is usually incorporated directly into the reward signal used during PPO training:

R_{total}(s, a) = R_{RM}(s, a) - \beta \, D_{KL}(\pi_\theta(\cdot|s) || \pi_{ref}(\cdot|s))

Or, more commonly in practice for per-token rewards:

R_{token}(s_t, a_t) = R_{RM, t} - \beta \log \frac{\pi_\theta(a_t|s_t)}{\pi_{ref}(a_t|s_t)}

Here:

$R_{RM}$ is the reward obtained from the Reward Model.
$\beta$ is a hyperparameter controlling the strength of the KL penalty.
The KL term penalizes the policy $\pi_\theta$ for diverging too much from the reference policy $\pi_{ref}$ on a token-by-token basis.

Why is this KL penalty so significant in RLHF?

Preserving Capabilities: It prevents the LLM, while optimizing for the human preference reward $R_{RM}$ , from drifting too far away from the general language modeling capabilities and knowledge embedded in the original SFT model. Without it, the model might learn to generate repetitive or nonsensical text that happens to score highly on the RM ("reward hacking") but is otherwise low quality.
Stabilization: Similar to the clipping in PPO, it acts as a regularization term, ensuring smoother and more stable training updates.

We will explore the practical implementation and tuning of this KL penalty within the PPO algorithm in Chapter 4.

Tying it Together for RLHF

This refresher highlights the core RL components adapted for RLHF:

We frame text generation as an MDP where the LLM acts as the policy.
The goal is to maximize a reward signal derived from human preferences (via the Reward Model).
PPO is employed to update the LLM's policy parameters stably.
A KL divergence penalty relative to the initial SFT model is incorporated to maintain language quality and prevent catastrophic forgetting or reward hacking.

Understanding these RL foundations, particularly the mechanics and rationale behind PPO and the KL constraint, is essential for effectively implementing and troubleshooting the RL fine-tuning stage of the RLHF pipeline, which we will detail in the upcoming chapters.

Was this section helpful?

References

Reinforcement Learning: An Introduction, Richard S. Sutton and Andrew G. Barto, 2018 (MIT Press) - Comprehensive introduction to reinforcement learning fundamentals, including Markov Decision Processes, policy gradient methods, and value functions. Essential for foundational understanding.
Proximal Policy Optimization Algorithms, John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, 2017 arXiv preprint arXiv:1707.06347 DOI: 10.48550/arXiv.1707.06347 - Introduces the Proximal Policy Optimization (PPO) algorithm, detailing its clipped surrogate objective and advantages for stable policy optimization, which is central to RLHF.
Training language models to follow instructions with human feedback, Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, Ryan Lowe, 2022 arXiv preprint arXiv:2203.02155 DOI: 10.48550/arXiv.2203.02155 - Describes the Reinforcement Learning from Human Feedback (RLHF) pipeline, specifically detailing the use of PPO with a KL divergence penalty for aligning language models with human preferences.