Policy Gradient Optimization

Chapter 1: Gradient Ascent

The simplest approach: compute the gradient, step in that direction, repeat.

Here α > 0 is the step factor (learning rate). With a decaying step size that satisfies the Robbins-Monro conditions (∑α_k = ∞, ∑α_k² < ∞), gradient ascent converges to a local optimum.

The step length problem

The actual distance moved in parameter space is α·||∇U(θ)||. This depends on both α and the gradient magnitude. In RL, gradient magnitudes can vary wildly across training: early on when policies are random, gradients are small; later, when policies find high-reward trajectories, gradients can explode. A fixed α that worked fine in epoch 1 may cause catastrophic overshooting in epoch 50.

Gradient scaling and clipping

Two cheap tricks tame large gradients before the update:

python
def gradient_ascent_step(theta, grad_U, alpha, L_max=None):
    """One step of (optionally norm-clipped) gradient ascent."""
    g = grad_U                              # gradient estimate from Ch 11
    if L_max is not None:
        norm = np.linalg.norm(g)
        if norm > L_max:
            g = g * (L_max / norm)            # scale: preserve direction
    return theta + alpha * g

Chapter 2: The Natural Gradient

Plain gradient ascent measures distance in parameter space: ||θ' − θ||. But what we really care about is how much the policy behavior changes. Two policies with nearby parameters can behave very differently; two policies with distant parameters can behave nearly identically. Parameter-space distance is the wrong ruler.

The natural gradient swaps in a better ruler: KL divergence between the action distributions.

The restricted gradient: fix the step length

Start simpler. Instead of fixing α, directly constrain how far θ can move in parameter space. Maximize the linear approximation to U(θ) subject to a spherical constraint:

This has a clean closed-form solution. The Lagrangian tells us θ' = θ + λ∇U. Plugging into the constraint gives λ = √(2ε/||∇U||²), so:

The step length is always exactly √(2ε), regardless of the gradient magnitude. Only the direction of the gradient matters. This decouples direction from magnitude — a major improvement.

The problem: parameter space is still the wrong metric

Suppose your policy has two parameters: θ₁ controls a softmax temperature (highly sensitive — tiny changes drastically shift the action distribution) and θ₂ controls a bias term (insensitive — large changes barely matter). A spherical constraint in parameter space treats both the same. You'll overstep θ₁ and understep θ₂.

The natural gradient update

Solve the constrained problem with the Fisher metric:

The solution: let u = F_θ⁻¹∇U(θ) be the natural gradient. Then:

The direction u = F⁻¹∇U rotates the standard gradient to compensate for parameter-space curvature. In parameter directions where the policy is insensitive (small Fisher eigenvalue), u is larger, taking a bigger step. In sensitive directions (large Fisher eigenvalue), u is smaller, being conservative.

The Fisher information matrix

The Fisher information matrix F_θ is the expected outer product of the score function (gradient of log policy):

For stochastic policies, environment transition terms cancel out (same argument as the likelihood-ratio gradient), so we only need the policy's log-probability gradient summed over the trajectory:

We estimate this from the same rollout trajectories used for the gradient estimate. No extra episodes required.

Chapter 3: Trust Region Policy Optimization

Natural gradient gives us the right direction. TRPO adds a guarantee: it searches along that direction for the best step that provably keeps the new policy close to the old one.

