SRA: MID/LED/MoFlow code + RUNNING.md instructions (code only, no data/ckpts)

LED/LED/cfg/nba/led_augment.yml

@@ -0,0 +1,38 @@
+# ------------------- General Options -------------------------
+description                  : LED
+results_root_dir             : results
+dataset                      : nba
+
+# ------------------- Dataset -------------------------
+past_frames                  : 10
+future_frames                : 20
+min_past_frames              : 10
+min_future_frames            : 20
+
+motion_dim                   : 2
+forecast_dim                 : 2
+
+traj_mean                    : [14, 7.5]
+traj_scale                   : 5
+
+# ------------------- Model -------------------------
+pretrained_core_denoising_model: './results/checkpoints/base_diffusion_model.p'
+debug                        : False # set to True for early stop in each epoch.
+
+diffusion                    : {
+  steps                      : 100,
+  beta_start                 : 1.e-4,
+  beta_end                   : 5.e-2,
+  beta_schedule              : 'linear'
+}
+
+# ------------------- Training Parameters -------------------------
+lr                           : 1.e-3
+train_batch_size             : 10
+test_batch_size              : 500
+num_epochs                   : 100
+
+lr_scheduler                 : 'step'
+decay_step                   : 8
+decay_gamma                  : 0.5
+

LED/LED/cfg/nba/led_augment_debug.yml

@@ -0,0 +1,38 @@
+# ------------------- General Options -------------------------
+description                  : LED
+results_root_dir             : results
+dataset                      : nba
+
+# ------------------- Dataset -------------------------
+past_frames                  : 10
+future_frames                : 20
+min_past_frames              : 10
+min_future_frames            : 20
+
+motion_dim                   : 2
+forecast_dim                 : 2
+
+traj_mean                    : [14, 7.5]
+traj_scale                   : 5
+
+# ------------------- Model -------------------------
+pretrained_core_denoising_model: './results/checkpoints/base_diffusion_model.p'
+debug                        : True # set to True for early stop in each epoch.
+
+diffusion                    : {
+  steps                      : 100,
+  beta_start                 : 1.e-4,
+  beta_end                   : 5.e-2,
+  beta_schedule              : 'linear'
+}
+
+# ------------------- Training Parameters -------------------------
+lr                           : 1.e-3
+train_batch_size             : 2
+test_batch_size              : 5
+num_epochs                   : 100
+test_interval                : 2
+lr_scheduler                 : 'step'
+decay_step                   : 8
+decay_gamma                  : 0.5
+

LED/LED/main_led_nba.py

@@ -0,0 +1,31 @@
+import argparse
+from trainer import train_led_trajectory_augment_input as led
+
+
+def parse_config():
+	parser = argparse.ArgumentParser()
+	parser.add_argument("--cuda", default=True)
+	parser.add_argument("--learning_rate", type=int, default=0.002)
+	parser.add_argument("--max_epochs", type=int, default=128)
+
+	parser.add_argument('--cfg', default='led_augment')
+	parser.add_argument('--gpu', type=int, default=0, help='Specify which GPU to use.')
+	parser.add_argument('--train', type=int, default=1, help='Whether train or evaluate.')
+	
+	parser.add_argument("--info", type=str, default='', help='Name of the experiment. '
+															 'It will be used in file creation.')
+	return parser.parse_args()
+
+
+def main(config):
+	t = led.Trainer(config)
+	if config.train==1:
+		t.fit()
+	else:
+		# t.save_data()
+		t.test_single_model()
+
+
+if __name__ == "__main__":
+	config = parse_config()
+	main(config)

LED/LED/models/layers.py

@@ -0,0 +1,169 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import Module, Linear
+
+
+class PositionalEncoding(nn.Module):
+	def __init__(self, d_model, dropout=0.1, max_len=5000):
+		super().__init__()
+
+		self.dropout = nn.Dropout(p=dropout)
+		pe = torch.zeros(max_len, d_model)
+		position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+		div_term = torch.exp(
+			torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
+		)
+		pe[:, 0::2] = torch.sin(position * div_term)
+		pe[:, 1::2] = torch.cos(position * div_term)
+		pe = pe.unsqueeze(0).transpose(0, 1)
+		self.register_buffer("pe", pe)
+
+	def forward(self, x):
+		x = x + self.pe[: x.size(0), :]
+		return self.dropout(x)
+
+
+class ConcatSquashLinear(Module):
+	def __init__(self, dim_in, dim_out, dim_ctx):
+		super(ConcatSquashLinear, self).__init__()
+		self._layer = Linear(dim_in, dim_out)
+		self._hyper_bias = Linear(dim_ctx, dim_out, bias=False)
+		self._hyper_gate = Linear(dim_ctx, dim_out)
+
+	def forward(self, ctx, x):
+		# ctx: (B, 1, F+3)
+		# x: (B, T, 2)
+		gate = torch.sigmoid(self._hyper_gate(ctx))
+		bias = self._hyper_bias(ctx)
+		# if x.dim() == 3:
+		#     gate = gate.unsqueeze(1)
+		#     bias = bias.unsqueeze(1)
+		ret = self._layer(x) * gate + bias
+		return ret
+	
+	def batch_generate(self, ctx, x):
+		# ctx: (B, n, 1, F+3)
+		# x: (B, n, T, 2)
+		gate = torch.sigmoid(self._hyper_gate(ctx))
+		bias = self._hyper_bias(ctx)
+		# if x.dim() == 3:
+		#     gate = gate.unsqueeze(1)
+		#     bias = bias.unsqueeze(1)
+		ret = self._layer(x) * gate + bias
+		return ret
+	
+
+class GAT(nn.Module):
+	def __init__(self, in_feat=2, out_feat=64, n_head=4, dropout=0.1, skip=True):
+		super(GAT, self).__init__()
+		self.in_feat = in_feat
+		self.out_feat = out_feat
+		self.n_head = n_head
+		self.skip = skip
+		self.w = nn.Parameter(torch.Tensor(n_head, in_feat, out_feat))
+		self.a_src = nn.Parameter(torch.Tensor(n_head, out_feat, 1))
+		self.a_dst = nn.Parameter(torch.Tensor(n_head, out_feat, 1))
+		self.bias = nn.Parameter(torch.Tensor(out_feat))
+
+		self.leaky_relu = nn.LeakyReLU(negative_slope=0.2)
+		self.softmax = nn.Softmax(dim=-1)
+		self.dropout = nn.Dropout(dropout)
+
+		nn.init.xavier_uniform_(self.w, gain=1.414)
+		nn.init.xavier_uniform_(self.a_src, gain=1.414)
+		nn.init.xavier_uniform_(self.a_dst, gain=1.414)
+		nn.init.constant_(self.bias, 0)
+
+	def forward(self, h, mask):
+		h_prime = h.unsqueeze(1) @ self.w
+		attn_src = h_prime @ self.a_src
+		attn_dst = h_prime @ self.a_dst
+		attn = attn_src @ attn_dst.permute(0, 1, 3, 2)
+		attn = self.leaky_relu(attn)
+		attn = self.softmax(attn)
+		attn = self.dropout(attn)
+		attn = attn * mask if mask is not None else attn
+		out = (attn @ h_prime).sum(dim=1) + self.bias
+		if self.skip:
+			out += h_prime.sum(dim=1)
+		return out, attn
+
+
+class MLP(nn.Module):
+	def __init__(self, in_feat, out_feat, hid_feat=(1024, 512), activation=None, dropout=-1):
+		super(MLP, self).__init__()
+		dims = (in_feat, ) + hid_feat + (out_feat, )
+
+		self.layers = nn.ModuleList()
+		for i in range(len(dims) - 1):
+			self.layers.append(nn.Linear(dims[i], dims[i + 1]))
+
+		self.activation = activation if activation is not None else lambda x: x
+		self.dropout = nn.Dropout(dropout) if dropout != -1 else lambda x: x
+
+	def forward(self, x):
+		for i in range(len(self.layers)):
+			x = self.activation(x)
+			x = self.dropout(x)
+			x = self.layers[i](x)
+		return x
+
+
+class social_transformer(nn.Module):
+	def __init__(self, past_len):
+		super(social_transformer, self).__init__()
+		self.encode_past = nn.Linear(past_len*6, 256, bias=False)
+		self.layer = nn.TransformerEncoderLayer(d_model=256, nhead=2, dim_feedforward=256)
+		self.transformer_encoder = nn.TransformerEncoder(self.layer, num_layers=2)
+
+	def forward(self, h, mask):
+		'''
+		h: batch_size, t, 2
+		'''
+		h_feat = self.encode_past(h.reshape(h.size(0), -1)).unsqueeze(1)
+		# print(h_feat.shape)
+		# n_samples, 1, 64
+		h_feat_ = self.transformer_encoder(h_feat, mask)
+		h_feat = h_feat + h_feat_
+
+		return h_feat
+
+
+class st_encoder(nn.Module):
+	def __init__(self):
+		super().__init__()
+		channel_in = 6
+		channel_out = 32
+		dim_kernel = 3
+		self.dim_embedding_key = 256
+		self.spatial_conv = nn.Conv1d(channel_in, channel_out, dim_kernel, stride=1, padding=1)
+		self.temporal_encoder = nn.GRU(channel_out, self.dim_embedding_key, 1, batch_first=True)
+
+		self.relu = nn.ReLU()
+
+		self.reset_parameters()
+
+	def reset_parameters(self):
+		nn.init.kaiming_normal_(self.spatial_conv.weight)
+		nn.init.kaiming_normal_(self.temporal_encoder.weight_ih_l0)
+		nn.init.kaiming_normal_(self.temporal_encoder.weight_hh_l0)
+		nn.init.zeros_(self.spatial_conv.bias)
+		nn.init.zeros_(self.temporal_encoder.bias_ih_l0)
+		nn.init.zeros_(self.temporal_encoder.bias_hh_l0)
+
+	def forward(self, X):
+		'''
+		X: b, T, 2
+
+		return: b, F
+		'''
+		X_t = torch.transpose(X, 1, 2)
+		X_after_spatial = self.relu(self.spatial_conv(X_t))
+		X_embed = torch.transpose(X_after_spatial, 1, 2)
+
+		output_x, state_x = self.temporal_encoder(X_embed)
+		state_x = state_x.squeeze(0)
+
+		return state_x
+

LED/LED/models/model_diffusion.py

@@ -0,0 +1,124 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import Module, Linear
+
+from models.layers import PositionalEncoding, ConcatSquashLinear
+
+class st_encoder(nn.Module):
+	def __init__(self):
+		super().__init__()
+		channel_in = 2
+		channel_out = 32
+		dim_kernel = 3
+		self.dim_embedding_key = 256
+		self.spatial_conv = nn.Conv1d(channel_in, channel_out, dim_kernel, stride=1, padding=1)
+		self.temporal_encoder = nn.GRU(channel_out, self.dim_embedding_key, 1, batch_first=True)
+
+		self.relu = nn.ReLU()
+
+		self.reset_parameters()
+
+	def reset_parameters(self):
+		nn.init.kaiming_normal_(self.spatial_conv.weight)
+		nn.init.kaiming_normal_(self.temporal_encoder.weight_ih_l0)
+		nn.init.kaiming_normal_(self.temporal_encoder.weight_hh_l0)
+		nn.init.zeros_(self.spatial_conv.bias)
+		nn.init.zeros_(self.temporal_encoder.bias_ih_l0)
+		nn.init.zeros_(self.temporal_encoder.bias_hh_l0)
+
+	def forward(self, X):
+		'''
+		X: b, T, 2
+
+		return: b, F
+		'''
+		X_t = torch.transpose(X, 1, 2)
+		X_after_spatial = self.relu(self.spatial_conv(X_t))
+		X_embed = torch.transpose(X_after_spatial, 1, 2)
+
+		output_x, state_x = self.temporal_encoder(X_embed)
+		state_x = state_x.squeeze(0)
+
+		return state_x
+
+
+class social_transformer(nn.Module):
+	def __init__(self):
+		super(social_transformer, self).__init__()
+		self.encode_past = nn.Linear(60, 256, bias=False)
+		# self.encode_past = nn.Linear(48, 256, bias=False)
+		self.layer = nn.TransformerEncoderLayer(d_model=256, nhead=2, dim_feedforward=256)
+		self.transformer_encoder = nn.TransformerEncoder(self.layer, num_layers=2)
+
+	def forward(self, h, mask):
+		'''
+		h: batch_size, t, 2
+		'''
+		# print(h.shape)
+		h_feat = self.encode_past(h.reshape(h.size(0), -1)).unsqueeze(1)
+		# print(h_feat.shape)
+		# n_samples, 1, 64
+		h_feat_ = self.transformer_encoder(h_feat, mask)
+		h_feat = h_feat + h_feat_
+
+		return h_feat
+
+
+class TransformerDenoisingModel(Module):
+
+	def __init__(self, context_dim=256, tf_layer=2):
+		super().__init__()
+		self.encoder_context = social_transformer()
+		self.pos_emb = PositionalEncoding(d_model=2*context_dim, dropout=0.1, max_len=24)
+		self.concat1 = ConcatSquashLinear(2, 2*context_dim, context_dim+3)
+		self.layer = nn.TransformerEncoderLayer(d_model=2*context_dim, nhead=2, dim_feedforward=2*context_dim)
+		self.transformer_encoder = nn.TransformerEncoder(self.layer, num_layers=tf_layer)
+		self.concat3 = ConcatSquashLinear(2*context_dim,context_dim,context_dim+3)
+		self.concat4 = ConcatSquashLinear(context_dim,context_dim//2,context_dim+3)
+		self.linear = ConcatSquashLinear(context_dim//2, 2, context_dim+3)
+
+
+	def forward(self, x, beta, context, mask):
+		batch_size = x.size(0)
+		beta = beta.view(batch_size, 1, 1)          # (B, 1, 1)
+		mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+		context = self.encoder_context(context, mask)
+		# context = context.view(batch_size, 1, -1)   # (B, 1, F)
+
+		time_emb = torch.cat([beta, torch.sin(beta), torch.cos(beta)], dim=-1)  # (B, 1, 3)
+		ctx_emb = torch.cat([time_emb, context], dim=-1)    # (B, 1, F+3)
+		
+		x = self.concat1(ctx_emb, x)
+		final_emb = x.permute(1,0,2)
+		final_emb = self.pos_emb(final_emb)
+		
+		trans = self.transformer_encoder(final_emb).permute(1,0,2)
+		trans = self.concat3(ctx_emb, trans)
+		trans = self.concat4(ctx_emb, trans)
+		return self.linear(ctx_emb, trans)
+	
+
+	def generate_accelerate(self, x, beta, context, mask):
+		batch_size = x.size(0)
+		beta = beta.view(beta.size(0), 1, 1)          # (B, 1, 1)
+		mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+		context = self.encoder_context(context, mask)
+		# context = context.view(batch_size, 1, -1)   # (B, 1, F)
+
+		time_emb = torch.cat([beta, torch.sin(beta), torch.cos(beta)], dim=-1)  # (B, 1, 3)
+		# time_emb: [11, 1, 3]
+		# context: [11, 1, 256]
+		ctx_emb = torch.cat([time_emb, context], dim=-1).repeat(1, 10, 1).unsqueeze(2)
+		# x: 11, 10, 20, 2
+		# ctx_emb: 11, 10, 1, 259
+		x = self.concat1.batch_generate(ctx_emb, x).contiguous().view(-1, 20, 512)
+		# x: 110, 20, 512
+		final_emb = x.permute(1, 0, 2)
+		final_emb = self.pos_emb(final_emb)
+		
+		trans = self.transformer_encoder(final_emb).permute(1, 0, 2).contiguous().view(-1, 10, 20, 512)
+		# trans: 11, 10, 20, 512
+		trans = self.concat3.batch_generate(ctx_emb, trans)
+		trans = self.concat4.batch_generate(ctx_emb, trans)
+		return self.linear.batch_generate(ctx_emb, trans)

LED/LED/models/model_led_initializer.py

@@ -0,0 +1,63 @@
+import torch
+import torch.nn as nn
+from models.layers import MLP, social_transformer, st_encoder
+
+class LEDInitializer(nn.Module):
+	def __init__(self, t_h: int=8, d_h: int=6, t_f: int=40, d_f: int=2, k_pred: int=20):
+		'''
+		Parameters
+		----
+		t_h: history timestamps,
+		d_h: dimension of each historical timestamp,
+		t_f: future timestamps,
+		d_f: dimension of each future timestamp,
+		k_pred: number of predictions.
+
+		'''
+		super(LEDInitializer, self).__init__()
+		self.n = k_pred
+		self.input_dim = t_h * d_h
+		self.output_dim = t_f * d_f * k_pred
+		self.fut_len = t_f
+
+		self.social_encoder = social_transformer(t_h)
+		self.ego_var_encoder = st_encoder()
+		self.ego_mean_encoder = st_encoder()
+		self.ego_scale_encoder = st_encoder()
+
+		self.scale_encoder = MLP(1, 32, hid_feat=(4, 16), activation=nn.ReLU())
+
+		self.var_decoder = MLP(256*2+32, self.output_dim, hid_feat=(1024, 1024), activation=nn.ReLU())
+		self.mean_decoder = MLP(256*2, t_f * d_f, hid_feat=(256, 128), activation=nn.ReLU())
+		self.scale_decoder = MLP(256*2, 1, hid_feat=(256, 128), activation=nn.ReLU())
+
+	
+	def forward(self, x, mask=None):
+		'''
+		x: batch size, t_p, 6
+		'''
+		mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+		social_embed = self.social_encoder(x, mask)
+		social_embed = social_embed.squeeze(1)
+		# B, 256
+		
+		ego_var_embed = self.ego_var_encoder(x)
+		ego_mean_embed = self.ego_mean_encoder(x)
+		ego_scale_embed = self.ego_scale_encoder(x)
+		# B, 256
+
+		mean_total = torch.cat((ego_mean_embed, social_embed), dim=-1)
+		
+		guess_mean = self.mean_decoder(mean_total).contiguous().view(-1, self.fut_len, 2)
+
+		scale_total = torch.cat((ego_scale_embed, social_embed), dim=-1)
+		guess_scale = self.scale_decoder(scale_total)
+
+		guess_scale_feat = self.scale_encoder(guess_scale)
+		var_total = torch.cat((ego_var_embed, social_embed, guess_scale_feat), dim=-1)
+		guess_var = self.var_decoder(var_total).reshape(x.size(0), self.n, self.fut_len, 2)
+
+		return guess_var, guess_mean, guess_scale
+
+
+

LED/LED/requirements.txt

@@ -0,0 +1,84 @@
+absl-py==1.4.0
+aiohttp==3.8.4
+aiosignal==1.3.1
+antlr4-python3-runtime==4.8
+async-timeout==4.0.2
+asynctest==0.13.0
+attrs==23.1.0
+cachetools==5.3.0
+charset-normalizer==3.1.0
+colour==0.1.5
+cycler==0.11.0
+descartes==1.1.0
+easydict==1.10
+fonttools==4.38.0
+frozenlist==1.3.3
+fsspec==2023.1.0
+future==0.18.3
+glob2==0.7
+google-auth==2.18.0
+google-auth-oauthlib==0.4.6
+googledrivedownloader==0.4
+grpcio==1.54.0
+h5py==3.8.0
+hydra-core==1.1.0
+imageio==2.27.0
+imgaug==0.4.0
+importlib-metadata==4.13.0
+importlib-resources==5.12.0
+isodate==0.6.1
+Jinja2==3.1.2
+joblib==1.2.0
+kiwisolver==1.4.4
+lapsolver==1.1.0
+llvmlite==0.39.1
+Markdown==3.4.3
+MarkupSafe==2.1.2
+matplotlib==3.5.3
+motmetrics==1.1.3
+multidict==6.0.4
+networkx==2.6.3
+numba==0.56.4
+numpy==1.19.0
+oauthlib==3.2.2
+omegaconf==2.1.0
+opencv-python==4.7.0.72
+packaging==23.1
+pandas==1.3.5
+Pillow==9.5.0
+polars==0.17.12
+protobuf==3.20.3
+pyasn1==0.5.0
+pyasn1-modules==0.3.0
+pyDeprecate==0.3.1
+pyntcloud==0.3.1
+pyparsing==3.0.9
+python-dateutil==2.8.2
+python-louvain==0.16
+pytorch-lightning==1.5.2
+pytz==2023.3
+PyWavelets==1.3.0
+PyYAML==6.0
+rdflib==6.3.2
+requests-oauthlib==1.3.1
+rsa==4.9
+scikit-image==0.19.3
+scikit-learn==1.0.2
+scipy==1.7.3
+shapely==2.0.1
+spconv==1.2.1
+tensorboard==2.11.2
+tensorboard-data-server==0.6.1
+tensorboard-plugin-wit==1.8.1
+tensorboardX==2.6
+threadpoolctl==3.1.0
+tifffile==2021.11.2
+torch==1.8.0+cu111
+torchaudio==0.8.0
+torchmetrics==0.11.4
+torchvision==0.9.0+cu111
+tqdm==4.65.0
+typing_extensions==4.5.0
+Werkzeug==2.2.3
+yarl==1.9.2
+zipp==3.15.0

LED/LED/trainer/train_led_trajectory_augment_input.py

@@ -0,0 +1,443 @@
+
+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 data.dataloader_nba import NBADataset, seq_collate
+
+
+from models.model_led_initializer import LEDInitializer as InitializationModel
+from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel
+
+import pdb
+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)
+		
+		# data normalization parameters
+		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 # define total diffusion steps
+
+		# make beta schedule and calculate the parameters used in denoising process.
+		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()
+		# load pretrained models
+		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.opt = torch.optim.AdamW(self.model_initializer.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)
+		
+		# ------------------------- prepare logs -------------------------
+		self.log = open(os.path.join(self.cfg.log_dir, 'log.txt'), 'a+')
+		self.print_model_param(self.model, name='Core Denoising Model')
+		self.print_model_param(self.model_initializer, name='Initialization Model')
+
+		# temporal reweight in the loss, it is not necessary.
+		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:
+		'''
+		Count the trainable/total parameters in `model`.
+		'''
+		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)
+		return None
+
+
+	def make_beta_schedule(self, schedule: str = 'linear', 
+			n_timesteps: int = 1000, 
+			start: float = 1e-5, end: float = 1e-2) -> torch.Tensor:
+		'''
+		Make beta schedule.
+
+		Parameters
+		----
+		schedule: str, in ['linear', 'quad', 'sigmoid'],
+		n_timesteps: int, diffusion steps,
+		start: float, beta start, `start<end`,
+		end: float, beta end,
+
+		Returns
+		----
+		betas: Tensor with the shape of (n_timesteps)
+
+		'''
+		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)
+
+	def noise_estimation_loss(self, x, y_0, mask):
+		batch_size = x.shape[0]
+		# Select a random step for each example
+		t = torch.randint(0, self.n_steps, size=(batch_size // 2 + 1,)).to(x.device)
+		t = torch.cat([t, self.n_steps - t - 1], dim=0)[:batch_size]
+		# x0 multiplier
+		a = self.extract(self.alphas_bar_sqrt, t, y_0)
+		beta = self.extract(self.betas, t, y_0)
+		# eps multiplier
+		am1 = self.extract(self.one_minus_alphas_bar_sqrt, t, y_0)
+		e = torch.randn_like(y_0)
+		# model input
+		y = y_0 * a + e * am1
+		output = self.model(y, beta, x, mask)
+		# batch_size, 20, 2
+		return (e - output).square().mean()
+
+
+
+	def p_sample(self, x, mask, cur_y, t):
+		if t==0:
+			z = torch.zeros_like(cur_y).to(x.device)
+		else:
+			z = torch.randn_like(cur_y).to(x.device)
+		t = torch.tensor([t]).cuda()
+		# Factor to the model output
+		eps_factor = ((1 - self.extract(self.alphas, t, cur_y)) / self.extract(self.one_minus_alphas_bar_sqrt, t, cur_y))
+		# Model output
+		beta = self.extract(self.betas, t.repeat(x.shape[0]), cur_y)
+		eps_theta = self.model(cur_y, beta, x, mask)
+		mean = (1 / self.extract(self.alphas, t, cur_y).sqrt()) * (cur_y - (eps_factor * eps_theta))
+		# Generate z
+		z = torch.randn_like(cur_y).to(x.device)
+		# Fixed sigma
+		sigma_t = self.extract(self.betas, t, cur_y).sqrt()
+		sample = mean + sigma_t * z
+		return (sample)
+	
+	def p_sample_accelerate(self, x, mask, cur_y, t):
+		if t==0:
+			z = torch.zeros_like(cur_y).to(x.device)
+		else:
+			z = torch.randn_like(cur_y).to(x.device)
+		t = torch.tensor([t]).cuda()
+		# Factor to the model output
+		eps_factor = ((1 - self.extract(self.alphas, t, cur_y)) / self.extract(self.one_minus_alphas_bar_sqrt, t, cur_y))
+		# Model output
+		beta = self.extract(self.betas, t.repeat(x.shape[0]), cur_y)
+		eps_theta = self.model.generate_accelerate(cur_y, beta, x, mask)
+		mean = (1 / self.extract(self.alphas, t, cur_y).sqrt()) * (cur_y - (eps_factor * eps_theta))
+		# Generate z
+		z = torch.randn_like(cur_y).to(x.device)
+		# Fixed sigma
+		sigma_t = self.extract(self.betas, t, cur_y).sqrt()
+		sample = mean + sigma_t * z * 0.00001
+		return (sample)
+
+
+
+	def p_sample_loop(self, x, mask, shape):
+		self.model.eval()
+		prediction_total = torch.Tensor().cuda()
+		for _ in range(20):
+			cur_y = torch.randn(shape).to(x.device)
+			for i in reversed(range(self.n_steps)):
+				cur_y = self.p_sample(x, mask, cur_y, i)
+			prediction_total = torch.cat((prediction_total, cur_y.unsqueeze(1)), dim=1)
+		return prediction_total
+	
+	def p_sample_loop_mean(self, x, mask, loc):
+		prediction_total = torch.Tensor().cuda()
+		for loc_i in range(1):
+			cur_y = loc
+			for i in reversed(range(NUM_Tau)):
+				cur_y = self.p_sample(x, mask, cur_y, i)
+			prediction_total = torch.cat((prediction_total, cur_y.unsqueeze(1)), dim=1)
+		return prediction_total
+
+	def p_sample_loop_accelerate(self, x, mask, loc):
+		'''
+		Batch operation to accelerate the denoising process.
+
+		x: [11, 10, 6]
+		mask: [11, 11]
+		cur_y: [11, 10, 20, 2]
+		'''
+		prediction_total = torch.Tensor().cuda()
+		cur_y = loc[:, :10]
+		for i in reversed(range(NUM_Tau)):
+			cur_y = self.p_sample_accelerate(x, mask, cur_y, i)
+		cur_y_ = loc[:, 10:]
+		for i in reversed(range(NUM_Tau)):
+			cur_y_ = self.p_sample_accelerate(x, mask, cur_y_, i)
+		# shape: B=b*n, K=10, T, 2
+		prediction_total = torch.cat((cur_y_, cur_y), dim=1)
+		return prediction_total
+
+
+
+	def fit(self):
+		# Training loop
+		for epoch in range(0, 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)
+			
+			if (epoch + 1) % self.cfg.test_interval == 0:
+				performance, samples = self._test_single_epoch()
+				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), self.log)
+				cp_path = self.cfg.model_path % (epoch + 1)
+				model_cp = {'model_initializer_dict': self.model_initializer.state_dict()}
+				torch.save(model_cp, cp_path)
+			self.scheduler_model.step()
+
+
+	def data_preprocess(self, data):
+		"""
+			pre_motion_3D: torch.Size([32, 11, 10, 2]), [batch_size, num_agent, past_frame, dimension]
+			fut_motion_3D: torch.Size([32, 11, 20, 2])
+			fut_motion_mask: torch.Size([32, 11, 20])
+			pre_motion_mask: torch.Size([32, 11, 10])
+			traj_scale: 1
+			pred_mask: None
+			seq: nba
+		"""
+		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:]
+		# augment input: absolute position, relative position, velocity
+		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()
+		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]
+			
+			generated_y = self.p_sample_loop_accelerate(past_traj, traj_mask, loc)
+			
+			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()
+			torch.nn.utils.clip_grad_norm_(self.model_initializer.parameters(), 1.)
+			self.opt.step()
+			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)
+		count = 0
+		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]
+			
+				pred_traj = self.p_sample_loop_accelerate(past_traj, traj_mask, loc)
+
+				fut_traj = fut_traj.unsqueeze(1).repeat(1, 20, 1, 1)
+				# b*n, K, T, 2
+				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]
+				count += 1
+				# if count==100:
+				# 	break
+		return performance, samples
+
+
+	def save_data(self):
+		'''
+		Save the visualization data.
+		'''
+		model_path = './results/checkpoints/led_vis.p'
+		model_dict = torch.load(model_path, map_location=torch.device('cpu'))['model_initializer_dict']
+		self.model_initializer.load_state_dict(model_dict)
+		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)
+		root_path = './visualization/data/'
+				
+		with torch.no_grad():
+			for data in self.test_loader:
+				_, traj_mask, past_traj, _ = self.data_preprocess(data)
+
+				sample_prediction, mean_estimation, variance_estimation = self.model_initializer(past_traj, traj_mask)
+				torch.save(sample_prediction, root_path+'p_var.pt')
+				torch.save(mean_estimation, root_path+'p_mean.pt')
+				torch.save(variance_estimation, root_path+'p_sigma.pt')
+
+				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]
+
+				pred_traj = self.p_sample_loop_accelerate(past_traj, traj_mask, loc)
+				pred_mean = self.p_sample_loop_mean(past_traj, traj_mask, mean_estimation)
+
+				torch.save(data['pre_motion_3D'], root_path+'past.pt')
+				torch.save(data['fut_motion_3D'], root_path+'future.pt')
+				torch.save(pred_traj, root_path+'prediction.pt')
+				torch.save(pred_mean, root_path+'p_mean_denoise.pt')
+
+				raise ValueError
+
+
+
+	def test_single_model(self):
+		model_path = './results/checkpoints/led_new.p'
+		model_dict = torch.load(model_path, map_location=torch.device('cpu'))['model_initializer_dict']
+		self.model_initializer.load_state_dict(model_dict)
+		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)
+		count = 0
+		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]
+			
+				pred_traj = self.p_sample_loop_accelerate(past_traj, traj_mask, loc)
+
+				fut_traj = fut_traj.unsqueeze(1).repeat(1, 20, 1, 1)
+				# b*n, K, T, 2
+				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]
+				count += 1
+					# if count==2:
+					# 	break
+		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)
+		
+	

LED/LED/utils/config.py

@@ -0,0 +1,60 @@
+import yaml
+import os
+import os.path as osp
+import glob
+import numpy as np
+from easydict import EasyDict
+from utils.utils import recreate_dirs
+
+
+class Config:
+
+    def __init__(self, cfg_id, info):
+        self.id = cfg_id
+        cfg_path = 'cfg/**/%s.yml' % cfg_id
+        files = glob.glob(cfg_path, recursive=True)
+        assert (len(files) == 1), 'YAML file [{}] does not exist!'.format(cfg_id)
+        self.yml_dict = EasyDict(yaml.safe_load(open(files[0], 'r')))
+
+        self.results_root_dir = os.path.expanduser(self.yml_dict['results_root_dir'])
+        # results dirs
+
+        self.cfg_dir = '%s/%s/%s' % (self.results_root_dir, cfg_id, info)
+        self.model_dir = '%s/models' % self.cfg_dir
+        self.log_dir = '%s/log' % self.cfg_dir
+        self.model_path = os.path.join(self.model_dir, 'model_%04d.p')
+        os.makedirs(self.model_dir, exist_ok=True)
+        os.makedirs(self.log_dir, exist_ok=True)
+
+    def get_last_epoch(self):
+        model_files = glob.glob(os.path.join(self.model_dir, 'model_*.p'))
+        if len(model_files) == 0:
+            return None
+        else:
+            model_file = osp.basename(model_files[0])
+            epoch = int(osp.splitext(model_file)[0].split('model_')[-1])
+            return epoch            
+
+    def __getattribute__(self, name):
+        yml_dict = super().__getattribute__('yml_dict')
+        if name in yml_dict:
+            return yml_dict[name]
+        else:
+            return super().__getattribute__(name)
+
+    def __setattr__(self, name, value):
+        try:
+            yml_dict = super().__getattribute__('yml_dict')
+        except AttributeError:
+            return super().__setattr__(name, value)
+        if name in yml_dict:
+            yml_dict[name] = value
+        else:
+            return super().__setattr__(name, value)
+
+    def get(self, name, default=None):
+        if hasattr(self, name):
+            return getattr(self, name)
+        else:
+            return default
+            

LED/LED/utils/utils.py

