""" 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