DiT4DiT — Video DiT hidden feature 条件化 Action DiT

背景：VLA 为什么需要视频动态先验

论文的出发点很明确：VLA 已经能用视觉和语言条件生成动作，但其 backbone 多数来自静态图文预训练。静态图文数据能提供物体、语义、关系和语言能力，却不直接提供“接触后会怎样”“物体如何随动作演化”“多阶段任务状态如何转换”等物理动态。

DiT4DiT 的核心假设是：通用视频生成模型为了预测未来帧，已经在互联网视频中学到大量时空结构和隐式物理规律。与其把这些能力用作外部可视化，不如把它们变成策略内部的动作条件。

意义：视频生成不是为了“看未来电影”

DiT4DiT 的贡献不是简单把未来视频生成出来，再让策略照着视频做。它明确避免依赖 fully reconstructed future frames，而是截取 Video DiT 去噪过程中的 hidden states 作为动作条件。这个差异非常关键：

DiT4DiT 基于哪个模型改进？

严格说，DiT4DiT 不是从零设计一个全新 transformer，而是把两个已经成熟的方向拼成一个 Video-Action Model：Video 端基于 Cosmos-Predict2.5-2B，Action 端基于 GR00T-N1 风格的 flow-matching Action DiT。它真正的创新不是“换了一个 backbone”，而是把 Video DiT 的中间去噪 hidden states 变成 Action DiT 的条件，并通过 video loss + action loss 联合训练。

模型结构总图：两个 DiT，但职责不同

输入:
  o_t: 当前/历史 ego 图像帧
  l:   语言指令
  s_t: 本体状态，可选
  a*:  expert action chunk，仅训练时使用

Video branch:
  frame_0 ──► Cosmos video processor / VAE ──► condition latent z_t^0
  frame_1..T ──► VAE ──► future latent supervision z_{t+1}^0
  instruction ──► tokenizer + text_encoder ──► prompt embeddings
  noisy future latent z^{τv} + condition latent + prompt
      └─► Cosmos Video DiT transformer blocks
             └─► hook block 17 (0-based) / paper layer 18
                    └─► hidden: [B, S=T×H×W, 2048]

Action branch:
  action noise / noised action chunk ──► ActionEncoder ─┐
  state s_t ──► StateEncoder MLP，可选 ────────────────┤
                                                       ├─► state/action tokens: [B, H(+state), 768]
  Video hidden [B,S,2048] ──────────────────────────────┘
      └─► 16-layer Action DiT
             - odd blocks: self-attention when interleave_self_attention=true
             - even blocks: cross-attention to Video hidden
             - timestep injected by AdaLN / timestep encoder
      └─► output projection 2560 ─► MLP action decoder ─► velocity / action_dim

训练目标:
  L_total = L_action_flow_matching + λ L_video_flow_matching

这张结构图的核心是：Cosmos Video DiT 不负责直接输出机器人动作，Action DiT 也不直接看原始图像。二者通过 Video hidden states 连接。这个连接点就是 DiT4DiT 和普通 VLA / 普通 video world model 的分界线。

论文结构 vs 本地代码结构

关键维度与配置

因此，如果要复现或迁移，不能只看论文 Table 4。LIBERO、RoboCasa、真机 G1 的 action/state 维度和 horizon 都会随本体改变；稳定不变的是 Video hidden dim=2048、hook 第 18 层、Action DiT cross-attend video hidden、动作 flow matching 4-step sampling 这些核心机制。

到底改进了什么？

1. 把 Cosmos-Predict2.5 从“视频生成器”改成“机器人动态特征提取器”。代码通过 forward hook 截取 transformer block hidden，不要求完整解码未来视频。
2. 把 GR00T-style Action DiT 的条件从静态 VLM token 改成 Video DiT hidden token。这让动作模型看到的是任务条件化的未来动态表征，而不是单帧图文语义。
3. 引入 video/action dual flow matching。Video branch 学未来 latent velocity field，Action branch 学动作 velocity field。
4. 引入 tri-timestep 解耦。视频去噪时间、hidden extraction 时间、动作去噪时间不强行绑在一起。
5. 通过 joint training 让 Video hidden 更适合控制。消融里 joint training 的任务阶段聚类更清楚，说明 action loss 会反向塑形 video hidden。

整体架构：Dual-DiT Video-Action Model

Video DiT 负责建模未来视觉动态；Action DiT 不是读取最终解码视频，而是通过 cross-attention 读取 Video DiT 某一层、某一去噪时刻的 hidden features。论文默认 hidden feature dim 为 2048，extract layer 为第 18 层；Action DiT 为 DiT-B 配置，inner dim 为 768，16 层，输出投影/decoder hidden size 为 2560，主设置动作 horizon 为 16。

