# Stage1 墙壁生长算法生成设计文档

## 背景与问题诊断

### 当前 Stage1 的 MaskGIT 采样机制

Stage1 负责生成地图的墙壁/地面骨架（floor/wall），使用标准 MaskGIT 迭代解码器：输入全 MASK 地图，通过余弦退火调度逐步揭开置信度最高的 token，共 18 步完成。每步流程如下：

```
全 MASK 地图 → 模型预测所有位置 → 选出 Top-K 高置信度位置 →
揭开这些位置 → 其余位置重新 MASK → 下一步
```

### 问题现象

当前 Stage1 生成的地图墙壁质量显著低于预期，具体表现为：

- 墙壁分布破碎，出现孤立墙块（block 不与外围墙壁连接）
- 房间边界不完整，未能形成闭合区域
- 走廊断断续续，不能形成连通通道
- 整体拓扑结构与训练集分布偏差大

### 根本原因：MaskGIT 缺乏空间连续约束

MaskGIT 的迭代解掩码（unmasking）机制本质上是**逐 token 独立决策**：每一步根据余弦退火比例选出全局置信度最高的 N 个位置进行揭开，选取标准仅依赖于该位置的类别预测置信度，不包含任何空间结构约束。

在 Stage1 墙壁生成中，墙壁的核心特性是**连通性**：

1. 所有墙壁必须与外围墙壁（地图边框）连通
2. 墙壁必须形成闭合或半闭合的房间边界
3. 墙壁的合理性完全由其邻域关系决定——一个孤立的墙块在游戏中毫无意义

MaskGIT 的全局置信度排序无视了这些空间邻接约束。一个远离外围墙壁的内部位置，即使模型以 99% 置信度预测为墙壁，也不应该被单独采出——因为游戏规则要求墙壁必须连通。当前机制却会欣然接受这个"高置信度的错误"。这本质上是全局独立采样与局部连续性要求之间的结构性矛盾。

类比：

| 机制               | 类比                                 | 问题             |
| ------------------ | ------------------------------------ | ---------------- |
| MaskGIT 全局 Top-K | 在所有空白处随机填色，先填最"确信"的 | 可能出现孤立色块 |
| 墙壁生长算法       | 从已有色块边缘逐步向外扩展           | 保证连通性       |

---

## 核心思路：生长式墙壁生成

### 基本思想

将 Stage1 从"全局置信度排序 + 逐步解掩码"改为**边界生长**（Frontier Growth）：

1. 初始状态：外围一圈固定为墙壁（tile=1），内部全部为 MASK
2. 每步：模型预测所有位置，但**只在已存在墙壁的邻接位置**接受墙壁预测
3. 新产生的墙壁成为下一轮的"种子"，继续向外生长
4. 直到没有新的邻接候选位置或达到最大步数

```
Step 0    Step 1       Step 2         Step 3
#######   #######      #######        #######
#.....#   #.....#      #...#.#        #.###.#
#.....#   #..#..#      #.###.#        #####.#
#.....#   #.....#      #.....#        #.#.#.#
#######   #######      #######        #######

# = 墙壁   . = MASK/空地
```

### 与原 MaskGIT 的对比

| 维度       | 原 MaskGIT（全局 Top-K） | 生长算法（邻接约束）                       |
| ---------- | ------------------------ | ------------------------------------------ |
| 选取机制   | 全局置信度排序取 Top-K   | 只在邻接已存在墙壁的位置候选中按置信度选取 |
| 空间约束   | 无，独立决策             | 强制连通：新墙壁必须"挨着"已有墙壁         |
| 墙壁连通性 | 不保证，依赖模型隐式学习 | 结构性保证                                 |
| 生成过程   | 并行：一步可揭开任意位置 | 扩散式：从边界逐步向内推进                 |
| 收敛步数   | 固定 18 步               | 动态，取决于地图墙壁密度                   |

---

## 算法详细设计

### 数据结构

