Make the WorldDiT model card demo first

#1
by bagellabs - opened
Files changed (1) hide show
  1. README.md +97 -146
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."
@@ -36,102 +34,95 @@ widget:
36
  url: "https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_10_frontview_task06_episode01.mp4"
37
  ---
38
 
39
- Bagel Labs is a physical AI research lab based in San Francisco and Toronto. It has developed an architecture that uses dramatically less data and compute to train while delivering a state-of-the-art model-size/performance tradeoff. WorldDiT demonstrates the model advantage; Paris 1.0 and Paris 2.0 demonstrate the training advantage.
40
-
41
  <p align="center">
42
  <img src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/paris2/model-card/v1/bagel_labs_logo.png" alt="Bagel Labs">
43
  </p>
44
 
45
- <h1 align="center">WorldDiT: A Unified Diffusion Backbone for<br>World and Action Modeling</h1>
 
 
46
 
47
  <p align="center">
48
- <a href="https://arxiv.org/abs/2607.23909">
49
- <img src="https://img.shields.io/badge/arXiv-2607.23909-B31B1B.svg?logo=arxiv&logoColor=white" alt="WorldDiT paper on arXiv">
 
 
 
50
  </a>
51
  </p>
52
 
53
- <p align="center">
54
- WorldDiT couples continuous action generation with auxiliary future normalized
55
- RGB patch prediction in one diffusion transformer. The architecture is designed
56
- as a general backbone for world and action modeling, while the current release
57
- evaluates it on LIBERO and provides four checkpoints, a self contained
58
- inference runtime, and an evaluator.
59
- </p>
60
 
61
  ## See WorldDiT act
62
 
63
  The four clips below show successful rollouts from the released checkpoints.
64
  Each clip covers a different LIBERO suite and camera view.
65
 
66
- <table>
67
- <tr>
68
- <td width="50%" valign="top">
69
- <video width="100%" controls muted loop playsinline preload="metadata"
70
- src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_spatial_frontview_task05_episode01.mp4"></video>
71
- <strong>LIBERO Spatial</strong><br>
72
- Task 5, front view.
73
- </td>
74
- <td width="50%" valign="top">
75
- <video width="100%" controls muted loop playsinline preload="metadata"
76
- src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_object_agentview_task08_episode01.mp4"></video>
77
- <strong>LIBERO Object</strong><br>
78
- Task 8, agent view.
79
- </td>
80
- </tr>
81
- <tr>
82
- <td width="50%" valign="top">
83
- <video width="100%" controls muted loop playsinline preload="metadata"
84
- src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_goal_sideview_task10_episode01.mp4"></video>
85
- <strong>LIBERO Goal</strong><br>
86
- Task 10, side view.
87
- </td>
88
- <td width="50%" valign="top">
89
- <video width="100%" controls muted loop playsinline preload="metadata"
90
- src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_10_frontview_task06_episode01.mp4"></video>
91
- <strong>LIBERO Long</strong><br>
92
- Task 6, front view.
93
- </td>
94
- </tr>
95
- </table>
96
-
97
- [Spatial MP4](https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_spatial_frontview_task05_episode01.mp4)
98
- · [Object MP4](https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_object_agentview_task08_episode01.mp4)
99
- · [Goal MP4](https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_goal_sideview_task10_episode01.mp4)
100
- · [Long MP4](https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_10_frontview_task06_episode01.mp4)
101
-
102
- ## Release snapshot
103
-
104
- | Reported LIBERO result | Released model |
105
- |---|---|
106
- | **94.9 percent** mean success<br>**1,898 of 2,000** successful episodes | **399.084 million** total parameters<br>**135.107 million** trainable parameters |
107
- | **98.0 percent** Spatial<br>**97.0 percent** Object | **Three** observation steps<br>**Seven** predicted actions |
108
- | **92.8 percent** Goal<br>**91.8 percent** Long | **Three** actions executed before replanning<br>**Seven** action dimensions |
109
 
110
- | Checkpoints | Runtime | Encoders and environment |
111
  |---|---|---|
112
- | 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 |
 
 
 
113
 
114
- The repository is self contained for WorldDiT inference. LIBERO provides the
115
- benchmark environments, assets, task definitions, and initial states.
116
 
117
- The released runtime and checkpoints were revalidated from a clean installation
118
- on eight RTX Pro 6000 Blackwell GPUs. The reported aggregate covers five
119
- hundred simulator episodes per suite. Three hundred episodes per suite informed
120
- staged checkpoint selection, while two hundred episodes per suite were disjoint
121
- from selection.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
122
 
