"""
=============================================================================
第 16 讲 · 02：机器人从自己的经验中改进 —— 模仿、过滤、优势条件化（Recap 式）
=============================================================================

对应讲义：lectures/lecture16_robotics_agents.md 第 5 节（Physical Intelligence 2025, π*0.6 与 Recap）

-----------------------------------------------------------------------------
【设定：做一杯咖啡】
-----------------------------------------------------------------------------
任务有 5 个阶段（取杯、磨豆、压粉、萃取、打奶）。每个阶段机器人选一种做法：
    稳：成功率高（0.97），但慢（2 个时间单位）；
    快：1 个时间单位，成功率**取决于阶段**：有的阶段快也很稳（0.98~0.99），有的阶段快就容易出错（0.55~0.60）；
    糙：0.7 个时间单位，成功率更低。
任何一个阶段失败，这杯咖啡就失败了。我们关心**吞吐量**：每 100 个时间单位做成几杯。
最好的策略大致是：在"快也稳"的阶段用快，在容易出错的阶段用稳（第一个阶段甚至可以用糙：那里失败了损失的时间很少）
—— 但机器人一开始并不知道哪些阶段是哪种。

数据的三个来源（和 Recap 一样）：
    · 示范：人遥控机器人做 100 次。人比较谨慎，但也不总是最优（稳 60%、快 30%、糙 10%，和阶段无关）。
    · 自主经验：机器人用当前的策略自己做，成功和失败都记录下来。
    · 纠正：机器人自己做时，人在旁边看着；看到它要用"糙"的做法，或者在容易出错的阶段用"快"，
      以 50% 的概率接管，改用"稳"（这一步记为专家的动作）。

四种从数据学策略的方法（都是表格式的：每个阶段，各种做法的出现次数归一化成概率）：
    1. 只模仿示范（BC）；
    2. 模仿所有数据：示范 + 自主经验 + 纠正，一视同仁；
    3. 过滤的模仿：只模仿**成功**的回合；
    4. 优势条件化（Recap 的核心）：先从所有数据估计每个阶段的价值 V(阶段) =
       "从这个阶段开始的平均回报"（回报 = 成功得 1 分，每个时间单位扣 0.1 分）；
       每个动作的优势 = 这一回合从这个阶段开始的实际回报 − V(阶段)。
       策略以"优势是否为正"为条件；部署时总是要求"好的那一类"，也就是只模仿优势为正的动作。
       失败的数据也有用：它让 V 更准确，并告诉模型哪些动作不好。
自我改进的循环：每一轮用当前的策略做 200 次（带纠正），加入数据，重新学习，共 5 轮。

所有数值都是**设定**的。演示的是"怎样利用混合质量的数据自我改进"的机制。

运行时间：约 1 秒（只用标准库）。
"""

import math
import random

STAGES = ("取杯", "磨豆", "压粉", "萃取", "打奶")
ACTIONS = ("稳", "快", "糙")
SUCCESS = {"稳": (0.97, 0.97, 0.97, 0.97, 0.97),
           "快": (0.99, 0.60, 0.99, 0.55, 0.98),
           "糙": (0.85, 0.35, 0.85, 0.30, 0.85)}
TIME = {"稳": 2.0, "快": 1.0, "糙": 0.7}
TIME_COST = 0.1             # 回报里每个时间单位扣多少分（≈ 最优吞吐量 / 100，让"回报高"和"吞吐量高"一致）
DEMO_MIX = (0.6, 0.3, 0.1)  # 人示范时三种做法的比例
RISKY = {1, 3}              # 人知道哪些阶段"快就容易出错"（用于纠正）
CORRECTION_P = 0.5


def sample(probs, rng):
    u, acc = rng.random(), 0.0
    for i, p in enumerate(probs):
        acc += p
        if u < acc:
            return i
    return len(probs) - 1


def run_episode(policy, rng, corrections=False):
    """
    做一杯咖啡。

    【参数】policy：每个阶段的动作概率 [[p稳, p快, p糙], ...]；corrections：人是否在旁边纠正
    【返回】(steps, success, total_time)：steps 是 [(阶段, 动作下标, 这一步之后剩余的时间), ...]
        —— 为了算"从这个阶段开始的回报"，先记录每一步的用时，回合结束后再整理。
    """
    steps, total, success = [], 0.0, True
    for s in range(len(STAGES)):
        a = sample(policy[s], rng)
        if corrections and (a == 2 or (a == 1 and s in RISKY)) and rng.random() < CORRECTION_P:
            a = 0                                                   # 人接管，改用"稳"
        name = ACTIONS[a]
        total += TIME[name]
        steps.append((s, a, TIME[name]))
        if rng.random() >= SUCCESS[name][s]:
            success = False
            break
    return steps, success, total


def returns_from_each_step(steps, success):
    """
    每一步"从这里开始的回报"：成功得 1 分（失败 0 分），减去从这一步到结束用的时间 × TIME_COST。

    【返回】和 steps 等长的列表
    """
    out, remaining = [], math.fsum(t for _, _, t in steps)
    for _, _, t in steps:
        out.append((1.0 if success else 0.0) - TIME_COST * remaining)
        remaining -= t
    return out


