重要性采样与PPO、TRPO

重要性采样与PPO、TRPO

本文是一篇学习笔记，主要参考李宏毅老师的强化学习课程。

on/off policy

on-policy 和 off-policy 的区分是强化学习中的重要概念。on-policy 是指我们要训练的 agent 与跟环境互动产生数据的 agent 是同一个 agent；off-policy 则是指二者是不同的 agent。以下围棋为例来直观地理解，on-policy 就是你自己下完棋之后，进行 ”复盘“，反思自己棋招中考虑欠妥的地方，是否导致了最终的失利；off-policy 则是 ”打谱“，研究别人的棋谱，学习高手制胜的妙招。

在我们之前介绍的强化学习方法中，Q-learning / DQN 是 off policy 的方法，而 policy gradient（如式 1 所示）每次更新要求其经验数据来自于当前要学习的 policy，因此是 on policy 的方法。
$$\begin{gathered} {\nabla }_{\theta }J({\pi }_{\theta })={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)R(\tau )\right] \\ \text{\hspace{1em}(1)} \end{gathered}$$
on policy 的问题非常明显：样本效率太低了。每次跑出来的经验数据样本只能用来更新一次 policy，之后就需要使用更新后的 policy 重新采样数据。因此，我们希望将 on-policy 的方法转换成 off-policy 的方法，用某个固定的 ${\pi }_{{\theta }^{\prime }}$ 来采样数据（做 demonstration，示范），样本可以多次用来更新目标 policy ${\pi }_{\theta }$，这样就能大大提高样本效率。

如何才能将 on-policy 的策略梯度转变成 off-policy 的呢？这就要用到重要性采样方法，下面先来介绍一下该方法。

重要性采样及其方差分析

我们想求函数 $f(x)$ 在概率 $p$ 下的期望：
$$E_{x\sim p}[f(x)]$$
按照期望的定义，应该是求积分：
$$E_{x\sim p}[f(x)]=\int p(x)f(x)dx$$
一般情况下，我们是没办法直接求上面的积分的。但是我们能够从分布 $p$ 中采样 $N$ 个样本 $x^i{|}_{i=1\dots N}$，从而近似地估计出这个积分：
$$E_{x\sim p}[f(x)]\approx \frac{1}{N}\sum _{i=1}^{N}f(x^i)$$
这就是蒙特卡洛估计，是一种无偏估计。

更进一步，假设我们甚至无法直接从 $p$ 分布采样数据或者从 $p$ 分布采样的效率不高，该怎么办呢？这时，就要用到重要性采样（Importance Samping）。考虑从更容易采样的 $q$ 分布中采样，并进行如下转换：
$$\begin{array}{rl}{E}_{x\sim p}[f(x)]& =\int p(x)f(x)dx\\ & =\int f(x)\frac{p(x)}{q(x)}q(x)dx\\ & =E_{x\sim q}\left[f(x)\frac{p(x)}{q(x)}\right]\end{array}$$
这里我们将想要求的期望展开成了积分的形式，然后上下同乘 $q(x)$，就把从 $p$ 分布下的期望转换成了 $q$ 分布下的期望，只需要多乘一个修正系数 $\frac{p(x)}{q(x)}$。重要性采样也是无偏的。

这就是重要性采样方法，通过简单的变换，乘一个修正权重 $\frac{p(x)}{q(x)}$ 将在对目标分布 $p$ 期望的估计，转换从 $q$ 分布中采样来估计。

虽然期望能够这样转换，但是方差呢？我们考察一下转换前后的方差：
$$\begin{gathered} 转换前方差：Va{r}_{x\sim p}[f(x)] \\ 转换后方差：Va{r}_{x\sim q}\left[f(x)\frac{p(x)}{q(x)}\right] \end{gathered}$$
我们知道，方差与期望的关系：
$$Var(X)=E[X^2]-(E[X]{)}^2$$

从而有：
$$\begin{array}{rl}转换前方差：Va{r}_{x\sim p}[f(x)]& =E_{x\sim p}[f(x{)}^2]-(E_{x\sim p}[f(x)]{)}^2\\ & \\ 转换后方差：Va{r}_{x\sim q}\left[f(x)\frac{p(x)}{q(x)}\right]& =E_{x\sim q}\left[{\left(f(x)\frac{p(x)}{q(x)}\right)}^2\right]-{\left(E_{x\sim q}\left[f(x)\frac{p(x)}{q(x)}\right]\right)}^2\\ & =E_{x\sim p}\left[f(x{)}^2\frac{p(x)}{q(x)}\right]-(E_{x\sim p}\left[f(x)\right]{)}^2把期望展开化简一下\end{array}$$
对比转换前后的方差，后面一项是相同的，区别在于前面一项，转换后的方差少了一个 $\frac{p(x)}{q(x)}$。

