10.Policy Gradients: REINFORCE → PPO

10.1 Why robotics chose policy search

Three reasons, all of which matter more on hardware than on benchmarks.

Continuous actions are natural. A policy that outputs the mean and standard deviation of a Gaussian over joint torques needs no argmax. Chapter 9's wall simply is not there.

Updates are smooth. A small parameter change produces a small change in behaviour. Value-based methods can flip the argmax at many states from one gradient step to the next, and a robot whose policy changes discontinuously mid-training is a robot that breaks something. This alone is worth a great deal of sample efficiency.

Prior structure fits naturally. You can constrain the policy class to whatever you know: a movement primitive with learnable weights (Chapter 17), a residual on top of a hand-tuned controller, an action space that respects joint limits by construction. Kober and colleagues make this argument at length, and it remains correct.

The cost is that policy gradient methods are usually on-policy — data must come from the current policy, so it is discarded after each update. That is expensive. It is also, per Chapter 8, exactly why they are stable: dropping the off-policy ingredient breaks the deadly triad.

10.2 The policy gradient theorem

Parameterize the policy as $\pi_\theta(a\mid s)$ and define the objective as expected return from the start-state distribution:

$J(\theta) = \mathbb{E}_{\pi_\theta}\left[\sum_{t=0}^{\infty} \gamma^t R_{t+1}\right].$

We want $\nabla_\theta J$. The difficulty is immediate: $J$ averages over trajectories whose distribution depends on $\theta$. Differentiating an expectation whose measure moves is not a routine operation.

The score-function identity resolves it. For any distribution $p_\theta$,

$\nabla_\theta p_\theta(x) = p_\theta(x)\, \nabla_\theta \log p_\theta(x),$

which is just the chain rule on $\log$, rearranged. It converts a gradient of a probability into an expectation of a gradient — and expectations we can sample.

REINFORCE turns the theorem into an algorithm: sample an episode, and for each step update

$\theta \leftarrow \theta + \alpha\, \gamma^t\, G_t\, \nabla_\theta \log \pi_\theta(A_t \mid S_t).$

It is unbiased and it is nearly unusable, because the variance of $G_t$ is enormous — the return accumulates every random event in the trajectory, and multiplying that by the score function does not help.

10.3 Baselines and advantage

Chapter 3 already gave the fix. Subtract any function $b(s)$ that does not depend on the action:

$\nabla_\theta J = \mathbb{E}\left[\big(q_\pi(S_t, A_t) - b(S_t)\big) \nabla_\theta \log \pi_\theta(A_t\mid S_t)\right].$

The best practical baseline is $v_\pi(s)$, which makes the multiplier the advantage:

$A_\pi(s,a) = q_\pi(s,a) - v_\pi(s).$

The interpretation is exactly right. Advantage asks "was this action better than what I typically do here?" rather than "was this outcome good?" — and only the first question carries information about which action to prefer. A state where everything goes well produces large returns for every action, and without a baseline all of them get reinforced.

Turn up the reward offset. REINFORCE's histogram spreads badly even though a constant added to every return says nothing about which action was good — the score function multiplies it regardless. The baseline subtracts it straight back out. All three estimators have the same mean throughout; only the spread differs, and spread is what you pay for in samples.

10.4 Actor–critic and GAE

We do not know $v_\pi$, so learn it: a critic $V_\phi(s)$ trained by the TD methods of Chapter 6, and an actor $\pi_\theta$ updated by the policy gradient using the critic's advantage estimate. This is generalized policy iteration again — Chapter 5's pattern, with a parameterized policy and a learned value.

How should the advantage be estimated? The one-step form $\delta_t = R_{t+1} + \gamma V(S_{t+1}) - V(S_t)$ has low variance and inherits the critic's bias. The Monte Carlo form $G_t - V(S_t)$ is unbiased and noisy. This is Chapter 7's dilemma verbatim, so it takes Chapter 7's answer.

Generalized advantage estimation is the $\lambda$-return applied to advantages:

$\hat A_t^{\text{GAE}(\gamma,\lambda)} = \sum_{l=0}^{\infty} (\gamma\lambda)^l\, \delta_{t+l}.$

At $\lambda = 0$ it is the one-step TD error; at $\lambda = 1$ it is the Monte Carlo advantage. In practice $\lambda \approx 0.95$ works across an unusually wide range of tasks, which is why it is one of the few hyperparameters practitioners rarely touch.

10.5 Trust regions and PPO

One problem remains, and it is the one that decides whether training survives.

