Proximal Policy Optimization (PPO)

Since 2017, PPO has been the de-facto standard for on-policy reinforcement learning. The trick: instead of solving a hard KL constraint like TRPO, PPO caps the update by simple clipping. We build up the clipped surrogate objective step by step, show the PPO update loop as a Mermaid diagram, implement a runnable mini-PPO on the numpy GridWorld and add a realistic PyTorch skeleton.

1. The Problem: Policy-Gradient Steps Can Be Fatal

Vanilla policy gradient has a safety problem: a single oversized update step can throw the policy into a region from which it never recovers. Learning collapses – even when the step is mathematically pointing in the right direction.

The mathematically clean answer: Trust Region Policy Optimization (TRPO, Schulman et al. 2015). TRPO formulates the update rule as constrained optimisation:

$$\max_\theta \mathbb{E}\left[\frac{\pi_\theta(a \mid s)}{\pi_{\theta_{\text{old}}}(a \mid s)} \hat{A}t\right] \quad \text{s.t.} \quad \mathbb{E}[\text{KL}(\pi{\theta_{\text{old}}} \,|\, \pi_\theta)] \le \delta$$

In words: maximise the expected advantage, but stay within a KL ball around the old policy. Theoretically elegant – practically, TRPO needs conjugate gradients and Hessian-vector products. Few frameworks implement it robustly.

2. The PPO Idea: Clipping Instead of a KL Constraint

Schulman et al. (2017) asked the pragmatic question: What if we replace the trust-region guarantee with a simple, hard cap on the ratio? The result is the clipped surrogate objective:

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

with the importance-sampling ratio

$$r_t(\theta) = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_{\text{old}}}(a_t \mid s_t)}.$$

The logic in plain words:

• If $\hat{A}t > 0$ (the action was better than average), we want to increase $\pi\theta(a)$. But only up to $r_t = 1 + \epsilon$. Beyond that: update capped at 0 – no further incentive to let the ratio explode.
• If $\hat{A}t < 0$, we want to decrease $\pi\theta(a)$. But only down to $r_t = 1 - \epsilon$.
• The min ensures that only "the worse" of the two options (clipped vs. unclipped) counts – PPO is conservative.

Typical value: $\epsilon = 0.2$.

3. Multiple Epochs per Batch: PPO's Sample Efficiency

Classical REINFORCE throws every batch away after exactly one gradient step. PPO is far more generous: on a batch of typically 2048–8192 steps you run $K$ epochs with minibatches ($K = 4 \dots 10$).

This works because the clipping kicks in after the first few updates: as soon as the new policy is far enough from $\pi_{\theta_\text{old}}$, the clipped path delivers a zero gradient – the remaining epochs no longer change anything. PPO is therefore semi-on-policy: not fully off-policy like SAC, but significantly more efficient than pure REINFORCE.

4. Runnable: Mini-PPO on GridWorld

We build a complete PPO – without PyTorch, with a pure numpy policy table. The steps follow the classic PPO recipe:

1. Collect a rollout under $\pi_{\theta_\text{old}}$.
2. Compute GAE advantages (simplified here: plain TD error).
3. Update with the clipped loss for $K$ epochs.

import numpy as np

rng = np.random.default_rng(0)
LENGTH = 7
GAMMA = 0.95
LR_A, LR_C = 0.05, 0.1
EPS = 0.2          # clip range
K_EPOCHS = 4
ITERATIONS = 200
STEPS_PER_ITER = 256

theta = np.zeros((LENGTH, 2))
V     = np.zeros(LENGTH)

def softmax(x):
    z = x - x.max(); e = np.exp(z); return e / e.sum()

def policy_probs_from(theta_arr, s): return softmax(theta_arr[s])

def step(s, a):
    s_next = max(0, min(LENGTH - 1, s + (1 if a == 1 else -1)))
    done = s_next in (0, LENGTH - 1)
    if done:
        r = 1.0 if s_next == LENGTH - 1 else -1.0
    else:
        r = 0.0
    return s_next, r, done

def collect_rollout(theta_old, n_steps):
    """Collect n_steps transitions, restart episodes in the middle."""
    data = []
    s = LENGTH // 2
    for _ in range(n_steps):
        probs = policy_probs_from(theta_old, s)
        a = int(rng.choice(2, p=probs))
        s_next, r, done = step(s, a)
        data.append((s, a, r, s_next, done, probs[a]))
        s = LENGTH // 2 if done else s_next
    return data

def compute_advantages(data):
    advs = []
    for (s, a, r, s_next, done, _) in data:
        target = r + (0.0 if done else GAMMA * V[s_next])
        advs.append(target - V[s])
    return np.array(advs)

