{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# REINFORCE（带价值基线）：CartPole-v1。\n",
    "\n",
    "\n",
    "\n",
    "配套教程：[策略梯度：REINFORCE](https://dqn.cn/advanced/reinforce)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/reinforce_cartpole.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-reinforce_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-reinforce_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",
    "GAMMA, LR_PI, LR_V, EPISODES = 0.99, 1e-3, 1e-3, 1000\n",
    "torch.manual_seed(0)\n",
    "env = gym.make(\"CartPole-v1\")\n",
    "obs_dim, n_actions = env.observation_space.shape[0], env.action_space.n\n",
    "\n",
    "policy = nn.Sequential(nn.Linear(obs_dim, 64), nn.Tanh(), nn.Linear(64, n_actions))   # 输出 logits\n",
    "value = nn.Sequential(nn.Linear(obs_dim, 64), nn.Tanh(), nn.Linear(64, 1))            # 基线 b(s) ≈ v(s)\n",
    "opt_pi = torch.optim.Adam(policy.parameters(), lr=LR_PI)\n",
    "opt_v = torch.optim.Adam(value.parameters(), lr=LR_V)\n",
    "\n",
    "returns_log = []\n",
    "for ep in range(EPISODES):\n",
    "    obs, _ = env.reset(seed=ep)\n",
    "    states, log_probs, rewards = [], [], []\n",
    "    done = False\n",
    "    while not done:\n",
    "        s = torch.as_tensor(obs, dtype=torch.float32)\n",
    "        dist = Categorical(logits=policy(s))\n",
    "        a = dist.sample()\n",
    "        obs, r, term, trunc, _ = env.step(a.item())\n",
    "        states.append(s); log_probs.append(dist.log_prob(a)); rewards.append(r)\n",
    "        done = term or trunc\n",
    "\n",
    "    # 计算每一步的回报 G_t（从后往前）\n",
    "    G, returns = 0.0, []\n",
    "    for r in reversed(rewards):\n",
    "        G = r + GAMMA * G\n",
    "        returns.append(G)\n",
    "    returns = torch.tensor(returns[::-1], dtype=torch.float32)\n",
    "    states = torch.stack(states)\n",
    "\n",
    "    v = value(states).squeeze(1)\n",
    "    adv = returns - v.detach()                         # 用基线降低方差\n",
    "    adv = (adv - adv.mean()) / (adv.std() + 1e-8)      # 标准化（常用技巧）\n",
    "\n",
    "    loss_pi = -(torch.stack(log_probs) * adv).sum()    # 梯度上升 => 最小化负目标\n",
    "    loss_v = ((returns - v) ** 2).mean()\n",
    "    opt_pi.zero_grad(); loss_pi.backward(); opt_pi.step()\n",
    "    opt_v.zero_grad(); loss_v.backward(); opt_v.step()\n",
    "\n",
    "    returns_log.append(sum(rewards))\n",
    "    if (ep + 1) % 100 == 0:\n",
    "        print(f\"episode {ep + 1:4d} | 最近100回合平均回报 {np.mean(returns_log[-100:]):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-reinforce_cartpole-py-2"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}