Policy gradients · from scalar to step

Reward doesn't point.
It scales.

In supervised learning, the loss tells you which direction to move. In reinforcement learning it can't — the environment isn't differentiable and never will be. So reward does something narrower and much stranger: it multiplies a gradient the policy already produced on its own.

$$\nabla_\theta J(\theta)=\mathbb{E}\Big[\underbrace{\nabla_\theta \log \pi_\theta(a\mid s)}_{\textstyle \text{direction}}\;\cdot\;\underbrace{A}_{\textstyle \text{scale}}\Big]$$
Direction — chosen by the policy Scale — one number from the world

You can't differentiate the world

The objective is expected return under your own policy:

$$J(\theta)=\mathbb{E}_{\tau\sim\pi_\theta}\big[R(\tau)\big]$$

Look at where $\theta$ appears. It is not inside $R$. It's in the distribution you're averaging over. $R$ is a game engine, a physics simulator, a compiler exit code, a human clicking thumbs-up — a black box that eats a trajectory and emits a number. There is no $\partial R/\partial\theta$ to compute, and no amount of autograd will invent one.

So the derivative has to be moved off the reward and onto the only differentiable thing in sight: the probability of the trajectory.

Moving the derivative

  1. Write the expectation as an integral, then push the gradient inside

    $$\nabla_\theta J=\nabla_\theta\!\int\! P(\tau;\theta)\,R(\tau)\,d\tau=\int\!\nabla_\theta P(\tau;\theta)\,R(\tau)\,d\tau$$

    $R(\tau)$ passes through untouched. It's a constant attached to each trajectory, not a function of the parameters.

  2. Use the log-derivative trick to get an expectation back

    $$\nabla_\theta P = P\,\nabla_\theta \log P \quad\Longrightarrow\quad \nabla_\theta J=\mathbb{E}_\tau\big[\nabla_\theta \log P(\tau;\theta)\cdot R(\tau)\big]$$

    Multiplying and dividing by $P$ restores the shape of an average, which means you can estimate it by sampling — by simply running the policy.

  3. Watch the environment cancel

    $$P(\tau;\theta)=\rho(s_0)\prod_t P(s_{t+1}\mid s_t,a_t)\,\pi_\theta(a_t\mid s_t)$$

    Take the log and the product becomes a sum. Differentiate, and every term that doesn't contain $\theta$ — the initial state distribution, all the transition dynamics — is a constant and dies. Only the policy survives:

    $$\nabla_\theta \log P(\tau;\theta)=\sum_t \nabla_\theta \log \pi_\theta(a_t\mid s_t)$$

    You never needed a model of the environment. That's the quiet miracle of this derivation.

  4. The estimator

    $$\nabla_\theta J=\mathbb{E}\Big[\sum_t \nabla_\theta \log \pi_\theta(a_t\mid s_t)\cdot R(\tau)\Big]$$

    Roll out, collect $(s,a,R)$, average. The reward was never differentiated — it arrived as a coefficient and stayed one.

What each factor actually does

$\nabla_\theta \log \pi_\theta(a\mid s)$ is the direction in parameter space that makes the action you already took more likely next time you see this state. It is computed entirely inside your network. It has no idea whether that action was a good one.

$A$ is the only quantity in the whole expression that has met reality, and it is a single number. Its sign decides whether you reinforce that action or suppress it. Its magnitude decides how hard.

The policy proposes the direction. The reward only signs it and sizes it.

The panel below makes this literal. A three-action softmax policy lives on a triangle — each corner is certainty about one action, the middle is a uniform policy. The arrow is the actual update. Change which action was sampled and the arrow swings to a new axis. Change the reward and the arrow only stretches or flips.

Interactive · one gradient stepstep 000
Policy on a probability simplex, with the gradient step drawn as an arrow a₁ a₂ a₃ CORNER = P(ACTION) → 1
Action sampled a₁
Advantage +1.00
Learning rate 0.30
The arrow always lies on the line joining the policy to the sampled action's corner. Positive advantage pulls toward it, negative pushes directly away, and the number only sets how far. Set the advantage to zero and the update vanishes — not because the action was bad, but because it carried no information.

It's weighted cross-entropy

Nobody implements the formula above by hand. In practice you write a scalar "loss," call .backward(), and let autograd produce the same thing:

