Trust Region Policy Optimization (TRPO) 公式推导

一. 优化问题的构建

1. 由期望累计奖励出发

记期望累计奖励$\eta (\pi )$为:
$$\eta (\pi )={\mathbb{E}}_{s_0\sim \rho (s_0),a_t\sim \pi (a|s),s_{t+1}\sim P(s_{t+1}|s_t,a_t)}[\sum _{t=0}^{\infty }{\gamma }^{t}r(s_t)],\text{\hspace{1em}(1)}$$其中$\rho (s_0)$为初始状态$s_0$的分布。可以证明：
$$\eta (\tilde{\pi })=\eta (\pi )+{\mathbb{E}}_{s_0\sim \rho (s_0),a_t\sim \tilde{\pi }(a|s),s_{t+1}\sim P(s_{t+1}|s_t,a_t)}[\sum _{t=0}^{\infty }{\gamma }^t{A}_{\pi }(s_t,a_t)],\text{\hspace{1em}(2)}$$具体证明过程参考TRPO论文Appendix A Lemma1。注意，公式(2)期望中$a_t$由新策略进行采样。记${\rho }_{\pi }(s)=P(s_0=s)+\gamma P(s_1=s)+...,$则公式(2)可以改写为：
$$\begin{gathered} \eta (\tilde{\pi })=\eta (\pi )+\sum _{t=0}^{\infty }\sum _{s}P(s_t=s|\tilde{\pi })\sum _a\tilde{\pi }(a|s){\gamma }^t{A}_{\pi }(s,a) \\ =\eta (\pi )+\sum _s\sum _{t=0}^{\infty }{\gamma }^{t}P(s_t=s|\tilde{\pi })\sum _a\tilde{\pi }(a|s)A_{\pi }(s,a) \\ =\eta (\pi )+\sum _s{\rho }_{\tilde{\pi }}(s)\sum _a\tilde{\pi }(a|s)A_{\pi }(s,a).\text{\hspace{1em}(3)} \end{gathered}$$公式(3)等号左边为新策略$\tilde{\pi }$的期望累计奖励$\eta (\tilde{\pi })$，右边为旧策略$\pi$的期望累计奖励$\eta (\pi )$加上某一项${\sum }_s{\rho }_{\tilde{\pi }}(s){\sum }_a\tilde{\pi }(a|s)A_{\pi }(s,a)$。TRPO的思想就是通过最大化${\sum }_s{\rho }_{\tilde{\pi }}(s){\sum }_a\tilde{\pi }(a|s)A_{\pi }(s,a)$，以提升策略的性能。但不能直接对${\sum }_s{\rho }_{\tilde{\pi }}(s){\sum }_a\tilde{\pi }(a|s)A_{\pi }(s,a)$进行优化，因为${\rho }_{\tilde{\pi }}(s)$导致了优化目标中含有未知项，即在未得到$\tilde{\pi }$之前，不知道${\rho }_{\tilde{\pi }}(s)$，直接求导十分复杂， 难以计算。

2. 对目标函数进行近似与简化

首先，由于上述提到的问题，我们将公式(3)简化为:
$$L_{\pi }(\tilde{\pi })=\eta (\pi )+\sum _s{\rho }_{\pi }(s)\sum _a\tilde{\pi }(a|s)A_{\pi }(s,a),\text{\hspace{1em}(4)}$$注意在公式(4)中，我们将${\rho }_{\tilde{\pi }}(s)$替换为了${\rho }_{\pi }(s)$。公式(4)具有以下的性质：
$$\begin{gathered} L_{{\pi }_{{\theta }_0}}({\pi }_{{\theta }_0})=\eta ({\pi }_{{\theta }_0}), \\ {\nabla }_{\theta }{L}_{{\pi }_{{\theta }_0}}({\pi }_{\theta }){|}_{\theta ={\theta }_0}={\nabla }_{\theta }{\eta }_{{\pi }_{{\theta }_0}}({\pi }_{\theta }){|}_{\theta ={\theta }_0},\text{\hspace{1em}(5)} \end{gathered}$$性质的证明见此公式： 公式(5)证明。下面的定理展示了$L_{\pi }(\tilde{\pi })$和$\eta (\tilde{\pi })$之间的关系:
$$\eta (\tilde{\pi })\ge L_{\pi }(\tilde{\pi })-C{D}_{TV}^{max}(\pi ,\tilde{\pi }),$$其中$C=\frac{4ϵ\gamma }{(1-\gamma {)}^2},ϵ={\max }_{s,a}{A}_{\pi }(s,a),$ $D_{TV}^{max}(\pi ,\tilde{\pi })={\max }_s{D}_{TV}(\pi ,\tilde{\pi })={\max }_s\frac{1}{2}{\sum }_i|{\pi }_i(\cdot |s)-{\tilde{\pi }}_i(\cdot |s)|$。证明见TRPO论文Appendix A Lemma2、3。文章用KL散度替换了TV散度，由于性质$D_{TV}(\pi ,\tilde{\pi }{)}^2\le D_{KL}(\pi ,\tilde{\pi })$，则可以得到：
$$\eta (\tilde{\pi })\ge L_{\pi }(\tilde{\pi })-C{D}_{KL}^{max}(\pi ,\tilde{\pi }),\text{\hspace{1em}(6)}$$其中$C=\frac{4ϵ\gamma }{(1-\gamma {)}^2},ϵ={\max }_{s,a}{A}_{\pi }(s,a)。$通过最大化等式右边$L_{\pi }(\tilde{\pi })-C{D}_{KL}^{max}(\pi ,\tilde{\pi })$，可以保证策略性能$\eta (\tilde{\pi })$的单调上升。

