# Baseline 动作臂契约（五臂公平消融）

主实验：同一个 Wan2.2-TI2V-5B、同一份数据、同一套训练配方，**唯一自由变量是动作以什么方式进模型**：

| 臂 | 代表作 | 机制（要按上游代码移植，不是凭印象） | 上游代码入口 |
|---|---|---|---|
| `arope` | 本项目 | 世界坐标 RoPE + known/new 掩码，0 参数 | `actionrope/arope.py` |
| `linear` | ReactiveGWM | 每个 DiT block 一个线性层，把动作映射成 hidden 上的偏置 | `vender/ReactiveGWM/inference/models/dit.py`、`training/data/action_utils.py`、`inference/utils/actions.py` |
| `xattn` | Matrix-Game 3 | 键盘/动作 token 与视频 token 沿时间轴做 cross-attention（前若干 block） | `vender/Matrix-Game/Matrix-Game-3/wan/modules/action_module.py`、`model.py`、`utils/conditions.py` |
| `prompt` | Incantation | 动作走文本，逐 latent（逐 cell）的 prompt 控制 | `vender/Incantation/modules/causal_model.py`、`inference.py`、`prompts/` |
| `adaln` | AlayaWorld | 动作嵌入进 adaLN 调制（与时间步嵌入同路） | `vender/AlayaWorld/alaya/control/action.py`、`alaya/model/components.py`、`ltx2/modules/model.py`、`timestep_embedding.py` |

旧报告的口径（要对齐）：四臂全部从零起训、同数据同 seed；**带动作的臂收到完全相同的控制信号（逐 cell 累计偏移）**；
文本一律剥除动作词（prompt 臂除外——文本就是它的通道）。Linear 移植自 ReactiveGWM 官方实现（30 个逐块线性偏置，+0.18M），
Cross-Attn 复刻 Matrix-Game 3.0 键盘通路（前 15 块时间轴交叉注意力，+98.6M）。

## 钩子接口（`actionrope/arope.py::arope_forward(..., arm=, action_inputs=)`）

`arm` 是一个 `nn.Module`（`baseline/base.py::ActionArm` 的子类），`arope_forward` 在固定位置调用它；
`arm=None` 时 `arope_forward` 与原实现逐位相同。调用顺序与形状：

```python
grid = dict(b, f=21, h=15, w=26, n_tok_per_frame=390)      # token 网格
feats   = arm.encode(action_inputs, grid)                    # 一次；返回任意张量/tuple（建议 [b, f, d]），会作为张量参数穿过梯度检查点
t_mod   = arm.modify_t_mod(t_mod, feats, grid)               # t_mod: [b, S, 6, dim]（逐 token，S = f·h·w；帧 0 的 t 已置 0）
context = arm.extra_context(context, feats, grid)            # 文本 context: [b, L, dim] 或 [(b f), L, dim]（逐 cell）
x       = arm.after_patch(x, feats, grid)                    # x: [b, S, dim]，patch embedding 之后、blocks 之前
for i, block in enumerate(dit.blocks):                       # 每个 block 内部（在梯度检查点里）：
    x = arm.block_pre(i, x, feats, grid)                     #   block 入口
    ... self_attn ...
    x = arm.block_mid(i, x, feats, grid)                     #   自注意力之后、文本 cross-attn 之前
    ... cross_attn ...
    x = arm.block_pre_ffn(i, x, feats, grid)                 #   文本 cross-attn 之后、FFN 之前（Matrix-Game 3 的注入位置）
    ... ffn ...
    x = arm.block_post(i, x, feats, grid)                    #   block 出口
```

token 排列是 `(f h w)`：token 序号 `s = f_idx * 390 + i_h * 26 + j_w`，所以 `[b, S, dim] → [b, f, h*w, dim]` 用 `x.view(b, f, -1, dim)`。

`ActionArm` 基类（`baseline/base.py`）所有钩子默认恒等、`encode` 默认返回 None；子类只覆盖自己用到的。
基类还提供：`install(dit)`（默认空；需要读 dit.dim / num_heads / 层数时用），`n_new_params()`，`zero_init_check()`。

### 动作输入 `action_inputs`（dataset 统一提供，臂自己选用）

