"""
=============================================================================
第 8 讲 · 代码 4：奖励过度优化 —— 代理奖励一路上升，真实质量先升后降
=============================================================================

对应讲义：lectures/lecture08_rlhf.md 第 5、6 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
Gao、Schulman、Hilton（2023）用一个"金标准奖励模型"代替人类，系统地测量了：策略离 SFT 越远（KL 越大），
代理奖励（训练用的奖励模型）一路上升，金标准奖励却先升后降。这个文件在玩具世界里**精确地**画出这条曲线：

    演示 1：KL 约束问题的解析最优解 π_β ∝ π_ref·exp(r/β)，β 从 10 扫到 0.02：KL、代理奖励、真实质量。
    演示 2：Best-of-n（从 SFT 采 n 个回答，挑奖励最高的）：精确的分布、KL 和真实质量；
            KL 公式 log n − (n−1)/n 什么时候准。
    演示 3：RL 训练的动态过程：KL 系数很小时，训练越久越糟 —— "提前停止"和"KL 惩罚"是同一个问题的两种答案。
    演示 4：四种奖励模型（线性 / 灵活 × 众包 / 专家）的过度优化曲线：两种误差来源各自造成多大的损失。
    演示 5：Gao 等的函数形式：真实质量 ≈ d(α − β·d)（d = √KL）拟合得怎样。

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

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

import math

import numpy as np

from toy_lm import (ALL, ALL_FEAT, ALL_GOLD, ALL_LEN, all_logprobs, clone_policy, exact_stats, make_reference,
                    make_reward_model, optimal_logprobs, pad, print_header, rl_train, show)

BETAS = [10.0, 5.0, 3.0, 2.0, 1.5, 1.0, 0.7, 0.5, 0.3, 0.2, 0.1, 0.05, 0.02]


def best_of_n(ref_lp, reward, n):
    """
    Best-of-n 的**精确**分布：从 π_ref 独立采 n 个回答，选奖励最高的一个（并列时随机选一个）。

    【思路】把回答按奖励分组（奖励相同的一组，例如只是词序不同），组按奖励从低到高排。
        记 q_g = 第 g 组在 π_ref 下的总概率，F_g = 前 g 组的累计概率。
        "n 个里最高的落在第 g 组"的概率 = F_g^n − F_{g−1}^n（n 个都 ≤ 第 g 组，减去 n 个都 < 第 g 组）。
        组内按 π_ref 的比例分配。
    【返回】(每个回答的概率数组, KL(BoN ‖ π_ref))
    【边界】KL 只需要按组算：组内的分布和 π_ref 成比例，组内那部分 KL 为 0。
    """
    key = np.round(reward, 9)                                   # 消除浮点误差造成的"假不相等"
    _, inv = np.unique(key, return_inverse=True)
    pr = np.exp(ref_lp)
    q = np.bincount(inv, weights=pr)
    F = np.cumsum(q)
    P = F ** n - (F - q) ** n
    p_y = P[inv] * pr / q[inv]
    m = P > 0
    return p_y, float(np.sum(P[m] * np.log(P[m] / q[m])))


def analytic_curve(ref_lp, reward, betas=BETAS):
    """对每个 β 计算解析最优解的统计量，返回 [(β, stats), ...]。"""
    return [(b, exact_stats(optimal_logprobs(ref_lp, reward, b), ref_lp, reward)) for b in betas]


# =============================================================================
# 演示 1
# =============================================================================
def demo_analytic(ref_lp, reward):
    print_header("演示 1：解析最优解 π_β ∝ π_ref·exp(r/β) —— β 从大到小")
    print("奖励模型：线性 + 众包（代码 02）。每一行都是精确值，不是采样估计。\n")
    print(f"  {pad('β', 6)} {pad('KL', 7, 'right')} {pad('√KL', 6, 'right')} {pad('代理奖励', 9, 'right')} "
          f"{pad('真实质量', 9, 'right')} {pad('平均长度', 9, 'right')} {pad('华丽', 6, 'right')} {pad('要点', 6, 'right')}")
    s0 = exact_stats(ref_lp, ref_lp, reward)
    print(f"  {pad('SFT', 6)} {0.0:>7.3f} {0.0:>6.2f} {s0['proxy']:>+9.3f} {s0['gold']:>+9.3f} {s0['length']:>9.2f} "
          f"{s0['fancy']:>6.2f} {s0['points']:>6.2f}")
    curve = analytic_curve(ref_lp, reward)
    best_beta, best = max(curve, key=lambda c: c[1]["gold"])
    for b, s in curve:
        mark = "  ← 真实质量最高" if b == best_beta else ""
        print(f"  {b:>6.2f} {s['kl']:>7.3f} {math.sqrt(s['kl']):>6.2f} {s['proxy']:>+9.3f} {s['gold']:>+9.3f} "
              f"{s['length']:>9.2f} {s['fancy']:>6.2f} {s['points']:>6.2f}{mark}")
    best_seq = int(np.argmax(reward))
    print(f"""
  β → 0 的极限：策略集中到奖励模型最喜欢的回答“{show(ALL[best_seq])}”（真实质量 {ALL_GOLD[best_seq]:+.2f}）。

