"""
=============================================================================
第 15 讲 · 代码 1：Dyna-Q——用学到的模型"想象"经验，提高样本效率
=============================================================================

对应讲义：lectures/lecture15_model_based_rl.md 第 1–2 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
Q-learning 每一条真实经验只用一次，然后就扔掉了。在真实环境中收集经验很昂贵（机器人会磨损、病人不能重来），
能不能让每一条经验发挥更大的作用？

Dyna 的想法：
    - 一边与环境交互，一边**记住**环境是怎样运作的（学一个模型：在 s 做 a 会得到 r、到达 s'）；
    - 每走一步真实的，就在脑子里"想象" n 步：从模型中取出以前的经验，再做 n 次 Q-learning 更新。
想象不花真实的代价，只花计算。

本文件（Sutton & Barto 例 8.1 的 Dyna 迷宫，6×9，最短路径 14 步）：
    1. 学习曲线：每个回合需要多少步？n = 0（普通 Q-learning）、5、50
    2. 为什么规划这么有效：前两个回合之后的贪心策略
    3. 真实经验 vs 计算量：n 越大，需要的真实经验越少，但计算越多

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    cd code/lecture15
    python 01_dyna_q.py      （约 20 秒）
"""

import numpy as np

from maze import DYNA_MAZE, run_dyna

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


def main():
    maze = DYNA_MAZE

    # ---------------------------------------------------------------------
    print("=" * 78)
    print("演示 1：学习曲线（每个回合用了多少步；30 个随机种子的平均；α = 0.1，ε = 0.1，γ = 0.95）")
    print("=" * 78)
    print(maze.render())
    print(f"  S = 起点，G = 终点，█ = 墙。最短路径 {maze.shortest_path()} 步（ε-贪心会多走几步）。")
    show = [1, 2, 3, 5, 10, 20, 50]
    print("\n   规划步数 n  " + "".join(f"  第{e:2d}回合" for e in show) + "    50 个回合的真实步数总和")
    for n in [0, 5, 50]:
        L = np.array([run_dyna(lambda t: maze, n, np.random.default_rng(sd), n_episodes=50)["episode_lengths"]
                      for sd in range(30)]).mean(axis=0)
        print(f"   {n:6d}      " + "".join(f"{L[e - 1]:9.1f}" for e in show) + f"        {L.sum():8.0f}")
    print("  观察：")
    print("  - 第 1 个回合，三者都要走七八百步：还不知道终点在哪里，只能随机乱走（有没有模型都一样）。")
    print("  - 之后，n = 50 只需 2～3 个回合就接近最短路径；n = 5 大约 5 个回合；")
    print("    普通 Q-learning（n = 0）要 20 个回合以上。这与 Sutton & Barto 图 8.2 的结果一致。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 2：为什么规划这么有效？——第 1 个回合结束、第 2 个回合刚走了 10 步时的贪心策略（· = Q 全为 0）")
    print("=" * 78)
    for n in [0, 50]:
        rng = np.random.default_rng(3)
        out = run_dyna(lambda t: maze, n, rng, n_episodes=1)
        # 用同一个随机种子重新跑一遍，这次在"第 1 个回合的步数 + 10 步"时停下，看看此时的 Q 表
        rng = np.random.default_rng(3)
        first_len = out["episode_lengths"][0]
        out2 = run_dyna(lambda t: maze, n, rng, max_steps=first_len + 10)
        print(f"\n  n = {n}（第 1 个回合用了 {first_len} 步）：")
        print("    " + maze.render(out2["Q"]).replace("\n", "\n    "))
    print("\n  观察：")
    print("  - n = 0：第 1 个回合结束时，只有\"终点旁边那一格\"的 Q 值被更新了（只有它的下一步得到了奖励）。")
    print("    奖励的信息每个回合只能往回传播一格——要传回起点，需要很多个回合。")
    print("  - n = 50：同样只走了一个回合，但规划把终点的奖励沿着模型中记住的转移反复地往回传播，")
    print("    一大片格子都已经知道该往哪走了。**真实经验只发生一次，想象可以重复很多次。**")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 3：真实经验 vs 计算量（30 个随机种子的平均，50 个回合）")
    print("=" * 78)
    print("   规划步数 n    真实步数总和    规划更新总次数    总更新次数（真实 + 规划）")
    for n in [0, 1, 5, 10, 50, 200]:
        real, plan = [], []
        for sd in range(30):
            out = run_dyna(lambda t: maze, n, np.random.default_rng(sd), n_episodes=50)
            real.append(sum(out["episode_lengths"]))
            plan.append(out["planning_updates"])
        print(f"   {n:6d}         {np.mean(real):9.0f}        {np.mean(plan):11.0f}          {np.mean(real) + np.mean(plan):11.0f}")
    print("  观察：")
    print("  - n 越大，需要的**真实**经验越少——但收益递减：真实步数主要花在第 1 个回合的\"找终点\"上，")
    print("    这部分是模型帮不上忙的（没见过的地方，模型里也没有）。")
    print("  - 总计算量随 n 迅速增长。基于模型的方法用**计算**换**数据**：")
    print("    当真实经验很贵（机器人、医疗）、而计算很便宜时，这笔交易非常划算。")
    print("  - 注意：这个迷宫是确定性的，模型只需要记住\"最近一次看到的结果\"就是完全准确的。")
    print("    这时 Dyna-Q 的规划就相当于第 5 讲 DQN 中的**经验回放**——从记忆中反复取出旧经验来学习。")
    print("    真正的挑战在于模型**不准确**的时候（02、03 文件）。")


if __name__ == "__main__":
    main()