强化学习 Proximal Policy Optimization (PPO)

参考: 李宏毅老师课件

PPO: Default reinforcement learning algorithm at OpenAI

PPO = Policy Gradient 从 On-policy 到 Off-policy, 再加一些constraint

Policy Gradient#

Basic Conception#

• Actor: 动作执行者(智能体)

• Env: 环境

• Reward Function: 奖励函数

• Policy $\pi$ : a network with parameter $\theta$.

  Input: 当前的 Env.

  Output: actor 要采取的下一个 action 的分布.

• Trajectory $\tau$: 一系列的 Env 和 Action, $\{s_1,a_1,s_2,a_2,\dots \}$
  
  在参数为 $\theta$ 情况下, 发生$\tau$的概率: $p_{\theta }(\tau )=p(s_1)p_{\theta }(a_1|s_1)p(s_2|s_1,a_1)p_{\theta }(a_2|s_2)\cdots$

Optimization#

Object#

给定 $\tau$, 可以计算 $\tau$ 的 reward, $R(\tau )$.

对于参数为 $\theta$ 的 Policy下, Trajectory $\tau$ 是采样得到的, 因此实际上需要计算的是 reward 的期望值$\overline{R_{\theta }}$. 我们希望 $\overline{R_{\theta }}$ 越大越好.

Policy Gradient#

Reward 的期望:

$$\begin{array}{}\text{(1)}& \begin{array}{r}\overline{R_{\theta }}=\sum _{\tau }R(\tau )p_{\theta }(\tau )\end{array}\end{array}$$

求 $\theta$ 的梯度:

$$\begin{array}{}\text{(2)}& \begin{array}{rl}\mathrm{\nabla }{\bar{R}}_{\theta }& =\sum _{\tau }R(\tau )\mathrm{\nabla }{p}_{\theta }(\tau )\\ & =\sum _{\tau }R(\tau )p_{\theta }(\tau )\frac{\mathrm{\nabla }{p}_{\theta }(\tau )}{p_{\theta }(\tau )}& & \text{分子分母同乘}{p}_{\theta }(\tau )\\ & =\sum _{\tau }R(\tau )p_{\theta }(\tau )\mathrm{\nabla }\log p_{\theta }(\tau )\\ & =E_{\tau \sim p_{\theta }(\tau )}[R(\tau )\mathrm{\nabla }\log p_{\theta }(\tau )]\\ & \approx \frac{1}{N}\sum _{n=1}^{N}R({\tau }^n)\mathrm{\nabla }\log p_{\theta }({\tau }^n)\\ & =\frac{1}{N}\sum _{n=1}^N\sum _{t=1}^{T_n}R({\tau }^n)\mathrm{\nabla }\log p_{\theta }(a_t^n|s_t^n)\end{array}\end{array}$$

由 $\mathrm{\nabla }\log p_{\theta }(\tau )=\frac{\mathrm{\nabla }{p}_{\theta }(\tau )}{p_{\theta }(\tau )}$, 可得到第三行公式.
此处可延伸出一个公式:

$$\begin{array}{}\text{(3)}& \mathrm{\nabla }f(x)=f(x)\mathrm{\nabla }\log f(x)\end{array}$$

由$\sum _{\tau }{p}_{\theta }(\tau )f(\tau )=E_{\tau \sim p_{\theta }(\tau )}[f(\tau )]$, 可得第四行

通过采样的方式估计期望值, 采样 $N$ 个 Trajectory, 既第五行公式

最后将 $p_{\theta }(\tau )$ 展开代入, 得第六行公式

Implementation#

最大化 Reward 的期望 $\overline{R_{\theta }}$, 由公式(2)中梯度的计算, 可以反推出目标函数在实现时定义如下:

$$\begin{array}{}\text{(4)}& \begin{array}{r}J(\theta )=\frac{1}{N}\sum _{n=1}^N\sum _{t=1}^{T_n}R({\tau }^n)\log p_{\theta }(a_t^n|s_t^n)\end{array}\end{array}$$

最大化 $object$ 等价于最小化 $loss$:

$$\begin{array}{}\text{(5)}& \begin{array}{r}loss=-\frac{1}{N}\sum _{n=1}^N\sum _{t=1}^{T_n}R({\tau }^n)\log p_{\theta }(a_t^n|s_t^n)\end{array}\end{array}$$

其中, $a_t^n,s_t^n$ 是在参数为 $\theta$ 的 policy 下采样得到的.

与交叉熵损失对比: 其实就是将采样得到的 $a_t^n$ 视作grand truth计算交叉熵, 区别在于针对不同的 Trajectory ${\tau }^n$, 要多乘了一个 $R({\tau }^n)$

