modeling_robottrack.py

"""Hugging Face model implementation for MiniCPM-RobotTrack."""

from contextlib import contextmanager
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import torch
from torch import nn
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import ModelOutput

from .configuration_minicpm import MiniCPMConfig
from .configuration_robottrack import MiniCPMRobotTrackConfig
from .modeling_minicpm import MiniCPMModel


@contextmanager
def _default_dtype(dtype: torch.dtype):
    previous = torch.get_default_dtype()
    torch.set_default_dtype(dtype)
    try:
        yield
    finally:
        torch.set_default_dtype(previous)


def _dtype_from_name(name: str) -> torch.dtype:
    try:
        dtype = getattr(torch, name)
    except AttributeError as exc:
        raise ValueError(f"unsupported backbone_dtype={name!r}") from exc
    if not isinstance(dtype, torch.dtype) or not dtype.is_floating_point:
        raise ValueError(f"backbone_dtype must name a floating-point torch dtype: {name!r}")
    return dtype


def _module_dtype(module: nn.Module) -> torch.dtype:
    try:
        return next(module.parameters()).dtype
    except StopIteration:
        return torch.float32


class VisionProjector(nn.Module):
    """Map concatenated DINOv3 and SigLIP features into MiniCPM space."""

    def __init__(self, input_dim: int, hidden_dim: int) -> None:
        super().__init__()
        self.layers = nn.Sequential(
            nn.LayerNorm(input_dim),
            nn.Linear(input_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim),
        )

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        return self.layers(features)


class TemporalMarkerEncoder(nn.Module):
    """Build one marker token for each frame represented in the sequence."""

    def __init__(self, hidden_dim: int, max_time_steps: int) -> None:
        super().__init__()
        self.time_embedding = nn.Embedding(max_time_steps, hidden_dim)
        self.stream_embedding = nn.Embedding(2, hidden_dim)
        self.camera_embedding = nn.Embedding(1, hidden_dim)

    def forward(self, time_step: int, stream_id: int, device: torch.device) -> torch.Tensor:
        time = torch.tensor([time_step], dtype=torch.long, device=device)
        stream = torch.tensor([stream_id], dtype=torch.long, device=device)
        camera = torch.zeros(1, dtype=torch.long, device=device)
        return (
            self.time_embedding(time)
            + self.stream_embedding(stream)
            + self.camera_embedding(camera)
        ).squeeze(0)


class FunnelTrajectoryHead(nn.Module):
    """Six-layer funnel MLP that predicts a fixed waypoint trajectory."""

    def __init__(
        self,
        hidden_dim: int,
        num_waypoints: int,
        action_dim: int,
        dropout: float,
        use_tanh: bool,
    ) -> None:
        super().__init__()
        output_dim = num_waypoints * action_dim
        self.num_waypoints = num_waypoints
        self.action_dim = action_dim
        self.use_tanh = use_tanh
        self.layers = nn.Sequential(
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, 4096),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(4096, 1024),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(1024, 512),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(512, 256),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(256, 128),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.LayerNorm(128),
            nn.Linear(128, output_dim),
        )

    def forward(self, control_state: torch.Tensor) -> torch.Tensor:
        trajectory = self.layers(control_state)
        if self.use_tanh:
            trajectory = torch.tanh(trajectory)
        return trajectory.view(-1, self.num_waypoints, self.action_dim)


@dataclass
class MiniCPMRobotTrackOutput(ModelOutput):
    """Output of MiniCPM-RobotTrack."""

    loss: Optional[torch.FloatTensor] = None
    trajectories: Optional[torch.FloatTensor] = None


