Proximal Policy Optimization from Absolute Zero

What You'll Master

The Setup & Big Picture

Same task as Problem 3, completely different algorithm. A 4-DOF Sawyer arm in MuJoCo simulation must learn to pick up a hammer and drive a nail into a board. The observation is ~39 numbers (positions of robot, hammer, nail). Action is a 4-D continuous vector in [-1, 1]. Reward: 0 every step, 1.0 the moment the nail is fully driven, then episode ends.

You'll implement PPO — Proximal Policy Optimization — the most widely-used policy-gradient algorithm in modern RL. The deliverable is a wandb plot of eval/episode_success over ~1 million environment steps.

Why PPO and not vanilla policy gradient

The simplest policy-gradient algorithm (REINFORCE) goes:

1. Sample a trajectory under π_θ.
2. For each (s_t, a_t) in the trajectory, increase log π_θ(a_t|s_t) in proportion to the return G_t.
3. Repeat.

This works in theory but is wildly unstable in practice. Two reasons:

1. High variance in the return G_t. Reduces convergence to a crawl.
2. Single big update can break the policy. The gradient says "move this way." But if you move too far, the next batch of rollouts comes from a totally different policy, and the gradient you just applied was based on the old policy. Now you're flying blind.

PPO addresses both problems. GAE reduces variance via a critic-blended advantage estimate. Clipping caps how far each update can push the policy. Together: stable, fast learning.

"On-policy" with caveats

Strictly on-policy means: collect data, do one gradient step, throw the data away. PPO does multiple epochs of gradient steps on the same batch. As long as the policy hasn't drifted too much (the clip enforces this), the same data is approximately valid for several updates. This is sometimes called "near-on-policy" or "approximately on-policy."

PyTorch Primer

If you know Python but PyTorch is fuzzy, read this chapter. You'll need every concept here when reading the code.

1. Tensors are NumPy arrays with two superpowers

A tensor is just a multi-dimensional array of numbers, like NumPy's ndarray. The two extras are:

• GPU residency. A tensor lives on either CPU or GPU memory. Calling .to(device) moves it. Math on tensors that live on a GPU is fast (parallelized across thousands of CUDA cores).
• Autograd. Each tensor remembers how it was computed. PyTorch builds a graph of operations behind the scenes so it can compute gradients automatically.

import torch

x = torch.tensor([1.0, 2.0, 3.0])    # shape [3], on CPU
x = x.to("cuda")                       # now on GPU
y = x * 2.0 + 1.0                       # [3, 5, 7], gradient graph remembered

2. Shapes are what bugs hide in

Most PyTorch bugs are shape bugs. The shape of a tensor is its dimensions: a vector of length 4 is shape [4]; a 32×4 matrix is shape [32, 4]. Operations broadcast across compatible shapes silently, which is convenient but dangerous.

a = torch.tensor([1., 2., 3.])             # shape [3]
b = torch.tensor([[10.], [20.], [30.]])     # shape [3, 1]
c = a + b   # shape [3, 3]! NumPy broadcasts. Probably not what you wanted.

3. Common shape operations

4. Autograd in 60 seconds

When you compute a tensor that has trainable parameters in its history, PyTorch tracks which operations made it. Calling .backward() on a scalar (a single number) tells PyTorch to walk backward through that history, computing gradients for every parameter along the way.

# Forward pass
w = torch.tensor([2.0], requires_grad=True)
x = torch.tensor([3.0])
loss = (w * x − 5.0) ** 2           # scalar, shape []

# Backward pass
loss.backward()                       # fills in w.grad
print(w.grad)                         # tensor([6.])

# Manual update (in practice the optimizer does this)
with torch.no_grad():
    w −= 0.01 * w.grad
    w.grad.zero_()                    # clear before next backward

Three rules to keep in your head:

1. You can only call .backward() on a scalar. If your loss is a vector, take .mean() or .sum() first.
2. Gradients accumulate. Always zero them before the next backward pass — optimizer.zero_grad() handles this.
3. Operations inside with torch.no_grad(): don't get tracked. Use this when computing values you don't want gradients for (like TD targets or log-probs at rollout time).

5. nn.Module: the building block

A network is a Python class that inherits from nn.Module. It has learnable parameters (typically the weights of nn.Linear layers it stores) and a forward() method that defines what happens when you call it.

