"""
=============================================================================
第 11 讲 · 代码 2：示范 vs 教学 —— 合作逆强化学习（CIRL）里的人和机器人
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
协助博弈（CIRL，Hadfield-Menell 等 2016）：人知道目标 θ，机器人不知道，两人共享同一个目标（人的回报）。
人先做一个动作（"示范"），机器人看了以后推断 θ，再替人做事。

一个简单的设定：目标有三种 θ1、θ2、θ3，人可以从 A、B、C 三个动作里选一个给机器人看。每种目标下各动作的价值：

              A      B      C
        θ1   1.0    0.0    0.6
        θ2   1.0    0.8    0.0
        θ3   0.0    0.6    1.0

注意：θ1 和 θ2 的最好动作都是 A。

两种人：
    · 示范者：做对自己最好的动作（玻尔兹曼理性，温度 τ = 0.1）—— "照常做事，让机器人看"
    · 教师  ：挑一个能让机器人最好地猜出 θ 的动作（Shafto 等 2014 的"教学"模型）
两种机器人：
    · 字面的机器人：假设人是示范者，用贝叶斯公式推断 θ（逆强化学习的标准假设）
    · 务实的机器人：知道人在教，按"教师"的模型推断 θ

    演示 1：四种组合下，机器人猜对的概率、对真实目标的平均信心、人为了"教"付出的代价。
    演示 2：机器人以为人在教、人其实只是在示范 —— "自信地猜错"。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_teaching_vs_demonstration.py        # 不到 1 秒

只需要 numpy。
"""

import numpy as np

from bandit_tools import pad, print_header

U = np.array([[1.0, 0.0, 0.6],
              [1.0, 0.8, 0.0],
              [0.0, 0.6, 1.0]])
GOALS = ["θ1", "θ2", "θ3"]
ACTIONS = ["A", "B", "C"]
TAU = 0.1          # 示范者的玻尔兹曼温度
ALPHA = 3.0        # 教师有多"用心"：P(动作 | θ) ∝ P_字面(θ | 动作)^α


def softmax_rows(x):
    """每一行做 softmax。"""
    x = x - x.max(axis=1, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=1, keepdims=True)


def posterior(likelihood):
    """
    均匀先验下的后验 P(θ | 动作)：把似然矩阵 P(动作 | θ)（行 = θ，列 = 动作）按列归一化。
    """
    return likelihood / likelihood.sum(axis=0, keepdims=True)


def build_models():
    """
    构造两种人、两种机器人。

    【返回】dict：
        demo ：示范者 P(动作 | θ) = softmax(U[θ] / τ)
        lit  ：字面的机器人 P(θ | 动作)，假设人是示范者
        teach：教师 P(动作 | θ) ∝ lit(θ | 动作)^α —— 选让字面的机器人最可能猜对的动作
        prag ：务实的机器人 P(θ | 动作)，假设人是教师
    【思路】这是"递归推理"的一层：教师在想"机器人会怎样理解我"，务实的机器人在想"人为什么选这个"。
        语言学里的"理性言语行为"模型（RSA）、认知科学里的"教学 vs 示范"（Ho 等 2016）都是同一个结构。
    """
    demo = softmax_rows(U / TAU)
    lit = posterior(demo)
    teach = softmax_rows(ALPHA * np.log(lit + 1e-300))
    prag = posterior(teach)
    return dict(demo=demo, lit=lit, teach=teach, prag=prag)


def evaluate(human, robot):
    """
    【返回】(猜对的概率, 对真实目标的平均后验概率, "自信地猜错"的概率, 人为此付出的平均代价)
        猜对：机器人后验最大的 θ 就是真实目标；
        自信地猜错：机器人把 ≥ 80% 的概率放在了一个错误的目标上；
        代价：人选的动作比"对自己最好的动作"少了多少价值（目标均匀分布）。
    """
    right = conf = wrong_conf = cost = 0.0
    for th in range(3):
        for a in range(3):
            p = human[th, a] / 3
            post = robot[:, a]
            right += p * (np.argmax(post) == th)
            conf += p * post[th]
            wrong_conf += p * (post.max() >= 0.8 and np.argmax(post) != th)
            cost += p * (U[th].max() - U[th, a])
    return right, conf, wrong_conf, cost


