{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模仿学习：行为克隆（BC）与 DAgger，CartPole-v1。\n",
    "\n",
    "专家是一个手写规则 a = 1[θ + 0.5·θ̇ > 0]（实测 20 个回合均达到满分 500）。\n",
    "\n",
    "配套教程：[模仿学习](https://dqn.cn/advanced/imitation-learning)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/bc_dagger_cartpole.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-bc_dagger_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-bc_dagger_cartpole-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "env = gym.make(\"CartPole-v1\")"
   ],
   "id": "advanced-bc_dagger_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-bc_dagger_cartpole-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rollout(policy_fn, seed, max_steps=500):\n",
    "    \"\"\"用 policy_fn 与环境交互，返回访问到的状态和回报。\"\"\"\n",
    "    o, _ = env.reset(seed=seed); states, G = [], 0.0\n",
    "    for _ in range(max_steps):\n",
    "        states.append(o)\n",
    "        o, r, te, tr, _ = env.step(policy_fn(o)); G += r\n",
    "        if te or tr:\n",
    "            break\n",
    "    return states, G"
   ],
   "id": "advanced-bc_dagger_cartpole-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit(X, y, steps=2000):\n",
    "    pi = nn.Sequential(nn.Linear(4, 64), nn.Tanh(), nn.Linear(64, 2))\n",
    "    opt = torch.optim.Adam(pi.parameters(), lr=1e-2)\n",
    "    X, y = torch.as_tensor(np.array(X), dtype=torch.float32), torch.as_tensor(y)\n",
    "    for _ in range(steps):\n",
    "        loss = F.cross_entropy(pi(X), y)\n",
    "        opt.zero_grad(); loss.backward(); opt.step()\n",
    "    return lambda o: int(pi(torch.as_tensor(o, dtype=torch.float32)).argmax())"
   ],
   "id": "advanced-bc_dagger_cartpole-py-5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(policy_fn, n=20):\n",
    "    return np.mean([rollout(policy_fn, 10_000 + i)[1] for i in range(n)])"
   ],
   "id": "advanced-bc_dagger_cartpole-py-6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_DEMO = 10      # 每轮只收集 10 个状态（刻意很少，以便观察复合误差）\n",
    "N_ROUNDS = 5"
   ],
   "id": "advanced-bc_dagger_cartpole-py-7"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run(seed):\n",
    "    torch.manual_seed(seed); np.random.seed(seed)\n",
    "    # ---------- 行为克隆：只用专家自己访问到的状态 ----------\n",
    "    demo_states, _ = rollout(expert, seed=seed, max_steps=N_DEMO)\n",
    "    X, y = list(demo_states), [expert(s) for s in demo_states]\n",
    "    policy = fit(X, y)\n",
    "    scores = [evaluate(policy)]\n",
    "    # ---------- DAgger：学生策略采样状态 → 专家标注 → 聚合数据 → 重新训练 ----------\n",
    "    for it in range(N_ROUNDS):\n",
    "        states, _ = rollout(policy, seed=seed * 100 + it, max_steps=N_DEMO)\n",
    "        X += states\n",
    "        y += [expert(s) for s in states]\n",
    "        policy = fit(X, y)\n",
    "        scores.append(evaluate(policy))\n",
    "    return scores"
   ],
   "id": "advanced-bc_dagger_cartpole-py-8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = np.array([run(seed) for seed in range(5)])     # 5 个随机种子取平均\n",
    "print(f\"BC（只用专家数据，{N_DEMO} 个样本）     平均评估回报 {results[:, 0].mean():6.1f}\")\n",
    "for it in range(1, N_ROUNDS + 1):\n",
    "    print(f\"DAgger 第{it}轮（共 {N_DEMO * (it + 1):3d} 个样本） 平均评估回报 {results[:, it].mean():6.1f}\")\n",
    "print(\"对照：BC 直接使用同样数量（60 个）的专家样本：\", end=\" \")\n",
    "bc60 = []\n",
    "for seed in range(5):\n",
    "    torch.manual_seed(seed)\n",
    "    st, _ = rollout(expert, seed=seed, max_steps=N_DEMO * (N_ROUNDS + 1))\n",
    "    bc60.append(evaluate(fit(st, [expert(s) for s in st])))\n",
    "print(f\"平均评估回报 {np.mean(bc60):.1f}\")"
   ],
   "id": "advanced-bc_dagger_cartpole-py-9"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}