In parts 1-3 we found that learning the values of different states (or state-action pairs) made it easy to define good polices; we simply selected high valued states and actions. Policy gradient methods use a different approach: learn policies directly by optimizing their parameters to maximize reward. These techniques allow us to tackle more interesting problems consisting of large or continuous action and state spaces. The math is a bit heavier , but so is the payoff.

policy gradient

Up to this point we have been learning value functions, which can be used to find good policies. Policy gradient methods directly learn the parameters for the policy. To do this we gradually tweak our policy parameters $\theta$ in the direction that increases expected reward. We therefore need to take the gradient of some performance metric $J(\theta)$ with respect to $\theta$ and update in that direction by some amount $\alpha \in (0,1]$:

policy gradient theorem

Let $J(\theta)$ be the amount of reward we expect under the policy. This is the thing we want to maximize. Differentiating $J(\theta)$ with respect to $\theta$ presents a challenge. The expected reward depends on the policy $\pi_\theta(a|s)$, which we take to be differentiable, but also on the distribution of states, which depends on potentially complex interactions between the policy and the environment.

The policy gradient theorem solves this problem. Consider the episodic case¹ in which we start at state $s_0$. Performance can be defined as the expected return for the episode, $v(s_0)$. The policy gradient theorem gives the following derivative, where $\mu(s)$ is the stationary distribution² of states under $\pi$:

The magic happens in the third line. The policy gradient theorem says we can take the gradient with respect to the policy without worrying about the gradient of the state distribution.

policy gradient theorem proof

We want to take the gradient of $J(\theta) = v(s_0)$. Let’s start by taking the gradient for any state $s$. First we’ll establish a recursive relationship between the value of a state and the subsequent state. For simplicity (but without loss of generality) we will not use discounting. Using the product rule we have:

Let’s expand the final term in this expression. Recall that in stochastic environments the same state-action pair $(s,a)$ can result in several possible subsequent state-reward pairs $(s’,r)$. Therefore, $q(s,a)$ can be expressed as an average over these possibilities, $q(s,a) = \sum_{s’,r}p(s’,r \mid s,a)[r + v(s’)]$, giving us:

When we distribute the rightmost $\nabla_\theta$ the $r$ disappears because it doesn’t depend on $\theta$ ($r$ behaves like a constant inside the summation):

Finally, to reduce the painfulness of the following steps we will set $\phi(s) = \sum_a \nabla \pi(a \mid s) q(s,a)$:

Okay! Get ready for some funny notation 📝📝📝. The probability of moving from state $s$ to state $x$ in exactly $k$ steps under the current policy will be denoted $\rho(s \rightarrow x, k)$. Notice that to compute the probability of moving from $s_0$ to $s_2$ in 2 steps, we add up the probabilities of getting there via all possible intermediate states $s_1$:

With this seemningly out-of-nowhere notation we are ready to prove the policy gradient theorem. First we shift around the summations a bit:

The probability of moving from $s$ to $s’$ in one step is the average of the probabilities for each action we could take from $s$. Mathematically, this means that $\rho(s \rightarrow s’, 1) = \sum_a p(a \mid s) p(s’ \mid s,a)$. Performing this substitution and then plugging in our formula recursively we have:

Now we distribute $\sum_{s’} \rho(s \rightarrow s’, 1)$ and use the fact that $\rho(s \rightarrow s’’, 2) = \sum_{s’} \rho(s \rightarrow s’, 1) \rho(s’ \rightarrow s’’, 1)$:

Continuing the pattern forever we get:

So 😬. Close 😬. I’m only going to introduce one more annoying thing. $\eta(s)$ will be the average number of steps we spend in state $s$ across episodes. It is the sum of the probabilities of being in state $s$ after zero steps, one steps, two steps…

Now let’s consider what happens to our gradient with $s = s_0$:

The proof is complete! In the last line we use the fact that the amount of time we spend in a state, $\eta(s)$, is proportional to the probability of being in that state, $\mu(s)$. Specifically, $\mu(s) = \eta(s) / \sum_{s’}\eta(s’)$, a fact we use in the fourth line.

algorithms

REINFORCE

Let’s put the policy gradient theorem to use. Although we don’t know the true distribution of states $\mu(s)$, we can continuously sample states under the current policy to approximate the gradient without bias. This Monte Carlo approach relies on the fact that:

There is a subtle shift here. In the first line we loop over specific states $s$, whereas in the second line we sample states $S_t$. Now the state is a random variable. This means that instead of averaging across states according to their (unknown!) probabilities, we can just behave in the world and average across the states we see!

Using similar logic, rather than performing updates by averaging across all actions, we can perform updates by sampling one action at a time because:

This leads to the stochastic update rule:

Because $\nabla \ln{f(x)} = \frac{\nabla f(x)}{f(x)}$ this is equivalent to:

This formula makes both mathematical and common sense. If the return $G_t$ is high, we want to make the actions that led to it more likely. In other words, we want to push $\theta$ in the direction of the gradient of $\pi$ for actions that are associated with large returns.

REINFORCE with baseline

REINFORCE has zero bias because it relies on sample returns, but it can suffer from very high variance due to stochasticity in these returns. We can reduce variance by subtracting a baseline from these returns. A common approach is to subtract an estimate of the state value function (or related functions):

This strategy does not introduce bias because the value function does not depend on the selected actions.

actor-critic

We can further reduce variance by constructing bootstrapped targets, e.g. by replacing $G_t$ with a one-step TD target:

In such actor-critic approaches the actor - the policy - selects actions whereas the critic - the learned value function - helps evaluate the goodness of the selected actions. Because we are now using a bootstrapped target estimate we have introduced bias into our updates.

Note that REINFORCE with baseline is not an actor-critic approach despite also using a learned value function. This is because using $\hat{v}(S_t)$ introduces bias when used to bootstrap the target, but not when used as a baseline. Even this (super useful) TensorFlow tutorial makes this mistake. See further discussion here.

policy parameterizations

Policy gradient methods present a natural way of dealing with large or continuous action spaces. Rather than learning probability mass functions over many different actions we can directly learn the parameters of probability distributions, for example the mean and standard deviation of Gaussians. We can then draw from these Gaussians when selecting actions in a continuous action space.

acknowledgements

Lilian Weng’s excellent policy gradient tutorial was very helpful, as was Sutton and Barto’s presentation of this proof.