class SimpleNet(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.layer1 = nn.Linear(in_dim, 64)    # weights + bias
        self.layer2 = nn.Linear(64, out_dim)

    def forward(self, x):
        x = torch.relu(self.layer1(x))
        return self.layer2(x)

net = SimpleNet(10, 4)
y = net(torch.randn(32, 10))         # calls .forward, returns shape [32, 4]

Notice net(x), not net.forward(x). The __call__ hook on Module triggers some PyTorch bookkeeping then calls forward.

nn.Sequential is a convenience for stacking layers without writing a forward method:

self.policy = nn.Sequential(
    nn.Linear(obs_dim, 256),
    nn.ReLU(),
    nn.Linear(256, 256),
    nn.ReLU(),
    nn.Linear(256, action_dim),
)

6. Optimizers

An optimizer takes a list of parameters and applies gradient updates to them. The most common: Adam.

optimizer = torch.optim.Adam(net.parameters(), lr=3e-4)

# Standard training loop:
optimizer.zero_grad()    # clear leftover grads
loss.backward()         # compute new grads
optimizer.step()         # take a gradient step

7. Distributions

For policy-gradient algorithms, you need to sample actions from a distribution and compute log-probabilities. PyTorch provides torch.distributions.

from torch.distributions import Normal

mu = torch.tensor([0.0])
sigma = torch.tensor([1.0])
dist = Normal(mu, sigma)

a = dist.sample()              # draws a random number, shape [1]
log_p = dist.log_prob(a)      # log-density at a, shape [1]

The homework uses utils.TruncatedNormal — same idea but the samples are clipped to [-1, 1] to match the action bounds. dist.log_prob(a) gives you log π(a|s), which is what the policy gradient theorem needs.

That's enough to read every line of on_policy.py. Onward.

The Policy Gradient Theorem

The whole reason PPO exists is to maximize one quantity: expected discounted return.

Read it: "the average over all possible trajectories τ that π_θ could produce, of the total discounted reward of that trajectory." We want to find the policy parameters θ that maximize J.

To do gradient ascent we need ∇θ J. The trajectory distribution depends on θ through the policy actions, so we can't trivially differentiate. The policy gradient theorem gives us a clever rewriting:

where Gt = Σk=0 γk rt+k is the return from time t onward.

In English: to make the policy better, take each action you just sampled and increase its log-probability by an amount proportional to how good the future turned out to be after that action. Good actions: push them up. Bad actions: push them down.

The naive estimator

Given a batch of N rollouts, we estimate the expectation by averaging:

If you implement just this, you have REINFORCE. It works, but it's slow. Two problems:

1. Variance: G_t can swing dramatically across rollouts. Same policy, totally different trajectories due to randomness in action sampling and environment.
2. No baseline: if rewards are always positive, every action's log-prob increases. The signal is "everything is good," which doesn't tell the policy what to focus on.

The fixes: baselines (advantage), GAE (multi-step variance reduction), clipping (stability). Each gets its own chapter.

Why We Need Advantage

Naive REINFORCE weights each action by the return G_t. Two issues:

Issue 1: Variance reduction via baselines

Subtract any state-dependent baseline b(st) from G_t. The expected gradient is unchanged (provable: cross-term has zero mean), but variance can drop dramatically.

The natural choice for b: the state-value function V(s) = 𝔼[G_t | s_t=s], i.e. how good is this state on average?

The advantage function

The cleaned-up policy gradient:

Now the gradient is "how much should we change θ per unit log-prob, weighted by how surprising-good this action was."

How do we estimate A(s, a)?

We don't know V exactly. We train a critic V_φ(s) in parallel with the policy, by regressing toward observed returns. With a critic, the simplest advantage estimator is the 1-step TD error:

This is low-variance (one step of randomness) but biased (uses imperfect V). The other extreme is the Monte Carlo advantage:

Unbiased but high variance (sums many random rewards). We want a middle ground that's tunable. That's exactly what GAE provides.

GAE: The Bias-Variance Trade

Generalized Advantage Estimation (Schulman et al., 2016) blends 1-step TD and Monte Carlo via a parameter λ ∈ [0, 1]. It's the algorithmic content of your compute_gae implementation.

The recursion

Computed backwards through a trajectory. (1 − dt) is the done-flag mask — same trick from Problem 1 to zero out bootstrap at terminal states.

The two ends of the spectrum

λ=0.95 is the standard choice and what this homework uses. It says "trust the actual rewards mostly, but blend in the critic's bootstrapped estimate to dampen variance."

Unrolled, GAE is a geometric sum

If you unroll the recursion, you get:

It's a geometrically-weighted sum of future TD errors. Each future TD error contributes less than the last by factor γλ. The recursion is just an efficient backwards pass to compute this sum.

The "returns" output

The function returns both advantages and returns:

Why? Recall A = Q − V, so returns = A + V is our estimate of Q (or G under the current policy, since the actor is ε-stochastic anyway). The critic should regress toward this — it's a low-variance estimate of "what does following the policy from here look like in terms of return."

Why PPO Clips

Vanilla policy gradient is unstable. One big update can move the policy so far that the rollouts you collected no longer represent it, and learning collapses.

We need a way to limit how much the policy can change per update. PPO does this with a clipped surrogate objective.

The importance ratio

Define the importance ratio between new and old policies for the same action:

If the policy hasn't changed, ρ = 1. If we doubled the probability of action a_t, ρ = 2. If we halved it, ρ = 0.5.

The clipped objective

Two terms inside the min:

1. ρ_t · Â_t — the unclipped surrogate. This is what vanilla policy gradient would use.
2. clip(ρ_t, 1−ε, 1+ε) · Â_t — the clipped surrogate. Caps ρ in [0.8, 1.2] for ε=0.2.

Then: take the elementwise min of these two, mean across the batch, negate (because PyTorch minimizes).

Why min — the asymmetric pessimistic cap

Walk through both signs of advantage:

The min always picks the worse outcome from the optimizer's perspective. It's a pessimistic bound. The optimizer can't game the surrogate by moving too far in either direction. That's what stabilizes PPO.

Value, Entropy, Reverse-KL

The full PPO loss has more than just the clipped surrogate. Three terms get added in on_policy.py:260-263:

loss = (policy_loss
        + self.value_coef    * value_loss
        − self.entropy_coef  * entropy
        + self.reverse_kl_coef * reverse_kl)

Value loss

Mean squared error between the critic and the GAE returns target. Trains V_φ alongside the policy. The value_coef hyperparameter weights how much we care about critic accuracy vs policy improvement (typically 0.5).

Entropy bonus

Entropy = how spread out is the action distribution. Without this, the policy can collapse to nearly deterministic during training, killing exploration. We add entropy to the objective (i.e. subtract from the loss, since loss = -objective). The entropy_coef is small (~0.001-0.01).

Reverse KL to a frozen reference

The reference_actor is a frozen snapshot of the policy taken right after BC pretraining. This term penalizes the RL policy for drifting too far from the BC initialization.

Why we need this: with sparse rewards, BC is the only thing telling the agent how to roughly do the task. Pure RL would forget BC during training. The KL anchor keeps it grounded. Same role as the BC-during-RL alternation in Problem 3 — same disease, different antibiotic.

The Full Algorithm

Note the discard step. PPO collects, updates a few times, throws data away. This is what makes it on-policy and what makes it sample-inefficient compared to off-policy methods.

Code Tour: on_policy.py

Three classes: Actor, Critic, PPOAgent. Same pattern as Problem 3 but the training loop differs.

The Actor

on_policy.py:13-39class Actor(nn.Module):
    def __init__(self, obs_shape, action_shape, hidden_dim, log_std_init=−2.0):
        super().__init__()
        self.policy = nn.Sequential(
            nn.Linear(obs_shape[0], hidden_dim), nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, hidden_dim),   nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, action_shape[0]))
        self.log_std = nn.Linear(hidden_dim, action_shape[0])

    def forward(self, obs):
        features = self.policy[:−1](obs)              # run all but final layer
        mu = torch.tanh(self.policy[−1](features))   # action mean in [-1, 1]
        std = torch.exp(self.log_std(features)).clamp(1e-5, 1.0)
        return utils.TruncatedNormal(mu, std)