3. 将优化问题进一步简化

记$\pi ={\pi }_{{\theta }_{old}},\tilde{\pi }={\pi }_{\theta }$，则优化问题可以构建为：
$${\text{maximize}}_{\theta }[L_{{\theta }_{old}}(\theta )-C{D}_{KL}^{max}({\theta }_{old},\theta )].$$文章中指出，如果按照理论中设置C的大小，则每次更新的步长非常小，效率很低。因此，文章用信赖域的约束优化问题代替上面的问题：
$$\begin{gathered} {\text{maximize}}_{\theta }{L}_{{\theta }_{old}}(\theta ), \\ \text{s.t.}{D}_{KL}^{max}({\theta }_{old},\theta )\le \delta . \end{gathered}$$文章进一步将约束中的KL散度由最大值约束简化为了平均值约束：
$$\begin{gathered} {\text{maximize}}_{\theta }{L}_{{\theta }_{old}}(\theta ), \\ \text{s.t.}{\bar{D}}_{KL}^{{\rho }_{{\theta }_{old}}}({\theta }_{old},\theta )\le \delta . \end{gathered}$$其中${\bar{D}}_{KL}^{\rho }({\theta }_1,{\theta }_2)={\mathbb{E}}_{s\sim \rho }[D_{KL}({\pi }_{{\theta }_1}(\cdot |s),{\pi }_{{\theta }_2}(\cdot |s))]$。由于在强化学习中，一般通过采样的方式收集样本，上述的优化目标还可以写成期望的形式：
$$\begin{gathered} {\text{maximize}}_{\theta }\sum _s{\rho }_{{\theta }_{old}}(s)\sum _{a\sim {\pi }_{\theta }}{\pi }_{\theta }(a|s)A_{{\theta }_{old}}(s,a), \\ \text{s.t.}{\bar{D}}_{KL}^{{\rho }_{{\theta }_{old}}}({\theta }_{old},\theta )\le \delta . \end{gathered}$$在上式中，${\text{maximize}}_{\theta }{\sum }_s{\rho }_{{\theta }_{old}}(s)[...]$代表根据分布$s\sim {\rho }_{{\theta }_{old}}(s)$对状态进行采样，可以写成${\mathbb{E}}_{s\sim {\rho }_{{\theta }_{old}}(s)}[...]$，同时，由于在收集样本时，一般是采用某些随机策略$q(\cdot |s)$，以增加对状态空间的探索，因此${\sum }_a{\pi }_{\theta }(a|s)$与数据中$a_t$来源于$q(\cdot |s)$不一致，通过重要性采样将其改写为${\sum }_{a\sim q}\frac{{\pi }_{\theta }(a|s)}{q(\cdot |s)}$。综上，优化问题可以改写为：
$$\begin{gathered} {\text{maximize}}_{\theta }{\mathbb{E}}_{s\sim {\rho }_{{\theta }_{old}},a\sim q}\frac{{\pi }_{\theta }(a|s)}{q(\cdot |s)}{A}_{{\theta }_{old}}(s,a), \\ \text{s.t.}{\bar{D}}_{KL}^{{\rho }_{{\theta }_{old}}}({\theta }_{old},\theta )\le \delta .\text{\hspace{1em}(7)} \end{gathered}$$