@@ -0,0 +1,336 @@
+"""
+Code borrowed from Xinshuo_PyToolbox: https://github.com/xinshuoweng/Xinshuo_PyToolbox
+"""
+
+import os
+import shutil
+import torch
+import numpy as np
+import random
+import time
+import copy
+import glob, glob2
+from torch import nn
+
+class AverageMeter(object):     
+  """Computes and stores the average and current value"""    
+  def __init__(self):   
+    self.reset()
+  
+  def reset(self):
+    self.val = 0
+    self.avg = 0
+    self.sum = 0
+    self.count = 0    
+    self.list = list()
+  
+  def update(self, val, n=1): 
+    self.val = val    
+    self.sum += val * n     
+    self.count += n
+    self.avg = self.sum / self.count  
+    self.list.append(val)
+
+
+def isnparray(nparray_test):
+	return isinstance(nparray_test, np.ndarray)
+
+
+def isinteger(integer_test):
+	if isnparray(integer_test): return False
+	try: return isinstance(integer_test, int) or int(integer_test) == integer_test
+	except ValueError: return False
+	except TypeError: return False
+
+
+def isfloat(float_test):
+	return isinstance(float_test, float)
+
+
+def isscalar(scalar_test):
+	try: return isinteger(scalar_test) or isfloat(scalar_test)
+	except TypeError: return False
+
+
+def islogical(logical_test):
+	return isinstance(logical_test, bool)
+
+    
+def isstring(string_test):
+    return isinstance(string_test, str)
+
+
+def islist(list_test):
+	return isinstance(list_test, list)
+
+
+def convert_secs2time(seconds):
+    '''
+    format second to human readable way
+    '''
+    assert isscalar(seconds), 'input should be a scalar to represent number of seconds'
+    m, s = divmod(int(seconds), 60)
+    h, m = divmod(m, 60)
+    return '[%d:%02d:%02d]' % (h, m, s)
+
+
+def get_timestring():
+    return time.strftime('%Y%m%d_%Hh%Mm%Ss')
+
+
+def recreate_dirs(*dirs):
+    for d in dirs:
+        if os.path.exists(d):
+            shutil.rmtree(d)
+        os.makedirs(d)
+
+
+def is_path_valid(pathname):
+	try:  
+		if not isstring(pathname) or not pathname: return False
+	except TypeError: return False
+	else: return True
+
+
+def is_path_creatable(pathname):
+	'''
+	if any previous level of parent folder exists, returns true
+	'''
+	if not is_path_valid(pathname): return False
+	pathname = os.path.normpath(pathname)
+	pathname = os.path.dirname(os.path.abspath(pathname))
+
+	# recursively to find the previous level of parent folder existing
+	while not is_path_exists(pathname):     
+		pathname_new = os.path.dirname(os.path.abspath(pathname))
+		if pathname_new == pathname: return False
+		pathname = pathname_new
+	return os.access(pathname, os.W_OK)
+
+
+def is_path_exists(pathname):
+	try: return is_path_valid(pathname) and os.path.exists(pathname)
+	except OSError: return False
+
+
+def is_path_exists_or_creatable(pathname):
+	try: return is_path_exists(pathname) or is_path_creatable(pathname)
+	except OSError: return False
+
+
+def isfile(pathname):
+	if is_path_valid(pathname):
+		pathname = os.path.normpath(pathname)
+		name = os.path.splitext(os.path.basename(pathname))[0]
+		ext = os.path.splitext(pathname)[1]
+		return len(name) > 0 and len(ext) > 0
+	else: return False
+
+
+def isfolder(pathname):
+	'''
+	if '.' exists in the subfolder, the function still justifies it as a folder. e.g., /mnt/dome/adhoc_0.5x/abc is a folder
+	if '.' exists after all slashes, the function will not justify is as a folder. e.g., /mnt/dome/adhoc_0.5x is NOT a folder
+	'''
+	if is_path_valid(pathname):
+		pathname = os.path.normpath(pathname)
+		if pathname == './': return True
+		name = os.path.splitext(os.path.basename(pathname))[0]
+		ext = os.path.splitext(pathname)[1]
+		return len(name) > 0 and len(ext) == 0
+	else: return False
+
+
+def mkdir_if_missing(input_path):
+    folder = input_path if isfolder(input_path) else os.path.dirname(input_path)
+    os.makedirs(folder, exist_ok=True)
+
+
+def safe_list(input_data, warning=True, debug=True):
+	'''
+	copy a list to the buffer for use
+	parameters:
+		input_data:		a list
+	outputs:
+		safe_data:		a copy of input data
+	'''
+	if debug: assert islist(input_data), 'the input data is not a list'
+	safe_data = copy.copy(input_data)
+	return safe_data
+
+
+def safe_path(input_path, warning=True, debug=True):
+    '''
+    convert path to a valid OS format, e.g., empty string '' to '.', remove redundant '/' at the end from 'aa/' to 'aa'
+    parameters:
+    	input_path:		a string
+    outputs:
+    	safe_data:		a valid path in OS format
+    '''
+    if debug: assert isstring(input_path), 'path is not a string: %s' % input_path
+    safe_data = copy.copy(input_path)
+    safe_data = os.path.normpath(safe_data)
+    return safe_data
+
+
+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)
+
+
+def initialize_weights(modules):
+    for m in modules:
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            if m.bias is not None: nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.BatchNorm2d):
+            nn.init.constant_(m.weight, 1)
+            if m.bias is not None: nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, 0, 0.01)
+            if m.bias is not None: nn.init.constant_(m.bias, 0)
+
+
+def print_log(print_str, log, same_line=False, display=True):
+	'''
+	print a string to a log file
+
+	parameters:
+		print_str:          a string to print
+		log:                a opened file to save the log
+		same_line:          True if we want to print the string without a new next line
+		display:            False if we want to disable to print the string onto the terminal
+	'''
+	if display:
+		if same_line: print('{}'.format(print_str), end='')
+		else: print('{}'.format(print_str))
+
+	if same_line: log.write('{}'.format(print_str))
+	else: log.write('{}\n'.format(print_str))
+	log.flush()
+
+
+def find_unique_common_from_lists(input_list1, input_list2, warning=True, debug=True):
+	'''
+	find common items from 2 lists, the returned elements are unique. repetitive items will be ignored
+	if the common items in two elements are not in the same order, the outputs follows the order in the first list
+
+	parameters:
+		input_list1, input_list2:		two input lists
+
+	outputs:
+		list_common:	a list of elements existing both in list_src1 and list_src2	
+		index_list1:	a list of index that list 1 has common items
+		index_list2:	a list of index that list 2 has common items
+	'''
+	input_list1 = safe_list(input_list1, warning=warning, debug=debug)
+	input_list2 = safe_list(input_list2, warning=warning, debug=debug)
+
+	common_list = list(set(input_list1).intersection(input_list2))
+	
+	# find index
+	index_list1 = []
+	for index in range(len(input_list1)):
+		item = input_list1[index]
+		if item in common_list:
+			index_list1.append(index)
+
+	index_list2 = []
+	for index in range(len(input_list2)):
+		item = input_list2[index]
+		if item in common_list:
+			index_list2.append(index)
+
+	return common_list, index_list1, index_list2
+
+
+def load_txt_file(file_path, debug=True):
+    '''
+    load data or string from text file
+    '''
+    file_path = safe_path(file_path)
+    if debug: assert is_path_exists(file_path), 'text file is not existing at path: %s!' % file_path
+    with open(file_path, 'r') as file: data = file.read().splitlines()
+    num_lines = len(data)
+    file.close()
+    return data, num_lines
+
+
+def load_list_from_folder(folder_path, ext_filter=None, depth=1, recursive=False, sort=True, save_path=None, debug=True):
+    '''
+    load a list of files or folders from a system path
+
+    parameters:
+        folder_path:    root to search 
+        ext_filter:     a string to represent the extension of files interested
+        depth:          maximum depth of folder to search, when it's None, all levels of folders will be searched
+        recursive:      False: only return current level
+                        True: return all levels till to the input depth
+
+    outputs:
+        fulllist:       a list of elements
+        num_elem:       number of the elements
+    '''
+    folder_path = safe_path(folder_path)
+    if debug: assert isfolder(folder_path), 'input folder path is not correct: %s' % folder_path
+    if not is_path_exists(folder_path): 
+        print('the input folder does not exist\n')
+        return [], 0
+    if debug:
+        assert islogical(recursive), 'recursive should be a logical variable: {}'.format(recursive)
+        assert depth is None or (isinteger(depth) and depth >= 1), 'input depth is not correct {}'.format(depth)
+        assert ext_filter is None or (islist(ext_filter) and all(isstring(ext_tmp) for ext_tmp in ext_filter)) or isstring(ext_filter), 'extension filter is not correct'
+    if isstring(ext_filter): ext_filter = [ext_filter]                               # convert to a list
+    # zxc
+
+    fulllist = list()
+    if depth is None:        # find all files recursively
+        recursive = True
+        wildcard_prefix = '**'
+        if ext_filter is not None:
+            for ext_tmp in ext_filter:
+                # wildcard = os.path.join(wildcard_prefix, '*' + string2ext_filter(ext_tmp))
+                wildcard = os.path.join(wildcard_prefix, '*' + ext_tmp)
+                curlist = glob2.glob(os.path.join(folder_path, wildcard))
+                if sort: curlist = sorted(curlist)
+                fulllist += curlist
+        else:
+            wildcard = wildcard_prefix
+            curlist = glob2.glob(os.path.join(folder_path, wildcard))
+            if sort: curlist = sorted(curlist)
+            fulllist += curlist
+    else:                    # find files based on depth and recursive flag
+        wildcard_prefix = '*'
+        for index in range(depth-1): wildcard_prefix = os.path.join(wildcard_prefix, '*')
+        if ext_filter is not None:
+            for ext_tmp in ext_filter:
+                # wildcard = wildcard_prefix + string2ext_filter(ext_tmp)
+                wildcard = wildcard_prefix + ext_tmp
+                curlist = glob.glob(os.path.join(folder_path, wildcard))
+                if sort: curlist = sorted(curlist)
+                fulllist += curlist
+            # zxc
+        else:
+            wildcard = wildcard_prefix
+            curlist = glob.glob(os.path.join(folder_path, wildcard))
+            # print(curlist)
+            if sort: curlist = sorted(curlist)
+            fulllist += curlist
+        if recursive and depth > 1:
+            newlist, _ = load_list_from_folder(folder_path=folder_path, ext_filter=ext_filter, depth=depth-1, recursive=True)
+            fulllist += newlist
+
+    fulllist = [os.path.normpath(path_tmp) for path_tmp in fulllist]
+    num_elem = len(fulllist)
+
+    # save list to a path
+    if save_path is not None:
+        save_path = safe_path(save_path)
+        if debug: assert is_path_exists_or_creatable(save_path), 'the file cannot be created'
+        with open(save_path, 'w') as file:
+            for item in fulllist: file.write('%s\n' % item)
+        file.close()
+
+    return fulllist, num_elem

LED/cfg/nba/led_augment.yml

@@ -0,0 +1,39 @@
+# ------------------- General Options -------------------------
+description                  : LED
+results_root_dir             : results
+dataset                      : nba
+
+# ------------------- Dataset -------------------------
+past_frames                  : 10
+future_frames                : 20
+min_past_frames              : 10
+min_future_frames            : 20
+
+motion_dim                   : 2
+forecast_dim                 : 2
+
+traj_mean                    : [14, 7.5]
+traj_scale                   : 5
+
+# ------------------- Model -------------------------
+pretrained_core_denoising_model: './results/checkpoints/base_diffusion_model.p'
+debug                        : False # set to True for early stop in each epoch.
+
+diffusion                    : {
+  steps                      : 100,
+  beta_start                 : 1.e-4,
+  beta_end                   : 5.e-2,
+  beta_schedule              : 'linear'
+}
+
+# ------------------- Training Parameters -------------------------
+lr                           : 1.e-3
+train_batch_size             : 10
+test_batch_size              : 500
+num_epochs                   : 100
+
+lr_scheduler                 : 'step'
+decay_step                   : 8
+decay_gamma                  : 0.5
+test_interval                : 4
+

LED/cfg/nba/led_augment_debug.yml

@@ -0,0 +1,38 @@
+# ------------------- General Options -------------------------
+description                  : LED
+results_root_dir             : results
+dataset                      : nba
+
+# ------------------- Dataset -------------------------
+past_frames                  : 10
+future_frames                : 20
+min_past_frames              : 10
+min_future_frames            : 20
+
+motion_dim                   : 2
+forecast_dim                 : 2
+
+traj_mean                    : [14, 7.5]
+traj_scale                   : 5
+
+# ------------------- Model -------------------------
+pretrained_core_denoising_model: './results/checkpoints/base_diffusion_model.p'
+debug                        : True # set to True for early stop in each epoch.
+
+diffusion                    : {
+  steps                      : 100,
+  beta_start                 : 1.e-4,
+  beta_end                   : 5.e-2,
+  beta_schedule              : 'linear'
+}
+
+# ------------------- Training Parameters -------------------------
+lr                           : 1.e-3
+train_batch_size             : 2
+test_batch_size              : 5
+num_epochs                   : 100
+test_interval                : 2
+lr_scheduler                 : 'step'
+decay_step                   : 8
+decay_gamma                  : 0.5
+

LED/cfg/sdd/sdd.yml

@@ -0,0 +1,34 @@
+description                  : LED — SDD (Stanford Drone Dataset)
+results_root_dir             : results_sdd
+dataset                      : sdd
+past_frames                  : 8
+future_frames                : 12
+motion_dim                   : 2
+forecast_dim                 : 2
+
+traj_mean                    : [0.0, 0.0]
+traj_scale                   : 1.0
+per_scene_norm               : false
+
+pretrained_core_denoising_model : './results_sdd/checkpoints/base_diffusion_model_sdd.p'
+
+diffusion:
+    steps                    : 100
+    beta_start               : 1.e-4
+    beta_end                 : 5.e-2
+    beta_schedule            : 'linear'
+
+train_batch_size             : 32
+test_batch_size              : 256
+num_epochs                   : 100
+decay_step                   : 8
+decay_gamma                  : 0.5
+test_interval                : 4
+
+pretrain:
+    num_epochs               : 150
+    lr                       : 1.e-3
+    train_batch_size         : 64
+    test_batch_size          : 256
+    decay_step               : 30
+    decay_gamma              : 0.5

LED/cfg/sport/football.yml

@@ -0,0 +1,53 @@
+# ------------------- General Options -------------------------
+description                  : LED-Sport-Football
+results_root_dir             : results
+dataset                      : football
+
+# ------------------- Dataset -------------------------
+data_dir                     : '/mnt/jaewoo4tb/srtp/srtp/raw_data/football'
+num_agents                   : 23
+past_frames                  : 10
+future_frames                : 20
+min_past_frames              : 10
+min_future_frames            : 20
+
+motion_dim                   : 2
+forecast_dim                 : 2
+
+# normalization (computed from football/train.npy)
+traj_mean                    : [-4.458, -2.909]
+traj_scale                   : 7.2
+
+# ------------------- Model -------------------------
+pretrained_core_denoising_model: './results/checkpoints/base_diffusion_model_football.p'
+debug                        : False
+
+diffusion                    : {
+  steps                      : 100,
+  beta_start                 : 1.e-4,
+  beta_end                   : 5.e-2,
+  beta_schedule              : 'linear'
+}
+
+# ------------------- Training Parameters -------------------------
+# train_batch_size must keep B*A*K_half ≤ ~1500 to fit in 48GB with the
+# 5-step leapfrog backprop chain; NBA used 10*11*10=1100. With A=23, a
+# batch of 6 gives 6*23*10=1380 — slightly under the NBA footprint.
+lr                           : 1.e-3
+train_batch_size             : 6
+test_batch_size              : 200
+num_epochs                   : 100
+
+lr_scheduler                 : 'step'
+decay_step                   : 8
+decay_gamma                  : 0.5
+test_interval                : 4
+
+# ------------------- Stage-1 Pretraining Parameters -------------------------
+pretrain                     : {
+  num_epochs                 : 100,
+  lr                         : 1.e-3,
+  train_batch_size           : 64,
+  decay_step                 : 10,
+  decay_gamma                : 0.8,
+}

LED/cfg/sport/soccer.yml

@@ -0,0 +1,56 @@
+# ------------------- General Options -------------------------
+description                  : LED-Sport-Soccer
+results_root_dir             : results
+dataset                      : soccer
+
+# ------------------- Dataset -------------------------
+data_dir                     : '/mnt/jaewoo4tb/srtp/srtp/raw_data/soccer'
+num_agents                   : 23
+past_frames                  : 10
+future_frames                : 20
+min_past_frames              : 10
+min_future_frames            : 20
+
+motion_dim                   : 2
+forecast_dim                 : 2
+
+# normalization (computed from soccer/train.npy)
+traj_mean                    : [-0.726, -0.278]
+traj_scale                   : 5.0
+per_scene_norm               : True
+
+# ------------------- Model -------------------------
+pretrained_core_denoising_model: './results/checkpoints/base_diffusion_model_soccer_psnorm.p'
+debug                        : False
+
+diffusion                    : {
+  steps                      : 100,
+  beta_start                 : 1.e-4,
+  beta_end                   : 5.e-2,
+  beta_schedule              : 'linear'
+}
+
+# ------------------- Training Parameters -------------------------
+# train_batch_size must keep B*A*K_half ≤ ~1500 to fit in 48GB with the
+# 5-step leapfrog backprop chain; NBA used 10*11*10=1100. With A=23, a
+# batch of 6 gives 6*23*10=1380 — slightly under the NBA footprint.
+# num_epochs lowered (from 100) because soccer has only 7.2K train scenes
+# and overfit hard in the first run; test_interval=1 to catch the peak.
+lr                           : 1.e-3
+train_batch_size             : 6
+test_batch_size              : 200
+num_epochs                   : 20
+
+lr_scheduler                 : 'step'
+decay_step                   : 4
+decay_gamma                  : 0.5
+test_interval                : 1
+
+# ------------------- Stage-1 Pretraining Parameters -------------------------
+pretrain                     : {
+  num_epochs                 : 150,
+  lr                         : 1.e-3,
+  train_batch_size           : 32,
+  decay_step                 : 10,
+  decay_gamma                : 0.8,
+}

LED/eval_sdd_led_allagents.py

@@ -0,0 +1,193 @@
+"""
+All-agents SDD evaluation for LED checkpoints — matches MoFlow's protocol.
+
+Protocol (mirrors MoFlow trainer's compute_ADE_FDE):
+  For each scene with A agents, run LED sampling → pred [A, K=20, T, 2].
+  Per-horizon buckets: 1.2s (frame 3), 2.4s (6), 3.6s (9), 4.8s (12).
+  ADE_min(H) = mean_{t=1..H} ‖pred - gt‖ → min over K → sum over A agents
+  FDE_min(H) = ‖pred[H-1] - gt[H-1]‖       → min over K → sum over A agents
+  ADE_avg(H) = mean over K instead of min
+  Report in pixels (× 50).
+
+Usage:
+  python eval_sdd_led_allagents.py --exp baseline_v2 --epoch 40
+  python eval_sdd_led_allagents.py --exp graph_sigma_v2 --epoch 28 --use_graph --use_v6_graph
+"""
+import argparse, os, sys, random, torch, numpy as np
+from torch.utils.data import DataLoader
+from data.dataloader_sdd import SDDDataset, sdd_seq_collate
+from models.model_led_initializer import LEDInitializer as InitializationModel
+from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel
+from trainer.train_sdd_led import NUM_Tau
+from utils.config import Config
+
+
+def build_models(cfg, use_graph, use_v6_graph, ckpt_path, device):
+    model = CoreDenoisingModel(past_len=cfg.past_frames).to(device)
+    core_ckpt = torch.load(cfg.pretrained_core_denoising_model, map_location='cpu')
+    model.load_state_dict(core_ckpt['model_dict'])
+    model.eval()
+
+    init = InitializationModel(
+        t_h=cfg.past_frames, d_h=6,
+        t_f=cfg.future_frames, d_f=2, k_pred=20).to(device)
+
+    ckpt = torch.load(ckpt_path, map_location='cpu')
+    init.load_state_dict(ckpt['model_initializer_dict'])
+    init.eval()
+
+    graph = None
+    if use_graph:
+        from models.future_interaction_graph_v6 import FutureInteractionGraphV6Wrapper
+        graph = FutureInteractionGraphV6Wrapper(
+            num_agents=64, future_steps=cfg.future_frames,
+            past_steps=cfg.past_frames, past_channels=6,
+            node_dim=128, top_n=5, num_denoise_steps=NUM_Tau).to(device)
+        sd = {k: v for k, v in ckpt['interaction_graph_dict'].items()
+              if '_single_edge_index' not in k}
+        graph.load_state_dict(sd, strict=False)
+        graph.eval()
+
+    return model, init, graph
+
+
+def make_beta_schedule(n=100, start=1e-4, end=5e-2):
+    return torch.linspace(start, end, n)
+
+
+def extract(a, t, x):
+    out = torch.gather(a, 0, t.to(a.device))
+    return out.reshape(t.shape[0], *([1] * (len(x.shape) - 1)))
+
+
+@torch.no_grad()
+def p_sample_accelerate(x, mask, cur_y, t, model, graph, use_v6_graph, sigma,
+                        betas, alphas, alphas_bar_sqrt, one_minus_alphas_bar_sqrt):
+    t_tensor = torch.tensor([int(t)]).to(x.device)
+    eps_factor = ((1 - extract(alphas, t_tensor, cur_y))
+                  / extract(one_minus_alphas_bar_sqrt, t_tensor, cur_y))
+    beta = extract(betas, t_tensor.repeat(x.shape[0]), cur_y)
+    eps_theta = model.generate_accelerate(cur_y, beta, x, mask)
+
+    if graph is not None:
+        abs_t = extract(alphas_bar_sqrt, t_tensor, cur_y)
+        am1_t = extract(one_minus_alphas_bar_sqrt, t_tensor, cur_y)
+        y0_hat = (cur_y - am1_t * eps_theta) / abs_t
+        delta = graph(y0_hat, x, int(t), sigma=sigma, A_override=x.size(0))
+        eps_theta = eps_theta - (abs_t / am1_t) * delta
+
+    mean = (1 / extract(alphas, t_tensor, cur_y).sqrt()) \
+           * (cur_y - eps_factor * eps_theta)
+    z = torch.randn_like(cur_y)
+    sigma_t = extract(betas, t_tensor, cur_y).sqrt()
+    return mean + sigma_t * z * 0.00001
+
+
+# Horizons in frames (assuming 2.5 fps → 3 = 1.2s, 6 = 2.4s, 9 = 3.6s, 12 = 4.8s)
+HORIZON_FRAMES = {'1.2s': 3, '2.4s': 6, '3.6s': 9, '4.8s': 12}
+
+
+@torch.no_grad()
+def run(args):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    cfg = Config(args.cfg, args.exp)
+
+    test_dset = SDDDataset(obs_len=cfg.past_frames,
+                           pred_len=cfg.future_frames, split='test')
+    loader = DataLoader(test_dset, batch_size=1, shuffle=False,
+                        num_workers=2, collate_fn=sdd_seq_collate)
+
+    ckpt_path = cfg.model_path % args.epoch
+    print(f'Loading checkpoint: {ckpt_path}')
+
+    model, init, graph = build_models(
+        cfg, use_graph=args.use_graph, use_v6_graph=args.use_v6_graph,
+        ckpt_path=ckpt_path, device=device)
+
+    betas = make_beta_schedule().to(device)
+    alphas = 1 - betas
+    alphas_prod = torch.cumprod(alphas, 0)
+    abs_sqrt = torch.sqrt(alphas_prod)
+    one_minus_abs_sqrt = torch.sqrt(1 - alphas_prod)
+
+    traj_mean = torch.FloatTensor(cfg.traj_mean).to(device).view(1, 1, 1, 2)
+    traj_scale = float(cfg.traj_scale)
+
+    np.random.seed(0); random.seed(0)
+    torch.manual_seed(0); torch.cuda.manual_seed_all(0)
+
+    # Accumulators for each horizon: sums of per-agent ADE/FDE (min over K and mean over K).
+    sums = {f'{k}_{h}': 0.0
+            for h in HORIZON_FRAMES for k in ['ADE_min', 'FDE_min', 'ADE_avg', 'FDE_avg']}
+    n_agents = 0
+    T = cfg.future_frames
+
+    for data in loader:
+        pre = data['pre_motion_3D'].to(device)       # [1, A, 8, 2]
+        fut = data['fut_motion_3D'].to(device)       # [1, A, 12, 2]
+        A = pre.size(1)
+        initial_pos = pre[:, :, -1:]
+        past_abs = ((pre - traj_mean) / traj_scale).contiguous().view(-1, cfg.past_frames, 2)
+        past_rel = ((pre - initial_pos) / traj_scale).contiguous().view(-1, cfg.past_frames, 2)
+        past_vel = torch.cat([past_rel[:, 1:] - past_rel[:, :-1],
+                              torch.zeros_like(past_rel[:, -1:])], dim=1)
+        past = torch.cat([past_abs, past_rel, past_vel], dim=-1)
+        fut_rel = ((fut - initial_pos) / traj_scale).contiguous().view(-1, T, 2)
+        mask = torch.ones(A, A).to(device)
+
+        sp, me, ve = init(past, mask)
+        ve = ve.clamp(min=-5, max=5)
+        sp = torch.exp(ve / 2)[..., None, None] * sp \
+             / (sp.std(dim=1).mean(dim=(1, 2))[:, None, None, None] + 1e-6)
+        loc = sp + me[:, None]
+        sigma_in = ve if args.use_v6_graph else None
+
+        # leapfrog: 20 modes = 10+10 two halves, each 5 reverse steps
+        cur_y = loc[:, :10]
+        for i in reversed(range(NUM_Tau)):
+            cur_y = p_sample_accelerate(
+                past, mask, cur_y, i, model, graph, args.use_v6_graph, sigma_in,
+                betas, alphas, abs_sqrt, one_minus_abs_sqrt)
+        cur_y_ = loc[:, 10:]
+        for i in reversed(range(NUM_Tau)):
+            cur_y_ = p_sample_accelerate(
+                past, mask, cur_y_, i, model, graph, args.use_v6_graph, sigma_in,
+                betas, alphas, abs_sqrt, one_minus_abs_sqrt)
+        pred = torch.cat((cur_y_, cur_y), dim=1)   # [A, K=20, T, 2]
+
+        # ALL agents in this scene, per-horizon ADE/FDE matching MoFlow
+        dist = torch.norm(pred - fut_rel.unsqueeze(1), dim=-1) * traj_scale  # [A, K, T]
+        for h_name, h_end in HORIZON_FRAMES.items():
+            # min over K
+            ade_min = dist[..., :h_end].mean(dim=-1).min(dim=-1)[0]          # [A]
+            fde_min = dist[..., h_end - 1].min(dim=-1)[0]                    # [A]
+            # avg over K (mean mode distance; for FVar/AVar later, unused here)
+            ade_avg = dist[..., :h_end].mean(dim=-1).mean(dim=-1)            # [A]
+            fde_avg = dist[..., h_end - 1].mean(dim=-1)                      # [A]
+            sums[f'ADE_min_{h_name}'] += ade_min.sum().item()
+            sums[f'FDE_min_{h_name}'] += fde_min.sum().item()
+            sums[f'ADE_avg_{h_name}'] += ade_avg.sum().item()
+            sums[f'FDE_avg_{h_name}'] += fde_avg.sum().item()
+        n_agents += A
+
+    # Report in pixels: multiply by 50.
+    print(f'\n{args.exp} @ epoch {args.epoch}  (n_agents={n_agents}, all-agents protocol)')
+    print('--- pixels ---')
+    for h in HORIZON_FRAMES:
+        am = sums[f'ADE_min_{h}'] / n_agents * 50.0
+        fm = sums[f'FDE_min_{h}'] / n_agents * 50.0
+        aa = sums[f'ADE_avg_{h}'] / n_agents * 50.0
+        fa = sums[f'FDE_avg_{h}'] / n_agents * 50.0
+        print(f'  ADE_min({h})={am:8.4f}  FDE_min({h})={fm:8.4f}  '
+              f'ADE_avg({h})={aa:8.4f}  FDE_avg({h})={fa:8.4f}')
+
+
+if __name__ == '__main__':
+    p = argparse.ArgumentParser()
+    p.add_argument('--cfg', default='sdd/sdd')
+    p.add_argument('--exp', required=True, help='info tag, e.g. baseline_v2')
+    p.add_argument('--epoch', type=int, required=True)
+    p.add_argument('--use_graph', action='store_true')
+    p.add_argument('--use_v6_graph', action='store_true')
+    args = p.parse_args()
+    run(args)

LED/eval_sdd_led_mid_protocol.py

@@ -0,0 +1,169 @@
+"""
+Re-evaluate an LED SDD checkpoint using the MID evaluation protocol:
+  per-pedestrian full-horizon (12-frame) ADE / final-frame FDE,
+  best_of_20 per pedestrian, averaged across all target pedestrians,
+  ×50 to report in pixels.
+
+Usage:
+  python eval_sdd_led_mid_protocol.py --exp baseline_v2 --epoch 60
+  python eval_sdd_led_mid_protocol.py --exp graph_sigma_v2 --epoch 60 --use_graph
+"""
+import argparse, os, sys, random, torch, numpy as np
+from torch.utils.data import DataLoader
+from data.dataloader_sdd import SDDDataset, sdd_seq_collate
+from models.model_led_initializer import LEDInitializer as InitializationModel
+from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel
+from trainer.train_sdd_led import NUM_Tau
+from utils.config import Config
+
+
+def build_models(cfg, use_graph, use_v6_graph, ckpt_path, device):
+    model = CoreDenoisingModel(past_len=cfg.past_frames).to(device)
+    core_ckpt = torch.load(cfg.pretrained_core_denoising_model, map_location='cpu')
+    model.load_state_dict(core_ckpt['model_dict'])
+    model.eval()
+
+    init = InitializationModel(
+        t_h=cfg.past_frames, d_h=6,
+        t_f=cfg.future_frames, d_f=2, k_pred=20).to(device)
+
+    ckpt = torch.load(ckpt_path, map_location='cpu')
+    init.load_state_dict(ckpt['model_initializer_dict'])
+    init.eval()
+
+    graph = None
+    if use_graph:
+        from models.future_interaction_graph_v6 import FutureInteractionGraphV6Wrapper
+        graph = FutureInteractionGraphV6Wrapper(
+            num_agents=64, future_steps=cfg.future_frames,
+            past_steps=cfg.past_frames, past_channels=6,
+            node_dim=128, top_n=5, num_denoise_steps=NUM_Tau).to(device)
+        sd = {k: v for k, v in ckpt['interaction_graph_dict'].items()
+              if '_single_edge_index' not in k}
+        graph.load_state_dict(sd, strict=False)
+        graph.eval()
+
+    return model, init, graph
+
+
+def make_beta_schedule(n=100, start=1e-4, end=5e-2):
+    return torch.linspace(start, end, n)
+
+
+def extract(a, t, x):
+    out = torch.gather(a, 0, t.to(a.device))
+    return out.reshape(t.shape[0], *([1] * (len(x.shape) - 1)))
+
+
+@torch.no_grad()
+def p_sample_accelerate(x, mask, cur_y, t, model, graph, use_v6_graph, sigma,
+                        betas, alphas, alphas_bar_sqrt, one_minus_alphas_bar_sqrt):
+    t_tensor = torch.tensor([int(t)]).to(x.device)
+    eps_factor = ((1 - extract(alphas, t_tensor, cur_y))
+                  / extract(one_minus_alphas_bar_sqrt, t_tensor, cur_y))
+    beta = extract(betas, t_tensor.repeat(x.shape[0]), cur_y)
+    eps_theta = model.generate_accelerate(cur_y, beta, x, mask)
+
+    if graph is not None:
+        abs_t = extract(alphas_bar_sqrt, t_tensor, cur_y)
+        am1_t = extract(one_minus_alphas_bar_sqrt, t_tensor, cur_y)
+        y0_hat = (cur_y - am1_t * eps_theta) / abs_t
+        delta = graph(y0_hat, x, int(t), sigma=sigma, A_override=x.size(0))
+        eps_theta = eps_theta - (abs_t / am1_t) * delta
+
+    mean = (1 / extract(alphas, t_tensor, cur_y).sqrt()) \
+           * (cur_y - eps_factor * eps_theta)
+    z = torch.randn_like(cur_y)
+    sigma_t = extract(betas, t_tensor, cur_y).sqrt()
+    return mean + sigma_t * z * 0.00001
+
+
+@torch.no_grad()
+def run(args):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    cfg = Config(args.cfg, args.exp)
+
+    test_dset = SDDDataset(obs_len=cfg.past_frames,
+                           pred_len=cfg.future_frames, split='test')
+    loader = DataLoader(test_dset, batch_size=1, shuffle=False,
+                        num_workers=2, collate_fn=sdd_seq_collate)
+
+    exp_dir = cfg.model_dir  # results_sdd/sdd/sdd/<info>/models
+    ckpt_path = cfg.model_path % args.epoch
+    print(f'Loading checkpoint: {ckpt_path}')
+
+    model, init, graph = build_models(
+        cfg, use_graph=args.use_graph, use_v6_graph=args.use_v6_graph,
+        ckpt_path=ckpt_path, device=device)
+
+    betas = make_beta_schedule().to(device)
+    alphas = 1 - betas
+    alphas_prod = torch.cumprod(alphas, 0)
+    abs_sqrt = torch.sqrt(alphas_prod)
+    one_minus_abs_sqrt = torch.sqrt(1 - alphas_prod)
+
+    traj_mean = torch.FloatTensor(cfg.traj_mean).to(device).view(1, 1, 1, 2)
+    traj_scale = float(cfg.traj_scale)
+
+    np.random.seed(0); random.seed(0)
+    torch.manual_seed(0); torch.cuda.manual_seed_all(0)
+
+    total_ade, total_fde, n = 0.0, 0.0, 0
+    T = cfg.future_frames
+    for data in loader:
+        pre = data['pre_motion_3D'].to(device)       # [1, A, 8, 2]
+        fut = data['fut_motion_3D'].to(device)
+        A = pre.size(1)
+        initial_pos = pre[:, :, -1:]
+        past_abs = ((pre - traj_mean) / traj_scale).contiguous().view(-1, cfg.past_frames, 2)
+        past_rel = ((pre - initial_pos) / traj_scale).contiguous().view(-1, cfg.past_frames, 2)
+        past_vel = torch.cat([past_rel[:, 1:] - past_rel[:, :-1],
+                              torch.zeros_like(past_rel[:, -1:])], dim=1)
+        past = torch.cat([past_abs, past_rel, past_vel], dim=-1)
+        fut_rel = ((fut - initial_pos) / traj_scale).contiguous().view(-1, T, 2)
+        mask = torch.ones(A, A).to(device)
+
+        sp, me, ve = init(past, mask)
+        ve = ve.clamp(min=-5, max=5)
+        sp = torch.exp(ve / 2)[..., None, None] * sp \
+             / (sp.std(dim=1).mean(dim=(1, 2))[:, None, None, None] + 1e-6)
+        loc = sp + me[:, None]
+        sigma_in = ve if args.use_v6_graph else None
+
+        # leapfrog: 20 modes = 10+10 two halves, each 5 reverse steps
+        cur_y = loc[:, :10]
+        for i in reversed(range(NUM_Tau)):
+            cur_y = p_sample_accelerate(
+                past, mask, cur_y, i, model, graph, args.use_v6_graph, sigma_in,
+                betas, alphas, abs_sqrt, one_minus_abs_sqrt)
+        cur_y_ = loc[:, 10:]
+        for i in reversed(range(NUM_Tau)):
+            cur_y_ = p_sample_accelerate(
+                past, mask, cur_y_, i, model, graph, args.use_v6_graph, sigma_in,
+                betas, alphas, abs_sqrt, one_minus_abs_sqrt)
+        pred = torch.cat((cur_y_, cur_y), dim=1)   # [A, 20, T, 2]
+
+        # target only (index 0) per scene, full-horizon ADE / final FDE, best_of_20
+        pred_0 = pred[0:1]
+        fut_0 = fut_rel[0:1]
+        dist = torch.norm(fut_0.unsqueeze(1) - pred_0, dim=-1) * traj_scale
+        ade = dist.mean(dim=-1).min(dim=-1)[0]
+        fde = dist[:, :, -1].min(dim=-1)[0]
+        total_ade += ade.sum().item()
+        total_fde += fde.sum().item()
+        n += 1
+
+    ade_px = total_ade / n * 50.0
+    fde_px = total_fde / n * 50.0
+    print(f'Epoch {args.epoch} MID-protocol: ADE={ade_px:.4f}  FDE={fde_px:.4f}  (n={n})')
+
+
+if __name__ == '__main__':
+    p = argparse.ArgumentParser()
+    p.add_argument('--cfg', default='sdd/sdd')
+    p.add_argument('--exp', required=True, help='info tag e.g. baseline_v2')
+    p.add_argument('--epoch', type=int, required=True)
+    p.add_argument('--use_graph', action='store_true')
+    p.add_argument('--use_v6_graph', action='store_true')
+    args = p.parse_args()
+    run(args)

LED/main_led_nba.py

@@ -0,0 +1,31 @@
+import argparse
+from trainer import train_led_trajectory_augment_input as led
+
+
+def parse_config():
+	parser = argparse.ArgumentParser()
+	parser.add_argument("--cuda", default=True)
+	parser.add_argument("--learning_rate", type=int, default=0.002)
+	parser.add_argument("--max_epochs", type=int, default=128)
+
+	parser.add_argument('--cfg', default='led_augment')
+	parser.add_argument('--gpu', type=int, default=0, help='Specify which GPU to use.')
+	parser.add_argument('--train', type=int, default=1, help='Whether train or evaluate.')
+	
+	parser.add_argument("--info", type=str, default='test', help='Name of the experiment. '
+															 'It will be used in file creation.')
+	return parser.parse_args()
+
+
+def main(config):
+	t = led.Trainer(config)
+	if config.train==1:
+		t.fit()
+	else:
+		# t.save_data()
+		t.test_single_model()
+
+
+if __name__ == "__main__":
+	config = parse_config()
+	main(config)

LED/main_led_nba_graph.py

