"""Self-contained DINOv3 temporal action-chunk policy definition. This module deliberately contains no model-hub access. A complete ``vision_config`` dictionary is embedded in the policy configuration, so constructing :class:`DinoV3ActionChunkModel` only creates modules. Callers are responsible for loading a local state dict afterwards. The image, text, phase, and token helpers are shared by feature-cache creation, training, and the submission adapter. Keeping those operations here prevents subtle train/deployment preprocessing drift. """ from __future__ import annotations import copy import hashlib import math import re from collections.abc import Mapping from typing import Any import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from transformers import DINOv3ViTConfig, DINOv3ViTModel ACTION_DIM = 7 DEFAULT_TEXT_DIM = 128 DEFAULT_PHASE_DIM = 4 DEFAULT_SPATIAL_HEADS = 8 DEFAULT_DIFFICULTIES = ("low", "medium", "hard", "very_high") TEXT_FEATURE_VERSION = "signed_hash_subtokens_v2" # Common paraphrases are collapsed before hashing. Underscore-delimited task # names retain both their complete token and their component tokens. TOKEN_ALIASES = { "build": "stack", "clean": "wipe", "container": "bin", "grab": "pick", "grasp": "pick", "move": "place", "put": "place", "sweep": "wipe", } _REQUIRED_CONFIG_KEYS = ( "vision_config", "image_size", "spatial_grid", "proprio_dim", "text_dim", "phase_dim", "hidden_dim", "history", "action_chunk", "ensemble_heads", "dropout", "task_to_id", "difficulty_to_id", ) def _require_plain_int(value: Any, name: str, *, minimum: int = 1) -> int: if isinstance(value, bool) or not isinstance(value, int) or value < minimum: raise ValueError(f"{name} must be an integer >= {minimum}, got {value!r}") return int(value) def _require_probability(value: Any, name: str) -> float: if isinstance(value, bool) or not isinstance(value, (int, float)): raise ValueError(f"{name} must be a number in [0, 1), got {value!r}") value = float(value) if not math.isfinite(value) or not 0.0 <= value < 1.0: raise ValueError(f"{name} must be a finite number in [0, 1), got {value!r}") return value def _validate_id_map(value: Any, name: str) -> dict[str, int]: if not isinstance(value, Mapping) or not value: raise ValueError(f"{name} must be a non-empty mapping of strings to IDs") result: dict[str, int] = {} for key, item in value.items(): if not isinstance(key, str) or not key.strip(): raise ValueError(f"{name} contains an invalid key: {key!r}") if isinstance(item, bool) or not isinstance(item, int) or item < 0: raise ValueError(f"{name}[{key!r}] must be a non-negative integer") if key in result: raise ValueError(f"{name} contains duplicate key {key!r}") result[key] = int(item) expected = list(range(len(result))) actual = sorted(result.values()) if actual != expected: raise ValueError( f"{name} IDs must be unique and contiguous 0..{len(result) - 1}, got {actual}" ) return result def validate_policy_config(config: Mapping[str, Any]) -> dict[str, Any]: """Validate and return an isolated copy of a policy configuration. Deployment/training metadata outside the architectural keys is retained, but all fields that affect tensor shapes or preprocessing are checked. The returned object can therefore be safely stored on a model without later mutations by the caller changing its behavior. """ if not isinstance(config, Mapping): raise TypeError(f"config must be a mapping, got {type(config).__name__}") missing = [key for key in _REQUIRED_CONFIG_KEYS if key not in config] if missing: raise ValueError(f"policy config is missing required keys: {missing}") validated = copy.deepcopy(dict(config)) vision_config = validated["vision_config"] if not isinstance(vision_config, Mapping) or not vision_config: raise ValueError("vision_config must be a non-empty mapping") vision_config = copy.deepcopy(dict(vision_config)) for key in ( "hidden_size", "intermediate_size", "image_size", "patch_size", "num_hidden_layers", "num_attention_heads", ): if key not in vision_config: raise ValueError(f"vision_config is missing required key {key!r}") _require_plain_int(vision_config[key], f"vision_config.{key}") if vision_config.get("model_type") != "dinov3_vit": raise ValueError( "vision_config.model_type must be 'dinov3_vit', got " f"{vision_config.get('model_type')!r}" ) register_tokens = vision_config.get("num_register_tokens") if ( isinstance(register_tokens, bool) or not isinstance(register_tokens, int) or register_tokens != 4 ): raise ValueError( "vision_config.num_register_tokens must be exactly 4, " f"got {register_tokens!r}" ) if ( int(vision_config["image_size"]) != 224 or int(vision_config["patch_size"]) != 16 ): raise ValueError( "The audited DINOv3 ViT contract requires image_size=224 and " f"patch_size=16, got {vision_config['image_size']!r} and " f"{vision_config['patch_size']!r}" ) if int(vision_config["hidden_size"]) % int(vision_config["num_attention_heads"]): raise ValueError( "vision_config.hidden_size must be divisible by " "vision_config.num_attention_heads" ) if "num_channels" in vision_config and int(vision_config["num_channels"]) != 3: raise ValueError( "only three-channel DINOv3 vision configurations are supported" ) validated["vision_config"] = vision_config image_size = _require_plain_int(validated["image_size"], "image_size") if image_size != int(vision_config["image_size"]): raise ValueError( "policy image_size must match vision_config.image_size, got " f"{image_size} and {vision_config['image_size']!r}" ) patch_size = int(vision_config["patch_size"]) if image_size % patch_size: raise ValueError( f"image_size={image_size} must be divisible by DINOv3 patch_size={patch_size}" ) patch_side = image_size // patch_size spatial_grid = _require_plain_int(validated["spatial_grid"], "spatial_grid") if spatial_grid > patch_side: raise ValueError( f"spatial_grid={spatial_grid} cannot exceed the {patch_side}x{patch_side} " "DINOv3 input patch grid" ) for key in ( "proprio_dim", "text_dim", "phase_dim", "hidden_dim", "history", "action_chunk", "ensemble_heads", ): validated[key] = _require_plain_int(validated[key], key) if validated["phase_dim"] != DEFAULT_PHASE_DIM: raise ValueError( f"phase_dim must be {DEFAULT_PHASE_DIM} for phase_vector(), " f"got {validated['phase_dim']}" ) if validated["hidden_dim"] % DEFAULT_SPATIAL_HEADS: raise ValueError( f"hidden_dim must be divisible by {DEFAULT_SPATIAL_HEADS} spatial heads" ) validated["dropout"] = _require_probability(validated["dropout"], "dropout") validated["task_to_id"] = _validate_id_map(validated["task_to_id"], "task_to_id") validated["difficulty_to_id"] = _validate_id_map( validated["difficulty_to_id"], "difficulty_to_id" ) if "action_dim" in validated and validated["action_dim"] != ACTION_DIM: raise ValueError( f"action_dim must be {ACTION_DIM}, got {validated['action_dim']!r}" ) if "rgb_dim" in validated and validated["rgb_dim"] != 3: raise ValueError(f"rgb_dim must be 3, got {validated['rgb_dim']!r}") if ( "text_feature_version" in validated and validated["text_feature_version"] != TEXT_FEATURE_VERSION ): raise ValueError( f"text_feature_version must be {TEXT_FEATURE_VERSION!r}, got " f"{validated['text_feature_version']!r}" ) return validated def text_vector(text: str, dim: int = DEFAULT_TEXT_DIM) -> np.ndarray: """Return a deterministic signed-hash bag-of-subtokens text feature.""" dim = _require_plain_int(dim, "text feature dimension") vec = np.zeros(dim, dtype=np.float32) raw_tokens = re.findall(r"[a-z0-9_]+", str(text).lower()) tokens: list[str] = [] for raw_token in raw_tokens: token = TOKEN_ALIASES.get(raw_token, raw_token) tokens.append(token) if "_" in token: tokens.extend(TOKEN_ALIASES.get(part, part) for part in token.split("_")) for token in tokens: digest = hashlib.blake2b(token.encode("utf-8"), digest_size=8).digest() value = int.from_bytes(digest, byteorder="little", signed=False) vec[value % dim] += 1.0 if value & 1 else -1.0 norm = float(np.linalg.norm(vec)) if norm > 0.0: vec /= norm return vec def phase_vector(step: int, horizon: int) -> np.ndarray: """Encode episode progress using the exact train/runtime four-vector.""" if isinstance(step, bool) or not isinstance(step, (int, np.integer)): raise ValueError(f"step must be an integer, got {step!r}") if isinstance(horizon, bool) or not isinstance(horizon, (int, np.integer)): raise ValueError(f"horizon must be an integer, got {horizon!r}") horizon_f = max(float(horizon), 1.0) progress = float(np.clip(float(step) / horizon_f, 0.0, 1.0)) return np.asarray( [ progress, 1.0 - progress, math.sin(math.pi * progress), math.cos(math.pi * progress), ], dtype=np.float32, ) def _images_to_nchw_rgb(images: torch.Tensor) -> torch.Tensor: if not isinstance(images, torch.Tensor): raise TypeError(f"images must be a torch.Tensor, got {type(images).__name__}") if images.ndim != 4: raise ValueError( f"expected a four-dimensional image tensor, got {tuple(images.shape)}" ) # Prefer an unambiguous channel-first interpretation, then NHWC. Normal # robotics images are 224x224, so both layouts are unambiguous in practice. if images.shape[1] in (1, 3, 4) and images.shape[-1] not in (1, 3, 4): nchw = images elif images.shape[-1] in (1, 3, 4): nchw = images.permute(0, 3, 1, 2) elif images.shape[1] in (1, 3, 4): nchw = images else: raise ValueError( f"cannot determine image channels for shape {tuple(images.shape)}" ) if nchw.shape[1] == 1: nchw = nchw.repeat(1, 3, 1, 1) elif nchw.shape[1] == 4: nchw = nchw[:, :3] if nchw.shape[1] != 3: raise ValueError( f"expected one, three, or four image channels, got {nchw.shape[1]}" ) return nchw def images_to_unit_rgb(images: torch.Tensor, image_size: int = 224) -> torch.Tensor: """Convert NHWC/NCHW uint8-like images to resized NCHW RGB in ``[0, 1]``.""" image_size = _require_plain_int(image_size, "image_size") rgb = _images_to_nchw_rgb(images).float() if rgb.numel() and float(rgb.detach().amax().cpu()) > 2.0: rgb = rgb / 255.0 if not bool(torch.isfinite(rgb).all().detach().cpu()): raise ValueError("images contain NaN or Inf") if tuple(rgb.shape[-2:]) != (image_size, image_size): rgb = F.interpolate( rgb, size=(image_size, image_size), mode="bilinear", align_corners=False, ) return rgb def normalize_images(images: torch.Tensor, image_size: int = 224) -> torch.Tensor: """Prepare image pixels for the DINOv3 vision encoder.""" rgb = images_to_unit_rgb(images, image_size=image_size) mean = rgb.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) std = rgb.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) return (rgb - mean) / std def rgb_grid_tokens( images: torch.Tensor, spatial_grid: int, image_size: int = 224 ) -> torch.Tensor: """Return global RGB plus a row-major spatial grid, shaped ``[B,1+G²,3]``.""" spatial_grid = _require_plain_int(spatial_grid, "spatial_grid") rgb = images_to_unit_rgb(images, image_size=image_size) global_rgb = rgb.mean(dim=(-2, -1)).unsqueeze(1) grid_rgb = F.adaptive_avg_pool2d(rgb, (spatial_grid, spatial_grid)) grid_rgb = grid_rgb.flatten(2).transpose(1, 2) return torch.cat([global_rgb, grid_rgb], dim=1) def pool_dinov3_tokens( hidden_states: torch.Tensor, spatial_grid: int, num_register_tokens: int, ) -> torch.Tensor: """Drop register tokens and pool DINOv3 patches to ``G x G`` plus CLS.""" spatial_grid = _require_plain_int(spatial_grid, "spatial_grid") if ( isinstance(num_register_tokens, bool) or not isinstance(num_register_tokens, int) or num_register_tokens < 0 ): raise ValueError( "num_register_tokens must be a non-negative integer, got " f"{num_register_tokens!r}" ) if not isinstance(hidden_states, torch.Tensor): raise TypeError("hidden_states must be a torch.Tensor") prefix_tokens = 1 + num_register_tokens if hidden_states.ndim != 3 or hidden_states.shape[1] <= prefix_tokens: raise ValueError( f"unexpected DINOv3 output shape: {tuple(hidden_states.shape)}" ) cls_token = hidden_states[:, :1] patches = hidden_states[:, prefix_tokens:] side = math.isqrt(int(patches.shape[1])) if side * side != int(patches.shape[1]): raise ValueError(f"DINOv3 patch count {patches.shape[1]} is not a square") patches = patches.transpose(1, 2).reshape( patches.shape[0], patches.shape[2], side, side ) patches = F.adaptive_avg_pool2d(patches, (spatial_grid, spatial_grid)) patches = patches.flatten(2).transpose(1, 2) return torch.cat([cls_token, patches], dim=1) class DinoV3ActionChunkHead(nn.Module): """Task-conditioned spatial pooling followed by a short temporal policy.""" def __init__( self, *, vision_dim: int, proprio_dim: int, text_dim: int, phase_dim: int, num_tasks: int, num_difficulties: int, hidden_dim: int = 256, history: int = 4, action_chunk: int = 8, ensemble_heads: int = 3, dropout: float = 0.10, ) -> None: super().__init__() self.vision_dim = _require_plain_int(vision_dim, "vision_dim") self.proprio_dim = _require_plain_int(proprio_dim, "proprio_dim") self.text_dim = _require_plain_int(text_dim, "text_dim") self.phase_dim = _require_plain_int(phase_dim, "phase_dim") self.hidden_dim = _require_plain_int(hidden_dim, "hidden_dim") self.history = _require_plain_int(history, "history") self.action_chunk = _require_plain_int(action_chunk, "action_chunk") self.ensemble_heads = _require_plain_int(ensemble_heads, "ensemble_heads") num_tasks = _require_plain_int(num_tasks, "num_tasks") num_difficulties = _require_plain_int(num_difficulties, "num_difficulties") dropout = _require_probability(dropout, "dropout") if self.hidden_dim % DEFAULT_SPATIAL_HEADS: raise ValueError( f"hidden_dim must be divisible by {DEFAULT_SPATIAL_HEADS} spatial heads" ) self.vision_proj = nn.Sequential( nn.LayerNorm(self.vision_dim), nn.Linear(self.vision_dim, self.hidden_dim) ) self.rgb_proj = nn.Sequential(nn.Linear(3, self.hidden_dim), nn.SiLU()) self.proprio_proj = nn.Sequential( nn.LayerNorm(self.proprio_dim + self.phase_dim), nn.Linear(self.proprio_dim + self.phase_dim, self.hidden_dim), nn.SiLU(), nn.Dropout(dropout), ) self.text_proj = nn.Sequential( nn.LayerNorm(self.text_dim), nn.Linear(self.text_dim, self.hidden_dim), nn.SiLU(), ) # The last row of each embedding is the trained unknown/fallback ID. self.task_embedding = nn.Embedding(num_tasks + 1, self.hidden_dim) self.difficulty_embedding = nn.Embedding(num_difficulties + 1, self.hidden_dim) self.task_scale = nn.Parameter(torch.tensor(0.5)) self.difficulty_scale = nn.Parameter(torch.tensor(0.25)) self.condition_norm = nn.LayerNorm(self.hidden_dim) self.spatial_attention = nn.MultiheadAttention( embed_dim=self.hidden_dim, num_heads=DEFAULT_SPATIAL_HEADS, dropout=dropout, batch_first=True, ) self.frame_fusion = nn.Sequential( nn.Linear(self.hidden_dim * 2, self.hidden_dim), nn.SiLU(), nn.LayerNorm(self.hidden_dim), nn.Dropout(dropout), ) self.temporal_gru = nn.GRU( input_size=self.hidden_dim, hidden_size=self.hidden_dim, num_layers=2, dropout=dropout, batch_first=True, ) self.output_heads = nn.ModuleList( [ nn.Sequential( nn.LayerNorm(self.hidden_dim), nn.Linear(self.hidden_dim, self.hidden_dim), nn.SiLU(), nn.Dropout(dropout), nn.Linear(self.hidden_dim, self.action_chunk * ACTION_DIM), ) for _ in range(self.ensemble_heads) ] ) def _validate_inputs( self, visual_tokens: torch.Tensor, rgb_tokens: torch.Tensor, proprio: torch.Tensor, phase: torch.Tensor, text_features: torch.Tensor, task_ids: torch.Tensor, difficulty_ids: torch.Tensor, ) -> tuple[int, int, int]: if visual_tokens.ndim != 4: raise ValueError( f"visual_tokens must have shape [B,T,V,D], got {tuple(visual_tokens.shape)}" ) batch, timesteps, token_count, vision_dim = visual_tokens.shape if timesteps != self.history: raise ValueError(f"expected history={self.history}, got {timesteps}") if vision_dim != self.vision_dim: raise ValueError(f"expected vision_dim={self.vision_dim}, got {vision_dim}") if rgb_tokens.shape != (batch, timesteps, token_count, 3): raise ValueError( "rgb_tokens must align with visual_tokens and end in RGB, got " f"{tuple(rgb_tokens.shape)}" ) if proprio.shape != (batch, timesteps, self.proprio_dim): raise ValueError( f"proprio must have shape {(batch, timesteps, self.proprio_dim)}, " f"got {tuple(proprio.shape)}" ) if phase.shape != (batch, timesteps, self.phase_dim): raise ValueError( f"phase must have shape {(batch, timesteps, self.phase_dim)}, " f"got {tuple(phase.shape)}" ) if text_features.shape != (batch, self.text_dim): raise ValueError( f"text_features must have shape {(batch, self.text_dim)}, " f"got {tuple(text_features.shape)}" ) for name, ids in (("task_ids", task_ids), ("difficulty_ids", difficulty_ids)): if ids.shape != (batch,): raise ValueError( f"{name} must have shape {(batch,)}, got {tuple(ids.shape)}" ) if ids.dtype not in (torch.int32, torch.int64): raise ValueError( f"{name} must contain integer IDs, got dtype={ids.dtype}" ) return batch, timesteps, token_count def forward_cached( self, visual_tokens: torch.Tensor, rgb_tokens: torch.Tensor, proprio: torch.Tensor, phase: torch.Tensor, text_features: torch.Tensor, task_ids: torch.Tensor, difficulty_ids: torch.Tensor, ) -> torch.Tensor: """Return raw action logits shaped ``[ensemble, batch, chunk, 7]``.""" batch, timesteps, token_count = self._validate_inputs( visual_tokens, rgb_tokens, proprio, phase, text_features, task_ids, difficulty_ids, ) visual = self.vision_proj(visual_tokens) + self.rgb_proj(rgb_tokens) state = self.proprio_proj(torch.cat([proprio, phase], dim=-1)) text = self.text_proj(text_features) task = self.task_embedding(task_ids) difficulty = self.difficulty_embedding(difficulty_ids) condition = self.condition_norm( state + text[:, None] + torch.tanh(self.task_scale) * task[:, None] + torch.tanh(self.difficulty_scale) * difficulty[:, None] ) flat_visual = visual.reshape(batch * timesteps, token_count, self.hidden_dim) flat_query = condition.reshape(batch * timesteps, 1, self.hidden_dim) attended, _ = self.spatial_attention( flat_query, flat_visual, flat_visual, need_weights=False ) attended = attended.reshape(batch, timesteps, self.hidden_dim) frames = self.frame_fusion(torch.cat([attended, condition], dim=-1)) temporal, _ = self.temporal_gru(frames) final = temporal[:, -1] outputs = [ head(final).reshape(batch, self.action_chunk, ACTION_DIM) for head in self.output_heads ] return torch.stack(outputs, dim=0) def forward( self, visual_tokens: torch.Tensor, rgb_tokens: torch.Tensor, proprio: torch.Tensor, phase: torch.Tensor, text_features: torch.Tensor, task_ids: torch.Tensor, difficulty_ids: torch.Tensor, ) -> torch.Tensor: return self.forward_cached( visual_tokens, rgb_tokens, proprio, phase, text_features, task_ids, difficulty_ids, ) class DinoV3ActionChunkModel(nn.Module): """Complete submission model containing frozen DINOv3 and the policy head.""" def __init__(self, config: Mapping[str, Any]) -> None: super().__init__() self.policy_config = validate_policy_config(config) self.spatial_grid = int(self.policy_config["spatial_grid"]) self.num_register_tokens = int( self.policy_config["vision_config"]["num_register_tokens"] ) # Offline construction only: this creates a model from the embedded # architecture. It never resolves a repository or downloads weights. vision_config = DINOv3ViTConfig(**self.policy_config["vision_config"]) self.vision = DINOv3ViTModel(vision_config) self.head = DinoV3ActionChunkHead( vision_dim=int(vision_config.hidden_size), proprio_dim=int(self.policy_config["proprio_dim"]), text_dim=int(self.policy_config["text_dim"]), phase_dim=int(self.policy_config["phase_dim"]), num_tasks=len(self.policy_config["task_to_id"]), num_difficulties=len(self.policy_config["difficulty_to_id"]), hidden_dim=int(self.policy_config["hidden_dim"]), history=int(self.policy_config["history"]), action_chunk=int(self.policy_config["action_chunk"]), ensemble_heads=int(self.policy_config["ensemble_heads"]), dropout=float(self.policy_config["dropout"]), ) def freeze_vision(self) -> None: """Freeze the backbone and keep it in deterministic inference mode.""" self.vision.requires_grad_(False) self.vision.eval() def train(self, mode: bool = True) -> "DinoV3ActionChunkModel": # A caller may train the complete wrapper for convenience. If the # backbone has been frozen, do not accidentally switch it back to train # mode through nn.Module.train() recursion. super().train(mode) if not any(parameter.requires_grad for parameter in self.vision.parameters()): self.vision.eval() return self def encode_images(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Encode an image batch into aligned DINOv3 and raw-RGB tokens.""" image_size = int(self.policy_config["image_size"]) rgb_tokens = rgb_grid_tokens( images, spatial_grid=self.spatial_grid, image_size=image_size ) pixels = normalize_images(images, image_size=image_size) vision_parameter = next(self.vision.parameters()) pixels = pixels.to(device=vision_parameter.device, dtype=vision_parameter.dtype) rgb_tokens = rgb_tokens.to( device=vision_parameter.device, dtype=vision_parameter.dtype ) hidden = self.vision(pixel_values=pixels).last_hidden_state return ( pool_dinov3_tokens( hidden, self.spatial_grid, self.num_register_tokens, ), rgb_tokens, ) def forward_cached( self, visual_tokens: torch.Tensor, rgb_tokens: torch.Tensor, proprio: torch.Tensor, phase: torch.Tensor, text_features: torch.Tensor, task_ids: torch.Tensor, difficulty_ids: torch.Tensor, ) -> torch.Tensor: return self.head.forward_cached( visual_tokens, rgb_tokens, proprio, phase, text_features, task_ids, difficulty_ids, ) def forward( self, visual_tokens: torch.Tensor, rgb_tokens: torch.Tensor, proprio: torch.Tensor, phase: torch.Tensor, text_features: torch.Tensor, task_ids: torch.Tensor, difficulty_ids: torch.Tensor, ) -> torch.Tensor: return self.forward_cached( visual_tokens, rgb_tokens, proprio, phase, text_features, task_ids, difficulty_ids, ) # Compatibility aliases make reference checkpoints/code easy to compare while # retaining descriptive names in the new trainer. CompetitivePolicyHead = DinoV3ActionChunkHead CompetitiveVLAModel = DinoV3ActionChunkModel __all__ = [ "ACTION_DIM", "DEFAULT_DIFFICULTIES", "DEFAULT_PHASE_DIM", "DEFAULT_SPATIAL_HEADS", "DEFAULT_TEXT_DIM", "TEXT_FEATURE_VERSION", "TOKEN_ALIASES", "CompetitivePolicyHead", "CompetitiveVLAModel", "DinoV3ActionChunkHead", "DinoV3ActionChunkModel", "images_to_unit_rgb", "normalize_images", "phase_vector", "pool_dinov3_tokens", "rgb_grid_tokens", "text_vector", "validate_policy_config", ]