MoFlow/models/backbone_graph_v2.py

"""
MotionTransformerGraphV2 — two-pass backbone with mode-specific graph edges.

Identical to MotionTransformerGraph (backbone_graph.py) except it uses
FutureInteractionGraphV2 instead of FutureInteractionGraph, so each of the
K denoising modes receives interaction edge features derived from its own
predicted future positions rather than a cross-mode average.
"""

from models.backbone_graph import MotionTransformerGraph
from models.graph_interaction_nba_v2 import FutureInteractionGraphV2


class MotionTransformerGraphV2(MotionTransformerGraph):
    """MotionTransformerGraph with mode-specific future interaction edges.

    Constructor arguments are identical to MotionTransformerGraph.
    Only __init__ is overridden to swap FutureInteractionGraph → V2.
    All forward logic (_forward_impl, forward) is inherited unchanged.
    """

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

        D        = self.dim
        time_dim = D

        # Replace the V1 graph module with V2 (same hyperparams)
        self.future_graph = FutureInteractionGraphV2(
            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,
        )

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