Deriving PPO from First Principles

PPO is the most widely used policy gradient algorithm in deep reinforcement learning. It trained OpenAI Five to beat the Dota 2 world champions in 2019, the robot hand that solved a Rubik’s cube that same year, and the RLHF step behind InstructGPT and ChatGPT in 2022.

Most explanations show the PPO objective in full and then work backwards. Here we build it up one term at a time, so that nothing appears without a reason.

PPO is really just a few additions to REINFORCE, an algorithm from 1992. REINFORCE is about as simple as a learning algorithm gets: a dozen lines of code.

The whole idea is one sentence: run your policy, see which episodes went well, and make those actions more likely.

1 · On-policy RL, in its simplest form

An agent acts, collects rewards, and the episode ends. We want whatever behavior collects the most reward.

All of that behavior lives in the policy $\pi_\theta$, one neural network with weights $\theta$. Handed a state it returns a probability distribution over actions, $\pi_\theta(a \mid s)$.

One pass through the environment leaves a trail of states, actions and rewards. Call it a trajectory:

How good was the episode? Add up the rewards, optionally shrinking later ones by a factor $\gamma$ per step so that reward now counts for more than reward later. That sum is the return $G$:

# rewards : torch.Tensor (T,)    reward at each step of one episode
# gamma   : float                discount factor
# ->        torch.Tensor scalar  the return of the whole episode
def trajectory_return(rewards, gamma):
  ...

The objective is the expected value of the return, for a trajectory sampled from the policy. Call it $J$:

We want to find the policy that maximizes $J$.

$\pi_\theta$ is a neural network, so maximizing $J$ means gradient ascent, which requires $\nabla_\theta J$. Luckily, this identity lets us compute it:

A gradient of an expectation becomes an expectation of a gradient, and those you can estimate by sampling. This is the score function estimator.

Apply it to $J$, with $f = G$:^*

A working algorithm, and a naive one.

Reading it carefully says why: every action in the episode is multiplied by the same number $G(\tau)$. An action taken at step $t = 200$ is credited with rewards collected at step $t = 3$.

2 · REINFORCE

An action cannot influence rewards that arrived before it. That is the principle of causality. Those earlier terms are not signal; they are noise, added to every update that action will ever receive. So we credit each action with only what came after it. That is the REINFORCE reward:

Swapping it in costs nothing: the discarded terms had expectation zero, so the estimator is still exactly unbiased, and strictly quieter. This is REINFORCE in the form people usually mean:

Notice what $G_t$ actually is. It is the observed return, a sum of rewards that really happened, not an estimate of anything. REINFORCE has no value function, no model of the environment, nothing to learn but the policy itself.

# rewards : torch.Tensor (T,)    reward at each step of one episode
# gamma   : float                discount factor
# ->        torch.Tensor (T,)    G[t] = return from step t onwards
def reward_to_go(rewards, gamma):
  G = torch.zeros_like(rewards)
  acc = 0.0
  for t in reversed(range(len(rewards))):
      acc = ...
      G[t] = acc
  return G

Whichever action you sample goes up

One state, three actions. We sample an action $a_t \sim \pi_\theta$ and collect the reward it pays. The episode is one step long, so $G_t = r_t$.

Set what each action returns and choose which one came out of the sample. The bars show where the probabilities end up.

Every action returns something positive, so only the one that gets sampled is reinforced, even the worst of the three.

3 · REINFORCE minus a baseline

If every $G_t$ is positive then every action is pushed up, and only the differences between them carry any information.

The fix is to subtract a reference point. Call it the baseline, $b(s_t)$:

The surprising part is that you may choose it freely: any such $b$ leaves the gradient’s expectation exactly where it was.

The reason is worth seeing without the algebra too. Probabilities in a state sum to one, so pushing all of them up by the same amount pushes none of them up relative to each other.

The best baseline is the one that predicts the return you were about to get: $b(s_t) = V^\pi(s_t)$, the value of being in state $s_t$.

The same three actions. Let’s replace $G_t$ with $G_t - V^\pi(s_t)$.

The return is measured against that average now, so a sampled action can be pushed down, not just up.

4 · One equation to rule them all

Look back at what actually changed between the last three formulas. The $\nabla \log \pi$ factor never moved. Only the number multiplying it did.

So give that coefficient a name and write the family once:

Every algorithm from here (A2C, TRPO, PPO) is a choice of $\Psi_t$.^*

5 · A2C

Three additions: bootstrapping, the critic (a second neural network, estimating $V(s)$), and GAE.

Bootstrapping

Using $G_t$ means waiting for the episode to end. But you do not need the whole episode to judge an action. The value of a state is the reward you collect plus the value of wherever you land. That is the Bellman equation:

