siglip2_action_chunk_model.py

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