"""
=============================================================================
第 6 讲 · 代码 2：从经验中学模型（model-based RL）—— 以及"只贪心不探索"会怎样
=============================================================================

对应讲义：lectures/lecture06_rl_control.md 第 2.4–2.5 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
代码 01 假设我们知道 P 和 R。现实中通常不知道，CS229 讲义给的办法最朴素：
    "数一数"——用经验估计 P̂_{sa}(s') = #(s,a→s') / #(s,a)，再在 P̂ 上做价值迭代。

    演示 1：在 4×3 网格里随机乱走 1、3、10 … 1000 个回合，学模型、规划、再拿到**真实**环境里评估。
            要多少数据才够？模型对自己的计划有多"自信"（规划出的价值 vs 真实价值）？
    演示 2：探索问题（exploration）。一条走廊：近处有 +1 的出口，远处有 +10 的出口。
            "学模型 → 贪心地执行计划"的 agent 会怎样？ε-greedy 能救吗？"乐观"（optimism）呢？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_model_based_rl.py          # 约 10 秒

需要 numpy。依赖本目录的 rl_tools.py。
"""

import numpy as np

from rl_tools import (estimate_model, greedy_policy, make_corridor, make_ng_grid, policy_evaluation,
                      policy_iteration, print_header, sample_episode, value_iteration)


# =============================================================================
# 演示 1：随机探索 → 学模型 → 规划
# =============================================================================
def demo_learn_model():
    print_header("演示 1：随机探索 → 估计 P̂、R̂ → 在 P̂ 上规划 → 在真实环境里评估（4x3 网格，γ = 0.99）")
    mdp = make_ng_grid()
    gamma = 0.99
    V_star, _ = value_iteration(mdp.P, mdp.R, gamma, tol=1e-12)
    live = ~mdp.terminal
    rng = np.random.default_rng(21)
    random_policy = lambda s, r: int(r.integers(mdp.A))          # 均匀随机选动作（lambda = 一行的小函数）
    reps = 50
    print(f"每一行重复 {reps} 次（每次重新采样），表中是平均值。V*(start) = {V_star[mdp.start]:.3f}。")
    print("V_hat = 学到的模型认为自己的计划值多少；V_true = 同一个计划在真实环境里实际值多少。\n")
    print(f"{'episodes':>9}{'transitions':>13}{'(s,a) tried':>13}{'TV error':>10}"
          f"{'V_hat(start)':>14}{'V_true(start)':>15}{'regret':>9}{'P(optimal)':>12}")
    summary = {}
    for n_ep in [1, 3, 10, 30, 100, 1000]:
        rows = []
        for _ in range(reps):
            trans, rews = [], []
            for _ in range(n_ep):
                tr, rw = sample_episode(mdp, random_policy, rng)
                trans += tr                                        # 列表相加 = 拼接
                rews += rw
            P_hat, R_hat, counts = estimate_model(trans, rews, mdp.S, mdp.A, mdp.terminal)
            V_hat, _ = value_iteration(P_hat, R_hat, gamma, tol=1e-10)
            pi_hat = greedy_policy(V_hat, P_hat, R_hat, gamma)    # 在**学到的**模型上贪心
            V_true = policy_evaluation(mdp.P, mdp.R, gamma, pi_hat)  # 在**真实**模型上评估同一个策略
            tried = counts[live] > 0                               # 非终止状态里，哪些 (s, a) 至少试过一次
            # total variation distance ½ Σ_s' |P̂ − P|：只在试过的 (s, a) 上算（没试过的是"均匀"，没有意义）
            tv = 0.5 * np.abs(P_hat[live] - mdp.P[live]).sum(axis=2)[tried]
            regret = V_star[mdp.start] - V_true[mdp.start]
            rows.append([len(trans), tried.mean(), tv.mean() if tv.size else np.nan,
                         V_hat[mdp.start], V_true[mdp.start], regret, regret < 1e-3])
        m = np.array(rows, dtype=float).mean(axis=0)
        summary[n_ep] = m
        print(f"{n_ep:>9}{m[0]:>13.0f}{m[1]:>13.2f}{m[2]:>10.3f}{m[3]:>14.3f}{m[4]:>15.3f}{m[5]:>9.3f}{m[6]:>12.2f}")
    gap3, gap10, gap100 = (summary[k][3] - summary[k][4] for k in (3, 10, 100))
    print(f"""
观察 Observations：
  · 数据越多，模型越准（TV error 从 {summary[1][2]:.2f} 降到 {summary[1000][2]:.3f}），计划越好（regret 从 {summary[1][5]:.2f} 降到
    {summary[1000][5]:.3f}）。1 个回合只试过 {summary[1][1]:.0%} 的 (s, a)；没试过的用"均匀分布"填，计划基本是瞎猜。
  · 不需要把模型学得很准才能得到好策略：100 个回合时 TV error 还有 {summary[100][2]:.3f}，P(optimal) 已经是
    {summary[100][6]:.2f}。和代码 01 演示 1 的"贪心策略比价值先收敛"是同一个道理：只要动作的排序对了就行。
  · 模型对自己的计划**过度自信**：V_hat − V_true = {gap3:+.3f}（3 个回合）、{gap10:+.3f}（10 个）、{gap100:+.3f}（100 个）。
    这是 optimizer's curse：在有误差的模型上取 max，会挑中那些"恰好被高估"的动作。数据越多，偏差越小。
    （CS234 第 15 讲的 03_model_error.py 专门讨论这个现象和模型误差的累积。）
  · 随机探索在这个小网格上够用：9 个非终止状态、36 个 (s, a)，乱走几十个回合就全试过了。
    状态一多、或者好东西藏在"随机走不到"的地方，随机探索就不行了 —— 看演示 2。""")


