| """Shared frame pooling, temporal pooling and reward heads for B1/C0/C1.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| @dataclass(frozen=True) |
| class SpanPoolResult: |
| frame_features: torch.Tensor |
| frame_mask: torch.Tensor |
| span_bounds: tuple[tuple[tuple[int, int], ...], ...] |
|
|
|
|
| def pool_image_token_spans( |
| hidden_states: torch.Tensor, |
| input_ids: torch.Tensor, |
| image_token_id: int, |
| image_grid_thw: torch.Tensor, |
| frames_per_sample: list[int], |
| spatial_merge_size: int, |
| ) -> SpanPoolResult: |
| """Mean-pool each image-token span and return padded per-frame features.""" |
|
|
| if hidden_states.ndim != 3 or input_ids.ndim != 2: |
| raise ValueError("hidden_states/input_ids must have shapes [B,S,D] and [B,S]") |
| if hidden_states.shape[:2] != input_ids.shape: |
| raise ValueError("hidden_states and input_ids batch/sequence dimensions differ") |
| if len(frames_per_sample) != hidden_states.shape[0] or any( |
| count <= 0 for count in frames_per_sample |
| ): |
| raise ValueError( |
| "frames_per_sample must contain one positive count per batch item" |
| ) |
| if sum(frames_per_sample) != image_grid_thw.shape[0]: |
| raise ValueError("frames_per_sample does not match image_grid_thw rows") |
| if spatial_merge_size <= 0: |
| raise ValueError("spatial_merge_size must be positive") |
|
|
| max_frames = max(frames_per_sample) |
| features = hidden_states.new_zeros( |
| (hidden_states.shape[0], max_frames, hidden_states.shape[-1]) |
| ) |
| frame_mask = torch.zeros( |
| (hidden_states.shape[0], max_frames), |
| dtype=torch.bool, |
| device=hidden_states.device, |
| ) |
| all_bounds: list[tuple[tuple[int, int], ...]] = [] |
| grid_offset = 0 |
| merge_area = spatial_merge_size**2 |
|
|
| for batch_index, frame_count in enumerate(frames_per_sample): |
| positions = torch.nonzero( |
| input_ids[batch_index] == image_token_id, as_tuple=False |
| ).flatten() |
| token_offset = 0 |
| bounds: list[tuple[int, int]] = [] |
| for frame_index in range(frame_count): |
| grid_t, grid_h, grid_w = ( |
| int(value) for value in image_grid_thw[grid_offset + frame_index] |
| ) |
| raw_tokens = grid_t * grid_h * grid_w |
| if raw_tokens % merge_area: |
| raise ValueError( |
| "image token grid is not divisible by spatial merge area" |
| ) |
| span_length = raw_tokens // merge_area |
| frame_positions = positions[token_offset : token_offset + span_length] |
| if frame_positions.numel() != span_length: |
| raise ValueError("image token span is empty or truncated") |
| if frame_positions.numel() > 1 and not torch.all( |
| frame_positions[1:] == frame_positions[:-1] + 1 |
| ): |
| raise ValueError("image token span is not contiguous") |
| start = int(frame_positions[0]) |
| stop = int(frame_positions[-1]) + 1 |
| features[batch_index, frame_index] = hidden_states[ |
| batch_index, frame_positions |
| ].mean(dim=0) |
| frame_mask[batch_index, frame_index] = True |
| bounds.append((start, stop)) |
| token_offset += span_length |
| if token_offset != positions.numel(): |
| raise ValueError("unassigned image tokens remain after span pooling") |
| all_bounds.append(tuple(bounds)) |
| grid_offset += frame_count |
|
|
| return SpanPoolResult(features, frame_mask, tuple(all_bounds)) |
|
|
|
|
| class MaskedQueryPool(nn.Module): |
| """Single-query attention pooling over valid ordered-frame features.""" |
|
|
| def __init__(self, hidden_size: int) -> None: |
| super().__init__() |
| self.query = nn.Parameter(torch.empty(hidden_size)) |
| nn.init.normal_(self.query, mean=0.0, std=hidden_size**-0.5) |
|
|
| def forward( |
| self, frame_features: torch.Tensor, frame_mask: torch.Tensor |
| ) -> torch.Tensor: |
| if frame_features.ndim != 3 or frame_mask.shape != frame_features.shape[:2]: |
| raise ValueError("expected frame_features [B,T,D] and frame_mask [B,T]") |
| if not torch.all(frame_mask.any(dim=1)): |
| raise ValueError("every trajectory must contain at least one valid frame") |
| scores = torch.einsum("btd,d->bt", frame_features, self.query) / math.sqrt( |
| frame_features.shape[-1] |
| ) |
| scores = scores.masked_fill(~frame_mask, torch.finfo(scores.dtype).min) |
| weights = torch.softmax(scores.float(), dim=1).to(frame_features.dtype) |
| return torch.einsum("bt,btd->bd", weights, frame_features) |
|
|
|
|
| class MLPHead(nn.Module): |
| def __init__( |
| self, hidden_size: int, output_size: int, dropout: float = 0.1 |
| ) -> None: |
| super().__init__() |
| bottleneck = hidden_size // 2 |
| self.net = nn.Sequential( |
| nn.Linear(hidden_size, bottleneck), |
| nn.LayerNorm(bottleneck), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(bottleneck, output_size), |
| ) |
|
|
| def forward(self, inputs: torch.Tensor) -> torch.Tensor: |
| return self.net(inputs) |
|
|
|
|
| class SharedRewardHeads(nn.Module): |
| """The identical topology used for frozen probes and matched-LoRA runs.""" |
|
|
| def __init__( |
| self, hidden_size: int, progress_bins: int = 10, dropout: float = 0.1 |
| ) -> None: |
| super().__init__() |
| self.temporal_pool = MaskedQueryPool(hidden_size) |
| self.progress_head = MLPHead(hidden_size, progress_bins, dropout) |
| self.success_head = MLPHead(hidden_size, 1, dropout) |
| self.preference_head = MLPHead(hidden_size, 1, dropout) |
|
|
| def forward( |
| self, frame_features: torch.Tensor, frame_mask: torch.Tensor |
| ) -> dict[str, torch.Tensor]: |
| trajectory = self.temporal_pool(frame_features, frame_mask) |
| return { |
| "progress_logits": self.progress_head(frame_features), |
| "success_logits": self.success_head(frame_features).squeeze(-1), |
| "trajectory_features": trajectory, |
| "trajectory_scores": self.preference_head(trajectory).squeeze(-1), |
| } |
|
|
| def preference_logit( |
| self, |
| frame_features_a: torch.Tensor, |
| frame_mask_a: torch.Tensor, |
| frame_features_b: torch.Tensor, |
| frame_mask_b: torch.Tensor, |
| ) -> torch.Tensor: |
| score_a = self.preference_head( |
| self.temporal_pool(frame_features_a, frame_mask_a) |
| ).squeeze(-1) |
| score_b = self.preference_head( |
| self.temporal_pool(frame_features_b, frame_mask_b) |
| ).squeeze(-1) |
| return score_a - score_b |
|
|
