"""
=============================================================================
第 10 讲 · 代码 3：在一个玩具策略上看清 RLHF、DPO 和过度优化
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
语言模型的"回答"有天文数字那么多种，很难直接看清 RLHF 和 DPO 在做什么。
这里把问题缩小到极致：只有**一个提示**、只有**几个可能的回答**，策略就是这几个回答上的一个概率分布。
在这个玩具世界里，一切都能精确计算：

    演示 1：RLHF 目标的最优解有闭式：π*(y) ∝ π_ref(y)·exp(r(y)/β) —— β 控制"离参考模型多远"
    演示 2：用策略梯度（REINFORCE）优化 RLHF 目标，会收敛到 π*
    演示 3：DPO 不用奖励模型、不用强化学习，只用偏好数据，也收敛到 π*
    演示 4：DPO 的隐式奖励 β·log(π/π_ref) 就是真实奖励（差一个常数）
    演示 5：DPO 的一个陷阱 —— 偏好"毫无噪声"时，β 管不住它
    演示 6：过度优化（古德哈特定律）—— 用一个不完美的奖励模型做 best-of-n，选得越狠，真实奖励反而越低

只用 numpy 和 PyTorch，不依赖前面任何讲次的代码，约 30 秒~1 分钟。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 03_rlhf_vs_dpo_toy.py
"""

import os

os.environ.setdefault("OMP_NUM_THREADS", "1")
import math

import numpy as np
import torch
import torch.nn.functional as F

torch.set_num_threads(1)

NAMES = ["正面、简洁", "正面、啰嗦", "中性", "有夸有骂", "负面", "跑题"]
R_TRUE = np.array([3.0, 2.0, 0.0, 0.5, -2.0, -1.0])        # 真实奖励 r*（人类心里的分数）
PI_REF = np.array([0.10, 0.15, 0.30, 0.25, 0.15, 0.05])    # 参考模型（SFT 模型）给每个回答的概率


def print_header(title):
    """打印一个带分隔线的标题。"""
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)


def optimal_policy(r, pi_ref, beta):
    """
    KL 约束下的最优策略：π*(y) = π_ref(y)·exp(r(y)/β) / Z。

    【思路】（讲义第 4.2、5.1 节）RLHF 的目标是
        max_π  E_{y~π}[r(y)] − β·KL(π ‖ π_ref)
        第一项：奖励越高越好；第二项：不要离参考模型太远。
        对这个目标求最优解（在概率单纯形上用拉格朗日乘子），得到上面的闭式解：
        在参考模型的基础上，按 exp(奖励 / β) 重新加权。β 越小，越"贪"奖励；β 越大，越接近 π_ref。

    【返回】np.ndarray，和 r 等长的概率分布
    """
    logits = np.log(pi_ref) + r / beta
    p = np.exp(logits - logits.max())
    return p / p.sum()


def kl(p, q):
    """KL(p ‖ q) = Σ p log(p/q)。"""
    m = p > 0
    return float(np.sum(p[m] * np.log(p[m] / q[m])))


def tv(p, q):
    """总变差距离 = ½ Σ |p − q|：两个分布差多远（0 = 完全相同，1 = 完全不重叠）。"""
    return 0.5 * float(np.abs(p - q).sum())


def show_policy(p, label):
    """打印一个策略：每个回答的概率。"""
    print(f"    {label:<14s}" + "".join(f"{x:>11.3f}" for x in p))


