docs/TREX_TRACK_FORCE_MODEL.md

# T-Rex Track-Force：16-step、20 Hz 动作与 5 Hz 触觉模型

本文描述独立的 `trex_track_force` 模型、数据契约、embedding、注意力、两阶段
flow matching、训练和在线执行。它不会修改原 DreamZero 的模型调用路径。

## 1. 固定时序契约

- 每个 action chunk 含 **16 个动作**，动作频率为 **20 Hz**。
- 一个 chunk 覆盖 `16 / 20 = 0.8 s`；动作时间戳跨度为 `0.75 s`。
- 触觉频率为 **5 Hz**，因此每 **4 个动作步**刷新一次。
- 一个 chunk 内的触觉刷新 offset 固定为 `[0, 4, 8, 12]`。
- 每次刷新只重新去噪尚未执行的 action suffix；已经执行的前缀保持不变。
- `max_chunk_size=4` 保留 DreamZero 原有的 autoregressive block memory。
  这里的 `max_chunk_size` 是跨 block 的记忆窗口，不是单个 action chunk 的长度。
- 一个训练 sample 仍由 4 个 autoregressive blocks 组成，因此共有 64 个动作；
  只有每 block 的 action 长度从 24 改为 16。
- 三路视频保持 DreamZero 的 33 帧布局：1 帧 clean conditioning observation，
  后接每 block 8 帧、共 32 帧 10 Hz future-video targets。
- 每个 block 都有 16 帧过去 track 和 16 步 target track；target 的第 0 步就是
  当前 anchor 帧，与 OpenPI 的 track window 契约一致。

## 2. 总体架构

```mermaid
flowchart LR
    V["三视角 RGB\nhead + left wrist + right wrist"] --> GRID["三视角拼图\n160 × 320"]
    GRID --> VAE["Wan2.2 VAE38\n48-channel latent"]
    GRID --> CLIP["Wan CLIP\n图像条件"]
    TXT["任务文本"] --> T5["UMT5\n文本条件"]

    TRK["250 点 track\npast clean + future noisy"] --> TE["TrackEncoder\n几何 + 可见性 + 身份 embedding"]
    ACT["16 × 62D delta-base action\n补零到 64D"] --> AE["Action encoder"]
    ST["62D 当前 EEF/hand state"] --> SE["State encoder"]

    VAE --> WAN["CausalWanTrackForceModel\nWan2.2-TI2V-5B · 30 blocks"]
    CLIP --> WAN
    T5 --> WAN
    TE --> WAN
    AE --> WAN
    SE --> WAN

    WAN --> COARSE["τ: 1 → 0.4\n6-step coarse action"]
    WAN --> TFLOW["未来 track flow"]
    WAN --> VFLOW["未来 video flow"]
    WAN --> MEM["τ=0.4 coarse memory"]

    RAW["10 fingers × 6D force\n16-sample history"] --> VQ["Per-finger force VQ-VAE\n64 codes · 256D"]
    FRESH["当前 10 × 6D force"] --> FT["Force-only Transformer\n6 layers · width 768"]
    VQ --> FT
    MEM --> FT
    COARSE --> FT
    FT --> FINE["τ: 0.4 → 0\n4-step tactile action suffix"]
```

粗阶段负责从视觉、状态、语言和 track 中得到动作的大尺度结构；独立的
force-only transformer 只使用力信号和粗阶段 memory 完成剩余去噪，使触觉
负责接触后的细粒度修正。

## 3. 62DoF 动作空间

每只手臂占 31 维，两侧共 62 维：

- 手腕/末端位姿：`xyz + rotation-6D`，共 9 维。
- 手部关节目标：22 维。

训练动作不是世界坐标绝对位姿。loader 以 chunk 起始状态为 reference，把每个
目标手腕位姿转换为 **delta-base**：

- 平移在 chunk 起始手腕坐标系中表示。
- 旋转为 `R_reference^-1 × R_target`，再编码为 rotation-6D。
- 手部 22 维仍是绝对目标。

模型内部把 62 个物理维补零为 64 维；采样的两个 padding 维在初始化和每个
Euler step 后都会重新置零。输出先按 relative-action 统计量反归一化，再用
chunk 起始 state 恢复绝对 EEF 位姿。

## 4. Track 数据和点身份

250 个点具有固定、可验证的顺序：

- `0:50`：头部视角左手及手臂。
- `50:100`：头部视角右手及手臂。
- `100:125`：左腕视角 5×5 背景点。
- `125:175`：左腕视角手掌点。
- `175:200`：右腕视角 5×5 背景点。
- `200:250`：右腕视角手掌点。

SAM2 只在 episode 第 0 帧根据固定 prompt 产生手/臂 mask；点从 mask 中采样，
随后由 CoTracker 跟踪整段视频。保存字段为：

- `observation.track_xy`：`[250, 2]`，坐标归一化到 `[0, 1]`。
- `observation.track_visibility`：`[250]`。
- metadata 中同时记录 view、hand、role、point index 和各 segment 边界。