考虑一个极端的情况，$p,q,f$ 如下图所示。很明显期望 $E_{x\sim p}[f(x)]$ 是负值，重要性采样是无偏的，因此修正后的期望 $E_{x\sim q}\left[f(x)\frac{p(x)}{q(x)}\right]$ 肯定也是负值。然而，我们从 $q$ 分布中采样，大概率采到的是正值，只有当采样足够充分，从 $q$ 中采到一次负值，这时 $f(x)$ 非常大，因此会乘上一个非常大的权重，把期望拉回成负值。

可以看到，当 $p,q$ 差异非常大时，必须要采样足够充分，这正是高方差的问题。因此，在重要性采样中，我们最好保证 $\frac{p(x)}{q(x)}$ 足够小，也就是 $p,q$ 两个分布尽可能接近，才能尽可能降低重要性采样的方差，在采样次数有限的情况下，保证正确性。

从 on policy 到 off policy

回到强化学习上，我们想把 on-policy 的策略梯度（式 1），转换成样本效率更高的 off-policy 算法，需要我们要学习的 policy ${\pi }_{\theta }$ 能够根据来自其他 policy ${\pi }_{{\theta }^{\prime }}$ 与环境互动取得的经验数据来学习，采用刚刚介绍的重要性采样，我们需要乘上一个修正系数：
$${\nabla }_{\theta }J({\pi }_{\theta })={\mathbb{E}}_{\tau \sim {\pi }_{{\theta }^{\prime }}}\left[\sum _{t=0}^T\frac{{\pi }_{\theta }(\tau )}{{\pi }_{{\theta }^{\prime }}(\tau )}{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)R(\tau )\right]$$
这样，${\pi }_{\theta }$ 就可以根据 ${\pi }_{{\theta }^{\prime }}$ 采样出来的数据进行多次更新，提高样本效率。

实际中，我们通常采用优势函数 $A^{\theta }(s_t,a_t)$ 来替代累积回报 $R(\tau )$，优势函数表示了状态 $s_t$ 下动作 $a_t$ 相较于其他动作 “好” 多少，一般认为它能比累积回报更好地反映动作的好坏。这样，使用重要性采样修正，将 on policy 转换为 off policy 的策略梯度可以写为：

$$\begin{array}{rl}{\nabla }_{\theta }J(\theta )& =E_{(s_t,a_t)\sim {\pi }_{\theta }}[A^{\theta }(s_t,a_t)\nabla \log {\pi }_{\theta }(a_t|s_t)]\\ & =E_{(s_t,a_t)\sim {\pi }_{{\theta }^{\prime }}}\left[\frac{{\pi }_{\theta }(s_t,a_t)}{{\pi }_{{\theta }^{\prime }}(s_t,a_t)}{A}^{{\theta }^{\prime }}(s_t,a_t){\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)\right]这里A^{\theta }需要替换成A^{{\theta }^{\prime }}\\ & =E_{(s_t,a_t)\sim {\pi }_{{\theta }^{\prime }}}\left[\frac{{\pi }_{\theta }(s_t|a_t)}{{\pi }_{{\theta }^{\prime }}(s_t|a_t)}\cancel{\frac{{\pi }_{\theta }(s_t)}{{\pi }_{{\theta }^{\prime }}(s_t)}}{A}^{{\theta }^{\prime }}(s_t,a_t){\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)\right]这里的\frac{{\pi }_{\theta }(s_t)}{{\pi }_{{\theta }^{\prime }}(s_t)}直接丢掉了\end{array}$$
根据这个策略梯度，我们可以写出对应目标函数：
$$\begin{gathered} J^{{\theta }^{\prime }}({\pi }_{\theta })=E_{(s_t,a_t)\sim {\pi }_{{\theta }^{\prime }}}\left[\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }^{\prime }}(a_t|s_t)}{A}^{{\theta }^{\prime }}(s_t,a_t)\right] \\ \text{\hspace{1em}(2)} \end{gathered}$$

PPO / TRPO

虽然通过应用重要性采样，将期望目标函数改写成了从 ${\pi }_{{\theta }^{\prime }}$ 采样的形式（式 2），但根据我们上面对方差的分析，要稳妥地应用重要性采样，需要保证两个分布 ${\pi }_{\theta }$ 和 ${\pi }_{{\theta }^{\prime }}$ 不能差太大。这正是 TRPO / PPO 要做的事情。

TRPO

