"""
=============================================================================
第 9 讲 · 代码 3：β 的作用；确定性的偏好与 IPO；离线 vs 在线（迭代）DPO
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
    演示 1：β 扫描：DPO 的 KL 和真实质量怎样随 β 变化？和"目标"π*(真实 BT 奖励) 比。
    演示 2：偏好是"确定性"的（标注没有噪声，例如多人投票取多数）时，DPO 的理论最优解会把 KL 推到无穷大，
            β 失去了作用（Azar 等 2023）。IPO 把目标从"概率比越大越好"换成"概率比等于一个有限的值"。
    演示 3：离线 DPO 只用 SFT 生成的数据；在线（迭代）DPO 每一轮用当前策略生成新回答、收集新偏好。
            同样多的标注，哪个更好？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 03_beta_offline_online.py        # 约 25 秒（CPU）

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

import numpy as np
import torch

from dpo_tools import (all_logprobs, exact_stats, labeler_utility, make_reference, make_reward_model,
                       optimal_logprobs, pad, preference_pairs, print_header, sample, train_preference,
                       true_bt_reward)


# =============================================================================
# 演示 1
# =============================================================================
def demo_beta_sweep(ref, ref_lp, data):
    print_header("演示 1：β 扫描 —— DPO 的 KL 与真实质量")
    chosen, rejected = data
    r_true = true_bt_reward()
    print("同样的 2000 对数据，每个 β 训练 20 遍。“目标”是 π*(真实 BT 奖励)（数据无限多时的极限）。\n")
    print(f"  {pad('β', 6)} {pad('DPO：KL', 9, 'right')} {pad('真实质量', 9, 'right')} {pad('平均长度', 9, 'right')} "
          f"{pad('华丽', 6, 'right')}   │ {pad('目标：KL', 9, 'right')} {pad('真实质量', 9, 'right')}")
    for beta in [5.0, 3.0, 1.0, 0.5, 0.3, 0.1, 0.05]:
        pol = train_preference("dpo", ref, chosen, rejected, beta=beta, epochs=20)
        s = exact_stats(all_logprobs(pol), ref_lp)
        t = exact_stats(optimal_logprobs(ref_lp, r_true, beta), ref_lp)
        print(f"  {beta:>6.2f} {s['kl']:>9.3f} {s['gold']:>+9.3f} {s['length']:>9.2f} {s['fancy']:>6.2f}   │ "
              f"{t['kl']:>9.3f} {t['gold']:>+9.3f}")
    print("""
观察：
  · β 大：DPO 离 SFT 很近，和目标几乎一样。β 减小：DPO 离 SFT 越来越远，但远不如目标走得远。
  · 真实质量在 β ≈ 0.5~1 附近最好，β 更小反而下降 —— 和 RLHF 的过度优化曲线形状相似，原因却不同：
    目标本身（π*）在 β 很小时真实质量很高（+2.4）；是**有限的离线数据**让 DPO 到不了那里，
    还在数据之外学到了“华丽越多越好”（代码 02 演示 2）。
  · 实践中 DPO 常用的 β 在 0.1 左右（Rafailov 等的实验用过 0.1 和 0.5）。和 RLHF 的 KL 系数一样，
    这个数字依赖于问题，没有通用的最佳值，要用留出数据或人工评估来选。""")


# =============================================================================
# 演示 2
# =============================================================================
def deterministic_pairs(ref, n, seed):
    """
    "确定性"的偏好：从 SFT 采两个回答，总是选标注员效用更高的那个（效用相同的对丢掉）。

    【偏好学习联系】现实中，如果每一对由很多人投票取多数，或者用一个很确定的评委模型打分，
        偏好就接近确定性（P(y_w ≻ y_l) ≈ 1）。
    """
    gen = torch.Generator().manual_seed(seed)
    a, b = sample(ref, n, gen), sample(ref, n, gen)
    chosen, rejected = [], []
    for x, y in zip(a, b):
        ux, uy = labeler_utility(x), labeler_utility(y)
        if ux > uy:
            chosen.append(x), rejected.append(y)
        elif uy > ux:
            chosen.append(y), rejected.append(x)
    return chosen, rejected


def demo_deterministic(ref, ref_lp, data):
    print_header("演示 2：确定性的偏好 —— DPO 的 KL 停不下来；IPO 停得下来")
    det = deterministic_pairs(ref, 2000, seed=1)
    print(f"确定性的偏好：{len(det[0])} 对（总是选效用更高的那个）。对照：演示 1 的带噪声的 2000 对。")
    print("报告训练不同遍数后，策略离 SFT 的 KL（以及真实质量）：\n")
    checkpoints = [5, 20, 60, 120]
    runs = [("DPO β=1，带噪声", "dpo", 1.0, data), ("DPO β=1，确定性", "dpo", 1.0, det),
            ("IPO τ=0.1，确定性", "ipo", 0.1, det)]
    print(f"  {pad('', 20)}" + "".join(pad(f"{c} 遍", 17, "right") for c in checkpoints))
    for name, kind, beta, (ch, rj) in runs:
        rec = {}

        def cb(ep, pol, rec=rec):
            if ep in checkpoints:
                s = exact_stats(all_logprobs(pol), ref_lp)
                rec[ep] = f"{s['kl']:.2f} / {s['gold']:+.2f}"
        train_preference(kind, ref, ch, rj, beta=beta, epochs=checkpoints[-1], callback=cb)
        print(f"  {pad(name, 20)}" + "".join(pad(rec[c], 17, "right") for c in checkpoints))
    print("""
