MoFlow/models/backbone_graph.py

"""
MotionTransformerGraph: extends MotionTransformer with a
FutureInteractionGraph module inserted after the per-agent self-attention.

Two-pass forward design:
  Pass 1 (torch.no_grad): run the full model with y_t as the graph edge
    source and no sigma weighting → produces y_0_hat and logvar_.
  Pass 2 (grad enabled): run the identical model again, now using y_0_hat
    from pass 1 as the graph edge source, and sigma derived from logvar_ to
    soft-weight the pairwise agent interaction (certain → uncertain).

Uncertainty:
  A logvar_head (same MLP structure as reg_head) predicts per-agent
  per-timestep log-variance [B, K, A, T*2].  From this, a per-agent
  uncertainty scalar sigma = sqrt(exp(logvar).mean()) is computed and
  passed to the graph as the directional weight:
      w_ij = sigmoid(γ * (σ_i − σ_j) + 0.5)
  so uncertain agents receive more information from certain neighbors.
  The NLL loss on logvar is applied in flow_matching.py.
"""

import math
import torch
import torch.nn as nn
from einops import rearrange, repeat

from models.backbone import MotionTransformer
from models.graph_interaction_nba_v6 import FutureInteractionGraphV6 as FutureInteractionGraph
from models.utils.common_layers import build_mlps
from utils.normalization import unnormalize_min_max, unnormalize_sqrt


