{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型预测控制（MPC）：在 Pendulum-v1 上比较三种基于“学到的模型”的规划方法\n",
    "\n",
    "1. 随机打靶（Random Shooting）：确定性神经网络模型 + 随机采样动作序列取最好的一条\n",
    "  2. CEM：确定性模型 + 交叉熵方法迭代优化动作序列\n",
    "  3. PETS：概率神经网络集成（5 个）+ 粒子轨迹采样（TS∞）+ CEM\n",
    "奖励函数视为已知（由预测的状态直接计算），只学习状态转移。\n",
    "用法：python mpc_pendulum.py；DQN_QUICK=1 python mpc_pendulum.py 只跑很短的版本用于检查代码。\n",
    "\n",
    "配套教程：[模型预测控制：随机打靶、CEM 与 PETS](https://dqn.cn/advanced/mpc)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/mpc_pendulum.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-mpc_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-mpc_pendulum-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import time\n",
    "\n",
    "import gymnasium as gym\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "QUICK = os.environ.get(\"DQN_QUICK\") == \"1\"\n",
    "torch.set_num_threads(1)\n",
    "H, POP, ITERS, ELITE = 25, 300, 5, 30          # 规划视界、候选序列数、CEM 迭代次数、精英数\n",
    "N_ENS, PARTICLES, POP_PETS = 5, 10, 200         # PETS：集成大小、每条候选序列的粒子数、候选数（计算量较大，略减）\n",
    "EPISODES = 3 if QUICK else 10                   # 第 1 个回合用随机动作收集初始数据\n",
    "A_MAX = 2.0"
   ],
   "id": "advanced-mpc_pendulum-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reward_fn(obs, a):\n",
    "    \"\"\"Pendulum 的奖励：-(θ² + 0.1·θ̇² + 0.001·u²)，θ 由 cos、sin 恢复。\"\"\"\n",
    "    th = torch.atan2(obs[..., 1], obs[..., 0])\n",
    "    return -(th ** 2 + 0.1 * obs[..., 2] ** 2 + 0.001 * a[..., 0] ** 2)"
   ],
   "id": "advanced-mpc_pendulum-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    \"\"\"输入 (s, a)，输出 Δs 的均值（概率模型还输出对数方差）。多个成员用一个批量维度并行计算。\"\"\"\n",
    "\n",
    "    def __init__(self, n_members, probabilistic):\n",
    "        super().__init__()\n",
    "        self.n, self.prob = n_members, probabilistic\n",
    "        out = 6 if probabilistic else 3\n",
    "        def P(*shape):\n",
    "            return nn.Parameter(torch.randn(n_members, *shape) * (1 / np.sqrt(shape[0])))\n",
    "        self.W = nn.ParameterList([P(4, 200), P(200, 200), P(200, out)])\n",
    "        self.b = nn.ParameterList([nn.Parameter(torch.zeros(n_members, 1, d)) for d in (200, 200, out)])\n",
    "        self.max_logvar = nn.Parameter(torch.full((1, 1, 3), 0.5)); self.min_logvar = nn.Parameter(torch.full((1, 1, 3), -10.0))\n",
    "        self.mu = self.sd = self.dmu = self.dsd = None\n",
    "\n",
    "    def set_norm(self, X, D):\n",
    "        self.mu, self.sd = X.mean(0), X.std(0) + 1e-6\n",
    "        self.dmu, self.dsd = D.mean(0), D.std(0) + 1e-6\n",
    "\n",
    "    def forward(self, x):                        # x: [n, B, 4]（已归一化）\n",
    "        for i, (W, b) in enumerate(zip(self.W, self.b)):\n",
    "            x = torch.baddbmm(b, x, W)\n",
    "            if i < 2:\n",
    "                x = x * torch.sigmoid(x)         # Swish 激活\n",
    "        if not self.prob:\n",
    "            return x, None\n",
    "        mean, logvar = x[..., :3], x[..., 3:]\n",
    "        logvar = self.max_logvar - nn.functional.softplus(self.max_logvar - logvar)\n",
    "        logvar = self.min_logvar + nn.functional.softplus(logvar - self.min_logvar)\n",
    "        return mean, logvar\n",
    "\n",
    "    def fit(self, X, D, rng, epochs=50, batch=256):\n",
    "        X, D = torch.as_tensor(X, dtype=torch.float32), torch.as_tensor(D, dtype=torch.float32)\n",
    "        self.set_norm(X, D); Xn, Dn = (X - self.mu) / self.sd, (D - self.dmu) / self.dsd\n",
    "        opt = torch.optim.Adam(self.parameters(), lr=1e-3, weight_decay=1e-5)\n",
    "        N = len(X)\n",
    "        boot = torch.as_tensor(rng.integers(0, N, (self.n, N)))          # 每个成员各自的自助采样\n",
    "        for _ in range(epochs):\n",
    "            perm = boot[:, torch.randperm(N)]\n",
    "            for i in range(0, N, batch):\n",
    "                j = perm[:, i:i + batch]\n",
    "                mean, logvar = self(Xn[j])\n",
    "                if self.prob:                    # 高斯负对数似然\n",
    "                    loss = (((mean - Dn[j]) ** 2) * torch.exp(-logvar) + logvar).mean()\n",
    "                    loss = loss + 0.01 * (self.max_logvar.sum() - self.min_logvar.sum())\n",
    "                else:\n",
    "                    loss = ((mean - Dn[j]) ** 2).mean()\n",
    "                opt.zero_grad(); loss.backward(); opt.step()\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def predict(self, s, a, gen):                # s: [n, B, 3], a: [n, B, 1]\n",
    "        x = (torch.cat([s, a], -1) - self.mu) / self.sd\n",
    "        mean, logvar = self(x)\n",
    "        if self.prob:\n",
    "            mean = mean + torch.exp(0.5 * logvar) * torch.randn(mean.shape, generator=gen)\n",
    "        return s + mean * self.dsd + self.dmu"
   ],
   "id": "advanced-mpc_pendulum-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def plan(model, s0, method, init_mean, gen):\n",
    "    \"\"\"返回规划出的动作序列（只执行第一个动作），以及下一步用于热启动的序列。\"\"\"\n",
    "    def rollout_return(seqs):                    # seqs: [POP, H]\n",
    "        P = PARTICLES if model.prob else 1\n",
    "        n_pop = seqs.shape[0]\n",
    "        # 粒子 p 固定使用集成成员 p % N_ENS（TS∞）；确定性模型只有 1 个成员\n",
    "        s = torch.as_tensor(s0, dtype=torch.float32).expand(model.n, n_pop * P // model.n if model.prob else n_pop, 3).clone()\n",
    "        a_all = seqs.repeat_interleave(P, 0) if model.prob else seqs\n",
    "        a_all = a_all.view(model.n, -1, H) if model.prob else a_all.view(1, -1, H)\n",
    "        ret = torch.zeros(a_all.shape[:2])\n",
    "        for t in range(H):\n",
    "            a = a_all[..., t:t + 1]\n",
    "            ret += reward_fn(s, a)\n",
    "            s = model.predict(s, a, gen)\n",
    "        return ret.reshape(n_pop, P).mean(1) if model.prob else ret.reshape(n_pop)\n",
    "\n",
    "    pop = POP_PETS if model.prob else POP\n",
    "    if method == \"rs\":\n",
    "        seqs = (torch.rand(POP, H, generator=gen) * 2 - 1) * A_MAX\n",
    "        return seqs[rollout_return(seqs).argmax()], None\n",
    "    mean, std = init_mean.clone(), torch.full((H,), A_MAX)\n",
    "    for _ in range(ITERS):\n",
    "        seqs = (mean + std * torch.randn(pop, H, generator=gen)).clamp(-A_MAX, A_MAX)\n",
    "        elite = seqs[rollout_return(seqs).topk(ELITE).indices]\n",
    "        mean, std = 0.9 * elite.mean(0) + 0.1 * mean, 0.9 * elite.std(0) + 0.1 * std\n",
    "    return mean, torch.cat([mean[1:], torch.zeros(1)])"
   ],
   "id": "advanced-mpc_pendulum-py-5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run(method, seed):\n",
    "    rng = np.random.default_rng(seed); torch.manual_seed(seed); gen = torch.Generator().manual_seed(seed)\n",
    "    env = gym.make(\"Pendulum-v1\"); env.action_space.seed(seed)\n",
    "    model = Model(N_ENS, True) if method == \"pets\" else Model(1, False)\n",
    "    X, D, returns = [], [], []\n",
    "    for ep in range(EPISODES):\n",
    "        s, _ = env.reset(seed=seed * 100 + ep); G = 0.0; warm = torch.zeros(H); t0 = time.time()\n",
    "        for t in range(200):\n",
    "            if ep == 0:\n",
    "                a = env.action_space.sample()\n",
    "            else:\n",
    "                seq, nxt = plan(model, s, method, warm, gen)\n",
    "                a = np.array([float(seq[0])], np.float32)\n",
    "                warm = nxt if nxt is not None else warm\n",
    "            s2, r, term, trunc, _ = env.step(a)\n",
    "            X.append(np.concatenate([s, a])); D.append(s2 - s); G += r; s = s2\n",
    "        returns.append(G)\n",
    "        model.fit(np.array(X), np.array(D), rng, epochs=20 if QUICK else 50)\n",
    "        print(f\"  {method:4s} 种子{seed} 回合{ep}  回报 {G:8.1f}  （{time.time() - t0:.0f}s）\", flush=True)\n",
    "    return returns"
   ],
   "id": "advanced-mpc_pendulum-py-6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    seeds = [0] if QUICK else [0, 1]\n",
    "    methods = (\"rs\", \"cem\", \"pets\")\n",
    "    if len(sys.argv) == 3:                        # 也可以单独运行：python mpc_pendulum.py pets 0\n",
    "        methods, seeds = (sys.argv[1],), [int(sys.argv[2])]\n",
    "    summary = {}\n",
    "    for method in methods:\n",
    "        summary[method] = [run(method, s) for s in seeds]\n",
    "    print(\"\\n最后 3 个回合的平均回报（每个种子）：\")\n",
    "    for m, rs in summary.items():\n",
    "        print(f\"{m:4s}: \" + \"  \".join(f\"{np.mean(r[-3:]):.1f}\" for r in rs))"
   ],
   "id": "advanced-mpc_pendulum-py-7"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}