Policy Gradients (REINFORCE)

Policy gradients skip the value table and optimize the policy directly. If a sampled action leads to a high return, raise its log-probability; if it leads to a low return, lower it. REINFORCE is that sentence written as a gradient, and it is the ancestor of PPO, GRPO, and every RLHF recipe in use today.

The idea

The objective is the expected return of the policy's own trajectories:

The difficulty is that the distribution being averaged over also depends on $\theta$. The log-derivative trick resolves it:

which turns the gradient of an expectation into an expectation of a gradient:

Weights no longer appear inside the distribution, so you can estimate the gradient from sampled episodes.

Three refinements make it usable:

• Baselines. Subtracting any function $b(s_t)$ leaves the expectation unchanged but shrinks variance: use $G_t - b(s_t)$ with $b \approx V(s_t)$.
• Reward-to-go. Rewards earned before step $t$ cannot depend on $a_t$, so replace the full return with $G_t^{\text{to-go}} = \sum_{k \ge t} \gamma^{k-t} r_{k+1}$.
• On-policy. The expectation is over samples from the current $\pi_\theta$, so one batch licenses exactly one update; the moment $\theta$ moves, the data is stale.

Worked example

A two-action softmax policy with logits $z = [1.0, 0.0]$ gives $\pi = [0.731, 0.269]$.

The gradient of $\log \pi(a_1)$ with respect to the logits is $(1 - \pi_1,\ -\pi_2) = (0.269,\ -0.269)$.

Suppose the sampled action was $a_1$ with return-to-go $G = 3$ and learning rate $\alpha = 0.1$. Ascending the objective gives

Re-normalizing, $\pi_1$ rises from 0.731 to about 0.762. Had the same action returned $G = -1$, the identical formula moves the logits the other way and $\pi_1$ drops to about 0.720. The action is not important; the sign and size of the return attached to it is the entire learning signal.

In code

import torch

logits = torch.tensor([[1.0, 0.0]], requires_grad=True)
log_probs = torch.log_softmax(logits, dim=-1)
action, G = 0, 3.0

loss = -(log_probs[0, action] * G)   # minimize the negative objective = ascend
loss.backward()

print(-0.1 * logits.grad)            # [[ 0.0807, -0.0807]] -> the ascent step

A baseline would replace G with G - value.detach(), and reward-to-go would replace the scalar with a per-step tensor.

Check yourself

1. REINFORCE is unbiased. Why subtract a baseline if the expected gradient does not change?
2. Why can you not reuse one REINFORCE batch for a second gradient step after $\theta$ has changed?
3. What is the difference between $G_t$ measured from the start of the episode and the reward-to-go, and which one does the update use?

Key takeaways

• The log-derivative trick turns "how well did this sample do" into a differentiable weight on log-probabilities.
• Baselines and reward-to-go reduce variance without introducing bias.
• Policy gradients are on-policy: the samples must come from the policy being updated.