{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 目标导向强化学习：DDPG 与 DDPG + HER（后见经验回放），在自建的二维“到达目标点”环境上对比。\n",
    "\n",
    "环境是稀疏奖励：没到达目标时每步 -1，到达（距离 < 0.1）时为 0。\n",
    "用法：python her_point_goal.py            # 默认 3 个种子 × 2 种方法\n",
    "      DQN_QUICK=1 python her_point_goal.py  # 缩短训练，仅用于检查代码能否运行\n",
    "\n",
    "配套教程：[目标导向强化学习与 HER](https://dqn.cn/advanced/goal-conditioned)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/advanced/her_point_goal.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "advanced-her_point_goal-py-0"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 环境准备：在 Colab 等缺少依赖的环境中自动安装（本地已装好依赖时什么也不做）\n",
    "import importlib.util, os, subprocess, sys\n",
    "sys.argv = sys.argv[:1]   # notebook 内核会带上 -f 等参数，清掉以免脚本把它们当成命令行参数\n",
    "# os.environ[\"DQN_QUICK\"] = \"1\"   # 取消注释：支持该开关的脚本只跑很短的版本，用于快速检查"
   ],
   "id": "advanced-her_point_goal-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import os\n",
    "\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)"
   ],
   "id": "advanced-her_point_goal-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PointGoalEnv:\n",
    "    \"\"\"二维平面 [-1, 1]²。智能体从原点出发，每步移动 0.1·a（a ∈ [-1, 1]²），\n",
    "    目标点在正方形内随机采样且离原点至少 0.5。一个回合固定 50 步，不提前终止。\"\"\"\n",
    "    T, STEP, TOL = 50, 0.1, 0.1\n",
    "\n",
    "    def __init__(self, seed):\n",
    "        self.rng = np.random.default_rng(seed)\n",
    "\n",
    "    def reset(self):\n",
    "        self.pos = np.zeros(2)\n",
    "        while True:\n",
    "            self.goal = self.rng.uniform(-1, 1, 2)\n",
    "            if np.linalg.norm(self.goal) >= 0.5:\n",
    "                break\n",
    "        return self.pos.copy(), self.goal.copy()\n",
    "\n",
    "    def step(self, a):\n",
    "        self.pos = np.clip(self.pos + self.STEP * np.clip(a, -1, 1), -1, 1)\n",
    "        return self.pos.copy()\n",
    "\n",
    "    @classmethod\n",
    "    def reward(cls, achieved, goal):\n",
    "        \"\"\"奖励只取决于“到达的位置”和“目标”，因此可以对任意目标重新计算——这正是 HER 需要的性质。\"\"\"\n",
    "        return -(np.linalg.norm(achieved - goal, axis=-1) >= cls.TOL).astype(np.float32)"
   ],
   "id": "advanced-her_point_goal-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mlp(i, o, out=None):\n",
    "    return nn.Sequential(nn.Linear(i, 256), nn.ReLU(), nn.Linear(256, 256), nn.ReLU(), nn.Linear(256, o),\n",
    "                         *([out] if out else []))"
   ],
   "id": "advanced-her_point_goal-py-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(use_her, seed, episodes, updates=40, batch=256, gamma=0.98, tau=0.05, k_future=4):\n",
    "    rng = np.random.default_rng(seed); torch.manual_seed(seed)\n",
    "    env, test_env = PointGoalEnv(seed), PointGoalEnv(10_000 + seed)\n",
    "    actor, critic = mlp(4, 2, nn.Tanh()), mlp(6, 1)\n",
    "    actor_t, critic_t = copy.deepcopy(actor), copy.deepcopy(critic)\n",
    "    opt_a = torch.optim.Adam(actor.parameters(), lr=1e-3)\n",
    "    opt_c = torch.optim.Adam(critic.parameters(), lr=1e-3)\n",
    "    cap = episodes * env.T * (1 + k_future)\n",
    "    buf = {k: torch.zeros(cap, d) for k, d in ((\"s\", 2), (\"g\", 2), (\"a\", 2), (\"r\", 1), (\"s2\", 2))}\n",
    "    n = 0\n",
    "    clip_q = -1.0 / (1.0 - gamma)          # 奖励 ∈ {-1, 0}，所以 Q ∈ [-1/(1-γ), 0]\n",
    "\n",
    "    def act(s, g, noise):\n",
    "        with torch.no_grad():\n",
    "            a = actor(torch.as_tensor(np.concatenate([s, g]), dtype=torch.float32)).numpy()\n",
    "        if noise:\n",
    "            a = np.clip(a + 0.2 * rng.standard_normal(2), -1, 1)\n",
    "            if rng.random() < 0.2:          # 20% 概率完全随机，保证探索\n",
    "                a = rng.uniform(-1, 1, 2)\n",
    "        return a\n",
    "\n",
    "    def evaluate(n=50):\n",
    "        ok = 0\n",
    "        for _ in range(n):\n",
    "            s, g = test_env.reset()\n",
    "            for _ in range(env.T):\n",
    "                s = test_env.step(act(s, g, noise=False))\n",
    "            ok += np.linalg.norm(s - g) < env.TOL\n",
    "        return ok / n\n",
    "\n",
    "    log = []\n",
    "    for ep in range(1, episodes + 1):\n",
    "        s, g = env.reset(); S, A = [s], []\n",
    "        for _ in range(env.T):\n",
    "            a = act(s, g, noise=True); s = env.step(a); S.append(s); A.append(a)\n",
    "        S, A = np.array(S), np.array(A)\n",
    "        # 原始目标的转移\n",
    "        goals = [np.repeat(g[None], env.T, 0)]\n",
    "        idx = [np.arange(env.T)]\n",
    "        if use_her:\n",
    "            # “future” 策略：对每个时刻 t，再从 t 之后实际到达的位置中随机取 k 个当作“假装的目标”\n",
    "            for _ in range(k_future):\n",
    "                fut = rng.integers(np.arange(env.T) + 1, env.T + 1)\n",
    "                goals.append(S[fut]); idx.append(np.arange(env.T))\n",
    "        for gs, ix in zip(goals, idx):\n",
    "            new = {\"s\": S[ix], \"a\": A[ix], \"s2\": S[ix + 1], \"g\": gs,\n",
    "                   \"r\": PointGoalEnv.reward(S[ix + 1], gs)[:, None]}\n",
    "            for k, v in new.items():\n",
    "                buf[k][n:n + len(ix)] = torch.as_tensor(v, dtype=torch.float32)\n",
    "            n += len(ix)\n",
    "        for _ in range(updates):\n",
    "            j = torch.as_tensor(rng.integers(0, n, batch))\n",
    "            s_, g_, a_, r_, s2_ = (buf[k][j] for k in (\"s\", \"g\", \"a\", \"r\", \"s2\"))\n",
    "            with torch.no_grad():\n",
    "                a2 = actor_t(torch.cat([s2_, g_], 1))\n",
    "                y = (r_ + gamma * critic_t(torch.cat([s2_, g_, a2], 1))).clamp(clip_q, 0.0)\n",
    "            loss_c = ((critic(torch.cat([s_, g_, a_], 1)) - y) ** 2).mean()\n",
    "            opt_c.zero_grad(); loss_c.backward(); opt_c.step()\n",
    "            pi = actor(torch.cat([s_, g_], 1))\n",
    "            loss_a = -critic(torch.cat([s_, g_, pi], 1)).mean() + (pi ** 2).mean()   # 小的动作惩罚，防止饱和\n",
    "            opt_a.zero_grad(); loss_a.backward(); opt_a.step()\n",
    "            for net, tgt in ((actor, actor_t), (critic, critic_t)):\n",
    "                for p, pt in zip(net.parameters(), tgt.parameters()):\n",
    "                    pt.data.mul_(1 - tau).add_(tau * p.data)\n",
    "        if ep % (50 if QUICK else 100) == 0:\n",
    "            log.append((ep, evaluate()))\n",
    "    return log"
   ],
   "id": "advanced-her_point_goal-py-5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    episodes = 100 if QUICK else 600\n",
    "    seeds = [0] if QUICK else [0, 1, 2]\n",
    "    # 参照：完全随机的动作能到达目标的比例\n",
    "    env = PointGoalEnv(123); rng = np.random.default_rng(123); hit = 0\n",
    "    for _ in range(1000):\n",
    "        s, g = env.reset()\n",
    "        for _ in range(env.T):\n",
    "            s = env.step(rng.uniform(-1, 1, 2))\n",
    "        hit += np.linalg.norm(s - g) < env.TOL\n",
    "    print(f\"随机策略：回合结束时位于目标附近的比例 {hit / 1000:.3f}\")\n",
    "    for use_her in (False, True):\n",
    "        name = \"DDPG + HER\" if use_her else \"DDPG      \"\n",
    "        for seed in seeds:\n",
    "            log = train(use_her, seed, episodes)\n",
    "            print(f\"{name} 种子{seed} 测试成功率：\" + \"  \".join(f\"{ep}:{sr:.2f}\" for ep, sr in log), flush=True)"
   ],
   "id": "advanced-her_point_goal-py-6"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}