DiT4DiT: Jointly Modeling Video Dynamics and Actions for Generalizable Robot Control

In plain words

Robots that act from camera images usually start from models trained on still pictures and text. Those models know what objects are, but not how a scene changes as things move. This paper starts from a video generator instead, a model trained to predict how a scene evolves, and uses it as the base of a robot policy.

The policy does not make the video model draw the next frames. It takes the video model's internal activations partway through its denoising process, at one fixed noise level, and passes them to a second, smaller network. That second network outputs the robot's next movements. Both networks are trained together, each with its own noise schedule, so the video part learns the scene's dynamics and the action part learns to read them.

The authors report that this gives high success in simulation and on a humanoid robot, with less robot data than the baselines they compare against, and faster training. These results come from their own experiments.

TL;DR

DiT4DiT joins a video diffusion transformer (Cosmos-Predict2.5-2B) to an action diffusion transformer and trains both with flow-matching losses. The action head reads intermediate features of one video block at a fixed timestep, so control does not need generated frames. The paper reports 98.6% average success on LIBERO and 50.8% over 24 RoboCasa-GR1 tasks, both trained from scratch and author-reported. The released code is MIT-licensed, but it does not pass the action loss gradient back into the video model (see Paper vs code).

Problem

Vision-language-action (VLA) models usually start from backbones pretrained on static image-text data. Physical dynamics are then learned only from the much smaller action dataset §1. Auxiliary objectives such as grounding and FLARE-style latent alignment are, according to the authors, sample-inefficient in their own tests §1 §3. Video-based policies that first generate future frames and then invert them are multi-stage, so control is indirect §1. The paper asks whether video generation can itself be the training objective, and how the representation learned by a video model should be taken out and coupled with action generation §1.

Method

The model is a Video-Action Model (VAM) with two diffusion transformers (DiTs), trained jointly §4 fig. 2.

Video DiT. The initialisation is Cosmos-Predict2.5-2B. It has a causal video VAE and a video DiT with flow-prediction parameterisation. Language conditions it through multi-layer embeddings from Cosmos-Reason1 §4.2. Instead of using the final denoised video, a forward hook takes the hidden activations of one block (or, the text says, an average over all blocks) at flow timestep τ_f §4.2 eq. 9.

Action DiT. It is adapted from GR00T-N1 §4.2. Each block uses AdaLN for the diffusion timestep and cross-attention to the video features h. The input sequence is the state embedding followed by the noisy action tokens. The text also mentions "learnable future tokens" in the sequence §4.2, but the released action model does not have them (see Paper vs code). The network outputs a velocity over the action chunk §4.2.

Tri-timestep scheme. Three independent timesteps are used §4.3 fig. 3:

• video: τ_v is drawn uniformly from [0, 1] §4.3;
• features: τ_f is a fixed, deterministic timestep for the hook §4.3;
• action: τ_a = 1 − σ with σ ~ Beta(α, β) §4.3, with α = 1.5 and β = 1.0 app. A.1.

Training. The VAE and the text encoder are described as frozen, and only the two DiTs are updated §4.3. Each step encodes the observation and the future frames, computes the video flow-matching loss, extracts features, computes the masked action loss, and sums them alg. 1.

Inference. Action-only inference runs one video-backbone pass with fresh noise at τ_f, then integrates the action velocity with Euler steps §4.4 alg. 2. Video generation for visualisation is a separate path with N_v Euler steps §4.4 fig. 10. The paper's ablation found one video step best for features §5.4.

Equations (numbering follows the LaTeX source order, see Pitfalls):

x_{\tau} = (1-\tau)\, x_0 + \tau\, z,\quad \tau \in [0,1] eq. 1

v^*(x_{\tau}, \tau) = z - x_0 eq. 2

\mathcal{L}_{\text{FM}} = \mathbb{E}\left[\left\| v_\theta(x_{\tau}, \tau) - (z - x_0) \right\|^2\right], \quad \tau \sim \mathcal{U}[0,1] eq. 3

x_{\tau-\Delta\tau} = x_{\tau} - \Delta\tau\, v_\theta(x_{\tau}, \tau) eq. 5

h_t^{\tau_f} = \mathcal{H}\big[v^{\text{video}}_\theta\big]\big(z_{t+1}^{\tau_f}, \tau_f \mid z_t^0, l\big) eq. 9

Joint loss, with λ a scalar that the paper does not give a value for:

\mathcal{L}^{\text{total}}_t = \mathcal{L}_{\text{action}} + \lambda\, \mathcal{L}_{\text{video}}