123
- ### Parameter count and reported success
 
124
 
125
- Among methods with complete four suite averages, WorldDiT lies on the reported
126
- Pareto frontier for total model parameters and mean LIBERO success.
 
127
 
128
- <p align="center">
129
- <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">
130
- </p>
131
 
132
- <p align="center">
133
- <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>
134
- </p>
 
 
 
 
 
 
 
 
 
135
 
136
  ## Run a smoke test
137
 
@@ -168,40 +159,27 @@ python eval.py \
168
  ```
169
 
170
  A successful smoke test confirms that the environment, checkpoint, visual
171
- encoders, simulator, and rendering path load together. Full benchmark reporting
172
- uses complete suite evaluations.
173
 
174
  ## How WorldDiT works
175
 
176
- Each of three recent observation steps contributes primary and wrist images
177
- together with robot state, while one language instruction conditions the
178
- sequence. During training, one diffusion transformer backbone learns a seven
179
- step action chunk together with an auxiliary future normalized RGB patch
180
- target. At deployment, the encoded history conditions the action velocity field
181
- directly. RGB patch token construction and RGB prediction head evaluation
182
- remain outside the inference graph, concentrating computation on action
183
- generation. The controller executes the first three predicted actions, observes
184
- again, and replans.
185
-
186
- <p align="center">
187
- <img src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v2/worlddit_inference_pipeline.png?v=09d6508" alt="WorldDiT inference pipeline">
188
- </p>
189
 
190
- <p align="center">
191
- <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>
192
- </p>
193
 
194
  | Training | Deployment |
195
  |---|---|
196
- | Action and future normalized RGB patch targets are learned by one backbone | Encoded history conditions the action velocity field |
197
  | Seven action steps are supervised | Seven actions are predicted |
198
- | Future normalized RGB patch supervision is present | RGB patch tokens and the RGB prediction head remain outside the inference graph |
199
  | The complete training objective is active | Three actions execute before replanning |
200
 
201
- ## Reference
202
-
203
- <details>
204
- <summary><strong>Full evaluation commands</strong></summary>
205
 
206
  ### One GPU
207
 
@@ -223,8 +201,8 @@ CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python eval.py \
223
 
224
  Each GPU receives an independent progress bar. After all workers finish, rank 0
225
  prints per task and overall success rates and writes a structured
226
- `results.json`. Use a new output directory for each evaluation to preserve
227
- earlier results.
228
 
229
  Supported suites.
230
 
@@ -235,10 +213,7 @@ libero_goal
235
  libero_10
236
  ```
237
 
238
- </details>
239
-
240
- <details>
241
- <summary><strong>Repository contents</strong></summary>
242
 
243
  ```text
244
  .
@@ -257,13 +232,9 @@ libero_10
257
  ```
258
 
259
  `dependencies/` contains the frozen visual and language encoder weights needed
260
- by the released WorldDiT runtime. The repository contains every model weight
261
- required for inference.
262
 
263
- </details>
264
-
265
- <details>
266
- <summary><strong>Inference API and tensor shapes</strong></summary>
267
 