@@ -0,0 +1,53 @@
+import argparse
+from trainer import train_led_graph as led
+
+
+def parse_config():
+	parser = argparse.ArgumentParser()
+	parser.add_argument("--cuda", default=True)
+	parser.add_argument("--learning_rate", type=int, default=0.002)
+	parser.add_argument("--max_epochs", type=int, default=128)
+
+	parser.add_argument('--cfg', default='led_augment')
+	parser.add_argument('--gpu', type=int, default=0, help='Specify which GPU to use.')
+	parser.add_argument('--train', type=int, default=1, help='Whether train or evaluate.')
+
+	parser.add_argument("--info", type=str, default='graph', help='Name of the experiment. '
+															 'It will be used in file creation.')
+
+	# Graph variant knobs.
+	parser.add_argument('--top_n', type=int, default=5,
+						help='Number of sparse neighbors per agent (max 10 for NBA with A=11).')
+	parser.add_argument('--residual_on', type=str, default='eps', choices=['eps', 'y0'],
+						help='Where to apply the graph residual: directly on epsilon, or '
+							 'on the implied y_0 estimate (then re-projected to epsilon).')
+	parser.add_argument('--use_sigma', action='store_true',
+						help='If set, pass the initializer variance_estimation to the graph '
+							 'as per-agent uncertainty (modulates node features and edges).')
+	parser.add_argument('--use_v6_graph', action='store_true',
+						help='If set, use MoFlow V6-style RAG-scoring graph (FutureInteractionGraphV6) '
+							 'instead of the default hand-crafted distance-based graph.')
+	parser.add_argument('--uncertainty_weight', type=float, default=1.0,
+						help='Weight for uncertainty NLL loss. Set to 0 for nosigma ablation.')
+	parser.add_argument('--edge_mode', type=str, default='full',
+						choices=['full', 'dist_only', 'relpos_only', 'heading_only', 'vel_only', 'full_relvel'],
+						help='Edge feature mode for RelTrajEncoder ablation.')
+	parser.add_argument('--resume_epoch', type=int, default=0,
+						help='Resume from this checkpoint epoch (0=start fresh).')
+	parser.add_argument('--neighbor_mode', type=str, default='rag',
+						choices=['rag', 'l2', 'semantic'],
+						help='Neighbor selection: rag (semantic+geo), l2 (closest), or semantic (learned only)')
+	return parser.parse_args()
+
+
+def main(config):
+	t = led.Trainer(config)
+	if config.train == 1:
+		t.fit()
+	else:
+		t.test_single_model()
+
+
+if __name__ == "__main__":
+	config = parse_config()
+	main(config)

LED/main_led_nba_grpo.py

@@ -0,0 +1,360 @@
+"""
+GRPO fine-tuning of LED + SRA graph on NBA.
+
+Formulation (single-step bandit on the initializer):
+  LED's 20-mode diversity comes entirely from the initializer; the 5-step
+  leapfrog denoising is near-deterministic (its DDPM noise is x1e-5). So we
+  treat the initializer's mode set `loc` [B*A, K, T, 2] as the ACTION:
+
+      loc      = initializer(past)                 # deterministic mean
+      loc_s    = loc + init_noise * z              # sampled action, z ~ N(0,I)
+      pred     = leapfrog_decode(loc_s)            # fixed decoder (graph+denoiser frozen)
+      reward   = per-agent accuracy (ADE/FDE) + joint (JADE/JFDE)
+
+  This sidesteps the long-chain credit-assignment / high-variance problem that
+  limited the MoFlow (10-step ODE) experiment: the whole trajectory-set is one
+  action, log-prob factorizes over (agent, mode), and GRPO's group-relative
+  advantage is taken over the K modes per agent.
+
+  Only the initializer is trained (graph + core denoiser frozen); KL anchor to a
+  frozen copy of the warm-start initializer. Eval uses the near-deterministic
+  decoder (matches the LED baseline) and reports ADE/FDE/JADE/JFDE.
+"""
+
+import os
+import sys
+import argparse
+import math
+import copy
+
+import numpy as np
+import torch
+
+from trainer.train_led_graph import Trainer as LEDTrainer, NUM_Tau
+
+
+# ---- GRPO reward (inlined from MoFlow grpo/rewards.py to avoid sys.path clash) ----
+def compute_reward_agentwise(pred, gt, init_pos, *, w_ade=1.0, w_fde=1.0,
+                             w_jade=1.0, w_jfde=1.0, w_col=0.0, w_kin=0.0,
+                             d_min=0.4, a_max=1.0, ball_idx=None):
+    B, K, A, T, _ = pred.shape
+    err = (pred - gt.unsqueeze(1)).norm(dim=-1)          # [B,K,A,T]
+    ade = err.mean(dim=-1)                                # [B,K,A]
+    fde = err[..., -1]                                    # [B,K,A]
+    r_marg = -(w_ade * ade + w_fde * fde)
+    jade = ade.mean(dim=2, keepdim=True)                 # [B,K,1]
+    jfde = fde.mean(dim=2, keepdim=True)
+    r_joint = -(w_jade * jade + w_jfde * jfde)
+    reward = r_marg + r_joint
+    info = {'ade': ade.detach(), 'jade': jade.squeeze(-1).detach(),
+            'ade_bestk': ade.min(dim=1).values.mean().detach(),
+            'jade_bestk': jade.squeeze(-1).min(dim=1).values.mean().detach()}
+    return reward, info
+
+
+def group_advantage(reward, eps=1e-4):
+    mean = reward.mean(dim=1, keepdim=True)
+    std = reward.std(dim=1, keepdim=True)
+    return (reward - mean) / (std + eps)
+
+
+def _player_mask(A, ball_idx, device):
+    pm = ~torch.eye(A, dtype=torch.bool, device=device)
+    if ball_idx is not None:
+        pm[ball_idx, :] = False
+        pm[:, ball_idx] = False
+    return pm
+
+
+def compute_reward_collision(pred, gt, init_pos, *, w_ade=0.3, w_col=1.0,
+                             d_min=0.4, ball_idx=10):
+    """Non-differentiable HARD collision-count reward (the objective the
+    supervised min-of-K loss cannot optimize) + a soft ADE term to hold accuracy.
+    pred/gt in RELATIVE metric (court) units; init_pos absolute [B,A,2].
+    Returns reward [B,K,A], info."""
+    B, K, A, T, _ = pred.shape
+    err = (pred - gt.unsqueeze(1)).norm(dim=-1)          # [B,K,A,T]
+    ade = err.mean(dim=-1)                                # [B,K,A]  (soft, accuracy)
+    abs_p = pred + init_pos[:, None, :, None, :]          # absolute positions
+    mind = (abs_p.unsqueeze(3) - abs_p.unsqueeze(2)).norm(dim=-1).min(dim=-1).values  # [B,K,A,A]
+    pm = _player_mask(A, ball_idx, pred.device)
+    hard = ((mind < d_min) & pm).float()                 # HARD indicator (non-diff)
+    coll_count = hard.sum(dim=-1)                         # [B,K,A]  #collisions of agent a
+    reward = -(w_ade * ade) - (w_col * coll_count)
+    info = {'ade_bestk': ade.min(dim=1).values.mean().detach(),
+            'jade_bestk': ade.mean(dim=2).min(dim=1).values.mean().detach(),
+            'coll_count': coll_count.mean().detach(),
+            'coll_rate': (coll_count > 0).float().mean().detach()}
+    return reward, info
+
+
+class LEDGRPOTrainer(LEDTrainer):
+    def __init__(self, config):
+        # graph config for the warm-start checkpoint (edge_relpos, v6, no sigma)
+        config.use_v6_graph = True
+        config.edge_mode = getattr(config, 'edge_mode', 'relpos_only')
+        config.neighbor_mode = getattr(config, 'neighbor_mode', 'rag')
+        config.top_n = getattr(config, 'top_n', 5)
+        config.use_sigma = False
+        config.residual_on = getattr(config, 'residual_on', 'eps')
+        super().__init__(config)
+
+        # ---- warm-start initializer + graph ----
+        ck = torch.load(config.warm_ckpt, map_location='cpu')
+        self.model_initializer.load_state_dict(ck['model_initializer_dict'])
+        self.interaction_graph.load_state_dict(ck['interaction_graph_dict'])
+        print(f'[LED-GRPO] warm-started from {config.warm_ckpt}')
+
+        # freeze graph + core denoiser; train ONLY the initializer
+        for p in self.interaction_graph.parameters():
+            p.requires_grad_(False)
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+        self.interaction_graph.eval()
+        self.model.eval()
+
+        # bigger rollout batch than LED's default (10) for stable GRPO advantages
+        if getattr(config, 'batch', 0):
+            from data.dataloader_nba import NBADataset, seq_collate
+            from torch.utils.data import DataLoader
+            tr = NBADataset(obs_len=self.cfg.past_frames, pred_len=self.cfg.future_frames, training=True)
+            self.train_loader = DataLoader(tr, batch_size=config.batch, shuffle=True,
+                                           num_workers=4, collate_fn=seq_collate, pin_memory=True, drop_last=True)
+
+        # frozen reference initializer (KL anchor)
+        self.ref_initializer = copy.deepcopy(self.model_initializer).cuda().eval()
+        for p in self.ref_initializer.parameters():
+            p.requires_grad_(False)
+
+        # optimizer over the initializer only
+        self.opt = torch.optim.AdamW(self.model_initializer.parameters(), lr=config.grpo_lr)
+
+        # GRPO hyperparams
+        self.G = 20
+        self.init_noise = float(config.init_noise)
+        self.kl_beta = float(config.kl_beta)
+        self.clip_eps = float(config.clip_eps)
+        self.inner_epochs = int(config.inner_epochs)
+        self.grpo_iters = int(config.grpo_iters)
+        self.eval_every = int(config.eval_every)
+        self.logratio_clip = 10.0
+        self.max_eval_batches = int(getattr(config, 'max_eval_batches', 0))
+        self.rw = dict(w_ade=config.w_ade, w_fde=config.w_fde,
+                       w_jade=config.w_jade, w_jfde=config.w_jfde,
+                       w_col=0.0, w_kin=0.0, ball_idx=None)
+        self.reward_mode = getattr(config, 'reward_mode', 'accuracy')
+        self.rw_coll = dict(w_ade=getattr(config, 'w_ade_soft', 0.3),
+                            w_col=getattr(config, 'w_col', 1.0),
+                            d_min=getattr(config, 'd_min', 0.4), ball_idx=10)
+        self.d_min_eval = getattr(config, 'd_min', 0.4)
+        self.ade_tol = getattr(config, 'ade_tol', 0.80)
+        self.best_sum = float('inf')
+        self.best_coll = float('inf')
+
+    # ------------------------------------------------------------------
+    def get_loc(self, past_traj, traj_mask):
+        """Initializer -> deterministic mode set loc [B*A, K, T, 2]."""
+        guess_var, guess_mean, guess_scale = self.model_initializer(past_traj, traj_mask)
+        sp = (torch.exp(guess_scale / 2)[..., None, None] * guess_var
+              / guess_var.std(dim=1).mean(dim=(1, 2))[:, None, None, None])
+        return sp + guess_mean[:, None]
+
+    def get_loc_from(self, initializer, past_traj, traj_mask):
+        guess_var, guess_mean, guess_scale = initializer(past_traj, traj_mask)
+        sp = (torch.exp(guess_scale / 2)[..., None, None] * guess_var
+              / guess_var.std(dim=1).mean(dim=(1, 2))[:, None, None, None])
+        return sp + guess_mean[:, None]
+
+    @staticmethod
+    def _logp(action, mean, std):
+        var = std * std
+        lp = -0.5 * (((action - mean) ** 2) / var + math.log(2 * math.pi * var))
+        return lp.sum(dim=(2, 3))           # [B*A, K]  sum over (T, 2)
+
+    def _to_bkat(self, x_ba_k, B, A):
+        """[B*A, K, T, 2] -> [B, K, A, T, 2]"""
+        K, T = x_ba_k.shape[1], x_ba_k.shape[2]
+        return x_ba_k.view(B, A, K, T, 2).permute(0, 2, 1, 3, 4)
+
+    # ------------------------------------------------------------------
+    def train(self):
+        A = 11
+        self.eval_grpo(-1)                     # same-subset baseline (before any update)
+        self.model_initializer.train()
+        dl = self._cycle(self.train_loader)
+        for it in range(self.grpo_iters):
+            data = next(dl)
+            B, traj_mask, past, fut = self.data_preprocess(data)
+
+            # ---- rollout: sample action loc_s, decode, reward ----
+            with torch.no_grad():
+                loc = self.get_loc(past, traj_mask)                  # [B*A,K,T,2]
+                z = torch.randn_like(loc)
+                loc_s = loc + self.init_noise * z
+                logp_old = self._logp(loc_s, loc, self.init_noise)   # [B*A,K]
+                pred = self.p_sample_loop_accelerate(past, traj_mask, loc_s)  # decode
+                loc_ref = self.get_loc_from(self.ref_initializer, past, traj_mask)
+
+            # reward in metric units ([B,K,A,T,2], scaled by traj_scale)
+            pred_m = self._to_bkat(pred, B, A) * self.traj_scale
+            gt_m = fut.view(B, A, fut.shape[1], 2) * self.traj_scale
+            if self.reward_mode == 'collision':
+                init_pos = data['pre_motion_3D'].cuda()[:, :, -1, :]     # [B,A,2] absolute
+                reward, info = compute_reward_collision(pred_m, gt_m, init_pos, **self.rw_coll)
+            else:
+                init_pos = torch.zeros(B, A, 2, device=pred.device)
+                reward, info = compute_reward_agentwise(pred_m, gt_m, init_pos, **self.rw)
+            # advantage per (agent): reshape reward [B,K,A] -> per-agent group over K
+            adv = group_advantage(reward)                             # [B,K,A]
+            # map advantage back to [B*A, K] to match logp layout
+            adv_bak = adv.permute(0, 2, 1).reshape(B * A, self.G)     # [B*A,K]
+
+            logp_old_flat = logp_old
+            loc_s_c = loc_s
+
+            # ---- PPO update (initializer only) ----
+            stats = {}
+            for _ in range(self.inner_epochs):
+                self.opt.zero_grad()
+                loc_new = self.get_loc(past, traj_mask)               # grad
+                logp_new = self._logp(loc_s_c, loc_new, self.init_noise)   # [B*A,K]
+                logratio = (logp_new - logp_old_flat).clamp(-self.logratio_clip, self.logratio_clip)
+                ratio = logratio.exp()
+                unclipped = ratio * adv_bak
+                clipped = ratio.clamp(1 - self.clip_eps, 1 + self.clip_eps) * adv_bak
+                pg = -torch.min(unclipped, clipped).mean()
+                kl = (0.5 * ((loc_new - loc_ref) ** 2) / (self.init_noise ** 2)).sum(dim=(2, 3)).mean()
+                loss = pg + self.kl_beta * kl
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(self.model_initializer.parameters(), 1.0)
+                self.opt.step()
+                stats = dict(pg=pg.item(), kl=kl.item(), ratio=ratio.mean().item(),
+                             clipfrac=((ratio - 1).abs() > self.clip_eps).float().mean().item())
+
+            if it % 10 == 0:
+                extra = (f'collrate={info["coll_rate"].item():.3f} cnt={info["coll_count"].item():.3f} '
+                         if 'coll_rate' in info else f'JADE*={info["jade_bestk"].item():.4f} ')
+                print(f'[LED-GRPO {it}/{self.grpo_iters}] R={reward.mean().item():.4f} '
+                      f'ADE*={info["ade_bestk"].item():.4f} {extra}'
+                      f'| pg={stats["pg"]:.4f} kl={stats["kl"]:.5f} '
+                      f'ratio={stats["ratio"]:.3f} clipfrac={stats["clipfrac"]:.3f}', flush=True)
+
+            if (it + 1) % self.eval_every == 0:
+                self.eval_grpo(it)
+                self.model_initializer.train()
+
+    # ------------------------------------------------------------------
+    @torch.no_grad()
+    def eval_grpo(self, it):
+        self.model_initializer.eval()
+        A = 11
+        perf = {'ADE': [0.]*4, 'FDE': [0.]*4, 'JADE': [0.]*4, 'JFDE': [0.]*4}
+        coll_thr = (0.2, 0.3, 0.4)
+        collP = {th: 0. for th in coll_thr}; collG = {th: 0. for th in coll_thr}
+        nP, nG = 0, 0
+        n_ag, n_sc = 0, 0
+        for bi, data in enumerate(self.test_loader):
+            if self.max_eval_batches and bi >= self.max_eval_batches:
+                break
+            B, traj_mask, past, fut = self.data_preprocess(data)
+            loc = self.get_loc(past, traj_mask)                       # deterministic
+            pred = self.p_sample_loop_accelerate(past, traj_mask, loc)
+            # --- collision: absolute positions, player-pairs, ball(10) excluded ---
+            ipos = data['pre_motion_3D'].cuda()[:, :, -1, :]          # [B,A,2]
+            Tf = fut.shape[1]
+            absP = self._to_bkat(pred, B, A) * self.traj_scale + ipos[:, None, :, None, :]   # [B,K,A,T,2]
+            absG = (fut.view(B, A, Tf, 2) * self.traj_scale + ipos[:, :, None, :]).unsqueeze(1)  # [B,1,A,T,2]
+            pm = _player_mask(A, 10, absP.device)
+            cpP = ((absP.unsqueeze(3) - absP.unsqueeze(2)).norm(dim=-1).min(dim=-1).values
+                   .masked_fill(~pm, 1e9).reshape(B, self.G, -1).min(-1).values)   # [B,K]
+            cpG = ((absG.unsqueeze(3) - absG.unsqueeze(2)).norm(dim=-1).min(dim=-1).values
+                   .masked_fill(~pm, 1e9).reshape(B, 1, -1).min(-1).values)        # [B,1]
+            for th in coll_thr:
+                collP[th] += (cpP < th).float().sum().item()
+                collG[th] += (cpG < th).float().sum().item()
+            nP += B * self.G; nG += B
+            fut_r = fut.unsqueeze(1).repeat(1, self.G, 1, 1)          # [B*A,K,T,2]
+            d = (fut_r - pred).norm(dim=-1) * self.traj_scale         # [B*A,K,T]
+            dB = d.view(B, A, self.G, d.shape[-1])                    # [B,A,K,T]
+            for ti in range(1, 5):
+                e = 5 * ti
+                # marginal: per-agent min over K
+                ade = d[..., :e].mean(-1).min(dim=1)[0].sum()
+                fde = d[..., e-1].min(dim=1)[0].sum()
+                # joint: per-scene, mean over agents then min over K
+                jade = dB[..., :e].mean(-1).mean(dim=1).min(dim=1)[0].sum()
+                jfde = dB[..., e-1].mean(dim=1).min(dim=1)[0].sum()
+                perf['ADE'][ti-1] += ade.item(); perf['FDE'][ti-1] += fde.item()
+                perf['JADE'][ti-1] += jade.item(); perf['JFDE'][ti-1] += jfde.item()
+            n_ag += B * A; n_sc += B
+        ade4 = perf['ADE'][3]/n_ag; fde4 = perf['FDE'][3]/n_ag
+        jade4 = perf['JADE'][3]/n_sc; jfde4 = perf['JFDE'][3]/n_sc
+        s = ade4 + fde4 + jade4 + jfde4
+        cstr = ' '.join(f'@{th}:{collP[th]/nP*100:.1f}%(GT{collG[th]/nG*100:.1f})' for th in coll_thr)
+        print(f'[LED-GRPO eval @ {it}] ADE={ade4:.4f} FDE={fde4:.4f} '
+              f'JADE={jade4:.4f} JFDE={jfde4:.4f} | coll[pred(GT)]: {cstr}', flush=True)
+        # checkpoint: collision mode -> best collision@d_min with ADE guard; else -> best sum
+        if self.reward_mode == 'collision':
+            c = collP[self.d_min_eval] / nP if self.d_min_eval in collP else collP[0.4] / nP
+            if ade4 <= getattr(self, 'ade_tol', 0.80) and c < self.best_coll:
+                self.best_coll = c
+                torch.save({'model_initializer_dict': self.model_initializer.state_dict(),
+                            'interaction_graph_dict': self.interaction_graph.state_dict()},
+                           os.path.join(self.cfg.log_dir, 'grpo_best.p'))
+                print(f'  new best coll@{self.d_min_eval}={c*100:.2f}% at ADE={ade4:.4f} -> grpo_best.p', flush=True)
+        elif s < self.best_sum:
+            self.best_sum = s
+            torch.save({'model_initializer_dict': self.model_initializer.state_dict(),
+                        'interaction_graph_dict': self.interaction_graph.state_dict()},
+                       os.path.join(self.cfg.log_dir, 'grpo_best.p'))
+            print(f'  new best sum={s:.4f} -> grpo_best.p', flush=True)
+
+    @staticmethod
+    def _cycle(dl):
+        while True:
+            for d in dl:
+                yield d
+
+
+def parse_config():
+    p = argparse.ArgumentParser()
+    p.add_argument('--cfg', default='led_augment')
+    p.add_argument('--info', default='grpo', type=str)
+    p.add_argument('--gpu', type=int, default=0)
+    p.add_argument('--cuda', default=True)
+    p.add_argument('--learning_rate', type=float, default=0.002)  # unused (grpo_lr used)
+    p.add_argument('--warm_ckpt', type=str,
+                   default='./results/led_augment/graph_v6_edge_relpos/models/model_0036.p')
+    p.add_argument('--edge_mode', default='relpos_only', type=str)
+    # GRPO
+    p.add_argument('--batch', type=int, default=64)
+    p.add_argument('--grpo_lr', type=float, default=1e-4)
+    p.add_argument('--init_noise', type=float, default=0.1)
+    p.add_argument('--kl_beta', type=float, default=0.0)
+    p.add_argument('--clip_eps', type=float, default=0.2)
+    p.add_argument('--inner_epochs', type=int, default=2)
+    p.add_argument('--grpo_iters', type=int, default=1000)
+    p.add_argument('--eval_every', type=int, default=50)
+    p.add_argument('--max_eval_batches', type=int, default=5)
+    p.add_argument('--w_ade', type=float, default=1.0)
+    p.add_argument('--w_fde', type=float, default=1.0)
+    p.add_argument('--w_jade', type=float, default=1.0)
+    p.add_argument('--w_jfde', type=float, default=1.0)
+    # collision (non-differentiable) reward
+    p.add_argument('--reward_mode', default='accuracy', choices=['accuracy', 'collision'])
+    p.add_argument('--w_ade_soft', type=float, default=0.3, help='soft ADE weight (hold accuracy)')
+    p.add_argument('--w_col', type=float, default=1.0, help='hard collision-count weight')
+    p.add_argument('--d_min', type=float, default=0.4)
+    p.add_argument('--ade_tol', type=float, default=0.80)
+    return p.parse_args()
+
+
+def main():
+    cfg = parse_config()
+    torch.cuda.set_device(cfg.gpu)
+    t = LEDGRPOTrainer(cfg)
+    t.train()
+
+
+if __name__ == '__main__':
+    main()

LED/main_sdd_led.py

@@ -0,0 +1,21 @@
+"""Stage-2: train LED initializer on SDD (baseline or +graph)."""
+import argparse
+from trainer.train_sdd_led import Trainer
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument('--cuda', default=True)
+    p.add_argument('--cfg', default='sdd/sdd')
+    p.add_argument('--gpu', type=int, default=0)
+    p.add_argument('--info', type=str, default='baseline')
+    p.add_argument('--learning_rate', type=float, default=0.002)
+    p.add_argument('--train', type=int, default=1)
+    p.add_argument('--grad_accum', type=int, default=16)
+    p.add_argument('--use_graph', action='store_true')
+    p.add_argument('--use_v6_graph', action='store_true')
+    p.add_argument('--residual_on', type=str, default='y0')
+    config = p.parse_args()
+    Trainer(config).fit()
+
+if __name__ == '__main__':
+    main()

LED/main_sdd_pretrain.py

@@ -0,0 +1,17 @@
+"""Stage-1: pretrain LED denoiser on SDD."""
+import argparse
+from trainer.train_sdd_pretrain import Trainer
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument('--cuda', default=True)
+    p.add_argument('--cfg', default='sdd/sdd')
+    p.add_argument('--gpu', type=int, default=0)
+    p.add_argument('--info', type=str, default='pretrain')
+    p.add_argument('--grad_accum', type=int, default=32)
+    config = p.parse_args()
+    config.train = 1
+    Trainer(config).fit()
+
+if __name__ == '__main__':
+    main()

LED/main_sport_led.py

@@ -0,0 +1,33 @@
+import argparse
+from trainer import train_sport_led as sl
+
+
+def parse_config():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--cuda', default=True)
+    parser.add_argument('--learning_rate', type=float, default=2e-3)
+    parser.add_argument('--cfg', default='soccer', help='Sport cfg name (soccer / football)')
+    parser.add_argument('--gpu', type=int, default=0)
+    parser.add_argument('--train', type=int, default=1)
+    parser.add_argument('--info', type=str, default='baseline')
+
+    parser.add_argument('--use_graph', action='store_true',
+                        help='If set, use the FutureInteractionGraph residual '
+                             '(default: baseline LED).')
+    parser.add_argument('--residual_on', type=str, default='y0',
+                        choices=['eps', 'y0'],
+                        help='Where the graph residual is applied (only used with --use_graph).')
+    parser.add_argument('--use_v6_graph', action='store_true',
+                        help='If set, use MoFlow V6-style RAG-scoring graph instead of '
+                             'the hand-crafted distance-based one.')
+    return parser.parse_args()
+
+
+def main(config):
+    t = sl.Trainer(config)
+    t.fit()
+
+
+if __name__ == '__main__':
+    config = parse_config()
+    main(config)

LED/main_sport_pretrain.py

@@ -0,0 +1,25 @@
+import argparse
+from trainer import train_sport_pretrain as sp
+
+
+def parse_config():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--cuda', default=True)
+    parser.add_argument('--learning_rate', type=float, default=1e-3,
+                        help='Ignored — actual lr comes from cfg.pretrain.lr.')
+    parser.add_argument('--cfg', default='soccer', help='Sport cfg name (soccer / football)')
+    parser.add_argument('--gpu', type=int, default=0)
+    parser.add_argument('--info', type=str, default='pretrain')
+    # unused but kept for interface symmetry with the LED mains
+    parser.add_argument('--train', type=int, default=1)
+    return parser.parse_args()
+
+
+def main(config):
+    t = sp.Trainer(config)
+    t.fit()
+
+
+if __name__ == '__main__':
+    config = parse_config()
+    main(config)

LED/models/future_interaction_graph.py

