Policy Gradient Methods
& RLHF

Reinforcement learning is the science of sequential decision making under uncertainty. An agent interacts with an environment, takes actions, and receives rewards. The goal is to learn a policy—a mapping from states to actions—that maximizes cumulative reward.

Why Policy Gradients?

Before policy gradients, RL relied primarily on value-based methods (Q-learning, SARSA). These methods learn the value of state-action pairs and derive a policy implicitly. But they face fundamental limitations:

Policy gradient methods are especially critical for LLM training: a language model's action space (vocabulary tokens) is discrete but enormous (~50,000 tokens), and we need stochastic policies that can represent nuanced, diverse outputs. Value-based methods cannot scale here.

A policy $\pi_\theta$ is a probability distribution over actions given states, parameterized by $\theta$ (e.g., neural network weights):

Common Parameterizations

Softmax policy (discrete actions):

where $h_\theta(s, a)$ is the preference (logit) of action $a$ in state $s$.

Gaussian policy (continuous actions):

For LLMs: The policy is a transformer with parameters $\theta$. Given context (state) $s = (x_1, \ldots, x_t)$, the policy outputs a distribution over the next token:

The core idea is strikingly simple: make good actions more probable, make bad actions less probable. We define a performance objective and take gradient steps to improve it.

The policy gradient The gradient of expected return w.r.t. policy parameters θ. Points in the direction of steepest increase in expected reward. tells us how to adjust $\theta$ to increase expected return. The key insight is that we can estimate this gradient using samples from the policy—no model of the environment is needed.

(A) Intuition: Increase Probability of Good Actions

(B) Formal Statement

Define the performance objective as expected cumulative discounted reward:

The Policy Gradient Theorem states:

