{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 离散动作 SAC：CartPole-v1。\n",
    "\n",
    "与连续 SAC 的区别：策略直接输出各动作概率，软价值、策略损失和温度损失都可以对动作“精确求期望”，不需要重参数化采样。\n",
    "用法：python sac_discrete_cartpole.py [种子]；DQN_QUICK=1 只跑很短的版本。\n",
    "\n",
    "配套教程：[SAC：柔性 Actor-Critic](https://dqn.cn/advanced/sac)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/sac_discrete_cartpole.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-sac_discrete_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-sac_discrete_cartpole-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import os\n",
    "import sys\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",
    "QUICK = os.environ.get(\"DQN_QUICK\") == \"1\"\n",
    "torch.set_num_threads(1)\n",
    "TOTAL_STEPS = 5_000 if QUICK else 40_000\n",
    "START_STEPS, BATCH, GAMMA, TAU, LR = 1_000, 128, 0.99, 0.005, 3e-4\n",
    "TARGET_ENT_RATIO = 0.5          # 目标熵 = 0.5 × log|A|。实测取 0.98（原论文）时 α 持续增大到上万，采样策略接近均匀随机"
   ],
   "id": "advanced-sac_discrete_cartpole-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mlp(i, o):\n",
    "    return nn.Sequential(nn.Linear(i, 128), nn.ReLU(), nn.Linear(128, 128), nn.ReLU(), nn.Linear(128, o))"
   ],
   "id": "advanced-sac_discrete_cartpole-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run(seed):\n",
    "    rng = np.random.default_rng(seed); torch.manual_seed(seed)\n",
    "    env, test_env = gym.make(\"CartPole-v1\"), gym.make(\"CartPole-v1\")\n",
    "    nS, nA = env.observation_space.shape[0], env.action_space.n\n",
    "    actor, q1, q2 = mlp(nS, nA), mlp(nS, nA), mlp(nS, nA)       # Q 网络一次输出所有动作的 Q 值\n",
    "    q1_t, q2_t = copy.deepcopy(q1), copy.deepcopy(q2)\n",
    "    opt_a = torch.optim.Adam(actor.parameters(), lr=LR)\n",
    "    opt_q = torch.optim.Adam(list(q1.parameters()) + list(q2.parameters()), lr=LR)\n",
    "    log_alpha = torch.zeros(1, requires_grad=True); opt_al = torch.optim.Adam([log_alpha], lr=LR)\n",
    "    target_ent = TARGET_ENT_RATIO * np.log(nA)\n",
    "    buf = {k: np.zeros((TOTAL_STEPS, d), np.float32) for k, d in ((\"s\", nS), (\"a\", 1), (\"r\", 1), (\"s2\", nS), (\"d\", 1))}\n",
    "\n",
    "    def evaluate(n=10):\n",
    "        tot = []\n",
    "        for k in range(n):\n",
    "            o, _ = test_env.reset(seed=10_000 + k); done, G = False, 0.0\n",
    "            while not done:\n",
    "                with torch.no_grad():\n",
    "                    a = int(actor(torch.as_tensor(o)).argmax())\n",
    "                o, r, te, tr, _ = test_env.step(a); G += r; done = te or tr\n",
    "            tot.append(G)\n",
    "        return float(np.mean(tot))\n",
    "\n",
    "    o, _ = env.reset(seed=seed); log, ep_ret, rets = [], 0.0, []\n",
    "    for step in range(TOTAL_STEPS):\n",
    "        if step < START_STEPS:\n",
    "            a = int(rng.integers(nA))\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                a = int(torch.distributions.Categorical(logits=actor(torch.as_tensor(o))).sample())\n",
    "        o2, r, te, tr, _ = env.step(a); ep_ret += r\n",
    "        for k, v in ((\"s\", o), (\"a\", [a]), (\"r\", [r]), (\"s2\", o2), (\"d\", [float(te)])):   # 只有真正终止才截断自举\n",
    "            buf[k][step] = v\n",
    "        o = o2\n",
    "        if te or tr:\n",
    "            rets.append(ep_ret); ep_ret = 0.0; o, _ = env.reset()\n",
    "        if step + 1 < START_STEPS:\n",
    "            continue\n",
    "        j = rng.integers(0, step + 1, BATCH)\n",
    "        s, a_, r_, s2, d = (torch.as_tensor(buf[k][j]) for k in (\"s\", \"a\", \"r\", \"s2\", \"d\"))\n",
    "        alpha = log_alpha.exp().detach()\n",
    "        with torch.no_grad():                    # 软价值 V(s') = Σ_a π(a|s') [min Q(s',a) - α log π(a|s')]\n",
    "            logp2 = F.log_softmax(actor(s2), -1); p2 = logp2.exp()\n",
    "            v2 = (p2 * (torch.min(q1_t(s2), q2_t(s2)) - alpha * logp2)).sum(-1, keepdim=True)\n",
    "            y = r_ + GAMMA * (1 - d) * v2\n",
    "        a_idx = a_.long()\n",
    "        lq = F.mse_loss(q1(s).gather(1, a_idx), y) + F.mse_loss(q2(s).gather(1, a_idx), y)\n",
    "        opt_q.zero_grad(); lq.backward(); opt_q.step()\n",
    "        logp = F.log_softmax(actor(s), -1); p = logp.exp()\n",
    "        with torch.no_grad():\n",
    "            qmin = torch.min(q1(s), q2(s))\n",
    "        la = (p * (alpha * logp - qmin)).sum(-1).mean()        # 策略损失：对动作精确求期望\n",
    "        opt_a.zero_grad(); la.backward(); opt_a.step()\n",
    "        ent = -(p.detach() * logp.detach()).sum(-1)\n",
    "        lal = (log_alpha * (ent - target_ent)).mean()           # 熵高于目标时降低 α，反之升高\n",
    "        opt_al.zero_grad(); lal.backward(); opt_al.step()\n",
    "        with torch.no_grad():\n",
    "            for n_, t_ in ((q1, q1_t), (q2, q2_t)):\n",
    "                for p_, tp in zip(n_.parameters(), t_.parameters()):\n",
    "                    tp.mul_(1 - TAU).add_(TAU * p_)\n",
    "        if (step + 1) % (1000 if QUICK else 5000) == 0:\n",
    "            log.append((step + 1, evaluate()))\n",
    "            print(f\"  种子{seed} 步数 {step + 1:6d}  最近10回合训练回报 {np.mean(rets[-10:]):6.1f}  贪心评估 {log[-1][1]:6.1f}  α={alpha.item():.3f}\", flush=True)\n",
    "    return log"
   ],
   "id": "advanced-sac_discrete_cartpole-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    seeds = [int(sys.argv[1])] if len(sys.argv) > 1 else ([0] if QUICK else [0, 1, 2])\n",
    "    for sd in seeds:\n",
    "        log = run(sd)\n",
    "        print(f\"种子{sd}: \" + \"  \".join(f\"{s}:{v:.0f}\" for s, v in log))"
   ],
   "id": "advanced-sac_discrete_cartpole-py-5"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}