Proof of the Policy Gradient Theorem

Here we prove the Policy gradient theorem, i.e., the gradient of an objective function $J(\theta )$ is $${\mathrm{\nabla }}_{\theta }J(\theta )=\sum _{s\in \mathcal{S}}\eta (s)\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s,\theta )q_{\pi }(s,a)$$ where $\eta$ is a state distribution and ${\mathrm{\nabla }}_{\theta }\pi$ is the gradient of $\pi$ with respect to $\theta$.

Moreover, this equation has a compact form expressed in terms of expectation: $${\mathrm{\nabla }}_{\theta }J(\theta )={\mathbb{E}}_{S\sim \eta ,A\sim \pi (S,\theta )}[{\mathrm{\nabla }}_{\theta }\ln \pi (A\mid S,\theta )q_{\pi }(S,A)],$$ where $\ln$ is the natural logarithm.

We prove this theorem in the discounted case and undiscounted cases separately. In each case, we prove it for 3 different metrics ${\overline{v}}_{\pi },{\overline{r}}_{\pi },{\overline{v}}_{\pi }^0$.

For simplicity, I only list the proof in the discounted case in the appendix. See the book for proof of the undiscounted case.

Sources:

1. Shiyu Zhao. Chapter 8: Policy Gradient Methods. Mathematical Foundations of Reinforcement Learning.

Derivation of the gradients in the discounted case

We next derive the gradients of the metrics in the discounted case where $\gamma \in (0,1)$. The state value and action value in the discounted case are defined as $$\begin{array}{rl}{v}_{\pi }(s)& =\mathbb{E}[R_{t+1}+\gamma R_{t+2}+{\gamma }^2{R}_{t+3}+\dots \mid S_t=s],\\ q_{\pi }(s,a)& =\mathbb{E}[R_{t+1}+\gamma R_{t+2}+{\gamma }^2{R}_{t+3}+\dots \mid S_t=s,A_t=a].\end{array}$$

It holds that $v_{\pi }(s)=\sum _{a\in \mathcal{A}}\pi (a\mid s,\theta )q_{\pi }(s,a)$ and the state value satisfies the Bellman equation.

First, we show that ${\overline{v}}_{\pi }(\theta )$ and ${\overline{r}}_{\pi }(\theta )$ are equivalent metrics.

Lemma 9.1

Lemma 9.1 (Equivalence between ${\overline{v}}_{\pi }(\theta )$ and ${\overline{r}}_{\pi }(\theta )$ ). In the discounted case where $\gamma \in$ $(0,1)$, it holds that $$\begin{array}{}\text{(1)}& {\overline{r}}_{\pi }=(1-\gamma ){\overline{v}}_{\pi }\end{array}$$

Proof: Note that ${\overline{v}}_{\pi }(\theta )=d_{\pi }^T{v}_{\pi }$ and ${\overline{r}}_{\pi }(\theta )=d_{\pi }^T{r}_{\pi }$, where $v_{\pi }$ and $r_{\pi }$ satisfy the Bellman equation $v_{\pi }=r_{\pi }+\gamma P_{\pi }{v}_{\pi }$. Multiplying $d_{\pi }^T$ on both sides of the Bellman equation yields $${\overline{v}}_{\pi }={\overline{r}}_{\pi }+\gamma d_{\pi }^T{P}_{\pi }{v}_{\pi }={\overline{r}}_{\pi }+\gamma d_{\pi }^T{v}_{\pi }={\overline{r}}_{\pi }+\gamma {\overline{v}}_{\pi }$$ which implies $\text{(1)}$. Second, the following lemma gives the gradient of $v_{\pi }(s)$ for any $s$.

Lemma 9.2

Lemma 9.2 (Gradient of $v_{\pi }(s)$ ). In the discounted case, it holds for any $s\in \mathcal{S}$ that $${\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)=\sum _{s^{\prime }\in \mathcal{S}}{\mathrm{Pr}}_{\pi }(s^{\prime }\mid s)\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s^{\prime },\theta )q_{\pi }(s^{\prime },a)$$ where $${\mathrm{Pr}}_{\pi }(s^{\prime }\mid s)\doteq \sum _{k=0}^{\mathrm{\infty }}{\gamma }^k{\left[P_{\pi }^k\right]}_{s{s}^{\prime }}={\left[{(I_n-\gamma P_{\pi })}^{-1}\right]}_{s{s}^{\prime }}$$ is the discounted total probability of transitioning from $s$ to $s^{\prime }$ under policy $\pi$. Here, $[\cdot {]}_{s{s}^{\prime }}$ denotes the entry in the sth row and $s^{\prime }$ th column, and ${\left[P_{\pi }^k\right]}_{s{s}^{\prime }}$ is the probability of transitioning from $s$ to $s^{\prime }$ using exactly $k$ steps under $\pi$.

