{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GAIL（生成对抗模仿学习）：CartPole-v1。\n",
    "\n",
    "专家与 bc_dagger_cartpole.py 相同：手写规则 a = 1[θ + 0.5·θ̇ > 0]。\n",
    "策略只能看到 GAIL 奖励 r = -log(1 - D(s, a))，看不到环境奖励；D 判别 (s, a) 是否来自专家。\n",
    "注意：该奖励恒为正，在 CartPole 这种“活得越久越好”的任务里自带“生存偏置”。\n",
    "为此再提供一个无符号偏置的对照奖励 r = log D - log(1 - D) = logit(D)（AIRL 形式），见教程正文。\n",
    "用法：python gail_cartpole.py [种子] [softplus|logit]；DQN_QUICK=1 只跑很短的版本。\n",
    "\n",
    "配套教程：[模仿学习](https://dqn.cn/advanced/imitation-learning)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/gail_cartpole.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-gail_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-gail_cartpole-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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 Categorical\n",
    "\n",
    "QUICK = os.environ.get(\"DQN_QUICK\") == \"1\"\n",
    "torch.set_num_threads(1)\n",
    "N_ENVS, N_STEPS = 4, 256\n",
    "TOTAL_STEPS = 20_480 if QUICK else 200_704\n",
    "GAMMA, LAM, LR, CLIP, EPOCHS, N_MB, ENT = 0.99, 0.95, 3e-4, 0.2, 10, 4, 0.01\n",
    "N_DEMO_TRAJ, SUBSAMPLE = 1, 20            # 只用 1 条专家轨迹，每 20 步取 1 个 (s, a)：共 25 对"
   ],
   "id": "advanced-gail_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-gail_cartpole-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def layer(i, o, std=np.sqrt(2)):\n",
    "    l = nn.Linear(i, o); nn.init.orthogonal_(l.weight, std); nn.init.zeros_(l.bias); return l"
   ],
   "id": "advanced-gail_cartpole-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(act_fn, n=20):\n",
    "    env, scores = gym.make(\"CartPole-v1\"), []\n",
    "    for i in range(n):\n",
    "        o, _ = env.reset(seed=10_000 + i); done, G = False, 0.0\n",
    "        while not done:\n",
    "            o, r, te, tr, _ = env.step(act_fn(o)); G += r; done = te or tr\n",
    "        scores.append(G)\n",
    "    return float(np.mean(scores))"
   ],
   "id": "advanced-gail_cartpole-py-5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def demos(seed):\n",
    "    env = gym.make(\"CartPole-v1\"); S, A = [], []\n",
    "    for k in range(N_DEMO_TRAJ):\n",
    "        o, _ = env.reset(seed=seed * 100 + k); done, t = False, 0\n",
    "        while not done:\n",
    "            a = expert(o)\n",
    "            if t % SUBSAMPLE == 0:\n",
    "                S.append(o); A.append(a)\n",
    "            o, _, te, tr, _ = env.step(a); done = te or tr; t += 1\n",
    "    return torch.as_tensor(np.array(S), dtype=torch.float32), torch.as_tensor(A)"
   ],
   "id": "advanced-gail_cartpole-py-6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bc(S, A, seed, steps=2000):\n",
    "    torch.manual_seed(seed)\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",
    "    for _ in range(steps):\n",
    "        loss = F.cross_entropy(pi(S), A); opt.zero_grad(); loss.backward(); opt.step()\n",
    "    return lambda o: int(pi(torch.as_tensor(o, dtype=torch.float32)).argmax())"
   ],
   "id": "advanced-gail_cartpole-py-7"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sa(s, a):\n",
    "    return torch.cat([s, F.one_hot(a.long(), 2).float()], -1)"
   ],
   "id": "advanced-gail_cartpole-py-8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gail(S_exp, A_exp, seed, reward=\"softplus\"):\n",
    "    torch.manual_seed(seed); np.random.seed(seed)\n",
    "    envs = gym.vector.SyncVectorEnv([lambda: gym.make(\"CartPole-v1\") for _ in range(N_ENVS)],\n",
    "                                    autoreset_mode=gym.vector.AutoresetMode.SAME_STEP)\n",
    "    actor = nn.Sequential(layer(4, 64), nn.Tanh(), layer(64, 64), nn.Tanh(), layer(64, 2, 0.01))\n",
    "    critic = nn.Sequential(layer(4, 64), nn.Tanh(), layer(64, 64), nn.Tanh(), layer(64, 1, 1.0))\n",
    "    disc = nn.Sequential(nn.Linear(6, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 1))\n",
    "    params = list(actor.parameters()) + list(critic.parameters())\n",
    "    opt = torch.optim.Adam(params, lr=LR, eps=1e-5)\n",
    "    d_opt = torch.optim.Adam(disc.parameters(), lr=3e-4)\n",
    "    exp_sa = sa(S_exp, A_exp)\n",
    "\n",
    "    obs = torch.as_tensor(envs.reset(seed=seed)[0], dtype=torch.float32)\n",
    "    O = torch.zeros(N_STEPS, N_ENVS, 4); A, LP, V, D = (torch.zeros(N_STEPS, N_ENVS) for _ in range(4))\n",
    "    FINAL = torch.zeros(N_STEPS, N_ENVS, 4); TR = torch.zeros(N_STEPS, N_ENVS)\n",
    "    env_returns, ep_ret = [], np.zeros(N_ENVS)\n",
    "    n_iters, log = TOTAL_STEPS // (N_ENVS * N_STEPS), []\n",
    "    for it in range(1, n_iters + 1):\n",
    "        for t in range(N_STEPS):\n",
    "            with torch.no_grad():\n",
    "                d = Categorical(logits=actor(obs)); a = d.sample(); V[t] = critic(obs).squeeze(-1)\n",
    "            o2, r, te, tr, info = envs.step(a.numpy())\n",
    "            O[t], A[t], LP[t] = obs, a, d.log_prob(a)\n",
    "            D[t] = torch.as_tensor(te | tr, dtype=torch.float32); TR[t] = torch.as_tensor(tr, dtype=torch.float32)\n",
    "            if tr.any():\n",
    "                FINAL[t] = torch.as_tensor(np.stack([info[\"final_obs\"][i] if tr[i] else o2[i] for i in range(N_ENVS)]), dtype=torch.float32)\n",
    "            ep_ret += r                                         # 环境回报只用于记录，不参与训练\n",
    "            for i in np.flatnonzero(te | tr):\n",
    "                env_returns.append(ep_ret[i]); ep_ret[i] = 0\n",
    "            obs = torch.as_tensor(o2, dtype=torch.float32)\n",
    "\n",
    "        # ---------- 1) 训练判别器：专家 (s,a) → 1，策略 (s,a) → 0 ----------\n",
    "        pol_sa = sa(O.reshape(-1, 4), A.reshape(-1))\n",
    "        for _ in range(5):\n",
    "            j = torch.randint(len(pol_sa), (256,)); k = torch.randint(len(exp_sa), (256,))\n",
    "            logits_p, logits_e = disc(pol_sa[j]), disc(exp_sa[k])\n",
    "            d_loss = F.binary_cross_entropy_with_logits(logits_e, torch.ones_like(logits_e)) + \\\n",
    "                     F.binary_cross_entropy_with_logits(logits_p, torch.zeros_like(logits_p))\n",
    "            d_opt.zero_grad(); d_loss.backward(); d_opt.step()\n",
    "\n",
    "        # ---------- 2) GAIL 奖励 r = -log(1 - D) = softplus(logit) ----------\n",
    "        with torch.no_grad():\n",
    "            logit = disc(pol_sa).reshape(N_STEPS, N_ENVS)\n",
    "            R = F.softplus(logit) if reward == \"softplus\" else logit\n",
    "            R = R + GAMMA * critic(FINAL).squeeze(-1) * TR   # 时间截断：自举\n",
    "            next_v = critic(obs).squeeze(-1); adv = torch.zeros_like(R); last = 0\n",
    "            for t in reversed(range(N_STEPS)):\n",
    "                nv = next_v if t == N_STEPS - 1 else V[t + 1]\n",
    "                delta = R[t] + GAMMA * nv * (1 - D[t]) - V[t]\n",
    "                last = delta + GAMMA * LAM * (1 - D[t]) * last; adv[t] = last\n",
    "            ret = adv + V\n",
    "\n",
    "        # ---------- 3) PPO 更新 ----------\n",
    "        b = [x.reshape(N_STEPS * N_ENVS, *x.shape[2:]) for x in (O, A, LP, adv, ret)]\n",
    "        for _ in range(EPOCHS):\n",
    "            for j in torch.randperm(N_STEPS * N_ENVS).chunk(N_MB):\n",
    "                s, a_, lp, ad, rt = (x[j] for x in b)\n",
    "                ad = (ad - ad.mean()) / (ad.std() + 1e-8)\n",
    "                d = Categorical(logits=actor(s)); ratio = (d.log_prob(a_) - lp).exp()\n",
    "                pg = -torch.min(ratio * ad, ratio.clamp(1 - CLIP, 1 + CLIP) * ad).mean()\n",
    "                loss = pg + 0.5 * ((critic(s).squeeze(-1) - rt) ** 2).mean() - ENT * d.entropy().mean()\n",
    "                opt.zero_grad(); loss.backward(); nn.utils.clip_grad_norm_(params, 0.5); opt.step()\n",
    "        if it % (5 if QUICK else 49) == 0:\n",
    "            score = evaluate(lambda o: int(actor(torch.as_tensor(o, dtype=torch.float32)).argmax()))\n",
    "            log.append((it * N_ENVS * N_STEPS, score))\n",
    "            print(f\"  GAIL[{reward}] 种子{seed} 步数 {it * N_ENVS * N_STEPS:6d} 最近20回合环境回报 {np.mean(env_returns[-20:]):6.1f}  贪心评估 {score:6.1f}\", flush=True)\n",
    "    return log"
   ],
   "id": "advanced-gail_cartpole-py-9"
  },
  {
   "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",
    "    rewards = [sys.argv[2]] if len(sys.argv) > 2 else [\"softplus\", \"logit\"]\n",
    "    for sd in seeds:\n",
    "        S, A = demos(sd)\n",
    "        print(f\"种子{sd}: 专家数据 {len(S)} 个 (s, a)；BC 贪心评估 {evaluate(bc(S, A, sd)):.1f}\")\n",
    "        for rw in rewards:\n",
    "            log = gail(S, A, sd, rw)\n",
    "            print(f\"种子{sd}: GAIL[{rw}] 最终贪心评估 {log[-1][1]:.1f}\")"
   ],
   "id": "advanced-gail_cartpole-py-10"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}