where $Q^{\pi_\theta}(s_t, a_t) = \mathbb{E}\left[\sum_{t'=t}^{T} \gamma^{t'-t} r_{t'}\right]$ is the action-value function (expected future reward from $(s_t, a_t)$).

(C) Detailed Step-by-Step Derivation

REINFORCE (Williams, 1992) is the simplest instantiation of the policy gradient theorem. It estimates $Q^{\pi}(s_t, a_t)$ using Monte Carlo returns—complete trajectory rollouts.

The REINFORCE update rule:

Baselines

A baseline $b(s)$ is any function of the current state (not the action) that we subtract from the return without introducing bias. The modified gradient:

Why is this unbiased? Because $\mathbb{E}_{a \sim \pi_\theta}[\nabla_\theta \log \pi_\theta(a|s) \cdot b(s)] = b(s) \cdot \nabla_\theta \sum_a \pi_\theta(a|s) = b(s) \cdot \nabla_\theta 1 = 0$.

The Advantage Function

Using the value function $V^\pi(s_t)$ as baseline gives the advantage function:

The advantage function A(s,a) measures how much better action a is compared to the average action in state s. Positive advantage = better than average; negative = worse than average. answers the question: "How much better is action $a$ compared to the average action in state $s$?"

Generalized Advantage Estimation (GAE)

GAE (Schulman et al., 2015) provides a bias-variance trade-off using temporal differences:

where $\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$ is the TD residual. Parameter $\lambda \in [0,1]$: $\lambda=0$ gives pure TD (low variance, high bias); $\lambda=1$ gives Monte Carlo (high variance, low bias).

Actor-Critic methods maintain two separate function approximators:

A2C / A3C Architecture

The actor loss (gradient ascent on policy):

The critic loss (value function regression):

Often combined with an entropy bonus to encourage exploration:

Natural Policy Gradient (NPG)

Standard gradient descent treats all parameter directions equally. But the policy space has a Riemannian geometry: equal changes in $\theta$ can have wildly different effects on $\pi_\theta$. The natural gradient accounts for this using the Fisher Information Matrix F(θ) = E[∇logπ · ∇logπᵀ] measures the curvature of the KL divergence between nearby policies. It encodes how policy-space distances relate to parameter-space distances. :

NPG makes updates that are invariant to the parameterization of $\theta$—a critical property for neural networks where the same policy can be represented by many different $\theta$ values.

Trust Region Policy Optimization (TRPO)

TRPO (Schulman et al., 2015) formalizes the intuition that policy updates should not be too large. It solves a constrained optimization problem:

The KL divergence constraint KL(p‖q) = Σ p(x) log(p(x)/q(x)). Here it measures how much the new policy deviates from the old one. TRPO requires this to be small (≤ δ ≈ 0.01). ensures the new policy stays close to the old one, maintaining the validity of the importance-weighted objective. TRPO uses conjugate gradients and line search, making it theoretically sound but computationally expensive.

Proximal Policy Optimization (Schulman et al., 2017) is arguably the most important modern policy gradient algorithm. It achieves TRPO's stability with a fraction of the computational cost.

(A) Starting Point: Policy Gradient Objective with Importance Sampling

When we collect data under policy $\pi_{\theta_\text{old}}$ but optimize $\theta \neq \theta_\text{old}$, we need importance sampling:

Define the probability ratio:

The surrogate objective (CPI — Conservative Policy Iteration):

(B) Introducing the KL Constraint (TRPO Approach)

TRPO adds an explicit KL constraint. The theoretical guarantee (from the surrogate objective bound):

This shows the true objective $J(\theta)$ is lower bounded by the surrogate minus a KL penalty. Maximizing this lower bound gives monotonic improvement guarantees.

(C) Trust Region Idea

(D) PPO Clipped Objective — Full Derivation

Policy gradient methods are instances of stochastic approximation (Robbins-Monro, 1951). We want to find $\theta^* = \arg\max_\theta J(\theta)$ where $J(\theta)$ cannot be computed exactly, only estimated.

where $\widehat{\nabla_\theta J}$ is an unbiased or consistent estimator of the true gradient. The Robbins-Monro conditions for convergence:

Typical choice: $\alpha_k = \frac{1}{k}$ or $\alpha_k = \frac{c}{k + c'}$. In practice, fixed learning rates with Adam often work better.

Modern deep RL combines policy gradients with deep neural networks as function approximators. Key developments:

Deep Deterministic Policy Gradient (DDPG)

For continuous action spaces, DDPG learns a deterministic policy $\mu_\theta(s)$ using the chain rule:

Soft Actor-Critic (SAC)

SAC maximizes a entropy-augmented objective, trading off reward for policy entropy:

SAC is off-policy (uses a replay buffer), entropy-regularized, and highly sample efficient—making it one of the most popular algorithms for continuous control.

Practical Tricks for Stability

(A) What is RLHF?

Reinforcement Learning from Human Feedback (RLHF) is the technique that transforms a pretrained language model into a helpful, harmless, and honest AI assistant. It aligns model behavior with human preferences using RL—specifically policy gradient methods.

(B) The RLHF Pipeline

(C) Role of Policy Gradient

The LLM is the policy $\pi_\theta$. Generating a response is a trajectory of token selections. The reward model provides a scalar reward at the end of generation:

The policy gradient update then pushes the LLM to generate responses that get higher reward from the reward model.

How GPT-style Models are Post-Trained

Starting from a pretrained model $\pi_\text{SFT}$, the RLHF loop works as follows:

Reward Modeling

The reward model $R_\phi$ is trained using the Bradley-Terry model for pairwise comparisons:

Training loss (negative log-likelihood of preferences):

KL Regularization: Why It's Critical

Without regularization, the policy would reward hack—find adversarial inputs that exploit weaknesses in the reward model while generating incoherent text. The KL penalty prevents this:

Per-Token Reward Formulation

In practice, the reward is applied per-token for the PPO advantage computation:

This distributes the terminal reward across tokens, enabling per-token advantage estimation.

DPO: Beyond PPO for LLMs

Direct Preference Optimization (Rafailov et al., 2023) bypasses the reward model entirely by showing that the optimal policy under KL-regularized RLHF satisfies:

DPO implicitly trains the reward model and policy simultaneously, dramatically simplifying the pipeline while achieving comparable results.

From RL to LLM Alignment: The Full Picture

Current Research Frontiers