范式层：RLHF (Reinforcement Learning from Human Feedback)
        ├── 需要reward model
        └── 需要RL算法来优化

算法层：PPO, TRPO, A3C, ...
        └── 传统RL算法，用于RLHF

新方法：DPO, IPO, KTO, SimPO
        └── 直接从偏好数据学习，不需要显式reward model

阶段1: 监督微调（SFT）
├── 输入：高质量示范数据
└── 输出：SFT模型（基础对齐）

阶段2: 训练Reward Model
├── 输入：人类偏好数据 (prompt, chosen, rejected)
└── 输出：独立的奖励模型（打分器）

阶段3: PPO强化学习
├── 输入：SFT模型 + Reward Model
├── 过程：
│   1. SFT模型生成回复
│   2. Reward Model给回复打分
│   3. 用PPO算法更新策略模型
│   4. 重复迭代
└── 输出：对齐后的模型

# 策略模型：决定生成什么
policy_output = policy_model.generate(prompt)

# 奖励模型：评价生成的好坏
reward = reward_model.score(prompt, policy_output)

# 价值模型：预测平均能得多少分（用于计算advantage）
baseline = value_model.predict(prompt)
advantage = reward - baseline  # 比平均好多少

# 参考模型：防止模型变化太激进
kl_penalty = KL(policy_model || reference_model)
final_reward = reward - 0.1 * kl_penalty

# 旧策略: P(好回答) = 0.1, P(坏回答) = 0.9
# 新策略: P(好回答) = 0.9, P(坏回答) = 0.1  # 变化太大！

→ 可能导致：
  - 模型突然"忘记"之前学的东西
  - 训练不稳定
  - 输出质量崩溃

epsilon = 0.2  # clip范围

# 概率比
ratio = π_new(action) / π_old(action)

# 限制在 [0.8, 1.2] 范围内
clipped_ratio = clip(ratio, 1-epsilon, 1+epsilon)

# 取保守的更新
objective = min(ratio * advantage, clipped_ratio * advantage)

→ 每次更新都是"小步前进"
→ 训练稳定

阶段1: 监督微调（SFT）
├── 与PPO相同
└── 得到基础模型

阶段2: 直接偏好优化
├── 输入：偏好对数据 (prompt, chosen, rejected)
├── 关键创新：跳过训练Reward Model
└── 直接从偏好数据学习

# DPO损失（数学形式）
loss = -log(σ(β * log(π_θ(y_w|x) / π_ref(y_w|x)) 
              - β * log(π_θ(y_l|x) / π_ref(y_l|x))))

# 其中：
# y_w: 被选中的回复（win）
# y_l: 被拒绝的回复（lose）
# π_θ: 当前模型（会更新）
# π_ref: 参考模型（固定）
# β: 温度参数

# 传统RLHF：
# 步骤1: 训练 reward_model: r(x,y) ≈ 人类偏好
# 步骤2: 优化策略: max E[reward_model(x, policy(x))]

# DPO：
# 一步到位：让 P(chosen) / P(rejected) 的比值变大
# 不需要显式的reward，偏好信息"隐式"编码在损失函数中

# ❌ 预训练模型：
User: "2+2等于多少？"
Model: "这是一个数学问题，让我来..." [可能输出各种格式]

# ✅ SFT后：
User: "2+2等于多少？"
Model: """
<start_working_out>
2 + 2 = 4
<end_working_out>
<SOLUTION>4</SOLUTION>
"""  # 严格遵循你定义的格式

# 预训练模型：有知识，但不知道"怎么用"
预训练模型："量子纠缠是..."[知道概念]

# SFT：教它如何组织答案、如何推理
SFT后："首先，量子纠缠的定义是...
        其次，它的物理原理包括...
        最后，应用场景有..."

# 如果直接RL：
初始模型 → 随机探索 → 99.99%的输出都得负分 → 无法学习

# 有SFT基础：
SFT模型 → 已经在"好答案"附近 → 探索空间小 → RL可以微调

