Release WorldDiT model and documentation

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+replay.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,354 @@
+---
+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>

checkpoints/libero_10/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c0fa3dd7597949d53dd98bdaec28a2267459be6bd1bebc6d9b2c0d82095529d8
+size 540446188

checkpoints/libero_goal/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5da34e96382fe831437459e29608d9db57356e2d1b9149cb20888cda72e6bdac
+size 540446196

checkpoints/libero_object/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c7d50e2038d7e7aea87268770007fc4f6a639158002b57ab573c32fb218b97d
+size 540446196

checkpoints/libero_spatial/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d2ddd556c1dc886dbc8927c6e53ad7c2df60dc6ff14b164753ca77cecbec9707
+size 540446196

config.json

@@ -0,0 +1,11 @@
+{
+  "model_type": "worlddit",
+  "architectures": ["WorldDiTPolicy"],
+  "context_frames": 3,
+  "action_horizon": 7,
+  "suites": ["libero_10", "libero_spatial", "libero_goal", "libero_object"],
+  "runtime": {
+    "model": "inference.py",
+    "evaluation": "eval.py"
+  }
+}

dependencies/ViT-B-32.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af
+size 353976522

dependencies/mae_pretrain_vit_base.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aec5f0b68e5f3193a00b07bc65a37440db549c15b36b8bea242606cc40c4bc5d
+size 343249461

eval.py