# =============================================================================
# 演示 1~4
# =============================================================================
def demo_closed_form():
    print_header("演示 1：RLHF 目标的最优解 π*(y) ∝ π_ref(y)·exp(r(y)/β)")
    print(f"\n    {'回答':<14s}" + "".join(f"{n:>9s}" for n in NAMES))
    print(f"    {'真实奖励 r*':<12s}" + "".join(f"{x:>11.1f}" for x in R_TRUE))
    show_policy(PI_REF, "参考模型 π_ref")
    print()
    print(f"    {'β':>6s}{'期望奖励':>10s}{'KL(π*‖π_ref)':>14s}   最优策略 π*")
    for beta in [10.0, 2.0, 1.0, 0.5, 0.2, 0.05]:
        p = optimal_policy(R_TRUE, PI_REF, beta)
        print(f"    {beta:>6.2f}{p @ R_TRUE:>10.2f}{kl(p, PI_REF):>14.3f}   " + " ".join(f"{x:.2f}" for x in p))
    print(f"\n    （参考模型本身的期望奖励是 {PI_REF @ R_TRUE:.2f}）")
    print("""
  观察：
    - β 很大时，π* 几乎就是 π_ref（不愿意离开参考模型）；β 越小，概率越集中到奖励最高的回答上，
      期望奖励越高，KL 也越大。β 控制着"追求奖励"和"保持原样"之间的权衡。
    - 即使 β 很小，π* 也不会给参考模型概率为 0 的回答分配概率：π* ∝ π_ref·exp(...)。
      RLHF 只能在参考模型"会说"的回答里重新分配概率，不能凭空创造新的回答。""")


def train_rlhf_pg(beta, steps=3000, batch=64, lr=0.05, seed=0):
    """
    用策略梯度（REINFORCE）优化 RLHF 目标 E[r] − β·KL(π‖π_ref)。

    【思路】（讲义第 4.3 节）
        对每个采样到的回答 y，把"奖励"定义为 r(y) − β·(log π(y) − log π_ref(y))（KL 惩罚摊到每个样本上），
        梯度 = E[ (奖励 − 基线) · ∇ log π(y) ]：奖励高于平均的回答，提高它的概率；低于平均的，降低它的概率。
        基线（这里用这一批的平均奖励）不改变梯度的期望，但能大大减小方差。PPO 在此基础上还加了"裁剪"等技巧。

    【返回】训练后的策略（概率分布）
    """
    torch.manual_seed(seed)
    logits = torch.tensor(np.log(PI_REF), dtype=torch.float32, requires_grad=True)      # 从参考模型出发
    log_ref = torch.tensor(np.log(PI_REF), dtype=torch.float32)
    r = torch.tensor(R_TRUE, dtype=torch.float32)
    opt = torch.optim.Adam([logits], lr=lr)
    for _ in range(steps):
        logp = torch.log_softmax(logits, dim=0)
        y = torch.multinomial(logp.exp(), batch, replacement=True)
        reward = r[y] - beta * (logp[y] - log_ref[y]).detach()
        adv = reward - reward.mean()
        loss = -(adv * logp[y]).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
    return torch.softmax(logits, dim=0).detach().numpy()


def make_pairs(n, seed, noisy=True):
    """
    从参考模型采样 n 对回答，按 Bradley–Terry 模型（noisy=True）或"永远选更好的"（noisy=False）标注。

    【返回】(winners, losers)：两个长度为 n 的整数数组（回答的编号）
    【边界】两个回答相同时跳过（没有偏好可言），所以实际对数可能略少于 n。
    """
    rng = np.random.default_rng(seed)
    a = rng.choice(len(PI_REF), n, p=PI_REF)
    b = rng.choice(len(PI_REF), n, p=PI_REF)
    keep = a != b
    a, b = a[keep], b[keep]
    if noisy:
        first = rng.random(len(a)) < 1 / (1 + np.exp(-(R_TRUE[a] - R_TRUE[b])))
    else:
        first = R_TRUE[a] > R_TRUE[b]
    return np.where(first, a, b), np.where(first, b, a)


def train_dpo(winners, losers, beta, steps=2000, lr=0.05, record=None):
    """
    DPO：直接在偏好数据上优化策略。

        loss = −mean log σ( β·[ (log π(y_w) − log π_ref(y_w)) − (log π(y_l) − log π_ref(y_l)) ] )

    【思路】（讲义第 5 节）把 RLHF 最优解反过来写：r(y) = β·log(π*(y)/π_ref(y)) + β·log Z，
        代入 Bradley–Terry 模型，Z 在"两个回答之差"里抵消了 —— 于是可以直接用策略 π 来表示奖励，
        对偏好数据做最大似然，不需要单独训练奖励模型，也不需要采样和强化学习。

    【参数】record：None 或一个列表；给了就每 200 步记录一次 (步数, 策略)
    【返回】训练后的策略
    """
    logits = torch.tensor(np.log(PI_REF), dtype=torch.float32, requires_grad=True)
    log_ref = torch.tensor(np.log(PI_REF), dtype=torch.float32)
    w, l = torch.tensor(winners), torch.tensor(losers)
    opt = torch.optim.Adam([logits], lr=lr)
    for step in range(steps + 1):
        if record is not None and step % 200 == 0:
            record.append((step, torch.softmax(logits, 0).detach().numpy()))
        if step == steps:
            break
        logp = torch.log_softmax(logits, dim=0)
        margin = beta * ((logp[w] - log_ref[w]) - (logp[l] - log_ref[l]))
        loss = -F.logsigmoid(margin).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
    return torch.softmax(logits, dim=0).detach().numpy()


