强化学习之自然梯度法

自然梯度法，即Natural Gradient，是一种优化算法，其他常见的优化算法比如一维搜索、牛顿法、最速下降法、共轭梯度法等。个人认为自然梯度的思想非常类似于最速下降法，只不过一个是寻找最优的标量$\alpha \in \mathbb{R}$，而NG是在一定约束下寻找最优的搜索方向$s^{\ast }=\alpha d$(严格来说搜索方向应该不包括学习率，但是为了便于说明，接下来的搜索方向都暗指包括学习率)。

另外NG算法还被广泛用于强化学习中，比如NPG、TRPO以及PPO算法都涉及到了自然梯度法思想的影响，因此对于强化学习领域的研究者来说，NG还是有必要了解的。

参考列表：
①费舍尔信息矩阵及自然梯度法或英文原版
②TRPO与NPG
③TRPO共轭梯度优化

1. 优化目标与约束

自然梯度法是梯度下降法的一个分支，因此参数的更新遵从：
$${\theta }^{(k+1)}\leftarrow {\theta }^{(k)}-{\alpha }^{(k)}\cdot g^{(k)}$$其中$g$为梯度，作为一个优化算法，这里设目标函数为$f(\theta )$，则目标函数为：
$${\theta }^{\ast }=\underset{\theta }{\mathrm{argmin}}f(\theta )$$那么由于在机器学习中，我们很难直接求得解析解，故一般都是通过迭代的方式去求得数值解，因此梯度下降法里面最重要的自然就是学习率以及搜索方向了，不同的梯度下降法也就是搜索方向或者学习率的选择不同。那么在Natural Gradient中的选择是这样的：
$$\begin{gathered} s^{\ast }=\underset{d}{\mathrm{argmin}}f({\theta }^k+s^{(k)}) \\ s.t.D_{KL}({\theta }^k,{\theta }^{(k+1)})\le ϵ \end{gathered}$$其表达的意思是：参数搜索范围在参数空间受限的情况下，找到最小化目标函数$f$的搜索方向$s^{\ast }$，这就是自然梯度法。在NG里约束参数搜索范围采用了$KL-divergence$法。

2. 必要知识点

在推导出自然梯度发最优搜索方向$d^{\ast }$的过程中，还需要了解一些必备知识点，比如①费雪信息矩阵FIM的定义②FIM与黑塞矩阵的关系③FIM与KL散度的关系等等。

2.1 费雪信息矩阵

什么是费雪信息矩阵(Fisher Information Matrix)？
定义某个分布关于模型$\theta$的分布为$p(x|\theta )$。
记得分函数$s_p(\theta )={\nabla }_{\theta }\log p(x|\theta ),s_p\in {\mathbb{R}}^d$，关于得分函数，在$KSD$中也出现过，只不过在$KSD$中是${\mathbb{E}}_p[{\mathcal{A}}_{p}f(x)]=0$，而在这里我们有${\mathbb{E}}_p[s_p(\theta )]=0$，证明如下：然后给出协方差的定义：$cov(X,Y)=\mathbb{E}[(X-\mathbb{E}X)\cdot (Y-\mathbb{E}Y)]$，因此得分函数的协方差表示为：
$$cov(s_p,s_p)=\mathbb{E}[(s_p-0)\cdot (s_p-0{)}^T]$$我们可以将其看成一种信息，因此可以说得分函数的协方差就是费雪信息。又因为费雪信息是以矩阵${\mathbb{R}}^{d\times d}$形式存在的，因此将得分函数展开：$$F={\mathbb{E}}_p[{\nabla }_{\theta }\log p(x|\theta )\cdot {\nabla }_{\theta }^T\log p(x|\theta )]$$
Note：

1. $F$就是费雪信息矩阵FIM。
2. 用样本估计的方式近似FIM的$\hat{F}$叫经验费雪信息矩阵：$$F\approx \hat{F}=\frac{1}{N}\sum _{i=1}^N{\nabla }_{\theta }\log p(x|\theta )\cdot {\nabla }_{\theta }^T\log p(x|\theta )$$

2.2 FIM&Hessian

