大模型RLHF中PPO强化学习代码学习笔记（二）

为了更详细的描述PPO伪代码中详细的数学过程，又进行了如下整理，图片来自 PPO算法讲解 。以下是PPO算法伪代码中 每一步的具体公式 、符号定义及作用解析，结合强化学习和大模型训练的场景：

一、符号约定

• $s$：状态（如batch_prompt，即模型的输入prompt）
• $a$：动作（如batch_response，即模型生成的回复）
• ${\pi }_{\theta }$：策略网络（active_model），输出动作分布和价值预测
• ${\pi }_{\text{old}}$：参考策略（ref_model，通常是更新前的${\pi }_{\theta }$，用于稳定训练）
• $V_{\theta }$：价值网络（与${\pi }_{\theta }$共享参数，预测状态价值）
• $r(s,a)$：奖励模型打分（如batch_scores，衡量$s\to a$的质量）
• $\text{KL}(p\Vert q)$：KL散度，衡量两个分布$p$和$q$的差异

二、逐行公式解析

1. 生成回复：模拟策略与环境交互

batch_response = active_model.generate(batch_prompt)  

• 公式：$a\sim {\pi }_{\theta }(\cdot |s)$（$s=\text{batch\_prompt}$，模型根据状态$s$采样生成动作$a$）
• 作用：获取模型的“行为样本”（prompt→response对），模拟强化学习中“智能体与环境交互”的过程。

2. 拼接数据：构建状态-动作对

batch_data = concat(batch_prompt, batch_response)  

• 公式：$\text{batch\_data}=(s,a)$（拼接状态$s$和动作$a$，作为后续模型输入）
• 作用：将“输入-输出”封装为训练单元，供策略、价值、奖励模型处理。

3. 奖励模型打分：初步评价动作质量

batch_scores = reward_model(batch_data)  

• 公式：$r=r(s,a)$（奖励模型对状态-动作对$(s,a)$打分，如内容质量、用户偏好）
• 作用：量化动作$a$的“内在价值”，是强化学习奖励的核心来源。

4. 策略与价值预测（当前模型）

batch_all_probs, batch_probs, batch_all_values = active_model.forward_pass(batch_data)  

• 公式：
  • 策略分布：${\pi }_{\theta }(\cdot |s)=\text{batch\_all\_probs}$（状态$s$下，所有可能动作的概率分布）
  • 动作概率：${\pi }_{\theta }(a|s)=\text{batch\_probs}$（生成当前动作$a$的概率）
  • 价值预测：$V_{\theta }(s)=\text{batch\_all\_values}$（预测状态$s$的期望回报）
• 作用：获取策略的概率分布（用于后续KL计算和策略更新）和价值预测（用于优势计算）。

5. 策略与价值预测（参考模型）

ref_all_probs, ref_probs, ref_all_values = ref_model.forward_pass(batch_data)  

• 公式：
  • 参考策略分布：${\pi }_{\text{old}}(\cdot |s)=\text{ref\_all\_probs}$
  • 参考动作概率：${\pi }_{\text{old}}(a|s)=\text{ref\_probs}$
• 作用：提供“基准分布”，通过KL散度惩罚策略的突变（如生成风格突然变化）。

6. 计算KL散度：惩罚策略突变

kls = compute_KL(batch_all_probs, ref_all_probs)  

• 公式：
  $$\text{KL}({\pi }_{\theta }\parallel {\pi }_{\text{old}})=\sum _x{\pi }_{\theta }(x|s)\log \frac{{\pi }_{\theta }(x|s)}{{\pi }_{\text{old}}(x|s)}$$
  （对所有可能的动作$x$，计算分布差异）
• 作用：避免策略更新过于剧烈（如生成结果突然从“正式”变“口语化”），保证训练稳定。

7. 最终奖励：结合KL惩罚

rewards = compute_rewards(batch_scores, kls)  

• 公式：$r_{\text{final}}=r-\alpha \cdot \text{kls}$（$\alpha$为KL惩罚系数，平衡“高奖励”和“策略稳定”）
• 作用：修正原始奖励，让模型既追求高得分，又不“乱变风格”。

8. 优势函数：衡量策略改进空间

advantages = compute_advantages(batch_all_values, rewards)  

• 公式：$\text{Adv}=r_{\text{final}}-V_{\theta }(s)$（实际回报与价值预测的差，回答“当前策略比预期好多少？”）
• 作用：指导策略更新——若$\text{Adv}>0$，强化当前动作；若$\text{Adv}<0$，弱化当前动作。