Policy gradient tells you a direction, not a step size. Take too large a step and the policy changes drastically; the data you collected is now from a policy that no longer exists, the next batch is collected under a worse policy, and performance collapses without recovering. Unlike supervised learning, a bad update poisons all future data.

The performance-difference lemma quantifies the risk: the improvement from $\pi_{\text{old}}$ to $\pi$ is exactly

$J(\pi) - J(\pi_{\text{old}}) = \mathbb{E}_{s \sim d_\pi, a\sim \pi}\left[A_{\pi_{\text{old}}}(s,a)\right],$

where the states are drawn from the new policy's distribution — which we cannot sample without running it. Approximating $d_\pi$ by $d_{\pi_{\text{old}}}$ gives a surrogate objective that is accurate only while the policies remain close. TRPO enforces closeness with a hard KL constraint and solves the resulting problem with conjugate gradients — principled, and heavy.

PPO achieves nearly the same effect with a clipped objective. With the probability ratio $r_t(\theta) = \frac{\pi_\theta(A_t\mid S_t)}{\pi_{\theta_{\text{old}}}(A_t\mid S_t)}$:

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

use burn::prelude::*;
 
pub struct PpoConfig {
    pub clip_eps: f32,
    pub entropy_coef: f32,
    pub value_coef: f32,
    pub target_kl: f32,
}
 
pub struct PpoDiagnostics {
    pub policy_loss: f32,
    pub value_loss: f32,
    pub entropy: f32,
    pub approx_kl: f32,
    pub clip_fraction: f32,
}
 
impl<B: AutodiffBackend> Ppo<B> {
    pub fn surrogate(&self, batch: &RolloutBatch<B>, cfg: &PpoConfig)
        -> (Tensor<B, 1>, PpoDiagnostics)
    {
        let (mean, log_std) = self.actor.forward(batch.states.clone());
        let log_probs = gaussian_log_prob(&mean, &log_std, &batch.actions);
 
        // r_t(θ) = π_θ(a|s) / π_old(a|s), computed in log space for stability.
        let ratio = (log_probs.clone() - batch.old_log_probs.clone()).exp();
 
        // Advantage normalization — not optional in practice.
        let adv = normalize(batch.advantages.clone());
 
        let unclipped = ratio.clone() * adv.clone();
        let clipped = ratio.clone().clamp(1.0 - cfg.clip_eps, 1.0 + cfg.clip_eps) * adv;
 
        // min() takes the pessimistic branch: the update never profits from
        // stepping outside the trust region in either direction.
        let policy_loss = -unclipped.min_pair(clipped).mean();
 
        let values = self.critic.forward(batch.states.clone()).squeeze(1);
        let value_loss = (values - batch.returns.clone()).powf_scalar(2.0).mean();
 
        let entropy = gaussian_entropy(&log_std).mean();
 
        let loss = policy_loss.clone()
            + value_loss.clone() * cfg.value_coef
            - entropy.clone() * cfg.entropy_coef;
 
        // Schulman's low-variance KL estimator; early-stop the epoch if it spikes.
        let log_ratio = log_probs - batch.old_log_probs.clone();
        let approx_kl = ((log_ratio.clone().exp() - 1.0) - log_ratio).mean();
 
        let clip_fraction = ratio
            .sub_scalar(1.0)
            .abs()
            .greater_elem(cfg.clip_eps)
            .float()
            .mean();
 
        (loss, PpoDiagnostics {
            policy_loss: policy_loss.into_scalar().elem(),
            value_loss: value_loss.into_scalar().elem(),
            entropy: entropy.into_scalar().elem(),
            approx_kl: approx_kl.into_scalar().elem(),
            clip_fraction: clip_fraction.into_scalar().elem(),
        })
    }
}

10.6 Chapter bridge

We changed what is learned. Instead of a value function with a policy extracted from it, the policy is the parameter vector, and the policy gradient theorem makes its gradient estimable from experience alone — with no model, and with the state-distribution shift already absorbed. Baselines and GAE tamed the variance; trust regions tamed the step size; PPO made the whole thing simple enough to run with Adam.

What we did not fix is sample efficiency. PPO is on-policy: every batch is used for a few epochs and thrown away. For a simulator with thousands of parallel environments that is acceptable, and Chapter 18 shows it is how locomotion is actually trained. For a robot learning on hardware, discarding data is a luxury nobody can afford.

Chapter 11 returns to off-policy learning — this time with continuous actions and a replay buffer, accepting the deadly triad's risks in exchange for reusing every transition many times. Reacher gets his proper debut as a continuous-control task, and maximum-entropy RL turns exploration from a schedule you tune into a quantity the objective optimizes.

References