"""
=============================================================================
第 9 讲 · 代码 3：RLHF 的 RL 阶段——KL 约束与奖励过度优化
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
RLHF 的第三步：用 RL 优化语言模型，让奖励模型的分数尽量高。但不是无约束地优化，而是

    max_π  E_{y~π}[ r̂(y) ]  -  β · KL(π ‖ π_ref)

多出来的 KL 项惩罚模型离参考模型（SFT 模型）太远。为什么需要它？02 文件已经给出了答案：
奖励模型只在参考模型的分布附近可信。本文件把这件事定量地展示出来：

    演示 1：验证这个目标的闭式最优解 π*(y) ∝ π_ref(y) exp(r̂(y)/β)
    演示 2：β 从大到小——奖励模型的分数一路上升，**真实**奖励却先升后降（过度优化）
    演示 3：用策略梯度（RL）真的去优化这个目标，看它是否收敛到闭式解

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    cd code/lecture09
    python 03_rlhf_kl.py      （约 1 秒）
"""

import numpy as np

from toy_llm import (N_RESP, LENGTH, FEATURES, true_reward, reference_policy, sample_preferences,
                     fit_linear_reward_model, optimal_kl_policy, softmax, kl)

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


def kl_objective(pi, reward, ref, beta):
    """KL 正则的 RL 目标 E_π[r] - β KL(π ‖ π_ref)。"""
    return float(np.sum(pi * reward) - beta * kl(pi, ref))


def policy_gradient_kl(reward, ref, beta, n_iters, batch, lr, rng):
    """
    用 REINFORCE（第 6 讲）优化 KL 正则目标，策略参数是 400 个回答的 logits θ（从 log π_ref 开始）。

    【RLHF 中的常见做法：把 KL 惩罚放进每个样本的"奖励"里】
        对每个采样到的回答 y：  R(y) = r̂(y) - β · [log π_θ(y) - log π_ref(y)]
        梯度估计：              (1/B) Σ_k (R(y_k) - b) ∇_θ log π_θ(y_k)，  b = 本批 R 的平均（基线）
    为什么可以这样做？因为 E_π[-β log(π/π_ref)] = -β KL(π‖π_ref)，用似然比技巧对它求导，
    多出来的一项 E[∇ log π] = 0（第 6 讲 2.4 节），所以"把 KL 当作奖励的一部分"正好给出正确的梯度。
    真实的 RLHF（用 PPO）中，这个 KL 惩罚是在**每个 token** 上计算的，道理相同。

    【返回】每隔 50 步记录一次 (步数, 当前策略与闭式解之间的 KL, 目标函数值)。
    """
    theta = np.log(ref).copy()
    target = optimal_kl_policy(reward, ref, beta)
    log = []
    for it in range(n_iters + 1):
        pi = softmax(theta)
        if it % 100 == 0:
            log.append((it, kl(target, pi), kl_objective(pi, reward, ref, beta)))
        ys = rng.choice(N_RESP, size=batch, p=pi)
        R = reward[ys] - beta * (np.log(pi[ys]) - np.log(ref[ys]))
        adv = R - R.mean()
        grad = np.zeros(N_RESP)
        np.add.at(grad, ys, adv)                  # Σ_k adv_k e_{y_k}
        grad -= adv.sum() * pi                     # Σ_k adv_k π（这一项因为 adv 的均值为 0，恰好为 0）
        theta += lr * grad / batch
    return softmax(theta), log


