metaclip_action_chunk_model.py

"""Self-contained MetaCLIP temporal action-chunk policy definition.

This module deliberately contains no model-hub access.  A complete
``clip_config`` dictionary is embedded in the policy configuration, so
constructing :class:`MetaCLIPActionChunkModel` only creates modules.  Callers are
responsible for loading a local state dict afterwards.

The image, 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 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 CLIPConfig, CLIPModel

ACTION_DIM = 7
DEFAULT_TEXT_DIM = 512
DEFAULT_PHASE_DIM = 4
DEFAULT_SPATIAL_HEADS = 8
DEFAULT_DIFFICULTIES = ("low", "medium", "hard", "very_high")
TEXT_FEATURE_VERSION = "metaclip_clip_bpe_projected_l2_text_v2"
METACLIP_IMAGE_MEAN = (0.48145466, 0.4578275, 0.40821073)
METACLIP_IMAGE_STD = (0.26862954, 0.26130258, 0.27577711)

_REQUIRED_CONFIG_KEYS = (
    "clip_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))
    clip_config = validated["clip_config"]
    if not isinstance(clip_config, Mapping) or not clip_config:
        raise ValueError("clip_config must be a non-empty mapping")
    clip_config = copy.deepcopy(dict(clip_config))
    if clip_config.get("model_type") != "clip":
        raise ValueError(
            "clip_config.model_type must be 'clip', got "
            f"{clip_config.get('model_type')!r}"
        )
    if (
        _require_plain_int(
            clip_config.get("projection_dim"), "clip_config.projection_dim"
        )
        != DEFAULT_TEXT_DIM
    ):
        raise ValueError(
            f"MetaCLIP projection_dim must be {DEFAULT_TEXT_DIM}, "
            f"got {clip_config.get('projection_dim')!r}"
        )

    vision_config = clip_config.get("vision_config")
    if not isinstance(vision_config, Mapping) or not vision_config:
        raise ValueError("clip_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"clip_config.vision_config is missing required key {key!r}"
            )
        _require_plain_int(vision_config[key], f"clip_config.vision_config.{key}")
    if vision_config.get("model_type") != "clip_vision_model":
        raise ValueError(
            "clip_config.vision_config.model_type must be 'clip_vision_model', got "
            f"{vision_config.get('model_type')!r}"
        )
    if (
        int(vision_config["image_size"]) != 224
        or int(vision_config["patch_size"]) != 16
    ):
        raise ValueError(
            "The audited MetaCLIP B/16 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 MetaCLIP vision configurations are supported"
        )

    text_config = clip_config.get("text_config")
    if not isinstance(text_config, Mapping) or not text_config:
        raise ValueError("clip_config.text_config must be a non-empty mapping")
    text_config = copy.deepcopy(dict(text_config))
    for key in (
        "hidden_size",
        "intermediate_size",
        "max_position_embeddings",
        "num_hidden_layers",
        "num_attention_heads",
        "vocab_size",
    ):
        if key not in text_config:
            raise ValueError(f"clip_config.text_config is missing required key {key!r}")
        _require_plain_int(text_config[key], f"clip_config.text_config.{key}")
    if text_config.get("model_type") != "clip_text_model":
        raise ValueError(
            "clip_config.text_config.model_type must be 'clip_text_model', got "
            f"{text_config.get('model_type')!r}"
        )
    if int(text_config["max_position_embeddings"]) != 77:
        raise ValueError(
            "MetaCLIP text max_position_embeddings must be 77, got "
            f"{text_config['max_position_embeddings']!r}"
        )
    clip_config["vision_config"] = vision_config
    clip_config["text_config"] = text_config
    validated["clip_config"] = clip_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 MetaCLIP 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} "
            "MetaCLIP 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(clip_config["projection_dim"]):
        raise ValueError(
            "text_dim must equal clip_config.projection_dim, got "
            f"{validated['text_dim']} and {clip_config['projection_dim']!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"
    )

    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 _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="bicubic",
            align_corners=False,
            antialias=True,
        )
    return rgb


def normalize_images(images: torch.Tensor, image_size: int = 224) -> torch.Tensor:
    """Prepare image pixels for the MetaCLIP vision encoder."""

    rgb = images_to_unit_rgb(images, image_size=image_size)
    mean = rgb.new_tensor(METACLIP_IMAGE_MEAN).view(1, 3, 1, 1)
    std = rgb.new_tensor(METACLIP_IMAGE_STD).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_metaclip_tokens(
    hidden_states: torch.Tensor,
    spatial_grid: int,
) -> torch.Tensor:
    """Pool MetaCLIP patch tokens to ``G x G`` and retain the CLS token."""

    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 MetaCLIP output shape: {tuple(hidden_states.shape)}"
        )
    cls_token = hidden_states[:, :1]
    patches = hidden_states[:, 1:]
    side = math.isqrt(int(patches.shape[1]))
    if side * side != int(patches.shape[1]):
        raise ValueError(f"MetaCLIP 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 MetaCLIPActionChunkHead(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 MetaCLIPActionChunkModel(nn.Module):
    """Complete submission model containing frozen MetaCLIP 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"])

        # Offline construction only: this creates a model from the embedded
        # architecture.  It never resolves a repository or downloads weights.
        clip_config = CLIPConfig.from_dict(self.policy_config["clip_config"])
        self.clip = CLIPModel(clip_config)
        self.head = MetaCLIPActionChunkHead(
            vision_dim=int(clip_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 MetaCLIP towers and keep them in inference mode."""

        self.clip.requires_grad_(False)
        self.clip.eval()

    def train(self, mode: bool = True) -> "MetaCLIPActionChunkModel":
        # 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.clip.parameters()):
            self.clip.eval()
        return self

    def encode_images(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Encode an image batch into aligned MetaCLIP 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.clip.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
        )
        hidden = self.clip.vision_model(pixel_values=pixels).last_hidden_state
        hidden = self.clip.vision_model.post_layernorm(hidden)
        return (
            pool_metaclip_tokens(hidden, self.spatial_grid),
            rgb_tokens,
        )

    def encode_text(
        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
    ) -> torch.Tensor:
        """Return normalized projected MetaCLIP instruction embeddings."""

        parameter = next(self.clip.text_model.parameters())
        input_ids = input_ids.to(device=parameter.device)
        attention_mask = attention_mask.to(device=parameter.device)
        outputs = self.clip.text_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True,
        )
        features = self.clip.text_projection(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,
        )


# Compatibility aliases make reference checkpoints/code easy to compare while
# retaining descriptive names in the new trainer.
CompetitivePolicyHead = MetaCLIPActionChunkHead
CompetitiveVLAModel = MetaCLIPActionChunkModel


__all__ = [
    "ACTION_DIM",
    "DEFAULT_DIFFICULTIES",
    "DEFAULT_PHASE_DIM",
    "DEFAULT_SPATIAL_HEADS",
    "DEFAULT_TEXT_DIM",
    "METACLIP_IMAGE_MEAN",
    "METACLIP_IMAGE_STD",
    "TEXT_FEATURE_VERSION",
    "CompetitivePolicyHead",
    "CompetitiveVLAModel",
    "MetaCLIPActionChunkHead",
    "MetaCLIPActionChunkModel",
    "images_to_unit_rgb",
    "normalize_images",
    "phase_vector",
    "pool_metaclip_tokens",
    "rgb_grid_tokens",
    "validate_policy_config",
]