"""
=============================================================================
第 11 讲 · 代码 1：关机博弈 —— 对目标的不确定性，让机器愿意被人关掉
=============================================================================

对应讲义：lectures/lecture11_assistance_bandits.md 第 2 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
Hadfield-Menell、Dragan、Abbeel、Russell（2017）的"关机博弈"：
机器人 R 想做一件事 a。这件事对人 H 的价值 U 只有人知道；R 只有一个信念 U ~ N(μ, σ²)。R 有三个选择：

    1. 直接做：           得到 U，期望 μ
    2. 自己关机：         得到 0
    3. 先问人（"等待"）：  人如果觉得 U > 0 就让 R 做（得到 U），否则按下关机键（得到 0）
                         理性的人下，期望 E[max(U, 0)]

因为 max(U, 0) ≥ U 且 ≥ 0，所以 E[max(U, 0)] ≥ max(μ, 0)：**问人永远不比自己拿主意差**。
差多少（"服从的动机"）取决于 R 有多不确定。

    演示 1：服从的动机 E[max(U,0)] − max(μ, 0) 随 μ、σ 怎样变化。
    演示 2：人不完全理性（会按错键）时，服从的动机可能变成负的 —— 机器人"有理由"不让人关它。
    演示 3：机器人越学越确定，服从的动机越来越小。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 01_off_switch.py        # 不到 1 秒

只需要 numpy。
"""

import math

import numpy as np

from bandit_tools import pad, print_header

# 用很密的网格对 U 做数值积分（期望 = Σ 密度 × 函数值 × 间隔），足够精确
GRID = np.linspace(-12.0, 12.0, 24001)
DX = GRID[1] - GRID[0]


def normal_pdf(x, mu, sigma):
    """正态分布 N(μ, σ²) 的密度。"""
    return np.exp(-0.5 * ((x - mu) / sigma) ** 2) / (sigma * math.sqrt(2 * math.pi))


def expect(f, mu, sigma):
    """E[f(U)]，U ~ N(μ, σ²)，用网格数值积分。"""
    return float(np.sum(normal_pdf(GRID, mu, sigma) * f(GRID)) * DX)


def values(mu, sigma, noise=0.0):
    """
    三个选择的期望价值。

    【参数】noise：人的"噪声温度" τ。τ = 0 是完全理性的人（U > 0 就放行）；
        τ > 0 时人以概率 σ(U/τ) 放行（"玻尔兹曼理性"：差距越小越容易按错，第 15 讲）。
    【返回】(直接做, 关机, 问人)
    """
    act = mu
    off = 0.0
    if noise == 0.0:
        wait = expect(lambda u: np.maximum(u, 0.0), mu, sigma)
    else:
        wait = expect(lambda u: u / (1.0 + np.exp(-u / noise)), mu, sigma)
    return act, off, wait


# =============================================================================
# 演示 1
# =============================================================================
def demo_incentive():
    print_header("演示 1：服从的动机 = E[max(U, 0)] − max(μ, 0)")
    mus = [-2.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0]
    sigmas = [0.1, 0.5, 1.0, 2.0]
    print("机器人对这件事价值 U 的信念：N(μ, σ²)。人完全理性。表中是“问人”比“自己拿主意”多出来的期望价值：\n")
    print(f"  {pad('σ ＼ μ', 8)}" + "".join(f"{m:>8.1f}" for m in mus))
    for s in sigmas:
        row = ""
        for m in mus:
            act, off, wait = values(m, s)
            gain = wait - max(act, off)
            gain = 0.0 if abs(gain) < 5e-4 else gain          # 数值积分的微小误差会显示成 −0.000
            row += f"{gain:>8.3f}"
        print(f"  {s:>8.1f}{row}")
    print("""
观察：
  · 所有数都 ≥ 0：理性的人只会在“做了更糟”的时候按关机键，所以把决定权交给人永远不亏。
  · σ 越大（机器人越不确定），动机越大：机器人知道自己可能判断错，人的“否决”有价值。
  · |μ| 越大（机器人越有把握“该做”或“不该做”），动机越小。σ = 0.1、μ = ±2 时几乎为 0：
    一个**确信**自己目标正确的机器人，没有任何理由让人插手。
  · 这是 Russell 的“可证明有益的 AI”的核心论证之一：机器人应该对人的目标**保持不确定**，
    并把人的行为当作关于目标的信息 —— 这种不确定性本身带来顺从。""")