对于RLHF来说, 就是：

1. 对于某个指令 $I$, 采样多条输出 $O=\{o_1,o_2,\cdots ,o_n\}$(即多个Trajectory, 其中每个 $o_i$ 表示一个Trajectory) 或者用多个指令，每个指令采样一个输出，不重要，不影响后续操作。
   2.把这些 $O$ 作为label, 用teacher forcing的方式计算交叉熵损失, 不过在对于每个样本 $o_i$, 其损失最终要额外乘一个奖励值$R(o_i)$

Tips#

Add a baseline#

$R({\tau }^n)$ 可能总为正数, 这样在 training时, 相当于告诉 model, 不论时什么action 都要将它的概率提升.

理想情况下, 这样是没有问题的, 因为 Reward 即使总是正的, 也有大有小.

当时实际上, action 是采样得到的, 这会导致如果有的 action 没有被采样到, 它的概率相对于被采样到的 action 就会下降, 而这时, 并不能表示当前环境下采取这个 action 不好. 改进: 减去一个 baseline, $b$.

直观上的解释是, 如果一个 $\tau$ 中的 action 导致reward 增加的不够多, 我们也认为这个 action 不够好. 改进后的目标函数如下:

$$\begin{array}{}\text{(6)}& \begin{array}{r}J(\theta )=\frac{1}{N}\sum _{n=1}^N\sum _{t=1}^{T_n}(R({\tau }^n)-b)\log p_{\theta }(a_t^n|s_t^n)\end{array}\end{array}$$

Assign Suitable Credit#

再来看一下目标函数:

$$\begin{array}{}\text{(7)}& \begin{array}{r}J(\theta )=\frac{1}{N}\sum _{n=1}^N\sum _{t=1}^{T_n}(R({\tau }^n)-b)\log p_{\theta }(a_t^n|s_t^n)\end{array}\end{array}$$

对于同一个 Trajectory $\tau$ 中, 针对每个状态 $s$ 下, 执行动作 $a$, 都有相同的 Reward 系数. 这是不合理的.
例如图的左边, 在 $s_b$ 执行 $a_2$ 不是一个好的选择, 他会导致接下来进入 $s_c$, 并执行 $a_3$, 得到 -2 分.
由此, 提出改改进1: 每个时刻的 reward 改为, 当前时刻到结束时刻的 reward 的总和.

某时刻的 action, 经过越长时间, 它的影响力就越小. 也就是与该 action 间隔很久的 reward 与该 action 的关系很小. 由此提出改进2: 加一个衰减系数 ${\gamma }^{t^{\prime }-t}$.

最后, 将整个系数项称为 Advantage Function, $A^{\theta }(s_t,a_t)$.其含义为, 在某 state 下, $a_t$ 相较于其他的 action, 有多好. (这个 $A$, 通常可以是用一个网络来预测，一般称作criic)

对于LLM来说，可以参考这篇文章。

最终得目标函数与梯度的公式如下:

$$\begin{array}{}\text{(8)}& \begin{array}{rl}J(\theta )& =\frac{1}{N}\sum _{n=1}^N\sum _{t=1}^{T_n}{A}^{\theta }(s_t,a_t)\log p_{\theta }(a_t^n|s_t^n)\\ \mathrm{\nabla }J(\theta )& =\frac{1}{N}\sum _{n=1}^N\sum _{t=1}^{T_n}{A}^{\theta }(s_t,a_t)\mathrm{\nabla }\log p_{\theta }(a_t^n|s_t^n)\end{array}\end{array}$$

*公式中的 $A^{\theta }(s_t,a_t)$ 不是指 $A$ 中含有参数 $\theta$, 只是代表从 $(s_t,a_t)$ 是从 $\theta$ 中采样来的.

On-policy $\to$ Off-policy#

On-policy#

梯度计算公式:

$$\begin{gathered} \begin{array}{}\text{(9)}& \mathrm{\nabla }{\bar{R}}_{\theta }=E_{\tau \sim p_{\theta }(\tau )}[R(\tau )\mathrm{\nabla }\log p_{\theta }(\tau )] \\ \end{array} \end{gathered}$$

目前为止的做法其实是一种 on-policy 的方法:

• 每次更新梯度前, 都需要从 ${\pi }_{\theta }$ 中采样 $\tau$.
• 参数更新后, 又需要用更新后的参数重新采样 $\tau$.

目标是: 从另一个 policy, ${\pi }_{{\theta }^{\prime }}$ 中采样数据, 用来训练 ${\pi }_{\theta }$. 这样就可以重复利用这些采样得到的数据.

Importance Sampling(重要性采样)#

