{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 离线强化学习玩具实验（CartPole-v1）：行为克隆（BC） vs 朴素离线 DQN vs 离散动作 CQL。\n",
    "\n",
    "专家规则 a = 1[θ + 0.5·θ̇ > 0]。对比两个数据集：\n",
    "  (1) 纯专家数据（5000 条，无噪声，动作覆盖极窄）\n",
    "  (2) 高噪声数据（20000 条，80% 概率随机动作）\n",
    "\n",
    "配套教程：[离线强化学习](https://dqn.cn/advanced/offline-rl)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/offline_cartpole.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-offline_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-offline_cartpole-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "random.seed(0); np.random.seed(0); torch.manual_seed(0)\n",
    "env = gym.make(\"CartPole-v1\")\n",
    "env.action_space.seed(0)  # action_space.sample() 使用独立的随机数生成器\n",
    "GAMMA = 0.99"
   ],
   "id": "advanced-offline_cartpole-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def expert(o):\n",
    "    return int(o[2] + 0.5 * o[3] > 0)"
   ],
   "id": "advanced-offline_cartpole-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ---------- 1) 收集固定数据集，之后不再与环境交互 ----------\n",
    "def collect(eps_data, n_transitions):\n",
    "    data, ep_returns = [], []\n",
    "    o, _ = env.reset(seed=0); G = 0.0\n",
    "    while len(data) < n_transitions:\n",
    "        a = env.action_space.sample() if random.random() < eps_data else expert(o)\n",
    "        o2, r, te, tr, _ = env.step(a)\n",
    "        data.append((o, a, r, o2, float(te))); G += r; o = o2\n",
    "        if te or tr:\n",
    "            ep_returns.append(G); G = 0.0; o, _ = env.reset()\n",
    "    S, A, R, S2, D = (torch.as_tensor(np.array(x), dtype=torch.float32) for x in zip(*data))\n",
    "    return S, A.long(), R, S2, D, np.mean(ep_returns)"
   ],
   "id": "advanced-offline_cartpole-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def net():\n",
    "    return nn.Sequential(nn.Linear(4, 128), nn.ReLU(), nn.Linear(128, 128), nn.ReLU(), nn.Linear(128, 2))"
   ],
   "id": "advanced-offline_cartpole-py-5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, n=20):\n",
    "    e = gym.make(\"CartPole-v1\"); tot = []\n",
    "    for i in range(n):\n",
    "        o, _ = e.reset(seed=10_000 + i); done, g = False, 0.0\n",
    "        while not done:\n",
    "            with torch.no_grad():\n",
    "                o, r, te, tr, _ = e.step(int(model(torch.as_tensor(o, dtype=torch.float32)).argmax()))\n",
    "            g += r; done = te or tr\n",
    "        tot.append(g)\n",
    "    return np.mean(tot)"
   ],
   "id": "advanced-offline_cartpole-py-6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_bc(steps=5_000):\n",
    "    pi = net(); opt = torch.optim.Adam(pi.parameters(), lr=1e-3)\n",
    "    for _ in range(steps):\n",
    "        i = torch.randint(0, len(S), (256,))\n",
    "        loss = F.cross_entropy(pi(S[i]), A[i])\n",
    "        opt.zero_grad(); loss.backward(); opt.step()\n",
    "    return pi"
   ],
   "id": "advanced-offline_cartpole-py-7"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_q(cql_alpha=0.0, steps=20_000):\n",
    "    q, q_t = net(), net(); q_t.load_state_dict(q.state_dict())\n",
    "    opt = torch.optim.Adam(q.parameters(), lr=3e-4)\n",
    "    for t in range(steps):\n",
    "        i = torch.randint(0, len(S), (256,))\n",
    "        qs = q(S[i])\n",
    "        q_sa = qs.gather(1, A[i, None]).squeeze(1)\n",
    "        with torch.no_grad():\n",
    "            y = R[i] + GAMMA * (1 - D[i]) * q_t(S2[i]).max(1).values\n",
    "        loss = F.mse_loss(q_sa, y)\n",
    "        if cql_alpha > 0:   # CQL 正则：压低所有动作的 Q（logsumexp），抬高数据中动作的 Q\n",
    "            loss = loss + cql_alpha * (torch.logsumexp(qs, dim=1) - q_sa).mean()\n",
    "        opt.zero_grad(); loss.backward(); opt.step()\n",
    "        if t % 500 == 0:\n",
    "            q_t.load_state_dict(q.state_dict())\n",
    "    return q"
   ],
   "id": "advanced-offline_cartpole-py-8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, eps_data, n in [(\"纯专家数据\", 0.0, 5_000), (\"高噪声数据\", 0.8, 20_000)]:\n",
    "    S, A, R, S2, D, behavior_return = collect(eps_data, n)\n",
    "    print(f\"== {name}：{n} 条转移，行为策略平均回报 {behavior_return:.1f}\")\n",
    "    print(f\"   行为克隆 BC  : 评估回报 {evaluate(train_bc()):6.1f}\")\n",
    "    print(f\"   朴素离线 DQN : 评估回报 {evaluate(train_q(0.0)):6.1f}\")\n",
    "    print(f\"   CQL (α=1.0)  : 评估回报 {evaluate(train_q(1.0)):6.1f}\")"
   ],
   "id": "advanced-offline_cartpole-py-9"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}