code/actionrope/SPEC.md

# ActionRoPE 实现契约（所有模块必须遵守）

复现目标（来自旧机器的两份报告 `恶意不息世界画布` / `世界RoPE·四科判决`）：
在 **Wan2.2-TI2V-5B** 上做动作条件世界模型，**动作不走文本、走 RoPE**：

1. **世界坐标 RoPE（ARoPE，零新增参数）**：每帧 token 的 (h, w) 位置编码 =
   窗口内格坐标 + 该帧相机累计偏移（换算成 token 单位）。同一世界位置 ⇒ 同一编码。
2. **known / new 支持掩码**：首帧足迹是全局唯一"真"区域；某帧某 token 的世界位置落在
   首帧足迹外 ⇒ `new`。new 区 loss ×2（`new_weight`），并可选把 mask 作为一个额外输入通道
   经 **零初始化 conv** 加到 patch embedding 上（`mask_channel`，默认开；关掉即严格零参数版）。
3. 文本只承载场景（动作从句已剥离，`data/latent/text_table_actionrope.pt`）；
   转场语料的 trigger 从句按 cell 变化 ⇒ **逐 cell 文本条件**（每帧 token 只 cross-attend
   自己 cell 的文本；所有 cell 文本相同时与原模型逐位等价）。
4. 对照臂 `plain`：普通 RoPE + 文本带动作，即报告里的 "Wan + Prompt"。**数据集的动作词是镜像的**
   （日志→屏幕矩阵 M 符号翻转，2026-09-04 人工看视频确认："moving left" 的 clip 玩家实际向右走），
   plain 臂训练/推理默认用 `prompts.mirror_action` 把动作词翻回真实屏幕方向，文本表用补齐了镜像串的
   `text_table_eybx_mirror.pt`（`actionrope/build_mirror_table.py` 生成，原 1,253 键 + 530 条镜像转场串）。

## 常量与几何

| 量 | 值 |
|---|---|
| 画幅 | 832×480（W×H），16 fps，81 帧 |
| latent | `[48, 21, 30, 52]`（C, F, H_lat, W_lat），latent 像素 = 16 px |
| token | patch (1,2,2) ⇒ 每帧 15×26 token，token = **32 px** |
| cell ↔ 帧 | cell 0 = 帧 0；cell k (k≥1) = 帧 4k−3 … 4k（含） |
| 屏幕坐标 | u 向右、v 向下，单位 px（832×480 画幅） |

### 相机 / 世界 → 屏幕

日志 `data/logs/<session>/state.npz` 是 20 Hz 玩家世界坐标 `(vt, x, y, z)`。
`offsets.json`：`video_t = log_vt + offset(log_vt)`（分段常数），`M_px1280` 是 3×2 矩阵，
行对应 (x, y, z)，列对应 1280×720 画幅下的 (u, v)：

```
p_px1280 = [x, y, z] @ M_px1280            # 玩家的"屏幕投影"，含高度项
p_px832  = p_px1280 * [832/1280, 480/720]
```

相机刚性钉在玩家上（无平滑无滞后），所以**背景相对屏幕的位移 = −(玩家屏幕投影位移)**。

### 相机偏移（sidecar 的核心量）

对 clip 的第 i 帧（video_t_i = video_t0 + i/16，日志按 video_t 最近邻采样，与 `make_clips.py` 一致）：

```
frame_offset_px[i] = p_px832(frame i) − p_px832(frame 0)        # (dx, dy)，frame 0 恒为 (0,0)
offset_px[k]       = mean_{i ∈ cell k} frame_offset_px[i]        # [21, 2]，cell 0 = (0,0)
```

**世界 token 坐标**：cell k 中屏幕格 (i_h, j_w) 的 token 的世界坐标为
`(i_h + offset_px[k].dy / 32,  j_w + offset_px[k].dx / 32)`（允许小数，RoPE 是连续的）。
理由：玩家向右走 D px ⇒ 背景整体左移 D px ⇒ 现在屏幕位置 s 上的内容，在首帧位于 s + D。

**known / new**（token 分辨率）：token 的世界坐标中心落在首帧足迹
`[−0.5, 14.5] × [−0.5, 25.5]`（h × w）内 ⇒ known，否则 new。
latent 分辨率同理（30×52，偏移除以 16）。

