Policy Gradient Theorem 的证明

如今，强化学习基本都采用参数化的神经网络来学习一个策略，而神经网络一般是通过梯度下降法或者各种变种来优化的，因此，获取累积回报关于策略的梯度至关重要。本节会推导策略梯度的表达式，并介绍实际训练中如何采样近似该表达式。

首先引入动作价值函数的定义：

$$Q^{\pi }(s_t,a_t)={\mathbb{E}}_{s_{t+1},a_{t+1},\dots }\left[\sum _{l=0}^{\mathrm{\infty }}{\gamma }^l{r}_{t+l}\right]$$

即在状态 $s_t$ 下采用动作 $a_t$ 后，后续动作服从策略 $\pi$ 的情况下的累积期望回报，其中 $\gamma \in (0,1)$ 是折扣因子。

接着，将策略梯度的计算过程详细展开：

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }J(\theta )& ={\mathrm{\nabla }}_{\theta }{V}^{\pi }(s_0)\\ & ={\mathrm{\nabla }}_{\theta }[\sum _{a_0}\pi (a_0\mid s_0)Q^{\pi }(s_0,a_0)]\\ & =\sum _{a_0}[{\mathrm{\nabla }}_{\theta }\pi (a_0\mid s_0)Q^{\pi }(s_0,a_0)+\pi (a_0\mid s_0){\mathrm{\nabla }}_{\theta }{Q}^{\pi }(s_0,a_0)]\\ & =\sum _{a_0}[{\mathrm{\nabla }}_{\theta }\pi (a_0\mid s_0)Q^{\pi }(s_0,a_0)+\pi (a_0\mid s_0){\mathrm{\nabla }}_{\theta }\sum _{s_1,r_1}p(s_1,r_1\mid s_0,a_0)(r_1+\gamma V^{\pi }(s_1))]\\ & =\sum _{a_0}{\mathrm{\nabla }}_{\theta }\pi (a_0\mid s_0)Q^{\pi }(s_0,a_0)+\sum _{a_0}\pi (a_0\mid s_0)\sum _{s_1}p(s_1\mid s_0,a_0)\gamma {\mathrm{\nabla }}_{\theta }{V}^{\pi }(s_1)\\ & =\sum _{a_0}{\mathrm{\nabla }}_{\theta }\pi (a_0\mid s_0)Q^{\pi }(s_0,a_0)\\ & +\sum _{a_0}\pi (a_0\mid s_0)\sum _{s_1}p(s_1\mid s_0,a_0)\gamma \sum _{a_1}{\mathrm{\nabla }}_{\theta }\pi (a_1\mid s_1)Q^{\pi }(s_1,a_1)\\ & +\sum _{a_0}\pi (a_0\mid s_0)\sum _{s_1}p(s_1\mid s_0,a_0)\gamma \sum _{a_1}\pi (a_1\mid s_1)\sum _{s_2}p(s_2\mid s_1,a_1)\gamma {\mathrm{\nabla }}_{\theta }{V}^{\pi }(s_2)\\ & =\sum _{a_0}{\mathrm{\nabla }}_{\theta }\pi (a_0\mid s_0)Q^{\pi }(s_0,a_0)\\ & +\sum _{a_0}\pi (a_0\mid s_0)\sum _{s_1}p(s_1\mid s_0,a_0)\gamma \sum _{a_1}{\mathrm{\nabla }}_{\theta }\pi (a_1\mid s_1)Q^{\pi }(s_1,a_1)+\cdots \\ & =\sum _{s_0}Pr(s_0\to s_0,0,\pi )\sum _{a_0}{\mathrm{\nabla }}_{\theta }\pi (a_0\mid s_0){\gamma }^0{Q}^{\pi }(s_0,a_0)\\ & +\sum _{s_1}Pr(s_0\to s_1,1,\pi )\sum _{a_1}{\mathrm{\nabla }}_{\theta }\pi (a_1\mid s_1){\gamma }^1{Q}^{\pi }(s_1,a_1)+\cdots \end{array}$$$$\begin{array}{rl}& =\sum _{s_0}Pr(s_0\to s_0,0,\pi )\sum _{a_0}\pi (a_0\mid s_0)[{\gamma }^0{Q}^{\pi }(s_0,a_0){\mathrm{\nabla }}_{\theta }\log \pi (a_0\mid s_0)]\\ & +\sum _{s_1}Pr(s_0\to s_1,1,\pi )\sum _{a_1}\pi (a_1\mid s_1)[{\gamma }^1{Q}^{\pi }(s_1,a_1){\mathrm{\nabla }}_{\theta }\log \pi (a_1\mid s_1)]+\cdots \\ & =\sum _{t=0}^{\mathrm{\infty }}\sum _{s_t}Pr(s_0\to s_t,t,\pi )\sum _{a_t}\pi (a_t\mid s_t)[{\gamma }^t{Q}^{\pi }(s_t,a_t){\mathrm{\nabla }}_{\theta }\log \pi (a_t\mid s_t)]\end{array}$$

其中 $Pr(s_0\to s_t,t,\pi )$ 代表从状态 $s_0$ 出发，按照策略 $\pi$ 与环境交互（rollout），在 $t$ 时刻到达状态 $s_t$ 的概率。