观察：
  · 代理奖励随 KL 单调上升（这是 π_β 的定义保证的：在给定 KL 下，它是代理奖励最高的分布）。
  · 真实质量先升后降：在 β ≈ {best_beta} 附近（KL 约 {best['kl']:.1f}）最高，之后一路下跌，最后变成负数 ——
    比什么都不做（SFT）还差得多。
  · 下跌的过程看得很清楚：要点越来越多（奖励模型以为“越多越好”）、回答越来越长、华丽一度增多；
    最后连华丽也被挤掉，只剩下奖励模型给分最高的“要点”。
  · **KL 约束的作用**：β 不是一个“技术细节”，它决定了策略在多大程度上相信奖励模型。β 太大，学不到东西；
    β 太小，奖励模型的错误被放大。InstructGPT 用的 KL 系数是 0.02（他们的奖励尺度和 KL 单位和这里不同，
    数字不能直接比较）。""")
    return curve


# =============================================================================
# 演示 2
# =============================================================================
def demo_best_of_n(ref_lp, reward, curve):
    print_header("演示 2：Best-of-n —— 不训练策略，只在采样时挑")
    print(f"  {pad('n', 6)} {pad('KL（精确）', 11, 'right')} {pad('log n − (n−1)/n', 16, 'right')} {pad('代理奖励', 9, 'right')} "
          f"{pad('真实质量', 9, 'right')} {pad('平均长度', 9, 'right')}")
    rows = []
    for n in [1, 2, 4, 8, 16, 32, 64, 128, 256, 1024, 4096, 16384]:
        p, k = best_of_n(ref_lp, reward, n)
        formula = math.log(n) - (n - 1) / n
        rows.append((n, k, float(p @ ALL_GOLD)))
        print(f"  {n:>6d} {k:>11.3f} {formula:>16.3f} {p @ reward:>+9.3f} {p @ ALL_GOLD:>+9.3f} {p @ ALL_LEN:>9.2f}")
    n_best, k_best, g_best = max(rows, key=lambda r: r[2])
    _, s_best = max(curve, key=lambda c: c[1]["gold"])
    print(f"""
  Best-of-n 真实质量最高：n = {n_best}，KL {k_best:.2f}，真实质量 {g_best:+.3f}。
  对照演示 1：解析最优的 RL 策略真实质量最高 {s_best['gold']:+.3f}（KL {s_best['kl']:.2f}）。