```
地图状态矩阵 state[13][13]，每格取值：
  0: FLOOR（空地/地面）
  1: WALL（墙壁）
  6: MASK（待生成）

外墙固化：state 中位于最外圈（row=0, row=12, col=0, col=12）的位置
          在初始化时设为 tile=1（WALL），并在整个过程中保持不可修改
```

### 主算法

```python
def wall_growth_sample(
    model, z, z_dist, struct, target_density,
    max_steps: int = 24
) -> np.ndarray:
    # 1. 初始化：外圈墙壁 + 内部 MASK
    state = torch.full((1, MAP_SIZE), MASK_TOKEN, dtype=torch.long)
    state[0, outer_wall_indices] = 1  # 固化外圈

    for step in range(max_steps):
        # 2. 模型前向：预测所有位置的 logits
        remain = compute_remaining(state, target_density, stage=1)
        logits = model(state, z, z_dist, struct, remain)
        probs = F.softmax(logits, dim=-1)  # [1, MAP_SIZE, NUM_CLASSES]
        wall_probs = probs[0, :, 1]         # 每个位置预测为墙壁的概率

        # 3. 寻找生长前沿：哪些 MASK 位置与已有墙壁四连通相邻
        mask_positions = (state[0] == MASK_TOKEN)
        adjacent_to_wall = compute_adjacency_mask(state[0])  # [MAP_SIZE] bool
        candidates = mask_positions & adjacent_to_wall

        if candidates.sum() == 0:
            break  # 没有可生长的位置，终止

        # 4. 在候选位置中选出置信度最高的墙壁预测
        cand_conf = wall_probs[candidates]
        num_to_reveal = compute_reveal_count(step, max_steps, candidates.sum())
        _, top_indices = torch.topk(cand_conf, k=num_to_reveal)

        cand_positions = candidates.nonzero(as_tuple=True)[0][top_indices]
        state[0, cand_positions] = 1  # 设为墙壁

    # 5. 剩余 MASK 位置 → 空地
    state[0, state[0] == MASK_TOKEN] = 0
    return state[0].cpu().numpy().reshape(MAP_H, MAP_W)
```

### 邻接计算

```python
def compute_adjacency_mask(flat_state: torch.Tensor) -> torch.Tensor:
    # 返回 [MAP_SIZE] 布尔张量，True 表示该位置与任意墙壁 4-邻接
    state_2d = flat_state.reshape(MAP_H, MAP_W)
    wall = (state_2d == 1)
    # 上/下/左/右 各 shift 一位取 OR
    adj = torch.zeros_like(wall, dtype=torch.bool)
    adj[:, 1:] |= wall[:, :-1]    # 上方有墙壁
    adj[:, :-1] |= wall[:, 1:]    # 下方有墙壁
    adj[1:, :] |= wall[:-1, :]    # 左方有墙壁
    adj[:-1, :] |= wall[1:, :]    # 右方有墙壁
    return adj.flatten()
```

### 每步揭开数量调度

与标准 MaskGIT 的余弦退火调度不同，生长算法在前期应积极揭开（扩大生长前沿），后期逐渐放缓（精细收束）：

```python
def compute_reveal_count(step: int, max_steps: int, num_candidates: int) -> int:
    # 前期多揭（扩大前沿），后期少揭（精细控制）
    t = step / max_steps
    ratio = 1.0 - t ** 0.5  # 从 1.0 降至 0.0，前期保持高位
    # 不能超过当前候选数量，至少揭 1 个（若有候选）
    return max(1, min(num_candidates, math.ceil(num_candidates * ratio)))
```

备选：沿用在生长前沿上的余弦退火——效果需通过实验验证。

### 候选类型扩展：FLOOR 的处理

上述算法仅在候选位置接受墙壁（tile=1）预测。另一种思路是同时接受 FLOOR（tile=0）预测，让"空地"也能与墙壁一同向外推进：

