Policy gradients and the log-derivative trick

Worked solutions

Problem 1

1. $J(\theta) = \sum_a p_\theta(a)\,r(a)$.The expectation over a finite set is a weighted sum.
2. $\nabla_\theta J = \sum_a r(a)\,\nabla_\theta p_\theta(a)$.Differentiate term by term; $r(a)$ does not depend on $\theta$.
3. $\nabla_\theta p_\theta(a) = p_\theta(a)\,\nabla_\theta\log p_\theta(a)$.The chain rule gives $\nabla\log p = \nabla p/p$; multiply through by $p > 0$.
4. $\nabla_\theta J = \sum_a p_\theta(a)\,r(a)\,\nabla_\theta\log p_\theta(a) = \mathbb{E}_{a\sim p_\theta}\big[r(a)\,\nabla_\theta\log p_\theta(a)\big]$.Substitute; a sum weighted by $p_\theta$ is an expectation under $p_\theta$.
5. $\nabla_\theta J = \mathbb{E}_{a\sim p_\theta}\big[r(a)\,\nabla_\theta\log p_\theta(a)\big]$The gradient of an expectation whose distribution depends on $\theta$ is again an expectation under that distribution, so it can be estimated by sampling: draw $a$, compute $r(a)\nabla_\theta\log p_\theta(a)$. It uses only the value $r(a)$, never a derivative of $r$, so the reward can be a game score or any black box, and $a$ can be discrete; that is what separates it from the ELBO page's reparameterisation, which needs $\nabla r$. For a continuous $a$ the sum becomes an integral and the argument is unchanged (Problem 6). Mistake 1 is the gradient moved inside instead.

Problem 2

1. $\mathbb{E}[\nabla_\theta\log p_\theta(a)] = \sum_a p_\theta(a)\,\nabla_\theta\log p_\theta(a) = \sum_a\nabla_\theta p_\theta(a)$.Problem 1's identity read from right to left.
2. $\sum_a\nabla_\theta p_\theta(a) = \nabla_\theta\sum_a p_\theta(a) = \nabla_\theta 1 = 0$.A finite sum of gradients is the gradient of the sum, and the probabilities sum to $1$ for every $\theta$.
3. $\mathbb{E}\big[(r(a) - b)\,\nabla_\theta\log p_\theta(a)\big] = \mathbb{E}\big[r(a)\,\nabla_\theta\log p_\theta(a)\big] - b\,\mathbb{E}[\nabla_\theta\log p_\theta(a)] = \nabla_\theta J - 0$.Linearity, then Problem 1 and step 2.
4. $\mathbb{E}_{a\sim p_\theta}[\nabla_\theta\log p_\theta(a)] = 0$, so $\mathbb{E}\big[(r(a) - b)\,\nabla_\theta\log p_\theta(a)\big] = \nabla_\theta J$ for every constant $b$The score has mean zero because probabilities sum to one, differentiated. For the softmax, $\sum_a p_a(e_a - p) = p - p = 0$ directly. A baseline leaves the mean of the estimator alone and changes only its variance, which Problem 5 minimises; it must not depend on the action drawn (Mistake 3).

Problem 3

