Improve the WorldDiT model card layout

#2
by bagellabs - opened
Files changed (1) hide show
  1. README.md +31 -83
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: A Unified Diffusion Backbone for<br>World and Action Modeling</h1>
 
 
44
 
45
  <p align="center">
46
- <a href="https://arxiv.org/abs/2607.23909">
47
- <img src="https://img.shields.io/badge/arXiv-2607.23909-B31B1B.svg?logo=arxiv&logoColor=white" alt="WorldDiT paper on arXiv">
48
- </a>
49
  </p>
50
 
51
  <p align="center">
52
- WorldDiT couples continuous action generation with auxiliary future normalized
53
- RGB patch prediction in one diffusion transformer. The architecture is designed
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** mean success<br>**1,898 of 2,000** successful episodes | **399.084 million** total parameters<br>**135.107 million** trainable parameters |
105
- | **98.0 percent** Spatial<br>**97.0 percent** Object | **Three** observation steps<br>**Seven** predicted actions |
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 reported aggregate covers five
117
- hundred simulator episodes per suite. Three hundred episodes per suite informed
118
- staged checkpoint selection, while two hundred episodes per suite were disjoint
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. Full benchmark reporting
170
- uses complete suite evaluations.
171
 
172
  ## How WorldDiT works
173
 
174
- Each of three recent observation steps contributes primary and wrist images
175
- together with robot state, while one language instruction conditions the
176
- sequence. During training, one diffusion transformer backbone learns a seven
177
- step action chunk together with an auxiliary future normalized RGB patch
178
- target. At deployment, the encoded history conditions the action velocity field
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
- <p align="center">
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 normalized RGB patch targets are learned by one backbone | Encoded history conditions the action velocity field |
195
  | Seven action steps are supervised | Seven actions are predicted |
196
- | Future normalized RGB patch supervision is present | RGB patch tokens and the RGB prediction head remain outside the inference graph |
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`. Use a new output directory for each evaluation to preserve
225
- earlier results.
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 WorldDiT runtime. The repository contains every model weight
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
- | Backbone | WorldDiT diffusion transformer |
291
- | Observation context | 3 observation steps |
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 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. |
307
- | 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. |
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