class MotionTransformerGraph(MotionTransformer):
    """MotionTransformer augmented with a future-interaction graph module
    and a per-agent uncertainty head.

    Extra constructor kwargs:
        graph_num_gnn_layers (int, default 2)
        graph_dropout        (float, default 0.1)
    """

    def __init__(self, model_config, logger, config,
                 graph_num_gnn_layers: int = 2,
                 graph_dropout: float = 0.1):
        super().__init__(model_config, logger, config)

        self.T_future  = config.future_frames  # 20
        self.A         = config.agents          # 11
        self.data_norm = config.get('data_norm', 'min_max')

        D        = self.dim   # 128
        time_dim = D

        self.future_graph = FutureInteractionGraph(
            embed_dim      = D,
            future_steps   = self.T_future,
            num_agents     = self.A,
            num_heads      = 4,
            dropout        = graph_dropout,
            num_gnn_layers = graph_num_gnn_layers,
            time_dim       = time_dim,
        )

        # Uncertainty head — same MLP structure as reg_head.
        # Predicts log-variance [B, K, A, T*2] from readout_token.
        self.logvar_head = build_mlps(
            c_in         = self.dim,
            mlp_channels = self.model_cfg.REGRESSION_MLPS,
            ret_before_act = True,
            without_norm   = True,
        )

        params_graph = sum(p.numel() for p in self.future_graph.parameters())
        params_logvar = sum(p.numel() for p in self.logvar_head.parameters())
        logger.info("FutureInteractionGraph parameters: {:,}".format(params_graph))
        logger.info("LogvarHead parameters: {:,}".format(params_logvar))

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _unnormalize_y(self, y_norm: torch.Tensor) -> torch.Tensor:
        """Unnormalize [B, K, A, T, 2] from training norm back to metres."""
        if self.data_norm == 'min_max':
            return unnormalize_min_max(
                y_norm,
                self.config.fut_traj_min,
                self.config.fut_traj_max,
                -1, 1,
            )
        elif self.data_norm == 'sqrt':
            sqrt_a_ = torch.tensor(
                [self.config.sqrt_x_a, self.config.sqrt_y_a],
                device=y_norm.device,
            )
            sqrt_b_ = torch.tensor(
                [self.config.sqrt_x_b, self.config.sqrt_y_b],
                device=y_norm.device,
            )
            return unnormalize_sqrt(y_norm, sqrt_a_, sqrt_b_)
        else:
            return y_norm

    # ------------------------------------------------------------------
    # Shared single-pass implementation
    # ------------------------------------------------------------------

    def _forward_impl(self, y, time, x_data,
                      y_0_for_graph=None,
                      sigma_for_graph=None,
                      skip_graph=False):
        """Single forward pass.

        y_0_for_graph:  [B, K, A, T*2] normalized, or None.
            None  → graph edge features built from y_t (noisy).
            Given → graph edge features built from this cleaner prediction.
        sigma_for_graph: [B, K, A] per-agent uncertainty scalar, or None.
            None  → graph uses scalar denoising tau for directional weight.
            Given → graph uses sigma to promote certain→uncertain flow.
        skip_graph: if True, skip the future interaction graph entirely.

        Returns: (denoiser_x [B,K,A,T*2], denoiser_cls [B,K,A],
                  logvar [B,K,A,T*2])
        """
        # ---- Shape normalisation (variable-A compatible) -----------------
        # Accept [B, K, A, T_future, 2] or [B, K, A, T_future*2] with A taken
        # from the input, not from self.A (supports SDD padded batches).
        if y.dim() == 5 and y.size(-1) == 2 and y.size(-2) == self.T_future:
            B, K, A = y.size(0), y.size(1), y.size(2)
            y = y.reshape(B, K, A, self.T_future * 2)
        else:
            assert y.size(-1) == self.T_future * 2, \
                f"Unexpected y shape: {y.shape}"
        device = y.device
        B, K, A, _ = y.shape

        # ---- Context encoder (past trajectories) -------------------------
        agent_type     = self.config.get('agent_type', 'sport')
        agent_mask_ctx = x_data.get('agent_mask', None) if isinstance(x_data, dict) else None
        encoder_out = self.context_encoder(
            x_data['past_traj_original_scale'],
            agent_type=agent_type,
            agent_mask=agent_mask_ctx,
        )                                               # [B, A, D]
        encoder_out_batch = repeat(
            encoder_out, 'b a d -> b k a d', k=K, a=A
        )                                               # [B, K, A, D]

        # ---- Noisy-y embedding -------------------------------------------
        y_emb = self.noisy_y_mlp(y)                     # [B, K, A, D]

        # ---- Time embedding (keep time_ ∈ [0,1] for tau) ----------------
        time_ = time
        if self.config.denoising_method == 'fm':
            time = time * 1000.0
        t_emb       = self.time_mlp(time)               # [B, D]
        t_emb_batch = repeat(t_emb, 'b d -> b k a d',
                             b=B, k=K, a=A)             # [B, K, A, D]

        # ---- Positional encodings ----------------------------------------
        k_pe = self.motion_query_embedding(
            torch.arange(self.model_cfg.NUM_PROPOSED_QUERY, device=device)
        )
        k_pe_batch = repeat(k_pe, 'k d -> b k a d', b=B, a=A)

        # Use actual A from input (supports variable-A SDD).
        a_pe = self.agent_order_embedding(torch.arange(A, device=device))
        a_pe_batch = repeat(a_pe, 'a d -> b k a d', b=B, k=K)

        # ---- K-level self-attention --------------------------------------
        y_emb_k = rearrange(
            self.apply_PE(y_emb, k_pe_batch, a_pe_batch),
            'b k a d -> (b a) k d',
        )
        y_emb_k = self.noisy_y_attn_k(y_emb_k)
        y_emb   = rearrange(y_emb_k, '(b a) k d -> b k a d', b=B, a=A)

        # ---- Agent-level self-attention (with padding mask for SDD) -----
        y_emb_a = rearrange(y_emb, 'b k a d -> (b k) a d')
        agent_mask_bka = x_data.get('agent_mask', None) if isinstance(x_data, dict) else None
        if agent_mask_bka is not None:
            kp_mask_a = ~agent_mask_bka.unsqueeze(1).expand(-1, K, -1).reshape(B * K, A)
            y_emb_a = self.noisy_y_attn_a(y_emb_a, src_key_padding_mask=kp_mask_a)
        else:
            y_emb_a = self.noisy_y_attn_a(y_emb_a)
        y_emb   = rearrange(y_emb_a, '(b k) a d -> b k a d', b=B, k=K)

        # ---- Embedding dropout (training only) ---------------------------
        if self.training and self.config.get('drop_method', None) == 'emb':
            m, k_drop = self.config.drop_logi_m, self.config.drop_logi_k
            p_m = 1 / (1 + torch.exp(-k_drop * (time_ - m)))
            p_m = p_m[:, None, None, None]
            y_emb = y_emb.masked_fill(torch.rand_like(p_m) < p_m, 0.)

        # ==================================================================
        # >>>  Future Euclidean Interaction Graph  <<<
        # ==================================================================
        if not skip_graph:
            # Edge source: y_0_hat from pass 1 (cleaner), or y_t (fallback).
            y_graph_src = (y_0_for_graph.view(B, K, A, self.T_future, 2)
                           if y_0_for_graph is not None
                           else y.view(B, K, A, self.T_future, 2))

            y_graph_unnorm = self._unnormalize_y(y_graph_src)
            init_pos = x_data['past_traj_original_scale'][:, :, -1, :2]  # [B, A, 2]
            y_abs    = y_graph_unnorm + init_pos.unsqueeze(1).unsqueeze(3)

            tau = time_  # [B] ∈ [0, 1]

            agent_mask = x_data.get('agent_mask', None) if isinstance(x_data, dict) else None
            y_emb_graph = self.future_graph(
                y_emb, y_abs, t_emb, tau,
                sigma_agent=sigma_for_graph,   # None in pass 1; [B,K,A] in pass 2
                agent_mask=agent_mask,          # [B,A] for padded SDD batches
            )
            y_emb = y_emb_graph
        # ==================================================================

        # ---- Context fusion + motion decoder ----------------------------
        emb_fusion = self.init_emb_fusion_mlp(
            torch.cat((encoder_out_batch, y_emb, t_emb_batch), dim=-1)
        )
        query_token  = self.post_pe_cat_mlp(
            self.apply_PE(emb_fusion, k_pe_batch, a_pe_batch)
        )
        readout_token = self.motion_decoder(query_token, t_emb)  # [B, K, A, D]

        # ---- Readout heads ----------------------------------------------
        denoiser_x   = self.reg_head(readout_token)               # [B, K, A, T*2]
        denoiser_cls = self.cls_head(readout_token).squeeze(-1)   # [B, K, A]
        logvar       = self.logvar_head(readout_token)            # [B, K, A, T*2]

        return denoiser_x, denoiser_cls, logvar

    # ------------------------------------------------------------------
    # Two-pass forward (overrides MotionTransformer.forward)
    # ------------------------------------------------------------------

    def forward(self, y, time, x_data, y_0_prev=None):
        """Two-pass forward — identical architecture both passes.

        Pass 1 (no_grad):
          - Training: graph uses GT future trajectory for edge features.
          - Inference: graph uses y_0_prev from previous sampling step,
            or skips graph if y_0_prev is None (first step).
          - Produces y_0_hat (clean prediction) and logvar_ (uncertainty).

        Pass 2 (grad):
          - Graph uses y_0_hat for cleaner edge features.
          - Graph uses sigma derived from logvar_ to weight agent messages:
              w_ij = sigmoid(γ * (σ_i − σ_j) + 0.5)
            promoting information flow from certain → uncertain agents.

        Returns: (denoiser_x, denoiser_cls, logvar)
          logvar is used by flow_matching.p_losses for the NLL uncertainty loss.
        """
        # Pass 1 — no gradient, get y_0_hat and preliminary logvar
        if self.training:
            # Use GT future as graph edge source during training.
            # A comes from input (variable for SDD), not from self.A.
            K = self.model_cfg.NUM_PROPOSED_QUERY
            gt_fut = x_data['fut_traj']                         # [B, A, T, 2]
            B_gt, A_gt = gt_fut.shape[0], gt_fut.shape[1]
            y_0_gt = gt_fut.unsqueeze(1).expand(-1, K, -1, -1, -1)  # [B, K, A, T, 2]
            y_0_for_pass1 = y_0_gt.reshape(B_gt, K, A_gt, -1)       # [B, K, A, T*2]
        else:
            y_0_for_pass1 = y_0_prev  # None at first step → skip graph

        with torch.no_grad():
            y_0_hat, _, logvar_ = self._forward_impl(
                y, time, x_data,
                y_0_for_graph=y_0_for_pass1,
                sigma_for_graph=None,
                skip_graph=(y_0_for_pass1 is None),
            )

        # Derive per-agent per-timestep uncertainty from logvar_:
        # logvar_ [B, K, A, T*2] → exp → mean over xy → sqrt → [B, K, A, T]
        B = y_0_hat.shape[0]
        K = self.model_cfg.NUM_PROPOSED_QUERY
        A = y_0_hat.shape[2]    # variable A from input (not self.A)
        T = self.T_future
        sigma_ = (logvar_.detach()
                         .view(B, K, A, T, 2)
                         .clamp(-10, 10)
                         .exp()
                         .mean(dim=-1)   # mean over xy only → [B, K, A, T]
                         .sqrt())        # [B, K, A, T]

        # Pass 2 — full gradient, graph uses y_0_hat edges + sigma weighting
        use_sigma = self.config.get('use_sigma_gating', True)
        return self._forward_impl(
            y, time, x_data,
            y_0_for_graph=y_0_hat.detach(),
            sigma_for_graph=sigma_ if use_sigma else None,
        )