Proximal Policy Optimisation

The policy-gradient surrogate is unbiased only near the sampling policy. PPO enforces that locality bluntly: it takes the minimum of the unclipped and clipped objectives, which removes the incentive to push the probability ratio outside $[1-\epsilon,\, 1+\epsilon]$.

In the RLHF setting the reward is a single scalar at the end of the sequence, shaped by a per-token KL penalty against the reference policy. Advantages come from GAE over that shaped reward.

import torch
from torch import Tensor


def ppo_clip_loss(
    logp: Tensor, logp_old: Tensor, adv: Tensor,
    eps: float = 0.2,
) -> Tensor:
    ratio = (logp - logp_old).exp()
    unclipped = ratio * adv
    clipped = ratio.clamp(1 - eps, 1 + eps) * adv
    return -torch.min(unclipped, clipped).mean()


def gae(rewards: Tensor, values: Tensor, gamma=1.0, lam=0.95) -> Tensor:
    """rewards, values: [T]. Returns advantages [T]."""
    adv = torch.zeros_like(rewards)
    running = 0.0
    for t in reversed(range(len(rewards))):
        next_v = values[t + 1] if t + 1 < len(values) else 0.0
        delta = rewards[t] + gamma * next_v - values[t]
        running = delta + gamma * lam * running
        adv[t] = running
    return adv

Normalise advantages over the batch, not over the sequence — per-sequence normalisation destroys the relative ranking between samples that the update depends on. Compute the ratio in log space as shown; exponentiating two probabilities and dividing loses precision exactly where the ratio matters.