standalone/c2f_nba_standalone.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
c2f_nba_standalone.py
================================================================================
FAITHFUL standalone re-implementation of the **Coarse-to-Fine** trajectory
predictor of ref [22] ("Towards Capturing the Temporal Dynamics for Trajectory
Prediction: A Coarse-to-Fine Approach") as a STANDALONE predictor for the NBA
basketball dataset.

It is a STANDALONE predictor: it consumes ONLY agent HISTORY (past trajectories)
and predicts K multi-modal futures via a native COARSE stage followed by an
autoregressive temporal FINE-refinement stage. It does NOT consume any external
denoiser / diffusion prediction (unlike the plug-in `CoarseToFineRefine` module,
which was handed the host's intermediate future estimate -> that borrowed SRA's
future-interaction signal and was therefore unfair; this standalone version
generates its own coarse trajectory, so it is faithful to the native method).

--------------------------------------------------------------------------------
WHAT THE NATIVE METHOD IS (and what we replicate)
--------------------------------------------------------------------------------
Coarse-to-fine = two-stage decoding that *captures temporal dynamics* by first
predicting a rough full trajectory and then refining it step-by-step:

  STAGE 0  (COARSE):  a multimodal decoder emits K rough full-horizon
    trajectories per agent from the social-encoded context (mode + agent query
    cross-attended to all agents' context, then an MLP trajectory head). This is
    the model's OWN coarse prediction -- it is NOT received from a host.

  STAGE 1..S (FINE):  the coarse trajectory is walked TEMPORALLY by a
    unidirectional (autoregressive) GRU -- the mechanism the paper uses to
    capture temporal dynamics -- conditioned on the mode's context, emitting a
    per-timestep residual correction  delta_t ; y_fine = y_coarse + delta.
    Repeated S times (progressive coarse -> fine -> finer). delta head is
    zero-initialised so training first fits the coarse stage, then the refiner
    engages (stable).

  Mode scores: a per-mode classification head; winner-takes-all training.

ENCODER (shared, IDENTICAL to the GameFormer standalone so the E4 comparison
isolates the DECODER mechanism, coarse-to-fine vs level-k, not the backbone):
  * AgentHistoryEncoder = 2-layer LSTM(6->256) over agent history + learned
    player/ball type embedding.
  * FusionEncoder = nn.TransformerEncoder (d=256, heads=8, ff=1024, gelu),
    `encoder_layers` deep -> agent<->agent social self-attention (no map: NBA).

LOSS (native coarse-to-fine supervision):
  * coarse WTA-L2 (variety loss over K modes)  +  fine WTA-L2  +  mode
    cross-entropy (label smoothing 0.2) on the fine-stage winning mode.
  * endpoint-emphasised distance (mean_t + checkpoint-step sum) for mode
    selection, marginal PER AGENT (NBA Table-1 metric is marginal min-ADE_20).

--------------------------------------------------------------------------------
NBA I/O + METRIC (matched to MoFlow / Table 1 -- identical to GameFormer standalone)
--------------------------------------------------------------------------------
  * Data: MoFlow's data/dataloader_nba.py::NBADatasetMinMax, same .npy, split,
    scaling (traj_scale=94/28, traj_mean=[14,7.5]). Past=10, Future=20 (4.0 s).
  * Predict in centered-abs frame (pos/scale - mean); cur_xy = last past step;
    gt_center = fut_traj_original_scale (displacement) + cur_xy.
  * Metric (identical to eval_perscene_moflow.py):
      d = ||pred_disp - gt_disp|| ; ADE4 = d[:,:20].mean_t.min_K ; FDE4 =
      d[:,19].min_K ; mean over 11 agents & scenes  (marginal min-of-K=20 @ 4s).

Prints:  [C2F-NBA] epoch N ADE4=.. FDE4=..
"""

import os
import sys
import argparse
import time

# ------------------------------------------------------------------ GPU FIRST
def _early_gpu():
    for i, a in enumerate(sys.argv):
        if a == '--gpu' and i + 1 < len(sys.argv):
            return sys.argv[i + 1]
        if a.startswith('--gpu='):
            return a.split('=', 1)[1]
    return None


_g = _early_gpu()
if _g is not None:
    os.environ['CUDA_VISIBLE_DEVICES'] = str(_g)
os.environ.setdefault('MPLBACKEND', 'Agg')

# ------------------------------------------------ reuse MoFlow's NBA pipeline
MOFLOW_ROOT = '/mnt/jaewoo4tb/srtp/MoFlow'
if MOFLOW_ROOT not in sys.path:
    sys.path.insert(0, MOFLOW_ROOT)

import types as _types
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

try:
    os.chdir(MOFLOW_ROOT)
except Exception:
    pass
from data.dataloader_nba import NBADatasetMinMax, seq_collate_nba


# ============================================================================
#  PRIMITIVES  (shared with the GameFormer standalone, verbatim)
# ============================================================================
D_MODEL = 256
N_HEADS = 8
DROPOUT = 0.1


class CrossTransformer(nn.Module):
    def __init__(self, dim=D_MODEL, heads=N_HEADS, dropout=DROPOUT):
        super().__init__()
        self.cross_attention = nn.MultiheadAttention(dim, heads, dropout, batch_first=True)
        self.norm_1 = nn.LayerNorm(dim)
        self.norm_2 = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim * 4), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(dim * 4, dim), nn.Dropout(dropout))

    def forward(self, query, key, value, mask=None):
        attn, _ = self.cross_attention(query, key, value, key_padding_mask=mask)
        attn = self.norm_1(attn)
        return self.norm_2(self.ffn(attn) + attn)


class AgentHistoryEncoder(nn.Module):
    """2-layer LSTM(6->256) over agent history + player/ball type embedding."""

    def __init__(self, in_dim=6, dim=D_MODEL, n_types=2):
        super().__init__()
        self.motion = nn.LSTM(in_dim, dim, 2, batch_first=True)
        self.type_emb = nn.Embedding(n_types, dim)

    def forward(self, hist, types):
        B, A, T, C = hist.shape
        traj, _ = self.motion(hist.reshape(B * A, T, C))
        out = traj[:, -1].reshape(B, A, -1)
        out = out + self.type_emb(types)[None]
        return out


class FusionEncoder(nn.Module):
    """Agent<->agent social self-attention (no map for NBA)."""

    def __init__(self, dim=D_MODEL, heads=N_HEADS, layers=6, dropout=DROPOUT):
        super().__init__()
        layer = nn.TransformerEncoderLayer(
            d_model=dim, nhead=heads, dim_feedforward=dim * 4,
            activation=F.gelu, dropout=dropout, batch_first=True)
        self.encoder = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False)

    def forward(self, tokens, mask=None):
        return self.encoder(tokens, src_key_padding_mask=mask)


# ============================================================================
#  COARSE-TO-FINE DECODER
# ============================================================================
class CoarseDecoder(nn.Module):
    """Stage-0: K multimodal rough full-horizon trajectories per agent.
    Mode + agent query added to the agent's context token, cross-attended to the
    full agent context (social), then an MLP trajectory head + a mode-score head."""

    def __init__(self, modalities, n_agents, future_len, dim=D_MODEL):
        super().__init__()
        self.M = modalities
        self.multi_modal_query_embedding = nn.Embedding(modalities, dim)
        self.agent_query_embedding = nn.Embedding(n_agents, dim)
        self.query_encoder = CrossTransformer(dim)
        self.traj_head = nn.Sequential(
            nn.Linear(dim, 512), nn.ELU(), nn.Dropout(0.1),
            nn.Linear(512, future_len * 2))
        self.score_head = nn.Sequential(
            nn.Linear(dim, 64), nn.ELU(), nn.Dropout(0.1), nn.Linear(64, 1))
        self.future_len = future_len
        self.register_buffer('modal', torch.arange(modalities).long())
        self.register_buffer('agent', torch.arange(n_agents).long())

    def forward(self, encoding, cur_xy, mask=None):
        B, A, D = encoding.shape
        M, T = self.M, self.future_len
        mm = self.multi_modal_query_embedding(self.modal)          # [M, D]
        ag = self.agent_query_embedding(self.agent)                # [A, D]
        query = encoding[:, :, None, :] + mm[None, None] + ag[None, :, None]  # [B,A,M,D]

        q = query.reshape(B * A, M, D)
        kv = encoding[:, None, :, :].expand(B, A, A, D).reshape(B * A, A, D)
        km = None
        if mask is not None:
            km = mask[:, None, :].expand(B, A, A).reshape(B * A, A)
        content = self.query_encoder(q, kv, kv, km)                # [B*A, M, D]

        coarse = self.traj_head(content).view(B, A, M, T, 2)
        coarse = coarse + cur_xy[:, :, None, None, :]              # centered-abs
        score = self.score_head(content).view(B, A, M)
        return content.view(B, A, M, D), coarse, score


class FineRefiner(nn.Module):
    """Stage-1..S: autoregressive temporal refinement of the coarse trajectory.
    A unidirectional GRU walks the (centered) coarse trajectory, conditioned on
    the mode context, emitting a per-timestep residual delta_t. Repeated S times."""

    def __init__(self, dim=D_MODEL, hidden=256, n_stages=2, dropout=DROPOUT):
        super().__init__()
        self.n_stages = n_stages
        self.pos_emb = nn.Linear(2, hidden)
        self.ctx_proj = nn.Linear(dim, hidden)
        self.gru = nn.GRU(hidden, hidden, 2, batch_first=True, dropout=dropout)
        self.delta = nn.Linear(hidden, 2)
        nn.init.zeros_(self.delta.weight)      # start as identity: fine == coarse at init
        nn.init.zeros_(self.delta.bias)

    def forward(self, coarse, content, cur_xy):
        # coarse:[B,A,M,T,2] content:[B,A,M,D] cur_xy:[B,A,2]
        B, A, M, T, _ = coarse.shape
        ctx = self.ctx_proj(content).reshape(B * A * M, 1, -1)     # [N,1,H]
        y = coarse
        for _ in range(self.n_stages):
            yc = (y - cur_xy[:, :, None, None, :]).reshape(B * A * M, T, 2)   # center
            seq = self.pos_emb(yc) + ctx                          # broadcast ctx over T
            h, _ = self.gru(seq)                                  # [N,T,H] autoregressive
            d = self.delta(h).view(B, A, M, T, 2)
            y = y + d
        return y


class CoarseToFineNBA(nn.Module):
    """Standalone coarse-to-fine predictor (NBA, map dropped)."""

    def __init__(self, n_agents=11, past_dim=6, future_len=20, modalities=20,
                 n_stages=2, dim=D_MODEL, heads=N_HEADS, enc_layers=6, hidden=256,
                 n_types=2, ball_idx=10):
        super().__init__()
        self.history_encoder = AgentHistoryEncoder(past_dim, dim, n_types)
        self.fusion_encoder = FusionEncoder(dim, heads, enc_layers)
        self.coarse_decoder = CoarseDecoder(modalities, n_agents, future_len, dim)
        self.fine_refiner = FineRefiner(dim, hidden, n_stages)
        types = torch.zeros(n_agents, dtype=torch.long)
        if 0 <= ball_idx < n_agents and n_types > 1:
            types[ball_idx] = 1
        self.register_buffer('agent_types', types)

    def forward(self, feats, cur_xy):
        enc = self.history_encoder(feats, self.agent_types)       # [B,A,D]
        enc = self.fusion_encoder(enc, None)                      # [B,A,D] social
        content, coarse, score = self.coarse_decoder(enc, cur_xy, None)
        fine = self.fine_refiner(coarse, content, cur_xy)         # [B,A,M,T,2]
        return {'coarse': coarse, 'fine': fine, 'scores': score}


# ============================================================================
#  LOSS  (coarse WTA + fine WTA + mode CE; marginal per agent)
# ============================================================================
def wta_l2(pred, gt, metric_idx):
    """Winner-takes-all L2 (variety loss). pred:[B,A,M,T,2] gt:[B,A,T,2].
    Mode selection uses endpoint-emphasised distance; returns best mode's mean
    displacement (reconstruction loss) and the winning mode index [B,A]."""
    B, A, M, T, _ = pred.shape
    d = torch.norm(pred - gt[:, :, None], dim=-1)                 # [B,A,M,T]
    sel = d.mean(-1) + d[..., metric_idx].sum(-1)                 # [B,A,M]
    best = sel.argmin(-1)                                         # [B,A]
    gi = best[..., None, None, None].expand(B, A, 1, T, 2)
    best_d = torch.norm(torch.gather(pred, 2, gi).squeeze(2) - gt, dim=-1)  # [B,A,T]
    reg = best_d.mean(-1) + best_d[..., metric_idx].sum(-1)       # [B,A] endpoint emphasis
    return reg.mean(), best


def c2f_loss(out, gt_center, metric_idx):
    coarse_reg, _ = wta_l2(out['coarse'], gt_center, metric_idx)
    fine_reg, best = wta_l2(out['fine'], gt_center, metric_idx)
    B, A, M = out['scores'].shape
    cls = F.cross_entropy(out['scores'].reshape(B * A, M), best.reshape(B * A),
                          label_smoothing=0.2)
    return coarse_reg + fine_reg + 2.0 * cls


# ============================================================================
#  DATA  (identical to the GameFormer standalone)
# ============================================================================
def build_loaders(args):
    cfg = _types.SimpleNamespace(
        traj_mean=[14, 7.5], data_norm='min_max',
        past_frames=args.past_len, future_frames=args.future_len, agents=args.agents)
    train_set = NBADatasetMinMax(
        obs_len=args.past_len, pred_len=args.future_len, training=True,
        num_scenes=args.n_train, cfg=cfg, data_dir=args.data_dir,
        rotate=False, data_norm='min_max')
    test_set = NBADatasetMinMax(
        obs_len=args.past_len, pred_len=args.future_len, training=False,
        test_scenes=args.n_test, cfg=cfg, data_dir=args.data_dir,
        rotate=False, data_norm='min_max')
    train_loader = DataLoader(
        train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.workers,
        collate_fn=seq_collate_nba, pin_memory=True, drop_last=True)
    test_loader = DataLoader(
        test_set, batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
        collate_fn=seq_collate_nba, pin_memory=True)
    return train_loader, test_loader


def unpack(data, device):
    feats = data['past_traj'].to(device)                       # [B,A,T,6] normalized
    past_orig = data['past_traj_original_scale'].to(device)    # [B,A,T,6] metric
    gt_disp = data['fut_traj_original_scale'].to(device)       # [B,A,Tf,2] displacement
    cur_xy = past_orig[:, :, -1, 0:2]
    gt_center = gt_disp + cur_xy[:, :, None, :]
    return feats, cur_xy, gt_disp, gt_center


# ============================================================================
#  EVAL  (marginal min-of-K=20 at 4.0 s; identical to eval_perscene_moflow.py)
# ============================================================================
@torch.no_grad()
def evaluate(model, loader, device, end):
    model.eval()
    ade_sum = fde_sum = 0.0
    n = 0
    for data in loader:
        feats, cur_xy, gt_disp, _ = unpack(data, device)
        out = model(feats, cur_xy)
        pred = out['fine'][..., :2]                               # [B,A,M,T,2] fine stage
        pred_disp = pred - cur_xy[:, :, None, None, :]
        d = torch.norm(pred_disp - gt_disp[:, :, None], dim=-1)   # [B,A,M,T]
        ade = d[..., :end].mean(-1).min(dim=-1).values            # [B,A]
        fde = d[..., end - 1].min(dim=-1).values                  # [B,A]
        ade_sum += ade.sum().item()
        fde_sum += fde.sum().item()
        n += ade.numel()
    return ade_sum / max(n, 1), fde_sum / max(n, 1)


# ============================================================================
#  TRAIN
# ============================================================================
def parse_args():
    p = argparse.ArgumentParser('Coarse-to-Fine standalone predictor for NBA')
    p.add_argument('--data_dir', default='/mnt/jaewoo4tb/srtp/MoFlow/data/nba', type=str)
    p.add_argument('--gpu', default='0', type=str)
    p.add_argument('--exp', default='c2f_nba', type=str)
    p.add_argument('--epochs', default=50, type=int)
    p.add_argument('--batch_size', default=128, type=int)
    p.add_argument('--test_batch', default=500, type=int)
    p.add_argument('--workers', default=4, type=int)
    p.add_argument('--seed', default=3407, type=int)

    p.add_argument('--modalities', default=20, type=int, help='K modes (NBA min-ADE_20)')
    p.add_argument('--stages', default=2, type=int, help='coarse->fine refinement stages')
    p.add_argument('--encoder_layers', default=6, type=int)
    p.add_argument('--hidden', default=256, type=int, help='fine-refiner GRU hidden')
    p.add_argument('--agents', default=11, type=int)
    p.add_argument('--past_len', default=10, type=int)
    p.add_argument('--future_len', default=20, type=int)
    p.add_argument('--ball_idx', default=10, type=int, help='-1 to disable type emb')

    p.add_argument('--lr', default=1e-4, type=float)
    p.add_argument('--weight_decay', default=1e-4, type=float)
    p.add_argument('--grad_clip', default=5.0, type=float)
    p.add_argument('--n_train', default=32500, type=int)
    p.add_argument('--n_test', default=12500, type=int)
    p.add_argument('--eval_every', default=2, type=int)
    p.add_argument('--save_dir', default=os.path.join(
        os.path.dirname(os.path.abspath(__file__)), 'c2f_nba_ckpt'), type=str)
    return p.parse_args()


def main():
    args = parse_args()
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    train_loader, test_loader = build_loaders(args)

    n_types = 2 if (0 <= args.ball_idx < args.agents) else 1
    model = CoarseToFineNBA(
        n_agents=args.agents, past_dim=6, future_len=args.future_len,
        modalities=args.modalities, n_stages=args.stages, dim=D_MODEL, heads=N_HEADS,
        enc_layers=args.encoder_layers, hidden=args.hidden, n_types=n_types,
        ball_idx=args.ball_idx).to(device)
    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f'[C2F-NBA] model params: {n_params/1e6:.2f} M | stages={args.stages} '
          f'K={args.modalities} enc_layers={args.encoder_layers} device={device}', flush=True)

    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    milestones = sorted(set(int(args.epochs * f) for f in (0.5, 0.65, 0.8, 0.9)))
    milestones = [m for m in milestones if 0 < m < args.epochs]
    sched = torch.optim.lr_scheduler.MultiStepLR(opt, milestones=milestones, gamma=0.5)

    end = args.future_len
    metric_idx = sorted(set([max(0, args.future_len // 2 - 1), args.future_len - 1]))

    os.makedirs(args.save_dir, exist_ok=True)
    best_ade = float('inf')

    for epoch in range(args.epochs):
        model.train()
        t0 = time.time()
        running = 0.0
        nb = 0
        for data in train_loader:
            feats, cur_xy, _, gt_center = unpack(data, device)
            out = model(feats, cur_xy)
            loss = c2f_loss(out, gt_center, metric_idx)
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            opt.step()
            running += float(loss.item())
            nb += 1
        sched.step()
        print(f'[C2F-NBA] epoch {epoch+1}/{args.epochs} loss={running/max(nb,1):.4f} '
              f'lr={opt.param_groups[0]["lr"]:.2e} ({time.time()-t0:.1f}s)', flush=True)

        if (epoch + 1) % args.eval_every == 0 or epoch == args.epochs - 1:
            ade, fde = evaluate(model, test_loader, device, end)
            print(f'[C2F-NBA] epoch {epoch+1} ADE4={ade:.4f} FDE4={fde:.4f}', flush=True)
            if ade < best_ade:
                best_ade = ade
                torch.save({'model': model.state_dict(), 'epoch': epoch + 1,
                            'ade4': ade, 'fde4': fde, 'args': vars(args)},
                           os.path.join(args.save_dir, f'{args.exp}_best.pt'))
                print(f'[C2F-NBA] saved best (ADE4={ade:.4f}) -> '
                      f'{os.path.join(args.save_dir, args.exp + "_best.pt")}', flush=True)

    print(f'[C2F-NBA] done. best ADE4={best_ade:.4f}', flush=True)


if __name__ == '__main__':
    main()