feat: 修改采样方式，改进码本

ginka/model.py

@@ -12,7 +12,7 @@ VQ_GAMMA = 0.1 # entropy loss 权重，鼓励码本使用均匀
 
 # 三通道 VQ 各自独立超参（L、K、层数、维度等均独立配置）
 # Stage1 墙壁骨架 — 结构最复杂，模型容量最大
-VQ_L1 = 24
+VQ_L1 = 16
 VQ_K1 = 32
 VQ_D_MODEL1 = 384
 VQ_NHEAD1 = 8
@@ -20,43 +20,43 @@ VQ_LAYERS1 = 6
 VQ_DIM_FF1 = 1536
 
 # Stage2 功能元素 — 中等复杂度
-VQ_L2 = 12
-VQ_K2 = 16
+VQ_L2 = 8
+VQ_K2 = 64
 VQ_D_MODEL2 = 256
 VQ_NHEAD2 = 4
-VQ_LAYERS2 = 6
+VQ_LAYERS2 = 4
 VQ_DIM_FF2 = 1024
 
 # Stage3 资源分布 — 最简单，模型容量最小
-VQ_L3 = 8
-VQ_K3 = 16
-VQ_D_MODEL3 = 192
-VQ_NHEAD3 = 4
+VQ_L3 = 24
+VQ_K3 = 24
+VQ_D_MODEL3 = 256
+VQ_NHEAD3 = 8
 VQ_LAYERS3 = 4
-VQ_DIM_FF3 = 768
+VQ_DIM_FF3 = 1024
 
 # 第一阶段 MaskGIT 超参
 STAGE1_MG_DMODEL = 512
-STAGE1_MG_NHEAD = 4
+STAGE1_MG_NHEAD = 8
 STAGE1_MG_NUM_LAYERS = 8
 STAGE1_MG_DIM_FF = 2048
 
 # 第二阶段 MaskGIT 超参
-STAGE2_MG_DMODEL = 256
-STAGE2_MG_NHEAD = 4
+STAGE2_MG_DMODEL = 384
+STAGE2_MG_NHEAD = 8
 STAGE2_MG_NUM_LAYERS = 6
-STAGE2_MG_DIM_FF = 1024
+STAGE2_MG_DIM_FF = 1536
 
 # 第三阶段 MaskGIT 超参
 STAGE3_MG_DMODEL = 256
-STAGE3_MG_NHEAD = 4
+STAGE3_MG_NHEAD = 8
 STAGE3_MG_NUM_LAYERS = 6
 STAGE3_MG_DIM_FF = 1024
 
 # 各阶段 VQ commit loss 权重（当前未单独使用，统一由 VQ_BETA 控制）
-STAGE1_VQ_WEIGHT = 0.5
-STAGE2_VQ_WEIGHT = 0.5
-STAGE3_VQ_WEIGHT = 0.5
+STAGE1_VQ_WEIGHT = 0.2
+STAGE2_VQ_WEIGHT = 0.2
+STAGE3_VQ_WEIGHT = 0.2
 
 # 全局参数
 NUM_CLASSES = 8 # 图块类型数
@@ -65,7 +65,7 @@ TOTAL_K = VQ_K1 + VQ_K2 + VQ_K3  # 预计算，供日志输出使用
 MAP_W = 13 # 地图宽度
 MAP_H = 13 # 地图高度
 
-LR = 1e-4 # AdamW 初始学习率
+LR = 2e-4 # AdamW 初始学习率
 MIN_LR = 1e-6 # 余弦退火最低学习率
 WEIGHT_DECAY = 1e-4 # L2 正则化系数
 EPOCHS = 400 # 总训练轮数

ginka/sample.py

@@ -1,52 +1,96 @@
+import heapq
 import math
+import random
+from collections import deque
 
 import numpy as np
 import torch
 import torch.nn.functional as F
 
 from .model import MASK_TOKEN, MAP_H, MAP_W, SeperatedModels
-from .utils import compute_remaining, MAP_SIZE, compute_adjacency_mask
+from .utils import (
+    compute_remaining, MAP_SIZE, compute_adjacency_mask,
+    WALL_DENSITY_IDX, DOOR_DENSITY_IDX, MONSTER_DENSITY_IDX,
+    ENTRANCE_DENSITY_IDX, RESOURCE_DENSITY_IDX
+)
 
