"""Small causal readiness head over frozen visual, state, and text features.""" from __future__ import annotations from typing import Any import torch from torch import nn class ResidualCrossAttentionBlock(nn.Module): def __init__(self, width: int, heads: int, ffn_width: int, dropout: float): super().__init__() self.query_norm = nn.LayerNorm(width) self.visual_norm = nn.LayerNorm(width) self.attention = nn.MultiheadAttention( width, heads, dropout=dropout, batch_first=True ) self.ffn_norm = nn.LayerNorm(width) self.ffn = nn.Sequential( nn.Linear(width, ffn_width), nn.GELU(), nn.Dropout(dropout), nn.Linear(ffn_width, width), nn.Dropout(dropout), ) def forward( self, query: torch.Tensor, visual: torch.Tensor, visual_mask: torch.Tensor ) -> torch.Tensor: normalized_visual = self.visual_norm(visual) update, _ = self.attention( self.query_norm(query), normalized_visual, normalized_visual, key_padding_mask=~visual_mask, need_weights=False, ) query = query + update return query + self.ffn(self.ffn_norm(query)) class CommitReadinessHead(nn.Module): """Predict evidence for committing away from the active semantic phase.""" def __init__( self, *, image_dim: int = 1152, text_dim: int = 4096, state_dim: int = 8, width: int = 256, state_transformer_layers: int = 1, cross_attention_layers: int = 2, attention_heads: int = 4, ffn_width: int = 1024, head_hidden_width: int = 128, dropout: float = 0.1, max_relative_age: int = 2048, state_recent_length: int = 16, ): super().__init__() self.image_dim = int(image_dim) self.text_dim = int(text_dim) self.state_dim = int(state_dim) self.width = int(width) self.state_recent_length = int(state_recent_length) self.visual_projection = nn.Linear(image_dim, width) self.visual_role_embedding = nn.Embedding(4, width) self.visual_view_embedding = nn.Embedding(2, width) self.visual_spatial_embedding = nn.Embedding(16, width) self.visual_age_embedding = nn.Embedding(max_relative_age + 1, width) self.state_projection = nn.Linear(state_dim, width) self.state_role_embedding = nn.Embedding(2, width) self.state_time_embedding = nn.Embedding(state_recent_length, width) state_layer = nn.TransformerEncoderLayer( d_model=width, nhead=attention_heads, dim_feedforward=ffn_width, dropout=dropout, activation="gelu", batch_first=True, norm_first=True, ) self.state_transformer = nn.TransformerEncoder( state_layer, num_layers=state_transformer_layers, enable_nested_tensor=False ) self.state_fusion = nn.Sequential( nn.Linear(2 * width, width), nn.GELU(), nn.Dropout(dropout) ) self.text_projection = nn.Linear(text_dim, width) self.query = nn.Parameter(torch.empty(1, 1, width)) nn.init.normal_(self.query, std=0.02) self.cross_attention = nn.ModuleList( [ ResidualCrossAttentionBlock(width, attention_heads, ffn_width, dropout) for _ in range(cross_attention_layers) ] ) self.output_norm = nn.LayerNorm(width) self.output = nn.Sequential( nn.Linear(width, head_hidden_width), nn.GELU(), nn.Dropout(dropout), nn.Linear(head_hidden_width, 1), ) @classmethod def from_config(cls, cfg: Any) -> "CommitReadinessHead": return cls( image_dim=int(cfg.features.image_dim), text_dim=int(cfg.features.text_dim), state_dim=int(cfg.features.state_dim), width=int(cfg.model.width), state_transformer_layers=int(cfg.model.state_transformer_layers), cross_attention_layers=int(cfg.model.cross_attention_layers), attention_heads=int(cfg.model.attention_heads), ffn_width=int(cfg.model.ffn_width), head_hidden_width=int(cfg.model.head_hidden_width), dropout=float(cfg.model.dropout), max_relative_age=int(cfg.features.max_relative_age), state_recent_length=int(cfg.features.state_recent_length), ) def forward(self, **features: torch.Tensor) -> dict[str, torch.Tensor]: visual_tokens = features["visual_tokens"] visual_mask = features["visual_mask"].bool() batch, token_count, image_dim = visual_tokens.shape if token_count != 608 or image_dim != self.image_dim: raise ValueError( f"Expected visual [B,608,{self.image_dim}], got {tuple(visual_tokens.shape)}" ) if ( visual_mask.shape != (batch, token_count) or not visual_mask.any(dim=1).all() ): raise ValueError("Each sample must have at least one valid visual token") visual = ( self.visual_projection(visual_tokens) + self.visual_role_embedding(features["visual_role_ids"].long()) + self.visual_view_embedding(features["visual_view_ids"].long()) + self.visual_spatial_embedding(features["visual_spatial_ids"].long()) + self.visual_age_embedding(features["visual_age_ids"].long()) ) active_state = features["active_start_state"] recent_states = features["recent_states"] recent_mask = features["recent_state_mask"].bool() if active_state.shape != (batch, self.state_dim): raise ValueError( f"Bad active-start state shape: {tuple(active_state.shape)}" ) if recent_states.shape != (batch, self.state_recent_length, self.state_dim): raise ValueError(f"Bad recent state shape: {tuple(recent_states.shape)}") if ( recent_mask.shape != (batch, self.state_recent_length) or not recent_mask.any(dim=1).all() ): raise ValueError("Each sample must have at least one valid recent state") active_summary = ( self.state_projection(active_state) + self.state_role_embedding.weight[0] ) time_ids = torch.arange(self.state_recent_length, device=recent_states.device) recent = ( self.state_projection(recent_states) + self.state_role_embedding.weight[1] + self.state_time_embedding(time_ids)[None] ) recent = self.state_transformer(recent, src_key_padding_mask=~recent_mask) recent_summary = (recent * recent_mask.unsqueeze(-1)).sum( dim=1 ) / recent_mask.sum(dim=1, keepdim=True).clamp_min(1) state_summary = self.state_fusion( torch.cat([active_summary, recent_summary], dim=-1) ) text = features["text_embedding"] if text.shape != (batch, self.text_dim): raise ValueError(f"Bad text embedding shape: {tuple(text.shape)}") text_summary = self.text_projection(text) query = ( self.query.expand(batch, -1, -1) + state_summary[:, None] + text_summary[:, None] ) for block in self.cross_attention: query = block(query, visual, visual_mask) commit_logit = self.output(self.output_norm(query[:, 0])).squeeze(-1) return { "commit_logit": commit_logit, "commit_score": torch.sigmoid(commit_logit), } def trainable_parameter_count(model: nn.Module) -> int: return sum( parameter.numel() for parameter in model.parameters() if parameter.requires_grad )