通过上述推导，就得到了无限长时间步下的策略梯度表达式。对于有限长时间步的环境，可以做一个简单的转化使其同样适用上述公式。假设时间步长度为 $T$，对于所有可能出现在最后一步的状态 $s_{T-1}$，定义：

1. 从 $s_{T-1}$ 出发，不论采取什么动作，一定会跳转到一个虚拟的吸收态 $s_T$，并返回奖励值 $0$。
2. 从 $s_T$ 出发，不论采取什么动作，一定会跳转回这个虚拟的吸收态 $s_T$，并返回奖励值 $0$。

由此将有限长的时间步扩展到了无限长，因为环境会陷入到 $s_T$ 的死循环中。

不过，上式实际上很难优化，因为它要求遍历整个状态空间和时间步空间，计算每个时间步上到达每个状态的概率。一方面，这在计算成本上无法容忍；另一方面，在绝大多数情况下无法获得环境的转移概率，因此无法计算特定时间步下整个状态空间上的概率分布。

为此，可以使用 Monte Carlo 方法，通过采样来逼近上面的策略梯度公式。先把上式转化为期望的形式：

$$\begin{array}{rl}& \sum _{t=0}^{\mathrm{\infty }}\sum _{s_t}Pr(s_0\to s_t,t,\pi )\sum _{a_t}\pi (a_t\mid s_t)[{\gamma }^t{Q}^{\pi }(s_t,a_t){\mathrm{\nabla }}_{\theta }\log \pi (a_t\mid s_t)]\\ & =\sum _{t=0}^{\mathrm{\infty }}{\mathbb{E}}_{s_t}\sum _{a_t}\pi (a_t\mid s_t)[{\gamma }^t{Q}^{\pi }(s_t,a_t){\mathrm{\nabla }}_{\theta }\log \pi (a_t\mid s_t)]\\ & =\sum _{t=0}^{\mathrm{\infty }}{\mathbb{E}}_{s_t,a_t}[{\gamma }^t{Q}^{\pi }(s_t,a_t){\mathrm{\nabla }}_{\theta }\log \pi (a_t\mid s_t)]\\ & ={\mathbb{E}}_{s_0,a_0,s_1,a_1,\dots }\sum _{t=0}^{\mathrm{\infty }}[{\gamma }^t{Q}^{\pi }(s_t,a_t){\mathrm{\nabla }}_{\theta }\log \pi (a_t\mid s_t)]\\ & ={\mathbb{E}}_{\tau }\sum _{t=0}^{\mathrm{\infty }}[{\gamma }^t{Q}^{\pi }(s_t,a_t){\mathrm{\nabla }}_{\theta }\log \pi (a_t\mid s_t)]\end{array}$$

其中 $\tau =[s_0,a_0,s_1,a_1,\dots ]$ 是按照策略 $\pi$ rollout 出来的状态动作轨迹。可以看出，将 ${\gamma }^t{Q}^{\pi }(s_t,a_t){\mathrm{\nabla }}_{\theta }\log \pi (a_t\mid s_t)$ 这一项先在时间步 $t$ 上求和，再关于轨迹 $\tau$ 取期望，就得到了策略梯度。至此，Monte Carlo 方法就可以很简单地结合进来：先将 ${\mathbb{E}}_{\tau }$ 替换为采样 $N$ 条轨迹 $[{\tau }^1,\dots ,{\tau }^N]$，并定义其中第 $n$ 条轨迹为

$${\tau }^n=⟨s_0^n,a_0^n,r_0^n,\dots ,s_{T_n-1}^n,a_{T_n-1}^n,r_{T_n-1}^n⟩,$$

轨迹长度为 $T_n$。最后对结果取平均：

$${\mathbb{E}}_{\tau }\sum _{t=0}^{\mathrm{\infty }}[{\gamma }^t{Q}^{\pi }(s_t,a_t){\mathrm{\nabla }}_{\theta }\log \pi (a_t\mid s_t)]=\frac{1}{N}\sum _{n=1}^N\sum _{t=0}^{T_n-1}[{\gamma }^t{Q}^{\pi }(s_t^n,a_t^n){\mathrm{\nabla }}_{\theta }\log \pi (a_t^n\mid s_t^n)]$$

注意到

$$Q^{\pi }(s_t^n,a_t^n)={\mathbb{E}}_{s_{t+1}^n,a_{t+1}^n,s_{t+2}^n,a_{t+2}^n,\dots \mid s_t^n,a_t^n}\left[\sum _{l=t}^{T_n-1}{\gamma }^{l-t}{r}_l^n\right],$$

因此从期望角度二者也是可以替换的。当算法没有显式地估计 $Q^{\pi }(s_t^n,a_t^n)$ 时，可以定义

$$G_t({\tau }^n)=\sum _{l=t}^{T_n-1}{\gamma }^{l-t}{r}_l^n,$$

即最朴素的回报，并用它替换 $Q^{\pi }(s_t^n,a_t^n)$，就得到实际使用的策略梯度公式：

$$\frac{1}{N}\sum _{n=1}^N\sum _{t=0}^{T_n-1}[{\gamma }^t{G}_t({\tau }^n){\mathrm{\nabla }}_{\theta }\log \pi (a_t^n\mid s_t^n)]$$