{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DQN 改进：Double DQN / Dueling DQN / 优先经验回放（PER），CartPole-v1。\n",
    "\n",
    "用法：python dqn_variants.py --double --dueling --per --seed 0\n",
    "\n",
    "配套教程：[DQN 改进：Double / Dueling / PER / Rainbow](https://dqn.cn/advanced/dqn-improvements)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/dqn_variants.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-dqn_variants-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": "advanced-dqn_variants-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import random\n",
    "\n",
    "import gymnasium as gym\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument(\"--double\", action=\"store_true\")\n",
    "parser.add_argument(\"--dueling\", action=\"store_true\")\n",
    "parser.add_argument(\"--per\", action=\"store_true\")\n",
    "parser.add_argument(\"--seed\", type=int, default=0)\n",
    "parser.add_argument(\"--steps\", type=int, default=50_000)\n",
    "args = parser.parse_args()\n",
    "\n",
    "GAMMA, LR, BATCH, BUFFER, START, TAU = 0.99, 5e-4, 64, 50_000, 1_000, 0.005\n",
    "EPS_START, EPS_END, EPS_DECAY = 1.0, 0.05, 10_000\n",
    "PER_ALPHA, PER_BETA0, PER_EPS = 0.6, 0.4, 1e-6\n",
    "\n",
    "random.seed(args.seed); np.random.seed(args.seed); torch.manual_seed(args.seed)\n",
    "env = gym.make(\"CartPole-v1\")\n",
    "obs_dim, n_actions = env.observation_space.shape[0], env.action_space.n"
   ],
   "id": "advanced-dqn_variants-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QNet(nn.Module):\n",
    "    def __init__(self, dueling):\n",
    "        super().__init__()\n",
    "        self.dueling = dueling\n",
    "        self.body = nn.Sequential(nn.Linear(obs_dim, 128), nn.ReLU(), nn.Linear(128, 128), nn.ReLU())\n",
    "        if dueling:\n",
    "            self.value = nn.Linear(128, 1)              # V(s)\n",
    "            self.adv = nn.Linear(128, n_actions)        # A(s, a)\n",
    "        else:\n",
    "            self.head = nn.Linear(128, n_actions)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = self.body(x)\n",
    "        if not self.dueling:\n",
    "            return self.head(h)\n",
    "        a = self.adv(h)\n",
    "        return self.value(h) + a - a.mean(dim=1, keepdim=True)   # Q = V + A - mean(A)"
   ],
   "id": "advanced-dqn_variants-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SumTree:\n",
    "    \"\"\"优先经验回放用的求和树：O(log N) 按优先级采样与更新。\"\"\"\n",
    "\n",
    "    def __init__(self, capacity):\n",
    "        self.capacity = capacity\n",
    "        self.tree = np.zeros(2 * capacity)\n",
    "        self.data = [None] * capacity\n",
    "        self.ptr, self.size = 0, 0\n",
    "\n",
    "    def add(self, priority, item):\n",
    "        self.data[self.ptr] = item\n",
    "        self.update(self.ptr, priority)\n",
    "        self.ptr = (self.ptr + 1) % self.capacity\n",
    "        self.size = min(self.size + 1, self.capacity)\n",
    "\n",
    "    def update(self, idx, priority):\n",
    "        i = idx + self.capacity\n",
    "        self.tree[i] = priority\n",
    "        i //= 2\n",
    "        while i >= 1:\n",
    "            self.tree[i] = self.tree[2 * i] + self.tree[2 * i + 1]\n",
    "            i //= 2\n",
    "\n",
    "    def find(self, value):\n",
    "        i = 1\n",
    "        while i < self.capacity:\n",
    "            left = 2 * i\n",
    "            if value <= self.tree[left]:\n",
    "                i = left\n",
    "            else:\n",
    "                value -= self.tree[left]\n",
    "                i = left + 1\n",
    "        return i - self.capacity\n",
    "\n",
    "    def total(self):\n",
    "        return self.tree[1]"
   ],
   "id": "advanced-dqn_variants-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_net, target_net = QNet(args.dueling), QNet(args.dueling)\n",
    "target_net.load_state_dict(q_net.state_dict())\n",
    "opt = torch.optim.Adam(q_net.parameters(), lr=LR)\n",
    "tree = SumTree(BUFFER)\n",
    "max_priority = 1.0\n",
    "\n",
    "obs, _ = env.reset(seed=args.seed)\n",
    "ep_ret, returns = 0.0, []\n",
    "for step in range(args.steps):\n",
    "    eps = EPS_START + min(step / EPS_DECAY, 1.0) * (EPS_END - EPS_START)\n",
    "    if random.random() < eps:\n",
    "        action = env.action_space.sample()\n",
    "    else:\n",
    "        with torch.no_grad():\n",
    "            action = int(q_net(torch.as_tensor(obs, dtype=torch.float32)[None]).argmax())\n",
    "    next_obs, r, term, trunc, _ = env.step(action)\n",
    "    # 新样本用当前最大优先级加入，保证至少被采样一次\n",
    "    tree.add(max_priority if args.per else 1.0, (obs, action, r, next_obs, float(term)))\n",
    "    obs, ep_ret = next_obs, ep_ret + r\n",
    "    if term or trunc:\n",
    "        returns.append(ep_ret)\n",
    "        obs, _ = env.reset()\n",
    "        ep_ret = 0.0\n",
    "\n",
    "    if step >= START:\n",
    "        # ---- 采样：均匀 或 按优先级 ----\n",
    "        if args.per:\n",
    "            seg = tree.total() / BATCH\n",
    "            idxs = [tree.find(random.uniform(seg * i, seg * (i + 1))) for i in range(BATCH)]\n",
    "            probs = np.array([tree.tree[i + BUFFER] for i in idxs]) / tree.total()\n",
    "            beta = PER_BETA0 + (1.0 - PER_BETA0) * step / args.steps   # β 线性退火到 1\n",
    "            w = (tree.size * probs) ** (-beta)\n",
    "            w = torch.as_tensor(w / w.max(), dtype=torch.float32)       # 重要性采样权重\n",
    "        else:\n",
    "            idxs = np.random.randint(0, tree.size, size=BATCH)\n",
    "            w = torch.ones(BATCH)\n",
    "        s, a, r_, s2, d = map(np.array, zip(*[tree.data[i] for i in idxs]))\n",
    "        s, s2 = torch.as_tensor(s, dtype=torch.float32), torch.as_tensor(s2, dtype=torch.float32)\n",
    "        a = torch.as_tensor(a).long().unsqueeze(1)\n",
    "        r_, d = torch.as_tensor(r_, dtype=torch.float32), torch.as_tensor(d, dtype=torch.float32)\n",
    "\n",
    "        q_sa = q_net(s).gather(1, a).squeeze(1)\n",
    "        with torch.no_grad():\n",
    "            if args.double:   # 在线网络选动作，目标网络评估\n",
    "                a2 = q_net(s2).argmax(dim=1, keepdim=True)\n",
    "                q_next = target_net(s2).gather(1, a2).squeeze(1)\n",
    "            else:\n",
    "                q_next = target_net(s2).max(dim=1).values\n",
    "            target = r_ + GAMMA * (1 - d) * q_next\n",
    "        td = target - q_sa\n",
    "        loss = (w * F.smooth_l1_loss(q_sa, target, reduction=\"none\")).mean()\n",
    "        opt.zero_grad(); loss.backward()\n",
    "        nn.utils.clip_grad_norm_(q_net.parameters(), 10.0)\n",
    "        opt.step()\n",
    "\n",
    "        if args.per:\n",
    "            new_p = (td.abs().detach().numpy() + PER_EPS) ** PER_ALPHA\n",
    "            for i, p in zip(idxs, new_p):\n",
    "                tree.update(i, p)\n",
    "            max_priority = max(max_priority, new_p.max())\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for p, tp in zip(q_net.parameters(), target_net.parameters()):\n",
    "                tp.mul_(1 - TAU).add_(TAU * p)"
   ],
   "id": "advanced-dqn_variants-py-5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(n=10):\n",
    "    e = gym.make(\"CartPole-v1\"); total = 0.0\n",
    "    for i in range(n):\n",
    "        o, _ = e.reset(seed=10_000 + i); done = False\n",
    "        while not done:\n",
    "            with torch.no_grad():\n",
    "                o, r, te, tr, _ = e.step(int(q_net(torch.as_tensor(o, dtype=torch.float32)[None]).argmax()))\n",
    "            total += r; done = te or tr\n",
    "    return total / n"
   ],
   "id": "advanced-dqn_variants-py-6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "name = \"+\".join([k for k in [\"double\", \"dueling\", \"per\"] if getattr(args, k)]) or \"vanilla\"\n",
    "print(f\"{name:20s} seed={args.seed} 最后10回合训练回报 {np.mean(returns[-10:]):6.1f} | 贪心评估(10回合) {evaluate():6.1f}\")"
   ],
   "id": "advanced-dqn_variants-py-7"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}