9. 回报计算：价值网络的学习目标

returns = advantages + batch_all_values  

• 公式：$\text{Returns}=V_{\theta }(s)+\text{Adv}=r_{\text{final}}$（等价于实际奖励，是价值网络的“理想预测值”）
• 作用：训练价值网络更准确预测状态$s$的期望回报。

10. 多轮优化循环：稳定更新策略与价值

for i in range(epoch):  

• 逻辑：在同一批数据上迭代训练，让模型充分学习，同时通过**策略比例（ratio）**限制更新幅度（避免过拟合）。

11. 重新计算策略与价值（更新后模型）

active_all_probs, active_probs, active_all_values = active_model.forward_pass(batch_data)  

• 公式：
  • 新策略分布：${\pi }_{\theta }^{\prime }(\cdot |s)=\text{active\_all\_probs}$（参数更新后，策略的动作分布）
  • 新动作概率：${\pi }_{\theta }^{\prime }(a|s)=\text{active\_probs}$（更新后生成$a$的概率）
  • 新价值预测：$V_{\theta }^{\prime }(s)=\text{active\_all\_values}$（更新后价值网络的预测）
• 作用：获取更新后的策略和价值预测，用于计算损失。

12. 价值损失：训练价值网络更准确

loss_state_value = torch.mean((returns - active_all_values) ** 2)  

• 公式：
  $${\mathcal{L}}_{\text{value}}=\frac{1}{N}\sum _{i=1}^N{\left({\text{Returns}}_i-V_{\theta }^{\prime }(s_i)\right)}^2$$
  （均方误差，最小化预测值与回报的差异）
• 作用：让价值网络更精准地预估“做这个动作能得多少分”。

13. 策略比例：衡量新旧策略差异

ratio = active_probs / batch_probs  

• 公式：$\rho =\frac{{\pi }_{\theta }^{\prime }(a|s)}{{\pi }_{\theta }(a|s)}$（新策略生成$a$的概率 ÷ 旧策略生成$a$的概率）
• 作用：判断策略变化幅度——$\rho >1$表示新策略更倾向于$a$，$\rho <1$则相反。

14. PPO策略损失：稳定更新策略

loss_ppo = torch.mean(-advantages * ratio)  

• 公式（简化版，实际含clip）：
  $${\mathcal{L}}_{\text{ppo}}=-\frac{1}{N}\sum _{i=1}^N{\text{Adv}}_i\cdot {\rho }_i$$
  （取负号，将“最大化优势×比例”转为梯度下降的“最小化问题”；实际工程会用clip(\rho, 1-\epsilon, 1+\epsilon)限制$\rho$，避免更新过猛）
• 作用：
  • 若$\text{Adv}>0$（动作比预期好）：鼓励$\rho$大（新策略更倾向于$a$）。
  • 若$\text{Adv}<0$（动作比预期差）：鼓励$\rho$小（新策略更规避$a$）。

15. 总损失：平衡策略与价值优化

loss = loss_ppo + value_loss_rate * loss_state_value  

• 公式：${\mathcal{L}}_{\text{total}}={\mathcal{L}}_{\text{ppo}}+\lambda \cdot {\mathcal{L}}_{\text{value}}$（$\lambda$为value_loss_rate，平衡两者权重）
• 作用：同时优化策略网络（生成更优回复）和价值网络（预测更准）。

16-18. 梯度更新：优化模型参数

loss.backward()    # 计算梯度：∇θℒ_total  
optimizer_step()   # 更新参数：θ ← θ - η·∇θℒ_total（η为学习率）  
optimizer_zero_grad()  # 清空梯度：∇θ ← 0  

• 作用：通过反向传播和梯度下降，逐步优化模型参数，让策略更优、价值预测更准。

三、核心逻辑总结

PPO的训练像一场 “策略微调游戏”：

1. 采样试错：模型生成回复，拿奖励（加KL惩罚）。
2. 自我评估：预测自己能得多少分（价值网络），对比实际得分算“优势”。
3. 稳定改进：通过“新旧策略比例”限制更新幅度，避免“画风突变”，同时训练价值网络更准。
4. 反复练习：同一批数据练多轮，逐步优化，最终让模型“又会写、又会估分”。

（实际工程中，还会用clip进一步限制策略变化，让训练更稳定！）