{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TD3：Pendulum-v1（双 Critic 取小 + 延迟策略更新 + 目标策略平滑）。\n",
    "\n",
    "\n",
    "\n",
    "配套教程：[TD3：双延迟深度确定性策略梯度](https://dqn.cn/advanced/td3)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/td3_pendulum.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-td3_pendulum-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-td3_pendulum-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\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",
    "GAMMA, TAU, LR_A, LR_C, BATCH = 0.99, 0.005, 1e-3, 1e-3, 256\n",
    "TOTAL_STEPS, START_STEPS, EXPL_NOISE = 20_000, 1_000, 0.1\n",
    "POLICY_NOISE, NOISE_CLIP, POLICY_DELAY = 0.2, 0.5, 2\n",
    "random.seed(0); np.random.seed(0); torch.manual_seed(0)\n",
    "\n",
    "env = gym.make(\"Pendulum-v1\")\n",
    "env.action_space.seed(0)  # 预热阶段的随机动作也要可复现\n",
    "obs_dim, act_dim = env.observation_space.shape[0], env.action_space.shape[0]\n",
    "max_action = float(env.action_space.high[0])        # Pendulum 的力矩范围 [-2, 2]"
   ],
   "id": "advanced-td3_pendulum-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mlp(i, o, out_act=None):\n",
    "    layers = [nn.Linear(i, 256), nn.ReLU(), nn.Linear(256, 256), nn.ReLU(), nn.Linear(256, o)]\n",
    "    return nn.Sequential(*layers, *( [out_act] if out_act else [] ))"
   ],
   "id": "advanced-td3_pendulum-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actor = mlp(obs_dim, act_dim, nn.Tanh())             # μ(s) ∈ [-1, 1]，再乘 max_action\n",
    "critic1, critic2 = mlp(obs_dim + act_dim, 1), mlp(obs_dim + act_dim, 1)   # 两个独立的 Q 网络\n",
    "actor_t, critic1_t, critic2_t = copy.deepcopy(actor), copy.deepcopy(critic1), copy.deepcopy(critic2)\n",
    "opt_a = torch.optim.Adam(actor.parameters(), lr=LR_A)\n",
    "opt_c = torch.optim.Adam(list(critic1.parameters()) + list(critic2.parameters()), lr=LR_C)"
   ],
   "id": "advanced-td3_pendulum-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Replay:\n",
    "    def __init__(self, size=100_000):\n",
    "        self.s = np.zeros((size, obs_dim), np.float32); self.a = np.zeros((size, act_dim), np.float32)\n",
    "        self.r = np.zeros((size, 1), np.float32); self.s2 = np.zeros((size, obs_dim), np.float32)\n",
    "        self.d = np.zeros((size, 1), np.float32); self.ptr, self.n, self.size = 0, 0, size\n",
    "\n",
    "    def add(self, s, a, r, s2, d):\n",
    "        self.s[self.ptr], self.a[self.ptr], self.r[self.ptr], self.s2[self.ptr], self.d[self.ptr] = s, a, r, s2, d\n",
    "        self.ptr = (self.ptr + 1) % self.size; self.n = min(self.n + 1, self.size)\n",
    "\n",
    "    def sample(self, k):\n",
    "        i = np.random.randint(0, self.n, k)\n",
    "        return [torch.as_tensor(x[i]) for x in (self.s, self.a, self.r, self.s2, self.d)]"
   ],
   "id": "advanced-td3_pendulum-py-5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def soft_update(net, target):\n",
    "    with torch.no_grad():\n",
    "        for p, tp in zip(net.parameters(), target.parameters()):\n",
    "            tp.mul_(1 - TAU).add_(TAU * p)"
   ],
   "id": "advanced-td3_pendulum-py-6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "buf = Replay()\n",
    "obs, _ = env.reset(seed=0)\n",
    "ep_ret, returns = 0.0, []\n",
    "for step in range(TOTAL_STEPS):\n",
    "    if step < START_STEPS:\n",
    "        a = env.action_space.sample()                                   # 预热：均匀随机动作\n",
    "    else:\n",
    "        with torch.no_grad():\n",
    "            a = max_action * actor(torch.as_tensor(obs, dtype=torch.float32)).numpy()\n",
    "        a = np.clip(a + np.random.normal(0, EXPL_NOISE * max_action, act_dim), -max_action, max_action)\n",
    "    obs2, r, term, trunc, _ = env.step(a)\n",
    "    buf.add(obs, a / max_action, r, obs2, float(term))                  # 存归一化后的动作\n",
    "    obs, ep_ret = obs2, ep_ret + r\n",
    "    if term or trunc:\n",
    "        returns.append(ep_ret); obs, _ = env.reset(); ep_ret = 0.0\n",
    "        if len(returns) % 10 == 0:\n",
    "            print(f\"step {step + 1:6d} | 最近10回合平均回报 {np.mean(returns[-10:]):8.1f}\")\n",
    "\n",
    "    if step >= START_STEPS:\n",
    "        s, a_b, r_b, s2, d = buf.sample(BATCH)\n",
    "        with torch.no_grad():\n",
    "            # 目标策略平滑：给目标动作加裁剪过的高斯噪声（动作已归一化到 [-1, 1]）\n",
    "            noise = (torch.randn_like(a_b) * POLICY_NOISE).clamp(-NOISE_CLIP, NOISE_CLIP)\n",
    "            a2 = (actor_t(s2) + noise).clamp(-1, 1)\n",
    "            sa2 = torch.cat([s2, a2], 1)\n",
    "            y = r_b + GAMMA * (1 - d) * torch.min(critic1_t(sa2), critic2_t(sa2))   # 取两个 Q 的较小值\n",
    "        sa = torch.cat([s, a_b], 1)\n",
    "        c_loss = F.mse_loss(critic1(sa), y) + F.mse_loss(critic2(sa), y)\n",
    "        opt_c.zero_grad(); c_loss.backward(); opt_c.step()\n",
    "\n",
    "        if step % POLICY_DELAY == 0:                                     # 延迟更新 Actor 与目标网络\n",
    "            a_loss = -critic1(torch.cat([s, actor(s)], 1)).mean()\n",
    "            opt_a.zero_grad(); a_loss.backward(); opt_a.step()\n",
    "            soft_update(actor, actor_t); soft_update(critic1, critic1_t); soft_update(critic2, critic2_t)\n",
    "\n",
    "# 确定性策略评估\n",
    "scores = []\n",
    "for i in range(10):\n",
    "    o, _ = env.reset(seed=1000 + i); done, total = False, 0.0\n",
    "    while not done:\n",
    "        with torch.no_grad():\n",
    "            o, r, te, tr, _ = env.step(max_action * actor(torch.as_tensor(o, dtype=torch.float32)).numpy())\n",
    "        total += r; done = te or tr\n",
    "    scores.append(total)\n",
    "print(f\"确定性策略评估（10 回合）平均回报: {np.mean(scores):.1f}\")"
   ],
   "id": "advanced-td3_pendulum-py-7"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}