1. $\log p_\theta(a) = \theta_a - \log\sum_b e^{\theta_b}$.The log of the softmax.
2. $\dfrac{\partial\log p_\theta(a)}{\partial\theta_k} = \delta_{ak} - \dfrac{e^{\theta_k}}{\sum_b e^{\theta_b}} = \delta_{ak} - p_k$, so $\nabla_\theta\log p_\theta(a) = e_a - p$.The derivative of the log-sum-exp is the softmax, as on the softmax page.
3. $\nabla_\theta J = \sum_a p_a r_a(e_a - p) = \sum_a p_a r_a e_a - \Big(\sum_a p_a r_a\Big)p = p\odot r - Jp$.Problem 1 with step 2; $\sum_a p_a r_a e_a$ is the vector with entries $p_a r_a$.
4. $p\odot r - Jp = p\odot(r - J\mathbf{1})$.Entry $k$ is $p_k r_k - Jp_k = p_k(r_k - J)$.
5. $\mathbf{1}^\top\nabla_\theta J = \sum_k p_k r_k - J\sum_k p_k = J - J = 0$.$\sum_k p_k = 1$.
6. $\nabla_\theta\log p_\theta(a) = e_a - p$; $\nabla_\theta J = p\odot(r - J\mathbf{1})$, whose entries sum to zeroAn action's logit rises when its reward is above the current average $J$ and falls when below, at a rate proportional to its probability, so a rarely chosen good action moves slowly, which is the exploration problem in one line. The entries must sum to zero because adding a constant to every logit leaves the softmax unchanged (the softmax page's shift invariance): $J$ is constant along $\mathbf{1}$, so its gradient is orthogonal to $\mathbf{1}$. Mistake 5 breaks this.

Problem 4

1. $p = (\tfrac13, \tfrac13, \tfrac13)$ and $J = \tfrac13(1 + 0 - 1) = 0$.The softmax of equal logits is uniform.
2. $\nabla_\theta J = p\odot(r - J\mathbf{1}) = \tfrac13(1, 0, -1)$.Problem 3 with $J = 0$.
3. $\theta_1 = \theta_0 + \eta\nabla_\theta J = (\tfrac13, 0, -\tfrac13)$.Gradient ascent on $J$ with $\eta = 1$.
4. $e^{1/3} \approx 1.3956$, $e^0 = 1$, $e^{-1/3} \approx 0.7165$, summing to $3.1121$, so $p \approx (0.4484, 0.3213, 0.2302)$.The softmax at $\theta_1$.
5. $J \approx 0.4484 - 0.2302 = 0.2182$.$p^\top r$ with $r_2 = 0$.
6. At $\theta_0$: $J = 0$ and $\nabla_\theta J = (\tfrac13, 0, -\tfrac13)$; $\theta_1 = (\tfrac13, 0, -\tfrac13)$ with $p \approx (0.4484, 0.3213, 0.2302)$ and $J \approx 0.2182$One step moved probability from the worst action to the best and raised $J$ from $0$ to $0.218$; the middle action, whose reward equals the mean, did not move. Repeating the steps sends $p \to e_1$ and $J \to 1$, ever more slowly, because the entry for action $3$ is $p_3(r_3 - J)$ and $p_3 \to 0$. The exact gradient here is what a REINFORCE step estimates from a sample: drawing $a = 1$ gives $(1 - 0)(e_1 - p) = (\tfrac23, -\tfrac13, -\tfrac13)$, and the $p$-weighted average of the three possible samples is the gradient above.

Problem 5

1. $\mathbb{E}\|\hat g_b\|^2 = \sum_a p_a\,(r_a - b)^2\,\|e_a - p\|^2$.The expectation over the finite set; the scalar $(r_a - b)$ comes out of the squared norm as its square.
2. $\operatorname{Var}(\hat g_b) = \mathbb{E}\|\hat g_b\|^2 - \|\nabla_\theta J\|^2$, and $\nabla_\theta J$ does not depend on $b$, so minimising the variance over $b$ is minimising $\mathbb{E}\|\hat g_b\|^2$.Problem 2: the mean of $\hat g_b$ is $\nabla_\theta J$ for every $b$.
3. $\dfrac{d}{db}\mathbb{E}\|\hat g_b\|^2 = -2\sum_a p_a(r_a - b)\|e_a - p\|^2 = 0$, so $b\sum_a p_a\|e_a - p\|^2 = \sum_a p_a\|e_a - p\|^2\,r_a$.Differentiate the square and set to zero.
4. $b^* = \dfrac{\sum_a p_a\|e_a - p\|^2\,r_a}{\sum_a p_a\|e_a - p\|^2}$, and the second derivative $2\sum_a p_a\|e_a - p\|^2 > 0$ makes it the minimum.A convex quadratic in $b$ has its minimum at its stationary point.
5. For uniform $p$, $\|e_a - p\|^2 = \big(1 - \tfrac1K\big)^2 + (K - 1)\tfrac{1}{K^2} = \dfrac{K - 1}{K}$ for every $a$, so the weights are equal and $b^* = \sum_a p_a r_a = J$.One entry of $e_a - p$ is $1 - \tfrac1K$ and the other $K - 1$ are $-\tfrac1K$; $(K - 1)^2 + (K - 1) = (K - 1)K$.
6. $\mathbb{E}\|\hat g_b\|^2 = \sum_a p_a(r_a - b)^2\|e_a - p\|^2$; $b^* = \dfrac{\sum_a p_a\|e_a - p\|^2\,r_a}{\sum_a p_a\|e_a - p\|^2}$; for uniform $p$, $b^* = J$The best baseline is a weighted mean of the rewards, weighted towards actions with large scores. The mean reward $J$, or a running estimate of it, which in reinforcement learning is the value function, is the usual choice: exactly optimal for a uniform policy and close otherwise. At Problem 4's $\theta_0$, $b = J = 0$ gives $\mathbb{E}\|\hat g\|^2 = \tfrac49$ against $\tfrac{10}{9}$ for $b = 1$, with the same mean.