# 典型的Reward Model结构
class RewardModel(nn.Module):
    def __init__(self, base_model):
        self.base_model = base_model  # 比如 Llama-7B
        self.value_head = nn.Linear(4096, 1)  # 输出标量分数

    def forward(self, input_ids):
        # 编码整个(prompt + response)
        hidden = self.base_model(input_ids).last_hidden_state

        # 取最后一个token的hidden state
        last_hidden = hidden[:, -1, :]

        # 输出一个分数
        score = self.value_head(last_hidden)
        return score  # 标量，比如 8.5

# 训练数据：人类偏好对
dataset = [
    {
        "prompt": "写一首关于春天的诗",
        "response_A": "春眠不觉晓，处处闻啼鸟...",  # 人类选择 👍
        "response_B": "春天来了..."                  # 人类拒绝 👎
    }
]

# 训练目标：Bradley-Terry模型
# P(A > B) = σ(r(A) - r(B))

loss = -log(σ(reward_model(prompt, A) - reward_model(prompt, B)))

# 训练后，reward_model能够预测人类偏好分数

def api_reward_function(prompt, completion):
    """用GPT-4作为reward model"""

    judge_prompt = f"""
请对以下回答打分（0-10分）：

问题：{prompt}
回答：{completion}

评分标准：
- 准确性（0-4分）
- 格式规范（0-3分）
- 推理清晰（0-3分）

只输出分数，如：8.5
"""

    response = openai.ChatCompletion.create(
        model="gpt-4",
        messages=[{"role": "user", "content": judge_prompt}]
    )

    score = float(response.choices[0].message.content.strip())
    return score

# 方式1：硬编码规则（适合有明确标准的任务）
def rule_based_reward(prompt, completion, answer):
    if extract_answer(completion) == answer:
        return 10.0  # 完全正确
    elif is_close(extract_answer(completion), answer):
        return 5.0   # 接近正确
    else:
        return 0.0   # 错误

# 方式2：训练Reward Model（适合主观评价）
reward_score = reward_model(prompt, completion)  # 学习人类偏好

# 方式3：LLM API（灵活但昂贵）
reward_score = gpt4_judge(prompt, completion)

# ========== PPO的一次迭代 ==========

# 初始化
prompt = "写一首关于春天的诗"
current_policy = your_model          # 会被更新
reference_policy = copy(your_model)  # 固定不变
value_network = separate_model       # 会被更新
reward_model = trained_rm            # 固定不变

# Step 1: 采样（用当前策略生成回答）
response = current_policy.generate(prompt)
# response = "春眠不觉晓，处处闻啼鸟"

# Step 2: 计算reward
reward = reward_model(prompt, response)  # 比如得到 8.5分

# Step 3: 计算value（预测的期望回报）
value = value_network(prompt)  # 预测这个prompt的平均回报，比如 7.0

# Step 4: 计算advantage（这个回答比平均好多少）
advantage = reward - value  # 8.5 - 7.0 = 1.5

# Step 5: 计算概率比（当前策略 vs 旧策略）
current_prob = current_policy.log_prob(response | prompt)  # -2.3
old_prob = reference_policy.log_prob(response | prompt)    # -2.5
ratio = exp(current_prob - old_prob)  # exp(0.2) ≈ 1.22

# Step 6: PPO的Clipped目标
epsilon = 0.2
unclipped = ratio * advantage          # 1.22 * 1.5 = 1.83
clipped = clip(ratio, 0.8, 1.2) * advantage  # 1.2 * 1.5 = 1.8
objective = min(unclipped, clipped)    # 取保守的：1.8

# Step 7: 反向传播更新
loss = -objective  # 负号因为要最大化objective
loss.backward()
optimizer.step()  # ✅ 更新 current_policy

PPO更新就像"小心翼翼地改进"：

1. 生成一个回答
2. 看看得了多少分
3. 跟平均水平比较（advantage）
4. 但不能改得太激进（clip）
5. 小步更新参数

就像学生改进答题策略：
- 这次比平均好 → 以后多这样答 ✅
- 但不能突然完全改变风格 ⚠️
- 一点点改进，保持稳定 👍