\mathcal{L}_{\text{action}} = \mathbb{E}\left[\left\| v^{\text{action}}_\phi(a_t^{\tau_a}, \tau_a \mid h_t^{\tau_f}, s) - (\epsilon - a_t^0) \right\|^2\right]

\mathcal{L}_{\text{video}} = \mathbb{E}\left[\left\| v^{\text{video}}_\theta(z_{t+1}^{\tau_v}, \tau_v \mid z_t^0, l) - (z - z_{t+1}^0) \right\|^2\right] eq. 10

Training step, as the paper's Algorithm 1 writes it:

# Joint training step  [alg. 1]
# Inputs:  obs (B, Tc, 3, H, W), next_frames (B, Tv, 3, H, W),
#          actions (B, 16, 32), state (B, 64), language, action_mask
# Output:  total_loss (scalar)

def training_step(obs, next_frames, actions, state, language, action_mask):
    # Encode the observation and the future frames with the VAE  [alg. 1]
    obs_latent = vae_encode(obs)
    future_latent = vae_encode(next_frames)

    # Video branch: uniform tau_v, interpolate, regress the velocity
    # [§4.3] [eq. 1] [eq. 2] [eq. 3]
    tau_video = uniform(0, 1)  # (B,)
    noise = normal_like(future_latent)
    noisy_future = (1 - tau_video) * future_latent + tau_video * noise
    video_velocity = video_dit(noisy_future, tau_video, obs_latent, language)
    video_loss = mean_squared_error(video_velocity, noise - future_latent)

    # Hidden features: fresh noise at a fixed feature timestep tau_f
    # Layer 18 in the paper's numbering [app. A.1] [§4.2] [eq. 9]
    feature_noise = normal_like(future_latent)
    hidden = hooked_block(video_dit, layer=18)(
        feature_noise, tau_f, obs_latent, language)

    # Action branch: tau_a = 1 - sigma, sigma ~ Beta(1.5, 1.0)  [§4.3]
    sigma = beta_sample(1.5, 1.0)  # (B,)
    tau_action = 1 - sigma
    action_noise = normal_like(actions)
    noisy_actions = (1 - tau_action) * actions + tau_action * action_noise
    action_velocity = action_dit(noisy_actions, tau_action, hidden, state)

    # Masked velocity error, normalised by the number of unmasked entries
    # [alg. 1] [eq. 10]
    action_loss = masked_mean_squared_error(
        action_velocity, action_noise - actions, action_mask)

    # Joint objective; lambda is not stated in the paper  [eq. 10]
    total_loss = action_loss + lam * video_loss
    return total_loss

Action-only inference, as the paper's Algorithm 2 writes it:

# Action inference  [alg. 2] [§4.4]
# Inputs:  obs (B, Tc, 3, H, W), state (B, 64), language
# Output:  actions (B, 16, 32)

def predict_actions(obs, state, language):
    # One video-backbone pass with fresh noise at tau_f, no sampling loop
    # [§4.4] [alg. 2]
    obs_latent = vae_encode(obs)
    feature_noise = normal(future_latent_shape)  # shape not stated
    hidden = hooked_block(video_dit, layer=18)(
        feature_noise, tau_f, obs_latent, language)

    # Start the action chunk from Gaussian noise  [alg. 2]
    actions = normal((B, 16, 32))
    num_steps = 4  # N_a, "Num Inference Timesteps"  [app. A.1]
    for step in range(num_steps):
        # Time runs from 1 (noise) toward 0 (clean)  [alg. 2] [§4.4]
        tau_action = 1 - step / num_steps
        velocity = action_dit(actions, tau_action, hidden, state)
        # Euler step  [eq. 5]
        actions = actions - (1 / num_steps) * velocity
    return actions

Hyperparameters

Values are from the paper's appendix unless marked as code-only. Shapes and units are as stated in the paper.

Results

Author-reported. Setup: LIBERO (Franka, four suites, §5.1), RoboCasa-GR1 tabletop with 24 tasks and 50 rollouts per task, horizon 720 §5.1, and a Unitree G1 with 7 real tasks and 20 rollouts each §5.1. The baselines in the simulation tables are other policies, with official open-source weights where the paper says so §5.1.

LIBERO (success rate, %) tab. 1. DiT4DiT was trained from scratch on action data from this benchmark only.

RoboCasa-GR1 (24 tasks, success rate, %) tab. 2. GR00T-N1.5 and GR00T-N1.6 are fine-tuned from their pretrained weights; Qwen3DiT and DiT4DiT are trained from scratch for the same number of steps tab. 2.