@@ -0,0 +1,155 @@
+"""
+Future-trajectory interaction graph for LED.
+
+Idea (ported from MoFlow's FutureInteractionGraphV6):
+  At each leapfrog denoising step, convert the frozen core denoiser's
+  epsilon prediction into an implied y_0 estimate, then run a small
+  agent-to-agent graph on that future geometry and emit a residual
+  correction added back to epsilon. The module is zero-initialized at
+  the output so it begins as a no-op and only nudges the reverse chain
+  as it learns.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class FutureInteractionGraph(nn.Module):
+    def __init__(
+        self,
+        num_agents: int = 11,
+        future_steps: int = 20,
+        past_steps: int = 10,
+        past_channels: int = 6,
+        node_dim: int = 128,
+        top_n: int = 5,
+        num_denoise_steps: int = 5,
+        use_sigma: bool = False,
+    ):
+        super().__init__()
+        self.A = num_agents
+        self.T = future_steps
+        self.D = node_dim
+        self.top_n = min(top_n, num_agents - 1)
+        self.use_sigma = use_sigma
+
+        self.past_proj = nn.Sequential(
+            nn.Linear(past_steps * past_channels, node_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(node_dim, node_dim),
+        )
+        self.y0_proj = nn.Sequential(
+            nn.Linear(future_steps * 2, node_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(node_dim, node_dim),
+        )
+        self.step_emb = nn.Embedding(num_denoise_steps, node_dim)
+
+        if use_sigma:
+            self.sigma_proj = nn.Sequential(
+                nn.Linear(1, node_dim // 4),
+                nn.ReLU(inplace=True),
+                nn.Linear(node_dim // 4, node_dim),
+            )
+
+        # edge feature: base 5 dims + 1 sigma_diff if use_sigma
+        edge_in_dim = 6 if use_sigma else 5
+        self.edge_mlp = nn.Sequential(
+            nn.Linear(edge_in_dim, 32),
+            nn.ReLU(inplace=True),
+            nn.Linear(32, node_dim),
+        )
+
+        self.attn_q = nn.Linear(node_dim, node_dim)
+        self.attn_k = nn.Linear(node_dim * 2, node_dim)
+        self.attn_v = nn.Linear(node_dim * 2, node_dim)
+        self.scale = node_dim ** -0.5
+
+        self.fuse = nn.Sequential(
+            nn.Linear(node_dim * 2, node_dim),
+            nn.ReLU(inplace=True),
+        )
+
+        self.decode = nn.Sequential(
+            nn.Linear(node_dim, node_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(node_dim, future_steps * 2),
+        )
+        # zero-init output so residual starts as a no-op
+        nn.init.zeros_(self.decode[-1].weight)
+        nn.init.zeros_(self.decode[-1].bias)
+
+    def forward(
+        self,
+        y0_hat: torch.Tensor,   # [B*A, K, T, 2]
+        past: torch.Tensor,     # [B*A, T_h, 6]
+        step_idx: int,          # current denoising step in [0, num_denoise_steps)
+        sigma: torch.Tensor = None,  # [B*A, 1] per-agent uncertainty (optional)
+    ) -> torch.Tensor:          # [B*A, K, T, 2]
+        BA, K, T, _ = y0_hat.shape
+        A = self.A
+        assert BA % A == 0, f"expected B*A agents with A={A}, got {BA}"
+        B = BA // A
+        device = y0_hat.device
+
+        past_feat = self.past_proj(past.reshape(BA, -1))                  # [BA, D]
+        past_feat = past_feat.unsqueeze(1).expand(BA, K, self.D)          # [BA, K, D]
+        y0_feat = self.y0_proj(y0_hat.reshape(BA, K, T * 2))              # [BA, K, D]
+        step = self.step_emb(torch.tensor(step_idx, device=device))       # [D]
+        nodes = past_feat + y0_feat + step                                # [BA, K, D]
+
+        if self.use_sigma and sigma is not None:
+            sigma_feat = self.sigma_proj(sigma)                            # [BA, D]
+            nodes = nodes + sigma_feat.unsqueeze(1)                        # broadcast over K
+
+        nodes = nodes.view(B, A, K, self.D).permute(0, 2, 1, 3).contiguous()
+        # nodes: [B, K, A, D]
+
+        pos = y0_hat.view(B, A, K, T, 2).permute(0, 2, 1, 3, 4).contiguous()
+        # pos: [B, K, A, T, 2]
+
+        # pairwise relative future trajectories: row i = receiver, col j = source
+        # rel[:, :, i, j] = pos_j - pos_i (where j is relative to i)
+        rel = pos.unsqueeze(2) - pos.unsqueeze(3)                          # [B,K,A_i,A_j,T,2]
+        mean_rel = rel.mean(dim=-2)                                        # [B,K,A,A,2]
+        std_rel = rel.std(dim=-2)                                          # [B,K,A,A,2]
+        dist = rel.norm(dim=-1)                                            # [B,K,A,A,T]
+        min_dist = dist.min(dim=-1).values.unsqueeze(-1)                   # [B,K,A,A,1]
+
+        if self.use_sigma and sigma is not None:
+            # Per-agent sigma → pairwise sigma difference as edge feature
+            sigma_bka = sigma.view(B, A, 1).permute(0, 2, 1).contiguous()  # [B, 1, A]
+            sigma_i = sigma_bka.unsqueeze(3).expand(-1, K, A, A)            # [B,K,A,A]
+            sigma_j = sigma_bka.unsqueeze(2).expand(-1, K, A, A)            # [B,K,A,A]
+            sigma_diff = (sigma_i - sigma_j).unsqueeze(-1)                  # [B,K,A,A,1]
+            edge_raw = torch.cat([mean_rel, std_rel, min_dist, sigma_diff], dim=-1)  # [B,K,A,A,6]
+        else:
+            edge_raw = torch.cat([mean_rel, std_rel, min_dist], dim=-1)    # [B,K,A,A,5]
+        edge_feat = self.edge_mlp(edge_raw)                                # [B,K,A,A,D]
+
+        # top-N neighbor selection by closest min-distance over future horizon
+        score = -min_dist.squeeze(-1)                                      # [B,K,A,A]
+        diag = torch.eye(A, dtype=torch.bool, device=device)
+        score = score.masked_fill(diag, float('-inf'))
+        _, top_idx = score.topk(self.top_n, dim=-1)                        # [B,K,A,N]
+
+        idx_node = top_idx.unsqueeze(-1).expand(-1, -1, -1, -1, self.D)
+        nodes_j = nodes.unsqueeze(2).expand(-1, -1, A, -1, -1)             # [B,K,A_i,A_j,D]
+        neigh_nodes = torch.gather(nodes_j, 3, idx_node)                   # [B,K,A,N,D]
+
+        idx_edge = top_idx.unsqueeze(-1).expand(-1, -1, -1, -1, self.D)
+        neigh_edges = torch.gather(edge_feat, 3, idx_edge)                 # [B,K,A,N,D]
+
+        q = self.attn_q(nodes).unsqueeze(-2)                               # [B,K,A,1,D]
+        kv_in = torch.cat([neigh_nodes, neigh_edges], dim=-1)              # [B,K,A,N,2D]
+        k = self.attn_k(kv_in)
+        v = self.attn_v(kv_in)
+        attn = (q * k).sum(dim=-1) * self.scale                            # [B,K,A,N]
+        attn = F.softmax(attn, dim=-1)
+        msg = (attn.unsqueeze(-1) * v).sum(dim=-2)                         # [B,K,A,D]
+
+        fused = self.fuse(torch.cat([nodes, msg], dim=-1))                 # [B,K,A,D]
+        residual = self.decode(fused).view(B, K, A, T, 2)
+        residual = residual.permute(0, 2, 1, 3, 4).contiguous().view(BA, K, T, 2)
+        return residual

LED/models/future_interaction_graph_v6.py

@@ -0,0 +1,173 @@
+"""
+V6-style future-trajectory interaction graph for LED.
+
+Re-uses MoFlow's `FutureInteractionGraphV6` verbatim (imported via sys.path)
+so that LED, MID, and MoFlow share the identical neighbor-scoring and
+message-passing component:
+
+  Per-node: q_i = W_q([y0_emb, sigma_mean]), k_j = W_k([y0_emb, sigma_mean])
+  Per-edge: semantic = (q_i · k_j)/sqrt(D_s), geo_bias = geo_mlp(mean_rel, std_rel, min_dist, heading_diff_mean)
+  Top-N by (semantic + geo_bias) → RelTrajEncoder([rel_pos, heading_diff], sigma_bias)
+  GNN layers → gated residual → refined node embedding
+
+This file is a thin adapter: it produces the V6 inputs from LED's native call
+signature and decodes the refined node embeddings into a per-agent, per-mode
+residual trajectory that LED adds to its frozen core denoiser's epsilon
+output (or to the implied y_0 estimate, depending on --residual_on).
+
+The decoder's final layer is zero-initialized so the overall residual starts
+as a no-op; V6's own gated residual inside is additionally safe at init
+because its gate_proj is randomly initialized but small.
+"""
+
+import os, sys
+import torch
+import torch.nn as nn
+
+# Import MoFlow's V6 graph module (and dependencies) directly.
+_MOFLOW_ROOT = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), '..', '..', 'MoFlow'))
+if _MOFLOW_ROOT not in sys.path:
+    sys.path.insert(0, _MOFLOW_ROOT)
+
+from models.graph_interaction_nba_v6 import FutureInteractionGraphV6
+from models.interaction_baselines import build_interaction_module  # E4: SRA_MODULE factory
+
+
+class FutureInteractionGraphV6Wrapper(nn.Module):
+    """Adapter that exposes MoFlow's V6 graph to LED's call signature.
+
+    LED call:
+        delta = graph(y0_hat, past, step_idx, sigma=sigma)
+        y0_hat [B*A, K, T, 2], past [B*A, T_h, 6], step_idx int, sigma [B*A, 1] or None
+        delta  [B*A, K, T, 2]
+
+    V6 call (underneath):
+        refined = future_graph(y_emb, y_abs, t_emb, tau, sigma_agent)
+        y_emb [B, K, A, D], y_abs [B, K, A, T, 2], t_emb [B, D], tau [B], sigma_agent [B, K, A, T] or None
+    """
+
+    def __init__(
+        self,
+        num_agents: int = 11,
+        future_steps: int = 20,
+        past_steps: int = 10,
+        past_channels: int = 6,
+        node_dim: int = 128,
+        top_n: int = 5,                  # matches MoFlow default
+        num_denoise_steps: int = 5,
+        num_gnn_layers: int = 2,
+        rel_traj_hidden: int = 32,
+        y0_score_dim: int = 32,
+        num_heads: int = 4,
+        dropout: float = 0.1,
+        edge_mode: str = 'full',
+        neighbor_mode: str = 'rag',
+    ):
+        super().__init__()
+        self.A = num_agents
+        self.T = future_steps
+        self.D = node_dim
+
+        # ---- Input encoders for the node embedding fed to V6 ----
+        self.past_proj = nn.Sequential(
+            nn.Linear(past_steps * past_channels, node_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(node_dim, node_dim),
+        )
+        self.y0_proj = nn.Sequential(
+            nn.Linear(future_steps * 2, node_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(node_dim, node_dim),
+        )
+        self.step_emb = nn.Embedding(num_denoise_steps, node_dim)
+
+        # ---- MoFlow V6 graph (scoring, edge encoding, GNN, gated residual) ----
+        self.future_graph = build_interaction_module(   # E4: sra|gameformer|c2f via env SRA_MODULE
+            embed_dim        = node_dim,
+            future_steps     = future_steps,
+            num_agents       = num_agents,
+            num_heads        = num_heads,
+            dropout          = dropout,
+            num_gnn_layers   = num_gnn_layers,
+            time_dim         = node_dim,
+            top_n_neighbors  = min(top_n, num_agents - 1),
+            rel_traj_hidden  = rel_traj_hidden,
+            y0_score_dim     = y0_score_dim,
+            edge_mode        = edge_mode,
+            neighbor_mode    = neighbor_mode,
+        )
+
+        # ---- Decoder: refined node embedding -> per-timestep trajectory residual ----
+        self.decode = nn.Sequential(
+            nn.Linear(node_dim, node_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(node_dim, future_steps * 2),
+        )
+        nn.init.zeros_(self.decode[-1].weight)
+        nn.init.zeros_(self.decode[-1].bias)
+
+    def forward(
+        self,
+        y0_hat: torch.Tensor,   # [B*A, K, T, 2]
+        past: torch.Tensor,     # [B*A, T_h, 6]
+        step_idx: int,
+        sigma: torch.Tensor = None,  # [B*A, 1] (per-agent scalar uncertainty, optional)
+        A_override: int = None,
+    ) -> torch.Tensor:          # [B*A, K, T, 2]
+        BA, K, T, _ = y0_hat.shape
+        # If A_override is given (variable-A scenes), use it; otherwise use init A.
+        if A_override is not None:
+            A = A_override
+        else:
+            A = self.A
+        if BA % A != 0:
+            # Fall back to treating whole batch as a single scene (B=1, A=BA)
+            A = BA
+        B = BA // A
+        device = y0_hat.device
+        # If A < 2, graph is meaningless — return zero residual
+        if A < 2:
+            return torch.zeros_like(y0_hat)
+        # Rebuild graph's edge index for this A
+        fg = self.future_graph
+        if fg.num_agents != A:
+            fg.num_agents = A
+            fg._E0 = A * (A - 1)
+            fg.top_n = max(1, min(fg.top_n, A - 1))
+            src, dst = [], []
+            for i in range(A):
+                for j in range(A):
+                    if i != j: src.append(j); dst.append(i)
+            fg._single_edge_index = torch.tensor([src, dst], dtype=torch.long, device=device)
+
+        # --- Build y_emb [B, K, A, D] from past + y0 + step ---
+        past_feat = self.past_proj(past.reshape(BA, -1))             # [BA, D]
+        past_feat = past_feat.unsqueeze(1).expand(BA, K, self.D)     # [BA, K, D]
+        y0_feat   = self.y0_proj(y0_hat.reshape(BA, K, T * 2))       # [BA, K, D]
+        step      = self.step_emb(torch.tensor(step_idx, device=device))  # [D]
+        nodes_bka = past_feat + y0_feat + step                       # [BA, K, D]
+        y_emb     = nodes_bka.view(B, A, K, self.D).permute(0, 2, 1, 3).contiguous()  # [B, K, A, D]
+
+        # --- y_abs [B, K, A, T, 2] (unnormalized future positions) ---
+        y_abs = y0_hat.view(B, A, K, T, 2).permute(0, 2, 1, 3, 4).contiguous()
+
+        # --- t_emb [B, D] and tau [B] ---
+        t_emb = self.step_emb(torch.tensor(step_idx, device=device)).unsqueeze(0).expand(B, -1)  # [B, D]
+        tau   = torch.full((B,), float(step_idx), dtype=torch.float32, device=device)
+
+        # --- sigma_agent [B, K, A, T] or None ---
+        # LED's variance_estimation is per-agent scalar [BA, 1]; broadcast over K and T.
+        sigma_agent = None
+        if sigma is not None:
+            sigma_ba = sigma.view(B, A)                              # [B, A]
+            sigma_agent = sigma_ba.view(B, 1, A, 1).expand(B, K, A, T)
+
+        # --- Run V6 graph: scoring, sparse edges, RelTrajEncoder, GNN, gated residual ---
+        y_emb_refined = self.future_graph(y_emb, y_abs, t_emb, tau, sigma_agent=sigma_agent)
+        # y_emb_refined: [B, K, A, D]
+
+        # --- Decode to per-mode per-agent trajectory residual ---
+        refined_bka = y_emb_refined.permute(0, 2, 1, 3).contiguous().view(BA, K, self.D)
+        residual = self.decode(refined_bka).view(BA, K, T, 2)
+        return residual

LED/models/interaction_baselines.py

@@ -0,0 +1,246 @@
+"""
+E4 plug-in baselines — alternative interaction/refinement modules that honor the EXACT
+FutureInteractionGraphV6 forward contract, for a protocol-matched comparison vs SRA:
+
+    forward(y_emb[B,K,A,D], y_abs[B,K,A,T,2], t_emb[B,D], tau[B],
+            sigma_agent[B,K,A,T]|None, agent_mask[B,A]|None=None) -> [B,K,A,D]
+    (gated residual, embedding->embedding, NO sigma out; robust to K in {1,10,20}, any A)
+
+Modules:
+  * GameFormerInteraction  [ref 20]: FAITHFUL GameFormer level-k InteractionDecoder —
+    host denoiser = InitialDecoder (level 0); denoiser K samples = GameFormer K modes;
+    FutureEncoder + score-weighted mode aggregation + interaction self-attention +
+    own-future-masked cross-attention + L-level GMM refinement. No neighbor selection,
+    no uncertainty (that is SRA). See the detailed faithfulness note above the class.
+  * CoarseToFineRefine     [ref 22]: Coarse-to-Fine (Jia et al. 2022) TEMPORAL refinement —
+    a per-agent temporal module (GRU / 1D-CNN over the future horizon). No interaction.
+    Orthogonal axis (temporal).
+
+Both accept the SAME __init__ kwargs as V6 (extras absorbed by **kwargs) and return the SAME
+gated-residual embedding, so they are true drop-ins at each host's V6 instantiation line.
+Selection is via env var SRA_MODULE in {sra, gameformer, c2f} (see build_interaction_module).
+"""
+import os
+import torch
+import torch.nn as nn
+
+
+class _GatedResidualHead(nn.Module):
+    """V6's output convention: out = input + gate(cat[input,refined]) * out_proj(refined).
+    Random-initialised (active at init), exactly like V6 — hosts that need no-op-at-init
+    (MID/LED) supply their own external zero-init projection; MoFlow uses the output directly."""
+    def __init__(self, embed_dim):
+        super().__init__()
+        self.gate_proj = nn.Sequential(nn.Linear(2 * embed_dim, embed_dim), nn.Sigmoid())
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        nn.init.zeros_(self.out_proj.weight)               # NO-OP AT INIT: module starts as identity
+        nn.init.zeros_(self.out_proj.bias)                 # (stabilizes MoFlow, which uses output directly)
+
+    def forward(self, orig, refined):                      # both [N, D]
+        gate = self.gate_proj(torch.cat([orig, refined], dim=-1))
+        res = self.out_proj(refined)
+        _cap = float(os.environ.get('MOFLOW_DAMP', 0) or 0)  # cap residual norm -> MoFlow flow-field stability
+        if _cap > 0:
+            rn = res.norm(dim=-1, keepdim=True)
+            res = res * (rn.clamp(max=_cap) / (rn + 1e-6))
+        return orig + gate * res
+
+
+# ============================================================================
+#  GameFormer [ref 20] -- FAITHFUL level-k InteractionDecoder as a denoiser plug-in
+# ============================================================================
+# Faithful to Liu et al. ICCV'23 (github.com/MCZhi/GameFormer, model/modules.py):
+#   * The host DENOISER plays the InitialDecoder (level 0): its current estimate
+#     (y_emb, y_abs) IS the level-0 prediction, and the denoiser's K SAMPLES are
+#     treated as GameFormer's K MODES (the modal set the level-k step reasons over).
+#   * Each InteractionDecoder level implements the real mechanism:
+#       - FutureEncoder: MLP on per-mode [x,y,heading,vx,vy], max-pooled over time
+#         (modules.FutureEncoder; box-size channels dropped -- no size in NBA/sport).
+#       - score-softmax-weighted aggregation of the K modes -> one future token/agent.
+#       - interaction SelfTransformer over agents (game-theoretic "condition on all
+#         other agents' level-(k-1) futures").
+#       - cross-attention of each mode's content (prev content + own future) to
+#         [interaction ; scene-context], with the agent's OWN future token MASKED
+#         (GameFormer's own-future masking).
+#       - a GMM-mu trajectory decoder emits the level-k refinement; iterate L levels.
+#   * The final level's content -> gated residual to the host embedding.
+# ONE gap vs the paper: per-level imitation SUPERVISION (GameFormer sums a GMM loss
+# over every level). A plug-in is trained only through the host denoising loss, so
+# the levels are learned end-to-end rather than level-wise-supervised.
+
+
+class _AgentSelfAttn(nn.Module):
+    """SelfTransformer over agents (modules.SelfTransformer). x: [B, A, D]."""
+    def __init__(self, dim, heads, dropout):
+        super().__init__()
+        self.heads = heads
+        self.attn = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
+        self.n1 = nn.LayerNorm(dim); self.n2 = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(nn.Linear(dim, dim * 4), nn.GELU(),
+                                 nn.Dropout(dropout), nn.Linear(dim * 4, dim))
+
+    def forward(self, x, key_padding_mask=None):
+        a, _ = self.attn(x, x, x, key_padding_mask=key_padding_mask, need_weights=False)
+        x = self.n1(a + x)
+        return self.n2(self.ffn(x) + x)
+
+
+class _CrossAttn(nn.Module):
+    """CrossTransformer (modules.CrossTransformer). q:[N,Lq,D] k/v:[N,Lk,D]."""
+    def __init__(self, dim, heads, dropout):
+        super().__init__()
+        self.heads = heads
+        self.attn = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
+        self.n1 = nn.LayerNorm(dim); self.n2 = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(nn.Linear(dim, dim * 4), nn.GELU(),
+                                 nn.Dropout(dropout), nn.Linear(dim * 4, dim))
+
+    def forward(self, q, k, v, attn_mask=None):
+        a, _ = self.attn(q, k, v, attn_mask=attn_mask, need_weights=False)
+        a = self.n1(a)
+        return self.n2(self.ffn(a) + a)
+
+
+class _FutureEncoder(nn.Module):
+    """modules.FutureEncoder: per-mode future -> vector via max-pool over an MLP on
+    [x, y, heading, vx, vy].  trajs [B,K,A,T,2] (K modes), cur_xy [B,K,A,2]."""
+    def __init__(self, dim, dt=0.2):
+        super().__init__()
+        self.dt = dt
+        self.mlp = nn.Sequential(nn.Linear(5, 64), nn.ReLU(inplace=True), nn.Linear(64, dim))
+
+    def forward(self, trajs, cur_xy):
+        pos = trajs - cur_xy.unsqueeze(-2)                             # centered position
+        xy = torch.cat([cur_xy.unsqueeze(-2), trajs], dim=-2)          # [B,K,A,T+1,2]
+        dxy = torch.diff(xy, dim=-2)
+        v = dxy / self.dt
+        theta = torch.atan2(dxy[..., 1], dxy[..., 0].clamp(min=1e-3)).unsqueeze(-1)
+        state = torch.cat([pos, theta, v], dim=-1)                     # [B,K,A,T,5]
+        h = self.mlp(state.detach())                                   # GameFormer DETACHES future feats
+        return h.max(dim=-2).values                                    # [B,K,A,D]
+
+
+class GameFormerInteraction(nn.Module):
+    """FAITHFUL GameFormer level-k InteractionDecoder as a denoiser plug-in (see header)."""
+
+    def __init__(self, embed_dim, future_steps, num_agents,
+                 num_heads=4, dropout=0.1, num_gnn_layers=2, time_dim=128,
+                 top_n_neighbors=5, rel_traj_hidden=32, y0_score_dim=32,
+                 edge_mode='full', neighbor_mode='rag', num_levels=None, **kwargs):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_agents = num_agents            # LED wrapper reads this
+        L = int(os.environ.get('GF_LEVELS', num_levels if num_levels is not None else 3))
+        self.num_levels = L
+        self.t_proj = nn.Linear(embed_dim, embed_dim)
+        self.future_encoder = _FutureEncoder(embed_dim)                # SHARED across levels
+        self.score_head = nn.ModuleList([                              # per-mode score -> softmax weight
+            nn.Sequential(nn.Linear(embed_dim, 64), nn.ELU(), nn.Linear(64, 1)) for _ in range(L)])
+        self.interaction_enc = nn.ModuleList([_AgentSelfAttn(embed_dim, num_heads, dropout) for _ in range(L)])
+        self.query_enc = nn.ModuleList([_CrossAttn(embed_dim, num_heads, dropout) for _ in range(L)])
+        self.traj_decode = nn.ModuleList([nn.Linear(embed_dim, future_steps * 2) for _ in range(L)])
+        for dec in self.traj_decode:                                   # each level starts as identity (stable)
+            nn.init.zeros_(dec.weight); nn.init.zeros_(dec.bias)
+        self.head = _GatedResidualHead(embed_dim)
+
+    def forward(self, y_emb, y_abs, t_emb, tau, sigma_agent=None, agent_mask=None):
+        B, K, A, D = y_emb.shape
+        T = y_abs.shape[3]
+        if A <= 1:
+            return y_emb
+        dev = y_emb.device
+        cur_xy = y_abs[..., 0, :]                            # [B,K,A,2] level-0 reference (1st future step)
+        y_traj = y_abs                                       # [B,K,A,T,2] current (level-0) future
+        t = self.t_proj(t_emb).view(B, 1, 1, D)              # denoiser timestep context
+        agent_ctx = y_emb.mean(dim=1) + t.squeeze(1)         # [B,A,D] fixed scene context (over modes)
+        kp = (~agent_mask.bool()) if agent_mask is not None else None   # [B,A] True=pad
+        ar = torch.arange(A, device=dev)
+        heads = self.query_enc[0].heads
+
+        content = y_emb                                      # level-0 content = host embedding
+        for l in range(self.num_levels):
+            multi_fut = self.future_encoder(y_traj, cur_xy)             # [B,K,A,D] per-mode future
+            w = self.score_head[l](multi_fut).softmax(dim=1)           # [B,K,A,1] over K modes
+            agg_fut = (multi_fut * w).sum(dim=1)                        # [B,A,D] aggregated future
+            interaction = self.interaction_enc[l](agg_fut, kp)         # [B,A,D] game-theoretic
+            ctx = torch.cat([interaction, agent_ctx], dim=1)           # [B,2A,D]
+
+            q = (content + multi_fut + t).reshape(B, K * A, D)         # prev content + own future + t
+            am = torch.zeros(B, K, A, 2 * A, dtype=torch.bool, device=dev)
+            am[:, :, ar, ar] = True                                     # mask OWN future (interaction block)
+            if agent_mask is not None:
+                pad = (~agent_mask.bool())[:, None, None, :]           # [B,1,1,A]
+                am[..., :A] = am[..., :A] | pad
+                am[..., A:] = am[..., A:] | pad
+            am = am.reshape(B, K * A, 2 * A)[:, None].expand(
+                B, heads, K * A, 2 * A).reshape(B * heads, K * A, 2 * A)
+            content = self.query_enc[l](q, ctx, ctx, am).reshape(B, K, A, D)
+            y_traj = y_traj + self.traj_decode[l](content).view(B, K, A, T, 2)   # level-k refinement
+        refined = content.reshape(B * K * A, D)
+        orig = y_emb.reshape(B * K * A, D)
+        return self.head(orig, refined).view(B, K, A, D)
+
+
+class CoarseToFineRefine(nn.Module):
+    """Coarse-to-Fine temporal refiner: per-agent GRU/1D-CNN over the future horizon (no interaction)."""
+
+    def __init__(self, embed_dim, future_steps, num_agents,
+                 num_heads=4, dropout=0.1, num_gnn_layers=2, time_dim=128,
+                 top_n_neighbors=5, rel_traj_hidden=32, y0_score_dim=32,
+                 edge_mode='full', neighbor_mode='rag', refine_type='gru', hidden=200, **kwargs):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_agents = num_agents            # LED wrapper reads this (skips variable-A rebuild when ==A)
+        self.refine_type = os.environ.get('C2F_REFINE', refine_type)
+        self.pos_emb = nn.Linear(2, hidden)
+        if self.refine_type == 'cnn':                                   # 1D-CNN temporal refiner (per-timestep output)
+            self.temporal = nn.Sequential(
+                nn.Conv1d(hidden, hidden, 3, padding=1), nn.ReLU(inplace=True),
+                nn.Conv1d(hidden, hidden, 3, padding=1), nn.ReLU(inplace=True),
+                nn.Conv1d(hidden, hidden, 3, padding=1), nn.ReLU(inplace=True))
+        else:                                                           # AUTOREGRESSIVE (unidirectional) GRU
+            self.temporal = nn.GRU(hidden, hidden, num_layers=2, batch_first=True, dropout=dropout)
+        self.delta_head = nn.Linear(hidden, 2)                          # per-timestep coarse->fine correction Δ_t
+        self.traj_encode = nn.Sequential(                              # re-encode the FINE trajectory -> host feature
+            nn.Linear(future_steps * 2, embed_dim), nn.ReLU(inplace=True),
+            nn.Linear(embed_dim, embed_dim))
+        self.t_proj = nn.Linear(embed_dim, embed_dim)
+        self.head = _GatedResidualHead(embed_dim)
+
+    def forward(self, y_emb, y_abs, t_emb, tau, sigma_agent=None, agent_mask=None):
+        # Faithful coarse-to-fine: host prediction = COARSE trajectory; walk it temporally
+        # (autoregressive GRU / 1D-CNN) emitting a per-timestep correction Δ_t -> FINE trajectory,
+        # then re-encode the fine trajectory into the host residual. Per-agent (no interaction).
+        B, K, A, D = y_emb.shape
+        T = y_abs.shape[3]
+        N = B * K * A
+        y_coarse = (y_abs - y_abs[..., :1, :]).reshape(N, T, 2)        # coarse trajectory (centered)
+        seq = self.pos_emb(y_coarse)                                    # [N, T, H]
+        if self.refine_type == 'cnn':
+            h = self.temporal(seq.transpose(1, 2)).transpose(1, 2)      # [N, T, H]
+        else:
+            h, _ = self.temporal(seq)                                   # [N, T, H] per-timestep (autoregressive)
+        delta = self.delta_head(h)                                      # [N, T, 2] coarse->fine correction
+        y_fine = y_coarse + delta                                       # refined (fine) trajectory
+        refined = self.traj_encode(y_fine.reshape(N, T * 2))            # re-encode fine traj -> feature
+        t = self.t_proj(t_emb).view(B, 1, 1, D).expand(B, K, A, D).reshape(N, D)
+        refined = refined + t
+        orig = y_emb.reshape(N, D)
+        return self.head(orig, refined).view(B, K, A, D)
+
+
+def build_interaction_module(name=None, **kwargs):
+    """Factory used at each host's V6 instantiation line. `name` defaults to env SRA_MODULE
+    (then 'sra'). SRA branch imports the real V6 and forwards only kwargs it accepts."""
+    name = (name or os.environ.get('SRA_MODULE') or 'sra').lower()
+    if name in ('sra', 'v6', 'graph'):
+        import inspect
+        from models.graph_interaction_nba_v6 import FutureInteractionGraphV6
+        allowed = set(inspect.signature(FutureInteractionGraphV6.__init__).parameters)
+        v6kw = {k: v for k, v in kwargs.items() if k in allowed}
+        return FutureInteractionGraphV6(**v6kw)
+    if name in ('gameformer', 'gf', 'gameformer_int'):
+        return GameFormerInteraction(**kwargs)
+    if name in ('c2f', 'coarse2fine', 'coarsetofine'):
+        return CoarseToFineRefine(**kwargs)
+    raise ValueError(f"unknown SRA_MODULE / interaction_module: {name!r}")

LED/models/layers.py

@@ -0,0 +1,169 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import Module, Linear
+
+
+class PositionalEncoding(nn.Module):
+	def __init__(self, d_model, dropout=0.1, max_len=5000):
+		super().__init__()
+
+		self.dropout = nn.Dropout(p=dropout)
+		pe = torch.zeros(max_len, d_model)
+		position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+		div_term = torch.exp(
+			torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
+		)
+		pe[:, 0::2] = torch.sin(position * div_term)
+		pe[:, 1::2] = torch.cos(position * div_term)
+		pe = pe.unsqueeze(0).transpose(0, 1)
+		self.register_buffer("pe", pe)
+
+	def forward(self, x):
+		x = x + self.pe[: x.size(0), :]
+		return self.dropout(x)
+
+
+class ConcatSquashLinear(Module):
+	def __init__(self, dim_in, dim_out, dim_ctx):
+		super(ConcatSquashLinear, self).__init__()
+		self._layer = Linear(dim_in, dim_out)
+		self._hyper_bias = Linear(dim_ctx, dim_out, bias=False)
+		self._hyper_gate = Linear(dim_ctx, dim_out)
+
+	def forward(self, ctx, x):
+		# ctx: (B, 1, F+3)
+		# x: (B, T, 2)
+		gate = torch.sigmoid(self._hyper_gate(ctx))
+		bias = self._hyper_bias(ctx)
+		# if x.dim() == 3:
+		#     gate = gate.unsqueeze(1)
+		#     bias = bias.unsqueeze(1)
+		ret = self._layer(x) * gate + bias
+		return ret
+	
+	def batch_generate(self, ctx, x):
+		# ctx: (B, n, 1, F+3)
+		# x: (B, n, T, 2)
+		gate = torch.sigmoid(self._hyper_gate(ctx))
+		bias = self._hyper_bias(ctx)
+		# if x.dim() == 3:
+		#     gate = gate.unsqueeze(1)
+		#     bias = bias.unsqueeze(1)
+		ret = self._layer(x) * gate + bias
+		return ret
+	
+
+class GAT(nn.Module):
+	def __init__(self, in_feat=2, out_feat=64, n_head=4, dropout=0.1, skip=True):
+		super(GAT, self).__init__()
+		self.in_feat = in_feat
+		self.out_feat = out_feat
+		self.n_head = n_head
+		self.skip = skip
+		self.w = nn.Parameter(torch.Tensor(n_head, in_feat, out_feat))
+		self.a_src = nn.Parameter(torch.Tensor(n_head, out_feat, 1))
+		self.a_dst = nn.Parameter(torch.Tensor(n_head, out_feat, 1))
+		self.bias = nn.Parameter(torch.Tensor(out_feat))
+
+		self.leaky_relu = nn.LeakyReLU(negative_slope=0.2)
+		self.softmax = nn.Softmax(dim=-1)
+		self.dropout = nn.Dropout(dropout)
+
+		nn.init.xavier_uniform_(self.w, gain=1.414)
+		nn.init.xavier_uniform_(self.a_src, gain=1.414)
+		nn.init.xavier_uniform_(self.a_dst, gain=1.414)
+		nn.init.constant_(self.bias, 0)
+
+	def forward(self, h, mask):
+		h_prime = h.unsqueeze(1) @ self.w
+		attn_src = h_prime @ self.a_src
+		attn_dst = h_prime @ self.a_dst
+		attn = attn_src @ attn_dst.permute(0, 1, 3, 2)
+		attn = self.leaky_relu(attn)
+		attn = self.softmax(attn)
+		attn = self.dropout(attn)
+		attn = attn * mask if mask is not None else attn
+		out = (attn @ h_prime).sum(dim=1) + self.bias
+		if self.skip:
+			out += h_prime.sum(dim=1)
+		return out, attn
+
+
+class MLP(nn.Module):
+	def __init__(self, in_feat, out_feat, hid_feat=(1024, 512), activation=None, dropout=-1):
+		super(MLP, self).__init__()
+		dims = (in_feat, ) + hid_feat + (out_feat, )
+
+		self.layers = nn.ModuleList()
+		for i in range(len(dims) - 1):
+			self.layers.append(nn.Linear(dims[i], dims[i + 1]))
+
+		self.activation = activation if activation is not None else lambda x: x
+		self.dropout = nn.Dropout(dropout) if dropout != -1 else lambda x: x
+
+	def forward(self, x):
+		for i in range(len(self.layers)):
+			x = self.activation(x)
+			x = self.dropout(x)
+			x = self.layers[i](x)
+		return x
+
+
+class social_transformer(nn.Module):
+	def __init__(self, past_len):
+		super(social_transformer, self).__init__()
+		self.encode_past = nn.Linear(past_len*6, 256, bias=False)
+		self.layer = nn.TransformerEncoderLayer(d_model=256, nhead=2, dim_feedforward=256)
+		self.transformer_encoder = nn.TransformerEncoder(self.layer, num_layers=2)
+
+	def forward(self, h, mask):
+		'''
+		h: batch_size, t, 2
+		'''
+		h_feat = self.encode_past(h.reshape(h.size(0), -1)).unsqueeze(1)
+		# print(h_feat.shape)
+		# n_samples, 1, 64
+		h_feat_ = self.transformer_encoder(h_feat, mask)
+		h_feat = h_feat + h_feat_
+
+		return h_feat
+
+
+class st_encoder(nn.Module):
+	def __init__(self):
+		super().__init__()
+		channel_in = 6
+		channel_out = 32
+		dim_kernel = 3
+		self.dim_embedding_key = 256
+		self.spatial_conv = nn.Conv1d(channel_in, channel_out, dim_kernel, stride=1, padding=1)
+		self.temporal_encoder = nn.GRU(channel_out, self.dim_embedding_key, 1, batch_first=True)
+
+		self.relu = nn.ReLU()
+
+		self.reset_parameters()
+
+	def reset_parameters(self):
+		nn.init.kaiming_normal_(self.spatial_conv.weight)
+		nn.init.kaiming_normal_(self.temporal_encoder.weight_ih_l0)
+		nn.init.kaiming_normal_(self.temporal_encoder.weight_hh_l0)
+		nn.init.zeros_(self.spatial_conv.bias)
+		nn.init.zeros_(self.temporal_encoder.bias_ih_l0)
+		nn.init.zeros_(self.temporal_encoder.bias_hh_l0)
+
+	def forward(self, X):
+		'''
+		X: b, T, 2
+
+		return: b, F
+		'''
+		X_t = torch.transpose(X, 1, 2)
+		X_after_spatial = self.relu(self.spatial_conv(X_t))
+		X_embed = torch.transpose(X_after_spatial, 1, 2)
+
+		output_x, state_x = self.temporal_encoder(X_embed)
+		state_x = state_x.squeeze(0)
+
+		return state_x
+

LED/models/model_diffusion.py

@@ -0,0 +1,125 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import Module, Linear
+
+from models.layers import PositionalEncoding, ConcatSquashLinear
+
+class st_encoder(nn.Module):
+	def __init__(self):
+		super().__init__()
+		channel_in = 2
+		channel_out = 32
+		dim_kernel = 3
+		self.dim_embedding_key = 256
+		self.spatial_conv = nn.Conv1d(channel_in, channel_out, dim_kernel, stride=1, padding=1)
+		self.temporal_encoder = nn.GRU(channel_out, self.dim_embedding_key, 1, batch_first=True)
+
+		self.relu = nn.ReLU()
+
+		self.reset_parameters()
+
+	def reset_parameters(self):
+		nn.init.kaiming_normal_(self.spatial_conv.weight)
+		nn.init.kaiming_normal_(self.temporal_encoder.weight_ih_l0)
+		nn.init.kaiming_normal_(self.temporal_encoder.weight_hh_l0)
+		nn.init.zeros_(self.spatial_conv.bias)
+		nn.init.zeros_(self.temporal_encoder.bias_ih_l0)
+		nn.init.zeros_(self.temporal_encoder.bias_hh_l0)
+
+	def forward(self, X):
+		'''
+		X: b, T, 2
+
+		return: b, F
+		'''
+		X_t = torch.transpose(X, 1, 2)
+		X_after_spatial = self.relu(self.spatial_conv(X_t))
+		X_embed = torch.transpose(X_after_spatial, 1, 2)
+
+		output_x, state_x = self.temporal_encoder(X_embed)
+		state_x = state_x.squeeze(0)
+
+		return state_x
+
+
+class social_transformer(nn.Module):
+	def __init__(self, past_len=10, in_channels=6):
+		super(social_transformer, self).__init__()
+		self.encode_past = nn.Linear(past_len * in_channels, 256, bias=False)
+		self.layer = nn.TransformerEncoderLayer(d_model=256, nhead=2, dim_feedforward=256)
+		self.transformer_encoder = nn.TransformerEncoder(self.layer, num_layers=2)
+
+	def forward(self, h, mask):
+		'''
+		h: batch_size, t, 2
+		'''
+		# print(h.shape)
+		h_feat = self.encode_past(h.reshape(h.size(0), -1)).unsqueeze(1)
+		# print(h_feat.shape)
+		# n_samples, 1, 64
+		h_feat_ = self.transformer_encoder(h_feat, mask)
+		h_feat = h_feat + h_feat_
+
+		return h_feat
+
+
+class TransformerDenoisingModel(Module):
+
+	def __init__(self, context_dim=256, tf_layer=2, past_len=10):
+		super().__init__()
+		self.encoder_context = social_transformer(past_len=past_len)
+		self.pos_emb = PositionalEncoding(d_model=2*context_dim, dropout=0.1, max_len=24)
+		self.concat1 = ConcatSquashLinear(2, 2*context_dim, context_dim+3)
+		self.layer = nn.TransformerEncoderLayer(d_model=2*context_dim, nhead=2, dim_feedforward=2*context_dim)
+		self.transformer_encoder = nn.TransformerEncoder(self.layer, num_layers=tf_layer)
+		self.concat3 = ConcatSquashLinear(2*context_dim,context_dim,context_dim+3)
+		self.concat4 = ConcatSquashLinear(context_dim,context_dim//2,context_dim+3)
+		self.linear = ConcatSquashLinear(context_dim//2, 2, context_dim+3)
+
+
+	def forward(self, x, beta, context, mask):
+		batch_size = x.size(0)
+		beta = beta.view(batch_size, 1, 1)          # (B, 1, 1)
+		mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+		context = self.encoder_context(context, mask)
+		# context = context.view(batch_size, 1, -1)   # (B, 1, F)
+
+		time_emb = torch.cat([beta, torch.sin(beta), torch.cos(beta)], dim=-1)  # (B, 1, 3)
+		ctx_emb = torch.cat([time_emb, context], dim=-1)    # (B, 1, F+3)
+		
+		x = self.concat1(ctx_emb, x)
+		final_emb = x.permute(1,0,2)
+		final_emb = self.pos_emb(final_emb)
+		
+		trans = self.transformer_encoder(final_emb).permute(1,0,2)
+		trans = self.concat3(ctx_emb, trans)
+		trans = self.concat4(ctx_emb, trans)
+		return self.linear(ctx_emb, trans)
+	
+
+	def generate_accelerate(self, x, beta, context, mask):
+		batch_size = x.size(0)
+		beta = beta.view(beta.size(0), 1, 1)          # (B, 1, 1)
+		mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+		context = self.encoder_context(context, mask)
+		# context = context.view(batch_size, 1, -1)   # (B, 1, F)
+
+		time_emb = torch.cat([beta, torch.sin(beta), torch.cos(beta)], dim=-1)  # (B, 1, 3)
+		# time_emb: [11, 1, 3]
+		# context: [11, 1, 256]
+		ctx_emb = torch.cat([time_emb, context], dim=-1).repeat(1, 10, 1).unsqueeze(2)
+		# x: 11, 10, 20, 2
+		# ctx_emb: 11, 10, 1, 259
+		K = x.size(1)
+		T = x.size(2)
+		D = 2 * 256
+		x = self.concat1.batch_generate(ctx_emb, x).contiguous().view(-1, T, D)
+		final_emb = x.permute(1, 0, 2)
+		final_emb = self.pos_emb(final_emb)
+
+		trans = self.transformer_encoder(final_emb).permute(1, 0, 2).contiguous().view(-1, K, T, D)
+		# trans: 11, 10, 20, 512
+		trans = self.concat3.batch_generate(ctx_emb, trans)
+		trans = self.concat4.batch_generate(ctx_emb, trans)
+		return self.linear.batch_generate(ctx_emb, trans)

LED/models/model_led_initializer.py

@@ -0,0 +1,63 @@
+import torch
+import torch.nn as nn
+from models.layers import MLP, social_transformer, st_encoder
+
+class LEDInitializer(nn.Module):
+	def __init__(self, t_h: int=8, d_h: int=6, t_f: int=40, d_f: int=2, k_pred: int=20):
+		'''
+		Parameters
+		----
+		t_h: history timestamps,
+		d_h: dimension of each historical timestamp,
+		t_f: future timestamps,
+		d_f: dimension of each future timestamp,
+		k_pred: number of predictions.
+
+		'''
+		super(LEDInitializer, self).__init__()
+		self.n = k_pred
+		self.input_dim = t_h * d_h
+		self.output_dim = t_f * d_f * k_pred
+		self.fut_len = t_f
+
+		self.social_encoder = social_transformer(t_h)
+		self.ego_var_encoder = st_encoder()
+		self.ego_mean_encoder = st_encoder()
+		self.ego_scale_encoder = st_encoder()
+
+		self.scale_encoder = MLP(1, 32, hid_feat=(4, 16), activation=nn.ReLU())
+
+		self.var_decoder = MLP(256*2+32, self.output_dim, hid_feat=(1024, 1024), activation=nn.ReLU())
+		self.mean_decoder = MLP(256*2, t_f * d_f, hid_feat=(256, 128), activation=nn.ReLU())
+		self.scale_decoder = MLP(256*2, 1, hid_feat=(256, 128), activation=nn.ReLU())
+
+	
+	def forward(self, x, mask=None):
+		'''
+		x: batch size, t_p, 6
+		'''
+		mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+		social_embed = self.social_encoder(x, mask)
+		social_embed = social_embed.squeeze(1)
+		# B, 256
+		
+		ego_var_embed = self.ego_var_encoder(x)
+		ego_mean_embed = self.ego_mean_encoder(x)
+		ego_scale_embed = self.ego_scale_encoder(x)
+		# B, 256
+
+		mean_total = torch.cat((ego_mean_embed, social_embed), dim=-1)
+		
+		guess_mean = self.mean_decoder(mean_total).contiguous().view(-1, self.fut_len, 2)
+
+		scale_total = torch.cat((ego_scale_embed, social_embed), dim=-1)
+		guess_scale = self.scale_decoder(scale_total)
+
+		guess_scale_feat = self.scale_encoder(guess_scale)
+		var_total = torch.cat((ego_var_embed, social_embed, guess_scale_feat), dim=-1)
+		guess_var = self.var_decoder(var_total).reshape(x.size(0), self.n, self.fut_len, 2)
+
+		return guess_var, guess_mean, guess_scale
+
+
+

LED/requirements.txt

@@ -0,0 +1,84 @@
+absl-py==1.4.0
+aiohttp==3.8.4
+aiosignal==1.3.1
+antlr4-python3-runtime==4.8
+async-timeout==4.0.2
+asynctest==0.13.0
+attrs==23.1.0
+cachetools==5.3.0
+charset-normalizer==3.1.0
+colour==0.1.5
+cycler==0.11.0
+descartes==1.1.0
+easydict==1.10
+fonttools==4.38.0
+frozenlist==1.3.3
+fsspec==2023.1.0
+future==0.18.3
+glob2==0.7
+google-auth==2.18.0
+google-auth-oauthlib==0.4.6
+googledrivedownloader==0.4
+grpcio==1.54.0
+h5py==3.8.0
+hydra-core==1.1.0
+imageio==2.27.0
+imgaug==0.4.0
+importlib-metadata==4.13.0
+importlib-resources==5.12.0
+isodate==0.6.1
+Jinja2==3.1.2
+joblib==1.2.0
+kiwisolver==1.4.4
+lapsolver==1.1.0
+llvmlite==0.39.1
+Markdown==3.4.3
+MarkupSafe==2.1.2
+matplotlib==3.5.3
+motmetrics==1.1.3
+multidict==6.0.4
+networkx==2.6.3
+numba==0.56.4
+numpy==1.19.0
+oauthlib==3.2.2
+omegaconf==2.1.0
+opencv-python==4.7.0.72
+packaging==23.1
+pandas==1.3.5
+Pillow==9.5.0
+polars==0.17.12
+protobuf==3.20.3
+pyasn1==0.5.0
+pyasn1-modules==0.3.0
+pyDeprecate==0.3.1
+pyntcloud==0.3.1
+pyparsing==3.0.9
+python-dateutil==2.8.2
+python-louvain==0.16
+pytorch-lightning==1.5.2
+pytz==2023.3
+PyWavelets==1.3.0
+PyYAML==6.0
+rdflib==6.3.2
+requests-oauthlib==1.3.1
+rsa==4.9
+scikit-image==0.19.3
+scikit-learn==1.0.2
+scipy==1.7.3
+shapely==2.0.1
+spconv==1.2.1
+tensorboard==2.11.2
+tensorboard-data-server==0.6.1
+tensorboard-plugin-wit==1.8.1
+tensorboardX==2.6
+threadpoolctl==3.1.0
+tifffile==2021.11.2
+torch==1.8.0+cu111
+torchaudio==0.8.0
+torchmetrics==0.11.4
+torchvision==0.9.0+cu111
+tqdm==4.65.0
+typing_extensions==4.5.0
+Werkzeug==2.2.3
+yarl==1.9.2
+zipp==3.15.0

LED/trainer/train_led_graph.py

@@ -0,0 +1,404 @@
+"""
+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)
+			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
+			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()
+
+			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)
+				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)
+				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)

