"""Self-contained SigLIP 2 So400m temporal action-chunk policy. This module deliberately performs no model-hub access. The exported policy configuration embeds the complete Hugging Face ``SiglipConfig`` dictionary, so constructing :class:`SigLIP2ActionChunkModel` only creates modules. Callers load the local state dict afterwards. The fixed-resolution SigLIP 2 checkpoint is represented by Transformers' ``SiglipConfig`` / ``SiglipModel`` classes (its serialized ``model_type`` is ``"siglip"``). Unlike CLIP, its vision sequence contains patch tokens only: there is no CLS token. We therefore retain the model's learned global attention-pooler output and add an 8x8 spatial grid made from every patch, producing the same 65-token policy contract used by the other action-chunk backbones. """ from __future__ import annotations import copy import math 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 SiglipConfig, SiglipModel ACTION_DIM = 7 DEFAULT_PHASE_DIM = 4 DEFAULT_SPATIAL_HEADS = 8 DEFAULT_TEXT_MAX_LENGTH = 64 DEFAULT_DIFFICULTIES = ("low", "medium", "hard", "very_high") TEXT_FEATURE_VERSION = "siglip2_sentencepiece_lowercase_projected_l2_text_v1" SIGLIP_IMAGE_MEAN = (0.5, 0.5, 0.5) SIGLIP_IMAGE_STD = (0.5, 0.5, 0.5) _REQUIRED_CONFIG_KEYS = ( "siglip_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", "text_max_length", "text_lowercase", ) 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") result[key] = int(item) actual = sorted(result.values()) expected = list(range(len(result))) 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 policy configuration. Shape and preprocessing fields are intentionally strict so feature-cache creation, training, export, and validator inference cannot silently drift. Dimensions are read from the embedded backbone config rather than hard-coded to 1152, while the audited So400m FixRes input contract remains fixed at 384 pixels with 14-pixel patches. """ 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)) siglip_config = validated["siglip_config"] if not isinstance(siglip_config, Mapping) or not siglip_config: raise ValueError("siglip_config must be a non-empty mapping") siglip_config = copy.deepcopy(dict(siglip_config)) if siglip_config.get("model_type") != "siglip": raise ValueError( "siglip_config.model_type must be 'siglip', got " f"{siglip_config.get('model_type')!r}" ) vision_config = siglip_config.get("vision_config") if not isinstance(vision_config, Mapping) or not vision_config: raise ValueError("siglip_config.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"siglip_config.vision_config is missing required key {key!r}" ) _require_plain_int( vision_config[key], f"siglip_config.vision_config.{key}" ) if vision_config.get("model_type") != "siglip_vision_model": raise ValueError( "siglip_config.vision_config.model_type must be " f"'siglip_vision_model', got {vision_config.get('model_type')!r}" ) if ( int(vision_config["image_size"]) != 384 or int(vision_config["patch_size"]) != 14 ): raise ValueError( "The audited SigLIP 2 So400m contract requires image_size=384 and " f"patch_size=14, 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 int(vision_config.get("num_channels", 3)) != 3: raise ValueError("only three-channel SigLIP vision configurations are supported") text_config = siglip_config.get("text_config") if not isinstance(text_config, Mapping) or not text_config: raise ValueError("siglip_config.text_config must be a non-empty mapping") text_config = copy.deepcopy(dict(text_config)) for key in ( "hidden_size", "intermediate_size", "num_hidden_layers", "num_attention_heads", "projection_size", "vocab_size", ): if key not in text_config: raise ValueError( f"siglip_config.text_config is missing required key {key!r}" ) _require_plain_int(text_config[key], f"siglip_config.text_config.{key}") if text_config.get("model_type") != "siglip_text_model": raise ValueError( "siglip_config.text_config.model_type must be 'siglip_text_model', got " f"{text_config.get('model_type')!r}" ) # Some early config.json files omit this value because 64 is the class # default. Store it explicitly in the exported contract. text_max_positions = _require_plain_int( text_config.get("max_position_embeddings", DEFAULT_TEXT_MAX_LENGTH), "siglip_config.text_config.max_position_embeddings", ) if text_max_positions != DEFAULT_TEXT_MAX_LENGTH: raise ValueError( "SigLIP 2 text max_position_embeddings must be " f"{DEFAULT_TEXT_MAX_LENGTH}, got {text_max_positions}" ) text_config["max_position_embeddings"] = text_max_positions siglip_config["vision_config"] = vision_config siglip_config["text_config"] = text_config validated["siglip_config"] = siglip_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_side = image_size // int(vision_config["patch_size"]) if patch_side < 1: raise ValueError("SigLIP patch size is larger than the configured image") 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} " "SigLIP 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["text_dim"] != int(text_config["projection_size"]): raise ValueError( "text_dim must equal text_config.projection_size, got " f"{validated['text_dim']} and {text_config['projection_size']!r}" ) 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" ) validated["text_max_length"] = _require_plain_int( validated["text_max_length"], "text_max_length" ) if validated["text_max_length"] != DEFAULT_TEXT_MAX_LENGTH: raise ValueError( f"text_max_length must be {DEFAULT_TEXT_MAX_LENGTH}, got " f"{validated['text_max_length']}" ) if validated["text_lowercase"] is not True: raise ValueError("text_lowercase must be true for the SigLIP 2 text contract") 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 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 canonicalize_instruction(text: str) -> str: """Apply the shared SigLIP 2 instruction normalization contract.""" if not isinstance(text, str): raise TypeError(f"instruction must be a string, got {type(text).__name__}") return text.strip().lower() 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)}" ) 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 = 384) -> torch.Tensor: """Convert NHWC/NCHW uint8-like images to resized 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), # The official checkpoint's preprocessor_config.json declares # PIL resample=2, i.e. bilinear interpolation. mode="bilinear", align_corners=False, antialias=True, ) return rgb def normalize_images(images: torch.Tensor, image_size: int = 384) -> torch.Tensor: """Prepare image pixels for the fixed-resolution SigLIP vision encoder.""" rgb = images_to_unit_rgb(images, image_size=image_size) mean = rgb.new_tensor(SIGLIP_IMAGE_MEAN).view(1, 3, 1, 1) std = rgb.new_tensor(SIGLIP_IMAGE_STD).view(1, 3, 1, 1) return (rgb - mean) / std def rgb_grid_tokens( images: torch.Tensor, spatial_grid: int, image_size: int = 384 ) -> torch.Tensor: """Return global RGB plus a row-major spatial grid, ``[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_siglip_tokens( hidden_states: torch.Tensor, pooler_output: torch.Tensor, spatial_grid: int, ) -> torch.Tensor: """Combine the learned global pooler with a pooled ``G x G`` patch grid. Fixed-resolution SigLIP has no CLS token. Its learned attention-pooling head is therefore the global token, while every item in ``last_hidden_state`` is retained when building the spatial grid. """ spatial_grid = _require_plain_int(spatial_grid, "spatial_grid") if not isinstance(hidden_states, torch.Tensor): raise TypeError("hidden_states must be a torch.Tensor") if hidden_states.ndim != 3 or hidden_states.shape[1] < 1: raise ValueError(f"unexpected SigLIP output shape: {tuple(hidden_states.shape)}") if not isinstance(pooler_output, torch.Tensor): raise TypeError("pooler_output must be a torch.Tensor") expected_pooler_shape = (hidden_states.shape[0], hidden_states.shape[2]) if tuple(pooler_output.shape) != expected_pooler_shape: raise ValueError( "unexpected SigLIP pooler shape: " f"{tuple(pooler_output.shape)}, expected {expected_pooler_shape}" ) side = math.isqrt(int(hidden_states.shape[1])) if side * side != int(hidden_states.shape[1]): raise ValueError(f"SigLIP patch count {hidden_states.shape[1]} is not a square") global_token = pooler_output.unsqueeze(1) patches = hidden_states.transpose(1, 2).reshape( hidden_states.shape[0], hidden_states.shape[2], side, side ) patches = F.adaptive_avg_pool2d(patches, (spatial_grid, spatial_grid)) patches = patches.flatten(2).transpose(1, 2) return torch.cat((global_token, patches), dim=1) class SigLIP2ActionChunkHead(nn.Module): """Task-conditioned spatial pooling followed by a temporal chunk 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(), ) 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( "visual_tokens must have shape [B,T,V,D], got " f"{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 actions 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 SigLIP2ActionChunkModel(nn.Module): """Complete offline submission model: frozen SigLIP 2 plus 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"]) siglip_config = SiglipConfig.from_dict(self.policy_config["siglip_config"]) self.siglip = SiglipModel(siglip_config) self.head = SigLIP2ActionChunkHead( vision_dim=int(siglip_config.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_backbone(self) -> None: """Freeze both SigLIP towers and keep them in inference mode.""" self.siglip.requires_grad_(False) self.siglip.eval() def train(self, mode: bool = True) -> "SigLIP2ActionChunkModel": super().train(mode) if not any(parameter.requires_grad for parameter in self.siglip.parameters()): self.siglip.eval() return self def encode_images(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Encode images into aligned 65-token SigLIP and raw-RGB grids.""" 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.siglip.vision_model.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 ) outputs = self.siglip.vision_model( pixel_values=pixels, interpolate_pos_encoding=False, return_dict=True, ) return ( pool_siglip_tokens( outputs.last_hidden_state, outputs.pooler_output, self.spatial_grid, ), rgb_tokens, ) def encode_text( self, input_ids: torch.Tensor, attention_mask: torch.Tensor ) -> torch.Tensor: """Return normalized projected SigLIP instruction embeddings.""" parameter = next(self.siglip.text_model.parameters()) input_ids = input_ids.to(device=parameter.device) attention_mask = attention_mask.to(device=parameter.device) outputs = self.siglip.text_model( input_ids=input_ids, attention_mask=attention_mask, return_dict=True, ) features = outputs.pooler_output return F.normalize(features.float(), dim=-1).to(dtype=parameter.dtype) 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, ) CompetitivePolicyHead = SigLIP2ActionChunkHead CompetitiveVLAModel = SigLIP2ActionChunkModel __all__ = [ "ACTION_DIM", "DEFAULT_DIFFICULTIES", "DEFAULT_PHASE_DIM", "DEFAULT_SPATIAL_HEADS", "DEFAULT_TEXT_MAX_LENGTH", "SIGLIP_IMAGE_MEAN", "SIGLIP_IMAGE_STD", "TEXT_FEATURE_VERSION", "CompetitivePolicyHead", "CompetitiveVLAModel", "SigLIP2ActionChunkHead", "SigLIP2ActionChunkModel", "canonicalize_instruction", "images_to_unit_rgb", "normalize_images", "phase_vector", "pool_siglip_tokens", "rgb_grid_tokens", "validate_policy_config", ]