{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 自定义 Gymnasium 环境：一维走廊。智能体从位置 0 出发，走到位置 N-1 得到奖励 1。\n",
    "\n",
    "\n",
    "\n",
    "配套教程：[Gymnasium 上手：CartPole 与 FrozenLake](https://dqn.cn/basics/gymnasium)（DQN.CN 强化学习中文教程）\n",
    "\n",
    "源码：`examples/basics/gym_custom_env.py`，本 notebook 由 `scripts/build-notebooks.mjs` 自动生成，代码与 .py 文件一致。代码默认在 CPU 上运行，固定了随机种子；部分脚本支持设置环境变量 `DQN_QUICK=1` 只跑很短的版本（见下一个单元）。"
   ],
   "id": "basics-gym_custom_env-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": "basics-gym_custom_env-py-1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "import numpy as np\n",
    "from gymnasium import spaces\n",
    "from gymnasium.utils.env_checker import check_env"
   ],
   "id": "basics-gym_custom_env-py-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Corridor(gym.Env):\n",
    "    metadata = {\"render_modes\": [\"ansi\"], \"render_fps\": 4}\n",
    "\n",
    "    def __init__(self, length=8, render_mode=None):\n",
    "        self.length = length\n",
    "        self.observation_space = spaces.Discrete(length)\n",
    "        self.action_space = spaces.Discrete(2)          # 0 = 向左，1 = 向右\n",
    "        self.render_mode = render_mode\n",
    "\n",
    "    def reset(self, *, seed=None, options=None):\n",
    "        super().reset(seed=seed)                        # 初始化 self.np_random\n",
    "        self.pos = 0\n",
    "        return self.pos, {}\n",
    "\n",
    "    def step(self, action):\n",
    "        self.pos = int(np.clip(self.pos + (1 if action == 1 else -1), 0, self.length - 1))\n",
    "        terminated = self.pos == self.length - 1\n",
    "        reward = 1.0 if terminated else 0.0\n",
    "        return self.pos, reward, terminated, False, {}\n",
    "\n",
    "    def render(self):\n",
    "        return \"\".join(\"A\" if i == self.pos else \".\" for i in range(self.length))"
   ],
   "id": "basics-gym_custom_env-py-3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "check_env(Corridor())                                   # 检查是否符合 Gymnasium API 规范\n",
    "gym.register(id=\"Corridor-v0\", entry_point=Corridor, max_episode_steps=50)\n",
    "env = gym.make(\"Corridor-v0\", render_mode=\"ansi\")\n",
    "\n",
    "# 常用 Wrapper：记录回合统计 + 奖励变换\n",
    "env = gym.wrappers.RecordEpisodeStatistics(env)\n",
    "env = gym.wrappers.TransformReward(env, lambda r: 10 * r)\n",
    "\n",
    "obs, info = env.reset(seed=0)\n",
    "done = False\n",
    "while not done:\n",
    "    obs, r, terminated, truncated, info = env.step(1)  # 一直向右\n",
    "    done = terminated or truncated\n",
    "print(env.render())\n",
    "print(\"回合统计:\", {k: float(v) for k, v in info[\"episode\"].items()})"
   ],
   "id": "basics-gym_custom_env-py-4"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}