"""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)