数学形式

基础 flow matching 路径：

Video-Action 联合分布写法：

Video hidden extraction：

总损失：

为什么 hidden feature 比最终帧更适合动作

训练：dual flow matching + tri-timestep

DiT4DiT 训练时同时优化视频动态与动作生成。关键不是“两个 loss 简单相加”，而是把三个时间步拆开：

推理：不需要完整视频采样

控制时的关键路径很短：

这就是论文能在真机达到 6Hz 的原因：它不是每个控制周期都生成完整未来视频，而是把视频生成模型当作“动态表征提取器”。

模型与训练配置

评测协议矩阵：作者到底测了什么

LIBERO：从 scratch 达到 98.6%

Long suite 是最有解释力的：DiT4DiT 97.6%，明显高于 Qwen3DiT 93.6 和 GR00T-N1.5 90.0。这个结果支持“视频动态 hidden 对长时状态转移有帮助”。

RoboCasa-GR1：24 个桌面任务，平均 50.8%

Unitree G1 真机：视频动态表征在长时任务上更稳

真机平台为 Unitree G1 humanoid，双 7-DoF 手臂 + ALOHA2 grippers，头部 Intel RealSense D435i ego RGB 640×480，PICO VR + XRoboToolkit 采集。七个任务每个 200 条 demo、20 次 rollout。

Qwen3DiT 在真机几乎崩溃，论文称所有任务不超过 10%，在 Drawer / Arrange / Box 上为 0%。这说明“静态 VLM + 同一动作头”不足以替代视频动态 backbone。

OOD / generalization

最关键消融 1：第 18 层最好

论文在 RoboCasa-GR1 的五个任务上比较 Video DiT 不同层的 hidden feature。结论是：早层 2–8 表现差，末层 24–28 也明显下降，默认第 18 层最佳。

最关键消融 2：hidden extraction 只做 1 step 最好

论文比较用于抽 hidden feature 的 denoising steps。结果是 1 step 最好，更多 step 单调下降。作者解释为：过多迭代会让 hidden states 过度承诺到像素级未来细节，反而损害动作需要的抽象动态先验。

最关键消融 3：joint training 改善时序结构

论文用 t-SNE / silhouette score 分析 hidden features。decoupled training 的 silhouette score 为 0.09，joint training 提升到 0.17，并表现出更顺滑的 Early → Middle → Late 任务阶段流动。

这说明 DiT4DiT 不是只把一个冻结视频模型当特征提取器，而是通过 action loss 反向塑形 Video DiT，使其 hidden feature 更适合动作。

效率

本地代码路径

代码印证 1：forward hook 抽 Video DiT 中间层

# Cosmos25.py
pipe = Cosmos2_5_PredictBasePipeline.from_pretrained(...)
self.text_encoder = pipe.text_encoder
self.transformer = pipe.transformer
self.vae = pipe.vae
self.scheduler = pipe.scheduler

target_layer = blocks[self.extract_layer]
target_layer.register_forward_hook(hook_fn)

def hook_fn(module, inp, out):
    if torch.is_tensor(out):
        self._cached_hidden.append(out.detach())
    elif isinstance(out, (tuple, list)) and torch.is_tensor(out[0]):
        self._cached_hidden.append(out[0].detach())

这说明 Video 端真实基座是 Cosmos-Predict2.5 / Cosmos2.5 pipeline。配置中 extract_layer: 17 是 0-based，对应论文表格里的第 18 层。这一点和论文消融完全对齐。

代码印证 2：只截第一次 transformer forward 的 hidden

# Cosmos25.py
self._cached_hidden.clear()
self._capture_hidden_enabled = True
...
for i, t in enumerate(timesteps):
    model_out = self.transformer(...)[0]
    if i == 0 and hidden_first is None and len(self._cached_hidden) > 0:
        hidden_first = self._cached_hidden[-1]
        self._capture_hidden_enabled = False
...
hidden = hidden_first if hidden_first is not None else self._cached_hidden[-1]

这就是“推理时不完整生成未来视频”的实现落点：默认 future_num_inference_steps: 1，只要一次 Video DiT forward 的中间 hidden。

代码印证 3：Action DiT 读取 video hidden

# DiT4DiT.py
backbone_outputs = self.backbone_interface(...)
last_hidden = backbone_outputs.hidden_states[-1]
action_loss = self.action_model(last_hidden_repeated, actions_target_repeated, action_mask, state_repeated)

# ActionDiT.py
model_output = self.model(
    hidden_states=sa_embs,
    encoder_hidden_states=vl_embs,
    timestep=t_discretized,
)

