PPO算法
要理解GRPO以及后续的改进,必须先理解PPO算法。
PPO算法也是继承自策略梯度的算法逻辑。
策略梯度算法
总体目标是训练一个神经网络,对于所有的状态给出对应的action,总的return值最大。
目标定义如下:
E(R(τ))τ∼Pθ(τ)=∑τR(τ)Pθ(τ)E(R(\tau)){\tau \sim P{\theta}(\tau)} = \sum_{\tau} R(\tau) P_{\theta}(\tau)E(R(τ))τ∼Pθ(τ)=τ∑R(τ)Pθ(τ)
求梯度,因为我们不可能计算所有的轨迹以及其概率,因此需要变形:
∇E(R(τ))τ∼Pθ(τ)=∇∑τR(τ)Pθ(τ)=∑τR(τ)∇Pθ(τ)=∑τR(τ)∇Pθ(τ)Pθ(τ)Pθ(τ)=∑τPθ(τ)R(τ)∇Pθ(τ)Pθ(τ)=∑τPθ(τ)R(τ)∇logPθ(τ)≈1N∑n=1NR(τn)∇logPθ(τn)\begin{aligned} \nabla \mathbb{E}(R(\tau)){\tau \sim P\theta(\tau)} &= \nabla \sum_{\tau} R(\tau)P_\theta(\tau) \\ &= \sum_{\tau} R(\tau)\nabla P_\theta(\tau) \\ &= \sum_{\tau} R(\tau)\nabla P_\theta(\tau) \frac{P_\theta(\tau)}{P_\theta(\tau)} \\ &= \sum_{\tau} P_\theta(\tau) R(\tau) \frac{\nabla P_\theta(\tau)}{P_\theta(\tau)} \\ &= \sum_{\tau} P_\theta(\tau) R(\tau) \nabla \log P_\theta(\tau) \\ &\approx \frac{1}{N} \sum_{n=1}^N R(\tau^n) \nabla \log P_\theta(\tau^n) \end{aligned}∇E(R(τ))τ∼Pθ(τ)=∇τ∑R(τ)Pθ(τ)=τ∑R(τ)∇Pθ(τ)=τ∑R(τ)∇Pθ(τ)Pθ(τ)Pθ(τ)=τ∑Pθ(τ)R(τ)Pθ(τ)∇Pθ(τ)=τ∑Pθ(τ)R(τ)∇logPθ(τ)≈N1n=1∑NR(τn)∇logPθ(τn)
继续推导:
Gradient≈1N∑n=1NR(τn)∇logPθ(τn)(对轨迹求梯度的采样近似)=1N∑n=1NR(τn)∇log∏t=1TnPθ(ant∣snt)(带入自回归连乘概率)=1N∑n=1NR(τn)∇∑t=1TnlogPθ(ant∣snt)(利用对数性质:log∏=∑log)=1N∑n=1NR(τn)∑t=1Tn∇logPθ(ant∣snt)(∇ 具有线性性质,可以点对点发给每个 Token)=1N∑n=1N∑t=1TnR(τn)∇logPθ(ant∣snt)(调整求和顺序,将奖励乘给每个 Token 的梯度)=∇1N∑n=1N∑t=1TnR(τn)logPθ(ant∣snt)(关键步:因为 R 与 θ 无关,把 ∇ 整体提出来)\begin{aligned} \text{Gradient} &\approx \frac{1}{N} \sum_{n=1}^N R(\tau^n) \mathbf{\nabla} \log P_\theta(\tau^n) \quad &\text{(对轨迹求梯度的采样近似)} \\ &= \frac{1}{N} \sum_{n=1}^N R(\tau^n) \mathbf{\nabla} \log \left \\prod_{t=1}\^{T_n} P_\\theta(a_n\^t \\mid s_n\^t) \\right \quad &\text{(带入自回归连乘概率)} \\ &= \frac{1}{N} \sum_{n=1}^N R(\tau^n) \mathbf{\nabla} \left \\sum_{t=1}\^{T_n} \\log P_\\theta(a_n\^t \\mid s_n\^t) \\right \quad &\text{(利用对数性质:}\log\prod = \sum\log\text{)} \\ &= \frac{1}{N} \sum_{n=1}^N R(\tau^n) \sum_{t=1}^{T_n} \mathbf{\nabla} \log P_\theta(a_n^t \mid s_n^t) \quad &\text{(}\nabla\text{ 具有线性性质,可以点对点发给每个 Token)} \\ &= \frac{1}{N} \sum_{n=1}^N \sum_{t=1}^{T_n} R(\tau^n) \mathbf{\nabla} \log P_\theta(a_n^t \mid s_n^t) \quad &\text{(调整求和顺序,将奖励乘给每个 Token 的梯度)} \\ &= \mathbf{\nabla} \left \\frac{1}{N} \\sum_{n=1}\^N \\sum_{t=1}\^{T_n} R(\\tau\^n) \\log P_\\theta(a_n\^t \\mid s_n\^t) \\right \quad &\text{(关键步:因为 } R \text{ 与 }\theta\text{ 无关,把 }\nabla\text{ 整体提出来)} \end{aligned}Gradient≈N1n=1∑NR(τn)∇logPθ(τn)=N1n=1∑NR(τn)∇logt=1∏TnPθ(ant∣snt)=N1n=1∑NR(τn)∇t=1∑TnlogPθ(ant∣snt)=N1n=1∑NR(τn)t=1∑Tn∇logPθ(ant∣snt)=N1n=1∑Nt=1∑TnR(τn)∇logPθ(ant∣snt)=∇N1n=1∑Nt=1∑TnR(τn)logPθ(ant∣snt)(对轨迹求梯度的采样近似)(带入自回归连乘概率)(利用对数性质:log∏=∑log)(∇ 具有线性性质,可以点对点发给每个 Token)(调整求和顺序,将奖励乘给每个 Token 的梯度)(关键步:因为 R 与 θ 无关,把 ∇ 整体提出来)
去掉梯度符号,最小化目标函数如下:
Loss(θ)=−1N∑n=1N∑t=1TnR(τn)logPθ(ant∣snt)\text{Loss}(\theta) = -\frac{1}{N} \sum_{n=1}^N \sum_{t=1}^{T_n} R(\tau^n) \log P_\theta(a_n^t \mid s_n^t)Loss(θ)=−N1n=1∑Nt=1∑TnR(τn)logPθ(ant∣snt)
Return-to-go
第 ttt 步写下的 Token,不应该为第 ttt 步之前已经拿到的奖励负责。因此,我们把整条轨迹的绝对总分 R(τn)R(\tau^n)R(τn) 扔掉,替换为从第 ttt 步开始向后累加的未来回报 GntG_n^tGnt。
因此,loss修改为:
Loss(θ)=−1N∑n=1N∑t=1TnGntlogPθ(ant∣snt)\text{Loss}(\theta) = -\frac{1}{N} \sum_{n=1}^N \sum_{t=1}^{T_n} \mathbf{G_n^t} \log P_\theta(a_n^t \mid s_n^t)Loss(θ)=−N1n=1∑Nt=1∑TnGntlogPθ(ant∣snt)
其中 Gnt=∑t′=tTnγt′−trnt′\mathbf{G_n^t} = \sum_{t'=t}^{T_n} \gamma^{t'-t} r_n^{t'}Gnt=∑t′=tTnγt′−trnt′。这里 rnt′r_n^{t'}rnt′ 是第 nnn 条轨迹在 t′t't′ 时刻环境给出的即时小奖励,γ\gammaγ 是时间折扣因子。
Actor-Critic架构
上面的设计依然存在一个巨大的问题,Reward的设计会导致高方差。
为了让模型只和"大盘平均表现"比输赢,我们在 GntG_n^tGnt 后面强行用减法扣掉一个只和当前状态 snts_n^tsnt 有关、和动作 anta_n^tant 无关的期望基线。我们用一个参数为 ϕ\phiϕ 的全新神经网络来预测这个基线,记为 Vϕ(snt)V_\phi(s_n^t)Vϕ(snt)。
Loss(θ)=−1N∑n=1N∑t=1Tn(Gnt−Vϕ(snt))logPθ(ant∣snt)\text{Loss}(\theta) = -\frac{1}{N} \sum_{n=1}^N \sum_{t=1}^{T_n} \mathbf{\left( G_n^t - V_\phi(s_n^t) \right)} \log P_\theta(a_n^t \mid s_n^t)Loss(θ)=−N1n=1∑Nt=1∑Tn(Gnt−Vϕ(snt))logPθ(ant∣snt)
把reward和状态之差称为优势,其实就是在某个状态的表现比这个状态的平均表现好多少,这才是评价这个action好坏的关键,损失函数修改为:
LossActor(θ)=−1N∑n=1N∑t=1TnAntlogPθ(ant∣snt)\text{Loss}{\text{Actor}}(\theta) = -\frac{1}{N} \sum{n=1}^N \sum_{t=1}^{T_n} A_n^t \log P_\theta(a_n^t \mid s_n^t)LossActor(θ)=−N1n=1∑Nt=1∑TnAntlogPθ(ant∣snt)
同时我们也需要去训练critic网络:
LossCritic(ϕ)=1N∑n=1N∑t=1Tn(Gnt−Vϕ(snt))2\text{Loss}{\text{Critic}}(\phi) = \frac{1}{N} \sum{n=1}^N \sum_{t=1}^{T_n} \left( G_n^t - V_\phi(s_n^t) \right)^2LossCritic(ϕ)=N1n=1∑Nt=1∑Tn(Gnt−Vϕ(snt))2
PPO算法
上述设计的核心问题是,我们采用当前策略采样的样本更新模型以后,策略就发生了变化,这时之前的样本就无法使用了。辛辛苦苦采样的样本就被使用了一次策略更新。
为了能让采出来的数据被反复连续训练好几个 Epoch(多步更新),论文引入了重要性采样。 我们设定一个"旧策略网络 πθold\pi_{\theta_{\text{old}}}πθold",专门用来死记硬背采样时的那一套概率 ()。当模型参数 θ\thetaθ 开始连续微调更新时,θold\theta_{\text{old}}θold 保持静止。
为此,论文定义了一个新旧模型的概率比值(Probability Ratio),记为 rnt(θ)\mathbf{r_n^t(\theta)}rnt(θ):
rnt(θ)=Pθ(ant∣snt)Pθold(ant∣snt)\mathbf{r_n^t(\theta)} = \frac{P_\theta(a_n^t \mid s_n^t)}{P_{\theta_{\text{old}}}(a_n^t \mid s_n^t)}rnt(θ)=Pθold(ant∣snt)Pθ(ant∣snt)
通过重要性采样的代数转换,原本公式里的 logPθ\log P_\thetalogPθ 顺理成章地被替换为了比值 rnt(θ)r_n^t(\theta)rnt(θ)。
LossCPI(θ)=−1N∑n=1N∑t=1Tnrnt(θ)Ant\text{Loss}{\text{CPI}}(\theta) = -\frac{1}{N} \sum{n=1}^N \sum_{t=1}^{T_n} \mathbf{r_n^t(\theta)} A_n^tLossCPI(θ)=−N1n=1∑Nt=1∑Tnrnt(θ)Ant
但是这样去更新多次,会发现一个很大的问题,比如一个token拿到正的advantage,模型会不断提高这个token的概率,导致概率比值被从1拉到10,甚至100,这样可能导致模型更新发生巨大的变化,摧毁模型预训练阶段的知识。因此我们必须对更新幅度进行限制。这里,PPO算法其实提出了两种不同的方法,即CLIP裁剪和KL散度限制,损失函数分别如下:
(1)CLIP版本
LossPPO-Clip(θ)=−1N∑n=1N∑t=1Tnmin(rnt(θ)Ant, clip(rnt(θ),1−ε,1+ε)Ant)\text{Loss}{\text{PPO-Clip}}(\theta) = -\frac{1}{N} \sum{n=1}^N \sum_{t=1}^{T_n} \min \left( r_n^t(\theta) A_n^t, \; \text{clip}(r_n^t(\theta), 1-\varepsilon, 1+\varepsilon) A_n^t \right)LossPPO-Clip(θ)=−N1n=1∑Nt=1∑Tnmin(rnt(θ)Ant,clip(rnt(θ),1−ε,1+ε)Ant)
(2)KL散度版本
LossPPO-Penalty(θ)=−1N∑n=1N∑t=1Tn(rnt(θ)Ant−βDKLπθold(⋅∣snt)∣∣πθ(⋅∣snt))\text{Loss}{\text{PPO-Penalty}}(\theta) = -\frac{1}{N} \sum{n=1}^N \sum_{t=1}^{T_n} \left( r_n^t(\theta) A_n^t - \beta \mathbb{D}_{\text{KL}}\\pi_{\\theta_{\\text{old}}}(\\cdot \\mid s_n\^t) \\mid\\mid \\pi_\\theta(\\cdot \\mid s_n\^t) \right)LossPPO-Penalty(θ)=−N1n=1∑Nt=1∑Tn(rnt(θ)Ant−βDKLπθold(⋅∣snt)∣∣πθ(⋅∣snt))