# 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`。