Policy gradient and PPO

Deep dive

How it works

Parts:

• policy π — action distribution from the state (actor)
• value function V — baseline and return target (critic)
• rollout buffer — s, a, log π_old, r, done, V for T steps × N environments
• advantage estimator — GAE
• optimizer loop — K epochs of clipped-surrogate minibatch steps

1. Rollout. Run N environments for T steps with π_old. Store s, a, log π_old(a|s), r, done, V(s).
2. GAE, backwards in time: δ_t = r_t + γ(1 − done_t)V(s_{t+1}) − V(s_t), Â_t = δ_t + γλ(1 − done_t)Â_{t+1}. Returns R_t = Â_t + V(s_t) are the value targets.
3. Optimize for K epochs over shuffled minibatches:
   • r = exp(log π(a|s) − log π_old(a|s)).
   • Normalize  within the minibatch (zero mean, unit std).
   • L_clip = −mean(min(r·Â, clip(r, 1−ε, 1+ε)·Â)).
   • L_V = c_v·mean((V(s) − R)²) (c_v 0.5-1), L_H = −c_e·mean(entropy) (c_e of order 0.01 for exploration, 0 for many continuous-control tasks).
   • Total loss → backprop → clip global gradient norm (0.5 is common) → Adam step.
4. Optionally stop the epochs early when the approximate KL mean((r − 1) − log r) exceeds a target such as 0.01-0.02.

Why clipping works: for  > 0 the objective stops growing once r > 1+ε; for  < 0, once r < 1−ε. In those regions the gradient is exactly zero, so no single sample can drag the policy far.

 > 0:  objective  ____/‾‾‾‾‾‾   flat after r = 1+ε          < 0:  ‾‾‾‾\____   flat before r = 1−ε
                   r →                                                r →

Code

/// PPO on a tabular softmax policy and tabular value function (chain env: N states, left/right,
/// +1 at the right end). The math per sample is exactly what a deep PPO computes per output head.
const N: usize = 8; const A: usize = 2;

struct Rng(u64);
impl Rng { fn f(&mut self) -> f32 { self.0 ^= self.0 << 13; self.0 ^= self.0 >> 7; self.0 ^= self.0 << 17;
                                     (self.0 >> 40) as f32 / (1u64 << 24) as f32 } }

fn softmax(l: &[f32]) -> [f32; A] {
    let m = l[0].max(l[1]);
    let e = [(l[0] - m).exp(), (l[1] - m).exp()];
    [e[0] / (e[0] + e[1]), e[1] / (e[0] + e[1])]
}

/// GAE(γ, λ) over one rollout; `values` has T+1 entries (bootstrap value last).
pub fn gae(rew: &[f32], values: &[f32], done: &[bool], gamma: f32, lam: f32) -> (Vec<f32>, Vec<f32>) {
    let t = rew.len();
    let (mut adv, mut last) = (vec![0.0; t], 0.0);
    for i in (0..t).rev() {
        let nonterm = if done[i] { 0.0 } else { 1.0 };
        let delta = rew[i] + gamma * nonterm * values[i + 1] - values[i];
        last = delta + gamma * lam * nonterm * last;
        adv[i] = last;
    }
    let ret = adv.iter().zip(values).map(|(a, v)| a + v).collect();
    (adv, ret)
}