encoder_hidden_states=vl_embs 就是 Video DiT 抽出来的 hidden token。Action DiT 内部交替 self-attention / cross-attention，使用 cross_attention_dim: 2048 对齐 video hidden dim。

代码印证 3.5：Action DiT 的 token 结构

# ActionDiT.py
state_features = self.state_encoder(state) if state is not None else None
action_features = self.action_encoder(noisy_trajectory, t_discretized)
if self.config.add_pos_embed:
    action_features = action_features + pos_embs

sa_embs = torch.cat((state_features, action_features), dim=1) \
    if state_features is not None else action_features

model_output = self.model(
    hidden_states=sa_embs,
    encoder_hidden_states=vl_embs,
    timestep=t_discretized,
)

这段代码把 Action DiT 的主输入明确成 state tokens + action tokens，Video hidden 作为 cross-attention memory。论文方法段提到 learnable “future tokens”，但本地 release 代码已核对主路径中没有看到显式 future token parameter。因此网页中把它标为论文描述与代码主路径的差异，而不是强行认定代码已经实现。

代码印证 4：动作 flow matching

# ActionDiT.py
noise = torch.randn(actions.shape, device=actions.device, dtype=actions.dtype)
t = self.sample_time(actions.shape[0], device=actions.device, dtype=actions.dtype)
noisy_trajectory = (1 - t) * actions + t * noise
velocity = noise - actions
...
loss = ((pred_actions - velocity) ** 2) * action_mask
loss = loss.sum() / action_mask.sum()

这与论文中的 action flow matching 对齐。推理时 num_inference_timesteps: 4，从随机动作噪声出发，按 flow ODE 积分得到 action chunk。

三篇工作的共同趋势

与 ImageWAM 的关系

ImageWAM 问的是：世界动作模型是否真的需要视频生成，还是只要图像编辑？它用图像编辑模型从 current → target 单帧变换，并在推理时只取编辑去噪 KV cache，不解码目标图。

两者不是互相否定。更准确的判断是：如果任务需要强时序动态，DiT4DiT 的视频 hidden 更有潜力；如果任务主要是对象状态/空间目标变化，ImageWAM 的编辑 cache 可能更划算。

与 DynaGuide 的关系

DynaGuide 不训练大型 VLA，而是训练一个轻量 dynamics model，预测 (obs, action_chunk) → future latent，再在 diffusion policy 的采样过程中对 action 求梯度，把动作往目标 latent 推近或远离负目标。

对 VLMGuide 的直接启发

1. 不要把目标设成“生成未来图/视频”。更应把目标设成“提取可用于动作选择的任务条件化中间表征”。
2. VLMGuide 的 guide 信号应落在 action selection。只要中间表征不能改变动作选择，就只是可视化或诊断。
3. 增加 action-conditioned reachability。ImageWAM/DiT4DiT 的 hidden/cache 可以表示“想象的目标变化”，但仍要判断候选动作是否能达到该变化。
4. 做 phase-aware guidance。reach / grasp / transport / place / recovery 不应共享同一种 guide；不同阶段需要不同 latent、不同权重。
5. 优先从轻量闭环开始。完整 DiT4DiT 需要大资源；VLMGuide 可以先用冻结 VLA + 中间表征 scorer/reranker 证明 policy-level 收益。

可落地方案：ImageWAM + DynaGuide + VLA repair

这个方案的论文叙事可以是：现有 VLA 会在 OOD、阶段错乱或视觉捷径下生成“看似合理但当前不该做”的动作；VLMGuide 不直接替换 policy，而是在运行时用生成模型中间表征 + action-conditioned dynamics 做动作一致性验证。

和触觉/世界模型课题的连接

DiT4DiT 的局限部分明确提到当前只有单个 ego camera，未来可加入 wrist camera 或 tactile feedback。对你们的触觉方向，这给出一个自然扩展：

选题建议

局限

• 单 ego camera。真机部署只用头部 ego RGB，相机被手臂或物体遮挡时会破坏时序连续性。
• 没有触觉/力反馈。论文没有在接触丰富、力敏感任务上系统验证。
• 训练资源高。论文配置为 32 GPUs、100k steps，完整复现门槛明显高于 DynaGuide / 普通 DP。
• 频率低于 VLA baseline。6Hz 可用但慢于 Qwen3DiT 9Hz 和 GR00T-N1.5 13Hz。
• 中间 hidden 可解释性弱。不能像最终视频那样直接肉眼看，需要消融、聚类或 probing 验证。

复现与代码使用注意

最终评价

DiT4DiT 是一篇很强的世界模型/VLA 交叉论文。它最值得学习的不是“视频模型很强”这个泛泛结论，而是一个更具体的技术判断：

自测题