# ========== 训练开始 ==========
# 初始模型（都一样）
policy_model = load_model("llama-7b-sft")      # 策略模型（会更新）
reference_model = load_model("llama-7b-sft")   # 参考模型（固定）
reference_model.eval()  # 设为评估模式
for param in reference_model.parameters():
    param.requires_grad = False  # 冻结所有参数

# ========== 迭代1 ==========
prompt = "如何学习编程？"

# 策略模型生成（会不断改进）
response_v1 = policy_model.generate(prompt)
# "学习编程需要多练习"

reward = reward_model.score(prompt, response_v1)  # 6.0分

# 用参考模型计算KL散度（防止偏离太远）
kl_divergence = KL(policy_model || reference_model)  # 0.05
final_reward = reward - 0.1 * kl_divergence  # 6.0 - 0.005 = 5.995

# 更新策略模型
update(policy_model)  # ✅ 更新
# reference_model 不变 ❌

# ========== 迭代100 ==========
# 策略模型已经改进很多
response_v100 = policy_model.generate(prompt)
# "学习编程建议：1)选择Python入门 2)每天练习 3)做项目..."

reward = 9.0

# 但要检查：是不是偏离原始模型太远？
kl_divergence = KL(policy_model || reference_model)  # 0.8 (比较大了)
final_reward = 9.0 - 0.1 * 0.8 = 8.92

# 如果KL太大，会被惩罚，防止模型"乱说"

# 没有参考模型的灾难：
# 假设reward model有bug，对"42"这个回答给高分

Iteration 1:
Q: "人生的意义是什么？"
A: "人生的意义在于..." → reward = 5.0

Iteration 50:
Q: "人生的意义是什么？"
A: "42" → reward = 10.0 (reward model的bug)
# 模型发现"42"能得高分

Iteration 100:
Q: "人生的意义是什么？"
A: "42"
Q: "如何学编程？"
A: "42"  # 什么都答"42"，模式崩溃！
Q: "今天天气？"
A: "42"

# 有参考模型的保护：
KL(policy || reference) 会变得非常大
→ final_reward = 10.0 - 0.1 * huge_KL = 负数
→ 模型不会学习这种"作弊"行为

# ========== DPO训练 ==========
policy_model = load_model("llama-7b-sft")
reference_model = load_model("llama-7b-sft")  # 复制一份
reference_model.eval()  # 冻结

for batch in dataset:
    prompt = batch['prompt']
    y_win = batch['chosen']
    y_lose = batch['rejected']

    # 计算当前策略的log概率
    log_prob_win_policy = policy_model.log_prob(y_win | prompt)
    log_prob_lose_policy = policy_model.log_prob(y_lose | prompt)

    # 计算参考策略的log概率（不更新）
    with torch.no_grad():
        log_prob_win_ref = reference_model.log_prob(y_win | prompt)
        log_prob_lose_ref = reference_model.log_prob(y_lose | prompt)

    # DPO损失（让chosen的相对概率更高）
    beta = 0.1
    loss = -log_sigmoid(
        beta * (log_prob_win_policy - log_prob_win_ref) -
        beta * (log_prob_lose_policy - log_prob_lose_ref)
    )

    # 更新
    loss.backward()
    optimizer.step()  # 只更新 policy_model
    # reference_model 保持不变！

# 把最大概率的输出作为label，然后计算loss？
# 不是这样的！

# DPO直接优化整个序列的log概率
# 不需要label，通过chosen vs rejected的相对概率来更新

# ========== Step 1: 准备输入 ==========
batch = {
    "prompt": "解释量子纠缠",
    "chosen": "量子纠缠是指两个粒子之间存在关联...",
    "rejected": "量子纠缠就是很神奇的现象"
}

# 拼接（注意：不是作为label，而是作为input）
chosen_input = tokenizer(batch["prompt"] + batch["chosen"])
rejected_input = tokenizer(batch["prompt"] + batch["rejected"])

