{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TRPO：CartPole-v1（共轭梯度 + 回溯线搜索 + KL 约束）。\n",
    "\n",
    "\n",
    "\n",
    "配套教程：[TRPO：信赖域策略优化](https://dqn.cn/advanced/trpo)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/trpo_cartpole.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-trpo_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-trpo_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",
    "from torch.distributions import Categorical, kl_divergence\n",
    "\n",
    "GAMMA, LAM, MAX_KL, CG_ITERS, DAMPING = 0.99, 0.97, 0.01, 10, 0.1\n",
    "N_ENVS, N_STEPS, ITERS = 8, 256, 50          # 每轮 2048 步，共约 10 万步\n",
    "torch.manual_seed(0)\n",
    "envs = gym.vector.SyncVectorEnv([lambda: gym.make(\"CartPole-v1\") for _ in range(N_ENVS)],\n",
    "                                autoreset_mode=gym.vector.AutoresetMode.SAME_STEP)\n",
    "envs = gym.wrappers.vector.RecordEpisodeStatistics(envs)\n",
    "obs_dim, n_actions = envs.single_observation_space.shape[0], envs.single_action_space.n\n",
    "\n",
    "policy = nn.Sequential(nn.Linear(obs_dim, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, n_actions))\n",
    "value = nn.Sequential(nn.Linear(obs_dim, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 1))\n",
    "v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)"
   ],
   "id": "advanced-trpo_cartpole-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def flat_grad(y, params, **kw):\n",
    "    return torch.cat([g.reshape(-1) for g in torch.autograd.grad(y, params, **kw)])"
   ],
   "id": "advanced-trpo_cartpole-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_flat(params, flat):\n",
    "    i = 0\n",
    "    for p in params:\n",
    "        n = p.numel(); p.data.copy_(flat[i:i + n].view_as(p)); i += n"
   ],
   "id": "advanced-trpo_cartpole-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conjugate_gradient(Avp, b, iters):\n",
    "    \"\"\"求解 F x = b，只需要 Fisher-向量积 Avp(v) = F v，无需显式构造 F。\"\"\"\n",
    "    x, r, p = torch.zeros_like(b), b.clone(), b.clone()\n",
    "    rr = r @ r\n",
    "    for _ in range(iters):\n",
    "        Ap = Avp(p)\n",
    "        alpha = rr / (p @ Ap)\n",
    "        x += alpha * p; r -= alpha * Ap\n",
    "        rr_new = r @ r\n",
    "        if rr_new < 1e-10:\n",
    "            break\n",
    "        p = r + (rr_new / rr) * p; rr = rr_new\n",
    "    return x"
   ],
   "id": "advanced-trpo_cartpole-py-5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = list(policy.parameters())\n",
    "obs = torch.as_tensor(envs.reset(seed=0)[0], dtype=torch.float32)\n",
    "ep_returns = []\n",
    "for it in range(ITERS):\n",
    "    # ---------- 采样 ----------\n",
    "    B_obs, B_act, B_rew, B_done, B_val = [], [], [], [], []\n",
    "    for t in range(N_STEPS):\n",
    "        with torch.no_grad():\n",
    "            a = Categorical(logits=policy(obs)).sample()\n",
    "            B_val.append(value(obs).squeeze(-1))\n",
    "        nobs, r, term, trunc, info = envs.step(a.numpy())\n",
    "        r = torch.as_tensor(r, dtype=torch.float32)\n",
    "        for i in np.flatnonzero(trunc):\n",
    "            with torch.no_grad():\n",
    "                r[i] += GAMMA * value(torch.as_tensor(info[\"final_obs\"][i], dtype=torch.float32)).item()\n",
    "        if \"episode\" in info:\n",
    "            ep_returns += list(info[\"episode\"][\"r\"][info[\"_episode\"]])\n",
    "        B_obs.append(obs); B_act.append(a); B_rew.append(r)\n",
    "        B_done.append(torch.as_tensor(term | trunc, dtype=torch.float32))\n",
    "        obs = torch.as_tensor(nobs, dtype=torch.float32)\n",
    "    with torch.no_grad():\n",
    "        next_v = value(obs).squeeze(-1)\n",
    "    adv, gae = [], 0\n",
    "    for t in reversed(range(N_STEPS)):\n",
    "        v_next = next_v if t == N_STEPS - 1 else B_val[t + 1]\n",
    "        delta = B_rew[t] + GAMMA * v_next * (1 - B_done[t]) - B_val[t]\n",
    "        gae = delta + GAMMA * LAM * (1 - B_done[t]) * gae\n",
    "        adv.insert(0, gae)\n",
    "    adv, vals = torch.stack(adv).reshape(-1), torch.stack(B_val).reshape(-1)\n",
    "    ret = adv + vals\n",
    "    adv = (adv - adv.mean()) / (adv.std() + 1e-8)\n",
    "    O, A = torch.stack(B_obs).reshape(-1, obs_dim), torch.stack(B_act).reshape(-1)\n",
    "\n",
    "    # ---------- 策略更新：max  g^T Δθ   s.t.  ½ Δθ^T F Δθ ≤ δ ----------\n",
    "    with torch.no_grad():\n",
    "        old_dist = Categorical(logits=policy(O))\n",
    "        old_logp = old_dist.log_prob(A)\n",
    "\n",
    "    def surrogate():\n",
    "        return ((policy_logp() - old_logp).exp() * adv).mean()\n",
    "\n",
    "    def policy_logp():\n",
    "        return Categorical(logits=policy(O)).log_prob(A)\n",
    "\n",
    "    def mean_kl():\n",
    "        return kl_divergence(old_dist, Categorical(logits=policy(O))).mean()\n",
    "\n",
    "    g = flat_grad(surrogate(), params)\n",
    "\n",
    "    def fvp(v):                                     # Fisher-向量积（KL 的 Hessian 乘向量）\n",
    "        kl_g = flat_grad(mean_kl(), params, create_graph=True)\n",
    "        return flat_grad(kl_g @ v, params) + DAMPING * v\n",
    "\n",
    "    step_dir = conjugate_gradient(fvp, g, CG_ITERS)\n",
    "    step_size = torch.sqrt(2 * MAX_KL / (step_dir @ fvp(step_dir)))\n",
    "    full_step = step_size * step_dir\n",
    "    old_params = torch.cat([p.data.reshape(-1) for p in params])\n",
    "    with torch.no_grad():\n",
    "        old_surr = surrogate()\n",
    "    for frac in 0.5 ** np.arange(10):               # 回溯线搜索：满足 KL 约束且目标提升\n",
    "        set_flat(params, old_params + frac * full_step)\n",
    "        with torch.no_grad():\n",
    "            if mean_kl() <= MAX_KL and surrogate() > old_surr:\n",
    "                break\n",
    "    else:\n",
    "        set_flat(params, old_params)\n",
    "\n",
    "    # ---------- 价值函数回归 ----------\n",
    "    for _ in range(10):\n",
    "        for idx in torch.randperm(len(O)).split(256):\n",
    "            v_loss = ((value(O[idx]).squeeze(-1) - ret[idx]) ** 2).mean()\n",
    "            v_opt.zero_grad(); v_loss.backward(); v_opt.step()\n",
    "\n",
    "    if (it + 1) % 10 == 0:\n",
    "        print(f\"iter {it + 1:3d} | steps {(it + 1) * N_ENVS * N_STEPS:6d} | 最近20回合平均回报 {np.mean(ep_returns[-20:]):6.1f}\")\n",
    "\n",
    "# ---------- 训练结束：用确定性（argmax）策略评估 10 个回合 ----------\n",
    "eval_env = gym.make(\"CartPole-v1\")\n",
    "scores = []\n",
    "for i in range(10):\n",
    "    o, _ = eval_env.reset(seed=10_000 + i)\n",
    "    done, total = False, 0.0\n",
    "    while not done:\n",
    "        with torch.no_grad():\n",
    "            act = int(policy(torch.as_tensor(o, dtype=torch.float32)).argmax())\n",
    "        o, r, te, tr, _ = eval_env.step(act)\n",
    "        total += r; done = te or tr\n",
    "    scores.append(total)\n",
    "print(f\"贪心评估（10 回合）平均回报: {np.mean(scores):.1f}\")"
   ],
   "id": "advanced-trpo_cartpole-py-6"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}