Extract shared maze code into common package (Gymnasium-style env)

2026-09-01 15:55:14 +08:00
parent d2dc1f0910
commit d54072ede4
5 changed files with 665 additions and 284 deletions
@@ -0,0 +1,303 @@
+"""Q-Learning 与 DQN 共用的迷宫环境、展示与绘图工具。"""
+from collections import deque
+
+import numpy as np
+
+
+START, END, WALL, ROAD = 0, 1, 2, 3                 # 分别是起点 终点 障碍 道路
+ACTION_DELTAS = ((-1, 0), (1, 0), (0, -1), (0, 1))  # 动作: 0上 1下 2左 3右
+ARROWS = "↑↓←→"
+CHARS = {START: "S", END: "G", WALL: "#", ROAD: "."}
+
+
+def generate_maze(rows, cols, seed=0, braid=0.0):
+    """按 seed 随机生成可解迷宫（递归回溯法），行列须为不小于 3 的奇数。
+
+    保证起点 (0,0) 到终点 (rows-1, cols-1) 连通。
+    braid: 0~1, 生成后按该概率拆掉剩余内墙形成环路。
+    0 = 完美迷宫（唯一通路, 长走廊多）；越大环路越多、可选路径越多。
+    """
+    if rows < 3 or cols < 3 or rows % 2 == 0 or cols % 2 == 0:
+        raise ValueError("行列必须为不小于 3 的奇数")
+    rng = np.random.default_rng(seed)
+    maze = np.full((rows, cols), WALL, dtype=np.int8)
+    maze[0, 0] = ROAD
+    stack = [(0, 0)]
+    while stack:
+        r, c = stack[-1]
+        neighbors = [
+            (r + dr, c + dc, r + dr // 2, c + dc // 2)
+            for dr, dc in ((-2, 0), (2, 0), (0, -2), (0, 2))
+            if 0 <= r + dr < rows and 0 <= c + dc < cols
+            and maze[r + dr, c + dc] == WALL
+        ]
+        if neighbors:
+            nr, nc, mr, mc = neighbors[int(rng.integers(len(neighbors)))]
+            maze[mr, mc] = maze[nr, nc] = ROAD
+            stack.append((nr, nc))
+        else:
+            stack.pop()
+    if braid > 0:
+        for r in range(1, rows - 1):
+            for c in range(1, cols - 1):
+                if maze[r, c] == WALL and rng.random() < braid and (
+                    (maze[r - 1, c] != WALL and maze[r + 1, c] != WALL)
+                    or (maze[r, c - 1] != WALL and maze[r, c + 1] != WALL)
+                ):
+                    maze[r, c] = ROAD   # 拆掉两侧都是路的墙, 形成环路
+    maze[0, 0], maze[rows - 1, cols - 1] = START, END
+    return maze
+
+
+class MazeEnv:
+    """Gymnasium 风格的迷宫环境（不依赖 gymnasium, 仅遵循其 API 约定）。
+
+    状态 obs 为格子坐标 (r, c)；动作: 0上 1下 2左 3右。
+    奖励: 每步 step_reward；撞墙/越界 wall_penalty（原地不动）；
+    到达终点 goal_reward 且 terminated。
+    step() 返回 (obs, reward, terminated, truncated, info)：
+      terminated = 到达终点（真正的终止状态）；
+      truncated  = 达到 max_steps 步数上限（超时截断, 非终止状态）。
+    """
+
+    def __init__(self, maze, start_pos=None,
+                 step_reward=-1.0, wall_penalty=-10.0, goal_reward=100.0, max_steps=200):
+        self.maze = np.asarray(maze)
+        self.shape = self.maze.shape
+        self.n_actions = len(ACTION_DELTAS)
+        self.start_pos = self._find(START) if start_pos is None else tuple(start_pos)
+        self.goal = self._find(END)
+
+        r, c = self.start_pos
+        if not (0 <= r < self.shape[0] and 0 <= c < self.shape[1]) or self.maze[r, c] == WALL:
+            raise ValueError(f"非法起点: {self.start_pos}")
+        if self.start_pos == self.goal:
+            raise ValueError("起点不能与终点重合")
+
+        self.step_reward = step_reward
+        self.wall_penalty = wall_penalty
+        self.goal_reward = goal_reward
+        self.max_steps = max_steps
+        self.steps = 0
+        self._pos = self.start_pos
+
+    def _find(self, cell):
+        hits = np.argwhere(self.maze == cell)
+        if len(hits) == 0:
+            raise ValueError(f"迷宫中找不到格子类型 {cell}")
+        return tuple(int(v) for v in hits[0])
+
+    def reset(self, *, seed=None, options=None):
+        """重置到起点, 返回 (obs, info)。seed/options 仅为 API 兼容保留。"""
+        self._pos = self.start_pos
+        self.steps = 0
+        return self._pos, {}
+
+    def step(self, action):
+        """执行动作, 返回 (obs, reward, terminated, truncated, info)。"""
+        if not 0 <= action < self.n_actions:
+            raise ValueError(f"非法动作: {action}")
+
+        dr, dc = ACTION_DELTAS[action]
+        nr, nc = self._pos[0] + dr, self._pos[1] + dc
+        self.steps += 1
+
+        terminated = False
+        if (
+            not (0 <= nr < self.shape[0] and 0 <= nc < self.shape[1])
+            or self.maze[nr][nc] == WALL
+        ):
+            # 遇到边界或者撞墙: 原地不动
+            reward = self.wall_penalty
+        else:
+            self._pos = (nr, nc)
+            terminated = self._pos == self.goal
+            reward = self.goal_reward if terminated else self.step_reward
+
+        truncated = not terminated and self.steps >= self.max_steps
+        return self._pos, reward, terminated, truncated, {}
+
+
+def shortest_path_len(env):
+    """BFS 求起点到终点的最短步数（不通返回 -1）。"""
+    dist = {env.start_pos: 0}
+    queue = deque([env.start_pos])
+    while queue:
+        pos = queue.popleft()
+        if pos == env.goal:
+            return dist[pos]
+        for dr, dc in ACTION_DELTAS:
+            nr, nc = pos[0] + dr, pos[1] + dc
+            nxt = (nr, nc)
+            if (0 <= nr < env.shape[0] and 0 <= nc < env.shape[1]
+                    and env.maze[nr, nc] != WALL and nxt not in dist):
+                dist[nxt] = dist[pos] + 1
+                queue.append(nxt)
+    return -1
+
+
+def evaluate(env, act, episodes=100):
+    """纯贪心策略（无探索）评测, 返回 (成功率, 成功局平均步数)。
+
+    act(pos) -> 动作编号, 由各算法提供（查表或前向推理）。
+    """
+    wins, steps_list = 0, []
+    for _ in range(episodes):
+        obs, _ = env.reset()
+        terminated = truncated = False
+        while not (terminated or truncated):
+            obs, _, terminated, truncated, _ = env.step(int(act(obs)))
+        if terminated:
+            wins += 1
+            steps_list.append(env.steps)
+    avg_steps = float(np.mean(steps_list)) if steps_list else float("nan")
+    return wins / episodes, avg_steps
+
+
+def show_maze(env):
+    print(f"迷宫 {env.shape[0]}x{env.shape[1]}（S起点 G终点 #墙）:")
+    for r in range(env.shape[0]):
+        print("".join(CHARS[env.maze[r][c]] for c in range(env.shape[1])))
+
+
+def show_policy(env, act):
+    print("学到的策略（每格最优动作）:")
+    for r in range(env.shape[0]):
+        row = ""
+        for c in range(env.shape[1]):
+            if env.maze[r][c] == WALL:
+                row += " # "
+            elif (r, c) == env.goal:
+                row += " G "
+            else:
+                row += f" {ARROWS[int(act((r, c)))]} "
+        print(row)
+
+
+def show_path(env, act):
+    print("\n贪心策略走出的一条路径:")
+    obs, path, terminated = env.reset()[0], [env.start_pos], False
+    truncated = False
+    while not (terminated or truncated):
+        obs, _, terminated, truncated, _ = env.step(int(act(obs)))
+        path.append(obs)
+    grid = [[CHARS[env.maze[r][c]] for c in range(env.shape[1])]
+            for r in range(env.shape[0])]
+    for i, (r, c) in enumerate(path):
+        if (r, c) not in (env.start_pos, env.goal):
+            grid[r][c] = str(i % 10)   # 若重复经过同一格，显示最后一步序号
+    print("\n".join(" ".join(row) for row in grid))
+    print(f"到达终点: {'是' if terminated else '否'}，共 {len(path) - 1} 步")
+
+
+# ─────────────────── 训练过程可视化（需 matplotlib） ───────────────────
+def _setup_plt():
+    import matplotlib
+
+    matplotlib.use("Agg")               # 不弹窗，直接保存图片
+    import matplotlib.pyplot as plt
+    plt.rcParams["font.sans-serif"] = ["Microsoft YaHei", "SimHei"]  # 中文字体
+    plt.rcParams["axes.unicode_minus"] = False                       # 正常显示负号
+    return plt
+
+
+def moving_avg(x, k=20):
+    x = np.asarray(x, dtype=float)
+    if len(x) < k:
+        return np.array([])
+    return np.convolve(x, np.ones(k) / k, mode="valid")
+
+
+def _plot_ma(ax, x, window, color):
+    x = np.asarray(x, dtype=float)
+    ax.plot(x, color="gray", alpha=0.35, label="逐局值")
+    ma = moving_avg(x, window)
+    if len(ma):
+        ax.plot(np.arange(window - 1, window - 1 + len(ma)), ma,
+                color=color, lw=2, label=f"滑动平均(窗口{window})")
+    ax.set_xlabel("训练局数 Episode")
+    ax.legend()
+
+
+def plot_training(env, rewards, successes, ep_steps, extra,
+                  window=20, out="training.png", title="走迷宫训练过程"):
+    """四联图: 回报曲线 / 滑动成功率 / 单局步数 / 第4格由 extra 提供。
+
+    extra: (数值序列, y轴标签, 子图标题, 是否对数y轴)。
+    """
+    plt = _setup_plt()
+    fig, axes = plt.subplots(2, 2, figsize=(12, 8))
+    fig.suptitle(title, fontsize=15)
+
+    ax = axes[0][0]
+    _plot_ma(ax, rewards, window, "C0")
+    ax.set_ylabel("单局回报")
+    ax.set_title("逐局回报（越高越好）")
+
+    ax = axes[0][1]
+    _plot_ma(ax, successes, window, "C1")
+    ax.set_ylim(-0.05, 1.05)
+    ax.set_ylabel("成功率")
+    ax.set_title(f"滑动成功率（窗口{window}）")
+
+    ax = axes[1][0]
+    _plot_ma(ax, ep_steps, window, "C2")
+    sp = shortest_path_len(env)
+    if sp >= 0:
+        ax.axhline(sp, color="red", ls="--", lw=1.5, label=f"最短路 {sp} 步")
+    ax.set_ylabel("单局步数")
+    ax.set_title("单局步数（降到红线 = 学会最短路）")
+    ax.legend()
+
+    values, ylabel, subtitle, log_scale = extra
+    ax = axes[1][1]
+    ax.plot(np.asarray(values, dtype=float), color="C3", lw=1.2 if log_scale else 2)
+    if log_scale:
+        ax.set_yscale("log")
+    ax.set_xlabel("训练局数 Episode")
+    ax.set_ylabel(ylabel)
+    ax.set_title(subtitle)
+
+    fig.tight_layout(rect=[0, 0, 1, 0.96])
+    fig.savefig(out, dpi=150)
+    print(f"\n图片已保存: {out}")
+
+
+def plot_snapshots(env, snapshots, v_final, out="snapshots.png",
+                   title="V(s)=max Q 的学习过程（价值从 G 回传）"):
+    """各训练阶段的价值热力图。
+
+    snapshots/v_final: (标签, V二维数组)，V 为该阶段每格的 max Q，墙格可为 nan。
+    """
+    plt = _setup_plt()
+    panels = list(snapshots) + [("最终", v_final)]
+    vals = [np.asarray(V, dtype=float)[~np.isnan(V)] for _, V in panels]
+    vmin = min(v.min() for v in vals)
+    vmax = max(v.max() for v in vals)
+    if vmin == vmax:
+        vmin, vmax = vmin - 1, vmax + 1
+
+    n = len(panels)
+    fig, axes = plt.subplots(1, n, figsize=(3 * n, 3.4))
+    axes = np.atleast_1d(axes)
+    masked_wall = env.maze == WALL
+    for ax, (tag, V) in zip(axes, panels):
+        V = np.asarray(V, dtype=float)
+        Vm = np.ma.masked_where(masked_wall, V)
+        im = ax.imshow(Vm, cmap="viridis", vmin=vmin, vmax=vmax, origin="upper")
+        fs = 8 if max(env.shape) <= 7 else 5
+        for r in range(env.shape[0]):
+            for c in range(env.shape[1]):
+                if masked_wall[r, c]:
+                    ax.text(c, r, "█", ha="center", va="center",
+                            color="black", fontsize=fs)
+                else:
+                    ax.text(c, r, f"{V[r, c]:.0f}", ha="center", va="center",
+                            color="white", fontsize=fs)
+        ax.set_xticks([])
+        ax.set_yticks([])
+        ax.set_title(f"第 {tag} 局后" if tag != "最终" else "训练结束")
+    fig.suptitle(title, fontsize=13)
+    fig.colorbar(im, ax=axes, fraction=0.03, pad=0.02)
+    fig.savefig(out, dpi=150, bbox_inches="tight")
+    print(f"图片已保存: {out}")