# =============================================================================
# 演示 2：探索问题 —— 近处 +1 vs 远处 +10
# =============================================================================
def plan(trans, rews, env, gamma, mode, r_max):
    """
    用至今为止的经验估计模型并规划，返回一个确定性策略（长度 S 的整数数组）。

    【参数】mode："model"（CS229 讲义的做法：没见过的 (s, a) 用均匀分布）或
        "optimistic"（R-max 式的乐观：没见过的 (s, a) 假设会通往一个"天堂"状态，之后每步都拿 r_max）
    【思路】乐观的实现：在 S 个真实状态之外加一个虚拟状态 heaven（下标 S），
        它的奖励是 r_max·(1 − γ)、永远停在自己身上，所以它的价值是 r_max。
        所有"没试过的 (s, a)"都以概率 1 通往 heaven。于是任何没试过的动作看起来都和 +10 出口一样好，
        agent 会主动去试 —— 试过之后它就变成"已知"的，乐观自动消失。
        规划用 policy_iteration（精确、几步就收敛，比 VI 快）。
    【边界】没有终止状态的 (s, a) 才需要乐观；终止状态的行保持全 0。
    """
    P_hat, R_hat, counts = estimate_model(trans, rews, env.S, env.A, env.terminal)
    if mode == "model":
        return policy_iteration(P_hat, R_hat, gamma)[0]
    S = env.S
    P2 = np.zeros((S + 1, env.A, S + 1))
    P2[:S, :, :S] = P_hat
    unseen = (counts == 0) & ~env.terminal[:, None]                # 没试过、且不是终止状态的 (s, a)
    P2[:S][unseen] = 0.0
    P2[:S, :, S][unseen] = 1.0                                     # 没试过 → 去 heaven
    P2[S, :, S] = 1.0                                              # heaven 永远停在自己身上
    R2 = np.append(R_hat, r_max * (1 - gamma))
    return policy_iteration(P2, R2, gamma)[0][:S]


def run_agent(env, gamma, mode, rng, n_episodes=30, eps=0.1, r_max=10.0, max_steps=100):
    """
    CS229 讲义的 model-based RL 循环：执行策略 → 更新 P̂、R̂ → 重新规划 → …

    【思路】第 0 个回合用均匀随机的动作（总得从某处开始）。之后每个回合都执行"当前计划"。
        每当试了一个**新的** (s, a)，立刻用全部经验重新估计模型、重新规划（R-max 的做法）。
        这个走廊是确定性的，同一个 (s, a) 第二次试不会带来新信息，所以这和"每一步都重新规划"结果完全一样，只是快得多。
    【参数】mode："greedy"（照计划走）、"egreedy"（以概率 eps 随机选动作）、"optimistic"（R-max 式乐观）
    【返回】dict：found_first（第 0 个随机回合是否到过 +10）、found_big（整个过程中是否到过 +10 出口）、goes_right（最终计划在起点是否向右）、
        last_returns（最后 10 个回合的平均折扣回报）、n_steps（总共走了多少步）
    """
    trans, rews = sample_episode(env, lambda s, r: int(r.integers(env.A)), rng, max_steps)
    big = env.S - 1
    found_first = any(s == big for s, _ in rews)                   # 第 0 个（随机）回合有没有碰巧走到 +10
    found_big = found_first
    plan_mode = "optimistic" if mode == "optimistic" else "model"
    returns = []
    for _ in range(n_episodes):
        pi = plan(trans, rews, env, gamma, plan_mode, r_max)
        tried = set((s, a) for s, a, _ in trans)                   # 集合（set）：查"试没试过"是 O(1)
        s = env.start
        rews.append((s, env.R[s]))
        ret, t = env.R[s], 0
        while not env.terminal[s] and t < max_steps:
            a = int(rng.integers(env.A)) if (mode == "egreedy" and rng.uniform() < eps) else int(pi[s])
            s_next = int(rng.choice(env.S, p=env.P[s, a]))
            trans.append((s, a, s_next))
            rews.append((s_next, env.R[s_next]))
            t += 1
            ret += gamma ** t * env.R[s_next]
            if (s, a) not in tried:                                # 新经验 → 重新规划
                tried.add((s, a))
                pi = plan(trans, rews, env, gamma, plan_mode, r_max)
            s = s_next
        found_big = found_big or s == big
        returns.append(ret)
    pi = plan(trans, rews, env, gamma, plan_mode, r_max)
    return {"found_first": found_first, "found_big": found_big, "goes_right": pi[env.start] == 1,
            "last_returns": float(np.mean(returns[-10:])), "n_steps": len(trans)}


