{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DQN 的 Atari 设置：标准预处理 + Nature DQN 卷积网络（只检查形状并随机交互几步，不包含训练）。\n",
    "\n",
    "依赖：pip install \"gymnasium[atari]\" opencv-python   （ale-py 0.9+ 已自带游戏 ROM）\n",
    "\n",
    "配套教程：[DQN：深度 Q 网络](https://dqn.cn/advanced/dqn)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/dqn_atari_setup.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-dqn_atari_setup-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 [\"ale_py\",\"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[atari]\", \"opencv-python\", \"gymnasium>=1.0\"], check=True)\n",
    "sys.argv = sys.argv[:1]   # notebook 内核会带上 -f 等参数，清掉以免脚本把它们当成命令行参数\n",
    "# os.environ[\"DQN_QUICK\"] = \"1\"   # 取消注释：支持该开关的脚本只跑很短的版本，用于快速检查"
   ],
   "id": "advanced-dqn_atari_setup-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ale_py\n",
    "import gymnasium as gym\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "gym.register_envs(ale_py)"
   ],
   "id": "advanced-dqn_atari_setup-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_env(env_id=\"ALE/Breakout-v5\"):\n",
    "    # frameskip=1：跳帧交给 AtariPreprocessing 处理；repeat_action_probability=0 对应原论文设置\n",
    "    env = gym.make(env_id, frameskip=1, repeat_action_probability=0.0)\n",
    "    env = gym.wrappers.AtariPreprocessing(\n",
    "        env,\n",
    "        noop_max=30,             # 开局随机执行 0~30 个 no-op，增加初始状态多样性\n",
    "        frame_skip=4,            # 每个动作重复 4 帧\n",
    "        screen_size=84,          # 缩放到 84×84\n",
    "        grayscale_obs=True,      # 灰度化\n",
    "        terminal_on_life_loss=False,\n",
    "        scale_obs=False,         # 保持 uint8 存储以节省回放内存，送入网络前再 /255\n",
    "    )\n",
    "    env = gym.wrappers.FrameStackObservation(env, stack_size=4)   # 堆叠最近 4 帧 → (4, 84, 84)\n",
    "    return env"
   ],
   "id": "advanced-dqn_atari_setup-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NatureDQN(nn.Module):\n",
    "    \"\"\"Mnih et al. (2015) 中的网络结构\"\"\"\n",
    "\n",
    "    def __init__(self, n_actions):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv2d(4, 32, kernel_size=8, stride=4), nn.ReLU(),\n",
    "            nn.Conv2d(32, 64, kernel_size=4, stride=2), nn.ReLU(),\n",
    "            nn.Conv2d(64, 64, kernel_size=3, stride=1), nn.ReLU(),\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(64 * 7 * 7, 512), nn.ReLU(),\n",
    "            nn.Linear(512, n_actions),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x.float() / 255.0)"
   ],
   "id": "advanced-dqn_atari_setup-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = make_env()\n",
    "obs, _ = env.reset(seed=0)\n",
    "q_net = NatureDQN(env.action_space.n)\n",
    "print(\"观测形状:\", np.array(obs).shape, \"| 动作数:\", env.action_space.n)\n",
    "print(\"网络参数量:\", sum(p.numel() for p in q_net.parameters()))\n",
    "total_reward = 0.0\n",
    "for t in range(200):\n",
    "    with torch.no_grad():\n",
    "        q = q_net(torch.as_tensor(np.array(obs))[None])\n",
    "    obs, r, term, trunc, _ = env.step(env.action_space.sample())\n",
    "    total_reward += np.sign(r)          # 训练时通常把奖励裁剪到 {-1, 0, +1}\n",
    "    if term or trunc:\n",
    "        obs, _ = env.reset()\n",
    "print(\"Q 值输出形状:\", tuple(q.shape), \"| 随机交互 200 步完成\")"
   ],
   "id": "advanced-dqn_atari_setup-py-5"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}