步步深入 TRPO

论文《Trust Region Policy Optimization》[1] 提出了大名鼎鼎的 TRPO 算法，这是 policy gradient 系列强化学习（RL）算法的里程碑之作。但原论文包含大量晦涩难懂的公式和定理，对于入门者并不友好。本文将详细讲解 TRPO 中关键公式的推导过程，希望能够理清 TRPO 作者想解决的问题以及采用的方法。

1. 引言

TRPO 和大多数 RL 算法一样，希望提升策略 $\pi$ 的期望累积回报 $\eta (\pi )$：

$$\eta (\pi )={\mathbb{E}}_{s_0,a_0,\dots }\left[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^t{r}_t\right],$$

其中每一步的动作 $a_t\sim \pi (a_t\mid s_t)$，服从策略 $\pi$ 所决定的动作概率分布。

在进一步分析 $\eta (\pi )$ 的性质之前，定义三个函数：动作价值函数 $Q^{\pi }(s_t,a_t)$、状态价值函数 $V^{\pi }(s_t)$ 和优势函数 $A^{\pi }(s_t,a_t)$。

$$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$ 的情况下的累积期望回报。

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

即在状态 $s_t$ 下，后续动作服从策略 $\pi$ 的情况下的累积期望回报。

$$A^{\pi }(s_t,a_t)=Q^{\pi }(s_t,a_t)-V^{\pi }(s_t),$$

表示在状态 $s_t$ 下，直接采用动作 $a_t$ 相比于按照 $a_t\sim \pi (a_t\mid s_t)$ 采样动作的优势。

如何提升 $\eta (\pi )$ 呢？或者说，如何找到一个新的策略 $\tilde{\pi }$ 使得 $\eta (\tilde{\pi })$ 高于 $\eta (\pi )$ 呢？这就需要分析 $\eta (\tilde{\pi })$ 和 $\eta (\pi )$ 的定量关系。这里引用一个 RL 领域的经典结论 [2]：

$$\begin{array}{}\text{(1)}& \eta (\tilde{\pi })=\eta (\pi )+{\mathbb{E}}_{s_0,a_0,\dots \sim \tilde{\pi }}[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^t{A}_{\pi }(s_t,a_t)].\end{array}$$

这里的 $A_{\pi }(s_t,a_t)$ 表示策略 $\pi$ 下的优势函数，也就是说：

$$A_{\pi }(s_t,a_t)=Q_{\pi }(s_t,a_t)-V_{\pi }(s_t),$$

其中动作价值函数和状态价值函数都对应策略 $\pi$。

式 (1) 可以这么理解：$\eta (\tilde{\pi })$ 与 $\eta (\pi )$ 的差等于按照策略 $\tilde{\pi }$ 采样动作 $a_t$、在走出的轨迹中每一步的策略 $\pi$ 下优势函数 $A_{\pi }(s_t,a_t)$ 的累积和。

下面给出简单证明：

$$\begin{array}{rl}\eta (\tilde{\pi })& ={\mathbb{E}}_{s_0,a_0,\dots \sim \tilde{\pi }}\left[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^t{r}_t\right]\\ & =\eta (\pi )+{\mathbb{E}}_{s_0,a_0,\dots \sim \tilde{\pi }}\left[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^t{r}_t\right]-\eta (\pi )\\ & =\eta (\pi )+{\mathbb{E}}_{s_0,a_0,\dots \sim \tilde{\pi }}[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^t{r}_t-V_{\pi }(s_0)]\\ & =\eta (\pi )+{\mathbb{E}}_{s_0,a_0,\dots \sim \tilde{\pi }}\left[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^t(r_t+\gamma V_{\pi }(s_{t+1})-V_{\pi }(s_t))\right]\\ & =\eta (\pi )+{\mathbb{E}}_{s_0,a_0,\dots \sim \tilde{\pi }}[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^t{A}_{\pi }(s_t,a_t)].\end{array}$$

注意论文原文的证明也是一样的裂项相消，区别只是这里写成了连等式的过程。实际上这个等式给了我们很强的指引：满足

$${\mathbb{E}}_{s_0,a_0,\dots \sim \tilde{\pi }}[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^t{A}_{\pi }(s_t,a_t)]\ge 0$$

的策略 $\tilde{\pi }$ 可以使得 $\eta (\tilde{\pi })\ge \eta (\pi )$。但是策略 $\tilde{\pi }$ 并没有显式出现在式中，因此接下来需要变形。

首先定义累积折扣状态访问频率：

$${\rho }_{\pi }(s)=P(s_0=s)+\gamma P(s_1=s)+{\gamma }^{2}P(s_2=s)+\cdots$$