def demo_exploration():
    print_header("演示 2：探索问题 —— 近处 +1、远处 +10，贪心地相信自己的模型会怎样？")
    n, gamma = 10, 0.95
    env = make_corridor(n, {0: 1.0, n - 1: 10.0}, start=2)
    V_star, _ = value_iteration(env.P, env.R, gamma, tol=1e-12)
    print(f"走廊：状态 0 … {n - 1}，从状态 2 出发，动作 左 / 右（确定性），γ = {gamma}，每步奖励 0。")
    print(f"状态 0 是 +1 出口（2 步），状态 {n - 1} 是 +10 出口（{n - 3} 步）。")
    print(f"最优：向右，V*(start) = {V_star[env.start]:.3f}；向左只能拿到 γ² × 1 = {gamma ** 2:.3f}。\n")
    agents = [("greedy (CS229 loop)", "greedy"), ("eps-greedy, eps=0.1", "egreedy"),
              ("optimistic (R-max)", "optimistic")]
    runs = 200
    print(f"每个 agent 独立跑 {runs} 次，每次 = 1 个随机回合 + 30 个按计划执行的回合。\n")
    print(f"{'agent':<22}{'random ep. hit +10':>20}{'found +10':>11}{'final plan: right':>19}{'mean return (last 10)':>23}{'total steps':>13}")
    res = {}
    for name, mode in agents:
        # 第 i 次运行用种子 (22, i)：三个 agent 的第 i 次运行有**相同的**第 0 个随机回合，差别只来自之后的探索策略。
        # 【numpy 小知识】default_rng 可以接收一个整数元组作种子，(22, 0)、(22, 1) … 给出互相独立的随机数流。
        out = [run_agent(env, gamma, mode, np.random.default_rng((22, i))) for i in range(runs)]
        r = {k: float(np.mean([o[k] for o in out])) for k in out[0]}
        res[mode] = r
        print(f"{name:<22}{r['found_first']:>20.3f}{r['found_big']:>11.3f}{r['goes_right']:>19.3f}{r['last_returns']:>23.3f}{r['n_steps']:>13.1f}")
    p_lucky = 2 / 9
    se = np.sqrt(p_lucky * (1 - p_lucky) / runs)                  # 二项比例的标准误
    extra = int(round((res['egreedy']['found_big'] - res['greedy']['found_big']) * runs))
    print(f"""
观察 Observations：
  · 第 0 个随机回合是从状态 2 出发的随机游走，先撞到 {n - 1} 而不是 0 的概率是 2/9 = {p_lucky:.3f}（赌徒破产问题）。
    greedy 找到 +10 的比例 {res['greedy']['found_big']:.3f} ≈ 2/9（{runs} 次运行的标准误约 {se:.2f}）：
    **只有第一个回合碰巧走到右边的那些运行学会了向右**。
    其余的运行里，模型只见过 +1，规划出的最优计划就是"向左"，执行"向左"又只会看到 +1 …… 永远不会知道右边有什么。
    这不是模型估计错了，而是**数据是 agent 自己的行为产生的**：不去试，就没有数据。
  · ε-greedy（ε = 0.1）没帮上忙：{runs} 次运行里，靠随机动作比 greedy 多找到 +10 的次数是 {extra}。
    要从状态 2 走到 {n - 1}，需要连续很多步"恰好"随机地向右，
    概率随距离指数下降。这种"抖动式"探索（dithering）在需要"深度探索"的问题上几乎无效。
  · 乐观（optimism in the face of uncertainty）：把没试过的动作当成"可能非常好"。agent 会**有计划地**去试每个没试过的
    (s, a)，试过之后乐观自动消失。{runs} 次全部找到 +10，最后 10 个回合的平均回报 {res['optimistic']['last_returns']:.3f}
    = V*(start)。代价是前期多走了一些路（total steps {res['optimistic']['n_steps']:.0f} vs {res['greedy']['n_steps']:.0f}）。
  · CS229 只点到为止。系统的探索理论（regret、UCB、Thompson sampling、R-max、PSRL）见 CS234 第 11–13 讲。""")


def main():
    demo_learn_model()
    demo_exploration()


if __name__ == "__main__":
    main()