# Created by BVH, Mar 2026. # Simplified dataloader for the UNIFIED (anydata) pipeline with action support. # Supports: dynamic view synthesis, forecast, world model. # Does NOT support: cross_modal, pose_estimation, inverse_dynamics, policy, trajectory. # 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_action_anydata, prepare_dense, random_choice, random_choice_k_unique, random_value, reanchor_cameras_to_view 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 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_action = anydata_dict.get('action', None) # dict mapping time to {'action', 'desired', 'actual': (B, D) tensor of float} (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_cams = sorted(set(key[1] for key in anydata_dict['timestep'].keys())) T_data = len(avail_times) V_data = len(avail_cams) # Task probabilities p_dvs = config.task_probs.get('dyn_view_synth', -0.1) p_fc = config.task_probs.get('forecast', -0.1) p_wm = config.task_probs.get('world_model', -0.1) bernoulli_fn = lambda p: random_value(device, generator) < p rnd_choice_fn = lambda choices: random_choice(choices, device, generator) # 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 cameras if 'used_cams' in directives: used_cams = directives['used_cams'] elif config.shuffle_cams2views: used_cams = random_choice_k_unique(avail_cams, V_used, device, generator) else: 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) # ---- Task selection: DVS, forecast, world model only ---- 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 is_world_model = False for dice_roll in range(1010): is_dyn_view_synth = (bernoulli_fn(p_dvs) and anydata_cams is not None and V_used >= 2) 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 is sampled independently (action conditioning task, orthogonal to video tasks) 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: 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: action is fully conditioned (input), not predicted if is_world_model: cond_frames_action = T_used total_frames_action = 0 else: cond_frames_action = 0 total_frames_action = 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, active_resolution=resolutions[used_cams[v]]) # 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, active_resolution=resolutions[used_cams[v]]) # Action 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] # Prefer 'desired' for input (= commanded EE state at current timestep). # Fallback to 'action' if desired is not available (e.g. DROID, RH20T, HumEv). # After collation, None values may become lists of Nones, so check for tensors. # TODO(bvh): investigate RH20T -- first_action_entry can be a list instead of dict # after collation, causing .get() to crash. Likely a data format issue in RH20T_d14. 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) # ^ 2x (B, T, D) tensor of float + masks. # 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) # Optional seed for per-sample variation 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) # Store results 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 += ['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( # str because variable length m for m, v in [('rgb', anydata_rgb or latent_rgb), ('cams', anydata_cams or latent_cams), ('action', anydata_action)] if v is not None), 'used_cams': ','.join(str(cam) for cam in used_cams), # str because variable length #### '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, 'total_frames_raw': total_frames_raw, 'cond_frames_latent': cond_frames_latent, 'total_frames_latent': total_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 return data_dict