# Created by BVH, Apr 2025. import os import torch from anydata.geometry.cameras.base import CameraBase from anydata.geometry.camera import Camera from custom.dataset.camera_utils import reanchor_cameras_to_view # noqa: F401 (re-export) from custom.utils.constants import VAE_RATIO_THW from imaginaire.utils import log def random_value(device, generator): ''' Returns a single random float in [0, 1) using the given torch generator for reproducibility. ''' val = torch.rand(1, device=device, generator=generator).item() return val def random_choice(choices, device, generator): ''' Picks and returns a single random element from the given list using the torch generator. ''' idx = torch.randint(0, len(choices), (1,), device=device, generator=generator).item() choice = choices[idx] return choice def random_choice_k_unique(choices, k, device, generator): ''' Picks k unique elements from the given list (without replacement) using the torch generator. Returns a list of the selected elements. ''' num_choices = len(choices) indices = torch.randperm(num_choices, device=device, generator=generator) selected_indices = indices[0:k] selected_choices = [choices[idx.item()] for idx in selected_indices] return selected_choices def maybe_clamp(data, clip_min, clip_max): ''' Clamps the tensor to [clip_min, clip_max] if either bound is not None. Returns data unchanged otherwise. ''' if not(clip_min is None and clip_max is None): data = torch.clamp(data, min=clip_min, max=clip_max) return data def biased_prng(x, p, seed=0): ''' Deterministic biased bit from an integer index. Uses a splitmix32-style finalizer to hash (x, seed) to a uniform 32-bit value, then thresholds at p. NOTE(bvh): This is designed to be used for randomness that has to be synchronized across ALL ranks (e.g. factors that noticeably affect forward pass speed). :param x: int, input index. :param p: float, probability of returning 1. :param seed: int, selects independent PRBS streams. :return: int, 0 or 1. ''' h = x * 0x45d9f3b ^ seed * 0x27d4eb2d h = ((h >> 16) ^ h) * 0x45d9f3b h = ((h >> 16) ^ h) * 0x45d9f3b h = (h >> 16) ^ h h &= 0xFFFFFFFF return int(h < p * 0x100000000) def prepare_dense(a4d_raw, a4d_latent, prefix, anydata_raw, anydata_latent, cam_idx, cond_frames_latent, pred_frames_latent=None, plucker_orig=False, scale_factor=1.0, clip_min=None, clip_max=None, clip_warn_thres=0.7, active_resolution=None): ''' This method can handle both raw and latent data, and typically updates one or both dictionaries as a function of the input information. Content gets moved from a4d_raw to a4d_latent later by a4d_model and a4d_vae. :param prefix (str): Any4D entry key name. :param anydata_raw (dict): data dict mapping (time, camera) to either (B, Cp, Hp, Wp) tensor of float, or CameraPinhole object. :param anydata_latent (dict): data dict mapping (time, camera) to (B, Cl, Hl, Wl) tensor of float. Use when the data has been pre-encoded / cached by the VAE, in which case the input & output masks become the same shape. :param cam_idx (int). :param cond_frames_latent: int or (B) tensor of int. Set to zero for output only. :param pred_frames_latent: int or (B) tensor of int. This is absolute, not relative to cond_frames. :param plucker_orig (bool): Whether to include the origin coordinates in the cams representation. :param scale_factor: float or (B) tensor of float. Applied for 3D / geometric data only. :param clip_min (float). :param clip_max (float). :param clip_warn_thres (float). :param active_resolution: raw resolution (BEFORE padding) as a post-collate [H[B], W[B]] list of 2 tensors :return (a4d_raw, a4d_latent). ''' assert prefix not in a4d_raw.keys(), f'{prefix} already in a4d_raw' assert prefix not in a4d_latent.keys(), f'{prefix} already in a4d_latent' assert (anydata_raw is None) != (anydata_latent is None), \ 'Must provide either anydata_raw or anydata_latent, but not both' actual_data = anydata_raw or anydata_latent avail_times = sorted(list(set([key[0] for key in actual_data.keys()]))) first_key = next(iter(actual_data.keys())) first_val = actual_data[first_key] # NOTE(bvh): use duck-typing to support both anydata's CameraBase and vidar's # vidar.geometry.cameras.base.CameraBase (different module path, same get_plucker API). convert_plucker = isinstance(first_val, CameraBase) or hasattr(first_val, 'get_plucker') is_latent = anydata_latent is not None if is_latent: # Pre-encoded / cached latent space assert avail_times == [0], f'Cached latent space data must have T=1, but got {avail_times}' a4d_data = actual_data[(0, cam_idx)] (B, Cl, Tl, Hl, Wl) = a4d_data.shape Tp = (Tl - 1) * VAE_RATIO_THW[0] + 1 Hp = Hl * VAE_RATIO_THW[1] Wp = Wl * VAE_RATIO_THW[2] else: # Raw / pixel space # NOTE(bvh): get_plucker T-loop ~= 52 ms/view; see 23.md. if convert_plucker: a4d_data = [actual_data[(t, cam_idx)].get_plucker(with_orig=plucker_orig) for t in avail_times] # ^ list of (B, 6/9, H, W) tensor of float32 with (rays, cross, orig). else: a4d_data = [actual_data[(t, cam_idx)] for t in avail_times] # ^ list of (B, 1/3/16/32/48, H, W) tensor of float32. a4d_data = torch.stack(a4d_data, dim=2) # (B, C, T, H, W) tensor of float32. (B, Cp, Tp, Hp, Wp) = a4d_data.shape # NOTE(yams_any4d 2026-07-09): windows sampled near episode tails can be # truncated to lengths violating the 1+4k VAE frame arithmetic (e.g. 23). # Trim to the largest valid 1+4k length instead of asserting out, so # short-episode datasets (raiden YAM) train without tail crashes. _Tp_valid = (Tp - 1) // VAE_RATIO_THW[0] * VAE_RATIO_THW[0] + 1 if _Tp_valid != Tp: a4d_data = a4d_data[:, :, :_Tp_valid] Tp = _Tp_valid Tl = (Tp - 1) // VAE_RATIO_THW[0] + 1 Hl = Hp // VAE_RATIO_THW[1] Wl = Wp // VAE_RATIO_THW[2] # Enforce nicely dividing dimensions for VAE processing assert Tp == (Tl - 1) * VAE_RATIO_THW[0] + 1, f'Tp = {Tp}, Tl = {Tl}' assert Hp == Hl * VAE_RATIO_THW[1], f'Hp = {Hp}, Hl = {Hl}' assert Wp == Wl * VAE_RATIO_THW[2], f'Wp = {Wp}, Wl = {Wl}' device = a4d_data.device tensor_kwargs = dict(device=device, dtype=torch.bfloat16) # WARNING(bvh): ^ we should take care to convert float32 to bfloat16 as late as possible, since # the mantissa of bfloat16 is only 7 bits, so we might run into precision issues otherwise # (especially when normalization steps are involved). if is_latent: a4d_data = a4d_data.detach().clone().to(**tensor_kwargs) else: # Scale & clip all involved depth / XYZ coordinate / pose translation values. # NOTE(bvh): This has a similar effect to scale_depth, scale_points, and scale_cams in # the augmentation utils, except it is done later to avoid changing the # underlying data dict (which could mess up visualizations / geometry etc). # For cams, we simulate scaling the translation component cams[k].pose[:, 0:3, 3], # which affects cross and orig the same way, but rays must stay the same. if isinstance(scale_factor, float): scale_factor = torch.ones(B, device=a4d_data.device, dtype=a4d_data.dtype) * scale_factor if convert_plucker: a4d_data[:, 3:] = a4d_data[:, 3:] * scale_factor[:, None, None, None, None] else: a4d_data = a4d_data * scale_factor[:, None, None, None, None] a4d_data = a4d_data.detach().clone().to(**tensor_kwargs) a4d_data = maybe_clamp(a4d_data, clip_min, clip_max) # ^ (B, C, T, H, W) tensor of bfloat16. min_oob_ratio = torch.sum(a4d_data <= clip_min) / a4d_data.numel() if clip_min is not None else -0.0 max_oob_ratio = torch.sum(a4d_data >= clip_max) / a4d_data.numel() if clip_max is not None else -0.0 if min_oob_ratio + max_oob_ratio > clip_warn_thres: log.warning(f'Many clamped values for {prefix}: ' f'min_oob_ratio = {min_oob_ratio:.3f}, max_oob_ratio = {max_oob_ratio:.3f}', rank0_only=False) if pred_frames_latent is None: pred_frames_latent = Tl if isinstance(cond_frames_latent, int): cond_frames_latent = torch.ones(B, device=device, dtype=torch.int32) * cond_frames_latent if isinstance(pred_frames_latent, int): pred_frames_latent = torch.ones(B, device=device, dtype=torch.int32) * pred_frames_latent # NOTE(bvh): For memory efficiency, we always keep all the masks in latent space. input_mask = torch.zeros(B, 1, Tl, Hl, Wl, **tensor_kwargs) output_mask = torch.zeros(B, 1, Tl, Hl, Wl, **tensor_kwargs) for b in range(B): # The first subsequence (if any) are for conditioning. input_mask[b, :, 0:cond_frames_latent[b]] = 1.0 # The remaining subsequence (if any) are for generation. if pred_frames_latent[b] > cond_frames_latent[b]: output_mask[b, :, cond_frames_latent[b]:pred_frames_latent[b]] = 1.0 # These could have been spatially padded by any4d_collate(). if active_resolution is not None: # TODO(bvh): debug/verify Hpa, Wpa = int(active_resolution[0][b]), int(active_resolution[1][b]) Hla = Hpa // VAE_RATIO_THW[1] Wla = Wpa // VAE_RATIO_THW[2] assert Hpa == Hla * VAE_RATIO_THW[1], f'Hpa = {Hpa}, Hla = {Hla}' assert Wpa == Wla * VAE_RATIO_THW[2], f'Wpa = {Wpa}, Wla = {Wla}' if Hla < Hl: input_mask[b, :, :, Hla:, :] = 0.0 output_mask[b, :, :, Hla:, :] = 0.0 if Wla < Wl: input_mask[b, :, :, :, Wla:] = 0.0 output_mask[b, :, :, :, Wla:] = 0.0 if is_latent: a4d_latent[prefix] = a4d_data else: a4d_raw[prefix] = a4d_data a4d_latent[f'{prefix}_input_mask'] = input_mask a4d_latent[f'{prefix}_output_mask'] = output_mask return (a4d_raw, a4d_latent) def prepare_action_vidar(a4d_raw, a4d_latent, action_data, cond_frames, pred_frames=None, entry_input='proprio', entry_output='action', src_input='proprio', src_output='action'): ''' Like prepare_dense, this puts content in a4d_raw and masks in a4d_latent. Content gets moved to a4d_latent later by vae.encode_a4d_batch in a4d_model.py. NOTE(bvh): This will generally include a mix of proprio and action into one set of tokens (with the same number of timesteps as frames), typically intended for conditioning and generation purposes respectively. :param action_data (dict): data dict mapping (time, 0) to: - LBM: {'action', 'proprio': (B, 16, 20) tensor of float}. - RoboCasa: {'action': (B, 7) tensor of float}. :param cond_frames: int or (B) tensor of int. Set to zero for output only. :param pred_frames: int or (B) tensor of int. This is absolute, not relative to cond_frames. :param src_input (str): 'proprio' or 'action'. :param src_output (str): 'action' or 'proprio'. :return (a4d_raw, a4d_latent). ''' first_action_key = next(iter(action_data.keys())) first_action_val = action_data[first_action_key][src_input] B = first_action_val.shape[0] device = first_action_val.device tensor_kwargs = dict(device=device, dtype=torch.bfloat16) if first_action_val.ndim == 3: # (B, T, C) per value so we need to take the first time step, e.g. LBM. collapse_time = True elif first_action_val.ndim == 2: # (B, C) per value so we not need to remove anything, e.g. RoboCasa. collapse_time = False else: raise ValueError(f'Unexpected input shape: {first_action_val.shape}') avail_times = sorted(list(set([key[0] for key in action_data.keys()]))) T = len(avail_times) if pred_frames is None: pred_frames = T if isinstance(cond_frames, int): cond_frames = torch.ones(B, device=device, dtype=torch.int32) * cond_frames if isinstance(pred_frames, int): pred_frames = torch.ones(B, device=device, dtype=torch.int32) * pred_frames if collapse_time: all_input = torch.stack([action_data[(i, 0)][src_input][:, 0] for i in range(T)], dim=1) all_output = torch.stack([action_data[(i, 0)][src_output][:, 0] for i in range(T)], dim=1) else: all_input = torch.stack([action_data[(i, 0)][src_input] for i in range(T)], dim=1) all_output = torch.stack([action_data[(i, 0)][src_output] for i in range(T)], dim=1) # ^ LBM: 2x (B, T, 20) tensor of float. # ^ RoboCasa: (B, T, 7) tensor of float. # all_input *= multiplier # all_output *= multiplier input_mask = torch.zeros(B, T, 1, **tensor_kwargs) output_mask = torch.zeros(B, T, 1, **tensor_kwargs) for b in range(B): input_mask[b, 0:cond_frames[b]] = 1.0 if pred_frames[b] > cond_frames[b]: output_mask[b, cond_frames[b]:pred_frames[b]] = 1.0 # ^ 2x (B, T, 1) tensor of bfloat16. if entry_input == entry_output: a4d_mixed = all_input * input_mask + all_output * output_mask a4d_mixed = a4d_mixed.detach().clone().to(**tensor_kwargs) # ^ (B, T, 20) tensor of float. a4d_raw[entry_input] = a4d_mixed a4d_raw[entry_output] = a4d_mixed a4d_latent[f'{entry_input}_input_mask'] = input_mask a4d_latent[f'{entry_output}_output_mask'] = output_mask else: a4d_raw[entry_input] = all_input a4d_raw[entry_output] = all_output a4d_latent[f'{entry_input}_input_mask'] = input_mask a4d_latent[f'{entry_input}_output_mask'] = output_mask * 0.0 a4d_latent[f'{entry_output}_input_mask'] = input_mask * 0.0 a4d_latent[f'{entry_output}_output_mask'] = output_mask return (a4d_raw, a4d_latent) # NOTE(bvh): T range in `video_entries` / `lowdim_*_entries` (a4d_config) are upper bounds, # not exact lengths. The dataloader is responsible for honoring T_actual <= T_max: # lowdim entries zero-pad to target_tokens here; highdim (video) entries can however vary in # tokens. def prepare_action_anydata(a4d_raw, a4d_latent, action_data, cond_frames, pred_frames=None, entry='action', src_input='actual', src_output='action', target_tokens=None): '''Prepare action entry for Any4D. Puts content in a4d_raw, masks in a4d_latent. Anydata action tensors are (B, D) per timestep. :param action_data: dict keyed by scalar time OR (time, cam) tuple, mapping to dicts with subkeys 'action' (B, D) and optionally 'desired' (B, D). :param cond_frames: int or (B) tensor of int. Set to zero for output only. :param pred_frames: int or (B) tensor of int. Absolute, not relative to cond_frames. :param entry: a4d entry name (both input and output share the same entry). :param src_input: key in action_data dicts for conditioning signal ('desired' or 'action'). :param src_output: key in action_data dicts for generation signal ('action', etc.). :param target_tokens: int or None, expected token count from adaln config. :return: (a4d_raw, a4d_latent). ''' sample_key = next(iter(action_data.keys())) keys_are_tuples = isinstance(sample_key, tuple) if keys_are_tuples: avail_times = sorted(set(key[0] for key in action_data.keys())) first_cam = sample_key[1] def _make_key(t): return (t, first_cam) else: avail_times = sorted(action_data.keys()) def _make_key(t): return t T = len(avail_times) # Find a valid entry to get B and device. first_entry = action_data[sample_key] first_val = first_entry[src_input] B = first_val.shape[0] device = first_val.device tensor_kwargs = dict(device=device, dtype=torch.bfloat16) if pred_frames is None: pred_frames = T if isinstance(cond_frames, int): cond_frames = torch.ones(B, device=device, dtype=torch.int32) * cond_frames if isinstance(pred_frames, int): pred_frames = torch.ones(B, device=device, dtype=torch.int32) * pred_frames # Stack across timesteps: each entry is (B, D), result is (B, T, D). # NOTE(bvh): occasionally `action_data[_make_key(t)]` was None for some t in # heterogeneous DROID batches (sample lacks first_cam at that t), causing # "'NoneType' object is not subscriptable". # Should hopefully be fixed once webbed DROID drops such samples # (TODO check with vitor/anydata though) D = next(v[src_input].shape[-1] for v in action_data.values() if v is not None and v.get(src_input) is not None) def _safe_get(t, key): entry = action_data.get(_make_key(t)) if entry is None or entry[key] is None: return torch.zeros(B, D, **tensor_kwargs) return entry[key] all_input = torch.stack([_safe_get(t, src_input) for t in avail_times], dim=1) all_output = torch.stack([_safe_get(t, src_output) for t in avail_times], dim=1) # Pad to target_tokens if the data has fewer timesteps than expected # (e.g. HumEv with 17 frames mixed with 41-frame datasets) if target_tokens is not None and T < target_tokens: pad_size = target_tokens - T all_input = torch.nn.functional.pad(all_input, (0, 0, 0, pad_size)) all_output = torch.nn.functional.pad(all_output, (0, 0, 0, pad_size)) T = target_tokens # Build masks: (B, T, 1) input_mask = torch.zeros(B, T, 1, **tensor_kwargs) output_mask = torch.zeros(B, T, 1, **tensor_kwargs) for b in range(B): input_mask[b, 0:cond_frames[b]] = 1.0 if pred_frames[b] > cond_frames[b]: output_mask[b, cond_frames[b]:pred_frames[b]] = 1.0 # Mix input (conditioning) and output (generation) into one tensor. a4d_mixed = all_input * input_mask + all_output * output_mask a4d_raw[entry] = a4d_mixed.detach().clone().to(**tensor_kwargs) # NOTE(yams_any4d): unmasked GT copy for auxiliary supervision (V4Head). The mixed # tensor above is ALL-ZERO for pure-forecast tasks (both masks 0), so aux heads # must never read a4d_raw[entry] as ground truth. a4d_raw[f'{entry}_gt_full'] = all_output.detach().clone().to(**tensor_kwargs) a4d_latent[f'{entry}_input_mask'] = input_mask a4d_latent[f'{entry}_output_mask'] = output_mask return (a4d_raw, a4d_latent) def prepare_traj_anydata(a4d_raw, a4d_latent, trajectory, cond_frames, pred_frames=None, entry='traj', target_tokens=None): '''Prepare trajectory entry for Any4D, analogous to prepare_action_anydata. Pads shorter clips to target_tokens and builds input/output masks. :param trajectory: (B, T, D) tensor of trajectory values (e.g. D=2 for driving). :param cond_frames: int, number of conditioned frames. :param pred_frames: int or None, number of predicted frames (0 = input-only). :param entry: str, a4d entry name. :param target_tokens: int or None, expected token count from adaln config. :return: (a4d_raw, a4d_latent). ''' B, T, D = trajectory.shape device = trajectory.device tensor_kwargs = dict(device=device, dtype=torch.bfloat16) if pred_frames is None: pred_frames = 0 # Pad to target_tokens if clip is shorter than model expects. T_avail = T if target_tokens is not None and T < target_tokens: pad_size = target_tokens - T trajectory = torch.nn.functional.pad(trajectory, (0, 0, 0, pad_size)) T = target_tokens # Clamp cond/pred to available frames (padded frames always get mask=0). cond_frames = min(cond_frames, T_avail) pred_frames = min(pred_frames, T_avail) # Build masks: (B, T, 1) input_mask = torch.zeros(B, T, 1, **tensor_kwargs) output_mask = torch.zeros(B, T, 1, **tensor_kwargs) input_mask[:, 0:cond_frames] = 1.0 if pred_frames > cond_frames: output_mask[:, cond_frames:pred_frames] = 1.0 a4d_raw[entry] = trajectory.detach().clone().to(**tensor_kwargs) a4d_latent[f'{entry}_input_mask'] = input_mask a4d_latent[f'{entry}_output_mask'] = output_mask return (a4d_raw, a4d_latent)