proof of the policy gradient theorem

The policy gradient theorem says that

For simplicity we'll assume a fixed initial state $s_0$ and fixed-length finite trajectories, but the result can be generalized to discounted-reward or average-reward notions of value in the continuing setting.

It seems like there are two proof approaches:

1. We can either recursively expand the value as a sequence of Q-values, beginning at the start state, or
2. We can work with rewards of entire trajectories $R(\tau)$, leading to a product of $T^2$ interactions between policy choices and rewards, and then remove the 'acausal' terms because they have expectation zero.

Sutton and Barto do the first, while John Schulman does the second in my deep RL notes.

Recursive expansion proof

Reproducing the Sutton and Barto approach for completeness: they expand out the value $V^{\pi_\theta}(s_0)$ to an expectation over q-values $\sum_{a_0} \pi_\theta(a_0 | s_0) Q^{\pi_\theta}(s_0, a_0)$, push the gradient inside and then apply the product rule

then recursively expand the gradient of the Q-value:

so that plugging this all back in we get the policy gradient as a sum in which each term $\nabla_\theta \pi_\theta(a | s) Q^{\pi_\theta}(s, a)$ has weight equal to the expected number of times that state $s$ is visited under policy $\pi_\theta$, which is just off from the expression at the top (which takes an expectation over states) by a constant factor of T.

Full trajectories derivation

This follows the derivation in deep RL notes.

Cleaning this up a bit we end up with the expression

which we can see as a variant of the theorem above in which the expectation over state visitations is split into an expectation over trajectories and the empirical sum over states within a trajectory.

Note that it's valid to move between sums of empirical rewards $\sum_{t'>t} \gamma_t r_{t}$ and their expectations---the Q values $Q^{\pi_\theta}(s_t, a_t)$---because this all happens inside of the expectation over $s_t, a_t$. I was worried that in general $r_{t'>t}$ is not independent of $\nabla_\theta \log \pi_\theta(a_t | s_t)$, but that doesn't matter because the latter is just a constant once we condition on $(s_t, a_t)$.