我去怎么不能渲染公式了。。。看看怎么解决。。。

从简单的公式推导介绍大致思想，然后从代码角度介绍具体实践。

Policy Gradient

对于当前的状态 $s$，我们需要采取一个行动 $a$，假设我们行动的策略由参数 $\theta$ 决定，那我们采取行动 $a$ 的概率为 $\pi_{\theta}(a|s) = P(a|s,\theta)$

目标是求得能最大化奖励函数 $J(\theta)$ 的参数 $\theta$，即求

我们把 $J(\theta)$ 看成是轨迹 $\tau$ 的期望奖励。每条轨迹的概率可以展开表示为：

为了求得 $\theta$ , 我们对 $\theta$ 求导，推导如下：

这里有一步巧妙的变换，将 $\nabla_{\theta}P(\tau;\theta)$ 变成了 $P(\tau;\theta) \nabla_{\theta}\text{log}P(\tau;\theta)$，log形式允许我们将连乘展开进行进一步推导。

将数据带回，我们可以得到奖励函数的最终形式

最后，对于最外层的期望，我们用多次采样模拟。假设采了 $m$ 次

但是采样非常昂贵，所以实际上我们可能只用Monte Carlo采样一次。甚至采样出来的样本还要被多次使用。

Loss

对于之前推导的结果，我们可以简单的积分回去:

常用的梯度下降用的是Loss，而这里是reward，所以给reward取个负号就是loss了。

上式比较具体，讨论公式时也常用简单版：$J(\theta)=\mathbb{E}_{\tau \sim \pi_{\theta}}\text{log}P(\tau;\theta)R(\tau)$

对于LLM而言，$\pi_{\theta}(a_t|s_t)$ 就等于 $P(t_i|t_{<i},\theta)$

还有一点就是，不要迷信Loss的绝对值大小，因为理论上你可以给Loss加上一个任意常数 $c$ 而不影响它的梯度。只需要关注合理的部分的loss即可。

Advantage & Reward

比起整条轨迹 $\tau$ 获得一个稀疏的奖励 $R(\tau)$，采用优势函数 $A(s,a)$ 是更被广泛采用的算法，定义如下：

其中，$Q(s,a)$ 是状态动作值函数，表示在当前状态 $s$ 下采取动作 $a$ 的期望奖励；$V(s)$ 是状态值函数，表示当前状态 $s$ 后续的平均收益

现在采用最多的策略是广义优势估计（Generalized Advantage Estimation）。

P.S. GRPO又回到了使用稀疏奖励

Baseline Adjust

先来严谨理解一下优势函数的选择

直接使用回报 $Q$ 可能导致方差过大，可以证明减去一个基线 $V(s_t)$，即在 $s_t$ 状态下的平均收益，可以保证梯度无偏而方差减小。

视 $s_t$ 为不变数，则

仅看第二项，有

所以是无偏的。条件是，value function与 $a_t$ 无关，才可以在第二步做一个分离。

$V(s_t)$ 取作 $\mathbb{E}_{a_t}[Q(s_t, a_t)]$ 是一个比较自然的想法，因此回报又被称为优势函数 $A(s_t, a_t) = Q(s_t, a_t) - V(s_t)$。或者还有像GAE中采用TD error

Generalized Advantage Estimation

状态值函数 $V(s)$

对于每条 $\tau$，我们可以应用 critic model 给出一个 value，在LLM里就是对rollout出来的数据再forward一次得到的logits。

状态动作值函数 $Q(s,a)$

状态动作值的取值为 $r_t+\gamma \cdot V(s_{t+1})$

直观理解就是，动作 $a_t$ 对是否拿到reward $r_t$ 有贡献，同时我们还期望下一步的value要尽量大。

由于 $V$ 是我们用 critic model 打分得到的logits，这里的 $V(s_{t+1})$ 实际上已经隐含了上一步取了 $a_t$ 的条件，甚至可以说在LLM里 $s_{t+1}$ 就是 $a_t$，因此 $Q$ 确实可以看作是 $s$ 与 $a$ 的函数。

广义优势估计 GAE

令 $\delta_t=r_t+\gamma \cdot V(s_{t+1}) - V(s_t)$，这一项被称为TD误差

然后，$A_t^\text{GAE}=\delta_t+\gamma\lambda\cdot A_{t+1}^\text{GAE}$

最后计算得到的奖励记为 $R_t=A_t+V(s_t)$

  • $\lambda$ 的作用是控制优势估计的步数。当 $\lambda=0$ 时，退化为仅用当前一步的 advantage；当 $\lambda=1$ 时，退化为使用无穷步的随步长衰减的 advantage 之和。

  • $\gamma$ 就是随step衰减，被称为 reward-to-go

  • 我还有一个很朴素的疑问暂时没有想明白，TD error 为什么只取一步？为什么不能是 $\delta_t = r_t + \gamma \cdot V(s_{t+1}) + \gamma^2 \cdot V(s_{t+2})- V(s_t)$。虽然可能很简单，比如说 $V(s_{t+2})$ 涉及另一次环境交互了，但我还是希望能得到一个数学上的解释。等我理解贝尔曼算子的性质后会再来补充这部分。肯定有一个原因是要保证 $A_t^{\text{GAE}}$ 收敛，不能随着 $n$ 发散

  • 并且，在LLM中，我觉得 $\delta_t$ 并不适合代表 $s_t$ 状态下选择动作（token） $a_t$ 的价值，因为 critic model 的 token level 的 value 方差太大，$V(s_{t+1})$ 并不能准确反映真正的价值。我举个例子，”I love eat appl”，那下一个词说 “e” 的概率很高，$V(s_{t+1})$ 非常大，但是这个合适吗？它适合作为基线吗？上面的无偏性证明是在期望的情况下。哪怕PPO batch size开的非常大，也没法估计这个位置取不同 $a_t$ 的期望，因为整条轨迹都不一样，所以必然是有偏的（而且我觉得偏的很大）。