$x$ 服从 $p$ 分布时, 计算 $f(x)$ 期望 $E_{x\sim p}[f(x)]$ 的做法: 一般是从 $p$ 中采样一些 $x$, 带入 $f(x)$ 求平均, 用这个值来估计所求期望.

现在, 假设无法从 $p$ 中直接采样 $x$, 但可以从另一个分布 $q$ 中采样 $x$. 可以对 $E_{x\sim p}[f(x)]$ 做如下变形:

$$\begin{array}{}\text{(10)}& \begin{array}{rl}{E}_{x\sim p}[f(x)]& =\int f(x)p(x)dx\\ & =\int f(x)\frac{p(x)}{q(x)}q(x)dx\\ & =E_{x\sim q}[f(x)\frac{p(x)}{q(x)}]\end{array}\end{array}$$

这样, 我们就可以用 $q$ 中采样的数据来估计期望值 $E_{x\sim p}[f(x)]$. 这就是 Importance Sampling.

Issue of Importance Sampling
理论上, 我们已经得出两个期望值是相等的:

$$\begin{array}{}\text{(11)}& E_{x\sim p}[f(x)]=E_{x\sim q}[f(x)\frac{p(x)}{q(x)}].\end{array}$$

那么它们的方差是否相等呢? $Va{r}_{x\sim p}[f(x)]==Va{r}_{x\sim q}[f(x)\frac{p(x)}{q(x)}]?$

由公式

$$\begin{array}{}\text{(12)}& Var[x]=E[x^2]-(E[x]{)}^2\end{array}$$

可以得出:

$$\begin{array}{}\text{(13)}& \begin{array}{rl}Va{r}_{x\sim p}[f(x)]& =E_{x\sim p}[f^2(x)]-(E_{x\sim p}[f(x)]{)}^2\\ Va{r}_{x\sim q}[f(x)\frac{p(x)}{q(x)}]& =E_{x\sim q}[(f(x)\frac{p(x)}{q(x)}{)}^2]-(E_{x\sim q}[f(x)\frac{p(x)}{q(x)}]{)}^2\\ & =\int (f(x)\frac{p(x)}{q(x)}{)}^{2}q(x)dx-(E_{x\sim p}[f(x)]{)}^2\\ & =\int f^2(x)\frac{p(x)}{q(x)}p(x)dx-(E_{x\sim p}[f(x)]{)}^2\\ & =E_{x\sim p}[f^2(x)\frac{p(x)}{q(x)}]-(E_{x\sim p}[f(x)]{)}^2\end{array}\end{array}$$

对比发现, 第一项中后者比前者多乘了一个 $\frac{p(x)}{q(x)}$, 也就是说当 $p$ 与 $q$ 相差很多时, 它们的方差也会差很多.

这样就会出现一问题: 理论上, 无论 $p,q$ 的分布是什么样的, 当我们从 $p$ 和 $q$ 采样足够多次时, 是可以得到 $E_{x\sim p}[f(x)]=E_{x\sim q}[f(x)\frac{p(x)}{q(x)}]$ 的.
但是当 $p,q$ 差距过大, 而我们采样的次数又不够多时, 因为它们之间的方差差距很大, 所以最后很可能导致期望差距很大.

一个直观的例子:

图中 $p,q$两个分布的差异很大.

当我们采样次数不够多, 导致没有采样到最左边那个样本时, 就会出现实际上 $E_{x\sim p}[f(x)]$ 应是一个负值, 但我们用 $E_{x\sim q}[f(x)\frac{p(x)}{q(x)}]$ 计算出来的却是一个正值.

而当我们采样到最左边那个样本时, 因为此时 $\frac{p(x)}{q(x)}$ 的值将会非常大, 所以可以把 $E_{x\sim q}[f(x)\frac{p(x)}{q(x)}]$ 拉回负值.

Off-policy#

将 Importance Sampling 用在 policy gradient 中, 我们就可以得到:

$$\begin{array}{}\text{(14)}& \begin{array}{rl}\mathrm{\nabla }{\bar{R}}_{\theta }& =E_{\tau \sim p_{\theta }(\tau )}[R(\tau )\mathrm{\nabla }\log p_{\theta }(\tau )]\\ & =E_{\tau \sim p_{{\theta }^{\prime }}(\tau )}[\frac{p_{\theta }(\tau )}{p_{{\theta }^{\prime }}(\tau )}R(\tau )\mathrm{\nabla }\log p_{\theta }(\tau )]\end{array}\end{array}$$

这样, 我们就可以从 ${\theta }^{\prime }$ 中采样数据, 然后多次利用这些数据来更新 $\theta$.

结合之前添加Tips后得到的公式(8), 有:

