mirror of
https://github.com/unanmed/ginka-generator.git
synced 2026-05-15 13:21:09 +08:00
323 lines
17 KiB
Markdown
323 lines
17 KiB
Markdown
# VQ 编码器三通道分拆预训练设计文档
|
||
|
||
## 背景与问题诊断
|
||
|
||
### 核心问题
|
||
|
||
在当前 VQ-VAE + MaskGIT 联合训练方案（以及方案 A 闭环约束、方案 D 全图预训练）中，一个根本性的类别不均衡问题始终未被解决：
|
||
|
||
地图中约 70–85% 的格子为墙壁（tile=1）和空地（tile=0），而怪物（9）、门（2）、入口（10）、宝石（4/5/6）、钥匙（3）、药水（7）、道具（8）等功能性 tile 仅占少数。当编码器以完整地图为输入、重建损失对所有位置平等计算时：
|
||
|
||
- 编码器的梯度信号被墙壁/空地主导，特征向量主要编码了空间结构，而非功能性内容；
|
||
- 解码（或生成）时，怪物/门/资源类 tile 的召回率远低于墙壁；
|
||
- 即使方案 D 进行了全图预训练，由于损失仍是全位置等权 Focal Loss，预训练阶段本身就强化了这种偏向；
|
||
- **已将损失由 CE 替换为 Focal Loss，但召回率仍无明显改善**——说明问题的根源在于梯度信号的来源比例，而非损失函数本身的形式；
|
||
- 方案 A 的闭环一致性约束在 z 本身质量不足时，仅能强化已有的偏向，难以从根本上改善。
|
||
|
||
### 改进目标
|
||
|
||
通过**按语义层次将地图拆分为三个独立通道**，采用**累积式输入 + 通道专属损失**的设计：每个通道的编码器输入包含当前通道及所有低等级 tile（提供空间上下文），而解码头的损失**仅在本通道专属 tile 的位置计算，低等级 tile 的损失权重为 0**。这样既保证了编码器能在有意义的空间结构中学习（避免功能性 tile 沦为孤立散点），又迫使解码器必须正确预测本通道 tile 才能降低损失——无法靠拟合高频的墙壁/空地来刷低损失值。
|
||
|
||
---
|
||
|
||
## Tile 类型划分
|
||
|
||
根据 `data/src/shared.ts` 中的定义，将 15 个有效 tile 类型（不含 MASK=15）按游戏语义分为三个通道：
|
||
|
||
| 通道 | 编码器输入（切片内容） | 解码损失计算范围 | 语义含义 |
|
||
| ------ | ---------------------------------------------------- | ------------------ | ------------------------------------ |
|
||
| 通道 1 | floor(0) + wall(1) | {1}（仅墙壁） | 空间骨架（地形结构） |
|
||
| 通道 2 | floor(0) + wall(1) + door(2) + mob(9) + entrance(10) | {2, 9, 10} | 关卡门控（交互元素，决定路径可达性） |
|
||
| 通道 3 | 完整地图（所有 tile） | {3, 4, 5, 6, 7, 8} | 收集资源（奖励与道具） |
|
||
|
||
**通道划分的关键设计原则**：
|
||
|
||
- 通道 1 仅包含 floor 和 wall，编码器集中学习地图的空间骨架结构；
|
||
- 通道 2 的输入**保留墙壁与空地作为空间上下文**，在此基础上叠加 door/mob/entrance，确保编码器能感知功能性 tile 在地图中的位置关系，而非将其视为无空间依附的孤立散点；
|
||
- 通道 3 的输入为**完整地图**，编码器在包含骨架与关卡结构的完整上下文中学习资源 tile 的空间分布；
|
||
- 预训练时每个解码头的损失**仅在本通道专属 tile 的位置计算，低等级 tile（floor、wall 及前级通道的 tile）损失权重为 0**——迫使编码器必须通过正确预测本通道 tile 来降低损失，无法靠拟合高频背景来规避优化压力。
|
||
|
||
---
|
||
|
||
## 整体架构
|
||
|
||
### 预训练阶段
|
||
|
||
```
|
||
完整地图 [B, H*W]
|
||
│
|
||
├──► 切片 1：floor(0)+wall(1)（这已是全部内容，无需替换）
|
||
│ │
|
||
│ ▼
|
||
│ Encoder_1 → VQ_1 → z_1 [B, L_1, d_z]
|
||
│ │
|
||
│ ▼
|
||
│ DecodeHead_1 → logits [B, H*W, C]
|
||
│ │
|
||
│ ▼
|
||
│ Loss_1：仅在 tile∈{1} 的位置计算 Focal Loss（仅墙壁，空地权重为 0）
|
||
│
|
||
├──► 切片 2：保留 floor(0)+wall(1)+door(2)+mob(9)+entry(10)，其余→floor(0)
|
||
│ │
|
||
│ ▼
|
||
│ Encoder_2 → VQ_2 → z_2 [B, L_2, d_z]
|
||
│ │
|
||
│ ▼
|
||
│ DecodeHead_2 → logits [B, H*W, C]
|
||
│ │
|
||
│ ▼
|
||
│ Loss_2：仅在 tile∈{2,9,10} 的位置计算 Focal Loss（floor/wall 权重为 0）
|
||
│
|
||
└──► 切片 3：完整地图（所有 tile，无需替换）
|
||
│
|
||
▼
|
||
Encoder_3 → VQ_3 → z_3 [B, L_3, d_z]
|
||
│
|
||
▼
|
||
DecodeHead_3 → logits [B, H*W, C]
|
||
│
|
||
▼
|
||
Loss_3：仅在 tile∈{3,4,5,6,7,8} 的位置计算 Focal Loss（其余 tile 权重为 0）
|
||
```
|
||
|
||
三路编码器**相互独立预训练**，每路的预训练损失：
|
||
|
||
$$\mathcal{L}_{pretrain}^{(k)} = \mathcal{L}_{FL}^{(k)} + \beta \cdot \mathcal{L}_{commit}^{(k)} + \gamma \cdot \mathcal{L}_{uniform}^{(k)}$$
|
||
|
||
其中 $\mathcal{L}_{FL}^{(k)}$ 为通道 $k$ 的通道专属掩码 Focal Loss（见下节）。
|
||
|
||
### 联合训练阶段
|
||
|
||
```
|
||
完整地图 ──► 三路切片 ──► [Enc_1, Enc_2, Enc_3] ──► [z_1, z_2, z_3]
|
||
│
|
||
z = Concat([z_1, z_2, z_3], dim=1)
|
||
│
|
||
▼
|
||
掩码地图 + z ──► MaskGIT (Cross-Attention) ──► 预测 logits ──► Focal Loss
|
||
```
|
||
|
||
联合训练总损失（不含预训练解码头）：
|
||
|
||
$$\mathcal{L}_{joint} = \mathcal{L}_{FL}^{MaskGIT} + \sum_{k=1}^{3} \left( \beta \cdot \mathcal{L}_{commit}^{(k)} + \gamma \cdot \mathcal{L}_{uniform}^{(k)} \right)$$
|
||
|
||
---
|
||
|
||
## 通道专属掩码 Focal Loss
|
||
|
||
这是方案的核心机制。对于通道 $k$，设其专属 tile 集合为 $\mathcal{T}_k$（不含低等级 tile），则损失计算为：
|
||
|
||
$$\mathcal{L}_{FL}^{(k)} = \frac{\sum_{i=1}^{H \times W} \mathbf{1}[y_i \in \mathcal{T}_k] \cdot \text{FL}(\hat{y}_i, y_i)}{\sum_{i=1}^{H \times W} \mathbf{1}[y_i \in \mathcal{T}_k] + \epsilon}$$
|
||
|
||
其中 $y_i$ 为真实 tile 类型，$\hat{y}_i$ 为解码头输出的 logits，$\text{FL}$ 为 Focal Loss。
|
||
|
||
**实现方式**（PyTorch 伪代码）：
|
||
|
||
```python
|
||
# 通道 2 示例（输入切片已包含 floor+wall 作为上下文）
|
||
CHANNEL_2_TILES = {2, 9, 10} # door, mob, entrance
|
||
|
||
# target: 完整地图 ground truth，[B, H*W]
|
||
# logits: DecodeHead_2 输出，[B, H*W, num_classes]
|
||
|
||
mask = torch.zeros_like(target, dtype=torch.bool)
|
||
for t in CHANNEL_2_TILES:
|
||
mask |= (target == t) # [B, H*W] bool，仅通道 2 专属 tile 的位置为 True
|
||
|
||
# focal_loss: reduction='none'，返回 [B * H*W]
|
||
fl = focal_loss(
|
||
logits.view(-1, num_classes),
|
||
target.view(-1),
|
||
) # [B * H*W]
|
||
fl = fl.view(B, -1) # [B, H*W]
|
||
|
||
loss_ch2 = (fl * mask).sum() / (mask.sum() + 1e-6)
|
||
```
|
||
|
||
**为什么输入包含墙壁、但损失不计算墙壁**：通道 2 的切片中保留了 floor+wall，是为了给编码器提供空间结构上下文，使门/怪/入口的位置有意义（否则孤立散点难以形成有效表示）。但损失仅在 `{2, 9, 10}` 位置计算，确保梯度信号完全来自这三类 tile——编码器如果只靠拟合高频的墙壁/空地，解码头在 `{2, 9, 10}` 位置的损失无法降低，从而被迫学习功能性 tile 的空间分布。
|
||
|
||
---
|
||
|
||
## 切片构造规则
|
||
|
||
| 通道 | 切片中保留的 tile | 其余位置替换为 | 解码损失计算的位置（专属 tile 集合 $\mathcal{T}_k$） |
|
||
| ---- | ----------------- | -------------- | ------------------------------------------------------- |
|
||
| 1 | 0, 1 | —（无需替换） | {1}（仅墙壁；空地是墙壁的补集，能预测墙壁即能区分空地） |
|
||
| 2 | 0, 1, 2, 9, 10 | 0（floor） | {2, 9, 10}（floor/wall 损失权重为 0） |
|
||
| 3 | 全部（完整地图） | —（无需替换） | {3, 4, 5, 6, 7, 8}（其余 tile 损失权重为 0） |
|
||
|
||
---
|
||
|
||
## 编码器架构设计
|
||
|
||
### 三路独立编码器
|
||
|
||
三个编码器均复用现有的 `GinkaVQVAE` 类，配置略有差异：
|
||
|
||
| 参数 | Encoder_1（结构骨架） | Encoder_2（关卡门控） | Encoder_3（收集资源） |
|
||
| --------------- | --------------------- | --------------------- | --------------------- |
|
||
| `L`（码字数） | 2 | 2 | 2 |
|
||
| `K`（码本大小） | 16 | 16 | 16 |
|
||
| `d_z` | 64 | 64 | 64 |
|
||
| `d_model` | 128 | 64 | 64 |
|
||
| `num_layers` | 2 | 2 | 2 |
|
||
|
||
- Encoder_1 处理高频 tile，适当加大 `d_model`；
|
||
- Encoder_2 和 Encoder_3 的功能性 tile 稀疏，信息量较小，可使用较小的 `d_model`；
|
||
- 三路 `d_z` 保持一致，以便拼接后维度齐整；
|
||
- 总参数量估算：Encoder_1 ~400K + Encoder_2 ~150K + Encoder_3 ~150K ≈ **700K**，在 1M 预算内。
|
||
|
||
> 考虑到训练集数量较少，可以考虑适当降低 K 和 L 的值，避免模型死记硬背，也可以防止训练集没有覆盖全部所有情况。1M 的参数量仅做估计，可以先尝试较大的参数量，如出现过拟合再降低。
|
||
|
||
### 联合训练时的 z 拼接
|
||
|
||
$$z = \text{Concat}([z_1, z_2, z_3], \dim=1) \in \mathbb{R}^{B \times (L_1+L_2+L_3) \times d_z}$$
|
||
|
||
以各通道 `L=2` 为例，总 memory 长度为 6，与当前 MaskGIT Cross-Attention 的 memory 规模（原来 L=2）相比略有增加，但绝对长度仍很小，不影响计算效率。
|
||
|
||
---
|
||
|
||
## 预训练流程
|
||
|
||
### 解码头复用
|
||
|
||
预训练时三路各自使用一个 `VQDecodeHead` 实例（现有类，`num_classes=16`），预训练结束后整体丢弃。解码头参数不迁移到联合训练阶段。
|
||
|
||
### 训练脚本
|
||
|
||
新增 `ginka/train_pretrain_split.py`（独立于现有的 `train_pretrain.py`）：
|
||
|
||
```python
|
||
# 伪代码结构
|
||
for epoch in ...:
|
||
for batch in dataloader:
|
||
raw_map = batch["raw_map"] # [B, H*W] 完整地图
|
||
|
||
# ─── 通道 1 ───
|
||
slice1 = make_slice(raw_map, keep={0, 1}) # floor+wall，切片即完整输入
|
||
z_q1, z_e1, _, vq_loss1, *_ = enc1(slice1)
|
||
logits1 = head1(z_q1)
|
||
fl1 = masked_focal(logits1, raw_map, tile_set={1}) # 仅对 wall 计损失
|
||
loss1 = fl1 + vq_loss1
|
||
|
||
# ─── 通道 2 ───
|
||
slice2 = make_slice(raw_map, keep={0, 1, 2, 9, 10}) # 保留 wall/floor 作为上下文
|
||
z_q2, z_e2, _, vq_loss2, *_ = enc2(slice2)
|
||
logits2 = head2(z_q2)
|
||
fl2 = masked_focal(logits2, raw_map, tile_set={2, 9, 10}) # 仅对专属 tile 计损失
|
||
loss2 = fl2 + vq_loss2
|
||
|
||
# ─── 通道 3 ───
|
||
slice3 = raw_map # 完整地图，无需切片
|
||
z_q3, z_e3, _, vq_loss3, *_ = enc3(slice3)
|
||
logits3 = head3(z_q3)
|
||
fl3 = masked_focal(logits3, raw_map, tile_set={3, 4, 5, 6, 7, 8}) # 仅对专属 tile 计损失
|
||
loss3 = fl3 + vq_loss3
|
||
|
||
total = loss1 + loss2 + loss3
|
||
total.backward()
|
||
optimizer.step()
|
||
```
|
||
|
||
三路编码器可以**同步训练**（同一 optimizer），也可以分别独立训练——独立训练更灵活，可以对收敛速度差异大的通道单独调参。
|
||
|
||
### 预训练监控指标
|
||
|
||
| 指标 | 说明 |
|
||
| ----------------------- | ---------------------------------------------------------------- |
|
||
| 通道 1 wall 位置准确率 | Encoder_1 能否正确重建墙壁分布 |
|
||
| 通道 2 功能 tile 召回率 | Encoder_2 对 door/mob/entrance 各类的召回，应 > 50%（稀疏 tile） |
|
||
| 通道 3 资源 tile 召回率 | Encoder_3 对各资源类的召回 |
|
||
| codebook 使用熵（各路） | 各通道 codebook 是否均匀使用，避免 collapse |
|
||
|
||
由于通道 2/3 的 tile 在每张地图中数量极少（典型地图中 door/mob/entrance 合计约 10–20 格，资源合计约 10–15 格），召回率指标比准确率更有意义。
|
||
|
||
---
|
||
|
||
## 联合训练流程
|
||
|
||
### 三阶段训练
|
||
|
||
| 阶段 | 模型状态 | 目标 | 建议轮数 |
|
||
| -------------------- | ----------------------------------------- | ------------------------------------------------ | ------------ |
|
||
| 阶段 0：分通道预训练 | 三路 Encoder + 三路 DecodeHead | 各通道 Focal Loss 收敛，功能 tile 召回率达到目标 | 30–60 epoch |
|
||
| 阶段 1：冻结热身 | 三路 Encoder 冻结 + MaskGIT 全参训练 | MaskGIT 适应三路 z 的联合分布 | 20–40 epoch |
|
||
| 阶段 2：完整联合训练 | 全部参数解冻，Encoder 使用较小 LR（×0.1） | 端到端联合优化 | 正常训练轮数 |
|
||
|
||
### 联合训练数据集
|
||
|
||
`GinkaJointDataset` 需扩展为同时提供三路切片：
|
||
|
||
```python
|
||
# 返回字典新增字段
|
||
{
|
||
"raw_map": ..., # [H*W] 完整地图（VQ 编码器输入）
|
||
"slice1": ..., # [H*W] 通道 1 切片
|
||
"slice2": ..., # [H*W] 通道 2 切片
|
||
"slice3": ..., # [H*W] 通道 3 切片
|
||
"masked_map": ..., # [H*W] MaskGIT 输入（掩码后地图）
|
||
"target_map": ..., # [H*W] MaskGIT CE ground truth
|
||
}
|
||
```
|
||
|
||
---
|
||
|
||
## 推理时的 z 采样
|
||
|
||
推理时三路编码器均独立采样，无需用户输入：
|
||
|
||
```python
|
||
# 完全随机生成
|
||
z1 = enc1.sample(B, device) # [B, L1, d_z]
|
||
z2 = enc2.sample(B, device) # [B, L2, d_z]
|
||
z3 = enc3.sample(B, device) # [B, L3, d_z]
|
||
z = torch.cat([z1, z2, z3], dim=1) # [B, L1+L2+L3, d_z]
|
||
```
|
||
|
||
**分通道条件控制**（可选扩展）：
|
||
|
||
| 场景 | 通道 1 z 来源 | 通道 2 z 来源 | 通道 3 z 来源 |
|
||
| ------------ | ------------- | ------------- | ------------- |
|
||
| 完全随机生成 | 随机采样 | 随机采样 | 随机采样 |
|
||
| 指定墙壁布局 | 用户地图编码 | 随机采样 | 随机采样 |
|
||
| 指定关卡结构 | 随机采样 | 参考图编码 | 随机采样 |
|
||
| 风格迁移 | 参考图编码 | 参考图编码 | 随机采样 |
|
||
|
||
---
|
||
|
||
## 与现有方案的关系
|
||
|
||
| 方案 | 与本方案的关系 |
|
||
| ------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||
| 方案 A | 兼容。由于通道 2/3 的编码器输入包含墙壁上下文（而非孤立散点），其 z 具备稳定的空间语义，一致性约束可同样作用于三个通道：通道 1 约束骨架结构，通道 2 约束关卡元素分布，通道 3 约束资源分布 |
|
||
| 方案 D | 本方案**替代并强化**方案 D 的预训练思路：方案 D 是全图等权 Focal Loss 预训练，本方案通过累积式输入 + 通道专属损失从根本上解决了类别不均衡问题 |
|
||
| 方案 C | 兼容。多阶段生成（先墙后门后资源）可以将通道 1/2/3 的 z 分别作为各阶段的生成条件 |
|
||
|
||
---
|
||
|
||
## 超参数建议
|
||
|
||
| 参数 | 建议初始值 | 备注 |
|
||
| -------------------------- | ---------- | ---------------------------------------- |
|
||
| 各通道 `L`（码字数） | 2 | 三路合计 6，可视效果适当扩大 |
|
||
| 各通道 `K`（码本大小） | 16 | 通道 2/3 可减小到 8（tile 种类少） |
|
||
| `d_z` | 64 | 三路保持一致，便于拼接 |
|
||
| `β`（commit loss 权重） | 0.25 | 同现有配置 |
|
||
| `γ`（uniform loss 权重） | 0.1 | 通道 2/3 码本小，可适当增大到 0.2 |
|
||
| 预训练 epoch | 30–60 | 以功能 tile 召回率达标为准，不以轮数为限 |
|
||
| 联合训练 Encoder LR 缩放比 | 0.1 | 阶段 2 解冻后使用较小 LR 微调 |
|
||
| z dropout 概率（联合训练） | 0.1–0.2 | 三路 z 各自独立 dropout |
|
||
|
||
---
|
||
|
||
## 实施步骤
|
||
|
||
- [ ] 在 `ginka/dataset.py` 中实现 `make_slice(map, keep_set)` 辅助函数，生成三路切片
|
||
- [ ] 扩展 `GinkaJointDataset.__getitem__`，新增 `slice1/slice2/slice3` 字段
|
||
- [ ] 在 `ginka/vqvae/model.py` 中确认 `GinkaVQVAE` 可独立实例化三次（无全局状态）
|
||
- [ ] 实现 `masked_focal(logits, target, tile_set)` 工具函数（`ginka/utils.py`）
|
||
- [ ] 新增 `ginka/train_pretrain_split.py` 预训练脚本（支持三路同步或分路训练）
|
||
- [ ] 修改 `ginka/train_vq.py`（联合训练脚本），支持加载三路编码器权重并拼接 z
|
||
- [ ] 修改 `GinkaMaskGIT.forward()` 以接受 `[B, L1+L2+L3, d_z]` 的拼接 z（Cross-Attention memory）
|
||
- [ ] 添加联合训练监控：各通道 codebook 使用熵、功能 tile 召回率
|