Text results §5.2: DiT4DiT is best on 16 of 24 tasks. It is 9.0 points above GR00T-N1.5 and 10.0 above GR00T-N1.6 on the average, and 14.6 points above Qwen3DiT.

Real Unitree G1 (7 tasks, 200 demos each, 20 rollouts per task) §5.2 fig. 5. Only some values are given in the text. Qwen3DiT never exceeds 10% on any task, and it yields a flat 0% on three of them: Drawer Interaction, Arrange Flower and Box Packing. DiT4DiT's pretraining uses 241,450 GR1 episodes, which the paper says is about 15% of the data scale of GR00T-N1.5 §5.1.

Zero-shot generalisation §5.3 fig. 6 fig. 7:

• Simulation (training on three bottle tasks, testing on Can, Cup, Milk and Wine): ToDrawerClose 54.5% for DiT4DiT vs 32.0% for Qwen3DiT (32.0% is implied by the stated 22.5-point gap); ToCabinetClose 34.0% vs 24.5%; ToMicrowaveClose 30.5% vs 17.0%.
• Real robot: Arrange Flower (Category) 70% for DiT4DiT, 0% for Qwen3DiT, 10% for GR00T-N1.5; Stack Cup (Number) 50%. Box Packing (Object) and Stack Cup (Category) values are not given in the text.

Ablations §5.4 fig. 8: on five RoboCasa tasks.

• Extraction layer: layers 2 to 8 perform poorly, layer 18 is best, and layers 24 to 28 collapse. Averaging all layers is close to but below layer 18.
• Denoising steps for features: a single step gives the highest success, and success falls monotonically as steps rise from 1 to 32.
• Joint versus decoupled training: silhouette score on the t-SNE of hidden features rises from 0.09 (decoupled) to 0.17 (joint).

Efficiency tab. 3: trainable parameters 2.2B (DiT4DiT), 2.3B (Qwen3DiT), 2.7B (GR00T-N1.5). Real-robot control frequency 6 Hz, 9 Hz and 13 Hz, respectively. The table caption says this was measured on one NVIDIA A100.

Paper vs code

Official code: Mondo-Robotics/DiT4DiT, MIT-licensed LICENSE. The paper's project page links it [dit4dit.github.io]. The organisation is the affiliation of several authors (Mondo Robotics) [paper: author block]. The commit checked is 66a6f3a12e2c8740157c1e478795952c040d31dd.

Each difference below has both sides cited:

• Gradient coupling (major). The paper says the joint objective optimises both DiTs, and Algorithm 1 updates θ and φ from the total loss §4.3 alg. 1. In the code, the hook returns out.detach() Cosmos25.py, and the extractor detaches the features before returning them Cosmos25.py. The action loss therefore does not reach the video DiT. The video DiT is trained only by its own flow-matching loss, and the two are coupled only through the shared training step.
• Feature timestep τ_f. The paper says τ_f is fixed and deterministic §4.3 §4.4, but Algorithm 1 samples τ_f uniformly over a grid of T steps for each training step alg. 1. The code does not sample τ_f. It uses the first timestep of a schedule with steps_for_hidden steps, which is 1 in the configs Cosmos25.py dit4dit_robocasa_gr1.yaml, and captures the hidden state from that first call Cosmos25.py. The noise level at that step is not given in the configs.
• Video timestep distribution. The paper says τ_v ~ U[0, 1] §4.3 alg. 1. The configs set time_distribution: logit_normal and high_sigma_ratio: 0.05, with 5% of samples drawn from [0.98, 1) dit4dit_robocasa_gr1.yaml. The sampling code is in Cosmos25.py.
• Action timestep. The paper says τ_a = 1 − σ with σ ~ Beta(1.5, 1.0) §4.3. The code sets the noise weight directly to the Beta sample divided by 0.999 ActionDiT.py ActionDiT.py. The mean noise weight is therefore about 0.6 in the code and 0.4 in the paper's formula.
• Learnable future tokens. The paper describes them in the action input §4.2. The action model concatenates state and action embeddings only ActionDiT.py. The config field num_target_vision_tokens is defined ActionDiT.py but I found no use in the forward path.
• Loss weight λ. The paper calls it a scalar and gives no value §4.3. The code uses trainer.loss_scale.future_video, defaulting to 1.0, and no config sets it train.py.
• Repeated diffusion steps. The appendix gives 4 app. A.1. The configs set 8 in framework.action_model dit4dit_robocasa_gr1.yaml, but the forward pass reads the trainer's value, which defaults to 4 DiT4DiT.py dit4dit_robocasa_gr1.yaml. The 8 does not appear in that forward path.
• Per-device batch size. The appendix gives 8 per device on 32 GPUs app. A.1. The configs give 16 dit4dit_robocasa_gr1.yaml.
• Weight decay. The appendix gives 1 × 10⁻⁸ app. A.1. The optimizer block has 1e-8 dit4dit_robocasa_gr1.yaml, but the trainer block has weight_decay: 0.0 dit4dit_robocasa_gr1.yaml. I did not trace which one reaches the optimizer.
• Frozen modules. The paper says the VAE and text encoder are frozen §4.3. All configs set freeze_modules: '' dit4dit_robocasa_gr1.yaml, and I found no requires_grad_(False) in the video backbone. The VAE is used in no_grad only for ground-truth latents Cosmos25.py, and prompt encoding runs under enable_grad Cosmos25.py. Whether these modules train depends on how they are loaded, and I did not verify it.
• Hardware for the frequency numbers. The efficiency table says 6 Hz on one A100 tab. 3. The appendix says real-time inference ran on an RTX 4090 app. A.3.
• Action execution at test time (code only). The RoboCasa interface returns the first n_action_steps=2 of the 16-step chunk model2robocasa_interface.py model2robocasa_interface.py. The LIBERO interface enables action ensembling by default, with horizon 3 and alpha 0.1 model2libero_interface.py. The paper states neither.
• Action ODE indexing. Algorithm 2 uses Δτ_a in the loop but defines only Δτ alg. 2. The code uses dt = 1/N and a time grid of 1 − i/N ActionDiT.py.