这一节讨论FIM与黑塞矩阵$H$之间的关系。
记$H_{\log p(x|\theta )}=\frac{{\partial }^2\log p(x|\theta )}{\partial {\theta }^2}$，则：
$$H_{\log p(x|\theta )}=J({\nabla }_{\theta }\log p(x|\theta ))=J(\frac{{\nabla }_{\theta }p(x|\theta )}{p(x|\theta )})$$Note：

1. $J$表示雅克比矩阵(Jocobin-matrix)。

对上式求期望可得：
$${\mathbb{E}}_{x\sim p}[H_{\log p(x|\theta )}]=-F$$证明如下：
记住这个结论：
$$F=-{\mathbb{E}}_p[H_{\log p(x|\theta )}]$$，意思就是说得分函数的一阶导数期望取反就是费雪信息矩阵$F$，接下来会用到。

2.3 FIM&KL-divergence

做强化学习的，对于KL散度应该不陌生了，之前在高斯平滑$Q$函数、TRPO、PPO这些论文中都出现过，这里还是啰嗦的给出其公式：
$$\begin{gathered} D_{KL}(p(x)||q(x))=\int p(x)\cdot \log \frac{p(x)}{q(x)}\mathrm{d}x={\mathbb{E}}_p[\log \frac{p(x)}{q(x)}] \\ ={\mathbb{E}}_p[\log p(x)]-{\mathbb{E}}_p[\log q(x)] \end{gathered}$$Note:

1. KL散度衡量的就是两个分布之间相似程度。这里不用距离的概念，是因为KL散度不是对称函数，因此不能算是距离。
2. KL散度的特性曲线如下：

记$\theta$为旧参数${\theta }_{old}$，${\theta }^{\prime }=\theta +s$为新参数${\theta }_{new}$，那么相对熵运用到NG算法中就是：
$$D_{KL}(p(x|\theta )||p(x|{\theta }^{\prime }))={\mathbb{E}}_{p(x|\theta )}[\log p(x|\theta )]-{\mathbb{E}}_{p(x|\theta )}[\log p(x|{\theta }^{\prime })]$$
接下来我们试着求取$KL$散度的一阶导数和二阶导数。
一阶导数和二阶导数所求如下：

也就是说$KL$散度在${\theta }^{\prime }=\theta$处的Hessian矩阵就是费雪信息矩阵$F$。

3. 最优搜索方向

接下来就开始推导出自然梯度法所优化出的搜索方向$s^{\ast }$了。
我们把目标函数和约束重新列一下：
$$\begin{gathered} \underset{\theta }{\min }f({\theta }^{\prime }) \\ s.t.kl({\theta }^{\prime })\le ϵ \end{gathered}$$

接下来我们对目标函数和约束用泰勒公式在${\theta }^{\prime }=\theta$处展开
对于目标函数：
$$f({\theta }^{\prime })\approx f(\theta )+g^T\cdot ({\theta }^{\prime }-\theta )+高阶项$$Note：

1. 这里我们忽略高阶项，直接做一阶近似。
2. 第一项是常数项，一般优化都省略掉。或者你也可以先留着，到后面就会发现常数项没有用。
3. $g=\frac{\partial f}{\partial \theta }{|}_{{\theta }^{\prime }=\theta }$。

对于约束项：
$$kl(\theta ||{\theta }^{\prime })\approx kl(\theta ||\theta )+{\nabla }_{{\theta }^{\prime }}^{T}kl(\theta ||{\theta }^{\prime }){|}_{{\theta }^{\prime }=\theta }({\theta }^{\prime }-\theta )+\frac{1}{2}({\theta }^{\prime }-\theta {)}^T{\nabla }_{{\theta }^{\prime }}^{2}kl(\theta ||{\theta }^{\prime }){|}_{{\theta }^{\prime }=\theta }({\theta }^{\prime }-\theta )+高阶项$$Note：

1. 第一项为0，第二项根据我们2.3节推导来看也是为0，因此只剩下二阶项，因此$KL$散度我们都是做二阶近似的。
2. 令$s={\theta }^{\prime }-\theta$，然后结合2.3节的重要结论，则KL散度可近似为：$$kl(\theta ||{\theta }^{\prime })\approx \frac{1}{2}{s}^{T}Fs$$这里的$F$就是费雪信息矩阵，是一个Hessian矩阵。从这里也可以看出$KL$散度的二阶导数就是FIM。

