"""SDD dataloader for MoFlow with MULTI-SCENE BATCHING via per-batch padding. Strategy: - Sort scenes by A (number of agents) at init time. - BatchSampler builds mini-batches from contiguous A-sorted indices so that all scenes in a batch have similar A → padding waste is small (≤ 2×). - Collate pads each batch to its own max_A and returns an `agent_mask` [B, max_A] (True = real agent, False = padding) that the backbone and graph module use to ignore padded slots. Scenes within a batch attend ONLY to agents in their own scene via a block-diagonal `traj_mask`; padded slots are zeroed out in the mask so no information leaks across scenes or into padding. """ import os, pickle, math, random import numpy as np import torch from torch.utils.data import Dataset, Sampler from utils.normalization import normalize_min_max def _sdd_to_scene(sdd, obs_len, pred_len): """Convert one raw scene [A, 20, 2] numpy array to torch tensors needed by MoFlow: (past_6ch, fut_rel, fut_vel, init_pos). Coords are already /50 and mean-centered per scene from preprocessing.""" traj = torch.from_numpy(sdd).float() # [A, 20, 2] pre = traj[:, :obs_len].contiguous() fut = traj[:, obs_len:].contiguous() init = pre[:, -1:] abs_ch = pre rel_ch = pre - init vel_ch = torch.cat([rel_ch[:, 1:] - rel_ch[:, :-1], torch.zeros_like(rel_ch[:, -1:])], dim=1) past_6 = torch.cat([abs_ch, rel_ch, vel_ch], dim=-1) # [A, T_past, 6] fut_r = fut - init # [A, T_future, 2] fut_vel = torch.cat([fut_r[:, 1:] - fut_r[:, :-1], torch.zeros_like(fut_r[:, -1:])], dim=1) return past_6, fut_r, fut_vel class SortedByASampler(Sampler): """Yield single indices in A-sorted order with bucket-local shuffling. Use with DataLoader(..., sampler=this, batch_size=B): consecutive indices land in the same mini-batch so batches end up A-similar. Bucket size W = 4 * batch_size controls how much A-shuffling happens per epoch. """ def __init__(self, a_counts, batch_size, shuffle=True, seed=0): self.a_counts = np.asarray(a_counts) self.batch_size = int(batch_size) self.shuffle = shuffle self.seed = seed self.epoch = 0 def set_epoch(self, epoch): self.epoch = int(epoch) def __iter__(self): order = np.argsort(self.a_counts, kind='stable') if self.shuffle: rng = np.random.RandomState(self.seed + self.epoch) W = 4 * self.batch_size for start in range(0, len(order), W): end = min(start + W, len(order)) rng.shuffle(order[start:end]) for idx in order.tolist(): yield int(idx) def __len__(self): return len(self.a_counts) def collate_padded(batch): """Pad a list of `(past_6, fut_rel, fut_vel)` scenes (each [A_i, T, *]) to a common max_A in this batch, stack, and emit an agent_mask. Returns the same dict structure as dataloader_nba's seq_collate_nba plus extra keys: 'agent_mask' [B, max_A] and 'A_per_scene' [B]. """ B = len(batch) A_list = [x[0].shape[0] for x in batch] max_A = max(A_list) T_past = batch[0][0].shape[1] T_fut = batch[0][1].shape[1] past_pad = torch.zeros(B, max_A, T_past, 6) fut_pad = torch.zeros(B, max_A, T_fut, 2) vel_pad = torch.zeros(B, max_A, T_fut, 2) agent_mask = torch.zeros(B, max_A, dtype=torch.bool) # True = real for i, (p6, fr, fv) in enumerate(batch): A_i = p6.shape[0] past_pad[i, :A_i] = p6 fut_pad[i, :A_i] = fr vel_pad[i, :A_i] = fv agent_mask[i, :A_i] = True # Block-diagonal traj_mask zeroed on padding N = B * max_A traj_mask = torch.zeros(N, N) flat_mask = agent_mask.view(-1) # [B*max_A] for i in range(B): s, e = i * max_A, (i + 1) * max_A block = agent_mask[i].unsqueeze(0) & agent_mask[i].unsqueeze(1) # [max_A, max_A] traj_mask[s:e, s:e] = block.float() return { 'batch_size': torch.tensor(B), 'past_traj': past_pad, # [B, max_A, T_past, 6] 'fut_traj': fut_pad, # [B, max_A, T_fut, 2] 'past_traj_original_scale': past_pad.clone(), # same as past_traj (normalize afterwards if needed) 'fut_traj_original_scale': fut_pad.clone(), 'traj_mask': traj_mask, # [B*max_A, B*max_A] 'fut_traj_vel': vel_pad, # [B, max_A, T_fut, 2] 'agent_mask': agent_mask, # [B, max_A] True=real } class SDDDatasetMinMax(Dataset): def __init__( self, obs_len=8, pred_len=12, training=True, cfg=None, data_dir='', split='train', data_norm='min_max', ): super().__init__() self.obs_len = obs_len self.pred_len = pred_len self.cfg = cfg self.data_norm = data_norm self.training = training path = os.path.join(data_dir, f'sdd_{split}_v2.pkl') with open(path, 'rb') as f: self.raw_scenes = pickle.load(f) # Pre-compute A per scene for the BatchSampler self.a_counts = [int(s.shape[0]) for s in self.raw_scenes] print(f'[SDDDatasetMinMax] {split}: {len(self.raw_scenes)} scenes, ' f'A min/mean/max = {min(self.a_counts)}/{np.mean(self.a_counts):.1f}/{max(self.a_counts)}') # Compute global min/max for min_max normalization (on training set) if training and data_norm == 'min_max': all_past, all_fut = [], [] for s in self.raw_scenes: past_6, fut_r, _ = _sdd_to_scene(s, obs_len, pred_len) all_past.append(past_6); all_fut.append(fut_r) past_cat = torch.cat(all_past, dim=0) fut_cat = torch.cat(all_fut, dim=0) cfg.past_traj_max = past_cat.max() cfg.past_traj_min = past_cat.min() cfg.fut_traj_max = fut_cat.max() cfg.fut_traj_min = fut_cat.min() def __len__(self): return len(self.raw_scenes) def __getitem__(self, idx): past_6, fut_r, fut_vel = _sdd_to_scene(self.raw_scenes[idx], self.obs_len, self.pred_len) if self.data_norm == 'min_max': past_6 = normalize_min_max(past_6, self.cfg.past_traj_min, self.cfg.past_traj_max, -1, 1).contiguous() fut_r = normalize_min_max(fut_r, self.cfg.fut_traj_min, self.cfg.fut_traj_max, -1, 1).contiguous() return past_6, fut_r, fut_vel # Alias kept for backwards-compat with import lines in fm_sdd.py / fm_sdd_graph.py seq_collate_sdd = collate_padded