**转场 clip**（`kind ∈ {burn, hardcut}`，帧 < n_pre 来自 `src_a` 的同下标帧，帧 ≥ n_pre 来自 `src_b`
的同下标帧，burn 段是两者混合）：

```
i <  n_pre : off[i] = offA[i]
i >= n_pre : off[i] = offA[n_pre−1] + (offB[i] − offB[n_pre−1])     # 累计位移保持连续
```

转场 clip 的 known / new 掩码**只按上面的几何定义**：B 段（帧 ≥ n_pre）的内容来自另一个场景，
首帧看不到，但其 token 若几何上落在首帧足迹内仍标 known（loss 权重 1、mask 通道 = 1）。
这是有意的简化（模型靠逐 cell 文本知道换了场景），不是缺陷；要按语义把 B 段整体标 new 需改此处。

### 增益校验（报告里的教训）

旧报告说 manifest 常数被逐帧小位移法的亚像素零偏压低 ~11%，真值应为横 27.5 / 纵 18.7 px/米；
本机 `offsets.json` 的 M 折算后为横 26.6 / 纵 16.9。**不要照抄任何一边**：sidecar 必须附带
对视频的实测校验（大基线相位相关/模板匹配，抽样 ≥200 个高速直行 clip），报告拟合增益比
`gain_fit = (实测位移) / (日志预测位移)` 及 R²。若两 session 一致且 R² ≥ 0.9，则把
`gain_correction`（分 x/y）写进 sidecar meta 并在生成 `offset_px` 时**已乘上**；否则保持 1.0 并在 meta 里说明。
（2026-09-04 人定例外：纵向 R² 0.92/0.84 未达门槛但 step-5000 推理证实 1.0 欠标定 27%，改用实测 1.12。）

## 文件契约

### `data/latent/arope_sidecar.pt`（由 `actionrope/sidecar.py` 生成）

```python
{
  "meta": {
    "version": "arope_sidecar_v1",
    "M_px832": [[..],[..],[..]],          # 3x2，已换算到 832x480
    "gain_correction": [gx, gy],          # 已应用到 offset_px 的修正倍率，默认 [1.0, 1.0]
    "gain_fit": {...},                    # 校验统计：每 session 的拟合斜率 / R² / 样本数
    "walk_speed_px_s": [sx, sy],          # 直行 clip 里的中位步行速度（|dx|/s, |dy|/s），供推理造指令用
    "n_clips": int, "n_valid": int, "notes": str,
  },
  "clips": {
    "<clip_id>": {
      "offset_px":       torch.float32[21, 2],   # (dx, dy)，cell 0 = 0
      "frame_offset_px": torch.float32[81, 2],
      "valid": bool,                             # 日志缺口 / 采样异常 ⇒ False（训练时剔除）
      "kind": "plain" | "burn" | "hardcut",
    }, ...
  }
}
```

覆盖 `data/latent/pool/` 全部 22,782 个 clip。

### 文本表

`data/latent/text_table_actionrope.pt`：`dict[str, bf16[512, 4096]]`，键 = `strip_action(prompt)`
（`actionrope/prompts.py`）。含空条件串 `NULL_PROMPT`。`plain` 臂用 `text_table_eybx_mirror.pt`
（键 = `mirror_action(原 prompt)`，是 `text_table_eybx.pt` 的超集）；`--no-mirror_labels` 旧约定才用原表、原串。

### 几何工具 `actionrope/geometry.py`（sidecar 与推理共用，避免两处各写一套）

```python
FRAMES, CELLS, W, H = 81, 21, 832, 480
LAT_H, LAT_W, TOK_H, TOK_W = 30, 52, 15, 26
PX_PER_TOKEN, PX_PER_LATENT = 32, 16
def cell_frames(k) -> list[int]
def frames_to_cells(frame_offset_px: np.ndarray[81,2]) -> np.ndarray[21,2]
def load_session_tracks(logs_dir, session) -> dict   # 排序去重后的 vt/x/y/z + offset 分段 + M_px832
def project_px832(tracks, x, y, z) -> (u, v)
def sample_frame_offsets(tracks, video_t0, n_frames=81) -> np.ndarray[81,2]   # 含 valid 判定
def known_mask_tokens(offset_px[F,2]) -> np.bool_[F, 15, 26]
def known_mask_latent(offset_px[F,2]) -> np.bool_[F, 30, 52]
```

## 模型契约 `actionrope/arope.py`

