MID/main_soccer_mid.py

"""
main_soccer_mid.py — MID baseline on the soccer dataset.

Adapted from main_nba_mid.py with:
  - NUM_AGENTS = 23 (soccer)
  - Per-scene normalization (abs channel centered by scene centroid at last obs)
  - No /= (94/28) court rescale (soccer data is already in field-normalized units)
  - val.npy used as both val and test
  - Adjusted batch sizes for 23 agents

Usage:
    python main_soccer_mid.py --gpu 0
"""

import os
import sys
import time
import logging
import argparse
import numpy as np
import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import Dataset, DataLoader
try:
    from tensorboardX import SummaryWriter
except Exception:
    from torch.utils.tensorboard import SummaryWriter
from tqdm.auto import tqdm

from models.diffusion import DiffusionTraj, VarianceSchedule, TransformerConcatLinear


OBS_LEN    = 10
PRED_LEN   = 20
NUM_AGENTS = 23
TRAJ_SCALE = 5.0
TRAJ_MEAN  = torch.FloatTensor([-0.726, -0.278])
K_EVAL     = 20
PER_SCENE_NORM = True


class SoccerDatasetMID(Dataset):
    def __init__(self, data_dir: str, split: str = 'train'):
        super().__init__()
        path = os.path.join(data_dir, f'{split}.npy')
        trajs = np.load(path).astype(np.float32)   # (N, 30, 23, 2)
        trajs = torch.from_numpy(trajs).permute(0, 2, 1, 3)  # (N, 23, 30, 2)
        self.pre = trajs[:, :, :OBS_LEN, :]
        self.fut = trajs[:, :, OBS_LEN:,  :]
        print(f'[SoccerDatasetMID] {split}: {path} → {trajs.shape}')

    def __len__(self):
        return len(self.pre)

    def __getitem__(self, idx):
        return self.pre[idx], self.fut[idx]


def collate_fn(batch):
    pre = torch.stack([b[0] for b in batch])
    fut = torch.stack([b[1] for b in batch])
    return pre, fut


def preprocess_batch(pre_motion, fut_motion, device):
    B, A = pre_motion.shape[:2]
    pre      = pre_motion.reshape(B * A, OBS_LEN, 2)
    fut      = fut_motion.reshape(B * A, PRED_LEN, 2)
    last_obs = pre[:, -1:, :]

    if PER_SCENE_NORM:
        scene_center = pre_motion[:, :, -1, :].mean(dim=1, keepdim=True)  # [B, 1, 2]
        scene_center = scene_center.unsqueeze(2)                           # [B, 1, 1, 2]
        abs_xy = ((pre_motion - scene_center) / TRAJ_SCALE).reshape(B * A, OBS_LEN, 2)
    else:
        traj_mean = TRAJ_MEAN.to(device)
        abs_xy = (pre - traj_mean) / TRAJ_SCALE

    rel_xy = (pre - last_obs) / TRAJ_SCALE
    vel_xy = torch.cat([rel_xy[:, 1:] - rel_xy[:, :-1],
                        torch.zeros_like(rel_xy[:, :1])], dim=1)

    past_6ch = torch.cat([abs_xy, rel_xy, vel_xy], dim=-1)
    fut_rel  = (fut - last_obs) / TRAJ_SCALE

    mask = torch.full((B * A, B * A), float('-inf'), device=device)
    for i in range(B):
        s, e = i * A, (i + 1) * A
        mask[s:e, s:e] = 0.0

    return past_6ch, fut_rel, mask, last_obs


