""" Multi-dataset data loader for multi-body dynamics. Loads sequences from multiple dataset directories, with train/val/test membership defined by a split file. """ import os import json import time import concurrent.futures as futures import h5py import numpy as np import torch from torch.utils.data import Dataset, DataLoader, Sampler, RandomSampler from torch.utils.data.distributed import DistributedSampler from typing import Dict, List, Tuple, Optional import glob from pathlib import Path import random from tqdm import tqdm import sys from functools import lru_cache import logging sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from utils.noise_utils import apply_random_walk_noise_to_batch, apply_independent_noise from data.object_scale import ObjectScale, check_object_counts from data.partial_view import PartialView, quat_to_rotmat from data.scene_frames import frame_dirs as scene_frame_dirs, is_packed, load_frame import zlib logger = logging.getLogger(__name__) # Padding value for positions: zero contributes nothing to # pooling/aggregation operations. PADDING_VALUE = 0.0 # Global LRU cache for frequently loaded numpy arrays # (~50000 entries at ~5KB each, roughly 250MB). @lru_cache(maxsize=50000) def _cached_numpy_load(filepath: str) -> np.ndarray: """Cached numpy array loading. Uses filepath as cache key.""" return np.load(filepath) def cached_load_npy(path: Path) -> np.ndarray: """Load numpy array with LRU caching.""" return _cached_numpy_load(str(path)) def check_scene_set(config: dict) -> str: """data.scene_set must name the scene set the split file declares ('benchmark' when either is absent). 'complete' splits (data_generation/build_complete_splits.py) leave out the MOVi scenes whose metadata omits simulated bodies; 'benchmark' is the published scene list. Returns the scene set. """ want = (config.get('data', {}) or {}).get('scene_set', 'benchmark') if want not in ('benchmark', 'complete'): raise ValueError(f"data.scene_set must be 'benchmark' or 'complete'; got {want!r}") with open(config['data']['split_file']) as f: have = json.load(f).get('scene_set', 'benchmark') if have != want: raise ValueError(f"data.scene_set is {want!r} but {config['data']['split_file']} declares scene_set {have!r}") return want class MultiObjectDynamics(Dataset): """MOVi multi-object dataset with per-object vertices and per-frame padding+mask. Expects dataset directories like movi_* with frames frame_00000/... and files: - meshes/obj_i_faces.npy (per-object faces, static per dataset) - frame_xxxxx/obj_i_pos.npy (per-object vertices in world space) """ def __init__( self, data_path: str, sequence_length: int = 3, stride: int = 1, split: str = 'train', train_ratio: float = 0.9, val_ratio: float = 0.1, split_file: Optional[str] = None, fast_index: bool = True, skip_first_frames: int = 0, skip_step_size: int = 1, max_vertices_per_object: Optional[int] = None, max_objects_per_frame: Optional[int] = None, anchor_k: Optional[int] = None, if_phys_param: bool = False, **kwargs ): self.data_path = Path(data_path) self.sequence_length = sequence_length self.stride = stride self.split = split self.skip_first_frames = skip_first_frames self.skip_step_size = skip_step_size self.max_vertices_per_object = max_vertices_per_object self.max_objects_per_frame = max_objects_per_frame self.if_phys_param = if_phys_param # Strict mode for physics parameters # Point Cloud Mode configuration # When enabled, load point cloud files (obj_i_pc{N}_pos.npy) instead of mesh vertices if_point_cloud_config = kwargs.get('if_point_cloud', {}) self.point_cloud_enable = if_point_cloud_config.get('enable', False) if self.point_cloud_enable: # the point-cloud path validates neither per-frame point counts nor per-resolution padding (a frame with a # missing point was accepted; resolutions above the first were truncated to it), and train.py does not # forward the setting: refused until implemented and validated raise NotImplementedError('point-cloud inputs (if_point_cloud.enable) are not supported by this loader') self.point_cloud_train_resolution = if_point_cloud_config.get('train_resolution', [512]) self.point_cloud_eval_resolution = if_point_cloud_config.get('eval_resolution', [512]) # Ensure resolutions are lists if isinstance(self.point_cloud_train_resolution, int): self.point_cloud_train_resolution = [self.point_cloud_train_resolution] if isinstance(self.point_cloud_eval_resolution, int): self.point_cloud_eval_resolution = [self.point_cloud_eval_resolution] # Select resolution based on split if self.point_cloud_enable: if split == 'train': self.point_cloud_resolutions = self.point_cloud_train_resolution else: self.point_cloud_resolutions = self.point_cloud_eval_resolution print(f"[{split}] Point Cloud Mode ENABLED:") print(f" - Resolutions: {self.point_cloud_resolutions}") print(f" - Data will be loaded from: frame_XXXXX/obj_i_pc{{N}}_pos.npy") else: self.point_cloud_resolutions = None # Anchor config (Nq): prefer explicit arg, else env ANCHOR_K, else default 32 if anchor_k is None: env_anchor = os.environ.get('ANCHOR_K', '').strip() if env_anchor: try: anchor_k = int(env_anchor) except Exception: raise ValueError(f"Invalid ANCHOR_K environment value: '{env_anchor}' (expected int)") else: anchor_k = 32 if anchor_k <= 0: raise ValueError(f"anchor_k must be positive; got {anchor_k}") self.anchor_k = int(anchor_k) # Object Permutation Augmentation (train only, at least 2 objects) augmentation_config = kwargs.get('augmentation', {}) # The random-walk / independent noise blocks were never applied on this path (the dataset never stored them and # the helpers expect body1_/body2_ arrays, not MOVi's obj_positions): enabling them is an error, not a no-op. for _blk in ('random_walk', 'independent'): for _cfg in (augmentation_config, kwargs.get('noise_config') or {}): if ((_cfg or {}).get(_blk) or {}).get('enable', False): raise NotImplementedError(f"data.augmentation.{_blk}.enable: true is not implemented for the MOVi loader " f"(it would be silently ignored); set it to false") # Anchor selection rule. 'index0' starts FPS at vertex 0 (order dependent); # 'centroid_farthest' starts at the vertex farthest from the centroid with lexicographic # tie-breaks, so the anchors depend on the vertex set only. self.anchor_fps_start = str(kwargs.get('anchor_fps_start', 'index0')) # Order-independent FPS ordering of every object's vertices (cached as obj_i_fpsorder_cf.npy); # the encoder's hierarchical levels take the first 50/25/12.5/... percent of this ordering. self.fps_order = bool(kwargs.get('fps_order', False)) # data.object_scale_file: apply each object's metadata size to the stored vertices (src/data/object_scale.py) _scale_file = kwargs.get('object_scale_file') self._object_scale = ObjectScale(_scale_file) if _scale_file else None if self._object_scale is not None: # the folder name is compared without regard to case _ds_real = os.path.basename(os.path.realpath(str(data_path))) if (self._object_scale.converted or '').lower() != _ds_real.lower(): raise ValueError(f"data.object_scale_file {_scale_file} was built for '{self._object_scale.converted}', " f"but the dataset path resolves to '{_ds_real}'") if self._object_scale is not None: print(f"[{split}] Object scale ACTIVE: vertices scaled about the object origin by the metadata size ({_scale_file})") # data.partial_view: partial point clouds of a simulated depth camera instead of the mesh vertices # (src/data/partial_view.py) _pv = kwargs.get('partial_view') or {} self._partial = None if _pv.get('enable', False): if self._object_scale is None: raise ValueError("data.partial_view needs data.object_scale_file for the objects' metadata sizes") self._partial = PartialView(_pv, Path(__file__).resolve().parents[2]) if self._partial.min_points < self.anchor_k: raise ValueError(f'data.partial_view.min_points ({self._partial.min_points}) must be at least ' f'anchor_k ({self.anchor_k})') self._partial_scenes = {} print(f"[{split}] {self._partial.describe(train=split == 'train')}") self._fps_order_by_dataset = {} if self.anchor_fps_start not in ('index0', 'centroid_farthest'): raise ValueError(f"data.anchor_fps_start must be 'index0' or 'centroid_farthest'; got '{self.anchor_fps_start}'") if self.anchor_fps_start != 'index0': print(f"[{split}] Anchor FPS start rule: {self.anchor_fps_start} (anchor cache suffix '_cf')") # Anchor resampling augmentation (train split only): with probability `prob` the FPS anchors # of every object are recomputed from a random start vertex. resample_config = augmentation_config.get('anchor_resampling', {}) self.anchor_resample_prob = float(resample_config.get('prob', 0.5)) if resample_config.get('enable', False) else 0.0 if self.anchor_resample_prob > 0 and split == 'train': print(f"[{split}] Anchor resampling augmentation enabled: prob={self.anchor_resample_prob:.2f}") permute_config = augmentation_config.get('object_permutation', {}) self.permute_enabled = permute_config.get('enable', False) and split == 'train' self.permute_prob = permute_config.get('prob_apply', 0.5) self.permute_debug = permute_config.get('if_debug', False) if self.permute_enabled: print(f"[{split}] Object permutation augmentation enabled: prob={self.permute_prob:.2f}, debug={self.permute_debug}") # Load the dataset split file (required; splits are never auto-generated). if not split_file: raise ValueError( "split_file is required. Set 'data.split_file' in the config to a split JSON, " "e.g. one of the dataset_split_*.json files in the repository root." ) if not os.path.exists(split_file): raise FileNotFoundError( f"Split file not found: {split_file}\n" "Check the 'data.split_file' path in the config. Split files " "(dataset_split_*.json) are provided in the repository root; new splits " "can be written for data produced by data_generation/movi_to_sim_data.py." ) with open(split_file, 'r') as f: s = json.load(f) names = s['train_datasets'] if split == 'train' else s['val_datasets'] if split == 'val' else s['test_datasets'] self.dataset_dirs = [self.data_path / n for n in names] if self._object_scale is not None: check_object_counts(self._object_scale, self.data_path, names) # Build sequences self.all_sequences = [] # each: { 'dataset': name, 'frames': [paths], 'frame_ids': [int] } # DDP-aware: only show progress bar on rank 0 import torch.distributed as dist show_progress = True if dist.is_available() and dist.is_initialized(): show_progress = (dist.get_rank() == 0) for ds in tqdm(self.dataset_dirs, desc=f"Indexing {split} datasets", unit="ds", disable=not show_progress): frames = scene_frame_dirs(ds) if self.skip_first_frames > 0: frames = frames[self.skip_first_frames:] if self.skip_step_size > 1: frames = frames[::self.skip_step_size] for start in range(0, len(frames) - sequence_length + 1, self.stride): seq = frames[start:start+sequence_length] if len(seq) == sequence_length: self.all_sequences.append({ 'dataset': ds.name, 'frames': [p.as_posix() for p in seq], 'frame_ids': [int(Path(p).name.split('_')[-1]) for p in seq] }) # Precompute/load per-object anchor indices and faces for each dataset from # the reference frame_00000; cached in memory. self._anchors_by_dataset = {} self._faces_by_dataset = {} # For point cloud mode, use the first resolution for __init__ validation init_pc_resolution = self.point_cloud_resolutions[0] if self.point_cloud_enable else None for ds in self.dataset_dirs: ds_name = ds.name frame0 = ds / 'frame_00000' if not (frame0.exists() or is_packed(ds)): raise FileNotFoundError( f"Missing reference frame for dataset '{ds_name}': {frame0}. " f"Expected a 'frame_00000' directory (or frames.npy) under {ds}.") # For point cloud mode, load point clouds for validation # For mesh mode, load mesh vertices ref_objs = self._load_frame_objects(frame0, pc_resolution=init_pc_resolution) if len(ref_objs) == 0: raise RuntimeError(f"Empty reference frame in {frame0}: no per-object vertex files found.") anchors_dir = ds / 'anchors' anchors_dir.mkdir(parents=True, exist_ok=True) ds_anchor_list = [] ds_faces_list = [] ds_order_list = [] for i, V in enumerate(ref_objs): nv = V.shape[0] if self.max_vertices_per_object is not None and nv > self.max_vertices_per_object: raise ValueError( f"Dataset '{ds_name}' reference object {i} has {nv} vertices, " f"exceeding max_vertices_per_object={self.max_vertices_per_object}.\n" f"Either reduce vertex count in dataset or increase data.max_vertices_per_object in config." ) # Point cloud mode: anchors are unused, so use sequential placeholder # indices rather than mesh anchors (whose indices may be out of range). if self.point_cloud_enable: idx = np.arange(min(self.anchor_k, nv), dtype=np.int32) if len(idx) < self.anchor_k: idx = np.pad(idx, (0, self.anchor_k - len(idx)), mode='constant') ds_anchor_list.append(idx) # Create empty faces (not used in point cloud mode) ds_faces_list.append(np.zeros((0, 3), dtype=np.int64)) continue # Mesh vertex mode: standard anchor and faces loading. if nv < self.anchor_k: raise ValueError( f"Dataset '{ds_name}' object {i} has only {nv} vertices (< anchor_k={self.anchor_k}). " f"Reduce ANCHOR_K or increase vertex cap (max_vertices_per_object).") _suffix = '' if self.anchor_fps_start == 'index0' else '_cf' anchor_path = anchors_dir / f"obj_{i}_anchors_fps{self.anchor_k}{_suffix}.npy" if anchor_path.exists(): idx = np.load(anchor_path) if idx.ndim != 1 or idx.shape[0] != self.anchor_k: raise ValueError( f"Invalid anchor file shape: {anchor_path} (expected ({self.anchor_k},))") if idx.dtype.kind not in ('i', 'u'): raise ValueError(f"Invalid anchor file dtype: {anchor_path} (expected integer)") if int(idx.max()) >= nv or int(idx.min()) < 0: raise ValueError( f"Anchor indices in {anchor_path} out of range [0,{nv-1}] for object {i}") if np.unique(idx).shape[0] != idx.shape[0]: raise ValueError(f"Anchor cache {anchor_path} repeats a vertex: {idx.tolist()}") ds_anchor_list.append(idx.astype(np.int32, copy=False)) else: # Compute deterministic FPS anchors on the STORED (unscaled) frame-0 vertices, as # pregen_cf_anchors.py does: FPS is invariant to the size scaling only in exact arithmetic (float32 # rounding of the scaled coordinates breaks the exact ties of symmetric meshes differently) V_ref = load_frame(ds, 0)[i].astype(np.float32, copy=False) idx = self._farthest_point_sample_np(V_ref, self.anchor_k, rule=self.anchor_fps_start) # Atomic write: concurrent DDP ranks must never read a partially written cache. tmp_path = anchor_path.with_name(f"{anchor_path.stem}.{os.getpid()}.tmp.npy") np.save(tmp_path, idx.astype(np.int32, copy=False)) os.replace(tmp_path, anchor_path) ds_anchor_list.append(idx.astype(np.int32, copy=False)) if self.fps_order: order_path = anchors_dir / f"obj_{i}_fpsorder_cf.npy" if order_path.exists(): order = np.load(order_path) if order.ndim != 1 or order.shape[0] != nv or order.dtype.kind not in ('i', 'u'): raise ValueError(f"Invalid FPS-order cache {order_path}: shape {order.shape}, expected ({nv},)") if not np.array_equal(np.sort(order), np.arange(nv)): raise ValueError(f"FPS-order cache {order_path} is not a permutation of the {nv} vertices") else: # stored (unscaled) frame-0 vertices, as pregen_fps_order.py does (see the anchors above) order = self._farthest_point_sample_np(load_frame(ds, 0)[i].astype(np.float32, copy=False), nv, rule='centroid_farthest') tmp_path = order_path.with_name(f"{order_path.stem}.{os.getpid()}.tmp.npy") np.save(tmp_path, order.astype(np.int32, copy=False)) os.replace(tmp_path, order_path) ds_order_list.append(order.astype(np.int32, copy=False)) # Load faces for this object (static per dataset), and filter to nv cap faces_path = ds / 'meshes' / f'obj_{i}_faces.npy' if not faces_path.exists(): raise FileNotFoundError( f"Missing faces file for dataset '{ds_name}' object {i}: {faces_path}") faces = np.load(faces_path) if faces.ndim != 2 or faces.shape[1] != 3: raise ValueError(f"Invalid faces array at {faces_path}: shape={faces.shape}, expected (F,3)") # every face must index this object's vertices (the vertex count is never capped: an object above # max_vertices_per_object is an error), so an out-of-range face is corrupt data, not a face to drop if faces.dtype.kind not in ('i', 'u'): raise ValueError(f"Faces of {faces_path} have dtype {faces.dtype}, expected integers") faces = faces.astype(np.int64, copy=False) if faces.size and (int(faces.min()) < 0 or int(faces.max()) >= nv): raise ValueError(f"Faces of {faces_path} index vertices outside [0, {nv - 1}]") ds_faces_list.append(faces) self._anchors_by_dataset[ds_name] = ds_anchor_list self._faces_by_dataset[ds_name] = ds_faces_list if self.fps_order: self._fps_order_by_dataset[ds_name] = ds_order_list # Cache reference frames (frame_00000) per dataset to avoid repeated loading. self._reference_frames_cache = {} # Per-dataset caches: frame lists (avoids glob in __getitem__), frame IDs # (avoids string parsing per sample), and per-object/max vertex counts. self._all_frames_by_dataset = {} self._frame_id_cache = {} self._max_nv_by_dataset = {} self._mesh_nv_by_dataset = {} for ds in self.dataset_dirs: ds_name = ds.name all_frames = scene_frame_dirs(ds) self._all_frames_by_dataset[ds_name] = all_frames for fp in all_frames: frame_id = int(fp.name.split('_')[-1]) # "frame_00123" -> 123 self._frame_id_cache[(ds_name, fp.name)] = frame_id # Frames are indexed by position: a missing or extra frame directory would silently change the time step # and misalign the ground truth (vis_traj.npy, which the evaluator indexes by frame number). _ids = [self._frame_id_cache[(ds_name, fp.name)] for fp in all_frames] _T_traj = np.load(ds / 'vis_traj.npy', mmap_mode='r').shape[1] if _ids != list(range(len(_ids))) or len(_ids) != _T_traj: raise ValueError(f"{ds}: frame directories must be frame_00000 .. frame_{_T_traj - 1:05d} without gaps " f"(vis_traj.npy has {_T_traj} frames); found {len(_ids)} directories, " f"first gap at {next((i for i, f in enumerate(_ids) if f != i), len(_ids))}") # Load frame_00000 once to record per-object vertex counts # (at point cloud resolution when enabled). frame0 = ds / 'frame_00000' ref_objs_data = self._load_frame_objects(frame0, pc_resolution=init_pc_resolution) nv_list = [obj.shape[0] for obj in ref_objs_data] self._mesh_nv_by_dataset[ds_name] = nv_list self._max_nv_by_dataset[ds_name] = max(nv_list) if nv_list else 0 # Load physics parameters (mass, friction, restitution) per dataset. # With if_phys_param=True the file and every object entry must exist. self._physics_by_dataset = {} for ds in self.dataset_dirs: ds_name = ds.name physics_path = ds / 'physics_params.json' num_objs = len(self._anchors_by_dataset[ds_name]) if self.if_phys_param: if not physics_path.exists(): raise FileNotFoundError( f"Physics parameters required but file missing: {physics_path}\n" f"Dataset: {ds_name} | if_phys_param=True requires physics_params.json\n" f"Please generate physics parameters or set if_phys_param=False" ) # Parse errors propagate deliberately. with open(physics_path, 'r') as f: physics_params = json.load(f) phys_list = [] for i in range(num_objs): obj_key = f"obj_{i}" if obj_key not in physics_params: raise KeyError( f"Missing physics parameters for {obj_key} in {physics_path}\n" f"Dataset: {ds_name} | Found objects: {list(physics_params.keys())}\n" f"Required: obj_0 through obj_{num_objs-1}" ) p = physics_params[obj_key] if 'mass' not in p: raise KeyError(f"Missing 'mass' field for {obj_key} in {physics_path}") if 'friction' not in p: raise KeyError(f"Missing 'friction' field for {obj_key} in {physics_path}") if 'restitution' not in p: raise KeyError(f"Missing 'restitution' field for {obj_key} in {physics_path}") phys_list.append([ float(p['mass']), float(p['friction']), float(p['restitution']) ]) self._physics_by_dataset[ds_name] = np.array(phys_list, dtype=np.float32) # (Nobj, 3) else: # Permissive mode: substitute defaults for missing entries. if physics_path.exists(): with open(physics_path, 'r') as f: physics_params = json.load(f) phys_list = [] for i in range(num_objs): obj_key = f"obj_{i}" if obj_key in physics_params: p = physics_params[obj_key] phys_list.append([ float(p.get('mass', 1.0)), float(p.get('friction', 0.5)), float(p.get('restitution', 0.5)) ]) else: phys_list.append([1.0, 0.5, 0.5]) self._physics_by_dataset[ds_name] = np.array(phys_list, dtype=np.float32) else: # File missing entirely: default values for all objects. self._physics_by_dataset[ds_name] = np.array([[1.0, 0.5, 0.5]] * num_objs, dtype=np.float32) # Multi-timescale training configuration; step sizes are expressed in # normalized time units. random_step_size_config = kwargs.get('random_step_size', {}) self.random_step_size_enable = random_step_size_config.get('enable', False) self.random_step_sizes = random_step_size_config.get('step_sizes', [1, 2, 5, 10]) self.random_start_frame = random_step_size_config.get('random_start_frame', True) self.sampling_weights = random_step_size_config.get('sampling_weights', None) if self.random_step_size_enable: if not isinstance(self.random_step_sizes, list) or len(self.random_step_sizes) == 0: raise ValueError( f"random_step_size.step_sizes must be a non-empty list, got: {self.random_step_sizes}" ) for ss in self.random_step_sizes: if not isinstance(ss, int) or ss <= 0: raise ValueError( f"All step_sizes must be positive integers, got: {self.random_step_sizes}" ) if self.sampling_weights is None: # Default to uniform sampling. self.sampling_weights = [1.0 / len(self.random_step_sizes)] * len(self.random_step_sizes) else: if len(self.sampling_weights) != len(self.random_step_sizes): raise ValueError( f"sampling_weights length ({len(self.sampling_weights)}) must match " f"step_sizes length ({len(self.random_step_sizes)})" ) weights_sum = sum(self.sampling_weights) if abs(weights_sum - 1.0) > 1e-6: raise ValueError( f"sampling_weights must sum to 1.0, got sum={weights_sum:.6f}\n" f"Weights: {self.sampling_weights}" ) print(f"\n{'='*60}") print(f"[{split}] Multi-timescale training ENABLED:") print(f" - step_sizes: {self.random_step_sizes}") print(f" - sampling_weights: {self.sampling_weights}") print(f" - random_start_frame: {self.random_start_frame}") print(f" - if_debug: {random_step_size_config.get('if_debug', False)}") print(f"{'='*60}\n") # Every dataset must have enough frames for every configured step size. print(f"[{split}] Validating dataset frame counts for multi-timescale training...") insufficient_datasets = [] for ds in self.dataset_dirs: frames = scene_frame_dirs(ds) if self.skip_first_frames > 0: frames = frames[self.skip_first_frames:] available_frames = len(frames) for step_size in self.random_step_sizes: frames_needed = (self.sequence_length - 1) * step_size + 1 if available_frames < frames_needed: insufficient_datasets.append({ 'dataset': ds.name, 'available': available_frames, 'needed': frames_needed, 'step_size': step_size, 'sequence_length': self.sequence_length }) if insufficient_datasets: error_msg = f"[{split}] Frame validation failed: {len(insufficient_datasets)} datasets have insufficient frames:\n" for info in insufficient_datasets[:5]: error_msg += ( f" - {info['dataset']}: available={info['available']}, " f"needed={info['needed']} (step_size={info['step_size']}, " f"sequence_length={info['sequence_length']})\n" ) if len(insufficient_datasets) > 5: error_msg += f" ... and {len(insufficient_datasets) - 5} more datasets\n" error_msg += ( f"\nSolutions:\n" f" 1. Reduce max step_size in config: data.random_step_size.step_sizes\n" f" 2. Reduce sequence_length in config: data.sequence_length\n" f" 3. Filter out datasets with insufficient frames\n" f" 4. Reduce skip_first_frames in config: data.skip_first_frames (currently {self.skip_first_frames})" ) raise ValueError(error_msg) print(f"[{split}] All {len(self.dataset_dirs)} datasets have sufficient frames for step_sizes {self.random_step_sizes}") else: # Disabled: use the fixed skip_step_size print(f"[{split}] Multi-timescale training disabled, using skip_step_size={self.skip_step_size}") def __len__(self) -> int: return len(self.all_sequences) def get_cache_stats(self) -> dict: """Return cache statistics for monitoring.""" return { 'numpy_cache_info': _cached_numpy_load.cache_info()._asdict(), 'reference_frames_cached': len(self._reference_frames_cache), 'reference_frames_datasets': list(self._reference_frames_cache.keys()) } @staticmethod def clear_numpy_cache(): """Clear the global numpy LRU cache to free memory.""" _cached_numpy_load.cache_clear() logger.info("Cleared numpy LRU cache") def _load_frame_objects(self, frame_dir: Path, pc_resolution: Optional[int] = None) -> List['np.ndarray']: """ Load per-object vertex/point data for a single frame. Args: frame_dir: Path to frame directory pc_resolution: If provided, load point cloud at this resolution instead of mesh vertices Returns: List of (Ni, 3) arrays, one per object """ # Data exceeding the configured limits is rejected rather than truncated. objs = [] idx = 0 if pc_resolution is not None: if self._object_scale is not None: raise NotImplementedError('data.object_scale_file is not supported with point-cloud inputs') # Point cloud mode: load from obj_i_pc{N}_pos.npy while (frame_dir / f'obj_{idx}_pc{pc_resolution}_pos.npy').exists(): pc_file = frame_dir / f'obj_{idx}_pc{pc_resolution}_pos.npy' V = cached_load_npy(pc_file) # (N_points, 3) # Point clouds have a fixed size; no vertex limit check needed. objs.append(V.astype('float32', copy=False)) idx += 1 # Mesh vertices exist but the requested point cloud resolution does not. if idx == 0 and (frame_dir / 'obj_0_pos.npy').exists(): raise FileNotFoundError( f"Point cloud files not found at resolution {pc_resolution} in {frame_dir}.\n" f"Expected: obj_i_pc{pc_resolution}_pos.npy\n" f"Mesh vertices exist but point cloud mode is enabled.\n" f"Please use a dataset with pre-computed point clouds." ) else: # Mesh vertex mode: load from obj_i_pos.npy (packed layout: the same arrays from frames.npy) packed = (load_frame(frame_dir.parent, int(frame_dir.name.split('_')[-1])) if is_packed(frame_dir.parent) else None) while (idx < len(packed)) if packed is not None else (frame_dir / f'obj_{idx}_pos.npy').exists(): V = packed[idx] if packed is not None else cached_load_npy(frame_dir / f'obj_{idx}_pos.npy') # (Ni, 3) if self._object_scale is not None: V = self._object_scale.apply(frame_dir / f'obj_{idx}_pos.npy', V) if self.max_vertices_per_object is not None and V.shape[0] > self.max_vertices_per_object: raise ValueError( f"Object {idx} in {frame_dir} has {V.shape[0]} vertices, " f"exceeding max_vertices_per_object={self.max_vertices_per_object}.\n" f"Either reduce vertex count in dataset or increase data.max_vertices_per_object in config." ) objs.append(V.astype('float32', copy=False)) idx += 1 if self.max_objects_per_frame is not None and len(objs) > self.max_objects_per_frame: raise ValueError( f"Frame {frame_dir} has {len(objs)} objects, " f"exceeding max_objects_per_frame={self.max_objects_per_frame}.\n" f"Either reduce object count in dataset or increase data.max_objects_per_frame in config." ) return objs def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: info = self.all_sequences[idx] # Point Cloud Mode: Sample resolution for this sample pc_resolution = None if self.point_cloud_enable: pc_resolution = np.random.choice(self.point_cloud_resolutions) # Multi-timescale training: sample a step size for this sequence. if self.random_step_size_enable: step_size = np.random.choice(self.random_step_sizes, p=self.sampling_weights) ds_name = info['dataset'] all_frames = self._all_frames_by_dataset[ds_name] # Apply skip_first_frames if self.skip_first_frames > 0: all_frames = all_frames[self.skip_first_frames:] # Total frames spanned by the sequence, e.g. sequence_length=4 with # step_size=10 uses frames [0, 10, 20, 30] and spans 31 frames. student_reach = (self.sequence_length - 1) * step_size total_frames_needed = student_reach + 1 # Frame counts were validated in __init__; assert as a sanity check. assert len(all_frames) >= total_frames_needed, \ f"Unexpected: insufficient frames in {ds_name} (should have been caught in __init__)" # Select a start frame within the valid range. if self.random_start_frame: max_start_frame = len(all_frames) - total_frames_needed start_frame = np.random.randint(0, max_start_frame + 1) else: # Deterministic start frame derived from idx (for reproducibility). max_start_frame = len(all_frames) - total_frames_needed start_frame = idx % (max_start_frame + 1) # Frame sequence [start_frame, start_frame + step_size, ...]. frame_indices = [start_frame + i * step_size for i in range(self.sequence_length)] frame_paths = [all_frames[i] for i in frame_indices] frame_ids = np.array([self._frame_id_cache[(ds_name, fp.name)] for fp in frame_paths], dtype=np.int32) selected_step_size = step_size else: # Fixed step size: use pre-built sequences with skip_step_size. ds_name = info['dataset'] frame_paths = [Path(p) for p in info['frames']] frame_ids = np.asarray(info.get('frame_ids', []), dtype=np.int32) selected_step_size = self.skip_step_size if self._partial is not None: return self._partial_item(ds_name, frame_ids, selected_step_size) # Load per-frame objects (ragged); point clouds are loaded instead of mesh # vertices when pc_resolution is set. frames_objs: List[List[np.ndarray]] = [self._load_frame_objects(fp, pc_resolution=pc_resolution) for fp in frame_paths] if pc_resolution is None: # every frame holds the reference frame's objects with the same vertex counts (a missing or truncated file # would otherwise become padding or a shorter object without notice) _ref_nv = self._mesh_nv_by_dataset[ds_name] for fp, fo in zip(frame_paths, frames_objs): if [V.shape[0] for V in fo] != _ref_nv: raise ValueError(f"{fp}: per-object vertex counts {[V.shape[0] for V in fo]} differ from frame_00000 " f"{_ref_nv}") num_objs = max(len(f) for f in frames_objs) if self.max_vertices_per_object is not None: max_nv = int(self.max_vertices_per_object) else: max_nv = self._max_nv_by_dataset[ds_name] T = len(frames_objs) # Allocate padded arrays. obj_positions = np.full((T, num_objs, max_nv, 3), PADDING_VALUE, dtype='float32') valid_obj_mask = np.zeros((num_objs,), dtype='bool') valid_vtx_mask = np.zeros((T, num_objs, max_nv), dtype='bool') for t, fo in enumerate(frames_objs): for i, V in enumerate(fo): nv = min(V.shape[0], max_nv) obj_positions[t, i, :nv, :] = V[:nv] valid_obj_mask[i] = True valid_vtx_mask[t, i, :nv] = True # Reference pose is always the first frame of the dataset (global frame_00000), # cached per dataset. Point cloud mode includes the resolution in the cache key. cache_key = f"{ds_name}_pc{pc_resolution}" if pc_resolution else ds_name if cache_key not in self._reference_frames_cache: ds_root = self.data_path / ds_name frame0 = ds_root / 'frame_00000' if not (frame0.exists() or is_packed(ds_root)): raise FileNotFoundError( f"Missing reference frame for dataset '{ds_name}': {frame0}. " f"Expected a 'frame_00000' directory (or frames.npy) under {ds_root}." ) ref_objs = self._load_frame_objects(frame0, pc_resolution=pc_resolution) if len(ref_objs) == 0: raise RuntimeError( f"Empty reference frame in {frame0}: no per-object vertex files found." ) self._reference_frames_cache[cache_key] = ref_objs else: ref_objs = self._reference_frames_cache[cache_key] # Every frame in the sequence must contain all reference objects, each stored # as a non-empty (N, 2|3) array. Point cloud mode is already validated inside # _load_frame_objects, so it is skipped here, and so is the packed layout (frames.npy holds every object of # every frame by construction; scene_frames checks its shape against frames_counts.npy). expected_num_objs = len(ref_objs) if pc_resolution is None and not is_packed(self.data_path / ds_name): for t, fdir in enumerate(frame_paths): for i_ref in range(expected_num_objs): npy_path = fdir / f"obj_{i_ref}_pos.npy" if not npy_path.exists(): raise RuntimeError( f"Missing expected file for dataset '{ds_name}': {npy_path} " f"(t={t}, obj={i_ref}). Ensure all frames contain the same set of objects as frame_00000.") try: arr_chk = np.load(npy_path) except Exception as e: raise RuntimeError( f"Failed to load file: {npy_path} (t={t}, obj={i_ref}). Error: {e}") if not (arr_chk.ndim == 2 and arr_chk.shape[1] in (2, 3) and arr_chk.shape[0] > 0): raise RuntimeError( f"Invalid/empty vertices at {npy_path}: shape={arr_chk.shape} " f"(t={t}, obj={i_ref}).") # Build the reference array with fixed Nv (pad/trim to max_nv). ref_obj_positions = np.full((num_objs, max_nv, 3), PADDING_VALUE, dtype='float32') for i, V in enumerate(ref_objs[:num_objs]): nv = min(V.shape[0], max_nv) ref_obj_positions[i, :nv, :] = V[:nv] # Load anchors (per object) for this dataset (computed in __init__) if ds_name not in self._anchors_by_dataset: raise KeyError(f"Missing anchors for dataset '{ds_name}'. Initialization may have failed.") ds_anchor_list = self._anchors_by_dataset[ds_name] if len(ds_anchor_list) < num_objs: raise ValueError( f"Anchors count ({len(ds_anchor_list)}) < objects in sequence ({num_objs}) for dataset '{ds_name}'.") anchor_idx = np.zeros((num_objs, self.anchor_k), dtype=np.int32) for i in range(num_objs): idx_i = ds_anchor_list[i] if idx_i.shape[0] != self.anchor_k: raise ValueError( f"Anchor length mismatch for dataset '{ds_name}' obj {i}: got {idx_i.shape[0]}, expected {self.anchor_k}") anchor_idx[i] = idx_i vtx_rank = None if self.fps_order: # rank[i, v] = position of vertex v in object i's FPS ordering (padding: max_nv) vtx_rank = np.full((num_objs, max_nv), max_nv, dtype=np.int32) ds_order_list = self._fps_order_by_dataset[ds_name] for i in range(num_objs): order_i = ds_order_list[i] nv_i = min(order_i.shape[0], max_nv) vtx_rank[i, order_i[:nv_i]] = np.arange(nv_i, dtype=np.int32) if (self.split == 'train' and self.anchor_resample_prob > 0 and not self.point_cloud_enable and np.random.rand() < self.anchor_resample_prob): _before = anchor_idx.copy() for i in range(num_objs): nv_i = min(ref_objs[i].shape[0], max_nv) if nv_i >= self.anchor_k: anchor_idx[i] = self._farthest_point_sample_np( ref_objs[i][:nv_i].astype(np.float32, copy=False), self.anchor_k, rule=self.anchor_fps_start, start=int(np.random.randint(nv_i))) if not getattr(self, '_anchor_resample_logged', False): self._anchor_resample_logged = True print(f"[train] Anchor resampling APPLIED (pid {os.getpid()}): obj0 {_before[0].tolist()} -> {anchor_idx[0].tolist()} " f"(rule={self.anchor_fps_start})", flush=True) # Object permutation must be applied before faces are gathered so that # faces_list[i] matches the permuted obj_i. Only faces for the current # sample's objects are taken. faces_list = self._faces_by_dataset[ds_name][:num_objs] if self.permute_debug: print(f"[object permutation] dataset={ds_name}, num_objs={num_objs}, faces_list length={len(faces_list)}, total_faces={len(self._faces_by_dataset[ds_name])}") perm = None # Set only if permutation is applied if self.permute_enabled and np.random.random() < self.permute_prob: if num_objs >= 2: # Shuffle all objects; there is no ground mesh to exclude. perm = np.arange(num_objs) np.random.shuffle(perm) if self.permute_debug: print(f"[object permutation] dataset={ds_name}, idx={idx}, num_objs={num_objs}") print(f" Original order: {list(range(num_objs))}") print(f" Shuffled order: {perm.tolist()}") print(f" Mapping: {' '.join([f'{i}->{perm[i]}' for i in range(num_objs)])}") # Permute all object-indexed data. obj_positions = obj_positions[:, perm] # (T, Nobj, Nv, 3) ref_obj_positions = ref_obj_positions[perm] # (Nobj, Nv, 3) ref_objs = [ref_objs[i] for i in perm] # Permute ref_objs list # Permute frames_objs (list of lists) frames_objs = [[fo[i] for i in perm] for fo in frames_objs] valid_obj_mask = valid_obj_mask[perm] # (Nobj,) valid_vtx_mask = valid_vtx_mask[:, perm] # (T, Nobj, Nv) anchor_idx = anchor_idx[perm] # (Nobj, Nq) if vtx_rank is not None: vtx_rank = vtx_rank[perm] # (Nobj, Nv) # Permute faces_list faces_list = [faces_list[i] for i in perm] # time_step_id: first global frame id of this sequence (scalar) ts_id_val = int(frame_ids[0]) if frame_ids.size > 0 else -1 # Get physics parameters for this dataset (permuted if augmentation applied) phys_params = self._physics_by_dataset[ds_name] # (Nobj_full, 3) phys_params = phys_params[:num_objs] # Trim to current sample's object count if perm is not None: phys_params = phys_params[perm] # Apply permutation to match permuted objects sample = { 'obj_positions': torch.from_numpy(obj_positions), 'valid_obj_mask': torch.from_numpy(valid_obj_mask), 'valid_vtx_mask': torch.from_numpy(valid_vtx_mask), 'ref_obj_positions': torch.from_numpy(ref_obj_positions), 'dt': torch.tensor(1.0, dtype=torch.float32), 'frame_ids': torch.from_numpy(frame_ids) if frame_ids.size > 0 else torch.tensor([], dtype=torch.int32), 'time_step_id': torch.tensor(ts_id_val, dtype=torch.int64), 'anchor_idx': torch.from_numpy(anchor_idx), **({'vtx_rank': torch.from_numpy(vtx_rank)} if vtx_rank is not None else {}), 'dataset_name': ds_name, 'faces_list': faces_list, # Store permuted faces (list of numpy arrays) 'permutation': perm if perm is not None else np.arange(num_objs), # Store permutation order (identity if not permuted) 'phys_params': torch.from_numpy(phys_params), # (Nobj, 3): [mass, friction, restitution] 'skip_step_size': torch.tensor(selected_step_size, dtype=torch.float32), # Time step size for temporal encoding (randomly sampled when multi-timescale enabled) 'is_point_cloud': pc_resolution is not None, # Whether this sample uses point cloud data } if pc_resolution is not None: sample['point_cloud_resolution'] = pc_resolution # Apply noise augmentation (train split only) if getattr(self, 'split', 'train') == 'train' and getattr(self, 'noise_config', None): rw_cfg = self.noise_config.get('random_walk', {}) if rw_cfg.get('enable', False): start_epoch = rw_cfg.get('start_epoch', 0) cur_epoch = getattr(self, 'current_epoch', 0) if cur_epoch >= start_epoch: sample, noise_info = apply_random_walk_noise_to_batch( sample, self.noise_config, is_training=True ) sample['noise_info'] = noise_info elif self.noise_config.get('independent', {}).get('enable', False): sample = apply_independent_noise(sample, self.noise_config, is_training=True) return sample def _partial_scene(self, ds_name: str) -> dict: """Shapes, metadata sizes and pose trajectory (T, N, 7: position, quaternion w x y z) of a scene.""" info = self._partial_scenes.get(ds_name) if info is None: ds = self.data_path / ds_name shapes = [o['asset_id'] for o in json.loads((ds / 'conversion.json').read_text())['objects']] traj = np.load(ds / 'vis_traj.npy')[0] sizes = self._object_scale.scales[ds_name] if not (len(shapes) == len(sizes) == traj.shape[1]): raise ValueError(f'{ds}: {len(shapes)} shapes, {len(sizes)} sizes, {traj.shape[1]} trajectories') info = {'shapes': shapes, 'sizes': np.asarray(sizes, dtype=np.float64), 'traj': traj} self._partial_scenes[ds_name] = info return info def _partial_item(self, ds_name: str, frame_ids: np.ndarray, step_size: int) -> Dict[str, torch.Tensor]: """Sample whose objects are the partial point clouds a depth camera observes at sample frame partial_view.observe_index, moved with the ground-truth poses to every frame (src/data/partial_view.py). Objects with fewer than min_points observed points are left out; `partial_obj_index` maps the sample's objects to the scene's. Points are in farthest-point order, so the anchors are the first anchor_k points and the FPS rank of a point is its index. Evaluation samples also carry the index of the observation frame in the sample (`partial_obs_index`) and the full collision meshes of all objects at that frame (`partial_full_obs`) for scoring. """ info = self._partial_scene(ds_name) traj = info['traj'] fids = np.asarray(frame_ids, dtype=np.int64) T = len(fids) k = min(self._partial.observe_index, T - 1) t_obs = int(fids[k]) train = self.split == 'train' rng = np.random.default_rng(None if train else zlib.crc32(f'{ds_name}:{t_obs}'.encode())) pos_obs = traj[t_obs, :, :3].astype(np.float64) quat_obs = traj[t_obs, :, 3:7].astype(np.float64) # The model's contact search needs at least two objects. A training view that observes fewer is redrawn; after # 8 such views the evaluation camera, which observes nearly all objects, is used instead, or the whole objects # for a scene without a recorded camera (augmented training scenes). for attempt in range(9 if train else 1): use_train_camera = train and attempt < 8 if train and not use_train_camera and ds_name not in self._partial.cameras: local = self._partial.whole_objects(info['shapes']) else: local = self._partial.observe(ds_name, info['shapes'], info['sizes'], pos_obs, quat_obs, use_train_camera, rng) observed = np.array([i for i, u in enumerate(local) if u is not None], dtype=np.int64) if observed.size >= 2: break if observed.size < 2: raise RuntimeError(f'{ds_name}: {observed.size} object(s) observed at frame {t_obs}, at least 2 needed ' f'({self._partial.describe(train=False)})') n_obj = int(observed.size) nmax = max(local[i].shape[0] for i in observed) P = traj[fids][:, observed, :3].astype(np.float64) # (T, n_obj, 3) Rq = quat_to_rotmat(traj[fids][:, observed, 3:7]) # (T, n_obj, 3, 3) P0 = traj[0, observed, :3].astype(np.float64) # reference pose: global frame 0 R0 = quat_to_rotmat(traj[0, observed, 3:7]) obj_positions = np.full((T, n_obj, nmax, 3), PADDING_VALUE, dtype='float32') valid_vtx_mask = np.zeros((T, n_obj, nmax), dtype='bool') ref_obj_positions = np.full((n_obj, nmax, 3), PADDING_VALUE, dtype='float32') vtx_rank = np.full((n_obj, nmax), nmax, dtype=np.int32) for j, i in enumerate(observed): U = local[i] * info['sizes'][i] # object frame, metres n = U.shape[0] obj_positions[:, j, :n] = P[:, j, None, :] + np.einsum('tab,nb->tna', Rq[:, j], U) valid_vtx_mask[:, j, :n] = True ref_obj_positions[j, :n] = P0[j] + U @ R0[j].T vtx_rank[j, :n] = np.arange(n, dtype=np.int32) anchor_idx = np.tile(np.arange(self.anchor_k, dtype=np.int32), (n_obj, 1)) valid_obj_mask = np.ones((n_obj,), dtype='bool') phys_params = self._physics_by_dataset[ds_name][observed] faces_list = [np.zeros((0, 3), dtype=np.int64) for _ in range(n_obj)] perm = None if self.permute_enabled and np.random.random() < self.permute_prob and n_obj >= 2: perm = np.random.permutation(n_obj) obj_positions = obj_positions[:, perm] valid_vtx_mask = valid_vtx_mask[:, perm] ref_obj_positions = ref_obj_positions[perm] vtx_rank = vtx_rank[perm] anchor_idx = anchor_idx[perm] phys_params = phys_params[perm] observed = observed[perm] faces_list = [faces_list[p] for p in perm] sample = { 'obj_positions': torch.from_numpy(obj_positions), 'valid_obj_mask': torch.from_numpy(valid_obj_mask), 'valid_vtx_mask': torch.from_numpy(valid_vtx_mask), 'ref_obj_positions': torch.from_numpy(ref_obj_positions), 'dt': torch.tensor(1.0, dtype=torch.float32), 'frame_ids': torch.from_numpy(fids.astype(np.int32)), 'time_step_id': torch.tensor(int(fids[0]), dtype=torch.int64), 'anchor_idx': torch.from_numpy(anchor_idx), **({'vtx_rank': torch.from_numpy(vtx_rank)} if self.fps_order else {}), 'dataset_name': ds_name, 'faces_list': faces_list, 'permutation': perm if perm is not None else np.arange(n_obj), 'phys_params': torch.from_numpy(np.ascontiguousarray(phys_params)), 'skip_step_size': torch.tensor(step_size, dtype=torch.float32), 'is_point_cloud': False, 'partial_obj_index': observed, } if not train: sample['partial_obs_index'] = k sample['partial_full_obs'] = self._load_frame_objects(self.data_path / ds_name / f'frame_{t_obs:05d}') return sample def set_epoch(self, epoch: int): """Set current epoch for noise scheduling.""" self.current_epoch = int(epoch) @staticmethod def _farthest_point_sample_np(vertices: 'np.ndarray', k: int, rule: str = 'index0', start: Optional[int] = None) -> 'np.ndarray': """Deterministic farthest point sampling on an (N,3) array. Returns indices of shape (k,). rule='index0': start at vertex 0, ties resolved by the lowest index. rule='centroid_farthest': start at the vertex farthest from the centroid and resolve every tie by the lexicographically smallest coordinates, so the result depends on the vertex set only, not on its ordering. An explicit `start` overrides the rule's start vertex. """ if vertices.ndim != 2 or vertices.shape[1] != 3: raise ValueError(f"Expected vertices shape (N,3); got {vertices.shape}") n = vertices.shape[0] if k > n: raise ValueError(f"FPS requested k={k} but only N={n} vertices available") V = vertices.astype(np.float64, copy=False) lex = rule == 'centroid_farthest' def argmax(d): if not lex: return int(np.argmax(d)) cand = np.flatnonzero(d >= d.max() - 1e-9 * max(1.0, abs(float(d.max())))) if cand.shape[0] == 1: return int(cand[0]) order = np.lexsort((V[cand, 2], V[cand, 1], V[cand, 0])) return int(cand[order[0]]) selected = np.empty((k,), dtype=np.int64) dists = np.full((n,), np.inf, dtype=np.float64) if start is not None: selected[0] = int(start) elif lex: selected[0] = argmax(((V - V.mean(axis=0)) ** 2).sum(axis=1)) else: selected[0] = 0 last = V[selected[0]] for i in range(1, k): delta = V - last dists = np.minimum(dists, (delta * delta).sum(axis=1)) # Never re-select a vertex: meshes may contain coincident duplicates, whose distance to an # already selected vertex is exactly zero. Excluding selected indices keeps the output a # permutation; duplicates are appended last (ties resolved by the rule above). dists[selected[:i]] = -np.inf idx = argmax(dists) selected[i] = idx last = V[idx] return selected.astype(np.int32, copy=False) def make_collate_fn(max_objects_per_frame: Optional[int] = None, max_vertices_per_object: Optional[int] = None): """Factory function to create a collate_fn with fixed padding dimensions. Args: max_objects_per_frame: Fixed maximum number of objects per frame (None = dynamic) max_vertices_per_object: Fixed maximum vertices per object (None = dynamic) Returns: Configured collate function """ def collate_fn(batch: List[Dict]) -> Dict[str, torch.Tensor]: """Custom collate function for batching sequences.""" # Detect multi-object sample: obj_positions present with per-sample padding if 'obj_positions' in batch[0]: # Determine dimensions T = batch[0]['obj_positions'].shape[0] # Use fixed padding if configured, otherwise compute from batch if max_objects_per_frame is not None: max_nobj = int(max_objects_per_frame) # Validate that no sample exceeds the fixed limit for s in batch: _, nobj, _, _ = s['obj_positions'].shape if nobj > max_nobj: raise ValueError( f"Sample has {nobj} objects, exceeding max_objects_per_frame={max_nobj}. " f"Increase data.max_objects_per_frame in config or reduce object count in dataset." ) else: # Dynamic padding max_nobj = 0 for s in batch: _, nobj, _, _ = s['obj_positions'].shape if nobj > max_nobj: max_nobj = nobj if max_vertices_per_object is not None: max_nv = int(max_vertices_per_object) # Validate that no sample exceeds the fixed limit for s in batch: _, _, nv, _ = s['obj_positions'].shape if nv > max_nv: raise ValueError( f"Sample has {nv} vertices, exceeding max_vertices_per_object={max_nv}. " f"Increase data.max_vertices_per_object in config or reduce vertex count in dataset." ) else: # Dynamic padding max_nv = 0 for s in batch: _, _, nv, _ = s['obj_positions'].shape if nv > max_nv: max_nv = nv B = len(batch) device = batch[0]['obj_positions'].device # Allocate padded tensors. obj_positions = torch.full((B, T, max_nobj, max_nv, 3), PADDING_VALUE, dtype=torch.float32, device=device) valid_obj_mask = torch.zeros((B, max_nobj), dtype=torch.bool, device=device) valid_vtx_mask = torch.zeros((B, T, max_nobj, max_nv), dtype=torch.bool, device=device) ref_obj_positions = torch.full((B, max_nobj, max_nv, 3), PADDING_VALUE, dtype=torch.float32, device=device) # Optional anchors, padded per object to max_nobj with fixed Nq. # anchor_idx pads with index 0; anchor_mask pads with False. has_anchors = 'anchor_idx' in batch[0] has_rank = 'vtx_rank' in batch[0] if has_rank: vtx_rank = torch.full((B, max_nobj, max_nv), int(max_nv), dtype=torch.int64, device=device) if has_anchors: anchor_k = int(batch[0]['anchor_idx'].shape[1]) anchor_idx = torch.zeros((B, max_nobj, anchor_k), dtype=torch.int64, device=device) anchor_mask = torch.zeros((B, max_nobj, anchor_k), dtype=torch.bool, device=device) # Fill per-sample for b, s in enumerate(batch): t_, nobj, nv, _ = s['obj_positions'].shape obj_positions[b, :, :nobj, :nv, :] = s['obj_positions'] valid_obj_mask[b, :nobj] = s['valid_obj_mask'] valid_vtx_mask[b, :, :nobj, :nv] = s['valid_vtx_mask'] ref_obj_positions[b, :nobj, :nv, :] = s['ref_obj_positions'] if has_rank: if 'vtx_rank' not in s: raise KeyError("Batch sample missing 'vtx_rank' while others provide it") r = s['vtx_rank'].to(device=device, dtype=torch.int64) vtx_rank[b, :nobj, :nv] = torch.where(r[:, :nv] >= nv, torch.full_like(r[:, :nv], int(max_nv)), r[:, :nv]) if has_anchors: if 'anchor_idx' not in s: raise KeyError("Batch sample missing 'anchor_idx' while others provide it") if s['anchor_idx'].dim() != 2: raise ValueError(f"anchor_idx must be (Nobj, Nq); got {tuple(s['anchor_idx'].shape)}") Nq = int(s['anchor_idx'].shape[1]) if Nq != anchor_k: raise ValueError(f"Inconsistent anchor_k across batch: {Nq} vs {anchor_k}") if s['anchor_idx'].shape[0] != nobj: raise ValueError(f"anchor_idx Nobj ({s['anchor_idx'].shape[0]}) != sample Nobj ({nobj})") anchor_idx[b, :nobj, :] = s['anchor_idx'].to(device=device, dtype=torch.int64) # With strict loader checks, all anchors are valid; set mask True anchor_mask[b, :nobj, :] = True collated = { 'obj_positions': obj_positions, 'valid_obj_mask': valid_obj_mask, 'valid_vtx_mask': valid_vtx_mask, 'ref_obj_positions': ref_obj_positions, 'dt': torch.stack([s['dt'] for s in batch], dim=0) } if has_anchors: collated['anchor_idx'] = anchor_idx collated['anchor_mask'] = anchor_mask if has_rank: collated['vtx_rank'] = vtx_rank # Handle phys_params (Nobj, 3) - pad to max_nobj with PADDING_VALUE if 'phys_params' in batch[0]: phys_params_padded = torch.full((B, max_nobj, 3), PADDING_VALUE, dtype=torch.float32) for b, s in enumerate(batch): nobj = s['phys_params'].shape[0] phys_params_padded[b, :nobj, :] = s['phys_params'] collated['phys_params'] = phys_params_padded # Handle skip_step_size (scalar) - can be stacked directly if 'skip_step_size' in batch[0]: collated['skip_step_size'] = torch.stack([s['skip_step_size'] for s in batch], dim=0) # Pass through any remaining keys; tensors are stacked and must agree in shape. for key in batch[0].keys(): if key in collated: continue if isinstance(batch[0][key], torch.Tensor): collated[key] = torch.stack([sample[key] for sample in batch], dim=0) else: collated[key] = [sample[key] for sample in batch] return collated return collate_fn def create_multi_dataloaders( data_path: str, train_batch_size: int = 32, val_batch_size: int = 32, eval_batch_size: int = 32, sequence_length: int = 10, num_workers: int = 4, train_ratio: float = 0.9, val_ratio: float = 0.1, split_file: Optional[str] = 'dataset_split.json', distributed: bool = False, rank: int = 0, world_size: int = 1, fast_index: bool = True, if_phys_param: bool = False, inference_mode: bool = False, **kwargs ) -> Tuple[DataLoader, DataLoader, DataLoader]: """ Create train, validation, and test dataloaders from multiple datasets. Args: data_path: Parent directory containing the dataset folders listed in the split file train_batch_size: Batch size for training loader val_batch_size: Batch size for validation loader eval_batch_size: Batch size for test/evaluation loader (can be larger, no gradients) sequence_length: Number of frames per sequence num_workers: Number of workers for data loading train_ratio: Ratio of datasets for training val_ratio: Ratio of datasets for validation **kwargs: Additional arguments for dataset Returns: train_loader, val_loader, test_loader """ # Sampler options (with-replacement support) replacement_sampling: bool = bool(kwargs.pop('replacement_sampling', False)) replacement_steps_per_epoch = kwargs.pop('replacement_steps_per_epoch', None) # Extract padding config for collate_fn (keep in kwargs for dataset creation) max_objects_per_frame = kwargs.get('max_objects_per_frame', None) max_vertices_per_object = kwargs.get('max_vertices_per_object', None) # Add if_phys_param to kwargs for dataset creation kwargs['if_phys_param'] = if_phys_param # Create configured collate function with fixed padding collate_fn = make_collate_fn( max_objects_per_frame=max_objects_per_frame, max_vertices_per_object=max_vertices_per_object ) # Default cap on vertices per object to bound memory usage. if 'max_vertices_per_object' not in kwargs: kwargs['max_vertices_per_object'] = 1024 # Inference mode skips train dataset indexing entirely. if inference_mode: logger.info("Inference mode: Skipping train dataset indexing") train_dataset = None else: train_dataset = MultiObjectDynamics( data_path, sequence_length=sequence_length, split='train', train_ratio=train_ratio, val_ratio=val_ratio, split_file=split_file, fast_index=fast_index, **kwargs ) # Only create val/test datasets if validation is enabled (val_ratio > 0) OR inference mode if val_ratio > 0.0 or inference_mode: val_dataset = MultiObjectDynamics( data_path, sequence_length=sequence_length, split='val', train_ratio=train_ratio, val_ratio=val_ratio, split_file=split_file, fast_index=fast_index, **kwargs ) test_dataset = MultiObjectDynamics( data_path, sequence_length=sequence_length, split='test', train_ratio=train_ratio, val_ratio=val_ratio, split_file=split_file, fast_index=fast_index, **kwargs ) else: # Validation disabled: create minimal dummy datasets to avoid errors logger.info("Validation disabled (val_ratio=0.0) - skipping val/test dataset creation") val_dataset = None test_dataset = None # Create distributed samplers when required train_sampler = None val_sampler = None test_sampler = None if distributed and train_dataset is not None: # For small datasets (overfit mode), disable drop_last to avoid empty loaders # drop_last=True only when dataset is large enough (>= world_size * train_batch_size) train_drop_last = len(train_dataset) >= world_size * train_batch_size if replacement_sampling: # Replacement-based distributed sampler class ReplacementDistributedSampler(Sampler): def __init__(self, dataset, num_replicas, rank, num_samples_per_rank, seed: int = 0): self.dataset = dataset self.num_replicas = num_replicas self.rank = rank self.num_samples = int(num_samples_per_rank) self.seed = int(seed) self.epoch = 0 def __iter__(self): g = torch.Generator() g.manual_seed(self.seed + self.epoch * 997 + self.rank * 23) total = len(self.dataset) # Draw with replacement indices = torch.randint(high=total, size=(self.num_samples,), generator=g, dtype=torch.int64) return iter(indices.tolist()) def __len__(self): return self.num_samples def set_epoch(self, epoch: int): self.epoch = int(epoch) # Determine per-rank target samples if replacement_steps_per_epoch is None: # Heuristic: ceil per-rank steps from dataset size steps = max(1, (len(train_dataset) + world_size * train_batch_size - 1) // (world_size * train_batch_size)) else: steps = int(replacement_steps_per_epoch) per_rank_samples = steps * train_batch_size train_sampler = ReplacementDistributedSampler( train_dataset, num_replicas=world_size, rank=rank, num_samples_per_rank=per_rank_samples, seed=42 ) else: train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True, drop_last=train_drop_last) # Create val/test samplers for distributed mode (separate from train_dataset condition) if distributed: if val_dataset is not None: val_sampler = DistributedSampler(val_dataset, num_replicas=world_size, rank=rank, shuffle=False, drop_last=False) if test_dataset is not None: test_sampler = DistributedSampler(test_dataset, num_replicas=world_size, rank=rank, shuffle=False, drop_last=False) # Create dataloaders # For DataLoader drop_last: use same logic as DistributedSampler if train_dataset is not None: dataloader_drop_last = len(train_dataset) >= (world_size if distributed else 1) * train_batch_size else: dataloader_drop_last = False # Handle None or empty train dataset (inference mode) if train_dataset is None: logger.info("Inference mode: Creating empty train DataLoader (no training data)") # Create a minimal empty dataset for the DataLoader class EmptyDataset(torch.utils.data.Dataset): def __len__(self): return 0 def __getitem__(self, idx): raise IndexError("Empty dataset") train_loader = DataLoader( EmptyDataset(), batch_size=1, shuffle=False, num_workers=0, collate_fn=collate_fn ) elif distributed: train_loader = DataLoader( train_dataset, batch_size=train_batch_size, shuffle=False, sampler=train_sampler, num_workers=num_workers, collate_fn=collate_fn, pin_memory=True, persistent_workers=True if num_workers and num_workers > 0 else False, prefetch_factor=2 if num_workers and num_workers > 0 else None, drop_last=dataloader_drop_last ) else: # Single GPU: optional with-replacement RandomSampler if replacement_sampling: if replacement_steps_per_epoch is None: steps = max(1, (len(train_dataset) + train_batch_size - 1) // train_batch_size) else: steps = int(replacement_steps_per_epoch) num_samples = steps * train_batch_size single_sampler = RandomSampler(train_dataset, replacement=True, num_samples=num_samples) train_loader = DataLoader( train_dataset, batch_size=train_batch_size, shuffle=False, sampler=single_sampler, num_workers=num_workers, collate_fn=collate_fn, pin_memory=True, persistent_workers=True if num_workers and num_workers > 0 else False, prefetch_factor=2 if num_workers and num_workers > 0 else None, drop_last=False ) else: train_loader = DataLoader( train_dataset, batch_size=train_batch_size, shuffle=True, sampler=None, num_workers=num_workers, collate_fn=collate_fn, pin_memory=True, persistent_workers=True if num_workers and num_workers > 0 else False, prefetch_factor=2 if num_workers and num_workers > 0 else None, drop_last=dataloader_drop_last ) # Only create val/test loaders if datasets exist if val_dataset is not None: val_loader = DataLoader( val_dataset, batch_size=val_batch_size, shuffle=False, sampler=val_sampler, num_workers=num_workers, collate_fn=collate_fn, pin_memory=True, persistent_workers=True if num_workers and num_workers > 0 else False, prefetch_factor=2 if num_workers and num_workers > 0 else None, drop_last=False ) else: val_loader = None if test_dataset is not None: test_loader = DataLoader( test_dataset, batch_size=eval_batch_size, shuffle=False, sampler=test_sampler, num_workers=num_workers, collate_fn=collate_fn, pin_memory=True, persistent_workers=True if num_workers and num_workers > 0 else False, prefetch_factor=2 if num_workers and num_workers > 0 else None, drop_last=False ) else: test_loader = None return train_loader, val_loader, test_loader