RLHF 到底在训什么？从 SFT 到 PPO的完整技术图谱

之前讲解了RLHF和PPO算法，后面的DPO和GRPO都是基于此方法的改进。但是我认为光是看懂原理是不够的，今天我打算来一篇完整的全流程工程实现，包括伪代码理清楚训练逻辑。

RLHF的核心在于如何将人类的偏好转化为数学上的奖励信号，并通过PPO这种On-policy算法来稳定地更新策略。

第一部分：RLHF 流程概览与 SFT (Supervised Fine-Tuning)

在进入复杂的PPO之前，我们需要构建地基。RLHF 的完整流水线通常被很多论文（如 InstructGPT）定义为三个步骤：

1. SFT (有监督微调)：让模型学会“像人一样说话”，掌握指令遵循的格式。

2. RM (奖励模型训练)：让模型学会“像人一样打分”，理解哪个回答更好。

3. PPO (强化学习微调)：利用 RM 的信号，优化生成策略。

如果直接跳过 SFT 进行 RL 训练，模型很难在庞大的动作空间中收敛到有意义的语言模式。

1. SFT 阶段伪代码

这一步本质上就是标准的语言模型训练（Next Token Prediction），但数据质量极高。

# ==========================================
# 阶段 1: Supervised Fine-Tuning (SFT)
# 目标: 获得一个能遵循指令的基础模型 (pi_sft)
# ==========================================

def train_sft(base_model, dataset, epochs, learning_rate):
    """
    base_model: 预训练的大语言模型 (Pretrained LLM)
    dataset: 高质量的 (Prompt, Response) 对
    """
    # 初始化优化器
    optimizer = AdamW(base_model.parameters(), lr=learning_rate)
    
    # 也就是最终我们将要在RL阶段使用的 Actor 的初始状态
    sft_policy = base_model 

    for epoch in range(epochs):
        for batch in dataset:
            # batch 包含 input_ids (prompt + completion) 和 labels
            # labels 中 prompt 部分通常被 mask 掉 (设为 -100)，只计算 response 的 loss
            input_ids, attention_mask, labels = batch
            
            # 前向传播
            outputs = sft_policy(input_ids, attention_mask=attention_mask)
            logits = outputs.logits
            
            # 计算交叉熵损失 (Cross Entropy Loss)
            # Loss = - sum(log(P(token_t | token_<t)))
            loss = cross_entropy_loss(logits, labels)
            
            # 反向传播与更新
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    
    print("SFT 模型训练完成。它现在能生成流畅的人类语言，但不知道哪种回答更符合人类偏好。")
    return sft_policy

关键点解析 (Expert Note):

• Masking: 在 SFT 中，非常关键的一点是 labels 的构建。我们只计算模型“回答”部分的 Loss，而不计算“提问”部分的 Loss。这是为了防止模型过度拟合 Prompt 的分布，而专注于生成 Response。

输入: "你好，请介绍一下Python。"  "Python是一种编程语言..."
标签:  [-100, -100, ... -100,   "Python", "是", "一种"...]
         ↑                        ↑
      prompt部分被mask        只计算response部分的loss

SFT的作用：

• 初始对齐：让模型学会"听从指令"

• 基础能力：掌握对话的基本格式和流畅性

• RL的起点：为后续强化学习阶段提供好的初始策略

假设数据集中的一个样本：

prompt = "解释一下量子计算"
response = "量子计算是一种利用量子力学原理进行计算的新型计算范式..."

# 模型输入
input_ids = encode(prompt + response)  # 拼接
labels = [-100]*len(prompt_tokens) + encode(response)  # prompt部分mask

SFT的目标是最简单的下一个token预测：

• 给定前文，预测下一个token

• 但通过mask机制，只优化response部分的预测

• 让模型学会根据prompt生成合适的response

实例逻辑：

第1步：准备标签

# token化后：
prompt_tokens  = [101, 123, 456, 789]  # "你好，介绍一下Python。"
response_tokens = [1025, 345, 678, 901]  # "Python是一种编程语言"

# 完整输入（拼接）：
input_ids = [101, 123, 456, 789, 1025, 345, 678, 901]

# 标签设置（关键！）：
labels    = [-100, -100, -100, -100, 1025, 345, 678, 901]
#         ↑ 忽略prompt部分的损失    ↑ 只计算response部分的损失

第2步：模型前向传播

# 模型输出（简化，实际是词汇表大小的概率分布）：
# 位置:    0      1      2      3      4        5      6      7
logits = [
  [0.1, 0.2, ...],  # 预测位置0的下一个token
  [0.3, 0.1, ...],  # 预测位置1的下一个token
  [0.2, 0.3, ...],  # 预测位置2的下一个token
  [0.4, 0.1, ...],  # 预测位置3的下一个token
  [0.7, 0.1, ...],  # 预测位置4的下一个token ← 这里应该预测"Python"(1025)
  [0.2, 0.6, ...],  # 预测位置5的下一个token ← 这里应该预测"是"(345)
  [0.1, 0.5, ...],  # 预测位置6的下一个token ← 这里应该预测"一种"(678)
  [0.3, 0.4, ...]   # 预测位置7的下一个token ← 这里应该预测"编程语言"(901)
]

第3步：计算损失（只看response部分）

假设正确token在模型预测中的概率：

位置4: 正确token=1025，模型给出的概率=0.1
位置5: 正确token=345，模型给出的概率=0.8
位置6: 正确token=678，模型给出的概率=0.6
位置7: 正确token=901，模型给出的概率=0.9

计算损失：

# 交叉熵损失 = -log(P(正确答案))

位置4的损失 = -log(0.1) ≈ 2.3026
位置5的损失 = -log(0.8) ≈ 0.2231
位置6的损失 = -log(0.6) ≈ 0.5108
位置7的损失 = -log(0.9) ≈ 0.1054

# 总损失 = 平均损失
总损失 = (2.3026 + 0.2231 + 0.5108 + 0.1054) / 4 ≈ 0.7855

第二部分：训练奖励模型（Reward Model, RM）

在 SFT 阶段，我们教会了模型“说话”。现在，我们需要构建一个“裁判”，告诉模型什么是“好话”。