| 策略               | 描述                                               | 优缺点                                     |
| ------------------ | -------------------------------------------------- | ------------------------------------------ |
| 仅接受墙壁         | 候选位置只接受 tile=1 预测                         | 简单，墙壁生长可控；但后期可能挤压空地空间 |
| 同时接受墙壁和空地 | 候选位置按模型预测的最高置信度类别决定是墙还是空地 | 更接近自然生长；但对模型预测质量要求更高   |

推荐先使用"仅接受墙壁"策略，原因是：空地本身没有连通性要求，一旦墙壁骨架确定，空地自然就是剩余位置。实验验证阶段可对比两种策略。

### 终止条件

生长在以下条件之一满足时终止：

1. **无候选位置**：所有 MASK 位置均不邻接现有墙壁——墙壁已完全生长，剩余位置就是空地
2. **达到最大步数**：`max_steps` 上限，防止极端情况无限循环（参考值：24 步，对应 13×13 地图）

---

## Stage1 训练方案：矩形掩码 + 连通损失

### 设计动机

标准 MaskGIT 训练使用 `std_mask()`（50% 随机散布掩码 / 50% 随机分块掩码），在整图上生成破碎的 MASK 图案。这种训练方式与生长式推理**完全不匹配**：

- **训练**：模型看到散落各处的孤立 MASK token，学习从全局上下文独立预测每个位置
- **推理**：MASK 区域是一个连续的边界生长前沿，模型需要在已有墙壁的基础上向外延伸

训练与推理的分布不一致（"train-test mismatch"）导致模型在推理时表现不佳：模型从未在训练中见过"全连通空白区域从边界开始生长"的场景。

**核心改造**：将 Stage1 训练改为矩形掩码（rectangular masking），且只在掩码区域内与已有墙壁连通的位置计算损失。这使训练任务与推理任务完全对齐——都是"从已有墙壁向外延伸填充一个连续区域"。

### 矩形掩码生成

直接使用 `train_seperated.py` 中已有的 `rect_mask` 函数（第 419 行），无需修改。调用方式：

```python
from ginka.train_seperated import rect_mask
mask = rect_mask(ratio=0.5)  # 反复放置随机矩形直到覆盖约 50% 面积
```

现有实现通过反复叠加随机小矩形（2×2 ~ 7×7）来达到目标遮盖比例，形状不规则但保证连续性。

### 连通损失计算

仅对掩码区域内**与外部墙壁 4-连通**的位置计算 CE Loss。原因：生长算法推理时，只有能从已有墙壁"触及"的位置才有机会被生长覆盖——模型不需要学习预测生长前沿触及不到的位置。

```
原图:                MASK（阴影区）:      连通位置（★）:
###########        ###########        ###########
#.........#        #.........#        #.........#
#.###.....#        #.###.....#        #.###.....#
#.#.#.....#   →    #▓▓▓▓▓▓▓..#   →    #.★▓▓▓▓..#  ← MASK 内左下角为墙，
#.....###.#        #▓▓▓▓▓▓▓##.#        #▓▓▓▓▓▓★#.#   因与外部墙连通故参与 loss
#######                #######                #######   MASK 内右上方为地板，不参与 loss
# = 墙  . = 地板  ▓ = MASK 区域  ★ = 参与 loss 的位置
```

具体算法：