-# MaskGIT 采样函数：通过迭代去掩码从离散隐变量 z 生成地图
+# 采样器超参
+DBC_GUMBEL = 2.0   # 揭开顺序的 Gumbel 噪声强度，0 关闭
+DBC_JITTER = 0.08  # 图块预算的随机浮动，0 关闭
+DBC_TRIES = 6      # 拒绝采样次数
+
+# 每个阶段负责的图块，在 target_density 里对应的下标
+STAGE_DENSITY_IDX = {
+    1: (WALL_DENSITY_IDX,),
+    2: (DOOR_DENSITY_IDX, MONSTER_DENSITY_IDX, ENTRANCE_DENSITY_IDX),
+    3: (RESOURCE_DENSITY_IDX,),
+}
+
+_NB4 = ((0, 1), (0, -1), (1, 0), (-1, 0))
 
 
-def wall_growth_sample(
-    model: torch.nn.Module,
-    inp: torch.Tensor,
-    z: torch.Tensor,
-    struct: torch.Tensor,
-    target_density: torch.Tensor,
-    max_steps: int = 24
-) -> np.ndarray:
-    # 墙壁生长算法：从 inp 中已有的墙壁出发，逐步向外生长
-    # 每步 MASK 位置决策（墙/非墙）后，新邻接面成为下一轮 MASK，逐步外扩
-    state = inp.clone()  # [B, MAP_SIZE]
-    # 初始 MASK：邻接已有墙壁的空地
-    init_adj = compute_adjacency_mask(state)
-    state[init_adj & (state == 0)] = MASK_TOKEN
+def floor_label(m: np.ndarray):
+    mask = (m != 1)
+    lab = -np.ones(m.shape, int)
+    n = 0
+    for i in range(MAP_H):
+        for j in range(MAP_W):
+            if mask[i, j] and lab[i, j] < 0:
+                q = deque([(i, j)])
+                lab[i, j] = n
+                while q:
+                    a, b = q.popleft()
+                    for da, db in _NB4:
+                        x, y = a + da, b + db
+                        if 0 <= x < MAP_H and 0 <= y < MAP_W and mask[x, y] and lab[x, y] < 0:
+                            lab[x, y] = n
+                            q.append((x, y))
+                n += 1
+    return lab, n
 
-    for step in range(max_steps):
-        mask_pos = (state == MASK_TOKEN)  # [B, MAP_SIZE]
-        if not mask_pos.any():
-            break
 