这个函数的含义是：在策略 $\pi$ 下，每一步状态 $s_t$ 等于 $s$ 的概率在折扣系数下的累积和。基于 ${\rho }_{\pi }(s)$ 的定义，对式 (1) 变形得到：

$$\text{begin{aligned} eta(tilde{pi}) \&= eta(pi) + sum\_{t=0}^{infty} sum\_s P(s\_t = s mid tilde{pi}) sum\_a tilde{pi}(a mid s) gamma^t A\_{pi}(s, a) \&= eta(pi) + sum\_s sum\_{t=0}^{infty} gamma^t P(s\_t = s mid tilde{pi}) sum\_a tilde{pi}(a mid s) A\_{pi}(s, a) \&= eta(pi) + sum\_s rho\_{tilde{pi}}(s) sum\_a tilde{pi}(a mid s) A\_{pi}(s, a). tag{2} end{aligned}}$$

通过式 (2) 不难看出，只要新策略 $\tilde{\pi }$ 满足：对于每个状态 $s$，

$$\sum _a\tilde{\pi }(a\mid s)A_{\pi }(s,a)\ge 0,$$

即可保证 $\eta (\tilde{\pi })\ge \eta (\pi )$。

很完美对不对？只要最大化 $\sum _a\tilde{\pi }(a\mid s)A_{\pi }(s,a)$ 就可以得到更优的策略 $\tilde{\pi }$。其实不然，如果在这里画上句号，就没 TRPO 什么事了。下面，真正的 TRPO 即将开始。

2. 前菜：为什么需要近似

在 RL 算法的实际应用中，我们通常通过神经网络来学习一个策略 $\pi$，即输入状态 $s$，输出这个状态下选择每个动作 $a$ 的概率 $\pi (a\mid s)$。既然是参数化的神经网络，就难免有误差。换句话说，“对于每个状态 $s$，$\sum _a\tilde{\pi }(a\mid s)A_{\pi }(s,a)\ge 0$” 这个完美条件难以成立，总有一些隐藏的坏点状态 $s$ 使得 $\sum _a\tilde{\pi }(a\mid s)A_{\pi }(s,a)<0$。

怎么办？其实也好办也不好办。

好办的是，根据式 (2)，

$$\eta (\tilde{\pi })-\eta (\pi )=\sum _s{\rho }_{\tilde{\pi }}(s)\sum _a\tilde{\pi }(a\mid s)A_{\pi }(s,a),$$

那么只要让整体

$$\sum _s{\rho }_{\tilde{\pi }}(s)\sum _a\tilde{\pi }(a\mid s)A_{\pi }(s,a)\ge 0$$

就行了，中间每个状态上的 $\sum _a\tilde{\pi }(a\mid s)A_{\pi }(s,a)$ 是正是负并不需要考虑。

不好办的是，想求 $\sum _s{\rho }_{\tilde{\pi }}(s)\sum _a\tilde{\pi }(a\mid s)A_{\pi }(s,a)$ 对策略 $\tilde{\pi }$ 的导数非常困难，因为 ${\rho }_{\tilde{\pi }}(s)$ 的导数我们搞不到。

3. 正篇 1：替代函数

这该怎么办？把难搞的东西给 ban 掉，换成相应的替代。一个很自然的想法就是把 ${\rho }_{\tilde{\pi }}(s)$ 换成 ${\rho }_{\pi }(s)$，于是定义一个近似的替代函数：

$$L_{\pi }(\tilde{\pi })=\eta (\pi )+\sum _s{\rho }_{\pi }(s)\sum _a\tilde{\pi }(a\mid s)A_{\pi }(s,a).$$

$L_{\pi }(\tilde{\pi })$ 比 $\eta (\tilde{\pi })$ 好处理多了，${\rho }_{\pi }(s)$ 中不包含策略 $\tilde{\pi }$，因此对策略 $\tilde{\pi }$ 的导数为零。但是，$L_{\pi }(\tilde{\pi })$ 既然是近似，必然有误差。这种近似的效果如何呢？

效果还真不错。观察发现，$L_{\pi }(\tilde{\pi })$ 和 $\eta (\tilde{\pi })$ 在 $\tilde{\pi }=\pi$ 处的值和对 $\tilde{\pi }$ 的导数都是一样的，也就是：

$$L_{\pi }(\pi )=\eta (\pi ),\text{且}{{\mathrm{\nabla }}_{\tilde{\pi }}{L}_{\pi }(\tilde{\pi })|}_{\tilde{\pi }=\pi }={{\mathrm{\nabla }}_{\tilde{\pi }}\eta (\tilde{\pi })|}_{\tilde{\pi }=\pi }.$$

第一个等式（值相等）一眼就可以看出来，第二个需要稍微证明一下：