@@ -0,0 +1,340 @@
+#!/usr/bin/env python3
+"""Evaluate a released WorldDiT checkpoint on LIBERO."""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import contextlib
+import json
+import math
+import os
+import random
+import subprocess
+import sys
+from collections import deque
+from pathlib import Path
+
+import clip
+import numpy as np
+import torch
+from PIL import Image
+from scipy.spatial.transform import Rotation
+from torch.nn.parallel import DistributedDataParallel as DDP
+from tqdm.auto import tqdm
+
+from inference import ACTION_HORIZON, CONTEXT_STEPS, SUITES, load_model
+
+
+def rank():
+    return int(os.environ.get("RANK", "0"))
+
+
+def world_size():
+    return int(os.environ.get("WORLD_SIZE", "1"))
+
+
+def launch(gpus: int):
+    if not 1 <= gpus <= torch.cuda.device_count():
+        raise ValueError(f"--gpus must be between 1 and {torch.cuda.device_count()}")
+    command = [
+        sys.executable,
+        "-m",
+        "torch.distributed.run",
+        "--standalone",
+        f"--nproc_per_node={gpus}",
+        str(Path(__file__).resolve()),
+        *sys.argv[1:],
+    ]
+    environment = os.environ.copy()
+    environment.setdefault("OMP_NUM_THREADS", "1")
+    return subprocess.run(command, check=False, env=environment).returncode
+
+
+def configure_libero(libero_root: Path, output: Path):
+    package = libero_root / "libero" / "libero"
+    required = (package / "bddl_files", package / "init_files", package / "assets")
+    if any(not path.is_dir() for path in required):
+        raise FileNotFoundError(f"invalid LIBERO checkout: {libero_root}")
+    config_dir = output / "libero_config"
+    if rank() == 0:
+        output.mkdir(parents=True, exist_ok=False)
+        config_dir.mkdir()
+        paths = {
+            "assets": str(package / "assets"),
+            "bddl_files": str(package / "bddl_files"),
+            "benchmark_root": str(package),
+            "datasets": str(package.parent / "datasets"),
+            "init_states": str(package / "init_files"),
+        }
+        (config_dir / "config.yaml").write_text(json.dumps(paths), encoding="utf-8")
+    torch.distributed.barrier()
+    os.environ["LIBERO_CONFIG_PATH"] = str(config_dir)
+    sys.path.insert(0, str(libero_root))
+
+
+def orientation(quaternion):
+    quaternion = np.asarray(quaternion, dtype=np.float64).copy()
+    quaternion[3] = np.clip(quaternion[3], -1.0, 1.0)
+    denominator = np.sqrt(1.0 - quaternion[3] ** 2)
+    axisangle = (
+        np.zeros(3)
+        if math.isclose(denominator, 0.0)
+        else quaternion[:3] * 2.0 * math.acos(quaternion[3]) / denominator
+    )
+    return Rotation.from_euler("xyz", axisangle).as_euler("xyz")
+
+
+def finish_action(action: torch.Tensor):
+    action = action.detach().cpu().numpy().copy()
+    action[-1] = 1.0 if action[-1] > 0.5 else -1.0
+    return action
+
+
+class PolicyRunner:
+    def __init__(
+        self, model, temperature: float, execution_horizon: int, max_steps: int
+    ):
+        self.model = model
+        self.processor = getattr(model, "module", model).image_processor
+        self.temperature = temperature
+        self.execution_horizon = execution_horizon
+        self.max_steps = max_steps
+
+    def reset(self):
+        self.primary = deque(maxlen=CONTEXT_STEPS)
+        self.wrist = deque(maxlen=CONTEXT_STEPS)
+        self.state = deque(maxlen=CONTEXT_STEPS)
+        self.pending = deque()
+        self.predictions = torch.zeros(
+            self.max_steps, self.max_steps + ACTION_HORIZON, 7, device="cuda"
+        )
+        self.valid = torch.zeros(
+            self.max_steps,
+            self.max_steps + ACTION_HORIZON,
+            dtype=torch.bool,
+            device="cuda",
+        )
+
+    def observe(self, observation):
+        primary = Image.fromarray(observation["agentview_image"][::-1])
+        wrist = Image.fromarray(observation["robot0_eye_in_hand_image"])
+        self.primary.append(self.processor(primary).unsqueeze(0).unsqueeze(0))
+        self.wrist.append(self.processor(wrist).unsqueeze(0).unsqueeze(0))
+        state = np.concatenate(
+            (
+                observation["robot0_eef_pos"],
+                orientation(observation["robot0_eef_quat"]),
+                observation["robot0_gripper_qpos"],
+            )
+        )
+        self.state.append(torch.from_numpy(state).float().view(1, 1, -1))
+
+    def prefill(self, observations):
+        for observation in observations[-CONTEXT_STEPS:-1]:
+            self.observe(observation)
+
+    def ensemble(self, chunk: torch.Tensor, timestep: int):
+        self.predictions[timestep, timestep : timestep + ACTION_HORIZON] = chunk
+        self.valid[timestep, timestep : timestep + ACTION_HORIZON] = True
+        selected = []
+        for target in range(timestep, timestep + self.execution_horizon):
+            mask = self.valid[: timestep + 1, target]
+            actions = self.predictions[: timestep + 1, target][mask]
+            weights = np.exp(-self.temperature * np.arange(len(actions)))
+            weights = torch.as_tensor(weights / weights.sum(), device="cuda").unsqueeze(
+                1
+            )
+            selected.append(finish_action((actions * weights).sum(0)))
+        self.pending.extend(selected[1:])
+        return selected[0]
+
+    @torch.inference_mode()
+    def act(self, observation, instruction: str, timestep: int):
+        self.observe(observation)
+        if self.pending:
+            return self.pending.popleft()
+        if len(self.primary) != CONTEXT_STEPS:
+            raise RuntimeError("evaluation requires three real context observations")
+        primary = torch.cat(tuple(self.primary), dim=1).cuda()
+        wrist = torch.cat(tuple(self.wrist), dim=1).cuda()
+        state = torch.cat(tuple(self.state), dim=1).cuda()
+        text = (
+            clip.tokenize([instruction] * CONTEXT_STEPS, truncate=True)
+            .view(1, CONTEXT_STEPS, -1)
+            .cuda()
+        )
+        chunk = self.model(primary, wrist, state, text)[0, -1]
+        return self.ensemble(chunk, timestep)
+
+
+def evaluate(args, model):
+    with (
+        open(os.devnull, "w") as quiet,
+        contextlib.redirect_stdout(quiet),
+        contextlib.redirect_stderr(quiet),
+    ):
+        from libero.libero import benchmark
+        from libero.libero.envs import OffScreenRenderEnv
+
+        suite = benchmark.get_benchmark_dict()[args.suite]()
+    runner = PolicyRunner(
+        model, args.temperature, args.execution_horizon, args.max_steps
+    )
+    total = args.tasks * args.episodes
+    assigned = list(range(total))[rank() :: world_size()]
+    local_results = []
+    progress = tqdm(
+        assigned,
+        desc=f"GPU {rank()}",
+        position=rank(),
+        dynamic_ncols=True,
+        leave=True,
+    )
+    for evaluation_id in progress:
+        task_id, episode_index = divmod(evaluation_id, args.episodes)
+        episode_id = args.episode_offset + episode_index
+        task = suite.get_task(task_id)
+        bddl = (
+            args.libero_path
+            / "libero"
+            / "libero"
+            / "bddl_files"
+            / task.problem_folder
+            / task.bddl_file
+        )
+        environment = OffScreenRenderEnv(
+            bddl_file_name=str(bddl),
+            camera_heights=128,
+            camera_widths=128,
+            render_gpu_device_id=int(os.environ["LOCAL_RANK"]),
+        )
+        try:
+            environment.reset()
+            environment.seed(66)
+            initial_states = torch.load(
+                args.libero_path
+                / "libero"
+                / "libero"
+                / "init_files"
+                / task.problem_folder
+                / task.init_states_file,
+                weights_only=False,
+            )
+            if episode_id >= len(initial_states):
+                raise IndexError(
+                    f"episode {episode_id} is unavailable for task {task_id}"
+                )
+            observation = environment.set_init_state(initial_states[episode_id])
+            warmup = []
+            for _ in range(5):
+                observation, _, _, _ = environment.step(np.zeros(7))
+                warmup.append(copy.deepcopy(observation))
+            runner.reset()
+            runner.prefill(warmup)
+            observation = warmup[-1]
+            success = 0
+            for steps in range(1, args.max_steps + 1):
+                action = runner.act(observation, task.language, steps - 1)
+                observation, _, done, _ = environment.step(action)
+                if done:
+                    success = 1
+                    break
+            local_results.append((evaluation_id, task_id, episode_id, success, steps))
+            progress.set_postfix(successes=sum(item[3] for item in local_results))
+        finally:
+            environment.close()
+
+    gathered = [None] * world_size() if rank() == 0 else None
+    torch.distributed.gather_object(local_results, gathered, dst=0)
+    if rank() != 0:
+        return
+    results = sorted((item for group in gathered for item in group), key=lambda x: x[0])
+    per_task = []
+    print()
+    for task_id in range(args.tasks):
+        values = [item[3] for item in results if item[1] == task_id]
+        rate = float(np.mean(values))
+        per_task.append(rate)
+        print(f"Task {task_id}: {sum(values)}/{len(values)} ({rate:.1%})")
+    successes = sum(item[3] for item in results)
+    print(f"Overall: {successes}/{len(results)} ({successes / len(results):.1%})")
+    report = {
+        "suite": args.suite,
+        "gpus": world_size(),
+        "episodes": len(results),
+        "successes": successes,
+        "success_rate": successes / len(results),
+        "per_task_success_rate": per_task,
+        "results": [
+            {"task": item[1], "episode": item[2], "success": item[3], "steps": item[4]}
+            for item in results
+        ],
+    }
+    (args.output_dir / "results.json").write_text(
+        json.dumps(report, indent=2) + "\n", encoding="utf-8"
+    )
+
+
+def parser():
+    value = argparse.ArgumentParser(description=__doc__)
+    value.add_argument("--suite", required=True, choices=SUITES)
+    value.add_argument("--gpus", type=int, default=1)
+    value.add_argument(
+        "--model-root", type=Path, default=Path(__file__).resolve().parent
+    )
+    value.add_argument(
+        "--libero-path", type=Path, default=Path("~/LIBERO").expanduser()
+    )
+    value.add_argument("--output-dir", type=Path, required=True)
+    value.add_argument("--tasks", type=int, default=10)
+    value.add_argument("--episodes", type=int, default=50)
+    value.add_argument("--episode-offset", type=int, default=0)
+    value.add_argument("--max-steps", type=int, default=600)
+    value.add_argument("--execution-horizon", type=int, choices=(1, 3), default=3)
+    value.add_argument("--temperature", type=float, default=0.01)
+    return value
+
+
+def main():
+    args = parser().parse_args()
+    if "RANK" not in os.environ:
+        return launch(args.gpus)
+    if args.gpus != world_size():
+        raise ValueError(
+            f"--gpus={args.gpus} but torchrun started {world_size()} workers"
+        )
+    args.model_root = args.model_root.expanduser().resolve()
+    args.libero_path = args.libero_path.expanduser().resolve()
+    args.output_dir = args.output_dir.expanduser().resolve()
+    os.environ.update(MUJOCO_GL="egl", PYOPENGL_PLATFORM="egl")
+    os.environ.setdefault("NCCL_IB_DISABLE", "1")
+    os.environ.setdefault("NCCL_P2P_DISABLE", "1")
+    os.environ.setdefault("NCCL_CUMEM_ENABLE", "0")
+    os.environ.setdefault("TORCH_NCCL_BLOCKING_WAIT", "1")
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+    torch.distributed.init_process_group(
+        backend="nccl", device_id=torch.device("cuda", int(os.environ["LOCAL_RANK"]))
+    )
+    try:
+        configure_libero(args.libero_path, args.output_dir)
+        model = load_model(args.model_root, args.suite, torch.device("cuda"))
+        seed = 66 + rank()
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        model = DDP(
+            model,
+            device_ids=[int(os.environ["LOCAL_RANK"])],
+            find_unused_parameters=True,
+        )
+        model.eval()
+        evaluate(args, model)
+        torch.distributed.barrier()
+    finally:
+        torch.distributed.destroy_process_group()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