$$\begin{array}{}\text{(15)}& \begin{array}{rl}\mathrm{\nabla }{\bar{R}}_{\theta }& =E_{(s_t,a_t)\sim {\pi }_{\theta }}[A^{\theta }(s_t,a_t)\mathrm{\nabla }\log p_{\theta }(a_t^n|s_t^n)]\\ & =E_{(s_t,a_t)\sim {\pi }_{{\theta }^{\prime }}}[\frac{p_{\theta }(s_t,a_t)}{p_{{\theta }^{\prime }}(s_t,a_t)}{A}^{{\theta }^{\prime }}(s_t,a_t)\mathrm{\nabla }\log p_{\theta }(a_t^n|s_t^n)]\\ & =E_{(s_t,a_t)\sim {\pi }_{{\theta }^{\prime }}}[\frac{p_{\theta }(a_t|s_t)p_{\theta }(s_t)}{p_{{\theta }^{\prime }}(a_t|s_t)p_{{\theta }^{\prime }}(s_t)}{A}^{{\theta }^{\prime }}(s_t,a_t)\mathrm{\nabla }\log p_{\theta }(a_t^n|s_t^n)]\\ & =E_{(s_t,a_t)\sim {\pi }_{{\theta }^{\prime }}}[\frac{p_{\theta }(a_t|s_t)}{p_{{\theta }^{\prime }}(a_t|s_t)}{A}^{{\theta }^{\prime }}(s_t,a_t)\mathrm{\nabla }\log p_{\theta }(a_t^n|s_t^n)]& & \text{假设}{p}_{\theta }(s_t)=p_{{\theta }^{\prime }}(s_t)\end{array}\end{array}$$

*为什么假设$p_{\theta }(s_t)=p_{{\theta }^{\prime }}(s_t)$, 因为难以计算?

再由公式(3)得:

$$\begin{array}{}\text{(16)}& \mathrm{\nabla }{\bar{R}}_{\theta }=E_{(s_t,a_t)\sim {\pi }_{{\theta }^{\prime }}}[\frac{\mathrm{\nabla }{p}_{\theta }(a_t|s_t)}{p_{{\theta }^{\prime }}(a_t|s_t)}{A}^{{\theta }^{\prime }}(s_t,a_t)]\end{array}$$

反推目标函数:

$$\begin{array}{}\text{(17)}& J^{{\theta }^{\prime }}(\theta )=E_{(s_t,a_t)\sim {\pi }_{{\theta }^{\prime }}}[\frac{p_{\theta }(a_t|s_t)}{p_{{\theta }^{\prime }}(a_t|s_t)}{A}^{{\theta }^{\prime }}(s_t,a_t)]\end{array}$$

Add constraint#

目前为止, 我们利用 Importance Sampling 完成了 Policy Gradient 从 On-policy 到 Off-policy 的优化.

但是 Importance Sampling 在实际应用中有一个不得不考虑的限制, 就是我们无法保证能采样足够多的数据, 这时当两个分布 $p_{\theta },p_{{\theta }^{\prime }}$差异过大时, 难以保证期望相等.

PPO做的事情, 简单说就是, 限制两个分布 $p_{\theta },p_{{\theta }^{\prime }}$ 不能差太多.

$$\begin{array}{}\text{(18)}& J_{PPO}^{{\theta }^{\prime }}(\theta )=J^{{\theta }^{\prime }}(\theta )-\beta KL(\theta ,{\theta }^{\prime })\end{array}$$

注: 此处 KL 散度指的不是将两个模型的参数看作分布,拉近两个模型的参数的距离. 而是两个模型行为上的距离, 就是当两个模型输入同样的 state 时, 希望输出的 action 的分布尽可能像

Conclusion#

PPO algorithm#

PPO2#

PPO2: 简化 PPO 的计算.

首先, 我们将横坐标 $x$ 设为 $\frac{p_{\theta }(a_t|s_t)}{p_{{\theta }^k}(a_t|s_t)}$, 则函数 $y=x$ 与 $y=clip(x,1-ϵ,1+ϵ)$ 的图像分别为图中的绿线和蓝线.
其中, $clip(x,a,b)=\{\begin{array}{rl}a,& x\le a\\ x,& a<x<b\\ b,& x\ge b\end{array}$

• 当 $A>0$ 时, $J_{PPO2}^{{\theta }^k}(\theta )$ 就是左图中红线, 我们要最大化目标函数, 也就希望 $x$ 越大越好, 但是当超过 $1+ϵ$ 后, 对目标函数就没有 benefit 了. 这样就阻止了进一步优化.
• 当 $A<0$ 时, 同理, 如右图.

目的依旧是保证两个分布 $p_{\theta },p_{{\theta }^k}$ 差距不能过大.

Experiment#