# ========== Step 2: 计算每个token的log概率 ==========
def compute_log_prob(model, input_ids, prompt_length):
    # 前向传播
    outputs = model(input_ids)
    logits = outputs.logits  # [batch, seq_len, vocab_size]

    # 只关注response部分（跳过prompt）
    response_logits = logits[:, prompt_length-1:-1, :]
    response_labels = input_ids[:, prompt_length:]

    # 计算每个token的log概率
    log_probs = F.log_softmax(response_logits, dim=-1)

    # 选出实际生成的token的log概率
    token_log_probs = torch.gather(
        log_probs, 
        dim=-1, 
        index=response_labels.unsqueeze(-1)
    ).squeeze(-1)

    # 整个response的log概率 = 所有token log概率之和
    sequence_log_prob = token_log_probs.sum(dim=-1)

    return sequence_log_prob

# ========== Step 3: 计算policy和reference的log概率 ==========
# Policy model（正在训练的）
policy_chosen_logprob = compute_log_prob(
    policy_model, chosen_input_ids, prompt_length
)
policy_rejected_logprob = compute_log_prob(
    policy_model, rejected_input_ids, prompt_length
)

# Reference model（固定不变的）
with torch.no_grad():  # 不需要梯度！
    ref_chosen_logprob = compute_log_prob(
        reference_model, chosen_input_ids, prompt_length
    )
    ref_rejected_logprob = compute_log_prob(
        reference_model, rejected_input_ids, prompt_length
    )

# ========== Step 4: 计算DPO损失 ==========
beta = 0.1

# 计算相对对数比率
policy_logratios = policy_chosen_logprob - policy_rejected_logprob
ref_logratios = ref_chosen_logprob - ref_rejected_logprob

# DPO损失：让chosen的相对概率更高
logits = beta * (policy_logratios - ref_logratios)
loss = -F.logsigmoid(logits).mean()

# ========== Step 5: 反向传播 ==========
loss.backward()  # 梯度只回传到policy_model！
optimizer.step()

loss → logits → policy_logratios 
             → policy_chosen_logprob, policy_rejected_logprob
             → policy_model的每个token的logits
             → policy_model的参数 ✅ 被更新

reference_model没有梯度 ❌ 不更新

# ========== SFT: Token-level的梯度 ==========
Input:  "天空是"
Target: "天空是蓝色的"

Position 0: Loss(predict="天" | input="")
Position 1: Loss(predict="空" | input="天")
Position 2: Loss(predict="是" | input="天空")
Position 3: Loss(predict="蓝" | input="天空是")
# → 每个token独立计算loss

# ========== DPO: Sequence-level的梯度 ==========
Chosen:   "天空是蓝色的"
Rejected: "天空是绿色的"

# 计算整个序列的log概率差
log_prob_chosen = Σ log P(token_i | previous)
log_prob_rejected = Σ log P(token_i | previous)

# Loss让chosen的概率相对更高
loss = -sigmoid(β * (log_prob_chosen - log_prob_rejected - ...))
# → 梯度回传到影响整个序列概率的参数

# 每次更新后，policy变成reference？
# 不是这样的！

# Reference model在整个DPO/PPO训练中保持固定
# 始终是训练开始时的SFT模型

# ========== DPO训练的完整流程 ==========

# 初始化（只在开始时做一次）
policy_model = load_model("llama-7b-sft")      # 会被更新
reference_model = load_model("llama-7b-sft")   # 克隆一份
reference_model.eval()                          # 设为评估模式
for param in reference_model.parameters():
    param.requires_grad = False                 # 冻结所有参数

# 训练循环
for epoch in range(num_epochs):
    for batch in dataloader:
        # Step 1: policy_model前向传播（需要梯度）
        policy_logprobs = compute_log_prob(policy_model, ...)

        # Step 2: reference_model前向传播（不需要梯度）
        with torch.no_grad():
            ref_logprobs = compute_log_prob(reference_model, ...)

        # Step 3: 计算loss并更新policy_model
        loss = dpo_loss(policy_logprobs, ref_logprobs)
        loss.backward()
        optimizer.step()  # ✅ 只更新policy_model

        # reference_model始终保持不变！❌ 不更新

如果reference_model也跟着更新：