With the results in Lemma 9.2, we are ready to derive the gradient of ${\overline{v}}_{\pi }^0$.

Theorem 9.2

Theorem 9.2 (Gradient of ${\overline{v}}_{\pi }^0$ in the discounted case). In the discounted case where $\gamma \in (0,1)$, the gradient of ${\overline{v}}_{\pi }^0=d_0^T{v}_{\pi }$ is $${\mathrm{\nabla }}_{\theta }{\overline{v}}_{\pi }^0=\mathbb{E}[{\mathrm{\nabla }}_{\theta }\ln \pi (A\mid S,\theta )q_{\pi }(S,A)]$$ where $S\sim {\rho }_{\pi }$ and $A\sim \pi (S,\theta )$. Here, the state distribution ${\rho }_{\pi }$ is $${\rho }_{\pi }(s)=\sum _{s^{\prime }\in \mathcal{S}}{d}_0\left(s^{\prime }\right){\mathrm{Pr}}_{\pi }(s\mid s^{\prime }),s\in \mathcal{S}$$ where ${\mathrm{Pr}}_{\pi }(s\mid s^{\prime })=\sum _{k=0}^{\mathrm{\infty }}{\gamma }^k{\left[P_{\pi }^k\right]}_{s^{\prime }s}={\left[{(I-\gamma P_{\pi })}^{-1}\right]}_{s^{\prime }s}$ is the discounted total probability of transitioning from $s^{\prime }$ to $s$ under policy $\pi$.

Theorem 9.3

Theorem 9.3 (Gradients of ${\overline{r}}_{\pi }$ and ${\overline{v}}_{\pi }$ in the discounted case). In the discounted case where $\gamma \in (0,1)$, the gradients of ${\overline{r}}_{\pi }$ and ${\overline{v}}_{\pi }$ are $$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }{\overline{r}}_{\pi }=(1-\gamma ){\mathrm{\nabla }}_{\theta }{\overline{v}}_{\pi }& \approx \sum _{s\in \mathcal{S}}{d}_{\pi }(s)\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s,\theta )q_{\pi }(s,a)\\ & =\mathbb{E}[{\mathrm{\nabla }}_{\theta }\ln \pi (A\mid S,\theta )q_{\pi }(S,A)],\end{array}$$ where $S\sim d_{\pi }$ and $A\sim \pi (S,\theta )$. Here, the approximation is more accurate when $\gamma$ is closer to 1 .

Appendix

Proof of Lemma 9.2

First, for any $s\in \mathcal{S}$, it holds that $$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)& ={\mathrm{\nabla }}_{\theta }[\sum _{a\in \mathcal{A}}\pi (a\mid s,\theta )q_{\pi }(s,a)]\\ & =\sum _{a\in \mathcal{A}}[{\mathrm{\nabla }}_{\theta }\pi (a\mid s,\theta )q_{\pi }(s,a)+\pi (a\mid s,\theta ){\mathrm{\nabla }}_{\theta }{q}_{\pi }(s,a)]\end{array}$$ where $q_{\pi }(s,a)$ is the action value given by $$q_{\pi }(s,a)=r(s,a)+\gamma \sum _{s^{\prime }\in \mathcal{S}}p(s^{\prime }\mid s,a)v_{\pi }\left(s^{\prime }\right)$$

