WorldDiT: A Unified Diffusion Backbone for
World and Action Modeling

WorldDiT couples continuous action generation with auxiliary future normalized RGB patch prediction in one diffusion transformer. The architecture is designed as a general backbone for world and action modeling, while the current release evaluates it on LIBERO and provides four checkpoints, a self contained inference runtime, and an evaluator.

See WorldDiT act

The four clips below show successful rollouts from the released checkpoints. Each clip covers a different LIBERO suite and camera view.

LIBERO Spatial Task 5, front view.	LIBERO Object Task 8, agent view.
LIBERO Goal Task 10, side view.	LIBERO Long Task 6, front view.

Spatial MP4 · Object MP4 · Goal MP4 · Long MP4

Release snapshot

Reported LIBERO result	Released model
94.9 percent mean success 1,898 of 2,000 successful episodes	399.084 million total parameters 135.107 million trainable parameters
98.0 percent Spatial 97.0 percent Object	Three observation steps Seven predicted actions
92.8 percent Goal 91.8 percent Long	Three actions executed before replanning Seven action dimensions

Checkpoints	Runtime	Encoders and environment
Spatial Object Goal Long	inference.py eval.py config.json	MAE ViT B OpenAI CLIP ViT B 32 SafeTensors and pinned requirements

The repository is self contained for WorldDiT inference. LIBERO provides the benchmark environments, assets, task definitions, and initial states.

The released runtime and checkpoints were revalidated from a clean installation on eight RTX Pro 6000 Blackwell GPUs. The reported aggregate covers five hundred simulator episodes per suite. Three hundred episodes per suite informed staged checkpoint selection, while two hundred episodes per suite were disjoint from selection.

Parameter count and reported success

Among methods with complete four suite averages, WorldDiT lies on the reported Pareto frontier for total model parameters and mean LIBERO success.

_{Reported LIBERO success against total model parameters for 24 methods. The line connects methods on the Pareto frontier with complete four suite averages. Because the methods follow different published evaluation protocols, the figure summarizes published results rather than a direct comparison under one common evaluation protocol.}

Run a smoke test

Download the repository and create a clean Python 3.12 environment.

hf download bageldotcom/worlddit --local-dir worlddit
cd worlddit

python3.12 -m venv venv
source venv/bin/activate
python -m pip install -r requirements.txt
python -m pip install --no-deps robosuite==1.4.1

LIBERO supplies the benchmark definitions, assets, and initial states. Keep the checkout at ~/LIBERO, which is the evaluator's default.

git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git ~/LIBERO

The released evaluation was validated with LIBERO commit 8f1084e3132a39270c3a13ebe37270a43ece2a01.

python eval.py \
  --suite libero_spatial \
  --gpus 1 \
  --tasks 1 \
  --episodes 1 \
  --max-steps 20 \
  --output-dir results/smoke

A successful smoke test confirms that the environment, checkpoint, visual encoders, simulator, and rendering path load together. Full benchmark reporting uses complete suite evaluations.

How WorldDiT works

Each of three recent observation steps contributes primary and wrist images together with robot state, while one language instruction conditions the sequence. During training, one diffusion transformer backbone learns a seven step action chunk together with an auxiliary future normalized RGB patch target. At deployment, the encoded history conditions the action velocity field directly. RGB patch token construction and RGB prediction head evaluation remain outside the inference graph, concentrating computation on action generation. The controller executes the first three predicted actions, observes again, and replans.

_{WorldDiT inference pipeline. The encoded observation history conditions action generation through twenty flow steps. The controller executes the first three actions from each seven action chunk, then updates the window and replans.}

Training	Deployment
Action and future normalized RGB patch targets are learned by one backbone	Encoded history conditions the action velocity field
Seven action steps are supervised	Seven actions are predicted
Future normalized RGB patch supervision is present	RGB patch tokens and the RGB prediction head remain outside the inference graph
The complete training objective is active	Three actions execute before replanning

Reference

Full evaluation commands

One GPU

python eval.py \
  --suite libero_spatial \
  --gpus 1 \
  --output-dir results/libero_spatial

Multiple GPUs

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python eval.py \
  --suite libero_spatial \
  --gpus 8 \
  --output-dir results/libero_spatial_8gpu

Each GPU receives an independent progress bar. After all workers finish, rank 0 prints per task and overall success rates and writes a structured results.json. Use a new output directory for each evaluation to preserve earlier results.

Supported suites.

libero_spatial
libero_object
libero_goal
libero_10

Repository contents

.
├── checkpoints/
│   ├── libero_10/model.safetensors
│   ├── libero_goal/model.safetensors
│   ├── libero_object/model.safetensors
│   └── libero_spatial/model.safetensors
├── dependencies/
│   ├── ViT-B-32.pt
│   └── mae_pretrain_vit_base.pth
├── eval.py
├── inference.py
├── config.json
└── requirements.txt

dependencies/ contains the frozen visual and language encoder weights needed by the released WorldDiT runtime. The repository contains every model weight required for inference.

Inference API and tensor shapes

from inference import load_model

model = load_model(".", suite="libero_spatial", device="cuda")
actions = model(primary_images, wrist_images, robot_state, text_tokens)

Input or output	Shape
Primary-camera images	[B, 3, 3, 224, 224]
Wrist-camera images	[B, 3, 3, 224, 224]
Robot state	[B, 3, 8]
OpenAI CLIP text tokens	[B, 3, 77]
Predicted action tensor	[B, 3, 7, 7]

Evaluation uses the final temporal slot of the predicted action tensor.

Architecture details

Component	Specification
Backbone	WorldDiT diffusion transformer
Observation context	3 observation steps
Action horizon	7 actions
Action dimension	7
Action aggregation	Temporal ensembling
Language encoder	OpenAI CLIP ViT-B/32
Visual encoder	MAE ViT-B
Evaluation	Headless LIBERO with EGL
Checkpoint format	SafeTensors

Use and scope

Intended use	Scope of the release
Research on world and action modeling for language conditioned robot manipulation. The architecture supports continuous action generation with auxiliary future normalized RGB patch prediction.	The current release evaluates WorldDiT in LIBERO simulation under the released protocol and provides checkpoints for all four suites.
Architecture research, reproduction, and evaluation of multimodal diffusion backbones for robot manipulation.	Real robot reliability, safety, and transfer across embodiments require dedicated future evaluation. The present experiments evaluate the integrated WorldDiT system. Targeted ablations are required to attribute performance to the future normalized RGB patch objective. Total instantiated parameter count characterizes model scale. Training cost, deployment latency, and runtime efficiency require dedicated measurements.

License

The WorldDiT checkpoints, model card, and original release materials are licensed under Creative Commons Attribution 4.0 International. You may copy, redistribute, and adapt them, including commercially, with appropriate credit to Bagel Labs and the WorldDiT authors, a link to the license, and an indication of any changes. Third party dependencies and assets remain governed by their upstream licenses.

Authors and contact

WorldDiT is developed by Sen Wang, Praveen Rajasekhar, Bidhan Roy, and Marcos Villagra at Bagel Labs. Questions can be sent to research@bagel.com.

Acknowledgments

This release builds on LIBERO, robosuite, OpenAI CLIP, and Masked Autoencoders. Third party components remain subject to their respective upstream terms.

---

Made with ❤️ by