然后我们将近似后的约束问题通过拉格朗日乘子法转为无约束问题：
$$minimizeL({\theta }^{\prime },\lambda )=f(\theta )+g^T\cdot s+\lambda (\frac{1}{2}{s}^{T}Fs-ϵ)$$整理加入$KKT$条件：
$$\begin{gathered} minimizeL({\theta }^{\prime },\lambda )=\frac{\lambda }{2}{s}^{T}Fs+g^T\cdot s+f(\theta )-\lambda ϵ \\ KKT：\{\begin{array}{l}\frac{1}{2}{s}^{T}Fs\le ϵ\\ \lambda =0\\ \lambda (\frac{1}{2}{s}^{T}Fs-ϵ)=0\end{array} \end{gathered}$$显然我们现在的目的在于优化出最好的$s^{\ast }$使得拉格朗日函数$L$达到最小值：因此我们求$\frac{\partial L}{\partial s}=0$。
易得：
$$\begin{gathered} \lambda Fs+g=0 \\ \iff s^{\ast }=-\frac{1}{\lambda }{F}^{-1}g \end{gathered}$$Note：

1. 至此我们就推导出了自然梯度法每次迭代的最佳搜索方向$s^{\ast }$。我们在TRPO里面也得出过这个式子，只不过那个是求$maximize$，因此不带“负号”。
2. 前面说过$s$包含了学习率和梯度$g$的乘积，因此常数$\frac{1}{\lambda }$可以吸收到学习率里面去，因此NG算法真正意义上的搜索方向为：$$F^{-1}g$$
3. 再来看这个$KKT$条件中的第三个，$KL$散度用于约束参数搜索范围在满足目标函数最小化的同时不要太大，因此在边界上的$KL$值是最好的，根据$KKT$性质所知，此时$\lambda \ge 0$，也就是说：
   $$\begin{gathered} \frac{1}{2}{s}^{T}Fs=ϵ \\ \iff \frac{1}{2}(\alpha g{)}^{T}F(\alpha g)=ϵ \\ \iff \alpha =\sqrt{\frac{2ϵ}{g^{T}Fg}} \end{gathered}$$

至此我们总算可以得出NG最重要的更新公式了：
$${\theta }^{(k+1)}\leftarrow {\theta }^k-\sqrt{\frac{2ϵ}{g^{T}Fg}}\cdot F^{-1}g$$

4. TRPO&NPG

这里给出了NG算法应用于强化学习中形成的Natural Policy Gradient算法，可以看到它最下面的更新公式，就是我们上述推导出来的。
NPG伪代码：
这一份是TRPO的伪代码，可以说TRPO进一步优化了NPG算法，关键在于TRPO可以推出这么一个式子：
$$\theta -{\theta }_{old}=s=\frac{1}{\lambda }{F}^{-1}g$$而TRPO采用共轭梯度优化算法可以通过迭代求出共轭搜索方向$d$，乘上学习率就可以求得$s$，因此可以用$d$来近似NPG算法中的搜索方向$F^{-1}g$，即$d\approx F^{-1}g$。这样最大的好处就是不用计算费雪信息矩阵，因为FIM的计算量是非常大的。
TRPO伪代码：

5. 总结

1. 将上述整合在一起，就是自然梯度下降法的伪代码：
2. 自然梯度法(NG)和最速下降法(steepest descent)很相似，只是两者优化的目标不一样，前者是为了求出使得$f(\theta +s)$最小的最佳搜索方向，后者是为了求出使得$f(\theta +s)$最小的最佳学习率。此外，共轭梯度算法(CG)可以对NG算法进行改进：CG通过迭代的方式去近似掉NG中的FIM。三者都是梯度下降法的分支，都可以应用到机器学习算法中去。
3. 此外，在推导出NG算法的过程中，认识到了FIM的定义、FIM与Hessian矩阵的关系、FIM与$KL$散度的关系。