基于 `diffsynth.models.wan_video_dit.WanModel`（TI2V-5B 配置：dim 3072, 24 头, head_dim 128,
30 层, patch (1,2,2), in/out 48, `seperated_timestep=True`, `fuse_vae_embedding_in_latents=True`）。
RoPE 分量：`precompute_freqs_cis_3d(128)` ⇒ f 轴 44 维、h 轴 42 维、w 轴 42 维（`dit.freqs` 实际是 complex128，theta 1e4）。

```python
def build_world_freqs(dit, f, h, w, offset_tok=None, device=None) -> torch.complex128
    # offset_tok: [b, f, 2] (dx_tok, dy_tok) 或 None（None ⇒ 与 dit.freqs 查表结果逐位相同）
    # 返回 [b, f*h*w, 1, 64]（offset_tok=None 时可返回 [f*h*w, 1, 64]）
    # 位置 = (frame_idx, i_h + dy_tok[f], j_w + dx_tok[f])，h/w 轴用 polar(1, pos * inv_freq) 现算

def known_mask(offset_tok, h, w) -> torch.bool [b, f, h, w]
def loss_weight_map(offset_px, lat_h=30, lat_w=52, new_weight=2.0) -> torch.float32 [b, 1, f, lat_h, lat_w]

def install_arope(dit, mask_channel: bool = True):
    # mask_channel ⇒ 新增 dit.arope_mask_embedding = nn.Conv3d(1, dim, patch, stride=patch)，零初始化
    # 不改动任何已有参数名；plain / 关闭时 ckpt 键集与原模型完全一致

def arope_forward(dit, latents, timestep, context, *,
                  offset_px=None,            # [b, f, 2]，None ⇒ 普通 RoPE
                  first_frame_cond=True,     # 帧 0 为干净条件帧：timestep 置 0（seperated_timestep 路径）
                  mask_input=None,           # [b, 1, f, 30, 52] float，known=1 / new=0；None ⇒ 不加
                  use_gradient_checkpointing=False,
                  use_gradient_checkpointing_offload=False) -> torch.Tensor [b, 48, f, 30, 52]
    # context: [b, 512, 4096]（整条一致）或 [b, f, 512, 4096]（逐 cell）。逐 cell 时 cross-attn 按帧分组：
    #   q: (b f) (h w) c，k/v: (b f) L c。所有 cell 相同 ⇒ 结果与整条模式逐位等价（bf16 容差）。
    #   text_embedding 只对去重后的串算一次。
    # 等价性要求：offset_px=None, mask_input=None, context 整条 ⇒ 与 diffsynth.pipelines.wan_video.model_fn_wan_video(
    #   dit=dit, latents=..., timestep=..., context=..., fuse_vae_embedding_in_latents=True) 输出一致（bf16 容差 ~1e-2 相对）。
```

时间步处理照抄 `model_fn_wan_video` 的 `seperated_timestep and fuse_vae_embedding_in_latents` 分支
（帧 0 的 t=0，其余帧 t=timestep，t_mod 逐 token）。block 调用照抄 `WanModel.forward`（含梯度检查点）。

## 训练契约 `actionrope/train.py`

* 只加载 DiT（`ModelConfig(path=[三个 safetensors 的绝对路径])`，`redirect_common_files=False`），
  bf16；不挂 T5/VAE。latent 与文本都是预计算的。
* 数据：`data/latent/train_eybx/*.pt`（软链）+ sidecar + 文本表；`valid=False` 的 clip 剔除；
  **转场 clip 的 `src_a` / `src_b` 若落在 val 目录也剔除**（val 的帧会原样出现在转场 clip 里，划分脚本没按 src 归属划）；
  sidecar 的 valid 规则：最近邻日志样本时间差 ≤ **0.5 s**（原 0.15 s；2026-09-04 视频实测 0.15–0.5 s 缺口的 clip
  偏移精度与 valid 基线一致，残差 5–7 px，人定放宽）。世界降速另用 meta 的 `rt_ratio` 闸：
  训练 `--min_rt_ratio 0.8`（默认），val / 评测 `--val_min_rt_ratio 0.9`（默认）。
  样本 = `{input_latents [48,21,30,52] bf16, context [21,512,4096] bf16, offset_px [21,2], clip_id}`。
  场景 dropout：以 `--scene_dropout p`（默认 0.1）把整条 21 个 cell 都换成 `NULL_PROMPT`
  （plain 臂换成 .pt 自带的 `prompts_bossdrop` 裸动作串并同样镜像，`NULL_PROMPT` 不在 eybx 表里）。
  随机数按 rank 错开（`set_seed(seed, device_specific=True)`）：噪声 ε、dropout 抽签各卡独立，数据分片由 accelerate 同步。