TRPO 可以算是 PPO 的前身，它将两个 policy 输出分布的 KL 散度作为了一个硬约束项，要求它小于 $\delta$，从而来保证两个 policy 差异不会太大。其训练目标为：
$$J_{\text{TRPO}}^{{\theta }^{\prime }}({\pi }_{\theta })=J^{{\theta }^{\prime }}({\pi }_{\theta }),\text{KL}[{\pi }_{\theta }(\cdot |s_t)||{\pi }_{{\theta }^{\prime }}(\cdot |s_t)]<\delta J^{{\theta }^{\prime }}({\pi }_{\theta })见式(2)$$
带约束的梯度上升是一个非常复杂的优化问题，因此 TRPO 论文中有大量困难的数学推导，理论理解和实际实现起来都比较麻烦。

PPO-KL

PPO 直接将两个 policy 输出分布的 KL 散度加到了目标函数中，作为一个类似正则项的软约束：
$$J_{\text{PPO-KL}}^{{\theta }^{\prime }}({\pi }_{\theta })=J^{{\theta }^{\prime }}({\pi }_{\theta })-\beta \text{KL}[{\pi }_{\theta }(\cdot |s_t)||{\pi }_{{\theta }^{\prime }}(\cdot |s_t)]J^{{\theta }^{\prime }}({\pi }_{\theta })见式(2)$$

这样实现起来就简单多了，并且实际性能也完全不比 TRPO 差。

PPO-KL 还提出了自适应的 KL penalty 调整，设置两个阈值 ${\text{KL}}_{\text{max}},{\text{KL}}_{\text{min}}$，根据训练过程中的 KL 值动态调整系数 $\beta$。如果训练过程中 KL 值超过了 ${\text{KL}}_{\text{max}}$，就说明 KL 的约束太弱了，两 policy 之间的差异过大，此时调高 $\beta$，增强约束；如果 KL 值小于 ${\text{KL}}_{\text{min}}$ 就说明 KL 约束太强了，两 policy 趋近相同也不是我们想要的，因为这样就没法优化策略 ${\pi }_{\theta }$ 了，此时调低 $\beta$，减弱约束。

注意这里的 KL 散度指的是两个 policy 输出分布 $\pi (\cdot |s)$ 的 KL 散度，而非参数 $\theta ,{\theta }^{\prime }$ 本身的 KL 散度。

PPO-CLIP

PPO 还有另一种形式 PPO-CLIP，也可以约束两 policy 之间的差异。其训练目标为：
$$J_{\text{PPO-CLIP}}^{{\theta }^{\prime }}({\pi }_{\theta })=E_{(s_t,a_t)\sim {\pi }^{\prime }}\left[\text{min}\left(\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }^{\prime }}(a_t|s_t)}{A}^{{\theta }^{\prime }}(s_t,a_t),\text{clip}\left(\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }^{\prime }}(a_t|s_t)},1-ϵ,1+ϵ\right)A^{{\theta }^{\prime }}(s_t,a_t)\right)\right]$$
结合下图来理解 PPO-CLIP 目标函数的意思。图中横轴表示 $\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }^{\prime }}(a_t|s_t)}$ 的值，纵轴表示输出值，绿色和蓝色虚线分别表示 min 函数中的第一、第二项，红线表示 min 函数的最终输出。

左图中，当优势函数 $A>0$ 时，说明当前动作是比较好的，此时我们想要增加状态 $s_t$ 下动作 $a_t$ 出现的概率，即增加 ${\pi }_{\theta }(a_t|s_t)$，但是，也不希望它增加到太大，一旦 $\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }^{\prime }}(a_t|s_t)}$ 超过了 $1+ϵ$，再增加就不再给它收益了。同理，当优势函数 $A<0$ 时，说明动作是比较差的，我们想要减小 ${\pi }_{\theta }(a_t|s_t)$，但是如果 $\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }^{\prime }}(a_t|s_t)}$ 小于了 $1-ϵ$，再减小也没有收益了。这样，我们就在尽可能优化 ${\pi }_{\theta }$ 的同时，保证 ${\pi }_{\theta }$ 和 ${\pi }_{{\theta }^{\prime }}$ 的差异不会太大，能够稳妥地应用重要性采样来提高样本效率。

总结

本文中，我们首先介绍了强化学习中 on/off policy 的概念，指出 off policy 的样本效率更高，因此我们想将 on policy 的策略梯度法转换为 off policy 的方法。我们借助重要性采样来实现这一转换，但同时也发现重要性采样的方差存在一定隐患，需要保证两个分布之间差异不要太大，这便是 TRPO / PPO 要解决的问题。TRPO 采用了 KL 硬约束来限制两分布的差异，PPO 则提出了 KL 软约束和 CLIP 两种方案。TRPO 和 PPO 在实际强化学习任务中表现差不太多，但是 PPO 从理论和实现上都更加简单，因此成为比较流行的强化学习方法。