fn main() {
    let (gamma, lam, eps, lr, ent_c, epochs, mb, horizon) = (0.99, 0.95, 0.2f32, 0.5f32, 0.01f32, 4, 32, 256);
    let mut logits = vec![0.0f32; N * A];
    let mut value = vec![0.0f32; N];
    let mut rng = Rng(7);
    let mut s = 0usize;
    for iter in 0..30 {
        let (mut ss, mut aa, mut lp_old, mut rr, mut dd, mut vv) = (vec![], vec![], vec![], vec![], vec![], vec![]);
        for _ in 0..horizon {                                       // 1. rollout with π_old
            let p = softmax(&logits[s * A..s * A + A]);
            let a = if rng.f() < p[0] { 0 } else { 1 };
            let s2 = if a == 1 { s + 1 } else { s.saturating_sub(1) };
            let done = s2 == N - 1;
            ss.push(s); aa.push(a); lp_old.push(p[a].ln()); vv.push(value[s]);
            rr.push(if done { 1.0 } else { 0.0 }); dd.push(done);
            s = if done { 0 } else { s2 };
        }
        vv.push(value[s]);                                          // bootstrap value
        let (adv, ret) = gae(&rr, &vv, &dd, gamma, lam);            // 2. advantages, value targets
        let mut idx: Vec<usize> = (0..horizon).collect();
        let mut clipfrac = 0.0;
        for _ in 0..epochs {                                        // 3. K epochs of minibatches
            for i in (1..horizon).rev() { idx.swap(i, (rng.f() * (i + 1) as f32) as usize % (i + 1)); }
            for chunk in idx.chunks(mb) {
                let mean = chunk.iter().map(|&i| adv[i]).sum::<f32>() / chunk.len() as f32;
                let std = (chunk.iter().map(|&i| (adv[i] - mean).powi(2)).sum::<f32>() / chunk.len() as f32).sqrt() + 1e-8;
                let (mut gl, mut gv) = (vec![0.0f32; N * A], vec![0.0f32; N]);
                for &i in chunk {
                    let (st, a) = (ss[i], aa[i]);
                    let p = softmax(&logits[st * A..st * A + A]);
                    let ratio = (p[a].ln() - lp_old[i]).exp();
                    let an = (adv[i] - mean) / std;                 // per-minibatch normalization
                    let clipped = (an > 0.0 && ratio > 1.0 + eps) || (an < 0.0 && ratio < 1.0 - eps);
                    if clipped { clipfrac += 1.0; }
                    let ent = -(p[0] * p[0].ln() + p[1] * p[1].ln());
                    for j in 0..A {
                        let onehot = if j == a { 1.0 } else { 0.0 };
                        // d(ratio·Â)/dlogit_j = ratio·Â·(1[j=a] − p_j); zero where the clip is active
                        let g_pg = if clipped { 0.0 } else { ratio * an * (onehot - p[j]) };
                        let g_ent = -p[j] * (p[j].ln() + ent);      // dH/dlogit_j
                        gl[st * A + j] += (g_pg + ent_c * g_ent) / chunk.len() as f32; // ascent direction
                    }
                    gv[st] += (value[st] - ret[i]) / chunk.len() as f32; // d/dV of ½(V − R)²
                }
                for k in 0..N * A { logits[k] += lr * gl[k]; }
                for k in 0..N { value[k] -= lr * gv[k]; }
            }
        }
        if iter % 10 == 9 {
            let p_right: Vec<String> = (0..N - 1).map(|s| format!("{:.2}", softmax(&logits[s * A..s * A + A])[1])).collect();
            println!("iter {iter}: P(right) per state {:?}, clip fraction {:.2}", p_right,
                     clipfrac / (epochs * horizon) as f32);
        }
    }
}

Complexity and constants

• Each sample is used K times (3-10). Wall-clock is usually dominated by environment stepping, so vectorize environments (N = 8-4096 in GPU simulators).
• Memory: the rollout buffer, T·N·(obs + 5 scalars). For LLM RLHF, the policy, a frozen reference, a reward model and a value model all live in memory: about 4 model copies.
• Healthy runs show clip fraction of order 0.1-0.3 and approximate KL of order 0.01 per update. Values far above mean the step is too large.

Efficiency