def demo_rlhf_and_dpo():
    beta = 0.5
    target = optimal_policy(R_TRUE, PI_REF, beta)
    print_header(f"演示 2：策略梯度（RLHF）优化，β = {beta}")
    print(f"\n    {'回答':<14s}" + "".join(f"{n:>9s}" for n in NAMES))
    show_policy(target, "闭式最优 π*")
    pg = train_rlhf_pg(beta)
    show_policy(pg, "策略梯度")
    print(f"\n    和 π* 的总变差距离：{tv(pg, target):.3f}")
    print("""
  观察：策略梯度只需要"采样 → 打分 → 好的回答提高概率"，就收敛到了闭式最优解附近。
    真实的 RLHF 就是这样做的，只是"回答"是整段文字、π 是一个大语言模型、奖励来自奖励模型，
    每一步都要生成大量文本并打分 —— 又慢又不稳定（讲义第 4.3 节）。""")

    print_header(f"演示 3：DPO —— 只用偏好数据，β = {beta}")
    print(f"\n    {'偏好数据（对）':<14s}{'和 π* 的总变差距离':>18s}{'期望奖励':>10s}{'KL(π‖π_ref)':>13s}")
    print(f"    {'（π* 本身）':<14s}{0:>18.3f}{target @ R_TRUE:>12.2f}{kl(target, PI_REF):>13.3f}")
    policies = {}
    for n in [100, 1000, 10000, 100000]:
        w, l = make_pairs(n, seed=1)
        p = train_dpo(w, l, beta)
        policies[n] = p
        print(f"    {len(w):<16d}{tv(p, target):>18.3f}{p @ R_TRUE:>12.2f}{kl(p, PI_REF):>13.3f}")
    print("""
  观察：偏好数据越多，DPO 得到的策略越接近 RLHF 的最优策略 π*。
    DPO 和 RLHF 优化的是**同一个目标**，只是 DPO 绕过了"先训练奖励模型、再做强化学习"这两步。""")

    print_header("演示 4：DPO 的隐式奖励 β·log(π/π_ref)")
    p = policies[100000]
    implicit = beta * np.log(p / PI_REF)
    print(f"\n    {'回答':<14s}" + "".join(f"{n:>9s}" for n in NAMES))
    print(f"    {'真实奖励 − 均值':<10s}" + "".join(f"{x:>11.2f}" for x in R_TRUE - R_TRUE.mean()))
    print(f"    {'隐式奖励 − 均值':<10s}" + "".join(f"{x:>11.2f}" for x in implicit - implicit.mean()))
    print("""
  观察：从 DPO 训练出的策略里，可以"读出"一个奖励：β·log(π(y)/π_ref(y))。它和真实奖励只差一个常数 ——
    这就是 DPO 论文标题的意思：*Your Language Model is Secretly a Reward Model*。""")


