README.md

metadata

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.

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 selection aware mean 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 frames 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 result is selection aware because three hundred episodes per suite informed staged checkpoint selection before the final five hundred episode score was assembled.

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

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. Output directories must be new so an earlier evaluation is never overwritten.

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 policy. No additional model downloads are required.

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
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

Use and scope

Intended use	Scope of the release
Research on language conditioned robot manipulation in the LIBERO simulator. The released checkpoints support reproduction, evaluation, and architecture research across the four released suites.	The results describe LIBERO simulation under the released evaluation protocol. They do not establish real robot reliability, safety, or transfer across embodiments.
The released checkpoints cover all four LIBERO suites.	The release does not isolate the causal contribution of the future visual target. Total parameter count 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.

---

Made with ❤️ by