ep_returns_per_iter = []
for it in range(ITERATIONS):
    theta_old = theta.copy()
    data = collect_rollout(theta_old, STEPS_PER_ITER)
    advs = compute_advantages(data)
    # Normalisation stabilises the loss.
    advs = (advs - advs.mean()) / (advs.std() + 1e-8)

    # K epochs of clipped update.
    for _ in range(K_EPOCHS):
        for (s, a, r, s_next, done, p_old_a), adv in zip(data, advs):
            probs_new = policy_probs_from(theta, s)
            p_new_a = probs_new[a]
            ratio = p_new_a / max(p_old_a, 1e-8)
            clipped = min(max(ratio, 1 - EPS), 1 + EPS)

            # Surrogate contribution: the min(...) decides which path is active.
            unclipped_obj = ratio * adv
            clipped_obj   = clipped * adv
            # Unclipped path provides the min → regular PG gradient.
            # Clipped path provides the min   → gradient = 0.
            if unclipped_obj <= clipped_obj:
                grad_log = -probs_new.copy(); grad_log[a] += 1.0
                # Gradient of ratio*adv w.r.t. theta equals ratio * grad_log * adv.
                theta[s] += LR_A * ratio * grad_log * adv

        # Critic update (TD).
        for (s, a, r, s_next, done, _) in data:
            target = r + (0.0 if done else GAMMA * V[s_next])
            V[s] += LR_C * (target - V[s])

    # Extract terminal returns from the done markers (approximation).
    done_returns = [r for (_, _, r, _, done, _) in data if done]
    if done_returns:
        ep_returns_per_iter.append(np.mean(done_returns))

print(f"Mean terminal reward, first 20 iter : {np.mean(ep_returns_per_iter[:20]):+.3f}")
print(f"Mean terminal reward, last  20 iter : {np.mean(ep_returns_per_iter[-20:]):+.3f}")
print("Learned P(right):")
for s in range(LENGTH):
    print(f"  s={s}: {policy_probs_from(theta, s)[1]:.2f}  V={V[s]:+.2f}")

The clip mechanism is visible in the if unclipped_obj <= clipped_obj branch: as soon as the ratio leaves the $[1-\epsilon, 1+\epsilon]$ corridor and that path produces the worse surrogate value, the gradient drops to 0. That is exactly PPO's safety net.

5. PyTorch Skeleton (Concept – Not Runnable in the Browser)

For a realistic impression of PPO in PyTorch – this variant requires gymnasium and runs locally only:

import torch
import torch.nn as nn
from torch.distributions import Categorical

class ActorCritic(nn.Module):
    def __init__(self, obs_dim, n_actions, hidden=64):
        super().__init__()
        self.shared = nn.Sequential(
            nn.Linear(obs_dim, hidden), nn.Tanh(),
            nn.Linear(hidden, hidden), nn.Tanh(),
        )
        self.actor_head  = nn.Linear(hidden, n_actions)
        self.critic_head = nn.Linear(hidden, 1)

    def forward(self, x):
        h = self.shared(x)
        return Categorical(logits=self.actor_head(h)), self.critic_head(h).squeeze(-1)


def ppo_update(model, optimizer, batch, eps=0.2, k_epochs=4):
    """batch = (obs, actions, old_log_probs, advantages, returns)."""
    obs, acts, old_log_probs, advs, returns = batch
    advs = (advs - advs.mean()) / (advs.std() + 1e-8)

    for _ in range(k_epochs):
        dist, values = model(obs)
        new_log_probs = dist.log_prob(acts)
        ratio = (new_log_probs - old_log_probs).exp()

        unclipped = ratio * advs
        clipped   = torch.clamp(ratio, 1 - eps, 1 + eps) * advs
        policy_loss = -torch.min(unclipped, clipped).mean()

        value_loss = ((returns - values) ** 2).mean()
        entropy    = dist.entropy().mean()

        loss = policy_loss + 0.5 * value_loss - 0.01 * entropy
        optimizer.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 0.5)
        optimizer.step()

Three production details missing from the numpy toy:

• Gradient-norm clipping (clip_grad_norm_(0.5)) – prevents spikes from outlier batches.
• Entropy bonus (-0.01 * entropy) – rewards exploration and prevents premature determinism.
• Value function and policy in one loss – with a fixed weight of 0.5 (PPO standard).

6. Hyperparameter Cheat Sheet

7. Discussion: Why PPO Became So Dominant

PPO is not the theoretically most elegant algorithm. It wins through engineering qualities:

• Few hyperparameters, all of them robust.
• Standard implementation fits into 200 lines of PyTorch.
• Scales to distributed setups (parallel workers, IMPALA-style architectures).
• Works in discrete and continuous action spaces alike (with a Gaussian policy).

These exact qualities made PPO the choice for RLHF in modern LLMs (ChatGPT, Llama, Claude) – with its own adjustments (reference-KL penalty, reward-model-based reward).

8. Wrap-up Quiz

9. Further Reading

• Next article: Soft Actor-Critic (SAC) (article 05 of this series).
• Cross reference: self-learning-agents/02-rl-fuer-code-agenten – PPO as the RLHF backbone.
• Knowledge base: masteringpytorch-secondedition.md – chapter on modern policy gradients.

Sources

• Schulman, J. et al. (2017) – Proximal Policy Optimization Algorithms, arXiv:1707.06347.
• Schulman, J. et al. (2015) – Trust Region Policy Optimization, ICML.
• Engstrom, L. et al. (2020) – Implementation Matters in Deep RL: A Case Study on PPO and TRPO – the famous reproducibility study on PPO tricks.