Problem 6

1. $\log p(a) = -\tfrac12\log(2\pi) - \log\sigma - \dfrac{(a - \mu)^2}{2\sigma^2}$.The Gaussian log-density.
2. $\partial_\mu\log p(a) = \dfrac{a - \mu}{\sigma^2}$ and $\partial_\sigma\log p(a) = -\dfrac1\sigma + \dfrac{(a - \mu)^2}{\sigma^3} = \dfrac{(a - \mu)^2 - \sigma^2}{\sigma^3}$.Chain rule on $(a - \mu)^2$, whose inner derivative in $\mu$ is $-1$; power rule on $\sigma^{-2}$ and the derivative of $-\log\sigma$.
3. $\partial_{\log\sigma}\log p(a) = \sigma\,\partial_\sigma\log p(a) = \dfrac{(a - \mu)^2}{\sigma^2} - 1$.$d/d(\log\sigma) = \sigma\,d/d\sigma$, since $d\sigma/d\log\sigma = \sigma$.
4. $J = \mathbb{E}[-(a - a^*)^2] = -\big(\operatorname{Var}(a) + (\mathbb{E}[a] - a^*)^2\big) = -\big((\mu - a^*)^2 + \sigma^2\big)$, so $\partial J/\partial\mu = -2(\mu - a^*)$ and $\partial J/\partial\sigma = -2\sigma$.The variance page's $\mathbb{E}[u^2] = \operatorname{Var}(u) + (\mathbb{E}[u])^2$ with $u = a - a^*$.
5. With $a = \mu + \sigma\varepsilon$: $\mathbb{E}\Big[r(a)\dfrac{a - \mu}{\sigma^2}\Big] = -\dfrac{1}{\sigma}\mathbb{E}\big[\big((\mu - a^*) + \sigma\varepsilon\big)^2\varepsilon\big] = -\dfrac1\sigma\mathbb{E}\big[(\mu - a^*)^2\varepsilon + 2(\mu - a^*)\sigma\varepsilon^2 + \sigma^2\varepsilon^3\big] = -2(\mu - a^*)$.$(a - \mu)/\sigma^2 = \varepsilon/\sigma$; expand the square; $\mathbb{E}[\varepsilon] = \mathbb{E}[\varepsilon^3] = 0$ and $\mathbb{E}[\varepsilon^2] = 1$.
6. $\partial_\mu\log p(a) = \dfrac{a - \mu}{\sigma^2}$, $\partial_\sigma\log p(a) = \dfrac{(a - \mu)^2 - \sigma^2}{\sigma^3}$, $\partial_{\log\sigma}\log p(a) = \dfrac{(a - \mu)^2}{\sigma^2} - 1$; $J = -\big((\mu - a^*)^2 + \sigma^2\big)$ and $\mathbb{E}\big[r(a)\,\partial_\mu\log p(a)\big] = -2(\mu - a^*) = \partial J/\partial\mu$The score in $\mu$ says that an action which landed above the mean pushes the mean up, weighted by its reward, and the score in $\sigma$ widens the policy when actions farther than $\sigma$ from the mean are rewarded; the check confirms $\mathbb{E}[r\,\partial_\sigma\log p] = -2\sigma$ too. The reparameterised route of the ELBO page would use $r'(a) = -2(a - a^*)$, the derivative of the reward, which an agent usually cannot evaluate; the log-derivative trick needs only values of $r$, at the cost of variance ($30$ against $4$ in the ELBO page's Problem 6). Parameterising by $\log\sigma$ keeps $\sigma > 0$ and makes the score scale-free. Mistake 2 differentiates in $a$ instead of $\mu$.

