Policy Gradients and REINFORCE

A policy samples an action. Sampling is not differentiable. So how does a gradient get from the reward back to the weights? The answer is one identity, it takes three lines, and every policy gradient method in use is built on it.

The thing that blocks everyone

The objective is the expected return under the policy: J(theta) = E[R(tau)], where tau is a trajectory drawn from the policy and R is its return. To improve the policy we need d J / d theta.

The obstacle is not the reward function, which is usually not differentiable either but does not need to be. The obstacle is that theta appears in the distribution being averaged over, not in the thing being averaged. Writing it out:

J(theta) = sum over tau of  p_theta(tau) * R(tau)

grad J   = sum over tau of  grad p_theta(tau) * R(tau)      <-- R has no theta in it

That expression is exact and useless. It is a sum over every possible trajectory, and it is not an expectation under anything you can sample from — grad p_theta(tau) is not a probability distribution, so there is no way to estimate the sum by drawing trajectories and averaging. Every attempt to get a usable policy gradient runs into this wall, and the standard treatments tend to skip straight over it.

The log-derivative trick

The fix is an identity from calculus, applied in the direction people rarely use. Since grad log f = grad f / f, it follows that grad f = f * grad log f. Substituting p_theta(tau) for f:

grad p_theta(tau)  =  p_theta(tau) * grad log p_theta(tau)

so:
grad J  =  sum over tau of  p_theta(tau) * grad log p_theta(tau) * R(tau)
        =  E over tau ~ p_theta of [ grad log p_theta(tau) * R(tau) ]

That is the whole move. The right-hand side is an expectation under the policy, so it can be estimated by the obvious procedure: run the policy, collect trajectories, average grad log p * R over them. Nothing had to be differentiated through the sampling step, because the sampling was absorbed into the expectation and only the log-probability of what was actually sampled gets differentiated.

There is a second, less advertised payoff. Expand the log-probability of a whole trajectory:

log p_theta(tau) = log p(s_0)
                 + sum over t of [ log pi_theta(a_t | s_t) + log P(s_t+1 | s_t, a_t) ]

grad log p_theta(tau) = sum over t of grad log pi_theta(a_t | s_t)

The initial state distribution and the environment dynamics do not depend on theta, so their gradients are zero and they drop out entirely. You never need a model of the environment. This is why policy gradient methods work on simulators, on games and on physical systems whose dynamics nobody has written down, and it falls out of the identity rather than being an extra assumption.

The estimator, term by term

grad J  ~=  (1/N) * sum over episodes of  sum over t of
                grad log pi_theta(a_t | s_t) * G_t

where G_t = r_t + gamma*r_(t+1) + gamma^2*r_(t+2) + ...

Read as an instruction, the update is: make everything you did in a good episode more likely, and everything you did in a bad episode less likely, in proportion to how good or bad. It has no way to distinguish the three good decisions in a successful episode from the twelve irrelevant ones, and it does not try. It relies on averaging over many episodes for the irrelevant actions to cancel, which is the mechanism behind the credit assignment problem.

One refinement is free and always applied: G_t should be the return from t onwards, not the return of the whole episode. Rewards collected before an action cannot have been caused by it, and dropping them removes variance without introducing bias.

REINFORCE in twelve lines

The gradient never appears explicitly in code. You write a scalar whose gradient equals the estimator, and let autograd do the rest. The negative sign is because optimisers minimise.

import torch

def discounted_returns(rewards, gamma=0.99):
    out, running = [], 0.0
    for r in reversed(rewards):
        running = r + gamma * running
        out.append(running)
    return list(reversed(out))

def reinforce_step(logits, actions, returns, optimizer, baseline=0.0):
    """logits: (T, A) from the policy. actions: (T,) long. returns: (T,) float."""
    logp = torch.log_softmax(logits, dim=-1)                  # (T, A)
    chosen = logp.gather(1, actions.unsqueeze(1)).squeeze(1)  # (T,)
    advantage = returns - baseline
    loss = -(chosen * advantage.detach()).mean()              # the surrogate
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

Two details in that snippet are load-bearing. The .detach() on the advantage is not optional: the return is data, not a function of the weights, and letting a gradient flow into it computes something that is not the policy gradient. And the printed loss value is meaningless as a training signal — it is a surrogate whose gradient is right, not an objective whose value means anything. Watch the average return instead.

Why subtracting a baseline is free

Suppose two actions in a state lead to returns of 100 and 102. Both are positive, so both get pushed up, with weights that differ by two per cent. The information you care about — that one is better — is a two per cent difference riding on a hundredfold common-mode signal, and it is buried in the sampling noise.

Subtract a baseline of 101 and the weights become -1 and +1. The two actions now push in opposite directions and the common mode is gone. That is the entire reason baselines exist, and it costs nothing in correctness, because the expected value of the extra term is exactly zero:

E over a ~ pi of [ grad log pi(a|s) * b(s) ]
  = b(s) * sum over a of  pi(a|s) * grad log pi(a|s)
  = b(s) * sum over a of  grad pi(a|s)
  = b(s) * grad ( sum over a of pi(a|s) )
  = b(s) * grad (1)
  = 0

The proof needs only that the baseline does not depend on the action. It may depend on the state, on the batch, or on anything else you like. Three baselines are in common use: the mean return of the batch (free, crude, and often enough); a learned value function V(s), which is what makes a method actor-critic; and the mean reward of a group of samples for the same prompt, which is what GRPO uses instead of a value network.

The same equation with a language model in it

Nothing in the derivation assumed a small action space. Substitute a language model and every term keeps its meaning.

• pi_theta(a_t | s_t) is the probability the model assigns to the next token given everything so far — the number you can already inspect through logprobs. The action space is the vocabulary.
• A trajectory is a completion. Its log-probability is the sum of the per-token log-probabilities, which is a quantity the training stack computes anyway.
• R(tau) is whatever scores the completion: a reward model, a unit test suite, or an exact-match check against a known answer.

The resulting update is a cross-entropy loss on the model’s own samples, weighted by how good each sample turned out to be. Seen that way, RL post-training is closer to supervised fine-tuning than the vocabulary suggests — the difference is that the training examples are generated by the model being trained, and their weights can be negative.

Where REINFORCE stops being enough

Plain REINFORCE is rarely used at scale, for three reasons that each have a named fix.

1. The variance grows with the horizon. The estimator multiplies a whole trajectory’s worth of gradient directions by one noisy scalar. A learned critic replaces the sampled return with an estimate and cuts the variance sharply, at the cost of some bias — that is actor-critic.
2. Every sample is used once. The derivation assumed trajectories drawn from the current policy, so the moment you update, your data is stale and must be thrown away. Importance sampling lets you reuse it for a few more steps, which is what PPO is built around.
3. Step size is unforgiving. There is no natural scale for the update, and a step that is too large can move the policy somewhere the collected data does not describe, from which it may not recover. Trust regions and clipping exist for this specific failure.

None of those changes the identity at the centre. Every method in this cluster that trains a policy directly is the log-derivative trick plus machinery for making its variance manageable.