Iteration 1:
policy: 远离初始点
reference: 也远离初始点
→ 失去了"初始分布"的参考

Iteration 1000:
policy: 可能已经偏离很远
reference: 也偏离了
→ KL惩罚失效，可能reward hacking

# 如果reference也更新（错误做法）
beta * (log(π_policy(chosen)) - log(π_policy(rejected))
      - log(π_ref(chosen)) + log(π_ref(rejected)))

# 如果policy和ref都朝一个方向更新，会导致：
# π_policy(chosen)/π_ref(chosen) → 1
# π_policy(rejected)/π_ref(rejected) → 1
# → 损失变为0，没有训练信号！

# ========== 动态更新reference（较少用）==========

# Iteration 1-100:
reference_model = initial_policy  # 固定

# Iteration 101-200:
reference_model = copy(policy_model_at_100)  # 更新一次
# 之后再固定100步

# Iteration 201-300:
reference_model = copy(policy_model_at_200)
# ...

# 这种做法的理由：
# - 避免policy偏离太远
# - 渐进式优化

# 但问题：
# - 容易过拟合
# - 失去初始分布的约束
# - 训练不稳定

概率链式法则 = 数学工具（概率论）
Policy Gradient = 优化方法（强化学习）

关系：Policy Gradient在计算时使用链式法则

# ========== 纯粹的概率论 ==========

# 定义：联合概率的分解
P(A, B, C) = P(A) · P(B|A) · P(C|A,B)

# 在文本生成中：
P("吃苹果") = P(吃) · P(苹|吃) · P(果|吃苹)

# 取对数：
log P("吃苹果") = log P(吃) + log P(苹|吃) + log P(果|吃苹)
                 = Σ log P(token_i | previous_tokens)

# 这是概率论的数学事实，与机器学习无关

# ========== 强化学习的优化方法 ==========

# 目标：最大化期望回报
J(θ) = E_τ~π_θ[R(τ)]
     = Σ_τ π_θ(τ) · R(τ)  # 对所有可能轨迹求期望

# Policy Gradient Theorem（核心定理）：
∇_θ J(θ) = E_τ~π_θ[R(τ) · ∇_θ log π_θ(τ)]

# 这是强化学习的优化理论

# ========== Policy Gradient需要计算 log π(τ) ==========

# 对于文本生成，τ = (token_1, token_2, ..., token_n)

# 这里用到概率链式法则：
log π_θ(τ) = log π_θ(token_1, token_2, ..., token_n)
           = log π_θ(token_1) 
             + log π_θ(token_2|token_1)
             + log π_θ(token_3|token_1,token_2)
             + ...
           = Σ_i log π_θ(token_i | token_<i)  # 链式法则！

# 所以梯度：
∇_θ log π_θ(τ) = Σ_i ∇_θ log π_θ(token_i | token_<i)

# 最终的Policy Gradient：
∇_θ J(θ) = E[R(τ) · Σ_i ∇_θ log π_θ(token_i | token_<i)]

概率链式法则 = "计算器"
- 告诉你如何计算P(整个序列)
- 数学工具

Policy Gradient = "优化策略"
- 告诉你如何调整参数让期望回报最大
- 优化目标

Policy Gradient在计算时使用链式法则
- 就像用计算器来执行优化策略

# ========== 自回归能力不仅没有损失，反而强化了 ==========

# 训练前：
P(好的回答) = 0.001  # log P = -6.9
P(差的回答) = 0.01   # log P = -4.6

# RL训练：优化整个序列的概率
# → 模型学会：在自回归的每一步，都倾向于选择能导向"好回答"的token

# 训练后：
P(好的回答) = 0.01   # log P = -4.6 ✅ 提高了
P(差的回答) = 0.001  # log P = -6.9 ✅ 降低了

# 自回归能力没有损失，反而学会了"长期规划"：
# - 在生成第1个token时，就考虑整个回答的质量
# - 避免生成"局部最优但全局差"的token

# 场景：解数学题
Question: "5 + 7 = ?"

# Response 1（推理清晰）：
"让我计算一下：5 + 7 = 12"
# Reward: 10.0（正确且有推理）