Problem 7

1. $p_\theta(\tau) = p(s_0)\prod_{t=0}^{T-1}\pi_\theta(a_t \mid s_t)\prod_{t=0}^{T-2}p(s_{t+1} \mid s_t, a_t)$.The chain rule of probability: each action depends on the current state through the policy and each next state on the current state and action through the dynamics.
2. $\log p_\theta(\tau) = \log p(s_0) + \sum_t\log\pi_\theta(a_t \mid s_t) + \sum_t\log p(s_{t+1} \mid s_t, a_t)$.The log of a product.
3. $\nabla_\theta\log p_\theta(\tau) = \sum_t\nabla_\theta\log\pi_\theta(a_t \mid s_t)$.Only the policy factors contain $\theta$; the start distribution and the dynamics do not.
4. $\nabla_\theta J = \mathbb{E}_\tau\big[R(\tau)\,\nabla_\theta\log p_\theta(\tau)\big] = \mathbb{E}_\tau\Big[R(\tau)\sum_t\nabla_\theta\log\pi_\theta(a_t \mid s_t)\Big]$.Problem 1 with trajectories as the outcomes and $R(\tau)$ as the reward, then step 3.
5. $\nabla_\theta\log p_\theta(\tau) = \sum_t\nabla_\theta\log\pi_\theta(a_t \mid s_t)$; $\nabla_\theta J = \mathbb{E}_\tau\Big[R(\tau)\sum_t\nabla_\theta\log\pi_\theta(a_t \mid s_t)\Big]$The dynamics drop out: the estimator never needs $p(s' \mid s, a)$, which is why policy gradients are model-free. The check enumerates every trajectory of a two-state, two-action, three-step environment with a tabular softmax policy and compares the expectation with finite differences of $J$. This is REINFORCE; it weights every action by the whole return, including actions taken after a reward was already collected, which Problem 8 removes.

Problem 8

1. Let $h_t = (s_0, a_0, \dots, s_{t-1}, a_{t-1}, s_t)$ be the history up to the state at time $t$. Given $h_t$, $a_t \sim \pi_\theta(\cdot \mid s_t)$, so $\mathbb{E}\big[\nabla_\theta\log\pi_\theta(a_t \mid s_t)\,\big|\,h_t\big] = \sum_a\pi_\theta(a \mid s_t)\nabla_\theta\log\pi_\theta(a \mid s_t) = 0$.Problem 2 applied to the distribution $\pi_\theta(\cdot \mid s_t)$.
2. For $t' < t$, $r_{t'} = r(s_{t'}, a_{t'})$ is a function of $h_t$, so $\mathbb{E}\big[r_{t'}\nabla_\theta\log\pi_\theta(a_t \mid s_t)\big] = \mathbb{E}\Big[r_{t'}\,\mathbb{E}\big[\nabla_\theta\log\pi_\theta(a_t \mid s_t)\,\big|\,h_t\big]\Big] = \mathbb{E}[r_{t'}\cdot 0] = 0$.Iterated expectations: condition on $h_t$, take $r_{t'}$ out of the inner expectation because it is fixed by $h_t$, then step 1.
3. $R(\tau)\sum_t\nabla_\theta\log\pi_\theta(a_t \mid s_t) = \sum_t\nabla_\theta\log\pi_\theta(a_t \mid s_t)\Big(\sum_{t' < t}r_{t'} + \sum_{t' \ge t}r_{t'}\Big)$.Split the return at $t$ inside each term.
4. Taking expectations, every past term $r_{t'}$ with $t' < t$ contributes $0$ by step 2, leaving $\nabla_\theta J = \mathbb{E}_\tau\Big[\sum_t\nabla_\theta\log\pi_\theta(a_t \mid s_t)\sum_{t' \ge t}r_{t'}\Big]$.Problem 7 and linearity.
5. $\mathbb{E}_\tau\big[r_{t'}\,\nabla_\theta\log\pi_\theta(a_t \mid s_t)\big] = 0$ for $t' < t$; $\nabla_\theta J = \mathbb{E}_\tau\Big[\sum_t\nabla_\theta\log\pi_\theta(a_t \mid s_t)\,\hat G_t\Big]$ with $\hat G_t = \sum_{t' \ge t}r_{t'}$An action cannot change rewards already received, and the algebra agrees: those terms are zero in expectation but not sample by sample, so dropping them lowers the variance without biasing the estimate, the same mechanism as Problem 2's baseline. A baseline may also be subtracted per state, $\hat G_t - b(s_t)$, because given $s_t$ it is a constant with respect to $a_t$; that difference is the advantage. Weighting by the past rewards instead is Mistake 4. The check verifies both the zero and the equality on the enumerated environment.

Problem 9

1. $\dfrac{\partial H}{\partial\theta_k} = -\sum_a\dfrac{\partial p_a}{\partial\theta_k}\,(\log p_a + 1)$.Product rule on $p_a\log p_a$: its derivative is $(\log p_a + 1)\,\partial p_a$.
2. $\dfrac{\partial H}{\partial\theta_k} = -\sum_a p_a(\delta_{ak} - p_k)(\log p_a + 1) = -p_k(\log p_k + 1) + p_k\sum_a p_a(\log p_a + 1)$.The softmax Jacobian; the $\delta_{ak}$ picks out $a = k$ and the $p_k$ factors out of the rest.
3. $\sum_a p_a(\log p_a + 1) = -H + 1$, so $\dfrac{\partial H}{\partial\theta_k} = -p_k(\log p_k + 1) + p_k(1 - H) = -p_k(\log p_k + H)$.$\sum_a p_a\log p_a = -H$ and $\sum_a p_a = 1$; the $\pm p_k$ cancel.
4. With $r_\theta(a) = -\log p_\theta(a)$, $H = \mathbb{E}_{a\sim p_\theta}[r_\theta(a)]$ and $\nabla_\theta H = \mathbb{E}\big[r_\theta(a)\nabla_\theta\log p_\theta(a)\big] + \mathbb{E}\big[\nabla_\theta r_\theta(a)\big]$; the second term is $-\mathbb{E}[\nabla_\theta\log p_\theta(a)] = 0$ and the first is $\sum_a p_a(-\log p_a)(e_a - p) = -p\odot\log p + \Big(\sum_a p_a\log p_a\Big)p = -p\odot\log p - Hp$.Product rule, because both the distribution and the reward depend on $\theta$; Problem 2 for the second term; Problem 3's score for the first.
5. $\nabla_\theta H(p) = -p\odot(\log p + H\mathbf{1}) = \mathbb{E}_{a\sim p}\big[(-\log p_\theta(a) - H)(e_a - p)\big]$The entropy gradient pushes down the logits of actions whose surprise $-\log p_a$ is below the average surprise $H$, the likely ones, and up the unlikely ones; it points towards the uniform policy, where it vanishes. Adding $\beta H$ to the objective adds $\beta$ times this to the update, or in sampled form adds the bonus $-\log p_\theta(a) - H$ to each sampled reward, with $H$ acting as its own baseline. The entries sum to zero again, as every softmax gradient must.

Problem 10