-        remain = compute_remaining(state, target_density, 1)
-        logits = model(state, z, struct, remain)
-        probs = F.softmax(logits, dim=-1)
-
-        wall_prob = probs[:, :, 1]  # [B, MAP_SIZE]
-        hits = (wall_prob > 0.5) & mask_pos
-        state[hits] = 1
-        state[mask_pos & ~hits] = 0
-
-        adj = compute_adjacency_mask(state)
-        new_mask = adj & (state == 0)
-        if not new_mask.any():
-            break
-        state[new_mask] = MASK_TOKEN
-
-    state[state == MASK_TOKEN] = 0
-    return state.cpu().numpy().reshape(state.size(0), MAP_H, MAP_W)
+def repair_connectivity(m: np.ndarray) -> np.ndarray:
+    # 0-1 Dijkstra 打通最少的墙，使所有非墙格连成一片
+    m = m.copy()
+    while True:
+        lab, n = floor_label(m)
+        if n <= 1:
+            return m
+        sizes = [int((lab == k).sum()) for k in range(n)]
+        main = int(np.argmax(sizes))
+        tgt = next(k for k in range(n) if k != main)
+        dist = np.full(m.shape, 1 << 30)
+        pq = []
+        prev = {}
+        end = None
+        for i in range(MAP_H):
+            for j in range(MAP_W):
+                if lab[i, j] == main:
+                    dist[i, j] = 0
+                    heapq.heappush(pq, (0, i, j))
+        while pq:
+            d, i, j = heapq.heappop(pq)
+            if d > dist[i, j]:
+                continue
+            if lab[i, j] == tgt:
+                end = (i, j)
+                break
+            for da, db in _NB4:
+                x, y = i + da, j + db
+                if not (0 <= x < MAP_H and 0 <= y < MAP_W):
+                    continue
+                nd = d + (1 if m[x, y] == 1 else 0)
+                if nd < dist[x, y]:
+                    dist[x, y] = nd
+                    prev[(x, y)] = (i, j)
+                    heapq.heappush(pq, (nd, x, y))
+        if end is None:
+            return m
+        cur = end
+        while cur in prev:
+            if m[cur] == 1:
+                m[cur] = 0
+            cur = prev[cur]
 
 
 def maskgit_sample(
@@ -61,6 +105,13 @@ def maskgit_sample(
     current = inp.clone()
     target_tensor = torch.tensor(target_tiles, dtype=torch.long, device=inp.device)
 
+    # 退火对象改成「本阶段还需要放多少个目标图块」
+    base = float(sum(target_density[0, i] for i in STAGE_DENSITY_IDX[stage])) * MAP_SIZE
+    if DBC_JITTER > 0:
+        base *= 1.0 + random.uniform(-DBC_JITTER, DBC_JITTER)
+    budget = int(round(base))
+    need0 = max(0, budget - int(torch.isin(current[:], target_tensor).sum()))
+
     # 迭代去掩码：每步根据置信度分数重新决定掩码位置
     for step in range(steps):
         remain = compute_remaining(current, target_density, stage)
@@ -72,20 +123,30 @@ def maskgit_sample(
 
         confidences = torch.gather(probs, -1, sampled.unsqueeze(-1)).squeeze(-1)
 
-        # 余弦退火调度：随步数推进，保留掩码的位置数量递减至 0
-        ratio = math.cos(((step + 1) / steps) * math.pi / 2)
-        num_to_mask = math.floor(ratio * MAP_SIZE)
+        # Gumbel 噪声退火：提高揭开顺序的多样性（MaskGIT 原论文做法）
+        if DBC_GUMBEL > 0:
+            u = torch.rand_like(confidences).clamp_(1e-9, 1 - 1e-9)
+            score = (torch.log(confidences.clamp_min(1e-9))
+                + DBC_GUMBEL * (1 - step / steps) * (-torch.log(-torch.log(u))))
+        else:
+            score = confidences
 
         # 结构位：current 中非空地、非掩码的位置（来自上一阶段，始终保留）
         struct_mask = (current[:] != MASK_TOKEN) & (current[:] != 0)
         # 候选位：sampled 为目标图块且不覆盖结构位
         candidate_mask = torch.isin(sampled[:], target_tensor) & ~struct_mask
         cand_count = candidate_mask.sum()
-        reveal_count = max(0, int(cand_count.item()) - num_to_mask)
+
+        # 预算感知揭开：按本阶段还需放置的图块数做余弦退火
+        ratio = math.cos(((step + 1) / steps) * math.pi / 2)
+        still = math.floor(ratio * need0)
+        now = max(0, budget - int(torch.isin(current[:], target_tensor).sum()))
+        reveal_count = min(max(0, now - still), int(cand_count.item()))
+
         next_state = current[:].clone()
         if reveal_count > 0 and cand_count > 0:
             cand_indices = candidate_mask.nonzero(as_tuple=False)
-            cand_conf = confidences[:][cand_indices[:, 0], cand_indices[:, 1]]
+            cand_conf = score[:][cand_indices[:, 0], cand_indices[:, 1]]
             top_k = min(reveal_count, cand_conf.size(0))
             _, top_idx = torch.topk(cand_conf, k=top_k, largest=True)
             reveal_rows = cand_indices[top_idx, 0]
@@ -107,21 +168,108 @@ def maskgit_sample(
 
     return current.cpu().numpy().reshape(current.size(0), MAP_H, MAP_W)
 
+
+def sample_with_retry(
+    model: torch.nn.Module, inp: torch.Tensor, z: torch.Tensor,
+    struct: torch.Tensor, target_density: torch.Tensor,
+    stage: int, steps: int, target_tiles: list[int],
+    tries: int = DBC_TRIES
+) -> np.ndarray:
+    # 拒绝采样：多次采样取地板连通的那张，都不连通则修复块数最少的
+    best = None
+    best_n = 1 << 30
+    for _ in range(tries):
+        m = maskgit_sample(
+            model, inp.clone(), z, struct, target_density,
+            stage, steps, target_tiles=target_tiles
+        )
+        m_2d = m.reshape(MAP_H, MAP_W)
+        _, n = floor_label(m_2d)
+        if n == 1:
+            return m
+        if n < best_n:
+            best = m
+            best_n = n
+    if len(best.shape) == 3:
+        repaired = repair_connectivity(best[0])
+        repaired = repaired[np.newaxis, ...]
+    else:
+        repaired = repair_connectivity(best.reshape(MAP_H, MAP_W))
+        repaired = repaired.reshape(best.shape)
+    return repaired
+
+
+def wall_growth_sample(
+    model: torch.nn.Module,
+    inp: torch.Tensor,
+    z: torch.Tensor,
+    struct: torch.Tensor,
+    target_density: torch.Tensor,
+    max_steps: int = 24
+) -> np.ndarray:
+    # 墙壁生长算法：从 inp 中已有的墙壁出发，逐步向外生长
+    # 所有未固定位置初始为 MASK，只有与墙壁邻接的 MASK 才参与评估
+    state = inp.clone()  # [B, MAP_SIZE]
+    state[state == 0] = MASK_TOKEN
+
+    for step in range(max_steps):
+        adj = compute_adjacency_mask(state)
+        mask_pos = adj & (state == MASK_TOKEN)
+        if not mask_pos.any():
+            break
+
+        remain = compute_remaining(state, target_density, 1)
+        logits = model(state, z, struct, remain)
+        probs = F.softmax(logits, dim=-1)
+
+        wall_prob = probs[:, :, 1]  # [B, MAP_SIZE]
+        hits = (wall_prob > 0.5) & mask_pos
+        state[hits] = 1
+        state[mask_pos & ~hits] = 0
+
+    state[state == MASK_TOKEN] = 0
+    return state.cpu().numpy().reshape(state.size(0), MAP_H, MAP_W)
+
+
 def full_generate(
     inp: torch.Tensor,
     z1: torch.Tensor, z2: torch.Tensor, z3: torch.Tensor,
     struct: torch.Tensor,
     target_density: torch.Tensor,
     models: SeperatedModels,
-    steps: int = 18
+    steps: int = 18,
+    seed_mode: bool = False,
+    stage1_method: str = "maskgit"
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-    # 三阶段级联生成：Stage1 墙壁生长 → Stage2 门/怪/入口 → Stage3 资源
+    # 三阶段级联生成
+    # seed_mode=True  从空白生成墙壁（自由生成）
+    # seed_mode=False 从 inp 中的已有结构出发（掩码补全）
+    # stage1_method: "maskgit"=MaskGIT+拒绝采样, "growth"=墙壁生长算法
     # 返回 (stage1结果, stage1+2合并, 最终完整地图)，形状均为 [B, H, W]
     device = inp.device
 
-    pred1_np = wall_growth_sample(
-        models.mg1, inp, z1, struct, target_density
-    )  # [B, H, W]
+    if seed_mode:
+        if stage1_method == "growth":
+            pred1_np = wall_growth_sample(
+                models.mg1, inp, z1, struct, target_density
+            )
+        else:
+            stage1_inp = torch.full((inp.size(0), MAP_SIZE), MASK_TOKEN, dtype=torch.long, device=device)
+            pred1_np = sample_with_retry(
+                models.mg1, stage1_inp, z1, struct, target_density, 1,
+                steps, target_tiles=[1]
+            )
+    else:
+        if stage1_method == "growth":
+            pred1_np = wall_growth_sample(
+                models.mg1, inp, z1, struct, target_density
+            )
+        else:
+            pred1_np = sample_with_retry(
+                models.mg1, inp, z1, struct, target_density, 1,
+                steps, target_tiles=[1]
+            )
+    # [B, H, W]
     inp2 = torch.tensor(pred1_np.reshape(pred1_np.shape[0], -1), dtype=torch.long, device=device)
     inp2[inp2 == 0] = MASK_TOKEN
 

ginka/train_seperated.py

@@ -20,8 +20,8 @@ from .utils import (
     compute_remaining, MAP_SIZE, summarize_codebook_hits
 )
 from .sample import full_generate
-from .dataset import GinkaSeperatedDataset
-from shared.image import matrix_to_image_cv
+from .dataset import GinkaSeperatedDataset, rect_mask, ensure_wall_connection
+from shared.image import matrix_to_image_cv, annotate
 
 # 图块 ID 定义：
 # 0. 空地   1. 墙壁   2. 普通门   3. 资源   4. 怪物   5. 入口   6. 机关门   7. 掩码（MASK_TOKEN）
@@ -107,133 +107,216 @@ def quantize_stage_latents(
     code_hits = (code_hits1, code_hits2, code_hits3)
     return (z_q1, z_q2, z_q3), commit_loss, code_hits, entropy_loss
 
-# 墙壁种子生成：随机放置墙壁后让模型生长
-def visualize_seed(
+# 每张图 2 行 x 4 列：MASKED | PREDICTED | MASKED | PREDICTED
+def make_compare_grid(samples_src, samples_pred, map_keys, tile_dict, TILE_SIZE):
+    SEP = 3
+    img_h = MAP_H * TILE_SIZE
+    img_w = MAP_W * TILE_SIZE
+    cols = 4
+    rows = 2
+    grid = np.ones((rows * img_h + (rows + 1) * SEP, cols * img_w + (cols + 1) * SEP, 3), dtype=np.uint8) * 255
+    for r in range(rows):
+        row_y = SEP + r * (img_h + SEP)
+        for c in range(cols):
+            is_src = (c % 2 == 0)
+            idx = r * 2 + c // 2
+            col_x = SEP + c * (img_w + SEP)
+            data = samples_src[idx] if is_src else samples_pred[idx]
+            grid[row_y:row_y + img_h, col_x:col_x + img_w] = annotate(
+                matrix_to_image_cv(data, tile_dict, TILE_SIZE), map_keys[idx], y=14
+            )
+    # 每两张图片之间的分隔线
+    for c in range(cols - 1):
+        line_x = (c + 1) * (img_w + SEP) + SEP // 2
+        grid[:, line_x:line_x + SEP, :] = 180
+    return grid
+
+def visualize_seed_growth(
     train_dataset: GinkaSeperatedDataset,
     models: SeperatedModels,
     device: torch.device,
     tile_dict,
     epoch: int
 ):
-    SEP = 3
-    TILE_SIZE = 32
-    img_h = MAP_H * TILE_SIZE
-    img_w = MAP_W * TILE_SIZE
-
-    def to_img(mat):
-        return matrix_to_image_cv(mat, tile_dict, TILE_SIZE)
-
     save_dir = f"result/seperated/e{epoch}"
     os.makedirs(save_dir, exist_ok=True)
+    samples_src = []
+    samples_pred = []
+    map_keys = []
+    for _ in range(20):
+        sample = train_dataset.random_sample_map()
+        struct_t = sample["struct_inject"].to(device).reshape(1, -1)
+        target_density_t = sample["target_density"].to(device).reshape(1, -1)
+        enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE)
+        enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE)
+        enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE)
+        enc1_np = enc1_t.cpu().numpy().reshape(MAP_H, MAP_W)
+        with torch.no_grad():
+            z_e1 = models.vq1(enc1_t)
+            z_e2 = models.vq2(enc2_t)
+            z_e3 = models.vq3(enc3_t)
+            z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3)
+            z1, z2, z3 = z_q
+            inp = torch.full((1, MAP_SIZE), 0, dtype=torch.long, device=device)
+            seed_count = random.randint(7, 14)
+            seed_idx = torch.randperm(MAP_SIZE, device=device)[:seed_count]
+            inp[0, seed_idx] = 1
+            inp_vis = inp.clone()
+            inp_vis[inp_vis == 0] = MASK_TOKEN
+            _, _, merged = full_generate(
+                inp, z1, z2, z3, struct_t, target_density_t, models,
+                steps=SEED_SAMPLE_STEPS, seed_mode=True, stage1_method="growth"
+            )
+        samples_src.append(inp_vis.cpu().numpy().reshape(MAP_H, MAP_W))
+        samples_pred.append(merged[0])
+        map_keys.append(sample["map_name"])
+    for i in range(5):
+        grid = make_compare_grid(
+            samples_src[i*4:i*4+4], samples_pred[i*4:i*4+4],
+            map_keys[i*4:i*4+4], tile_dict, 32
+        )
+        cv2.imwrite(f"{save_dir}/seed_growth_{i}.png", grid)
 
-    for img_i in range(5):
-        samples_data = []
-        for _ in range(4):
-            sample = train_dataset.random_sample_map()
-            struct_t = sample["struct_inject"].to(device).reshape(1, -1)
-            target_density_t = sample["target_density"].to(device).reshape(1, -1)
-            enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE)
-            enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE)
-            enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE)
-
-            with torch.no_grad():
-                z_e1 = models.vq1(enc1_t)
-                z_e2 = models.vq2(enc2_t)
-                z_e3 = models.vq3(enc3_t)
-                z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3)
-                z1, z2, z3 = z_q
-
-                inp = torch.full((1, MAP_SIZE), 0, dtype=torch.long, device=device)
-                seed_count = random.randint(7, 14)
-
-                # 分层采样：将 13x13 地图划分为 4x4 网格，从不同格中随机取种子点
-                grid_h, grid_w = 4, 4
-                cell_h = MAP_H // grid_h + 1
-                cell_w = MAP_W // grid_w + 1
-                all_cells = [(r, c) for r in range(grid_h) for c in range(grid_w)]
-                chosen_cells = random.sample(all_cells, seed_count)
-                for cell_r, cell_c in chosen_cells:
-                    r_min = cell_r * cell_h
-                    r_max = min(r_min + cell_h, MAP_H)
-                    c_min = cell_c * cell_w
-                    c_max = min(c_min + cell_w, MAP_W)
-                    sr = random.randint(r_min, max(r_min, r_max - 1))
-                    sc = random.randint(c_min, max(c_min, c_max - 1))
-                    flat_idx = sr * MAP_W + sc
-                    inp[0, flat_idx] = 1
-
-                _, _, merged = full_generate(
-                    inp, z1, z2, z3,
-                    struct_t, target_density_t, models,
-                    steps=SEED_SAMPLE_STEPS
-                )
-
-            samples_data.append(merged[0])
-
-        grid = np.ones((2 * img_h + 3 * SEP, 2 * img_w + 3 * SEP, 3), dtype=np.uint8) * 255
-        for r in range(2):
-            for c in range(2):
-                y = SEP + r * (img_h + SEP)
-                x = SEP + c * (img_w + SEP)
-                grid[y:y + img_h, x:x + img_w] = to_img(samples_data[r * 2 + c])
-        cv2.imwrite(f"{save_dir}/seed{img_i}.png", grid)
-
-# 矩形掩码生成：对真实地图随机掩码后让模型补全
-def visualize_mask(
+def visualize_seed_maskgit(
     train_dataset: GinkaSeperatedDataset,
     models: SeperatedModels,
     device: torch.device,
     tile_dict,
     epoch: int
 ):
-    SEP = 3
-    TILE_SIZE = 32
-    img_h = MAP_H * TILE_SIZE
-    img_w = MAP_W * TILE_SIZE
-
-    def to_img(mat):
-        return matrix_to_image_cv(mat, tile_dict, TILE_SIZE)
-
     save_dir = f"result/seperated/e{epoch}"
     os.makedirs(save_dir, exist_ok=True)
+    samples_src = []
+    samples_pred = []
+    map_keys = []
+    for _ in range(20):
+        sample = train_dataset.random_sample_map()
+        struct_t = sample["struct_inject"].to(device).reshape(1, -1)
+        target_density_t = sample["target_density"].to(device).reshape(1, -1)
+        enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE)
+        enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE)
+        enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE)
+        enc1_np = enc1_t.cpu().numpy().reshape(MAP_H, MAP_W)
+        with torch.no_grad():
+            z_e1 = models.vq1(enc1_t)
+            z_e2 = models.vq2(enc2_t)
+            z_e3 = models.vq3(enc3_t)
+            z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3)
+            z1, z2, z3 = z_q
+            inp = torch.full((1, MAP_SIZE), MASK_TOKEN, dtype=torch.long, device=device)
+            _, _, merged = full_generate(
+                inp, z1, z2, z3, struct_t, target_density_t, models,
+                steps=GENERATE_STEP, seed_mode=True, stage1_method="maskgit"
+            )
+        samples_src.append(inp.cpu().numpy().reshape(MAP_H, MAP_W))
+        samples_pred.append(merged[0])
+        map_keys.append(sample["map_name"])
+    for i in range(5):
+        grid = make_compare_grid(
+            samples_src[i*4:i*4+4], samples_pred[i*4:i*4+4],
+            map_keys[i*4:i*4+4], tile_dict, 32
+        )
+        cv2.imwrite(f"{save_dir}/seed_maskgit_{i}.png", grid)
 
