从头推导了 Policy Gradient 算法相关的一些数学细节，并结合了 verl 中对应的代码实现。我觉得还是比较详细的。

Policy Gradient

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

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

$$\mathop{\arg\max}\limits_{\theta}J(\theta) = \mathop{\arg\max}\limits_{\theta}\mathbb{E}_{\tau \sim \pi_{\theta}}R(\tau)=\mathop{\arg\max}\limits_{\theta}\sum_{\tau}P(\tau;\theta)R(\tau)$$

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

$$P(\tau|\theta) = \prod_{t=0}^T P(s_{t+1}|s_t, a_t) \, \cdot \, \pi_{\theta}(a_t | s_t)$$

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

$$\begin{aligned} \nabla_{\theta}J(\theta) &= \sum_{\tau}\nabla_{\theta}P(\tau;\theta)R(\theta) \\ &= \sum_{\tau} P(\tau;\theta) \frac{\nabla_{\theta}P(\tau;\theta)}{P(\tau;\theta)} R(\tau) \\ &= \sum_{\tau} P(\tau;\theta) \nabla_{\theta}\text{log}P(\tau;\theta)R(\tau) \\ &= \mathbb{E}_{\tau \sim \pi_{\theta}} \nabla_{\theta}\text{log}P(\tau;\theta)R(\tau) \end{aligned}$$

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

$$\begin{aligned} \nabla_{\theta}\text{log}P(\tau;\theta) &= \nabla_{\theta}\sum_{t=0}^T\text{log}P(s_{t+1}|s_t, a_t) + \nabla_{\theta}\sum_{t=0}^T\text{log}\pi_{\theta}(a_t|s_t) \\ &= \nabla_{\theta}\sum_{t=0}^T\text{log}\pi_{\theta}(a_t|s_t) \end{aligned}$$

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

$$\begin{aligned} \nabla_{\theta}J(\theta) &= \mathbb{E}_{\tau \sim \pi_{\theta}} [\nabla_{\theta}\sum_{t=0}^T\text{log}\pi_{\theta}(a_t|s_t)R(\tau)] \\ &= \mathbb{E}_{\tau \sim \pi_{\theta}} [\sum_{t=0}^T\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t)R(\tau)] \end{aligned}$$

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

$$\begin{aligned} \nabla_{\theta}J(\theta) &\approx \frac{1}{m}\sum_{i=1}^mR(\tau^{(i)})[\sum_{t^{(i)}=0}^{T^{(i)}} \nabla_{\theta}\text{log}\pi_{\theta}(a_{t^{(i)}}|s_{t^{(i)}})] \end{aligned}$$

Loss

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

$$\nabla_{\theta}J(\theta) = \mathbb{E}_{\tau \sim \pi_{\theta}} [\sum_{t=0}^T\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t)A(s_t, a_t)]$$ $$J(\theta) = \mathbb{E}_{\tau \sim \pi_{\theta}} [\sum_{t=0}^T\text{log}\pi_{\theta}(a_t|s_t)A(s_t, a_t)]$$

常用的梯度下降用的是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)$ 是更被广泛采用的算法，定义如下：

$$A(s,a)=Q(s,a) - V(s)$$

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

采用较多的策略有PPO中使用的广义优势估计（Generalized Advantage Estimation），或是GRPO中使用的组内均值作为基线。

Baseline Adjust

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

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

$$\mathbb{E}_{\tau}[\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t)\cdot V(s_t)] = \mathbb{E}_{s_t}[\mathbb{E}_{a_t \sim \pi_{\theta}(\cdot|s_t)}[\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t)\cdot V(s_t)| s_t]]$$

视 $s_t$ 为不变数，则

$$\mathbb{E}_{a_t \sim \pi_{\theta}(\cdot|s_t)}[\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t)\cdot V(s_t)| s_t] = V(s_t) \cdot \mathbb{E}_{a_t}[\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t)|s_t]$$

仅看第二项，有

$$\begin{aligned} \mathbb{E}_{a_t\sim \pi_{\theta}}[\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t)|s_t] &= \int \pi_{\theta}(a_t|s_t)\cdot \frac{\nabla_{\theta}\pi_{\theta}(a_t|s_t)}{\pi_{\theta}(a_t|s_t)} da_t \\ &= \nabla_{\theta} \int \pi_{\theta}(a_t|s_t)da_t \\ &= \nabla_{\theta}1 \\ &= 0 \end{aligned}$$

所以是无偏的。条件是，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)$，表示采取动作 $a_t$ 相比于其它操作能带来多少收益，带有一点中心化的思想。

Optimal Baseline

我们来求解一下理论的最优基线 $b(s_t)$，已知梯度为：

