Policy gradients: the policy gradient theorem, REINFORCE and baselines

Before this lesson: returns, $V^\pi$, $Q^\pi$ and the advantage (lesson 01); Monte Carlo returns (lesson 04); the score-function gradient estimator (calculus 12); the chain rule (calculus 04); softmax (foundations 05); stochastic gradient ascent (optimization 05). Time: about 75 minutes, plus the notebook.

---

Why this matters #

Value-based methods (lesson 05 and lesson 06) need a max over actions. That is hard when the actions are continuous (the torques of a robot arm) or very many (all next tokens of a language model). Policy gradients work in both cases, they can learn stochastic policies, and they need no derivative of the reward: the reward can come from a simulator, a game result, a unit test or a human rating.

• Actor-critic methods and PPO (lesson 08 and lesson 09) are REINFORCE plus a learned baseline and a better estimate of the advantage.
• RL for language models. RLHF with PPO, and the critic-free methods RLOO and GRPO, all start from the score-function gradient of this lesson. They differ in the baseline.
• AlphaGo trained its policy network by self-play with a policy gradient (Silver et al., 2016). Neural architecture search (Zoph and Le, 2017) and sequence-level training of captioning models (self-critical sequence training, Rennie et al., 2017) use REINFORCE with a baseline.

1. Predict first #

2. The picture #

Policies with parameters #

A parameterized policy $\pi_\theta(a \mid s)$ gives a probability for every action, and it is differentiable in the parameters $\theta$. Two families cover most uses.

• Softmax policy (discrete actions). A function $h_\theta(s, a)$ gives a score, called a logit or a preference, to every action. The policy is the softmax of the logits:

In the tabular case $h_\theta(s, a) = \theta(s, a)$: one number per state–action pair. With features, $h_\theta(s, a) = \theta^\top x(s, a)$; in deep RL, a network outputs the logits.

• Gaussian policy (continuous actions). A function $\mu_\theta(s)$ gives the mean action, and the action is $a \sim N(\mu_\theta(s), \sigma^2)$. The standard deviation $\sigma$ is fixed, or it is a learned parameter (usually through $\log \sigma$, so that $\sigma > 0$).

The objective is the expected return from the start, $J(\theta)$. We do gradient ascent: $\theta \leftarrow \theta + \eta \, \widehat{\nabla J}$, with a learning rate $\eta$ and an estimate $\widehat{\nabla J}$ of the gradient.

One step first: the log-derivative trick #

Take one state and $K$ actions with fixed rewards $r_a$ (a bandit, lesson 03). The objective is the expected reward

Its gradient is $\nabla J = \sum_a \nabla \pi_\theta(a)\, r_a$. This sum is not an expectation, so we cannot estimate it from samples of $a$. The log-derivative trick fixes this. Since $\nabla \log \pi = \nabla \pi / \pi$, we have $\nabla \pi = \pi\, \nabla \log \pi$, and

Now the gradient is an expectation over actions from the policy itself. Sample $a \sim \pi_\theta$, look at the reward, and $\nabla \log \pi_\theta(a)\, r_a$ is an unbiased estimate of $\nabla J$. We never need to know how $r$ depends on $a$.

The vector $\nabla_\theta \log \pi_\theta(a)$ is the score of action $a$. It is the direction in parameter space that makes $a$ more likely, fastest. The reward says how much we want that.

The softmax score. For $\pi = \operatorname{softmax}(\theta)$, write $\log \pi_a = \theta_a - \log \sum_b e^{\theta_b}$. Differentiate with respect to $\theta_i$:

1. the first term gives $\partial \theta_a / \partial \theta_i = 1$ if $i = a$ and 0 otherwise;
2. the second term gives $e^{\theta_i} / \sum_b e^{\theta_b} = \pi_i$;
3. so $\partial \log \pi_a / \partial \theta_i = \mathbb 1[i = a] - \pi_i$, that is,

the one-hot vector of $a$ minus the probability vector. It raises the logit of $a$ and lowers every logit in proportion to its probability. Its entries add up to 0.

Interactive figure: loading…

Continuous actions: the Gaussian score #

For a Gaussian policy $a \sim N(\mu, \sigma^2)$,

so the scores of the two parameters are