LED/trainer/train_led_trajectory_augment_input.py

@@ -0,0 +1,443 @@
+
+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 data.dataloader_nba import NBADataset, seq_collate
+
+
+from models.model_led_initializer import LEDInitializer as InitializationModel
+from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel
+
+import pdb
+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)
+		
+		# data normalization parameters
+		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 # define total diffusion steps
+
+		# make beta schedule and calculate the parameters used in denoising process.
+		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()
+		# load pretrained models
+		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.opt = torch.optim.AdamW(self.model_initializer.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)
+
+		# ------------------------- prepare logs -------------------------
+		self.log = open(os.path.join(self.cfg.log_dir, 'log.txt'), 'a+')
+		self.print_model_param(self.model, name='Core Denoising Model')
+		self.print_model_param(self.model_initializer, name='Initialization Model')
+
+		# temporal reweight in the loss, it is not necessary.
+		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:
+		'''
+		Count the trainable/total parameters in `model`.
+		'''
+		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)
+		return None
+
+
+	def make_beta_schedule(self, schedule: str = 'linear', 
+			n_timesteps: int = 1000, 
+			start: float = 1e-5, end: float = 1e-2) -> torch.Tensor:
+		'''
+		Make beta schedule.
+
+		Parameters
+		----
+		schedule: str, in ['linear', 'quad', 'sigmoid'],
+		n_timesteps: int, diffusion steps,
+		start: float, beta start, `start<end`,
+		end: float, beta end,
+
+		Returns
+		----
+		betas: Tensor with the shape of (n_timesteps)
+
+		'''
+		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)
+
+	def noise_estimation_loss(self, x, y_0, mask):
+		batch_size = x.shape[0]
+		# Select a random step for each example
+		t = torch.randint(0, self.n_steps, size=(batch_size // 2 + 1,)).to(x.device)
+		t = torch.cat([t, self.n_steps - t - 1], dim=0)[:batch_size]
+		# x0 multiplier
+		a = self.extract(self.alphas_bar_sqrt, t, y_0)
+		beta = self.extract(self.betas, t, y_0)
+		# eps multiplier
+		am1 = self.extract(self.one_minus_alphas_bar_sqrt, t, y_0)
+		e = torch.randn_like(y_0)
+		# model input
+		y = y_0 * a + e * am1
+		output = self.model(y, beta, x, mask)
+		# batch_size, 20, 2
+		return (e - output).square().mean()
+
+
+
+	def p_sample(self, x, mask, cur_y, t):
+		if t==0:
+			z = torch.zeros_like(cur_y).to(x.device)
+		else:
+			z = torch.randn_like(cur_y).to(x.device)
+		t = torch.tensor([t]).cuda()
+		# Factor to the model output
+		eps_factor = ((1 - self.extract(self.alphas, t, cur_y)) / self.extract(self.one_minus_alphas_bar_sqrt, t, cur_y))
+		# Model output
+		beta = self.extract(self.betas, t.repeat(x.shape[0]), cur_y)
+		eps_theta = self.model(cur_y, beta, x, mask)
+		mean = (1 / self.extract(self.alphas, t, cur_y).sqrt()) * (cur_y - (eps_factor * eps_theta))
+		# Generate z
+		z = torch.randn_like(cur_y).to(x.device)
+		# Fixed sigma
+		sigma_t = self.extract(self.betas, t, cur_y).sqrt()
+		sample = mean + sigma_t * z
+		return (sample)
+	
+	def p_sample_accelerate(self, x, mask, cur_y, t):
+		if t==0:
+			z = torch.zeros_like(cur_y).to(x.device)
+		else:
+			z = torch.randn_like(cur_y).to(x.device)
+		t = torch.tensor([t]).cuda()
+		# Factor to the model output
+		eps_factor = ((1 - self.extract(self.alphas, t, cur_y)) / self.extract(self.one_minus_alphas_bar_sqrt, t, cur_y))
+		# Model output
+		beta = self.extract(self.betas, t.repeat(x.shape[0]), cur_y)
+		eps_theta = self.model.generate_accelerate(cur_y, beta, x, mask)
+		mean = (1 / self.extract(self.alphas, t, cur_y).sqrt()) * (cur_y - (eps_factor * eps_theta))
+		# Generate z
+		z = torch.randn_like(cur_y).to(x.device)
+		# Fixed sigma
+		sigma_t = self.extract(self.betas, t, cur_y).sqrt()
+		sample = mean + sigma_t * z * 0.00001
+		return (sample)
+
+
+
+	def p_sample_loop(self, x, mask, shape):
+		self.model.eval()
+		prediction_total = torch.Tensor().cuda()
+		for _ in range(20):
+			cur_y = torch.randn(shape).to(x.device)
+			for i in reversed(range(self.n_steps)):
+				cur_y = self.p_sample(x, mask, cur_y, i)
+			prediction_total = torch.cat((prediction_total, cur_y.unsqueeze(1)), dim=1)
+		return prediction_total
+	
+	def p_sample_loop_mean(self, x, mask, loc):
+		prediction_total = torch.Tensor().cuda()
+		for loc_i in range(1):
+			cur_y = loc
+			for i in reversed(range(NUM_Tau)):
+				cur_y = self.p_sample(x, mask, cur_y, i)
+			prediction_total = torch.cat((prediction_total, cur_y.unsqueeze(1)), dim=1)
+		return prediction_total
+
+	def p_sample_loop_accelerate(self, x, mask, loc):
+		'''
+		Batch operation to accelerate the denoising process.
+
+		x: [11, 10, 6]
+		mask: [11, 11]
+		cur_y: [11, 10, 20, 2]
+		'''
+		prediction_total = torch.Tensor().cuda()
+		cur_y = loc[:, :10]
+		for i in reversed(range(NUM_Tau)):
+			cur_y = self.p_sample_accelerate(x, mask, cur_y, i)
+		cur_y_ = loc[:, 10:]
+		for i in reversed(range(NUM_Tau)):
+			cur_y_ = self.p_sample_accelerate(x, mask, cur_y_, i)
+		# shape: B=b*n, K=10, T, 2
+		prediction_total = torch.cat((cur_y_, cur_y), dim=1)
+		return prediction_total
+
+
+
+	def fit(self):
+		# Training loop
+		for epoch in range(0, 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)
+			
+			if (epoch + 1) % self.cfg.test_interval == 0:
+				performance, samples = self._test_single_epoch()
+				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), self.log)
+				cp_path = self.cfg.model_path % (epoch + 1)
+				model_cp = {'model_initializer_dict': self.model_initializer.state_dict()}
+				torch.save(model_cp, cp_path)
+			self.scheduler_model.step()
+
+
+	def data_preprocess(self, data):
+		"""
+			pre_motion_3D: torch.Size([32, 11, 10, 2]), [batch_size, num_agent, past_frame, dimension]
+			fut_motion_3D: torch.Size([32, 11, 20, 2])
+			fut_motion_mask: torch.Size([32, 11, 20])
+			pre_motion_mask: torch.Size([32, 11, 10])
+			traj_scale: 1
+			pred_mask: None
+			seq: nba
+		"""
+		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:]
+		# augment input: absolute position, relative position, velocity
+		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()
+		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]
+			
+			generated_y = self.p_sample_loop_accelerate(past_traj, traj_mask, loc)
+			
+			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()
+			torch.nn.utils.clip_grad_norm_(self.model_initializer.parameters(), 1.)
+			self.opt.step()
+			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)
+		count = 0
+		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]
+			
+				pred_traj = self.p_sample_loop_accelerate(past_traj, traj_mask, loc)
+
+				fut_traj = fut_traj.unsqueeze(1).repeat(1, 20, 1, 1)
+				# b*n, K, T, 2
+				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]
+				count += 1
+				# if count==100:
+				# 	break
+		return performance, samples
+
+
+	def save_data(self):
+		'''
+		Save the visualization data.
+		'''
+		model_path = './results/checkpoints/led_vis.p'
+		model_dict = torch.load(model_path, map_location=torch.device('cpu'))['model_initializer_dict']
+		self.model_initializer.load_state_dict(model_dict)
+		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)
+		root_path = './visualization/data/'
+				
+		with torch.no_grad():
+			for data in self.test_loader:
+				_, traj_mask, past_traj, _ = self.data_preprocess(data)
+
+				sample_prediction, mean_estimation, variance_estimation = self.model_initializer(past_traj, traj_mask)
+				torch.save(sample_prediction, root_path+'p_var.pt')
+				torch.save(mean_estimation, root_path+'p_mean.pt')
+				torch.save(variance_estimation, root_path+'p_sigma.pt')
+
+				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]
+
+				pred_traj = self.p_sample_loop_accelerate(past_traj, traj_mask, loc)
+				pred_mean = self.p_sample_loop_mean(past_traj, traj_mask, mean_estimation)
+
+				torch.save(data['pre_motion_3D'], root_path+'past.pt')
+				torch.save(data['fut_motion_3D'], root_path+'future.pt')
+				torch.save(pred_traj, root_path+'prediction.pt')
+				torch.save(pred_mean, root_path+'p_mean_denoise.pt')
+
+				raise ValueError
+
+
+
+	def test_single_model(self):
+		model_path = './results/checkpoints/led_new.p'
+		model_dict = torch.load(model_path, map_location=torch.device('cpu'))['model_initializer_dict']
+		self.model_initializer.load_state_dict(model_dict)
+		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)
+		count = 0
+		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]
+			
+				pred_traj = self.p_sample_loop_accelerate(past_traj, traj_mask, loc)
+
+				fut_traj = fut_traj.unsqueeze(1).repeat(1, 20, 1, 1)
+				# b*n, K, T, 2
+				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]
+				count += 1
+					# if count==2:
+					# 	break
+		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)
+		
+	

LED/trainer/train_sdd_led.py

