library_name: pytorch
pipeline_tag: robotics
language:
- en
tags:
- worlddit
- world-action-model
- world-models
- libero
- robot-learning
- robotic-manipulation
- imitation-learning
- diffusion-transformer
- diffusion-policy
- flow-matching
inference: false
widget:
- example_title: LIBERO Spatial, task 5
text: Successful rollout, front view.
output:
url: >-
https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_spatial_frontview_task05_episode01.mp4
- example_title: LIBERO Object, task 8
text: Successful rollout, agent view.
output:
url: >-
https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_object_agentview_task08_episode01.mp4
- example_title: LIBERO Goal, task 10
text: Successful rollout, side view.
output:
url: >-
https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_goal_sideview_task10_episode01.mp4
- example_title: LIBERO Long, task 6
text: Successful rollout, front view.
output:
url: >-
https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_10_frontview_task06_episode01.mp4
WorldDiT
One diffusion backbone learns what to do and what comes next.
WorldDiT learns continuous robot action chunks and a future visual target through one shared diffusion transformer. Deployment keeps only the action path.
This release includes four LIBERO checkpoints, a self contained inference runtime, and an evaluator for reproducing the reported suite results.
See WorldDiT act
The four clips below show successful rollouts from the released checkpoints. Each clip covers a different LIBERO suite and camera view.
- Prompt
- Successful rollout, front view.
- Prompt
- Successful rollout, agent view.
- Prompt
- Successful rollout, side view.
- Prompt
- Successful rollout, front view.
| Suite | View | Video |
|---|---|---|
| LIBERO Spatial | Front view | Open MP4 |
| LIBERO Object | Agent view | Open MP4 |
| LIBERO Goal | Side view | Open MP4 |
| LIBERO Long | Front view | Open MP4 |
What is in this release
| Release component | Included artifact |
|---|---|
| LIBERO Spatial policy | SafeTensors checkpoint |
| LIBERO Object policy | SafeTensors checkpoint |
| LIBERO Goal policy | SafeTensors checkpoint |
| LIBERO Long policy | SafeTensors checkpoint |
| Model runtime | inference.py |
| Evaluation runtime | eval.py |
| Frozen encoders | CLIP ViT B 32 and MAE ViT B |
| Configuration | config.json |
| Environment | Pinned Python requirements |
The repository is self contained for WorldDiT inference. LIBERO still provides the benchmark environments, assets, task definitions, and initial states.
Reported LIBERO results
Across the four released suite checkpoints, WorldDiT records 1,898 successful episodes out of 2,000 under the selection aware evaluation protocol.
| Suite | Successful episodes | Success rate |
|---|---|---|
| LIBERO Spatial | 490 of 500 | 98.0 percent |
| LIBERO Object | 485 of 500 | 97.0 percent |
| LIBERO Goal | 464 of 500 | 92.8 percent |
| LIBERO Long | 459 of 500 | 91.8 percent |
| Selection aware mean | 1,898 of 2,000 | 94.9 percent |
The released runtime and checkpoints were revalidated from a clean installation on eight RTX Pro 6000 Blackwell GPUs.
The result is selection aware because three hundred episodes per suite informed staged checkpoint selection before the final five hundred episode score was assembled.
Model at a glance
| Property | Released configuration |
|---|---|
| Total parameters | 399.084 million |
| Trainable parameters | 135.107 million |
| Observation context | Three frames |
| Predicted action horizon | Seven actions |
| Executed before replanning | Three actions |
| Action dimension | Seven |
| Visual encoder | MAE ViT B |
| Language encoder | OpenAI CLIP ViT B 32 |
| Checkpoint format | SafeTensors |
| Evaluation environment | Headless LIBERO with EGL |
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. It is not a benchmark result.
How WorldDiT works
WorldDiT uses three recent observations, robot state, and language as context. During training, one diffusion transformer learns a seven step action chunk and an auxiliary future visual target. During deployment, the future visual path is absent. The policy executes the first three predicted actions, observes again, and replans.
Future visual prediction is a training signal, not a deployment path.
| Training | Deployment |
|---|---|
| Action and future visual targets share one backbone | Only the action path remains |
| Seven action steps are supervised | Seven actions are predicted |
| Future visual supervision is present | No future visual output is requested |
| The complete training objective is active | Three actions execute before replanning |
Evaluation
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. Output directories must be new so an earlier evaluation is
never overwritten.
Supported suites.
libero_spatial
libero_object
libero_goal
libero_10
What this repository contains
.
βββ 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 policy. No additional model downloads are required.
Inference API
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 |
|---|---|
| Policy | WorldDiT diffusion transformer |
| Observation context | 3 frames |
| 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 |
Intended use
WorldDiT is intended for research on language conditioned robot manipulation in the LIBERO simulator. The released checkpoints support reproduction, evaluation, and architecture research across the four released suites.
Scope of the release
The reported results describe LIBERO simulation under the released evaluation protocol. They do not establish real robot reliability, safety, or transfer across embodiments.
The present release does not isolate the causal contribution of the future visual target. Total parameter count also does not measure training cost, deployment latency, or runtime efficiency.
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.