Since $r(s,a)=\sum _{r}rp(r\mid s,a)$ is independent of $\theta$, we have $${\mathrm{\nabla }}_{\theta }{q}_{\pi }(s,a)=0+\gamma \sum _{s^{\prime }\in \mathcal{S}}p(s^{\prime }\mid s,a){\mathrm{\nabla }}_{\theta }{v}_{\pi }\left(s^{\prime }\right)$$ Substituting this result into the policy gradient ${\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)$ yields $$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)& =\sum _{a\in \mathcal{A}}[{\mathrm{\nabla }}_{\theta }\pi (a\mid s,\theta )q_{\pi }(s,a)+\pi (a\mid s,\theta )\gamma \sum _{s^{\prime }\in \mathcal{S}}p(s^{\prime }\mid s,a){\mathrm{\nabla }}_{\theta }{v}_{\pi }\left(s^{\prime }\right)]\\ & =\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s,\theta )q_{\pi }(s,a)+\gamma \sum _{a\in \mathcal{A}}\pi (a\mid s,\theta )\sum _{s^{\prime }\in \mathcal{S}}p(s^{\prime }\mid s,a){\mathrm{\nabla }}_{\theta }{v}_{\pi }\left(s^{\prime }\right)\end{array}$$

It is notable that ${\mathrm{\nabla }}_{\theta }{v}_{\pi }$ appears on both sides of the above equation. Here, we use the matrix-vector form to calculate it. In particular, let $$u(s)\doteq \sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s,\theta )q_{\pi }(s,a)$$ Since $$\sum _{a\in \mathcal{A}}\pi (a\mid s,\theta )\sum _{s^{\prime }\in \mathcal{S}}p(s^{\prime }\mid s,a){\mathrm{\nabla }}_{\theta }{v}_{\pi }\left(s^{\prime }\right)=\sum _{s^{\prime }\in \mathcal{S}}p(s^{\prime }\mid s){\mathrm{\nabla }}_{\theta }{v}_{\pi }\left(s^{\prime }\right)=\sum _{s^{\prime }\in \mathcal{S}}{\left[P_{\pi }\right]}_{s{s}^{\prime }}{\mathrm{\nabla }}_{\theta }{v}_{\pi }\left(s^{\prime }\right)$$ equation $${\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)=\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s,\theta )q_{\pi }(s,a)+\gamma \sum _{a\in \mathcal{A}}\pi (a\mid s,\theta )\sum _{s^{\prime }\in \mathcal{S}}p(s^{\prime }\mid s,a){\mathrm{\nabla }}_{\theta }{v}_{\pi }\left(s^{\prime }\right)$$ can be written in matrix-vector form as $$\underset{{\mathrm{\nabla }}_{\theta }{v}_{\pi }\in {\mathbb{R}}^{mn}}{\underbrace{\left[\begin{array}{c}⋮\\ {\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)\\ ⋮\end{array}\right]}}=\underset{u\in {\mathbb{R}}^{mn}}{\underbrace{\left[\begin{array}{c}⋮\\ u(s)\\ ⋮\end{array}\right]}}+\gamma (P_{\pi }\otimes I_m)\underset{{\mathrm{\nabla }}_{\theta }{v}_{\pi }\in {\mathbb{R}}^{mn}}{\underbrace{\left[\begin{array}{c}⋮\\ {\mathrm{\nabla }}_{\theta }{v}_{\pi }\left(s^{\prime }\right)\\ ⋮\end{array}\right]}},$$ which can be written concisely as $${\mathrm{\nabla }}_{\theta }{v}_{\pi }=u+\gamma (P_{\pi }\otimes I_m){\mathrm{\nabla }}_{\theta }{v}_{\pi }.$$ Here, $n=|\mathcal{S}|$, and $m$ is the dimension of the parameter vector $\theta$. The reason that the Kronecker product $\otimes$ emerges in the equation is that ${\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)$ is a vector. The above equation is a linear equation of ${\mathrm{\nabla }}_{\theta }{v}_{\pi }$, which can be solved as $$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }{v}_{\pi }& ={(I_{nm}-\gamma P_{\pi }\otimes I_m)}^{-1}u\\ & ={(I_n\otimes I_m-\gamma P_{\pi }\otimes I_m)}^{-1}u\\ & =[{(I_n-\gamma P_{\pi })}^{-1}\otimes I_m]u.\end{array}$$