So stop after one step: take the reward you got and let the value of the next state stand in for the rest.

We do not have $V^\pi$. In its place put a second network, $V_\phi$, with its own parameters $\phi$, trained to predict what a state is worth.

The gap between what $V_\phi$ predicted and what one step of reality says is the TD error, $\delta_t$.

We are improving an estimate using another estimate. That is bootstrapping, and it is the one idea separating A2C from everything before it.

# rewards    : torch.Tensor (T,)  reward collected at each step of the rollout
# values     : torch.Tensor (T,)  the critic's V(s_t) at each step
# last_value : torch.Tensor ()    the critic's V(s_T), where the rollout stopped
# gamma      : float
# ->           torch.Tensor (T,)  one TD error per step
def td_error(rewards, values, last_value, gamma):
  ...

The critic

There are two networks now, doing different jobs. $\pi_\theta$ picks actions: it is the actor, and the only part that gets used once training is over. $V_\phi$ never chooses anything. It predicts, and its prediction is what the actor’s choices get measured against. It is the critic.

And the quantity being measured has a name. How much better an action turned out than what the policy would have done on average is the advantage:

It is zero-mean under the policy, $\mathbb{E}_{a\sim\pi}[A^\pi(s,a)] = 0$, so better-than-average actions get positive weight and worse-than-average get negative. Section 3’s baseline was already estimating it with a whole episode; $\delta_t$ estimates it with one step.

Three words, three pieces. A2C is advantage actor-critic.

The gain is less variance: one transition’s worth of noise instead of an entire episode’s, and no waiting for the episode to end, since a single step is now enough to learn from.

The cost is correctness. $\delta_t$ is an advantage estimate only when $V_\phi = V^\pi$, and $V_\phi$ is never quite right, so where every estimator so far was unbiased and merely noisy, this one is biased.

GAE

We now have two estimates of the same advantage, and they disagree about how far to look before trusting the critic.

• REINFORCE (Monte Carlo) waits for the entire episode: no bias, a great deal of noise.
• The TD error looks one step ahead: very little noise, some bias.

Nothing forces that choice. You could take two real rewards before falling back on $V_\phi$, or three, or ten:

So rather than picking one, take a weighted average of all of them, with the weight decaying by a factor $\lambda$ each step further out. That average collapses into a single sum over the TD errors we already have:

This is the generalized advantage estimator, and $\lambda$ is the dial.

At $\lambda = 1$ nothing decays: Monte Carlo again.

At $\lambda = 0$ only the first term survives: the single TD error. The two algorithms we have built so far are the same one, at opposite ends of this dial.

So which setting? In practice almost everyone picks $\lambda \approx 0.95$ — close to the Monte Carlo end, but not quite at it. The plot below shows why.

# rewards    : torch.Tensor (T,)     reward at each step
# values     : torch.Tensor (T,)     V(s_t), the critic at each step
# last_value : torch.Tensor scalar   V(s_T), to bootstrap the final step
# gamma, lam : float
# ->           torch.Tensor (T,)     the GAE advantage at each step
def gae(rewards, values, last_value, gamma, lam):
  adv = torch.zeros_like(rewards)
  acc = 0.0
  for t in reversed(range(len(rewards))):
      v_next = last_value if t == len(rewards) - 1 else values[t + 1]
      delta = rewards[t] + gamma * v_next - values[t]
      acc = ...
      adv[t] = acc
  return adv

6 · PPO

A2C uses each transition once and throws it away. After a single gradient step $\pi_\theta \ne \pi_{\text{old}}$, so the batch was produced by a policy we no longer have, and the on-policy expectation from section 1, which assumed the data came from the policy being updated, no longer applies.

Weight each sample by how much more likely the new policy is to have taken it:

and the objective is usable again:

Reweighting samples drawn from one distribution in order to estimate an expectation under another is called importance sampling; $\rho_t$ is the importance weight. Here it is what makes a second pass over the same batch legitimate.

A quick check that nothing was smuggled in: at the very start of the update $\pi_\theta = \pi_{\text{old}}$, so $\rho \equiv 1$ and the gradient is exactly A2C’s.

# logp_new : torch.Tensor (N,)  log pi_theta(a|s), the policy being updated
# logp_old : torch.Tensor (N,)  log pi_old(a|s), the one that collected the batch
# ->         torch.Tensor (N,)  one ratio per sample
def importance_ratio(logp_new, logp_old):
  ...

There is a catch, though. Nothing bounds $\rho_t$. With $\hat A_t > 0$ the objective grows without limit as $\rho_t \to \infty$, so gradient ascent will happily drive one sampled action’s probability to 1 on the strength of a single batch.