$$\begin{aligned} \nabla_{\theta}J(\theta) &= \mathbb{E}_{\tau} [\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t)(Q(s_t, a_t) - b(s_t))] \\ \end{aligned}$$

由方差的定义，梯度的方差为 $\text{Var}(\nabla_{\theta}J) = \mathbb{E}([\nabla_{\theta}J]^2) - [\mathbb{E}(\nabla_{\theta}J)]^2$，其中后一项我们已经证明过了无偏，所以只需考虑 $\mathbb{E}([\nabla_{\theta}J]^2)$ 项的变化会怎么影响方差。

考虑优化问题，目标是找到 $b(s,t)$ 来最小化：

$$L(b) = \mathbb{E}_\tau [(Q(s_t, a_t) - b(s_t)) \nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t)]^2$$

还是一样，我们先固定状态 $s_t$，此时 $b(s_t)$ 是一个常数：

$$L_{s_t}(b) = \mathbb{E}_{a_t \sim \pi_{\theta}(\cdot|s_t)}[((Q(s_t, a_t) - b(s_t)) \nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t))^2|s_t]$$

对 $b(s_t)$ 求导，并令导函数为0：

$$\begin{aligned} \frac{\partial L_{s_t}}{\partial b(s_t)} &= \mathbb{E}_{a_t \sim \pi_{\theta}(\cdot|s_t)}[-2 \cdot (Q(s_t, a_t) - b(s_t)) (\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t))^2|s_t] = 0 \\ \mathbb{E}_{a_t \sim \pi_{\theta}(\cdot|s_t)}[b(s_t) (\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t))^2|s_t] &= \mathbb{E}_{a_t \sim \pi_{\theta}(\cdot|s_t)}[Q(s_t, a_t) (\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t))^2|s_t] \\ b(s_t) \mathbb{E}_{a_t \sim \pi_{\theta}(\cdot|s_t)}[ (\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t))^2|s_t] &= \mathbb{E}_{a_t \sim \pi_{\theta}(\cdot|s_t)}[Q(s_t, a_t) (\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t))^2|s_t] \\ b(s_t) &= \frac{\mathbb{E}_{a_t \sim \pi_{\theta}(\cdot|s_t)}[Q(s_t, a_t) (\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t))^2|s_t]}{\mathbb{E}_{a_t \sim \pi_{\theta}(\cdot|s_t)}[ (\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t))^2|s_t]} \end{aligned}$$

整体来看，这个最优基线 $b(s,t)$ 的计算相当复杂，它既依赖于我们正在优化的策略 $\theta$，又依赖于回报 $Q_t$ 的真实分布，还要计算每一步的策略梯度最后加权。因此，我们需要一个好的近似，而历史选择了忽略加权项：

$$V(s_t) = \mathbb{E}_{a_t \sim \pi_{\theta}(\cdot|s_t)}[Q(s_t, a_t)|s_t]$$

$V(s_t)$ 是在状态 $s_t$ 下，遵循当前策略 $\pi_\theta$ 所能获得的期望回报 $\mathbb{E}_{a_t \sim \pi_{\theta}(\cdot|s_t)}[Q(s_t, a_t)]$。它不依赖于任何具体的动作 $a_t$，因此满足梯度无偏的条件。在实际中我们不会真的去采样 $a_t$，而是额外学习一个价值网络 critic model，来想办法获得近似的 $V(s_t)$

删去加权项的影响

仅仅只是好奇，随想一下相关的影响。

原本加权项会给梯度较大的动作 $a_t$ 对应的回报 $Q(s_t,a_t)$ 更高的权重，也就是说，基线 $b(s_t)$ 会偏大，让那些梯度更大的动作选择造成的影响小一点。我们的近似求解实际上唯一的影响就是让方差变大，感觉其实很难说会有什么影响。

原式其实可以抽象为一个问题，求解 $\frac{\mathbb{E}(AB)}{\mathbb{E}(A)}$

注意到协方差的定义是 $\text{Cov}(A,B) = \mathbb{E}(AB) - \mathbb{E}(A)\mathbb{E}(B)$，我们可以得到：

$$\mathbb{E}(AB) = \mathbb{E}(A)\mathbb{E}(B) + \text{Cov}(A,B)$$

因此原问题可以改写为：

$$\frac{\mathbb{E}(AB)}{\mathbb{E}(A)} = \frac{\mathbb{E}(A)\mathbb{E}(B) + \text{Cov}(A,B)}{\mathbb{E}(A)} = \mathbb{E}(B) + \frac{\text{Cov}(A,B)}{\mathbb{E}(A)}$$

