{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SAC（自动调节温度 α）：Pendulum-v1。\n",
    "\n",
    "\n",
    "\n",
    "配套教程：[SAC：柔性 Actor-Critic](https://dqn.cn/advanced/sac)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/sac_pendulum.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-sac_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-sac_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",
    "from torch.distributions import Normal\n",
    "\n",
    "GAMMA, TAU, LR, BATCH = 0.99, 0.005, 3e-4, 256\n",
    "TOTAL_STEPS, START_STEPS = 20_000, 1_000\n",
    "random.seed(0); np.random.seed(0); torch.manual_seed(0)\n",
    "\n",
    "env = gym.make(\"Pendulum-v1\")\n",
    "obs_dim, act_dim = env.observation_space.shape[0], env.action_space.shape[0]\n",
    "max_action = float(env.action_space.high[0])"
   ],
   "id": "advanced-sac_pendulum-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mlp(i, o):\n",
    "    return nn.Sequential(nn.Linear(i, 256), nn.ReLU(), nn.Linear(256, 256), nn.ReLU(), nn.Linear(256, o))"
   ],
   "id": "advanced-sac_pendulum-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Actor(nn.Module):\n",
    "    \"\"\"高斯策略 + tanh 压缩：a = tanh(u), u ~ N(μ(s), σ(s))\"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.net = mlp(obs_dim, 2 * act_dim)\n",
    "\n",
    "    def forward(self, s):\n",
    "        mu, log_std = self.net(s).chunk(2, dim=-1)\n",
    "        log_std = log_std.clamp(-20, 2)\n",
    "        dist = Normal(mu, log_std.exp())\n",
    "        u = dist.rsample()                                  # 重参数化采样，梯度可回传\n",
    "        a = torch.tanh(u)\n",
    "        # tanh 变换的对数概率修正：log π(a|s) = log N(u) - Σ log(1 - tanh(u)^2)\n",
    "        log_prob = dist.log_prob(u).sum(-1) - torch.log(1 - a.pow(2) + 1e-6).sum(-1)\n",
    "        return a, log_prob, torch.tanh(mu)"
   ],
   "id": "advanced-sac_pendulum-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actor = Actor()\n",
    "q1, q2 = mlp(obs_dim + act_dim, 1), mlp(obs_dim + act_dim, 1)\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)          # 学习 log α 保证 α > 0\n",
    "opt_alpha = torch.optim.Adam([log_alpha], lr=LR)\n",
    "target_entropy = -act_dim                               # 常用经验值：-|A|\n",
    "\n",
    "S = np.zeros((TOTAL_STEPS, obs_dim), np.float32); A = np.zeros((TOTAL_STEPS, act_dim), np.float32)\n",
    "R = np.zeros((TOTAL_STEPS, 1), np.float32); S2 = np.zeros_like(S); D = np.zeros_like(R)\n",
    "\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 = np.random.uniform(-1, 1, act_dim).astype(np.float32)\n",
    "    else:\n",
    "        with torch.no_grad():\n",
    "            a = actor(torch.as_tensor(obs, dtype=torch.float32))[0].numpy()\n",
    "    obs2, r, term, trunc, _ = env.step(a * max_action)\n",
    "    S[step], A[step], R[step], S2[step], D[step] = obs, a, 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} | α={log_alpha.exp().item():.3f}\")\n",
    "\n",
    "    if step >= START_STEPS:\n",
    "        i = np.random.randint(0, step + 1, BATCH)\n",
    "        s, a_b, r_b, s2, d = (torch.as_tensor(x[i]) for x in (S, A, R, S2, D))\n",
    "        alpha = log_alpha.exp().detach()\n",
    "\n",
    "        # ---- Soft Q 更新：y = r + γ (min Q_targ(s',a') - α log π(a'|s')) ----\n",
    "        with torch.no_grad():\n",
    "            a2, logp2, _ = actor(s2)\n",
    "            sa2 = torch.cat([s2, a2], 1)\n",
    "            y = r_b + GAMMA * (1 - d) * (torch.min(q1_t(sa2), q2_t(sa2)) - alpha * logp2.unsqueeze(1))\n",
    "        sa = torch.cat([s, a_b], 1)\n",
    "        q_loss = F.mse_loss(q1(sa), y) + F.mse_loss(q2(sa), y)\n",
    "        opt_q.zero_grad(); q_loss.backward(); opt_q.step()\n",
    "\n",
    "        # ---- 策略更新：min E[α log π(a|s) - min Q(s,a)] ----\n",
    "        a_new, logp, _ = actor(s)\n",
    "        sa_new = torch.cat([s, a_new], 1)\n",
    "        a_loss = (alpha * logp - torch.min(q1(sa_new), q2(sa_new)).squeeze(1)).mean()\n",
    "        opt_a.zero_grad(); a_loss.backward(); opt_a.step()\n",
    "\n",
    "        # ---- 温度更新：让策略熵接近 target_entropy ----\n",
    "        alpha_loss = -(log_alpha * (logp.detach() + target_entropy)).mean()\n",
    "        opt_alpha.zero_grad(); alpha_loss.backward(); opt_alpha.step()\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for net, tgt in ((q1, q1_t), (q2, q2_t)):\n",
    "                for p, tp in zip(net.parameters(), tgt.parameters()):\n",
    "                    tp.mul_(1 - TAU).add_(TAU * p)\n",
    "\n",
    "scores = []\n",
    "for k in range(10):\n",
    "    o, _ = env.reset(seed=1000 + k); done, total = False, 0.0\n",
    "    while not done:\n",
    "        with torch.no_grad():\n",
    "            act = actor(torch.as_tensor(o, dtype=torch.float32))[2].numpy()   # 评估时用均值动作 tanh(μ)\n",
    "        o, r, te, tr, _ = env.step(act * max_action)\n",
    "        total += r; done = te or tr\n",
    "    scores.append(total)\n",
    "print(f\"确定性策略评估（10 回合）平均回报: {np.mean(scores):.1f}\")"
   ],
   "id": "advanced-sac_pendulum-py-5"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}