{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 连续动作离线强化学习：Pendulum-v1 上的 BC、朴素离线 SAC 与 CQL(H)（基于 SAC）。\n",
    "\n",
    "数据集：在线 SAC 训练前 DATA_STEPS 步的完整回放缓冲区（前 1000 步为随机动作），质量从随机到较好混杂，\n",
    "类似 D4RL 的 \"medium-replay\"。之后三种方法只用这份固定数据训练，不再与环境交互。\n",
    "用法：python cql_pendulum.py [种子]；DQN_QUICK=1 只跑很短的版本。\n",
    "\n",
    "配套教程：[离线强化学习](https://dqn.cn/advanced/offline-rl)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/cql_pendulum.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-cql_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-cql_pendulum-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",
    "from torch.distributions import Normal\n",
    "\n",
    "QUICK = os.environ.get(\"DQN_QUICK\") == \"1\"\n",
    "torch.set_num_threads(1)\n",
    "DATA_STEPS = 2_000 if QUICK else 6_000\n",
    "OFFLINE_STEPS = 2_000 if QUICK else 20_000\n",
    "GAMMA, TAU, LR, BATCH, A_MAX = 0.99, 0.005, 3e-4, 256, 2.0\n",
    "CQL_ALPHA, CQL_N = 5.0, 10            # CQL 正则系数、每个状态采样的动作数"
   ],
   "id": "advanced-cql_pendulum-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mlp(i, o, h=256):\n",
    "    return nn.Sequential(nn.Linear(i, h), nn.ReLU(), nn.Linear(h, h), nn.ReLU(), nn.Linear(h, o))"
   ],
   "id": "advanced-cql_pendulum-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Actor(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__(); self.net = mlp(3, 2)\n",
    "\n",
    "    def forward(self, s):\n",
    "        mu, log_std = self.net(s).chunk(2, -1)\n",
    "        d = Normal(mu, log_std.clamp(-20, 2).exp()); u = d.rsample(); a = torch.tanh(u)\n",
    "        logp = d.log_prob(u).sum(-1) - torch.log(1 - a.pow(2) + 1e-6).sum(-1)\n",
    "        return a, logp, torch.tanh(mu)"
   ],
   "id": "advanced-cql_pendulum-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SAC:\n",
    "    def __init__(self):\n",
    "        self.actor, self.q1, self.q2 = Actor(), mlp(4, 1), mlp(4, 1)\n",
    "        self.q1_t, self.q2_t = copy.deepcopy(self.q1), copy.deepcopy(self.q2)\n",
    "        self.opt_a = torch.optim.Adam(self.actor.parameters(), lr=LR)\n",
    "        self.opt_q = torch.optim.Adam(list(self.q1.parameters()) + list(self.q2.parameters()), lr=LR)\n",
    "        self.log_alpha = torch.zeros(1, requires_grad=True); self.opt_al = torch.optim.Adam([self.log_alpha], lr=LR)\n",
    "\n",
    "    def update(self, s, a, r, s2, cql=0.0):\n",
    "        alpha = self.log_alpha.exp().detach()\n",
    "        with torch.no_grad():\n",
    "            a2, lp2, _ = self.actor(s2); sa2 = torch.cat([s2, a2], 1)\n",
    "            y = r + GAMMA * (torch.min(self.q1_t(sa2), self.q2_t(sa2)) - alpha * lp2[:, None])\n",
    "        sa = torch.cat([s, a], 1); q1, q2 = self.q1(sa), self.q2(sa)\n",
    "        lq = F.mse_loss(q1, y) + F.mse_loss(q2, y)\n",
    "        if cql > 0:\n",
    "            # CQL(H)：压低 logsumexp_a Q(s,a)（用均匀动作 + 当前策略动作做重要性采样估计），抬高数据中的 Q(s,a)\n",
    "            B = len(s); s_rep = s.repeat_interleave(CQL_N, 0)\n",
    "            a_rand = torch.rand(B * CQL_N, 1) * 2 - 1\n",
    "            with torch.no_grad():\n",
    "                a_pi, lp_pi, _ = self.actor(s_rep)\n",
    "            for q_net, q_data in ((self.q1, q1), (self.q2, q2)):\n",
    "                q_rand = q_net(torch.cat([s_rep, a_rand], 1)).view(B, CQL_N) - np.log(0.5)   # 均匀分布密度 1/2\n",
    "                q_pi = q_net(torch.cat([s_rep, a_pi], 1)).view(B, CQL_N) - lp_pi.view(B, CQL_N)\n",
    "                lse = torch.logsumexp(torch.cat([q_rand, q_pi], 1), 1) - np.log(2 * CQL_N)\n",
    "                lq = lq + cql * (lse.mean() - q_data.mean())\n",
    "        self.opt_q.zero_grad(); lq.backward(); self.opt_q.step()\n",
    "        an, lp, _ = self.actor(s); san = torch.cat([s, an], 1)\n",
    "        la = (alpha * lp - torch.min(self.q1(san), self.q2(san)).squeeze(1)).mean()\n",
    "        self.opt_a.zero_grad(); la.backward(); self.opt_a.step()\n",
    "        lal = -(self.log_alpha * (lp.detach() - 1.0)).mean()      # 目标熵 = -|A| = -1\n",
    "        self.opt_al.zero_grad(); lal.backward(); self.opt_al.step()\n",
    "        with torch.no_grad():\n",
    "            for n, t in ((self.q1, self.q1_t), (self.q2, self.q2_t)):\n",
    "                for p, tp in zip(n.parameters(), t.parameters()):\n",
    "                    tp.mul_(1 - TAU).add_(TAU * p)"
   ],
   "id": "advanced-cql_pendulum-py-5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(act, n=10):\n",
    "    env, tot = gym.make(\"Pendulum-v1\"), []\n",
    "    for k in range(n):\n",
    "        o, _ = env.reset(seed=10_000 + k); done, G = False, 0.0\n",
    "        while not done:\n",
    "            with torch.no_grad():\n",
    "                a = act(torch.as_tensor(o, dtype=torch.float32))\n",
    "            o, r, te, tr, _ = env.step(a.numpy() * A_MAX); G += r; done = te or tr\n",
    "        tot.append(G)\n",
    "    return float(np.mean(tot))"
   ],
   "id": "advanced-cql_pendulum-py-6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collect(seed):\n",
    "    \"\"\"在线 SAC 训练 DATA_STEPS 步，把整个回放缓冲区作为离线数据集。\"\"\"\n",
    "    rng = np.random.default_rng(seed); env = gym.make(\"Pendulum-v1\"); agent = SAC()\n",
    "    D = {k: torch.zeros(DATA_STEPS, d) for k, d in ((\"s\", 3), (\"a\", 1), (\"r\", 1), (\"s2\", 3))}\n",
    "    o, _ = env.reset(seed=seed); rets, G = [], 0.0\n",
    "    for t in range(DATA_STEPS):\n",
    "        if t < 1000:\n",
    "            a = rng.uniform(-1, 1, 1).astype(np.float32)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                a = agent.actor(torch.as_tensor(o, dtype=torch.float32))[0].numpy()\n",
    "        o2, r, te, tr, _ = env.step(a * A_MAX); G += r\n",
    "        for k, v in ((\"s\", o), (\"a\", a), (\"r\", [r]), (\"s2\", o2)):\n",
    "            D[k][t] = torch.as_tensor(np.asarray(v, np.float32))\n",
    "        o = o2\n",
    "        if te or tr:\n",
    "            rets.append(G); G = 0.0; o, _ = env.reset()\n",
    "        if t >= 1000:\n",
    "            j = torch.as_tensor(rng.integers(0, t + 1, BATCH))\n",
    "            agent.update(*(D[k][j] for k in (\"s\", \"a\", \"r\", \"s2\")))\n",
    "    return D, rets"
   ],
   "id": "advanced-cql_pendulum-py-7"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_offline(D, method, seed):\n",
    "    torch.manual_seed(seed); rng = np.random.default_rng(seed); N = len(D[\"s\"])\n",
    "    if method == \"bc\":\n",
    "        pi = mlp(3, 1); opt = torch.optim.Adam(pi.parameters(), lr=1e-3)\n",
    "        for _ in range(OFFLINE_STEPS):\n",
    "            j = torch.as_tensor(rng.integers(0, N, BATCH))\n",
    "            loss = F.mse_loss(torch.tanh(pi(D[\"s\"][j])), D[\"a\"][j]); opt.zero_grad(); loss.backward(); opt.step()\n",
    "        return evaluate(lambda o: torch.tanh(pi(o)))\n",
    "    agent = SAC()\n",
    "    for _ in range(OFFLINE_STEPS):\n",
    "        j = torch.as_tensor(rng.integers(0, N, BATCH))\n",
    "        agent.update(*(D[k][j] for k in (\"s\", \"a\", \"r\", \"s2\")), cql=CQL_ALPHA if method == \"cql\" else 0.0)\n",
    "    with torch.no_grad():\n",
    "        q_data = torch.min(agent.q1(torch.cat([D[\"s\"], D[\"a\"]], 1)), agent.q2(torch.cat([D[\"s\"], D[\"a\"]], 1))).mean().item()\n",
    "    return evaluate(lambda o: agent.actor(o)[2]), q_data"
   ],
   "id": "advanced-cql_pendulum-py-8"
  },
  {
   "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])\n",
    "    for sd in seeds:\n",
    "        torch.manual_seed(sd)\n",
    "        D, rets = collect(sd)\n",
    "        print(f\"种子{sd}: 数据集 {len(D['s'])} 条转移，{len(rets)} 个回合，回合回报 最差 {min(rets):.0f} / 平均 {np.mean(rets):.0f} / 最好 {max(rets):.0f}\", flush=True)\n",
    "        print(f\"  BC          评估回报 {train_offline(D, 'bc', sd):8.1f}\", flush=True)\n",
    "        for m in (\"sac\", \"cql\"):\n",
    "            score, q = train_offline(D, m, sd)\n",
    "            name = \"朴素离线 SAC\" if m == \"sac\" else f\"CQL(α={CQL_ALPHA:g})\"\n",
    "            print(f\"  {name:11s} 评估回报 {score:8.1f}   数据上的平均 Q 估计 {q:9.1f}\", flush=True)"
   ],
   "id": "advanced-cql_pendulum-py-9"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}