# Response 2（直接答案）：
"12"
# Reward: 8.0（正确但无推理）

# RL训练目标：
# 让 Response 1 的整体概率提高
# → 模型学会：生成"让"时，要考虑后续能形成完整推理
# → 这需要考虑整个序列，而不是单个token

# ========== DPO的多样性问题 ==========

Question: "描述一下春天"

# 训练数据中的 chosen response：
chosen = "春天是万物复苏的季节，鸟语花香，生机勃勃。"

# 其他同样好的表达（但不在训练数据中）：
alternative_1 = "春暖花开，大地回春，处处洋溢着新生的气息。"
alternative_2 = "春风拂面，柳绿花红，是一年中最美好的时光。"

# ========== DPO训练后会发生什么？ ==========

# 训练前：
P(chosen) = 0.01
P(alternative_1) = 0.01  # 概率相似
P(alternative_2) = 0.01

# DPO训练后：
P(chosen) = 0.15         # ✅ 显著提高
P(alternative_1) = 0.005 # ❌ 可能降低！
P(alternative_2) = 0.004 # ❌ 可能降低！

# 问题：虽然chosen的质量提高了，但多样性丧失了

# ========== DPO优化的目标 ==========

# DPO鼓励：
# 1. 提高 chosen 的概率 ↑↑  (直接优化)
# 2. 降低 rejected 的概率 ↓↓ (直接优化)

# 但没有机制鼓励：
# 3. 保持 alternatives 的概率

# 结果（概率归一化）：
# - chosen ↑↑
# - rejected ↓↓
# - alternatives ↓ (间接受影响)

# ========== 训练进展 ==========

# Epoch 0 (初始SFT模型)
>>> generate("推荐一本书", temperature=0.8, num_samples=5)
[
    "推荐《三体》，这是刘慈欣的科幻巨著...",
    "我建议你读一读《百年孤独》...",
    "《人类简史》是一本很好的历史书...",
    "如果喜欢推理，可以试试《白夜行》...",
    "《活着》是余华的代表作，非常感人..."
]
# → 5个不同的回答 ✅

# Epoch 5 (DPO训练后)
>>> generate("推荐一本书", temperature=0.8, num_samples=5)
[
    "推荐《三体》，这是刘慈欣的科幻巨著...",
    "推荐《三体》，这是刘慈欣的科幻巨著...",
    "推荐《三体》，刘慈欣的科幻小说...",
    "推荐《三体》，这是刘慈欣的科幻巨著...",
    "我推荐《三体》，这是刘慈欣的..."
]
# → 都是三体！模式崩溃 ❌

# KL散度的作用：
KL(π_policy || π_ref) = E[log(π_policy(y)/π_ref(y))]

# 如果 π_policy 过度拟合某个回答：
# π_policy("三体...") = 0.9
# π_ref("三体...") = 0.01
# → KL = log(0.9/0.01) = 4.5  # 非常大！

# DPO loss 会惩罚这种过度集中
# → 一定程度上保持多样性

# 但为什么不完美？
# 因为KL约束是"全局"的，不是"针对alternative"的

# β = 0.01 (小β)：
# → 对齐效果好，但多样性丧失严重

# β = 1.0 (大β)：
# → 保持多样性好，但对齐效果差

# 实践中的选择：β ≈ 0.1 ~ 0.5
# → 在对齐和多样性之间平衡

# ❌ 差的偏好数据：
[
    {"prompt": "推荐书", "chosen": "推荐《三体》...", "rejected": "随便看看"},
    {"prompt": "推荐书", "chosen": "推荐《三体》...", "rejected": "不知道"},
]
# → 所有chosen都是三体

# ✅ 好的偏好数据：
[
    {"prompt": "推荐书", "chosen": "推荐《三体》...", "rejected": "随便看看"},
    {"prompt": "推荐书", "chosen": "推荐《百年孤独》...", "rejected": "不知道"},
    {"prompt": "推荐书", "chosen": "推荐《人类简史》...", "rejected": "我不确定"},
    {"prompt": "推荐科幻书", "chosen": "推荐《三体》...", "rejected": "随便"},
]
# → 多样的chosen，模型学会"根据情况选择"

