{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 动态规划：在 FrozenLake-v1（4x4，打滑）上实现策略迭代与价值迭代。\n",
    "\n",
    "\n",
    "\n",
    "配套教程：[动态规划](https://dqn.cn/basics/dynamic-programming)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/basics/dp.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "basics-dp-py-0"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 环境准备：在 Colab 等缺少依赖的环境中自动安装（本地已装好依赖时什么也不做）\n",
    "import importlib.util, os, subprocess, sys\n",
    "_need = [m for m in [\"gymnasium\"] if importlib.util.find_spec(m) is None]\n",
    "if _need or \"google.colab\" in sys.modules:\n",
    "    subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"gymnasium>=1.0\"], check=True)\n",
    "sys.argv = sys.argv[:1]   # notebook 内核会带上 -f 等参数，清掉以免脚本把它们当成命令行参数\n",
    "# os.environ[\"DQN_QUICK\"] = \"1\"   # 取消注释：支持该开关的脚本只跑很短的版本，用于快速检查"
   ],
   "id": "basics-dp-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "import numpy as np\n",
    "\n",
    "env = gym.make(\"FrozenLake-v1\", is_slippery=True)\n",
    "P = env.unwrapped.P\n",
    "nS, nA = env.observation_space.n, env.action_space.n\n",
    "gamma, theta = 0.99, 1e-8"
   ],
   "id": "basics-dp-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def q_from_v(v, s):\n",
    "    \"\"\"一步前瞻：q(s,a) = Σ p(s',r|s,a)[r + γ v(s')]\"\"\"\n",
    "    q = np.zeros(nA)\n",
    "    for a in range(nA):\n",
    "        for prob, s2, r, done in P[s][a]:\n",
    "            q[a] += prob * (r + gamma * v[s2] * (not done))\n",
    "    return q"
   ],
   "id": "basics-dp-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def policy_evaluation(policy, v=None):\n",
    "    v = np.zeros(nS) if v is None else v.copy()\n",
    "    while True:\n",
    "        delta = 0.0\n",
    "        for s in range(nS):\n",
    "            new = q_from_v(v, s)[policy[s]]       # 确定性策略\n",
    "            delta = max(delta, abs(new - v[s]))\n",
    "            v[s] = new                           # 原地更新（in-place）\n",
    "        if delta < theta:\n",
    "            return v"
   ],
   "id": "basics-dp-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def policy_iteration():\n",
    "    policy = np.zeros(nS, dtype=int)\n",
    "    v = np.zeros(nS)\n",
    "    for i in range(100):\n",
    "        v = policy_evaluation(policy, v)\n",
    "        new_policy = np.array([np.argmax(q_from_v(v, s)) for s in range(nS)])\n",
    "        if np.array_equal(new_policy, policy):\n",
    "            return policy, v, i + 1\n",
    "        policy = new_policy"
   ],
   "id": "basics-dp-py-5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def value_iteration():\n",
    "    v = np.zeros(nS)\n",
    "    for i in range(10_000):\n",
    "        delta = 0.0\n",
    "        for s in range(nS):\n",
    "            best = np.max(q_from_v(v, s))\n",
    "            delta = max(delta, abs(best - v[s]))\n",
    "            v[s] = best\n",
    "        if delta < theta:\n",
    "            break\n",
    "    policy = np.array([np.argmax(q_from_v(v, s)) for s in range(nS)])\n",
    "    return policy, v, i + 1"
   ],
   "id": "basics-dp-py-6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(policy, episodes=10_000, seed=0):\n",
    "    rng_env = gym.make(\"FrozenLake-v1\", is_slippery=True)\n",
    "    wins = 0\n",
    "    for ep in range(episodes):\n",
    "        s, _ = rng_env.reset(seed=seed + ep)\n",
    "        done = False\n",
    "        while not done:\n",
    "            s, r, term, trunc, _ = rng_env.step(int(policy[s]))\n",
    "            done = term or trunc\n",
    "        wins += r\n",
    "    return wins / episodes"
   ],
   "id": "basics-dp-py-7"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pi_pi, v_pi, n_pi = policy_iteration()\n",
    "pi_vi, v_vi, n_vi = value_iteration()\n",
    "arrows = np.array([\"←\", \"↓\", \"→\", \"↑\"])\n",
    "print(f\"策略迭代：{n_pi} 轮收敛；价值迭代：{n_vi} 次扫描收敛\")\n",
    "print(\"两种方法得到的策略是否一致：\", np.array_equal(pi_pi, pi_vi))\n",
    "print(\"最优策略：\\n\", arrows[pi_vi].reshape(4, 4))\n",
    "print(\"v*(起点) =\", round(v_vi[0], 4))\n",
    "print(f\"实际成功率（10000 回合）：{evaluate(pi_vi):.3f}\")"
   ],
   "id": "basics-dp-py-8"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}