Policy Gradient from Scratch

Before going into the policy gradient we need to have some understanding/review of the following term $G(t)$ is the discounted future return from time t or $G_t=r_t+\gamma r_{t+1}+{\gamma }^2{r}_{t+2}+\dots$ , ${\pi }_{\theta }(a|s)$ is the policy, note that it is a distribution we will discuss that if the policy is deterministic later (hopefully). we also have a state transition probability $P(s_{t+1}\mid s_t,a_t)$, a state value function $v(s_t)$ and a action value function $Q(s_t,a_t)=\mathbb{E}\left[G_t\mid s_t,a_t\right]$.

In reinforcement learning, the objective is defined as

$$J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[G_0(\tau )]$$

Then we can expand it to the following integration over all possible trajectory

$${\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[G_0(\tau )]=\int p_{\theta }(\tau )G_0(\tau )d\tau$$

where

$$p_{\theta }(\tau )=p(s_0)\prod _t{\pi }_{\theta }(a_t\mid s_t)p(s_{t+1}\mid s_t,a_t).$$

Above we see that the expected return of a trajectory is the integration over the product of initial state probability, the policy output distribution and the state transition probability and the collected return from all possible trajectory.

The gradient with respect to the policy parameter

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

Use the log gradient trick

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

Expand we have

$$\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))$$

Taking gradient w.r.t. $\theta$ gives

$${\nabla }_{\theta }\log p_{\theta }(\tau )=\sum _t{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t\mid s_t))$$

since $\log p(s_0)$ and $\log p(s_{t+1}\mid s_t,a_t)$ do not depend on $\theta$.

Then we could write

$$\begin{array}{rl}{\nabla }_{\theta }J(\theta )& =\int p_{\theta }(\tau )\sum _t{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t\mid s_t))G_0(\tau )d\tau \\ & ={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _t{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t\mid s_t))G_0(\tau )\right]\end{array}$$

Note that $G_0(\tau )$ contains all time step’s future return, however for Expectation at time step $t$, by Markov property, rewards before time $t$ do not depend on action $a_t$, so we can replace $G_0(\tau )$ with $G_t$, and write that

$${\nabla }_{\theta }J(\theta )=\sum _t{\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t\mid s_t))G_t\right].$$

Which is the standard REINFORCE formulation , by law of total expectation

$$\begin{array}{rl}{E}_Y[E[X|Y]]& =E[X]\\ & ={\int }_{y}P(Y)E[X|Y]dy\\ & ={\int }_{y}P(Y)\left({\int }_{x}P(X|Y)Xdx\right)dy\\ & ={\int }_{y}P(Y)\left({\int }_x\frac{P(X,Y)}{P(Y)}Xdx\right)dy\\ & ={\int }_y{\int }_{x}P(X,Y)Xdxdy\\ & ={\int }_{x}X\left({\int }_{y}P(X,Y)dy\right)dx\\ & ={\int }_{x}P(X)Xdx\\ & =E[X].\end{array}$$

and that $G(t)$ is a nosier realization of $Q(s_t,a_t)$ if Q is exact (conditional expectation always reduce variance ⇐) , with law of total expectation, and conditioning on the state and action, we get the following

$${\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\sum _t{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t|s_t))E(G_t|s_t,a_t)]={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\sum _t{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t|s_t))Q(s_t,a_t)]$$

which is policy gradient, below is a detailed derivation, we look at a specific time-step

$$\begin{array}{rl}{\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t|s_t))G(t)]& ={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\mathbb{E}[{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t|s_t))G(t)|s_t,a_t]]\\ & ={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t|s_t))\mathbb{E}[G(t)|s_t,a_t]]\\ & ={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t|s_t))Q(s_t,a_t)]\end{array}$$

Note that $G(t)$ is a Monte-Carlo sample of the trajectory and is ill-posed with high variance, now our mission is to find ways to reduce the variance. first, we introduce a baseline $b(s)$ that is only state dependent, now we proceed to prove that such baseline does not create bias on the policy gradient.

$${\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t|s_t))(Q(s_t,a_t)-b(s_t))]$$

Then

$$\begin{array}{rl}{\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\mathbb{E}[{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t|s_t))b(s_t)|s_t]]& ={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[b(s_t)\mathbb{E}[{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t|s_t))]]\\ \mathbb{E}[{\nabla }_{\theta }\log ({\pi }_{\theta }(a_t|s_t))]& =\sum _a{\pi }_{\theta }(a_t|s_t){\nabla }_{\theta }\log ({\pi }_{\theta }(a_t|s_t))\end{array}$$

Note that here we use the log trick again

$${\nabla }_{\theta }\log ({\pi }_{\theta }(a_t|s_t))=\frac{\nabla {\pi }_{\theta }(a_t|s_t)}{{\pi }_{\theta }(a_t|s_t)}$$

The equation above become

$$\sum _a{\pi }_{\theta }(a_t|s_t){\nabla }_{\theta }\log ({\pi }_{\theta }(a_t|s_t))=\sum _a{\nabla }_{\theta }{\pi }_{\theta }(a_t|s_t)={\nabla }_{\theta }\sum _a{\pi }_{\theta }(a_t|s_t)={\nabla }_{\theta }1=0$$

So as long as we have a non action dependent baseline, we can arrive have an unbiased estimation of policy gradient.