Say $\pi_{\text{old}}$ took an action with probability $0.1$ and $\pi_\theta$ now takes it with probability $0.9$. Then $\rho_t = 9$, and that single sample pulls nine times as hard as it did when it was collected.

We need to constrain how far the policy may move in one update. The principled way is to bound the KL divergence between $\pi_{\text{old}}$ and $\pi_\theta$ — a measure of how far apart two distributions are. An improvement guarantee can be proved mathematically when using KL divergence.

TRPO, for trust region policy optimization, is the 2015 algorithm PPO grew out of, and it uses the KL directly: maximize the objective subject to $\mathrm{KL}(\pi_{\text{old}} \,\|\, \pi_\theta) \le \delta$.

That works, but it is expensive.

PPO approximates that constraint inside the objective itself, by clipping the ratio so that moving too far stops paying:

If that looks intimidating, there are only two cases — one per sign of $\hat A_t$:

# ratio     : torch.Tensor (N,)      pi_theta / pi_old, one per sample
# adv       : torch.Tensor (N,)      advantage estimate, one per sample
# clip_coef : float                  epsilon, e.g. 0.2
# ->          torch.Tensor scalar    the LOSS (a negated objective)
def clipped_surrogate(ratio, adv, clip_coef):
  ...

Clipping allows PPO to reuse data

What PPO adds is data reuse: several gradient steps on each batch instead of one. The ratio $\rho_t$ is what makes that legitimate, since a batch collected under $\pi_{\text{old}}$ still counts for the policy we have now. The clip is what keeps those extra steps safe. Neither data reuse nor clipping does anything alone. Let’s run both on the same environment.

Cliff Walking is a small gridworld: $-1$ per step, $-100$ for stepping off the cliff, and an optimal route worth $-13$.

Run at $\gamma = 1$ for 400k environment steps, it defeats A2C completely. Every row below sees that same amount of experience; only the update changes:

Sixteen times the learning per unit of experience is what lets the policy find the goal, and the clip is what stops sixteen unconstrained steps on one batch from blowing up.

Why $\gamma = 1$ is what breaks A2C here is worth its own look. Bootstrapping needs the Bellman backup to be a contraction, and $\gamma$ is the contraction modulus. At $\gamma = 1$ there is no fixed point to fall into at all:

The whole loss

Everything so far has been the policy term. A real implementation adds two more:

The critic term trains $V_\phi$: plain regression onto the returns the rollout actually produced, $L^{V} = \mathbb{E}\big[(V_\phi(s_t) - R_t)^2\big]$.

The entropy term pays the policy to stay spread out. $\bar H$ is largest when it still gives every action a chance and zero once it always picks the same one, so rewarding it delays the moment exploration stops.

# pg_loss  : torch.Tensor scalar   clipped surrogate, already negated
# v_loss   : torch.Tensor scalar   mean squared error of the critic
# entropy  : torch.Tensor (N,)     H(pi) at each state in the batch
# vf_coef  : float                 weight on the critic
# ent_coef : float                 weight on exploration
# ->         torch.Tensor scalar   the one number you call .backward() on
def ppo_loss(pg_loss, v_loss, entropy, vf_coef, ent_coef):
  policy_term  = pg_loss
  value_term   = vf_coef * v_loss
  entropy_term = ent_coef * entropy.mean()
  return ...

And that is PPO — the whole objective, nothing held back. You derived it: the return, the reward-to-go, the baseline, the advantage, bootstrapping, GAE, the importance ratio, the clip. Every term above is something you built a reason for.

Every step, in one table

The first two rows are free: variance down, bias unchanged. Only the move to A2C costs something, and it is the same thing every time: a target that refers to its own estimate. That is what buys the sample efficiency, and it is what breaks when $\gamma = 1$.

Where to go next

The same path, further on. TRPO is the step PPO approximates: a hard KL constraint with a monotonic-improvement guarantee inherited from Conservative Policy Iteration. GAE is the $\lambda$ dial in its own paper. GRPO drops the critic and returns to a Monte Carlo estimator with an empirical group-mean baseline — section 3, rediscovered for language models.

Exploration is its own field. RND and ICM supply the directed exploration sampling cannot; Go-Explore attacks it by remembering promising states; the noisy-TV problem is where curiosity itself breaks.

Off-policy avoids throwing data away at all. SAC for continuous control, and the deadly triad for why that is harder than it sounds.

Canonical sources. Sutton & Barto, Reinforcement Learning: An Introduction. Ch. 13 for policy gradients. The 37 Implementation Details of PPO for everything the equations leave out, and CleanRL for single-file implementations worth reading start to finish.