Source-linked AI summary
WorldDiT: A Unified Diffusion Architecture for World and Action Modeling
Sen Wang, R. Gnana Praveen, Bidhan Roy, Marcos Villagra
TL;DR
Robot policies often rely on large pretrained vision-language backbones, making the roles of scale and architecture difficult to separate. WorldDiT unifies diffusion-based action generation with future visual prediction and, across four LIBERO suites, places on the reported parameter–success Pareto frontier without a large VLM action backbone.
Problem
Large pretrained action backbones make it difficult to isolate whether strong robot control comes from architectural design, pretrained representations, or their combination.
Method
WorldDiT uses a shared diffusion transformer to generate continuous action chunks while predicting normalized RGB patches from future camera frames during training.
Results
Across four LIBERO suites, WorldDiT lies on the reported parameter–success Pareto frontier and records a 94.9% mean success rate.
Takeaways & Limitations
The reported results support unified world-and-action modeling as a compact basis for future scaling studies.
Takeaways & Limitations
The evaluation is not fully held out, and incomplete public artifacts prevent reproducing every comparison under a shared protocol.
Abstract
from arXiv · showhide
Many recent robot policies pursue stronger control by using large pretrained vision-language models (VLMs) as the action backbone. We introduce WorldDiT, a unified diffusion transformer architecture that couples action generation with visual world modeling and achieves strong performance without a large pretrained VLM action backbone. During training, a single diffusion transformer generates continuous action chunks and predicts normalized RGB patch targets from future camera frames. Across four LIBERO simulation suites, WorldDiT lies on the reported Pareto frontier for total model parameters and mean success among methods reporting all four suites. These results provide a strong sub-billion-parameter baseline for future scaling studies.
1. Introduction
WorldDiT asks whether unified diffusion can provide strong robot control without conflating action design with a large pretrained backbone. It couples action generation with future RGB patch prediction during training, omits that auxiliary prediction at inference, and reports a sub-billion-parameter baseline for future scaling studies.
- Large pretrained action backbones make the contributions of model scale, architecture, and pretraining difficult to separate.The paper motivates testing a unified diffusion transformer without relying on a large pretrained action backbone.
- WorldDiT uses a shared diffusion transformer to generate continuous seven-step action chunks and predict normalized RGB patches from future camera frames.The backbone conditions on recent visual observations, robot state, and language; at inference it generates actions, executes the first three, and replans.
- Below one billion parameters, reported results characterize the tradeoff between parameter count and mean success across four LIBERO suites.One WorldDiT configuration is evaluated under the specified protocol across all 24 methods, providing a baseline for future scaling studies rather than evidence of scaling behavior.
- RGB patch prediction supplies training supervision but is absent at inference, forming a unified action-generation and future-world-modeling design.This auxiliary target is predicted from future primary-camera and wrist-camera frames during training.
2. Method
WorldDiT uses one shared diffusion transformer backbone to generate continuous robot action chunks while predicting normalized RGB patches from future camera frames as auxiliary supervision. It conditions on multimodal observation histories and trains with flow matching to jointly regress action and RGB-patch velocities.
- Unified architecture: WorldDiT uses one shared DiT backbone to model continuous robot actions and normalized RGB patches selected from future primary-camera and wrist-camera frames.This couples action generation with an auxiliary future normalized RGB patch objective without relying on a large vision-language model to generate action tokens autoregressively.
- Multimodal conditioning: Given language instructions, multi-view observations, and robot states, frozen visual and language encoders produce context tokens for the shared backbone.The backbone receives corrupted action and future RGB patch tokens, timestep embeddings, and learned register tokens, then predicts their flow velocities.
- Optimization: Flow matching trains velocities along a straight path from Gaussian noise to clean action or RGB-patch targets, with total loss weighted across both velocity losses.The coefficients w_action and w_rgb specify the corresponding loss weights.
- Training targets: WorldDiT predicts one H-step action chunk and one future world target of normalized RGB patches from primary-camera and wrist-camera frames at i+H.Only the final temporal slot contributes to the loss.
- Deployment: At deployment, WorldDiT starts action generation from Gaussian noise, integrates the learned action velocity field, executes a prefix of the predicted chunk, and replans.No future RGB patch targets or future action labels are provided during deployment.
3. Experiments
Experiments evaluate WorldDiT across four LIBERO suites using a unified action-and-RGB-patch training setup. It achieves a 94.9% mean success rate and lies on the reported parameter–success Pareto frontier, though its aggregate includes checkpoint-selection episodes.
- Experimental setup: WorldDiT uses libero_90 for pretraining and fine-tunes independently on libero_spatial, libero_object, libero_goal, and libero_10, reported as LIBERO Long.Training windows contain multi-view RGB observations, robot states, language instructions, action sequences, and future normalized RGB patch targets.
- Evaluation limitations: 94.9% should not be interpreted as an unbiased test estimate because the aggregate includes 300 episodes per suite used during staged checkpoint selection.The reported WorldDiT score aggregates 500 simulator episodes per suite.
- Results: 98.0% on Spatial, 97.0% on Object, 92.8% on Goal, and 91.8% on Long produce a 94.9% mean success rate.Each suite uses 500 simulator episodes; Long is the hardest suite.
- Parameter efficiency: 399.084 million total parameters and 135.107 million trainable parameters place WorldDiT on the reported Pareto frontier.Total parameters include frozen modules required at inference.
- Parameter efficiency: Among methods reporting means across all four suites, higher reported success requires more total parameters, while methods at or below WorldDiT’s count report lower success.This comparison uses the published results included in the paper.
4. Discussion
WorldDiT couples continuous action generation with future visual prediction in a single diffusion transformer while retaining an action-only deployment path. Its 399M-parameter system lies on the reported LIBERO parameter–success Pareto frontier without a large pretrained VLM action backbone.
- A single diffusion transformer couples continuous action generation with future visual prediction while retaining an action-only deployment path.
- 399M-parameter system lies on the reported LIBERO parameter–success Pareto frontier.The result indicates strong benchmark performance without placing a large pretrained vision-language model in the action backbone.
- Shared flow matching for action and visual targets may support independently trained experts for heterogeneous hardware and compute-constrained robot platforms.This direction follows configurations explored in Paris and Paris 2.0.