@@ -0,0 +1,239 @@
+"""
+Stage-2 LED trainer for SDD. Variable-A scenes, batch_size=1, grad_accum.
+Optional graph module (--use_graph --use_v6_graph).
+Eval: only target agent (index 0) counts toward ADE/FDE.
+"""
+
+import os, time, torch, random, 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_sdd import SDDDataset, sdd_seq_collate
+from models.model_led_initializer import LEDInitializer as InitializationModel
+from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel
+
+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.use_v6_graph = bool(getattr(config, 'use_v6_graph', False))
+        self.residual_on = getattr(config, 'residual_on', 'y0')
+        self.grad_accum  = getattr(config, 'grad_accum', 16)
+
+        train_dset = SDDDataset(obs_len=self.cfg.past_frames,
+                                pred_len=self.cfg.future_frames, split='train')
+        test_dset  = SDDDataset(obs_len=self.cfg.past_frames,
+                                pred_len=self.cfg.future_frames, split='test')
+        self.train_loader = DataLoader(train_dset, batch_size=1, shuffle=True,
+            num_workers=2, collate_fn=sdd_seq_collate)
+        self.test_loader = DataLoader(test_dset, batch_size=1, shuffle=False,
+            num_workers=2, collate_fn=sdd_seq_collate)
+
+        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.n_steps = self.cfg.diffusion.steps
+        self.betas = self._make_beta_schedule(
+            self.cfg.diffusion.beta_schedule, self.n_steps,
+            self.cfg.diffusion.beta_start, 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)
+
+        self.model = CoreDenoisingModel(past_len=self.cfg.past_frames).cuda()
+        ckpt_path = self.cfg.pretrained_core_denoising_model
+        if not os.path.isfile(ckpt_path):
+            raise FileNotFoundError(f'Missing pretrained denoiser: {ckpt_path}')
+        self.model.load_state_dict(torch.load(ckpt_path, map_location='cpu')['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
+        if self.use_graph:
+            from models.future_interaction_graph_v6 import FutureInteractionGraphV6Wrapper
+            self.interaction_graph = FutureInteractionGraphV6Wrapper(
+                num_agents=64, future_steps=self.cfg.future_frames,
+                past_steps=self.cfg.past_frames, past_channels=6,
+                node_dim=128, top_n=5, num_denoise_steps=NUM_Tau).cuda()
+            params += list(self.interaction_graph.parameters())
+
+        self.opt = torch.optim.AdamW(params, lr=config.learning_rate)
+        self.scheduler = torch.optim.lr_scheduler.StepLR(
+            self.opt, step_size=self.cfg.decay_step, gamma=self.cfg.decay_gamma)
+
+        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_param(self.model, 'Core Denoiser')
+        self._print_param(self.model_initializer, 'Initializer')
+        if self.interaction_graph:
+            self._print_param(self.interaction_graph, 'Graph')
+
+        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_param(self, m, name):
+        t = sum(p.numel() for p in m.parameters())
+        tr = sum(p.numel() for p in m.parameters() if p.requires_grad)
+        print_log(f'[{name}] {tr}/{t}', self.log)
+
+    def _make_beta_schedule(self, schedule, n, start, end):
+        if schedule == 'linear': return torch.linspace(start, end, n)
+        return torch.linspace(start, end, n)
+
+    def _extract(self, a, t, x):
+        out = torch.gather(a, 0, t.to(a.device))
+        return out.reshape(t.shape[0], *([1] * (len(x.shape) - 1)))
+
+    def p_sample_accelerate(self, x, mask, cur_y, t, sigma=None):
+        t_tensor = torch.tensor([int(t)]).cuda()
+        eps_factor = ((1 - self._extract(self.alphas, t_tensor, cur_y))
+                      / self._extract(self.one_minus_alphas_bar_sqrt, t_tensor, cur_y))
+        beta = self._extract(self.betas, t_tensor.repeat(x.shape[0]), cur_y)
+        eps_theta = self.model.generate_accelerate(cur_y, beta, x, mask)
+
+        if self.interaction_graph is not None:
+            abs_t = self._extract(self.alphas_bar_sqrt, t_tensor, cur_y)
+            am1_t = self._extract(self.one_minus_alphas_bar_sqrt, t_tensor, cur_y)
+            y0_hat = (cur_y - am1_t * eps_theta) / abs_t
+            delta = self.interaction_graph(
+                y0_hat, x, int(t), sigma=sigma, A_override=x.size(0))
+            eps_theta = eps_theta - (abs_t / am1_t) * delta
+
+        mean = (1 / self._extract(self.alphas, t_tensor, cur_y).sqrt()) \
+               * (cur_y - eps_factor * eps_theta)
+        z = torch.randn_like(cur_y)
+        sigma_t = self._extract(self.betas, t_tensor, 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)
+
+    def data_preprocess(self, data):
+        pre = data['pre_motion_3D'].cuda()
+        fut = data['fut_motion_3D'].cuda()
+        A = pre.size(1)
+        initial_pos = pre[:, :, -1:]
+        past_abs = ((pre - self.traj_mean) / self.traj_scale).contiguous().view(-1, self.cfg.past_frames, 2)
+        past_rel = ((pre - initial_pos) / self.traj_scale).contiguous().view(-1, self.cfg.past_frames, 2)
+        past_vel = torch.cat([past_rel[:, 1:] - past_rel[:, :-1],
+                              torch.zeros_like(past_rel[:, -1:])], dim=1)
+        past_traj = torch.cat([past_abs, past_rel, past_vel], dim=-1)
+        fut_traj  = ((fut - initial_pos) / self.traj_scale).contiguous().view(-1, self.cfg.future_frames, 2)
+        mask = torch.ones(A, A).cuda()
+        return A, mask, past_traj, fut_traj
+
+    def fit(self):
+        for epoch in range(self.cfg.num_epochs):
+            lt, ld, lu = self._train_epoch(epoch)
+            print_log(f'[{time.strftime("%Y-%m-%d %H:%M:%S")}] Epoch: {epoch}\t'
+                      f'Loss: {lt:.6f}\tDist: {ld:.6f}\tUnc: {lu:.6f}', self.log)
+            self.tb.add_scalar('train/loss', lt, epoch)
+            self.tb.add_scalar('train/loss_dist', ld, epoch)
+
+            if (epoch + 1) % self.cfg.test_interval == 0:
+                perf, n = self._test_epoch()
+                # MID protocol: scale normalized ADE/FDE by 50 to report in pixels
+                ade_px = perf['ADE'] / n * 50.0
+                fde_px = perf['FDE'] / n * 50.0
+                print_log(f'Epoch {epoch} Best Of 20: ADE: {ade_px:.4f} FDE: {fde_px:.4f}', self.log)
+                self.tb.add_scalar('val/ADE_px', ade_px, epoch)
+                self.tb.add_scalar('val/FDE_px', fde_px, epoch)
+
+                cp = {'model_initializer_dict': self.model_initializer.state_dict()}
+                if self.interaction_graph:
+                    cp['interaction_graph_dict'] = self.interaction_graph.state_dict()
+                torch.save(cp, self.cfg.model_path % (epoch + 1))
+            self.scheduler.step()
+        self.tb.flush(); self.tb.close()
+
+    def _train_epoch(self, epoch):
+        self.model.train(); self.model_initializer.train()
+        if self.interaction_graph: self.interaction_graph.train()
+        lt, ld, lu, cnt = 0, 0, 0, 0
+        self.opt.zero_grad()
+        for i, data in enumerate(self.train_loader):
+            A, mask, past, fut = self.data_preprocess(data)
+            sp, me, ve = self.model_initializer(past, mask)
+            ve = ve.clamp(min=-5, max=5)
+            sp = torch.exp(ve / 2)[..., None, None] * sp \
+                 / (sp.std(dim=1).mean(dim=(1, 2))[:, None, None, None] + 1e-6)
+            loc = sp + me[:, None]
+            sigma_in = ve if self.use_v6_graph else None
+            gen = self.p_sample_loop_accelerate(past, mask, loc, sigma=sigma_in)
+            loss_d = ((gen - fut.unsqueeze(1)).norm(p=2, dim=-1)
+                      * self.temporal_reweight).mean(dim=-1).min(dim=1)[0].mean()
+            loss_u = (torch.exp(-ve)
+                      * (gen - fut.unsqueeze(1)).norm(p=2, dim=-1).mean(dim=(1, 2))
+                      + ve).mean()
+            loss = loss_d * 50 + loss_u
+            (loss / self.grad_accum).backward()
+            if (i + 1) % self.grad_accum == 0:
+                params = list(self.model_initializer.parameters())
+                if self.interaction_graph: params += list(self.interaction_graph.parameters())
+                nn.utils.clip_grad_norm_(params, 1.0)
+                self.opt.step(); self.opt.zero_grad()
+            lt += loss.item(); ld += loss_d.item() * 50; lu += loss_u.item(); cnt += 1
+            self.global_step += 1
+        self.opt.step(); self.opt.zero_grad()
+        return lt / cnt, ld / cnt, lu / cnt
+
+    def _test_epoch(self):
+        """MID-style SDD protocol:
+        per-pedestrian full-horizon ADE (mean L2 over 12 future frames) and
+        FDE (L2 at final frame), best_of_20 per pedestrian, then average
+        across all evaluated pedestrians. Coordinates are already in
+        MID's ÷50 mean-centered space, so the final ADE/FDE is multiplied
+        by 50 to report in pixels.
+        Each scene in the SDD dataloader corresponds to one target
+        pedestrian (index 0) + its neighbors; we evaluate only the target
+        per scene so each pedestrian is counted exactly once (matches
+        MID's get_timesteps_data qualification intent).
+        """
+        T = self.cfg.future_frames
+        perf = {'ADE': 0.0, 'FDE': 0.0}
+        n = 0
+        np.random.seed(0); random.seed(0)
+        torch.manual_seed(0); torch.cuda.manual_seed_all(0)
+        self.model_initializer.eval()
+        if self.interaction_graph: self.interaction_graph.eval()
+        with torch.no_grad():
+            for data in self.test_loader:
+                A, mask, past, fut = self.data_preprocess(data)
+                sp, me, ve = self.model_initializer(past, mask)
+                ve = ve.clamp(min=-5, max=5)
+                sp = torch.exp(ve / 2)[..., None, None] * sp \
+                     / (sp.std(dim=1).mean(dim=(1, 2))[:, None, None, None] + 1e-6)
+                loc = sp + me[:, None]
+                sigma_in = ve if self.use_v6_graph else None
+                pred = self.p_sample_loop_accelerate(past, mask, loc, sigma=sigma_in)
+                # MID protocol: only target agent (index 0) per scene
+                pred_0 = pred[0:1]              # [1, 20, T, 2]
+                fut_0  = fut[0:1]               # [1, T, 2]
+                dist = torch.norm(fut_0.unsqueeze(1) - pred_0, dim=-1) * self.traj_scale  # [1, 20, T]
+                # best_of_20 per pedestrian, then full-horizon ADE / final FDE
+                ade_per_ped = dist.mean(dim=-1).min(dim=-1)[0]       # [1]
+                fde_per_ped = dist[:, :, -1].min(dim=-1)[0]           # [1]
+                perf['ADE'] += ade_per_ped.sum().item()
+                perf['FDE'] += fde_per_ped.sum().item()
+                n += 1
+        return perf, n

LED/trainer/train_sdd_pretrain.py

@@ -0,0 +1,137 @@
+"""
+Stage-1 pretraining for LED on SDD.
+Variable-A scenes with batch_size=1 + gradient accumulation.
+past_frames=8, future_frames=12. Pixel coordinates.
+"""
+
+import os, time, torch, 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_sdd import SDDDataset, sdd_seq_collate
+from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel
+
+
+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.grad_accum = getattr(config, 'grad_accum', 32)
+
+        train_dset = SDDDataset(obs_len=self.cfg.past_frames,
+                                pred_len=self.cfg.future_frames, split='train')
+        test_dset  = SDDDataset(obs_len=self.cfg.past_frames,
+                                pred_len=self.cfg.future_frames, split='test')
+        self.train_loader = DataLoader(train_dset, batch_size=1, shuffle=True,
+            num_workers=2, collate_fn=sdd_seq_collate)
+        self.val_loader = DataLoader(test_dset, batch_size=1, shuffle=False,
+            num_workers=2, collate_fn=sdd_seq_collate)
+
+        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.n_steps = self.cfg.diffusion.steps
+        self.betas = self._make_beta_schedule(
+            self.cfg.diffusion.beta_schedule, self.n_steps,
+            self.cfg.diffusion.beta_start, 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)
+
+        self.model = CoreDenoisingModel(past_len=self.cfg.past_frames).cuda()
+
+        pre_lr = float(self.cfg.pretrain['lr'])
+        self.pre_epochs = int(self.cfg.pretrain['num_epochs'])
+        self.opt = torch.optim.AdamW(self.model.parameters(), lr=pre_lr)
+        self.scheduler = torch.optim.lr_scheduler.StepLR(
+            self.opt, step_size=int(self.cfg.pretrain.get('decay_step', 30)),
+            gamma=float(self.cfg.pretrain.get('decay_gamma', 0.5)))
+
+        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.ckpt_path = self.cfg.pretrained_core_denoising_model
+        total = sum(p.numel() for p in self.model.parameters())
+        print_log(f'Core Denoiser params: {total:,}', self.log)
+
+    def _make_beta_schedule(self, schedule, n, start, end):
+        if schedule == 'linear': return torch.linspace(start, end, n)
+        elif schedule == 'quad': return torch.linspace(start**0.5, end**0.5, n)**2
+        return torch.linspace(start, end, n)
+
+    def _extract(self, a, t, x):
+        out = torch.gather(a, 0, t.to(a.device))
+        return out.reshape(t.shape[0], *([1] * (len(x.shape) - 1)))
+
+    def data_preprocess(self, data):
+        pre = data['pre_motion_3D'].cuda()       # [1, A, 8, 2]
+        fut = data['fut_motion_3D'].cuda()       # [1, A, 12, 2]
+        A = pre.size(1)
+        initial_pos = pre[:, :, -1:]              # [1, A, 1, 2]
+        past_abs = ((pre - self.traj_mean) / self.traj_scale).contiguous().view(-1, self.cfg.past_frames, 2)
+        past_rel = ((pre - initial_pos) / self.traj_scale).contiguous().view(-1, self.cfg.past_frames, 2)
+        past_vel = torch.cat([past_rel[:, 1:] - past_rel[:, :-1],
+                              torch.zeros_like(past_rel[:, -1:])], dim=1)
+        past_traj = torch.cat([past_abs, past_rel, past_vel], dim=-1)  # [A, 8, 6]
+        fut_traj  = ((fut - initial_pos) / self.traj_scale).contiguous().view(-1, self.cfg.future_frames, 2)
+        mask = torch.ones(A, A).cuda()
+        return A, mask, past_traj, fut_traj
+
+    def noise_estimation_loss(self, x, y_0, mask):
+        B = x.shape[0]
+        t = torch.randint(0, self.n_steps, size=(B // 2 + 1,)).to(x.device)
+        t = torch.cat([t, self.n_steps - t - 1], dim=0)[:B]
+        a   = self._extract(self.alphas_bar_sqrt, t, y_0)
+        beta = self._extract(self.betas, t, y_0)
+        am1 = self._extract(self.one_minus_alphas_bar_sqrt, t, y_0)
+        e   = torch.randn_like(y_0)
+        y   = y_0 * a + e * am1
+        out = self.model(y, beta, x, mask)
+        return (e - out).square().mean()
+
+    def fit(self):
+        best_val = float('inf')
+        for epoch in range(self.pre_epochs):
+            self.model.train()
+            loss_sum, n = 0.0, 0
+            self.opt.zero_grad()
+            for i, data in enumerate(self.train_loader):
+                _, mask, past, fut = self.data_preprocess(data)
+                loss = self.noise_estimation_loss(past, fut, mask)
+                (loss / self.grad_accum).backward()
+                if (i + 1) % self.grad_accum == 0:
+                    nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
+                    self.opt.step(); self.opt.zero_grad()
+                loss_sum += loss.item(); n += 1
+                self.tb.add_scalar('pretrain_step/loss', loss.item(), self.global_step)
+                self.global_step += 1
+            self.opt.step(); self.opt.zero_grad()
+
+            train_loss = loss_sum / max(1, n)
+            print_log(f'[{time.strftime("%Y-%m-%d %H:%M:%S")}] Pretrain Epoch {epoch}  train_mse={train_loss:.6f}', self.log)
+            self.tb.add_scalar('pretrain_epoch/train_loss', train_loss, epoch)
+
+            if (epoch + 1) % 5 == 0:
+                self.model.eval()
+                val_sum, val_n = 0.0, 0
+                with torch.no_grad():
+                    for data in self.val_loader:
+                        _, mask, past, fut = self.data_preprocess(data)
+                        val_sum += self.noise_estimation_loss(past, fut, mask).item()
+                        val_n += 1
+                val_loss = val_sum / max(1, val_n)
+                print_log(f'  val_mse={val_loss:.6f}', self.log)
+                self.tb.add_scalar('pretrain_epoch/val_loss', val_loss, epoch)
+                if val_loss < best_val:
+                    best_val = val_loss
+                    os.makedirs(os.path.dirname(self.ckpt_path), exist_ok=True)
+                    torch.save({'model_dict': self.model.state_dict(), 'epoch': epoch}, self.ckpt_path)
+                    print_log(f'  -> saved {self.ckpt_path}', self.log)
+
+            self.scheduler.step()
+        self.tb.flush(); self.tb.close()

LED/trainer/train_sport_led.py

@@ -0,0 +1,352 @@
+"""
+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

LED/trainer/train_sport_pretrain.py

@@ -0,0 +1,210 @@
+"""
+Stage-1 pretraining for LED on sport datasets (soccer / football).
+
+This is the training stage that LED's public README leaves on its TODO list:
+trains the core 100-step DDPM denoiser (TransformerDenoisingModel) from
+scratch on a sport dataset using the standard eps-MSE loss. The resulting
+checkpoint is saved to `cfg.pretrained_core_denoising_model` so that
+train_sport_led.py (stage 2) can load it as a frozen refiner.
+
+Optimizes: self.model (core denoiser) only.
+Loss: noise_estimation_loss — predict eps at a uniformly-sampled diffusion
+      step t ∈ [0, n_steps) and regress via MSE.
+"""
+
+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_diffusion import TransformerDenoisingModel as CoreDenoisingModel
+
+
+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)
+
+        # ------------------------- 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',
+        )
+
+        pre_bs = self.cfg.pretrain['train_batch_size']
+        self.train_loader = DataLoader(
+            train_dset, batch_size=pre_bs, shuffle=True,
+            num_workers=4, collate_fn=sport_seq_collate, pin_memory=True)
+        self.val_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)
+
+        # ------------------------- model (the only thing being trained) -------------------------
+        self.model = CoreDenoisingModel().cuda()
+
+        pre_lr          = float(self.cfg.pretrain['lr'])
+        pre_decay_step  = int(self.cfg.pretrain['decay_step'])
+        pre_decay_gamma = float(self.cfg.pretrain['decay_gamma'])
+        self.pre_epochs = int(self.cfg.pretrain['num_epochs'])
+
+        self.opt = torch.optim.AdamW(self.model.parameters(), lr=pre_lr)
+        self.scheduler = torch.optim.lr_scheduler.StepLR(
+            self.opt, step_size=pre_decay_step, gamma=pre_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.ckpt_path = self.cfg.pretrained_core_denoising_model
+
+    def print_model_param(self, model: nn.Module, name: str = 'Model'):
+        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)
+
+    # ------------------------------------------------------------------
+    # Data preprocess — identical math to LED NBA trainer but with
+    # num_agents parameterized so the block-diagonal mask is A×A, not 11×11.
+    # ------------------------------------------------------------------
+    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
+
+    # ------------------------------------------------------------------
+    # Pretraining loss (LED's noise_estimation_loss, unchanged)
+    # ------------------------------------------------------------------
+    def noise_estimation_loss(self, x, y_0, mask):
+        batch_size = x.shape[0]
+        t = torch.randint(0, self.n_steps, size=(batch_size // 2 + 1,)).to(x.device)
+        t = torch.cat([t, self.n_steps - t - 1], dim=0)[:batch_size]
+        a   = self.extract(self.alphas_bar_sqrt, t, y_0)
+        beta = self.extract(self.betas, t, y_0)
+        am1 = self.extract(self.one_minus_alphas_bar_sqrt, t, y_0)
+        e   = torch.randn_like(y_0)
+        y   = y_0 * a + e * am1
+        out = self.model(y, beta, x, mask)
+        return (e - out).square().mean()
+
+    # ------------------------------------------------------------------
+    # Training loop
+    # ------------------------------------------------------------------
+    def fit(self):
+        best_val = float('inf')
+        for epoch in range(self.pre_epochs):
+            self.model.train()
+            loss_sum, n_batches = 0.0, 0
+            for data in self.train_loader:
+                _, mask, past, fut = self.data_preprocess(data)
+                loss = self.noise_estimation_loss(past, fut, mask)
+                self.opt.zero_grad()
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
+                self.opt.step()
+                loss_sum += float(loss.item())
+                n_batches += 1
+                self.tb.add_scalar('pretrain_step/loss', loss.item(), self.global_step)
+                self.global_step += 1
+
+            train_loss = loss_sum / max(1, n_batches)
+            print_log(f'[{time.strftime("%Y-%m-%d %H:%M:%S")}] Pretrain Epoch {epoch}  train_mse={train_loss:.6f}', self.log)
+            self.tb.add_scalar('pretrain_epoch/train_loss', train_loss, epoch)
+            self.tb.add_scalar('pretrain_epoch/lr', self.opt.param_groups[0]['lr'], epoch)
+
+            # Quick val MSE to monitor
+            self.model.eval()
+            val_sum, val_n = 0.0, 0
+            with torch.no_grad():
+                for data in self.val_loader:
+                    _, mask, past, fut = self.data_preprocess(data)
+                    val_sum += float(self.noise_estimation_loss(past, fut, mask).item())
+                    val_n += 1
+            val_loss = val_sum / max(1, val_n)
+            print_log(f'                            val_mse={val_loss:.6f}', self.log)
+            self.tb.add_scalar('pretrain_epoch/val_loss', val_loss, epoch)
+
+            if val_loss < best_val:
+                best_val = val_loss
+                os.makedirs(os.path.dirname(self.ckpt_path), exist_ok=True)
+                torch.save({'model_dict': self.model.state_dict(), 'epoch': epoch}, self.ckpt_path)
+                print_log(f'                            -> saved {self.ckpt_path} (best val {best_val:.6f})', self.log)
+
+            self.scheduler.step()
+
+        self.tb.flush()
+        self.tb.close()

LED/utils/config.py

@@ -0,0 +1,60 @@
+import yaml
+import os
+import os.path as osp
+import glob
+import numpy as np
+from easydict import EasyDict
+from utils.utils import recreate_dirs
+
+
+class Config:
+
+    def __init__(self, cfg_id, info):
+        self.id = cfg_id
+        cfg_path = 'cfg/**/%s.yml' % cfg_id
+        files = glob.glob(cfg_path, recursive=True)
+        assert (len(files) == 1), 'YAML file [{}] does not exist!'.format(cfg_id)
+        self.yml_dict = EasyDict(yaml.safe_load(open(files[0], 'r')))
+
+        self.results_root_dir = os.path.expanduser(self.yml_dict['results_root_dir'])
+        # results dirs
+
+        self.cfg_dir = '%s/%s/%s' % (self.results_root_dir, cfg_id, info)
+        self.model_dir = '%s/models' % self.cfg_dir
+        self.log_dir = '%s/log' % self.cfg_dir
+        self.model_path = os.path.join(self.model_dir, 'model_%04d.p')
+        os.makedirs(self.model_dir, exist_ok=True)
+        os.makedirs(self.log_dir, exist_ok=True)
+
+    def get_last_epoch(self):
+        model_files = glob.glob(os.path.join(self.model_dir, 'model_*.p'))
+        if len(model_files) == 0:
+            return None
+        else:
+            model_file = osp.basename(model_files[0])
+            epoch = int(osp.splitext(model_file)[0].split('model_')[-1])
+            return epoch            
+
+    def __getattribute__(self, name):
+        yml_dict = super().__getattribute__('yml_dict')
+        if name in yml_dict:
+            return yml_dict[name]
+        else:
+            return super().__getattribute__(name)
+
+    def __setattr__(self, name, value):
+        try:
+            yml_dict = super().__getattribute__('yml_dict')
+        except AttributeError:
+            return super().__setattr__(name, value)
+        if name in yml_dict:
+            yml_dict[name] = value
+        else:
+            return super().__setattr__(name, value)
+
+    def get(self, name, default=None):
+        if hasattr(self, name):
+            return getattr(self, name)
+        else:
+            return default
+            

LED/utils/utils.py

@@ -0,0 +1,336 @@
+"""
+Code borrowed from Xinshuo_PyToolbox: https://github.com/xinshuoweng/Xinshuo_PyToolbox
+"""
+
+import os
+import shutil
+import torch
+import numpy as np
+import random
+import time
+import copy
+import glob, glob2
+from torch import nn
+
+class AverageMeter(object):     
+  """Computes and stores the average and current value"""    
+  def __init__(self):   
+    self.reset()
+  
+  def reset(self):
+    self.val = 0
+    self.avg = 0
+    self.sum = 0
+    self.count = 0    
+    self.list = list()
+  
+  def update(self, val, n=1): 
+    self.val = val    
+    self.sum += val * n     
+    self.count += n
+    self.avg = self.sum / self.count  
+    self.list.append(val)
+
+
+def isnparray(nparray_test):
+	return isinstance(nparray_test, np.ndarray)
+
+
+def isinteger(integer_test):
+	if isnparray(integer_test): return False
+	try: return isinstance(integer_test, int) or int(integer_test) == integer_test
+	except ValueError: return False
+	except TypeError: return False
+
+
+def isfloat(float_test):
+	return isinstance(float_test, float)
+
+
+def isscalar(scalar_test):
+	try: return isinteger(scalar_test) or isfloat(scalar_test)
+	except TypeError: return False
+
+
+def islogical(logical_test):
+	return isinstance(logical_test, bool)
+
+    
+def isstring(string_test):
+    return isinstance(string_test, str)
+
+
+def islist(list_test):
+	return isinstance(list_test, list)
+
+
+def convert_secs2time(seconds):
+    '''
+    format second to human readable way
+    '''
+    assert isscalar(seconds), 'input should be a scalar to represent number of seconds'
+    m, s = divmod(int(seconds), 60)
+    h, m = divmod(m, 60)
+    return '[%d:%02d:%02d]' % (h, m, s)
+
+
+def get_timestring():
+    return time.strftime('%Y%m%d_%Hh%Mm%Ss')
+
+
+def recreate_dirs(*dirs):
+    for d in dirs:
+        if os.path.exists(d):
+            shutil.rmtree(d)
+        os.makedirs(d)
+
+
+def is_path_valid(pathname):
+	try:  
+		if not isstring(pathname) or not pathname: return False
+	except TypeError: return False
+	else: return True
+
+
+def is_path_creatable(pathname):
+	'''
+	if any previous level of parent folder exists, returns true
+	'''
+	if not is_path_valid(pathname): return False
+	pathname = os.path.normpath(pathname)
+	pathname = os.path.dirname(os.path.abspath(pathname))
+
+	# recursively to find the previous level of parent folder existing
+	while not is_path_exists(pathname):     
+		pathname_new = os.path.dirname(os.path.abspath(pathname))
+		if pathname_new == pathname: return False
+		pathname = pathname_new
+	return os.access(pathname, os.W_OK)
+
+
+def is_path_exists(pathname):
+	try: return is_path_valid(pathname) and os.path.exists(pathname)
+	except OSError: return False
+
+
+def is_path_exists_or_creatable(pathname):
+	try: return is_path_exists(pathname) or is_path_creatable(pathname)
+	except OSError: return False
+
+
+def isfile(pathname):
+	if is_path_valid(pathname):
+		pathname = os.path.normpath(pathname)
+		name = os.path.splitext(os.path.basename(pathname))[0]
+		ext = os.path.splitext(pathname)[1]
+		return len(name) > 0 and len(ext) > 0
+	else: return False
+
+
+def isfolder(pathname):
+	'''
+	if '.' exists in the subfolder, the function still justifies it as a folder. e.g., /mnt/dome/adhoc_0.5x/abc is a folder
+	if '.' exists after all slashes, the function will not justify is as a folder. e.g., /mnt/dome/adhoc_0.5x is NOT a folder
+	'''
+	if is_path_valid(pathname):
+		pathname = os.path.normpath(pathname)
+		if pathname == './': return True
+		name = os.path.splitext(os.path.basename(pathname))[0]
+		ext = os.path.splitext(pathname)[1]
+		return len(name) > 0 and len(ext) == 0
+	else: return False
+
+
+def mkdir_if_missing(input_path):
+    folder = input_path if isfolder(input_path) else os.path.dirname(input_path)
+    os.makedirs(folder, exist_ok=True)
+
+
+def safe_list(input_data, warning=True, debug=True):
+	'''
+	copy a list to the buffer for use
+	parameters:
+		input_data:		a list
+	outputs:
+		safe_data:		a copy of input data
+	'''
+	if debug: assert islist(input_data), 'the input data is not a list'
+	safe_data = copy.copy(input_data)
+	return safe_data
+
+
+def safe_path(input_path, warning=True, debug=True):
+    '''
+    convert path to a valid OS format, e.g., empty string '' to '.', remove redundant '/' at the end from 'aa/' to 'aa'
+    parameters:
+    	input_path:		a string
+    outputs:
+    	safe_data:		a valid path in OS format
+    '''
+    if debug: assert isstring(input_path), 'path is not a string: %s' % input_path
+    safe_data = copy.copy(input_path)
+    safe_data = os.path.normpath(safe_data)
+    return safe_data
+
+
+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)
+
+
+def initialize_weights(modules):
+    for m in modules:
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            if m.bias is not None: nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.BatchNorm2d):
+            nn.init.constant_(m.weight, 1)
+            if m.bias is not None: nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, 0, 0.01)
+            if m.bias is not None: nn.init.constant_(m.bias, 0)
+
+
+def print_log(print_str, log, same_line=False, display=True):
+	'''
+	print a string to a log file
+
+	parameters:
+		print_str:          a string to print
+		log:                a opened file to save the log
+		same_line:          True if we want to print the string without a new next line
+		display:            False if we want to disable to print the string onto the terminal
+	'''
+	if display:
+		if same_line: print('{}'.format(print_str), end='')
+		else: print('{}'.format(print_str))
+
+	if same_line: log.write('{}'.format(print_str))
+	else: log.write('{}\n'.format(print_str))
+	log.flush()
+
+
+def find_unique_common_from_lists(input_list1, input_list2, warning=True, debug=True):
+	'''
+	find common items from 2 lists, the returned elements are unique. repetitive items will be ignored
+	if the common items in two elements are not in the same order, the outputs follows the order in the first list
+
+	parameters:
+		input_list1, input_list2:		two input lists
+
+	outputs:
+		list_common:	a list of elements existing both in list_src1 and list_src2	
+		index_list1:	a list of index that list 1 has common items
+		index_list2:	a list of index that list 2 has common items
+	'''
+	input_list1 = safe_list(input_list1, warning=warning, debug=debug)
+	input_list2 = safe_list(input_list2, warning=warning, debug=debug)
+
+	common_list = list(set(input_list1).intersection(input_list2))
+	
+	# find index
+	index_list1 = []
+	for index in range(len(input_list1)):
+		item = input_list1[index]
+		if item in common_list:
+			index_list1.append(index)
+
+	index_list2 = []
+	for index in range(len(input_list2)):
+		item = input_list2[index]
+		if item in common_list:
+			index_list2.append(index)
+
+	return common_list, index_list1, index_list2
+
+
+def load_txt_file(file_path, debug=True):
+    '''
+    load data or string from text file
+    '''
+    file_path = safe_path(file_path)
+    if debug: assert is_path_exists(file_path), 'text file is not existing at path: %s!' % file_path
+    with open(file_path, 'r') as file: data = file.read().splitlines()
+    num_lines = len(data)
+    file.close()
+    return data, num_lines
+
+
+def load_list_from_folder(folder_path, ext_filter=None, depth=1, recursive=False, sort=True, save_path=None, debug=True):
+    '''
+    load a list of files or folders from a system path
+
+    parameters:
+        folder_path:    root to search 
+        ext_filter:     a string to represent the extension of files interested
+        depth:          maximum depth of folder to search, when it's None, all levels of folders will be searched
+        recursive:      False: only return current level
+                        True: return all levels till to the input depth
+
+    outputs:
+        fulllist:       a list of elements
+        num_elem:       number of the elements
+    '''
+    folder_path = safe_path(folder_path)
+    if debug: assert isfolder(folder_path), 'input folder path is not correct: %s' % folder_path
+    if not is_path_exists(folder_path): 
+        print('the input folder does not exist\n')
+        return [], 0
+    if debug:
+        assert islogical(recursive), 'recursive should be a logical variable: {}'.format(recursive)
+        assert depth is None or (isinteger(depth) and depth >= 1), 'input depth is not correct {}'.format(depth)
+        assert ext_filter is None or (islist(ext_filter) and all(isstring(ext_tmp) for ext_tmp in ext_filter)) or isstring(ext_filter), 'extension filter is not correct'
+    if isstring(ext_filter): ext_filter = [ext_filter]                               # convert to a list
+    # zxc
+
+    fulllist = list()
+    if depth is None:        # find all files recursively
+        recursive = True
+        wildcard_prefix = '**'
+        if ext_filter is not None:
+            for ext_tmp in ext_filter:
+                # wildcard = os.path.join(wildcard_prefix, '*' + string2ext_filter(ext_tmp))
+                wildcard = os.path.join(wildcard_prefix, '*' + ext_tmp)
+                curlist = glob2.glob(os.path.join(folder_path, wildcard))
+                if sort: curlist = sorted(curlist)
+                fulllist += curlist
+        else:
+            wildcard = wildcard_prefix
+            curlist = glob2.glob(os.path.join(folder_path, wildcard))
+            if sort: curlist = sorted(curlist)
+            fulllist += curlist
+    else:                    # find files based on depth and recursive flag
+        wildcard_prefix = '*'
+        for index in range(depth-1): wildcard_prefix = os.path.join(wildcard_prefix, '*')
+        if ext_filter is not None:
+            for ext_tmp in ext_filter:
+                # wildcard = wildcard_prefix + string2ext_filter(ext_tmp)
+                wildcard = wildcard_prefix + ext_tmp
+                curlist = glob.glob(os.path.join(folder_path, wildcard))
+                if sort: curlist = sorted(curlist)
+                fulllist += curlist
+            # zxc
+        else:
+            wildcard = wildcard_prefix
+            curlist = glob.glob(os.path.join(folder_path, wildcard))
+            # print(curlist)
+            if sort: curlist = sorted(curlist)
+            fulllist += curlist
+        if recursive and depth > 1:
+            newlist, _ = load_list_from_folder(folder_path=folder_path, ext_filter=ext_filter, depth=depth-1, recursive=True)
+            fulllist += newlist
+
+    fulllist = [os.path.normpath(path_tmp) for path_tmp in fulllist]
+    num_elem = len(fulllist)
+
+    # save list to a path
+    if save_path is not None:
+        save_path = safe_path(save_path)
+        if debug: assert is_path_exists_or_creatable(save_path), 'the file cannot be created'
+        with open(save_path, 'w') as file:
+            for item in fulllist: file.write('%s\n' % item)
+        file.close()
+
+    return fulllist, num_elem

LED/viz_denoising_process.py

@@ -0,0 +1,329 @@
+"""
+Visualize the full LED denoising process with/without uncertainty.
+
+Generates side-by-side plots showing:
+- Left: Without uncertainty (graph only)
+- Right: With uncertainty (graph + sigma coloring)
+
+For each sample:
+- Past trajectories (black lines)
+- GT future (green dashed)
+- Denoising steps τ=4,3,2,1,0 with progressively refined predictions
+- For sigma version: trajectory color indicates uncertainty (red=high, blue=low)
+"""
+
+import os
+import sys
+import torch
+import random
+import numpy as np
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import matplotlib.cm as cm
+from matplotlib.colors import Normalize
+
+sys.path.insert(0, os.path.dirname(__file__))
+
+from utils.config import Config
+from data.dataloader_nba import NBADataset, seq_collate
+from torch.utils.data import DataLoader
+from models.model_led_initializer import LEDInitializer as InitializationModel
+from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel
+from models.future_interaction_graph_v6 import FutureInteractionGraphV6Wrapper
+
+NUM_Tau = 5
+TRAJ_SCALE = 94.0 / 28.0  # to convert back to feet
+
+
+def load_models(ckpt_path, use_sigma=True, use_v6_graph=True, edge_mode='relpos_only',
+                top_n=5, device='cuda'):
+    """Load LED models from checkpoint."""
+    cfg = Config('led_augment', 'viz')
+
+    model = CoreDenoisingModel().to(device)
+    model_cp = torch.load(cfg.pretrained_core_denoising_model, map_location='cpu')
+    model.load_state_dict(model_cp['model_dict'])
+    model.eval()
+
+    model_init = InitializationModel(t_h=10, d_h=6, t_f=20, d_f=2, k_pred=20).to(device)
+
+    graph = FutureInteractionGraphV6Wrapper(
+        num_agents=11, future_steps=20, past_steps=10,
+        past_channels=6, node_dim=128, top_n=top_n,
+        num_denoise_steps=NUM_Tau, edge_mode=edge_mode,
+    ).to(device)
+
+    ckpt = torch.load(ckpt_path, map_location='cpu')
+    model_init.load_state_dict(ckpt['model_initializer_dict'])
+    graph.load_state_dict(ckpt['interaction_graph_dict'])
+    model_init.eval()
+    graph.eval()
+
+    return cfg, model, model_init, graph
+
+
+def make_beta_schedule(n_timesteps=100, start=1e-5, end=1e-2):
+    return torch.linspace(start, end, n_timesteps).cuda()
+
+
+def denoise_with_intermediates(model, graph, past_traj, traj_mask, loc,
+                                betas, alphas_prod, alphas_bar_sqrt,
+                                one_minus_alphas_bar_sqrt, alphas,
+                                use_sigma=False, sigma=None):
+    """Run denoising and return intermediate predictions at each step."""
+    intermediates = []  # list of (y0_hat, sigma_val) at each step
+
+    cur_y = loc[:, :10]
+    for i in reversed(range(NUM_Tau)):
+        t = torch.tensor([i]).cuda()
+        eps_factor = ((1 - alphas[i]) / one_minus_alphas_bar_sqrt[i])
+        beta = betas[i].repeat(past_traj.shape[0]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+
+        eps_theta = model.generate_accelerate(cur_y, beta.squeeze(-1).squeeze(-1), past_traj, traj_mask)
+
+        alpha_bar_sqrt_t = alphas_bar_sqrt[i]
+        one_minus_abs_t = one_minus_alphas_bar_sqrt[i]
+        y0_hat = (cur_y - one_minus_abs_t * eps_theta) / alpha_bar_sqrt_t
+
+        sigma_input = sigma if use_sigma else None
+        delta = graph(y0_hat, past_traj, i, sigma=sigma_input)
+        eps_theta = eps_theta + delta
+
+        # Store y0_hat after graph correction
+        y0_corrected = (cur_y - one_minus_abs_t * eps_theta) / alpha_bar_sqrt_t
+        intermediates.append({
+            'step': i,
+            'y0_hat': y0_corrected.detach().cpu(),
+            'sigma': sigma.detach().cpu() if sigma is not None else None,
+        })
+
+        mean = (1 / alphas[i].sqrt()) * (cur_y - eps_factor * eps_theta)
+        z = torch.randn_like(cur_y)
+        sigma_t = betas[i].sqrt()
+        cur_y = mean + sigma_t * z * 0.00001
+
+    # Second half (modes 10-19)
+    cur_y_ = loc[:, 10:]
+    for i in reversed(range(NUM_Tau)):
+        t = torch.tensor([i]).cuda()
+        eps_factor = ((1 - alphas[i]) / one_minus_alphas_bar_sqrt[i])
+        beta = betas[i].repeat(past_traj.shape[0]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+        eps_theta = model.generate_accelerate(cur_y_, beta.squeeze(-1).squeeze(-1), past_traj, traj_mask)
+        alpha_bar_sqrt_t = alphas_bar_sqrt[i]
+        one_minus_abs_t = one_minus_alphas_bar_sqrt[i]
+        y0_hat = (cur_y_ - one_minus_abs_t * eps_theta) / alpha_bar_sqrt_t
+        sigma_input = sigma if use_sigma else None
+        delta = graph(y0_hat, past_traj, i, sigma=sigma_input)
+        eps_theta = eps_theta + delta
+        mean = (1 / alphas[i].sqrt()) * (cur_y_ - eps_factor * eps_theta)
+        z = torch.randn_like(cur_y_)
+        sigma_t = betas[i].sqrt()
+        cur_y_ = mean + sigma_t * z * 0.00001
+
+    final_pred = torch.cat((cur_y_, cur_y), dim=1)
+    return intermediates, final_pred
+
+
+def plot_single_sample(ax, past, gt_fut, intermediates, final_pred,
+                        initial_pos, traj_mean, traj_scale,
+                        sigma_vals=None, title='', show_uncertainty=False):
+    """Plot one sample's denoising process on an axis."""
+    A = 11
+    K = 20
+
+    # Convert past to absolute court positions
+    past_abs = past.reshape(A, 10, 6)[:, :, :2]  # abs_xy channels
+    past_abs = past_abs * traj_scale + traj_mean  # back to court coords
+
+    # GT future
+    gt_abs = gt_fut.reshape(A, 20, 2) * traj_scale + initial_pos.reshape(A, 1, 2)
+
+    # Colors for agents
+    agent_colors = plt.cm.tab10(np.linspace(0, 1, A))
+
+    # Draw court
+    ax.set_xlim(-2, 30)
+    ax.set_ylim(-2, 17)
+    ax.set_aspect('equal')
+    ax.set_facecolor('#2d5016')  # dark green court
+
+    # Draw past trajectories
+    for a in range(A):
+        ax.plot(past_abs[a, :, 0], past_abs[a, :, 1], '-',
+                color='white', alpha=0.5, linewidth=1)
+        ax.plot(past_abs[a, -1, 0], past_abs[a, -1, 1], 'o',
+                color='white', markersize=4)
+
+    # Draw GT future
+    for a in range(A):
+        ax.plot(gt_abs[a, :, 0], gt_abs[a, :, 1], '--',
+                color='lime', alpha=0.4, linewidth=1)
+
+    # Draw denoising steps (from noisy to clean)
+    step_alphas = [0.15, 0.25, 0.35, 0.5, 0.7]
+    for idx, inter in enumerate(intermediates):
+        y0 = inter['y0_hat']  # [B*A, K, T, 2]
+        step = inter['step']
+        alpha = step_alphas[min(idx, len(step_alphas) - 1)]
+
+        # Take best mode (mode 0 for simplicity)
+        y0_mode0 = y0[:A, 0, :, :]  # [A, T, 2]
+        y0_abs = y0_mode0 * traj_scale + initial_pos.reshape(A, 1, 2)
+
+        if show_uncertainty and inter['sigma'] is not None:
+            sigma_a = inter['sigma'][:A, 0].numpy()
+            norm = Normalize(vmin=sigma_a.min(), vmax=sigma_a.max())
+            cmap = cm.coolwarm  # blue=certain, red=uncertain
+
+            for a in range(A):
+                color = cmap(norm(sigma_a[a]))
+                ax.plot(y0_abs[a, :, 0], y0_abs[a, :, 1], '-',
+                        color=color, alpha=alpha, linewidth=1.5)
+        else:
+            for a in range(A):
+                ax.plot(y0_abs[a, :, 0], y0_abs[a, :, 1], '-',
+                        color=agent_colors[a], alpha=alpha, linewidth=1)
+
+    # Draw final prediction (best mode)
+    if final_pred is not None:
+        pred = final_pred[:A]  # [A, K, T, 2]
+        pred_abs = pred.numpy() * traj_scale + initial_pos.reshape(A, 1, 1, 2)
+        gt_exp = np.expand_dims(gt_abs, 1).repeat(pred_abs.shape[1], axis=1)
+        ade_per_mode = np.linalg.norm(pred_abs - gt_exp, axis=-1).mean(axis=-1)  # [A, K]
+        best_modes = ade_per_mode.argmin(axis=-1)  # [A]
+
+        for a in range(A):
+            best = pred_abs[a, best_modes[a]]
+            ax.plot(best[:, 0], best[:, 1], '-',
+                    color=agent_colors[a], alpha=0.9, linewidth=2)
+            ax.plot(best[-1, 0], best[-1, 1], '*',
+                    color=agent_colors[a], markersize=6)
+
+    ax.set_title(title, fontsize=10, color='white')
+    ax.tick_params(colors='gray')
+
+
+def main():
+    device = 'cuda:0'  # CUDA_VISIBLE_DEVICES remaps to 0
+    torch.cuda.set_device(0)
+
+    # Paths
+    ckpt_sigma = '/mnt/jaewoo4tb/srtp/LED/results/led_augment/graph_v6_edge_relpos/models/model_0052.p'
+    ckpt_nosigma = '/mnt/jaewoo4tb/srtp/LED/results/led_augment/graph_v6_nosigma_n3/models/model_0084.p'
+
+    out_dir = '/mnt/jaewoo4tb/srtp/LED/visualizations'
+    os.makedirs(out_dir, exist_ok=True)
+
+    # Load models
+    cfg_s, model_s, init_s, graph_s = load_models(
+        ckpt_sigma, use_sigma=True, edge_mode='relpos_only', top_n=5, device=device)
+    cfg_n, model_n, init_n, graph_n = load_models(
+        ckpt_nosigma, use_sigma=False, edge_mode='full', top_n=3, device=device)
+
+    # Diffusion schedule
+    betas = torch.linspace(1e-5, 1e-2, 100).to(device)
+    alphas = 1 - betas
+    alphas_prod = torch.cumprod(alphas, 0)
+    alphas_bar_sqrt = torch.sqrt(alphas_prod)
+    one_minus_alphas_bar_sqrt = torch.sqrt(1 - alphas_prod)
+
+    traj_mean = torch.FloatTensor(cfg_s.traj_mean).to(device).unsqueeze(0).unsqueeze(0).unsqueeze(0)
+    traj_scale = cfg_s.traj_scale
+
+    # Load test data
+    test_dset = NBADataset(obs_len=10, pred_len=20, training=False)
+    test_loader = DataLoader(test_dset, batch_size=1, shuffle=False, collate_fn=seq_collate)
+
+    # Set seed for reproducibility
+    np.random.seed(42)
+    random.seed(42)
+    torch.manual_seed(42)
+
+    num_samples = 5
+    sample_indices = sorted(random.sample(range(len(test_dset)), num_samples))
+
+    with torch.no_grad():
+        for sample_idx, data in enumerate(test_loader):
+            if sample_idx not in sample_indices:
+                continue
+            if sample_idx > max(sample_indices):
+                break
+
+            batch_size = 1
+            traj_mask = torch.ones(11, 11).to(device)
+            initial_pos = data['pre_motion_3D'].to(device)[:, :, -1:]
+            past_traj_abs = ((data['pre_motion_3D'].to(device) - traj_mean) / traj_scale).view(-1, 10, 2)
+            past_traj_rel = ((data['pre_motion_3D'].to(device) - initial_pos) / traj_scale).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'].to(device) - initial_pos) / traj_scale).view(-1, 20, 2)
+
+            # --- With sigma ---
+            sample_pred_s, mean_s, var_s = init_s(past_traj, traj_mask)
+            sample_pred_s = (torch.exp(var_s / 2)[..., None, None]
+                             * sample_pred_s / sample_pred_s.std(dim=1).mean(dim=(1, 2))[:, None, None, None])
+            loc_s = sample_pred_s + mean_s[:, None]
+
+            intermediates_s, final_s = denoise_with_intermediates(
+                model_s, graph_s, past_traj, traj_mask, loc_s,
+                betas, alphas_prod, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, alphas,
+                use_sigma=True, sigma=var_s)
+
+            # --- Without sigma ---
+            sample_pred_n, mean_n, var_n = init_n(past_traj, traj_mask)
+            sample_pred_n = (torch.exp(var_n / 2)[..., None, None]
+                             * sample_pred_n / sample_pred_n.std(dim=1).mean(dim=(1, 2))[:, None, None, None])
+            loc_n = sample_pred_n + mean_n[:, None]
+
+            intermediates_n, final_n = denoise_with_intermediates(
+                model_n, graph_n, past_traj, traj_mask, loc_n,
+                betas, alphas_prod, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, alphas,
+                use_sigma=False, sigma=None)
+
+            # --- Plot ---
+            fig, axes = plt.subplots(1, 2, figsize=(20, 8))
+            fig.patch.set_facecolor('#1a1a1a')
+
+            init_pos_cpu = initial_pos.cpu().squeeze(0)
+            traj_mean_cpu = traj_mean.cpu().squeeze(0).squeeze(0)
+
+            plot_single_sample(
+                axes[0], past_traj.cpu().numpy(), fut_traj.cpu().numpy(),
+                intermediates_n, final_n.cpu(),
+                init_pos_cpu.numpy(), traj_mean_cpu.numpy(), traj_scale,
+                title=f'Without Uncertainty (sample {sample_idx})')
+
+            plot_single_sample(
+                axes[1], past_traj.cpu().numpy(), fut_traj.cpu().numpy(),
+                intermediates_s, final_s.cpu(),
+                init_pos_cpu.numpy(), traj_mean_cpu.numpy(), traj_scale,
+                sigma_vals=var_s.cpu(), show_uncertainty=True,
+                title=f'With Uncertainty (sample {sample_idx})')
+
+            # Add colorbar for uncertainty
+            sm = plt.cm.ScalarMappable(cmap=cm.coolwarm)
+            sm.set_array([])
+            cbar = fig.colorbar(sm, ax=axes[1], shrink=0.6, pad=0.02)
+            cbar.set_label('Uncertainty (σ)', color='white')
+            cbar.ax.yaxis.set_tick_params(color='white')
+            plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color='white')
+
+            plt.suptitle(f'LED Denoising Process — Sample {sample_idx}\n'
+                         f'Fading lines: τ=4→0 (noisy→clean). '
+                         f'Green dashed: GT. Stars: final endpoints.',
+                         color='white', fontsize=12)
+
+            plt.tight_layout()
+            save_path = os.path.join(out_dir, f'denoising_sample_{sample_idx:04d}.png')
+            plt.savefig(save_path, dpi=150, bbox_inches='tight',
+                        facecolor=fig.get_facecolor())
+            plt.close()
+            print(f'Saved: {save_path}')
+
+    print(f'\nAll visualizations saved to {out_dir}/')
+
+
+if __name__ == '__main__':
+    main()

LED/viz_denoising_steps.py

@@ -0,0 +1,299 @@
+"""
+Visualize LED denoising: initializer → refinement → final, on basketball court.
+
+LED's denoising is a 2-stage process:
+  1. Initializer: produces 20 diverse trajectory modes (noisy, multimodal)
+  2. Leapfrog refinement: 5 DDPM steps that clean up each mode (subtle changes)
+
+We visualize 5 stages:
+  Step 0: Raw initializer output (mean_estimation only, no variance)
+  Step 1: Initializer + variance scaling (diverse modes)
+  Step 2: After τ=4,3 refinement (2 DDPM steps)
+  Step 3: After τ=2,1 refinement (4 DDPM steps)
+  Step 4: Final prediction τ=0 (5 DDPM steps, fully denoised)
+
+For sigma version: trajectory color = uncertainty (green=certain, red=uncertain)
+For nosigma version: trajectory color = team (blue=home, orange=away, green=ball)
+"""
+
+import os, sys, random
+import numpy as np
+import torch
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import matplotlib.cm as cm
+from matplotlib.colors import Normalize
+
+sys.path.insert(0, os.path.dirname(__file__))
+
+from utils.config import Config
+from data.dataloader_nba import NBADataset, seq_collate
+from torch.utils.data import DataLoader
+from models.model_led_initializer import LEDInitializer as InitializationModel
+from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel
+from models.future_interaction_graph_v6 import FutureInteractionGraphV6Wrapper
+
+NUM_Tau = 5
+COURT_IMG = '/mnt/jaewoo4tb/srtp/srtp/raw_data/nba/court.png'
+COURT_W, COURT_H = 28.0, 15.0
+A = 11
+
+HOME_PAST, HOME_FUT = '#a8d5ff', '#187bff'
+AWAY_PAST, AWAY_FUT = '#ffc9a8', '#ff5b2e'
+BALL_PAST, BALL_FUT = '#b9f2b9', '#1f9d1f'
+GT_COLOR = '#333333'
+
+def agent_colors(idx):
+    if idx < 5:    return HOME_PAST, HOME_FUT
+    elif idx < 10: return AWAY_PAST, AWAY_FUT
+    else:          return BALL_PAST, BALL_FUT
+
+_court_cache = None
+def court_img():
+    global _court_cache
+    if _court_cache is None:
+        _court_cache = plt.imread(COURT_IMG)
+    return _court_cache
+
+
+def load_models(ckpt_path, use_sigma, edge_mode, top_n, device):
+    cfg = Config('led_augment', 'viz')
+    model = CoreDenoisingModel().to(device)
+    cp = torch.load(cfg.pretrained_core_denoising_model, map_location='cpu', weights_only=False)
+    model.load_state_dict(cp['model_dict']); model.eval()
+
+    model_init = InitializationModel(t_h=10, d_h=6, t_f=20, d_f=2, k_pred=20).to(device)
+    graph = FutureInteractionGraphV6Wrapper(
+        num_agents=11, future_steps=20, past_steps=10,
+        past_channels=6, node_dim=128, top_n=top_n,
+        num_denoise_steps=NUM_Tau, edge_mode=edge_mode).to(device)
+
+    ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)
+    model_init.load_state_dict(ckpt['model_initializer_dict'])
+    graph.load_state_dict(ckpt['interaction_graph_dict'])
+    model_init.eval(); graph.eval()
+    return cfg, model, model_init, graph
+
+
+def get_all_stages(model, graph, model_init, past_traj, traj_mask,
+                    betas, alphas, abs_sqrt, oma_sqrt,
+                    use_sigma, traj_scale, init_pos_np):
+    """Return trajectory predictions at each meaningful stage."""
+    sample_pred, mean_est, var_est = model_init(past_traj, traj_mask)
+
+    stages = []
+
+    # Stage 0: Mean estimation only (single mode per agent)
+    mean_abs = mean_est.detach().cpu().numpy()  # [B*A, T, 2]
+    mean_abs = mean_abs * traj_scale + init_pos_np.reshape(-1, 1, 2)
+    stages.append({
+        'trajs': mean_abs[:A, np.newaxis],  # [A, 1, T, 2]
+        'sigma': None,
+        'title': 'Stage 1: Mean estimation',
+    })
+
+    # Stage 1: Initializer with variance scaling (20 diverse modes)
+    sample_pred_scaled = (torch.exp(var_est / 2)[..., None, None]
+                          * sample_pred
+                          / sample_pred.std(dim=1).mean(dim=(1, 2))[:, None, None, None])
+    loc = sample_pred_scaled + mean_est[:, None]
+
+    loc_abs = loc.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2)
+    sigma_np = var_est.detach().cpu().numpy() if use_sigma else None
+    stages.append({
+        'trajs': loc_abs[:A],  # [A, K=20, T, 2]
+        'sigma': sigma_np,
+        'title': 'Stage 2: Initializer (20 modes)',
+    })
+
+    # Run denoising and capture intermediate states
+    sigma_input = var_est if use_sigma else None
+    cur_y = loc[:, :10]
+
+    checkpoints = {3: 'Stage 3: After 2 DDPM steps',
+                   1: 'Stage 4: After 4 DDPM steps',
+                   -1: 'Stage 5: Final (5 DDPM steps)'}
+
+    for i in reversed(range(NUM_Tau)):
+        ef = (1 - alphas[i]) / oma_sqrt[i]
+        beta = betas[i].repeat(past_traj.shape[0]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+        eps = model.generate_accelerate(cur_y, beta.squeeze(-1).squeeze(-1), past_traj, traj_mask)
+        y0h = (cur_y - oma_sqrt[i] * eps) / abs_sqrt[i]
+        delta = graph(y0h, past_traj, i, sigma=sigma_input)
+        eps = eps + delta
+        mean = (1 / alphas[i].sqrt()) * (cur_y - ef * eps)
+        z = torch.randn_like(cur_y)
+        cur_y = mean + betas[i].sqrt() * z * 0.00001
+
+        if i in checkpoints:
+            # Get y0 estimate (clean prediction) at this point
+            y0_est = (cur_y - oma_sqrt[max(0, i-1)] * eps) / abs_sqrt[max(0, i-1)] if i > 0 else cur_y
+            y0_abs = y0_est.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2)
+            stages.append({
+                'trajs': y0_abs[:A],  # [A, K=10, T, 2]
+                'sigma': sigma_np,
+                'title': checkpoints[i],
+            })
+
+    # Stage 5 (final): use cur_y directly
+    # Also do the second half (modes 10-19)
+    cur_y2 = loc[:, 10:]
+    for i in reversed(range(NUM_Tau)):
+        ef = (1 - alphas[i]) / oma_sqrt[i]
+        beta = betas[i].repeat(past_traj.shape[0]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+        eps = model.generate_accelerate(cur_y2, beta.squeeze(-1).squeeze(-1), past_traj, traj_mask)
+        y0h = (cur_y2 - oma_sqrt[i] * eps) / abs_sqrt[i]
+        delta = graph(y0h, past_traj, i, sigma=sigma_input)
+        eps = eps + delta
+        mean = (1 / alphas[i].sqrt()) * (cur_y2 - ef * eps)
+        z = torch.randn_like(cur_y2)
+        cur_y2 = mean + betas[i].sqrt() * z * 0.00001
+
+    final = torch.cat((cur_y2, cur_y), dim=1)
+    final_abs = final.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2)
+    stages[-1] = {
+        'trajs': final_abs[:A],  # [A, K=20, T, 2]
+        'sigma': sigma_np,
+        'title': 'Stage 5: Final (all 20 modes)',
+    }
+
+    return stages
+
+
+def draw_stage(ax, past_abs, gt_abs, trajs, sigma_vals, title, show_sigma):
+    """Draw one stage on court."""
+    ax.imshow(court_img(), extent=[0, COURT_W, COURT_H, 0], zorder=0, alpha=0.5)
+    ax.set_xlim(0, COURT_W); ax.set_ylim(COURT_H, 0)
+    ax.axis('off')
+    ax.set_title(title, fontsize=9, pad=3)
+
+    # Past
+    for a in range(A):
+        pc, _ = agent_colors(a)
+        ax.plot(past_abs[a, :, 0], past_abs[a, :, 1], color=pc, lw=0.8,
+                marker='o', ms=1.5, alpha=0.6, zorder=2)
+
+    # GT
+    for a in range(A):
+        gt = np.concatenate([past_abs[a, -1:], gt_abs[a]], axis=0)
+        ax.plot(gt[:, 0], gt[:, 1], color=GT_COLOR, lw=0.7, marker='o',
+                ms=1.0, alpha=0.4, zorder=3, linestyle='--')
+
+    K = trajs.shape[1]
+    K_show = min(K, 10)
+
+    if show_sigma and sigma_vals is not None:
+        sigma_std = np.exp(sigma_vals[:A, 0] / 2)
+        norm = Normalize(vmin=sigma_std.min() - 0.01, vmax=sigma_std.max() + 0.01)
+        cmap = cm.RdYlGn_r
+
+        for a in range(A):
+            color = cmap(norm(sigma_std[a]))
+            for k in range(K_show):
+                pred = np.concatenate([past_abs[a, -1:], trajs[a, k]], axis=0)
+                ax.plot(pred[:, 0], pred[:, 1], color=color, lw=0.5, alpha=0.3, zorder=4)
+            # Best mode
+            dists = np.linalg.norm(trajs[a, :K_show] - gt_abs[a:a+1], axis=-1).mean(axis=-1)
+            best_k = dists.argmin()
+            best = np.concatenate([past_abs[a, -1:], trajs[a, best_k]], axis=0)
+            ax.plot(best[:, 0], best[:, 1], color=color, lw=2.0, alpha=0.9,
+                    zorder=5, marker='o', ms=2.0)
+    else:
+        for a in range(A):
+            _, fc = agent_colors(a)
+            for k in range(K_show):
+                pred = np.concatenate([past_abs[a, -1:], trajs[a, k]], axis=0)
+                ax.plot(pred[:, 0], pred[:, 1], color=fc, lw=0.5, alpha=0.25, zorder=4)
+            dists = np.linalg.norm(trajs[a, :K_show] - gt_abs[a:a+1], axis=-1).mean(axis=-1)
+            best_k = dists.argmin()
+            best = np.concatenate([past_abs[a, -1:], trajs[a, best_k]], axis=0)
+            ax.plot(best[:, 0], best[:, 1], color=fc, lw=2.0, alpha=0.9,
+                    zorder=5, marker='o', ms=2.0)
+
+
+def main():
+    device = 'cuda:0'
+
+    ckpt_sigma = '/mnt/jaewoo4tb/srtp/LED/results/led_augment/graph_v6_edge_relpos/models/model_0052.p'
+    ckpt_nosigma = '/mnt/jaewoo4tb/srtp/LED/results/led_augment/graph_v6_nosigma_n3/models/model_0084.p'
+
+    out_dir = '/mnt/jaewoo4tb/srtp/LED/visualizations/denoising_steps'
+    os.makedirs(out_dir, exist_ok=True)
+
+    cfg, model_s, init_s, graph_s = load_models(
+        ckpt_sigma, True, 'relpos_only', 5, device)
+    _, model_n, init_n, graph_n = load_models(
+        ckpt_nosigma, False, 'full', 3, device)
+
+    betas = torch.linspace(1e-5, 1e-2, 100).to(device)
+    alphas = 1 - betas
+    alphas_prod = torch.cumprod(alphas, 0)
+    abs_sqrt = torch.sqrt(alphas_prod)
+    oma_sqrt = torch.sqrt(1 - alphas_prod)
+
+    traj_scale = cfg.traj_scale
+
+    test_dset = NBADataset(obs_len=10, pred_len=20, training=False)
+    test_loader = DataLoader(test_dset, batch_size=1, shuffle=False, collate_fn=seq_collate)
+
+    np.random.seed(42); random.seed(42); torch.manual_seed(42)
+    sample_indices = sorted(random.sample(range(len(test_dset)), 3))
+
+    with torch.no_grad():
+        for sample_idx, data in enumerate(test_loader):
+            if sample_idx not in sample_indices:
+                continue
+            if sample_idx > max(sample_indices):
+                break
+
+            traj_mean_t = torch.FloatTensor(cfg.traj_mean).cuda().unsqueeze(0).unsqueeze(0).unsqueeze(0)
+            initial_pos = data['pre_motion_3D'].cuda()[:, :, -1:]
+            past_abs_np = data['pre_motion_3D'].numpy().squeeze(0) / (94.0 / 28.0)
+            fut_abs_np = data['fut_motion_3D'].numpy().squeeze(0) / (94.0 / 28.0)
+            init_pos_np = initial_pos.cpu().numpy().squeeze(0) / (94.0 / 28.0)
+
+            past_traj_abs = ((data['pre_motion_3D'].cuda() - traj_mean_t) / traj_scale).view(-1, 10, 2)
+            past_traj_rel = ((data['pre_motion_3D'].cuda() - initial_pos) / traj_scale).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)
+            traj_mask = torch.ones(11, 11).cuda()
+
+            for version, model, init_model, graph, use_sigma, label in [
+                ('nosigma', model_n, init_n, graph_n, False, 'Without Uncertainty'),
+                ('sigma', model_s, init_s, graph_s, True, 'With Uncertainty'),
+            ]:
+                stages = get_all_stages(
+                    model, graph, init_model, past_traj, traj_mask,
+                    betas, alphas, abs_sqrt, oma_sqrt,
+                    use_sigma, traj_scale, init_pos_np)
+
+                for si, stage in enumerate(stages):
+                    fig, ax = plt.subplots(1, 1, figsize=(7, 5), dpi=200)
+                    title = f'{label} — {stage["title"]}'
+                    draw_stage(ax, past_abs_np, fut_abs_np, stage['trajs'],
+                               stage['sigma'], title, show_sigma=use_sigma)
+
+                    if use_sigma and stage['sigma'] is not None:
+                        sigma_std = np.exp(stage['sigma'][:A, 0] / 2)
+                        sm = cm.ScalarMappable(cmap=cm.RdYlGn_r,
+                                               norm=Normalize(vmin=sigma_std.min() - 0.01,
+                                                              vmax=sigma_std.max() + 0.01))
+                        sm.set_array([])
+                        cbar = fig.colorbar(sm, ax=ax, shrink=0.5, pad=0.02)
+                        cbar.set_label('σ (uncertainty)', fontsize=7)
+
+                    plt.tight_layout()
+                    fname = f'sample_{sample_idx:04d}_{version}_stage{si}.png'
+                    fig.savefig(os.path.join(out_dir, fname),
+                                bbox_inches='tight', pad_inches=0.02, dpi=200)
+                    plt.close(fig)
+
+                print(f'  Sample {sample_idx} {version}: {len(stages)} stages saved')
+
+    print(f'\nAll saved to {out_dir}/')
+
+
+if __name__ == '__main__':
+    main()