$${{\mathrm{\nabla }}_{\tilde{\pi }}\eta (\tilde{\pi })|}_{\tilde{\pi }=\pi }={\sum _s{\rho }_{\pi }(s){\mathrm{\nabla }}_{\tilde{\pi }}|}_{\tilde{\pi }=\pi }\sum _a\tilde{\pi }(a\mid s)A_{\pi }(s,a)+{\sum _s{\mathrm{\nabla }}_{\tilde{\pi }}|}_{\tilde{\pi }=\pi }{\rho }_{\tilde{\pi }}(s)\sum _a\pi (a\mid s)A_{\pi }(s,a).$$

注意到 $\sum _a\pi (a\mid s)A_{\pi }(s,a)=0$，代入可得：

$${{\mathrm{\nabla }}_{\tilde{\pi }}\eta (\tilde{\pi })|}_{\tilde{\pi }=\pi }={\sum _s{\rho }_{\pi }(s){\mathrm{\nabla }}_{\tilde{\pi }}|}_{\tilde{\pi }=\pi }\sum _a\tilde{\pi }(a\mid s)A_{\pi }(s,a)={{\mathrm{\nabla }}_{\tilde{\pi }}{L}_{\pi }(\tilde{\pi })|}_{\tilde{\pi }=\pi }.$$

这种在 $\tilde{\pi }=\pi$ 处的值和梯度都相等的情况，叫做 $L_{\pi }(\tilde{\pi })$ 是对 $\eta (\tilde{\pi })$ 的一阶近似。结合数学分析的知识可以知道，当

$${{\mathrm{\nabla }}_{\tilde{\pi }}{L}_{\pi }(\tilde{\pi })|}_{\tilde{\pi }=\pi }={{\mathrm{\nabla }}_{\tilde{\pi }}\eta (\tilde{\pi })|}_{\tilde{\pi }=\pi }\ne 0$$

时，$\tilde{\pi }=\pi$ 处必然存在一个邻域，域内的 $\tilde{\pi }$ 满足：若 $L_{\pi }(\tilde{\pi })$ 增大，则 $\eta (\tilde{\pi })$ 也增大。这说明，在一定步长内优化 $L_{\pi }(\tilde{\pi })$，会使得 $\eta (\tilde{\pi })$ 也得到优化。

4. 正篇 2：信赖域

仅凭替代函数对优化函数的一阶近似性质，我们只能知道在一定步长内提升 $L_{\pi }(\tilde{\pi })$ 会使得 $\eta (\tilde{\pi })$ 也提升，但还不知道步长要选多大。文章标题中的关键词 trust region（信赖域）正是探讨优化的步长要在什么范围（域）内选择。

沿着这个思路出发，文章的核心贡献点之一就是进一步量化了 $L_{\pi }(\tilde{\pi })$ 和 $\eta (\tilde{\pi })$ 的关系，提出了下面的不等式：

$$\eta (\tilde{\pi })\ge L_{\pi }(\tilde{\pi })-\frac{4ϵ\gamma }{(1-\gamma {)}^2}\alpha ,$$

其中

$$\alpha =\underset{s}{max}{D}_{\mathrm{KL}}(\pi (\cdot \mid s)\parallel \tilde{\pi }(\cdot \mid s)),ϵ=\underset{s,a}{max}|A_{\pi }(s,a)|.$$

这一步基本解决了步长的问题。因为我们得到了 $\eta (\tilde{\pi })$、$L_{\pi }(\tilde{\pi })$ 以及步长（这里表现为 $\tilde{\pi }$ 与 $\pi$ 之间的 KL 散度）的定量关系。这个关系表现为 $\frac{4ϵ\gamma }{(1-\gamma {)}^2}\alpha$ 这个惩罚项：步长越大，惩罚就越大，此时 $\eta (\tilde{\pi })$ 越难以享受到提升 $L_{\pi }(\tilde{\pi })$ 所带来的优化效果。

有了定量关系就好办了，可以直接把优化目标从 $L_{\pi }(\tilde{\pi })$ 改为：

$$M_{\pi }(\tilde{\pi })=L_{\pi }(\tilde{\pi })-\frac{4ϵ\gamma }{(1-\gamma {)}^2}\alpha .$$

注意到 $M_{\pi }(\tilde{\pi })$ 是 $\eta (\tilde{\pi })$ 的下界，我们希望优化 $M_{\pi }(\tilde{\pi })$ 来提升 $\eta (\tilde{\pi })$。下面证明：优化 $M_{\pi }(\tilde{\pi })$ 得到的最优解 $\overline{\pi }$ 一定是更好的策略，即 $\eta (\overline{\pi })\ge \eta (\pi )$。