二. 优化问题求解

文章采用顺序二次规划的思想，将目标函数线性化、将约束进行二阶Taylor展开，通过迭代的方式逐步求解问题(7)。

1. 目标函数的线性化

记$f(\theta )=\frac{{\pi }_{\theta }(a|s)}{q(\cdot |s)}{A}_{{\theta }_{old}}(s,a)$，则$f(\theta )=f({\theta }_{old})+{\nabla }_{\theta }f(\theta ){|}_{\theta ={\theta }_{old}}(\theta -{\theta }_{old})$。其中${\nabla }_{\theta }f(\theta ){|}_{\theta ={\theta }_{old}}$
可以通过自动微分计算。

2. 约束的二阶Taylor展开

KL散度的二阶Taylor展开可以写为以下形式:
$$g(\theta )=D_{KL}(p_{{\theta }_0},p_{\theta })=D_{KL}(p_{{\theta }_0},p_{{\theta }_0})+{\nabla }_{\theta }{D}_{KL}(p_{{\theta }_0},p_{\theta }){|}_{\theta ={\theta }_0}(\theta -{\theta }_0)+\frac{1}{2}(\theta -{\theta }_0{)}^{T}F(\theta -{\theta }_0)+O({\theta }^3).$$其中$D_{KL}(p_{{\theta }_0},p_{{\theta }_0})=0,$
$$\begin{gathered} {\nabla }_{\theta }{D}_{KL}(p_{{\theta }_0},p_{\theta }){|}_{\theta ={\theta }_0}={\nabla }_{\theta }\int p_{{\theta }_0}(x)[\log p_{{\theta }_0}(x)-\log p_{\theta }(x)]dx{|}_{\theta ={\theta }_0} \\ =-\int p_{{\theta }_0}(x){\nabla }_{\theta }\log p_{\theta }(x)dx{|}_{\theta ={\theta }_0} \\ =-\int p_{{\theta }_0}(x)\frac{{\nabla }_{\theta }{p}_{\theta }(x)}{p_{\theta }(x)}dx{|}_{\theta ={\theta }_0} \\ =-{\nabla }_{\theta }\int p_{\theta }(x)dx{|}_{\theta ={\theta }_0} \\ =-{\nabla }_{\theta }1 \\ =0. \end{gathered}$$$F$为Fisher信息矩阵，定义可见Fisher信息矩阵的定义。具体推导过程可见KL散度的二阶Taylor展开。

综上，优化问题可以写为：
$$\begin{gathered} {\text{maximize}}_{\theta }{\mathbb{E}}_{s\sim {\rho }_{{\theta }_{old}},a\sim q}[f({\theta }_{old})+{\nabla }_{\theta }f(\theta ){|}_{\theta ={\theta }_{old}}(\theta -{\theta }_{old})], \\ s.t.\frac{1}{2}(\theta -{\theta }_0{)}^{T}F(\theta -{\theta }_0)\le \delta .\text{\hspace{1em}(8)} \end{gathered}$$

3.高效求解优化问题

由问题(8)，根据Language乘子法，可知${\nabla }_{\theta }f(\theta ){|}_{\theta ={\theta }_{old}}+\lambda F(\theta -{\theta }_0)=0,$即更新方向$g=(\theta -{\theta }_0)$满足
$$Fg=b,\text{\hspace{1em}(9)}$$其中$F$为Fisher信息矩阵，$g=(\theta -{\theta }_0),b={\nabla }_{\theta }f(\theta ){|}_{\theta ={\theta }_{old}}$。文中采用共轭梯度法对(9)式进行快速计算。通过(9)式求得更新方向$g$后，假设步长为$\beta$,需满足KL散度的约束，即
$$\delta =D_{KL}\approx \frac{1}{2}(\beta g{)}^{T}A(\beta g)=\frac{1}{2}{\beta }^2{g}^{T}Ag.$$解得$\beta =\sqrt{2\delta /g^{T}Ag}$。文中从$\beta =\sqrt{2\delta /g^{T}Ag}$开始进行线搜索，求得使$L_{{\theta }_{old}}(\theta )-\mathbb{I}(D_{KL}({\theta }_{old},\theta )\le \delta )$取得最大值的${\beta }^{\ast }$，其中$\mathbb{I}(\cdot )$为示性函数，当条件满足时取0，不满足时取$+\infty$。