WorldDiT: A Unified Diffusion Architecture for World and Action Modeling

2026-07-27Machine Learning

Machine LearningRobotics
AI summary

The authors present WorldDiT, a new robot control model that doesn't rely on big pretrained vision-language models. Instead, it uses a single 'diffusion transformer' that both plans robot actions and predicts future images from the robot's camera. They tested WorldDiT in four robot simulation environments and found it performed well compared to other methods, using fewer parameters. This work offers a compact baseline model for future improvements in robot control.

robot policiespretrained vision-language modelsdiffusion transformeraction generationvisual world modelingRGB patchescamera framessimulation suitesmodel parametersPareto frontier
Authors
Sen Wang, R. Gnana Praveen, Bidhan Roy, Marcos Villagra
Abstract
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.