16. PPO's Clipped Surrogate Loss

Problem

Proximal Policy Optimization (PPO) updates a policy using an importance ratio between the new and old policy's probability of the action taken, ratio = exp(new_log_prob - old_log_prob), multiplied by the advantage estimate. Left unconstrained, a large ratio can push the policy arbitrarily far in one update and destabilize training. PPO's fix is the clipped surrogate objective: take the minimum of the unclipped ratio * advantage and a version of it with the ratio clamped to [1-eps, 1+eps], which removes the incentive to move the ratio far outside that trust region while still allowing beneficial small updates.

Given new_log_probs, old_log_probs, and advantages (each shape (batch,)), compute the clipped surrogate loss as a scalar.

Source: src/16_ppo_clipped_surrogate_loss.py

def ppo_clipped_surrogate_loss(new_log_probs, old_log_probs, advantages, clip_eps=0.2) -> torch.Tensor

>>> import torch
>>> new_log_probs = torch.tensor([1.0])
>>> old_log_probs = torch.tensor([0.0])
>>> advantages = torch.tensor([1.0])
>>> ppo_clipped_surrogate_loss(new_log_probs, old_log_probs, advantages, clip_eps=0.2)
tensor(-1.2000)

Step-by-Step Approach

1. Compute the importance ratio in log-space first and exponentiate: ratio = torch.exp(new_log_probs - old_log_probs). Working from log-probs (rather than dividing raw probabilities, which most policy heads don't even expose directly) is both more natural given what policies output and numerically safer.
2. Compute the unclipped surrogate: surr1 = ratio * advantages.
3. Compute the clipped surrogate: clamp the ratio to [1-clip_eps, 1+clip_eps] with torch.clamp, then multiply by advantages to get surr2.
4. Take the elementwise minimum of the two surrogates with torch.min(surr1, surr2) — this is the "pessimistic bound" that gives PPO its name-adjacent stability property.
5. Average over the batch and negate, since PPO's objective is something to maximize but training loops minimize a loss: -torch.min(surr1, surr2).mean().

The key insight is why min (not just always using the clipped version) matters: clipping alone would still let a large *negative* advantage's surrogate move arbitrarily favorably if the ratio happened to fall outside the clip range in the "helpful" direction, but the ratio is not clamped at all, so it would understate the pessimistic case. Taking the min of clipped and unclipped means the clip range only ever removes the incentive to over-improve the objective, never adds an incentive that wasn't already there — the objective is only ever more conservative than the unclipped one, never more optimistic.

Reference solution

import torch


def ppo_clipped_surrogate_loss(new_log_probs, old_log_probs, advantages, clip_eps=0.2):
    # working in log-space and exponentiating the difference is the
    # numerically stable way to compute a probability ratio -- never divide
    # raw probabilities directly
    ratio = torch.exp(new_log_probs - old_log_probs)  # (batch,)
    surr1 = ratio * advantages  # unclipped surrogate
    clipped_ratio = torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps)
    surr2 = clipped_ratio * advantages  # clipped surrogate
    # taking the min (not just using the clipped version) is what makes this
    # a *pessimistic* bound: it only clips away improvement, never penalty
    loss = -torch.min(surr1, surr2).mean()
    return loss

Key Functions & Tricks

• torch.exp(new_log_probs - old_log_probs) — computes a probability ratio via log-space subtraction, avoiding ever dividing two small raw probabilities directly.
• torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps) — the trust-region clip; both bounds are supplied explicitly rather than clamping only one side.
• torch.min(surr1, surr2) — elementwise minimum, the mechanism that makes the clip a one-directional (pessimistic) constraint rather than a hard cap.
• .mean() then negate — converts a per-sample objective into a scalar loss suitable for loss.backward(), flipping sign since PPO's objective is maximized but optimizers minimize.

How to Recognize This Pattern