Policy Gradients and REINFORCE

The score-function gradient turned "differentiate an expectation" into "sample, then weight the score by the value." Policy gradients apply that same trick to control: the distribution is a policy $\pi_\theta(a\mid s)$ over actions, and the value being maximized is the expected return. REINFORCE is the estimator that falls out, and a baseline is the one line that makes it usable.

The policy gradient theorem

Let a policy $\pi_\theta(a\mid s)$ induce a trajectory distribution with expected return $J(\theta) = \mathbb{E}_{\pi_\theta}\!\big[\sum_t \gamma^t r_t\big]$. The policy gradient theorem says the gradient is an expectation over the states the policy visits and the actions it takes,¹

$$\nabla_\theta J(\theta) = \mathbb{E}_{\pi_\theta}\!\left[\, \nabla_\theta \log \pi_\theta(a\mid s)\; Q^{\pi}(s,a) \right].$$

The environment dynamics never appear in the gradient. Only the score $\nabla_\theta \log \pi_\theta(a\mid s)$, the part we control, is weighted by $Q^\pi(s,a)$, the value of taking that action. This is the score-function identity over actions instead of over an abstract sample $x$.

REINFORCE: the sampled-return estimator

$Q^\pi$ is unknown, but along a sampled trajectory the return $G_t = \sum_{k\ge t}\gamma^{k-t} r_k$ is an unbiased sample of it. Substituting $G_t$ for $Q^\pi(s_t,a_t)$ gives the Monte Carlo estimator and its update,²

$$\nabla_\theta J \approx \sum_t G_t\, \nabla_\theta \log \pi_\theta(a_t\mid s_t), \qquad \theta \leftarrow \theta + \alpha\, G_t\, \nabla_\theta \log \pi_\theta(a_t\mid s_t).$$

Actions that preceded a high return get their log-probability pushed up. It is unbiased but noisy: the return $G_t$ carries the variance of the whole future.

Baselines cut variance for free

Subtract any function $b(s)$ that does not depend on the action. The expected update is unchanged, because the expected score is zero:

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

So $\mathbb{E}\big[b(s)\,\nabla_\theta \log \pi_\theta\big] = 0$, and replacing $G_t$ with $G_t - b(s_t)$ keeps the estimator unbiased while changing its variance. Choosing $b(s)\approx \mathbb{E}[G\mid s]$ centers the weight near zero, shrinking $\mathbb{E}\big[(G-b)^2\,\lVert\nabla\log\pi\rVert^2\big]$. The mean return is the simplest such baseline and already removes most of the noise.³

Worked example: softmax over K arms

Drop the state. With logits $\theta\in\mathbb{R}^K$ the policy is $\pi = \operatorname{softmax}(\theta)$, and the score of a chosen arm $a$ is the one-hot vector minus the probability vector,

$$\nabla_\theta \log \pi(a) = e_a - \pi .$$

Start uniform, $\theta = 0$ so $\pi = (\tfrac13,\tfrac13,\tfrac13)$ for $K=3$. Sample arm $a = 2$ and observe reward $G = 2$. The score is

$$\nabla_\theta \log \pi(2) = \begin{bmatrix}0\\1\\0\end{bmatrix} - \begin{bmatrix}1/3\\1/3\\1/3\end{bmatrix} = \begin{bmatrix}-1/3\\2/3\\-1/3\end{bmatrix}.$$

With no baseline and $\alpha = 0.1$, the REINFORCE step is $\Delta\theta = \alpha\,G\,(e_2 - \pi)$:

$$\Delta\theta = 0.1 \cdot 2 \cdot \begin{bmatrix}-1/3\\2/3\\-1/3\end{bmatrix} = \begin{bmatrix}-0.067\\0.133\\-0.067\end{bmatrix} \;\Rightarrow\; \pi \to (0.310,\, 0.379,\, 0.310).$$

The probability of arm 2 rose from $0.333$ to $0.379$. Now subtract a baseline $b = 1.0$, the average reward seen so far. The weight becomes $G - b = 1$:

$$\Delta\theta = 0.1 \cdot (2 - 1) \cdot \begin{bmatrix}-1/3\\2/3\\-1/3\end{bmatrix} = \begin{bmatrix}-0.033\\0.067\\-0.033\end{bmatrix} \;\Rightarrow\; \pi \to (0.322,\, 0.356,\, 0.322).$$

Same direction, smaller step, because arm 2 only beat the average by one. The crucial case is an arm that pays exactly the average, $G = b$: its update is $0$. Below-average arms are pushed down. Without a baseline every arm here has positive reward, so every sampled arm is pushed up, and the policy has to separate the good from the mediocre out of noise alone. Centering by $b$ turns "how large is the reward" into "how much better than usual," which is the signal that actually matters.

REINFORCE on a 4-armed banditinteractive

steps

0

π(a3) best

25.0%

avg reward

0.000

Var[∇ estimate]

b = 00.000

b = mean0.000

ratio0%

α = 0.10baseline: on (b = mean)

PauseStepReset

Turn the baseline off and the probability bars still march toward the green arm, but the variance bar jumps and the reward curve gets choppier: the update is being knocked around by the reward's mean level rather than its differences. This is the whole story of modern policy-gradient methods in miniature. Replace the running-mean baseline with a learned value function $V_\phi(s)$ and the weight $G_t - V_\phi(s_t)$ becomes the advantage, the object at the center of actor-critic and PPO.

Footnotes

1. Sutton, McAllester, Singh, Mansour, Policy Gradient Methods for Reinforcement Learning with Function Approximation (2000). ↩

2. Williams, Simple statistical gradient-following algorithms for connectionist reinforcement learning (1992). ↩

3. Sutton & Barto, Reinforcement Learning: An Introduction (2018). ↩