```python
{
  "offset_px":  float32 [b, 21, 2],   # 逐 cell 累计玩家屏幕投影位移 (dx, dy) px，cell 0 = 0；与 ARoPE 同一个信号（sidecar）
  "offset_tok": float32 [b, 21, 2],   # = offset_px / 32
  "delta_tok":  float32 [b, 21, 2],   # 逐 cell 增量 = offset_tok[k] − offset_tok[k−1]，cell 0 = 0（速度）
  "action_idx": int64   [b, 21],      # 9 类离散标签（0 idle，1 up … 按 prompts.ACTIONS），已翻回真实屏幕方向
}
```

**主信号是 `offset_px` / `offset_tok`（累计偏移），与旧报告一致。** 上游代码用离散按键的（Matrix-Game 键盘）
按其原样用 `action_idx`（one-hot / embedding），并在 docstring 里写明；若上游同时接受连续量（鼠标/相机），
连续量用 `delta_tok`。每个臂在 `baseline/<name>.py` 顶部 docstring 写清：上游是怎么做的（文件:行）、这里怎么对应、
用了哪个信号、新增参数量、初始化方式（**新增路径必须零初始化或等效地保证"装上瞬间输出与原模型相同"**）。

### 文本

`linear` / `xattn` / `adaln` / `arope`：`text_mode=scene`（动作词剥离，`text_table_actionrope.pt`）。
`prompt`：`text_mode=scene_action`（动作词翻回真实方向，`text_table_eybx_mirror.pt`），逐 cell 文本已经是
`arope_forward` 的原生能力（每帧 token 只 attend 自己 cell 的文本）。Incantation 若有别的逐 latent 做法
（例如逐帧 prompt 掩码、专门的 attention），按上游移植并说明差异。

### 训练 / 推理接入（integration 阶段做，臂的作者不要改这些文件）

* `actionrope/train.py --arm {arope,plain,linear,xattn,prompt,adaln}`：`plain` 保留（= prompt 臂的旧名，行为不变）；
  新臂：`install_arope(dit, mask_channel=False)`（普通 RoPE、无 mask、权重全 1）+ `arm = baseline.build_arm(name, dit)`，
  arm 的参数与 DiT 一起训练（同 lr）；ckpt 导出时 arm 参数以 `arm.` 前缀一起写进 safetensors；config.json 记 arm 名与新增参数量。
* `actionrope/infer.py --arm <name>`：按 ckpt 元数据/键集自动装 arm；动作脚本 → `action_inputs`（`offset_px` 同 ARoPE；
  `action_idx` 由每 cell 速度方向得到，同 `velocity_to_label(mirror=False)`）。
* `baseline/__init__.py::build_arm(name, dit) -> ActionArm`、`ARMS = {...}`。

### 测试（每个臂 `tests/test_arm_<name>.py`，单卡，模型只加载一次）

1. 装上后 **输出与原模型逐位相同**（零初始化）——`arope_forward(arm=arm)` vs `arope_forward()`。
2. 给非零动作、把新增参数随机扰动后输出**有变化**且梯度能回到新增参数与 DiT（`use_gradient_checkpointing=True`）。
3. 新增参数量打印并与上游/旧报告对照（linear ≈ 0.18M，xattn ≈ 98.6M；差得远要解释）。
4. ckpt 键集：`arm.state_dict()` 键名无冲突、可 strict 加载。
5. 单卡 2 步真实训练 smoke（`--limit 8 --max_steps 2`）在 integration 阶段跑。

## 纪律

只在 `/opt/dlami/nvme/zhiyangdeng/ActionRoPE` 内读写；新文件只放 `baseline/`、`tests/`；`vender/` 只读（不改、不删、不跑它们的脚本、不下载它们的权重）；
不改 `actionrope/*.py`（钩子已加好；确需改动写进 report 交 integration 决定）；不动 data/、models/、DiffSynth-Studio/；
禁止 rm -rf、sudo、pip 卸载/升降级、杀非自己启动的进程、任何网络上传/下载（vender 已克隆好）；不读凭据。
GPU：正在跑 8 卡训练（每卡 24 GB 已用，143 GB 总），单卡前向+反向 smoke 用分配给你的那一张，跑完立刻释放。