Two heads sharing the same backbone: a mu head (action mean) and a log_std head (action variability). Crucially, std is learned here (unlike Problem 3 where std=0.1 is fixed). The slice self.policy[:-1] runs all layers except the final one, so we get the shared 256-D feature vector that both heads branch from.

The Critic

on_policy.py:42-58class Critic(nn.Module):
    def __init__(self, obs_shape, hidden_dim):
        super().__init__()
        self.value_net = nn.Sequential(
            nn.Linear(obs_shape[0], hidden_dim), nn.LayerNorm(hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, 1))                # scalar V(s)

    def forward(self, obs):
        return self.value_net(obs)                  # shape [batch, 1]

V(s) not Q(s, a) — the critic only takes obs as input. PPO doesn't need Q because the policy gradient theorem works with V as a baseline.

The PPOAgent.update method — structure

This is where everything happens. Roughly:

1. Concatenate the rollout into flat tensors (lines 168-189).
2. Compute GAE advantages and returns — your Task 2(a) (lines 192-205).
3. Loop ppo_epochs times, with shuffled minibatches (lines 218-275).
4. Inside each minibatch: compute clipped surrogate — your Task 2(b) — plus value loss, entropy, reverse-KL. Combine. Accumulate gradients.
5. Step optimizer once per epoch (lines 274-275).