-    for img_i in range(5):
-        samples_data = []
-        for _ in range(4):
-            sample = train_dataset.random_sample_map()
-            raw_map = sample["raw_map"].cpu().numpy().reshape(MAP_H, MAP_W)
-            ratio = random.uniform(0.2, 0.8)
+def visualize_mask_growth(
+    train_dataset: GinkaSeperatedDataset,
+    models: SeperatedModels,
+    device: torch.device,
+    tile_dict,
+    epoch: int
+):
+    save_dir = f"result/seperated/e{epoch}"
+    os.makedirs(save_dir, exist_ok=True)
+    samples_src = []
+    samples_pred = []
+    map_keys = []
+    for _ in range(20):
+        sample = train_dataset.random_sample_map()
+        raw_map = sample["raw_map"].cpu().numpy().reshape(MAP_H, MAP_W)
+        enc1_np = sample["encoder_stage1"].cpu().numpy().reshape(MAP_H, MAP_W)
+        enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE)
+        enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE)
+        enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE)
+        struct_t = sample["struct_inject"].to(device).reshape(1, -1)
+        target_density_t = sample["target_density"].to(device).reshape(1, -1)
 
-            enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE)
-            enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE)
-            enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE)
-            struct_t = sample["struct_inject"].to(device).reshape(1, -1)
-            target_density_t = sample["target_density"].to(device).reshape(1, -1)
+        # 矩形分块掩码 + 确保墙壁连通性，与训练时一致
+        ratio = random.uniform(0.3, 0.9)
+        rmask = rect_mask(ratio).reshape(MAP_H, MAP_W)
+        rmask = ensure_wall_connection(rmask, enc1_np)
+        inp = torch.tensor(raw_map.flatten(), dtype=torch.long, device=device).reshape(1, MAP_SIZE)
+        inp[0, rmask.flatten()] = MASK_TOKEN
+        masked_np = inp.cpu().numpy().reshape(MAP_H, MAP_W)
 
