Improve the WorldDiT model card layout
#2
by bagellabs - opened
README.md
CHANGED
|
@@ -12,11 +12,9 @@ tags:
|
|
| 12 |
- robotic-manipulation
|
| 13 |
- imitation-learning
|
| 14 |
- diffusion-transformer
|
|
|
|
| 15 |
- flow-matching
|
| 16 |
-
- pareto-frontier
|
| 17 |
-
- parameter-efficiency
|
| 18 |
inference: false
|
| 19 |
-
license: cc-by-4.0
|
| 20 |
widget:
|
| 21 |
- example_title: "LIBERO Spatial, task 5"
|
| 22 |
text: "Successful rollout, front view."
|
|
@@ -40,20 +38,19 @@ widget:
|
|
| 40 |
<img src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/paris2/model-card/v1/bagel_labs_logo.png" alt="Bagel Labs">
|
| 41 |
</p>
|
| 42 |
|
| 43 |
-
<h1 align="center">WorldDiT
|
|
|
|
|
|
|
| 44 |
|
| 45 |
<p align="center">
|
| 46 |
-
|
| 47 |
-
|
| 48 |
-
|
| 49 |
</p>
|
| 50 |
|
| 51 |
<p align="center">
|
| 52 |
-
|
| 53 |
-
|
| 54 |
-
as a general backbone for world and action modeling, while the current release
|
| 55 |
-
evaluates it on LIBERO and provides four checkpoints, a self contained
|
| 56 |
-
inference runtime, and an evaluator.
|
| 57 |
</p>
|
| 58 |
|
| 59 |
## See WorldDiT act
|
|
@@ -101,8 +98,8 @@ Each clip covers a different LIBERO suite and camera view.
|
|
| 101 |
|
| 102 |
| Reported LIBERO result | Released model |
|
| 103 |
|---|---|
|
| 104 |
-
| **94.9 percent**
|
| 105 |
-
| **98.0 percent** Spatial<br>**97.0 percent** Object | **Three** observation
|
| 106 |
| **92.8 percent** Goal<br>**91.8 percent** Long | **Three** actions executed before replanning<br>**Seven** action dimensions |
|
| 107 |
|
| 108 |
| Checkpoints | Runtime | Encoders and environment |
|
|
@@ -113,23 +110,9 @@ The repository is self contained for WorldDiT inference. LIBERO provides the
|
|
| 113 |
benchmark environments, assets, task definitions, and initial states.
|
| 114 |
|
| 115 |
The released runtime and checkpoints were revalidated from a clean installation
|
| 116 |
-
on eight RTX Pro 6000 Blackwell GPUs. The
|
| 117 |
-
hundred
|
| 118 |
-
|
| 119 |
-
from selection.
|
| 120 |
-
|
| 121 |
-
### Parameter count and reported success
|
| 122 |
-
|
| 123 |
-
Among methods with complete four suite averages, WorldDiT lies on the reported
|
| 124 |
-
Pareto frontier for total model parameters and mean LIBERO success.
|
| 125 |
-
|
| 126 |
-
<p align="center">
|
| 127 |
-
<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">
|
| 128 |
-
</p>
|
| 129 |
-
|
| 130 |
-
<p align="center">
|
| 131 |
-
<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>
|
| 132 |
-
</p>
|
| 133 |
|
| 134 |
## Run a smoke test
|
| 135 |
|
|
@@ -166,34 +149,24 @@ python eval.py \
|
|
| 166 |
```
|
| 167 |
|
| 168 |
A successful smoke test confirms that the environment, checkpoint, visual
|
| 169 |
-
encoders, simulator, and rendering path load together.
|
| 170 |
-
|
| 171 |
|
| 172 |
## How WorldDiT works
|
| 173 |
|
| 174 |
-
|
| 175 |
-
|
| 176 |
-
|
| 177 |
-
|
| 178 |
-
|
| 179 |
-
directly. RGB patch token construction and RGB prediction head evaluation
|
| 180 |
-
remain outside the inference graph, concentrating computation on action
|
| 181 |
-
generation. The controller executes the first three predicted actions, observes
|
| 182 |
-
again, and replans.
|
| 183 |
|
| 184 |
-
|
| 185 |
-
<img src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v2/worlddit_inference_pipeline.png?v=09d6508" alt="WorldDiT inference pipeline">
|
| 186 |
-
</p>
|
| 187 |
-
|
| 188 |
-
<p align="center">
|
| 189 |
-
<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>
|
| 190 |
-
</p>
|
| 191 |
|
| 192 |
| Training | Deployment |
|
| 193 |
|---|---|
|
| 194 |
-
| Action and future
|
| 195 |
| Seven action steps are supervised | Seven actions are predicted |
|
| 196 |
-
| Future
|
| 197 |
| The complete training objective is active | Three actions execute before replanning |
|
| 198 |
|
| 199 |
## Reference
|
|
@@ -221,8 +194,8 @@ CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python eval.py \
|
|
| 221 |
|
| 222 |
Each GPU receives an independent progress bar. After all workers finish, rank 0
|
| 223 |
prints per task and overall success rates and writes a structured
|
| 224 |
-
`results.json`.
|
| 225 |
-
|
| 226 |
|
| 227 |
Supported suites.
|
| 228 |
|
|
@@ -255,8 +228,7 @@ libero_10
|
|
| 255 |
```
|
| 256 |
|
| 257 |
`dependencies/` contains the frozen visual and language encoder weights needed
|
| 258 |
-
by the released
|
| 259 |
-
required for inference.
|
| 260 |
|
| 261 |
</details>
|
| 262 |
|
|
@@ -287,8 +259,8 @@ Evaluation uses the final temporal slot of the predicted action tensor.
|
|
| 287 |
|
| 288 |
| Component | Specification |
|
| 289 |
|---|---|
|
| 290 |
-
|
|
| 291 |
-
| Observation context | 3
|
| 292 |
| Action horizon | 7 actions |
|
| 293 |
| Action dimension | 7 |
|
| 294 |
| Action aggregation | Temporal ensembling |
|
|
@@ -303,32 +275,8 @@ Evaluation uses the final temporal slot of the predicted action tensor.
|
|
| 303 |
|
| 304 |
| Intended use | Scope of the release |
|
| 305 |
|---|---|
|
| 306 |
-
| Research on
|
| 307 |
-
|
|
| 308 |
-
|
| 309 |
-
## Citation
|
| 310 |
-
|
| 311 |
-
If you use WorldDiT in your research, please cite the paper.
|
| 312 |
-
|
| 313 |
-
```bibtex
|
| 314 |
-
@article{260723909,
|
| 315 |
-
title={{WorldDiT: A Unified Diffusion Architecture for World and Action Modeling}},
|
| 316 |
-
author={Sen Wang and R. Gnana Praveen and Bidhan Roy and Marcos Villagra},
|
| 317 |
-
year={{2026}},
|
| 318 |
-
eprint={{2607.23909}},
|
| 319 |
-
archivePrefix={{arXiv}}
|
| 320 |
-
}
|
| 321 |
-
```
|
| 322 |
-
|
| 323 |
-
## License
|
| 324 |
-
|
| 325 |
-
The WorldDiT checkpoints, model card, and original release materials are
|
| 326 |
-
licensed under [Creative Commons Attribution 4.0
|
| 327 |
-
International](https://creativecommons.org/licenses/by/4.0/). You may copy,
|
| 328 |
-
redistribute, and adapt them, including commercially, with appropriate credit
|
| 329 |
-
to Bagel Labs and the WorldDiT authors, a link to the license, and an indication
|
| 330 |
-
of any changes. Third party dependencies and assets remain governed by their
|
| 331 |
-
upstream licenses.
|
| 332 |
|
| 333 |
## Authors and contact
|
| 334 |
|
|
|
|
| 12 |
- robotic-manipulation
|
| 13 |
- imitation-learning
|
| 14 |
- diffusion-transformer
|
| 15 |
+
- diffusion-policy
|
| 16 |
- flow-matching
|
|
|
|
|
|
|
| 17 |
inference: false
|
|
|
|
| 18 |
widget:
|
| 19 |
- example_title: "LIBERO Spatial, task 5"
|
| 20 |
text: "Successful rollout, front view."
|
|
|
|
| 38 |
<img src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/paris2/model-card/v1/bagel_labs_logo.png" alt="Bagel Labs">
|
| 39 |
</p>
|
| 40 |
|
| 41 |
+
<h1 align="center">WorldDiT</h1>
|
| 42 |
+
|
| 43 |
+
<h2 align="center">One diffusion backbone learns what to do and what comes next.</h2>
|
| 44 |
|
| 45 |
<p align="center">
|
| 46 |
+
WorldDiT learns continuous robot action chunks and a future visual target
|
| 47 |
+
through one shared diffusion transformer. Deployment keeps only the action
|
| 48 |
+
path.
|
| 49 |
</p>
|
| 50 |
|
| 51 |
<p align="center">
|
| 52 |
+
This release includes four LIBERO checkpoints, a self contained inference
|
| 53 |
+
runtime, and an evaluator for reproducing the reported suite results.
|
|
|
|
|
|
|
|
|
|
| 54 |
</p>
|
| 55 |
|
| 56 |
## See WorldDiT act
|
|
|
|
| 98 |
|
| 99 |
| Reported LIBERO result | Released model |
|
| 100 |
|---|---|
|
| 101 |
+
| **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 |
|
| 102 |
+
| **98.0 percent** Spatial<br>**97.0 percent** Object | **Three** observation frames<br>**Seven** predicted actions |
|
| 103 |
| **92.8 percent** Goal<br>**91.8 percent** Long | **Three** actions executed before replanning<br>**Seven** action dimensions |
|
| 104 |
|
| 105 |
| Checkpoints | Runtime | Encoders and environment |
|
|
|
|
| 110 |
benchmark environments, assets, task definitions, and initial states.
|
| 111 |
|
| 112 |
The released runtime and checkpoints were revalidated from a clean installation
|
| 113 |
+
on eight RTX Pro 6000 Blackwell GPUs. The result is selection aware because
|
| 114 |
+
three hundred episodes per suite informed staged checkpoint selection before
|
| 115 |
+
the final five hundred episode score was assembled.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 116 |
|
| 117 |
## Run a smoke test
|
| 118 |
|
|
|
|
| 149 |
```
|
| 150 |
|
| 151 |
A successful smoke test confirms that the environment, checkpoint, visual
|
| 152 |
+
encoders, simulator, and rendering path load together. It is not a benchmark
|
| 153 |
+
result.
|
| 154 |
|
| 155 |
## How WorldDiT works
|
| 156 |
|
| 157 |
+
WorldDiT uses three recent observations, robot state, and language as context.
|
| 158 |
+
During training, one diffusion transformer learns a seven step action chunk and
|
| 159 |
+
an auxiliary future visual target. During deployment, the future visual path is
|
| 160 |
+
absent. The policy executes the first three predicted actions, observes again,
|
| 161 |
+
and replans.
|
|
|
|
|
|
|
|
|
|
|
|
|
| 162 |
|
| 163 |
+
> Future visual prediction is a training signal, not a deployment path.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 164 |
|
| 165 |
| Training | Deployment |
|
| 166 |
|---|---|
|
| 167 |
+
| Action and future visual targets share one backbone | Only the action path remains |
|
| 168 |
| Seven action steps are supervised | Seven actions are predicted |
|
| 169 |
+
| Future visual supervision is present | No future visual output is requested |
|
| 170 |
| The complete training objective is active | Three actions execute before replanning |
|
| 171 |
|
| 172 |
## Reference
|
|
|
|
| 194 |
|
| 195 |
Each GPU receives an independent progress bar. After all workers finish, rank 0
|
| 196 |
prints per task and overall success rates and writes a structured
|
| 197 |
+
`results.json`. Output directories must be new so an earlier evaluation is
|
| 198 |
+
never overwritten.
|
| 199 |
|
| 200 |
Supported suites.
|
| 201 |
|
|
|
|
| 228 |
```
|
| 229 |
|
| 230 |
`dependencies/` contains the frozen visual and language encoder weights needed
|
| 231 |
+
by the released policy. No additional model downloads are required.
|
|
|
|
| 232 |
|
| 233 |
</details>
|
| 234 |
|
|
|
|
| 259 |
|
| 260 |
| Component | Specification |
|
| 261 |
|---|---|
|
| 262 |
+
| Policy | WorldDiT diffusion transformer |
|
| 263 |
+
| Observation context | 3 frames |
|
| 264 |
| Action horizon | 7 actions |
|
| 265 |
| Action dimension | 7 |
|
| 266 |
| Action aggregation | Temporal ensembling |
|
|
|
|
| 275 |
|
| 276 |
| Intended use | Scope of the release |
|
| 277 |
|---|---|
|
| 278 |
+
| 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. |
|
| 279 |
+
| 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. |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 280 |
|
| 281 |
## Authors and contact
|
| 282 |
|