策略梯度小计

RL 好难啊，RL 好难啊。

这里记录（抄一些） cs336 和 spinning up 的推导和训练流程，主要遵循 cs336 的一些精简定义。以及做作业的一些坑。

定义

• Categorical policy：$a_t\sim {\pi }_{\theta }(\cdot |s_t)$ , 状态和环境 $o,s$ 在这里等价。
• ${\pi }_{\theta }(a_t|s_t)=\mathrm{softmax}(\cdots )$
• Trajectories: $\tau =(s_0,a_0,\cdots ,s_T,a_T),s_0\sim {\rho }_0(s_0)$, 在语言模型中，简单认为拼接 $s_{t+1}=s_t||a_t$
• reward: $r_T=R(s_T,a_T)$ ，在语言模型中，我们认为预测了正确的文本算 $R=1$ 奖励，否则为 $0$
• 一个轨迹的奖励为 $R(\tau )=\sum _{t=0}^T{r}_t$ , 我们暂时只考虑又穷轨迹，且不引入折扣因子。
• objective of agent: maximize $J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[R(\tau )]$
• problem: ${\theta }^{\ast }={\mathrm{argmax}}_{\theta }J(\theta )$

方程

简单的梯度上升为：${\theta }_{k+1}={\theta }_k+\alpha {\mathrm{\nabla }}_{\theta }J({\theta }_k)$

一个重要的结论为 ${\mathrm{\nabla }}_{\theta }J({\theta }_k)=\mathbb{E}[{\mathrm{\nabla }}_{\theta }\log P(\tau |\theta )R(\tau )]=\mathbb{E}[\sum _t{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t|s_t)R(\tau )]$

因此从一个批次样本 $D=\{{\tau }^i{\}}_{i=1}^N$ 可以获得无偏估计:

这里的 $b(s_t)$ 可以是一个训练的模型 （PPO），例如 $V^{\pi }(s)=\mathbb{E}[R(\tau )|s_t=s]$ ，即 $s_t=s$ 开始策略预测的期望回报，也可以是不用训练的值，例如后续的 GRPO。

一个重要结论为：

考虑将期望展开，将对数导数展开得到常数导数恒定为 $0$ ，因此添加 $b(s_t)$ 仍为无偏估计。

应该注意的是，我们在此重新定义了一个梯度，这不是损失，不应将其作为评估指标。

• 将来添加是否含基线的消融测试。

Off-Policy

我们可以将先前的一次采样结果缓存，重复利用，得到离轨策略梯度：

这仍然是一个无偏估计。

off-policy 在采样困难的场景有较大的用处，因为可以多次复用初始采样，但由于还要维护历史值，其计算量比在线验证更大。

注意：这里是对概率的比值，而实际编写时通常直接得到的为对数概率，需要做一个转化。

GRPO

GRPO 在 baseline 上不训练一个模型，而是采用直接计算的方法，对于一个问题 $q$ ，采样 $G$ 个输出 $o^{1,2\cdots ,G}$ ，并计算对应奖励 $r^{1,2,\cdots ,G}$ ，其 $\mathrm{advantage}$ 为：

一个输出 $o^i$ 内部的每个 token，都用同一个 $A^i$ 作为基线。

clip

clip 的用处在于限制 off-policy 乘积因子的改变幅度：

通常 $ϵ$ 取 $0.2$ 。

当该输出奖励高，$A^i>0$ ，clip 会限制梯度的增长，当 $A^i<0$ ，clip 也会限制模型的下降。