Read the first one as a picture. A sample $a$ to the right of $\mu$ has a positive score: making it more likely means moving $\mu$ to the right. The REINFORCE term $\frac{a - \mu}{\sigma^2}\, r(a)$ moves $\mu$ toward samples with high reward. If $\mu = \mu_\theta(s)$ is the output of a network, the chain rule gives $\nabla_\theta \log \pi = \frac{a - \mu}{\sigma^2} \nabla_\theta \mu_\theta(s)$, and autodiff does this for you.

Interactive figure: loading…

3. Definitions and main results #

Episodes and the objective #

An episode, or trajectory, is $\tau = (s_0, a_0, r_0, s_1, a_1, r_1, \dots)$, where $r_t$ is the reward after $(s_t, a_t)$ (Sutton and Barto write $R_{t+1}$; see lesson 01). Under the policy $\pi_\theta$ its probability is

with the start-state distribution $p_0$. The objective is the expected discounted return from the start:

Two more objects:

• the reward-to-go (the return from step $t$), $G_t = \sum_{k \ge t} \gamma^{k-t} r_k$;
• the discounted state visitation $d^\pi(s) = \sum_{t \ge 0} \gamma^t \Pr(s_t = s)$. It is not normalized: it adds up to $\mathbb E[\sum_t \gamma^t]$, not to 1. The normalized version is $(1-\gamma) d^\pi$ when $\gamma < 1$ and episodes never end.

Assumptions. Finite state and action sets, bounded rewards, and $\pi_\theta(a \mid s)$ differentiable in $\theta$. To keep the proof short, every episode ends after at most $H$ steps ($H$ may be huge); then all sums are finite, and we may exchange $\nabla$ and $\sum$. A limit argument extends the result to episodes of unbounded length when $\gamma < 1$ or when the episode ends with probability 1 and has a finite expected length.

The policy gradient theorem #

Theorem (policy gradient; Sutton, McAllester, Singh and Mansour, NeurIPS 1999). Under the assumptions above,

Proof.

1. Write $J$ as a sum over trajectories: $J = \sum_\tau p_\theta(\tau)\, R(\tau)$ with $R(\tau) = \sum_t \gamma^t r_t$.

2. Likelihood ratio. $R(\tau)$ does not depend on $\theta$ once $\tau$ is fixed. Use $\nabla p_\theta = p_\theta \nabla \log p_\theta$:

3. The dynamics drop out. The log of the product is

Only the middle sum depends on $\theta$, so $\nabla \log p_\theta(\tau) = \sum_t \nabla \log \pi_\theta(a_t \mid s_t)$. No model of $P$ is needed. Together with step 2: $\nabla J = \mathbb E\big[\sum_t \nabla \log \pi_\theta(a_t \mid s_t) \sum_k \gamma^k r_k\big]$.

4. Causality. Take a pair $k < t$: the reward $r_k$ came before the action $a_t$. Condition on the history $h_t = (s_0, a_0, r_0, \dots, s_t)$. Given $h_t$, the reward $r_k$ is fixed, and

So $\mathbb E[\nabla \log \pi_\theta(a_t \mid s_t)\, r_k] = \mathbb E\big[r_k\, \mathbb E[\nabla \log \pi_\theta(a_t \mid s_t) \mid h_t]\big] = 0$. Drop all these pairs. What remains for step $t$ is $\sum_{k \ge t} \gamma^k r_k = \gamma^t G_t$:

5. Replace $G_t$ by $Q^\pi$. Condition on $h_t$ and $a_t$. By the Markov property, $\mathbb E[G_t \mid h_t, a_t] = Q^\pi(s_t, a_t)$, and the score is fixed given $h_t, a_t$. The tower rule gives the first form of the theorem.

6. The visitation form. Write the expectation at step $t$ as a sum over $s_t = s$ and $a_t = a$: $\mathbb E[\gamma^t \nabla \log \pi(a_t \mid s_t)\, Q^\pi(s_t, a_t)] = \gamma^t \sum_s \Pr(s_t = s) \sum_a \pi(a \mid s) \nabla \log \pi(a \mid s)\, Q^\pi(s, a)$. Use $\pi \nabla \log \pi = \nabla \pi$ and sum over $t$: $\sum_t \gamma^t \Pr(s_t = s) = d^\pi(s)$. $\blacksquare$

