LED/trainer/train_led_graph.py

"""
LED trainer variant that augments the leapfrog denoising chain with a
future-trajectory interaction graph (ported in spirit from MoFlow's
FutureInteractionGraphV6 / MID graphv6_v5 ideas).

At each of the NUM_Tau leapfrog reverse steps, the frozen core denoiser
produces an epsilon prediction; we recover the implied y_0 estimate,
run a small inter-agent graph on the predicted future trajectories, and
add its output as a residual correction to epsilon. The graph is the
only new trainable module besides the existing LED initializer.

This file mirrors trainer/train_led_trajectory_augment_input.py so that
the baseline script stays untouched and both variants can be run
side-by-side.
"""

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_nba import NBADataset, 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)

		# ------------------------- prepare train/test data loader -------------------------
		train_dset = NBADataset(
			obs_len=self.cfg.past_frames,
			pred_len=self.cfg.future_frames,
			training=True)

		self.train_loader = DataLoader(
			train_dset,
			batch_size=self.cfg.train_batch_size,
			shuffle=True,
			num_workers=4,
			collate_fn=seq_collate,
			pin_memory=True)

		test_dset = NBADataset(
			obs_len=self.cfg.past_frames,
			pred_len=self.cfg.future_frames,
			training=False)

		self.test_loader = DataLoader(
			test_dset,
			batch_size=self.cfg.test_batch_size,
			shuffle=False,
			num_workers=4,
			collate_fn=seq_collate,
			pin_memory=True)

		self.traj_mean = torch.FloatTensor(self.cfg.traj_mean).cuda().unsqueeze(0).unsqueeze(0).unsqueeze(0)
		self.traj_scale = self.cfg.traj_scale

		# ------------------------- define 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)

		# ------------------------- define models -------------------------
		self.model = CoreDenoisingModel().cuda()
		model_cp = torch.load(self.cfg.pretrained_core_denoising_model, map_location='cpu')
		self.model.load_state_dict(model_cp['model_dict'])

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

		self.residual_on = getattr(config, 'residual_on', 'eps')
		self.use_sigma = bool(getattr(config, 'use_sigma', False))
		self.use_v6_graph = bool(getattr(config, 'use_v6_graph', False))
		self.uncertainty_weight = float(getattr(config, 'uncertainty_weight', 1.0))
		top_n = getattr(config, 'top_n', 5)
		edge_mode = getattr(config, 'edge_mode', 'full')
		neighbor_mode = getattr(config, 'neighbor_mode', 'rag')
		if self.use_v6_graph:
			self.interaction_graph = FutureInteractionGraphV6Wrapper(
				num_agents=11,
				future_steps=self.cfg.future_frames,
				past_steps=self.cfg.past_frames,
				past_channels=6,
				node_dim=128,
				top_n=top_n,
				num_denoise_steps=NUM_Tau,
				edge_mode=edge_mode,
				neighbor_mode=neighbor_mode,
			).cuda()
		else:
			self.interaction_graph = FutureInteractionGraph(
				num_agents=11,
				future_steps=self.cfg.future_frames,
				past_steps=self.cfg.past_frames,
				past_channels=6,
				node_dim=128,
				top_n=top_n,
				num_denoise_steps=NUM_Tau,
				use_sigma=self.use_sigma,
			).cuda()

		self.opt = torch.optim.AdamW(
			list(self.model_initializer.parameters())
			+ list(self.interaction_graph.parameters()),
			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)

		self.resume_epoch = int(getattr(config, 'resume_epoch', 0))
		if self.resume_epoch > 0:
			cp_path = self.cfg.model_path % self.resume_epoch
			cp = torch.load(cp_path, map_location='cpu')
			self.model_initializer.load_state_dict(cp['model_initializer_dict'])
			self.interaction_graph.load_state_dict(cp['interaction_graph_dict'])
			for _ in range(self.resume_epoch):
				self.scheduler_model.step()

		# ------------------------- prepare 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')
		self.print_model_param(self.interaction_graph, name='Future Interaction Graph')

		self.temporal_reweight = torch.FloatTensor(
			[21 - i for i in range(1, 21)]).cuda().unsqueeze(0).unsqueeze(0) / 10

	def print_model_param(self, model: nn.Module, name: str = 'Model') -> None:
		total_num = sum(p.numel() for p in model.parameters())
		trainable_num = sum(p.numel() for p in model.parameters() if p.requires_grad)
		print_log("[{}] Trainable/Total: {}/{}".format(name, trainable_num, total_num), self.log)

	def make_beta_schedule(self, schedule: str = 'linear',
			n_timesteps: int = 1000,
			start: float = 1e-5, end: float = 1e-2) -> torch.Tensor:
		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, input, t, x):
		shape = x.shape
		out = torch.gather(input, 0, t.to(input.device))
		reshape = [t.shape[0]] + [1] * (len(shape) - 1)
		return out.reshape(*reshape)

	# ------------------------------------------------------------------
	# Leapfrog reverse step with graph residual on epsilon
	# ------------------------------------------------------------------
	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)

		# Implied y_0 estimate for the graph input.
		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

		delta = self.interaction_graph(y0_hat, x, step_idx, sigma=sigma)
		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()
		sample = mean + sigma_t * z * 0.00001
		return sample

	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)
		prediction_total = torch.cat((cur_y_, cur_y), dim=1)
		return prediction_total

	# ------------------------------------------------------------------
	# Training / evaluation
	# ------------------------------------------------------------------
	def fit(self):
		for epoch in range(self.resume_epoch, self.cfg.num_epochs):
			loss_total, loss_distance, loss_uncertainty = self._train_single_epoch(epoch)
			print_log('[{}] Epoch: {}\t\tLoss: {:.6f}\tLoss Dist.: {:.6f}\tLoss Uncertainty: {:.6f}'.format(
				time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
				epoch, loss_total, loss_distance, loss_uncertainty), self.log)

			self.tb.add_scalar('train_epoch/loss_total', loss_total, epoch)
			self.tb.add_scalar('train_epoch/loss_dist_x50', loss_distance, epoch)
			self.tb.add_scalar('train_epoch/loss_uncertainty', loss_uncertainty, 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 time_i in range(4):
					ade = performance['ADE'][time_i] / samples
					fde = performance['FDE'][time_i] / samples
					print_log('--ADE({}s): {:.4f}\t--FDE({}s): {:.4f}'.format(
						time_i + 1, ade, time_i + 1, fde), self.log)
					self.tb.add_scalar('val/ADE_{}s'.format(time_i + 1), ade, epoch)
					self.tb.add_scalar('val/FDE_{}s'.format(time_i + 1), fde, epoch)
				cp_path = self.cfg.model_path % (epoch + 1)
				model_cp = {
					'model_initializer_dict': self.model_initializer.state_dict(),
					'interaction_graph_dict': self.interaction_graph.state_dict(),
				}
				torch.save(model_cp, cp_path)
			self.scheduler_model.step()
		self.tb.flush()
		self.tb.close()

	def data_preprocess(self, data):
		batch_size = data['pre_motion_3D'].shape[0]
		traj_mask = torch.zeros(batch_size * 11, batch_size * 11).cuda()
		for i in range(batch_size):
			traj_mask[i * 11:(i + 1) * 11, i * 11:(i + 1) * 11] = 1.

		initial_pos = data['pre_motion_3D'].cuda()[:, :, -1:]
		past_traj_abs = ((data['pre_motion_3D'].cuda() - self.traj_mean) / self.traj_scale).contiguous().view(-1, 10, 2)
		past_traj_rel = ((data['pre_motion_3D'].cuda() - initial_pos) / self.traj_scale).contiguous().view(-1, 10, 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 = ((data['fut_motion_3D'].cuda() - initial_pos) / self.traj_scale).contiguous().view(-1, 20, 2)
		return batch_size, traj_mask, past_traj, fut_traj

	def _train_single_epoch(self, epoch):
		self.model.train()
		self.model_initializer.train()
		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)

			# --- σ 안정화 가드 (기본 전부 off => 원본과 비트 단위로 동일) --------------
			# 배경: --use_sigma 로 logvar 를 그래프에 넣으면 NLL 이외의 gradient 경로가
			#       하나 더 생긴다. logvar 가 음수로 밀리면 exp(-logvar) 가 폭주하고
			#       exp(logvar/2)*x / std(x) 의 분모도 0 으로 가서 NaN 이 된다.
			#       (실제로 LED full SRA σ ON 이 epoch 4 에서 이렇게 죽었다)
			_clamp = float(os.environ.get('LED_SIGMA_CLAMP', 0.0) or 0.0)
			if _clamp > 0:
				variance_estimation = variance_estimation.clamp(-_clamp, _clamp)

			sample_prediction = (torch.exp(variance_estimation / 2)[..., None, None]
								 * sample_prediction
								 / sample_prediction.std(dim=1).mean(dim=(1, 2))[:, None, None, None]
								   .clamp_min(float(os.environ.get('LED_STD_EPS', 0.0) or 0.0)))
			loc = sample_prediction + mean_estimation[:, None]

			sigma_input = variance_estimation if self.use_sigma else None
			# σ 를 '게이트 신호'로만 쓰고 분산 헤드로 역전파하지 않는 옵션
			if sigma_input is not None and os.environ.get('LED_SIGMA_DETACH', '') not in ('', '0', 'false', 'False'):
				sigma_input = sigma_input.detach()
			generated_y = self.p_sample_loop_accelerate(past_traj, traj_mask, loc, sigma=sigma_input)

			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()

			# logvar 추이를 남긴다 (NaN 이 나면 원인을 사후에 알 수 있게)
			if count % 200 == 0:
				with torch.no_grad():
					_v = variance_estimation.detach()
					self.tb.add_scalar('sigma/logvar_min',  float(_v.min()),  self.global_step)
					self.tb.add_scalar('sigma/logvar_max',  float(_v.max()),  self.global_step)
					self.tb.add_scalar('sigma/logvar_mean', float(_v.mean()), self.global_step)
					self.tb.add_scalar('sigma/pred_std_min',
									   float(sample_prediction.std(dim=1).mean(dim=(1, 2)).min()),
									   self.global_step)

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

			self.opt.zero_grad()
			loss.backward()
			grad_norm = torch.nn.utils.clip_grad_norm_(
				list(self.model_initializer.parameters())
				+ list(self.interaction_graph.parameters()),
				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()
		self.interaction_graph.eval()
		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)
				_c = float(os.environ.get('LED_SIGMA_CLAMP', 0.0) or 0.0)
				if _c > 0:
					variance_estimation = variance_estimation.clamp(-_c, _c)
				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_input = variance_estimation if self.use_sigma else None
				pred_traj = self.p_sample_loop_accelerate(past_traj, traj_mask, loc, sigma=sigma_input)

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

	def test_single_model(self):
		model_path = './results/checkpoints/led_graph.p'
		ckpt = torch.load(model_path, map_location=torch.device('cpu'))
		self.model_initializer.load_state_dict(ckpt['model_initializer_dict'])
		if 'interaction_graph_dict' in ckpt:
			self.interaction_graph.load_state_dict(ckpt['interaction_graph_dict'])
		else:
			print_log('WARNING: checkpoint has no interaction_graph_dict; '
					  'using zero-initialized graph (equivalent to baseline LED).',
					  log=self.log)

		performance = {'FDE': [0, 0, 0, 0], 'ADE': [0, 0, 0, 0]}
		samples = 0
		print_log(model_path, log=self.log)

		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()
		self.interaction_graph.eval()
		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)
				_c = float(os.environ.get('LED_SIGMA_CLAMP', 0.0) or 0.0)
				if _c > 0:
					variance_estimation = variance_estimation.clamp(-_c, _c)
				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_input = variance_estimation if self.use_sigma else None
				pred_traj = self.p_sample_loop_accelerate(past_traj, traj_mask, loc, sigma=sigma_input)

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

		for time_i in range(4):
			print_log('--ADE({}s): {:.4f}\t--FDE({}s): {:.4f}'.format(
				time_i + 1, performance['ADE'][time_i] / samples,
				time_i + 1, performance['FDE'][time_i] / samples), log=self.log)