# states, actions, advantages come from rollouts
logits = policy(states)
logp   = -F.cross_entropy(logits, actions, reduction="none")  # = log π(a|s)
loss   = -(advantages.detach() * logp).mean()
loss.backward()

That is cross-entropy where the label is whatever action you happened to sample, weighted by a number the environment handed you. The advantage is detach()ed — it's outside the computation graph on purpose. Autograd only ever differentiates $\log\pi$.

One thing worth internalising: this is not a loss. Its value means nothing — it isn't a distance from anything, it doesn't measure error, and watching it go down tells you very little. Only its gradient is meaningful, and only because we constructed it to match $\nabla_\theta J$ in expectation. It's a surrogate, and it earns its name honestly.

Then everything else is variance

The estimator is unbiased, and that's about all it has going for it. Rolling one scalar backward across a whole trajectory is an extremely noisy way to learn. Almost every technique layered on top is a variance fix that leaves the expectation alone.

Reward-to-go

The naive form multiplies every action by the return of the whole trajectory — including reward collected before that action was taken. An action cannot influence the past. Dropping those terms changes nothing in expectation and removes a pile of pure noise:

$$\hat{g}=\sum_t \nabla_\theta \log \pi_\theta(a_t\mid s_t)\underbrace{\sum_{t'\ge t}\gamma^{t'-t} r_{t'}}_{\text{reward-to-go}}$$

Baselines are free

Subtract any $b(s)$ that doesn't depend on the action. The correction term vanishes exactly:

$$\mathbb{E}_{a\sim\pi}\big[\nabla_\theta \log \pi_\theta(a\mid s)\,b(s)\big]=b(s)\sum_a \nabla_\theta \pi_\theta(a\mid s)=b(s)\,\nabla_\theta 1=0$$

Because probabilities sum to one, their gradients sum to zero. So you may subtract whatever you like, for free, as long as it's constant across actions. Choose $b(s)=V^\pi(s)$ and the coefficient becomes the advantage $A=Q-V$, which changes the question being asked at every step — from "did this go well?" to "did this go better than I expected from here?"

The difference is not cosmetic. With all-positive rewards and no baseline, every single sampled action gets reinforced. You're pushing probability up everywhere and learning only from the differences in how hard — differences that sit buried inside the sampling noise.

Interactive · same expectation, different variance3-arm bandit
Rewards are drawn around means of 8, 10 and 14 — all comfortably positive. Both agents use the identical learning rate and have the same expected gradient; the only difference is that one subtracts a running mean of past reward. Averaged over 60 runs. The gap is entirely variance.

What this looks like in modern RL

PPO's clipping. Once you take more than one gradient step on a batch, the data was drawn from $\pi_{\text{old}}$, not $\pi_\theta$, so you correct with an importance ratio $r_t(\theta)=\pi_\theta/\pi_{\text{old}}$. PPO then clips it:

$$L=\mathbb{E}\big[\min\big(r_t(\theta)\hat{A}_t,\;\text{clip}(r_t(\theta),1-\epsilon,1+\epsilon)\hat{A}_t\big)\big]$$

The min makes the objective flat once the ratio has moved too far in the direction the advantage wanted, so the gradient becomes exactly zero there. It's a trust region enforced by making the surrogate stop rewarding you, rather than by a constraint solver.

RLHF. The reward model scores an entire completion and returns one scalar. That single number then weights the log-probability gradient of every token in the response. Credit assignment across hundreds of tokens from one number is exactly as hard as it sounds, and it's why so much RLHF engineering is really variance engineering.

GRPO. Drops the value network entirely. Sample a group of $G$ responses to the same prompt, and use the group's mean reward as the baseline — normalising by the group's standard deviation too. It's legitimate for precisely the reason above: the group mean is constant across the actions being compared, so subtracting it costs no bias and buys a large amount of variance.

The mental model

Reward is a coefficient

It multiplies a gradient; it never produces one. Sign chooses reinforce or suppress, magnitude chooses how much.

It needn't be differentiable

Or continuous, or smooth, or even computable in closed form. It only has to be comparable across samples from the same state.

One scalar vs. a whole vector

Supervised learning gets a full gradient per example. RL gets a number per trajectory. That gap is most of why RL is so sample-hungry.

The rest is variance reduction

Reward-to-go, baselines, GAE, clipping, group normalisation — all the same identity, estimated less noisily.