The surprising part: changing $\theta$ also changes which states the agent visits, but the gradient has no term $\nabla d^\pi$. All of that effect is inside $Q^\pi$. This is why the gradient can be estimated from episodes without a model.

Interactive figure: loading…

REINFORCE #

Replace the expectation in step 4 by one sampled episode. This is REINFORCE (Williams, 1992).

The factor $\gamma^t$. It comes from the proof and is needed for an unbiased gradient of the discounted objective. Most implementations drop it and use $\sum_t (G_t - b(s_t)) \nabla \log \pi_\theta(a_t \mid s_t)$. Then late steps count as much as early ones. This update is not the gradient of $J$; Nota and Thomas (2020) show that it is in general not the gradient of any function. It still works well in practice, and it is closer to what we usually want: good behaviour in every state, not only near the start.

Baselines #

Lemma (a baseline adds no bias). For any function $b(s)$ of the state only,

So $G_t - b(s_t)$ can replace $G_t$ at every step: condition on $s_t$ as in step 4 of the proof, and the extra term has mean 0. Only the variance changes.

Why the variance falls. Suppose every return is about $-50 \pm 2$. Without a baseline, every sampled action gets the weight $\approx -50$: each update pushes the sampled actions down with a big weight, and only the small differences between episodes tell good from bad. The noise in "which action happened to be sampled" is multiplied by 50. Subtract $b \approx -50$, and the weights become $\approx \pm 2$: an action is pushed up if it did better than usual and down if it did worse.

The natural choice is $b(s) = V^\pi(s)$. Then $G_t - V^\pi(s_t)$ is an estimate of the advantage $A^\pi(s_t, a_t) = Q^\pi(s_t, a_t) - V^\pi(s_t)$: how much better the action was than the policy's average in that state. Given the state, the advantage has mean 0. Learning $V^\pi$ with a second model gives the actor-critic methods of lesson 08.

A baseline must not depend on the action. With $b(s, a)$ the lemma fails: $\sum_a \nabla \pi(a \mid s)\, b(s, a)$ is in general not 0, and the estimate is biased. (Methods that use action-dependent control variates add a correction term to remove this bias.)

The variance-optimal constant baseline

Take the one-step problem with score $s_a = \nabla \log \pi(a)$ and a constant baseline $b$. The total variance of $\hat g = (r_a - b) s_a$ is $\mathbb E[(r_a - b)^2 \|s_a\|^2] - \|\nabla J\|^2$, because the mean does not depend on $b$. Set the derivative in $b$ to zero: $-2\, \mathbb E[(r_a - b) \|s_a\|^2] = 0$, so

$$b^* = \frac{\mathbb E\big[\|s_a\|^2\, r_a\big]}{\mathbb E\big[\|s_a\|^2\big]} .$$

It is a weighted average of the rewards, with weights $\|s_a\|^2$. It equals $J = \mathbb E[r_a]$ when every action has a score of the same length (as in the worked example below), and it differs from $J$ in general. In practice $b(s) \approx V^\pi(s)$ is used: it is close to optimal and easy to learn. The "best constant b" button in the first figure computes $b^*$.

The exact gradient of a tabular softmax policy #

Corollary. For $\pi_\theta(a \mid s) = \operatorname{softmax}(\theta(s, \cdot))_a$,

Proof.