观察：
  · Best-of-n 也会过度优化：n 越大，越会挑中奖励模型高估的回答，真实质量先升后降。
  · KL(BoN ‖ π_ref) 的公式 log n − (n−1)/n 在 n 不大时很准；n 很大时，精确值比公式小 ——
    公式假设“奖励没有并列”，而这里很多回答奖励相同（只是词序不同）。一般情况下这个公式是上界（Beirami 等 2024）。
  · 在这个玩具世界里，Best-of-n 的真实质量峰值比 RL 的解析最优解还高一些。解析最优解在每个 KL 下把**代理奖励**推到最高
    （它对代理奖励是“最有效率”的），也就最有效率地利用了奖励模型的错误；Best-of-n 只按名次挑，不那么“激进”。
    Gao 等（2023）在真实的奖励模型上也发现两种方法的过度优化曲线形状不同。
  · Best-of-n 简单、不需要训练，但每次生成要采 n 个回答、算 n 次奖励，推理成本高；
    它常被用作 RLHF 的对照，也被用来生成“拒绝采样微调”的数据（Llama 2）。""")
    return rows


# =============================================================================
# 演示 3
# =============================================================================
def demo_dynamics(ref, ref_lp, reward, beta=0.05, steps=500):
    print_header(f"演示 3：RL 训练的动态过程 —— KL 系数很小（β = {beta}）时，训练越久越糟")
    print(f"RLOO，每步 128 个回答，学习率 2e-4（比代码 03 小，为了看清过程）。每 25 步精确地评估一次：\n")
    print(f"  {pad('步数', 6)} {pad('KL', 7, 'right')} {pad('代理奖励', 9, 'right')} {pad('真实质量', 9, 'right')} "
          f"{pad('平均长度', 9, 'right')}")
    rec = []

    def cb(step, pol):
        s = exact_stats(all_logprobs(pol), ref_lp, reward)
        rec.append((step, s))
        print(f"  {step:>6d} {s['kl']:>7.3f} {s['proxy']:>+9.3f} {s['gold']:>+9.3f} {s['length']:>9.2f}")
    rl_train(clone_policy(ref), ref, reward, beta, steps, lr=2e-4, method="rloo", seed=5, log_every=25, callback=cb)
    peak_step, peak = max(rec, key=lambda r: r[1]["gold"])
    print(f"""
  真实质量在第 {peak_step} 步最高（{peak['gold']:+.3f}，KL {peak['kl']:.2f}），之后下降。

观察：
  · KL 惩罚很弱时，训练的每一步都在让代理奖励更高、离 SFT 更远；真实质量在某个时刻达到最高，然后下滑。
  · 两种控制过度优化的办法：(1) 用较大的 β，让最优解本身就停在“安全”的地方（演示 1）；
    (2) 用较小的 β，但在真实质量最高的时候**提前停止**。两者都需要一个“真实质量”的信号来做选择 ——
    现实中就是定期的人工评估（或一个独立的、更可信的评估模型）。
  · 两种办法在这里并不完全等价：第 100 步提前停止的策略（KL 0.89，真实质量 +1.15），不如同样 KL 下的
    解析最优解（演示 1，β = 1：KL 0.92，真实质量 +1.34）—— 训练轨迹上的中间策略，并不是那个 KL 下的最优分布。
    Gao 等在他们的设定里发现 KL 惩罚的效果大致相当于提前停止；不同设定下结论可能不同。
  · 只看训练用的奖励模型，永远发现不了这个问题：它一直在涨。这就是为什么 RLHF 的论文都报告人工评估的胜率，
    而不只是奖励模型的分数。""")


# =============================================================================
# 演示 4
# =============================================================================
def demo_four_rms(ref, ref_lp):
    print_header("演示 4：四种奖励模型的过度优化曲线")
    print("每种奖励模型用 2000 对数据训练；对每个 β 求解析最优解，报告**真实质量**：\n")
    betas = [3.0, 1.0, 0.3, 0.1, 0.02]
    print(f"  {pad('奖励模型', 14)}" + "".join(pad(f"β={b}", 9, "right") for b in betas)
          + pad("最高", 8, "right") + pad("在 β =", 8, "right"))
    for kind, kname in [("linear", "线性"), ("rich", "灵活")]:
        for lab, lname in [("crowd", "众包"), ("expert", "专家")]:
            _, r, _ = make_reward_model(ref, kind=kind, labeler=lab)
            curve = analytic_curve(ref_lp, r)
            best_b, best_s = max(curve, key=lambda c: c[1]["gold"])
            golds = {b: s["gold"] for b, s in curve}
            print(f"  {pad(kname + ' + ' + lname, 14)}" + "".join(f"{golds[b]:>+9.3f}" for b in betas)
                  + f"{best_s['gold']:>+8.3f}{best_b:>8.2f}")
    p_ref = np.exp(ref_lp)
    print(f"""
  参照：SFT 策略的真实质量 {p_ref @ ALL_GOLD:+.3f}；真实质量的最大可能值 +2.700。

