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. 总体架构

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。时间序列输入特征为：

[x, y, visibility, Δx, Δy]

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

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

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：

[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 顺序为：

[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

使用线性插值：

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. 入口

训练：

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 推理：

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。