## 5. Embedding 设计

### 5.1 Video、语言和图像条件

- 三视角拼入 2×2 canvas 的三个有效格，再统一缩放到 `160×320`。
- Wan2.2 VAE38 生成 48-channel latent。
- `patch_size=[1,2,2]`，每 latent frame 形成 50 个 video tokens。
- video token 使用 Wan 原生 3D RoPE，编码时间、高度和宽度。
- UMT5 输出 4096D 文本条件。
- Wan CLIP 输出 1280D 图像条件，并经 `img_emb` 投影。
- CLIP/T5 cross-attention 只作用于 observation/action query，不向 track query
  泄漏视觉或语言信息。

### 5.2 Action 和 state

- action：每个 62D 动作补零到 64D，经线性层投影到 Wan hidden dim 3072。
- 一个 block 有 16 个 action tokens；位置由 1D action RoPE 编码。
- state：62D 当前状态补到 64D，经线性层得到一个 state token。
- flow 时间 `τ` 经 sinusoidal embedding 和 MLP 后参与 Wan modulation。

### 5.3 TrackEncoder

每个点分别产生一个 `past token` 和一个 `future token`。时间序列输入特征为：

```text
[x, y, visibility, Δx, Δy]
```

不可见时刻的坐标和 motion 先清零，随后完整的 `16×5` 时序按固定顺序展平，再经
`Linear → SiLU → Linear` 投影。不能在时间维求平均：future flow 的每个时刻含有
独立噪声，平均会丢失“哪一个噪声属于哪一个 timestep”，使 16-step 重建不可解。

最终每个点 token 是下列 embedding 的和，再经 LayerNorm：

```text
trajectory
+ view(head / left_wrist / right_wrist)
+ hand(none / left / right)
+ role(head_hand / wrist_background / wrist_hand)
+ canonical point id(0...249)
+ autoregressive block id
+ temporal role(past / future)
```

左右手不只通过点序号区分，还具有显式 hand embedding。背景点使用
`hand=none` 和独立 role embedding。

与 OpenPI 一致，track target 始终是 `[0,1]` 内的绝对归一化 XY；target window
从当前帧开始，因此第 0 步等于 GT anchor。`TrackDecoder` 对 250 个
future-track hidden token 分别执行 `LayerNorm → Linear(16×2)`，预测 CFM
velocity `noise-clean_xy`。Euler 积分后的结果已经是绝对坐标，不做空间
`cumsum`，也不再额外加 anchor。

### 5.4 Force-only VQ-VAE

输入只包含 force/wrench，不包含 deformation map：

```text
[batch, history=16, fingers=10, wrench=6]
```

- 两只手共享同一套时序卷积 encoder/decoder。
- 每只手内部加入 5 个 finger identity embeddings。
- 输出每根手指一个 256D latent，共 10 个 tactile history tokens。
- EMA codebook 大小为 64，带 commitment loss、perplexity 统计和 dead-code
  revival。
- 左右手通过 side embedding 区分，手指位置通过 finger embedding 区分。
- 训练可直接输入 raw 16-step history；推理也可输入预计算的 10 个离散 codes。
- episode 前缀缺失的历史填为归一化中性值，并用 validity mask 从 VQ 重建
  loss 中排除 padding 步。

### 5.5 Force-only Transformer

输入 token 包括：

- 16 个 noisy action tokens。
- 10 个当前 force tokens。
- 10 个 VQ history tokens。
- 从 Wan `τ=0.4` hidden state 提取的 coarse memory tokens。

action token 叠加 action position、5Hz force slot、token type、当前 refresh
offset 和 `τ` embedding。force/history token叠加 finger position、token type 和
refresh offset。独立 transformer 为 6 层、12 heads、hidden dim 768。

## 6. 非对称 attention 契约

每个 autoregressive block 的 packed 顺序为：

```text
[obs, action, state, track_past, track_future]
```

其可见性严格为：

- observation query 和 action query 在同一 block 内互相可见。
- observation/action query 可读取同 block 的 state、past track 和 noisy future
  track，也可读取有限 AR 窗口内的历史 obs/action/track。
- state query 只读取自己的 state token，避免把多模态信息反向带给 track。
- past-track query 读取历史 block 的 track 和当前 past track，但不能读取当前
  noisy future track。
- future-track query 只读取当前及历史 track。
- **任何 track query 都不能读取 observation、action、state、CLIP 或文本。**
- 任何 query 都不能读取未来 block。

因此实现了要求的方向性：`obs/action → track` 表示 obs/action 可以把 track
作为条件；反方向被 mask 禁止。

## 7. 两阶段 flow matching

使用线性插值：

```text
x_τ = τ · noise + (1 - τ) · clean
target flow = noise - clean
x_next = x_τ + (τ_next - τ) · predicted_flow
```

### 粗阶段：Wan，`τ=1 → 0.4`

- 训练时 action/Wan expert 遵循原始 T-Rex，在完整 `(0,1]` 上按
  `Beta(1.5,1.0)` 采样；其均值约为 `0.6`，但这不是 split timestep。