def demo_four_combos(models):
    print_header("演示 1：两种人 × 两种机器人")
    print("每种目标下各动作的价值（θ1 和 θ2 的最好动作都是 A）：")
    print("        " + "".join(f"{a:>6s}" for a in ACTIONS))
    for g, row in zip(GOALS, U):
        print(f"    {g:>3s} " + "".join(f"{v:>6.1f}" for v in row))
    print("\n人选每个动作的概率：")
    for name, key in [("示范者", "demo"), ("教师", "teach")]:
        print(f"  {name}：" + "；".join(f"{g} → " + " ".join(f"{a}{p:.2f}" for a, p in zip(ACTIONS, models[key][i]))
                                      for i, g in enumerate(GOALS)))
    print(f"\n  {pad('人', 8)} {pad('机器人', 10)} {pad('猜对的概率', 11, 'right')} {pad('对真实目标的平均信心', 21, 'right')} "
          f"{pad('自信地猜错', 11, 'right')} {pad('人的代价', 9, 'right')}")
    for hname, hkey in [("示范者", "demo"), ("教师", "teach")]:
        for rname, rkey in [("字面的", "lit"), ("务实的", "prag")]:
            right, conf, wrong, cost = evaluate(models[hkey], models[rkey])
            print(f"  {pad(hname, 8)} {pad(rname, 10)} {right:>11.3f} {conf:>21.3f} {wrong:>11.3f} {cost:>9.3f}")
    print("""
观察：
  · 示范者：目标是 θ1 还是 θ2，它都选 A（对自己最好）。机器人看到 A，最多只能猜“θ1 或 θ2 各一半”。
    人做得最好的事，不一定是最能让机器人学会的事。
  · 教师：目标是 θ2 时，它故意选 B —— 对自己略差（0.8 < 1.0），却是 θ2 独有的“信号”；
    目标是 θ1 时选 A。字面的机器人猜对的概率从 0.69 升到 0.95，代价只是人平均少拿 0.06 的价值。
  · 务实的机器人更进一步：它知道“如果人的目标是 θ2，人多半会选 B”，所以看到 A 就相当确信是 θ1
    （对真实目标的平均信心从 0.77 升到 0.92）。
  · Hadfield-Menell 等（2016）证明：在协助博弈里，人的最优策略一般不是“照常示范”，而是“教”；
    把逆强化学习（只是观察人怎样做事）换成协助博弈，人和机器人都会改变行为。""")


def demo_mismatch(models):
    print_header("演示 2：机器人以为人在教，人其实只是在示范")
    right, conf, wrong, _ = evaluate(models["demo"], models["prag"])
    r2, c2, w2, _ = evaluate(models["demo"], models["lit"])
    print(f"示范者 + 务实的机器人：猜对的概率 {right:.3f}，自信地猜错（≥ 80% 的信心放在错误目标上）的概率 {wrong:.3f}")
    print(f"示范者 + 字面的机器人：猜对的概率 {r2:.3f}，自信地猜错的概率 {w2:.3f}")
    print("""
观察：
  · 目标是 θ2 的示范者选了 A（对自己最好）。务实的机器人想：“如果是 θ2，人会选 B 来提示我；
    人选了 A，所以多半是 θ1。”—— 于是把约 88% 的信心放在了错误的目标上。字面的机器人看到 A 只给 θ1 约一半的
    信心，至少还知道自己不确定。两种机器人猜对的概率一样，差别在于“错的时候有多自信”。
  · 对人的行为“读得越深”，越依赖于对人的假设。假设错了，推断就会自信地出错。
    这和第 15 讲的“有限理性”是同一个问题：从行为推断偏好，必须假设人是怎样从偏好产生行为的。
  · 对大语言模型的启示：标注员写示范、写偏好时，可能在“完成任务”，也可能在“教模型”（例如故意写得更规范）；
    模型从数据里推断“人想要什么”时，隐含着对这一点的假设。""")


def main():
    models = build_models()
    demo_four_combos(models)
    demo_mismatch(models)


if __name__ == "__main__":
    main()