PPO 算法

算法推导

给定一系列的动作 $s_1\rightarrow a_1\rightarrow s_2 \rightarrow a_2\rightarrow …\rightarrow s_T$，那么其轨迹可以记作 $\tau={s_1,a_1,s_2,a_2,…,s_T,a_T}$。利用 policy p 有：
$$
\begin{aligned}
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) p(s_3|s_2,a_2) \ldots \\
&= p(s_1) \prod_{t=1}^{T} p_\theta(a_t|s_t) p(s_{t+1}|s_t,a_t)
\end{aligned} \tag{1}
$$
奖励记作：
$$
R(\tau)=\sum_{t=1}^{T}{r_t} \tag{2}
$$
期望得到的奖励：
$$
\bar{R_\theta}=\sum_{\tau}^{}{R(\tau)p_\theta(\tau)}=E_{\tau\sim p_\theta(\tau)}[R(\tau)] \tag{3}
$$
对其参数 $\theta$ 求梯度：
$$
\nabla\bar{R_\theta}=\Delta \sum_{\tau}^{}{R(\tau)p_\theta(\tau)}=\sum_{\tau}^{}{R(\tau)\Delta p_\theta(\tau)} \tag{4}
$$

转换一下，下面第二步到第三步是因为：对数概率的微分等于概率本身乘以概率的微分。
$$
\begin{aligned}
\nabla\bar{R_\theta} & =\sum_{\tau}^{}{R(\tau)\Delta p_\theta(\tau)} \\ &=\sum_{\tau}^{}{R(\tau)p_\theta(\tau)\Delta p_\theta(\tau)/p_\theta(\tau)} \\ &=\sum_{\tau}^{}{R(\tau)p_\theta(\tau)\Delta logp_\theta(\tau)}
\end{aligned}
\tag{5}
$$
即用下式替换可以得到：
$$
\frac{d}{dx} \log(p(x)) = \frac{1}{p(x)} \frac{dp(x)}{dx} \tag{6}
$$

其中 $\tau$ 是从概率分布 $p_\theta(\tau)$ 中抽样得到的，即 $\tau \sim p_\theta(\tau)$。所以，期望值可以用样本平均来估计：
$$
\nabla\bar{R_\theta}=E_{\tau\sim p_\theta(\tau)}[R(\tau)\Delta \log p_\theta(\tau)] \approx \frac{1}{N}\sum_{n=1}^{N}{R(\tau^{n})\Delta \log(p_\theta(\tau^{n}))} \tag{7}
$$
将式 1 带入得：
$$
\nabla\bar{R_\theta} \approx \frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T_n}{R(\tau^{n})\Delta log(p_\theta(a_{t}^{n}|s_{t}^{n}))} \tag{8}
$$
那么 loss：
$$
\mathcal{L} = -\frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T_n} R(\tau^{(n)})\Delta \log(p_\theta(a_{t}^{(n)}|s_{t}^{(n)}))
\tag{9}
$$
这里，

$\mathcal{L}$ 表示损失函数，负号表示我们希望最小化该损失，
$N$ 是样本数，
$T_n$ 是第 $n$ 个样本的轨迹长度，
$R(\tau^{(n)})$ 是回报，
$\nabla \log(p_\theta(a_{t}^{(n)}|s_{t}^{(n)}))$ 是对数概率的变化。

为了确保奖励有正有负，不然会一直是正的，引入 advantage function.
$$
\nabla\bar{R_\theta} \approx \frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T_n}{(R(\tau^{n})-b)\Delta log(p_\theta(a_{t}^{n}|s_{t}^{n}))} \tag{10}
$$

可以看到，对于每个 action，都有相同的 $R(\tau_n) – b$，可以看做是这个动作的权重，但是每个序列的所有动作都有相同的权重看起来并不是特别合理，可以直觉地考虑，每个动作只会对对后续动作产生影响，则权重应该来自于后续的 reward。所以修改奖励：
$$
\nabla\bar{R_\theta} \approx \frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T_n}{(\sum_{t’=t}^{T_n}{r_{t’}^{n}}-b)\Delta log(p_\theta(a_{t}^{n}|s_{t}^{n}))} \tag{11}
$$
对于每一个 action，他对当前奖励影响较大，随着时间推移，这个动作的影响会越来越小，所以应该添加一个修正因子 $\gamma < 1$
$$
\nabla\bar{R_\theta} \approx \frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T_n}{(\sum_{t’=t}^{T_n}{r_{t’}^{n}\gamma^{t’-t}}-b)\Delta \log(p_\theta(a_{t}^{n}|s_{t}^{n}))} \tag{12}
$$
得到奖励期望的梯度。

由上面可知，
$$
\text{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}))} \tag{13}
$$
$\log (p_\theta(a_{t}^{n}|s_{t}^{n}))$ 对应模型输出 token 的 log_prob，生成的句子的整体 reward 搭配每个 token 的 log_prob 构成了整体 loss。