# Created by BVH & VG, Jul 2025. # Flexible dataloader for the UNIFIED (anydata) pipeline. # Reads from data_dict['anydata'] and assigns Any4D entries based on availability and config. import numpy as np import torch from rich import print from custom.dataloader.utils import prepare_dense, random_choice, random_choice_k_unique, random_value, reanchor_cameras_to_view # DEPRECATED(bvh): action conditioning moved to unified_anyact / unified_anydrive: # from custom.dataloader.utils import prepare_action_anydata # from custom.dataloader.wm_utils import perturb_action_basile1, perturb_action_basile2 from custom.utils.constants import VAE_RATIO_THW from imaginaire.utils import log # Warn at most once per process that action conditioning is deprecated in unified_flex (avoid per-batch spam). _warned_action_deprecation = False def anydata_to_any4d(data_dict, phase=None, config=None, directives=None, generator=None, train_step=-1, val_step=-1, **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. # Check which highdim modalities are actually in the architecture config. _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_depth_entries = any(x.startswith('depth') for x in _entries_no_mask) has_points_entries = any(x.startswith('points') for x in _entries_no_mask) has_cams_entries = any(x.startswith('cams') for x in _entries_no_mask) # 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 anydata dictionary (this already includes some model-side transforms) 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 relevant raw / pixel modalities. anydata_rgb = anydata_dict.get('rgb', None) if load_rgb else None # dict mapping (time, camera) to (B, 3, H, W) tensor of float in [0, 1]. anydata_depth = anydata_dict.get('depth', None) if load_depth else None # dict mapping (time, camera) to (B, 3, H, W) tensor of float in [0, +inf). anydata_points = anydata_dict.get('points', None) if load_points else None # dict mapping (time, camera) to (B, 3, H, W) tensor of float in (-inf, +inf). anydata_cams = anydata_dict.get('cams', None) if load_cams else None # dict mapping (time, camera) to CameraPinhole object (re-anchored if reference_camera is set). anydata_action = anydata_dict.get('action', None) if load_action else None # dict mapping time to {'action', 'desired', 'actual': (B, D) tensor of float} (cam dim stripped in post_process_sample). # Extract relevant latent modalities. latent_rgb = anydata_dict.get('lat_rgb', None) if load_rgb else None # dict mapping (time, camera) to (B, 16, H, W) tensor of float. latent_depth = anydata_dict.get('lat_depth', None) if load_depth else None # dict mapping (time, camera) to (B, 16, H, W) tensor of float. latent_points = anydata_dict.get('lat_points', None) if load_points else None # dict mapping (time, camera) to (B, 16, H, W) tensor of float. latent_cams = anydata_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 = (anydata_rgb or anydata_depth or anydata_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 anydata_dict['timestep'].keys()]))) avail_cams = sorted(list(set([key[1] for key in anydata_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 anydata_rgb is not None or latent_rgb is not None else [] avail_modals += ['depth'] if anydata_depth is not None or latent_depth is not None else [] avail_modals += ['points'] if anydata_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 (modalities, views, frames) used either for input, or output, or neither M_used = M_data # Always use all available # View count (every phase); directives['use_views'] (int or mode) overrides config. V_upto = min(V_data, V_max) use_views = directives.get('use_views', config.use_views) if isinstance(use_views, int): V_used = max(1, min(use_views, V_upto)) elif use_views == 'rand1': V_used = rnd_choice_fn(list(range(1, V_upto + 1))) elif use_views == 'rand2': V_used = rnd_choice_fn(list(range(2, V_upto + 1))) if V_upto >= 2 else 1 elif use_views == 'all': V_used = V_upto else: raise ValueError(f"unknown use_views '{use_views}' (expected int / 'rand1' / 'rand2' / 'all')") 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] # 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) 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_est' 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 anydata_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 anydata_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: # 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') # 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: # 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 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])) 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 # cond_frames_raw: 13, total_frames_raw: 41, cond_frames_latent: 4, total_frames_latent: 11 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}") # DEPRECATED(bvh): action conditioning forced off in unified_flex (use unified_anyact/anydrive). global _warned_action_deprecation if anydata_action is not None and not _warned_action_deprecation: _warned_action_deprecation = True log.warning('unified_flex no longer conditions on actions (inv_dyn/policy/world_model); ' 'use unified_anyact or unified_anydrive for action/trajectory conditioning') (is_inv_dyn, is_policy, is_world_model) = (False, False, False) # 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: # 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 total_frames_action = T_used elif is_policy: cond_frames_action = cond_frames_raw total_frames_action = T_used elif is_world_model: # Aka forward dynamics cond_frames_action = T_used total_frames_action = 0 else: cond_frames_action = 0 total_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 entries ---- # 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 has_rgb_entries and '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}', anydata_rgb, latent_rgb, used_cams[v], cond_frames_rgb) # ^ (B, 3/16, T, H, W) tensor of bfloat16 in [0, 1] + masks. if has_depth_entries and '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}', anydata_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 has_points_entries and '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}', anydata_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 has_cams_entries and (anydata_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 = 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) # ^ (B, 6/9/32/48, T, H, W) tensor of bfloat16 in [-1, 1] + masks. # NOTE(bvh): Trajectory (driving) deprecated/removed from unified_flex to keep dataloader complexity manageable, # see unified_anydrive.py instead. # Action # DEPRECATED(bvh): action population removed along with action conditioning above; # see unified_anyact.py for the maintained version of this block. orig_action = None # if anydata_action is not None: # entry_input = 'action' # first_action_key = next(iter(anydata_action.keys())) # first_action_entry = anydata_action[first_action_key] # desired_val = first_action_entry.get('desired', None) # has_desired = isinstance(desired_val, torch.Tensor) # anydata_input = 'desired' if has_desired else 'action' # # Get expected token count from adaln config to pad shorter clips # adaln_action = config.lowdim_adaln_entries.get(entry_input, # config.lowdim_adaln_entries.get('action', None)) # target_tokens = (adaln_action[1] - adaln_action[0]) if adaln_action else None # (a4d_raw, a4d_latent) = prepare_action_anydata( # a4d_raw, a4d_latent, anydata_action, cond_frames_action, pred_frames=total_frames_action, # entry=entry_input, src_input=anydata_input, src_output='action', # target_tokens=target_tokens) # # Save original action before perturbation (for visualization) # orig_action = a4d_raw['action'].clone() # if 'perturb_action' in directives: # perturb_action = directives['perturb_action'] # perturb_seed = directives.get('perturb_seed', None) # 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) # 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_est'] 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, 'total_frames_raw': total_frames_raw, 'cond_frames_latent': cond_frames_latent, 'total_frames_latent': total_frames_latent, 'cond_frames_action': cond_frames_action, 'total_frames_action': total_frames_action, } if orig_action is not None: data_dict['meta']['orig_action'] = orig_action # (B, T, 20) tensor of float return data_dict