{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# n 步 SARSA（on-policy，ε-贪心）在 CliffWalking-v1 上的对比：n = 1, 3, 5。\n",
    "\n",
    "n 步目标：G_{t:t+n} = r_{t+1} + γ r_{t+2} + ... + γ^{n-1} r_{t+n} + γ^n Q(s_{t+n}, a_{t+n})。\n",
    "\n",
    "配套教程：[Q-learning 与 SARSA](https://dqn.cn/basics/q-learning-sarsa)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/basics/nstep_sarsa_cliff.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "basics-nstep_sarsa_cliff-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-nstep_sarsa_cliff-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(\"CliffWalking-v1\")\n",
    "nS, nA = env.observation_space.n, env.action_space.n\n",
    "EPISODES, SEEDS, GAMMA, EPS = 500, 10, 1.0, 0.1"
   ],
   "id": "basics-nstep_sarsa_cliff-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eps_greedy(Q, s, eps, rng):\n",
    "    if rng.random() < eps:\n",
    "        return int(rng.integers(nA))\n",
    "    q = Q[s]\n",
    "    return int(rng.choice(np.flatnonzero(q == q.max())))"
   ],
   "id": "basics-nstep_sarsa_cliff-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nstep_sarsa(n, alpha, seed):\n",
    "    rng = np.random.default_rng(seed)\n",
    "    Q = np.zeros((nS, nA)); returns = []\n",
    "    for ep in range(EPISODES):\n",
    "        s, _ = env.reset(seed=seed * 10_000 + ep)\n",
    "        S, A, R = [s], [eps_greedy(Q, s, EPS, rng)], [0.0]      # R[0] 占位，使 R[t+1] 对应第 t 步的奖励\n",
    "        T, t = np.inf, 0\n",
    "        while True:\n",
    "            if t < T:\n",
    "                s2, r, term, trunc, _ = env.step(A[t])\n",
    "                S.append(s2); R.append(r)\n",
    "                if term or trunc:\n",
    "                    T = t + 1\n",
    "                else:\n",
    "                    A.append(eps_greedy(Q, s2, EPS, rng))\n",
    "            tau = t - n + 1                                     # 本步要更新的时刻\n",
    "            if tau >= 0:\n",
    "                G = sum(GAMMA ** (i - tau - 1) * R[i] for i in range(tau + 1, int(min(tau + n, T)) + 1))\n",
    "                if tau + n < T:\n",
    "                    G += GAMMA ** n * Q[S[tau + n], A[tau + n]]\n",
    "                Q[S[tau], A[tau]] += alpha * (G - Q[S[tau], A[tau]])\n",
    "            if tau == T - 1:\n",
    "                break\n",
    "            t += 1\n",
    "        returns.append(sum(R))\n",
    "    return np.array(returns)"
   ],
   "id": "basics-nstep_sarsa_cliff-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    print(f\"CliffWalking，ε={EPS}，{EPISODES} 回合 × {SEEDS} 个种子；回报越接近 0 越好（最优路径为 -13）\")\n",
    "    for n in (1, 3, 5):\n",
    "        for alpha in (0.1, 0.5):\n",
    "            R = np.stack([nstep_sarsa(n, alpha, sd) for sd in range(SEEDS)])\n",
    "            print(f\"n={n}  α={alpha:.1f}  前50回合平均回报 {R[:, :50].mean():7.1f}  \"\n",
    "                  f\"最后100回合平均回报 {R[:, -100:].mean():6.1f}\")"
   ],
   "id": "basics-nstep_sarsa_cliff-py-5"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}