Pitfalls

• The joint training gradient is one-way in the released code. An implementation that follows Algorithm 1 literally (gradients through h into the video DiT) differs from the released code. One that follows the code differs from the paper's stated objective. Pick one, and say which in your results.
• τ_f is not pinned down. The paper gives no value and two different procedures (fixed versus sampled per step). The ablation covers the number of denoising steps and the layer, not the timestep value. Choose τ_f, test it, and report it.
• Layer indexing. The paper says layer 18 app. A.1. The config uses index 17 on a 0-based list Cosmos25.py. The indexing convention is not stated in the paper.
• Ablation design. The ablations use five RoboCasa tasks §5.4, and the step ablation is described as a strict monotonic decline in success. Those are author-reported numbers from a small set of tasks.
• Not like-for-like comparisons. The simulation baselines GR00T-N1.5 and N1.6 are fine-tuned from pretrained weights, while DiT4DiT is trained from scratch tab. 2 §5.1. For real-robot tests, GR00T-N1.5 starts from official weights that the paper says use far more data, while DiT4DiT is pretrained on GR1 episodes §5.1.
• Data counts. LIBERO is described as having exactly 500 demonstrations per suite §5.1, but the appendix says the official dataset has 1,693 episodes app. A.2. The paper does not reconcile them.
• Equation numbers. Equations are numbered in LaTeX source order, and the source has one unlabeled display (eq. 4). Check the numbers against the compiled PDF before relying on them.
• Cost per control step. The video backbone adds per-step cost. The authors report 6 Hz, against 9 Hz and 13 Hz for the baselines, and they note that LLM features could be cached for fixed tasks §5.5.
• Action normalisation and state preprocessing. The paper does not state how actions and states are normalised, or how the 16-step chunk maps to the robot's control rate. The repository has a state/action transform module state_action.py; I did not read it, so take the exact scheme from there rather than from this card.
• Frame choices differ by config. The stride and frame count vary between the LIBERO config (9 indices, stride 2), the RoboCasa and real-robot configs (17 indices, stride 2), and the decoupled-WBC config (stride 4) dit4dit_libero.yaml dit4dit_g1_decoupled_wbc.yaml.

Lineage

• Builds on: Cosmos-Predict2.5-2B video model and its causal VAE §4.2; the GR00T-N1 action DiT with AdaLN §4.2; Cosmos-Reason1 text embeddings §4.2; flow matching §4.1; DiT (Peebles and Xie) §4.2.
• Compared against: Diffusion Policy, Dita, π0, UniVLA, π0.5, OpenVLA-OFT, CogVLA, GR00T-N1.5 and N1.6 tab. 1 tab. 2; grounding and FLARE-style latent alignment as proxy objectives §3; Qwen3DiT, a parameter-matched VLM-based baseline with the same action DiT §5.1.
• Superseded by: not stated.