#!/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()