{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gymnasium 上手：FrozenLake 表格型 Q-learning，以及 CartPole 随机策略基线。\n",
    "\n",
    "\n",
    "\n",
    "配套教程：[Gymnasium 上手：CartPole 与 FrozenLake](https://dqn.cn/basics/gymnasium)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/basics/gym_frozenlake_q.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "basics-gym_frozenlake_q-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-gym_frozenlake_q-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "import numpy as np\n",
    "\n",
    "# ---------- 1. FrozenLake：表格型 Q-learning ----------\n",
    "env = gym.make(\"FrozenLake-v1\", is_slippery=True)\n",
    "nS, nA = env.observation_space.n, env.action_space.n\n",
    "Q = np.zeros((nS, nA))\n",
    "alpha, gamma = 0.1, 0.99\n",
    "eps, eps_min, eps_decay = 1.0, 0.05, 0.9995\n",
    "rng = np.random.default_rng(0)\n",
    "\n",
    "for ep in range(20_000):\n",
    "    s, _ = env.reset(seed=ep)\n",
    "    done = False\n",
    "    while not done:\n",
    "        a = int(rng.integers(nA)) if rng.random() < eps else int(np.argmax(Q[s]))\n",
    "        s2, r, terminated, truncated, _ = env.step(a)\n",
    "        target = r + gamma * np.max(Q[s2]) * (not terminated)\n",
    "        Q[s, a] += alpha * (target - Q[s, a])\n",
    "        s, done = s2, terminated or truncated\n",
    "    eps = max(eps_min, eps * eps_decay)\n",
    "\n",
    "wins = 0\n",
    "for ep in range(1000):\n",
    "    s, _ = env.reset(seed=100_000 + ep)\n",
    "    done = False\n",
    "    while not done:\n",
    "        s, r, terminated, truncated, _ = env.step(int(np.argmax(Q[s])))\n",
    "        done = terminated or truncated\n",
    "    wins += r\n",
    "print(f\"FrozenLake 贪心策略成功率（1000 回合）: {wins / 1000:.3f}\")\n",
    "\n",
    "# ---------- 2. CartPole：随机策略基线 ----------\n",
    "env = gym.make(\"CartPole-v1\")\n",
    "env.action_space.seed(0)                     # action_space 有独立的随机数生成器，需单独设种子\n",
    "returns = []\n",
    "for ep in range(100):\n",
    "    obs, _ = env.reset(seed=ep)\n",
    "    G, done = 0.0, False\n",
    "    while not done:\n",
    "        obs, r, terminated, truncated, _ = env.step(env.action_space.sample())\n",
    "        G += r\n",
    "        done = terminated or truncated\n",
    "    returns.append(G)\n",
    "print(f\"CartPole 随机策略：平均回报 {np.mean(returns):.1f} ± {np.std(returns):.1f}（100 回合）\")"
   ],
   "id": "basics-gym_frozenlake_q-py-2"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}