```python
def compute_stage1_loss_mask(
    rect_mask: np.ndarray, target1: np.ndarray, inp1: np.ndarray
) -> np.ndarray:
    # rect_mask:  [H, W] bool，True=被矩形掩码覆盖
    # target1:    [H, W] int，GT 墙壁骨架（仅含 0=floor, 1=wall）
    # inp1:       [H, W] int，输入地图（MASK 区域为 MASK_ID，其余保留原值）
    # 返回: [H, W] bool，True=该位置需要计算 loss
    # 策略：在 rect_mask 内，仅对从"外部墙壁"可 4-连通触及的 WALL 位置计算 loss
    H, W = rect_mask.shape
    MASK_ID = 7
    loss_mask = np.zeros((H, W), dtype=bool)

    # 1. 找到外部墙壁（不在 rect_mask 内且值为 WALL=1 的位置）
    external_wall = (~rect_mask) & (inp1 == 1)

    # 2. 从外部墙壁出发 BFS，沿 target1==1 传播，找到所有连通墙壁
    visited = np.zeros((H, W), dtype=bool)
    # BFS/DFS 种子：外部墙壁位置
    seeds = np.argwhere(external_wall)
    for r, c in seeds:
        if visited[r, c]:
            continue
        stack = [(r, c)]
        visited[r, c] = True
        while stack:
            cr, cc = stack.pop()
            for nr, nc in [(cr-1, cc), (cr+1, cc), (cr, cc-1), (cr, cc+1)]:
                if 0 <= nr < H and 0 <= nc < W:
                    if not visited[nr, nc] and target1[nr, nc] == 1:
                        visited[nr, nc] = True
                        stack.append((nr, nc))

    # 3. 连通位置 = rect_mask 内 & 被 BFS 访问到的位置（即与外部墙连通的墙壁）
    loss_mask = rect_mask & visited
    return loss_mask
```

关键逻辑：

1. **外部墙壁** = 不在掩码区域内且值为 WALL 的位置
2. **BFS 传播** = 从外部墙壁出发，沿 GT 中值为 WALL(1) 的格子 4-连通扩散，标记所有可达的墙壁位置
3. **Loss 掩码** = 掩码区域内、且被 BFS 标记为可达的位置

注意：BFS 使用 GT（target1）而非输入（inp1），因为需要标注的是"GT 中与外部墙壁连通的墙壁分布"，而非输入中已存在的墙壁。即使掩码区域内的某个墙壁位置在输入中被 MASK 覆盖（值为 MASK_ID），只要它在 GT 中与外部墙壁连通，就应该参与 loss。

loss_mask 中：

- **True** = 掩码区域内、与外部墙壁连通的墙壁位置 → 模型应该预测 WALL
- **False** = 掩码区域内的其他位置（地板、或与外部墙壁不连通的孤立墙壁） → 不参与 loss

### 对比：新旧训练方案

| 维度 | 旧方案（MaskGIT std_mask） | 新方案（矩形掩码 + 连通损失） |
| ---- | ------------------------- | ---------------------------- |
| 掩码形状 | 随机散布 + 随机方块叠加 | 连续矩形块（复用 `rect_mask`） |
| MASK 覆盖比例 | Beta(2,2)，约 5%~95% | 可由 `ratio` 参数控制 |
| loss 计算范围 | 所有 MASK 位置 | 仅 MASK 区域中与外部墙壁连通的位置 |
| 训练-推理一致性 | 低：散点 MASK vs 连续生长 | 高：矩形空缺 vs 连续生长边界 |
| 模型学习目标 | 从破碎上下文预测单个位置 | 从已有墙壁向外延伸填补连续区域 |

### 生长式推断

生长算法仅影响**Stage1 的推理采样方式**。推理时：

1. 初始状态：外围一圈固定为墙壁（tile=1），内部全部为 MASK
2. 每步：模型预测所有位置，但只在已存在墙壁的邻接位置接受墙壁预测
3. 新产生的墙壁成为下一轮的"种子"，继续向外生长
4. 直到没有新的邻接候选位置或达到最大步数

训练时使用矩形掩码 + 连通损失，推理时使用生长算法——两者在概念上完全一致：都是"从已有墙壁向外延伸"。训练为推理做了准备的准备（teacher forcing），推理则自回归地展开（autoregressive rollout）。

### 推理管线影响

```
当前推理管线:
  Stage1 (MaskGIT Top-K) → floor/wall 骨架
  Stage2 (MaskGIT)       → +门/怪物/入口
  Stage3 (MaskGIT)       → +资源

修改后推理管线:
  Stage1 (生长算法)      → floor/wall 骨架       [改动点]
  Stage2 (MaskGIT)       → +门/怪物/入口         [不变]
  Stage3 (MaskGIT)       → +资源                 [不变]
```

