"""Query-conditioned continuous-time pointer models for P4D-Belief.""" from __future__ import annotations import math from dataclasses import dataclass import torch from torch import nn from readyagent.p4d_belief.data import Catalog @dataclass(frozen=True) class ModelConfig: hidden_dim: int = 128 layers: int = 3 heads: int = 4 dropout: float = 0.15 use_instance_identity: bool = False use_compatibility: bool = False class SemanticLocationEncoder(nn.Module): """Shared location encoder with no scene-specific state-ID embedding.""" def __init__(self, catalog: Catalog, hidden_dim: int) -> None: super().__init__() region_count = len(catalog.schema["region_category_to_id"]) receptacle_count = len(catalog.schema["receptacle_category_to_id"]) semantic_dim = hidden_dim // 4 self.region_embedding = nn.Embedding(region_count, semantic_dim) self.receptacle_embedding = nn.Embedding(receptacle_count, semantic_dim) self.project = nn.Sequential( nn.Linear(semantic_dim * 2 + 4, hidden_dim), nn.GELU(), nn.LayerNorm(hidden_dim), ) center = catalog.center_xyz.float() known = ~catalog.is_unknown.bool() mean = center[known].mean(dim=0) std = center[known].std(dim=0).clamp_min(1e-3) self.register_buffer("region_category", catalog.region_category.long()) self.register_buffer("receptacle_category", catalog.receptacle_category.long()) self.register_buffer("center", center) self.register_buffer("center_mean", mean) self.register_buffer("center_std", std) self.register_buffer("is_unknown", catalog.is_unknown.float()) def base_nodes(self) -> torch.Tensor: normalized_center = (self.center - self.center_mean) / self.center_std return self.project( torch.cat( ( self.region_embedding(self.region_category), self.receptacle_embedding(self.receptacle_category), normalized_center, self.is_unknown[:, None], ), dim=-1, ) ) class RelativeTimeEncoderLayer(nn.Module): """Transformer layer with a learned, per-head continuous-time bias.""" def __init__(self, hidden_dim: int, heads: int, dropout: float) -> None: super().__init__() self.heads = heads self.time_bias = nn.Linear(2, heads, bias=False) self.attention = nn.MultiheadAttention( hidden_dim, heads, dropout=dropout, batch_first=True ) self.norm1 = nn.LayerNorm(hidden_dim) self.norm2 = nn.LayerNorm(hidden_dim) self.feed_forward = nn.Sequential( nn.Linear(hidden_dim, hidden_dim * 4), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden_dim * 4, hidden_dim), nn.Dropout(dropout), ) self.dropout = nn.Dropout(dropout) def forward( self, tokens: torch.Tensor, times_days: torch.Tensor, valid: torch.Tensor ) -> torch.Tensor: delta_hours = (times_days[:, :, None] - times_days[:, None, :]) * 24.0 time_features = torch.stack( (torch.sign(delta_hours), torch.log1p(delta_hours.abs())), dim=-1 ) bias = self.time_bias(time_features).permute(0, 3, 1, 2) key_mask = (~valid)[:, None, None, :] bias = bias.masked_fill(key_mask, -1e4) batch, length = valid.shape attention_mask = bias.reshape(batch * self.heads, length, length) attended, _ = self.attention( tokens, tokens, tokens, attn_mask=attention_mask, need_weights=False ) tokens = self.norm1(tokens + self.dropout(attended)) return self.norm2(tokens + self.feed_forward(tokens)) class HistoryInputs(nn.Module): def __init__(self, catalog: Catalog, config: ModelConfig) -> None: super().__init__() hidden = config.hidden_dim category_count = len(catalog.schema["category_to_id"]) instance_count = len(catalog.schema.get("instance_uuid_to_id", {})) self.location = SemanticLocationEncoder(catalog, hidden) self.object_embedding = nn.Embedding(category_count, hidden) self.instance_embedding = ( nn.Embedding(instance_count, hidden) if config.use_instance_identity and instance_count else None ) self.event_embedding = nn.Embedding(3, hidden) self.weekday_embedding = nn.Embedding(7, hidden // 4) self.quality_projection = nn.Sequential( nn.Linear(6, hidden // 2), nn.GELU(), nn.Linear(hidden // 2, hidden) ) self.time_projection = nn.Linear(5, hidden) self.query_projection = nn.Sequential( nn.Linear(hidden + hidden // 4 + 4, hidden), nn.GELU(), nn.LayerNorm(hidden), ) self.history_summary_projection = nn.Sequential( nn.Linear(7, hidden), nn.GELU(), nn.Linear(hidden, hidden) ) self.observable_context_projection = nn.Sequential( nn.Linear(category_count + 5, hidden), nn.GELU(), nn.Linear(hidden, hidden), ) self.query_context_norm = nn.LayerNorm(hidden) self.candidate_geometry = nn.Sequential( nn.Linear(5, hidden), nn.GELU(), nn.LayerNorm(hidden) ) self.token_norm = nn.LayerNorm(hidden) def forward( self, batch: dict[str, torch.Tensor] ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: base_nodes = self.location.base_nodes() event_type = batch["event_type"].long() history_mask = batch["history_mask"].bool() observed = batch["observed_state_id"].long() inspected = batch["candidate_state_id"].long() location_id = torch.where(event_type == 1, observed, inspected).clamp_min(0) location_token = base_nodes[location_id] query_days = batch["query_time_days"].float() event_days = batch["event_time_days"].float() age_hours = ((query_days[:, None] - event_days) * 24.0).clamp_min(0.0) time_features = torch.stack( ( torch.log1p(age_hours), torch.sin(2 * math.pi * event_days), torch.cos(2 * math.pi * event_days), torch.sin(2 * math.pi * event_days / 7.0), torch.cos(2 * math.pi * event_days / 7.0), ), dim=-1, ) object_token = self.object_embedding(batch["target_category_id"].long()) if self.instance_embedding is not None: object_token = object_token + self.instance_embedding( batch["target_instance_id"].long() ) tokens = self.token_norm( location_token + self.event_embedding(event_type) + self.quality_projection(batch["evidence_features"].float()) + self.time_projection(time_features) + object_token[:, None, :] ) query_features = torch.cat( ( object_token, self.weekday_embedding(batch["query_weekday_id"].long()), batch["query_time_of_day_sin_cos"].float(), torch.log1p(batch["elapsed_since_last_positive_days"].float())[:, None], torch.ones_like(query_days)[:, None], ), dim=-1, ) positive = (event_type == 1) & history_mask inspection = (event_type == 2) & history_mask positive_count = positive.sum(dim=1).float() inspection_count = inspection.sum(dim=1).float() state_one_hot = torch.nn.functional.one_hot( observed.clamp_min(0), num_classes=base_nodes.shape[0] ).bool() distinct_states = (state_one_hot & positive[:, :, None]).any(dim=1).sum(dim=1).float() previous = torch.full( (event_type.shape[0],), -1, dtype=torch.long, device=event_type.device ) change_count = torch.zeros_like(positive_count) for position in range(event_type.shape[1]): current = observed[:, position] is_positive = positive[:, position] change_count += ( is_positive & (previous >= 0) & (current != previous) ).float() previous = torch.where(is_positive, current, previous) inspected_time = torch.where( inspection, event_days, torch.full_like(event_days, -1e9) ).max(dim=1).values time_since_inspection = torch.where( inspection.any(dim=1), (query_days - inspected_time).clamp_min(0.0), batch["elapsed_since_last_positive_days"].float(), ) negative_strength = batch["evidence_features"].float()[:, :, 5] negative_sum = (negative_strength * inspection.float()).sum(dim=1) negative_mean = negative_sum / inspection_count.clamp_min(1.0) negative_max = negative_strength.masked_fill(~inspection, 0.0).max(dim=1).values summary = torch.stack( ( torch.log1p(positive_count), torch.log1p(inspection_count), torch.log1p(distinct_states), torch.log1p(change_count), torch.log1p(time_since_inspection * 24.0), negative_mean, negative_max, ), dim=-1, ) last_node = base_nodes[batch["last_state"].long()] observable_context = torch.zeros_like(last_node) if "context_mask" in batch: context_mask = batch["context_mask"].bool() context_age_hours = ( (query_days[:, None] - batch["context_time_days"].float()) * 24.0 ).clamp_min(0.0) context_features = torch.cat( ( batch["context_category_counts"].float(), batch["context_observation_features"].float(), torch.log1p(context_age_hours)[:, :, None], ), dim=-1, ) projected_context = self.observable_context_projection(context_features) observable_context = ( projected_context * context_mask[:, :, None].float() ).sum(dim=1) / context_mask.sum(dim=1).clamp_min(1)[:, None] query_token = self.query_context_norm( self.query_projection(query_features) + last_node + self.history_summary_projection(summary) + observable_context ) return tokens, query_token, base_nodes, history_mask def candidate_nodes( self, base_nodes: torch.Tensor, batch: dict[str, torch.Tensor] ) -> torch.Tensor: candidate_ids = batch["candidate_state_ids"].long().clamp_min(0) nodes = base_nodes[candidate_ids] centers = self.location.center[candidate_ids] last_center = self.location.center[batch["last_state"].long()] relative = (centers - last_center[:, None, :]) / self.location.center_std distance = torch.linalg.vector_norm(relative, dim=-1, keepdim=True) is_last = ( candidate_ids == batch["last_state"].long()[:, None] ).float()[:, :, None] # Geometry is added without introducing a scene-specific state embedding. geometry = torch.cat((relative, distance, is_last), dim=-1) return nodes + self.candidate_geometry(geometry) class CTTransformerBackbone(nn.Module): def __init__(self, catalog: Catalog, config: ModelConfig) -> None: super().__init__() self.inputs = HistoryInputs(catalog, config) self.layers = nn.ModuleList( RelativeTimeEncoderLayer(config.hidden_dim, config.heads, config.dropout) for _ in range(config.layers) ) def forward( self, batch: dict[str, torch.Tensor] ) -> tuple[torch.Tensor, torch.Tensor]: history, query, base_nodes, history_mask = self.inputs(batch) tokens = torch.cat((history, query[:, None, :]), dim=1) query_valid = torch.ones( (history_mask.shape[0], 1), dtype=torch.bool, device=history_mask.device ) valid = torch.cat((history_mask, query_valid), dim=1) times = torch.cat( (batch["event_time_days"].float(), batch["query_time_days"].float()[:, None]), dim=1, ) for layer in self.layers: tokens = layer(tokens, times, valid) candidates = self.inputs.candidate_nodes(base_nodes, batch) return tokens[:, -1], candidates class GRUBackbone(nn.Module): def __init__(self, catalog: Catalog, config: ModelConfig) -> None: super().__init__() self.inputs = HistoryInputs(catalog, config) self.gru = nn.GRU( config.hidden_dim, config.hidden_dim, num_layers=2, dropout=config.dropout, batch_first=True, ) def forward( self, batch: dict[str, torch.Tensor] ) -> tuple[torch.Tensor, torch.Tensor]: history, query, base_nodes, history_mask = self.inputs(batch) lengths = history_mask.sum(dim=1).long() sequence = torch.zeros( history.shape[0], history.shape[1] + 1, history.shape[2], device=history.device, dtype=history.dtype, ) sequence[:, : history.shape[1]] = history sequence[torch.arange(len(lengths), device=history.device), lengths] = query packed = nn.utils.rnn.pack_padded_sequence( sequence, (lengths + 1).cpu(), batch_first=True, enforce_sorted=False ) _, hidden = self.gru(packed) candidates = self.inputs.candidate_nodes(base_nodes, batch) return hidden[-1], candidates class PointerBase(nn.Module): def __init__(self, backbone: nn.Module, hidden_dim: int) -> None: super().__init__() self.backbone = backbone self.query_pointer = nn.Linear(hidden_dim, hidden_dim, bias=False) self.candidate_pointer = nn.Linear(hidden_dim, hidden_dim, bias=False) self.scale = hidden_dim**-0.5 def pointer_logits( self, query: torch.Tensor, candidates: torch.Tensor, mask: torch.Tensor ) -> torch.Tensor: logits = torch.einsum( "bd,bmd->bm", self.query_pointer(query), self.candidate_pointer(candidates) ) * self.scale return logits.masked_fill(~mask.bool(), -1e4) class DirectPointer(PointerBase): """Direct softmax comparator using the same CT Transformer and pointer.""" def __init__(self, catalog: Catalog, config: ModelConfig) -> None: super().__init__(CTTransformerBackbone(catalog, config), config.hidden_dim) def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: query, candidates = self.backbone(batch) logits = self.pointer_logits(query, candidates, batch["candidate_mask"]) probabilities = torch.softmax(logits, dim=-1) last_probability = probabilities.gather( 1, batch["last_candidate_index"].long()[:, None] ).squeeze(1) return {"probabilities": probabilities, "rho": last_probability, "logits": logits} class GRUDirectPointer(PointerBase): def __init__(self, catalog: Catalog, config: ModelConfig) -> None: super().__init__(GRUBackbone(catalog, config), config.hidden_dim) def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: query, candidates = self.backbone(batch) logits = self.pointer_logits(query, candidates, batch["candidate_mask"]) probabilities = torch.softmax(logits, dim=-1) rho = probabilities.gather( 1, batch["last_candidate_index"].long()[:, None] ).squeeze(1) return {"probabilities": probabilities, "rho": rho, "logits": logits} class P4DBelief(PointerBase): """Persistence-relocation factorized pointer with a single final-state NLL.""" def __init__(self, catalog: Catalog, config: ModelConfig) -> None: super().__init__(CTTransformerBackbone(catalog, config), config.hidden_dim) self.persistence = nn.Linear(config.hidden_dim, 1) self.compatibility = ( nn.Sequential( nn.Linear(config.hidden_dim * 2, config.hidden_dim), nn.GELU(), nn.Linear(config.hidden_dim, 1), ) if config.use_compatibility else None ) def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: query, candidates = self.backbone(batch) rho = torch.sigmoid(self.persistence(query)).squeeze(-1) relocation_mask = batch["candidate_mask"].bool().clone() relocation_mask.scatter_( 1, batch["last_candidate_index"].long()[:, None], False ) relocation_logits = self.pointer_logits(query, candidates, relocation_mask) if self.compatibility is not None: paired = torch.cat( (query[:, None, :].expand_as(candidates), candidates), dim=-1 ) relocation_logits = relocation_logits + self.compatibility(paired).squeeze(-1) relocation_logits = relocation_logits.masked_fill(~relocation_mask, -1e4) relocation = torch.softmax(relocation_logits, dim=-1) probabilities = relocation * (1.0 - rho[:, None]) probabilities.scatter_( 1, batch["last_candidate_index"].long()[:, None], rho[:, None] ) probabilities = probabilities * batch["candidate_mask"].float() return { "probabilities": probabilities, "rho": rho, "relocation_probabilities": relocation, "relocation_mask": relocation_mask, "context": query, "candidates": candidates, } def state_nll(probabilities: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: selected = probabilities.gather(1, targets.long()[:, None]).squeeze(1) return -torch.log(selected.clamp_min(1e-9)).mean() def bayesian_update( prior: torch.Tensor, checked_candidate: torch.Tensor, *, observed: torch.Tensor, confidence: torch.Tensor, coverage: torch.Tensor, ) -> torch.Tensor: """Apply one parameter-free visibility-aware observation update.""" likelihood = torch.ones_like(prior) strength = (confidence * coverage).clamp(0.0, 1.0) rows = torch.arange(prior.shape[0], device=prior.device) not_seen_likelihood = (1.0 - strength).clamp_min(1e-6) likelihood[rows, checked_candidate.long()] = torch.where( observed.bool(), torch.ones_like(strength), not_seen_likelihood, ) if observed.any(): detected_rows = rows[observed.bool()] likelihood[detected_rows] = 1e-6 likelihood[detected_rows, checked_candidate[observed.bool()].long()] = strength[ observed.bool() ].clamp_min(1e-6) posterior = prior * likelihood return posterior / posterior.sum(dim=-1, keepdim=True).clamp_min(1e-12)