LED/trainer/train_sport_led.py

"""
Stage-2 LED trainer for sport datasets (soccer / football).

Loads a sport-specific pretrained core denoiser (produced by
train_sport_pretrain.py), then trains the leapfrog initializer in the
standard LED way. If --use_graph is set, also instantiates
FutureInteractionGraph (top_n=5, residual_on='y0' — the winning NBA
variant) and adds its output as a residual correction inside each
leapfrog reverse step.

Agent count, data path, traj_mean and traj_scale all come from the
sport config, so the same trainer runs on soccer and football with
different ymls.
"""

import os
import time
import torch
import random
import numpy as np
import torch.nn as nn

from utils.config import Config
from utils.utils import print_log

from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from data.dataloader_sport import SportDataset, sport_seq_collate
from models.model_led_initializer import LEDInitializer as InitializationModel
from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel
from models.future_interaction_graph import FutureInteractionGraph
from models.future_interaction_graph_v6 import FutureInteractionGraphV6Wrapper


NUM_Tau = 5


class Trainer:
    def __init__(self, config):
        if torch.cuda.is_available():
            torch.cuda.set_device(config.gpu)
        self.device = torch.device('cuda') if config.cuda else torch.device('cpu')
        self.cfg = Config(config.cfg, config.info)
        self.use_graph   = bool(getattr(config, 'use_graph', False))
        self.residual_on = getattr(config, 'residual_on', 'y0')

        # ------------------------- data -------------------------
        self.num_agents = self.cfg.num_agents
        train_dset = SportDataset(
            data_dir   = self.cfg.data_dir,
            num_agents = self.num_agents,
            obs_len    = self.cfg.past_frames,
            pred_len   = self.cfg.future_frames,
            split      = 'train',
        )
        val_dset = SportDataset(
            data_dir   = self.cfg.data_dir,
            num_agents = self.num_agents,
            obs_len    = self.cfg.past_frames,
            pred_len   = self.cfg.future_frames,
            split      = 'val',
        )
        self.train_loader = DataLoader(
            train_dset, batch_size=self.cfg.train_batch_size, shuffle=True,
            num_workers=4, collate_fn=sport_seq_collate, pin_memory=True)
        self.test_loader = DataLoader(
            val_dset, batch_size=self.cfg.test_batch_size, shuffle=False,
            num_workers=4, collate_fn=sport_seq_collate, pin_memory=True)

        self.traj_mean  = torch.FloatTensor(self.cfg.traj_mean).cuda().unsqueeze(0).unsqueeze(0).unsqueeze(0)
        self.traj_scale = float(self.cfg.traj_scale)
        self.per_scene_norm = bool(self.cfg.get('per_scene_norm', False))

        # ------------------------- diffusion parameters -------------------------
        self.n_steps = self.cfg.diffusion.steps
        self.betas = self.make_beta_schedule(
            schedule=self.cfg.diffusion.beta_schedule, n_timesteps=self.n_steps,
            start=self.cfg.diffusion.beta_start, end=self.cfg.diffusion.beta_end).cuda()
        self.alphas = 1 - self.betas
        self.alphas_prod = torch.cumprod(self.alphas, 0)
        self.alphas_bar_sqrt = torch.sqrt(self.alphas_prod)
        self.one_minus_alphas_bar_sqrt = torch.sqrt(1 - self.alphas_prod)

        # ------------------------- models -------------------------
        self.model = CoreDenoisingModel().cuda()
        ckpt_path = self.cfg.pretrained_core_denoising_model
        if not os.path.isfile(ckpt_path):
            raise FileNotFoundError(
                f'Missing sport-specific pretrained denoiser: {ckpt_path}. '
                'Run train_sport_pretrain.py first.')
        core_cp = torch.load(ckpt_path, map_location='cpu')
        self.model.load_state_dict(core_cp['model_dict'])

        self.model_initializer = InitializationModel(
            t_h=self.cfg.past_frames, d_h=6,
            t_f=self.cfg.future_frames, d_f=2,
            k_pred=20).cuda()

        params = list(self.model_initializer.parameters())
        self.interaction_graph = None
        self.use_v6_graph = bool(getattr(config, 'use_v6_graph', False))
        if self.use_graph:
            if self.use_v6_graph:
                self.interaction_graph = FutureInteractionGraphV6Wrapper(
                    num_agents        = self.num_agents,
                    future_steps      = self.cfg.future_frames,
                    past_steps        = self.cfg.past_frames,
                    past_channels     = 6,
                    node_dim          = 128,
                    top_n             = min(int(__import__('os').environ.get('LED_TOP_N', 5)), self.num_agents - 1),
                    num_denoise_steps = NUM_Tau,
                ).cuda()
            else:
                self.interaction_graph = FutureInteractionGraph(
                    num_agents        = self.num_agents,
                    future_steps      = self.cfg.future_frames,
                    past_steps        = self.cfg.past_frames,
                    past_channels     = 6,
                    node_dim          = 128,
                    top_n             = min(int(__import__('os').environ.get('LED_TOP_N', 5)), self.num_agents - 1),
                    num_denoise_steps = NUM_Tau,
                ).cuda()
            params += list(self.interaction_graph.parameters())

        self.opt = torch.optim.AdamW(params, lr=config.learning_rate)
        self.scheduler_model = torch.optim.lr_scheduler.StepLR(
            self.opt, step_size=self.cfg.decay_step, gamma=self.cfg.decay_gamma)

        # ------------------------- logs -------------------------
        self.log = open(os.path.join(self.cfg.log_dir, 'log.txt'), 'a+')
        self.tb  = SummaryWriter(log_dir=os.path.join(self.cfg.log_dir, 'tb'))
        self.global_step = 0
        self.print_model_param(self.model, name='Core Denoising Model')
        self.print_model_param(self.model_initializer, name='Initialization Model')
        if self.use_graph:
            self.print_model_param(self.interaction_graph, name='Future Interaction Graph')

        # temporal reweight: [T, T-1, ..., 1] / (T/2)
        T = self.cfg.future_frames
        self.temporal_reweight = torch.FloatTensor(
            [(T + 1) - i for i in range(1, T + 1)]).cuda().unsqueeze(0).unsqueeze(0) / (T / 2)

    def print_model_param(self, model: nn.Module, name: str):
        total = sum(p.numel() for p in model.parameters())
        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print_log(f'[{name}] Trainable/Total: {trainable}/{total}', self.log)

    def make_beta_schedule(self, schedule='linear', n_timesteps=1000, start=1e-5, end=1e-2):
        if schedule == 'linear':
            betas = torch.linspace(start, end, n_timesteps)
        elif schedule == 'quad':
            betas = torch.linspace(start ** 0.5, end ** 0.5, n_timesteps) ** 2
        elif schedule == 'sigmoid':
            betas = torch.linspace(-6, 6, n_timesteps)
            betas = torch.sigmoid(betas) * (end - start) + start
        return betas

    def extract(self, inp, t, x):
        shape = x.shape
        out = torch.gather(inp, 0, t.to(inp.device))
        reshape = [t.shape[0]] + [1] * (len(shape) - 1)
        return out.reshape(*reshape)

    # ------------------------------------------------------------------
    # Leapfrog reverse step (+ optional graph residual)
    # ------------------------------------------------------------------
    def p_sample_accelerate(self, x, mask, cur_y, t, sigma=None):
        step_idx = int(t)
        t = torch.tensor([t]).cuda()
        eps_factor = ((1 - self.extract(self.alphas, t, cur_y))
                      / self.extract(self.one_minus_alphas_bar_sqrt, t, cur_y))
        beta = self.extract(self.betas, t.repeat(x.shape[0]), cur_y)
        eps_theta = self.model.generate_accelerate(cur_y, beta, x, mask)

        if self.interaction_graph is not None:
            alpha_bar_sqrt_t = self.extract(self.alphas_bar_sqrt, t, cur_y)
            one_minus_abs_t  = self.extract(self.one_minus_alphas_bar_sqrt, t, cur_y)
            y0_hat = (cur_y - one_minus_abs_t * eps_theta) / alpha_bar_sqrt_t
            if self.use_v6_graph:
                delta = self.interaction_graph(y0_hat, x, step_idx, sigma=sigma)
            else:
                delta = self.interaction_graph(y0_hat, x, step_idx)
            if self.residual_on == 'eps':
                eps_theta = eps_theta + delta
            else:
                eps_theta = eps_theta - (alpha_bar_sqrt_t / one_minus_abs_t) * delta

        mean = (1 / self.extract(self.alphas, t, cur_y).sqrt()) \
               * (cur_y - (eps_factor * eps_theta))
        z = torch.randn_like(cur_y).to(x.device)
        sigma_t = self.extract(self.betas, t, cur_y).sqrt()
        return mean + sigma_t * z * 0.00001

    def p_sample_loop_accelerate(self, x, mask, loc, sigma=None):
        cur_y = loc[:, :10]
        for i in reversed(range(NUM_Tau)):
            cur_y = self.p_sample_accelerate(x, mask, cur_y, i, sigma=sigma)
        cur_y_ = loc[:, 10:]
        for i in reversed(range(NUM_Tau)):
            cur_y_ = self.p_sample_accelerate(x, mask, cur_y_, i, sigma=sigma)
        return torch.cat((cur_y_, cur_y), dim=1)

    # ------------------------------------------------------------------
    # Data preprocess (num_agents parameterized)
    # ------------------------------------------------------------------
    def data_preprocess(self, data):
        A = self.num_agents
        batch_size = data['pre_motion_3D'].shape[0]

        traj_mask = torch.zeros(batch_size * A, batch_size * A).cuda()
        for i in range(batch_size):
            traj_mask[i * A:(i + 1) * A, i * A:(i + 1) * A] = 1.

        pre = data['pre_motion_3D'].cuda()
        fut = data['fut_motion_3D'].cuda()
        initial_pos = pre[:, :, -1:]

        if self.per_scene_norm:
            scene_center = pre[:, :, -1, :].mean(dim=1, keepdim=True).unsqueeze(2)
            past_traj_abs = ((pre - scene_center) / self.traj_scale).contiguous().view(-1, self.cfg.past_frames, 2)
        else:
            past_traj_abs = ((pre - self.traj_mean) / self.traj_scale).contiguous().view(-1, self.cfg.past_frames, 2)
        past_traj_rel = ((pre - initial_pos)     / self.traj_scale).contiguous().view(-1, self.cfg.past_frames, 2)
        past_traj_vel = torch.cat(
            (past_traj_rel[:, 1:] - past_traj_rel[:, :-1],
             torch.zeros_like(past_traj_rel[:, -1:])), dim=1)
        past_traj = torch.cat((past_traj_abs, past_traj_rel, past_traj_vel), dim=-1)

        fut_traj = ((fut - initial_pos) / self.traj_scale).contiguous().view(-1, self.cfg.future_frames, 2)
        return batch_size, traj_mask, past_traj, fut_traj

    # ------------------------------------------------------------------
    # Training / validation
    # ------------------------------------------------------------------
    def fit(self):
        for epoch in range(self.cfg.num_epochs):
            loss_total, loss_dt, loss_dc = self._train_single_epoch(epoch)
            print_log(
                f'[{time.strftime("%Y-%m-%d %H:%M:%S")}] Epoch: {epoch}\t\tLoss: {loss_total:.6f}\t'
                f'Loss Dist.: {loss_dt:.6f}\tLoss Uncertainty: {loss_dc:.6f}', self.log)
            self.tb.add_scalar('train_epoch/loss_total', loss_total, epoch)
            self.tb.add_scalar('train_epoch/loss_dist_x50', loss_dt, epoch)
            self.tb.add_scalar('train_epoch/loss_uncertainty', loss_dc, epoch)
            self.tb.add_scalar('train_epoch/lr', self.opt.param_groups[0]['lr'], epoch)

            if (epoch + 1) % self.cfg.test_interval == 0:
                performance, samples = self._test_single_epoch()
                for i in range(4):
                    ade = performance['ADE'][i] / samples
                    fde = performance['FDE'][i] / samples
                    print_log(f'--ADE({i+1}s): {ade:.4f}\t--FDE({i+1}s): {fde:.4f}', self.log)
                    self.tb.add_scalar(f'val/ADE_{i+1}s', ade, epoch)
                    self.tb.add_scalar(f'val/FDE_{i+1}s', fde, epoch)

                cp_path = self.cfg.model_path % (epoch + 1)
                cp = {'model_initializer_dict': self.model_initializer.state_dict()}
                if self.interaction_graph is not None:
                    cp['interaction_graph_dict'] = self.interaction_graph.state_dict()
                torch.save(cp, cp_path)
            self.scheduler_model.step()
        self.tb.flush(); self.tb.close()

    def _train_single_epoch(self, epoch):
        self.model.train()
        self.model_initializer.train()
        if self.interaction_graph is not None:
            self.interaction_graph.train()

        loss_total, loss_dt, loss_dc, count = 0, 0, 0, 0
        for data in self.train_loader:
            batch_size, traj_mask, past_traj, fut_traj = self.data_preprocess(data)

            sample_prediction, mean_estimation, variance_estimation = self.model_initializer(past_traj, traj_mask)
            sample_prediction = torch.exp(variance_estimation / 2)[..., None, None] \
                              * sample_prediction \
                              / sample_prediction.std(dim=1).mean(dim=(1, 2))[:, None, None, None]
            loc = sample_prediction + mean_estimation[:, None]

            sigma_in = None if __import__('os').environ.get('LED_NO_SIGMA') else (variance_estimation if self.use_v6_graph else None)
            generated_y = self.p_sample_loop_accelerate(past_traj, traj_mask, loc, sigma=sigma_in)

            loss_dist = ((generated_y - fut_traj.unsqueeze(dim=1)).norm(p=2, dim=-1)
                         * self.temporal_reweight).mean(dim=-1).min(dim=1)[0].mean()
            loss_uncertainty = (torch.exp(-variance_estimation)
                                * (generated_y - fut_traj.unsqueeze(dim=1)).norm(p=2, dim=-1).mean(dim=(1, 2))
                                + variance_estimation).mean()

            loss = loss_dist * 50 + loss_uncertainty
            loss_total += loss.item()
            loss_dt    += loss_dist.item() * 50
            loss_dc    += loss_uncertainty.item()

            self.opt.zero_grad()
            loss.backward()
            params = list(self.model_initializer.parameters())
            if self.interaction_graph is not None:
                params += list(self.interaction_graph.parameters())
            grad_norm = torch.nn.utils.clip_grad_norm_(params, 1.)
            self.opt.step()

            self.tb.add_scalar('train_step/loss_total', loss.item(), self.global_step)
            self.tb.add_scalar('train_step/loss_dist_x50', loss_dist.item() * 50, self.global_step)
            self.tb.add_scalar('train_step/loss_uncertainty', loss_uncertainty.item(), self.global_step)
            self.tb.add_scalar('train_step/grad_norm', float(grad_norm), self.global_step)
            self.global_step += 1

            count += 1
            if self.cfg.debug and count == 2:
                break
        return loss_total / count, loss_dt / count, loss_dc / count

    def _test_single_epoch(self):
        performance = {'FDE': [0, 0, 0, 0], 'ADE': [0, 0, 0, 0]}
        samples = 0

        def prepare_seed(rand_seed):
            np.random.seed(rand_seed); random.seed(rand_seed)
            torch.manual_seed(rand_seed); torch.cuda.manual_seed_all(rand_seed)
        prepare_seed(0)

        self.model_initializer.eval()
        if self.interaction_graph is not None:
            self.interaction_graph.eval()

        # validation horizon: 4 checkpoints evenly across future_frames
        T_fut = self.cfg.future_frames
        step = max(1, T_fut // 4)
        horizons = [min(T_fut, step * (i + 1)) for i in range(4)]

        with torch.no_grad():
            for data in self.test_loader:
                batch_size, traj_mask, past_traj, fut_traj = self.data_preprocess(data)

                sample_prediction, mean_estimation, variance_estimation = self.model_initializer(past_traj, traj_mask)
                sample_prediction = torch.exp(variance_estimation / 2)[..., None, None] \
                                  * sample_prediction \
                                  / sample_prediction.std(dim=1).mean(dim=(1, 2))[:, None, None, None]
                loc = sample_prediction + mean_estimation[:, None]

                sigma_in = None if __import__('os').environ.get('LED_NO_SIGMA') else (variance_estimation if self.use_v6_graph else None)
                pred_traj = self.p_sample_loop_accelerate(past_traj, traj_mask, loc, sigma=sigma_in)

                fut_traj_k = fut_traj.unsqueeze(1).repeat(1, 20, 1, 1)
                distances = torch.norm(fut_traj_k - pred_traj, dim=-1) * self.traj_scale
                for i, h in enumerate(horizons):
                    ade = distances[:, :, :h].mean(dim=-1).min(dim=-1)[0].sum()
                    fde = distances[:, :, h - 1].min(dim=-1)[0].sum()
                    performance['ADE'][i] += ade.item()
                    performance['FDE'][i] += fde.item()
                samples += distances.shape[0]
        return performance, samples