def main():
    rng = np.random.default_rng(1)
    ref = reference_policy()
    r_true = true_reward()
    chosen, rejected = sample_preferences(ref, 3000, rng)
    w = fit_linear_reward_model(chosen, rejected)
    r_hat = FEATURES @ w
    print(f"奖励模型（与 02 文件相同的训练方式）：w = {w}（质量, 长度）")
    print(f"参考策略：平均长度 {np.sum(ref * LENGTH):.2f}，平均真实奖励 {np.sum(ref * r_true):.3f}\n")

    # ---------------------------------------------------------------------
    print("=" * 78)
    print("演示 1：闭式解 π*(y) ∝ π_ref(y) exp(r̂(y)/β) 真的是最优的吗？（β = 0.5）")
    print("=" * 78)
    beta = 0.5
    pi_star = optimal_kl_policy(r_hat, ref, beta)
    best = kl_objective(pi_star, r_hat, ref, beta)
    print(f"  闭式解的目标值：{best:.4f}")
    worse = 0
    for _ in range(1000):
        perturbed = softmax(np.log(pi_star) + rng.normal(0, 0.3, N_RESP))
        worse += kl_objective(perturbed, r_hat, ref, beta) < best
    print(f"  随机扰动 1000 次，目标值变差的次数：{worse} / 1000")
    print("  → 闭式解确实是最优的（讲义 3.3 节有推导：目标等于 -β KL(π ‖ π*) 加一个常数）。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 2：β 越小，优化越激进——奖励模型的分数 vs 真实奖励")
    print("=" * 78)
    print("     β     KL(π‖π_ref)   平均长度   奖励模型分数   真实奖励   │ 若直接优化真实奖励：真实奖励")
    ref_rhat, ref_true = np.sum(ref * r_hat), np.sum(ref * r_true)
    rows = []
    for beta in [10, 3, 1, 0.5, 0.3, 0.2, 0.1, 0.05, 0.02]:
        pi = optimal_kl_policy(r_hat, ref, beta)
        gold = optimal_kl_policy(r_true, ref, beta)
        rows.append((beta, kl(pi, ref), np.sum(pi * r_true)))
        print(f"  {beta:6.2f}     {kl(pi, ref):7.2f}      {np.sum(pi * LENGTH):6.2f}      "
              f"{np.sum(pi * r_hat) - ref_rhat:+7.2f}       {np.sum(pi * r_true) - ref_true:+7.2f}  │        "
              f"{np.sum(gold * r_true) - ref_true:+7.2f}")
    best_beta = max(rows, key=lambda x: x[2])
    print(f"  （奖励模型分数和真实奖励都表示为相对于参考策略的提升。）")
    print(f"\n  真实奖励在 β ≈ {best_beta[0]}（KL ≈ {best_beta[1]:.1f}）时最高，之后反而下降。")
    print("  观察：")
    print("  - 奖励模型的分数随 β 减小**单调上升**——从它的角度看，优化越激进越好。")
    print("  - 真实奖励却**先升后降**：优化到一定程度后，模型开始利用奖励模型的缺陷——")
    print("    回答越写越长（平均长度从 2.6 涨到 9 以上），奖励模型给高分，真实质量却在下降。")
    print("    这就是**奖励过度优化**（reward over-optimization），也叫 Goodhart 定律：")
    print("    \"当一个指标变成了优化目标，它就不再是一个好指标。\"")
    print("  - 最右一列是\"如果我们能直接优化真实奖励\"：它随 β 减小单调上升。")
    print("    两者的差距完全来自奖励模型在分布外的错误。")
    print("  - KL 项的作用：把模型限制在奖励模型\"还可信\"的范围内。β 是需要仔细调的超参数。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 3：用策略梯度（RL）优化 KL 正则目标，会收敛到闭式解吗？（β = 0.5）")
    print("=" * 78)
    beta = 0.5
    pi_rl, log = policy_gradient_kl(r_hat, ref, beta, n_iters=20_000, batch=256, lr=1.0, rng=rng)
    target = optimal_kl_policy(r_hat, ref, beta)
    print("   步数    KL(闭式解 ‖ 当前策略)    目标值")
    for it, d, obj in log[::25]:
        print(f"   {it:5d}        {d:8.4f}           {obj:7.4f}")
    print(f"  闭式解的目标值：{kl_objective(target, r_hat, ref, beta):.4f}")
    print(f"  RL 学到的策略：平均长度 {np.sum(pi_rl * LENGTH):.2f}（闭式解 {np.sum(target * LENGTH):.2f}），"
          f"真实奖励 {np.sum(pi_rl * r_true):.3f}（闭式解 {np.sum(target * r_true):.3f}）")
    print("  → RL 一步步逼近闭式解。剩下的一点差距主要来自那些在 π_ref 下概率极小的回答：")
    print("    它们很少被采样到，梯度信号也就极少——这是 RL 样本效率问题的一个缩影。")
    print("    在真实的语言模型中，回答有天文数字那么多，没法直接写出闭式解，")
    print("    只能用 PPO 这样的 RL 算法去逼近它。**但既然知道了最优解的形式，能不能绕过 RL？**")
    print("    这正是下一个文件 DPO 的出发点。")


if __name__ == "__main__":
    main()