def demo_dpo_trap():
    print_header("演示 5：DPO 的陷阱 —— 偏好毫无噪声时，β 管不住它")
    beta = 0.5
    target = optimal_policy(R_TRUE, PI_REF, beta)
    w, l = make_pairs(10000, seed=2, noisy=False)
    rec = []
    train_dpo(w, l, beta, steps=2000, record=rec)
    print(f"\n    偏好数据：永远选真实奖励更高的那个（没有任何噪声）。π* 给\"正面、简洁\"的概率是 {target[0]:.2f}，"
          f"KL = {kl(target, PI_REF):.2f}")
    print(f"    {'训练步数':>8s}{'P(正面、简洁)':>14s}{'KL(π‖π_ref)':>13s}")
    for step, p in rec[::2]:
        print(f"    {step:>8d}{p[0]:>14.3f}{kl(p, PI_REF):>13.2f}")
    print("""
  观察：在演示 3 里，DPO 停在 π* 附近；在这里，它却一路把"正面、简洁"的概率推向 1，KL 越来越大。
    原因：如果数据里 A 总是赢 B，Bradley–Terry 的最大似然估计会认为 A 比 B 好"无穷多"，
    DPO 的损失只有在 log(π(A)/π(B)) → ∞ 时才降到 0 —— β 只改变"走多快"，挡不住"走多远"。
    真实数据里每一对回答通常只被标注一次，"永远是这个赢"的情况很常见，所以 DPO 可能过拟合偏好数据、
    偏离参考模型太远（Azar 等 2023 的 IPO 等方法就是为了修补这一点）。实践中常用早停、更大的 β、
    或者在 DPO 损失之外再加上对参考模型的约束。""")


# =============================================================================
# 演示 6：过度优化
# =============================================================================
def best_of_n_distribution(scores, pi_ref, n):
    """
    best-of-n 的精确分布：从 π_ref 独立采样 n 个回答，选 scores 最高的那个。

    【思路】按 scores 从低到高排序，F_i = 前 i 个回答的总概率（"抽到的回答不比第 i 个好"的概率）。
        n 个样本里最好的那个恰好是第 i 个 ⟺ 所有样本都不比它好、且不是所有样本都比它差：
            P(选中第 i 个) = F_i^n − F_{i−1}^n
    【返回】np.ndarray，best-of-n 策略给每个回答的概率
    """
    order = np.argsort(scores)
    F_cum = np.cumsum(pi_ref[order])
    F_prev = np.concatenate([[0.0], F_cum[:-1]])
    p = np.zeros_like(pi_ref)
    p[order] = F_cum ** n - F_prev ** n
    return p


N_LIST = [1, 2, 4, 8, 16, 32, 64, 256, 1024, 4096, 16384]


def overoptimization_experiment(seed, K=2000, n_pairs=1000, l2=1e-5, steps=1000):
    """
    一次过度优化实验（给定随机种子，之后的计算都是精确的）。

    【步骤】
        1. 造 K 种可能的回答：真实奖励 r* ~ N(0, 1)；参考模型 π_ref 是一个"长尾"分布（少数回答很常见，大多数很罕见）
        2. 从 π_ref 采样 n_pairs 对回答，按 Bradley–Terry 模型标注，用最大似然（很弱的 L2 正则）拟合奖励模型 r̂
        3. 对每个 n，精确计算 best-of-n（用 r̂ 选）的分布，以及它的 KL、r̂ 的期望、r* 的期望

    【返回】dict：kl、kl_formula、proxy、gold、gold_select（用 r* 本身来选时的真实奖励），都是长度为 len(N_LIST) 的数组；
            以及 n_seen（至少出现过一次的回答种数）、corr_w / corr（奖励模型和真实奖励的加权 / 不加权相关系数）
    """
    rng = np.random.default_rng(seed)
    r_true = rng.normal(size=K)
    pi_ref = rng.dirichlet(np.full(K, 0.3))
    a = rng.choice(K, n_pairs, p=pi_ref)
    b = rng.choice(K, n_pairs, p=pi_ref)
    first = rng.random(n_pairs) < 1 / (1 + np.exp(-(r_true[a] - r_true[b])))
    win, lose = np.where(first, a, b), np.where(first, b, a)
    r_hat = torch.zeros(K, requires_grad=True)
    opt = torch.optim.Adam([r_hat], lr=0.1)
    for _ in range(steps):
        loss = -F.logsigmoid(r_hat[win] - r_hat[lose]).mean() + l2 * (r_hat ** 2).sum()
        opt.zero_grad()
        loss.backward()
        opt.step()
    r_hat = r_hat.detach().numpy()
    c = np.cov(r_hat, r_true, aweights=pi_ref)                       # 按参考模型的概率加权的协方差矩阵
    out = {"kl": [], "kl_formula": [], "proxy": [], "gold": [], "gold_select": [],
           "n_seen": len(set(a) | set(b)), "corr_w": c[0, 1] / np.sqrt(c[0, 0] * c[1, 1]),
           "corr": np.corrcoef(r_hat, r_true)[0, 1]}
    for n in N_LIST:
        p = best_of_n_distribution(r_hat, pi_ref, n)
        out["kl"].append(kl(p, pi_ref))
        out["kl_formula"].append(math.log(n) - (n - 1) / n)
        out["proxy"].append(p @ r_hat)
        out["gold"].append(p @ r_true)
        out["gold_select"].append(best_of_n_distribution(r_true, pi_ref, n) @ r_true)
    return {k: np.array(v) if isinstance(v, list) else v for k, v in out.items()}


