KOH 2v2 规则蒸馏方案

这个仓库实现的是一个面向 KOH 2v2 对抗环境的训练与调试方案。目标是分别训练进攻方 attack 和防守方 defense 两个固定权重模型，并导出比赛要求的 .safetensors 文件；最终采用的主路线是基于规则 teacher 的策略蒸馏。

1. 问题介绍

题目本质上是一个小规模、部分可观测、双智能体协同的对抗问题：

• 对局形式：2v2
• 阵营：
  • T / attack：进攻方，下包并等待爆炸
  • CT / defense：防守方，阻止下包或在下包后拆弹
• 地图：25 x 25 ASCII 地图，墙体同时阻挡移动、视线和弹道
• 动作：每个角色每步 9 个动作，包含上下左右移动、停留、四方向开火
• 机制难点：
  • 射击先于移动结算，不能靠本步移动躲枪
  • 下包/拆弹都要求连续 stay 3 步
  • 双人协同需要避免抢格、对穿和站位冲突
  • 胜负既包含歼灭，也包含时间、下包和拆弹条件

除此之外，这个题目还有一个很关键的比赛约束：提交时不能自由设计大模型，必须使用赛事规定的固定网络结构。当前实现里使用的是固定 MLP：

77 -> 256 -> 256 -> 18

也就是说：

• 输入维度固定为团队观测 77 维
• 只能通过比较小的隐藏层容量表达策略
• 输出固定为两名角色的联合动作，共 18 维

这会带来一个额外挑战：模型参数量非常有限，无法依赖“大网络自己学出复杂战术”，必须尽量把任务结构、阶段信息和协同规则压缩进有限容量的策略中。因此，这个问题不只是“训练一个能打的策略”，而是“在极小模型容量下，把下包、拆弹、交战和双人协同这些关键行为稳定编码进去”。

环境定义和规则细节见 KOH_rules.md 与 koh_env.py。

2. 最终使用的方法

最终采用的是“规则蒸馏”为主的方案，而不是直接从零做纯强化学习。

2.1 为什么选择规则蒸馏

这个任务的奖励稀疏且带明显阶段性：

• 进攻方既要走位接近包点，又要在合适时机 stay 下包
• 防守方既要守点，又要在下包后切换成 retake/defuse 模式
• 双人协同如果没有先验，很容易出现长时间空转、抢位、无效移动

直接从零开始 RL 在这个环境里很容易遇到：

• 学不到稳定的下包/拆弹流程
• 样本效率低
• 双人动作冲突多
• 对地图变化泛化较差

规则 teacher 则天然能把关键结构先编码进去，再用监督学习压进固定网络中，训练更稳定，起点也更高。

2.2 规则蒸馏的核心思路

整体流程是：

1. 用规则 teacher 在真实环境中 rollout，收集 (state, action) 样本。
2. 可选补充随机合法状态样本，扩大覆盖面。
3. 对防守方可额外补充 post-plant 状态，强化拆弹阶段。
4. 用固定 MLP 77 -> 256 -> 256 -> 18 做监督学习，分别预测两名角色的动作。

对应代码：

• 训练入口：koh_baseline_template.py
• 蒸馏训练：koh_baseline_train.py
• teacher 规则：koh_baseline_policies.py
• 数据采样：koh_baseline_data.py

2.3 teacher 规则概述

teacher 不是纯脚本乱走，而是显式利用环境信息做联合决策，核心规则包括：

• 射击优先：若可见敌人在同一行/列且在射程内，优先开火
• BFS 走位：利用环境中的距离表，朝 A/B 点或炸弹点走最短路
• 双人冲突消解：避免两人走向同一格或对穿换位
• 阶段性目标：
  • 进攻方默认朝包点推进，到点后连续 stay 下包
  • 防守方未下包时分点防守，下包后共同转向炸弹点拆弹
  • 防守方看见敌人但暂时打不到时，会先追击而不是盲目守点

更详细的 teacher 说明见 KOH_teacher_policy_distill.md。

