# Created by BVH, Apr 2026. # Dataloader for DRIVING datasets in the UNIFIED (anydata) pipeline. # Based on unified_flex (trajectory) and unified_anyact (simplified task logic). # Supports: forecast, world model (trajectory conditioning), dyn_view_synth (incl. cams-free). # Does NOT support: cross_modal, pose_estimation, inverse_dynamics, policy, robot action. # Handles different driving datasets with different # cameras via dataset-specific cam mapping. import numpy as np import torch from rich import print from custom.dataloader.utils import prepare_dense, prepare_traj_anydata, random_choice, random_value, reanchor_cameras_to_view from custom.dataloader.wm_utils import compute_trajectory_from_extrinsics, compute_trajectory_from_ego_pose, perturb_traj_basile1, perturb_traj_basile2, perturb_traj_basile3 from custom.utils.constants import VAE_RATIO_THW from imaginaire.utils import log # Index 0 = ego/front; remaining entries are whatever directional spread we want next. # We walk this list in order, look up each name in the per-sample cam_names from metadata, # and emit the matching int cam_idx; missing names are skipped. # TODO(bvh): verify per dataset CAM_ORDER = { 'argo': [ 'ring_front_center', # front 'ring_front_left', # left 'ring_front_right', # right 'ring_side_left', # left 'ring_side_right', # right 'ring_rear_left', # left 'ring_rear_right', # right ], 'ddad': [ 'CAMERA_01', # front 'CAMERA_05', # left 'CAMERA_06', # right 'CAMERA_09', # back 'CAMERA_07', # left 'CAMERA_08', # right ], 'gaia': [ # NOTE(bvh): I tried to align the first 5 views with TSS4 MDC by visual inspection, # based on the sample data shared by Philip Tsai and Yiting Liu. # => the "Front FCM" camera appears most closely aligned with "front_wide". # See also https://docs.google.com/document/d/1llYJmzku7xfcTeqVjdnZvd7Tqla_f4gjUfepqXZNwUs/edit # The idea is to enable better finetuning later when only these 5 cameras are available. # BUT we also need to align with OTHER datasets as well as possible, # hence first 4 views = front > left > right > back, even if not all PVM. # NOTE(bvh): this is also used by viz_anydrive_gaia() to map camera names to view indices. 'front_wide', # front 'left_side_pvm', # left 'right_side_pvm', # right 'rear_pvm', # back 'front_pvm', # front 'left_side_forward', # left 'right_side_forward', # right 'rear_tele', # back 'front_tele', # front 'left_side_rearward', # left 'right_side_rearward', # right ], 'lyftl5': [ '0', # front '5', # left '1', # right '4', # back '3', # left '2', # right '6', # back ], 'pd4d': [ 'yaw-0', # front 'yaw-neg-60', # left 'yaw-60', # right ], 'pdv2': [ 'camera_01', # front 'camera_05', # left 'camera_06', # right 'camera_09', # back 'camera_07', # left 'camera_08', # right ], 'waymo': [ '0', # front (FRONT) '1', # left (FRONT_LEFT) '2', # right (FRONT_RIGHT) '3', # left (SIDE_LEFT) '4', # right (SIDE_RIGHT) ], } # Cameras whose translation vectors get weighted-averaged to define the ego trajectory. # Midpoint approximates vehicle center, which is a better trajectory proxy than any single camera origin. # Weights normalize internally so absolute scale is irrelevant; entries with missing cam_name are dropped. # TODO(bvh): verify per dataset EGO_TRAJ_PROXY = { 'argo': [('ring_front_center', 0.34), ('ring_rear_left', 0.33), ('ring_rear_right', 0.33)], # no rear cam 'ddad': [('CAMERA_01', 0.5), ('CAMERA_09', 0.5)], # gaia: union of wide and pvm configs; drop-missing + renormalize picks the right subset per config. 'gaia': [('front_wide', 0.5), ('left_side_rearward', 0.25), ('right_side_rearward', 0.25), ('front_pvm', 0.5), ('rear_pvm', 0.5)], 'lyftl5': [('0', 0.5), ('4', 0.5)], 'pd4d': [('yaw-0', 1.0)], # no rear cam 'pdv2': [('camera_01', 0.5), ('camera_09', 0.5)], 'waymo': [('0', 0.34), ('3', 0.33), ('4', 0.33)], # no rear cam } def _dset_key(dset_name): my_key = dset_name.lower().split('_')[0] match = [k for k in CAM_ORDER.keys() if my_key.startswith(k)] assert len(match) == 1, (f'None or multiple CAM_ORDER entries match {dset_name!r} ' f'among {list(CAM_ORDER.keys())}: {match}') return match[0] def _used_cam_inds(dset_name, cam_names): ''' Walk CAM_ORDER for this dataset; map names to int indices via cam_names. Skip missing. :param dset_name: str, e.g. 'DDAD', 'LyftL5'. :param cam_names: list of str, names in dict-key (= cam_idx) order, from metadata['cameras'][b]. :return: list of int cam indices in priority order. Falls back to all cam_names indices on miss. ''' name_to_idx = {name: idx for idx, name in enumerate(cam_names)} order = CAM_ORDER.get(_dset_key(dset_name), None) used = [name_to_idx[name] for name in order if name in name_to_idx] return used def _ego_traj_proxy_inds(dset_name, cam_names): ''' Translate EGO_TRAJ_PROXY (cam_name, weight) entries to (cam_idx, weight); drop missing names. :return: list of (int, float) tuples; falls back to [(0, 1.0)] if nothing matched. ''' name_to_idx = {name: idx for idx, name in enumerate(cam_names)} proxy = EGO_TRAJ_PROXY.get(_dset_key(dset_name), None) resolved = [(name_to_idx[name], w) for name, w in proxy if name in name_to_idx] assert resolved, (f'No EGO_TRAJ_PROXY cams matched cam_names={cam_names} for {dset_name!r}. ' f'Expected at least one of {[n for n, _ in proxy]} to be present.') return resolved def anydata_to_any4d(data_dict, phase=None, config=None, directives=None, generator=None, train_step=-1, val_step=-1, **leftover): # NOTE(bvh): temporal stride/trim (incl. cross-rank synchronized lengths via # sync_temporal_transforms) is now handled post-collate in DataTransforms, not here. assert phase in ['train', 'val', 'test'] if directives is None: directives = dict() # V_max from entries (not config.num_views, since concatenation might be involved) _entries_no_mask = [x for x in config.all_highdim_entries if '_mask' not in x and '_ref' not in x] # _ref is for Wan V_max = 1 + max([int(x[3:]) for x in _entries_no_mask if x.startswith('rgb')]) has_rgb_entries = any(x.startswith('rgb') for x in _entries_no_mask) has_cams_entries = any(x.startswith('cams') for x in _entries_no_mask) a4d_raw = dict() a4d_latent = dict() anydata_dict = data_dict['anydata'] meta_ds_dict = anydata_dict['metadata'] resolutions = meta_ds_dict['resolution'] # dict-V {cam_idx: [H[B], W[B]]} scale_factors = data_dict['scale_factor'] # (B) tensor of float dset_names = data_dict['dset_name'] # list (B) of str dset_name = dset_names[0] # str # Extract whatever modalities are available from anydata_dict. anydata_rgb = anydata_dict.get('rgb', None) # dict mapping (time, camera) to (B, 3, H, W) tensor of float in [0, 1]. anydata_cams = anydata_dict.get('cams', None) # dict mapping (time, camera) to CameraPinhole object (re-anchored if reference_camera is set). anydata_ego_pose = anydata_dict.get('ego_pose', None) # dict mapping time to (B, 4, 4) ego "Tcw" pose (driving only; cam dim stripped in post_process_sample). latent_rgb = anydata_dict.get('lat_rgb', None) # dict mapping (time, camera) to (B, 16, H, W) tensor of float. latent_cams = anydata_dict.get('lat_cams', None) # dict mapping (time, camera) to (B, 32/48, H, W) tensor of float. # Available keys & device first_valid = anydata_rgb or latent_rgb assert first_valid is not None, 'Need at least rgb or lat_rgb' keys = list(first_valid.keys()) B = first_valid[keys[0]].shape[0] device = first_valid[keys[0]].device # Available timesteps & camera indices. avail_times = sorted(set(key[0] for key in anydata_dict['timestep'].keys())) avail_cam_inds = sorted(set(key[1] for key in anydata_dict['timestep'].keys())) T_data = len(avail_times) V_data = len(avail_cam_inds) # NOTE(bvh): using cam_names from batch[0] only. Ordering may differ across batch elements; # in that case ordering will be "correct" for sample 0 and only approximately correct for the # rest. Mostly an issue if upstream randomly subsamples cameras per sample. first_cam_names = anydata_dict['metadata']['cameras'][0] # Get task probabilities # Driving WM configs run forecast (always, by default) + optional world model (trajectory # conditioning); DVS configs (r[n]cdvs11_gaia) set dyn_view_synth=1.0 and forecast=0.0. p_dvs = config.task_probs.get('dyn_view_synth', -0.1) p_fc = config.task_probs.get('forecast', 1.0) p_wm = config.task_probs.get('world_model', -0.1) # Prepare to roll the dice bernoulli_fn = lambda p: random_value(device, generator) < p if config.anydrive_cams2views: # NOTE(bvh): We apply a fixed order per dataset; # this replaces the logic from robotics dataloaders (e.g. shuffle_cams2views, use_views). used_cams = _used_cam_inds(dset_name, first_cam_names) used_cams = used_cams[0:V_max] V_used = len(used_cams) T_used = T_data # Always use all available else: # Leave order loaded by dataset (following YAML) intact. # e.g. for infotech 5V->11V, we want: # output (rgb0-5) = [front_tele,rear_tele,left_side_forward,right_side_forward,left_side_rearward,right_side_rearward] # input (rgb6-10) = [front_wide,front_pvm,rear_pvm,left_side_pvm,right_side_pvm] used_cams = avail_cam_inds[0:V_max] V_used = len(used_cams) # derive from actual cams (avoid over-index if fewer than V_max available) T_used = T_data # Always use all available # Re-anchor cameras to a reference view if configured. # This makes plucker rays relative to whichever camera was assigned to the reference view. ref_view = getattr(config, 'reference_view', False) if ref_view is not False and ref_view is not None and anydata_cams is not None: ref_cam_idx = used_cams[ref_view] per_timestep = getattr(config, 'ref_view_rel_per_timestep', True) anydata_cams = reanchor_cameras_to_view(anydata_cams, ref_cam_idx, per_frame=per_timestep) # ---- Task selection: DVS, forecast, world model ---- if 'tasks' in directives: is_dyn_view_synth = 'dyn_view_synth' in directives['tasks'] is_forecast = 'forecast' in directives['tasks'] is_world_model = 'world_model' in directives['tasks'] else: is_dyn_view_synth = False is_forecast = False for dice_roll in range(1010): # NOTE(bvh): camera poses are only required for DVS if the architecture actually # conditions on them; with fixed rigs (poses constant across samples), cams-free # configs (rncdvs11_gaia) learn the geometry implicitly from the view index. is_dyn_view_synth = (bernoulli_fn(p_dvs) and V_used >= 2 and (anydata_cams is not None or latent_cams is not None or not has_cams_entries)) is_forecast = bernoulli_fn(p_fc) if is_dyn_view_synth or is_forecast: break if dice_roll >= 1000: # Do not change this to a warning, it indicates a user config mistake that needs to be fixed raise ValueError('Failed to select a valid set of tasks after 1000 dice rolls') # World model (trajectory conditioning) sampled independently is_world_model = bernoulli_fn(p_wm) # DVS: condition on subset of views, predict the rest if is_dyn_view_synth: if 'num_pred_views' in directives: cond_views = V_used - directives.get('num_pred_views', 1) else: cond_choices_views = directives.get('cond_choices_views', list(range(1, V_used))) cond_views = random_choice(cond_choices_views, device, generator) cond_views_cams = V_used cond_views_notcams = cond_views else: cond_views_cams = V_used cond_views_notcams = V_used # Forecast: condition on first N frames, predict the rest if is_forecast: if 'cond_frames_raw' in directives: cond_frames_raw = directives['cond_frames_raw'] else: if 'cond_frames_choices' in directives: cond_frames_choices = directives['cond_frames_choices'] else: if phase == 'train': max_cond_frames = (T_used * 2) // 3 else: max_cond_frames = (T_used + 1) // 3 cond_frames_choices = list(range(1, max_cond_frames, VAE_RATIO_THW[0])) if not cond_frames_choices: # very short clips -> condition on 1 frame cond_frames_choices = [1] cond_frames_raw = random_choice(cond_frames_choices, device, generator) else: cond_frames_raw = T_used total_frames_raw = T_used cond_frames_latent = (cond_frames_raw - 1) // VAE_RATIO_THW[0] + 1 total_frames_latent = (total_frames_raw - 1) // VAE_RATIO_THW[0] + 1 log.debug(f"cond_frames_raw: {cond_frames_raw}, total_frames_raw: {total_frames_raw}, " f"cond_frames_latent: {cond_frames_latent}, total_frames_latent: {total_frames_latent}") # World model: trajectory is fully conditioned (input), not predicted if is_world_model: cond_frames_traj = T_used total_frames_traj = 0 else: cond_frames_traj = 0 total_frames_traj = 0 assert is_dyn_view_synth or is_forecast, \ 'At least one video task (DVS or forecast) must be active' # ---- Populate Any4D entries ---- # TODO(bvh): detect which tensors are padded (or even all zero, since we can now have # different # loaded cameras WITHIN the same batch too), and zero out all 3 masks accordingly! # This info must be loaded from metadata['resolution'] (not 'resolution_original') # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ for v in range(V_used): # RGB if has_rgb_entries and (anydata_rgb is not None or latent_rgb is not None): cur_input = (v >= V_used - cond_views_notcams) cond_frames_rgb = cond_frames_latent if cur_input else 0 (a4d_raw, a4d_latent) = prepare_dense( a4d_raw, a4d_latent, f'rgb{v}', anydata_rgb, latent_rgb, used_cams[v], cond_frames_rgb) # Cameras (plucker) if has_cams_entries and (anydata_cams is not None or latent_cams is not None): cur_input = (v >= V_used - cond_views_cams) cond_frames_cams = total_frames_latent if cur_input else 0 (a4d_raw, a4d_latent) = prepare_dense( a4d_raw, a4d_latent, f'cams{v}', anydata_cams, latent_cams, used_cams[v], cond_frames_cams, plucker_orig=False, scale_factor=scale_factors, clip_min=-1.0, clip_max=1.0) # ---- Ego vehicle trajectory parsing ---- # Only parse a trajectory when the architecture has a traj slot (DVS configs do not). adaln_traj = config.lowdim_adaln_entries.get('traj', None) # In some datasets, e.g. GAIA, we now have ego_pose and ego_vehicle keys available, so try this first traj_from_ego = None if adaln_traj is not None and anydata_ego_pose is not None: traj_from_ego = compute_trajectory_from_ego_pose( anydata_dict, avail_times) # (B, T, 2) # Fallback to weighted mean of camera extrinsics through EGO_TRAJ_PROXY traj_from_cams = None if adaln_traj is not None and anydata_cams is not None: ego_traj_proxy = _ego_traj_proxy_inds(dset_name, first_cam_names) # ^ list of (cam_idx: int, weight: float) tuples traj_from_cams = compute_trajectory_from_extrinsics( anydata_dict, avail_times, ego_traj_proxy) # (B, T, 2) orig_traj = traj_from_ego if traj_from_ego is not None else traj_from_cams if orig_traj is not None: # Bake in multiplication since we want downstream modifications to be directly comparable traj_multiplier = getattr(config, 'traj_multiplier', 1.0) orig_traj = orig_traj.clone() * traj_multiplier # ^ (B, T_used, 2) # Get expected token count from adaln config to pad shorter clips target_tokens = (adaln_traj[1] - adaln_traj[0]) if adaln_traj else None trajectory = orig_traj.clone() if 'perturb_traj' in directives: perturb_traj = directives['perturb_traj'] perturb_seed = directives.get('perturb_seed', None) if perturb_traj == 'basile1': trajectory = perturb_traj_basile1(trajectory, cond_frames_raw, seed=perturb_seed) elif perturb_traj == 'basile2': trajectory = perturb_traj_basile2(trajectory, cond_frames_raw, seed=perturb_seed) elif perturb_traj == 'basile3': trajectory = perturb_traj_basile3(trajectory, cond_frames_raw, seed=perturb_seed) # NOTE(bvh): trajectory is input-only for now (pred_frames=0). (a4d_raw, a4d_latent) = prepare_traj_anydata( a4d_raw, a4d_latent, trajectory, cond_frames_traj, pred_frames=total_frames_traj, target_tokens=target_tokens) # Store results data_dict['a4d_raw'] = a4d_raw data_dict['a4d_latent'] = a4d_latent # Metadata meta_tasks = [] meta_tasks += ['dyn_view_synth'] if is_dyn_view_synth else [] meta_tasks += ['forecast'] if is_forecast else [] meta_tasks += ['world_model'] if is_world_model else [] data_dict['meta'] = { 'used_modals': ','.join( m for m, v in [('rgb', anydata_rgb or latent_rgb), ('cams', anydata_cams or latent_cams), ('traj', orig_traj)] if v is not None), 'used_cams': ','.join(str(cam) for cam in used_cams), # a bit ugly but we should avoid changing standards 'tasks': ','.join(meta_tasks), 'cond_modals': 'rgb', # TODO update 'cond_views_cams': cond_views_cams, 'cond_views_notcams': cond_views_notcams, 'cond_frames_raw': cond_frames_raw, 'total_frames_raw': total_frames_raw, 'cond_frames_latent': cond_frames_latent, 'total_frames_latent': total_frames_latent, 'cond_frames_traj': cond_frames_traj, 'total_frames_traj': total_frames_traj, } if orig_traj is not None: data_dict['meta']['orig_traj'] = orig_traj return data_dict