-            with torch.no_grad():
-                z_e1 = models.vq1(enc1_t)
-                z_e2 = models.vq2(enc2_t)
-                z_e3 = models.vq3(enc3_t)
-                z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3)
-                z1, z2, z3 = z_q
+        with torch.no_grad():
+            z_e1 = models.vq1(enc1_t)
+            z_e2 = models.vq2(enc2_t)
+            z_e3 = models.vq3(enc3_t)
+            z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3)
+            z1, z2, z3 = z_q
+            _, _, merged = full_generate(
+                inp, z1, z2, z3, struct_t, target_density_t, models,
+                stage1_method="growth"
+            )
+        samples_src.append(masked_np)
+        samples_pred.append(merged[0])
+        map_keys.append(sample["map_name"])
+    for i in range(5):
+        grid = make_compare_grid(
+            samples_src[i*4:i*4+4], samples_pred[i*4:i*4+4],
+            map_keys[i*4:i*4+4], tile_dict, 32
+        )
+        cv2.imwrite(f"{save_dir}/mask_growth_{i}.png", grid)
 
-                mask = torch.rand(MAP_SIZE, device=device) < ratio
-                inp = torch.tensor(raw_map.flatten(), dtype=torch.long, device=device).reshape(1, MAP_SIZE)
-                inp[0, mask] = MASK_TOKEN
+def visualize_mask_maskgit(
+    train_dataset: GinkaSeperatedDataset,
+    models: SeperatedModels,
+    device: torch.device,
+    tile_dict,
+    epoch: int
+):
+    save_dir = f"result/seperated/e{epoch}"
+    os.makedirs(save_dir, exist_ok=True)
+    samples_src = []
+    samples_pred = []
+    map_keys = []
+    for _ in range(20):
+        sample = train_dataset.random_sample_map()
+        raw_map = sample["raw_map"].cpu().numpy().reshape(MAP_H, MAP_W)
+        enc1_np = sample["encoder_stage1"].cpu().numpy().reshape(MAP_H, MAP_W)
+        enc1_t = sample["encoder_stage1"].to(device).reshape(1, MAP_SIZE)
+        enc2_t = sample["encoder_stage2"].to(device).reshape(1, MAP_SIZE)
+        enc3_t = sample["encoder_stage3"].to(device).reshape(1, MAP_SIZE)
+        struct_t = sample["struct_inject"].to(device).reshape(1, -1)
+        target_density_t = sample["target_density"].to(device).reshape(1, -1)
 