2.4 当前落地版本

当前仓库里可以直接看到最终方案的几个实际变体：

• attack: 蒸馏为主，后续可选 PPO 微调
• defense: 蒸馏为主，并额外强化 post-plant / defuse 阶段

checkpoints/ 中保留了若干实验产物，例如：

• defense_stage1_postplant.safetensors
• defense_stage1_postplant20k.safetensors
• defense_stage1_redistill_current_teacher.safetensors
• attack_stage2.safetensors
• defense_stage2.safetensors

从这些命名可以看出，最终稳定有效的主线不是“纯 PPO”，而是“先蒸馏拿到可靠行为，再考虑少量 RL 微调”。

3. 运行方法

3.1 环境准备

建议使用 uv：

uv sync

如果只想临时运行脚本，也可以直接：

uv run python -V

依赖定义见 pyproject.toml。

3.2 训练 attack / defense

仓库里已经给了一个可直接复现的训练脚本 train.sh。

典型的 stage 1 规则蒸馏命令如下：

uv run \
  koh_baseline_template.py \
  --train-mode distill \
  --role attack \
  --maps-dir maps \
  --teacher-episodes 3000 \
  --random-state-samples 0 \
  --epochs 50 \
  --batch-size 64 \
  --output checkpoints/attack_stage1.safetensors

uv run \
  koh_baseline_template.py \
  --train-mode distill \
  --role defense \
  --maps-dir maps \
  --teacher-episodes 3000 \
  --random-state-samples 0 \
  --postplant-state-samples 20000 \
  --epochs 50 \
  --batch-size 64 \
  --output checkpoints/defense_stage1.safetensors

参数说明：

• --train-mode distill：使用规则蒸馏
• --role：训练 attack 或 defense
• --maps-dir maps：使用地图池训练
• --teacher-episodes：teacher rollout 的局数
• --random-state-samples：补充随机合法状态样本，当前默认实验里常设为 0
• --postplant-state-samples：补充 post-plant 状态，主要对 defense 有帮助
• --epochs / --batch-size / --learning-rate：监督训练超参数
• --output：输出权重路径

3.3 可选的 PPO 微调

如果想在蒸馏权重上继续做对抗微调，可以使用 --train-mode ppo，并通过 --init-weights 和 --enemy-checkpoint 指定初始化模型和固定对手。

train.sh 里已经给了 stage 2 示例，不过当前是注释状态：

uv run \
  koh_baseline_template.py \
  --train-mode ppo \
  --role attack \
  --maps-dir maps \
  --episodes 1000 \
  --init-weights checkpoints/attack_stage1.safetensors \
  --enemy-checkpoint checkpoints/defense_stage1.safetensors \
  --output checkpoints/attack_stage2.safetensors

uv run \
  koh_baseline_template.py \
  --train-mode ppo \
  --role defense \
  --maps-dir maps \
  --episodes 1000 \
  --init-weights checkpoints/defense_stage1.safetensors \
  --enemy-checkpoint checkpoints/attack_stage1.safetensors \
  --output checkpoints/defense_stage2.safetensors

3.4 测试与对局回放

仓库里给了一个简单的测试入口 test.sh：

uv run python koh_view_replay.py \
  --attack-checkpoint checkpoints/attack_stage1.safetensors \
  --defense-checkpoint checkpoints/defense_stage1.safetensors \
  --map maps/sunset.txt

4. 可视化脚本使用

可视化脚本是 koh_view_replay.py，它会在终端里播放对局回放。

4.1 直接生成并播放一局

uv run python koh_view_replay.py \
  --attack-checkpoint checkpoints/attack_stage1.safetensors \
  --defense-checkpoint checkpoints/defense_stage1.safetensors \
  --map maps/sunset.txt

常用参数：

• --attack-checkpoint / --defense-checkpoint：指定双方模型
• --map：指定单张地图
• --maps-dir：如果不指定 --map，可从地图目录读取
• --seed：控制生成的 round id / 随机种子
• --attack-policy / --defense-policy：不加载 checkpoint 时，可选 teacher 或 scripted