$\text{Cov}(A,B)$ 衡量的是 $A$ 和 $B$ 是否统计独立与不相关。这里其实也牵涉到一个容易混淆的问题，回报 $Q(s_t, a_t)$ 与策略梯度 $\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t)$ 是否相关呢？

$Q(s_t, a_t)$ 衡量的是在状态 $s_t$ 执行动作 $a_t$ 后，遵循当前策略 $\pi_\theta$ 继续执行下去所能获得的期望回报，它衡量的是 $a_t$ 本身的好坏，与环境的奖励函数有关，也和状态转移的策略 $\theta$ 有关

$\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t)$ 是一个向量，告诉我们最快调整策略网络权重的方向，能让网络倾向于输出动作 $a_t$。模长大，则意味着微调参数 $\theta$ 会对 $a_t$ 的概率产生巨大影响；反之则影响不大。它完全由策略 $\theta$ 决定，这一步是不需要reward function参与的。

一言以蔽之，前者衡量动作的外部价值，后者衡量策略对动作的内部敏感度。二者有关系但关系不大。

因此，省略掉原公式中的策略梯度权重项，直接用 $\mathbb{E}_{a_t \sim \pi_{\theta}(\cdot|s_t)}[Q(s_t, a_t)|s_t]$ 来近似似乎变得更有道理了。

Generalized Advantage Estimation

现在，我们有了一个价值网络 $V$ 来估计每个step的 $V(s_t)$，但是我们的 $Q$ 却是稀疏的，可能只有一条轨迹的最后才能获得一个离散分数。怎么衡量每个 step 的期望回报 $Q(s_t, a_t)$ 呢？

Monte Carlo Return

一种方法叫 MC return， 把 $t$ 时刻到结尾的所有真实奖励 $r$ 加起来，得到的回报 $R_t$ 作为当前的 $Q$

$$\begin{aligned} R_t &= r_t + r_{t+1} + \cdots\\ A_t &= R_t - V(s_t) \end{aligned}$$

它的优点是，$R_t$ 本身就是无偏的奖励；但问题在于对于稀疏奖励，一段轨迹只有一个奖励，它没法更细粒度地区分 $R_t$ 的贡献究竟来源于哪个step，可能会造成错误的归因。

当然，也可以根据对奖励乘以一个距离衰减 $\lambda^{i-j}$，虽然还是不够实际，太远的话奖励可能已经稀释没了

$$R_t = r_t + \lambda r_{t+1} + \lambda^2 r_{t+2} + \cdots$$

TD error

另一种方法叫 TD error，它只看一步，用 当前选择的即时奖励 加上 下一步的价值估计

$$Q(s_t, a_t) = r_t + \gamma V(s_{t+1})$$

所以，优势函数的估计为：

$$A_t = r_t + \gamma V(s_{t+1}) - V(s_t)$$

它的优点在于，只依赖于下一步的随机性，方差低；但是它的 $Q$ 严重依赖于 $V$ 的准确性，是有偏的，而且在训练早期会偏得很厉害（毕竟 critic model 最后会接一个随机初始化的线性层来将 hidden state 映射到一个 scalar）

TD error 也可以多看几步，面临的问题同 MC return，奖励越无偏，但是方差越大。