1. The softmax score in state $s'$ only involves the row $\theta(s', \cdot)$: $\partial \log \pi(a' \mid s') / \partial \theta(s, a) = \mathbb 1[s' = s]\,\big(\mathbb 1[a' = a] - \pi(a \mid s)\big)$.
2. Insert this in the visitation form, written with $\nabla \pi = \pi \nabla \log \pi$: $\partial J / \partial \theta(s, a) = d^\pi(s) \sum_{a'} \pi(a' \mid s) \big(\mathbb 1[a' = a] - \pi(a \mid s)\big) Q^\pi(s, a')$.
3. The sum is $\pi(a \mid s) Q^\pi(s, a) - \pi(a \mid s) \sum_{a'} \pi(a' \mid s) Q^\pi(s, a') = \pi(a \mid s)\big(Q^\pi(s, a) - V^\pi(s)\big)$. $\blacksquare$

Read the three factors: how often (and how early) we visit $s$; how likely $a$ is now; how much better $a$ is than the average action in $s$. Two consequences. The gradient of row $s$ adds up to 0 (the advantages, weighted by $\pi$, add up to 0). And when $\pi(a \mid s)$ is almost 0, the gradient for a good action $a$ is almost 0 too: a softmax policy that is nearly deterministic in the wrong direction escapes very slowly (Mei et al., 2020). The notebook checks this formula against finite differences.

An entropy bonus #

Policy gradients can make the policy deterministic too early, before it has explored. A common fix adds the entropy $\mathcal H(\pi(\cdot \mid s)) = -\sum_a \pi(a \mid s) \log \pi(a \mid s)$ (information theory 01) of the visited states to the objective, with a small weight $c_e$. For a softmax, $\partial \mathcal H / \partial \theta_a = -\pi_a(\log \pi_a + \mathcal H)$: it raises the logits of rare actions and lowers those of frequent ones. A2C uses this bonus, and SAC makes entropy part of the objective itself (lesson 08).

Why the best policy can be stochastic: the short corridor #

In an MDP where the agent sees the state, some optimal policy is deterministic (lesson 01). With function approximation the agent may not be able to tell states apart, and then this is no longer true. The short corridor (Sutton and Barto, Example 13.1) shows it.

• Three cells S, X, Y and a goal G on the right; the reward is $-1$ per step, $\gamma = 1$.
• In S and Y, "right" moves right and "left" moves left (left in S stays in S).
• In X the actions are switched: "right" moves to S and "left" moves to Y.
• All three cells have the same features, so the policy is the same in every cell: right with probability $p$.

"Always right" goes S → X → S → … and never ends. "Always left" stays in S forever. Both have value $-\infty$. A stochastic policy reaches the goal. Write the Bellman equations for the three cells (lesson 01):

1. The first gives $V_S = V_X - 1/p$. Put it into the second: $(1-p) V_X = -2 + (1-p) V_Y$, so $V_X = V_Y - 2/(1-p)$.
2. Put that into the third: $V_Y = -1 + (1-p) V_Y - 2$, so $p V_Y = -3$ and $V_Y = -3/p$.
3. Back-substitute: $V_X = -3/p - 2/(1-p)$ and

4. Set $J'(p) = 0$: $2p(1-p) - (2p - 4)(1 - 2p) = 0$ simplifies to $p^2 - 4p + 2 = 0$, so $p^* = 2 - \sqrt 2 \approx 0.586$ (the other root is larger than 1). The value is $J(p^*) = -(6 + 4\sqrt 2) \approx -11.66$.

A greedy method over action values cannot represent this policy, but a policy-gradient method finds it.

Interactive figure: loading…

Answers to section 1. (1) No. A constant is a state-independent baseline (with $b = -100$), so its term has mean 0. The variance changes: here it grows, because every score is multiplied by a larger number. (2) No. The transition terms vanish from $\nabla \log p_\theta(\tau)$ (step 3). We need only the sampled rewards and the scores of the sampled actions. (3) No. Both deterministic policies never reach the goal; the best policy goes right with probability $2 - \sqrt 2 \approx 0.586$.

4. Worked example #

One step, three actions. $\theta = (0, 0, 0)$, so $\pi = (1/3, 1/3, 1/3)$; the rewards are $r = (1, 0, 2)$.

Step 1: the objective and the exact gradient. $J = \frac13(1 + 0 + 2) = 1$. From the corollary with one state ($d = 1$, $A = r - J$): $\partial J / \partial \theta_a = \pi_a (r_a - J)$, so

Step 2: the REINFORCE samples without a baseline. The sample for action $a$ is $r_a (e_a - \pi)$, drawn with probability $\frac13$:

action	$e_a - \pi$	$r_a$	sample $g_a$	$\|g_a\|^2$
1	$(\frac23, -\frac13, -\frac13)$	1	$(\frac23, -\frac13, -\frac13)$	$\frac69$
2	$(-\frac13, \frac23, -\frac13)$	0	$(0, 0, 0)$	$0$
3	$(-\frac13, -\frac13, \frac23)$	2	$(-\frac23, -\frac23, \frac43)$	$\frac{24}{9}$

Mean: $\frac13\big[(\frac23, -\frac13, -\frac13) + (0, 0, 0) + (-\frac23, -\frac23, \frac43)\big] = \frac13 (0, -1, 1) = (0, -\frac13, \frac13)$ ✓. The total variance (the sum of the variances of the three entries) is $\mathbb E\|g\|^2 - \|\nabla J\|^2 = \frac13 \cdot \frac{30}{9} - \frac29 = \frac{10}{9} - \frac29 = \frac89$.

Step 3: with the baseline $b = J = 1$. The samples become $(r_a - 1)(e_a - \pi)$:

action	$r_a - 1$	sample $g_a$	$\|g_a\|^2$
1	0	$(0, 0, 0)$	0
2	$-1$	$(\frac13, -\frac23, \frac13)$	$\frac69$
3	1	$(-\frac13, -\frac13, \frac23)$	$\frac69$

Mean: $\frac13 (0, -1, 1)$, the same ✓. Total variance: $\frac13 \cdot \frac{12}{9} - \frac29 = \frac49 - \frac29 = \frac29$. The baseline removed $75\%$ of the variance. Here every action has $\|e_a - \pi\|^2 = \frac23$, so the variance-optimal constant $b^*$ of the deeper box equals $J = 1$.

Step 4: one gradient step. With $\eta = 1$, $\theta \leftarrow (0, -\frac13, \frac13)$, so $\pi \approx (0.321, 0.230, 0.448)$ and $J \approx 0.321 + 2 \cdot 0.448 = 1.218$. The best action (reward 2) gained probability, and the worst lost it.

The short corridor. From section 3, $J(0.5) = \frac{1 - 4}{0.25} = -12$, $J(0.1) = \frac{-3.8}{0.09} \approx -42.2$, and the best value is $J(2 - \sqrt 2) \approx -11.66$. A run that starts at $p = 0.1$ must climb a steep slope first; near $p^*$ the curve is flat, so the noise of the estimates matters most there.

5. In code #

NumPy checks the one-step example: the exact gradient against finite differences, and the mean and variance of the REINFORCE samples with and without a baseline.

import numpy as np

def softmax(z):
    e = np.exp(z - z.max())
    return e / e.sum()

r = np.array([1.0, 0.0, 2.0])
theta = np.zeros(3)
pi = softmax(theta)
J = pi @ r
grad = pi * (r - J)                       # exact gradient of J = sum_a pi_a r_a
print(grad.round(4))                      # [ 0.     -0.3333  0.3333]

# check with central finite differences
fd = np.array([(softmax(theta + 1e-6 * e) @ r - softmax(theta - 1e-6 * e) @ r) / 2e-6
               for e in np.eye(3)])
print(np.allclose(fd, grad))              # True

# REINFORCE samples g_a = (r_a - b)(e_a - pi): exact mean and total variance
for b in (0.0, 1.0):
    g = np.array([(r[a] - b) * (np.eye(3)[a] - pi) for a in range(3)])
    mean = pi @ g
    var = pi @ (g**2).sum(axis=1) - mean @ mean
    print(b, mean.round(4), round(var, 4))
# 0.0 [ 0.     -0.3333  0.3333] 0.8889
# 1.0 [ 0.     -0.3333  0.3333] 0.2222

# short corridor: J(p) has its maximum at p* = 2 - sqrt(2)
p = np.linspace(0.01, 0.99, 9801)
J_p = (2 * p - 4) / (p * (1 - p))
print(p[J_p.argmax()].round(4), J_p.max().round(2))   # 0.5858 -11.66

In PyTorch the same REINFORCE loss is one line: loss = -(gamma_t * (G - b).detach() * dist.log_prob(actions)).sum(), where dist is a Categorical(logits=...) or Normal(mu, sigma) distribution.

6. Where it shows up #

• Language models. A language model is a softmax policy over tokens. RLHF fine-tunes it with PPO (a policy gradient with a learned critic and a trust region); RLOO and GRPO use REINFORCE with a baseline built from other samples for the same prompt (lesson 09).
• Robotics and control. Gaussian policies over joint torques, trained with PPO or SAC (lesson 08) in simulators such as MuJoCo.
• Games. AlphaGo's policy network was improved by self-play with a policy gradient before it was used in tree search (lesson 10).
• Non-differentiable objectives. Whenever a sampled discrete choice is scored by a black box: neural architecture search (a controller samples architectures, the reward is validation accuracy), caption and translation models trained on BLEU or CIDEr (self-critical training uses the greedy output's score as the baseline).
• ML for science. Molecule generators that sample a molecule token by token and are fine-tuned with RL toward a property predictor (for example REINVENT, Olivecrona et al., 2017), and chip floor-planning with a policy network trained by PPO (Mirhoseini et al., 2021).

7. Common mistakes #

• Wrong sign. Optimizers minimize. The loss is $-\sum_t (G_t - b)\log \pi_\theta(a_t \mid s_t)$; without the minus sign you make good actions less likely.
• Gradient through the weight. $G_t - b(s_t)$ is a number in the policy loss. If $b$ comes from a network, detach it, or the policy loss also trains the baseline in a wrong way.
• Action-dependent baseline. $b(s, a)$ adds bias. Only functions of the state (or constants) are safe.
• Using the full return for every step. It is unbiased, but rewards before $t$ only add noise. Use the reward-to-go.
• The log-probability of the wrong action. The score must be of the action that was actually sampled and executed, from the policy that sampled it. Recomputing it after the update, or with a different policy, is a different estimator (lesson 09 handles the ratio $\pi_\theta / \pi_{\text{old}}$).
• Reward timing. Here $r_t$ follows $(s_t, a_t)$. With Sutton and Barto's $R_{t+1}$, the reward-to-go is $G_t = R_{t+1} + \gamma R_{t+2} + \dots$; mixing the two conventions shifts every return by one step.
• "A baseline changes the result." It changes only the variance. If an algorithm improves when you add a baseline, it is because the updates are less noisy, not because the target moved.

8. Check yourself #

1. Why do the transition probabilities not appear in the policy gradient?
   Answer

   $\log p_\theta(\tau)$ is a sum of $\log p_0(s_0)$, the log-policy terms and the log-transition terms. Only the log-policy terms depend on $\theta$, so the gradient of the other terms is 0.
2. Prove that a state-dependent baseline adds no bias.
   Answer

   Condition on $s_t$. Then $\mathbb E_{a_t}[\nabla \log \pi(a_t \mid s_t)\, b(s_t)] = b(s_t) \sum_a \pi(a \mid s_t)\nabla \log \pi(a \mid s_t) = b(s_t)\nabla \sum_a \pi(a \mid s_t) = b(s_t) \nabla 1 = 0$. The tower rule gives 0 for the unconditional expectation too.
3. A softmax policy over 2 actions has logits $(0, \ln 3)$. Give $\pi$ and the score of action 1.
   Answer

   $\pi = (1, 3)/4 = (0.25, 0.75)$. The score of action 1 is $e_1 - \pi = (0.75, -0.75)$.
4. Give $\partial J / \partial \theta(s, a)$ for a tabular softmax policy and explain each factor.
   Answer

   $d^\pi(s)\,\pi(a \mid s)\,A^\pi(s, a)$: how much (discounted) time the policy spends in $s$, how likely $a$ is there, and how much better $a$ is than the average action in $s$.
5. Why can a stochastic policy beat every deterministic policy in the short corridor, although lesson 01 says that an optimal deterministic policy always exists?
   Answer

   Lesson 01 assumes that the policy can depend on the true state. In the corridor all cells look the same to the agent, so the policy must use the same action probabilities in every cell. Among such policies, randomness is the only way to sometimes go right and sometimes go left. The best one has $p = 2 - \sqrt 2$.
6. A batch of returns is $(10, 12, 8, 10)$. Which constant baseline would you subtract, and why is it safe?
   Answer

   About their mean, 10 (a running average or a learned $V(s)$ in practice). A constant does not depend on the action, so its term has mean 0; it removes the large common part of the returns, so the weights become $(0, 2, -2, 0)$ and the variance falls. (Using the mean of the same batch adds a tiny bias; leave-one-out means avoid it, see lesson 09.)

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. What is the log-derivative trick?

   Show answer

2. State the policy gradient theorem (expectation form).

   Show answer

3. Why does the policy gradient need no model of the transitions?

   Show answer

4. Write the REINFORCE estimate from one episode, with a baseline.

   Show answer

5. Why does a baseline $b(s)$ add no bias?

   Show answer

6. What is the score $\nabla_\theta \log \pi(a)$ of a softmax policy?

   Show answer

7. What is the score of a Gaussian policy $N(\mu, \sigma^2)$ with respect to $\mu$?

   Show answer

8. Give the exact gradient of a tabular softmax policy.

   Show answer

9. Define the advantage, and give its mean under the policy.

   Show answer

10. Three actions, uniform policy, rewards $(1, 0, 2)$. Total variance of the REINFORCE sample with $b = 0$ and with $b = 1$?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Derive the policy gradient." — Likelihood ratio on trajectories; the dynamics drop out of $\nabla \log p_\theta(\tau)$; causality gives the reward-to-go; conditioning gives $Q^\pi$.
• "Why subtract a baseline? Can it depend on the action?" — Variance. It is unbiased because the score has mean 0 in every state. An action-dependent baseline is biased unless you add a correction term.
• "Why use policy gradients instead of Q-learning?" — Continuous or huge action spaces, stochastic policies (partial observability, games), and smooth changes of the policy. The cost: high variance, on-policy data (low sample efficiency), and only a local optimum.
• "What does REINFORCE do in a language model?" — Sample a response, score it, and raise the log-probability of its tokens in proportion to (score − baseline).
• "Is the usual discounted policy gradient really a gradient?" — Without the $\gamma^t$ factor, no (Nota and Thomas, 2020); it still works well in practice.

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution:

1. softmax_policy and score_function for a tabular softmax policy, checked with finite differences.
2. exact_values, discounted_visitation and exact_policy_gradient on a small 3-state MDP; check $d^\pi\, \pi\, A^\pi$ against finite differences of $J$.
3. reinforce_estimates: one estimate per sampled episode; the mean matches the exact gradient, and the baseline $b = V^\pi$ cuts the variance.
4. REINFORCE on the short corridor: learn $p$ and compare it with $2 - \sqrt 2$.
5. A Gaussian policy on a one-dimensional continuous bandit: learn the mean with the score $(a - \mu)/\sigma^2$.

11. Watch and read #

Policy Gradient Methods | Reinforcement Learning Part 6Mutual Information · 29 min

A clear derivation of the policy gradient theorem and REINFORCE, with good pictures.

An introduction to Policy Gradient methods - Deep Reinforcement LearningArxiv Insights · 20 min

The intuition first: why policy gradients, what the score does, and why the estimates are noisy.

L3 Policy Gradients and Advantage Estimation (Foundations of Deep RL Series)Pieter Abbeel · 41 min

The likelihood-ratio derivation, baselines and the step from REINFORCE to actor-critic.

RL Course by David Silver - Lecture 7: Policy Gradient MethodsGoogle DeepMind · David Silver's RL course · 1 h 34 min

Softmax and Gaussian policies, the policy gradient theorem, REINFORCE and actor-critic.

• Sutton and Barto, Reinforcement Learning: An Introduction, 2nd edition, free PDF: §13.1–13.4 (policy approximation, the short corridor, the policy gradient theorem, REINFORCE and baselines).
• OpenAI Spinning Up, Part 3: Intro to Policy Optimization: the derivation with reward-to-go and baselines, and a short PyTorch implementation.
• Williams, Simple statistical gradient-following algorithms for connectionist reinforcement learning (1992): the original REINFORCE paper.
• Sutton, McAllester, Singh and Mansour, Policy gradient methods for reinforcement learning with function approximation (NeurIPS 1999): the policy gradient theorem.
• Nota and Thomas, Is the policy gradient a gradient? (AAMAS 2020): what happens when the $\gamma^t$ factor is dropped.
• Mei, Xiao, Szepesvári and Schuurmans, On the global convergence rates of softmax policy gradient methods (ICML 2020): why near-deterministic softmax policies learn slowly.