The widget at the end of §15.3 showed something strange. Adding the same constant to every reward left the expected policy gradient exactly where it was, yet it made each individual estimate much noisier, and with enough of an offset the learning curves dissolved into noise. The cure in the widget was to subtract the average reward before multiplying by the score. That raises an obvious worry. If we change the weights on the scores, are we still estimating the gradient of the thing we want to maximize?

We are, and the reason is a one-line lemma that does a surprising amount of work. This section proves it, uses it to justify both baselines and reward-to-go, works out which baseline removes the most variance, and arrives at the advantage function, the quantity that every modern policy gradient method, PPO and GRPO included, multiplies its scores by. Then we learn the baseline with a second model, the critic, which gives actor-critic methods.

The expected score is zero

Fix a state $s$ and let $\pi_\theta(\cdot\mid s)$ be any differentiable policy. Then

In plain words, probabilities always sum to one, so you cannot raise all of them at once, and on average the directions that raise each sampled action's probability cancel out. For a softmax policy you can see it directly, $\mathbb{E}_a[\mathbf{e}_a - \boldsymbol\pi] = \boldsymbol\pi - \boldsymbol\pi = \mathbf{0}$. Spinning Up calls this the expected grad-log-prob lemma. It is the score-function analogue of the fact that a log-likelihood's gradient (the score in statistics) has mean zero under the model.

Baselines add no bias

Take the REINFORCE estimator with reward-to-go from §15.3, $\sum_t\nabla_\theta\log\pi_\theta(a_t\mid s_t)\,G_t$, where $G_t$ is the (discounted) return from step $t$. (We drop the $\gamma^t$ factor, as practice does. Everything below holds with it too.) Now subtract a baseline $b(s_t)$ from each return,

The change in the expected estimate is $-\sum_t\mathbb{E}[\nabla_\theta\log\pi_\theta(a_t\mid s_t)\,b(s_t)]$. Compute each term by first conditioning on the state $s_t$ (the law of iterated expectations), so that $b(s_t)$ is a constant inside the inner expectation,

So subtracting any function of the state leaves the expected gradient exactly unchanged. The same argument works if the baseline depends on anything known before $a_t$ is drawn, the whole history $h_t = (s_0, a_0, \dots, s_t)$, the time step $t$, a random number, or the parameters themselves. Condition on that information instead of on $s_t$ alone and the inner expectation is still zero.

The one thing a baseline must not depend on is the action $a_t$ (or anything after it). Then the inner expectation becomes $\sum_a\pi_\theta(a\mid s)\,b(s,a)\nabla\log\pi_\theta(a\mid s) = \sum_a b(s,a)\nabla\pi_\theta(a\mid s)$, which is not zero in general. The extreme case makes the point. If you chose $b(s, a) = G$, the "perfect" baseline that equals the return, every estimate would be exactly zero and nothing would ever be learned.

Reward-to-go comes from the same lemma

In §15.3 we replaced the trajectory form, $R(\tau)\sum_t\nabla\log\pi_\theta(a_t\mid s_t)$, by the reward-to-go form, using the policy gradient theorem. The lemma gives a direct proof that needs no theorem. Expand $R(\tau) = \sum_{t'}\gamma^{t'}r_{t'}$ and look at one cross term, the reward at step $t'$ times the score at step $t$, with $t' < t$. The reward $r_{t'}$ was earned before $a_t$ was chosen, so it is a function of the history $h_t$, and