268
  ```python
269
  from inference import load_model
@@ -282,15 +253,12 @@ actions = model(primary_images, wrist_images, robot_state, text_tokens)
282
 
283
  Evaluation uses the final temporal slot of the predicted action tensor.
284
 
285
- </details>
286
-
287
- <details>
288
- <summary><strong>Architecture details</strong></summary>
289
 
290
  | Component | Specification |
291
  |---|---|
292
- | Backbone | WorldDiT diffusion transformer |
293
- | Observation context | 3 observation steps |
294
  | Action horizon | 7 actions |
295
  | Action dimension | 7 |
296
  | Action aggregation | Temporal ensembling |
@@ -299,38 +267,21 @@ Evaluation uses the final temporal slot of the predicted action tensor.
299
  | Evaluation | Headless LIBERO with EGL |
300
  | Checkpoint format | SafeTensors |
301
 
302
- </details>
303
 
304
- ## Use and scope
305
-
306
- | Intended use | Scope of the release |
307
- |---|---|
308
- | 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. |
309
- | 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. |
310
 
311
- ## Citation
312
-
313
- If you use WorldDiT in your research, please cite the paper.
314
-
315
- ```bibtex
316
- @article{260723909,
317
- title={{WorldDiT: A Unified Diffusion Architecture for World and Action Modeling}},
318
- author={Sen Wang and R. Gnana Praveen and Bidhan Roy and Marcos Villagra},
319
- year={{2026}},
320
- eprint={{2607.23909}},
321
- archivePrefix={{arXiv}}
322
- }
323
- ```
324
 
325
- ## License
 
 
326
 
327
- The WorldDiT checkpoints, model card, and original release materials are
328
- licensed under [Creative Commons Attribution 4.0
329
- International](https://creativecommons.org/licenses/by/4.0/). You may copy,
330
- redistribute, and adapt them, including commercially, with appropriate credit
331
- to Bagel Labs and the WorldDiT authors, a link to the license, and an indication
332
- of any changes. Third party dependencies and assets remain governed by their
333
- upstream licenses.
334
 
335
  ## Authors and contact
336
 
 
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."
 
34
  url: "https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_10_frontview_task06_episode01.mp4"
35
  ---
36
 
 
 
37
  <p align="center">
38
  <img src="https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/paris2/model-card/v1/bagel_labs_logo.png" alt="Bagel Labs">
39
  </p>
40
 
41
+ # WorldDiT
42
+
43
+ ## One diffusion backbone learns what to do and what comes next.
44
 
45
  <p align="center">
46
+ <a href="https://huggingface.co/bageldotcom/worlddit" target="_blank">
47
+ <img src="https://img.shields.io/badge/🤗_DOWNLOAD_WORLDDIT_WEIGHTS-FFD21E?style=for-the-badge&logoColor=000000" alt="Download WorldDiT Weights">
48
+ </a>
49
+ <a href="https://github.com/Lifelong-Robot-Learning/LIBERO" target="_blank">
50
+ <img src="https://img.shields.io/badge/🤖_LIBERO_BENCHMARK-FF6B6B?style=for-the-badge&logoColor=white" alt="LIBERO Benchmark">
51
  </a>
52
  </p>
53
 
54
+ WorldDiT learns continuous robot action chunks and a future visual target
55
+ through one shared diffusion transformer. Deployment keeps only the action
56
+ path.
57
+
58
+ This release includes four LIBERO checkpoints, a self contained inference
59
+ runtime, and an evaluator for reproducing the reported suite results.
 
60
 
61
  ## See WorldDiT act
62
 
63
  The four clips below show successful rollouts from the released checkpoints.
64
  Each clip covers a different LIBERO suite and camera view.
65
 
66
+ <Gallery />
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
67
 
68
+ | Suite | View | Video |
69
  |---|---|---|
70
+ | LIBERO Spatial | Front view | [Open MP4](https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_spatial_frontview_task05_episode01.mp4) |
71
+ | LIBERO Object | Agent view | [Open MP4](https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_object_agentview_task08_episode01.mp4) |
72
+ | LIBERO Goal | Side view | [Open MP4](https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_goal_sideview_task10_episode01.mp4) |
73
+ | LIBERO Long | Front view | [Open MP4](https://pub-2c09ae97630f4932a23e622b450076e0.r2.dev/worlddit/model-card/v1/worlddit_libero_10_frontview_task06_episode01.mp4) |
74
 
75
+ ## What is in this release
 
76
 
77
+ | Release component | Included artifact |
78
+ |---|---|
79
+ | LIBERO Spatial policy | SafeTensors checkpoint |
80
+ | LIBERO Object policy | SafeTensors checkpoint |
81
+ | LIBERO Goal policy | SafeTensors checkpoint |
82
+ | LIBERO Long policy | SafeTensors checkpoint |
83
+ | Model runtime | `inference.py` |
84
+ | Evaluation runtime | `eval.py` |
85
+ | Frozen encoders | CLIP ViT B 32 and MAE ViT B |
86
+ | Configuration | `config.json` |
87
+ | Environment | Pinned Python requirements |
88
+
89
+ The repository is self contained for WorldDiT inference. LIBERO still provides
90
+ the benchmark environments, assets, task definitions, and initial states.
91
+
92
+ ## Reported LIBERO results
93
+
94
+ Across the four released suite checkpoints, WorldDiT records 1,898 successful
95
+ episodes out of 2,000 under the selection aware evaluation protocol.
96
+
97
+ | Suite | Successful episodes | Success rate |
98
+ |---|---:|---:|
99
+ | LIBERO Spatial | 490 of 500 | 98.0 percent |
100
+ | LIBERO Object | 485 of 500 | 97.0 percent |
101
+ | LIBERO Goal | 464 of 500 | 92.8 percent |
102
+ | LIBERO Long | 459 of 500 | 91.8 percent |
103
+ | Selection aware mean | 1,898 of 2,000 | 94.9 percent |
104
 
105
+ The released runtime and checkpoints were revalidated from a clean installation
106
+ on eight RTX Pro 6000 Blackwell GPUs.
107
 
108
+ The result is selection aware because three hundred episodes per suite informed
109
+ staged checkpoint selection before the final five hundred episode score was
110
+ assembled.
111
 
112
+ ## Model at a glance
 
 
113
 
114
+ | Property | Released configuration |
115
+ |---|---|
116
+ | Total parameters | 399.084 million |
117
+ | Trainable parameters | 135.107 million |
118
+ | Observation context | Three frames |
119
+ | Predicted action horizon | Seven actions |
120
+ | Executed before replanning | Three actions |
121
+ | Action dimension | Seven |
122
+ | Visual encoder | MAE ViT B |
123
+ | Language encoder | OpenAI CLIP ViT B 32 |
124
+ | Checkpoint format | SafeTensors |
125
+ | Evaluation environment | Headless LIBERO with EGL |
126
 
127
  ## Run a smoke test
128
 
 
159
  ```