def fit_policy(episodes, method):
    """
    从数据学一个表格式的策略。

    【参数】episodes：[(steps, success), ...]；method："bc" / "filtered" / "advantage"
    【返回】每个阶段的动作概率
    【思路】
        bc        ：每个阶段，统计所有数据里三种动作出现的次数；
        filtered  ：只统计成功的回合；
        advantage ：先算 V(阶段) = 所有数据里"从这个阶段开始的回报"的平均；再只统计优势 > 0 的动作。
    【边界】某个阶段没有任何被统计的动作（例如所有回合都在前面失败了），退回均匀分布。
        每个计数都加 0.5 的平滑，让策略保留一点探索。
    """
    counts = [[0.5] * len(ACTIONS) for _ in STAGES]
    if method == "advantage":
        sums = [[] for _ in STAGES]
        for steps, success in episodes:
            for (s, _, _), g in zip(steps, returns_from_each_step(steps, success)):
                sums[s].append(g)
        value = [math.fsum(x) / len(x) if x else 0.0 for x in sums]
    for steps, success in episodes:
        if method == "filtered" and not success:
            continue
        gs = returns_from_each_step(steps, success)
        for (s, a, _), g in zip(steps, gs):
            if method == "advantage" and g - value[s] <= 0:
                continue
            counts[s][a] += 1
    return [[c / math.fsum(row) for c in row] for row in counts]


def evaluate(policy, rng, n=3000):
    """不带纠正地做 n 次，返回 (成功率, 每 100 个时间单位做成几杯)。"""
    wins, time = 0, 0.0
    for _ in range(n):
        _, success, t = run_episode(policy, rng)
        wins += success
        time += t
    return wins / n, 100 * wins / time


def best_policy():
    """上帝视角的最优策略（每个阶段选使总吞吐量最高的做法），只用来做参照。"""
    best, best_tp = None, -1.0
    for combo in range(len(ACTIONS) ** len(STAGES)):
        choice, c = [], combo
        for _ in STAGES:
            choice.append(c % len(ACTIONS))
            c //= len(ACTIONS)
        p_ok, exp_time, alive = 1.0, 0.0, 1.0
        for s, a in enumerate(choice):
            exp_time += alive * TIME[ACTIONS[a]]
            alive *= SUCCESS[ACTIONS[a]][s]
        tp = 100 * alive / exp_time
        if tp > best_tp:
            best, best_tp = choice, tp
    return best, best_tp


def main():
    rng = random.Random(1610)
    demos = []
    demo_policy = [list(DEMO_MIX) for _ in STAGES]
    for _ in range(100):
        steps, success, _ = run_episode(demo_policy, rng)
        demos.append((steps, success))

    best, best_tp = best_policy()
    print("=" * 78)
    print("最优策略（上帝视角）：" + "、".join(f"{STAGES[s]}{ACTIONS[a]}" for s, a in enumerate(best))
          + f"，吞吐量 {best_tp:.2f} 杯 / 100 时间单位")
    print("=" * 78)
    eval_rng = random.Random(1611)
    sr, tp = evaluate(demo_policy, eval_rng)
    print(f"人的示范本身：成功率 {sr:.3f}，吞吐量 {tp:.2f}")

    for corrections in (True, False):
        print("\n" + "=" * 78)
        print("自我改进的循环（每轮自主做 200 次" + ("，人在旁边纠正）" if corrections else "，没有人纠正）"))
        print("=" * 78)
        print(f"{'':26}" + "".join(f"{'第 ' + str(r) + ' 轮':>14}" for r in range(6)))
        methods = (("模仿所有数据", "bc"), ("过滤的模仿（只学成功的）", "filtered"), ("优势条件化（Recap 式）", "advantage"))
        for name, method in methods:
            loop_rng = random.Random(1620 + len(method) + (100 if corrections else 0))
            data = list(demos)
            policy = fit_policy(data, "bc")                      # 第 0 轮：只模仿示范
            cells = []
            for r in range(6):
                sr, tp = evaluate(policy, eval_rng)
                cells.append(f"{sr:.2f} / {tp:5.2f}")
                if r == 5:
                    break
                for _ in range(200):
                    steps, success, _ = run_episode(policy, loop_rng, corrections)
                    data.append((steps, success))
                policy = fit_policy(data, method)
            width = sum(2 if ord(ch) > 127 else 1 for ch in name)
            print(name + " " * max(1, 26 - width) + "".join(f"{c:>14}" for c in cells))
        print("（每格：成功率 / 吞吐量）")

    final = fit_policy(data, "advantage")
    print("\n最后一轮优势条件化的策略（没有纠正的那组）：")
    for s, row in enumerate(final):
        print(f"  {STAGES[s]}：" + "  ".join(f"{ACTIONS[a]} {p:.2f}" for a, p in enumerate(row)))

    print("\n观察：")
    print("  · 第 0 轮三种方法用的是同一个策略（只模仿示范），几个数字的差别是评估的随机波动（约 ±0.1）。")
    print("  · 不加区分地模仿所有数据：没有纠正时原地踏步 —— 模仿自己的数据，学到的还是自己；")
    print("    有纠正时缓慢变好，因为纠正把\"糙\"和\"在危险的阶段求快\"换成了\"稳\"。")
    print("  · 过滤的模仿（只学成功的回合）成功率提高得快，但吞吐量提高有限：它只看\"成没成\"，不看\"快不快\"，")
    print("    成功的回合里大量是又稳又慢的动作。")
    print("  · 优势条件化同时利用了成功、失败和用时，吞吐量提高得最多；学到的策略在\"快也稳\"的打奶阶段")
    print("    几乎总是用快，在容易出错的磨豆、萃取阶段主要用稳。失败的数据也有用：它让价值估计更准。")
    print("  · 离最优（13.74）还有距离：表格式的策略保留着早期示范的影子，每轮只加 200 次数据，改进是渐进的。")
    print("    Recap 报告，在最难的任务上吞吐量提高一倍以上、失败率大约减半。")
    print("  · 人的纠正对所有方法都有帮助：它在关键时刻提供\"正确的做法\"，正是自主经验里最缺的数据。\n")


if __name__ == "__main__":
    main()