{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dyna-Q：在 Sutton & Barto 第 8 章的迷宫（6×9）上比较不同规划步数 n。\n",
    "\n",
    "\n",
    "\n",
    "配套教程：[基于模型的强化学习](https://dqn.cn/advanced/model-based)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/dyna_q_maze.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-dyna_q_maze-py-0"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 环境准备：在 Colab 等缺少依赖的环境中自动安装（本地已装好依赖时什么也不做）\n",
    "import importlib.util, os, subprocess, sys\n",
    "sys.argv = sys.argv[:1]   # notebook 内核会带上 -f 等参数，清掉以免脚本把它们当成命令行参数\n",
    "# os.environ[\"DQN_QUICK\"] = \"1\"   # 取消注释：支持该开关的脚本只跑很短的版本，用于快速检查"
   ],
   "id": "advanced-dyna_q_maze-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "MAZE = [\n",
    "    \".......#G\",\n",
    "    \"..#....#.\",\n",
    "    \"S.#....#.\",\n",
    "    \"..#....#.\",\n",
    "    \".....#...\",\n",
    "    \".........\",\n",
    "]\n",
    "ROWS, COLS = len(MAZE), len(MAZE[0])\n",
    "START = next((r, c) for r in range(ROWS) for c in range(COLS) if MAZE[r][c] == \"S\")\n",
    "GOAL = next((r, c) for r in range(ROWS) for c in range(COLS) if MAZE[r][c] == \"G\")\n",
    "ACTIONS = [(-1, 0), (1, 0), (0, -1), (0, 1)]"
   ],
   "id": "advanced-dyna_q_maze-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def step(state, a):\n",
    "    r, c = state[0] + ACTIONS[a][0], state[1] + ACTIONS[a][1]\n",
    "    if not (0 <= r < ROWS and 0 <= c < COLS) or MAZE[r][c] == \"#\":\n",
    "        r, c = state                          # 撞墙则原地不动\n",
    "    if (r, c) == GOAL:\n",
    "        return (r, c), 1.0, True\n",
    "    return (r, c), 0.0, False"
   ],
   "id": "advanced-dyna_q_maze-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dyna_q(n_planning, episodes=50, alpha=0.1, gamma=0.95, eps=0.1, seed=0):\n",
    "    rng = np.random.default_rng(seed)\n",
    "    Q = np.zeros((ROWS, COLS, 4))\n",
    "    model = {}                                # 确定性模型：(s, a) -> (r, s', done)\n",
    "    steps_per_episode = []\n",
    "    for _ in range(episodes):\n",
    "        s, done, n_steps = START, False, 0\n",
    "        while not done:\n",
    "            q = Q[s]\n",
    "            a = int(rng.integers(4)) if rng.random() < eps else int(rng.choice(np.flatnonzero(q == q.max())))\n",
    "            s2, r, done = step(s, a)\n",
    "            # (a) 直接强化学习：用真实经验做一次 Q-learning 更新\n",
    "            Q[s][a] += alpha * (r + gamma * np.max(Q[s2]) * (not done) - Q[s][a])\n",
    "            # (b) 模型学习\n",
    "            model[(s, a)] = (r, s2, done)\n",
    "            # (c) 规划：从模型中随机抽取 n 个已见过的 (s, a) 做模拟更新\n",
    "            keys = list(model.keys())\n",
    "            for _ in range(n_planning):\n",
    "                ps, pa = keys[rng.integers(len(keys))]\n",
    "                pr, ps2, pdone = model[(ps, pa)]\n",
    "                Q[ps][pa] += alpha * (pr + gamma * np.max(Q[ps2]) * (not pdone) - Q[ps][pa])\n",
    "            s, n_steps = s2, n_steps + 1\n",
    "        steps_per_episode.append(n_steps)\n",
    "    return np.array(steps_per_episode)"
   ],
   "id": "advanced-dyna_q_maze-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in [0, 5, 50]:\n",
    "    curves = np.mean([dyna_q(n, seed=i) for i in range(30)], axis=0)\n",
    "    print(f\"n={n:2d} | 第2回合平均步数 {curves[1]:7.1f} | 第5回合 {curves[4]:6.1f} | 第50回合 {curves[-1]:5.1f}\")"
   ],
   "id": "advanced-dyna_q_maze-py-5"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}