-                _, _, merged = full_generate(
-                    inp, z1, z2, z3, struct_t, target_density_t, models
-                )
+        # 矩形分块掩码 + 确保墙壁连通性，与训练时一致
+        ratio = random.uniform(0.3, 0.9)
+        rmask = rect_mask(ratio).reshape(MAP_H, MAP_W)
+        rmask = ensure_wall_connection(rmask, enc1_np)
+        inp = torch.tensor(raw_map.flatten(), dtype=torch.long, device=device).reshape(1, MAP_SIZE)
+        inp[0, rmask.flatten()] = MASK_TOKEN
+        masked_np = inp.cpu().numpy().reshape(MAP_H, MAP_W)
 
-            samples_data.append(merged[0])
-
-        grid = np.ones((2 * img_h + 3 * SEP, 2 * img_w + 3 * SEP, 3), dtype=np.uint8) * 255
-        for r in range(2):
-            for c in range(2):
-                y = SEP + r * (img_h + SEP)
-                x = SEP + c * (img_w + SEP)
-                grid[y:y + img_h, x:x + img_w] = to_img(samples_data[r * 2 + c])
-        cv2.imwrite(f"{save_dir}/mask{img_i}.png", grid)
+        with torch.no_grad():
+            z_e1 = models.vq1(enc1_t)
+            z_e2 = models.vq2(enc2_t)
+            z_e3 = models.vq3(enc3_t)
+            z_q, _, _, _ = quantize_stage_latents(models, z_e1, z_e2, z_e3)
+            z1, z2, z3 = z_q
+            _, _, merged = full_generate(
+                inp, z1, z2, z3, struct_t, target_density_t, models,
+                stage1_method="maskgit"
+            )
+        samples_src.append(masked_np)
+        samples_pred.append(merged[0])
+        map_keys.append(sample["map_name"])
+    for i in range(5):
+        grid = make_compare_grid(
+            samples_src[i*4:i*4+4], samples_pred[i*4:i*4+4],
+            map_keys[i*4:i*4+4], tile_dict, 32
+        )
+        cv2.imwrite(f"{save_dir}/mask_maskgit_{i}.png", grid)
 
 def train(device: torch.device):
     args = parse_arguments()