160
 
161
  A successful smoke test confirms that the environment, checkpoint, visual
162
+ encoders, simulator, and rendering path load together. It is not a benchmark
163
+ result.
164
 
165
  ## How WorldDiT works
166
 
167
+ WorldDiT uses three recent observations, robot state, and language as context.
168
+ During training, one diffusion transformer learns a seven step action chunk and
169
+ an auxiliary future visual target. During deployment, the future visual path is
170
+ absent. The policy executes the first three predicted actions, observes again,
171
+ and replans.
 
 
 
 
 
 
 
 
172
 
173
+ > Future visual prediction is a training signal, not a deployment path.
 
 
174
 
175
  | Training | Deployment |
176
  |---|---|
177
+ | Action and future visual targets share one backbone | Only the action path remains |
178
  | Seven action steps are supervised | Seven actions are predicted |
179
+ | Future visual supervision is present | No future visual output is requested |
180
  | The complete training objective is active | Three actions execute before replanning |
181
 
182
+ ## Evaluation
 
 
 
183
 
184
  ### One GPU
185
 
 
201
 
202
  Each GPU receives an independent progress bar. After all workers finish, rank 0
203
  prints per task and overall success rates and writes a structured
204
+ `results.json`. Output directories must be new so an earlier evaluation is
205
+ never overwritten.
206
 
207
  Supported suites.
208
 
 
213
  libero_10
214
  ```
215
 
216
+ ## What this repository contains
 
 
 
217
 
218
  ```text
219
  .
 
232
  ```
233
 
234
  `dependencies/` contains the frozen visual and language encoder weights needed
235
+ by the released policy. No additional model downloads are required.
 
236
 
237
+ ## Inference API
 
 
 
238
 
239
  ```python
240
  from inference import load_model
 
253
 
254
  Evaluation uses the final temporal slot of the predicted action tensor.
255
 
256
+ ## Architecture details
 
 
 
257
 
258
  | Component | Specification |
259
  |---|---|
260
+ | Policy | WorldDiT diffusion transformer |
261
+ | Observation context | 3 frames |
262
  | Action horizon | 7 actions |
263
  | Action dimension | 7 |
264
  | Action aggregation | Temporal ensembling |
 
267
  | Evaluation | Headless LIBERO with EGL |
268
  | Checkpoint format | SafeTensors |
269
 
270
+ ## Intended use
271
 
272
+ WorldDiT is intended for research on language conditioned robot manipulation in
273
+ the LIBERO simulator. The released checkpoints support reproduction,
274
+ evaluation, and architecture research across the four released suites.
 
 
 
275
 
276
+ ## Scope of the release
 
 
 
 
 
 
 
 
 
 
 
 
277
 
278
+ The reported results describe LIBERO simulation under the released evaluation
279
+ protocol. They do not establish real robot reliability, safety, or transfer
280
+ across embodiments.
281
 
282
+ The present release does not isolate the causal contribution of the future
283
+ visual target. Total parameter count also does not measure training cost,
284
+ deployment latency, or runtime efficiency.
 
 
 
 
285
 
286
  ## Authors and contact
287