Code

在Verl中这部分逻辑的代码实现为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
with torch.no_grad():
    lastgaelam = 0
    advantages_reversed = []
    gen_len = token_level_rewards.shape[-1]

    for t in reversed(range(gen_len)):
        nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
        delta = token_level_rewards[:, t] + gamma * nextvalues - values[:, t]
        lastgaelam = delta + gamma * lam * lastgaelam
        advantages_reversed.append(lastgaelam)
    advantages = torch.stack(advantages_reversed[::-1], dim=1)

    returns = advantages + values
    advantages = verl_F.masked_whiten(advantages, eos_mask)
return advantages, returns

显然，我们可以倒序计算来处理TD error的后向依赖关系。

KL Loss Constraint

简单的使用 $\Delta\theta = \alpha\nabla_{\theta}J(\theta)$，模型可能会崩掉，需要增加限制。对于语言模型而言，这个限制就是输出的KL散度。

特别注意，KL散度是非对称的，需要分析一下。我们假设分布 $P$ 是已知的，$Q$ 是待优化的 \
1°) 正向KL散度：$\mathcal{D}_{KL}(P||Q) = \int_{x} P(x)\text{log}\frac{P(x)}{Q(x)}dx$ \
当 $P(x)$ 偏大时，KL散度会比较大，所以会更关注在 $P(x)$ 较大的地方，$Q(x)$也要尽量大，这样才能让KL散度偏小。因此，分布会尽可能覆盖所有的概率峰。\
2°) 反向KL散度：$\mathcal{D}_{KL}(Q||P) = \int_{x} Q(x)\text{log}\frac{Q(x)}{P(x)}dx$ \
这时恰好相反，当 $P(x)$ 偏小的时候，KL散度会非常大，所以会更关注在 $P(x)$ 较小的地方，$Q(x)$ 一定要很小，这会导致最后的分布集中在单个峰上，避免跨越峰之间的低概率区域。\
总结一下，正向KL散度能比较均衡地贴合原始分布。而反向KL散度比较极端，会坍缩到少数合理分布。\
参考：https://blog.csdn.net/mch2869253130/article/details/108998463

所以，我们要选择一个能最大化 $J(\theta)$ 的变分 $\Delta\theta$，同时又要保证输出分布不要漂移太多，写成数学形式是：

利用凸优化理论中的拉格朗日松弛（虽然我已经忘了，依稀记得还有什么KKT条件），可以将约束加入优化目标中作为惩罚项

$\lambda\epsilon$ 这一项可以忽略，它不影响梯度的方向只影响梯度的大小，等价于调节学习率的大小，因为 $\alpha \Delta \theta^{\ast} = \alpha’ (\Delta \theta^{\ast} - \lambda \epsilon )$

Importance Sampling

假设我们想估计 $f(x)$ 的期望，$x\sim p(x)$，则 $E_{x\sim p}[f(x)] \approx \frac{1}{N}\sum_{i=1}^Nf(x^i)$

但是如果 $p(x)$ 不可采样，我们可以用另一个可采样的分布 $q(x)$ 来近似采样。

其中 $\frac{p(x)}{q(x)}$ 被称为 importance weight

虽然期望 $E_{x\sim p}[f(x)]=E_{x\sim q}[f(x)\frac{p(x)}{q(x)}]$ 构成了一个无偏估计，但是考虑对应的方差：

可以发现，当 $q(x)$ 和 $p(x)$ 相差很大的时候，重要性权重 $\frac{p(x)}{q(x)}$ 就会偏离1，进而导致方差差异变大。因此，重要性采样在采样不充分的时候会逐渐失真。

之前也说过，采样很昂贵。一条数据采出来可能被训练多次。那在训过一次之后，这条数据就不能看作是由当前的分布 $\pi_{\theta}$ 采出来的了，而是过去的分布 $\pi_{\theta’}$ 采样的结果。

有了重要性采样，我们可以先off-policy地用 $\pi_{\theta’}$ 采样很多样本，然后再进行训练：

你可能会想问 $\frac{\pi_{\theta}}{\pi_{\theta’}}$ 怎么算，让我们展开来理解一下具体的意思。

这样就清楚多了，这两个值显然都是可求的。再进一步，我们回到梯度形式

太Amazing了。我们可以进一步，把采样分解到每个时间步，得到

PPO

TODO

PPO penalty

给reward添加一个 per token 的KL control，我没搞明白它数学上的正当性是哪里来的。不过，可以简单当成对reward的一个修正，只做了一点点改变就拿到reward肯定比做了很多改变才拿到reward好；做了很多改变还拿不到reward的策略更是要狠狠惩罚。

参考策略可以是自己，也可以是sft启动后的ref model，不过现在都不这么做了。而且从后者角度理解的话就会觉得有点奇怪，让模型去学ref model，可能是找一条更平滑的接近？

PPO Clip

Reference

【Proximal Policy Optimization (PPO) 算法理解：从策略梯度开始】https://zhuanlan.zhihu.com/p/614115887

【从Importance Sampling到PPO】https://zhuanlan.zhihu.com/p/388707220

【Why does the “reward to go” trick in policy gradient methods work?】https://ai.stackexchange.com/questions/9614/why-does-the-reward-to-go-trick-in-policy-gradient-methods-work