{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 多智能体 PPO：IPPO 与 MAPPO 在 MPE simple_spread（3 个智能体合作覆盖 3 个地标）上的对比。\n",
    "\n",
    "两者都让 3 个智能体共享同一个策略网络（输入加上智能体编号的 one-hot），区别只在评论家（价值网络）：\n",
    "  IPPO：V(o_i)，只看自己的局部观测；MAPPO：V(o_1, o_2, o_3, i)，训练时看到所有智能体的观测（CTDE）。\n",
    "依赖：pip install mpe2（MPE 已从 PettingZoo 中独立出来）。\n",
    "用法：python ippo_mappo_spread.py [ippo|mappo] [种子]；DQN_QUICK=1 时只跑很短的版本用于检查代码。\n",
    "\n",
    "配套教程：[多智能体强化学习](https://dqn.cn/advanced/multi-agent)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/ippo_mappo_spread.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-ippo_mappo_spread-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 [\"mpe2\"] 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\", \"mpe2\"], check=True)\n",
    "sys.argv = sys.argv[:1]   # notebook 内核会带上 -f 等参数，清掉以免脚本把它们当成命令行参数\n",
    "# os.environ[\"DQN_QUICK\"] = \"1\"   # 取消注释：支持该开关的脚本只跑很短的版本，用于快速检查"
   ],
   "id": "advanced-ippo_mappo_spread-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from mpe2 import simple_spread_v3\n",
    "from torch.distributions import Categorical\n",
    "\n",
    "QUICK = os.environ.get(\"DQN_QUICK\") == \"1\"\n",
    "torch.set_num_threads(1)\n",
    "N, T = 3, 25                                   # 智能体数量、每回合步数\n",
    "EPISODES_PER_ITER, ITERS = 32, (10 if QUICK else 1000)\n",
    "GAMMA, LAM, LR, CLIP, EPOCHS, N_MB, ENT = 0.99, 0.95, 7e-4, 0.2, 5, 4, 0.01"
   ],
   "id": "advanced-ippo_mappo_spread-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_env():\n",
    "    return simple_spread_v3.parallel_env(N=N, max_cycles=T, continuous_actions=False)"
   ],
   "id": "advanced-ippo_mappo_spread-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mlp(i, o, std=1.0):\n",
    "    net = nn.Sequential(nn.Linear(i, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, o))\n",
    "    for m in net:\n",
    "        if isinstance(m, nn.Linear):\n",
    "            nn.init.orthogonal_(m.weight, np.sqrt(2)); nn.init.zeros_(m.bias)\n",
    "    nn.init.orthogonal_(net[-1].weight, std)\n",
    "    return net"
   ],
   "id": "advanced-ippo_mappo_spread-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run(method, seed):\n",
    "    torch.manual_seed(seed); np.random.seed(seed)\n",
    "    env = make_env(); agents = env.possible_agents\n",
    "    obs_dim, n_act = env.observation_space(agents[0]).shape[0], env.action_space(agents[0]).n\n",
    "    eye = torch.eye(N)\n",
    "    actor = mlp(obs_dim + N, n_act, 0.01)\n",
    "    critic_in = (obs_dim * N if method == \"mappo\" else obs_dim) + N\n",
    "    critic = mlp(critic_in, 1)\n",
    "    opt = torch.optim.Adam(list(actor.parameters()) + list(critic.parameters()), lr=LR, eps=1e-5)\n",
    "\n",
    "    def critic_input(o):                       # o: [..., N, obs_dim]\n",
    "        if method == \"mappo\":\n",
    "            g = o.reshape(*o.shape[:-2], 1, N * obs_dim).expand(*o.shape[:-2], N, N * obs_dim)\n",
    "            return torch.cat([g, eye.expand(*o.shape[:-2], N, N)], -1)\n",
    "        return torch.cat([o, eye.expand(*o.shape[:-2], N, N)], -1)\n",
    "\n",
    "    t0, log, ep_count = time.time(), [], 0\n",
    "    for it in range(1, ITERS + 1):\n",
    "        B = EPISODES_PER_ITER\n",
    "        O = torch.zeros(B, T, N, obs_dim); A = torch.zeros(B, T, N, dtype=torch.long)\n",
    "        LP = torch.zeros(B, T, N); R = torch.zeros(B, T, N); V = torch.zeros(B, T + 1, N)\n",
    "        for b in range(B):\n",
    "            o, _ = env.reset(seed=seed * 1_000_000 + ep_count); ep_count += 1\n",
    "            for t in range(T):\n",
    "                ot = torch.as_tensor(np.stack([o[a] for a in agents]), dtype=torch.float32)\n",
    "                with torch.no_grad():\n",
    "                    d = Categorical(logits=actor(torch.cat([ot, eye], -1)))\n",
    "                    act = d.sample(); V[b, t] = critic(critic_input(ot)).squeeze(-1)\n",
    "                O[b, t], A[b, t], LP[b, t] = ot, act, d.log_prob(act)\n",
    "                o, r, term, trunc, _ = env.step({a: int(act[i]) for i, a in enumerate(agents)})\n",
    "                R[b, t] = torch.as_tensor([r[a] for a in agents], dtype=torch.float32)\n",
    "            ot = torch.as_tensor(np.stack([o[a] for a in agents]), dtype=torch.float32)\n",
    "            with torch.no_grad():               # 回合是被 max_cycles 截断的，最后一步用 V(s_T) 自举\n",
    "                V[b, T] = critic(critic_input(ot)).squeeze(-1)\n",
    "        adv = torch.zeros(B, T, N); last = torch.zeros(B, N)\n",
    "        for t in reversed(range(T)):\n",
    "            delta = R[:, t] + GAMMA * V[:, t + 1] - V[:, t]\n",
    "            last = delta + GAMMA * LAM * last; adv[:, t] = last\n",
    "        ret = adv + V[:, :T]\n",
    "        flat = lambda x: x.reshape(B * T, *x.shape[2:])\n",
    "        Of, Af, LPf, ADVf, RETf = map(flat, (O, A, LP, adv, ret))\n",
    "        for _ in range(EPOCHS):\n",
    "            perm = torch.randperm(B * T)\n",
    "            for j in perm.chunk(N_MB):\n",
    "                o_, a_, lp_, ad_, rt_ = Of[j], Af[j], LPf[j], ADVf[j], RETf[j]\n",
    "                ad_ = (ad_ - ad_.mean()) / (ad_.std() + 1e-8)\n",
    "                d = Categorical(logits=actor(torch.cat([o_, eye.expand(len(j), N, N)], -1)))\n",
    "                ratio = (d.log_prob(a_) - lp_).exp()\n",
    "                pg = -torch.min(ratio * ad_, ratio.clamp(1 - CLIP, 1 + CLIP) * ad_).mean()\n",
    "                v_loss = ((critic(critic_input(o_)).squeeze(-1) - rt_) ** 2).mean()\n",
    "                loss = pg + 0.5 * v_loss - ENT * d.entropy().mean()\n",
    "                opt.zero_grad(); loss.backward()\n",
    "                nn.utils.clip_grad_norm_(list(actor.parameters()) + list(critic.parameters()), 0.5); opt.step()\n",
    "        # 每个智能体一个回合的累计奖励（3 个智能体取平均）\n",
    "        log.append(R.sum(1).mean().item())\n",
    "        if it % (5 if QUICK else 25) == 0:\n",
    "            print(f\"  {method} 种子{seed} 迭代{it:4d} 环境步数 {it * B * T:7d}  最近25次迭代平均回合回报 {np.mean(log[-25:]):7.2f}  （{time.time() - t0:.0f}s）\", flush=True)\n",
    "    return log"
   ],
   "id": "advanced-ippo_mappo_spread-py-5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_baseline(episodes=200):\n",
    "    env = make_env(); rng = np.random.default_rng(0); tot = []\n",
    "    for e in range(episodes):\n",
    "        env.reset(seed=e); G = 0.0\n",
    "        for _ in range(T):\n",
    "            _, r, _, _, _ = env.step({a: int(rng.integers(5)) for a in env.agents}); G += np.mean(list(r.values()))\n",
    "        tot.append(G)\n",
    "    return np.mean(tot)"
   ],
   "id": "advanced-ippo_mappo_spread-py-6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    methods, seeds = (\"ippo\", \"mappo\"), ([0] if QUICK else [0, 1])\n",
    "    if len(sys.argv) == 3:\n",
    "        methods, seeds = (sys.argv[1],), [int(sys.argv[2])]\n",
    "    else:\n",
    "        print(f\"随机策略的平均回合回报（200 回合）：{random_baseline():.2f}\")\n",
    "    for m in methods:\n",
    "        for sd in seeds:\n",
    "            log = run(m, sd)\n",
    "            print(f\"{m:5s} 种子{sd}: 最后 25 次迭代平均回合回报 {np.mean(log[-25:]):.2f}\")"
   ],
   "id": "advanced-ippo_mappo_spread-py-7"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}