{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A2C（同步 Advantage Actor-Critic + GAE）：CartPole-v1，8 个并行环境。\n",
    "\n",
    "\n",
    "\n",
    "配套教程：[Actor-Critic：A2C / A3C](https://dqn.cn/advanced/actor-critic)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/a2c_cartpole.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-a2c_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-a2c_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 = 8, 5, 200_000\n",
    "GAMMA, LAM, LR, ENT_COEF, VF_COEF = 0.99, 0.95, 7e-4, 0.01, 0.5\n",
    "torch.manual_seed(0)\n",
    "\n",
    "envs = gym.vector.SyncVectorEnv(\n",
    "    [lambda: gym.make(\"CartPole-v1\") for _ in range(N_ENVS)],\n",
    "    autoreset_mode=gym.vector.AutoresetMode.SAME_STEP,   # 回合结束时同一步自动 reset\n",
    ")\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-a2c_cartpole-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ActorCritic(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.actor = nn.Sequential(nn.Linear(obs_dim, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, n_actions))\n",
    "        self.critic = nn.Sequential(nn.Linear(obs_dim, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 1))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return Categorical(logits=self.actor(x)), self.critic(x).squeeze(-1)"
   ],
   "id": "advanced-a2c_cartpole-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ActorCritic()\n",
    "opt = torch.optim.Adam(model.parameters(), lr=LR)\n",
    "obs, _ = envs.reset(seed=0)\n",
    "obs = torch.as_tensor(obs, dtype=torch.float32)\n",
    "ep_returns = []\n",
    "\n",
    "for update in range(TOTAL_STEPS // (N_ENVS * N_STEPS)):\n",
    "    log_probs, values, rewards, dones, entropies = [], [], [], [], []\n",
    "    for _ in range(N_STEPS):\n",
    "        dist, v = model(obs)\n",
    "        a = dist.sample()\n",
    "        next_obs, r, term, trunc, info = envs.step(a.numpy())\n",
    "        r = torch.as_tensor(r, dtype=torch.float32)\n",
    "        # 被截断（truncated）的回合：把最终状态的价值自举进奖励，避免把时间上限误当成失败\n",
    "        if trunc.any():\n",
    "            final = torch.as_tensor(np.stack([info[\"final_obs\"][i] if trunc[i] else next_obs[i] for i in range(N_ENVS)]), dtype=torch.float32)\n",
    "            with torch.no_grad():\n",
    "                r = r + GAMMA * model(final)[1] * torch.as_tensor(trunc, dtype=torch.float32)\n",
    "        if \"episode\" in info:\n",
    "            ep_returns += list(info[\"episode\"][\"r\"][info[\"_episode\"]])\n",
    "        log_probs.append(dist.log_prob(a)); values.append(v); entropies.append(dist.entropy())\n",
    "        rewards.append(r); dones.append(torch.as_tensor(term | trunc, dtype=torch.float32))\n",
    "        obs = torch.as_tensor(next_obs, dtype=torch.float32)\n",
    "\n",
    "    # GAE(λ) 计算优势\n",
    "    with torch.no_grad():\n",
    "        next_v = model(obs)[1]\n",
    "    adv, gae = [], torch.zeros(N_ENVS)\n",
    "    for t in reversed(range(N_STEPS)):\n",
    "        v_next = next_v if t == N_STEPS - 1 else values[t + 1].detach()\n",
    "        delta = rewards[t] + GAMMA * v_next * (1 - dones[t]) - values[t].detach()\n",
    "        gae = delta + GAMMA * LAM * (1 - dones[t]) * gae\n",
    "        adv.insert(0, gae)\n",
    "    adv = torch.stack(adv)\n",
    "    values, log_probs, entropies = torch.stack(values), torch.stack(log_probs), torch.stack(entropies)\n",
    "    returns = adv + values.detach()\n",
    "\n",
    "    loss = -(log_probs * adv).mean() + VF_COEF * ((returns - values) ** 2).mean() - ENT_COEF * entropies.mean()\n",
    "    opt.zero_grad(); loss.backward()\n",
    "    nn.utils.clip_grad_norm_(model.parameters(), 0.5)\n",
    "    opt.step()\n",
    "\n",
    "    steps = (update + 1) * N_ENVS * N_STEPS\n",
    "    if steps % 20_000 == 0 and ep_returns:\n",
    "        print(f\"steps {steps:7d} | 最近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(model.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-a2c_cartpole-py-4"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}