code/src/readyagent/p4d_belief/models.py

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