"""
=============================================================================
第 9 讲 · 代码 2：同样的偏好数据，DPO vs RLHF；DPO 把概率放到了哪里；似然位移
=============================================================================

对应讲义：lectures/lecture09_dpo.md 第 3、4 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
用第 8 讲代码 02 的同一批 2000 对偏好数据：
    · RLHF：先训练奖励模型（线性 / 灵活），再求 KL 约束的最优策略（第 8 讲：RL 收敛到解析解，这里直接用解析解）；
    · DPO：直接在偏好对上训练策略。

    演示 1：不同 β 下，两条路线的真实质量和 KL。
    演示 2：β 小的时候，DPO 把概率放到了哪里？—— 放到训练数据里没出现过的回答上的，是堆满"华丽"的回答。
    演示 3：训练过程：被选中回答的概率反而下降（似然位移）；训练越久 KL 越大。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_dpo_vs_rlhf.py        # 约 20 秒（CPU）

需要 numpy 和 PyTorch。依赖本目录的 dpo_tools.py。
"""

import math

import numpy as np
import torch

from dpo_tools import (ALL, ALL_GOLD, INDEX, all_logprobs, encode, exact_stats, make_reference, make_reward_model,
                       optimal_logprobs, pad, preference_pairs, print_header, sequence_logprob, show,
                       train_preference, true_bt_reward)

EPOCHS = 20


# =============================================================================
# 演示 1
# =============================================================================
def demo_compare(ref, ref_lp, data, r_lin, r_rich):
    print_header("演示 1：同样的 2000 对偏好数据 —— RLHF（两种奖励模型）vs DPO")
    chosen, rejected = data
    r_true = true_bt_reward()
    print(f"RLHF 用解析最优解（第 8 讲代码 03 验证过 RL 会收敛到它）；DPO 训练 {EPOCHS} 遍。")
    print("“目标”一列是 π*(真实 BT 奖励)：数据无限多时 DPO（和足够灵活的 RLHF）会到达的地方。\n")
    print(f"  {pad('β', 5)} {pad('', 14)}" + "".join(pad(n, 13, "right") for n in ["RLHF 线性", "RLHF 灵活", "DPO", "目标"]))
    for beta in [3.0, 1.0, 0.3, 0.1]:
        stats = [exact_stats(optimal_logprobs(ref_lp, r, beta), ref_lp) for r in (r_lin, r_rich)]
        pol = train_preference("dpo", ref, chosen, rejected, beta=beta, epochs=EPOCHS)
        stats.append(exact_stats(all_logprobs(pol), ref_lp))
        stats.append(exact_stats(optimal_logprobs(ref_lp, r_true, beta), ref_lp))
        print(f"  {beta:>5.1f} {pad('真实质量', 14)}" + "".join(f"{s['gold']:>+13.3f}" for s in stats))
        print(f"  {'':>5s} {pad('KL(π ‖ π_ref)', 14)}" + "".join(f"{s['kl']:>13.3f}" for s in stats))
    print("""
观察：
  · β = 3、1：DPO 的真实质量明显高于“RLHF 线性”，和“RLHF 灵活”相近。DPO 的隐式奖励由策略网络表示，
    比线性奖励模型灵活，没有第 8 讲那种“要点越多越好”的设定错误。
  · β 小（0.3、0.1）：RLHF 灵活继续变好（接近它的极限 +2.4），DPO 却停在 +1.3~+1.6，离“目标”很远。
    KL 也远小于目标：2000 对数据、20 遍训练，DPO 没有到达它“应该”到达的 π*。演示 2 看它到底去了哪里。
  · 两条路线的取舍：RLHF 要训练奖励模型、跑 RL（贵、不稳定），但奖励模型可以给**策略新生成的回答**打分；
    DPO 简单、稳定、便宜（只是一个分类损失），但它只从**固定的数据**学。""")