LED/viz_uncertainty_denoising.py

@@ -0,0 +1,371 @@
+"""
+Visualize LED denoising with/without uncertainty — matching qual_uncertainty.png style.
+
+For each denoising step:
+  - Past: blue dots+lines
+  - Future GT: red dots+lines
+  - Prediction: green, with darkness proportional to uncertainty
+    (dark green = certain, light green = uncertain)
+
+Left column: No uncertainty (all predictions same green)
+Right column: With uncertainty (per-agent green intensity from σ)
+
+Produces one image per sample with rows = denoising stages,
+columns = [no uncertainty, with uncertainty].
+"""
+
+import os, sys, random
+import numpy as np
+import torch
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+from matplotlib.lines import Line2D
+
+sys.path.insert(0, '/mnt/jaewoo4tb/srtp/LED')
+
+from utils.config import Config
+from data.dataloader_nba import NBADataset, seq_collate
+from torch.utils.data import DataLoader
+from models.model_led_initializer import LEDInitializer as InitializationModel
+from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel
+from models.future_interaction_graph_v6 import FutureInteractionGraphV6Wrapper
+
+NUM_Tau = 5
+COURT_IMG = '/mnt/jaewoo4tb/srtp/srtp/raw_data/nba/court.png'
+COURT_W, COURT_H = 28.0, 15.0
+A = 11
+OUT_DIR = '/mnt/jaewoo4tb/srtp/NeurIPS_2026_SRT/Viz_uncertainty'
+
+# Colors matching the reference figure
+PAST_COLOR = '#2962FF'      # blue
+GT_COLOR = '#E91E63'        # red/pink
+PRED_BASE = np.array([0.2, 0.6, 0.1])  # base green for predictions
+
+_court_cache = None
+def court_img():
+    global _court_cache
+    if _court_cache is None:
+        _court_cache = plt.imread(COURT_IMG)
+    return _court_cache
+
+
+def load_models(ckpt_path, use_sigma, edge_mode, top_n, device):
+    os.chdir('/mnt/jaewoo4tb/srtp/LED')
+    cfg = Config('led_augment', 'viz')
+    model = CoreDenoisingModel().to(device)
+    cp = torch.load(cfg.pretrained_core_denoising_model, map_location='cpu', weights_only=False)
+    model.load_state_dict(cp['model_dict']); model.eval()
+
+    model_init = InitializationModel(t_h=10, d_h=6, t_f=20, d_f=2, k_pred=20).to(device)
+    graph = FutureInteractionGraphV6Wrapper(
+        num_agents=11, future_steps=20, past_steps=10,
+        past_channels=6, node_dim=128, top_n=top_n,
+        num_denoise_steps=NUM_Tau, edge_mode=edge_mode).to(device)
+
+    ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)
+    model_init.load_state_dict(ckpt['model_initializer_dict'])
+    graph.load_state_dict(ckpt['interaction_graph_dict'])
+    model_init.eval(); graph.eval()
+    return cfg, model, model_init, graph
+
+
+def get_stages(model, graph, model_init, past_traj, traj_mask,
+               betas, alphas, abs_sqrt, oma_sqrt,
+               use_sigma, traj_scale, init_pos_np):
+    """Get prediction at each meaningful stage."""
+    sample_pred, mean_est, var_est = model_init(past_traj, traj_mask)
+    sample_pred_scaled = (torch.exp(var_est / 2)[..., None, None]
+                          * sample_pred
+                          / sample_pred.std(dim=1).mean(dim=(1, 2))[:, None, None, None])
+    loc = sample_pred_scaled + mean_est[:, None]
+
+    sigma_input = var_est if use_sigma else None
+    sigma_np = var_est.detach().cpu().numpy() if use_sigma else None
+
+    stages = []
+
+    # Stage 0: Initializer output (before any denoising)
+    loc_abs = loc.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2)
+    stages.append({'trajs': loc_abs[:A], 'sigma': sigma_np, 'label': 'Initializer'})
+
+    # Run denoising, capture at key steps
+    cur_y = loc[:, :10]
+    capture_at = {3: 'After 2 steps', 1: 'After 4 steps'}
+
+    for i in reversed(range(NUM_Tau)):
+        ef = (1 - alphas[i]) / oma_sqrt[i]
+        beta = betas[i].repeat(past_traj.shape[0]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+        eps = model.generate_accelerate(cur_y, beta.squeeze(-1).squeeze(-1), past_traj, traj_mask)
+        y0h = (cur_y - oma_sqrt[i] * eps) / abs_sqrt[i]
+        delta = graph(y0h, past_traj, i, sigma=sigma_input)
+        eps = eps + delta
+        mean = (1 / alphas[i].sqrt()) * (cur_y - ef * eps)
+        z = torch.randn_like(cur_y)
+        cur_y = mean + betas[i].sqrt() * z * 0.00001
+
+        if i in capture_at:
+            cur_abs = cur_y.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2)
+            stages.append({'trajs': cur_abs[:A], 'sigma': sigma_np, 'label': capture_at[i]})
+
+    # Final: run second half too
+    cur_y2 = loc[:, 10:]
+    for i in reversed(range(NUM_Tau)):
+        ef = (1 - alphas[i]) / oma_sqrt[i]
+        beta = betas[i].repeat(past_traj.shape[0]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+        eps = model.generate_accelerate(cur_y2, beta.squeeze(-1).squeeze(-1), past_traj, traj_mask)
+        y0h = (cur_y2 - oma_sqrt[i] * eps) / abs_sqrt[i]
+        delta = graph(y0h, past_traj, i, sigma=sigma_input)
+        eps = eps + delta
+        mean = (1 / alphas[i].sqrt()) * (cur_y2 - ef * eps)
+        z = torch.randn_like(cur_y2)
+        cur_y2 = mean + betas[i].sqrt() * z * 0.00001
+
+    final = torch.cat((cur_y2, cur_y), dim=1)
+    final_abs = final.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2)
+    stages.append({'trajs': final_abs[:A], 'sigma': sigma_np, 'label': 'Final'})
+
+    return stages
+
+
+def uncertainty_to_green(sigma_vals, a):
+    """Convert per-agent uncertainty to green color intensity.
+    High certainty (low σ) → dark green, Low certainty (high σ) → light green.
+    """
+    if sigma_vals is None:
+        return (0.2, 0.65, 0.1, 0.8)  # default green
+
+    sigma_std = np.exp(sigma_vals[:A, 0] / 2)  # actual std
+    # Normalize to [0, 1] range
+    s_min, s_max = sigma_std.min(), sigma_std.max()
+    if s_max - s_min < 1e-6:
+        norm_val = 0.5
+    else:
+        norm_val = (sigma_std[a] - s_min) / (s_max - s_min)
+
+    # Dark green (certain) to light yellow-green (uncertain)
+    # Interpolate: dark (0.1, 0.4, 0.05) ↔ light (0.6, 0.85, 0.3)
+    dark = np.array([0.05, 0.35, 0.0])
+    light = np.array([0.55, 0.82, 0.25])
+    color = dark + norm_val * (light - dark)
+    return (*color, 0.85)
+
+
+def draw_stage(ax, past_abs, gt_abs, trajs, sigma_vals, show_sigma,
+               xlim=None, ylim=None):
+    """Draw one stage on court."""
+    ax.imshow(court_img(), extent=[0, COURT_W, COURT_H, 0], zorder=0, alpha=0.45)
+    if xlim:
+        ax.set_xlim(*xlim)
+    else:
+        ax.set_xlim(0, COURT_W)
+    if ylim:
+        ax.set_ylim(*ylim)
+    else:
+        ax.set_ylim(COURT_H, 0)
+    ax.axis('off')
+
+    K = trajs.shape[1]
+    K_show = min(5, K)
+
+    # Find best mode per agent
+    best_modes = []
+    for a in range(A):
+        if K > 1:
+            dists = np.linalg.norm(trajs[a, :K_show] - gt_abs[a:a+1], axis=-1).mean(axis=-1)
+            best_modes.append(dists.argmin())
+        else:
+            best_modes.append(0)
+
+    # Draw predictions (green, intensity by uncertainty)
+    for a in range(A):
+        if show_sigma:
+            color = uncertainty_to_green(sigma_vals, a)
+        else:
+            color = (0.2, 0.65, 0.1, 0.7)
+
+        # Light modes
+        for k in range(K_show):
+            pred = trajs[a, k]
+            ax.plot(pred[:, 0], pred[:, 1], color=color, lw=0.4, alpha=0.25, zorder=3)
+
+        # Best mode (thicker)
+        best = trajs[a, best_modes[a]]
+        pred_line = np.concatenate([past_abs[a, -1:], best], axis=0)
+        ax.plot(pred_line[:, 0], pred_line[:, 1], color=color, lw=1.8,
+                marker='o', ms=1.8, markevery=3, zorder=5)
+
+    # Draw past (blue)
+    for a in range(A):
+        ax.plot(past_abs[a, :, 0], past_abs[a, :, 1], color=PAST_COLOR,
+                lw=1.2, marker='o', ms=2.0, markevery=2, alpha=0.85, zorder=6)
+
+    # Draw GT future (red)
+    for a in range(A):
+        gt = np.concatenate([past_abs[a, -1:], gt_abs[a]], axis=0)
+        ax.plot(gt[:, 0], gt[:, 1], color=GT_COLOR, lw=1.0,
+                marker='o', ms=1.5, markevery=3, alpha=0.7, zorder=4)
+
+
+def make_figure(stages_nosigma, stages_sigma, past_abs, gt_abs, sample_idx, zoom_region=None):
+    """Create the full comparison figure."""
+    n_stages = len(stages_nosigma)
+
+    # If zoom region provided, add zoomed row at bottom
+    n_rows = n_stages + (1 if zoom_region else 0)
+
+    fig, axes = plt.subplots(n_rows, 2, figsize=(12, 3.0 * n_rows), dpi=200)
+    if n_rows == 1:
+        axes = axes.reshape(1, 2)
+
+    for row in range(n_stages):
+        # Left: no uncertainty
+        draw_stage(axes[row, 0], past_abs, gt_abs,
+                   stages_nosigma[row]['trajs'], None, show_sigma=False)
+        if row == 0:
+            axes[row, 0].set_title('No uncertainty', fontsize=11, fontweight='bold')
+
+        # Right: with uncertainty
+        draw_stage(axes[row, 1], past_abs, gt_abs,
+                   stages_sigma[row]['trajs'], stages_sigma[row]['sigma'], show_sigma=True)
+        if row == 0:
+            axes[row, 1].set_title('Using uncertainty', fontsize=11, fontweight='bold')
+
+        # Row label
+        axes[row, 0].text(0.02, 0.95, stages_nosigma[row]['label'],
+                          transform=axes[row, 0].transAxes, fontsize=8,
+                          verticalalignment='top', fontweight='bold',
+                          bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.7))
+
+    # Zoomed bottom row
+    if zoom_region and n_rows > n_stages:
+        xl, yl = zoom_region
+        for col in range(2):
+            stages = stages_nosigma if col == 0 else stages_sigma
+            sigma = None if col == 0 else stages[-1]['sigma']
+            draw_stage(axes[-1, col], past_abs, gt_abs,
+                       stages[-1]['trajs'], sigma, show_sigma=(col == 1),
+                       xlim=xl, ylim=yl)
+            # Draw zoom box on the row above
+            from matplotlib.patches import Rectangle
+            rect = Rectangle((xl[0], yl[1]), xl[1]-xl[0], yl[0]-yl[1],
+                              linewidth=1.5, edgecolor='black', facecolor='none', zorder=10)
+            axes[-2, col].add_patch(rect)
+
+    # Add noise level indicator on the right side
+    # Gray gradient arrow from τ^K (top) to 0 (bottom)
+    noise_ax = fig.add_axes([0.92, 0.15, 0.03, 0.7])
+    gradient = np.linspace(0.3, 1.0, 256).reshape(256, 1)
+    noise_ax.imshow(gradient, aspect='auto', cmap='Greys_r', extent=[0, 1, 0, 1])
+    noise_ax.set_xticks([])
+    noise_ax.set_yticks([0, 1])
+    noise_ax.set_yticklabels(['0', 'τ$^K$'], fontsize=8)
+    noise_ax.set_ylabel('Noise level', fontsize=8, rotation=270, labelpad=12)
+    noise_ax.yaxis.set_label_position('right')
+
+    # Legend
+    legend_elements = [
+        Line2D([0], [0], color=PAST_COLOR, lw=2, marker='o', ms=4, label='Past'),
+        Line2D([0], [0], color=GT_COLOR, lw=2, marker='o', ms=4, label='Future'),
+        Line2D([0], [0], color=(0.2, 0.65, 0.1), lw=2, marker='o', ms=4, label='Prediction'),
+    ]
+    fig.legend(handles=legend_elements, loc='upper center', ncol=3,
+               fontsize=9, frameon=True, fancybox=True, shadow=True,
+               bbox_to_anchor=(0.45, 0.98))
+
+    # Uncertainty colorbar for the right column
+    from matplotlib.cm import ScalarMappable
+    from matplotlib.colors import LinearSegmentedColormap, Normalize
+    dark = (0.05, 0.35, 0.0)
+    light = (0.55, 0.82, 0.25)
+    cmap_unc = LinearSegmentedColormap.from_list('unc', [dark, light])
+    sm = ScalarMappable(cmap=cmap_unc, norm=Normalize(0, 1))
+    sm.set_array([])
+    cbar_ax = fig.add_axes([0.52, 0.96, 0.15, 0.012])
+    cbar = fig.colorbar(sm, cax=cbar_ax, orientation='horizontal')
+    cbar.set_ticks([0, 1])
+    cbar.set_ticklabels(['Certain', 'Uncertain'], fontsize=7)
+    cbar_ax.set_title('Uncertainty level', fontsize=7, pad=2)
+
+    plt.subplots_adjust(hspace=0.05, wspace=0.02, right=0.90, top=0.93)
+
+    save_path = os.path.join(OUT_DIR, f'denoising_uncertainty_sample_{sample_idx:04d}.png')
+    fig.savefig(save_path, bbox_inches='tight', pad_inches=0.05, dpi=200)
+    plt.close(fig)
+    print(f'Saved: {save_path}')
+
+
+def main():
+    device = 'cuda:0'
+    os.environ['CUDA_VISIBLE_DEVICES'] = '3'  # use GPU 3
+    torch.cuda.set_device(0)
+
+    ckpt_sigma = '/mnt/jaewoo4tb/srtp/LED/results/led_augment/graph_v6_edge_relpos/models/model_0052.p'
+    ckpt_nosigma = '/mnt/jaewoo4tb/srtp/LED/results/led_augment/graph_v6_nosigma_n3/models/model_0084.p'
+
+    cfg, model_s, init_s, graph_s = load_models(
+        ckpt_sigma, True, 'relpos_only', 5, device)
+    _, model_n, init_n, graph_n = load_models(
+        ckpt_nosigma, False, 'full', 3, device)
+
+    betas = torch.linspace(1e-5, 1e-2, 100).to(device)
+    alphas = 1 - betas
+    alphas_prod = torch.cumprod(alphas, 0)
+    abs_sqrt = torch.sqrt(alphas_prod)
+    oma_sqrt = torch.sqrt(1 - alphas_prod)
+
+    traj_scale = cfg.traj_scale
+
+    test_dset = NBADataset(obs_len=10, pred_len=20, training=False)
+    test_loader = DataLoader(test_dset, batch_size=1, shuffle=False, collate_fn=seq_collate)
+
+    np.random.seed(42); random.seed(42); torch.manual_seed(42)
+    sample_indices = sorted(random.sample(range(len(test_dset)), 5))
+
+    with torch.no_grad():
+        for sample_idx, data in enumerate(test_loader):
+            if sample_idx not in sample_indices:
+                continue
+            if sample_idx > max(sample_indices):
+                break
+
+            traj_mean_t = torch.FloatTensor(cfg.traj_mean).cuda().unsqueeze(0).unsqueeze(0).unsqueeze(0)
+            initial_pos = data['pre_motion_3D'].cuda()[:, :, -1:]
+            # Data is already in court units (0-28 x 0-15), no further scaling needed
+            past_abs_np = data['pre_motion_3D'].numpy().squeeze(0)
+            fut_abs_np = data['fut_motion_3D'].numpy().squeeze(0)
+            init_pos_np = initial_pos.cpu().numpy().squeeze(0)
+
+            past_traj_abs = ((data['pre_motion_3D'].cuda() - traj_mean_t) / traj_scale).view(-1, 10, 2)
+            past_traj_rel = ((data['pre_motion_3D'].cuda() - initial_pos) / traj_scale).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)
+            traj_mask = torch.ones(11, 11).cuda()
+
+            # Get stages for both versions
+            stages_nosigma = get_stages(
+                model_n, graph_n, init_n, past_traj, traj_mask,
+                betas, alphas, abs_sqrt, oma_sqrt,
+                False, traj_scale, init_pos_np)
+
+            stages_sigma = get_stages(
+                model_s, graph_s, init_s, past_traj, traj_mask,
+                betas, alphas, abs_sqrt, oma_sqrt,
+                True, traj_scale, init_pos_np)
+
+            # Compute zoom region around the action
+            all_pos = np.concatenate([past_abs_np.reshape(-1, 2), fut_abs_np.reshape(-1, 2)])
+            cx, cy = all_pos.mean(axis=0)
+            span = max(all_pos.max(axis=0) - all_pos.min(axis=0)) * 0.6
+            zoom = ([cx - span, cx + span], [cy + span, cy - span])
+
+            make_figure(stages_nosigma, stages_sigma, past_abs_np, fut_abs_np,
+                        sample_idx, zoom_region=zoom)
+
+    print(f'\nAll saved to {OUT_DIR}/')
+
+
+if __name__ == '__main__':
+    main()

LED/viz_uncertainty_individual.py