- Wan 联合预测 action flow、future-track flow 和 future-video flow。
- 推理按总计 10 个、`Δτ=-0.1` 的 Euler 网格运行前 6 步，到达 `τ=0.4`。
- 在精确的 `τ=0.4` 再执行一次 Wan，生成给触觉 transformer 使用的 detached
  coarse memory。

### 精阶段：触觉，`τ=0.4 → 0`

- 训练时 `τ_tactile = 0.4 × Beta(1.5,1.0)`，覆盖 `(0,0.4]`。
- 每个训练 chunk 对 offset `0/4/8/12` 分别提供对应的 force 与 16-sample
  history。
- offset 为 `k` 时，loss 只覆盖 action `k:16`。
- 推理运行剩余 4 个 Euler updates；每次 5Hz 刷新只写入未执行 suffix。
- runtime controller 保存之前已经发出的 prefix，确保后续触觉刷新不能改写历史
  command。

## 8. 训练 loss

总 loss 由以下部分加权求和：

- `dynamics_loss`：未来视频 latent flow MSE。
- `action_loss`：粗阶段 62 个物理动作维的 flow MSE。
- `track_loss`：可见性 mask 后的绝对 XY CFM velocity MSE。
- `force_loss`：按 refresh offset mask 后的动作 suffix flow MSE。
- `vq_loss`：仅在有效历史步计算的 force history reconstruction loss。
- `commitment_loss`：VQ commitment loss，默认系数 0.25。

训练使用 Wan2.2-TI2V-5B 共享权重和 LoRA。LoRA 注入
`q,k,v,o,k_img,v_img,ffn.0,ffn.2`；新建的 action/state/track/force、video token
投影和 decoder 完整训练。checkpoint 保存 LoRA、新模块参数以及 VQ EMA buffers。

## 9. 在线执行

1. 收集三视角 conditioning observation（在线接口也可编码多帧历史）、当前
   62D state、16 帧 past track、当前 force 和 16-sample force history。
2. Wan 对 action、future track 和 future video 运行 6 步，从 `τ=1` 到
   `τ=0.4`。
3. offset 0 的触觉运行剩余 4 步，把完整 16-step action 去噪到 `τ=0`。
4. 以 20 Hz 执行动作。
5. 执行 4、8、12 步后，各接收一次新的 5Hz 触觉；每次从缓存的 coarse state
   重新去噪剩余 suffix，同时保留已执行 prefix。
6. 输出 delta-base 动作反归一化并恢复为绝对手腕 `xyz + rotation-6D`。

## 10. 入口

训练：

```bash
cd /scratch1/home/zhicao/dreamzero
bash scripts/train/trex_track_force_training_wan22.sh
```

训练脚本直接使用已经构建完成的 `data/trex_small`，启动前只执行 schema
完整性校验，不会再次运行 SAM2/CoTracker。

日志、checkpoint 与原 T-Rex 训练共用同一套 Trainer 回调。每隔
`wandb_video_reconstruction_steps` 个 global step 会分别保存两个监控视频：
不带标记的原始重建视频位于
`OUTPUT_DIR/eval_videos/train_step_XXXXXX.mp4`，20Hz 预测 track motion
叠加到 10Hz head/left-wrist/right-wrist 三面板后的版本位于
`OUTPUT_DIR/eval_track_videos/train_step_XXXXXX.mp4`。原始预测与目标
track 同时写入 `OUTPUT_DIR/eval_tracks/train_step_XXXXXX.npz`。两个视频
分别记录到 W&B `eval/predicted_video` 和 `eval/predicted_track_video`；
`dynamics/action/track/force/VQ/commitment` loss 和触觉 codebook
perplexity/active-code 指标同时写入 W&B 与 `OUTPUT_DIR/loss_log.jsonl`。
两个视频的首帧都是精确 GT conditioning frame；track target/prediction 的第 0
步在模型链路内就是当前 GT anchor。可视化不再对整段预测做事后平移。

训练可视化默认使用
`wandb_video_reconstruction_inference_steps=1`：先运行一次 coarse WAN
Euler update，再在 `tau=0.4` 用一次边界预测恢复 clean video/track；不会直接
解码半噪声状态，并跳过不会改变这两个输出的 action-only 触觉精修。
这不会改变训练或正式在线推理的 6+4 两阶段 schedule；如需更高质量的监控视频，
可以单独提高该值，但耗时近似按 inference steps 线性增长。

离线 NPZ 推理：

```bash
python scripts/eval/trex_track_force_inference.py \
  --checkpoint checkpoints/trex_track_force_wan22_lora/checkpoint-8000 \
  --dataset-root data/trex_small \
  --input sample_input.npz \
  --output prediction.npz
```

核心实现位于 `groot/vla/model/trex_track_force/`；独立 Hydra 配置为
`model=trex_track_force/vla` 和 `data=dreamzero/trex_track_force_wan22`。