（每格是“KL / 真实质量”。）

观察：
  · 带噪声的偏好：DPO 的 KL 增长得越来越慢 —— 有噪声时，BT 似然的最优点是有限的。
  · 确定性的偏好：DPO 的 KL 一直在涨。原因（Azar 等 2023）：如果数据里 y_w 总是赢 y_l，BT 似然要求
    P(y_w ≻ y_l) = σ(β(h_w − h_l)) → 1，也就是 h_w − h_l → ∞ —— 无论 β 多大，最优解都在无穷远处，
    β 的“正则化”作用消失了。真实的偏好数据集里，大部分回答对只被标注过一次，从数据的角度看就是“确定性”的。
  · IPO 把目标换成“h_w − h_l 等于一个有限的值 1/(2τ)”（平方损失），KL 很快就稳定下来。
  · 注意：IPO 的 τ 和 DPO 的 β 不是同一个尺度（τ = 1 时目标间隔只有 0.5，几乎不动），要分别调。
  · 这里的确定性偏好带来了更好的真实质量 —— 因为它去掉了标注噪声。“KL 停不下来”本身不一定有害，
    有害的是它不再受 β 控制：模型会离 SFT 多远，取决于训练多久。""")


# =============================================================================
# 演示 3
# =============================================================================
def demo_online(ref, ref_lp, data, beta=0.3, rounds=4, per_round=1000):
    print_header(f"演示 3：离线 vs 在线（迭代）DPO（β = {beta}）")
    chosen, rejected = data
    print(f"在线：第 0 轮用演示 1 的 2000 对（SFT 生成）训练；之后每一轮从**当前策略**采 {per_round} 对新回答、")
    print("让众包标注员比较，把新数据加进去，从 π_ref 重新训练（参考策略始终是 SFT）。")
    print("离线：同样总数的偏好对，全部从 SFT 采样。\n")
    print(f"  {pad('轮次', 5)} {pad('偏好对总数', 11, 'right')}   {pad('在线：KL', 9, 'right')} {pad('真实质量', 9, 'right')} "
          f"{pad('华丽', 6, 'right')}   │ {pad('离线：KL', 9, 'right')} {pad('真实质量', 9, 'right')} {pad('华丽', 6, 'right')}")
    all_c, all_r = list(chosen), list(rejected)
    off_c, off_r = list(chosen), list(rejected)
    policy = ref
    for rnd in range(rounds):
        if rnd > 0:
            gen = torch.Generator().manual_seed(100 + rnd)
            c, r = preference_pairs(policy, per_round, gen, np.random.default_rng(100 + rnd))
            all_c += c
            all_r += r
            gen = torch.Generator().manual_seed(200 + rnd)
            c, r = preference_pairs(ref, per_round, gen, np.random.default_rng(200 + rnd))
            off_c += c
            off_r += r
        policy = train_preference("dpo", ref, all_c, all_r, beta=beta, epochs=20, seed=rnd)
        offline = train_preference("dpo", ref, off_c, off_r, beta=beta, epochs=20, seed=rnd)
        s_on = exact_stats(all_logprobs(policy), ref_lp)
        s_off = exact_stats(all_logprobs(offline), ref_lp)
        print(f"  {rnd:>5d} {len(all_c):>11d}   {s_on['kl']:>9.3f} {s_on['gold']:>+9.3f} {s_on['fancy']:>6.2f}   │ "
              f"{s_off['kl']:>9.3f} {s_off['gold']:>+9.3f} {s_off['fancy']:>6.2f}")
    t = exact_stats(optimal_logprobs(ref_lp, true_bt_reward(), beta), ref_lp)
    print(f"""
  目标 π*(真实 BT 奖励)：KL {t['kl']:.3f}，真实质量 {t['gold']:+.3f}。

观察：
  · 在线 DPO 每一轮都在**策略自己会生成的回答**上收集偏好：上一轮学歪了（例如堆华丽）的地方，
    这一轮就会有数据告诉它“六个华丽并不好”。几轮之后，真实质量接近目标。
  · 离线 DPO 增加同样多的数据，但数据都来自 SFT，覆盖不到策略要去的地方，改进有限。
  · 这和第 8 讲 RLHF 的“迭代地更新奖励模型”是同一个思想：**让数据跟上策略**。代价是每一轮都要标注新数据
    （或者用一个评委模型代替人，即“在线 AI 反馈”）。Llama 3 等模型的后训练就是多轮迭代的。
  · 在线 DPO 和 RLHF 的界限因此变得模糊：两者都是“生成 → 评价 → 更新”的循环，区别只在“评价”是一个
    显式的奖励模型还是新的偏好标注、“更新”用 RL 还是 DPO 损失。""")


def main():
    ref = make_reference()
    ref_lp = all_logprobs(ref)
    _, _, data = make_reward_model(ref)
    demo_beta_sweep(ref, ref_lp, data)
    demo_deterministic(ref, ref_lp, data)
    demo_online(ref, ref_lp, data)


if __name__ == "__main__":
    main()