DiT4DiT: Joint Video & Action Diffusion

• The paper introduces DiT4DiT, an end-to-end framework that couples video and action diffusion transformers to jointly model video dynamics and control actions in robotic tasks.
• It leverages intermediate denoising features from a generative video process as temporally grounded conditions to predict accurate robot actions.
• Empirical results demonstrate up to 7× faster convergence and 10× improved data efficiency compared to conventional vision-language-action models.

DiT4DiT refers to an end-to-end framework that couples a video Diffusion Transformer and an action Diffusion Transformer to model both video dynamics and control actions in generalizable robot manipulation tasks. This architecture leverages intermediate denoising features—extracted from the video generative process—as temporally grounded conditions for action prediction. DiT4DiT directly addresses the representational bottleneck of vision-language-action models pre-trained only on static image–text corpora by enabling joint, efficient learning of physical dynamics and policy control from video and action demonstrations. The approach establishes state-of-the-art performance on both simulated and real-world robotic benchmarks, delivering order-of-magnitude improvements in data efficiency and convergence speed compared to prior VLA and policy learning methods (Ma et al., 11 Mar 2026).

1. Motivation for Joint Video-Action Diffusion Transformers

Conventional VLA models primarily inherit rich image-language representations from large static corpora, forcing downstream policy learning to account for all physical dynamics solely via limited labeled episodes. Semantic auxiliary objectives (e.g., object-centric or future-feature-alignment losses) often fail to internalize the temporal and physics structure critical for high-precision, long-horizon robotic control. In contrast, generative video models—particularly diffusion-based architectures—are trained to generate coherent frame sequences, capturing spatiotemporal continuity, causality, and implicit physics priors.

Empirical analysis demonstrates that casting world modeling as a generative video prediction task serves as an unsupervised proxy for faster, more sample-efficient policy learning. Specifically, DiT4DiT accelerates convergence by up to 7× and achieves at least 10× greater sample efficiency compared to semantic-centric grounding approaches (Ma et al., 11 Mar 2026).

2. Cascaded Architecture and Feature Coupling

DiT4DiT composes two flow-matching diffusion transformers in a cascaded structure:

• Video Diffusion Transformer (Video DiT): Encodes the observation sequence and goal (via a spatial-temporal VAE) as latents $z_t^0$. It performs guided denoising to produce video features $z_{t+1}^0$.
• Intermediate Feature Extraction: At a fixed denoising step $\tau_f$, an internal activation $h_t^{\tau_f}$ is extracted from a specific transformer layer (e.g., layer 18).
• Action Diffusion Transformer (Action DiT): Receives the intermediate video features $h_t^{\tau_f}$, the robot state, and a noisy action vector. It performs denoising towards the correct action trajectory.

This setup allows the action module to be temporally conditioned on the latent video dynamics, rather than on fully reconstructed frames, serving as richer, more abstract and temporally grounded policy input.

3. Dual Flow-Matching Training Objective

The overall training procedure is orchestrated by a dual flow-matching objective, with decoupled diffusion timesteps and noise scales for the video and action components:

• Video Diffusion: At sampled $\tau_v$, the Video DiT regresses the velocity field between noisy and clean video latents.
• Feature Extraction: Hidden state $h_t^{\tau_f}$ is extracted at a fixed intermediary denoising step.
• Action Diffusion: Action DiT, cross-attending to $h_t^{\tau_f}$, regresses the action velocity field at sampled $\tau_a$.

The total loss per batch is given by

$$L_\text{total} = \mathbb{E}_{\tau_v, z}\left[\| v_\theta^\text{video}(z_{t+1}^{\tau_v}, \tau_v | z_t^0, l) - (z - z_{t+1}^0) \|^2 \right] + \lambda\,\mathbb{E}_{\tau_a, \epsilon}\left[ \| v_\phi^\text{action}(a_t^{\tau_a}, \tau_a | h_t^{\tau_f}, s) - (\epsilon - a_t^0)\|^2 \right].$$

Importantly, the video and policy modules are trained end-to-end; extracted video features act as a continuous, differentiable policy anchor.

4. Training Regime, Implementation Details, and Hyperparameters

Training is conducted across several robotic benchmarks:

• LIBERO: 4 suites × 500 demonstrations per suite (7-DoF Panda arm).
• RoboCasa-GR1: 24 tasks × 1,000 trajectories (29-DoF dual-arm humanoid).
• Unitree G1: 7 real-world tasks × 200 episodes (16-DoF; VR teleoperation).

Key architecture and optimization parameters:

• Video DiT initialized from Cosmos-Predict 2.5-2B (hidden size 2048).
• Action DiT: 16 layers, hidden size 2560, cross-attention 2048, horizon 16, diffusion steps 4.
• Joint batch size 256, trained for 100k steps with AdamW (distinct learning rates per module), and cosine LR decay.
• Noise schedules draw $\tau_v$ uniformly, action noise via $1-\sigma$ for $\sigma \sim \text{Beta}(1.5, 1.0)$.

5. Empirical Results and Data Efficiency

Performance is consistently state-of-the-art in both simulation and real-robot tasks:

LIBERO Suite

Method	Spatial	Object	Goal	Long	Avg
DiT4DiT (scratch)	98.4	99.6	98.6	97.6	98.6

RoboCasa-GR1 (average across 24 tasks):

• DiT4DiT: 50.8% vs. Qwen3DiT: 36.2%, GR00T-N1.5: 41.8%.

Unitree G1 (real-world, selected tasks):

• Arrange Flower: 75%, Stack Cups: 60%, Drawer Interaction: 90%, Box Packing: 50%. Baselines often fall below 25% on these tasks.

DiT4DiT demonstrates pronounced zero-shot generalization—for instance, achieving a 54.5% success rate in Bottle→Close with unseen objects, versus 32.0% for Qwen3DiT.

Sample efficiency and convergence plots indicate that DiT4DiT reaches ≥80% success approximately 7× faster and achieves more than 10× performance in data-limited regimes compared to strong semantic and policy-learning baselines.

6. Ablation Studies and Analysis

• Feature extraction layer: Peak control transfer occurs when extracting features from layer 18; earlier layers yield <30% success, and later layers degrade as representations specialize to pixel-space.
• Conditioning steps: Conditioning on a single denoising step is optimal. More steps—corresponding to more complete reconstructions—cause monotonic performance decline, suggesting abstraction rather than low-level fidelity is most beneficial for policy grounding.
• Joint vs. decoupled training: Joint end-to-end optimization yields much smoother temporal feature trajectories (2× silhouette score improvement in t-SNE), while decoupled training produces fragmented scene representations.

These findings underline the value of intermediate denoising features and the necessity of joint optimization for generalizable robot policies.

7. Limitations and Prospective Directions

Current limitations include reliance on a single egocentric camera (risking occlusion), marginally slower deployment rates (6 Hz) compared to non-generative VLA baselines (9–13 Hz), and a still-modest scale of pretraining data relative to some recent generalist models. Future research may incorporate multi-modal sensory fusion (e.g., wrist cameras, tactile arrays), embodiment diversity, hierarchical LLM-based planning for compositional tasks, and multi-agent coordination via action DiT cross-attending to multiple egocentric video streams.

A plausible implication is that generative modeling of pixel-level dynamics—rather than semantic alignment alone—acts as a potent scaling proxy for complex policy learning. DiT4DiT showcases that abstracted video denoising features, extracted from a generative world model, can be productively coupled to robot action inference, jointly modeling "what will happen" and "what to do" in a single, end-to-end framework (Ma et al., 11 Mar 2026).