In plain words, an action cannot be credited or blamed for rewards that arrived before it, because on average those rewards are uncorrelated with which action was taken. Deleting all of these zero-mean terms leaves $\sum_t\nabla\log\pi_\theta(a_t\mid s_t)\sum_{t'\ge t}\gamma^{t'}r_{t'} = \sum_t\gamma^t\,G_t\,\nabla\log\pi_\theta(a_t\mid s_t)$, exactly the theorem's sample form. The deleted terms had zero mean but nonzero variance, so removing them can only help (in expectation, removing a zero-mean term that is uncorrelated with the rest always lowers the variance, and in practice it helps even when the terms are correlated). Reward-to-go and baselines are the same idea. Each subtracts something whose expected product with the score is zero.

Computing reward-to-go for a batch of variable-length episodes stored as a padded array is one of the most common pieces of code in RL, and in RLHF it runs over batches of token sequences of different lengths.

import numpy as np

def reward_to_go(rewards, mask, gamma):
    raise NotImplementedError

How much variance a baseline removes

Unbiasedness says a baseline is free. Variance says which baseline to pick. Look at a single decision, so that the estimator is $\mathbf{g} = (G - b)\,\mathbf{s}$ with $\mathbf{s} = \nabla_\theta\log\pi_\theta(a\mid s)$ the score of the sampled action and $G$ its (random) return, and let $b$ be a constant. The mean of $\mathbf{g}$ does not depend on $b$, so minimizing the total variance $\mathbb{E}\lVert\mathbf{g}\rVert^2 - \lVert\mathbb{E}\mathbf{g}\rVert^2$ is the same as minimizing the second moment,

a parabola in $b$. Setting its derivative to zero gives the optimal constant baseline,

In plain words, the best baseline is an average of the returns, weighted by how large each sample's score is. Samples whose scores are large contribute the most variance, so the baseline should match their returns best.

If the score's size is unrelated to the return, $b^\star$ is just $\mathbb{E}[G]$, the expected return, which in a state $s$ is $V^\pi(s)$. The two usually differ a little (the weighting favors rare actions, which have larger scores, since $\lVert\mathbf{e}_a - \boldsymbol\pi\rVert$ grows as $\pi_a$ shrinks), but almost nobody bothers with $b^\star$ in practice. $V^\pi(s)$ captures nearly all of the gain in absolute terms, and it has a second virtue that $b^\star$ lacks, which we will see in a moment. (Greensmith, Bartlett, and Baxter (2004) give the full analysis of optimal baselines, including state-dependent ones.)

For this bandit, whose rewards share a large offset, the expected reward as a baseline cuts the variance of a single estimate from 44 to 0.07, a factor of more than 600, and that ratio is the number of samples you save. The optimal constant shrinks what is left by another factor of about 2.4, because it leans toward the returns of the two rarer actions, whose scores are larger. In absolute terms almost all of the gain came from the first step. Drag the rewards and the policy around below and the parabola moves, but its minimum always sits close to the expected reward.

The same effect shows up in a real episode. Here is the short corridor from §15.3, with the single-episode gradient estimate of $\partial J/\partial\theta$ (the logit of "right") sampled 3,000 times at $p = 0.8$, with and without the average return as a baseline.

Both estimators are centered on the exact gradient, $-7$. (The policy goes right too often, so the gradient says to lower $\theta$.) The baseline cuts the standard deviation from about 34 to about 21, which means $(34/21)^2 \approx 2.6$ times fewer episodes for the same precision. The gain is smaller than for the bandit because a single constant cannot follow the return as it changes along the episode. In this corridor every state looks the same to the agent, so a state-dependent baseline is the same single constant. A baseline that depends on the time step, which the lemma also allows, could follow the shrinking returns and do better.

import numpy as np

def pg_estimates(z, actions, rewards, b):
    raise NotImplementedError

def baseline_check(z, r_mean, b, n, seed):
    raise NotImplementedError

import numpy as np

def optimal_baseline(scores, G):
    return float(np.mean(G))   # placeholder: the plain mean

The advantage function

Choose the baseline $b(s_t) = V^\pi(s_t)$. Since $\mathbb{E}[G_t\mid s_t, a_t] = Q^\pi(s_t,a_t)$, the expected weight on each score becomes

the advantage of §15.1, and the policy gradient takes the form used by every algorithm from here on,

In plain words, raise the probability of actions that did better than the policy usually does from that state, and lower the probability of actions that did worse. The size of the push is how much better or worse.

This is what makes $V^\pi$ the natural baseline, beyond its near-optimal variance. The advantage is a relative judgment, and relative judgments are what a policy update needs. Three properties make it pleasant to work with.

• It averages to zero over the policy's own actions, $\sum_a\pi(a\mid s)A^\pi(s,a) = V^\pi(s) - V^\pi(s) = 0$, so in every state some actions are pushed up and some down, and the policy is not dragged in one direction by the overall level of reward.
• Its sign answers the question an update needs answered, better or worse than the current behavior. A state where everything is bad (every return is $-40$) still has actions with positive advantage, the least bad ones.
• It is the quantity that policy improvement is about. The greedy policy of §15.2 picks $\arg\max_aQ^\pi(s,a) = \arg\max_aA^\pi(s,a)$, and §15.6 will show that the improvement of one policy over another is exactly an average of advantages.

The rest of this chapter is mostly about estimating $A^\pi(s_t, a_t)$ well. Sampling $G_t - V(s_t)$ is one estimate. The next two sections give better ones.

Learning the baseline, the critic

We do not know $V^\pi$, so we learn it. A value network (or a table) $V_\psi(s)$ with parameters $\psi$ is trained by regression on the returns the policy actually produces,

The minimizer of the expected squared error is the conditional mean, $\mathbb{E}[G_t\mid s_t] = V^\pi(s_t)$, so with enough capacity and data the regression recovers the right baseline. Two facts keep this safe. First, if the targets are Monte Carlo returns, an inaccurate $V_\psi$ only costs variance, never bias, because the unbiasedness proof above worked for any function of the state. Second, the policy changes after every update, so $V_\psi$ is always chasing a moving target, and in practice it is updated a few gradient steps per batch, alongside the policy.

The value model is called the critic and the policy the actor. In RLHF the critic is a whole second language model, usually initialized from the reward model, that outputs a value for every prefix of a response (§16.4). It is one of the four models that make RLHF with PPO expensive, and dropping it in favor of a baseline computed from a group of sampled responses is GRPO's main saving (§18.4).

Actor-critic methods

Once we have a critic, we can use it for more than a baseline. Replace the sampled return $G_t$, which needs the whole rest of the episode, with a one-step bootstrapped estimate, as TD(0) did in §15.2,

the TD error, with $V_\psi(s_{t+1}) = 0$ if $s_{t+1}$ is terminal. Is that a reasonable estimate of the advantage? If the critic is exact, $V_\psi = V^\pi$, then yes, exactly, because

using the one-step relation $Q^\pi(s,a) = r(s,a) + \gamma\sum_{s'}P(s'\mid s,a)V^\pi(s')$ from §15.1. In plain words, the TD error is a one-sample advantage. It says how much better this step turned out than the critic expected, measured with the critic's own estimate of where it landed.

If the critic is wrong, $\delta_t$ is biased, and unlike the baseline case the bias now reaches the gradient, because the error in $\gamma V_\psi(s_{t+1})$ depends on where the action led. In exchange, its variance is far lower than $G_t - V_\psi(s_t)$, since it involves one reward and one transition instead of a whole future. This is the Monte Carlo versus TD trade-off of §15.2 again, now for advantages, and §15.5 turns it into a dial.

An actor-critic algorithm runs both updates side by side. After each transition $(s_t, a_t, r_t, s_{t+1})$,

where $\alpha_c$ and $\alpha_a$ are the critic's and actor's step sizes. (The critic update is a semi-gradient step on $\frac12\delta_t^2$ that treats the target $r_t + \gamma V_\psi(s_{t+1})$ as a constant, exactly TD(0).) It is generalized policy iteration from §15.2 run continuously, a critic evaluating the current policy and an actor improving against the critic. A2C and A3C (Mnih et al. 2016) are this algorithm with many environments running in parallel, $n$-step returns instead of one-step ones, and an entropy bonus, and they were the standard deep RL baselines before PPO.

import numpy as np

def actor_critic_step(theta, v, s, a, r, s_next, terminated, lr_actor, lr_critic, gamma):
    raise NotImplementedError

Here is actor-critic on the gridworld, in the synchronous style of A2C. Sixteen copies of the world run side by side, and after every single step of all sixteen, the critic and the actor each take one averaged update. Episodes are cut off after 50 steps, and a cut-off is not a terminal state, so the critic still bootstraps from where the agent was. We track the exact value of the start cell under the current policy, which value iteration told us is at most $V^\star = 0.351$.

Actor-critic reaches 0.333 after 8,000 transitions and 0.340 after 16,000, against the optimal 0.351, learning from every single step without waiting for any episode to finish. On this small world REINFORCE with a batch of complete episodes does about as well for the same number of transitions, because its returns are short and well scaled. Actor-critic's real advantages appear when episodes are long or never end, where waiting for the full return is slow or impossible, and when the reward scale makes raw returns noisy. Its real weakness is the critic's bias early in training, which you can see in the first numbers, when the critic still knows nothing.

Normalizing advantages

Almost every implementation standardizes the advantages within each batch before using them, replacing $\hat A_i$ by $(\hat A_i - \bar A)/(\operatorname{std}(\hat A) + \epsilon)$. Subtracting the batch mean is one more baseline (with the small batch-dependence bias discussed above). Dividing by the standard deviation rescales the whole gradient by one positive number per batch, which does not change its direction, only its length, and acts as an automatic learning-rate adjustment that makes the step size insensitive to the reward scale. It is one of the "implementation details" that turns out to matter in practice (§15.6).

In PyTorch

The A2C loss in production code combines the actor's surrogate (with advantages detached, so the policy gradient does not flow into the critic), the critic's regression loss, and usually an entropy bonus. Here it is for one batch of gridworld transitions, with separate actor and critic networks.

Read more

• OpenAI Spinning Up. Part 3: Intro to Policy Optimization. The expected grad-log-prob lemma, reward-to-go, and baselines, with proofs in the same style.
• Greensmith, Bartlett & Baxter (2004). Variance Reduction Techniques for Gradient Estimates in Reinforcement Learning. The careful analysis of baselines, including the optimal one.
• Mnih et al. (2016). Asynchronous Methods for Deep Reinforcement Learning. A3C and its synchronous cousin A2C, the standard actor-critic recipe with parallel environments and entropy regularization.
• Konda & Tsitsiklis (2000). Actor-Critic Algorithms. The two-timescale analysis behind running a fast critic and a slow actor.
• Sutton & Barto (2018). Reinforcement Learning: An Introduction, 2nd edition. Sections 13.4 and 13.5 cover REINFORCE with baseline and one-step actor-critic.
• Levine. CS 285, Lecture 6, Actor-Critic Algorithms. Critics, bootstrapped advantages, and the bias-variance trade-off, from the lecture we follow.

Watch next

▶

▶