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

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

<h2 align="center">One diffusion backbone learns what to do and what comes next.</h2>

<p align="center">
WorldDiT learns continuous robot action chunks and a future visual target
through one shared diffusion transformer. Deployment keeps only the action
path.
</p>

<p align="center">
This release includes four LIBERO checkpoints, a self contained inference
runtime, and an evaluator for reproducing the reported suite results.
</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** selection aware mean<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 frames<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 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.

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

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

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

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

</details>

## 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](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>