Publish WorldDiT LIBERO release

README.md

@@ -0,0 +1,210 @@
+---
+library_name: pytorch
+pipeline_tag: robotics
+tags:
+  - worlddit
+  - libero
+  - robot-learning
+  - imitation-learning
+  - diffusion-policy
+---
+
+<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: Context-3 / Action-7 LIBERO Policies</h1>
+
+<p align="center">
+  <a href="https://huggingface.co/bageldotcom/worlddit" target="_blank">
+    <img src="https://img.shields.io/badge/🤗_DOWNLOAD_WORLDDIT_WEIGHTS-FFD21E?style=for-the-badge&logoColor=000000" alt="Download WorldDiT Weights">
+  </a>
+  <a href="https://github.com/Lifelong-Robot-Learning/LIBERO" target="_blank">
+    <img src="https://img.shields.io/badge/🤖_LIBERO_BENCHMARK-FF6B6B?style=for-the-badge&logoColor=white" alt="LIBERO Benchmark">
+  </a>
+</p>
+
+WorldDiT is a diffusion-transformer policy for language-conditioned robotic
+manipulation. This public release provides suite-specific checkpoints for all
+four LIBERO benchmark suites together with a compact, self-contained inference
+and evaluation runtime.
+
+The runtime is intentionally minimal: `inference.py` constructs the policy and
+loads a checkpoint, while `eval.py` performs headless single- or multi-GPU
+LIBERO evaluation. The research training code is not required.
+
+# Results
+
+All results use 10 tasks × 50 episodes with the released evaluation protocol.
+The public runtime and checkpoints were revalidated from a clean installation
+across 2,000 episodes on eight GPUs.
+
+| Suite | Successes | Success rate |
+|---|---:|---:|
+| LIBERO-Spatial | 490/500 | **98.0%** |
+| LIBERO-Object | 485/500 | **97.0%** |
+| LIBERO-Goal | 464/500 | **92.8%** |
+| LIBERO-10 | 459/500 | **91.8%** |
+
+# Key Characteristics
+
+- Three-frame observation context
+- Seven-step action prediction horizon
+- Three actions executed between policy replans
+- Temporally ensembled action predictions
+- Suite-specific checkpoints for all four LIBERO benchmarks
+- Headless evaluation on one or more GPUs
+- Compact two-file inference and evaluation runtime
+
+---
+
+# What This Repository Contains
+
+```text
+.
+├── checkpoints/
+│   ├── libero_10/model.safetensors
+│   ├── libero_goal/model.safetensors
+│   ├── libero_object/model.safetensors
+│   └── libero_spatial/model.safetensors
+├── dependencies/
+│   ├── ViT-B-32.pt
+│   └── mae_pretrain_vit_base.pth
+├── eval.py
+├── inference.py
+├── config.json
+└── requirements.txt
+```
+
+`dependencies/` contains the frozen visual and language encoder weights needed
+by the released policy. No additional model downloads are required.
+
+---
+
+# Installation
+
+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`.
+
+---
+
+# Evaluation
+
+## One GPU
+
+```bash
+python eval.py \
+  --suite libero_spatial \
+  --gpus 1 \
+  --output-dir results/libero_spatial
+```
+
+## Multiple GPUs
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python eval.py \
+  --suite libero_spatial \
+  --gpus 8 \
+  --output-dir results/libero_spatial_8gpu
+```
+
+Each GPU receives an independent progress bar. After all workers finish, rank 0
+prints per-task and overall success rates and writes a structured
+`results.json`. Output directories must be new so an earlier evaluation is
+never overwritten.
+
+Supported suites:
+
+```text
+libero_spatial
+libero_object
+libero_goal
+libero_10
+```
+
+For a short installation smoke test:
+
+```bash
+python eval.py \
+  --suite libero_spatial \
+  --gpus 1 \
+  --tasks 1 \
+  --episodes 1 \
+  --max-steps 20 \
+  --output-dir results/smoke
+```
+
+---
+
+# Inference API
+
+```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.
+
+---
+
+# Architecture Details
+
+| Component | Specification |
+|---|---|
+| Policy | WorldDiT diffusion transformer |
+| Observation context | 3 frames |
+| Action horizon | 7 actions |
+| Action dimension | 7 |
+| Language encoder | OpenAI CLIP ViT-B/32 |
+| Visual encoder | MAE ViT-B |
+| Evaluation | Headless LIBERO with EGL |
+| Checkpoint format | SafeTensors |
+
+---
+
+# 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

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