{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 贝尔曼期望方程：用矩阵形式求解 FrozenLake 上均匀随机策略的状态价值，并与迭代法对比。\n",
    "\n",
    "\n",
    "\n",
    "配套教程：[贝尔曼方程](https://dqn.cn/basics/bellman)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/basics/bellman_solve.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "basics-bellman_solve-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-bellman_solve-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            # P[s][a] = [(prob, s_next, reward, terminated), ...]\n",
    "nS, nA = env.observation_space.n, env.action_space.n\n",
    "gamma = 0.9\n",
    "pi = np.full((nS, nA), 1.0 / nA)   # 均匀随机策略 π(a|s)\n",
    "\n",
    "# 构造 P_π（nS×nS）和 r_π（nS）\n",
    "P_pi = np.zeros((nS, nS))\n",
    "r_pi = np.zeros(nS)\n",
    "for s in range(nS):\n",
    "    for a in range(nA):\n",
    "        for prob, s2, r, done in P[s][a]:\n",
    "            P_pi[s, s2] += pi[s, a] * prob * (0.0 if done else 1.0)  # 终止后不再有未来价值\n",
    "            r_pi[s] += pi[s, a] * prob * r\n",
    "\n",
    "# 方法一：解析解 v = (I - γ P_π)^{-1} r_π\n",
    "v_exact = np.linalg.solve(np.eye(nS) - gamma * P_pi, r_pi)\n",
    "\n",
    "# 方法二：反复应用贝尔曼期望算子直到收敛\n",
    "v = np.zeros(nS)\n",
    "for it in range(1000):\n",
    "    v_new = r_pi + gamma * P_pi @ v\n",
    "    if np.max(np.abs(v_new - v)) < 1e-10:\n",
    "        break\n",
    "    v = v_new\n",
    "\n",
    "print(\"解析解：\\n\", v_exact.reshape(4, 4).round(4))\n",
    "print(f\"迭代 {it} 次后与解析解的最大误差: {np.max(np.abs(v - v_exact)):.2e}\")"
   ],
   "id": "basics-bellman_solve-py-2"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}