首先注意到两个事实：

1. $M_{\pi }(\pi )=L_{\pi }(\pi )-0=\eta (\pi )-0=\eta (\pi )$；
2. $M_{\pi }(\overline{\pi })\ge M_{\pi }(\pi )$。

第一个由于 $\tilde{\pi }=\pi$ 时惩罚项为零，第二个则利用了 $\overline{\pi }$ 为最优解这个特性。那么证明就一目了然：

$$\eta (\overline{\pi })\ge M_{\pi }(\overline{\pi })\ge M_{\pi }(\pi )=\eta (\pi ).$$

到此，证明了直接优化 $M_{\pi }(\tilde{\pi })$ 就可以得到更优的策略。

5. 正篇 3：优化

$M_{\pi }(\tilde{\pi })$ 的优化面临两个难点：

1. $\underset{s}{max}{D}_{\mathrm{KL}}(\pi (\cdot \mid s)\parallel \tilde{\pi }(\cdot \mid s))$ 难以计算；
2. $\underset{s,a}{max}|A_{\pi }(s,a)|$ 也难以计算。

对于第一个难点，TRPO 作者将 $\underset{s}{max}{D}_{\mathrm{KL}}(\pi (\cdot \mid s)\parallel \tilde{\pi }(\cdot \mid s))$ 用 ${\mathbb{E}}_{s\sim {\rho }_{\pi }}{D}_{\mathrm{KL}}(\pi (\cdot \mid s)\parallel \tilde{\pi }(\cdot \mid s))$ 进行了近似替换，因为 $s\sim {\rho }_{\pi }$ 这个分布上的期望可以用蒙特卡洛法解决。

第二个难点其实是惩罚项系数的问题，这个系数中的 $\underset{s,a}{max}|A_{\pi }(s,a)|$ 也是个难确定的东西。所以 TRPO 作者直接就把最大化 $M_{\pi }(\tilde{\pi })$ 换成了它的对偶问题，也就是从：

$$\underset{\tilde{\pi }}{max}{M}_{\pi }(\tilde{\pi })=L_{\pi }(\tilde{\pi })-\frac{4ϵ\gamma }{(1-\gamma {)}^2}\alpha$$

变成：

$$\begin{array}{c}\underset{\tilde{\pi }}{max}{L}_{\pi }(\tilde{\pi })\\ \text{s.t.}{\mathbb{E}}_{s\sim {\rho }_{\pi }}{D}_{\mathrm{KL}}(\pi (\cdot \mid s)\parallel \tilde{\pi }(\cdot \mid s))\le \delta .\end{array}$$

这是一种偷懒方式：不管 $\underset{s,a}{max}|A_{\pi }(s,a)|$ 是多大，总存在对应的 $\delta$ 使得对偶问题和原问题一致。那么 $\delta$ 具体取多大？那就是调参的事了，哪个值性能好用哪个。

实际训练时，策略是用参数化的网络来实现的。用 $\theta$ 和 $\tilde{\theta }$ 来表示更新前后策略 $\pi$ 和 $\tilde{\pi }$ 的参数，因此优化目标可以表示为 $L_{\theta }(\tilde{\theta })$。文章中优化 $L_{\theta }(\tilde{\theta })$ 时，采用了通过 Fisher information matrix 来计算自然梯度（natural gradient）的方法，具体建模为：

$$\begin{array}{c}\underset{\tilde{\theta }}{max}{L}_{\theta }(\tilde{\theta })\\ \text{s.t.}\frac{1}{2}(\tilde{\theta }-\theta {)}^{\mathsf{T}}A(\tilde{\theta },\theta )(\tilde{\theta }-\theta )\le \delta ,\end{array}$$

其中 $A(\tilde{\theta },\theta )$ 是 KL 散度关于 $\tilde{\theta }$ 的 Hessian 矩阵，也就是 Fisher information matrix。求解这个问题可以直接套用自然梯度的公式。令 $g={\mathrm{\nabla }}_{\theta }{L}_{\theta }(\tilde{\theta })$，则自然梯度的方向为 $A^{-1}g$。考虑到 $\delta$ 对步长的约束，实际更新使用的自然梯度为：

$$\sqrt{\frac{2\delta }{g^{\mathsf{T}}{A}^{-1}g}}{A}^{-1}g.$$

参考文献

1. Schulman, J., Levine, S., Abbeel, P., Jordan, M. and Moritz, P., 2015. Trust region policy optimization. In International conference on machine learning (pp. 1889-1897). PMLR.
2. Kakade, S. and Langford, J., 2002. Approximately optimal approximate reinforcement learning. In Proceedings of the 19th International Conference on Machine Learning.