def mixed_loss(policy_model, batch):
    # DPO损失（对齐）
    dpo_loss = compute_dpo_loss(policy_model, batch)

    # SFT损失（保持多样性）
    sft_loss = compute_sft_loss(policy_model, diverse_responses)

    # 组合
    total_loss = dpo_loss + 0.1 * sft_loss

    return total_loss

# DPO：只比较1对(chosen, rejected)
# → 容易过拟合chosen

# GRPO：比较一组responses
responses = [
    "推荐《三体》，科幻巨著...",     # reward: 9.0
    "推荐《百年孤独》，文学经典...", # reward: 8.5
    "推荐《人类简史》，历史佳作...", # reward: 8.0
    "随便看看",                       # reward: 2.0
]

# GRPO优势：根据相对质量调整，保持多样性
advantages = normalize(rewards)
loss = -Σ advantages[i] * log P(response[i])

# ========== 代码生成：DPO效果好 ✅ ==========
Question: "写一个快速排序"
# 正确的实现就那么几种，多样性不重要
# → DPO效果excellent!

# ========== 自然语言：DPO可能受限 ⚠️ ==========
Question: "描述春天"
# 有无数种优美的表达
# → DPO可能过度拟合一种表达

# ========== 解决方案 ==========
# 对于代码/JSON：aggressive DPO
beta_code = 0.1  # 小β，强对齐

# 对于创意写作：保守DPO或其他方法
beta_creative = 0.5  # 大β，保持多样性

# 数学题：答案是对还是错 → 完美的reward
def math_reward(answer, ground_truth):
    return 1.0 if answer == ground_truth else 0.0

# 编程题：代码能否通过测试用例 → 明确的reward
def code_reward(code, test_cases):
    passed = run_tests(code, test_cases)
    return len(passed) / len(test_cases)

传统SFT：只能学到人类的思考方式（天花板）
RL：可以探索出人类不会用的解题方法

RL框架促进了高级推理模式的自然涌现：
- 自我反思
- 验证
- 动态策略适应

# 决策树：
if 任务有明确的对错标准:
    if 资源充足:
        使用PPO/GRPO（可能效果略好）
    else:
        使用DPO（更简单稳定）
elif 任务需要主观评价:
    if 有大量偏好数据:
        使用DPO
    else:
        使用LLM-as-a-Judge + GRPO
else:
    继续SFT，慎用RL

# ✅ 高质量偏好数据的特征：
1. 多样性：chosen response要覆盖不同风格
2. 明确性：chosen和rejected的差异要明显
3. 一致性：相似问题的偏好要一致
4. 数量：足够的数据量（建议>10k对）

# ❌ 避免：
1. 所有chosen都是同一个答案
2. chosen和rejected差异不明显
3. 偏好标准不一致

# 1. 超参数设置
beta = 0.1  # DPO: 代码任务可以更小(0.05)，创意任务更大(0.5)
learning_rate = 2e-4  # LoRA可以更大，全参数更小(1e-6)

# 2. 监控指标
- KL散度：不要超过5.0
- Reward平均值：应该持续增加
- 生成多样性：定期采样检查

# 3. 早停策略
if validation_reward下降 or KL过大:
    停止训练，使用前一个checkpoint

# ❌ 错误1：Reference model也在更新
# ✅ 正确：Reference model保持固定

# ❌ 错误2：只用一对偏好数据训练
# ✅ 正确：使用多样化的大规模数据

# ❌ 错误3：β设置过小导致模式崩溃
# ✅ 正确：从较大的β开始，逐渐调小

# ❌ 错误4：忽略KL散度监控
# ✅ 正确：实时监控KL，设置阈值

# 当前代码的问题：
def compute_grpo_loss_simple(model, tokenizer, experience):
    outputs = model(**inputs, labels=inputs['input_ids'])
    loss = outputs.loss  # ⚠️ 这是SFT的loss
    weighted_loss = loss * (1.0 - reward / 10.0)  # ⚠️ 不够准确