For any state $s$, it follows from this equation that $$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)& =\sum _{s^{\prime }\in \mathcal{S}}{\left[{(I_n-\gamma P_{\pi })}^{-1}\right]}_{s{s}^{\prime }}u\left(s^{\prime }\right)\\ & =\sum _{s^{\prime }\in \mathcal{S}}{\left[{(I_n-\gamma P_{\pi })}^{-1}\right]}_{s{s}^{\prime }}\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s^{\prime },\theta )q_{\pi }(s^{\prime },a).\end{array}$$

The quantity ${\left[{(I_n-\gamma P_{\pi })}^{-1}\right]}_{s{s}^{\prime }}$ has a clear probabilistic interpretation. In particular, since $${(I_n-\gamma P_{\pi })}^{-1}=I+\gamma P_{\pi }+{\gamma }^2{P}_{\pi }^2+\cdots ,$$ we have $${\left[{(I_n-\gamma P_{\pi })}^{-1}\right]}_{s{s}^{\prime }}=[I{]}_{s{s}^{\prime }}+\gamma {\left[P_{\pi }\right]}_{s{s}^{\prime }}+{\gamma }^2{\left[P_{\pi }^2\right]}_{s{s}^{\prime }}+\cdots =\sum _{k=0}^{\mathrm{\infty }}{\gamma }^k{\left[P_{\pi }^k\right]}_{s{s}^{\prime }}.$$ Note that ${\left[P_{\pi }^k\right]}_{s{s}^{\prime }}$ is the probability of transitioning from $s$ to $s^{\prime }$ using exactly $k$ steps (see my previous post). Therefore, $${\left[{(I_n-\gamma P_{\pi })}^{-1}\right]}_{s{s}^{\prime }}$$ is the discounted total probability of transitioning from $s$ to $s^{\prime }$ using any number of steps. By denoting $${\left[{(I_n-\gamma P_{\pi })}^{-1}\right]}_{s{s}^{\prime }}\doteq {\mathrm{Pr}}_{\pi }(s^{\prime }\mid s),$$ we obtain $$\sum _{s^{\prime }\in \mathcal{S}}{\left[{(I_n-\gamma P_{\pi })}^{-1}\right]}_{s{s}^{\prime }}\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s^{\prime },\theta )q_{\pi }(s^{\prime },a)=\sum _{s^{\prime }\in \mathcal{S}}{\mathrm{Pr}}_{\pi }(s^{\prime }\mid s)\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s^{\prime },\theta )q_{\pi }(s^{\prime },a)={\mathrm{\nabla }}_{\theta }{v}_{\pi }(s).$$

Q. E. D.

Proof of Theorem 9.2

Since $d_0(s)$ is independent of $\pi$, we have $${\mathrm{\nabla }}_{\theta }{\overline{v}}_{\pi }^0={\mathrm{\nabla }}_{\theta }\sum _{s\in \mathcal{S}}{d}_0(s)v_{\pi }(s)=\sum _{s\in \mathcal{S}}{d}_0(s){\mathrm{\nabla }}_{\theta }{v}_{\pi }(s).$$

Substituting the expression of ${\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)$ given in Lemma 9.2 into the above equation yields $$ $$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }{\overline{v}}_{\pi }^0=\sum _{s\in \mathcal{S}}{d}_0(s){\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)& =\sum _{s\in \mathcal{S}}{d}_0(s)\sum _{s^{\prime }\in \mathcal{S}}{\mathrm{Pr}}_{\pi }(s^{\prime }\mid s)\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s^{\prime },\theta )q_{\pi }(s^{\prime },a)\\ & =\sum _{s^{\prime }\in \mathcal{S}}(\sum _{s\in \mathcal{S}}{d}_0(s){\mathrm{Pr}}_{\pi }(s^{\prime }\mid s))\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s^{\prime },\theta )q_{\pi }(s^{\prime },a)\\ & \doteq \sum _{s^{\prime }\in \mathcal{S}}{\rho }_{\pi }\left(s^{\prime }\right)\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s^{\prime },\theta )q_{\pi }(s^{\prime },a)\\ & =\sum _{s\in \mathcal{S}}{\rho }_{\pi }(s)\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s,\theta )q_{\pi }(s,a)\text{ (change }{s}^{\prime }\text{ to }s\text{ ) }\\ & =\sum _{s\in \mathcal{S}}{\rho }_{\pi }(s)\sum _{a\in \mathcal{A}}\pi (a\mid s,\theta ){\mathrm{\nabla }}_{\theta }\ln \pi (a\mid s,\theta )q_{\pi }(s,a)\\ & =\mathbb{E}[{\mathrm{\nabla }}_{\theta }\ln \pi (A\mid S,\theta )q_{\pi }(S,A)],\end{array}$$

