"""
=============================================================================
第 9 讲 · 代码 2：奖励模型——在数据分布内很准，在数据分布外会出错
=============================================================================

对应讲义：lectures/lecture09_rlhf_dpo.md 第 2、4 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
RLHF 的第二步：用人类的偏好数据训练一个**奖励模型** r̂(y)，让它模仿人类的判断。
之后的 RL 阶段不再询问人类，而是用 r̂ 来给模型的回答打分。

问题在于：奖励模型只在"人类比较过的那些回答"附近见过数据。
RL 会把模型推向 r̂ 认为好的地方——很可能是奖励模型**从没见过**的地方。
本文件先训练奖励模型，然后测量它在"数据分布内"和"数据分布外"的准确率。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    cd code/lecture09
    python 02_reward_model.py      （约 2 秒）
"""

import numpy as np

from toy_llm import (N_RESP, LENGTH, QUALITY, FEATURES, true_reward, reference_policy,
                     sample_preferences, fit_linear_reward_model, sigmoid)

np.set_printoptions(precision=3, suppress=True)


def pairwise_accuracy(reward_hat, idx_a, idx_b):
    """
    在一组回答对上，奖励模型判断"哪个更好"与**真实奖励**的判断一致的比例。
    （用真实奖励的判断作为标准，而不是用带噪声的人类标签，这样能看出模型本身的对错。）
    """
    r = true_reward()
    agree = np.sign(reward_hat[idx_a] - reward_hat[idx_b]) == np.sign(r[idx_a] - r[idx_b])
    return float(agree.mean())


def main():
    rng = np.random.default_rng(1)
    ref = reference_policy()
    r = true_reward()

    # ---------------------------------------------------------------------
    print("=" * 78)
    print("演示 1：参考策略（SFT 模型）会生成什么样的回答？")
    print("=" * 78)
    bins = np.arange(0, 11, 1)
    print("   长度区间    π_ref 的概率     该区间回答的平均真实奖励")
    for lo, hi in zip(bins[:-1], bins[1:]):
        m = (LENGTH >= lo) & (LENGTH < hi)
        print(f"   [{lo:2d}, {hi:2d})      {ref[m].sum():7.3f}          {r[m].mean():6.2f}   "
              + "█" * int(ref[m].sum() * 60))
    print(f"\n  π_ref 的平均长度 {np.sum(ref * LENGTH):.2f}，平均真实奖励 {np.sum(ref * r):.3f}")
    print(f"  所有候选回答中真实奖励最高的是长度 {LENGTH[np.argmax(r)]:.1f} 的回答（奖励 {r.max():.2f}）")
    print("  → SFT 模型几乎只写长度 1~4 的回答。在这个区间里，更长的回答通常更好。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 2：用 3000 对人类比较训练线性奖励模型 r̂(y) = w·[质量, 长度]")
    print("=" * 78)
    chosen, rejected = sample_preferences(ref, 3000, rng)
    w = fit_linear_reward_model(chosen, rejected)
    r_hat = FEATURES @ w
    human_agree = np.mean(r[chosen] > r[rejected])
    print(f"  学到的权重 w = {w}（质量, 长度）")
    print(f"  训练数据中，人类的选择与真实奖励一致的比例：{human_agree * 100:.1f}%（人类偏好有噪声）")
    print("  → 奖励模型学到了\"质量越高越好\"（正确）和\"越长越好\"（只在短回答的范围内正确）。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 3：奖励模型在哪里准、在哪里不准？")
    print("=" * 78)
    test_rng = np.random.default_rng(2)
    # (a) 数据分布内：从 π_ref 中抽的回答对
    a = test_rng.choice(N_RESP, size=5000, p=ref)
    b = test_rng.choice(N_RESP, size=5000, p=ref)
    print(f"  (a) 从 SFT 模型中抽的回答对（训练分布内）：准确率 {pairwise_accuracy(r_hat, a, b) * 100:.1f}%")
    # (b) 长回答之间的比较：两个回答都长于 6
    long_idx = np.flatnonzero(LENGTH > 6)
    a = test_rng.choice(long_idx, size=5000)
    b = test_rng.choice(long_idx, size=5000)
    print(f"  (b) 两个回答都长于 6（训练分布外）：         准确率 {pairwise_accuracy(r_hat, a, b) * 100:.1f}%")
    # (c) 专门比较"长度差别大、质量差不多"的回答
    a = test_rng.choice(long_idx, size=5000)
    mid_idx = np.flatnonzero((LENGTH > 4) & (LENGTH < 5.5))
    b = test_rng.choice(mid_idx, size=5000)
    print(f"  (c) 长回答（>6） vs 适中回答（4~5.5）：          准确率 {pairwise_accuracy(r_hat, a, b) * 100:.1f}%")
    print("  → 在训练分布内，奖励模型相当准。一旦到了分布外，尤其是需要判断\"太长是不是不好\"时，")
    print("    它的准确率跌破 50%——比随机猜还差：它系统性地认为更长的更好，而真相恰恰相反。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 4：沿着长度看，奖励模型与真实奖励的差距（固定质量 = 0）")
    print("=" * 78)
    print("   长度    真实奖励    奖励模型     差距")
    for length in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]:
        true_v = 0.8 * length - 0.08 * length ** 2
        hat_v = w[1] * length
        # 奖励模型只能确定到一个常数（01 文件），这里把两者在长度 2.5 处对齐
        offset = (0.8 * 2.5 - 0.08 * 2.5 ** 2) - w[1] * 2.5
        print(f"   {length:4d}    {true_v:7.2f}    {hat_v + offset:7.2f}    {hat_v + offset - true_v:+6.2f}")
    print("  → 在长度 1~4（数据最多的地方）两者几乎重合；长度越大，奖励模型高估得越厉害。")
    print("    下一个文件会看到：如果不加约束地用 RL 去最大化这个奖励模型，模型会学会\"写得越长越好\"。")


if __name__ == "__main__":
    main()