{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 连续动作 PPO-Clip：Pendulum-v1（对角高斯策略，状态无关的可学习 log σ）。\n",
    "\n",
    "与离散版 ppo_cartpole.py 相比，主要变化：高斯策略、动作裁剪、截断时的价值自举、较小的 γ 和较小的小批量。\n",
    "用法：python ppo_pendulum.py [种子]；DQN_QUICK=1 只跑很短的版本。\n",
    "\n",
    "配套教程：[PPO：近端策略优化](https://dqn.cn/advanced/ppo)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/ppo_pendulum.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-ppo_pendulum-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_pendulum-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",
    "from torch.distributions import Normal\n",
    "\n",
    "QUICK = os.environ.get(\"DQN_QUICK\") == \"1\"\n",
    "torch.set_num_threads(1)\n",
    "N_ENVS, N_STEPS = 4, 1024\n",
    "TOTAL_STEPS = 40_960 if QUICK else 204_800\n",
    "GAMMA, LAM, LR = 0.9, 0.95, 1e-3            # Pendulum 的奖励很密集，较小的 γ 学得更快\n",
    "CLIP, EPOCHS, MINIBATCH, VF_COEF, MAX_GRAD = 0.2, 10, 64, 0.5, 0.5"
   ],
   "id": "advanced-ppo_pendulum-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); nn.init.orthogonal_(l.weight, std); nn.init.zeros_(l.bias); return l"
   ],
   "id": "advanced-ppo_pendulum-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run(seed):\n",
    "    torch.manual_seed(seed); np.random.seed(seed)\n",
    "    envs = gym.vector.SyncVectorEnv([lambda: gym.make(\"Pendulum-v1\") for _ in range(N_ENVS)],\n",
    "                                    autoreset_mode=gym.vector.AutoresetMode.SAME_STEP)\n",
    "    envs = gym.wrappers.vector.RecordEpisodeStatistics(envs)\n",
    "    obs_dim, act_dim = envs.single_observation_space.shape[0], envs.single_action_space.shape[0]\n",
    "    a_low, a_high = envs.single_action_space.low, envs.single_action_space.high\n",
    "\n",
    "    actor_mean = nn.Sequential(layer(obs_dim, 64), nn.Tanh(), layer(64, 64), nn.Tanh(), layer(64, act_dim, 0.01))\n",
    "    actor_logstd = nn.Parameter(torch.zeros(act_dim))          # 与状态无关的 log σ，初始 σ = 1\n",
    "    critic = nn.Sequential(layer(obs_dim, 64), nn.Tanh(), layer(64, 64), nn.Tanh(), layer(64, 1, 1.0))\n",
    "    params = list(actor_mean.parameters()) + [actor_logstd] + list(critic.parameters())\n",
    "    opt = torch.optim.Adam(params, lr=LR, eps=1e-5)\n",
    "\n",
    "    def dist(o):\n",
    "        return Normal(actor_mean(o), actor_logstd.exp().expand(o.shape[0], act_dim))\n",
    "\n",
    "    obs_buf = torch.zeros(N_STEPS, N_ENVS, obs_dim); act_buf = torch.zeros(N_STEPS, N_ENVS, act_dim)\n",
    "    logp_buf, rew_buf, done_buf, val_buf = (torch.zeros(N_STEPS, N_ENVS) for _ in range(4))\n",
    "    o, _ = envs.reset(seed=seed); o = torch.as_tensor(o, dtype=torch.float32)\n",
    "    returns, n_iters = [], TOTAL_STEPS // (N_ENVS * N_STEPS)\n",
    "    for it in range(1, n_iters + 1):\n",
    "        for g in opt.param_groups:\n",
    "            g[\"lr\"] = LR * (1 - (it - 1) / n_iters)               # 学习率线性退火\n",
    "        for t in range(N_STEPS):\n",
    "            with torch.no_grad():\n",
    "                d = dist(o); a = d.sample(); val_buf[t] = critic(o).squeeze(1)\n",
    "            obs_buf[t], act_buf[t], logp_buf[t] = o, a, d.log_prob(a).sum(1)   # 存未裁剪的动作，保证概率比正确\n",
    "            o2, r, term, trunc, info = envs.step(np.clip(a.numpy(), a_low, a_high))\n",
    "            rew_buf[t] = torch.as_tensor(r, dtype=torch.float32)\n",
    "            # Pendulum 只有时间截断（200 步）：被截断时把 γ·V(s_T) 加进奖励，而不是当作“价值为 0 的终止”\n",
    "            if trunc.any():\n",
    "                final = torch.as_tensor(np.stack([info[\"final_obs\"][i] if trunc[i] else o2[i] for i in range(N_ENVS)]), dtype=torch.float32)\n",
    "                with torch.no_grad():\n",
    "                    rew_buf[t] += GAMMA * critic(final).squeeze(1) * torch.as_tensor(trunc, dtype=torch.float32)\n",
    "            done_buf[t] = torch.as_tensor(term | trunc, dtype=torch.float32)\n",
    "            if \"episode\" in info:\n",
    "                returns += list(info[\"episode\"][\"r\"][info[\"_episode\"]])\n",
    "            o = torch.as_tensor(o2, dtype=torch.float32)\n",
    "\n",
    "        with torch.no_grad():                                      # GAE\n",
    "            next_v = critic(o).squeeze(1); adv = torch.zeros_like(rew_buf); last = 0\n",
    "            for t in reversed(range(N_STEPS)):\n",
    "                nv = next_v if t == N_STEPS - 1 else val_buf[t + 1]\n",
    "                delta = rew_buf[t] + GAMMA * nv * (1 - done_buf[t]) - val_buf[t]\n",
    "                last = delta + GAMMA * LAM * (1 - done_buf[t]) * last; adv[t] = last\n",
    "            ret = adv + val_buf\n",
    "\n",
    "        b = [x.reshape(N_STEPS * N_ENVS, *x.shape[2:]) for x in (obs_buf, act_buf, logp_buf, adv, ret)]\n",
    "        for _ in range(EPOCHS):\n",
    "            perm = torch.randperm(N_STEPS * N_ENVS)\n",
    "            for i in range(0, len(perm), MINIBATCH):\n",
    "                j = perm[i:i + MINIBATCH]\n",
    "                s, a, lp_old, A, R = (x[j] for x in b)\n",
    "                A = (A - A.mean()) / (A.std() + 1e-8)\n",
    "                d = dist(s); ratio = (d.log_prob(a).sum(1) - lp_old).exp()\n",
    "                pg = -torch.min(ratio * A, ratio.clamp(1 - CLIP, 1 + CLIP) * A).mean()\n",
    "                v_loss = ((critic(s).squeeze(1) - R) ** 2).mean()\n",
    "                loss = pg + VF_COEF * v_loss                        # 连续控制中通常不加熵奖励\n",
    "                opt.zero_grad(); loss.backward(); nn.utils.clip_grad_norm_(params, MAX_GRAD); opt.step()\n",
    "        if it % (2 if QUICK else 10) == 0:\n",
    "            print(f\"  种子{seed} 步数 {it * N_ENVS * N_STEPS:7d} | 最近20回合平均回报 {np.mean(returns[-20:]):8.1f} | σ={actor_logstd.exp().mean().item():.2f}\", flush=True)\n",
    "\n",
    "    # 评估：用高斯策略的均值作为确定性动作\n",
    "    test, scores = gym.make(\"Pendulum-v1\"), []\n",
    "    for k in range(10):\n",
    "        x, _ = test.reset(seed=1000 + k); done, G = False, 0.0\n",
    "        while not done:\n",
    "            with torch.no_grad():\n",
    "                a = actor_mean(torch.as_tensor(x, dtype=torch.float32)[None])[0].numpy()\n",
    "            x, r, te, tr, _ = test.step(np.clip(a, a_low, a_high)); G += r; done = te or tr\n",
    "        scores.append(G)\n",
    "    return float(np.mean(scores))"
   ],
   "id": "advanced-ppo_pendulum-py-4"
  },
  {
   "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",
    "    for sd in seeds:\n",
    "        print(f\"种子{sd}: 确定性策略评估（10 回合）平均回报 {run(sd):.1f}\", flush=True)"
   ],
   "id": "advanced-ppo_pendulum-py-5"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}