* loss（与 DiffSynth `FlowMatchSFTLoss` 一致，加权重）：
  `timestep` 从 `FlowMatchScheduler("Wan").set_timesteps(1000, training=True)` 的 1000 个时间步里均匀抽；
  `x_t = (1−σ) x0 + σ ε`，目标 `ε − x0`；帧 0 用干净 x0 替换（TI2V 首帧条件）并在 loss 中剔除；
  `loss = Σ w·(pred−target)² / Σ w`（w 来自 `loss_weight_map`，帧 0 不计），再乘 `scheduler.training_weight(timestep)`。
* 分布式：`accelerate launch --config_file configs/accelerate_zero2.yaml`，DeepSpeed ZeRO-2，bf16，
  每卡 batch 1，梯度检查点开，AdamW（lr 默认 1e-5, wd 0.01），线性 warmup 可选。
  `--grad_clip` 写进 DeepSpeed 的 `gradient_clipping`（日志 `grad_norm` 是裁剪前的全局范数）。
  `--save_state` 写 `state_A` / `state_B` 两槽轮转，`<output>/state` 软链指向最新完整槽。
* 参数：`--arm {arope,plain}`、`--mask_channel/--no_mask_channel`、`--new_weight`、`--text_table`、
  `--scene_dropout`、`--lr`、`--max_steps`、`--save_every`、`--val_every --val_n`、`--output`、
  `--resume`、`--limit`（smoke 用）、`--seed`、`--num_workers`。
* 输出 `outputs/<run>/`：`step-N.safetensors`（只含 DiT 权重，键名与原模型一致 + 可选 `arope_mask_embedding.*`），
  `train_log.csv` / tensorboard，`config.json`（含 arm、gain_correction 等复现信息）。
* plain 臂 = `offset_px=None` + 无 mask + 权重全 1 + 文本带动作从句（默认 `--mirror_labels`：动作词翻回真实方向）。

## 推理契约 `actionrope/infer.py`

* 输入：首帧（png / mp4 首帧 / `clip_id`，取 clip mp4 的第 0 帧）、场景从句、动作脚本、ckpt。
* 动作脚本 → `frame_offset_px[81,2]`（用 sidecar meta 的 `walk_speed_px_s` × 倍率 × 方向单位向量，逐帧累计），
  再 `frames_to_cells` ⇒ `offset_px[21,2]`。`--replay <clip_id>` 直接用 sidecar 真值（可与 GT clip 逐帧比较 PSNR）。
* 采样：DiffSynth `FlowMatchScheduler("Wan")`，50 步，shift 5，帧 0 = VAE 编码的首帧 latent（fp32 编码后存 bf16，
  与数据集 latent 同精度）并每步保持；CFG 可选（负条件：arope 臂 `NULL_PROMPT`；plain 臂逐 cell 裸动作串
  `In a 2.5D top down view, the player is {action}.`，即该臂训练 dropout 用的串）。
  VAE：`models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth`（`WanVideoVAE38`）。
* 输出 mp4（832×480，16 fps）+ 一行文本（实测位移：对生成视频做相位相关得到的累计 dx/dy，与指令对比）。

## 环境与纪律

* 解释器：`/opt/dlami/nvme/zhiyangdeng/ActionRoPE/.venv/bin/python`（3.12，torch 2.11+cu128，8×H200）。
  允许往 venv 里 `pip install` 缺的包（如 `opencv-python-headless`、`scipy`）。
* 所有路径用绝对路径；只在 `/opt/dlami/nvme/zhiyangdeng/ActionRoPE` 内读写；大文件（sidecar、ckpt、样例视频）
  只能落在 `data/` / `outputs/`。临时文件放 scratchpad，不要写到 `/`、`/root`、`/tmp` 之外。
* **不要读取、修改任何 AWS 鉴权信息（ak/sk、`~/.aws`、环境变量凭据）**。
* 不要启动超过 smoke 规模的训练（≤ 数十步）。GPU 分配以任务指令为准，避免互相抢卡。