$$\left\{ \begin{aligned} A_t^{(1)} &= r_t + \gamma V(s_{t+1}) - V(s_t) \\ A_t^{(2)} &= r_t + \gamma r_{t+1} +\gamma^2 V(s_{t+2}) - V(s_t) \\ A_t^{(3)} &= r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \gamma^3 V(s_{t+3}) - V(s_t) \\ & \vdots \end{aligned} \right.$$

当 TD error 取到无限步的时候，也就变成了带距离衰减的 MC return 的形式。

综合一下

GAE就是把两者综合一下，用一个额外的参数 $\lambda$ 来加权几步的 TD error。我们令 $\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$，则

$$\begin{aligned} A_t^{GAE(\gamma, \lambda)} &= A_t^{(1)} + \lambda A_t^{(2)} + \lambda^2 A_t^{(3)} + \cdots \\ &= \delta_t + \lambda(\delta_t + \gamma \delta_{t+1}) + \lambda^2(\delta_t + \gamma \delta_{t+1} + \gamma^2 \delta_{t+2}) + \cdots \\ &= \delta_t (1+ \lambda + \lambda^2 + \cdots) + \gamma \delta_{t+1}(\lambda + \lambda^2 + \lambda^3 + \cdots) + \cdots \\ &= \frac{1}{1-\lambda} \delta_t + \gamma \frac{\lambda}{1-\lambda}\delta_{t+1} + \gamma^2 \frac{\lambda^2}{1-\lambda}\delta_{t+2} + \cdots \\ &= \frac{1}{1-\lambda} \sum_{l=0}^\infin (\gamma \lambda)^l \delta_{t+l} \end{aligned}$$

我们可以直接把 $\frac{1}{1-\lambda}$ 丢掉，毕竟这种 scalar 的乘数可以融到 learning rate 里

而且这里的级数求和，实际上只有在无穷时间步的时候才能取，也引入了偏差。

$\gamma$ 是时间衰减系数，衡量的是对长期或短期奖励的权重。$\gamma$ 越接近 1，则模型更关注长期回报，否则越接近 0，则更关注即时奖励。

$\lambda$ 是对 TD 求和的权重。$\lambda=0$ 就是 TD error（低方差，高偏差），$\lambda=1$ 就是 MC return（高方差，低偏差），目的是为了在偏差和方差之间取得平衡。

Code

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

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散度。

$$\mathcal{D}_{KL}(\pi_{\theta}||\pi_{\theta+\Delta\theta}) = \int_{x} \pi_{\theta}(x)\text{log}\frac{\pi_{\theta}(x)}{\pi_{\theta+\Delta\theta}(x)}dx$$

特别注意，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$，同时又要保证输出分布不要漂移太多，写成数学形式是：

$$\Delta\theta^* = \mathop{\arg\max}\limits_{\mathcal{D}_{KL}(\pi_{\theta}||\pi_{\theta+\Delta\theta}) \leq \epsilon}J(\theta+\Delta\theta)$$

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

$$\Delta\theta^* = \mathop{\arg\max}\limits_{\Delta\theta}J(\theta+\Delta\theta)-\lambda (\mathcal{D}_{KL}(\pi_{\theta} || \pi_{\theta+\Delta\theta})-\epsilon)$$

$\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)$ 来近似采样。

$$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)}]$$

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

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

$$\text{Var}_{x\sim p}[f(x)]=E_{x\sim p}[f(x)^2]-(E_{x\sim p}[f(x)])^2$$ $$\begin{aligned} \text{Var}_{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 \\ &= E_{x\sim p}[f(x)^2\frac{p(x)}{q(x)}]-(E_{x\sim q}[f(x)])^2 \end{aligned}$$

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

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

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

$$\begin{aligned} J(\theta) &=\mathbb{E}_{\tau \sim \pi_{\theta}}[\text{log}P(\tau;\theta)R(\tau)] \\ &=\mathbb{E}_{\tau \sim \pi_{\theta'}}[\frac{\pi_{\theta}}{\pi_{\theta'}}\text{log}P(\tau;\theta)R(\tau)] \end{aligned}$$

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

$$\begin{aligned} J(\theta) &= \mathbb{E}_{\tau \sim \pi_{\theta}} [\sum_{t=0}^T\text{log}\pi_{\theta}(a_t|s_t)A(s_t, a_t)] \\ &= \mathbb{E}_{\tau \sim \pi_{\theta'}} [\sum_{t=0}^T\frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta'}(a_t|s_t)}\text{log}\pi_{\theta}(a_t|s_t)A(s_t, a_t)] \end{aligned}$$

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

$$\begin{aligned} \nabla_{\theta}J(\theta) &= \mathbb{E}_{\tau \sim \pi_{\theta'}} [\sum_{t=0}^T\frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta'}(a_t|s_t)}\nabla_{\theta}\text{log}\pi_{\theta}(a_t|s_t)A(s_t, a_t)] \\ &= \mathbb{E}_{\tau \sim \pi_{\theta'}} [\sum_{t=0}^T\frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta'}(a_t|s_t)}\frac{\nabla_{\theta}\pi_{\theta}(a_t|s_t)}{\pi_{\theta}(a_t|s_t)}A(s_t, a_t)] \\ &= \mathbb{E}_{\tau \sim \pi_{\theta'}} [\sum_{t=0}^T\frac{\nabla_{\theta}\pi_{\theta}(a_t|s_t)}{\pi_{\theta'}(a_t|s_t)}A(s_t, a_t)] \\ J(\theta) &= \mathbb{E}_{\tau \sim \pi_{\theta'}} [\sum_{t=0}^T\frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta'}(a_t|s_t)}A(s_t, a_t)] \end{aligned}$$

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

$$J(\theta) = \mathbb{E}_{(s_t, a_t) \sim \pi_{\theta'}} [\frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta'}(a_t|s_t)}A(s_t, a_t)]$$

PPO

TODO

PPO penalty

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

$$r_t = r_t - \beta \text{log} \frac{\pi_{\theta}}{\pi_{\theta'}}$$

参考策略可以是自己，也可以是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