One subtle point: gradients accumulate across minibatches inside an epoch, with each minibatch loss scaled by (batch_size / total_size). The optimizer steps once per epoch. This is equivalent to averaging the loss across the entire epoch, but avoids OOM from holding the whole rollout in one forward pass.

Your Three Changes, Decoded

Now the centerpiece: every line you wrote, explained at the level of "what is this PyTorch operation actually doing." This chapter assumes you've read the PyTorch primer in Chapter 02.

for t in reversed(range(T)):
    delta = rewards[t] + discounts[t] * next_values[t] − values[t]
    gae = delta + discounts[t] * self.gae_lambda * gae
    advantages[t] = gae

returns = advantages + values
return advantages, returns

values_all = self.critic(obs_all).squeeze(−1)
next_values_all = self.critic(next_obs_all).squeeze(−1)
advantages_all, returns_all = self.compute_gae(rew_all, values_all, next_values_all, disc_all, done_all)

adv_std = advantages_all.std(unbiased=False)
advantages_all_old = advantages_all
advantages_all = (advantages_all − advantages_all.mean()) / (adv_std + 1e-8)

ratio = torch.exp(new_log_prob − olp_ep)
surrogate1 = ratio * adv_ep
surrogate2 = torch.clamp(ratio, 1.0 − self.clip_eps, 1.0 + self.clip_eps) * adv_ep
policy_loss = −torch.min(surrogate1, surrogate2).mean()

with torch.no_grad():
    ref_dist = self.reference_actor(obs_ep)
    ref_log_prob = ref_dist.log_prob(act_ep).sum(−1)
reverse_kl = (ratio * (new_log_prob − ref_log_prob)).mean()

values_pred = self.critic(obs_ep).squeeze(−1)
value_loss = F.mse_loss(values_pred, ret_ep)

loss = (policy_loss
        + self.value_coef    * value_loss
        − self.entropy_coef  * entropy
        + self.reverse_kl_coef * reverse_kl)

(loss * (batch_size / T)).backward()

nn.utils.clip_grad_norm_(params, max_norm=1)
self.opt.step()

Cheat Sheet & Self-Quiz

Equations to memorize

Tensor shape cheat-sheet

Self-quiz

1. Why does the policy gradient theorem need only log-probs of actions taken — not over all possible actions?
2. What does subtracting the value baseline V(s) accomplish? Why doesn't it bias the gradient?
3. What does λ control in GAE? What are λ=0 and λ=1 corner cases called?
4. Why does GAE iterate backwards through time?
5. Why is the importance ratio computed as exp of a difference, not as a direct division?
6. What happens to the gradient when ratio > 1+ε and  > 0?
7. Why is the policy loss negated before backward?
8. Why is GAE computation wrapped in with torch.no_grad():?
9. What's the difference between old_log_prob and new_log_prob in the same minibatch?
10. Why does the entropy term get subtracted from the loss?
11. Why does PPO need a frozen reference actor for KL?
12. Why is the actor's std learned here but fixed at 0.1 in Problem 3?

Running it

# Pre-flight (one-time)
modal secret create wandb-secret WANDB_API_KEY=<your_key>

# Launch training
cd /Users/ozyphus/Documents/GitHub/cs224r/homework/hw2/hw2
modal run --detach modal_on_policy.py

~45 minutes to 1.5 hours on Modal A10. The wandb URL prints at startup. Watch eval/episode_success, policy_loss, and actor_std — if std collapses to 0 in the first 50k steps, the entropy_coef is too low.

You can now teach this

Three big ideas, in order of importance:

1. Policy gradient with a baseline. Push log-prob of action up by amount A(s,a). The baseline V(s) reduces variance without changing the gradient's expectation.
2. GAE. A geometrically-weighted sum of TD errors that interpolates between 1-step TD and Monte Carlo. λ tunes the bias-variance trade. The backwards recursion is just an efficient way to compute the geometric sum.
3. PPO clipping. Compute the importance ratio ρ = π_new/π_old. Take min of (ρ·A, clipped(ρ)·A). The min creates an asymmetric pessimistic cap that prevents the policy from drifting too far per update. This is the entire stability mechanism.

If a friend asks: "Why does PPO use min instead of average?" — you say: "Because the optimizer is adversarial. If we used average, the optimizer would happily blow past the clip in directions where the unclipped term is favorable, defeating the whole point. The min is a pessimistic bound: the optimizer always sees the worse of the two surrogates, so it can't game the clip."

You can teach this. Go finish the training.