Policy Gradients

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. Our best CartPole policy was a hand-coded if statement scoring \(481.5\), and every new environment would need a new rule. We instead parameterize the policy with a neural network and optimize its expected return. The objective is an expectation over the policy’s own randomness, run through an environment whose physics we cannot differentiate directly.

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)}}. \] This is the classifier architecture used since Logistic Regression: features in, logits out, and a softmax over classes, with actions taking the place of classes. 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. The same randomness also lets us differentiate the expected return.

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.

We need \(\nabla_{\boldsymbol\theta} J\), but \(\boldsymbol\theta\) changes which trajectories are sampled rather than appearing inside a deterministic computation. Sampling \(a_t \sim \pi_{\boldsymbol\theta}(\cdot|s_t)\) is not differentiable, and the next state comes from environment dynamics \(\Pr(s_{t+1} | s_t, a_t)\) that may also be unavailable for differentiation. The log-derivative identity handles both cases.

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 proof uses the log-derivative identity, 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). \] The right-hand form expresses the gradient of a probability as the probability 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 environment’s dynamics enter as additive constants in log-space and vanish under the gradient. We need only the derivative of the policy, which backpropagation computes 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. The next section derives \(\nabla_{\boldsymbol\theta} \log \pi_{\boldsymbol\theta}(a_t|s_t)\) for a softmax policy.

The Softmax Policy’s Gradient

For our softmax policy, that per-step gradient has a closed form we have already met. The Logistic Regression lecture showed that the cross-entropy loss \(-\log [\mathrm{softmax}(\mathbf{z})]_c\) has gradient \(\mathbf{p} - \mathbf{y}\) with respect to the logits. A policy gradient uses the negative of this expression, scaled by the return.

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.

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.

The demo runs this loop on CartPole and shows the variability of its learning curve.

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. Replaying a fixed trained policy shows that a single episode’s return can span nearly its whole possible range. Discounting only partly reduces that variation before the return multiplies the gradient.

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. With high-variance updates, a large learning rate can damage the policy, while a small learning rate requires more episodes.

Problem 24 audits reward-shift invariance in a two-armed bandit. The exact objective ignores a common shift, but the one-sample REINFORCE noise does not; optimizing the shift reveals both removable offset noise and irreducible random-reward noise.

The next lecture subtracts a baseline from the return before it multiplies the gradient. This leaves the estimator unbiased, and the variance-minimizing baseline is the control-variate coefficient from Problem 2.