code/stereo_core/pair_residual_model.py

"""PAIR: Predictive Agent-Interaction Representation.

The deployment policy is strictly local.  Synchronised peer actions and future
local effects are available only to the training-time teacher.  Unlike the
previous predictability pilot, PAIR never changes Local-ARCA's action-role
router.  A zero-initialised low-rank residual adapter can only modulate the
decoded action tokens, so update zero is exactly the measured Local-ARCA policy.
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F

from stereo_decoder_variants import StereoARCA


class StereoPAIRResidual(StereoARCA):
    """PAIR-conditioned interaction residual over a Local-ARCA policy.

    ``joint_training=False`` reproduces the conservative frozen-checkpoint
    probe.  ``joint_training=True`` is the formal from-scratch PAIR policy:
    the same local observation/action path, local interaction student and
    residual adapter are optimized jointly while the synchronized teacher
    remains training-only.
    """

    def __init__(
        self, *args, event_dim=96, residual_rank=32, joint_training=False, **kwargs
    ):
        super().__init__(*args, **kwargs)
        d = self.query.shape[-1]
        self.event_dim = event_dim
        self.residual_rank = residual_rank
        self.local_event_head = nn.Sequential(
            nn.LayerNorm(d),
            nn.Linear(d, d),
            nn.GELU(),
            nn.Linear(d, event_dim),
        )
        self.confidence_head = nn.Sequential(
            nn.LayerNorm(d),
            nn.Linear(d, d // 2),
            nn.GELU(),
            nn.Linear(d // 2, 1),
        )
        self.residual_down = nn.Linear(d, residual_rank, bias=False)
        self.event_scale = nn.Linear(event_dim, residual_rank)
        self.query_scale = nn.Parameter(torch.zeros(self.query.shape[1], residual_rank))
        self.residual_up = nn.Linear(residual_rank, d, bias=False)
        # Hard baseline-preservation property: before training, PAIR produces
        # exactly Local-ARCA's decoded hidden state and action chunk.
        nn.init.zeros_(self.residual_up.weight)
        nn.init.zeros_(self.event_scale.weight)
        nn.init.zeros_(self.event_scale.bias)
        self.last_local_event = None
        self.last_confidence = None
        self.last_gates = None
        self.last_residual_norm = None
        self.joint_training = bool(joint_training)

    def freeze_local_arca(self):
        """Freeze every parameter inherited from Local-ARCA.

        New PAIR modules remain trainable.  This makes role-route preservation
        structural rather than a soft regularisation preference.
        """
        new_prefixes = (
            "local_event_head.",
            "confidence_head.",
            "residual_down.",
            "event_scale.",
            "residual_up.",
        )
        for name, parameter in self.named_parameters():
            parameter.requires_grad = name == "query_scale" or name.startswith(new_prefixes)

    def forward(self, image, depth_mm, qpos, actions=None):
        def base_forward():
            observation = self._rgbd_tokens(image, depth_mm)
            state_vector = self.state(qpos)
            gates, auxiliary = self._route(state_vector, observation, image.shape[0])
            if actions is not None:
                hidden = self.posterior(self.action(actions) + self.pos)
                mu, logvar = self.latent(hidden.mean(1)).chunk(2, -1)
                logvar = logvar.clamp(-10.0, 5.0)
                latent = mu + torch.randn_like(mu) * torch.exp(0.5 * logvar)
            else:
                mu = logvar = None
                latent = torch.zeros(
                    (image.shape[0], self.z_proj.in_features), device=image.device
                )
            memory = torch.cat(
                (
                    state_vector.unsqueeze(1),
                    self.z_proj(latent).unsqueeze(1),
                    observation,
                ),
                dim=1,
            )
            decoded = self.decoder(
                self.query.expand(image.shape[0], -1, -1),
                memory,
                observation,
                gates,
            )
            base_action = self.out(decoded)

            return observation, state_vector, gates, auxiliary, decoded, base_action, mu, logvar

        # The conservative probe structurally freezes Local-ARCA.  The formal
        # model executes the identical graph with gradients enabled.
        if self.joint_training:
            observation, state_vector, gates, auxiliary, decoded, base_action, mu, logvar = base_forward()
        else:
            with torch.no_grad():
                observation, state_vector, gates, auxiliary, decoded, base_action, mu, logvar = base_forward()

        context = (
            self.route_state(state_vector if self.joint_training else state_vector.detach())
            + self.route_observation(
                (observation if self.joint_training else observation.detach()).mean(1)
            )
        )
        local_event = F.normalize(
            self.local_event_head(context).float(), dim=-1, eps=1e-6
        )
        confidence = torch.sigmoid(self.confidence_head(context).float())
        scale = torch.sigmoid(
            self.event_scale(local_event).unsqueeze(1)
            + self.query_scale.unsqueeze(0)
        )
        residual_hidden = self.residual_up(
            F.gelu(
                self.residual_down(decoded if self.joint_training else decoded.detach())
            )
            * scale.to(decoded.dtype)
        )
        residual_hidden = confidence.to(decoded.dtype).unsqueeze(1) * residual_hidden
        action = self.out(
            (decoded if self.joint_training else decoded.detach()) + residual_hidden
        )

        self.last_local_event = local_event
        self.last_confidence = confidence
        self.last_gates = gates
        # Epsilon avoids the undefined d(sqrt(x))/dx at the deliberately
        # zero-initialised residual while preserving the reported RMS value.
        self.last_residual_norm = (
            residual_hidden.float().square().mean() + 1e-12
        ).sqrt()
        return (
            action,
            mu,
            logvar,
            auxiliary,
            local_event,
            confidence,
            base_action,
            gates,
        )


class PairwiseLagInteractionTeacher(nn.Module):
    """Training-only pairwise, directed and lag-selective interaction teacher.

    Each target arm attends over all (peer, lag) candidates in its synchronised
    unordered team.  The target's future *local* effect is privileged training
    supervision, never a deployment input.
    """

    def __init__(
        self,
        action_dim,
        effect_dim,
        lags,
        event_dim=96,
        width=192,
    ):
        super().__init__()
        self.effect_dim = effect_dim
        self.lags = tuple(int(value) for value in lags)
        self.action_step = nn.Sequential(
            nn.Linear(action_dim, width),
            nn.GELU(),
            nn.Linear(width, width),
        )
        self.action_feature = nn.Sequential(
            nn.LayerNorm(2 * width),
            nn.Linear(2 * width, width),
            nn.GELU(),
            nn.Linear(width, width),
        )
        self.effect_feature = nn.Sequential(
            nn.LayerNorm(effect_dim),
            nn.Linear(effect_dim, width),
            nn.GELU(),
            nn.Linear(width, width),
        )
        self.lag_embedding = nn.Embedding(len(self.lags), width)
        self.pair_value = nn.Sequential(
            nn.LayerNorm(4 * width),
            nn.Linear(4 * width, width),
            nn.GELU(),
            nn.Linear(width, event_dim),
        )
        self.pair_score = nn.Sequential(
            nn.LayerNorm(4 * width),
            nn.Linear(4 * width, width),
            nn.GELU(),
            nn.Linear(width, 1),
        )
        self.own_effect = nn.Sequential(
            nn.Linear(width, width),
            nn.GELU(),
            nn.Linear(width, len(self.lags) * effect_dim),
        )
        self.interaction_effect = nn.Sequential(
            nn.Linear(event_dim, width),
            nn.GELU(),
            nn.Linear(width, len(self.lags) * effect_dim),
        )

    def action_features(self, actions):
        delta = actions[:, 1:] - actions[:, :-1]
        encoded = self.action_step(delta.float())
        return self.action_feature(
            torch.cat((encoded.mean(1), encoded.square().mean(1).sqrt()), dim=-1)
        )

    def forward(self, actions, future_effects, groups):
        """Return event, predictive-gain confidence target and diagnostics."""
        own = self.action_features(actions)
        effect_features = self.effect_feature(future_effects.float())
        events = torch.zeros(
            (len(actions), self.interaction_effect[0].in_features),
            device=actions.device,
            dtype=torch.float32,
        )
        attention_entropy = []
        for group in groups.unique(sorted=True):
            ids = (groups == group).nonzero(as_tuple=False).flatten()
            values = own.index_select(0, ids)
            group_effects = effect_features.index_select(0, ids)
            for local_index, target_id in enumerate(ids):
                candidates, scores = [], []
                target = values[local_index]
                peer_indices = [
                    index for index in range(len(ids)) if index != local_index
                ]
                if not peer_indices:
                    continue
                for peer_index in peer_indices:
                    peer = values[peer_index]
                    for lag_index in range(len(self.lags)):
                        effect = group_effects[local_index, lag_index]
                        lag = self.lag_embedding.weight[lag_index]
                        pair = torch.cat(
                            (target, peer, target - peer, effect + lag), dim=-1
                        )
                        candidates.append(self.pair_value(pair))
                        scores.append(self.pair_score(pair))
                score = torch.cat(scores, dim=0)
                weight = score.softmax(dim=0)
                candidate = torch.stack(candidates, dim=0)
                event = (weight.unsqueeze(-1) * candidate).sum(0)
                events[target_id] = event.float()
                attention_entropy.append(
                    -(weight.clamp_min(1e-8).log() * weight).sum()
                )

        teacher_event = F.normalize(events, dim=-1, eps=1e-6)
        target = future_effects.float().flatten(1)
        own_prediction = self.own_effect(own)
        full_prediction = own_prediction + self.interaction_effect(teacher_event)
        own_error = (own_prediction - target).square().mean(-1)
        full_error = (full_prediction - target).square().mean(-1)
        # Only the locally consequential part receives a strong deployment
        # confidence target.  Detaching prevents the teacher from gaming it.
        gain = ((own_error - full_error) / own_error.clamp_min(1e-6)).clamp(0, 1).detach()
        zero = own_error.new_zeros(())
        entropy = torch.stack(attention_entropy).mean() if attention_entropy else zero
        return (
            teacher_event,
            gain.unsqueeze(-1),
            own_error.mean(),
            full_error.mean(),
            entropy,
        )


def interaction_alignment(student_event, teacher_event, temperature=0.15):
    student = F.normalize(student_event.float(), dim=-1, eps=1e-6)
    teacher = F.normalize(teacher_event.detach().float(), dim=-1, eps=1e-6)
    logits = student @ teacher.t() / temperature
    target = torch.arange(len(student), device=student.device)
    return 0.5 * (
        F.cross_entropy(logits, target) + F.cross_entropy(logits.t(), target)
    )