对于 LLM 这种生成任务，直接定义一个数学上的 Reward 函数（类似于强化学习玩游戏时的得分）是非常困难的。因此，我们利用人类的偏好数据训练一个神经网络来充当 Reward Function。

RM 的核心任务是：给定一个 Prompt 和一个 Response，输出一个标量分数（Scalar Score）。

数据形式：通常是成对比较数据。标注员不会直接给一个句子打 7.5 分，而是看两个回答 A 和 B，判断“A 比 B 好”。 数据结构通常为三元组：(Prompt, Chosen_Response, Rejected_Response)。

2. RM 阶段伪代码

这里的核心在于损失函数的设计。我们使用 Bradley-Terry 模型假设下的 Pairwise Ranking Loss。

# ==========================================
# 阶段 2: Reward Model Training
# 目标: 训练一个能模拟人类偏好的打分模型
# ==========================================

import torch.nn.functional as F

class RewardModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        # 通常基于 SFT 模型初始化，保留其语言理解能力
        self.backbone = base_model
        # 将原本的 LM Head (vocab_size) 替换为 Scalar Head (1)
        # 也就是把输出维度从 [Batch, Seq_Len, Vocab] 映射到 [Batch, 1]
        self.scalar_head = nn.Linear(hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        # 获取 transformer 的最后一层隐藏状态
        hidden_states = self.backbone(input_ids, attention_mask).last_hidden_state
        
        # 通常取最后一个 token (EOS token) 的 embedding 或者是所有 token 的 average pooling
        # 来代表整个句子的语义表示
        sentence_representation = get_pooling_or_eos_embedding(hidden_states)
        
        # 输出标量奖励值
        reward = self.scalar_head(sentence_representation)
        return reward

def train_reward_model(sft_model, comparison_dataset, epochs, lr):
    """
    comparison_dataset: 包含 (prompt, chosen, rejected) 的数据
    """
    rm = RewardModel(sft_model)
    optimizer = Adam(rm.parameters(), lr=lr)

    for epoch in range(epochs):
        for batch in comparison_dataset:
            # batch 包含两组输入：
            # 1. prompt + chosen response (input_ids_chosen)
            # 2. prompt + rejected response (input_ids_rejected)
            
            # 分别通过 RM 计算奖励分值
            # 注意：同一个 RM 同时处理“胜者”和“败者”
            r_chosen = rm(batch.input_ids_chosen, batch.mask_chosen)
            r_rejected = rm(batch.input_ids_rejected, batch.mask_rejected)
            
            # === 核心：Pairwise Ranking Loss ===
            # 我们希望 r_chosen > r_rejected
            # Loss = -log(sigmoid(r_chosen - r_rejected))
            # 当差值越大，sigmoid 越接近 1，log(1)=0，loss 越小
            loss = -torch.log(torch.sigmoid(r_chosen - r_rejected)).mean()
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
    print("Reward Model 训练完成。它现在是人类偏好的代理 (Proxy)。")
    print("注意：在接下来的 PPO 阶段，RM 的参数将被冻结 (Frozen)，仅作为环境反馈信号。")
    return rm

关键点解析 (Expert Note):

Scalar Head: 原本的 LLM 是为了预测下一个 Token 的概率分布，输出维度巨大（Vocab Size，如 32k 或 128k）。RM 不需要预测下一个词，只需要评估整体质量，所以必须把 Head 换成输出维度为 1 的线性层。

数据一致性: 这里的 prompt 必须和 SFT 阶段以及后续 PPO 阶段的分布保持一致，否则 RM 会出现 OOD (Out of Distribution) 问题，导致给出的奖励不准确。

代码细节：

1. 奖励模型的作用

核心目标：

训练一个能区分回答质量的模型，为RL阶段提供量化反馈。

为什么需要RM？

人类反馈: "这个回答好，那个回答不好"
↓
奖励模型: r_chosen=8.5, r_rejected=2.3
↓
RL优化: 让模型输出能获得更高奖励的回答

2. 奖励模型架构详解

class RewardModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        # 基于SFT模型初始化
        self.backbone = base_model
        # 替换输出头：从词汇表预测 → 标量分数
        self.scalar_head = nn.Linear(hidden_size, 1)

架构变化：

SFT模型输出: [batch, seq_len, vocab_size]  # 每个位置预测词
奖励模型输出: [batch, 1]                   # 整个回答的打分

3. 数据处理流程

数据集格式：

每个样本包含三部分：

{
    "prompt": "解释一下量子计算",
    "chosen_response": "量子计算是一种...（清晰准确的解释）",  # 人类优选
    "rejected_response": "量子计算就是...（含糊不清或错误的解释）"  # 人类不选
}

输入拼接：

chosen_input = prompt + chosen_response
rejected_input = prompt + rejected_response

# 例子：
chosen:   "解释一下量子计算[SEP]量子计算是一种利用量子力学原理..."
rejected: "解释一下量子计算[SEP]量子计算就是使用量子计算机..."

实例逻辑：

数据收集

需要收集：提示（Prompt） + 两个回答（一个优选，一个次选）

例子：
提示："解释什么是黑洞"
优选回答："黑洞是宇宙中的一个区域，引力极强，连光也无法逃脱..."
次选回答："黑洞就是太空中的一个洞，能把东西都吸进去..."

# 每条数据包含：
{
    "prompt": "问题文本",
    "chosen": "人类偏好的回答",
    "rejected": "人类不偏好的回答"
}

1. 基础架构

原始SFT模型：
输入：[Batch, Seq_Len] → 模型 → 输出：[Batch, Seq_Len, Vocab_Size]

奖励模型改造：
输入：[Batch, Seq_Len] → 模型 → 输出：[Batch, 1]（单个分数）

2. 模型初始化

• 主干网络：使用训练好的SFT模型

• 输出头替换：将语言模型头替换为回归头

SFT模型：最后一个线性层输出维度=词汇表大小
奖励模型：最后一个线性层输出维度=1

3. 特征提取策略

如何从序列中提取特征得到单个分数？

策略A：EOS Token表示法

输入：你好[SEP]黑洞是一种...<EOS>
                  ↑
              取EOS位置的特征向量 → 线性层 → 分数

策略B：平均池化

你好[SEP]黑[SEP]洞[SEP]是[SEP]...
  ↓     ↓   ↓   ↓   ↓
[向量] [向量] [向量] [向量] ...  → 求所有token向量的平均值 → 线性层 → 分数

训练流程详解

步骤1：数据批处理

每个batch包含：
batch = {
    "chosen_input_ids": [batch_size, seq_len],
    "chosen_attention_mask": [batch_size, seq_len],
    "rejected_input_ids": [batch_size, seq_len], 
    "rejected_attention_mask": [batch_size, seq_len]
}

步骤2：前向传播

对于每个样本：
1. chosen数据 → 奖励模型 → 得到分数 r_chosen
2. rejected数据 → 奖励模型 → 得到分数 r_rejected

步骤3：计算损失

核心思想：让 r_chosen > r_rejected

计算差值：diff = r_chosen - r_rejected
计算sigmoid：p = 1 / (1 + exp(-diff))  # 将差值映射到(0,1)
计算损失：loss = -log(p)

理解：
- 如果 r_chosen >> r_rejected → diff很大 → p接近1 → -log(p)接近0
- 如果 r_chosen ≈ r_rejected → diff接近0 → p=0.5 → -log(0.5)=0.69
- 如果 r_chosen << r_rejected → diff为负 → p接近0 → -log(p)很大

步骤4：反向传播更新

loss.backward() → 计算梯度
optimizer.step() → 更新模型参数

高级技巧：

不只有一个总分，而是多个维度：
- 准确性：0-10分
- 有用性：0-10分
- 安全性：0-10分
- 风格匹配：0-10分

不仅输出分数，还输出置信度：
分数：8.5 ± 0.3
用于：低置信度时请求人工评估

部署后继续收集人类反馈：
1. 用户对回答的评分
2. 用户选择（A/B测试）
3. 人工审核结果
用于持续改进奖励模型

第三部分：PPO 的初始化与采样 (Rollout)。

如果说前两个阶段是准备食材和厨具，那么从这里开始，真正的“烹饪”（强化学习迭代）才刚刚开始。这也是工程实现上最痛苦、显存爆炸的阶段。

在 PPO 这一步，我们通常需要在显存中同时维护四个模型（或者说逻辑上的四个实体）：

1. Actor (策略模型): 我们要训练的主角（初始权重来自 SFT）。

2. Reference Model (参考模型): SFT 的一个冻结副本，用于计算 KL 散度，防止 Actor 跑偏。

3. Reward Model (奖励模型): 阶段 2 训练好的，冻结，用于给分。

4. Critic (价值模型): 在 PPO 中必须的一个组件，用于估计 Value Function $V(s)$，通常用 SFT 权重初始化。

这一部分的核心逻辑是：让 Actor 去“玩游戏”（生成文本），然后收集环境反馈（Reward + KL Penalty）。

3. PPO Rollout 阶段伪代码

# ==========================================
# 阶段 3: PPO Rollout (采样与奖励计算)
# 目标: 生成数据，构建经验池 (Experience Buffer)
# ==========================================

class PPOTrainer:
    def __init__(self, sft_model, reward_model, kl_coef=0.1):
        # 1. Actor: 可训练，也就是当前的 Policy
        self.actor = copy.deepcopy(sft_model) 
        
        # 2. Critic: 可训练，用于估计 Value(State)
        # 结构与 Actor 类似，但输出层改为 scalar (预测当前序列的长期期望回报)
        self.critic = init_value_head(copy.deepcopy(sft_model))
        
        # 3. Reference Model: 冻结，用于约束 Actor
        self.ref_model = copy.deepcopy(sft_model)
        self.ref_model.eval() # Freeze
        
        # 4. Reward Model: 冻结，提供原始分数
        self.reward_model = reward_model
        self.reward_model.eval() # Freeze
        
        self.kl_coef = kl_coef # KL 惩罚系数 beta

    def generate_experience(self, prompts):
        """
        采样步骤 (Rollout Phase)
        prompts: 一个 batch 的提示词
        """
        with torch.no_grad(): # 采样阶段不需要反向传播
            # 1. Action: Actor 根据 Prompt 生成 Response
            # 这里的 sequence 包含 prompt + generated_tokens
            sequence, action_mask = self.actor.generate(prompts)
            
            # 2. Log Probabilities (Old Policy): 
            # 计算当前 Actor 生成这些 token 的对数概率
            # shape: [batch, seq_len]
            old_log_probs = self.actor.get_log_probs(sequence)
            
            # 3. Reference Log Probabilities:
            # 计算原始 SFT 模型生成同样 token 的对数概率
            ref_log_probs = self.ref_model.get_log_probs(sequence)
            
            # 4. Raw Reward: 
            # 把它丢给 Reward Model 打分 (只看完整句子)
            # shape: [batch] (每个句子一个分)
            raw_rewards = self.reward_model(sequence)
            
            # 5. 计算 KL 散度惩罚 (KL Divergence Penalty)
            # 我们希望 Actor 不要偏离 SFT 太多
            # KL(P || Q) ≈ log(P(x)) - log(Q(x))
            # 这是一个 token-level 的计算
            kl_divergence = old_log_probs - ref_log_probs
            
            # 6. 综合奖励 (Total Reward)
            # R_total = R_model - beta * KL
            # 注意：Raw Reward 通常只加在最后一个 token 上
            # 而 KL Penalty 是施加在生成的每一个 token 上的
            rewards = torch.zeros_like(sequence)
            
            # 在生成过程的每一步扣除 KL 惩罚
            rewards = -self.kl_coef * kl_divergence 
            
            # 在序列的最后一个 token 加上 Reward Model 的打分
            for i in range(batch_size):
                end_idx = find_end_index(action_mask[i])
                rewards[i, end_idx] += raw_rewards[i]
            
            # 7. Value Estimation:
            # Critic 预测每个 token 状态下的价值 V(s)
            values = self.critic(sequence)
            
        return {
            "prompts": prompts,
            "sequence": sequence,
            "old_log_probs": old_log_probs,
            "rewards": rewards, # 包含了 KL 的最终奖励
            "values": values,   # Critic 的预测值
            "mask": action_mask
        }

关键点解析 (Expert Note):

• Reward Hacking (奖励黑客): 为什么要引入 ref_model 和 KL 散度？因为 RM 只是人类偏好的一个有偏代理。如果你完全不加约束地训练 Actor 最大化 RM 的分数，Actor 很快就会学会输出一些人类看不懂但 RM 认为分很高的乱码（Adversarial Attacks），或者输出极端的、讨好的废话。KL 惩罚项强迫 Actor 保持在“说人话”的范围内。

• Token-level Reward: 这是一个很反直觉的地方。虽然 RM 只给整个句子打一个分，但在强化学习视角下，每个生成的 Token 都是一个 Action。

  • 中间 Token 的 Reward = - beta * KL (只有惩罚，没有 RM 分数)。

  • 最后一个 Token 的 Reward = RM_Score - beta * KL。

• 显存地狱: 这一步你可以看到，显存里塞了 4 个 LLM（虽然其中两个是 inference only）。实际工程中，通常会使用 LoRA (仅微调 Adapter)、Offloading (把 Ref Model 放到 CPU) 或 Shared Backbone (Actor 和 Critic 共享底层参数) 来优化。

PPO Rollout（采样与奖励计算）阶段详解

这一段比较重要，我会讲的细致一些。

一、PPO训练框架概述

1. 强化学习在RLHF中的应用

传统RL: 智能体 ← 环境 → 奖励
RLHF:   语言模型 ← 奖励模型 → 人类偏好

2. PPO在RLHF中的四个角色

# 四个关键模型：
1. Actor（演员）：当前要训练的语言模型策略
2. Critic（评论家）：价值函数，评估状态好坏
3. Reference Model（参考模型）：原始的SFT模型，作为锚点
4. Reward Model（奖励模型）：人类偏好的代理

二、Rollout阶段详解

阶段目标：生成经验数据

输入：一批提示（prompts）
输出：完整的交互轨迹（状态、动作、奖励、价值估计）

三、代码逐行详解

1. 初始化PPOTrainer

def __init__(self, sft_model, reward_model, kl_coef=0.1):
    # 1. Actor: 可训练的策略网络
    self.actor = copy.deepcopy(sft_model) 
    # 为什么要复制？因为我们要在sft模型基础上优化
    
    # 2. Critic: 价值函数，评估状态的好坏
    self.critic = init_value_head(copy.deepcopy(sft_model))
    # 结构类似Actor，但输出单个标量（价值估计）
    
    # 3. Reference Model: 冻结的参考模型
    self.ref_model = copy.deepcopy(sft_model)
    self.ref_model.eval()  # 冻结，不训练
    # 作用：防止Actor偏离原始SFT模型太远
    
    # 4. Reward Model: 冻结的奖励模型
    self.reward_model = reward_model
    self.reward_model.eval()  # 冻结，不训练
    
    self.kl_coef = kl_coef  # KL惩罚系数

2. 生成经验（generate_experience）

sequence, action_mask = self.actor.generate(prompts)

输入: ["解释一下机器学习", "如何泡茶"]
输出: 
sequence: [
  "解释一下机器学习[SEP]机器学习是人工智能的一个分支...",
  "如何泡茶[SEP]首先准备茶叶和热水，然后..."
]
action_mask: 标记哪些token是生成的（非prompt部分）

prompt: "解释一下机器学习"
生成过程（自回归）：
1. 输入"解释一下机器学习" → 输出"机器"
2. 输入"解释一下机器学习机器" → 输出"学习"
3. 输入"解释一下机器学习学习" → 输出"是"
...
直到生成<EOS>或达到最大长度

步骤2：计算旧策略的对数概率

old_log_probs = self.actor.get_log_probs(sequence)

计算：每个生成token在Actor当前策略下的对数概率
例子：
生成的token: ["机器", "学习", "是", ...]
概率:       [0.8, 0.7, 0.9, ...]
对数概率:   [log(0.8), log(0.7), log(0.9), ...] = [-0.223, -0.357, -0.105, ...]

对数概率的意义：

• 高对数概率：模型对这个token很确定

• 低对数概率：模型对这个token不确定

• 用于后续计算重要性采样比率

  将完整的prompt+response输入奖励模型
  输出：每个序列的总体质量分数
  例如：raw_rewards = [8.5, 6.2]  # 两个序列的分数

步骤3：计算参考模型的对数概率

ref_log_probs = self.ref_model.get_log_probs(sequence)

用冻结的参考模型计算相同序列的概率
保持模型"不忘初心"的锚点

步骤4：计算原始奖励

raw_rewards = self.reward_model(sequence)

将完整的prompt+response输入奖励模型
输出：每个序列的总体质量分数
例如：raw_rewards = [8.5, 6.2]  # 两个序列的分数

输入完整序列 → 特征提取 → 线性层 → 标量分数
           ↓
      [CLS]解释...机器学习是...[EOS]
                        ↑
                   取[EOS]特征

步骤5：计算KL散度惩罚

kl_divergence = old_log_probs - ref_log_probs

KL散度计算公式：
KL(P || Q) = Σ P(x) * log(P(x)/Q(x))
在token级别近似为：
KL ≈ log(P(token)) - log(Q(token))

其中：
P：Actor当前策略
Q：Reference模型策略

步骤6：计算综合奖励

# 初始化奖励矩阵
rewards = torch.zeros_like(sequence)  # 形状：[batch, seq_len]

# 每一步扣除KL惩罚
rewards = -self.kl_coef * kl_divergence

# 在序列末尾加上奖励模型的分数
for i in range(batch_size):
    end_idx = find_end_index(action_mask[i])  # 找到生成部分的末尾
    rewards[i, end_idx] += raw_rewards[i]

序列: [P1, P2, P3, R1, R2, R3, R4, R5]
奖励: [ 0,  0,  0, k1, k2, k3, k4, k5+raw]
      ↑           ↑                 ↑
    prompt部分   生成token        末尾token
    奖励为0    KL惩罚         KL惩罚+RM分数
    
其中：
k1..k5 = -kl_coef * KL(token_i)
raw = 奖励模型给出的分数

为什么这样设计？

1. 每一步KL惩罚：防止每一步都偏离参考模型

2. 只在最后给奖励：因为只有完整回答才能评估质量

3. prompt部分不计算奖励：因为我们不评估输入的好坏

步骤7：价值估计

values = self.critic(sequence)

Critic的任务：预测每个状态（token位置）的长期回报

例子：
序列: "解释一下机器学习[SEP]机器[SEP]学习[SEP]是..."
价值:  [3.1, 3.2, 3.3, 5.1, 6.2, 7.3, 8.0, ...]
      ↑              ↑
    开始状态        中间状态

Critic的输出结构：

输入: [batch, seq_len, hidden_size]
输出: [batch, seq_len]  # 每个位置的价值估计

四、经验缓冲区结构

experience_buffer = {
    "prompts": prompts,          # 原始提示
    "sequence": sequence,        # 完整序列
    "old_log_probs": old_log_probs,  # 旧策略的对数概率
    "rewards": rewards,          # 每一步的即时奖励
    "values": values,           # Critic的价值估计
    "mask": action_mask         # 哪些位置是生成的
}

一个完整轨迹示例：

# 单个样本的轨迹
prompt = "解释一下机器学习"

# Actor生成的回答
sequence = [
    "解释", "一下", "机器", "学习",  # prompt
    "机器", "学习", "是", "AI", "分支"  # 生成的回答
]

# 生成部分的掩码
action_mask = [0, 0, 0, 0, 1, 1, 1, 1, 1]  # 1表示生成部分

# Actor的对数概率
old_log_probs = [0.0, 0.0, 0.0, 0.0, -0.2, -0.3, -0.1, -0.4, -0.2]

# 参考模型的对数概率
ref_log_probs = [0.0, 0.0, 0.0, 0.0, -0.5, -0.6, -0.3, -0.7, -0.4]

# 奖励（假设kl_coef=0.1, raw_reward=8.0）
kl_penalties = -0.1 * (old_log_probs - ref_log_probs)
# 计算：第5个token: -0.1*(-0.2+0.5) = -0.1 * 0.3 = -0.03
# ... 以此类推
# 最后一个token加上RM奖励
rewards = [0, 0, 0, 0, -0.03, -0.03, -0.02, -0.03, 7.98]

# Critic的价值估计
values = [5.0, 5.1, 5.2, 5.3, 6.0, 6.5, 7.0, 7.5, 8.0]

没有KL惩罚的问题：
- 模型可能"走捷径"：找到奖励模型漏洞
- 过度优化：生成不自然的文本
- 灾难性遗忘：忘记语言基本能力

高KL惩罚（β=0.2）：
- 保守：变化小，收敛慢
- 安全：不易产生奇怪输出

低KL惩罚（β=0.01）：
- 激进：变化大，收敛快
- 风险：可能产生异常文本

六、采样策略细节

# 多种生成策略
def generate_responses(prompts, generation_params):
    # 策略1：贪心解码（确定性）
    if strategy == "greedy":
        return greedy_decode(prompts)
    
    # 策略2：采样（带温度）
    elif strategy == "sampling":
        return sample_with_temperature(prompts, temperature=0.7)
    
    # 策略3：集束搜索
    elif strategy == "beam_search":
        return beam_search(prompts, beam_width=4)

旧策略：采样时使用的策略
新策略：更新后的策略
重要性采样比率：r(θ) = π_new(a|s) / π_old(a|s)

作用：用旧策略的数据评估新策略

完整数据流：

输入提示
    ↓
Actor生成回答
    ↓
计算各项指标：
1. Actor对数概率
2. Reference对数概率  
3. 奖励模型分数
4. KL散度惩罚
5. Critic价值估计
    ↓
组合成经验轨迹
    ↓
存入经验缓冲区
    ↓
等待PPO更新阶段使用

第四部分：优势计算 (GAE) 与 PPO 目标函数

在上一部分，我们收集了一堆数据：sequence, rewards (包含 KL 惩罚), values (Critic 的预测), 和 old_log_probs。

但我们不能直接拿这些数据去更新模型。强化学习的核心难题在于信用分配 (Credit Assignment)：如果模型生成了 100 个词，最后得了个高分，那么这 100 个词里，哪一个词贡献最大？是第一个词选得好，还是第 99 个词选得好？

为了解决这个问题，我们需要计算优势函数 (Advantage Function)。而 PPO 算法的精髓，在于如何利用这个优势值安全地更新策略。

这部分包含两个核心数学逻辑：

1. GAE (Generalized Advantage Estimation)：一种平滑预测误差的方法，用于平衡方差和偏差。

2. PPO Clipping Loss：通过限制策略更新的幅度，保证训练稳定。

4. GAE 与 Loss 计算伪代码

# ==========================================
# 阶段 4: GAE Calculation & PPO Update
# 目标: 计算每个 Token 到底有多"好" (Advantage)，并构建 Loss
# ==========================================

def compute_gae_and_returns(rewards, values, gamma=0.99, lambda_=0.95):
    """
    计算广义优势估计 (GAE)。
    这是在时间轴上倒序计算的 (Dynamic Programming)。
    """
    advantages = torch.zeros_like(rewards)
    last_advantage = 0
    
    # 倒序遍历序列 (Time step t from T down to 0)
    # 因为当前的动作好坏取决于未来的收益
    for t in reversed(range(len(rewards))):
        # 1. 计算 TD Error (时序差分误差)
        # delta = r_t + gamma * V(s_{t+1}) - V(s_t)
        # 这里的 values 是 Critic 预测的
        if t == len(rewards) - 1:
            next_value = 0 # 结束状态价值为 0
        else:
            next_value = values[t + 1]
            
        delta = rewards[t] + gamma * next_value - values[t]
        
        # 2. 递归计算 Advantage (GAE 公式)
        # A_t = delta + (gamma * lambda) * A_{t+1}
        advantages[t] = delta + (gamma * lambda_) * last_advantage
        last_advantage = advantages[t]
        
    # Returns = Advantage + Value (用于训练 Critic)
    returns = advantages + values
    return advantages, returns

def ppo_loss_function(batch, actor, critic, epsilon=0.2):
    """
    计算 PPO 的最终 Loss
    batch: 包含 rollout 阶段收集的所有数据以及计算好的 advantages
    """
    # 1. 重新前向传播 (Re-forward)
    # 在 PPO 的内层循环中，Actor 参数在变，所以要重新计算 log_probs
    new_logits = actor(batch.prompts) 
    new_log_probs = get_log_probs(new_logits, batch.sequence)
    
    new_values = critic(batch.prompts) # Critic 也要更新
    
    # 2. 计算概率比率 (Ratio)
    # ratio = pi_new(a|s) / pi_old(a|s)
    # 为了数值稳定性，通常在 log 域计算：ratio = exp(log_new - log_old)
    ratio = torch.exp(new_log_probs - batch.old_log_probs)
    
    # 3. 计算 Actor Loss (Policy Loss) - PPO 的灵魂
    # 目标是最大化 Advantage，但不能偏离 old policy 太多
    surr1 = ratio * batch.advantages
    surr2 = torch.clamp(ratio, 1 - epsilon, 1 + epsilon) * batch.advantages
    
    # 取最小值（悲观估计），然后取负号因为我们要 Minimize Loss
    policy_loss = -torch.min(surr1, surr2).mean()
    
    # 4. 计算 Critic Loss (Value Loss)
    # 让 Critic 预测的 value 越接近真实的 returns 越好
    # 有时这里也会对 value 进行 clip
    value_loss = F.mse_loss(new_values, batch.returns)
    
    # 5. 总 Loss
    # 通常还会加一个 Entropy Bonus 鼓励探索 (loss -= coeff * entropy)
    total_loss = policy_loss + 0.5 * value_loss
    
    return total_loss

关键点解析 (Expert Note):

• 为什么要倒序 (Reversed Loop): GAE 的核心思想是当前动作的好坏不仅取决于当下的奖励，还取决于未来的价值。就像下棋，这一步走得好不好，得看它给后续几步创造了多大的优势。倒序计算可以方便地把未来的信号传播回当前时刻。

• Advantage Normalization (白化): 在实际代码中，在计算 Loss 之前，通常会对一个 batch 内的 advantages 进行归一化 (adv - mean) / (std + 1e-8)。这对于训练的稳定性至关重要，否则梯度的尺度会剧烈波动。

• Ratio 的物理意义:

  • 如果 ratio > 1，说明新策略比旧策略更倾向于采取这个动作。

  • 如果这个动作的 advantage > 0 (是好动作)，我们鼓励 ratio 变大（直到被 clip）。

  • 如果这个动作的 advantage < 0 (是坏动作)，我们鼓励 ratio 变小。

• Critic 的角色: 注意代码中的 returns = advantages + values。Actor 的目标是最大化 Advantage，而 Critic 的目标是逼近真实的 Return。Critic 越准，Advantage 的估计就越准，Actor 学得就越快。

PPO更新阶段详解：GAE计算与损失函数

一、PPO更新阶段概述

Rollout阶段（采样） → 收集经验
    ↓
GAE计算（分析） → 计算每个动作的好坏程度
    ↓
PPO更新（学习） → 更新Actor和Critic
    ↓
重复直到收敛

经验缓冲区包含：
1. 状态序列 (sequence)
2. 动作 (生成的token)
3. 奖励 (reward)
4. 旧策略概率 (old_log_probs)
5. Critic预测值 (values)

目标：用这些数据更新策略，使其生成更好的回答

二、广义优势估计（GAE）详解

1. 为什么需要GAE？

问题：在语言模型中，奖励只在序列最后给出，如何评估每个token的好坏？

例子：

序列：我喜欢吃[SEP]苹果[SEP]很甜[EOS]
奖励：只在EOS位置有RM分数=8.5
问题：如何分配这个8.5的分数到每个token？

解决方案：

• 蒙特卡洛方法：用完整序列的回报

• 时间差分：用相邻状态的价值差

• GAE：两者结合，平衡偏差和方差

2. GAE数学原理

2.1 时序差分误差（TD Error）

δ_t = r_t + γ * V(s_{t+1}) - V(s_t)

其中：
- r_t: 时间步t的即时奖励
- V(s_t): 状态s_t的价值
- γ: 折扣因子（通常0.99）

δ_t 衡量了Critic预测的误差
- δ_t > 0: 实际结果比预测好
- δ_t < 0: 实际结果比预测差
- δ_t = 0: 预测准确

2.2 GAE定义

A_t^{GAE(γ,λ)} = Σ_{l=0}^{∞} (γλ)^l δ_{t+l}

其中λ是权衡参数：
- λ=0: 只看一步TD误差
- λ=1: 看完整蒙特卡洛回报
- 通常λ=0.95: 平衡两者

具体计算示例

序列长度: 4 (prompt不算)
即时奖励: [-0.01, -0.01, -0.01, 8.5]  # 前三步KL惩罚，最后一步RM奖励
Critic预测: [7.0, 7.5, 8.0, 8.5]  # 每个位置的价值预测
γ=0.99, λ=0.95

计算过程：

步骤3 (t=3):
next_value = 0
δ_3 = 8.5 + 0.99 * 0 - 8.5 = 0
A_3 = 0 + (0.99 * 0.95)*0 = 0

步骤2 (t=2):
next_value = values[3] = 8.5
δ_2 = -0.01 + 0.99 * 8.5 - 8.0 = -0.01 + 8.415 - 8.0 = 0.405
A_2 = 0.405 + (0.99 * 0.95)*0 = 0.405

步骤1 (t=1):
next_value = values[2] = 8.0
δ_1 = -0.01 + 0.99 * 8.0 - 7.5 = -0.01 + 7.92 - 7.5 = 0.41
A_1 = 0.41 + (0.99 * 0.95)*0.405 = 0.41 + 0.381 ≈ 0.791

步骤0 (t=0):
next_value = values[1] = 7.5
δ_0 = -0.01 + 0.99 * 7.5 - 7.0 = -0.01 + 7.425 - 7.0 = 0.415
A_0 = 0.415 + (0.99 * 0.95)*0.791 = 0.415 + 0.744 ≈ 1.159

结果：
优势: [1.159, 0.791, 0.405, 0.0]
回报: 优势 + 价值 = [8.159, 8.291, 8.405, 8.5]

实例逻辑

1. 单个token的PPO更新

场景：

提示："我喜欢吃"
生成token："苹果"
旧策略概率：log_prob_old = -0.5 (概率≈0.6065)
Critic预测：value = 7.0
优势：advantage = 1.0
回报：return = 8.0
ε=0.2

计算：

1. 新策略概率：log_prob_new = -0.3 (概率≈0.7408)
2. 比率：ratio = exp(-0.3 - (-0.5)) = exp(0.2) ≈ 1.221
3. 策略损失：
   surr1 = 1.221 * 1.0 = 1.221
   clipped_ratio = clamp(1.221, 0.8, 1.2) = 1.2
   surr2 = 1.2 * 1.0 = 1.2
   policy_loss = -min(1.221, 1.2) = -1.2
4. Critic损失：
   新Critic预测：value_new = 7.5
   value_loss = (7.5 - 8.0)^2 = 0.25
5. 总损失：total_loss = -1.2 + 0.5 * 0.25 = -1.075

模型学会了：
1. 提高生成"苹果"的概率（从0.6065到0.7408）
2. Critic调整预测（从7.0到7.5，更接近实际回报8.0）

第五部分：RLHF 完整训练循环 (The Main Loop)

RLHF 的训练通常是嵌套循环结构：

• 外层循环 (Step)：采样数据 (Rollout)。

• 内层循环 (PPO Epoch)：利用同一批数据多次更新梯度 (这就是 PPO 相比于 A2C 等算法 sample efficiency 高的原因)。

5. 主循环伪代码

# ==========================================
# 阶段 5: The Full RLHF Training Loop
# 目标: 整合所有组件，开始迭代
# ==========================================

def run_rlhf_training(args):
    # 1. 初始化四个模型 (通常 Actor/Critic 为 Train模式, Ref/RM 为 Eval模式)
    # 实际工程中，Ref 和 RM 常常使用 Offloading 技术放在 CPU 以节省显存
    ppo_trainer = PPOTrainer(sft_model, reward_model, kl_coef=0.1)
    
    # 优化器只优化 Actor 和 Critic 的参数
    optimizer = AdamW(ppo_trainer.get_learnable_parameters(), lr=1e-6)
    
    # === 外层循环：Steps / Episodes ===
    for step in range(args.total_steps):
        
        # 1. 获取 Batch Prompts
        prompts = prompt_dataset.sample(args.batch_size)
        
        # 2. 采样 (Rollout Phase) - 显存消耗峰值 1
        # Actor 生成数据，Critic 预测价值，RM 打分
        experience_data = ppo_trainer.generate_experience(prompts)
        
        # 3. 计算优势 (GAE Phase) - CPU/GPU 计算
        advantages, returns = compute_gae_and_returns(
            experience_data['rewards'], 
            experience_data['values']
        )
        # 将计算结果存回 experience_data
        experience_data['advantages'] = advantages
        experience_data['returns'] = returns
        
        # 4. PPO 更新 (Update Phase) - 显存消耗峰值 2
        # 我们使用这批数据更新多次 (PPO Epochs)
        for ppo_epoch in range(args.ppo_epochs):
            
            # Shuffle 数据并切分为 Mini-batches
            # 这一步是为了 SGD 的随机性，防止过拟合最近的数据
            data_loader = create_minibatch_loader(experience_data, args.mini_batch_size)
            
            for mini_batch in data_loader:
                # 计算 Loss (包含 Policy Loss, Value Loss, Entropy Bonus)
                # 这一步会再次前向传播 Actor 和 Critic
                loss = ppo_loss_function(mini_batch, ppo_trainer.actor, ppo_trainer.critic)
                
                # 反向传播
                optimizer.zero_grad()
                loss.backward()
                
                # 梯度裁剪 (Gradient Clipping) - 防止梯度爆炸，稳定训练
                torch.nn.utils.clip_grad_norm_(ppo_trainer.get_learnable_parameters(), 1.0)
                
                optimizer.step()
        
        # 5. 监控与日志 (Logging)
        # 观察 Mean Reward 和 Mean KL 是最重要的指标
        mean_reward = experience_data['rewards'].mean()
        mean_kl = (experience_data['old_log_probs'] - experience_data['new_log_probs']).mean()
        
        print(f"Step {step}: Mean Reward = {mean_reward:.4f}, Mean KL = {mean_kl:.4f}")
        
        # 6. 动态调整 KL 系数 (可选 Trick)
        # 如果 KL 太大，说明 Actor 跑偏了，增大惩罚系数
        # 如果 KL 太小，说明 Actor 没学到东西，减小惩罚系数
        ppo_trainer.update_kl_coef(mean_kl, target_kl=0.1)

    print("RLHF 训练结束！")

如果只照着教科书写 PPO，在几十亿参数的模型上通常是训不起来的。以下是支撑现代 LLM RLHF 的核心工程细节：

1. 显存优化 (The VRAM Bottleneck):                ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​问题: 同时加载 4 个模型（Actor, Critic, Ref, RM）太占显存。        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​解法：1.Offloading: 既然 Ref Model 和 RM 不需要反向传播，推理完一次后直接把权重卸载到 CPU，PPO 更新时只留 Actor 和 Critic 在 GPU。        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​         Hydra Head: Actor 和 Critic 共享同一个 Backbone，只是最后一层 Head 不同（但这会导致训练不稳定，通常不推荐）。        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​        ​​​​​​​ LoRA: Actor 和 Critic 仅训练 LoRA 权重，冻结 Backbone。

2. Reward Normalization (奖励归一化):RM 输出的分数可能范围很大（比如 -10 到 10），这会导致 Value Loss 难以收敛。通常会对 Reward 进行归一化（减均值除方差），或者使用 Running Mean/Std 缩放，这对训练稳定性至关重要。

3. Token-level vs Sentence-level KL:虽然 KL 是 Token 级别的，但为了计算方便，有时只统计 Average KL。但在代码实现中，必须确保 KL Penalty 是逐个 Token 施加在 Advantage 上的，否则无法进行细粒度的动作修正。

4. 初始化陷阱:Critic 的初始化非常关键。如果 Critic 一开始预测的值离真实 Reward 太远，GAE 就会有巨大的方差，导致 Actor 在第一步就“崩”了。通常建议先对 Critic 进行几轮 Warm-up 训练，或者将 Critic 的输出层 Bias 初始化为 Reward 的均值。

5. DPO (Direct Preference Optimization) 的兴起:既然 PPO 这么麻烦（4 个模型、超参数敏感、显存大），现在学术界和工业界（如 Llama 3）开始转向 DPO。DPO 从数学上推导证明了可以直接通过优化 SFT 模型的 Cross-Entropy Loss 来隐式地最大化 Reward，从而去掉了 RM 和 Critic，甚至去掉了 PPO 循环，把 RL 问题变成了 Supervised Learning 问题。

精华总结，全流程解释：

一、RLHF训练全貌

训练的三个阶段

阶段1：SFT监督微调
输入：基础大模型 + 高质量的问答对
输出：能够理解并遵循指令的模型

阶段2：奖励模型训练
输入：SFT模型 + 人类标注的"好/坏"回答对比
输出：能够评价回答质量的奖励模型

阶段3：RLHF强化学习训练
输入：SFT模型 + 奖励模型
输出：优化后的最终模型

二、PPO训练循环的六个核心步骤

步骤1：初始化准备

目标：准备好所有需要的模型和工具

四个关键角色：

1. Actor（演员）：要训练的语言模型，负责生成回答
2. Critic（评论家）：价值网络，评估生成的质量
3. Reference（参考模型）：原始的SFT模型，作为基准
4. Reward Model（奖励模型）：人类偏好的代理，负责打分

- Actor和Critic：可训练，不断更新
- Reference和Reward Model：冻结的，作为稳定的标准
- 优化器：只优化Actor和Critic的参数

步骤2：采样（生成回答）

目标：用当前的Actor模型生成一批回答

具体过程：

输入：一批提示（如"解释什么是量子计算"）
过程：
1. Actor读取提示
2. 逐个token生成回答
3. 记录每个token的生成概率
输出：完整的prompt+response序列

提示："解释什么是机器学习"

Actor生成："机器学习是人工智能的一个分支，让计算机能够从数据中学习，而无需明确编程。"

同时记录：
- 生成"机器"的概率：0.85
- 生成"学习"的概率：0.90
- 生成"是"的概率：0.95
- ...（每个token的概率）

1. 用Critic评估序列价值：每个位置的价值
2. 用Reward Model评估整体质量：一个总分
3. 用Reference模型计算参考概率
4. 计算KL散度惩罚
5. 计算总奖励

步骤3：计算优势（好坏程度）

目标：计算每个动作（生成每个token）到底有多"好"

为什么需要优势：

问题：奖励只在最后给出，如何评估中间每个token的贡献？

举例：
回答："机器学习是人工智能的一个重要分支"
奖励：最后得到8.5分

我们需要知道：
- 生成"机器"对这个8.5分贡献了多少？
- 生成"学习"贡献了多少？
- 每个token应该得到多少"功劳"？

GAE（广义优势估计）计算：

核心思想：综合考虑即时奖励和未来潜力

从后往前计算：
1. 最后一个token：只看即时奖励
2. 倒数第二个token：即时奖励 + 一部分未来价值
3. 倒数第三个token：即时奖励 + 更多未来价值
4. ...一直往前推

结果：
- 早期token的优势较大（对最终结果影响大）
- 后期token的优势较小（不确定性小）
- 优势可以是正数（好动作）或负数（坏动作）

数学原理：

优势 = 当前实际回报 - 预期回报

其中：
- 当前实际回报：实际获得的奖励
- 预期回报：Critic预测的价值
- 正优势：比预期好
- 负优势：比预期差

步骤4：PPO更新（学习优化）

目标：用收集到的经验数据更新Actor和Critic

PPO的核心技巧：

关键问题：如何用旧数据学习新策略？

解决方案：重要性采样 + 裁剪

比喻：用去年的考试题练习今年的考试
- 如果题目变化不大，可以重用
- 如果变化很大，需要谨慎使用

具体更新过程：

A. 策略更新（Actor）：

目标：让好动作的概率增加，坏动作的概率减少

计算步骤：
1. 比较新旧概率：新概率 ÷ 旧概率
2. 如果优势为正（好动作）：增加这个token的概率
3. 如果优势为负（坏动作）：减少这个token的概率
4. 但要有度：限制每次更新的幅度

裁剪机制：
- 如果概率变化太大（比如超过20%），就裁剪到20%
- 防止一步迈太大，破坏现有能力

生成token"机器"：
- 旧概率：0.85
- 新概率：0.90
- 比率：0.90/0.85 = 1.06
- 优势：+1.5（好动作）

计算：
- 无裁剪目标：1.06 × 1.5 = 1.59
- 裁剪后（如果超过1.2）：1.2 × 1.5 = 1.8
- 实际目标：min(1.59, 1.8) = 1.59

优化方向：让这个token的概率继续提高

B. 价值更新（Critic）：

目标：让Critic的预测更准确

计算：预测值应该接近实际回报
损失 = (预测值 - 实际回报)²

比喻：
- Critic是股票分析师
- 预测值是分析师的股价预测
- 实际回报是真实股价
- 目标：让预测越来越准

用同一批数据更新多次（通常是4-10次）
原因：
1. 生成数据成本高，要充分利用
2. 小步多次更新更稳定
3. 每次用不同的数据子集，增加随机性

RLHF训练本质上是在多个目标间寻找平衡：

目标1：最大化人类偏好（奖励）
目标2：最小化模型偏离（KL散度）
目标3：保持探索能力（熵）
目标4：稳定训练过程（裁剪）

成功的RLHF训练需要：
1. 高质量的奖励模型
2. 合理的超参数设置
3. 细致的监控调整
4. 充分的耐心迭代

通过这个精密的训练过程，原始的预训练大模型最终转变为能够理解人类意图、遵循人类价值观、提供有用帮助的智能助手。这个过程虽然复杂，但正是它让人工智能真正"理解"了什么是人类眼中的"好"。