Implementing PPO from scratch in PyTorch

Proximal Policy Optimization is the default on-policy algorithm for a reason: it is simple to implement, stable to train, and forgiving of hyperparameters. But most people meet it through a framework that hides the parts that actually matter. When I built it from scratch for my MSc work, the value was in writing every component by hand. This post walks through each one and explains why it is there.

The shape of the problem

PPO is a policy-gradient method. We have a policy network that maps a state to a distribution over actions, and a value network that estimates the expected return from a state. We collect experience by running the current policy in the environment, then use that experience to take a few careful improvement steps before throwing it away and collecting more. The whole game is taking the largest improvement step we safely can without the new policy drifting so far from the old one that the data we collected becomes meaningless.

Generalised Advantage Estimation

Before we can improve the policy we need a signal for how much better or worse an action was than the value network expected. That signal is the advantage. The naive estimate, the full Monte Carlo return minus the value baseline, is unbiased but very noisy. A one-step temporal-difference estimate is low variance but biased. GAE gives us a dial between the two.

For each timestep we first compute the TD residual, often written delta:

delta_t = r_t + gamma * V(s_{t+1}) * (1 - done_t) - V(s_t)

The GAE advantage is then an exponentially weighted sum of these residuals, with the lambda parameter controlling the decay. In practice you compute it with a single backward pass over the trajectory, which is far cleaner than the summation suggests:

gae = 0
for t in reversed(range(T)):
    next_value = values[t + 1] if t + 1 < T else last_value
    nonterminal = 1.0 - dones[t]
    delta = rewards[t] + gamma * next_value * nonterminal - values[t]
    gae = delta + gamma * lam * nonterminal * gae
    advantages[t] = gae
returns = advantages + values

Two details people get wrong here. First, the nonterminal mask must zero out the bootstrap across episode boundaries, otherwise reward leaks backwards across resets. Second, the value targets for the critic are advantages + values, computed before any normalisation. After computing advantages it is standard to normalise them across the batch to zero mean and unit variance, which stabilises the gradient scale.

The clipped surrogate objective

This is the heart of PPO. We define the probability ratio between the new and old policies for the action that was actually taken:

ratio = exp(new_log_prob - old_log_prob)

The plain policy-gradient objective is ratio * advantage. The problem is that nothing stops a single update from pushing the ratio far from one, which can collapse the policy. PPO clips the ratio into a small interval around one, typically [1 - eps, 1 + eps] with eps around 0.2, and takes the pessimistic minimum of the clipped and unclipped terms:

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

The minimum is what makes this work. When the advantage is positive we want to increase the action's probability, but the clip caps how much credit a single update can take. When the advantage is negative the clip caps how far we push the probability down. Taking the minimum means the objective only ever gives a conservative estimate of the improvement, so the policy cannot exploit a large ratio to make a reckless step. We negate because optimisers minimise.

The value loss

The critic is trained by regression onto the returns we computed earlier. A plain mean-squared error works:

value_loss = 0.5 * (new_values - returns).pow(2).mean()

Many implementations also clip the value update, mirroring the policy clip, so the new value estimate cannot move too far from the old one in a single batch. It is optional and its benefit is environment-dependent, but it is worth knowing it exists.

The entropy bonus

Left alone, the policy will happily become deterministic too early and stop exploring. To counter this we add the entropy of the policy distribution as a small bonus, which rewards keeping some uncertainty in the action choice:

entropy = dist.entropy().mean()
loss = policy_loss + vf_coef * value_loss - ent_coef * entropy

Note the sign: we subtract the entropy term from the loss because higher entropy is good and we are minimising. The coefficient is small, often around 0.01, and it decays in importance naturally as the advantages start to dominate.

The training loop

With the pieces in place the loop is straightforward, and its structure is what makes PPO on-policy yet sample-efficient:

1. Run the current policy for a fixed number of steps, storing states, actions, log-probabilities, rewards, dones and value estimates.
2. Compute advantages and returns with GAE, then normalise the advantages.
3. For several epochs, shuffle the batch and iterate over minibatches. For each minibatch, recompute log-probabilities and values under the current network, form the three loss terms, and take an optimiser step.
4. Discard the data and go back to step one.

What I took away

Building PPO by hand made clear how few of its parts are exotic. The clip is two lines. GAE is one backward pass. The losses are a sum of three familiar terms. What is easy to get wrong is the plumbing around them: the terminal masks in GAE, the order of advantage normalisation, the sign on the entropy term, and reusing the old log-probabilities rather than recomputing them. Get those right and PPO is remarkably robust. The full implementation is on my GitHub as ppo-from-scratch, with a multi-GPU version using DistributedDataParallel alongside it.