class MiniCPMRobotTrackModel(PreTrainedModel):
    """MiniCPM visual tracking policy with a funnel trajectory head."""

    config_class = MiniCPMRobotTrackConfig
    base_model_prefix = "backbone"
    main_input_name = "input_ids"
    supports_gradient_checkpointing = True
    _supports_sdpa = True
    _supports_flash_attn_2 = True
    _no_split_modules = ["MiniCPMDecoderLayer"]

    def __init__(self, config: MiniCPMRobotTrackConfig) -> None:
        super().__init__(config)

        backbone_config = MiniCPMConfig(**config.backbone_config)
        attention_implementation = getattr(config, "_attn_implementation", None)
        if attention_implementation is not None:
            backbone_config._attn_implementation = attention_implementation
        backbone_config.use_cache = False

        backbone_dtype = _dtype_from_name(config.backbone_dtype)
        with _default_dtype(backbone_dtype):
            self.backbone = MiniCPMModel(backbone_config)

        hidden_dim = int(backbone_config.hidden_size)
        self.vision_projector = VisionProjector(config.vision_feature_dim, hidden_dim)
        self.temporal_markers = TemporalMarkerEncoder(hidden_dim, config.max_time_steps)
        self.control_query = nn.Parameter(torch.empty(1, 1, hidden_dim))
        nn.init.normal_(self.control_query, mean=0.0, std=0.02)
        self.trajectory_head = FunnelTrajectoryHead(
            hidden_dim=hidden_dim,
            num_waypoints=config.num_waypoints,
            action_dim=config.action_dim,
            dropout=config.trajectory_dropout,
            use_tanh=config.use_tanh_actions,
        )

        output_scale = torch.ones(1, 1, config.action_dim, dtype=torch.float32)
        output_scale[..., :2] = config.xy_scale
        self.register_buffer("output_scale", output_scale)

    def get_input_embeddings(self) -> nn.Module:
        return self.backbone.get_input_embeddings()

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.backbone.set_input_embeddings(value)

    def _insert_temporal_markers(
        self,
        tokens: torch.Tensor,
        time_indices: torch.Tensor,
        stream_id: int,
    ) -> torch.Tensor:
        if tokens.ndim != 3 or time_indices.ndim != 2:
            raise ValueError("visual tokens and time indices must have shapes [B, N, C] and [B, N]")
        if tokens.shape[:2] != time_indices.shape:
            raise ValueError("visual token and time-index shapes do not match")
        if tokens.size(1) == 0:
            return tokens

        packed_rows: List[torch.Tensor] = []
        time_rows = time_indices.detach().to("cpu").tolist()
        for batch_index, time_row in enumerate(time_rows):
            pieces: List[torch.Tensor] = []
            start = 0
            while start < len(time_row):
                time_step = int(time_row[start])
                if not 0 <= time_step < self.config.max_time_steps:
                    raise ValueError(f"time index {time_step} is outside the configured range")
                end = start + 1
                while end < len(time_row) and int(time_row[end]) == time_step:
                    end += 1
                marker = self.temporal_markers(time_step, stream_id, tokens.device)
                pieces.extend((marker.unsqueeze(0), tokens[batch_index, start:end]))
                start = end
            packed_rows.append(torch.cat(pieces, dim=0))

        packed_lengths = {row.size(0) for row in packed_rows}
        if len(packed_lengths) != 1:
            raise ValueError("each batch item must contain the same number of represented frames")
        return torch.stack(packed_rows, dim=0)

    def _build_sequence(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor],
        coarse_tokens: torch.Tensor,
        coarse_time_indices: torch.Tensor,
        fine_tokens: torch.Tensor,
        fine_time_indices: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        device = self.control_query.device
        input_ids = input_ids.to(device)
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids, dtype=torch.long)
        else:
            attention_mask = attention_mask.to(device)

        batch_size = coarse_tokens.size(0)
        if fine_tokens.size(0) != batch_size or input_ids.size(0) != batch_size:
            raise ValueError("batch dimensions do not match")

        projector_dtype = _module_dtype(self.vision_projector)
        history = self.vision_projector(coarse_tokens.to(device=device, dtype=projector_dtype))
        current = self.vision_projector(fine_tokens.to(device=device, dtype=projector_dtype))
        history = self._insert_temporal_markers(
            history, coarse_time_indices.to(device), stream_id=0
        )
        current = self._insert_temporal_markers(
            current, fine_time_indices.to(device), stream_id=1
        )
        text = self.backbone.get_input_embeddings()(input_ids)

        control_query = self.control_query.expand(batch_size, -1, -1)
        sequence = torch.cat((text, history, current, control_query), dim=1)
        sequence = sequence.to(dtype=_module_dtype(self.backbone))

        full_attention_mask = torch.cat(
            (
                attention_mask,
                torch.ones(batch_size, history.size(1), dtype=torch.long, device=device),
                torch.ones(batch_size, current.size(1), dtype=torch.long, device=device),
                torch.ones(batch_size, 1, dtype=torch.long, device=device),
            ),
            dim=1,
        )
        return sequence, full_attention_mask

    def normalize_trajectory(self, trajectory: torch.Tensor) -> torch.Tensor:
        return trajectory / self.output_scale.to(device=trajectory.device, dtype=trajectory.dtype)

    def forward(
        self,
        input_ids: torch.LongTensor,
        coarse_tokens: torch.Tensor,
        coarse_time_indices: torch.Tensor,
        fine_tokens: torch.Tensor,
        fine_time_indices: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        valid_mask: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[MiniCPMRobotTrackOutput, Tuple[torch.Tensor, ...]]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        sequence, full_attention_mask = self._build_sequence(
            input_ids=input_ids,
            attention_mask=attention_mask,
            coarse_tokens=coarse_tokens,
            coarse_time_indices=coarse_time_indices,
            fine_tokens=fine_tokens,
            fine_time_indices=fine_time_indices,
        )
        output = self.backbone(
            inputs_embeds=sequence,
            attention_mask=full_attention_mask,
            use_cache=False,
            return_dict=True,
        )
        control_state = output.last_hidden_state[:, -1].to(
            dtype=_module_dtype(self.trajectory_head)
        )
        normalized_trajectory = self.trajectory_head(control_state)
        trajectories = normalized_trajectory * self.output_scale.to(
            normalized_trajectory.dtype
        )

        loss = None
        if labels is not None:
            labels = labels.to(device=trajectories.device, dtype=trajectories.dtype)
            if labels.shape != trajectories.shape:
                raise ValueError("labels and predicted trajectories must have identical shapes")
            normalized_labels = self.normalize_trajectory(labels)
            squared_error = (normalized_trajectory - normalized_labels).square()
            if valid_mask is None:
                loss = squared_error.mean()
            else:
                mask = valid_mask.to(
                    device=trajectories.device, dtype=trajectories.dtype
                ).unsqueeze(-1)
                denominator = mask.sum() * trajectories.size(-1)
                loss = (
                    squared_error.mul(mask).sum() / denominator
                    if denominator.item() > 0
                    else trajectories.sum() * 0.0
                )

        if not return_dict:
            return (trajectories,) if loss is None else (loss, trajectories)
        return MiniCPMRobotTrackOutput(loss=loss, trajectories=trajectories)