{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PPO-Clip：CartPole-v1，4 个并行环境（参考 CleanRL 的单文件实现风格）。\n",
    "\n",
    "\n",
    "\n",
    "配套教程：[PPO：近端策略优化](https://dqn.cn/advanced/ppo)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/ppo_cartpole.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-ppo_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-ppo_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\n",
    "\n",
    "N_ENVS, N_STEPS, TOTAL_STEPS = 4, 128, 200_000\n",
    "GAMMA, LAM, LR = 0.99, 0.95, 2.5e-4\n",
    "CLIP, EPOCHS, N_MINIBATCH, ENT_COEF, VF_COEF, MAX_GRAD = 0.2, 10, 4, 0.01, 0.5, 0.5\n",
    "torch.manual_seed(0)\n",
    "\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"
   ],
   "id": "advanced-ppo_cartpole-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def layer(i, o, std=np.sqrt(2)):\n",
    "    l = nn.Linear(i, o)\n",
    "    nn.init.orthogonal_(l.weight, std); nn.init.zeros_(l.bias)   # 正交初始化\n",
    "    return l"
   ],
   "id": "advanced-ppo_cartpole-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actor = nn.Sequential(layer(obs_dim, 64), nn.Tanh(), layer(64, 64), nn.Tanh(), layer(64, n_actions, 0.01))\n",
    "critic = nn.Sequential(layer(obs_dim, 64), nn.Tanh(), layer(64, 64), nn.Tanh(), layer(64, 1, 1.0))\n",
    "params = list(actor.parameters()) + list(critic.parameters())\n",
    "opt = torch.optim.Adam(params, lr=LR, eps=1e-5)\n",
    "\n",
    "buf_obs = torch.zeros(N_STEPS, N_ENVS, obs_dim)\n",
    "buf_act, buf_logp, buf_rew, buf_done, buf_val = (torch.zeros(N_STEPS, N_ENVS) for _ in range(5))\n",
    "obs = torch.as_tensor(envs.reset(seed=0)[0], dtype=torch.float32)\n",
    "ep_returns, n_updates = [], TOTAL_STEPS // (N_ENVS * N_STEPS)\n",
    "\n",
    "for update in range(n_updates):\n",
    "    for g in opt.param_groups:                     # 学习率线性退火\n",
    "        g[\"lr\"] = LR * (1 - update / n_updates)\n",
    "    # ---------- 1) 用旧策略采样 ----------\n",
    "    for t in range(N_STEPS):\n",
    "        with torch.no_grad():\n",
    "            dist = Categorical(logits=actor(obs))\n",
    "            a = dist.sample()\n",
    "            buf_val[t] = critic(obs).squeeze(-1)\n",
    "        next_obs, r, term, trunc, info = envs.step(a.numpy())\n",
    "        r = torch.as_tensor(r, dtype=torch.float32)\n",
    "        if trunc.any():                            # 截断时自举最终状态价值\n",
    "            for i in np.flatnonzero(trunc):\n",
    "                with torch.no_grad():\n",
    "                    r[i] += GAMMA * critic(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",
    "        buf_obs[t], buf_act[t], buf_logp[t] = obs, a, dist.log_prob(a)\n",
    "        buf_rew[t], buf_done[t] = r, torch.as_tensor(term | trunc, dtype=torch.float32)\n",
    "        obs = torch.as_tensor(next_obs, dtype=torch.float32)\n",
    "\n",
    "    # ---------- 2) GAE ----------\n",
    "    with torch.no_grad():\n",
    "        next_v = critic(obs).squeeze(-1)\n",
    "        adv, gae = torch.zeros(N_STEPS, N_ENVS), 0\n",
    "        for t in reversed(range(N_STEPS)):\n",
    "            v_next = next_v if t == N_STEPS - 1 else buf_val[t + 1]\n",
    "            delta = buf_rew[t] + GAMMA * v_next * (1 - buf_done[t]) - buf_val[t]\n",
    "            gae = delta + GAMMA * LAM * (1 - buf_done[t]) * gae\n",
    "            adv[t] = gae\n",
    "        ret = adv + buf_val\n",
    "\n",
    "    # ---------- 3) 多轮小批量优化裁剪目标 ----------\n",
    "    b_obs, b_act, b_logp = buf_obs.reshape(-1, obs_dim), buf_act.reshape(-1), buf_logp.reshape(-1)\n",
    "    b_adv, b_ret = adv.reshape(-1), ret.reshape(-1)\n",
    "    batch = N_ENVS * N_STEPS\n",
    "    mb = batch // N_MINIBATCH\n",
    "    for _ in range(EPOCHS):\n",
    "        perm = torch.randperm(batch)\n",
    "        for start in range(0, batch, mb):\n",
    "            idx = perm[start:start + mb]\n",
    "            dist = Categorical(logits=actor(b_obs[idx]))\n",
    "            logp = dist.log_prob(b_act[idx])\n",
    "            ratio = (logp - b_logp[idx]).exp()                       # r_t(θ)\n",
    "            a_mb = (b_adv[idx] - b_adv[idx].mean()) / (b_adv[idx].std() + 1e-8)\n",
    "            pg_loss = -torch.min(ratio * a_mb, ratio.clamp(1 - CLIP, 1 + CLIP) * a_mb).mean()\n",
    "            v_loss = ((critic(b_obs[idx]).squeeze(-1) - b_ret[idx]) ** 2).mean()\n",
    "            loss = pg_loss + VF_COEF * v_loss - ENT_COEF * dist.entropy().mean()\n",
    "            opt.zero_grad(); loss.backward()\n",
    "            nn.utils.clip_grad_norm_(params, MAX_GRAD)\n",
    "            opt.step()\n",
    "\n",
    "    steps = (update + 1) * batch\n",
    "    if (update + 1) % 40 == 0 or update == n_updates - 1:\n",
    "        print(f\"steps {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(actor(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-ppo_cartpole-py-4"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}