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,202 @@
+import argparse
+from collections import deque
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from common.maze import (
+    WALL,
+    MazeEnv,
+    evaluate,
+    generate_maze,
+    plot_snapshots,
+    plot_training,
+    show_maze,
+    show_path,
+    show_policy,
+)
+
+
+class DQN(nn.Module):
+    """输入: 状态编码(one-hot 位置向量); 输出: 4 个动作的 Q 值。"""
+
+    def __init__(self, n_states, n_actions, hidden=128):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(n_states, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, n_actions),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class ReplayBuffer:
+    """经验回放池: 存储 (s, a, r, s', done), 随机抽样打破时序相关性。"""
+
+    def __init__(self, capacity, seed=0):
+        self.buf = deque(maxlen=capacity)
+        self.rng = np.random.default_rng(seed)
+
+    def push(self, *transition):
+        self.buf.append(transition)
+
+    def sample(self, batch_size):
+        idx = self.rng.choice(len(self.buf), batch_size, replace=False)
+        batch = [self.buf[i] for i in idx]
+        return tuple(np.stack(col) for col in zip(*batch))
+
+    def __len__(self):
+        return len(self.buf)
+
+
+def encode_state(pos, n_states, n_cols):
+    """把格子坐标 (r, c) 编码成 one-hot 向量。"""
+    one = np.zeros(n_states, dtype=np.float32)
+    one[pos[0] * n_cols + pos[1]] = 1.0
+    return one
+
+
+def optimize(net, target, opt, buffer, batch_size, gamma):
+    """一次梯度步: 最小化 (r + γ·max Q(s',·;θ⁻) − Q(s,a;θ))²"""
+    s, a, r, s2, done = buffer.sample(batch_size)
+    s = torch.as_tensor(s)
+    a = torch.as_tensor(a, dtype=torch.long).unsqueeze(1)
+    r = torch.as_tensor(r, dtype=torch.float32)
+    s2 = torch.as_tensor(s2, dtype=torch.float32)
+    done = torch.as_tensor(done, dtype=torch.float32)
+
+    q = net(s).gather(1, a).squeeze(1)
+    with torch.no_grad():
+        y = r + gamma * target(s2).max(1).values * (1.0 - done)
+    loss = F.mse_loss(q, y)
+
+    opt.zero_grad()
+    loss.backward()
+    opt.step()
+    return float(loss.item())
+
+
+def value_map(env, net, encode):
+    """计算每个格子的 V(s)=max_a Q(s,a), 墙为 nan。"""
+    rows, cols = env.shape
+    V = np.full((rows, cols), np.nan)
+    with torch.no_grad():
+        for r in range(rows):
+            for c in range(cols):
+                if env.maze[r, c] == WALL:
+                    continue
+                x = torch.as_tensor(encode((r, c))).unsqueeze(0)
+                V[r, c] = float(net(x).max().item())
+    return V
+
+
+def train(env, episodes=800, hidden=128, batch_size=64, buffer_capacity=20000,
+          gamma=0.99, lr=1e-3, epsilon=1.0, eps_min=0.02, eps_decay=0.995,
+          target_update=200, seed=0, snap_at=()):
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    n_states = env.shape[0] * env.shape[1]
+    net = DQN(n_states, env.n_actions, hidden).to(device)
+    target = DQN(n_states, env.n_actions, hidden).to(device)
+    target.load_state_dict(net.state_dict())
+    opt = torch.optim.Adam(net.parameters(), lr=lr)
+    buffer = ReplayBuffer(buffer_capacity, seed=seed)
+    rng = np.random.default_rng(seed + 1)
+
+    def encode(pos):
+        return encode_state(pos, n_states, env.shape[1])
+
+    losses, rewards, successes, ep_steps, eps_hist = [], [], [], [], []
+    snapshots = []
+    snap_at = sorted({int(e) for e in snap_at if 0 < int(e) <= episodes})
+    si = 0
+    total_steps = 0
+
+    for ep in range(1, episodes + 1):
+        pos, _ = env.reset()
+        total_r, ep_loss, n_loss = 0.0, 0.0, 0
+        terminated = truncated = False
+
+        while not (terminated or truncated):
+            # ε-greedy
+            if rng.random() < epsilon:
+                a = int(rng.integers(env.n_actions))
+            else:
+                with torch.no_grad():
+                    x = torch.as_tensor(encode(pos)).unsqueeze(0).to(device)
+                    a = int(net(x).argmax().item())
+
+            nxt, r, terminated, truncated, _ = env.step(a)
+            # 超时截断(truncated)不代表环境终止, TD 目标不该被截断, 故只存 terminated
+            buffer.push(encode(pos), a, r, encode(nxt), terminated)
+
+            if len(buffer) >= batch_size:
+                ep_loss += optimize(net, target, opt, buffer, batch_size, gamma)
+                n_loss += 1
+                total_steps += 1
+                if total_steps % target_update == 0:
+                    target.load_state_dict(net.state_dict())
+
+            pos, total_r = nxt, total_r + r
+
+        epsilon = max(eps_min, epsilon * eps_decay)
+        losses.append(ep_loss / n_loss if n_loss else float("nan"))
+        rewards.append(total_r)
+        successes.append(terminated)
+        ep_steps.append(env.steps)
+        eps_hist.append(epsilon)
+        if si < len(snap_at) and ep == snap_at[si]:
+            snapshots.append((ep, value_map(env, net, encode)))
+            si += 1
+
+    return net, device, encode, (losses, rewards, successes, ep_steps, eps_hist, snapshots)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="DQN 走迷宫（PyTorch 深度强化学习实现）")
+    parser.add_argument("--seed", type=int, default=0, help="随机种子")
+    parser.add_argument("--rows", type=int, default=9, help="迷宫行数（奇数）")
+    parser.add_argument("--cols", type=int, default=9, help="迷宫列数（奇数）")
+    parser.add_argument("--episodes", type=int, default=800, help="训练局数")
+    parser.add_argument("--braid", type=float, default=0.0, help="拆墙成环概率 0~1")
+    parser.add_argument("--hidden", type=int, default=128, help="隐藏层宽度")
+    parser.add_argument("--plot", action="store_true",
+                        help="生成 dqn_training.png 与 dqn_snapshots.png")
+    args = parser.parse_args()
+
+    env = MazeEnv(generate_maze(args.rows, args.cols, seed=args.seed, braid=args.braid))
+    show_maze(env)
+    snap_at = (1, 20, 100, 300, 500) if args.plot else ()
+    net, device, encode, hist = train(
+        env, episodes=args.episodes, hidden=args.hidden, seed=args.seed, snap_at=snap_at)
+    losses, rewards, successes, ep_steps, eps_hist, snapshots = hist
+
+    def act(pos):
+        with torch.no_grad():
+            x = torch.as_tensor(encode(pos)).unsqueeze(0).to(device)
+            return int(net(x).argmax().item())
+
+    show_policy(env, act)
+    show_path(env, act)
+    rate, avg_steps = evaluate(env, act)
+    print(f"\n前50局平均回报: {np.mean(rewards[:50]):6.1f}   "
+          f"后50局平均回报: {np.mean(rewards[-50:]):6.1f}")
+    print(f"后50局成功率: {np.mean(successes[-50:]):.0%}")
+    print(f"评测(100局贪心): 成功率 {rate:.0%}，平均 {avg_steps:.1f} 步到达终点")
+    if args.plot:
+        plot_training(env, rewards, successes, ep_steps,
+                      extra=(losses, "TD 损失 (MSE, log)", "TD 损失（下降即逼近 Q*）", True),
+                      out="dqn_training.png", title="DQN 走迷宫训练过程")
+        plot_snapshots(env, snapshots, value_map(env, net, encode),
+                       out="dqn_snapshots.png",
+                       title="DQN 学到的 V(s)=max Q 随训练的变化（价值从 G 回传）")