def demo_overoptimization():
    print_header("演示 6：过度优化 —— 用不完美的奖励模型做 best-of-n")
    e = overoptimization_experiment(seed=0)
    print(f"\n    2000 种回答；奖励模型用 1000 对偏好数据训练，只有 {e['n_seen']} 种回答至少出现过一次。")
    print(f"    奖励模型和真实奖励的相关系数：在参考模型常说的回答上（按 π_ref 加权）{e['corr_w']:.2f}，不加权 {e['corr']:.2f}")
    print(f"\n    （随机种子 0）")
    print(f"    {'n':>6s}{'KL（精确）':>11s}{'log n − (n−1)/n':>16s}{'奖励模型的分数':>14s}{'真实奖励':>9s}{'用真实奖励选':>12s}")
    for i, n in enumerate(N_LIST):
        print(f"    {n:>6d}{e['kl'][i]:>11.2f}{e['kl_formula'][i]:>16.2f}{e['proxy'][i]:>14.2f}"
              f"{e['gold'][i]:>11.2f}{e['gold_select'][i]:>12.2f}")
    runs = [overoptimization_experiment(seed=s) for s in range(20)]
    gold = np.array([r["gold"] for r in runs])
    proxy = np.array([r["proxy"] for r in runs])
    declined = int(np.sum(gold[:, -1] < gold.max(axis=1) - 0.1))
    print(f"\n    20 个随机种子的平均（每个种子：重新造回答、重新采样偏好数据、重新训练奖励模型）：")
    print(f"    {'n':>6s}" + "".join(f"{n:>7d}" for n in N_LIST))
    print(f"    {'奖励模型':<6s}" + "".join(f"{x:>7.2f}" for x in proxy.mean(0)))
    print(f"    {'真实奖励':<6s}" + "".join(f"{x:>7.2f}" for x in gold.mean(0)))
    print(f"\n    20 个种子里，有 {declined} 个种子的真实奖励在 n = 16384 时比它的最高点低了 0.1 以上。")
    print("""
  观察（每个种子内部是精确计算；不同种子之间差别很大）：
    - best-of-n：从参考模型采样 n 个回答，用奖励模型挑分数最高的。n 越大，"选得越狠"，离参考模型越远：
      KL 大约是 log n − (n−1)/n（回答的种类很多时，精确值和公式很接近；n 太大时离散效应开始显现）。
    - 奖励模型的分数随 n **一直上升** —— 我们在优化的就是它。
    - 真实奖励一开始也上升（奖励模型大体上是对的），但涨得比奖励模型的分数慢得多，很快就停滞了；
      两者的差距越拉越大。在不少种子里（比如种子 0），选得太狠之后真实奖励甚至**不升反降**：
      被选中的往往是奖励模型"看走眼"的回答 —— 很罕见、在偏好数据里只出现过一两次、估计误差很大。
      也有一些种子里真实奖励一直上升（奖励模型恰好没有在最高分的地方看走眼）：过度优化是一种风险，不是必然。
    - 最后一列（种子 0）：如果直接用真实奖励来选，真实奖励一直上升 —— 问题不在 best-of-n 本身，而在奖励模型不完美。
    这就是**古德哈特定律**："当一个指标变成目标，它就不再是好的指标。"Gao 等（2022）在真实的语言模型上
    系统地测量了这种曲线（讲义第 6.1 节）；04 号文件在迷你 GPT 上做同样的实验。""")


def main():
    demo_closed_form()
    demo_rlhf_and_dpo()
    demo_dpo_trap()
    demo_overoptimization()


if __name__ == "__main__":
    main()