# =============================================================================
# 演示 2
# =============================================================================
def demo_where_mass(ref, ref_lp, data, r_rich, beta=0.1):
    print_header(f"演示 2：β = {beta} 时，DPO 把概率放到了哪里？")
    chosen, rejected = data
    seen = np.zeros(len(ALL), dtype=bool)
    for s in chosen + rejected:
        seen[INDEX[s]] = True
    pol = train_preference("dpo", ref, chosen, rejected, beta=beta, epochs=EPOCHS)
    lp = all_logprobs(pol)
    p = np.exp(lp)
    p_rich = np.exp(optimal_logprobs(ref_lp, r_rich, beta))
    p_ref = np.exp(ref_lp)
    print(f"训练数据里出现过的不同回答：{int(seen.sum())} 个（全部 {len(ALL)} 个回答的 {seen.mean():.1%}）。\n")
    print(f"  {pad('策略', 14)} {pad('在“出现过的”上的概率', 21, 'right')} {pad('出现过的部分的真实质量', 23, 'right')} "
          f"{pad('没出现过的部分的真实质量', 25, 'right')} {pad('总的真实质量', 13, 'right')}")
    for name, q in [("SFT（π_ref）", p_ref), ("DPO", p), ("RLHF 灵活", p_rich)]:
        g_seen = float(q[seen] @ ALL_GOLD[seen] / q[seen].sum())
        g_unseen = float(q[~seen] @ ALL_GOLD[~seen] / q[~seen].sum())
        print(f"  {pad(name, 14)} {q[seen].sum():>21.3f} {g_seen:>+23.3f} {g_unseen:>+25.3f} {q @ ALL_GOLD:>+13.3f}")
    print("\n  DPO 放在“没出现过的回答”上概率最高的 6 个：")
    order = np.argsort(-np.where(seen, -1.0, p), kind="stable")[:6]
    for k in order:
        print(f"    {pad(show(ALL[k]), 34)} 概率 {p[k]:.3f}   真实质量 {ALL_GOLD[k]:+.2f}")
    best_rich = int(np.argmax(p_rich))
    print(f"\n  对照：RLHF 灵活概率最高的回答是“{show(ALL[best_rich])}”（{'出现过' if seen[best_rich] else '没出现过'}，"
          f"真实质量 {ALL_GOLD[best_rich]:+.2f}）。")
    print("""
观察：
  · 两种方法都把大部分概率放在了训练数据里**没出现过**的回答上 —— 这本身不是问题：最好的回答
    不一定恰好被采样到过。RLHF 灵活放到那里的，是真实质量很高的回答。
  · 问题在于 DPO 放到“没出现过”的回答上的是什么：真实质量明显更低，其中有堆满“华丽”的回答
    （例如六个华丽，真实质量为负）。
  · 为什么？DPO 只在“被选中 vs 被拒绝”的那些对上约束了 β·log(π/π_ref) 的**差值**。数据里华丽多的回答
    常常被选中（标注偏差），策略网络把“华丽越多越好”推广到了数据之外 —— 没有任何数据告诉它
    “六个华丽”不好。β 越小，同样的偏好差值需要越大的 log 概率比，策略被推得越远。
  · RLHF 的奖励模型也会外推（第 8 讲），但灵活的奖励模型是一个单独的、在所有回答上都有定义的函数，
    RL 阶段在策略**新生成**的回答上用它打分；DPO 的“隐式奖励”是策略网络本身，它在数据之外怎样推广，
    取决于策略网络的结构。Xu 等（2024）在真实的大模型上也发现：DPO 可能给分布外的回答分配很高的概率，
    仔细调过的 PPO 往往更好。
  · 解决办法之一是**在线 / 迭代 DPO**：用当前策略生成新回答、收集新偏好，让数据跟上策略（代码 03 演示 3）。""")


def demo_dynamics(ref, ref_lp, data, beta=1.0):
    print_header("演示 3：训练过程 —— 似然位移，以及训练越久 KL 越大")
    chosen, rejected = data
    enc_c, enc_r = encode(chosen), encode(rejected)
    gen = torch.Generator().manual_seed(42)
    val_c, val_r = preference_pairs(ref, 2000, gen, np.random.default_rng(42))
    with torch.no_grad():
        ref_c = sequence_logprob(ref, enc_c).mean().item()
        ref_r = sequence_logprob(ref, enc_r).mean().item()
    print(f"β = {beta}。训练数据上的平均 log π：被选中的 {ref_c:.3f}，被拒绝的 {ref_r:.3f}（训练前，即 π_ref）。\n")
    print(f"  {pad('遍数', 5)} {pad('log π(被选中)', 14, 'right')} {pad('log π(被拒绝)', 14, 'right')} "
          f"{pad('训练集准确率', 13, 'right')} {pad('留出准确率', 11, 'right')} {pad('KL', 7, 'right')} {pad('真实质量', 9, 'right')}")

    def cb(ep, pol):
        if ep in (1, 2, 5, 10, 20, 40, 80):
            lp = all_logprobs(pol)
            r_hat = lp - ref_lp
            with torch.no_grad():
                lc = sequence_logprob(pol, enc_c).mean().item()
                lr_ = sequence_logprob(pol, enc_r).mean().item()
            tr = float(np.mean(r_hat[[INDEX[s] for s in chosen]] > r_hat[[INDEX[s] for s in rejected]]))
            va = float(np.mean(r_hat[[INDEX[s] for s in val_c]] > r_hat[[INDEX[s] for s in val_r]]))
            st = exact_stats(lp, ref_lp)
            print(f"  {ep:>5d} {lc:>14.3f} {lr_:>14.3f} {tr:>13.3f} {va:>11.3f} {st['kl']:>7.3f} {st['gold']:>+9.3f}")
    train_preference("dpo", ref, chosen, rejected, beta=beta, epochs=80, callback=cb)
    print("""
观察：
  · **似然位移**：训练中“被拒绝回答”的概率大幅下降（这是想要的），可“被选中回答”的概率**也在下降**。
    DPO 只要求两者的**差**变大；两个都下降、但被拒绝的降得更多，也能做到。被挪走的概率去了别的回答 ——
    包括数据里没有的回答（演示 2）。Razin 等（2024）把这叫“似然位移”，它可能把概率挪到意想不到的地方。
    一些变体（例如在损失里加上被选中回答的 SFT 损失）试图阻止被选中回答的概率下降。
  · 训练集准确率一直上升，留出准确率很快就不怎么变了 —— 后面的训练主要在“背”训练数据。
  · KL 随训练遍数一直增大：β 只出现在损失里，并不直接约束训练过程中的 KL（对比 RLHF 的 KL 惩罚）。
    所以 DPO 的“训练多久”也是一个需要用留出数据 / 人工评估来选的超参数。""")


def main():
    ref = make_reference()
    ref_lp = all_logprobs(ref)
    _, r_lin, data = make_reward_model(ref)
    _, r_rich, _ = make_reward_model(ref, kind="rich")
    demo_compare(ref, ref_lp, data, r_lin, r_rich)
    demo_where_mass(ref, ref_lp, data, r_rich)
    demo_dynamics(ref, ref_lp, data)


if __name__ == "__main__":
    main()