$$ where $S\sim {\rho }_{\pi }$ and $A\sim \pi (S,\theta )$. The proof is complete.

Proof of Theorem 9.3

It follows from the definition of ${\overline{v}}_{\pi }$ that $$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }{\overline{v}}_{\pi }& ={\mathrm{\nabla }}_{\theta }\sum _{s\in \mathcal{S}}{d}_{\pi }(s)v_{\pi }(s)\\ & =\sum _{s\in \mathcal{S}}{\mathrm{\nabla }}_{\theta }{d}_{\pi }(s)v_{\pi }(s)+\sum _{s\in \mathcal{S}}{d}_{\pi }(s){\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)\end{array}$$

This equation contains two terms. On the one hand, substituting the expression of ${\mathrm{\nabla }}_{\theta }{v}_{\pi }$ given in (9.17) into the second term gives $$\begin{array}{rl}\sum _{s\in \mathcal{S}}{d}_{\pi }(s){\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)& =(d_{\pi }^T\otimes I_m){\mathrm{\nabla }}_{\theta }{v}_{\pi }\\ & =(d_{\pi }^T\otimes I_m)[{(I_n-\gamma P_{\pi })}^{-1}\otimes I_m]u\\ & =\left[d_{\pi }^T{(I_n-\gamma P_{\pi })}^{-1}\right]\otimes I_{m}u\end{array}$$

It is noted that $$d_{\pi }^T{(I_n-\gamma P_{\pi })}^{-1}=\frac{1}{1-\gamma }{d}_{\pi }^T$$ which can be easily verified by multiplying $(I_n-\gamma P_{\pi })$ on both sides of the equation. Therefore, (9.21) becomes $$\begin{array}{rl}\sum _{s\in \mathcal{S}}{d}_{\pi }(s){\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)& =\frac{1}{1-\gamma }{d}_{\pi }^T\otimes I_{m}u\\ & =\frac{1}{1-\gamma }\sum _{s\in \mathcal{S}}{d}_{\pi }(s)\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s,\theta )q_{\pi }(s,a)\end{array}$$ On the other hand, the first term of $$\sum _{s\in \mathcal{S}}{\mathrm{\nabla }}_{\theta }{d}_{\pi }(s)v_{\pi }(s)+\sum _{s\in \mathcal{S}}{d}_{\pi }(s){\mathrm{\nabla }}_{\theta }{v}_{\pi }(s)$$ involves ${\mathrm{\nabla }}_{\theta }{d}_{\pi }$. However, since the second term contains $\frac{1}{1-\gamma }$, the second term becomes dominant, and the first term becomes negligible (#TODO) when $\gamma \to 1$. Therefore, $${\mathrm{\nabla }}_{\theta }{\overline{v}}_{\pi }\approx \frac{1}{1-\gamma }\sum _{s\in \mathcal{S}}{d}_{\pi }(s)\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s,\theta )q_{\pi }(s,a)$$

Furthermore, it follows from ${\overline{r}}_{\pi }=(1-\gamma ){\overline{v}}_{\pi }$ that $$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }{\overline{r}}_{\pi }=(1-\gamma ){\mathrm{\nabla }}_{\theta }{\overline{v}}_{\pi }& \approx \sum _{s\in \mathcal{S}}{d}_{\pi }(s)\sum _{a\in \mathcal{A}}{\mathrm{\nabla }}_{\theta }\pi (a\mid s,\theta )q_{\pi }(s,a)\\ & =\sum _{s\in \mathcal{S}}{d}_{\pi }(s)\sum _{a\in \mathcal{A}}\pi (a\mid s,\theta ){\mathrm{\nabla }}_{\theta }\ln \pi (a\mid s,\theta )q_{\pi }(s,a)\\ & =\mathbb{E}[{\mathrm{\nabla }}_{\theta }\ln \pi (A\mid S,\theta )q_{\pi }(S,A)]\end{array}$$ The approximation in the above equation requires that the first term does not go to infinity when $\gamma \to 1$.