@@ -0,0 +1,197 @@
+"""
+Draw each denoising step as a separate image.
+Past=blue, GT=red, Prediction=green (darkness by uncertainty).
+"""
+
+import os, sys, random
+import numpy as np
+import torch
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+
+sys.path.insert(0, '/mnt/jaewoo4tb/srtp/LED')
+
+from utils.config import Config
+from data.dataloader_nba import NBADataset, seq_collate
+from torch.utils.data import DataLoader
+from models.model_led_initializer import LEDInitializer
+from models.model_diffusion import TransformerDenoisingModel
+from models.future_interaction_graph_v6 import FutureInteractionGraphV6Wrapper
+
+NUM_Tau = 5
+COURT_IMG = '/mnt/jaewoo4tb/srtp/srtp/raw_data/nba/court.png'
+COURT_W, COURT_H = 28.0, 15.0
+A = 11
+OUT_DIR = '/mnt/jaewoo4tb/srtp/NeurIPS_2026_SRT/Viz_uncertainty'
+
+PAST_COLOR = '#2962FF'
+GT_COLOR = '#E91E63'
+
+_court_cache = None
+def court_img():
+    global _court_cache
+    if _court_cache is None:
+        _court_cache = plt.imread(COURT_IMG)
+    return _court_cache
+
+
+def load_models(ckpt_path, use_sigma, edge_mode, top_n, device):
+    os.chdir('/mnt/jaewoo4tb/srtp/LED')
+    cfg = Config('led_augment', 'viz')
+    model = TransformerDenoisingModel().to(device)
+    cp = torch.load(cfg.pretrained_core_denoising_model, map_location='cpu', weights_only=False)
+    model.load_state_dict(cp['model_dict']); model.eval()
+    init_m = LEDInitializer(t_h=10, d_h=6, t_f=20, d_f=2, k_pred=20).to(device)
+    graph = FutureInteractionGraphV6Wrapper(
+        num_agents=11, future_steps=20, past_steps=10,
+        past_channels=6, node_dim=128, top_n=top_n,
+        num_denoise_steps=NUM_Tau, edge_mode=edge_mode).to(device)
+    ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)
+    init_m.load_state_dict(ckpt['model_initializer_dict'])
+    graph.load_state_dict(ckpt['interaction_graph_dict'])
+    init_m.eval(); graph.eval()
+    return cfg, model, init_m, graph
+
+
+def get_all_steps(model, graph, init_m, past_traj, traj_mask,
+                  betas, alphas, abs_sqrt, oma_sqrt,
+                  use_sigma, traj_scale, init_pos_np):
+    sample_pred, mean_est, var_est = init_m(past_traj, traj_mask)
+    sample_pred = (torch.exp(var_est / 2)[..., None, None]
+                   * sample_pred / sample_pred.std(dim=1).mean(dim=(1, 2))[:, None, None, None])
+    loc = sample_pred + mean_est[:, None]
+    sigma_input = var_est if use_sigma else None
+    sigma_np = var_est.detach().cpu().numpy() if use_sigma else None
+
+    steps = []
+
+    # Initializer
+    loc_abs = loc.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2)
+    steps.append({'trajs': loc_abs[:A], 'sigma': sigma_np, 'label': 'Initializer'})
+
+    # Denoise first 10 modes
+    cur_y = loc[:, :10]
+    for i in reversed(range(NUM_Tau)):
+        ef = (1 - alphas[i]) / oma_sqrt[i]
+        beta = betas[i].repeat(past_traj.shape[0]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+        eps = model.generate_accelerate(cur_y, beta.squeeze(-1).squeeze(-1), past_traj, traj_mask)
+        y0h = (cur_y - oma_sqrt[i] * eps) / abs_sqrt[i]
+        delta = graph(y0h, past_traj, i, sigma=sigma_input)
+        eps = eps + delta
+        mean = (1 / alphas[i].sqrt()) * (cur_y - ef * eps)
+        z = torch.randn_like(cur_y)
+        cur_y = mean + betas[i].sqrt() * z * 0.00001
+
+        cur_abs = cur_y.detach().cpu().numpy() * traj_scale + init_pos_np.reshape(-1, 1, 1, 2)
+        steps.append({'trajs': cur_abs[:A], 'sigma': sigma_np, 'label': f'Step {NUM_Tau - i}/{NUM_Tau}'})
+
+    return steps
+
+
+def unc_color(sigma_vals, a):
+    if sigma_vals is None:
+        return (0.2, 0.65, 0.1, 0.8)
+    sigma_std = np.exp(sigma_vals[:A, 0] / 2)
+    s_min, s_max = sigma_std.min(), sigma_std.max()
+    norm = (sigma_std[a] - s_min) / (s_max - s_min + 1e-8)
+    dark = np.array([0.05, 0.35, 0.0])
+    light = np.array([0.55, 0.82, 0.25])
+    c = dark + norm * (light - dark)
+    return (*c, 0.85)
+
+
+def draw(ax, past, gt, trajs, sigma, show_sigma):
+    ax.imshow(court_img(), extent=[0, COURT_W, COURT_H, 0], zorder=0, alpha=0.45)
+    ax.set_xlim(0, COURT_W); ax.set_ylim(COURT_H, 0)
+    ax.axis('off')
+
+    K = min(5, trajs.shape[1])
+
+    for a in range(A):
+        color = unc_color(sigma, a) if show_sigma else (0.2, 0.65, 0.1, 0.8)
+        dists = np.linalg.norm(trajs[a, :K] - gt[a:a+1], axis=-1).mean(axis=-1)
+        best_k = dists.argmin()
+
+        best = np.concatenate([past[a, -1:], trajs[a, best_k]], axis=0)
+        ax.plot(best[:, 0], best[:, 1], color=color, lw=1.8,
+                marker='o', ms=1.8, markevery=3, zorder=5)
+
+    for a in range(A):
+        ax.plot(past[a, :, 0], past[a, :, 1], color=PAST_COLOR,
+                lw=1.2, marker='o', ms=2.0, markevery=2, alpha=0.85, zorder=6)
+
+    for a in range(A):
+        g = np.concatenate([past[a, -1:], gt[a]], axis=0)
+        ax.plot(g[:, 0], g[:, 1], color=GT_COLOR, lw=1.0,
+                marker='o', ms=1.5, markevery=3, alpha=0.7, zorder=4)
+
+
+def main():
+    device = 'cuda:0'
+    os.environ['CUDA_VISIBLE_DEVICES'] = '3'
+    torch.cuda.set_device(0)
+
+    ckpt_s = '/mnt/jaewoo4tb/srtp/LED/results/led_augment/graph_v6_edge_relpos/models/model_0052.p'
+    ckpt_n = '/mnt/jaewoo4tb/srtp/LED/results/led_augment/graph_v6_nosigma_n3/models/model_0084.p'
+
+    cfg, model_s, init_s, graph_s = load_models(ckpt_s, True, 'relpos_only', 5, device)
+    _, model_n, init_n, graph_n = load_models(ckpt_n, False, 'full', 3, device)
+
+    betas = torch.linspace(1e-5, 1e-2, 100).to(device)
+    alphas = 1 - betas
+    alphas_prod = torch.cumprod(alphas, 0)
+    abs_sqrt = torch.sqrt(alphas_prod)
+    oma_sqrt = torch.sqrt(1 - alphas_prod)
+    traj_scale = cfg.traj_scale
+
+    test_dset = NBADataset(obs_len=10, pred_len=20, training=False)
+    test_loader = DataLoader(test_dset, batch_size=1, shuffle=False, collate_fn=seq_collate)
+
+    np.random.seed(42); random.seed(42); torch.manual_seed(42)
+    sample_indices = sorted(random.sample(range(len(test_dset)), 3))
+
+    with torch.no_grad():
+        for idx, data in enumerate(test_loader):
+            if idx not in sample_indices:
+                continue
+            if idx > max(sample_indices):
+                break
+
+            traj_mean_t = torch.FloatTensor(cfg.traj_mean).cuda().unsqueeze(0).unsqueeze(0).unsqueeze(0)
+            initial_pos = data['pre_motion_3D'].cuda()[:, :, -1:]
+            # Data is already in court units (0-28 x 0-15), no further scaling needed
+            past_np = data['pre_motion_3D'].numpy().squeeze(0)
+            fut_np = data['fut_motion_3D'].numpy().squeeze(0)
+            init_np = initial_pos.cpu().numpy().squeeze(0)
+
+            pa = ((data['pre_motion_3D'].cuda() - traj_mean_t) / traj_scale).view(-1, 10, 2)
+            pr = ((data['pre_motion_3D'].cuda() - initial_pos) / traj_scale).view(-1, 10, 2)
+            pv = torch.cat((pr[:, 1:] - pr[:, :-1], torch.zeros_like(pr[:, :1])), dim=1)
+            past_traj = torch.cat((pa, pr, pv), dim=-1)
+            mask = torch.ones(11, 11).cuda()
+
+            for version, model, init_m, graph, use_sigma, tag in [
+                ('nosigma', model_n, init_n, graph_n, False, 'no_uncertainty'),
+                ('sigma', model_s, init_s, graph_s, True, 'with_uncertainty'),
+            ]:
+                steps = get_all_steps(model, graph, init_m, past_traj, mask,
+                                      betas, alphas, abs_sqrt, oma_sqrt,
+                                      use_sigma, traj_scale, init_np)
+
+                for si, step in enumerate(steps):
+                    fig, ax = plt.subplots(1, 1, figsize=(7, 5), dpi=200)
+                    draw(ax, past_np, fut_np, step['trajs'], step['sigma'], use_sigma)
+                    plt.tight_layout()
+                    fname = f'sample_{idx:04d}_{tag}_step{si}.png'
+                    fig.savefig(os.path.join(OUT_DIR, fname),
+                                bbox_inches='tight', pad_inches=0.02, dpi=200)
+                    plt.close(fig)
+
+                print(f'  Sample {idx} {tag}: {len(steps)} images')
+
+    print(f'All saved to {OUT_DIR}/')
+
+
+if __name__ == '__main__':
+    main()

MID/configs/baseline.yaml

@@ -0,0 +1,46 @@
+############ MID Config #############
+# optimizer
+lr: 0.001
+# dataset
+data_dir: processed_data
+  # the path with the directory with XXX.pkl
+# Training Prarmeters:
+diffnet: TransformerConcatLinear #TransformerConcatLinear
+encoder_dim: 256
+tf_layer: 3
+epochs: 90
+batch_size: 256
+eval_batch_size: 256
+k_eval: 25
+seed: 123
+eval_every: 30
+# Testing
+eval_at: 70
+eval_mode: False
+# set to "ddim" to enable fast sampling
+sampling: ddpm
+############### Trajectron++ Config #################
+# misc
+conf: None
+debug: False
+preprocess_workers: 0
+# model parameters
+offline_scene_graph: yes
+dynamic_edges: yes
+edge_state_combine_method: sum
+edge_influence_combine_method: attention
+edge_addition_filter: [0.25, 0.5, 0.75, 1.0]
+edge_removal_filter: [1.0, 0.0]
+override_attention_radius: []
+incl_robot_node: False
+map_encoding: False
+augment: True
+node_freq_mult_train: False
+node_freq_mult_eval: False
+scene_freq_mult_train: False
+scene_freq_mult_eval: False
+scene_freq_mult_viz: False
+no_edge_encoding: False
+# Data Parameters:
+device: cuda
+eval_device: None

MID/configs/baseline_sdd.yaml

@@ -0,0 +1,35 @@
+lr: 0.001
+data_dir: processed_data
+diffnet: TransformerConcatLinear
+encoder_dim: 256
+tf_layer: 3
+epochs: 90
+batch_size: 256
+eval_batch_size: 256
+k_eval: 25
+seed: 123
+eval_every: 4
+eval_at: 70
+eval_mode: False
+sampling: ddpm
+conf: None
+debug: False
+preprocess_workers: 0
+offline_scene_graph: yes
+dynamic_edges: yes
+edge_state_combine_method: sum
+edge_influence_combine_method: attention
+edge_addition_filter: [0.25, 0.5, 0.75, 1.0]
+edge_removal_filter: [1.0, 0.0]
+override_attention_radius: []
+incl_robot_node: False
+map_encoding: False
+augment: True
+node_freq_mult_train: False
+node_freq_mult_eval: False
+scene_freq_mult_train: False
+scene_freq_mult_eval: False
+scene_freq_mult_viz: False
+no_edge_encoding: False
+device: cuda
+eval_device: None

MID/configs/baseline_sdd_eval.yaml

@@ -0,0 +1,35 @@
+lr: 0.001
+data_dir: processed_data
+diffnet: TransformerConcatLinear
+encoder_dim: 256
+tf_layer: 3
+epochs: 90
+batch_size: 256
+eval_batch_size: 256
+k_eval: 25
+seed: 123
+eval_every: 4
+eval_at: 48
+eval_mode: True
+sampling: ddim
+conf: None
+debug: False
+preprocess_workers: 0
+offline_scene_graph: yes
+dynamic_edges: yes
+edge_state_combine_method: sum
+edge_influence_combine_method: attention
+edge_addition_filter: [0.25, 0.5, 0.75, 1.0]
+edge_removal_filter: [1.0, 0.0]
+override_attention_radius: []
+incl_robot_node: False
+map_encoding: False
+augment: True
+node_freq_mult_train: False
+node_freq_mult_eval: False
+scene_freq_mult_train: False
+scene_freq_mult_eval: False
+scene_freq_mult_viz: False
+no_edge_encoding: False
+device: cuda
+eval_device: None

MID/dataset/__init__.py

@@ -0,0 +1,2 @@
+from .dataset import EnvironmentDataset, NodeTypeDataset
+from .preprocessing import collate, get_node_timestep_data, get_timesteps_data, restore

MID/dataset/dataset.py

@@ -0,0 +1,76 @@
+from torch.utils import data
+import numpy as np
+from .preprocessing import get_node_timestep_data
+
+
+class EnvironmentDataset(object):
+    def __init__(self, env, state, pred_state, node_freq_mult, scene_freq_mult, hyperparams, **kwargs):
+        self.env = env
+        self.state = state
+        self.pred_state = pred_state
+        self.hyperparams = hyperparams
+        self.max_ht = self.hyperparams['maximum_history_length']
+        self.max_ft = kwargs['min_future_timesteps']
+        self.node_type_datasets = list()
+        self._augment = False
+        for node_type in env.NodeType:
+            if node_type not in hyperparams['pred_state']:
+                continue
+            self.node_type_datasets.append(NodeTypeDataset(env, node_type, state, pred_state, node_freq_mult,
+                                                           scene_freq_mult, hyperparams, **kwargs))
+
+    @property
+    def augment(self):
+        return self._augment
+
+    @augment.setter
+    def augment(self, value):
+        self._augment = value
+        for node_type_dataset in self.node_type_datasets:
+            node_type_dataset.augment = value
+
+    def __iter__(self):
+        return iter(self.node_type_datasets)
+
+
+class NodeTypeDataset(data.Dataset):
+    def __init__(self, env, node_type, state, pred_state, node_freq_mult,
+                 scene_freq_mult, hyperparams, augment=False, **kwargs):
+        self.env = env
+        self.state = state
+        self.pred_state = pred_state
+        self.hyperparams = hyperparams
+        self.max_ht = self.hyperparams['maximum_history_length']
+        self.max_ft = kwargs['min_future_timesteps']
+
+        self.augment = augment
+
+        self.node_type = node_type
+        self.index = self.index_env(node_freq_mult, scene_freq_mult, **kwargs)
+        self.len = len(self.index)
+        self.edge_types = [edge_type for edge_type in env.get_edge_types() if edge_type[0] is node_type]
+
+    def index_env(self, node_freq_mult, scene_freq_mult, **kwargs):
+        index = list()
+        for scene in self.env.scenes:
+            present_node_dict = scene.present_nodes(np.arange(0, scene.timesteps), type=self.node_type, **kwargs)
+            for t, nodes in present_node_dict.items():
+                for node in nodes:
+                    index += [(scene, t, node)] *\
+                             (scene.frequency_multiplier if scene_freq_mult else 1) *\
+                             (node.frequency_multiplier if node_freq_mult else 1)
+
+        return index
+
+    def __len__(self):
+        return self.len
+
+    def __getitem__(self, i):
+        (scene, t, node) = self.index[i]
+
+        if self.augment:
+            scene = scene.augment()
+            node = scene.get_node_by_id(node.id)
+
+        return get_node_timestep_data(self.env, scene, t, node, self.state, self.pred_state,
+                                      self.edge_types, self.max_ht, self.max_ft, self.hyperparams)

MID/dataset/homography_warper.py

@@ -0,0 +1,471 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Tuple, Optional
+
+
+pi = torch.tensor(3.14159265358979323846)
+
+
+def deg2rad(tensor: torch.Tensor) -> torch.Tensor:
+    r"""Function that converts angles from degrees to radians.
+    Args:
+        tensor (torch.Tensor): Tensor of arbitrary shape.
+    Returns:
+        torch.Tensor: tensor with same shape as input.
+    """
+    if not isinstance(tensor, torch.Tensor):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
+            type(tensor)))
+
+    return tensor * pi.to(tensor.device).type(tensor.dtype) / 180.
+
+
+def angle_to_rotation_matrix(angle: torch.Tensor) -> torch.Tensor:
+    """
+    Creates a rotation matrix out of angles in degrees
+    Args:
+        angle: (torch.Tensor): tensor of angles in degrees, any shape.
+    Returns:
+        torch.Tensor: tensor of *x2x2 rotation matrices.
+    Shape:
+        - Input: :math:`(*)`
+        - Output: :math:`(*, 2, 2)`
+    Example:
+        >>> input = torch.rand(1, 3)  # Nx3
+        >>> output = kornia.angle_to_rotation_matrix(input)  # Nx3x2x2
+    """
+    ang_rad = deg2rad(angle)
+    cos_a: torch.Tensor = torch.cos(ang_rad)
+    sin_a: torch.Tensor = torch.sin(ang_rad)
+    return torch.stack([cos_a, sin_a, -sin_a, cos_a], dim=-1).view(*angle.shape, 2, 2)
+
+
+def get_rotation_matrix2d(
+        center: torch.Tensor,
+        angle: torch.Tensor,
+        scale: torch.Tensor) -> torch.Tensor:
+    r"""Calculates an affine matrix of 2D rotation.
+    The function calculates the following matrix:
+    .. math::
+        \begin{bmatrix}
+            \alpha & \beta & (1 - \alpha) \cdot \text{x}
+            - \beta \cdot \text{y} \\
+            -\beta & \alpha & \beta \cdot \text{x}
+            + (1 - \alpha) \cdot \text{y}
+        \end{bmatrix}
+    where
+    .. math::
+        \alpha = \text{scale} \cdot cos(\text{radian}) \\
+        \beta = \text{scale} \cdot sin(\text{radian})
+    The transformation maps the rotation center to itself
+    If this is not the target, adjust the shift.
+    Args:
+        center (Tensor): center of the rotation in the source image.
+        angle (Tensor): rotation radian in degrees. Positive values mean
+            counter-clockwise rotation (the coordinate origin is assumed to
+            be the top-left corner).
+        scale (Tensor): isotropic scale factor.
+    Returns:
+        Tensor: the affine matrix of 2D rotation.
+    Shape:
+        - Input: :math:`(B, 2)`, :math:`(B)` and :math:`(B)`
+        - Output: :math:`(B, 2, 3)`
+    Example:
+        >>> center = torch.zeros(1, 2)
+        >>> scale = torch.ones(1)
+        >>> radian = 45. * torch.ones(1)
+        >>> M = kornia.get_rotation_matrix2d(center, radian, scale)
+        tensor([[[ 0.7071,  0.7071,  0.0000],
+                 [-0.7071,  0.7071,  0.0000]]])
+    """
+    if not torch.is_tensor(center):
+        raise TypeError("Input center type is not a torch.Tensor. Got {}"
+                        .format(type(center)))
+    if not torch.is_tensor(angle):
+        raise TypeError("Input radian type is not a torch.Tensor. Got {}"
+                        .format(type(angle)))
+    if not torch.is_tensor(scale):
+        raise TypeError("Input scale type is not a torch.Tensor. Got {}"
+                        .format(type(scale)))
+    if not (len(center.shape) == 2 and center.shape[1] == 2):
+        raise ValueError("Input center must be a Bx2 tensor. Got {}"
+                         .format(center.shape))
+    if not len(angle.shape) == 1:
+        raise ValueError("Input radian must be a B tensor. Got {}"
+                         .format(angle.shape))
+    if not len(scale.shape) == 1:
+        raise ValueError("Input scale must be a B tensor. Got {}"
+                         .format(scale.shape))
+    if not (center.shape[0] == angle.shape[0] == scale.shape[0]):
+        raise ValueError("Inputs must have same batch size dimension. Got {}"
+                         .format(center.shape, angle.shape, scale.shape))
+    # convert radian and apply scale
+    scaled_rotation: torch.Tensor = angle_to_rotation_matrix(angle) * scale.view(-1, 1, 1)
+    alpha: torch.Tensor = scaled_rotation[:, 0, 0]
+    beta: torch.Tensor = scaled_rotation[:, 0, 1]
+
+    # unpack the center to x, y coordinates
+    x: torch.Tensor = center[..., 0]
+    y: torch.Tensor = center[..., 1]
+
+    # create output tensor
+    batch_size: int = center.shape[0]
+    M: torch.Tensor = torch.zeros(
+        batch_size, 2, 3, device=center.device, dtype=center.dtype)
+    M[..., 0:2, 0:2] = scaled_rotation
+    M[..., 0, 2] = (torch.tensor(1.) - alpha) * x - beta * y
+    M[..., 1, 2] = beta * x + (torch.tensor(1.) - alpha) * y
+    return M
+
+def convert_points_to_homogeneous(points: torch.Tensor) -> torch.Tensor:
+    r"""Function that converts points from Euclidean to homogeneous space.
+    Examples::
+        >>> input = torch.rand(2, 4, 3)  # BxNx3
+        >>> output = kornia.convert_points_to_homogeneous(input)  # BxNx4
+    """
+    if not isinstance(points, torch.Tensor):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
+            type(points)))
+    if len(points.shape) < 2:
+        raise ValueError("Input must be at least a 2D tensor. Got {}".format(
+            points.shape))
+
+    return torch.nn.functional.pad(points, [0, 1], "constant", 1.0)
+
+
+def convert_points_from_homogeneous(
+        points: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+    r"""Function that converts points from homogeneous to Euclidean space.
+    Examples::
+        >>> input = torch.rand(2, 4, 3)  # BxNx3
+        >>> output = kornia.convert_points_from_homogeneous(input)  # BxNx2
+    """
+    if not isinstance(points, torch.Tensor):
+        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
+            type(points)))
+
+    if len(points.shape) < 2:
+        raise ValueError("Input must be at least a 2D tensor. Got {}".format(
+            points.shape))
+
+    # we check for points at infinity
+    z_vec: torch.Tensor = points[..., -1:]
+
+    # set the results of division by zeror/near-zero to 1.0
+    # follow the convention of opencv:
+    # https://github.com/opencv/opencv/pull/14411/files
+    mask: torch.Tensor = torch.abs(z_vec) > eps
+    scale: torch.Tensor = torch.ones_like(z_vec).masked_scatter_(
+        mask, torch.tensor(1.0).to(points.device) / z_vec[mask])
+
+    return scale * points[..., :-1]
+
+def transform_points(trans_01: torch.Tensor,
+                     points_1: torch.Tensor) -> torch.Tensor:
+    r"""Function that applies transformations to a set of points.
+    Args:
+        trans_01 (torch.Tensor): tensor for transformations of shape
+          :math:`(B, D+1, D+1)`.
+        points_1 (torch.Tensor): tensor of points of shape :math:`(B, N, D)`.
+    Returns:
+        torch.Tensor: tensor of N-dimensional points.
+    Shape:
+        - Output: :math:`(B, N, D)`
+    Examples:
+        >>> points_1 = torch.rand(2, 4, 3)  # BxNx3
+        >>> trans_01 = torch.eye(4).view(1, 4, 4)  # Bx4x4
+        >>> points_0 = kornia.transform_points(trans_01, points_1)  # BxNx3
+    """
+    if not torch.is_tensor(trans_01) or not torch.is_tensor(points_1):
+        raise TypeError("Input type is not a torch.Tensor")
+    if not trans_01.device == points_1.device:
+        raise TypeError("Tensor must be in the same device")
+    if not trans_01.shape[0] == points_1.shape[0] and trans_01.shape[0] != 1:
+        raise ValueError("Input batch size must be the same for both tensors or 1")
+    if not trans_01.shape[-1] == (points_1.shape[-1] + 1):
+        raise ValueError("Last input dimensions must differe by one unit")
+    # to homogeneous
+    points_1_h = convert_points_to_homogeneous(points_1)  # BxNxD+1
+    # transform coordinates
+    points_0_h = torch.matmul(
+        trans_01.unsqueeze(1), points_1_h.unsqueeze(-1))
+    points_0_h = torch.squeeze(points_0_h, dim=-1)
+    # to euclidean
+    points_0 = convert_points_from_homogeneous(points_0_h)  # BxNxD
+    return points_0
+
+
+def multi_linspace(a, b, num, endpoint=True, device='cpu', dtype=torch.float):
+    """This function is just like np.linspace, but will create linearly
+    spaced vectors from a start to end vector.
+    Inputs:
+        a - Start vector.
+        b - End vector.
+        num - Number of samples to generate. Default is 50. Must be above 0.
+        endpoint - If True, b is the last sample.
+                   Otherwise, it is not included. Default is True.
+    """
+
+    return a[..., None] + (b-a)[..., None]/(num-endpoint) * torch.arange(num, device=device, dtype=dtype)
+
+
+def create_batched_meshgrid(
+        x_min: torch.Tensor,
+        y_min: torch.Tensor,
+        x_max: torch.Tensor,
+        y_max: torch.Tensor,
+        height: int,
+        width: int,
+        device: Optional[torch.device] = torch.device('cpu')) -> torch.Tensor:
+    """Generates a coordinate grid for an image.
+    When the flag `normalized_coordinates` is set to True, the grid is
+    normalized to be in the range [-1,1] to be consistent with the pytorch
+    function grid_sample.
+    http://pytorch.org/docs/master/nn.html#torch.nn.functional.grid_sample
+    Args:
+        height (int): the image height (rows).
+        width (int): the image width (cols).
+        normalized_coordinates (Optional[bool]): whether to normalize
+          coordinates in the range [-1, 1] in order to be consistent with the
+          PyTorch function grid_sample.
+    Return:
+        torch.Tensor: returns a grid tensor with shape :math:`(1, H, W, 2)`.
+    """
+    # generate coordinates
+    xs = multi_linspace(x_min, x_max, width, device=device, dtype=torch.float)
+    ys = multi_linspace(y_min, y_max, height, device=device, dtype=torch.float)
+
+    # generate grid by stacking coordinates
+    bs = x_min.shape[0]
+    batched_grid_i_list = list()
+    for i in range(bs):
+        batched_grid_i_list.append(torch.stack(torch.meshgrid([xs[i], ys[i]])).transpose(1, 2))  # 2xHxW
+    batched_grid: torch.Tensor = torch.stack(batched_grid_i_list, dim=0)
+    return batched_grid.permute(0, 2, 3, 1)  # BxHxWx2
+
+
+def homography_warp(patch_src: torch.Tensor,
+                    centers: torch.Tensor,
+                    dst_homo_src: torch.Tensor,
+                    dsize: Tuple[int, int],
+                    mode: str = 'bilinear',
+                    padding_mode: str = 'zeros') -> torch.Tensor:
+    r"""Function that warps image patchs or tensors by homographies.
+    See :class:`~kornia.geometry.warp.HomographyWarper` for details.
+    Args:
+        patch_src (torch.Tensor): The image or tensor to warp. Should be from
+                                  source of shape :math:`(N, C, H, W)`.
+        dst_homo_src (torch.Tensor): The homography or stack of homographies
+                                     from source to destination of shape
+                                     :math:`(N, 3, 3)`.
+        dsize (Tuple[int, int]): The height and width of the image to warp.
+        mode (str): interpolation mode to calculate output values
+          'bilinear' | 'nearest'. Default: 'bilinear'.
+        padding_mode (str): padding mode for outside grid values
+          'zeros' | 'border' | 'reflection'. Default: 'zeros'.
+    Return:
+        torch.Tensor: Patch sampled at locations from source to destination.
+    Example:
+        >>> input = torch.rand(1, 3, 32, 32)
+        >>> homography = torch.eye(3).view(1, 3, 3)
+        >>> output = kornia.homography_warp(input, homography, (32, 32))
+    """
+
+    out_height, out_width = dsize
+    image_height, image_width = patch_src.shape[-2:]
+    x_min = 2. * (centers[..., 0] - out_width/2) / image_width - 1.
+    y_min = 2. * (centers[..., 1] - out_height/2) / image_height - 1.
+    x_max = 2. * (centers[..., 0] + out_width/2) / image_width - 1.
+    y_max = 2. * (centers[..., 1] + out_height/2) / image_height - 1.
+    warper = HomographyWarper(x_min, y_min, x_max, y_max, out_height, out_width, mode, padding_mode)
+    return warper(patch_src, dst_homo_src)
+
+
+def normal_transform_pixel(height, width):
+
+    tr_mat = torch.Tensor([[1.0, 0.0, -1.0],
+                           [0.0, 1.0, -1.0],
+                           [0.0, 0.0, 1.0]])  # 1x3x3
+
+    tr_mat[0, 0] = tr_mat[0, 0] * 2.0 / (width - 1.0)
+    tr_mat[1, 1] = tr_mat[1, 1] * 2.0 / (height - 1.0)
+
+    tr_mat = tr_mat.unsqueeze(0)
+
+    return tr_mat
+
+
+def src_norm_to_dst_norm(dst_pix_trans_src_pix: torch.Tensor,
+                         dsize_src: Tuple[int, int], dsize_dst: Tuple[int, int]) -> torch.Tensor:
+    # source and destination sizes
+    src_h, src_w = dsize_src
+    dst_h, dst_w = dsize_dst
+    # the devices and types
+    device: torch.device = dst_pix_trans_src_pix.device
+    dtype: torch.dtype = dst_pix_trans_src_pix.dtype
+    # compute the transformation pixel/norm for src/dst
+    src_norm_trans_src_pix: torch.Tensor = normal_transform_pixel(
+        src_h, src_w).to(device, dtype)
+    src_pix_trans_src_norm = torch.inverse(src_norm_trans_src_pix)
+    dst_norm_trans_dst_pix: torch.Tensor = normal_transform_pixel(
+        dst_h, dst_w).to(device, dtype)
+    # compute chain transformations
+    dst_norm_trans_src_norm: torch.Tensor = (
+        dst_norm_trans_dst_pix @ (dst_pix_trans_src_pix @ src_pix_trans_src_norm)
+    )
+    return dst_norm_trans_src_norm
+
+
+def transform_warp_impl(src: torch.Tensor, centers: torch.Tensor, dst_pix_trans_src_pix: torch.Tensor,
+                        dsize_src: Tuple[int, int], dsize_dst: Tuple[int, int],
+                        grid_mode: str, padding_mode: str) -> torch.Tensor:
+    """Compute the transform in normalized cooridnates and perform the warping.
+    """
+    dst_norm_trans_src_norm: torch.Tensor = src_norm_to_dst_norm(
+        dst_pix_trans_src_pix, dsize_src, dsize_src)
+
+    src_norm_trans_dst_norm = torch.inverse(dst_norm_trans_src_norm)
+    return homography_warp(src, centers, src_norm_trans_dst_norm, dsize_dst, grid_mode, padding_mode)
+
+
+class HomographyWarper(nn.Module):
+    r"""Warps image patches or tensors by homographies.
+    .. math::
+        X_{dst} = H_{src}^{\{dst\}} * X_{src}
+    Args:
+        height (int): The height of the image to warp.
+        width (int): The width of the image to warp.
+        mode (str): interpolation mode to calculate output values
+          'bilinear' | 'nearest'. Default: 'bilinear'.
+        padding_mode (str): padding mode for outside grid values
+          'zeros' | 'border' | 'reflection'. Default: 'zeros'.
+    """
+
+    def __init__(
+            self,
+            x_min: torch.Tensor,
+            y_min: torch.Tensor,
+            x_max: torch.Tensor,
+            y_max: torch.Tensor,
+            height: int,
+            width: int,
+            mode: str = 'bilinear',
+            padding_mode: str = 'zeros') -> None:
+        super(HomographyWarper, self).__init__()
+        self.width: int = width
+        self.height: int = height
+        self.mode: str = mode
+        self.padding_mode: str = padding_mode
+
+        # create base grid to compute the flow
+        self.grid: torch.Tensor = create_batched_meshgrid(x_min, y_min, x_max, y_max, height, width)
+
+    def warp_grid(self, dst_homo_src: torch.Tensor) -> torch.Tensor:
+        r"""Computes the grid to warp the coordinates grid by an homography.
+        Args:
+            dst_homo_src (torch.Tensor): Homography or homographies (stacked) to
+                              transform all points in the grid. Shape of the
+                              homography has to be :math:`(N, 3, 3)`.
+        Returns:
+            torch.Tensor: the transformed grid of shape :math:`(N, H, W, 2)`.
+        """
+        batch_size: int = dst_homo_src.shape[0]
+        device: torch.device = dst_homo_src.device
+        dtype: torch.dtype = dst_homo_src.dtype
+        # expand grid to match the input batch size
+        grid: torch.Tensor = self.grid
+        if len(dst_homo_src.shape) == 3:  # local homography case
+            dst_homo_src = dst_homo_src.view(batch_size, 1, 3, 3)  # NxHxWx3x3
+        # perform the actual grid transformation,
+        # the grid is copied to input device and casted to the same type
+        flow: torch.Tensor = transform_points(
+            dst_homo_src, grid.to(device).to(dtype))  # NxHxWx2
+        return flow.view(batch_size, self.height, self.width, 2)  # NxHxWx2
+
+    def forward(  # type: ignore
+            self,
+            patch_src: torch.Tensor,
+            dst_homo_src: torch.Tensor) -> torch.Tensor:
+        r"""Warps an image or tensor from source into reference frame.
+        Args:
+            patch_src (torch.Tensor): The image or tensor to warp.
+                                      Should be from source.
+            dst_homo_src (torch.Tensor): The homography or stack of homographies
+             from source to destination. The homography assumes normalized
+             coordinates [-1, 1].
+        Return:
+            torch.Tensor: Patch sampled at locations from source to destination.
+        Shape:
+            - Input: :math:`(N, C, H, W)` and :math:`(N, 3, 3)`
+            - Output: :math:`(N, C, H, W)`
+        Example:
+            >>> input = torch.rand(1, 3, 32, 32)
+            >>> homography = torch.eye(3).view(1, 3, 3)
+            >>> warper = kornia.HomographyWarper(32, 32)
+            >>> output = warper(input, homography)  # NxCxHxW
+        """
+        if not dst_homo_src.device == patch_src.device:
+            raise TypeError("Patch and homography must be on the same device. \
+                            Got patch.device: {} dst_H_src.device: {}."
+                            .format(patch_src.device, dst_homo_src.device))
+
+        return F.grid_sample(patch_src, self.warp_grid(dst_homo_src),  # type: ignore
+                             mode=self.mode, padding_mode=self.padding_mode, align_corners=True)
+
+
+def warp_affine_crop(src: torch.Tensor, centers: torch.Tensor, M: torch.Tensor,
+                dsize: Tuple[int, int], flags: str = 'bilinear',
+                padding_mode: str = 'zeros') -> torch.Tensor:
+    r"""Applies an affine transformation to a tensor.
+
+    The function warp_affine transforms the source tensor using
+    the specified matrix:
+
+    .. math::
+        \text{dst}(x, y) = \text{src} \left( M_{11} x + M_{12} y + M_{13} ,
+        M_{21} x + M_{22} y + M_{23} \right )
+
+    Args:
+        src (torch.Tensor): input tensor of shape :math:`(B, C, H, W)`.
+        M (torch.Tensor): affine transformation of shape :math:`(B, 2, 3)`.
+        dsize (Tuple[int, int]): size of the output image (height, width).
+        mode (str): interpolation mode to calculate output values
+          'bilinear' | 'nearest'. Default: 'bilinear'.
+        padding_mode (str): padding mode for outside grid values
+          'zeros' | 'border' | 'reflection'. Default: 'zeros'.
+
+    Returns:
+        torch.Tensor: the warped tensor.
+
+    Shape:
+        - Output: :math:`(B, C, H, W)`
+
+    .. note::
+       See a working example `here <https://kornia.readthedocs.io/en/latest/
+       tutorials/warp_affine.html>`__.
+    """
+    if not torch.is_tensor(src):
+        raise TypeError("Input src type is not a torch.Tensor. Got {}"
+                        .format(type(src)))
+
+    if not torch.is_tensor(M):
+        raise TypeError("Input M type is not a torch.Tensor. Got {}"
+                        .format(type(M)))
+
+    if not len(src.shape) == 4:
+        raise ValueError("Input src must be a BxCxHxW tensor. Got {}"
+                         .format(src.shape))
+
+    if not (len(M.shape) == 3 or M.shape[-2:] == (2, 3)):
+        raise ValueError("Input M must be a Bx2x3 tensor. Got {}"
+                         .format(src.shape))
+
+    # we generate a 3x3 transformation matrix from 2x3 affine
+    M_3x3: torch.Tensor = F.pad(M, [0, 0, 0, 1, 0, 0],
+                                mode="constant", value=0)
+    M_3x3[:, 2, 2] += 1.0
+
+    # launches the warper
+    h, w = src.shape[-2:]
+    return transform_warp_impl(src, centers, M_3x3, (h, w), dsize, flags, padding_mode)

MID/dataset/preprocessing.py

@@ -0,0 +1,233 @@
+import torch
+import numpy as np
+import collections.abc
+from torch.utils.data._utils.collate import default_collate
+import dill
+container_abcs = collections.abc
+
+
+def restore(data):
+    """
+    In case we dilled some structures to share between multiple process this function will restore them.
+    If the data input are not bytes we assume it was not dilled in the first place
+
+    :param data: Possibly dilled data structure
+    :return: Un-dilled data structure
+    """
+    if type(data) is bytes:
+        return dill.loads(data)
+    return data
+
+
+def collate(batch):
+    if len(batch) == 0:
+        return batch
+    elem = batch[0]
+    if elem is None:
+        return None
+    elif isinstance(elem, container_abcs.Sequence):
+        if len(elem) == 4: # We assume those are the maps, map points, headings and patch_size
+            scene_map, scene_pts, heading_angle, patch_size = zip(*batch)
+            if heading_angle[0] is None:
+                heading_angle = None
+            else:
+                heading_angle = torch.Tensor(heading_angle)
+            map = scene_map[0].get_cropped_maps_from_scene_map_batch(scene_map,
+                                                                     scene_pts=torch.Tensor(scene_pts),
+                                                                     patch_size=patch_size[0],
+                                                                     rotation=heading_angle)
+            return map
+        transposed = zip(*batch)
+        return [collate(samples) for samples in transposed]
+    elif isinstance(elem, container_abcs.Mapping):
+        # We have to dill the neighbors structures. Otherwise each tensor is put into
+        # shared memory separately -> slow, file pointer overhead
+        # we only do this in multiprocessing
+        neighbor_dict = {key: [d[key] for d in batch] for key in elem}
+        return dill.dumps(neighbor_dict) if torch.utils.data.get_worker_info() else neighbor_dict
+    return default_collate(batch)
+
+
+def get_relative_robot_traj(env, state, node_traj, robot_traj, node_type, robot_type):
+    # TODO: We will have to make this more generic if robot_type != node_type
+    # Make Robot State relative to node
+    _, std = env.get_standardize_params(state[robot_type], node_type=robot_type)
+    std[0:2] = env.attention_radius[(node_type, robot_type)]
+    robot_traj_st = env.standardize(robot_traj,
+                                    state[robot_type],
+                                    node_type=robot_type,
+                                    mean=node_traj,
+                                    std=std)
+    robot_traj_st_t = torch.tensor(robot_traj_st, dtype=torch.float)
+
+    return robot_traj_st_t
+
+
+def get_node_timestep_data(env, scene, t, node, state, pred_state,
+                           edge_types, max_ht, max_ft, hyperparams,
+                           scene_graph=None):
+    """
+    Pre-processes the data for a single batch element: node state over time for a specific time in a specific scene
+    as well as the neighbour data for it.
+
+    :param env: Environment
+    :param scene: Scene
+    :param t: Timestep in scene
+    :param node: Node
+    :param state: Specification of the node state
+    :param pred_state: Specification of the prediction state
+    :param edge_types: List of all Edge Types for which neighbours are pre-processed
+    :param max_ht: Maximum history timesteps
+    :param max_ft: Maximum future timesteps (prediction horizon)
+    :param hyperparams: Model hyperparameters
+    :param scene_graph: If scene graph was already computed for this scene and time you can pass it here
+    :return: Batch Element
+    """
+
+    # Node
+    timestep_range_x = np.array([t - max_ht, t])
+    timestep_range_y = np.array([t + 1, t + max_ft])
+
+    x = node.get(timestep_range_x, state[node.type])
+    y = node.get(timestep_range_y, pred_state[node.type])
+    first_history_index = (max_ht - node.history_points_at(t)).clip(0)
+
+    _, std = env.get_standardize_params(state[node.type], node.type)
+    std[0:2] = env.attention_radius[(node.type, node.type)]
+    rel_state = np.zeros_like(x[0])
+    rel_state[0:2] = np.array(x)[-1, 0:2]
+    x_st = env.standardize(x, state[node.type], node.type, mean=rel_state, std=std)
+    if list(pred_state[node.type].keys())[0] == 'position':  # If we predict position we do it relative to current pos
+        y_st = env.standardize(y, pred_state[node.type], node.type, mean=rel_state[0:2])
+    else:
+        y_st = env.standardize(y, pred_state[node.type], node.type)
+
+    x_t = torch.tensor(x, dtype=torch.float)
+    y_t = torch.tensor(y, dtype=torch.float)
+    x_st_t = torch.tensor(x_st, dtype=torch.float)
+    y_st_t = torch.tensor(y_st, dtype=torch.float)
+
+    # Neighbors
+    neighbors_data_st = None
+    neighbors_edge_value = None
+    if hyperparams['edge_encoding']:
+        # Scene Graph
+        scene_graph = scene.get_scene_graph(t,
+                                            env.attention_radius,
+                                            hyperparams['edge_addition_filter'],
+                                            hyperparams['edge_removal_filter']) if scene_graph is None else scene_graph
+
+        neighbors_data_st = dict()
+        neighbors_edge_value = dict()
+        for edge_type in edge_types:
+            neighbors_data_st[edge_type] = list()
+            # We get all nodes which are connected to the current node for the current timestep
+            connected_nodes = scene_graph.get_neighbors(node, edge_type[1])
+
+            if hyperparams['dynamic_edges'] == 'yes':
+                # We get the edge masks for the current node at the current timestep
+                edge_masks = torch.tensor(scene_graph.get_edge_scaling(node), dtype=torch.float)
+                neighbors_edge_value[edge_type] = edge_masks
+
+            for connected_node in connected_nodes:
+                neighbor_state_np = connected_node.get(np.array([t - max_ht, t]),
+                                                       state[connected_node.type],
+                                                       padding=0.0)
+
+                # Make State relative to node where neighbor and node have same state
+                _, std = env.get_standardize_params(state[connected_node.type], node_type=connected_node.type)
+                std[0:2] = env.attention_radius[edge_type]
+                equal_dims = np.min((neighbor_state_np.shape[-1], x.shape[-1]))
+                rel_state = np.zeros_like(neighbor_state_np)
+                rel_state[:, ..., :equal_dims] = x[-1, ..., :equal_dims]
+                neighbor_state_np_st = env.standardize(neighbor_state_np,
+                                                       state[connected_node.type],
+                                                       node_type=connected_node.type,
+                                                       mean=rel_state,
+                                                       std=std)
+
+                neighbor_state = torch.tensor(neighbor_state_np_st, dtype=torch.float)
+                neighbors_data_st[edge_type].append(neighbor_state)
+
+    # Robot
+    robot_traj_st_t = None
+    timestep_range_r = np.array([t, t + max_ft])
+    if hyperparams['incl_robot_node']:
+        x_node = node.get(timestep_range_r, state[node.type])
+        if scene.non_aug_scene is not None:
+            robot = scene.get_node_by_id(scene.non_aug_scene.robot.id)
+        else:
+            robot = scene.robot
+        robot_type = robot.type
+        robot_traj = robot.get(timestep_range_r, state[robot_type], padding=0.0)
+        robot_traj_st_t = get_relative_robot_traj(env, state, x_node, robot_traj, node.type, robot_type)
+
+    # Map
+    map_tuple = None
+    if hyperparams['use_map_encoding']:
+        if node.type in hyperparams['map_encoder']:
+            if node.non_aug_node is not None:
+                x = node.non_aug_node.get(np.array([t]), state[node.type])
+            me_hyp = hyperparams['map_encoder'][node.type]
+            if 'heading_state_index' in me_hyp:
+                heading_state_index = me_hyp['heading_state_index']
+                # We have to rotate the map in the opposit direction of the agent to match them
+                if type(heading_state_index) is list:  # infer from velocity or heading vector
+                    heading_angle = -np.arctan2(x[-1, heading_state_index[1]],
+                                                x[-1, heading_state_index[0]]) * 180 / np.pi
+                else:
+                    heading_angle = -x[-1, heading_state_index] * 180 / np.pi
+            else:
+                heading_angle = None
+
+            scene_map = scene.map[node.type]
+            map_point = x[-1, :2]
+
+
+            patch_size = hyperparams['map_encoder'][node.type]['patch_size']
+            map_tuple = (scene_map, map_point, heading_angle, patch_size)
+
+    return (first_history_index, x_t, y_t, x_st_t, y_st_t, neighbors_data_st,
+            neighbors_edge_value, robot_traj_st_t, map_tuple)
+
+
+def get_timesteps_data(env, scene, t, node_type, state, pred_state,
+                       edge_types, min_ht, max_ht, min_ft, max_ft, hyperparams):
+    """
+    Puts together the inputs for ALL nodes in a given scene and timestep in it.
+
+    :param env: Environment
+    :param scene: Scene
+    :param t: Timestep in scene
+    :param node_type: Node Type of nodes for which the data shall be pre-processed
+    :param state: Specification of the node state
+    :param pred_state: Specification of the prediction state
+    :param edge_types: List of all Edge Types for which neighbors are pre-processed
+    :param max_ht: Maximum history timesteps
+    :param max_ft: Maximum future timesteps (prediction horizon)
+    :param hyperparams: Model hyperparameters
+    :return:
+    """
+    nodes_per_ts = scene.present_nodes(t,
+                                       type=node_type,
+                                       min_history_timesteps=min_ht,
+                                       min_future_timesteps=max_ft,
+                                       return_robot=not hyperparams['incl_robot_node'])
+    batch = list()
+    nodes = list()
+    out_timesteps = list()
+    for timestep in nodes_per_ts.keys():
+            scene_graph = scene.get_scene_graph(timestep,
+                                                env.attention_radius,
+                                                hyperparams['edge_addition_filter'],
+                                                hyperparams['edge_removal_filter'])
+            present_nodes = nodes_per_ts[timestep]
+            for node in present_nodes:
+                nodes.append(node)
+                out_timesteps.append(timestep)
+                batch.append(get_node_timestep_data(env, scene, timestep, node, state, pred_state,
+                                                    edge_types, max_ht, max_ft, hyperparams,
+                                                    scene_graph=scene_graph))
+    if len(out_timesteps) == 0:
+        return None
+    return collate(batch), nodes, out_timesteps