{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# DQN：CartPole-v1 上的最小可运行实现（经验回放 + 目标网络软更新）。\n", "\n", "\n", "\n", "配套教程：[DQN：深度 Q 网络](https://dqn.cn/advanced/dqn)（DQN.CN 强化学习中文教程）\n", "\n", "源码：`examples/advanced/dqn_cartpole.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。" ], "id": "advanced-dqn_cartpole-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_cartpole-py-1" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import random\n", "from collections import deque\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", "# ---------- 超参数 ----------\n", "GAMMA = 0.99\n", "LR = 5e-4\n", "BATCH_SIZE = 128\n", "BUFFER_SIZE = 50_000\n", "LEARNING_STARTS = 1_000\n", "TAU = 0.005 # 目标网络软更新系数\n", "EPS_START, EPS_END, EPS_DECAY_STEPS = 1.0, 0.05, 10_000\n", "TOTAL_STEPS = 50_000\n", "\n", "random.seed(0); np.random.seed(0); torch.manual_seed(0)\n", "env = gym.make(\"CartPole-v1\")\n", "obs_dim = env.observation_space.shape[0]\n", "n_actions = env.action_space.n\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")" ], "id": "advanced-dqn_cartpole-py-2" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class QNet(nn.Module):\n", " def __init__(self, obs_dim, n_actions):\n", " super().__init__()\n", " self.net = nn.Sequential(\n", " nn.Linear(obs_dim, 128), nn.ReLU(),\n", " nn.Linear(128, 128), nn.ReLU(),\n", " nn.Linear(128, n_actions),\n", " )\n", "\n", " def forward(self, x):\n", " return self.net(x)" ], "id": "advanced-dqn_cartpole-py-3" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "q_net = QNet(obs_dim, n_actions).to(device)\n", "target_net = QNet(obs_dim, n_actions).to(device)\n", "target_net.load_state_dict(q_net.state_dict())\n", "optimizer = torch.optim.Adam(q_net.parameters(), lr=LR)\n", "buffer = deque(maxlen=BUFFER_SIZE)" ], "id": "advanced-dqn_cartpole-py-4" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def epsilon(step):\n", " frac = min(step / EPS_DECAY_STEPS, 1.0)\n", " return EPS_START + frac * (EPS_END - EPS_START)" ], "id": "advanced-dqn_cartpole-py-5" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "obs, _ = env.reset(seed=0)\n", "env.action_space.seed(0) # 动作空间的随机数生成器需要单独设置种子\n", "episode_return, returns = 0.0, []\n", "\n", "for step in range(TOTAL_STEPS):\n", " # 1) ε-贪心选择动作\n", " if random.random() < epsilon(step):\n", " action = env.action_space.sample()\n", " else:\n", " with torch.no_grad():\n", " q = q_net(torch.as_tensor(obs, dtype=torch.float32, device=device))\n", " action = int(q.argmax().item())\n", "\n", " # 2) 与环境交互，存入回放缓冲区\n", " next_obs, reward, terminated, truncated, _ = env.step(action)\n", " # 注意：只有真正终止（terminated）时才不自举；被截断（truncated）时仍应自举\n", " buffer.append((obs, action, reward, next_obs, float(terminated)))\n", " obs = next_obs\n", " episode_return += reward\n", " if terminated or truncated:\n", " returns.append(episode_return)\n", " obs, _ = env.reset()\n", " episode_return = 0.0\n", "\n", " # 3) 从缓冲区采样并更新网络\n", " if step >= LEARNING_STARTS:\n", " batch = random.sample(buffer, BATCH_SIZE)\n", " s, a, r, s2, d = map(np.array, zip(*batch))\n", " s = torch.as_tensor(s, dtype=torch.float32, device=device)\n", " a = torch.as_tensor(a, dtype=torch.int64, device=device).unsqueeze(1)\n", " r = torch.as_tensor(r, dtype=torch.float32, device=device)\n", " s2 = torch.as_tensor(s2, dtype=torch.float32, device=device)\n", " d = torch.as_tensor(d, dtype=torch.float32, device=device)\n", "\n", " q_sa = q_net(s).gather(1, a).squeeze(1)\n", " with torch.no_grad():\n", " target = r + GAMMA * (1 - d) * target_net(s2).max(dim=1).values\n", " loss = F.smooth_l1_loss(q_sa, target) # Huber 损失\n", "\n", " optimizer.zero_grad()\n", " loss.backward()\n", " nn.utils.clip_grad_norm_(q_net.parameters(), 10.0)\n", " optimizer.step()\n", "\n", " # 4) 软更新目标网络：θ⁻ ← τθ + (1-τ)θ⁻\n", " if step >= LEARNING_STARTS:\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)\n", "\n", " if step % 5_000 == 0 and returns:\n", " print(f\"step {step:6d} | ε={epsilon(step):.2f} | 最近10回合平均回报 {np.mean(returns[-10:]):.1f}\")\n", "\n", "env.close()\n", "\n", "# ---------- 训练结束：贪心策略评估 10 个回合 ----------\n", "eval_env = gym.make(\"CartPole-v1\")\n", "scores = []\n", "for i in range(10):\n", " o, _ = eval_env.reset(seed=10_000 + i)\n", " done, total = False, 0.0\n", " while not done:\n", " with torch.no_grad():\n", " a = int(q_net(torch.as_tensor(o, dtype=torch.float32, device=device)).argmax())\n", " o, r, te, tr, _ = eval_env.step(a)\n", " total += r; done = te or tr\n", " scores.append(total)\n", "print(f\"贪心评估（10 回合）平均回报: {np.mean(scores):.1f}\")" ], "id": "advanced-dqn_cartpole-py-6" } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python" }, "colab": { "provenance": [] } }, "nbformat": 4, "nbformat_minor": 5 }