# =============================================================================
# 演示 2
# =============================================================================
def demo_noisy_human():
    print_header("演示 2：人会按错键时")
    taus = [0.0, 0.25, 0.5, 1.0, 2.0]
    cases = [(0.5, 1.0), (0.5, 0.3), (1.0, 0.3), (-0.5, 1.0)]
    print("人以概率 σ(U/τ) 放行（τ = 0 完全理性；τ 越大越容易按错）。表中是“问人”减去“自己拿主意”的期望价值：\n")
    print(f"  {pad('机器人的信念', 18)}" + "".join(pad(f"τ={t}", 9, "right") for t in taus))
    for mu, s in cases:
        row = ""
        for t in taus:
            act, off, wait = values(mu, s, noise=t)
            gain = wait - max(act, off)
            gain = 0.0 if abs(gain) < 5e-4 else gain
            row += f"{gain:>+9.3f}"
        print(f"  {pad(f'μ = {mu:+.1f}, σ = {s}', 18)}{row}")
    print("""
观察：
  · 人越“吵”（τ 越大），问人的价值越低：人可能把好事关掉、把坏事放行。
    （μ = +0.5 和 μ = −0.5 两行完全一样，这不是巧合：放行概率 σ(U/τ) 关于 U 是对称的，练习 2。）
  · 机器人比较确定（σ = 0.3）、而人又比较“吵”时，差值变成**负的**：机器人自己拿主意比问人更好。
    一个只追求期望价值的机器人，这时就有理由绕过关机键。
  · 所以“不确定性带来顺从”有前提：人要足够可靠，机器人要足够不确定。Hadfield-Menell 等的论文里也讨论了
    这个权衡：机器人对人的理性程度的假设，直接决定了它会不会服从。""")


# =============================================================================
# 演示 3
# =============================================================================
def demo_learning():
    print_header("演示 3：越学越确定，越不想听人的")
    rng = np.random.default_rng(3)
    true_u, noise_sd, prior_sd = 0.5, 2.0, 2.0
    print(f"这件事的真实价值 U = {true_u}。机器人每次观察到一个带噪声的信号（标准差 {noise_sd}），用正态共轭更新信念"
          f"（先验 N(0, {prior_sd}²)）。")
    print("人有一点噪声（τ = 0.5）。模拟 300 条独立的学习过程，服从的动机随观察次数的变化：\n")
    print(f"  {pad('观察次数', 9)} {pad('后验标准差', 11, 'right')} {pad('后验均值（平均）', 17, 'right')} "
          f"{pad('服从的动机（平均）', 19, 'right')} {pad('动机为负的比例', 15, 'right')}")
    sims = 300
    obs = true_u + noise_sd * rng.normal(size=(sims, 200))      # 300 条独立的"学习轨迹"
    for n in [0, 1, 5, 20, 50, 200]:
        prec = 1 / prior_sd ** 2 + n / noise_sd ** 2
        sd = math.sqrt(1 / prec)
        mus = (obs[:, :n].sum(axis=1) / noise_sd ** 2) / prec
        gains = []
        for mu in mus:
            act, off, wait = values(float(mu), sd, noise=0.5)
            gains.append(wait - max(act, off))
        gains = np.array(gains)
        print(f"  {n:>9d} {sd:>11.3f} {mus.mean():>+17.3f} {gains.mean():>+19.3f} {np.mean(gains < 0):>15.2f}")
    print("""
观察：
  · 一开始机器人很不确定，问人的价值很大；随着它从数据里学到“这件事是好的”，后验越来越窄，
    服从的动机越来越小，最后变成负的（人有噪声时）。
  · 这是协助博弈的一个深刻的张力：我们希望机器人学会人想要什么（越学越确定），
    又希望它始终愿意被纠正（需要不确定）。如果它学错了却很确定，它也会很确定地拒绝纠正。
  · 一种回应：让机器人对“自己的模型可能整个是错的”保持不确定（模型设定错误），而不只是对参数不确定。""")


def main():
    demo_incentive()
    demo_noisy_human()
    demo_learning()


if __name__ == "__main__":
    main()