""" MID + Graph on SDD — v4 (output-level). Combines v3's correct fixes with v2's stable training recipe: From v3 (bug fixes): 1. Velocity→position: y_pos = cumsum(vel)*dt + init_pos 2. Position normalization: y_abs / pos_scale (SDD pixels ~±500, scale=100) 3. Trajectory-aware nodes: node_proj([context, y_vel_flat]) 4. No sigma/logvar (removed untrained uncertainty head) 5. Cosine LR schedule From v2 (stability): 1. Per-agent independent diffusion t (not shared — MID batches mix agents from different timesteps, shared t degrades base training) 2. Two-pass training: pass1 base-only → y_0_hat, pass2 base+graph with y_0_hat edges. Avoids GT train/eval gap. 3. delta_scale=0.1, gate_init=0.1 (moderate — v2's 0.05 too weak, v3's 0.3 too strong) 4. graph_warmup=5 epochs Target: beat previous graph best ADE=8.53 (baseline=8.27). """ import os, sys, time, logging, argparse, math, random import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch import optim from torch.utils.tensorboard import SummaryWriter # tbX-broken import dill from dataset import EnvironmentDataset, collate, get_timesteps_data, restore from models.autoencoder import AutoEncoder from models.trajectron import Trajectron from utils.model_registrar import ModelRegistrar from utils.trajectron_hypers import get_traj_hypers from models.diffusion import DiffusionTraj, VarianceSchedule, TransformerConcatLinear import evaluation MOFLOW_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'MoFlow')) sys.path.insert(0, MOFLOW_ROOT) from models.graph_interaction_nba_v6 import FutureInteractionGraphV6 from models.context_encoder.mtr_encoder import SinusoidalPosEmb class GraphDenoiserWrapperV4(nn.Module): def __init__(self, base_net, encoder_dim=256, pred_len=12, graph_hidden=128, top_n=5, num_gnn_layers=2, graph_dropout=0.1, dt=0.4, pos_scale=100.0): super().__init__() self.base_net = base_net self.pred_len = pred_len self.graph_hidden = graph_hidden self.max_top_n = top_n self.dt = dt self.pos_scale = pos_scale self.node_proj = nn.Sequential( nn.Linear(encoder_dim + pred_len * 2, graph_hidden), nn.ReLU(inplace=True)) self.time_mlp = nn.Sequential( SinusoidalPosEmb(graph_hidden), nn.Linear(graph_hidden, graph_hidden), nn.ReLU(), nn.Linear(graph_hidden, graph_hidden)) self.future_graph = FutureInteractionGraphV6( embed_dim=graph_hidden, future_steps=pred_len, num_agents=64, num_heads=4, dropout=graph_dropout, num_gnn_layers=num_gnn_layers, time_dim=graph_hidden, top_n_neighbors=min(top_n, 63), rel_traj_hidden=32, y0_score_dim=32) self.graph_out_proj = nn.Sequential( nn.Linear(graph_hidden, graph_hidden), nn.ReLU(inplace=True), nn.Linear(graph_hidden, pred_len * 2), nn.Tanh()) nn.init.zeros_(self.graph_out_proj[-2].weight) nn.init.zeros_(self.graph_out_proj[-2].bias) gi = 0.1 self.raw_gate = nn.Parameter(torch.tensor(math.log(gi / (1.0 - gi)))) self.register_buffer('delta_scale', torch.tensor(0.1)) def forward(self, x_t, beta, context, y_vel=None, init_pos=None, skip_graph=False): """ x_t: [N, T, 2] noisy trajectory (velocity space) beta: [N] diffusion beta context: [N, enc_dim] encoder output y_vel: [N, T, 2] velocity for graph edge features (y_0_hat from pass1) init_pos: [N, 2] last observed absolute position """ eps_pred = self.base_net(x_t, beta=beta, context=context) N = x_t.size(0) T = self.pred_len if (not skip_graph) and y_vel is not None and init_pos is not None and N >= 2: D = self.graph_hidden y_vel_flat = y_vel.view(N, T * 2) node_input = torch.cat([context, y_vel_flat], dim=-1) node_emb = self.node_proj(node_input).view(1, 1, N, D) y_pos = (torch.cumsum(y_vel.view(N, T, 2), dim=1) * self.dt + init_pos.unsqueeze(1)) y_abs = (y_pos / self.pos_scale).view(1, 1, N, T, 2) beta_scene = beta.mean().unsqueeze(0) tau = (beta_scene / 0.05).clamp(0, 1) t_emb = self.time_mlp(beta_scene) y_emb_out = self.future_graph( node_emb, y_abs, t_emb, tau, sigma_agent=None) graph_out = self.graph_out_proj( y_emb_out.squeeze(1).squeeze(0)) delta = graph_out.view(N, T, 2) * self.delta_scale gate = torch.sigmoid(self.raw_gate) eps_pred = eps_pred + gate * delta return eps_pred class MIDGraphV4: def __init__(self, config): self.config = config torch.backends.cudnn.benchmark = True self._build() def _build(self): self._skip_graph_override = False self.model_dir = os.path.join("./experiments", self.config.exp_name) self.log_writer = SummaryWriter(log_dir=self.model_dir) os.makedirs(self.model_dir, exist_ok=True) log_name = f"sdd_{time.strftime('%Y-%m-%d-%H-%M')}.log" self.log = logging.getLogger(self.config.exp_name) self.log.setLevel(logging.INFO) self.log.addHandler(logging.FileHandler(os.path.join(self.model_dir, log_name))) self.log.addHandler(logging.StreamHandler()) self.log.info(f"Config: {self.config}") self.train_data_path = os.path.join(self.config.data_dir, "sdd_train.pkl") self.eval_data_path = os.path.join(self.config.data_dir, "sdd_test.pkl") self.hyperparams = get_traj_hypers() self.hyperparams['enc_rnn_dim_edge'] = self.config.encoder_dim // 2 self.hyperparams['enc_rnn_dim_edge_influence'] = self.config.encoder_dim // 2 self.hyperparams['enc_rnn_dim_history'] = self.config.encoder_dim // 2 self.hyperparams['enc_rnn_dim_future'] = self.config.encoder_dim // 2 self.registrar = ModelRegistrar(self.model_dir, "cuda") with open(self.train_data_path, 'rb') as f: self.train_env = dill.load(f, encoding='latin1') with open(self.eval_data_path, 'rb') as f: self.eval_env = dill.load(f, encoding='latin1') self.encoder = Trajectron(self.registrar, self.hyperparams, "cuda") self.encoder.set_environment(self.train_env) self.encoder.set_annealing_params() base_net = TransformerConcatLinear( point_dim=2, context_dim=self.config.encoder_dim, tf_layer=self.config.tf_layer, residual=False) self.graph_net = GraphDenoiserWrapperV4( base_net, encoder_dim=self.config.encoder_dim, pred_len=12, graph_hidden=128, top_n=self.config.top_n_neighbors, num_gnn_layers=self.config.graph_gnn_layers, graph_dropout=self.config.graph_dropout, dt=self.config.dt, pos_scale=self.config.pos_scale).cuda() if hasattr(self.config, 'graph_gate_init') and self.config.graph_gate_init is not None: gi = float(max(min(self.config.graph_gate_init, 0.999), 1e-4)) with torch.no_grad(): self.graph_net.raw_gate.fill_(math.log(gi / (1.0 - gi))) self.var_sched = VarianceSchedule(num_steps=100, beta_T=5e-2, mode='linear') graph_keys = ('future_graph', 'node_proj', 'time_mlp', 'graph_out_proj', 'raw_gate') self._graph_params = [p for n, p in self.graph_net.named_parameters() if any(k in n for k in graph_keys)] self._base_params = [p for n, p in self.graph_net.named_parameters() if not any(k in n for k in graph_keys)] self.optimizer = optim.Adam([ {'params': self.registrar.get_all_but_name_match('map_encoder').parameters()}, {'params': self.graph_net.parameters()}, ], lr=self.config.lr) warm = max(1, int(getattr(self.config, 'lr_warmup_epochs', 2))) from torch.optim.lr_scheduler import LambdaLR, CosineAnnealingLR, SequentialLR warm_sched = LambdaLR(self.optimizer, lr_lambda=lambda e: min(1.0, (e + 1) / warm)) cosine_sched = CosineAnnealingLR( self.optimizer, T_max=self.config.epochs - warm, eta_min=1e-5) self.scheduler = SequentialLR( self.optimizer, schedulers=[warm_sched, cosine_sched], milestones=[warm]) self.train_scenes = self.train_env.scenes self.eval_scenes = self.eval_env.scenes self.log.info(f"Train scenes: {len(self.train_scenes)}, " f"Eval scenes: {len(self.eval_scenes)}") def _get_loss(self, batch, node_type): (first_history_index, x_t_raw, y_t, x_st_t, y_st_t, neighbors_data_st, neighbors_edge_value, robot_traj_st_t, map_) = batch context = self.encoder.get_latent(batch, node_type) y_0 = y_t.cuda() # [N, 12, 2] velocities N = y_0.size(0) init_pos = x_t_raw[:, -1, 0:2].cuda() # [N, 2] absolute position # Per-agent independent diffusion t (v2 style — stable for MID) t = self.var_sched.uniform_sample_t(N) alpha_bar = self.var_sched.alpha_bars[t].cuda() beta = self.var_sched.betas[t].cuda() c0 = alpha_bar.sqrt().view(N, 1, 1) c1 = (1 - alpha_bar).sqrt().view(N, 1, 1) e_rand = torch.randn_like(y_0) x_noisy = c0 * y_0 + c1 * e_rand if self._skip_graph_override or N < 2: eps_pred = self.graph_net(x_noisy, beta, context, skip_graph=True) else: # Two-pass: pass1 base-only → y_0_hat, pass2 base+graph with torch.no_grad(): eps_base = self.graph_net(x_noisy, beta, context, skip_graph=True) y_0_hat = (x_noisy - c1 * eps_base) / c0 # [N, T, 2] velocity eps_pred = self.graph_net(x_noisy, beta, context, y_vel=y_0_hat, init_pos=init_pos) return F.mse_loss(eps_pred.reshape(-1, 2), e_rand.reshape(-1, 2)) def train(self): node_type = "PEDESTRIAN" ph = self.hyperparams['prediction_horizon'] max_hl = self.hyperparams['maximum_history_length'] graph_warm = int(getattr(self.config, 'graph_warmup_epochs', 5)) for epoch in range(1, self.config.epochs + 1): self.graph_net.train() total_loss, n_batches = 0.0, 0 self._skip_graph_override = (epoch <= graph_warm) for scene in self.train_scenes: for t_start in range(0, scene.timesteps, 10): timesteps = np.arange(t_start, t_start + 10) batch = get_timesteps_data( env=self.train_env, scene=scene, t=timesteps, node_type=node_type, state=self.hyperparams['state'], pred_state=self.hyperparams['pred_state'], edge_types=self.train_env.get_edge_types(), min_ht=1, max_ht=max_hl, min_ft=12, max_ft=12, hyperparams=self.hyperparams) if batch is None: continue loss = self._get_loss(batch[0], node_type) self.optimizer.zero_grad() loss.backward() nn.utils.clip_grad_norm_(self._graph_params, 0.1) nn.utils.clip_grad_norm_(self._base_params, 1.0) self.optimizer.step() total_loss += loss.item() n_batches += 1 self.scheduler.step() avg = total_loss / max(1, n_batches) self.log.info(f"Epoch {epoch} train_loss={avg:.4f}") self.log_writer.add_scalar('loss/train', avg, epoch) if epoch % self.config.eval_every == 0: ade, fde = self._eval(node_type, ph, max_hl) ade *= 50; fde *= 50 self.log.info(f"Epoch {epoch} Best Of 20: " f"ADE: {ade:.4f} FDE: {fde:.4f}") self.log_writer.add_scalar('metric/ADE', ade, epoch) self.log_writer.add_scalar('metric/FDE', fde, epoch) torch.save({ 'encoder': self.registrar.model_dict, 'graph_net': self.graph_net.state_dict(), }, os.path.join(self.model_dir, f"sdd_epoch{epoch}.pt")) @torch.no_grad() def _eval(self, node_type, ph, max_hl): self.graph_net.eval() ade_errors, fde_errors = [], [] for scene in self.eval_scenes: for t_start in range(0, scene.timesteps, 10): timesteps = np.arange(t_start, t_start + 10) batch = get_timesteps_data( env=self.eval_env, scene=scene, t=timesteps, node_type=node_type, state=self.hyperparams['state'], pred_state=self.hyperparams['pred_state'], edge_types=self.eval_env.get_edge_types(), min_ht=7, max_ht=max_hl, min_ft=12, max_ft=12, hyperparams=self.hyperparams) if batch is None: continue test_batch, nodes, timesteps_o = batch context = self.encoder.get_latent(test_batch, node_type) dynamics = self.encoder.node_models_dict[node_type].dynamic N = context.size(0) _, x_t_raw, *_ = test_batch init_pos = x_t_raw[:, -1, 0:2].cuda() preds = self._sample_with_graph( context, N, init_pos, num_points=12, K=20) predicted_y_pos = dynamics.integrate_samples(preds) predictions = predicted_y_pos.cpu().numpy() predictions_dict = {} for i, ts in enumerate(timesteps_o): if ts not in predictions_dict: predictions_dict[ts] = {} predictions_dict[ts][nodes[i]] = np.transpose( predictions[:, [i]], (1, 0, 2, 3)) batch_error = evaluation.compute_batch_statistics( predictions_dict, scene.dt, max_hl=max_hl, ph=ph, node_type_enum=self.eval_env.NodeType, kde=False, map=None, best_of=True, prune_ph_to_future=True) ade_errors = np.hstack( (ade_errors, batch_error[node_type]['ade'])) fde_errors = np.hstack( (fde_errors, batch_error[node_type]['fde'])) return np.mean(ade_errors), np.mean(fde_errors) def _sample_with_graph(self, context, N, init_pos, num_points=12, K=20): traj_list = [] stride = 5 for _ in range(K): x_t = torch.randn(N, num_points, 2, device=context.device) y_0_prev = None for t in range(self.var_sched.num_steps, 0, -stride): alpha_bar = self.var_sched.alpha_bars[t] alpha_bar_next = self.var_sched.alpha_bars[t - stride] beta = self.var_sched.betas[[t] * N].cuda() if y_0_prev is not None and N >= 2: eps = self.graph_net(x_t, beta, context, y_vel=y_0_prev, init_pos=init_pos) else: eps = self.graph_net(x_t, beta, context, skip_graph=True) x0_pred = ((x_t - (1 - alpha_bar).sqrt() * eps) / alpha_bar.sqrt()) y_0_prev = x0_pred x_t = (alpha_bar_next.sqrt() * x0_pred + (1 - alpha_bar_next).sqrt() * eps) traj_list.append(x_t) return torch.stack(traj_list) def main(): p = argparse.ArgumentParser() p.add_argument('--data_dir', default='processed_data') p.add_argument('--exp_name', default='mid_sdd_graph_v4') p.add_argument('--gpu', type=int, default=0) p.add_argument('--epochs', type=int, default=100) p.add_argument('--lr', type=float, default=1e-3) p.add_argument('--eval_every', type=int, default=3) p.add_argument('--encoder_dim', type=int, default=256) p.add_argument('--tf_layer', type=int, default=3) p.add_argument('--top_n_neighbors', type=int, default=5) p.add_argument('--graph_gnn_layers', type=int, default=2) p.add_argument('--graph_dropout', type=float, default=0.1) p.add_argument('--graph_gate_init', type=float, default=0.1) p.add_argument('--graph_warmup_epochs', type=int, default=5) p.add_argument('--lr_warmup_epochs', type=int, default=2) p.add_argument('--dt', type=float, default=0.4, help='Scene timestep (SDD: 0.4s at 2.5Hz)') p.add_argument('--pos_scale', type=float, default=100.0, help='Divide positions by this before graph') config = p.parse_args() torch.cuda.set_device(config.gpu) MIDGraphV4(config).train() if __name__ == '__main__': main()