@@ -283,6 +366,7 @@ def train(device: torch.device):
         loss2_total = torch.Tensor([0]).to(device)
         loss3_total = torch.Tensor([0]).to(device)
         commit_total = torch.Tensor([0]).to(device)
+        entropy_total = torch.Tensor([0]).to(device)
         code_hits_total = (torch.zeros(result.quantizer1.K, device=device), torch.zeros(result.quantizer2.K, device=device), torch.zeros(result.quantizer3.K, device=device))  # validate
 
         for batch in tqdm(dataloader, leave=False, desc="Epoch Progress", disable=disable_tqdm):
@@ -330,10 +414,12 @@ def train(device: torch.device):
             logits2 = result.mg2(inp2, z_q2, struct, remain2)
             logits3 = result.mg3(inp3, z_q3, struct, remain3)
 
-            # 三阶段 Cross Entropy：仅对输入中为 MASK_TOKEN 的位置计算 loss
+            # 三阶段 Cross Entropy
+            # Stage1: 仅对掩码位置计算 loss
+            # Stage2/3: 掩码 + 空地位置均计算 loss，模型需学习哪里不能放东西
             mask1 = (inp1 == MASK_TOKEN)
-            mask2 = (inp2 == MASK_TOKEN)
-            mask3 = (inp3 == MASK_TOKEN)
+            mask2 = (inp2 == MASK_TOKEN) | (inp2 == 0)
+            mask3 = (inp3 == MASK_TOKEN) | (inp3 == 0)
             loss1 = cross_entropy_loss(logits1, target1, mask1)
             loss2 = cross_entropy_loss(logits2, target2, mask2)
             loss3 = cross_entropy_loss(logits3, target3, mask3)
