Policy Gradients
Last lecture posed the problem of the course’s final unit. An agent acts, the environment replies with a state and a reward, and the objective is the expected return \(\mathbb{E}[G_0]\), with no labels, no answer key, and no dataset the agent did not generate itself. We also confessed a failure: our best CartPole policy was a hand-coded if statement scoring \(481.5\), and every new environment would need a human to sit down and write a new one. Today the policy becomes a neural network with parameters \(\boldsymbol\theta\), and we do what this course always does with parameters: compute a gradient and climb. But the objective is an expectation over the policy’s own randomness, run through an environment whose physics we cannot differentiate. Ordinary backpropagation, the tool that trained every model since week one, does not even apply.
A Neural Network Policy
Recall from last lecture that a policy \(\pi(a|s)\) is a probability distribution over actions for each state. To learn one, we parameterize it with a policy network. A network with parameters \(\boldsymbol\theta\) (all its weights and biases, flattened into one vector) maps the state to logits \(\mathbf{z}(s) \in \mathbb{R}^{|\mathcal{A}|}\), one score per action, and a softmax turns those scores into a distribution over actions: \[ \pi_{\boldsymbol\theta}(a | s) = [\mathrm{softmax}(\mathbf{z}(s))]_a = \frac{e^{z_a(s)}}{\sum_{l \in \mathcal{A}} e^{z_l(s)}}. \] Look familiar? This is exactly the architecture of every classifier we have built since Logistic Regression: features in, logits out, softmax over the classes, except that the “classes” are now the actions the agent can take. At each step the agent samples \(a_t \sim \pi_{\boldsymbol\theta}(\cdot | s_t)\). Last lecture gave one reason to keep the policy random: sampling means every action gets tried, which is exploration. Today we will see a second: that same randomness is the handle that lets us differentiate the expected return at all.
The Objective and the Obstacle
Running the policy produces a trajectory \(\tau = (s_0, a_0, r_0, s_1, a_1, r_1, \ldots)\), and we write \(G(\tau) = \sum_{t \geq 0} \gamma^t r_t\) for its total discounted return, which is last lecture’s \(G_0\) written as a function of the trajectory that produced it. Rewriting last lecture’s objective with the parameters in place of the policy gives what we want to maximize: \[ \max_{\boldsymbol\theta} \; J(\boldsymbol\theta) = \mathbb{E}_{\tau \sim \pi_{\boldsymbol\theta}}\left[G(\tau)\right]. \] The plan is gradient ascent: \(\boldsymbol\theta^{(k+1)} = \boldsymbol\theta^{(k)} + \alpha \nabla_{\boldsymbol\theta} J(\boldsymbol\theta^{(k)})\), the same update as always with the sign flipped because we are maximizing.
So we need \(\nabla_{\boldsymbol\theta} J\), and here the supervised recipe breaks: \(\boldsymbol\theta\) no longer sits inside a formula the chain rule can march through; it sits inside the probability distribution, changing which trajectories get sampled. Two concrete obstacles:
- Sampling is not differentiable. The step \(a_t \sim \pi_{\boldsymbol\theta}(\cdot|s_t)\) rolls dice; there is no derivative of a die roll with respect to the weights that loaded it.
- The environment is a black box. The next state comes from \(\Pr(s_{t+1} | s_t, a_t)\), which is real physics, or a simulator we can run but not differentiate. Even if we could backpropagate through our own network, the chain would dead-end at the environment.
A three-line identity, derived in class, gets us past both obstacles.
The Log-Derivative Trick
We state the identity for a general distribution, because it is used far beyond reinforcement learning. Let \(p_{\boldsymbol\theta}(x)\) be a probability distribution over outcomes \(x\) whose shape depends on parameters \(\boldsymbol\theta\), and let \(f(x)\) be any function of the outcome that does not depend on \(\boldsymbol\theta\).
Claim: \(\displaystyle \nabla_{\boldsymbol\theta} \, \mathbb{E}_{x \sim p_{\boldsymbol\theta}}\left[f(x)\right] = \mathbb{E}_{x \sim p_{\boldsymbol\theta}}\left[f(x) \, \nabla_{\boldsymbol\theta} \log p_{\boldsymbol\theta}(x)\right]\).
Proof of Claim
The engine of the proof is the log-derivative identity, which is just the chain rule applied to a logarithm: \[ \nabla_{\boldsymbol\theta} \log p_{\boldsymbol\theta}(x) = \frac{\nabla_{\boldsymbol\theta} \, p_{\boldsymbol\theta}(x)}{p_{\boldsymbol\theta}(x)} \qquad\Longleftrightarrow\qquad \nabla_{\boldsymbol\theta} \, p_{\boldsymbol\theta}(x) = p_{\boldsymbol\theta}(x) \, \nabla_{\boldsymbol\theta} \log p_{\boldsymbol\theta}(x). \] Read the right-hand version as a trade: the gradient of a probability equals the probability itself times the gradient of its log.
Now write the expectation as a sum over outcomes and compute: \[ \begin{align*} \nabla_{\boldsymbol\theta} \, \mathbb{E}_{x \sim p_{\boldsymbol\theta}}\left[f(x)\right] &= \nabla_{\boldsymbol\theta} \sum_{x} p_{\boldsymbol\theta}(x) \, f(x) \\&= \sum_{x} f(x) \, \nabla_{\boldsymbol\theta} \, p_{\boldsymbol\theta}(x) \\&= \sum_{x} f(x) \, p_{\boldsymbol\theta}(x) \, \nabla_{\boldsymbol\theta} \log p_{\boldsymbol\theta}(x) \\&= \mathbb{E}_{x \sim p_{\boldsymbol\theta}}\left[f(x) \, \nabla_{\boldsymbol\theta} \log p_{\boldsymbol\theta}(x)\right]. \end{align*} \] Here, we used the definition of expectation, then moved the gradient inside the sum (gradients are linear, so they pass through finite sums; for infinite sums or integrals, replace \(\sum_x\) with \(\int dx\) and every step survives, though the swap then needs mild regularity conditions we will not worry about), then applied the log-derivative identity to \(\nabla_{\boldsymbol\theta} \, p_{\boldsymbol\theta}(x)\), and finally recognized the sum of \(p_{\boldsymbol\theta}(x)\) times a quantity as an expectation again.On the left is a gradient of an expectation, the thing we cannot compute, because the expectation sums over every possible outcome. On the right is an expectation of a quantity we can compute for any single sampled outcome: \(f(x)\) is the function evaluated at the sample, and \(\nabla_{\boldsymbol\theta} \log p_{\boldsymbol\theta}(x)\) is the gradient of our own model’s log-probability, which backpropagation delivers. Moving the gradient inside the expectation is what puts it within reach of sampling.
The REINFORCE Gradient
Now apply the claim to reinforcement learning: the outcome \(x\) is a whole trajectory \(\tau\), the function is the return \(f(\tau) = G(\tau)\), and the distribution \(p_{\boldsymbol\theta}(\tau)\) is the probability that running policy \(\pi_{\boldsymbol\theta}\) in the environment produces exactly the trajectory \(\tau\). Substituting these into the claim gives the gradient of the objective: \[ \nabla_{\boldsymbol\theta} J(\boldsymbol\theta) = \mathbb{E}_{\tau \sim \pi_{\boldsymbol\theta}}\left[G(\tau) \, \nabla_{\boldsymbol\theta} \log p_{\boldsymbol\theta}(\tau)\right]. \] The return \(G(\tau)\) is a number we read off the episode, so the only piece left to understand is \(\log p_{\boldsymbol\theta}(\tau)\). A trajectory is built step by step (start somewhere, act, transition, act, transition), so by the chain rule of probability and the Markov property, its probability is a product: \[ p_{\boldsymbol\theta}(\tau) = \Pr(s_0) \prod_{t \geq 0} \pi_{\boldsymbol\theta}(a_t | s_t) \Pr(s_{t+1} | s_t, a_t), \] one factor per decision the policy made and one per roll of the environment’s dice. Taking the logarithm turns the product into a sum, and the gradient then deletes every term that does not contain \(\boldsymbol\theta\): \[ \begin{align*} \nabla_{\boldsymbol\theta} \log p_{\boldsymbol\theta}(\tau) &= \nabla_{\boldsymbol\theta} \left[\log \Pr(s_0) + \sum_{t \geq 0} \log \pi_{\boldsymbol\theta}(a_t|s_t) + \sum_{t \geq 0} \log \Pr(s_{t+1}|s_t, a_t)\right] \\&= \sum_{t \geq 0} \nabla_{\boldsymbol\theta} \log \pi_{\boldsymbol\theta}(a_t | s_t), \end{align*} \] since \(\log \Pr(s_0)\) and each \(\log \Pr(s_{t+1}|s_t,a_t)\) are constants as far as \(\boldsymbol\theta\) is concerned. The last line settles the black-box obstacle. The environment’s dynamics entered as additive constants in log-space and vanished under the gradient, so we never need the environment’s derivative, only our own policy’s, and backpropagation computes that in one pass. Substituting back gives the policy gradient: \[ \nabla_{\boldsymbol\theta} J(\boldsymbol\theta) = \mathbb{E}_{\tau \sim \pi_{\boldsymbol\theta}}\left[G(\tau) \sum_{t \geq 0} \nabla_{\boldsymbol\theta} \log \pi_{\boldsymbol\theta}(a_t | s_t)\right]. \] Gradient ascent on this expression is the REINFORCE algorithm of Ronald Williams (1992). The expression is still an expectation over every trajectory the policy might produce, and we can no more enumerate those than we could before.
Monte Carlo Estimation
We cannot compute the expectation, but we can sample from it, which is precisely what lecture two’s Monte Carlo estimator was built for: sample, evaluate, average. Run \(N\) episodes with the current policy to get trajectories \(\tau^{(1)}, \ldots, \tau^{(N)}\), and average the quantity inside the expectation: \[ \widehat{\nabla J} = \frac{1}{N} \sum_{i=1}^N G(\tau^{(i)}) \sum_{t \geq 0} \nabla_{\boldsymbol\theta} \log \pi_{\boldsymbol\theta}\big(a_t^{(i)} \big| s_t^{(i)}\big). \] By the same argument that made the sample mean unbiased in lecture two, \(\mathbb{E}[\widehat{\nabla J}] = \nabla_{\boldsymbol\theta} J(\boldsymbol\theta)\) for every \(N\), even \(N = 1\), a single episode. Stochastic gradient descent made the same move, estimating the full-dataset gradient from a batch; REINFORCE swaps batches for episodes, and one full loop of the algorithm fits in four lines:
- Run one episode with the current policy, recording states, actions, and rewards.
- Compute the return \(G(\tau)\) back to front with last lecture’s recursion \(G_t = r_t + \gamma G_{t+1}\), one multiply and add per step.
- Compute \(\widehat{\nabla J} = G(\tau) \sum_t \nabla_{\boldsymbol\theta} \log \pi_{\boldsymbol\theta}(a_t|s_t)\) by backpropagation.
- Ascend: \(\boldsymbol\theta \leftarrow \boldsymbol\theta + \alpha \widehat{\nabla J}\), and go back to step 1.
Each update costs one episode of interaction plus one backward pass, so the agent learns from experience it generated seconds ago. Step 3 is the only line we have not made concrete: what does \(\nabla_{\boldsymbol\theta} \log \pi_{\boldsymbol\theta}(a_t|s_t)\) actually look like?
The Softmax Policy’s Gradient
For our softmax policy, that per-step gradient has a closed form we have already met. Recall the Logistic Regression lecture’s central computation: the cross-entropy loss \(-\log [\mathrm{softmax}(\mathbf{z})]_c\) has gradient \(\mathbf{p} - \mathbf{y}\) with respect to the logits, “softmax minus one-hot”. That lecture promised the same shape would come back, scaled by a reward, when we trained an agent’s policy. It does.
Claim: For a softmax policy with logits \(\mathbf{z}(s_t)\), probabilities \(\mathbf{p}_t = \mathrm{softmax}(\mathbf{z}(s_t))\), and sampled action \(a_t\) with one-hot vector \(\mathbf{y}_t\), the REINFORCE gradient with respect to the logits at step \(t\) is \[ G(\tau)\,\nabla_{\mathbf{z}} \log \pi_{\boldsymbol\theta}(a_t | s_t) = G(\tau) \left(\mathbf{y}_t - \mathbf{p}_t\right). \]
Proof of Claim
Write the log-probability of the sampled action directly from the softmax definition: \[ \log \pi_{\boldsymbol\theta}(a_t | s_t) = \log \frac{e^{z_{a_t}}}{\sum_l e^{z_l}} = z_{a_t} - \log \sum_{l} e^{z_l}. \] Differentiate with respect to logit \(z_j\), one term at a time. The first term \(z_{a_t}\) contributes \(1\) if \(j = a_t\) and \(0\) otherwise, which is exactly the \(j\)-th entry \([\mathbf{y}_t]_j\) of the one-hot vector. The second term is the log-sum-exp we differentiated in the Logistic Regression lecture, and the chain rule gives: \[ \frac{\partial}{\partial z_j} \log \sum_{l} e^{z_l} = \frac{1}{\sum_{l} e^{z_l}} \cdot \frac{\partial}{\partial z_j} \sum_{l} e^{z_l} = \frac{e^{z_j}}{\sum_{l} e^{z_l}} = [\mathbf{p}_t]_j. \] Subtracting the second contribution from the first gives the \(j\)-th entry of the gradient: \[ \frac{\partial}{\partial z_j} \log \pi_{\boldsymbol\theta}(a_t | s_t) = [\mathbf{y}_t]_j - [\mathbf{p}_t]_j. \] Stacking these entries over \(j\) gives \(\nabla_{\mathbf{z}} \log \pi_{\boldsymbol\theta}(a_t|s_t) = \mathbf{y}_t - \mathbf{p}_t\), and multiplying by the return gives the claim. (The chain rule then carries this gradient from the logits back through the network’s weights, exactly as backpropagation always has.)So REINFORCE’s per-step update, seen from the logits, is the familiar cross-entropy gradient with the sampled action playing the role of the label, scaled by the return: ascending \(G(\tau)(\mathbf{y}_t - \mathbf{p}_t)\) is descending the weighted cross-entropy \(G(\tau) \cdot \left(-\log \pi_{\boldsymbol\theta}(a_t|s_t)\right)\). A positive return pushes probability toward the actions the agent took, in proportion to how good the episode was; a negative return pushes probability away from them.
One sentence now describes the whole algorithm, and it closes the maximum-likelihood arc that has run since the Mean Squared Error lecture: REINFORCE is return-weighted maximum likelihood on the agent’s own behavior. Maximum likelihood up-weights the log-probability of the observed data, every observation equally; REINFORCE up-weights the log-probability of the actions the agent observed itself taking, each weighted by the return that followed.
Today’s demo runs exactly this loop on CartPole, and the trained policy clears the pole far more often than the untrained one. Its learning curve is jagged rather than smooth, though, lurching through progress, collapse, and recovery, a symptom of the estimator we derived, not of a poor choice of hyperparameters.
Why the Variance Is High
Our single-episode gradient estimate is one noisy scalar, \(G(\tau)\), multiplying the entire vector \(\sum_t \nabla_{\boldsymbol\theta} \log \pi_{\boldsymbol\theta}(a_t|s_t)\), so every coordinate is scaled by the same roll of the dice. And \(G(\tau)\) is genuinely noisy. The demo confirms this by freezing a trained policy and replaying it many times: a single episode’s return swings across nearly its whole possible range, and discounting only partly compresses that swing before it multiplies the gradient.
There is a second problem, and it is the one Problem 24 is built around. CartPole’s rewards are all positive, so every return is positive, and every update pushes probability toward whatever the agent just did, good episodes and bad ones alike, differing only in how hard they push. An action is favored because it was followed by a bigger number, not because it beat the alternatives, and the relative differences that carry the real signal ride on top of the returns’ common bulk. The demo makes this concrete: a tenfold increase in the learning rate collapses training to worse than the random policy it started from, because a single noisy update becomes large enough to destroy the network. The variance is the reason the learning rate had to be timid, and timid learning rates are why REINFORCE needs far more episodes than an if statement needed lines.
Problem 24 puts this under a microscope, on the smallest environment where every quantity can be computed exactly by enumeration: a two-armed bandit with one state, two actions, and one softmax parameter. You will compute the exact gradient, verify that the REINFORCE estimator is unbiased, and then shift both rewards by the same constant \(c\), which changes nothing about which arm is better, and watch the estimator’s variance grow like \(c^2\) while the gradient stands still.
Looking Forward
Today we got what last lecture lacked: learning. The cost was variance, and the lesson to carry forward is that an unbiased estimator is only as good as its variance, so the next question is always what we can subtract from it. The final lecture of the course asks exactly that. We will subtract a baseline from the return before it multiplies the gradient, prove that this leaves the estimator unbiased for any baseline, and then choose the baseline that kills the most variance, which turns out to be exactly the control-variate coefficient you derived on Problem 2, thirteen weeks ago. The variance-reduction thread that opened the course closes it.