Define the surrogate objective — an estimate of U(θ') computed from data collected under θ:

The ratio π_θ'/π_θ is importance weighting: it corrects for the fact that actions were sampled from the old policy, not the new one. Q_θ(s,a) is the action-value function estimate. This quantity can be evaluated for any θ' without running new episodes.

The optimization problem

TRPO solves a constrained optimization at each step:

The KL divergence constraint is in policy space, not parameter space. Two policy parameterizations that produce identical distributions are treated identically — TRPO is truly parameterization-invariant.

The algorithm: conjugate gradient + line search

The line search evaluates L(θ') and the KL for each candidate by computing importance ratios π_θ'/π_θ on the stored rollout data. No new episodes needed. Typically only 10-20 function evaluations per step.

Why a hard constraint, not a penalty?

A natural alternative: add β·KL to the objective as a penalty. The problem: the right β varies wildly across environments and even across training stages of the same environment. TRPO's hard constraint δ has a consistent semantic — "never let the policies diverge more than this" — making it a robust hyperparameter choice.

python (pseudocode)
def trpo_step(policy, rollouts, delta=0.01, max_backtracks=10):
    grad = compute_policy_gradient(rollouts)          # Ch 11
    u = conjugate_gradient(fisher_vector_product, grad) # Fu = grad
    step_size = np.sqrt(2 * delta / np.dot(u, fisher_vector_product(u)))
    full_step = step_size * u                          # natural gradient

    theta_old = policy.get_params()
    for k in range(max_backtracks):                     # line search
        theta_new = theta_old + (0.5 ** k) * full_step
        policy.set_params(theta_new)
        kl = compute_kl(policy, rollouts)
        improvement = compute_surrogate(policy, rollouts) - compute_surrogate_old(rollouts)
        if kl <= delta and improvement > 0:
            return                                       # accept step
    policy.set_params(theta_old)                         # reject all, keep old

Chapter 4: Proximal Policy Optimization

TRPO works. But it's complicated: you need a conjugate gradient solver, a Fisher-vector-product computation, and a custom line search. In 2017, Schulman et al. asked: what if you could get most of TRPO's benefits with a single clipped objective that standard optimizers (Adam, SGD) can handle directly?

The answer: PPO. Define the probability ratio:

r_t = 1 when θ = θ_old. r_t > 1 means θ now assigns more probability to this action than θ_old did. r_t < 1 means less probability.

The clipped surrogate objective is:

where A_t is the advantage estimate at time t, and ε is typically 0.2.

What the clipping actually does

The clip function forces r_t to stay in [1−ε, 1+ε]. The min takes the pessimistic bound between clipped and unclipped. Let's work through the two cases:

SHOWCASE: Interactive PPO Clipping Explorer

What to observe in the demo

With positive advantage: Set A=1.5. Move r from 0 to 2. Notice the orange curve rises linearly until r=1+ε, then goes flat. The optimizer gets no gradient signal to push r above 1+ε. That's the safety mechanism.

With negative advantage: Set A=−1.5. The objective becomes a flipped curve. Now the orange curve falls as r decreases, but flattens at r=1−ε. No benefit to making r very small (overly suppressing the action).

With zero advantage: Set A=0. The objective is flat everywhere — if an action is neither good nor bad, PPO makes no update regardless of r. This is correct.

python
import torch

def ppo_loss(log_pi_new, log_pi_old, advantages, eps=0.2):
    """
    log_pi_new: [B] log probabilities under new policy
    log_pi_old: [B] log probabilities under old policy (fixed)
    advantages: [B] advantage estimates (e.g. from GAE)
    Returns: scalar PPO clipped objective (to be maximized)
    """
    ratio = (log_pi_new - log_pi_old).exp()  # [B] — r_t(theta)
    # Unclipped surrogate
    obj1 = ratio * advantages                  # [B]
    # Clipped surrogate
    obj2 = ratio.clamp(1 - eps, 1 + eps) * advantages  # [B]
    # Pessimistic min, then mean over batch
    loss = torch.min(obj1, obj2).mean()        # scalar
    return -loss  # negate: PyTorch minimizes, we want to maximize

# Training loop (one iteration)
for epoch in range(n_epochs):             # e.g. 10 epochs on same data
    log_pi = policy.log_prob(states, actions)
    adv = compute_advantages(rewards, values, gamma, lam)
    loss = ppo_loss(log_pi, log_pi_old, adv)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Chapter 5: TRPO vs PPO

Both TRPO and PPO are trying to solve the same problem: take the biggest safe step toward a better policy. They differ dramatically in how they define and enforce safety.

When TRPO is still worth it

TRPO's hard KL constraint provides stronger guarantees in some situations. If you need:

• Monotonic improvement guarantees (safety-critical deployments)
• Invariance to policy parameterization (comparing across different network architectures)
• Understanding why a method fails (TRPO's diagnostics are cleaner)

Then TRPO is worth the extra implementation effort. In research, TRPO remains a useful baseline precisely because its guarantees are mathematically precise.

When PPO clips too aggressively

PPO's clipping is an approximation. There are edge cases where:

• The clip threshold ε is wrong for the current training stage (too tight early on, too loose later)
• Multiple epochs corrupt the old log-probs assumption (the importance ratio drifts)
• The advantage estimator has high variance, making the clip region the wrong shape

Modern practice: monitor the fraction of timesteps where r_t is clipped (the "clip fraction"). If it's >20-30%, reduce ε or the number of epochs. If it's near 0%, ε may be too tight.

Chapter 6: Implementation Details

The gap between understanding PPO conceptually and getting it to work is filled with engineering decisions that papers often omit. Here's what actually matters.

The full PPO objective

The paper's actual objective combines three terms:

Three components:

Advantage estimation: GAE

The advantage A_t appears in every term of L^CLIP. PPO almost always uses Generalized Advantage Estimation (GAE) from Schulman et al. (2016):

where δ_t = r_t + γV(s_t+1) − V(s_t) is the TD error. The λ parameter interpolates between:

• λ=0: pure TD(0) — low variance, high bias
• λ=1: Monte Carlo returns — high variance, low bias
• λ=0.95: the sweet spot in most benchmarks

Normalization: the hidden magic

These details are rarely highlighted but often make or break PPO:

Observation normalization: Maintain running statistics of observations and normalize them. Many environments have wildly different observation ranges. A neural network trained on unnormalized observations can spend most of training learning the scale, not the task.

Reward normalization: Scale rewards by a running std estimate of the discounted return. Prevents the value function from having to represent astronomically different value scales.

python
class PPOBuffer:
    """Collect rollouts, compute GAE advantages, normalize."""
    def compute_advantages(self, gamma=0.99, lam=0.95):
        adv = np.zeros_like(self.rewards)
        last_gae = 0
        for t in reversed(range(len(self.rewards))):
            if t == len(self.rewards) - 1:
                next_val = 0  # terminal
            else:
                next_val = self.values[t + 1]
            delta = self.rewards[t] + gamma * next_val - self.values[t]
            last_gae = delta + gamma * lam * last_gae
            adv[t] = last_gae
        # Normalize advantages (critical!)
        adv = (adv - adv.mean()) / (adv.std() + 1e-8)
        return adv

Data flow: shapes and types

Chapter 7: Hyperparameter Sensitivity

PPO is widely praised for robustness, but "robust" is relative. Understanding which hyperparameters matter — and why — separates practitioners who tune PPO correctly from those who cargo-cult settings from Atari papers onto robotics tasks.

The clip threshold ε

ε = 0.2 is the Schulman et al. default. It says: "never let the policy probability ratio deviate more than 20% from 1." In practice:

Number of epochs per rollout

More epochs = more gradient updates from the same data = more sample efficiency. But:

• After each epoch, θ changes, making log_pi_old stale
• The importance ratio r_t drifts away from 1
• Eventually the advantage estimates are wrong for the current policy

Empirical finding: 3-10 epochs is a good range. Beyond 10, you're usually hurting more than helping. Monitor the average KL between old and new policy per batch — if it exceeds ~0.01-0.05, reduce epochs.

Rollout length T and batch size

The rollout length T determines how many environment steps you collect before each update. Tradeoffs:

Value function clipping

A subtle detail: some PPO implementations also clip the value function update. The clipped value loss is:

The idea: prevent the value function from changing too rapidly. In practice, value clipping is often omitted (and Andrychowicz et al. 2021 showed it can sometimes hurt). Use it if you observe value function instability, skip it otherwise.

Chapter 8: Modern Extensions

PPO (2017) is now a foundation layer, not an endpoint. The field has built a series of extensions that address its remaining weaknesses.

PPO with KL penalty (PPO-KL)

Instead of clipping, add a KL penalty to the objective. Adaptively adjust the penalty coefficient β based on whether the current KL is above or below a target:

python (adaptive KL penalty)
if kl_mean < target_kl / 1.5:
    beta /= 2      # KL too small: relax penalty
elif kl_mean > target_kl * 1.5:
    beta *= 2      # KL too large: tighten penalty
loss = -L_clip + beta * kl_mean

This variant is preferred in some RLHF pipelines because it provides smoother gradients than hard clipping.

DAPO: Decoupled Clip and Dynamic Sampling (2025)

Recent work from Byte et al. identifies a subtle problem: in Group Relative Policy Optimization (GRPO, used in DeepSeek-R1), the entropy can collapse because tokens that are already highly probable get clipped and receive zero gradient. DAPO fixes this by:

• Using separate clip thresholds ε_low and ε_high for the lower and upper bounds
• Removing the lower clip entirely (i.e., only clipping upward) when advantage is positive
• Adding a token-level entropy loss to prevent collapse
• Dynamically oversampling prompts where the policy has improved significantly

REINFORCE Leave-One-Out (RLOO)

For language model fine-tuning, RLOO simplifies the PPO update: for each prompt, sample K responses, use the mean of the other K−1 as a baseline. This gives an unbiased advantage estimate without a learned value function. Recent work (Ahmadian et al. 2024) shows RLOO matches PPO on RLHF benchmarks with less code and memory.

Safe RL and constrained policy optimization

TRPO's constrained optimization framework naturally extends to constrained RL: maximize expected return subject to safety constraints (e.g., collision rate ≤ 5%). CPO (Constrained Policy Optimization, Achiam et al. 2017) extends TRPO to handle multiple safety constraints simultaneously using a primal-dual approach. PPO-based safe RL methods use Lagrangian relaxation to add constraint costs to the objective.

Offline-to-online fine-tuning

A modern pattern: pre-train a policy with behavior cloning (supervised learning on demonstrations), then fine-tune with PPO. The clip mechanism is essential here — it prevents PPO from immediately forgetting the learned behaviors. The initial r_t values are all near 1, and the clip keeps updates conservative, allowing the policy to improve while retaining good priors.

Chapter 9: Connections and What Comes Next

Policy gradient optimization doesn't exist in isolation. It's the optimization backbone that every modern RL system plugs into. Here's where Ch 12 sits in the larger picture.

What Ch 12 assumes (from Ch 11)

Every method here takes ∇U(θ) as input — the policy gradient estimate. Chapter 11 covers the likelihood-ratio trick, REINFORCE, baseline subtraction, and actor-critic gradient estimators. The quality of ∇U(θ) directly limits what optimization can achieve: a noisy gradient with high variance will cause unstable training regardless of how sophisticated the optimizer is.

What plugs into Ch 12 (Ch 13 — Actor-Critic)

Chapter 13 addresses a major weakness: the advantage estimate A_t used in L^CLIP. Using raw Monte Carlo returns has high variance. Actor-critic methods replace Monte Carlo with a learned value function (the critic), dramatically reducing variance. The critic is trained to minimize Bellman error while the actor uses PPO's clipped objective. This is the standard architecture for modern RL.

Comparison to value-based methods (Ch 7-9)

Connection to RLHF (Ch 18 — Imitation Learning)

RLHF pipelines use PPO as the optimization step, but with a reward model trained from human preferences instead of an environment reward. The policy is initialized from supervised fine-tuning (SFT) — which is imitation learning (Ch 18). The clip mechanism prevents the policy from drifting too far from the SFT policy, which acts as a regularizer against reward hacking.

Recommended reading

“To understand is to know what to do.”
— Wittgenstein, Zettel