Deriving Policy Gradient

In this post, I derive the gradient of a policy. This tutorial follows the steps in the first part of lecture 5 of CS285 at UC Berkeley.

Terminology

In reinforcement learning, a trajectory $\tau$ is a sequence of states and actions ${\mathbf{s}}_1,{\mathbf{a}}_1,\dots ,{\mathbf{s}}_T,{\mathbf{a}}_T$ collected by a policy. And it can be defined as

$$p_{\theta }(\tau )=p_{\theta }\left({\mathbf{s}}_1,{\mathbf{a}}_1,\dots ,{\mathbf{s}}_T,{\mathbf{a}}_T\right)=p\left({\mathbf{s}}_1\right)\prod _{t=1}^T{\pi }_{\theta }\left({\mathbf{a}}_t\mid {\mathbf{s}}_t\right)p\left({\mathbf{s}}_{t+1}\mid {\mathbf{s}}_t,{\mathbf{a}}_t\right).$$

A trajectory distribution is a probability distribution over a sequence of states and actions. In the equation above, it is represented by the chain rule of probability by multiplying the initial state distribution $p\left({\mathbf{s}}_1\right)$ by the product of policy probability ${\pi }_{\theta }\left({\mathbf{a}}_t\mid {\mathbf{s}}_t\right)$ and transition probability $p\left({\mathbf{s}}_{t+1}\mid {\mathbf{s}}_t,{\mathbf{a}}_t\right)$ over all time steps.

The reinforcement learning objective can be written as the optimal parameter $\theta$ that maximizes the expected reward under a trajectory.

$${\theta }^{\star }=\arg \underset{\theta }{\max }{E}_{\tau \sim p_{\theta }(\tau )}\left[\sum _{t}r\left({\mathbf{s}}_t,{\mathbf{a}}_t\right)\right]$$

In the following derivation, I will use $J(\theta )$ to represent the expected reward ${\sum }_{t}r\left({\mathbf{s}}_t,{\mathbf{a}}_t\right)$, and replace ${\sum }_{t=1}^{T}r\left({\mathbf{s}}_t,{\mathbf{a}}_t\right)$ with $r(\tau )$ for simplicity.

Differentiating the Policy Directly

Now, we have our objective

$${\theta }^{\star }=\arg \underset{\theta }{\max }J(\theta )$$

and expectation of rewards

$$J(\theta )=E_{\tau \sim p_{\theta }(\tau )}[r(\tau )].$$

We expand the expectation for continuous variables to its integral form.

$$J(\theta )=E_{\tau \sim p_{\theta }(\tau )}[r(\tau )]=\int p_{\theta }(\tau )r(\tau )d\tau$$

Then, the gradient (or the derivative) of the expected reward can be written as ${\nabla }_{\theta }J(\theta )$ by directly putting the differentiation operator ${\nabla }_{\theta }$ inside the integral because it is linear.

$${\nabla }_{\theta }J(\theta )=\int {\nabla }_{\theta }{p}_{\theta }(\tau )r(\tau )d\tau$$

Now, we need to use the log-derivative identity

$$\frac{d}{dx}\log (x)=\frac{dx}{x}$$

to help us expand the ${\nabla }_{\theta }{p}_{\theta }(\tau )$ by applying it inversely as

$$p_{\theta }(\tau ){\nabla }_{\theta }\log p_{\theta }(\tau )=p_{\theta }(\tau )\frac{{\nabla }_{\theta }{p}_{\theta }(\tau )}{p_{\theta }(\tau )}={\nabla }_{\theta }{p}_{\theta }(\tau ).$$

If we replace the ${\nabla }_{\theta }{p}_{\theta }(\tau )$ term in the gradient of expectation by the left hand side $p_{\theta }(\tau ){\nabla }_{\theta }\log p_{\theta }(\tau )$ and convert it back into the expectation form, we will get

$$\begin{array}{rl}{\nabla }_{\theta }J(\theta )& =\int {\nabla }_{\theta }{p}_{\theta }(\tau )r(\tau )d\tau \\ & =\int p_{\theta }(\tau ){\nabla }_{\theta }\log p_{\theta }(\tau )r(\tau )d\tau \\ & =E_{\tau \sim p_{\theta }(\tau )}\left[{\nabla }_{\theta }\log p_{\theta }(\tau )r(\tau )\right].\end{array}$$

Recall that the trajectory distribution is

$$p_{\theta }(\tau )=p\left({\mathbf{s}}_1\right)\prod _{t=1}^T{\pi }_{\theta }\left({\mathbf{a}}_t\mid {\mathbf{s}}_t\right)p\left({\mathbf{s}}_{t+1}\mid {\mathbf{s}}_t,{\mathbf{a}}_t\right),$$

and if we take the logarithm of both sides, we will get

$$\log p_{\theta }(\tau )=\log p\left({\mathbf{s}}_1\right)+\sum _{t=1}^T\log {\pi }_{\theta }\left({\mathbf{a}}_t\mid {\mathbf{s}}_t\right)+\log p\left({\mathbf{s}}_{t+1}\mid {\mathbf{s}}_t,{\mathbf{a}}_t\right).$$

We proceed to substitute the right hand side of this equation for $\log p_{\theta }(\tau )$ inside the expectation.

$$\begin{array}{rl}{\nabla }_{\theta }J(\theta )& =E_{\tau \sim p_{\theta }(\tau )}\left[{\nabla }_{\theta }\log p_{\theta }(\tau )r(\tau )\right]\\ & =E_{\tau \sim p_{\theta }(\tau )}\left[{\nabla }_{\theta }\left[\log p\left({\mathbf{s}}_1\right)+\sum _{t=1}^T\log {\pi }_{\theta }\left({\mathbf{a}}_t\mid {\mathbf{s}}_t\right)+\log p\left({\mathbf{s}}_{t+1}\mid {\mathbf{s}}_t,{\mathbf{a}}_t\right)\right]r(\tau )\right]\end{array}$$

Both $\log p\left({\mathbf{s}}_1\right)$ and $\log p\left({\mathbf{s}}_{t+1}\mid {\mathbf{s}}_t,{\mathbf{a}}_t\right)$ do not depend on $\theta$, so we cancel them out and simplify it to

$${\nabla }_{\theta }J(\theta )=E_{\tau \sim p_{\theta }(\tau )}\left[\left(\sum _{t=1}^T{\nabla }_{\theta }\log {\pi }_{\theta }\left({\mathbf{a}}_t\mid {\mathbf{s}}_t\right)\right)\left(\sum _{t=1}^{T}r\left({\mathbf{s}}_t,{\mathbf{a}}_t\right)\right)\right].$$

Reference

[1] CS 285 Deep Reinforcement Learning.