{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MBPO（基于模型的策略优化）与 SAC 的样本效率对比：Pendulum-v1，相同的真实交互步数。\n",
    "\n",
    "MBPO = SAC + 概率神经网络集成模型 + 从真实状态出发的短视界模型展开（这里 k = 1）。\n",
    "对照组 sac10：同样每个真实步做 10 次更新，但只用真实数据，用来区分“模型数据”与“更新次数”的作用。\n",
    "用法：python mbpo_pendulum.py [mbpo|sac|sac10] [种子]；DQN_QUICK=1 时只跑很短的版本用于检查代码。\n",
    "\n",
    "配套教程：[MBPO：基于模型的策略优化](https://dqn.cn/advanced/mbpo)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/mbpo_pendulum.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-mbpo_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-mbpo_pendulum-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "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",
    "import torch.nn.functional as F\n",
    "from torch.distributions import Normal\n",
    "\n",
    "QUICK = os.environ.get(\"DQN_QUICK\") == \"1\"\n",
    "torch.set_num_threads(1)\n",
    "GAMMA, TAU, LR, BATCH = 0.99, 0.005, 3e-4, 256\n",
    "TOTAL_STEPS = 800 if QUICK else 5_000\n",
    "START_STEPS = 500\n",
    "MODEL_EVERY, ROLLOUT_STARTS, ROLLOUT_LEN = 250, 2_000, 1   # 每 250 步重训模型并生成一批模型数据\n",
    "UTD, REAL_RATIO = 10, 0.05                                  # MBPO：每个真实步做 10 次 SAC 更新，5% 真实数据\n",
    "N_ENS, N_ELITE = 7, 5\n",
    "OBS, ACT, A_MAX = 3, 1, 2.0"
   ],
   "id": "advanced-mbpo_pendulum-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mlp(i, o, h=256):\n",
    "    return nn.Sequential(nn.Linear(i, h), nn.ReLU(), nn.Linear(h, h), nn.ReLU(), nn.Linear(h, o))"
   ],
   "id": "advanced-mbpo_pendulum-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Actor(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.net = mlp(OBS, 2 * ACT)\n",
    "\n",
    "    def forward(self, s):\n",
    "        mu, log_std = self.net(s).chunk(2, -1)\n",
    "        dist = Normal(mu, log_std.clamp(-20, 2).exp())\n",
    "        u = dist.rsample(); a = torch.tanh(u)\n",
    "        logp = dist.log_prob(u).sum(-1) - torch.log(1 - a.pow(2) + 1e-6).sum(-1)\n",
    "        return a, logp, torch.tanh(mu)"
   ],
   "id": "advanced-mbpo_pendulum-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Ensemble(nn.Module):\n",
    "    \"\"\"N 个概率网络并行：输入 (s, a)，输出 (Δs, r) 的高斯分布。\"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        def P(*shape):\n",
    "            return nn.Parameter(torch.randn(N_ENS, *shape) / np.sqrt(shape[0]))\n",
    "        self.W = nn.ParameterList([P(OBS + ACT, 200), P(200, 200), P(200, 200), P(200, 2 * (OBS + 1))])\n",
    "        self.b = nn.ParameterList([nn.Parameter(torch.zeros(N_ENS, 1, d)) for d in (200, 200, 200, 2 * (OBS + 1))])\n",
    "        self.max_lv = nn.Parameter(torch.full((1, 1, OBS + 1), 0.5)); self.min_lv = nn.Parameter(torch.full((1, 1, OBS + 1), -10.0))\n",
    "        self.elites = list(range(N_ELITE))\n",
    "\n",
    "    def forward(self, x):\n",
    "        for i, (W, b) in enumerate(zip(self.W, self.b)):\n",
    "            x = torch.baddbmm(b, x, W)\n",
    "            if i < 3:\n",
    "                x = x * torch.sigmoid(x)\n",
    "        mean, lv = x.chunk(2, -1)\n",
    "        lv = self.max_lv - F.softplus(self.max_lv - lv); lv = self.min_lv + F.softplus(lv - self.min_lv)\n",
    "        return mean, lv\n",
    "\n",
    "    def train_on(self, X, Y, rng, epochs):\n",
    "        X, Y = torch.as_tensor(X), torch.as_tensor(Y)\n",
    "        self.mu, self.sd = X.mean(0), X.std(0) + 1e-6; self.ymu, self.ysd = Y.mean(0), Y.std(0) + 1e-6\n",
    "        Xn, Yn = (X - self.mu) / self.sd, (Y - self.ymu) / self.ysd\n",
    "        n_val = min(500, len(X) // 5); perm = torch.as_tensor(rng.permutation(len(X)))\n",
    "        val, tr = perm[:n_val], perm[n_val:]\n",
    "        opt = torch.optim.Adam(self.parameters(), lr=1e-3, weight_decay=1e-5)\n",
    "        boot = tr[torch.as_tensor(rng.integers(0, len(tr), (N_ENS, len(tr))))]\n",
    "        for _ in range(epochs):\n",
    "            order = boot[:, torch.randperm(len(tr))]\n",
    "            for i in range(0, len(tr), 256):\n",
    "                j = order[:, i:i + 256]\n",
    "                mean, lv = self(Xn[j])\n",
    "                loss = (((mean - Yn[j]) ** 2) * torch.exp(-lv) + lv).mean() + 0.01 * (self.max_lv.sum() - self.min_lv.sum())\n",
    "                opt.zero_grad(); loss.backward(); opt.step()\n",
    "        with torch.no_grad():                    # 按验证集 MSE 选出最好的 N_ELITE 个成员\n",
    "            mean, _ = self(Xn[val].expand(N_ENS, -1, -1))\n",
    "            mse = ((mean - Yn[val]) ** 2).mean((1, 2))\n",
    "            self.elites = mse.argsort()[:N_ELITE].tolist()\n",
    "        return mse.sort().values[:N_ELITE].mean().item()\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def step(self, s, a, gen):\n",
    "        x = ((torch.cat([s, a], -1) - self.mu) / self.sd).expand(N_ENS, -1, -1)\n",
    "        mean, lv = self(x)\n",
    "        sample = mean + torch.exp(0.5 * lv) * torch.randn(mean.shape, generator=gen)\n",
    "        pick = torch.as_tensor(self.elites)[torch.randint(0, N_ELITE, (len(s),), generator=gen)]\n",
    "        out = sample[pick, torch.arange(len(s))] * self.ysd + self.ymu    # 每条样本随机选一个精英成员\n",
    "        return s + out[:, :OBS], out[:, OBS:]"
   ],
   "id": "advanced-mbpo_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, test_env = gym.make(\"Pendulum-v1\"), gym.make(\"Pendulum-v1\")\n",
    "    actor = Actor(); q1, q2 = mlp(OBS + ACT, 1), mlp(OBS + ACT, 1)\n",
    "    q1_t, q2_t = copy.deepcopy(q1), copy.deepcopy(q2)\n",
    "    opt_a = torch.optim.Adam(actor.parameters(), lr=LR)\n",
    "    opt_q = torch.optim.Adam(list(q1.parameters()) + list(q2.parameters()), lr=LR)\n",
    "    log_alpha = torch.zeros(1, requires_grad=True); opt_al = torch.optim.Adam([log_alpha], lr=LR)\n",
    "    real = {k: torch.zeros(TOTAL_STEPS, d) for k, d in ((\"s\", OBS), (\"a\", ACT), (\"r\", 1), (\"s2\", OBS))}\n",
    "    cap = ROLLOUT_STARTS * ROLLOUT_LEN * 8                     # 模型数据只保留最近若干批\n",
    "    fake = {k: torch.zeros(cap, d) for k, d in ((\"s\", OBS), (\"a\", ACT), (\"r\", 1), (\"s2\", OBS))}\n",
    "    n_fake = ptr = 0\n",
    "    model = Ensemble() if method == \"mbpo\" else None\n",
    "\n",
    "    def sac_update(s, a, r, s2):                # Pendulum 没有终止状态（只有时间截断），所以不需要 done\n",
    "        alpha = log_alpha.exp().detach()\n",
    "        with torch.no_grad():\n",
    "            a2, lp2, _ = actor(s2); sa2 = torch.cat([s2, a2], 1)\n",
    "            y = r + GAMMA * (torch.min(q1_t(sa2), q2_t(sa2)) - alpha * lp2[:, None])\n",
    "        sa = torch.cat([s, a], 1)\n",
    "        lq = F.mse_loss(q1(sa), y) + F.mse_loss(q2(sa), y)\n",
    "        opt_q.zero_grad(); lq.backward(); opt_q.step()\n",
    "        an, lp, _ = actor(s); san = torch.cat([s, an], 1)\n",
    "        la = (alpha * lp - torch.min(q1(san), q2(san)).squeeze(1)).mean()\n",
    "        opt_a.zero_grad(); la.backward(); opt_a.step()\n",
    "        lal = -(log_alpha * (lp.detach() - ACT)).mean()\n",
    "        opt_al.zero_grad(); lal.backward(); opt_al.step()\n",
    "        with torch.no_grad():\n",
    "            for n, t in ((q1, q1_t), (q2, q2_t)):\n",
    "                for p, tp in zip(n.parameters(), t.parameters()):\n",
    "                    tp.mul_(1 - TAU).add_(TAU * p)\n",
    "\n",
    "    def evaluate(n=5):\n",
    "        tot = 0.0\n",
    "        for k in range(n):\n",
    "            o, _ = test_env.reset(seed=10_000 + k); done = False\n",
    "            while not done:\n",
    "                with torch.no_grad():\n",
    "                    a = actor(torch.as_tensor(o, dtype=torch.float32))[2].numpy()\n",
    "                o, r, te, tr, _ = test_env.step(a * A_MAX); tot += r; done = te or tr\n",
    "        return tot / n\n",
    "\n",
    "    o, _ = env.reset(seed=seed); env.action_space.seed(seed); t0 = time.time(); log = []\n",
    "    for step in range(TOTAL_STEPS):\n",
    "        if step < START_STEPS:\n",
    "            a = rng.uniform(-1, 1, ACT).astype(np.float32)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                a = actor(torch.as_tensor(o, dtype=torch.float32))[0].numpy()\n",
    "        o2, r, te, tr, _ = env.step(a * A_MAX)\n",
    "        for k, v in ((\"s\", o), (\"a\", a), (\"r\", [r]), (\"s2\", o2)):\n",
    "            real[k][step] = torch.as_tensor(np.asarray(v, np.float32))\n",
    "        o = o2\n",
    "        if te or tr:\n",
    "            o, _ = env.reset()\n",
    "        n_real = step + 1\n",
    "        if step + 1 < START_STEPS:\n",
    "            continue\n",
    "        if method == \"mbpo\":\n",
    "            if (step + 1 - START_STEPS) % MODEL_EVERY == 0:\n",
    "                X = torch.cat([real[\"s\"][:n_real], real[\"a\"][:n_real]], 1).numpy()\n",
    "                Y = torch.cat([real[\"s2\"][:n_real] - real[\"s\"][:n_real], real[\"r\"][:n_real]], 1).numpy()\n",
    "                mse = model.train_on(X, Y, rng, epochs=5 if QUICK else 30)\n",
    "                s = real[\"s\"][torch.as_tensor(rng.integers(0, n_real, ROLLOUT_STARTS))]\n",
    "                for _ in range(ROLLOUT_LEN):          # 从真实状态出发，用当前策略在模型里走 k 步\n",
    "                    with torch.no_grad():\n",
    "                        a_m = actor(s)[0]\n",
    "                    s2_m, r_m = model.step(s, a_m, gen)\n",
    "                    m = len(s); idx = (torch.arange(m) + ptr) % cap\n",
    "                    for k, v in ((\"s\", s), (\"a\", a_m), (\"r\", r_m), (\"s2\", s2_m)):\n",
    "                        fake[k][idx] = v\n",
    "                    ptr = (ptr + m) % cap; n_fake = min(n_fake + m, cap); s = s2_m\n",
    "            for _ in range(UTD):\n",
    "                n_r = int(BATCH * REAL_RATIO)\n",
    "                i = torch.as_tensor(rng.integers(0, n_real, n_r)); j = torch.as_tensor(rng.integers(0, n_fake, BATCH - n_r))\n",
    "                sac_update(*(torch.cat([real[k][i], fake[k][j]]) for k in (\"s\", \"a\", \"r\", \"s2\")))\n",
    "        else:\n",
    "            for _ in range(UTD if method == \"sac10\" else 1):\n",
    "                i = torch.as_tensor(rng.integers(0, n_real, BATCH))\n",
    "                sac_update(*(real[k][i] for k in (\"s\", \"a\", \"r\", \"s2\")))\n",
    "        if (step + 1) % (200 if QUICK else 1000) == 0:\n",
    "            log.append((step + 1, evaluate()))\n",
    "            print(f\"  {method} 种子{seed} 真实步数 {step + 1:5d}  评估回报 {log[-1][1]:8.1f}  （{time.time() - t0:.0f}s）\", flush=True)\n",
    "    return log"
   ],
   "id": "advanced-mbpo_pendulum-py-6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    methods, seeds = (\"sac\", \"sac10\", \"mbpo\"), ([0] if QUICK else [0, 1])\n",
    "    if len(sys.argv) == 3:\n",
    "        methods, seeds = (sys.argv[1],), [int(sys.argv[2])]\n",
    "    for m in methods:\n",
    "        for sd in seeds:\n",
    "            log = run(m, sd)\n",
    "            print(f\"{m:4s} 种子{sd}: \" + \"  \".join(f\"{s}:{v:.0f}\" for s, v in log))"
   ],
   "id": "advanced-mbpo_pendulum-py-7"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}