| """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:] |
| |
| 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 |
|
|