4.2 先保存 replay，再回放

生成 replay JSON：

uv run python koh_view_replay.py \
  --attack-checkpoint checkpoints/attack_stage1.safetensors \
  --defense-checkpoint checkpoints/defense_stage1.safetensors \
  --map maps/sunset.txt \
  --save-replay replay.json

播放已有 replay：

uv run python koh_view_replay.py --replay replay.json

4.3 播放控制

viewer 是 TTY 交互式脚本，按键如下：

• n：下一帧
• N：上一帧
• q：退出

如果不是在真实终端中运行，脚本只能生成 replay，不能交互播放。

5. 探索过的方法对比

这里把仓库里实际存在过、或代码中已经实现的路线做一个简要对比。

5.1 纯规则策略

特点：

• 不需要训练
• 行为可解释
• 起点强，尤其适合这个规则复杂的环境

问题：

• 最终提交要求是固定神经网络权重
• 规则上限明显，泛化和细节优化能力有限

结论：

• 适合作为 teacher，不适合作为最终提交形态

5.2 纯 DQN / 纯 RL

代码里保留了 legacy rl 路线和更完整的 ppo 路线。

优点：

• 理论上可超越手工规则
• 能从对抗反馈里继续修正策略

问题：

• 从零训练样本效率低
• 很容易卡在“会走但不会赢”的阶段
• 对下包/拆弹这种多步阶段目标学习不稳定
• 双人协同难，动作冲突多

结论：

• 单独作为主路线不稳
• 更适合接在蒸馏之后做小步微调

5.3 规则蒸馏

优点：

• 收敛快，训练稳定
• 可以直接把“射击优先、最短路走位、分点防守、到点下包/拆弹、双人避碰”等结构灌进模型
• 更适合作为比赛里固定网络的可靠起点

问题：

• 上限受 teacher 质量限制
• teacher 没覆盖到的状态，学生也容易学不到

结论：

• 这是当前仓库最终采用的主方案

5.4 规则蒸馏 + 特定状态增强

这是当前 defense 路线里比较关键的一步。通过 --postplant-state-samples 人工补充“已下包”的合法状态，显式提升 retake/defuse 阶段覆盖。

优点：

• 对防守方很有效
• 直接补齐普通 rollout 中占比偏少但很关键的后期局面

问题：

• 需要自己识别哪个阶段是性能瓶颈
• 属于任务特定的数据分布修正

结论：

• 是比单纯加长 teacher rollout 更有针对性的增强手段

5.5 规则蒸馏 + PPO 微调

优点：

• 以蒸馏策略为稳定初始化，再从真实对抗结果中继续修正
• 有机会学到 teacher 没写出的细节

问题：

• 训练成本更高
• 对对手选择和训练顺序敏感
• 如果微调过头，可能反而破坏蒸馏得到的结构化行为

结论：

• 适合作为增量实验路线
• 但最终稳定版本仍应以蒸馏效果为主，不应完全依赖 PPO

6. 仓库结构

• koh_env.py：环境定义与规则结算
• koh_baseline_template.py：训练入口
• koh_baseline_train.py：蒸馏、DQN、PPO 训练逻辑
• koh_baseline_data.py：地图、数据采样、训练环境包装
• koh_baseline_policies.py：teacher 和 scripted policy
• koh_view_replay.py：终端回放可视化
• train.sh：训练命令示例
• test.sh：回放命令示例

7. 推荐结论

如果以“稳定出成绩”为目标，推荐顺序是：

1. 先用规则蒸馏训练 attack 和 defense
2. 对 defense 额外加入 post-plant 状态增强
3. 通过 koh_view_replay.py 检查实际行为是否符合预期
4. 只有在蒸馏策略已经稳定的前提下，再尝试少量 PPO 微调

这条路线和当前仓库保留下来的实验痕迹是一致的，也是这个环境里最务实的方案。