# Created by BVH & VG, Jul 2025. # Flexible dataloader for the VIDAR (legacy) pipeline. # Reads from data_dict['vidar'] and assigns Any4D entries based on availability and config. # NOTE: This file is deprecated in favor of unified_flex.py for new datasets. import numpy as np import torch from rich import print from custom.dataloader.utils import prepare_action_vidar, prepare_dense, random_choice, random_choice_k_unique, random_value from custom.dataloader.wm_utils import perturb_traj_basile1, perturb_traj_basile2, perturb_traj_basile3, perturb_action_basile1, perturb_action_basile2 from custom.utils.constants import VAE_RATIO_THW from imaginaire.utils import log def vidar_to_any4d(data_dict, phase=None, config=None, directives=None, generator=None, **leftover): # Valid phases assert phase in ['train', 'val', 'test'] if directives is None: directives = dict() # to avoid crashes later # load_modals represents M_max, but always use all # V_max = config.num_views # Maximum number of views supported by the current model config. # underlying data represents T_max, but we always use all # Redefine based on ENTRIES rather than STREAMS since concatenation might be involved: V_max = 1 + max([int(x[3:]) for x in config.all_highdim_entries if x.startswith('rgb') and '_mask' not in x]) # Initialize dictionaries to populate & return. a4d_raw = dict() a4d_latent = dict() load_rgb = ('rgb' in config.load_modals) load_depth = ('depth' in config.load_modals) load_points = ('points' in config.load_modals) load_cams = ('cams' in config.load_modals) load_action = ('action' in config.load_modals) load_traj = ('traj' in config.load_modals) # Get vidar dictionary (this already includes some model-side transforms) vidar_dict = data_dict['vidar'] scale_factors = data_dict['scale_factor'] # (B) tensor of float dset_names = [x.split('__')[0].lower() for x in vidar_dict['tag']] # list (B) of str # Extract relevant raw / pixel modalities. vidar_rgb = vidar_dict.get('rgb', None) if load_rgb else None # dict mapping (time, camera) to (B, 3, H, W) tensor of float in [0, 1]. vidar_depth = vidar_dict.get('depth', None) if load_depth else None # dict mapping (time, camera) to (B, 3, H, W) tensor of float in [0, +inf). vidar_points = vidar_dict.get('points', None) if load_points else None # dict mapping (time, camera) to (B, 3, H, W) tensor of float in (-inf, +inf). vidar_cams = vidar_dict.get('cams', None) if load_cams else None # dict mapping (time, camera) to CameraPinhole object. vidar_action = vidar_dict.get('action', None) if load_action else None # dict mapping (time, 0) to {'action', 'proprio': (B, 16, 20) tensor of float}. # Extract relevant latent modalities. latent_rgb = vidar_dict.get('lat_rgb', None) if load_rgb else None # dict mapping (time, camera) to (B, 16, H, W) tensor of float. latent_depth = vidar_dict.get('lat_depth', None) if load_depth else None # dict mapping (time, camera) to (B, 16, H, W) tensor of float. latent_points = vidar_dict.get('lat_points', None) if load_points else None # dict mapping (time, camera) to (B, 16, H, W) tensor of float. latent_cams = vidar_dict.get('lat_cams', None) if load_cams else None # dict mapping (time, camera) to (B, 32/48, H, W) tensor of float. # Available keys & device first_valid = (vidar_rgb or vidar_depth or vidar_points or \ latent_rgb or latent_depth or latent_points) assert first_valid is not None, 'Insufficient input information' keys = list(first_valid.keys()) B = first_valid[keys[0]].shape[0] device = first_valid[keys[0]].device # Available timesteps & cameras avail_times = sorted(list(set([key[0] for key in vidar_dict['timestep'].keys()]))) avail_cams = sorted(list(set([key[1] for key in vidar_dict['timestep'].keys()]))) T_data = len(avail_times) V_data = len(avail_cams) # Available modalities (except special cases: camera, language, action) avail_modals = [] avail_modals += ['rgb'] if vidar_rgb is not None or latent_rgb is not None else [] avail_modals += ['depth'] if vidar_depth is not None or latent_depth is not None else [] avail_modals += ['points'] if vidar_points is not None or latent_points is not None else [] M_data = len(avail_modals) # Get task probabilities p_xm = config.task_probs.get('cross_modal', -0.1) p_dvs = config.task_probs.get('dyn_view_synth', -0.1) p_fc = config.task_probs.get('forecast', -0.1) p_pe = config.task_probs.get('pose_est', -0.1) p_id = config.task_probs.get('inv_dyn', -0.1) p_p = config.task_probs.get('policy', -0.1) p_wm = config.task_probs.get('world_model', -0.1) # Prepare to roll the dice bernoulli_fn = lambda p: random_value(device, generator) < p rnd_choice_fn = lambda choices: random_choice(choices, device, generator) # Determine tokens used either for input or output M_used = M_data # Always use all available if phase == 'train': if config.use_views == 'rand1': V_used = rnd_choice_fn(list(range(1, min(V_data, V_max) + 1))) # Sometimes ignore some elif config.use_views == 'rand2': V_used = rnd_choice_fn(list(range(2, min(V_data, V_max) + 1))) # Sometimes ignore some elif config.use_views == 'all': V_used = min(V_data, V_max) else: V_used = min(V_data, V_max) T_used = T_data # Always use all available # Determine active modalities in this batch (either for input or output) if 'used_modals' in directives: used_modals = directives['used_modals'] else: used_modals = random_choice_k_unique(avail_modals, M_used, device, generator) # Determine active cameras in this batch (either for input or output) # NOTE(bvh): Make sure to respect these indices when we visualize stuff later! if 'used_cams' in directives: used_cams = directives['used_cams'] elif config.shuffle_cams2views: # NOTE(bvh): Even if V_used == V_data, this random subsampling step is important to ensure # debiasing between Any4D viewpoints and original dataset cameras. used_cams = random_choice_k_unique(avail_cams, V_used, device, generator) else: # NOTE(bvh): This is the deterministic variant, useful if we are using all cameras # in the webdataset and want to link them to Any4D views in a consistent way. used_cams = avail_cams[0:V_used] if 'tasks' in directives: # supports both str or list of str is_cross_modal = 'cross_modal' in directives['tasks'] is_dyn_view_synth = 'dyn_view_synth' in directives['tasks'] is_forecast = 'forecast' in directives['tasks'] is_pose_estim = 'pose_estim' in directives['tasks'] else: # Imagine everything is given as input as a starting point. # Then we start gradually masking out random things. # We must perform at least one task, so if none are selected, we keep retrying. is_cross_modal = False is_dyn_view_synth = False is_forecast = False is_pose_estim = False for dice_roll in range(1010): # Remove one or more modalities (except camera)? is_cross_modal = (bernoulli_fn(p_xm) and M_used >= 2) # Remove one entire viewpoint (except camera)? is_dyn_view_synth = (bernoulli_fn(p_dvs) and vidar_cams is not None and V_used >= 2) # Remove most future frames (except camera)? is_forecast = bernoulli_fn(p_fc) # Remove camera modality (if not DVS)? is_pose_estim = (bernoulli_fn(p_pe) and vidar_cams is not None and V_used >= 2) valid_tasks = ((is_cross_modal or is_dyn_view_synth or is_forecast or is_pose_estim) \ and not(is_dyn_view_synth and is_pose_estim)) if valid_tasks: break if dice_roll >= 1000: raise ValueError('Failed to select a valid set of tasks after 1000 dice rolls') # Then determine the INPUT modalities and viewpoints specifically if is_cross_modal: # Select strict, non-empty subset of modalities as input, the complement is output. if 'cond_modals' in directives: cond_modals = directives['cond_modals'] # list of str else: cond_modals = [] for dice_roll in range(1010): cond_modals = [modality for modality in used_modals if bernoulli_fn(0.5)] if 0 < len(cond_modals) < M_used: break if dice_roll >= 1000: raise ValueError('Failed to select a valid set of modalities after 1000 dice rolls') else: cond_modals = used_modals if is_dyn_view_synth or is_pose_estim: # Select strict, non-empty subset of viewpoints as input, the complement is output. # NOTE(bvh): Output means either cams is still input but predict all modalities (DVS), # or some modalities are input but predict cams (PE). 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) if is_dyn_view_synth: cond_views_cams = V_used cond_views_notcams = cond_views elif is_pose_estim: cond_views_cams = cond_views cond_views_notcams = V_used else: cond_views_cams = V_used cond_views_notcams = V_used 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': # Select strict, non-empty subset of frames as input, but at most ~66%. # e.g. [1, 5, 9, 13, 17, 21, 25] if Tp = 41 max_cond_frames = (T_used * 2) // 3 else: # Select strict, non-empty subset of frames as input, but at most ~33%. # e.g. [1, 5, 9, 13] if Tp = 41 # Use fewer conditioning frames during validation/test to better stress the model's forecasting capability. max_cond_frames = (T_used + 1) // 3 cond_frames_choices = list(range(1, max_cond_frames, VAE_RATIO_THW[0])) cond_frames_raw = random_choice(cond_frames_choices, device, generator) else: cond_frames_raw = T_used pred_frames_raw = T_used cond_frames_latent = (cond_frames_raw - 1) // VAE_RATIO_THW[0] + 1 pred_frames_latent = (pred_frames_raw - 1) // VAE_RATIO_THW[0] + 1 # This action part works quite differently from the above: # Only one of these can be active at a time, but nothing is also possible! if 'tasks' in directives: is_inv_dyn = 'inv_dyn' in directives['tasks'] is_policy = 'policy' in directives['tasks'] is_world_model = 'world_model' in directives['tasks'] else: # is_inv_dyn = bernoulli_fn(p_id) # is_policy = bernoulli_fn(p_p / (1.0 - p_id) if p_id < 1.0 else 0.0) # is_world_model = bernoulli_fn(p_wm / (1.0 - p_id - p_p) if p_id + p_p < 1.0 else 0.0) if bernoulli_fn(p_id): (is_inv_dyn, is_policy, is_world_model) = (True, False, False) elif bernoulli_fn(p_p / (1.0 - p_id) if p_id < 1.0 else 0.0): (is_inv_dyn, is_policy, is_world_model) = (False, True, False) elif bernoulli_fn(p_wm / (1.0 - p_id - p_p) if p_id + p_p < 1.0 else 0.0): (is_inv_dyn, is_policy, is_world_model) = (False, False, True) else: # actions will become DC (don't care) (is_inv_dyn, is_policy, is_world_model) = (False, False, False) assert not(is_inv_dyn and is_policy), \ 'Cannot have both inverse dynamics and policy active at the same time' assert not(is_inv_dyn and is_world_model), \ 'Cannot have both inverse dynamics and world model active at the same time' assert not(is_policy and is_world_model), \ 'Cannot have both policy and world model active at the same time' if is_inv_dyn: cond_frames_action = 0 pred_frames_action = T_used elif is_policy: cond_frames_action = cond_frames_raw pred_frames_action = T_used elif is_world_model: # Aka forward dynamics cond_frames_action = T_used pred_frames_action = 0 else: cond_frames_action = 0 pred_frames_action = 0 # NOTE(bvh): ^ these values are also used for trajectory masks # Sanity check assert is_cross_modal or is_dyn_view_synth or is_forecast or is_pose_estim or is_inv_dyn or is_policy, \ 'No tasks are selected (or more specifically, no tokens are set as output), probably a bug / mistake in directives?' # Populate Any4D dictionaries with entries (either raw or latent). # Loop over all video viewpoints & modalities to process them. # NOTE(bvh): Conditional (input) views intentionally start from the LAST one. for v in range(V_used): if 'rgb' in used_modals: cur_input = ('rgb' in cond_modals and 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}', vidar_rgb, latent_rgb, used_cams[v], cond_frames_rgb) # ^ (B, 3/16, T, H, W) tensor of bfloat16 in [0, 1] + masks. if 'depth' in used_modals: cur_input = ('depth' in cond_modals and v >= V_used - cond_views_notcams) cond_frames_depth = cond_frames_latent if cur_input else 0 (a4d_raw, a4d_latent) = prepare_dense( a4d_raw, a4d_latent, f'depth{v}', vidar_depth, latent_depth, used_cams[v], cond_frames_depth, scale_factor=scale_factors, clip_min=0.0, clip_max=1.0) # ^ (B, 1/16, T, H, W) tensor of bfloat16 in [0, 1] + masks. if 'points' in used_modals: cur_input = ('points' in cond_modals and v >= V_used - cond_views_notcams) cond_frames_points = cond_frames_latent if cur_input else 0 (a4d_raw, a4d_latent) = prepare_dense( a4d_raw, a4d_latent, f'points{v}', vidar_points, latent_points, used_cams[v], cond_frames_points, scale_factor=scale_factors, clip_min=-1.0, clip_max=1.0) # ^ (B, 3/16, T, H, W) tensor of bfloat16 in [-1, 1] + masks. if vidar_cams is not None or latent_cams is not None: # NOTE(bvh): Three important differences here: # Do not consider cond_modals, and apply special cond_views and cond_frames! cur_input = (v >= V_used - cond_views_cams) cond_frames_cams = pred_frames_latent if cur_input else 0 (a4d_raw, a4d_latent) = prepare_dense( a4d_raw, a4d_latent, f'cams{v}', vidar_cams, latent_cams, used_cams[v], cond_frames_cams, plucker_orig=False, scale_factor=scale_factors, clip_min=-1.0, clip_max=1.0) # ^ (B, 6/9/32/48, T, H, W) tensor of bfloat16 in [-1, 1] + masks. orig_action = None # Will be set if perturb_action is used if vidar_action is not None: entry_input = 'proprio' if 'proprio' in config.all_lowdim_entries else 'action' vidar_input = 'proprio' if 'proprio' in vidar_action[keys[0]] else 'action' (a4d_raw, a4d_latent) = prepare_action_vidar( a4d_raw, a4d_latent, vidar_action, cond_frames_action, pred_frames=pred_frames_action, entry_input=entry_input, entry_output='action', src_input=vidar_input, src_output='action') # ^ 2x (B, Tp, 20) tensor of float + masks. if 'perturb_action' in directives: perturb_action = directives['perturb_action'] perturb_seed = directives.get('perturb_seed', None) # Optional seed for per-sample variation # Save original action before perturbation (for visualization) orig_action = a4d_raw['action'].clone() if perturb_action == 'basile1': a4d_raw['action'] = perturb_action_basile1(a4d_raw['action'], cond_frames_raw, seed=perturb_seed) elif perturb_action == 'basile2': a4d_raw['action'] = perturb_action_basile2(a4d_raw['action'], cond_frames_raw, seed=perturb_seed) if load_traj: # all_positions = torch.stack([vidar_dict['cams'][(t, 0)].pose[:, 0:3, 3] # for t in avail_times], dim=1) # (B, T, 3) tensor of float # relative_positions = all_positions - all_positions[:, 0:1] # (B, T, 3) tensor of float all_positions = torch.stack([vidar_dict['pose'][(t, 'ego')][:, 0:3, 3] for t in avail_times], dim=1) # (B, T, 3) tensor of float relative_positions = all_positions.clone() relative_positions[:, 0, :] = 0.0 # convert absolute to relative # Y is pointing up in camera space, so take (x, z) components of forward facing camera # to represent (approximate) vehicle pose over time # NOTE(bvh): this is only valid for CAMERA_01 in DDAD trajectory = relative_positions[:, :, [0, 2]] # (B, T, 2) tensor of float trajectory[:, :, 1] *= -1.0 # now z = positive forward (instead of negative) orig_traj = trajectory.clone() if 'perturb_traj' in directives: perturb_traj = directives['perturb_traj'] perturb_seed = directives.get('perturb_seed', None) # Optional seed for per-sample variation if perturb_traj == 'basile1': trajectory = perturb_traj_basile1(orig_traj, cond_frames_raw, seed=perturb_seed) elif perturb_traj == 'basile2': trajectory = perturb_traj_basile2(orig_traj, cond_frames_raw, seed=perturb_seed) elif perturb_traj == 'basile3': trajectory = perturb_traj_basile3(orig_traj, cond_frames_raw, seed=perturb_seed) # TODO: figure out if this is float16 or float32? # TODO: consider scaling traj values by 3D scale_factor and/or traj_multiplier etc.. # TODO: always set T_traj = 41 because lowdim cannot change shape like video can # and simply pad everything with zeroes a4d_raw['traj'] = trajectory # (B, T, 2) tensor of float a4d_latent['traj_input_mask'] = torch.zeros_like(trajectory[:, :, 0:1]) # (B, T, 1) tensor of float a4d_latent['traj_input_mask'][:, 0:cond_frames_action] = 1.0 a4d_latent['traj_output_mask'] = torch.zeros_like(trajectory[:, :, 0:1]) # (B, T, 1) tensor of float a4d_latent['traj_output_mask'][:, cond_frames_action:pred_frames_action] = 1.0 # if 'scale_factor' in config.all_lowdim_entries: # a4d_latent['scale_factor'] = scale_factors.view(-1, 1, 1) # (B, 1) # a4d_latent['scale_factor_input_mask'] = torch.ones_like(scale_factors).view(-1, 1, 1) # (B, 1) # This will be processed by Any4D VAE as needed later: data_dict['a4d_raw'] = a4d_raw data_dict['a4d_latent'] = a4d_latent # Keep track of metadata, including selected modalities, viewpoints, tasks, frames, etc. # This is informative only (not used by the model). meta_tasks = [] meta_tasks += ['cross_modal'] if is_cross_modal else [] meta_tasks += ['dyn_view_synth'] if is_dyn_view_synth else [] meta_tasks += ['forecast'] if is_forecast else [] meta_tasks += ['pose_estim'] if is_pose_estim else [] meta_tasks += ['inv_dyn'] if is_inv_dyn else [] meta_tasks += ['policy'] if is_policy else [] meta_tasks += ['world_model'] if is_world_model else [] data_dict['meta'] = { 'used_modals': ','.join(used_modals), # str because variable length 'used_cams': ','.join(str(cam) for cam in used_cams), # str because variable length #### 'tasks': ','.join(meta_tasks), #### 'cond_modals': ','.join(cond_modals), 'cond_views_cams': cond_views_cams, 'cond_views_notcams': cond_views_notcams, 'cond_frames_raw': cond_frames_raw, 'pred_frames_raw': pred_frames_raw, 'cond_frames_latent': cond_frames_latent, 'pred_frames_latent': pred_frames_latent, 'cond_frames_action': cond_frames_action, 'pred_frames_action': pred_frames_action, } if load_traj: data_dict['meta']['orig_traj'] = orig_traj # (B, T, 2) tensor of float if orig_action is not None: data_dict['meta']['orig_action'] = orig_action # (B, T, 20) tensor of float return data_dict