inference.py

@@ -0,0 +1,431 @@
+"""Compact inference-only WorldDiT runtime for the released LIBERO checkpoints."""
+
+from __future__ import annotations
+
+import math
+from functools import partial
+from pathlib import Path
+
+import clip
+import torch
+from einops import rearrange, repeat
+from einops_exts import rearrange_many
+from safetensors import safe_open
+from timm.models.vision_transformer import Block, PatchEmbed
+from torch import einsum, nn
+
+SUITES = ("libero_10", "libero_spatial", "libero_goal", "libero_object")
+CONTEXT_STEPS = 3
+ACTION_HORIZON = 7
+HIDDEN_DIM = 1024
+
+
+class VisionEncoder(nn.Module):
+    """The encoder half of the frozen MAE dependency."""
+
+    def __init__(self):
+        super().__init__()
+        self.patch_embed = PatchEmbed(224, 16, 3, 768)
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, 768))
+        self.pos_embed = nn.Parameter(torch.zeros(1, 197, 768), requires_grad=False)
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    768,
+                    12,
+                    4,
+                    qkv_bias=True,
+                    norm_layer=partial(nn.LayerNorm, eps=1e-6),
+                )
+                for _ in range(12)
+            ]
+        )
+        self.norm = nn.LayerNorm(768, eps=1e-6)
+
+    def forward(self, images: torch.Tensor) -> torch.Tensor:
+        patches = self.patch_embed(images) + self.pos_embed[:, 1:]
+        # Retain the released encoder's mask_ratio=0 behavior, including RNG use.
+        order = torch.rand(patches.shape[:2], device=patches.device).argsort(dim=1)
+        patches = torch.gather(patches, 1, order.unsqueeze(-1).expand_as(patches))
+        cls = (self.cls_token + self.pos_embed[:, :1]).expand(images.shape[0], -1, -1)
+        tokens = torch.cat((cls, patches), dim=1)
+        for block in self.blocks:
+            tokens = block(tokens)
+        return self.norm(tokens)
+
+
+class PerceiverAttention(nn.Module):
+    def __init__(self, dim: int = 768, dim_head: int = 64, heads: int = 8):
+        super().__init__()
+        self.scale = dim_head**-0.5
+        self.heads = heads
+        inner = dim_head * heads
+        self.norm_media = nn.LayerNorm(dim)
+        self.norm_latents = nn.LayerNorm(dim)
+        self.to_q = nn.Linear(dim, inner, bias=False)
+        self.to_kv = nn.Linear(dim, inner * 2, bias=False)
+        self.to_out = nn.Linear(inner, dim, bias=False)
+
+    def forward(self, media: torch.Tensor, latents: torch.Tensor) -> torch.Tensor:
+        media, latents = self.norm_media(media), self.norm_latents(latents)
+        query = self.to_q(latents)
+        key, value = self.to_kv(torch.cat((media, latents), dim=-2)).chunk(2, dim=-1)
+        query, key, value = rearrange_many(
+            (query, key, value), "b t n (h d) -> b h t n d", h=self.heads
+        )
+        scores = einsum("... i d, ... j d -> ... i j", query * self.scale, key)
+        weights = (scores - scores.amax(dim=-1, keepdim=True).detach()).softmax(-1)
+        output = einsum("... i j, ... j d -> ... i d", weights, value)
+        return self.to_out(rearrange(output, "b h t n d -> b t n (h d)"))
+
+
+class PerceiverResampler(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.latents = nn.Parameter(torch.randn(16, 768))
+        self.layers = nn.ModuleList()
+        for _ in range(3):
+            feed_forward = nn.Sequential(
+                nn.LayerNorm(768),
+                nn.Linear(768, 3072, bias=False),
+                nn.GELU(),
+                nn.Linear(3072, 768, bias=False),
+            )
+            self.layers.append(nn.ModuleList((PerceiverAttention(), feed_forward)))
+        self.norm = nn.LayerNorm(768)
+
+    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
+        batch, steps = tokens.shape[:2]
+        tokens = rearrange(tokens, "b t f v d -> b t (f v) d")
+        latents = repeat(self.latents, "n d -> b t n d", b=batch, t=steps)
+        for attention, feed_forward in self.layers:
+            latents = attention(tokens, latents) + latents
+            latents = feed_forward(latents) + latents
+        return self.norm(latents)
+
+
+def _time_embedding(timestep: torch.Tensor, dim: int) -> torch.Tensor:
+    half = dim // 2
+    frequency = torch.exp(
+        -math.log(10000.0)
+        * torch.arange(half, device=timestep.device, dtype=timestep.dtype)
+        / max(half - 1, 1)
+    )
+    phase = timestep.unsqueeze(-1) * frequency.unsqueeze(0) * 1000.0
+    embedding = torch.cat((torch.sin(phase), torch.cos(phase)), dim=-1)
+    return embedding if dim % 2 == 0 else torch.nn.functional.pad(embedding, (0, 1))
+
+
+class MLP(nn.Module):
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, output_dim),
+        )
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        return self.net(value)
+
+
+def _modulate(value: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor):
+    return value * (1.0 + scale) + shift
+
+
+class DFDiTBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(HIDDEN_DIM, elementwise_affine=False)
+        self.attn = nn.MultiheadAttention(HIDDEN_DIM, 16, batch_first=True)
+        self.norm2 = nn.LayerNorm(HIDDEN_DIM, elementwise_affine=False)
+        self.mlp = nn.Sequential(
+            nn.Linear(HIDDEN_DIM, HIDDEN_DIM * 4),
+            nn.GELU(approximate="tanh"),
+            nn.Dropout(0.0),
+            nn.Linear(HIDDEN_DIM * 4, HIDDEN_DIM),
+        )
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(HIDDEN_DIM, 6 * HIDDEN_DIM)
+        )
+
+    def forward(self, tokens, modulation, mask):
+        shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = self.adaLN_modulation(
+            modulation
+        ).chunk(6, dim=-1)
+        attention_input = _modulate(self.norm1(tokens), shift_a, scale_a)
+        attention = self.attn(
+            attention_input,
+            attention_input,
+            attention_input,
+            attn_mask=mask,
+            need_weights=False,
+        )[0]
+        tokens = tokens + gate_a * attention
+        return tokens + gate_m * self.mlp(
+            _modulate(self.norm2(tokens), shift_m, scale_m)
+        )
+
+
+class ActionSampler(nn.Module):
+    """Action sampler used by the released checkpoints."""
+
+    def __init__(self):
+        super().__init__()
+        self.context_norm = nn.LayerNorm(HIDDEN_DIM)
+        self.context_proj = nn.Linear(HIDDEN_DIM, HIDDEN_DIM)
+        self.action_tokenizer = MLP(7, HIDDEN_DIM * 4, HIDDEN_DIM)
+        self.action_decoder = MLP(HIDDEN_DIM, HIDDEN_DIM * 4, 7)
+        self.action_pos = nn.Parameter(
+            torch.randn(1, ACTION_HORIZON, HIDDEN_DIM) * 0.02
+        )
+        self.step_pos = nn.Parameter(
+            torch.randn(1, CONTEXT_STEPS, 1, HIDDEN_DIM) * 0.02
+        )
+        self.context_type = nn.Parameter(torch.randn(1, 1, HIDDEN_DIM) * 0.02)
+        self.action_type = nn.Parameter(torch.randn(1, 1, HIDDEN_DIM) * 0.02)
+        self.register_type = nn.Parameter(torch.randn(1, 1, HIDDEN_DIM) * 0.02)
+        self.register_tokens = nn.Parameter(torch.randn(1, 4, HIDDEN_DIM) * 0.02)
+        self.time_proj = nn.Sequential(
+            nn.Linear(HIDDEN_DIM, HIDDEN_DIM * 4),
+            nn.SiLU(),
+            nn.Linear(HIDDEN_DIM * 4, HIDDEN_DIM),
+        )
+        self.blocks = nn.ModuleList([DFDiTBlock() for _ in range(4)])
+        self.final_norm = nn.LayerNorm(HIDDEN_DIM)
+
+    def _times(self, timestep: torch.Tensor, count: int, dtype: torch.dtype):
+        embedded = self.time_proj(_time_embedding(timestep, HIDDEN_DIM).to(dtype))
+        return embedded.unsqueeze(1).expand(-1, count, -1)
+
+    @staticmethod
+    def _mask(steps: int, context_count: int, dtype, device):
+        targets, registers = ACTION_HORIZON, 4
+        block = context_count + targets + registers
+        mask = torch.full(
+            (steps * block, steps * block), -torch.inf, dtype=dtype, device=device
+        )
+        for step in range(steps):
+            start = step * block
+            context = slice(start, start + context_count)
+            action = slice(start + context_count, start + context_count + targets)
+            register = slice(start + context_count + targets, start + block)
+            for source_step in range(step + 1):
+                source = source_step * block
+                visible_context = slice(source, source + context_count)
+                mask[context, visible_context] = 0
+                mask[action, visible_context] = 0
+                mask[register, visible_context] = 0
+            mask[action, action] = 0
+            mask[action, register] = 0
+            mask[register, action] = 0
+            mask[register, register] = 0
+        return mask
+
+    def _velocity(self, context: torch.Tensor, noisy_action: torch.Tensor, timestep):
+        batch, steps, context_count, _ = context.shape
+        flat_batch = batch * steps
+        dtype, device = context.dtype, context.device
+        encoded_context = self.context_proj(self.context_norm(context))
+        encoded_context = encoded_context + self.context_type.to(dtype)
+        encoded_context = encoded_context + self.step_pos[:, :steps].to(dtype).expand(
+            -1, -1, context_count, -1
+        )
+        zero_time = torch.zeros(flat_batch, dtype=dtype, device=device)
+        context_mod = self._times(zero_time, context_count, dtype).view(
+            batch, steps, context_count, HIDDEN_DIM
+        ) + self.context_type.to(dtype)
+
+        flat_action = noisy_action.reshape(flat_batch, ACTION_HORIZON, 7)
+        action_time = self._times(timestep, ACTION_HORIZON, dtype)
+        action = self.action_tokenizer(flat_action.to(dtype))
+        action = (
+            action
+            + self.action_pos.to(dtype)
+            + self.action_type.to(dtype)
+            + action_time
+        )
+        action = action.view(batch, steps, ACTION_HORIZON, HIDDEN_DIM)
+        action = action + self.step_pos[:, :steps].to(dtype).expand(
+            -1, -1, ACTION_HORIZON, -1
+        )
+        action_mod = (action_time + self.action_type.to(dtype)).view(
+            batch, steps, ACTION_HORIZON, HIDDEN_DIM
+        )
+
+        registers = self.register_tokens.to(dtype).expand(batch, steps, -1, -1)
+        registers = registers + self.register_type.to(dtype)
+        registers = registers + self.step_pos[:, :steps].to(dtype).expand(-1, -1, 4, -1)
+        register_mod = self._times(zero_time, 4, dtype).view(
+            batch, steps, 4, HIDDEN_DIM
+        ) + self.register_type.to(dtype)
+
+        tokens = torch.cat((encoded_context, action, registers), dim=2).flatten(1, 2)
+        modulation = torch.cat((context_mod, action_mod, register_mod), dim=2).flatten(
+            1, 2
+        )
+        mask = self._mask(steps, context_count, tokens.dtype, device)
+        for block in self.blocks:
+            tokens = block(tokens, modulation, mask)
+        tokens = self.final_norm(tokens).view(batch, steps, -1, HIDDEN_DIM)
+        return self.action_decoder(
+            tokens[:, :, context_count : context_count + ACTION_HORIZON]
+        )
+
+    def forward(self, context: torch.Tensor, sampling_steps: int = 20):
+        batch, steps = context.shape[:2]
+        dtype, device = self.context_norm.weight.dtype, context.device
+        context = context.to(dtype)
+        action = torch.randn(
+            batch, steps, ACTION_HORIZON, 7, dtype=dtype, device=device
+        )
+        for time in torch.linspace(0.0, 1.0, sampling_steps + 1, device=device)[:-1]:
+            timestep = torch.full(
+                (batch * steps,), float(time.item()), dtype=dtype, device=device
+            )
+            action = action + self._velocity(context, action, timestep) / sampling_steps
+        return action
+
+
+class WorldDiTPolicy(nn.Module):
+    """Released policy with only modules used by inference."""
+
+    def __init__(self, mae_path: Path, clip_path: Path):
+        super().__init__()
+        self.text_projector = nn.Linear(512, HIDDEN_DIM)
+        self.arm_state_encoder = nn.Linear(6, HIDDEN_DIM)
+        self.gripper_state_encoder = nn.Linear(2, HIDDEN_DIM)
+        self.state_projector = nn.Linear(HIDDEN_DIM * 2, HIDDEN_DIM)
+        self.vision_encoder = VisionEncoder()
+        self.perceiver_resampler = PerceiverResampler()
+        self.image_primary_projector = nn.Linear(768, HIDDEN_DIM)
+        self.cls_token_primary_projector = nn.Linear(768, HIDDEN_DIM)
+        self.image_wrist_projector = nn.Linear(768, HIDDEN_DIM)
+        self.cls_token_wrist_projector = nn.Linear(768, HIDDEN_DIM)
+        self.embedding_layer_norm = nn.LayerNorm(HIDDEN_DIM)
+        self.transformer_backbone_position_embedding = nn.Parameter(
+            torch.zeros(1, CONTEXT_STEPS, 1, HIDDEN_DIM)
+        )
+        self.unified_action_world_head = ActionSampler()
+
+        mae = torch.load(mae_path, map_location="cpu", weights_only=False)
+        self.vision_encoder.load_state_dict(mae["model"], strict=False)
+        self.clip_model, self.image_processor = clip.load(str(clip_path), device="cpu")
+        self.vision_encoder.requires_grad_(False)
+        self.clip_model.requires_grad_(False)
+
+    def _images(self, primary: torch.Tensor, wrist: torch.Tensor):
+        batch, steps = primary.shape[:2]
+        vision_dtype = next(self.vision_encoder.parameters()).dtype
+        with torch.no_grad():
+            primary_tokens = self.vision_encoder(primary.flatten(0, 1).to(vision_dtype))
+            wrist_tokens = self.vision_encoder(wrist.flatten(0, 1).to(vision_dtype))
+        cls_primary = primary_tokens[:, :1]
+        cls_wrist = wrist_tokens[:, :1]
+        resampler_dtype = next(self.perceiver_resampler.parameters()).dtype
+        cls_primary = cls_primary.to(resampler_dtype)
+        cls_wrist = cls_wrist.to(resampler_dtype)
+        primary_tokens = primary_tokens[:, 1:].to(resampler_dtype)
+        wrist_tokens = wrist_tokens[:, 1:].to(resampler_dtype)
+        primary_latents = self.perceiver_resampler(
+            primary_tokens.reshape(batch * steps, 196, 768).unsqueeze(1).unsqueeze(1)
+        )
+        wrist_latents = self.perceiver_resampler(
+            wrist_tokens.reshape(batch * steps, 196, 768).unsqueeze(1).unsqueeze(1)
+        )
+        images = torch.cat(
+            (
+                self.image_primary_projector(primary_latents.flatten(0, 2)).view(
+                    batch, steps, 16, HIDDEN_DIM
+                ),
+                self.image_wrist_projector(wrist_latents.flatten(0, 2)).view(
+                    batch, steps, 16, HIDDEN_DIM
+                ),
+            ),
+            dim=2,
+        )
+        cls = torch.cat(
+            (
+                self.cls_token_primary_projector(cls_primary).view(
+                    batch, steps, 1, HIDDEN_DIM
+                ),
+                self.cls_token_wrist_projector(cls_wrist).view(
+                    batch, steps, 1, HIDDEN_DIM
+                ),
+            ),
+            dim=2,
+        )
+        return images, cls
+
+    def forward(self, image_primary, image_wrist, state, text_token):
+        batch, steps = state.shape[:2]
+        if steps != CONTEXT_STEPS:
+            raise ValueError(f"expected {CONTEXT_STEPS} context frames, got {steps}")
+        with torch.no_grad():
+            text = self.clip_model.encode_text(text_token.flatten(0, 1)).type_as(state)
+        text = self.text_projector(text).view(batch, steps, 1, HIDDEN_DIM)
+        flat_state = state.flatten(0, 1)
+        state_token = self.state_projector(
+            torch.cat(
+                (
+                    self.arm_state_encoder(flat_state[:, :6]),
+                    self.gripper_state_encoder(flat_state[:, 6:]),
+                ),
+                dim=1,
+            )
+        ).view(batch, steps, 1, HIDDEN_DIM)
+        images, cls = self._images(image_primary, image_wrist)
+        context = torch.cat((text, state_token, images, cls), dim=2)
+        context = context + self.transformer_backbone_position_embedding[:, :steps].to(
+            context.dtype
+        )
+        context = self.embedding_layer_norm(
+            context.to(self.embedding_layer_norm.weight.dtype)
+        )
+        return self.unified_action_world_head(context, sampling_steps=20)
+
+
+def load_model(
+    model_root: str | Path,
+    suite: str,
+    device: str | torch.device = "cuda",
+    *,
+    vision_bfloat16: bool = True,
+) -> WorldDiTPolicy:
+    """Load one suite checkpoint from a downloaded model-repository directory."""
+    if suite not in SUITES:
+        raise ValueError(f"unknown suite {suite!r}; choose one of {SUITES}")
+    root = Path(model_root).expanduser().resolve()
+    paths = {
+        "checkpoint": root / "checkpoints" / suite / "model.safetensors",
+        "mae": root / "dependencies" / "mae_pretrain_vit_base.pth",
+        "clip": root / "dependencies" / "ViT-B-32.pt",
+    }
+    missing = [str(path) for path in paths.values() if not path.is_file()]
+    if missing:
+        raise FileNotFoundError(f"missing model files: {missing}")
+
+    model = WorldDiTPolicy(paths["mae"], paths["clip"])
+    with safe_open(paths["checkpoint"], framework="pt", device="cpu") as source:
+        released = {
+            name.removeprefix("module."): source.get_tensor(name)
+            for name in source.keys()
+        }
+    current = model.state_dict()
+    retained = {name: value for name, value in released.items() if name in current}
+    expected = {
+        name
+        for name in current
+        if not name.startswith(("vision_encoder.", "clip_model."))
+    }
+    missing_parameters = sorted(expected - retained.keys())
+    if missing_parameters:
+        raise RuntimeError(
+            f"checkpoint is missing inference parameters: {missing_parameters}"
+        )
+    model.load_state_dict(retained, strict=False)
+    model.float()
+    if vision_bfloat16:
+        model.vision_encoder.bfloat16()
+    model.to(torch.device(device)).eval()
+    return model

replay.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2178af4051b31c3faa1d7a590cfd5395ca91ec90f4c2170cd14f4bdcd31050aa
+size 386846

requirements.txt

@@ -0,0 +1,32 @@
+--extra-index-url https://download.pytorch.org/whl/cu128
+
+torch==2.9.1+cu128
+torchvision==0.24.1+cu128
+numpy==1.26.4
+scipy==1.17.1
+pillow==12.3.0
+tqdm==4.68.4
+opencv-python==4.10.0.84
+matplotlib==3.11.1
+omegaconf==2.3.0
+timm==0.9.16
+einops==0.8.2
+einops-exts==0.0.4
+safetensors==0.8.0
+huggingface-hub==0.36.0
+ftfy==6.2.0
+regex==2026.7.10
+packaging==24.0
+git+https://github.com/openai/CLIP.git
+
+# Headless LIBERO / robosuite runtime
+numba==0.66.0
+termcolor==3.3.0
+bddl==3.6.0
+mujoco==3.3.2
+gym==0.26.2
+cloudpickle==3.1.2
+easydict==1.13
+glfw==2.10.0
+PyOpenGL==3.1.10
+PyYAML==6.0.3