| """`linear` 臂：ReactiveGWM 的逐块线性动作偏置（30 个 bias-free Linear，加在每个 DiT block 入口）。 |
| |
| 上游机制（vender/ReactiveGWM，只读） |
| ----------------------------------- |
| * 模型 `inference/models/dit.py::WanModelAction`（同一个 Wan2.2-TI2V-5B 基座，dim 3072 / 30 层）： |
| - L207-209 `self.action_embedders = nn.ModuleList([nn.Linear(num_buttons, dim, bias=False) for _ in range(num_layers)])` |
| —— 每个 block 一个 **无偏置** 线性层，把动作向量投到 hidden 维；没有门控（docstring："No gates"）。 |
| - L222-226 `_bin_action`：原始逐帧键盘 one-hot `[B, T_raw, K]` 用 `adaptive_max_pool1d` 压到 latent 帧数 `[B, f, K]` |
| （每个 latent 帧 = 4 个视频帧的"按过即 1"）。 |
| - L228-234 `_inject_action`：`emb = action_embedders[i](action[B,f,K].to(x.dtype))` → `[B, f, C]`， |
| 沿 h·w 展开成逐 token 偏置 `[B, f·h·w, C]`，`x = x + bias`。 |
| - L284-286 前向：`for i, block: x = _inject_action(x, action, i, f, h, w); x = block(x, ctx, t_mod, freqs)` |
| —— 偏置加在 **block 入口的残差流**（进 norm1/self-attn 之前），不是 adaLN、不是 cross-attn。 |
| * 动作向量 `training/data/action_utils.py` / `inference/utils/actions.py`：parquet 的 10 个按键列 |
| （`inference/constants.py::SF_BUTTON_COLS`，UP/DOWN/LEFT/RIGHT + 6 攻击键）0/1 值，`hold_last_upsample`（10 帧窗） |
| 补成逐视频帧 dense one-hot，不做归一化。 |
| * 初始化：`training/bidirectional/train.py` L17-18 / L134-137："ActionModule keys stay at their default (zero / xavier) init"， |
| 其余权重从 Wan2.2-TI2V-5B 按形状拷贝（`_transfer_weights`）。训练时 action_embedders 与 DiT 全参数同 lr 训练； |
| 推理有 `action_cfg_scale`（`inference/pipeline.py` L218-224，动作置零做无动作分支）。 |
| * 上游参数量：30 × 10 × 3072 = 921,600（10 键）。 |
| |
| 本实现的对应（baseline/SPEC.md 的钩子） |
| -------------------------------------- |
| * `encode`：取 `action_inputs["offset_tok"]`（float32 `[b, 21, 2]`，逐 cell 累计屏幕位移 (dx, dy)，单位 token， |
| 与 ARoPE 臂**完全相同的控制信号**——旧报告口径"带动作的臂收到相同的逐 cell 累计偏移"）。 |
| 上游的"逐帧 one-hot → adaptive_max_pool 到 f 个 latent 帧"这一步，在我们这边由 sidecar 的 |
| `frames_to_cells`（cell 内逐帧偏移取均值）已经做完，dataset 直接给逐 cell 的 `[b, 21, 2]`， |
| 所以这里不再池化，只做形状校验。不归一化（上游也不归一化；零初始化下量纲只影响学习动态，不影响装上瞬间的等价性）。 |
| cell 0 的偏移恒为 0 ⇒ 无偏置 Linear 在条件帧上的偏置恒为 0，与"首帧是干净条件帧"一致。 |
| * `block_pre(i, x)`：`x + expand(action_embedders[i](offset_tok))`，与上游 `_inject_action` 同位置同算式。 |
| 小矩阵乘在 fp32 里算再转回 x.dtype：offset_tok 最大约 ±11 token，bf16 在该量级的分辨率是 1/16 token = 2 px， |
| fp32 保住亚像素信息；上游输入是 0/1 one-hot，bf16 本来就无损，所以它直接 `.to(x.dtype)`。 |
| * 参数量：30 × 2 × 3072 = **184,320 ≈ 0.18M**（与旧报告 +0.18M 一致：旧报告就是 2 维偏移输入）。 |
| * 初始化：**全零**（SPEC 要求"装上瞬间输出与原模型逐位相同"；上游 docstring 也把 zero 列为默认之一）。 |
| 零权重 ⇒ 偏置张量逐位为 0 ⇒ `x + 0 == x` 逐位相等。 |
| * ckpt 键：`action_embedders.{i}.weight`（与上游键布局同名；integration 阶段导出时加 `arm.` 前缀，与 DiT 键无冲突）。 |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from baseline.base import ActionArm |
|
|
| |
| ACTION_DIM = 2 |
|
|
|
|
| class LinearArm(ActionArm): |
| name = "linear" |
|
|
| def __init__(self, dim: int = 3072, num_layers: int = 30, in_dim: int = ACTION_DIM): |
| super().__init__() |
| self.dim, self.num_layers, self.in_dim = dim, num_layers, in_dim |
| self.action_embedders = self._build(dim, num_layers, in_dim) |
|
|
| @staticmethod |
| def _build(dim: int, num_layers: int, in_dim: int) -> nn.ModuleList: |
| |
| embedders = nn.ModuleList([nn.Linear(in_dim, dim, bias=False) for _ in range(num_layers)]) |
| for lin in embedders: |
| nn.init.zeros_(lin.weight) |
| return embedders |
|
|
| def install(self, dit) -> None: |
| """按 dit 的 dim / 层数建参数，并放到 dit 的 device / dtype（与 DiT 一起以 bf16 训练）。""" |
| dim, num_layers = int(dit.dim), len(dit.blocks) |
| if (dim, num_layers) != (self.dim, self.num_layers): |
| self.dim, self.num_layers = dim, num_layers |
| self.action_embedders = self._build(dim, num_layers, self.in_dim) |
| ref = dit.patch_embedding.weight |
| self.to(device=ref.device, dtype=ref.dtype) |
|
|
| |
| def encode(self, action_inputs: dict | None, grid: dict): |
| """offset_tok [b, f, 2] → 同一张量（float32）。None ⇒ 不注入（对应上游 keyboard_action=None 的分支）。""" |
| if action_inputs is None: |
| return None |
| off = action_inputs["offset_tok"] |
| b, f = grid["b"], grid["f"] |
| assert off.shape == (b, f, self.in_dim), f"offset_tok 形状应为 [{b}, {f}, {self.in_dim}]，得到 {tuple(off.shape)}" |
| return off.to(torch.float32) |
|
|
| def block_pre(self, idx: int, x: torch.Tensor, feats, grid: dict) -> torch.Tensor: |
| """上游 `_inject_action`：逐帧线性偏置沿帧内 h·w 个 token 广播后加到残差流。""" |
| if feats is None: |
| return x |
| b, f, n = grid["b"], grid["f"], grid["n_tok_per_frame"] |
| weight = self.action_embedders[idx].weight |
| emb = F.linear(feats.to(device=x.device), weight.to(torch.float32)).to(x.dtype) |
| bias = emb.unsqueeze(2).expand(b, f, n, self.dim).reshape(b, f * n, self.dim) |
| return x + bias |
|
|
| |
| def zero_init_check(self) -> bool: |
| return all(bool((lin.weight == 0).all()) for lin in self.action_embedders) |
|
|
