README.md

---
library_name: pytorch
pipeline_tag: robotics
language:
  - en
tags:
  - worlddit
  - world-action-model
  - world-models
  - libero
  - robot-learning
  - robotic-manipulation
  - imitation-learning
  - diffusion-transformer
  - flow-matching
  - pareto-frontier
  - parameter-efficiency
inference: false
license: cc-by-4.0
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"
---

<p align="center">
  <img src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/paris2/model-card/v1/bagel_labs_logo.png" alt="Bagel Labs">
</p>

<h1 align="center">WorldDiT: A Unified Diffusion Backbone for<br>World and Action Modeling</h1>

<p align="center">
  <a href="https://arxiv.org/abs/2607.23909">
    <img src="https://img.shields.io/badge/arXiv-2607.23909-B31B1B.svg?logo=arxiv&logoColor=white" alt="WorldDiT paper on arXiv">
  </a>
</p>

<p align="center">
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.
</p>

## See WorldDiT act

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

<table>
  <tr>
    <td width="50%" valign="top">
      <video width="100%" controls muted loop playsinline preload="metadata"
        src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_spatial_frontview_task05_episode01.mp4"></video>
      <strong>LIBERO Spatial</strong><br>
      Task 5, front view.
    </td>
    <td width="50%" valign="top">
      <video width="100%" controls muted loop playsinline preload="metadata"
        src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_object_agentview_task08_episode01.mp4"></video>
      <strong>LIBERO Object</strong><br>
      Task 8, agent view.
    </td>
  </tr>
  <tr>
    <td width="50%" valign="top">
      <video width="100%" controls muted loop playsinline preload="metadata"
        src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_goal_sideview_task10_episode01.mp4"></video>
      <strong>LIBERO Goal</strong><br>
      Task 10, side view.
    </td>
    <td width="50%" valign="top">
      <video width="100%" controls muted loop playsinline preload="metadata"
        src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_10_frontview_task06_episode01.mp4"></video>
      <strong>LIBERO Long</strong><br>
      Task 6, front view.
    </td>
  </tr>
</table>

[Spatial MP4](https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_spatial_frontview_task05_episode01.mp4)
· [Object MP4](https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_object_agentview_task08_episode01.mp4)
· [Goal MP4](https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_goal_sideview_task10_episode01.mp4)
· [Long MP4](https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_10_frontview_task06_episode01.mp4)

## Release snapshot

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

| Checkpoints | Runtime | Encoders and environment |
|---|---|---|
| Spatial<br>Object<br>Goal<br>Long | `inference.py`<br>`eval.py`<br>`config.json` | MAE ViT B<br>OpenAI CLIP ViT B 32<br>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.

<p align="center">
  <img src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v2/worlddit_parameter_efficiency.png" alt="Reported mean LIBERO success against total instantiated model parameters">
</p>

<p align="center">
  <sub>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.</sub>
</p>

## Run a smoke test

Download the repository and create a clean Python 3.12 environment.

```bash
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.

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

The released evaluation was validated with LIBERO commit
`8f1084e3132a39270c3a13ebe37270a43ece2a01`.

```bash
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.

<p align="center">
  <img src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v2/worlddit_inference_pipeline.png?v=09d6508" alt="WorldDiT inference pipeline">
</p>

<p align="center">
  <sub>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.</sub>
</p>

| 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

<details>
<summary><strong>Full evaluation commands</strong></summary>

### One GPU

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

### Multiple GPUs

```bash
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.

```text
libero_spatial
libero_object
libero_goal
libero_10
```

</details>

<details>
<summary><strong>Repository contents</strong></summary>

```text
.
├── 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.

</details>

<details>
<summary><strong>Inference API and tensor shapes</strong></summary>

```python
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.

</details>

<details>
<summary><strong>Architecture details</strong></summary>

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

</details>

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

## Citation

If you use WorldDiT in your research, please cite the paper.

```bibtex
@article{260723909,
  title={{WorldDiT: A Unified Diffusion Architecture for World and Action Modeling}},
  author={Sen Wang and R. Gnana Praveen and Bidhan Roy and Marcos Villagra},
  year={{2026}},
  eprint={{2607.23909}},
  archivePrefix={{arXiv}}
}
```

## License

The WorldDiT checkpoints, model card, and original release materials are
licensed under [Creative Commons Attribution 4.0
International](https://creativecommons.org/licenses/by/4.0/). 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](https://github.com/Lifelong-Robot-Learning/LIBERO),
[robosuite](https://github.com/ARISE-Initiative/robosuite),
[OpenAI CLIP](https://github.com/openai/CLIP), and
[Masked Autoencoders](https://github.com/facebookresearch/mae). Third party
components remain subject to their respective upstream terms.

---

<div style="display: flex; align-items: center; gap: 8px;">
  <span>Made with ❤️ by</span>
  <a href="https://twitter.com/bageldotcom" target="_blank">
    <img src="https://img.shields.io/badge/Bagel_Labs-1DA1F2?style=for-the-badge&logo=twitter&logoColor=white" alt="Follow Bagel Labs on Twitter" height="28">
  </a>
</div>