Stage2 和 Stage3 的训练与推理**保持不变**——它们继续使用标准 MaskGIT 方法。这是因为门/怪物/资源是稀疏离散元素，不涉及全局连通性约束，MaskGIT 的独立 token 预测对此类任务足够有效。

### 模型参数

生长算法**不改变模型结构**（GinkaMaskGIT 的 Transformer 架构不变），仅修改：

- `dataset.py`：Stage1 的 `apply_subset*` 方法，将 MaskGIT 掩码策略改为矩形掩码 + 连通损失掩码
- `train_seperated.py`：新增 `wall_growth_sample` 函数替代 `maskgit_sample` 的 Stage1 分支
- 训练 loss 计算中，Stage1 的 CE Loss 需要乘以 `loss_mask` 进行筛选

模型检查点格式与 Stage2/3 训练流程不受影响。

---

## 潜在风险与应对

| 风险           | 描述                                                                       | 应对策略                                                                 |
| -------------- | -------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| 墙壁密度过高   | 生长不退让，可能填满所有可达区域                                           | 增加 target_density 的约束权重；在后期步数中降低揭开速率                 |
| 生长形状单一   | 从外圈均匀向内心扩散，产生"洋葱圈"形墙壁                                   | 引入随机扰动（每步以一定概率跳过某些候选位置）；调整 target_density 分布 |
| 屋顶封闭       | 生长过早将空地包围封闭，导致内部无法再生长墙壁                             | 在揭开判断中加入"是否会导致孤岛"的后处理检查                             |
| 模型能力不匹配 | 生长算法依赖模型在邻接位置的预测质量，若模型本身不能捕捉邻接关系则效果有限 | 可结合已有的二维因式位置嵌入，确保模型具备足够的局部感知能力             |
| 步数不均匀     | 不同地图墙壁密度不同，生长步数的固定上限可能不够或过早收敛                 | 将 max_steps 设为动态参数，或改为"无新墙即终止"的收敛策略                |

### 墙壁"洋葱圈"问题专项说明

若每步以固定速率在所有邻接生长前沿均匀揭开，且模型倾向于预测墙壁，则墙壁可能从外围逐层向内均匀扩散，形成洋葱圈状结构。应对策略：

1. **target_density 约束**：靶向墙壁密度作为硬约束，当已生成的墙壁数量接近目标密度时，大幅降低每步揭开速率
2. **非均匀揭开**：每步仅在候选位置中随机选取子集，而非全部可选中的最高置信度位置，增加生长路径的随机性
3. **后处理修剪**：生成完成后，对墙壁进行一次 BFS 修剪，移除不改变连通性的"冗余墙壁"分支

---

## 实施计划

### 数据集修改

- [ ] Stage1 训练中直接调用已有 `rect_mask` 替代 `std_mask`
- [ ] 新增 `compute_stage1_loss_mask` 函数：BFS 计算连通损失掩码
- [ ] 在 `__getitem__` 返回的 sample 中新增 `"loss_mask_stage1"` 字段

### 推理修改

- [ ] 在 `train_seperated.py` 中实现 `wall_growth_sample` 函数
- [ ] 修改 `full_generate_specific_z` 中 Stage1 的调用，替换为生长算法
- [ ] 实现 `compute_adjacency_mask` 工具函数（四邻接检测）
- [ ] 实现生长步数调度 `compute_reveal_count`

### 训练脚本修改

- [ ] Stage1 的 CE Loss 改为 `(F.cross_entropy(..., reduction='none') * loss_mask_stage1.float()).sum() / loss_mask_stage1.sum()`
- [ ] Stage2/3 的训练与原有逻辑保持一致，不做修改

### 验证

- [ ] 单步测试：加载已有检查点，运行生长算法一次，输出墙壁骨架地图
- [ ] 添加可视化对比：同一组 z 下，MaskGIT Top-K vs 生长算法的 Stage1 输出
- [ ] 对比指标：墙壁连通率、孤立墙块数、房间闭合度