# 改进建议：
def compute_grpo_loss_proper(model, tokenizer, experience):
    # 1. 手动计算log概率（不用labels）
    logits = model(**inputs).logits
    log_probs = compute_sequence_log_prob(logits, completion_tokens)

    # 2. 计算group相对advantage
    advantages = (rewards - rewards.mean()) / (rewards.std() + 1e-8)

    # 3. GRPO损失
    loss = -(log_probs * advantages).mean()

    return loss

# DPO实现：
# https://github.com/huggingface/trl

# GRPO实现：
# https://github.com/deepseek-ai/DeepSeek-Math

# OpenAI的RLHF论文：
# InstructGPT: Training language models to follow instructions

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import torch.nn.functional as F

def compute_dpo_loss(policy_model, reference_model, batch, beta=0.1):
    """完整的DPO损失计算"""

    # 准备输入
    prompt = batch['prompt']
    chosen = batch['chosen']
    rejected = batch['rejected']

    chosen_input = tokenizer(prompt + chosen, return_tensors="pt")
    rejected_input = tokenizer(prompt + rejected, return_tensors="pt")

    # Policy model log概率
    policy_chosen_logprob = compute_log_prob(
        policy_model, chosen_input, len(tokenizer(prompt)['input_ids'])
    )
    policy_rejected_logprob = compute_log_prob(
        policy_model, rejected_input, len(tokenizer(prompt)['input_ids'])
    )

    # Reference model log概率
    with torch.no_grad():
        ref_chosen_logprob = compute_log_prob(
            reference_model, chosen_input, len(tokenizer(prompt)['input_ids'])
        )
        ref_rejected_logprob = compute_log_prob(
            reference_model, rejected_input, len(tokenizer(prompt)['input_ids'])
        )

    # DPO损失
    policy_logratios = policy_chosen_logprob - policy_rejected_logprob
    ref_logratios = ref_chosen_logprob - ref_rejected_logprob

    logits = beta * (policy_logratios - ref_logratios)
    loss = -F.logsigmoid(logits).mean()

    return loss

def compute_log_prob(model, input_ids, prompt_length):
    """计算序列的log概率"""
    outputs = model(input_ids['input_ids'])
    logits = outputs.logits

    # 只关注response部分
    response_logits = logits[:, prompt_length-1:-1, :]
    response_labels = input_ids['input_ids'][:, prompt_length:]

    # 计算log概率
    log_probs = F.log_softmax(response_logits, dim=-1)
    token_log_probs = torch.gather(
        log_probs, dim=-1, index=response_labels.unsqueeze(-1)
    ).squeeze(-1)

    return token_log_probs.sum(dim=-1)

def compute_grpo_loss(model, experience, reference_model=None, beta=0.1):
    """GRPO损失计算"""

    prompts = experience['prompts']
    completions = experience['completions']
    rewards = experience['rewards']

    total_loss = 0.0

    for i, prompt in enumerate(prompts):
        log_probs = []

        # 计算每个completion的log概率
        for j, completion in enumerate(completions[i]):
            full_text = tokenizer.apply_chat_template(
                prompt + completion, tokenize=False
            )
            inputs = tokenizer(full_text, return_tensors="pt")

            log_prob = compute_log_prob(
                model, inputs, len(tokenizer(prompt)['input_ids'])
            )
            log_probs.append(log_prob)

        log_probs_tensor = torch.stack(log_probs)
        rewards_tensor = torch.tensor(rewards[i])

        # Group相对advantage
        advantages = (rewards_tensor - rewards_tensor.mean()) / \
                     (rewards_tensor.std() + 1e-8)

        # GRPO损失
        loss = -(log_probs_tensor * advantages).mean()

        # 可选：添加KL惩罚
        if reference_model is not None:
            with torch.no_grad():
                ref_log_probs = [compute_log_prob(
                    reference_model, inputs, prompt_len
                ) for inputs in all_inputs]
            kl = (log_probs_tensor - torch.stack(ref_log_probs)).mean()
            loss += beta * kl

        total_loss += loss

    return total_loss / len(prompts)