@@ -353,12 +439,16 @@ def train(device: torch.device):
             loss2_total += loss2.detach()
             loss3_total += loss3.detach()
             commit_total += commit_loss.detach()
+            entropy_total += entropy_loss.detach()
             code_hits_total = (code_hits_total[0] + code_hits[0].detach(), code_hits_total[1] + code_hits[1].detach(), code_hits_total[2] + code_hits[2].detach())  # accumulate train
 
         # 每个 epoch 结束后更新学习率
         result.scheduler.step()
 
         data_length = len(dataloader)
+        # 总 loss 去掉 entropy_loss 以便与 CE + VQ 对齐
+        loss_display = (loss_total.item() - VQ_GAMMA * entropy_total.item()) / data_length
+        entropy_display = entropy_total.item() / data_length
         stats = summarize_codebook_hits(code_hits_total)
         parts = []
         for name in ["q1(stage1)", "q2(stage2)", "q3(stage3)"]:
@@ -368,9 +458,9 @@ def train(device: torch.device):
         tqdm.write(
             f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] "
             f"E: {epoch + 1} | "
-            f"Loss: {loss_total.item() / data_length:.4f} | "
+            f"Loss: {loss_display:.4f} | "
             f"CE: {loss1_total.item() / data_length:.4f}, {loss2_total.item() / data_length:.4f}, {loss3_total.item() / data_length:.4f} | "
-            f"VQ: {commit_total.item() / data_length:.4f} | "
+            f"cmt: {commit_total.item() / data_length:.4f} ent: {entropy_display:+.4f} | "
             f"VQ: {' | '.join(parts)} | "
             f"Total: {c['active']}/{c['K']} ppl={c['ppl']:.1f} | "
             f"LR: {result.scheduler.get_last_lr()[0]:.6f}"
@@ -378,8 +468,10 @@ def train(device: torch.device):
 
         # 每 CHECKPOINT 个 epoch 执行可视化并保存检查点
         if (epoch + 1) % CHECKPOINT == 0:
-            visualize_seed(dataset, result, device, tile_dict, epoch + 1)
-            visualize_mask(dataset, result, device, tile_dict, epoch + 1)
+            visualize_seed_growth(dataset, result, device, tile_dict, epoch + 1)
+            visualize_seed_maskgit(dataset, result, device, tile_dict, epoch + 1)
+            visualize_mask_growth(dataset, result, device, tile_dict, epoch + 1)
+            visualize_mask_maskgit(dataset, result, device, tile_dict, epoch + 1)
             ckpt_path = f"result/seperated/sep-{epoch + 1}.pth"
             result.save(ckpt_path, epoch + 1)
             tqdm.write(f"Saved checkpoint: {ckpt_path}")