• Cost: each sample is used K = 3-10 times, so training is K backward passes per rollout. Wall-clock is usually environment stepping, not the network. Memory is the rollout buffer, T·N·(obs + 5 scalars); LLM RLHF holds about 4 model copies (policy, reference, reward, value).
• Levers:
  • Vectorized environments, N = 8-4096 in GPU simulators — amortizes policy forward passes and stepping (above) — simulators on GPU.
  • Compute GAE in one backward sweep over the buffer — O(T·N) — always.
  • Early stop an update when approximate KL exceeds the target — saves wasted epochs and divergence (above: healthy KL of order 0.01) — unstable tasks.
  • Share one trunk between actor and critic — one forward serves both heads — pixel or large-observation inputs.
  • Group-relative baseline (GRPO) — removes the value model, one of about 4 copies (Shao et al. 2024) — LLM RL.
• Naive — REINFORCE, one gradient step per rollout, then discard · 1 use per sample, high variance
• Standard — PPO with GAE, K epochs of minibatches · K uses per sample
• State of the art — GPU-resident simulation with thousands of environments; critic-free GRPO for LLMs · stepping and generation dominate

Implementation notes

• Separate policy and value heads (or networks) for continuous control. Orthogonal init with a small gain (0.01) on the policy output layer, 1 on the value head (37-details write-up).
• Continuous actions: diagonal Gaussian with a state-independent learned log-std. Compute log-probs on the unsquashed sample, or include the tanh Jacobian correction if you squash.
• Observation normalization (running mean/var) and reward scaling by a running std of returns matter more than most hyperparameters on MuJoCo.
• Anneal lr linearly to 0. Clip gradient global norm at 0.5.
• Truncated episodes: bootstrap with V(s_T) on time limits, zero only on true termination.
• RLHF: per-token reward is −β·KL(π‖π_ref) plus the reward-model score at the last token. GRPO drops the value network and uses group-normalized rewards of several samples per prompt as the advantage.
• In PyTorch, compute ratio = (logp − logp_old).exp() from stored log-probs, never from stored probabilities, and keep logp_old detached.

Pitfalls

• Sign and normalization bugs: minimizing the surrogate instead of maximizing it, or normalizing the value targets instead of the advantages.
• Forgetting (1 − done) in GAE across episode boundaries inside a rollout.
• Too many epochs or too large a minibatch lr: the clip only blocks per-sample gradients; the policy can still drift through samples that are not clipped. Watch approximate KL and early-stop.
• Recomputing log π_old after the first update instead of storing it at rollout time.
• Entropy bonus too high: the policy stays uniform. Too low: premature collapse in discrete action spaces.
• Value loss on a different scale than the policy loss with a shared network: one head dominates the shared features.

Variants and alternatives

In production

OpenAI Five (Dota 2), InstructGPT/ChatGPT RLHF, OpenAI's dexterous in-hand robot hand, DeepMind robotics, CleanRL ppo.py, Stable-Baselines3 PPO, RLlib, Hugging Face TRL (PPOTrainer, GRPOTrainer), Isaac Gym/Lab (thousands of GPU environments).

Implementations

Sources

• OpenAI et al., Dota 2 with Large Scale Deep Reinforcement Learning, 2019. paper
• OpenAI et al., Learning Dexterous In-Hand Manipulation, 2018. paper
• Schulman, Wolski, Dhariwal, Radford, Klimov, Proximal Policy Optimization Algorithms, 2017. paper
• Schulman, Moritz, Levine, Jordan, Abbeel, High-Dimensional Continuous Control Using Generalized Advantage Estimation, ICLR 2016. paper
• Schulman, Levine, Moritz, Jordan, Abbeel, Trust Region Policy Optimization, ICML 2015. paper
• Williams, Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning, Machine Learning 1992. paper
• Huang et al., The 37 Implementation Details of Proximal Policy Optimization, ICLR Blog Track 2022. paper
• Engstrom et al., Implementation Matters in Deep Policy Gradients: A Case Study on PPO and TRPO, ICLR 2020. paper
• Ouyang et al., Training Language Models to Follow Instructions with Human Feedback, NeurIPS 2022. paper
• Shao et al., DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models, 2024 (GRPO). paper