观察：
  · 灵活 + 专家（没有设定错误、没有标注偏差）：β 越小越好，最后到达真实最优 +2.7 —— 没有过度优化。
    过度优化不是 RL 的“本性”，而是奖励模型的错误被优化放大。
  · 灵活 + 众包（只有标注偏差）：真实质量接近最高后稳定在 +2.4 —— 偏差让它“填满华丽”，损失 0.3，但不会崩溃。
  · 线性 + 专家（只有设定错误）：先升后降，β 很小时跌到负数 —— 外推的错误是无界的（“要点越多越好”）。
  · 线性 + 众包（两种都有）：峰值和“线性 + 专家”差不多，但 β 更小时反而跌得慢一些：众包标注员对华丽的偏爱，
    让策略把一部分位置填上华丽（每个只扣 0.1），而不是更多的要点（第三个起每个扣 0.6）—— 两种错误碰巧部分抵消。
    错误之间会相互作用，不能简单地相加。
  · 启示：(1) 奖励模型的容量和训练数据的覆盖范围决定了“能安全地走多远”；
    (2) 标注偏差决定了“最终会偏到哪里”；(3) KL 约束只是一个保险，不能修正奖励模型本身的错误。""")


# =============================================================================
# 演示 5
# =============================================================================
def demo_functional_form(ref_lp, reward, bon_rows, curve):
    print_header("演示 5：Gao 等的函数形式")
    g0 = float(np.exp(ref_lp) @ ALL_GOLD)

    def fit(points, form):
        """对 (KL, 真实质量) 点拟合 真实质量 − 起点 = d(α − β·f(d))，d = √KL；返回 (α, β, R²)。"""
        d = np.sqrt(np.array([max(k, 0.0) for k, _ in points]))
        y = np.array([g for _, g in points]) - g0
        keep = d > 1e-6
        d, y = d[keep], y[keep]
        X = np.stack([d, -d * d if form == "bon" else -d * np.log(d)], axis=1)
        coef, *_ = np.linalg.lstsq(X, y, rcond=None)
        r2 = 1 - np.sum((y - X @ coef) ** 2) / np.sum((y - y.mean()) ** 2)
        return coef[0], coef[1], r2

    bon_pts = [(k, g) for _, k, g in bon_rows]
    rl_pts = [(s["kl"], s["gold"]) for _, s in curve]
    print("Gao、Schulman、Hilton（2023）发现，以 d = √KL 为横轴：")
    print("  Best-of-n：真实质量 ≈ d(α − β·d)；RL：真实质量 ≈ d(α − β·log d)。在这里拟合（R² 越接近 1 越好）：\n")
    print(f"  {pad('曲线', 22)} {pad('形式 d(α − β·d)', 20, 'right')} {pad('形式 d(α − β·log d)', 22, 'right')}")
    for name, pts in [("Best-of-n", bon_pts), ("解析最优（RL 的极限）", rl_pts)]:
        a1, b1, r1 = fit(pts, "bon")
        a2, b2, r2 = fit(pts, "rl")
        print(f"  {pad(name, 22)} {pad(f'R² = {r1:.3f}', 20, 'right')} {pad(f'R² = {r2:.3f}', 22, 'right')}")
    print("""
观察：
  · Best-of-n 的曲线很符合 d(α − β·d)（R² 约 0.99），和 Gao 等的发现一致。
  · 解析最优解的曲线在这里也更符合 d(α − β·d)，而不是论文里 RL 的 d(α − β·log d)。
    可能的原因：论文里的 RL 是真实的 PPO 训练轨迹（每个点是训练中途的检查点，不是某个 β 的最优解），
    奖励模型是神经网络；这里是一个 19531 个回答的小世界和线性奖励模型。函数形式是经验规律，不是定理。
  · 论文更重要的发现是这些曲线的参数怎样随奖励模型的大小、数据量变化：奖励模型越大、数据越多，
    峰值越高、出现得越晚 —— 但峰值总是存在（演示 4 的“线性 + 众包”说明了为什么：偏差和外推错误不会因为
    数据多就消失）。""")


def main():
    ref = make_reference()
    ref_lp = all_logprobs(ref)
    _, reward, _ = make_reward_model(ref, kind="linear", labeler="crowd")
    curve = demo_analytic(ref_lp, reward)
    bon_rows = demo_best_of_n(ref_lp, reward, curve)
    demo_dynamics(ref, ref_lp, reward)
    demo_four_rms(ref, ref_lp)
    demo_functional_form(ref_lp, reward, bon_rows, curve)


if __name__ == "__main__":
    main()