class _STEncoder(nn.Module):
    def __init__(self, in_channels=6, hidden=256):
        super().__init__()
        self.conv = nn.Conv1d(in_channels, 32, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        self.gru  = nn.GRU(32, hidden, num_layers=1, batch_first=True)
        nn.init.kaiming_normal_(self.conv.weight)
        nn.init.kaiming_normal_(self.gru.weight_ih_l0)
        nn.init.kaiming_normal_(self.gru.weight_hh_l0)
        nn.init.zeros_(self.conv.bias)
        nn.init.zeros_(self.gru.bias_ih_l0)
        nn.init.zeros_(self.gru.bias_hh_l0)

    def forward(self, x):
        h = self.relu(self.conv(x.transpose(1, 2)))
        _, state = self.gru(h.transpose(1, 2))
        return state.squeeze(0)


class _SocialTransformer(nn.Module):
    def __init__(self, past_len=OBS_LEN, hidden=256):
        super().__init__()
        self.proj    = nn.Linear(past_len * 6, hidden, bias=False)
        layer        = nn.TransformerEncoderLayer(
                           d_model=hidden, nhead=2,
                           dim_feedforward=hidden, batch_first=False)
        self.encoder = nn.TransformerEncoder(layer, num_layers=2)

    def forward(self, x_flat, mask):
        h = self.proj(x_flat).unsqueeze(1)
        h = h + self.encoder(h, mask)
        return h.squeeze(1)


class SoccerEncoder(nn.Module):
    def __init__(self, encoder_dim=256, past_len=OBS_LEN):
        super().__init__()
        self.ego_encoder    = _STEncoder(in_channels=6, hidden=256)
        self.social_encoder = _SocialTransformer(past_len=past_len, hidden=256)
        self.fusion         = nn.Linear(512, encoder_dim)

    def forward(self, past_6ch, social_mask):
        ego    = self.ego_encoder(past_6ch)
        social = self.social_encoder(
                     past_6ch.reshape(past_6ch.size(0), -1), social_mask)
        return self.fusion(torch.cat([ego, social], dim=-1))


class Trainer:
    def __init__(self, args):
        self.args   = args
        self.device = torch.device(f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu')
        self._build_dirs()
        self._build_data()
        self._build_model()
        self._build_optimizer()

    def _build_dirs(self):
        self.exp_dir = os.path.join('experiments', self.args.exp_name)
        os.makedirs(self.exp_dir, exist_ok=True)
        self.tb_log  = SummaryWriter(log_dir=self.exp_dir)
        log_path = os.path.join(
            self.exp_dir,
            'soccer_{}.log'.format(time.strftime('%Y-%m-%d-%H-%M')))
        self.log = logging.getLogger(self.args.exp_name)
        self.log.setLevel(logging.INFO)
        self.log.addHandler(logging.FileHandler(log_path))
        self.log.addHandler(logging.StreamHandler(sys.stdout))
        self.log.info(f"Args: {self.args}")

    def _build_data(self):
        train_dset = SoccerDatasetMID(self.args.data_dir, split='train')
        test_dset  = SoccerDatasetMID(self.args.data_dir, split='val')
        self.train_loader = DataLoader(
            train_dset, batch_size=self.args.batch_size,
            shuffle=True, num_workers=4,
            collate_fn=collate_fn, pin_memory=True)
        self.test_loader = DataLoader(
            test_dset, batch_size=self.args.eval_batch_size,
            shuffle=False, num_workers=4,
            collate_fn=collate_fn, pin_memory=True)
        self.log.info(f"Train: {len(train_dset)}  Val/Test: {len(test_dset)}")

    def _build_model(self):
        self.encoder = SoccerEncoder(
            encoder_dim=self.args.encoder_dim, past_len=OBS_LEN,
        ).to(self.device)
        net = TransformerConcatLinear(
            point_dim=2, context_dim=self.args.encoder_dim,
            tf_layer=self.args.tf_layer, residual=False)
        self.diffusion = DiffusionTraj(
            net=net,
            var_sched=VarianceSchedule(num_steps=100, beta_T=5e-2, mode='linear'),
        ).to(self.device)
        n_enc  = sum(p.numel() for p in self.encoder.parameters())
        n_diff = sum(p.numel() for p in self.diffusion.parameters())
        self.log.info(f"Encoder: {n_enc:,}  Diffusion: {n_diff:,}")

    def _build_optimizer(self):
        params = list(self.encoder.parameters()) + list(self.diffusion.parameters())
        self.optimizer = optim.Adam(params, lr=self.args.lr)
        self.scheduler = optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=0.98)

    def train(self):
        best_ade = float('inf')
        for epoch in range(1, self.args.epochs + 1):
            self.encoder.train()
            self.diffusion.train()
            total_loss, count = 0.0, 0
            pbar = tqdm(self.train_loader, ncols=90)
            for pre, fut in pbar:
                pre, fut = pre.to(self.device), fut.to(self.device)
                past_6ch, fut_rel, mask, _ = preprocess_batch(pre, fut, self.device)
                context = self.encoder(past_6ch, mask)
                loss    = self.diffusion.get_loss(fut_rel, context)
                self.optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(
                    list(self.encoder.parameters()) +
                    list(self.diffusion.parameters()), 1.0)
                self.optimizer.step()
                total_loss += loss.item()
                count += 1
                pbar.set_description(f"Epoch {epoch} loss={total_loss/count:.4f}")

            self.scheduler.step()
            avg_loss = total_loss / count
            self.tb_log.add_scalar('loss/train', avg_loss, epoch)
            self.log.info(f"Epoch {epoch}  train_loss={avg_loss:.4f}")

            if epoch % self.args.eval_every == 0:
                ade, fde = self.evaluate()
                self.tb_log.add_scalar('metric/ADE', ade, epoch)
                self.tb_log.add_scalar('metric/FDE', fde, epoch)
                self.log.info(f"Epoch {epoch}  ADE={ade:.4f}  FDE={fde:.4f}")
                if ade < best_ade:
                    best_ade = ade
                    torch.save({
                        'encoder': self.encoder.state_dict(),
                        'diffusion': self.diffusion.state_dict(),
                        'epoch': epoch, 'ade': ade, 'fde': fde,
                    }, os.path.join(self.exp_dir, 'best.pt'))
                    self.log.info(f"  ** New best ADE={ade:.4f} FDE={fde:.4f}")

    @torch.no_grad()
    def evaluate(self):
        self.encoder.eval()
        self.diffusion.eval()
        ade_sum, fde_sum, n_agents = 0.0, 0.0, 0
        for pre, fut in tqdm(self.test_loader, ncols=90, desc='Eval'):
            pre, fut = pre.to(self.device), fut.to(self.device)
            B = pre.size(0)
            past_6ch, _, mask, last_obs = preprocess_batch(pre, fut, self.device)
            context = self.encoder(past_6ch, mask)
            pred_rel = self.diffusion.sample(
                num_points=PRED_LEN, context=context,
                sample=K_EVAL, bestof=True,
                sampling=self.args.sampling, step=self.args.sampling_step)
            pred_abs = pred_rel * TRAJ_SCALE + last_obs.unsqueeze(0)
            fut_abs  = fut.reshape(B * NUM_AGENTS, PRED_LEN, 2)
            dist     = (pred_abs - fut_abs.unsqueeze(0)).norm(dim=-1)
            ade_sum += dist.mean(dim=-1).min(dim=0).values.sum().item()
            fde_sum += dist[:, :, -1].min(dim=0).values.sum().item()
            n_agents += B * NUM_AGENTS
        return ade_sum / n_agents, fde_sum / n_agents


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument('--data_dir', type=str,
                   default='/mnt/jaewoo4tb/srtp/srtp/raw_data/soccer')
    p.add_argument('--exp_name',     type=str, default='mid_soccer_baseline')
    p.add_argument('--gpu',          type=int, default=0)
    p.add_argument('--epochs',       type=int, default=100)
    p.add_argument('--batch_size',   type=int, default=32)
    p.add_argument('--eval_batch_size', type=int, default=64)
    p.add_argument('--lr',           type=float, default=1e-3)
    p.add_argument('--eval_every',   type=int, default=5)
    p.add_argument('--encoder_dim',  type=int, default=256)
    p.add_argument('--tf_layer',     type=int, default=3)
    p.add_argument('--sampling',     type=str, default='ddim', choices=['ddpm', 'ddim'])
    p.add_argument('--sampling_step', type=int, default=10)
    return p.parse_args()


if __name__ == '__main__':
    args = parse_args()
    trainer = Trainer(args)
    trainer.train()