From 7b68092a0b94366142eece42507236e856fee7e7 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Wed, 16 Sep 2026 17:09:18 -0400 Subject: [PATCH 01/79] keypoints: the 144-D wrist-first hand-keypoint action becomes the human default Human data trains on hand keypoints instead of a palm-origin cartesian pose. The action is 144-D wrist-first, per hand [wrist xyz (3) | wrist rot6d (6) | 21 MANO keypoints (63)] x {left, right}, built by the keypoints_{wrist,head}frame_6d modes: the existing ypr keypoint builders produce the 138-D layout and KeypointsYPRToRot6D re-encodes just the wrist rotation, leaving the keypoint blocks untouched. Both the action and the observations.state.keypoints proprio are converted, so neither carries a wrapping euler angle. MANO keypoint 0 IS the wrist, so per hand the wrist position rides at both ch 0-2 and ch 9-11; that is kept deliberately (21 keypoints is the MANO interface every consumer reshapes to) and written down in the bimanual_keypoint_layout docstring. include_ee_pose additionally builds the cartesian head-frame ee_pose next to the keypoint action (keymap + transform list), so a keypoint pi0.5 run discretizes the same State: prompt block as the cartesian runs and differs from them only in the action space. The cartesian builder runs with delete_target_world=False to keep obs_head_pose alive for the keypoint builder, and its actions_cartesian by-product is dropped. 144 does not fit openpi's action vector, whose PyTorch port hard-codes 32-wide action projections regardless of Pi0Config.action_dim. PI now reads action_dim and, when it is not 32, swaps in action_dim-wide action_in_proj / action_out_proj AFTER the strict base-weight load: only those two layers start fresh, everything else comes from the checkpoint, and sample_actions already draws noise at config.action_dim. pad_to_width zero-pads narrower embodiments (eva's 20-D) into the same vector and raises rather than truncating a wider one. _reduce_loss then restricts the mean to the slots the converter actually packs, so eva does not average its loss over 112 zero-padded dims and contribute a fraction of the gradient signal it had at width 32. HumanBimanualKeypoints is identity-in, slice-out. HPT human configs move to a 144-D state_keypoints stem. hpt_bc_flow_human holds the body -- there is no vendor notion at the embodiment level, so the canonical config is the embodiment one -- and hpt_bc_flow_{mecka,aria,scale} are two-line aliases of it. HPTModel.add_stem drops null stem entries so a child config can delete a stem its base declares, which hydra merges cannot express, and hpt_cotrain_flow_seperate_head uses it to drop state_ee_pose. With the keypoints_*_6d modes in place, the guard from the previous commit widens to the keypoints ypr and quat modes, so every remaining discontinuous rotation target needs an explicit allow_legacy_rotation=True (the aria->MANO and mecka keypoint viz tutorials, which read pre-6D data). The cartesian configs that were still pinned to the ypr widths move with it: hpt_bc_flow_human_cartesian and the cotrain base to human 20/18, the shared-head cotrains to a 20-wide shared action. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/algo/hpt.py | 3 + egomimic/algo/pi.py | 74 ++++- egomimic/hydra_configs/data/aria.yaml | 9 +- .../hydra_configs/data/aria_keypoints.yaml | 2 +- .../data/aria_keypoints_wrist.yaml | 2 +- egomimic/hydra_configs/data/aria_qwen.yaml | 12 +- .../hydra_configs/data/cotrain_pi_base.yaml | 22 +- egomimic/hydra_configs/data/human.yaml | 18 +- egomimic/hydra_configs/data/mecka.yaml | 17 +- egomimic/hydra_configs/data/scale.yaml | 11 +- .../hydra_configs/model/hpt_bc_flow_aria.yaml | 128 +------- .../model/hpt_bc_flow_human.yaml | 34 ++- .../model/hpt_bc_flow_human_cartesian.yaml | 134 +++++++++ .../model/hpt_bc_flow_mecka.yaml | 128 +------- .../model/hpt_bc_flow_scale.yaml | 128 +------- .../model/hpt_bc_keypoints_base.yaml | 6 +- .../model/hpt_cotrain_enc_dec_base.yaml | 4 +- .../model/hpt_cotrain_flow_seperate_head.yaml | 39 ++- .../model/hpt_cotrain_flow_shared_head.yaml | 3 +- .../hpt_cotrain_mecka_flow_shared_head.yaml | 3 +- .../hpt_cotrain_scale_flow_shared_head.yaml | 3 +- egomimic/hydra_configs/model/pi0.5_base.yaml | 5 + .../hydra_configs/model/pi0.5_bc_aria.yaml | 10 +- .../hydra_configs/model/pi0.5_bc_mecka.yaml | 13 +- .../hydra_configs/model/pi0.5_bc_scale.yaml | 9 +- .../model/pi0.5_cotrain_eva_aria.yaml | 14 +- egomimic/rldb/embodiment/embodiment.py | 4 + egomimic/rldb/embodiment/human.py | 173 ++++++++--- egomimic/rldb/zarr/action_chunk_transforms.py | 69 +++++ .../scripts/tutorials/aria_to_mano_convert.py | 4 +- .../scripts/tutorials/mecka_keypoints_viz.py | 4 +- .../scripts/tutorials/zarr_data_viz.ipynb | 4 +- egomimic/utils/action_utils.py | 69 ++++- egomimic/utils/pose_utils.py | 62 +++- egomimic/utils/viz_utils.py | 6 +- tests/fixtures/synthetic_episodes.py | 10 + tests/fixtures/train_harness.py | 23 +- tests/unit/test_infra_port_wiring.py | 258 ++++++++++++++++ tests/unit/test_keypoints_6d.py | 276 ++++++++++++++++++ tests/unit/test_model_dims.py | 24 +- tests/unit/test_pi05_norm_rot6d.py | 56 ++++ tests/unit/test_train_step.py | 37 ++- 42 files changed, 1397 insertions(+), 513 deletions(-) create mode 100644 egomimic/hydra_configs/model/hpt_bc_flow_human_cartesian.yaml create mode 100644 tests/unit/test_infra_port_wiring.py create mode 100644 tests/unit/test_keypoints_6d.py diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index ad567b0fb6..8184f5ac74 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -145,6 +145,9 @@ def init_domain_stem(self, domain_name, stem_spec): A specification containing configurations for each modality's stem. """ + # A ``null`` entry lets a child config drop a stem its base declares + # (hydra merges cannot delete keys). + stem_spec = {k: v for k, v in stem_spec.items() if v is not None} self.stem_spec[domain_name] = stem_spec self.modalities[domain_name] = list(stem_spec.keys()) diff --git a/egomimic/algo/pi.py b/egomimic/algo/pi.py index d055eab546..7e26e352e5 100644 --- a/egomimic/algo/pi.py +++ b/egomimic/algo/pi.py @@ -25,7 +25,7 @@ gather_pi_images, ) from egomimic.rldb.embodiment.embodiment import get_embodiment, get_embodiment_id -from egomimic.utils.action_utils import ConverterRegistry +from egomimic.utils.action_utils import ConverterRegistry, pad_to_width logger = logging.getLogger(__name__) # Ensure logger propagates to root logger and has appropriate level @@ -142,6 +142,7 @@ def __init__( self.action_converters = action_converters self.action_registry = ConverterRegistry() + self._packed_widths: dict[tuple, int] = {} arcfg = self.action_converters default_ac_key = getattr(arcfg, "ac_key", "actions_cartesian") @@ -195,9 +196,45 @@ def __init__( "No pytorch_weight_path: no base weights loaded. The weights must " "come from a checkpoint; otherwise this trains from scratch." ) + # openpi's PyTorch port hard-codes 32-wide action projections + # regardless of ``Pi0Config.action_dim``. Wider action spaces (the + # 144-D hand-keypoint action, or a cotrain that pads eva's 20-D into + # the same vector) need ``action_dim``-wide ones, swapped in AFTER the + # (strict) base-weight load. Only these two layers start fresh; + # ``sample_actions`` already draws its noise at ``config.action_dim``. + self.action_dim = int(self.config.model.action_dim) + if self.action_dim != 32: + self._resize_action_projections(self.action_dim) self.nets = nn.ModuleDict() self.nets["policy"] = self.model + def _resize_action_projections(self, action_dim: int) -> None: + target = ( + self.model.module + if isinstance(self.model, torch.nn.parallel.DistributedDataParallel) + else self.model + ) + old_in, old_out = target.action_in_proj, target.action_out_proj + width = old_in.out_features + new_in = nn.Linear(action_dim, width).to( + dtype=old_in.weight.dtype, device=old_in.weight.device + ) + new_out = nn.Linear(width, action_dim).to( + dtype=old_out.weight.dtype, device=old_out.weight.device + ) + target.action_in_proj = new_in + target.action_out_proj = new_out + logger.warning( + "action_dim=%d != 32: re-initialized action_in_proj (%d->%d) and " + "action_out_proj (%d->%d); the pretrained 32-wide projections are " + "discarded, everything else comes from the base checkpoint.", + action_dim, + action_dim, + width, + width, + action_dim, + ) + def _control_mode_for(self, emb_name: str | None) -> str: if self.control_mode and emb_name is not None: for key, val in self.control_mode.items(): @@ -479,19 +516,40 @@ def forward_eval(self, batch): return unnorm_preds def _reduce_loss(self, losses, action, embodiment_id, ac_key) -> torch.Tensor: - """Mean flow-matching loss for the embodiment.""" + """Mean flow-matching loss over the embodiment's packed action width. + + openpi returns the per-element loss ``(B, H, action_dim)``. With + ``action_dim`` widened for the 144-D keypoint action, a narrower + embodiment (eva's 32-slot cartesian layout) would otherwise average + its loss over 112 zero-padded "predict the noise" dims and contribute + a fraction of the gradient signal it had at width 32; restrict the + mean to the slots the converter actually packs. + """ if isinstance(losses, (list, tuple)): losses = torch.stack(losses) elif not isinstance(losses, torch.Tensor): losses = torch.tensor(losses, device=action.device, dtype=torch.float32) + if losses.ndim == 3: + losses = losses[..., : self._packed_width(action, embodiment_id, ac_key)] return losses.mean() + def _packed_width(self, action, embodiment_id, ac_key) -> int: + """Width the embodiment's converter packs into, cached per pairing: it + is a property of the converter, not of the batch.""" + key = (embodiment_id, ac_key) + if key not in self._packed_widths: + converter = self.action_registry.get(embodiment_id, ac_key) + self._packed_widths[key] = int( + converter.to32_norm_6d(action[:1, :1]).shape[-1] + ) + return self._packed_widths[key] + def _postprocess_sampled_actions(self, pred_actions, _batch, embodiment_id, ac_key): """Raw ``sample_actions`` output -> unnormalized native action dict. - Unpack the normalized native action (xyz+6D(+gripper)) from the - model's action vector, then unnormalize via the standard pipeline - (stats were computed on the native layout).""" + Unpack the normalized native action (xyz+6D(+gripper) or the 144-D + keypoint vector) from the model's action vector, then unnormalize via + the standard pipeline (stats were computed on the native layout).""" ref = _batch[ac_key] _, T, D = ref.shape converter = self.action_registry.get(embodiment_id, ac_key) @@ -549,7 +607,7 @@ def log_info(self, info): return log def _robomimic_to_pi_data(self, batch, proprio_keys, lang_keys, ac_key, embodiment): - """Dataset batch -> (openpi observation, 32-D action chunk).""" + """Dataset batch -> (openpi observation, ``action_dim``-wide action chunk).""" if ac_key not in batch: raise KeyError(f"Missing action key '{ac_key}' in batch") @@ -561,8 +619,8 @@ def _robomimic_to_pi_data(self, batch, proprio_keys, lang_keys, ac_key, embodime converter = self.action_registry.get(emb_id, ac_key) # The action is already normalized and in its native layout (the # ypr->6D conversion happened in the data transforms). Pack it into - # the canonical 32-slot block layout. - action32 = converter.to32_norm_6d(action) + # the canonical block layout and zero-pad up to the model width. + action32 = pad_to_width(converter.to32_norm_6d(action), self.action_dim) # OpenPI expects a fixed camera tuple under its own names. Human data # only has the front camera; the missing wrist slots get a copy of it diff --git a/egomimic/hydra_configs/data/aria.yaml b/egomimic/hydra_configs/data/aria.yaml index ef1550dd1f..dbd5f47d17 100644 --- a/egomimic/hydra_configs/data/aria.yaml +++ b/egomimic/hydra_configs/data/aria.yaml @@ -1,3 +1,6 @@ +# One pinned aria episode (stored as human_bimanual) on the 144-D wrist-frame +# hand-keypoint action; see data/human.yaml for the layout and the cartesian +# override. _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: @@ -8,10 +11,12 @@ train_datasets: folder_path: ${paths.dataset_dir} key_map: _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian + keymap_mode: keypoints + include_ee_pose: true transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d + mode: keypoints_wristframe_6d + include_ee_pose: true # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 # State: prompt bins align with the robot 20-D layout on every config. pad_proprio_gripper: true diff --git a/egomimic/hydra_configs/data/aria_keypoints.yaml b/egomimic/hydra_configs/data/aria_keypoints.yaml index 59d0a3ab62..7ed22d6eb7 100644 --- a/egomimic/hydra_configs/data/aria_keypoints.yaml +++ b/egomimic/hydra_configs/data/aria_keypoints.yaml @@ -11,7 +11,7 @@ train_datasets: keymap_mode: keypoints transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: keypoints_headframe_ypr + mode: keypoints_headframe_6d stride: 3 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml b/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml index 4c22356e61..1a4edb25d3 100644 --- a/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml +++ b/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml @@ -11,7 +11,7 @@ train_datasets: keymap_mode: keypoints transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: keypoints_wristframe_ypr + mode: keypoints_wristframe_6d stride: 3 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/aria_qwen.yaml b/egomimic/hydra_configs/data/aria_qwen.yaml index 9612b56990..cd976a643d 100644 --- a/egomimic/hydra_configs/data/aria_qwen.yaml +++ b/egomimic/hydra_configs/data/aria_qwen.yaml @@ -1,7 +1,7 @@ # Aria data with raw text annotations in the batch, so the HPT Qwen stem can -# sample one annotation per item. Mirror of `aria.yaml` plus -# `annotation_key: annotations` in the keymap (which adds a per-frame raw-text -# annotation key — see Embodiment.get_keymap). +# sample one annotation per item. Mirror of `aria.yaml` (144-D wrist-frame +# hand-keypoint action) plus `annotation_key: annotations` in the keymap +# (which adds a per-frame raw-text annotation key; see Embodiment.get_keymap). _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: @@ -12,11 +12,13 @@ train_datasets: folder_path: ${paths.dataset_dir} key_map: _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian + keymap_mode: keypoints + include_ee_pose: true annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d + mode: keypoints_wristframe_6d + include_ee_pose: true # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 # State: prompt bins align with the robot 20-D layout on every config. pad_proprio_gripper: true diff --git a/egomimic/hydra_configs/data/cotrain_pi_base.yaml b/egomimic/hydra_configs/data/cotrain_pi_base.yaml index bfa45b5bba..7fb6122886 100644 --- a/egomimic/hydra_configs/data/cotrain_pi_base.yaml +++ b/egomimic/hydra_configs/data/cotrain_pi_base.yaml @@ -1,3 +1,19 @@ +# Base for every pi0.5 data config. +# +# eva: wrist-relative cartesian actions (`cartesian_wristframe_6d`: each arm's +# chunk in its current EEF frame, xyz+6D+gripper per arm, 20D) with the +# cam-frame proprio ee_pose in the same 6D layout. Cam-frame runs override +# the mode to `cartesian_6d` on the command line. +# human: the 144-D wrist-first hand-keypoint action (`keypoints_wristframe_6d`: +# per hand [wrist xyz | wrist rot6d | 21 MANO keypoints] in that wrist's +# frame) with the head-frame `observations.state.keypoints` proprio, plus +# (`include_ee_pose`) the cartesian head-frame `observations.state.ee_pose` +# that the pi0.5 prompt discretizes into its `State:` block, so a keypoint +# run shares its prompt state with the cartesian runs. The cartesian human +# modes stay available by overriding +# `resolver.key_map.keymap_mode=cartesian` + `resolver.transform_list.mode=cartesian_6d` +# (and dropping include_ee_pose) together with a cartesian PI model config. +# Norm stats must be computed on these layouts. _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: @@ -23,11 +39,13 @@ train_datasets: folder_path: ${paths.dataset_dir} key_map: _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian + keymap_mode: keypoints + include_ee_pose: true annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d + mode: keypoints_wristframe_6d + include_ee_pose: true # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 # State: prompt bins align with the robot 20-D layout on every config. pad_proprio_gripper: true diff --git a/egomimic/hydra_configs/data/human.yaml b/egomimic/hydra_configs/data/human.yaml index 59a472a346..9c4a0b27eb 100644 --- a/egomimic/hydra_configs/data/human.yaml +++ b/egomimic/hydra_configs/data/human.yaml @@ -1,8 +1,14 @@ _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper -# Canonical human (ex-aria) cartesian data config, keyed by the collapsed -# embodiment human_bimanual (episodes reprocessed to embodiment="human_bimanual" -# + intrinsics in zarr.json). +# Canonical human cartesian data config, keyed by the collapsed embodiment +# human_bimanual (episodes reprocessed to embodiment="human_bimanual" + +# intrinsics in zarr.json). +# Action: the 144-D wrist-frame hand-keypoint action (`keypoints_wristframe_6d`) +# with the head-frame `observations.state.keypoints` proprio and, via +# include_ee_pose, the cartesian head-frame `observations.state.ee_pose` the +# pi0.5 prompt discretizes (HPT ignores it: no stem consumes it). Override +# `resolver.key_map.keymap_mode=cartesian resolver.transform_list.mode=cartesian_6d` +# (and include_ee_pose=false on both) for the cartesian action. # NOTE: the filter selects on the SQL value 'human_bimanual'. This ASSUMES the SQL # embodiment/robot_name column has been migrated from the old vendor values # ('aria_bimanual'). Until that migration lands, this dataset matches zero rows. @@ -14,10 +20,12 @@ train_datasets: folder_path: ${paths.dataset_dir} key_map: _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian + keymap_mode: keypoints + include_ee_pose: true transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d + mode: keypoints_wristframe_6d + include_ee_pose: true # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 # State: prompt bins align with the robot 20-D layout on every config. pad_proprio_gripper: true diff --git a/egomimic/hydra_configs/data/mecka.yaml b/egomimic/hydra_configs/data/mecka.yaml index 3fe7f6c6ca..b904d4dde0 100644 --- a/egomimic/hydra_configs/data/mecka.yaml +++ b/egomimic/hydra_configs/data/mecka.yaml @@ -1,3 +1,12 @@ +# One pinned mecka episode (stored as human_bimanual): the smoke / overfit +# config for HPT and Pi on the 144-D wrist-frame hand-keypoint action. +# fix_left_wrist_convention: this episode predates the mecka_to_zarr rot_left +# fix (LEFT wrist frame double-mirrored); Rz(180) corrects the raw left +# pose before any frame math. With include_ee_pose it corrects both the +# wrist_pose that the 144-D keypoint action is built from and the cartesian +# ee_pose proprio. +# pad_proprio_gripper: 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) +# so the pi0.5 State: prompt bins align with the robot 20-D layout. _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: @@ -8,13 +17,13 @@ train_datasets: folder_path: ${paths.dataset_dir} key_map: _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian + keymap_mode: keypoints + include_ee_pose: true transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d + mode: keypoints_wristframe_6d + include_ee_pose: true fix_left_wrist_convention: true - # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 - # State: prompt bins align with the robot 20-D layout on every config. pad_proprio_gripper: true stride: 1 filters: diff --git a/egomimic/hydra_configs/data/scale.yaml b/egomimic/hydra_configs/data/scale.yaml index 81a486e8cc..5b97388514 100644 --- a/egomimic/hydra_configs/data/scale.yaml +++ b/egomimic/hydra_configs/data/scale.yaml @@ -1,3 +1,8 @@ +# One pinned scale episode (stored as human_bimanual) on the 144-D wrist-frame +# hand-keypoint action; see data/human.yaml for the layout and the cartesian +# override. fix_left_wrist_convention is on: scale runs its own converter, +# but it writes the same pre-fix left-wrist convention as mecka (measured on +# all 20 production episodes). _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper train_datasets: @@ -8,10 +13,12 @@ train_datasets: folder_path: ${paths.dataset_dir} key_map: _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian + keymap_mode: keypoints + include_ee_pose: true transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d + mode: keypoints_wristframe_6d + include_ee_pose: true # Scale exports were written by the pre-fix hand-pose converter too: # measured on all 20 production episodes, the left wrist matches # rot_left to <= 2.06 deg with the right-hand control clean. diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_aria.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_aria.yaml index eb345dd1e9..3389976760 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_aria.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_aria.yaml @@ -1,122 +1,6 @@ -_target_: egomimic.pl_utils.pl_model.ModelWrapper -robomimic_model: - _target_: egomimic.algo.hpt.HPT - camera_transforms: {} - - dims: - # cartesian_6d + pad_proprio_gripper: action [xyz | rot6d] x 2 = 18, - # proprio ee_pose grip-padded to the robot 20-D layout. - human_bimanual: {proprio: 20, action: 18} - - diffusion: true - 6dof: true - - ac_keys: - human_bimanual: "actions_cartesian" - - trunk: - embed_dim: 256 - num_blocks: 16 - num_heads: 8 - token_postprocessing: "action_token" - observation_horizon: 1 - action_horizon: 64 - no_trunk: false - use_domain_embedding: true - drop_path: 0.1 - weight_init_style: "pytorch" - - multitask: false - pretrained: false - pretrained_checkpoint: "" # TODO - reverse_kl_samples: 8 - - domains: ["human_bimanual"] - shared_obs_keys: ["front_img_1"] - - shared_stem_specs: - front_img_1: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: 256 - output_dim: 256 - widths: [256] - specs: - random_horizon_masking: false - cross_attn: - crossattn_latent: 16 - crossattn_heads: 8 - crossattn_dim_head: 64 - crossattn_modality_dropout: 0.1 - modality_embed_dim: 256 - - stem_specs: - human_bimanual: - state_ee_pose: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} - output_dim: 256 - widths: [256] - specs: - random_horizon_masking: false - cross_attn: - crossattn_latent: 16 - crossattn_heads: 8 - crossattn_dim_head: 64 - crossattn_modality_dropout: 0.1 - modality_embed_dim: 256 - - head_specs: - human_bimanual: - _target_: egomimic.models.fm_policy.FMPolicy - action_horizon: 100 - num_inference_steps: 50 - pooling: null - time_dist: "beta" - infer_ac_dims: - human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} - model: - _target_: egomimic.models.denoising_nets.CrossTransformer - nblocks: 6 - cond_dim: 256 - hidden_dim: 128 - act_dim: ${model.robomimic_model.dims.human_bimanual.action} - act_seq: 100 - n_heads: 4 - dropout: 0.1 - mlp_layers: 4 - mlp_ratio: 4 - - encoder_specs: - front_img_1: - _target_: egomimic.models.hpt_nets.ResNet - output_dim: 256 - - train_image_augs: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.ColorJitter - brightness: 0.1 - contrast: 0.1 - saturation: 0.1 - hue: 0.05 - - _target_: torchvision.transforms.Normalize - mean: [0.485, 0.456, 0.406] - std: [0.229, 0.224, 0.225] - eval_image_augs: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.Normalize - mean: [0.485, 0.456, 0.406] - std: [0.229, 0.224, 0.225] - -optimizer: - _target_: torch.optim.AdamW - _partial_: true - lr: 3e-4 - weight_decay: 0.0001 - -scheduler: - _target_: torch.optim.lr_scheduler.CosineAnnealingLR - _partial_: true - T_max: 1400 - eta_min: 1e-5 +# HPT flow-matching BC on human data (aria, stored as human_bimanual) with +# the 144-D wrist-frame hand-keypoint action (data/aria.yaml); identical to +# hpt_bc_flow_human.yaml apart from this note. +defaults: + - hpt_bc_flow_human + - _self_ diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml index 21c3f7fb01..c8daab5660 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml @@ -1,23 +1,25 @@ +# HPT flow-matching BC on human data with the 144-D wrist-frame hand-keypoint +# action (`keypoints_wristframe_6d`): per hand [wrist xyz | wrist rot6d | 21 +# MANO keypoints] in that wrist's frame. Proprio = the 144-D head-frame +# `observations.state.keypoints` through the `state_keypoints` stem (the +# ee_pose a data config may also emit for the pi0.5 prompt has no stem here +# and is ignored). There is no vendor notion at the embodiment level, so this +# is the canonical human keypoint config; hpt_bc_flow_{mecka,aria,scale} are +# aliases of it, one per data config. For the cartesian action pair a +# cartesian data override with model=hpt_bc_flow_human_cartesian. _target_: egomimic.pl_utils.pl_model.ModelWrapper robomimic_model: _target_: egomimic.algo.hpt.HPT - # camera_transforms is a dead HPT.__init__ param after refactor 80e5523a - # ("move intrinsics/extrinsics into embodiment classes" — CameraTransforms was - # removed; HPT never stored/used it). Per-episode intrinsics now travel in - # batch["intrinsics"] (zarr.json) and viz reads them via the embodiment classes. - # HPT still takes the positional arg, so pass an empty mapping (it is ignored). camera_transforms: {} dims: - # cartesian_6d + pad_proprio_gripper: action [xyz | rot6d] x 2 = 18, - # proprio ee_pose grip-padded to the robot 20-D layout. - human_bimanual: {proprio: 20, action: 18} + human_bimanual: {proprio: 144, action: 144} - diffusion: true + diffusion: true 6dof: true ac_keys: - human_bimanual: "actions_cartesian" + human_bimanual: "actions_keypoints" trunk: embed_dim: 256 @@ -30,12 +32,12 @@ robomimic_model: use_domain_embedding: true drop_path: 0.1 weight_init_style: "pytorch" - + multitask: false pretrained: false - pretrained_checkpoint: "" + pretrained_checkpoint: "" # TODO reverse_kl_samples: 8 - + domains: ["human_bimanual"] shared_obs_keys: ["front_img_1"] @@ -53,10 +55,10 @@ robomimic_model: crossattn_dim_head: 64 crossattn_modality_dropout: 0.1 modality_embed_dim: 256 - + stem_specs: human_bimanual: - state_ee_pose: + state_keypoints: _target_: egomimic.models.hpt_nets.MLPPolicyStem input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} output_dim: 256 @@ -71,7 +73,7 @@ robomimic_model: modality_embed_dim: 256 head_specs: - human_bimanual: + human_bimanual: _target_: egomimic.models.fm_policy.FMPolicy action_horizon: 100 num_inference_steps: 50 diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_human_cartesian.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_human_cartesian.yaml new file mode 100644 index 0000000000..bdb2775b90 --- /dev/null +++ b/egomimic/hydra_configs/model/hpt_bc_flow_human_cartesian.yaml @@ -0,0 +1,134 @@ +# HPT flow-matching BC on human data with the CARTESIAN head-frame action +# (18-D xyz+rot6d per arm pair) and the 20-D gripper-padded ee_pose proprio. +# The repo default for human data is the 144-D hand-keypoint action +# (hpt_bc_flow_human); use this with a cartesian data override, e.g. +# data=mecka data.train_datasets.human_bimanual.resolver.key_map.keymap_mode=cartesian \ +# data.train_datasets.human_bimanual.resolver.key_map.include_ee_pose=false \ +# data.train_datasets.human_bimanual.resolver.transform_list.mode=cartesian_6d \ +# data.train_datasets.human_bimanual.resolver.transform_list.include_ee_pose=false +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.hpt.HPT + # camera_transforms is a dead HPT.__init__ param after refactor 80e5523a + # ("move intrinsics/extrinsics into embodiment classes" — CameraTransforms was + # removed; HPT never stored/used it). Per-episode intrinsics now travel in + # batch["intrinsics"] (zarr.json) and viz reads them via the embodiment classes. + # HPT still takes the positional arg, so pass an empty mapping (it is ignored). + camera_transforms: {} + + # Per-embodiment widths, declared once. Stems and heads below interpolate + # these; override e.g. model.robomimic_model.dims..proprio=16. + dims: + human_bimanual: {proprio: 20, action: 18} + + diffusion: true + 6dof: true + + ac_keys: + human_bimanual: "actions_cartesian" + + trunk: + embed_dim: 256 + num_blocks: 16 + num_heads: 8 + token_postprocessing: "action_token" + observation_horizon: 1 + action_horizon: 64 + no_trunk: false + use_domain_embedding: true + drop_path: 0.1 + weight_init_style: "pytorch" + + multitask: false + pretrained: false + pretrained_checkpoint: "" + + domains: ["human_bimanual"] + shared_obs_keys: ["front_img_1"] + + shared_stem_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: 256 + output_dim: 256 + widths: [256] + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 16 + crossattn_heads: 8 + crossattn_dim_head: 64 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 256 + + stem_specs: + human_bimanual: + state_ee_pose: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} + output_dim: 256 + widths: [256] + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 16 + crossattn_heads: 8 + crossattn_dim_head: 64 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 256 + + head_specs: + human_bimanual: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 50 + pooling: null + time_dist: "beta" + infer_ac_dims: + human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} + model: + _target_: egomimic.models.denoising_nets.CrossTransformer + nblocks: 6 + cond_dim: 256 + hidden_dim: 128 + act_dim: ${model.robomimic_model.dims.human_bimanual.action} + act_seq: 100 + n_heads: 4 + dropout: 0.1 + mlp_layers: 4 + mlp_ratio: 4 + + encoder_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.ResNet + output_dim: 256 + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.ColorJitter + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.05 + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + eval_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 3e-4 + weight_decay: 0.0001 + +scheduler: + _target_: torch.optim.lr_scheduler.CosineAnnealingLR + _partial_: true + T_max: 1400 + eta_min: 1e-5 diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_mecka.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_mecka.yaml index eb345dd1e9..30c9af5165 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_mecka.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_mecka.yaml @@ -1,122 +1,6 @@ -_target_: egomimic.pl_utils.pl_model.ModelWrapper -robomimic_model: - _target_: egomimic.algo.hpt.HPT - camera_transforms: {} - - dims: - # cartesian_6d + pad_proprio_gripper: action [xyz | rot6d] x 2 = 18, - # proprio ee_pose grip-padded to the robot 20-D layout. - human_bimanual: {proprio: 20, action: 18} - - diffusion: true - 6dof: true - - ac_keys: - human_bimanual: "actions_cartesian" - - trunk: - embed_dim: 256 - num_blocks: 16 - num_heads: 8 - token_postprocessing: "action_token" - observation_horizon: 1 - action_horizon: 64 - no_trunk: false - use_domain_embedding: true - drop_path: 0.1 - weight_init_style: "pytorch" - - multitask: false - pretrained: false - pretrained_checkpoint: "" # TODO - reverse_kl_samples: 8 - - domains: ["human_bimanual"] - shared_obs_keys: ["front_img_1"] - - shared_stem_specs: - front_img_1: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: 256 - output_dim: 256 - widths: [256] - specs: - random_horizon_masking: false - cross_attn: - crossattn_latent: 16 - crossattn_heads: 8 - crossattn_dim_head: 64 - crossattn_modality_dropout: 0.1 - modality_embed_dim: 256 - - stem_specs: - human_bimanual: - state_ee_pose: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} - output_dim: 256 - widths: [256] - specs: - random_horizon_masking: false - cross_attn: - crossattn_latent: 16 - crossattn_heads: 8 - crossattn_dim_head: 64 - crossattn_modality_dropout: 0.1 - modality_embed_dim: 256 - - head_specs: - human_bimanual: - _target_: egomimic.models.fm_policy.FMPolicy - action_horizon: 100 - num_inference_steps: 50 - pooling: null - time_dist: "beta" - infer_ac_dims: - human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} - model: - _target_: egomimic.models.denoising_nets.CrossTransformer - nblocks: 6 - cond_dim: 256 - hidden_dim: 128 - act_dim: ${model.robomimic_model.dims.human_bimanual.action} - act_seq: 100 - n_heads: 4 - dropout: 0.1 - mlp_layers: 4 - mlp_ratio: 4 - - encoder_specs: - front_img_1: - _target_: egomimic.models.hpt_nets.ResNet - output_dim: 256 - - train_image_augs: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.ColorJitter - brightness: 0.1 - contrast: 0.1 - saturation: 0.1 - hue: 0.05 - - _target_: torchvision.transforms.Normalize - mean: [0.485, 0.456, 0.406] - std: [0.229, 0.224, 0.225] - eval_image_augs: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.Normalize - mean: [0.485, 0.456, 0.406] - std: [0.229, 0.224, 0.225] - -optimizer: - _target_: torch.optim.AdamW - _partial_: true - lr: 3e-4 - weight_decay: 0.0001 - -scheduler: - _target_: torch.optim.lr_scheduler.CosineAnnealingLR - _partial_: true - T_max: 1400 - eta_min: 1e-5 +# HPT flow-matching BC on human data (mecka, stored as human_bimanual) with +# the 144-D wrist-frame hand-keypoint action (data/mecka.yaml); identical to +# hpt_bc_flow_human.yaml apart from this note. +defaults: + - hpt_bc_flow_human + - _self_ diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_scale.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_scale.yaml index 6ae2b6d1d4..68235751ab 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_scale.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_scale.yaml @@ -1,122 +1,6 @@ -_target_: egomimic.pl_utils.pl_model.ModelWrapper -robomimic_model: - _target_: egomimic.algo.hpt.HPT - camera_transforms: {} - - dims: - # cartesian_6d + pad_proprio_gripper: action [xyz | rot6d] x 2 = 18, - # proprio ee_pose grip-padded to the robot 20-D layout. - human_bimanual: {proprio: 20, action: 18} - - diffusion: true - 6dof: true - - ac_keys: - human_bimanual: "actions_cartesian" - - trunk: - embed_dim: 256 - num_blocks: 64 - num_heads: 8 - token_postprocessing: "action_token" - observation_horizon: 1 - action_horizon: 64 - no_trunk: false - use_domain_embedding: true - drop_path: 0.1 - weight_init_style: "pytorch" - - multitask: false - pretrained: false - pretrained_checkpoint: "" # TODO - reverse_kl_samples: 8 - - domains: ["human_bimanual"] - shared_obs_keys: ["front_img_1"] - - shared_stem_specs: - front_img_1: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: 256 - output_dim: 256 - widths: [256] - specs: - random_horizon_masking: false - cross_attn: - crossattn_latent: 16 - crossattn_heads: 8 - crossattn_dim_head: 64 - crossattn_modality_dropout: 0.1 - modality_embed_dim: 256 - - stem_specs: - human_bimanual: - state_ee_pose: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} - output_dim: 256 - widths: [256] - specs: - random_horizon_masking: false - cross_attn: - crossattn_latent: 16 - crossattn_heads: 8 - crossattn_dim_head: 64 - crossattn_modality_dropout: 0.1 - modality_embed_dim: 256 - - head_specs: - human_bimanual: - _target_: egomimic.models.fm_policy.FMPolicy - action_horizon: 100 - num_inference_steps: 50 - pooling: null - time_dist: "beta" - infer_ac_dims: - human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} - model: - _target_: egomimic.models.denoising_nets.CrossTransformer - nblocks: 6 - cond_dim: 256 - hidden_dim: 128 - act_dim: ${model.robomimic_model.dims.human_bimanual.action} - act_seq: 100 - n_heads: 4 - dropout: 0.1 - mlp_layers: 4 - mlp_ratio: 4 - - encoder_specs: - front_img_1: - _target_: egomimic.models.hpt_nets.ResNet - output_dim: 256 - - train_image_augs: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.ColorJitter - brightness: 0.1 - contrast: 0.1 - saturation: 0.1 - hue: 0.05 - - _target_: torchvision.transforms.Normalize - mean: [0.485, 0.456, 0.406] - std: [0.229, 0.224, 0.225] - eval_image_augs: - _target_: torchvision.transforms.Compose - transforms: - - _target_: torchvision.transforms.Normalize - mean: [0.485, 0.456, 0.406] - std: [0.229, 0.224, 0.225] - -optimizer: - _target_: torch.optim.AdamW - _partial_: true - lr: 3e-4 - weight_decay: 0.0001 - -scheduler: - _target_: torch.optim.lr_scheduler.CosineAnnealingLR - _partial_: true - T_max: 1400 - eta_min: 1e-5 +# HPT flow-matching BC on human data (scale, stored as human_bimanual) with +# the 144-D wrist-frame hand-keypoint action (data/scale.yaml); identical to +# hpt_bc_flow_human.yaml apart from this note. +defaults: + - hpt_bc_flow_human + - _self_ diff --git a/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml b/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml index 51dfa890bc..b4d1465fcc 100644 --- a/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml @@ -3,8 +3,12 @@ robomimic_model: _target_: egomimic.algo.hpt.HPT camera_transforms: {} + # 144 = per hand [wrist xyz (3) | wrist rot6d (6) | 21 MANO keypoints (63)] + # x 2, the `keypoints_*_6d` transform modes (the 138-D ypr layout of the + # `keypoints_*_ypr` modes is no longer the default). dims: - human_bimanual: {proprio: 138, action: 138} + human_bimanual: {proprio: 144, action: 144} + ac_keys: human_bimanual: "actions_keypoints" diff --git a/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml b/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml index e515212854..1c10c3480e 100644 --- a/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml +++ b/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml @@ -4,8 +4,10 @@ robomimic_model: camera_transforms: {} dims: + # eva proprio is joint positions (7/arm), not a pose. eva_bimanual: {proprio: 14, action: 14} - human_bimanual: {proprio: 12, action: 12} + # human cartesian_6d: 18-D action, 20-D gripper-padded ee_pose proprio. + human_bimanual: {proprio: 20, action: 18} ac_keys: human_bimanual: "actions_cartesian" diff --git a/egomimic/hydra_configs/model/hpt_cotrain_flow_seperate_head.yaml b/egomimic/hydra_configs/model/hpt_cotrain_flow_seperate_head.yaml index 499472930d..a3d73c5697 100644 --- a/egomimic/hydra_configs/model/hpt_cotrain_flow_seperate_head.yaml +++ b/egomimic/hydra_configs/model/hpt_cotrain_flow_seperate_head.yaml @@ -1,12 +1,41 @@ defaults: - hpt_cotrain_enc_dec_base +# Separate-head cotrain: eva on the base's 14-D joint-position action +# (`actions_joints` ac_key), human on the 144-D wrist-frame hand-keypoint +# action through a `state_keypoints` stem; the base's `state_ee_pose` stem is +# dropped (null) so a keypoint data config's extra ee_pose proprio is ignored. +# No shipped HPT cotrain data config emits this pairing yet; build one from a +# per-vendor keypoint config plus data/eva.yaml (whose eva action is the 20-D +# wrist-frame 6D layout, so override dims.eva_bimanual to match). The +# shared-head cotrain configs (hpt_cotrain_*_shared_head) need one action +# space and stay cartesian on both sides. robomimic_model: domains: ["eva_bimanual", "human_bimanual"] - diffusion: true + diffusion: true + dims: + human_bimanual: {proprio: 144, action: 144} + ac_keys: + human_bimanual: "actions_keypoints" + stem_specs: + human_bimanual: + state_ee_pose: null + state_keypoints: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} + output_dim: 256 + widths: [256] + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 16 + crossattn_heads: 8 + crossattn_dim_head: 64 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 256 head_specs: shared: null - human_bimanual: + human_bimanual: _target_: egomimic.models.fm_policy.FMPolicy action_horizon: 100 num_inference_steps: 50 @@ -26,7 +55,7 @@ robomimic_model: mlp_layers: 4 mlp_ratio: 4 - eva_bimanual: + eva_bimanual: _target_: egomimic.models.fm_policy.FMPolicy action_horizon: 100 num_inference_steps: 50 @@ -43,9 +72,9 @@ robomimic_model: act_seq: 100 n_heads: 4 dropout: 0.1 - mlp_layers: 4 + mlp_layers: 4 mlp_ratio: 4 - + optimizer: _target_: torch.optim.AdamW _partial_: true diff --git a/egomimic/hydra_configs/model/hpt_cotrain_flow_shared_head.yaml b/egomimic/hydra_configs/model/hpt_cotrain_flow_shared_head.yaml index 56fde2ad4e..fad676946b 100644 --- a/egomimic/hydra_configs/model/hpt_cotrain_flow_shared_head.yaml +++ b/egomimic/hydra_configs/model/hpt_cotrain_flow_shared_head.yaml @@ -13,7 +13,8 @@ robomimic_model: # Shared head: every embodiment's actions are zero-padded to action_width, so # the per-embodiment action widths inherited from the base follow it. dims: - action_width: 14 + # eva cartesian_6d is 20-D; the human 18-D action zero-pads up to it. + action_width: 20 eva_bimanual: action: ${model.robomimic_model.dims.action_width} human_bimanual: diff --git a/egomimic/hydra_configs/model/hpt_cotrain_mecka_flow_shared_head.yaml b/egomimic/hydra_configs/model/hpt_cotrain_mecka_flow_shared_head.yaml index 0598b67e83..b2811a2000 100644 --- a/egomimic/hydra_configs/model/hpt_cotrain_mecka_flow_shared_head.yaml +++ b/egomimic/hydra_configs/model/hpt_cotrain_mecka_flow_shared_head.yaml @@ -32,7 +32,8 @@ robomimic_model: # Shared head: every embodiment's actions are zero-padded to action_width, so # the per-embodiment action widths inherited from the base follow it. dims: - action_width: 14 + # eva cartesian_6d is 20-D; the human 18-D action zero-pads up to it. + action_width: 20 eva_bimanual: action: ${model.robomimic_model.dims.action_width} human_bimanual: diff --git a/egomimic/hydra_configs/model/hpt_cotrain_scale_flow_shared_head.yaml b/egomimic/hydra_configs/model/hpt_cotrain_scale_flow_shared_head.yaml index 42dc191dff..f927f22cd2 100644 --- a/egomimic/hydra_configs/model/hpt_cotrain_scale_flow_shared_head.yaml +++ b/egomimic/hydra_configs/model/hpt_cotrain_scale_flow_shared_head.yaml @@ -29,7 +29,8 @@ robomimic_model: # Shared head: every embodiment's actions are zero-padded to action_width, so # the per-embodiment action widths inherited from the base follow it. dims: - action_width: 14 + # eva cartesian_6d is 20-D; the human 18-D action zero-pads up to it. + action_width: 20 eva_bimanual: action: ${model.robomimic_model.dims.action_width} human_bimanual: diff --git a/egomimic/hydra_configs/model/pi0.5_base.yaml b/egomimic/hydra_configs/model/pi0.5_base.yaml index e94ff5c095..88f2a7947e 100644 --- a/egomimic/hydra_configs/model/pi0.5_base.yaml +++ b/egomimic/hydra_configs/model/pi0.5_base.yaml @@ -7,6 +7,11 @@ robomimic_model: pytorch_weight_path: ${oc.env:EGOVERSE_PI05_WEIGHTS,/storage/home/hcoda1/4/paphiwetsa3/r-dxu345-0/projects/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch} model: pi05: true + # Width of the model's action vector. 32 is openpi's pretrained width + # (the cartesian block layout fits in it). Configs with a human domain + # raise it to 144 for the wrist-first hand-keypoint action; PI then + # re-initializes only openpi's two action projections at that width + # and zero-pads narrower embodiments (eva's 20-D) into the same vector. action_dim: 32 action_horizon: 100 max_token_len: 180 diff --git a/egomimic/hydra_configs/model/pi0.5_bc_aria.yaml b/egomimic/hydra_configs/model/pi0.5_bc_aria.yaml index 8782b49304..7dcfe5dc02 100644 --- a/egomimic/hydra_configs/model/pi0.5_bc_aria.yaml +++ b/egomimic/hydra_configs/model/pi0.5_bc_aria.yaml @@ -2,17 +2,21 @@ defaults: - pi0.5_base - _self_ +# pi0.5 BC on human data (aria, stored as human_bimanual) with the 144-D +# wrist-frame hand-keypoint action (data/aria.yaml); see pi0.5_bc_mecka.yaml. robomimic_model: camera_transforms: {} + config: + model: + action_dim: 144 ac_keys: - human_bimanual: "actions_cartesian" + human_bimanual: "actions_keypoints" domains: ["human_bimanual"] action_converters: rules: HUMAN_BIMANUAL: - _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + _target_: egomimic.utils.action_utils.HumanBimanualKeypoints # optional fallback if no match is found fallback: _target_: egomimic.utils.action_utils.BaseActionConverter - diff --git a/egomimic/hydra_configs/model/pi0.5_bc_mecka.yaml b/egomimic/hydra_configs/model/pi0.5_bc_mecka.yaml index 8782b49304..fa4d5a0deb 100644 --- a/egomimic/hydra_configs/model/pi0.5_bc_mecka.yaml +++ b/egomimic/hydra_configs/model/pi0.5_bc_mecka.yaml @@ -2,17 +2,24 @@ defaults: - pi0.5_base - _self_ +# pi0.5 BC on human data (mecka, stored as human_bimanual) with the 144-D +# wrist-frame hand-keypoint action (data/mecka.yaml, `keypoints_wristframe_6d`). +# The action vector is widened to 144: PI re-initializes openpi's two action +# projections at that width, everything else loads from pi05_base. The prompt +# state is the cartesian head-frame ee_pose (proprio_keys_for_prompt default). robomimic_model: camera_transforms: {} + config: + model: + action_dim: 144 ac_keys: - human_bimanual: "actions_cartesian" + human_bimanual: "actions_keypoints" domains: ["human_bimanual"] action_converters: rules: HUMAN_BIMANUAL: - _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + _target_: egomimic.utils.action_utils.HumanBimanualKeypoints # optional fallback if no match is found fallback: _target_: egomimic.utils.action_utils.BaseActionConverter - diff --git a/egomimic/hydra_configs/model/pi0.5_bc_scale.yaml b/egomimic/hydra_configs/model/pi0.5_bc_scale.yaml index 79ed54e8f2..dae132236f 100644 --- a/egomimic/hydra_configs/model/pi0.5_bc_scale.yaml +++ b/egomimic/hydra_configs/model/pi0.5_bc_scale.yaml @@ -2,16 +2,21 @@ defaults: - pi0.5_base - _self_ +# pi0.5 BC on human data (scale, stored as human_bimanual) with the 144-D +# wrist-frame hand-keypoint action (data/scale.yaml); see pi0.5_bc_mecka.yaml. robomimic_model: camera_transforms: {} + config: + model: + action_dim: 144 ac_keys: - human_bimanual: "actions_cartesian" + human_bimanual: "actions_keypoints" domains: ["human_bimanual"] action_converters: rules: HUMAN_BIMANUAL: - _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + _target_: egomimic.utils.action_utils.HumanBimanualKeypoints # optional fallback if no match is found fallback: _target_: egomimic.utils.action_utils.BaseActionConverter diff --git a/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml index d9c2695b68..c1ef9350b1 100644 --- a/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml +++ b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml @@ -2,11 +2,19 @@ defaults: - pi0.5_base - _self_ +# pi0.5 cotrain: eva on the cam-frame cartesian 6D action (20-D, packed into +# the canonical 32-slot block layout and zero-padded), human on the 144-D +# wrist-frame hand-keypoint action (data/cotrain_pi_base.yaml). The action +# vector is 144 wide; PI re-initializes openpi's two action projections at +# that width, everything else loads from pi05_base. robomimic_model: camera_transforms: {} + config: + model: + action_dim: 144 ac_keys: eva_bimanual: "actions_cartesian" - human_bimanual: "actions_cartesian" + human_bimanual: "actions_keypoints" domains: ["eva_bimanual", "human_bimanual"] action_converters: @@ -14,7 +22,7 @@ robomimic_model: EVA_BIMANUAL: _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler HUMAN_BIMANUAL: - _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + _target_: egomimic.utils.action_utils.HumanBimanualKeypoints # optional fallback if no match is found fallback: - _target_: egomimic.utils.action_utils.BaseActionConverter \ No newline at end of file + _target_: egomimic.utils.action_utils.BaseActionConverter diff --git a/egomimic/rldb/embodiment/embodiment.py b/egomimic/rldb/embodiment/embodiment.py index c6fb9d1870..83a00db10b 100644 --- a/egomimic/rldb/embodiment/embodiment.py +++ b/egomimic/rldb/embodiment/embodiment.py @@ -106,6 +106,10 @@ def _strip_pi_keymap_mode(cls, keymap_mode: str) -> str: "cartesian_padded": "cartesian_6d", "cartesian_wristframe_ypr": "cartesian_wristframe_6d", "cartesian_wristframe_quat": "cartesian_wristframe_6d", + "keypoints_headframe_ypr": "keypoints_headframe_6d", + "keypoints_headframe_quat": "keypoints_headframe_6d", + "keypoints_wristframe_ypr": "keypoints_wristframe_6d", + "keypoints_wristframe_quat": "keypoints_wristframe_6d", } diff --git a/egomimic/rldb/embodiment/human.py b/egomimic/rldb/embodiment/human.py index cadef181c6..d30e701b43 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -17,6 +17,8 @@ ConcatKeys, DeleteKeys, InterpolatePose, + KeypointsRot6DToYPR, + KeypointsYPRToRot6D, PadGripperZeros, PoseCoordinateFrameTransform, QuaternionPoseToYPR, @@ -226,18 +228,24 @@ def get_keymap( keymap_mode: str, has_head_pose: bool = True, include_aria_keypoints: bool = False, + include_ee_pose: bool = False, norm_mode: bool = False, annotation_key: str = None, ): """Build the keymap. Per-vendor knobs are explicit args from the data config: ``has_head_pose`` (False only for contributors that truly omit - ``obs_head_pose``; incompatible with cartesian modes, which pivot on it) and ``include_aria_keypoints`` - (Aria=True). ``norm_mode``/``annotation_key`` behave as in the base. + ``obs_head_pose``; incompatible with cartesian modes, which pivot on + it), ``include_aria_keypoints`` (Aria=True) and ``include_ee_pose`` + (keypoints mode only: also read the palm-origin ``obs_ee_pose`` keys so + the keypoint transform modes can build the cartesian + ``observations.state.ee_pose`` proprio, the pi0.5 prompt state). + ``norm_mode``/``annotation_key`` behave as in the base. """ key_map = cls._get_keymap( keymap_mode, has_head_pose=has_head_pose, include_aria_keypoints=include_aria_keypoints, + include_ee_pose=include_ee_pose, ) if annotation_key is not None and not norm_mode: key_map[annotation_key] = { @@ -260,10 +268,13 @@ def _get_keymap( keymap_mode: str, has_head_pose: bool = True, include_aria_keypoints: bool = False, + include_ee_pose: bool = False, ): """Canonical MANO keymap. ``include_aria_keypoints`` additionally exposes - the raw Aria-layout proprio keypoints alongside the MANO ones. The front - image is always ``VIZ_IMAGE_KEY``; Pi renames it onto its own slot. + the raw Aria-layout proprio keypoints alongside the MANO ones; + ``include_ee_pose`` adds the palm-origin ee_pose keys to the keypoints + keymap. The front image is always ``VIZ_IMAGE_KEY``; Pi renames it + onto its own slot. """ base_mode = _strip_pi_keymap_mode(cls, keymap_mode) front_key = cls.VIZ_IMAGE_KEY @@ -338,6 +349,23 @@ def _get_keymap( "zarr_key": "right.obs_wrist_pose", }, } + if include_ee_pose: + # Palm-origin ee_pose alongside the keypoints: lets the + # keypoint transform modes build the SAME head-frame + # ``observations.state.ee_pose`` proprio the cartesian pi runs + # put in the prompt (left-wrist fix, 6D, grip-padded), so a + # keypoint-action run differs from them only in the action + # space. Pair with ``get_transform_list(include_ee_pose=True)``. + for side in ("left", "right"): + key_map[f"{side}.action_ee_pose"] = { + "key_type": "action_keys", + "zarr_key": f"{side}.obs_ee_pose", + "horizon": horizon, + } + key_map[f"{side}.obs_ee_pose"] = { + "key_type": "proprio_keys", + "zarr_key": f"{side}.obs_ee_pose", + } if include_aria_keypoints: # Raw Aria-layout keypoints exposed alongside the canonical MANO # ones (proprio, no horizon: no transform consumes them). @@ -370,12 +398,15 @@ def get_transform_list( "cartesian_wristframe_6d", "keypoints_headframe_ypr", "keypoints_headframe_quat", + "keypoints_headframe_6d", "keypoints_wristframe_ypr", "keypoints_wristframe_quat", + "keypoints_wristframe_6d", ], stride: int = 3, fix_left_wrist_convention: bool = False, pad_proprio_gripper: bool = False, + include_ee_pose: bool = False, allow_legacy_rotation: bool = False, ) -> list[Transform]: """Transform pipeline. ``stride`` is the per-vendor action stride @@ -389,15 +420,15 @@ def get_transform_list( wrist-frame convention of zarrs written before the ``rot_left`` fix in ``mecka_to_zarr.compute_hand_pose_xyzquat`` (which double-mirrored the left hand onto the right hand's spatial convention): the raw left - pose keys (ee_pose in cartesian modes, wrist_pose in keypoints modes; - the converter writes the same rotation into both) are right-multiplied - by Rz(180 deg) before any frame math, exactly equivalent to - reconverting. It is a property of the DATA, not of the vendor: enable - it for any zarr written by the pre-fix hand-pose converter (measured: - every mecka and scale episode in the production set), and disable it - for data reconverted after the fix, which would double-flip. Aria and - LightWheel are not covered by that measurement -- no such episodes are - in the production set -- so leave them unset. + pose keys (ee_pose in cartesian modes, wrist_pose in keypoints modes, + both when ``include_ee_pose``; the converter writes the same rotation + into both) are right-multiplied by Rz(180 deg) before any frame math, + exactly equivalent to reconverting. It is a property of the DATA, not + of the vendor: enable it for any zarr written by the pre-fix hand-pose + converter (measured: every mecka and scale episode in the production + set), and disable it for data reconverted after the fix, which would + double-flip. Aria and LightWheel are not covered by that measurement -- + no such episodes are in the production set -- so leave them unset. ``left_wrist_convention`` re-derives the verdict from a zarr's own keypoints when a setting has to be checked. @@ -405,21 +436,31 @@ def get_transform_list( ``observations.state.ee_pose`` 18 -> 20 (zero grip slots at 9/19) so the pi0.5 ``State:`` prompt bins align positionally with the robot 20D layout. 6D modes only. + + ``include_ee_pose`` (keypoints modes only; pair with + ``get_keymap(include_ee_pose=True)``) additionally builds the + cartesian pipeline's head-frame ``observations.state.ee_pose`` proprio + (left-wrist fix, 6D-encoded, optionally grip-padded) next to the + keypoint action, so a keypoint-action pi0.5 run shares its prompt + state with the cartesian runs and differs only in the action space. """ _reject_legacy_rotation(cls, mode, allow_legacy_rotation) is_keypoints = mode.startswith("keypoints") prefix: list[Transform] = [] if fix_left_wrist_convention: - fix_keys = ( - ["left.action_wrist_pose", "left.obs_wrist_pose"] - if is_keypoints - else ["left.action_ee_pose", "left.obs_ee_pose"] - ) + fix_keys = [] + if is_keypoints: + fix_keys += ["left.action_wrist_pose", "left.obs_wrist_pose"] + if not is_keypoints or include_ee_pose: + fix_keys += ["left.action_ee_pose", "left.obs_ee_pose"] prefix = [RotateLocalFrame(keys=fix_keys)] - if pad_proprio_gripper and not mode.endswith("_6d"): + if include_ee_pose and not is_keypoints: + raise ValueError("include_ee_pose only applies to the keypoints modes") + has_6d_ee_pose = (not is_keypoints and mode.endswith("_6d")) or include_ee_pose + if pad_proprio_gripper and not has_6d_ee_pose: raise ValueError( "pad_proprio_gripper needs a 6D-encoded ee_pose proprio: a " - "cartesian *_6d mode" + "cartesian *_6d mode, or a keypoints mode with include_ee_pose" ) proprio_6d: list[Transform] = [ @@ -464,22 +505,47 @@ def get_transform_list( + [CartesianYPRToRot6D(action_key="actions_cartesian")] + proprio_6d ) - if mode == "keypoints_headframe_ypr": - return prefix + _build_human_keypoints_bimanual_transform_list( - stride=stride, is_quat=False - ) - if mode == "keypoints_headframe_quat": - return prefix + _build_human_keypoints_bimanual_transform_list( - stride=stride, is_quat=True - ) - if mode == "keypoints_wristframe_ypr": - return prefix + _build_human_keypoints_eef_frame_transform_list( - stride=stride, is_quat=False - ) - if mode == "keypoints_wristframe_quat": - return prefix + _build_human_keypoints_eef_frame_transform_list( - stride=stride, is_quat=True + + if is_keypoints: + frame, _, rot = mode[len("keypoints_") :].partition("_") + if frame not in ("headframe", "wristframe") or rot not in ( + "ypr", + "quat", + "6d", + ): + raise ValueError( + f"Unsupported transform_list mode '{mode}' for {cls.__name__}" + ) + ee_pose: list[Transform] = [] + if include_ee_pose: + # The cartesian builder keeps ``obs_head_pose`` alive for the + # keypoint builder (delete_target_world=False); its + # ``actions_cartesian`` by-product is dropped. + ee_pose = ( + _build_human_cartesian_bimanual_transform_list( + stride=stride, delete_target_world=False + ) + + proprio_6d + + [DeleteKeys(keys_to_delete=["actions_cartesian"])] + ) + build = ( + _build_human_keypoints_eef_frame_transform_list + if frame == "wristframe" + else _build_human_keypoints_bimanual_transform_list ) + if rot == "quat": + return prefix + ee_pose + build(stride=stride, is_quat=True) + keypoints = build(stride=stride, is_quat=False) + if rot == "6d": + # 138-D ypr layout -> 144-D: per hand [wrist xyz (3) | wrist + # rot6d (6) | 21 keypoints (63)], for both the action and the + # proprio ``observations.state.keypoints`` (head-frame wrist + # pose + keypoints), continuous like the cartesian 6d modes. + keypoints = keypoints + [ + KeypointsYPRToRot6D(action_key="actions_keypoints"), + KeypointsYPRToRot6D(action_key="observations.state.keypoints"), + ] + return prefix + ee_pose + keypoints raise ValueError(f"Unsupported transform_list mode '{mode}' for {cls.__name__}") @@ -1083,6 +1149,43 @@ def _build_human_cartesian_revert_6d_wristframe_transform_list( ] +def _build_human_keypoints_revert_6d_transform_list( + *, + action_key: str = "actions_keypoints", + obs_key: str = "observations.state.keypoints", +) -> list[Transform]: + """Revert head-frame 144-D keypoint actions (``keypoints_headframe_6d``) + to the 138-D ypr layout the keypoint overlay draws (the keypoint blocks + are already in head frame; only the wrist rotation representation + changes). The proprio is reverted the same way. + """ + return [ + KeypointsRot6DToYPR(action_key=action_key), + KeypointsRot6DToYPR(action_key=obs_key), + ] + + +def _build_human_keypoints_revert_6d_wristframe_transform_list( + *, + action_key: str = "actions_keypoints", + obs_key: str = "observations.state.keypoints", +) -> list[Transform]: + """Revert wrist-frame 144-D keypoint actions (``keypoints_wristframe_6d``) + to head-frame keypoints for the overlay: (1) wrist rot6d -> ypr on the + action and the proprio (144 -> 138 each); (2) the standard keypoint eef + revert projects each hand's 63-D keypoint chunk back to head frame via the + head-frame wrist pose in the proprio, giving a 126-D head-frame + ``actions_keypoints`` (the wrist-pose slices are consumed by the revert). + """ + return [ + KeypointsRot6DToYPR(action_key=action_key), + KeypointsRot6DToYPR(action_key=obs_key), + *_build_human_keypoints_revert_eef_frame_transform_list( + action_key=action_key, obs_key=obs_key, is_quat=False + ), + ] + + def _build_human_cartesian_eef_frame_transform_list( *, target_world: str = "obs_head_pose", diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/zarr/action_chunk_transforms.py index 62a0f91452..6a1352afbe 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -33,6 +33,7 @@ _xyzwxyz_to_matrix, _xyzypr_to_matrix, _ypr_to_rot6d, + bimanual_keypoint_layout, wxyz_to_xyzw, xyzw_to_wxyz, ) @@ -492,6 +493,74 @@ def transform(self, batch: dict) -> dict: return batch +class KeypointsYPRToRot6D(Transform): + """Convert a wrist-first bimanual keypoint vector from a ypr wrist rotation + to the continuous 6D one: per hand + ``[wrist xyz (3) | wrist ypr (3) | 21 keypoints (63)]`` (138 total) -> + ``[wrist xyz (3) | wrist rot6d (6) | 21 keypoints (63)]`` (144 total). + The keypoint blocks are copied through untouched. Works on action chunks + ``(T, 138)`` and single proprio vectors ``(138,)``. + """ + + def __init__( + self, action_key: str = "actions_keypoints", output_key: str | None = None + ): + self.action_key = action_key + self.output_key = output_key or action_key + + def transform(self, batch: dict) -> dict: + arr, is_tensor = _as_array(batch[self.action_key]) + if arr.shape[-1] != 138: + raise ValueError( + f"KeypointsYPRToRot6D expects last-dim 138, got {arr.shape} for " + f"'{self.action_key}'" + ) + layout = bimanual_keypoint_layout(138) + per_hand, n_rot = layout["per_hand"], len(layout["rot"]) // 2 + blocks = [] + for hand in range(2): + o = hand * per_hand + xyz = len(layout["wrist_xyz"]) // 2 + blocks += [ + arr[..., o : o + xyz], + _ypr_to_rot6d(arr[..., o + xyz : o + xyz + n_rot]), + arr[..., o + xyz + n_rot : o + per_hand], + ] + batch[self.output_key] = _like_input(np.concatenate(blocks, axis=-1), is_tensor) + return batch + + +class KeypointsRot6DToYPR(Transform): + """Inverse of :class:`KeypointsYPRToRot6D`: 144 -> 138 (wrist rot6d -> + ypr via Gram-Schmidt, keypoints untouched).""" + + def __init__( + self, action_key: str = "actions_keypoints", output_key: str | None = None + ): + self.action_key = action_key + self.output_key = output_key or action_key + + def transform(self, batch: dict) -> dict: + arr, is_tensor = _as_array(batch[self.action_key]) + layout = bimanual_keypoint_layout(arr.shape[-1]) + if layout is None or arr.shape[-1] != 144: + raise ValueError( + f"KeypointsRot6DToYPR expects last-dim 144, got {arr.shape} for " + f"'{self.action_key}'" + ) + per_hand = layout["per_hand"] + blocks = [] + for hand in range(2): + o = hand * per_hand + blocks += [ + arr[..., o : o + 3], + _rot6d_to_ypr(arr[..., o + 3 : o + 9]), + arr[..., o + 9 : o + per_hand], + ] + batch[self.output_key] = _like_input(np.concatenate(blocks, axis=-1), is_tensor) + return batch + + class RotateLocalFrame(Transform): """Right-multiply listed xyz+quat(wxyz) pose keys by a constant LOCAL rotation: ``R_new = R_old @ R_fix``. Relabels the pose's own axes without diff --git a/egomimic/scripts/tutorials/aria_to_mano_convert.py b/egomimic/scripts/tutorials/aria_to_mano_convert.py index 7047e06967..305aa9fbe0 100644 --- a/egomimic/scripts/tutorials/aria_to_mano_convert.py +++ b/egomimic/scripts/tutorials/aria_to_mano_convert.py @@ -78,7 +78,9 @@ def fetch_episode_loader(): print(f"Aria episode: {episode_hash} (of {len(aria_df)})") key_map = Human.get_keymap(keymap_mode="keypoints") - transform_list = Human.get_transform_list(mode="keypoints_headframe_ypr", stride=3) + transform_list = Human.get_transform_list( + mode="keypoints_headframe_ypr", stride=3, allow_legacy_rotation=True + ) resolver = S3EpisodeResolver( str(CACHE_DIR), key_map=key_map, transform_list=transform_list ) diff --git a/egomimic/scripts/tutorials/mecka_keypoints_viz.py b/egomimic/scripts/tutorials/mecka_keypoints_viz.py index 65786380ce..2364e04efc 100644 --- a/egomimic/scripts/tutorials/mecka_keypoints_viz.py +++ b/egomimic/scripts/tutorials/mecka_keypoints_viz.py @@ -34,7 +34,9 @@ def main() -> None: print(f"Using mecka episode: {episode_hash} (of {len(mecka_df)} candidates)") key_map = Human.get_keymap(keymap_mode="keypoints") - transform_list = Human.get_transform_list(mode="keypoints_headframe_ypr", stride=1) + transform_list = Human.get_transform_list( + mode="keypoints_headframe_ypr", stride=1, allow_legacy_rotation=True + ) resolver = S3EpisodeResolver( str(CACHE_DIR), key_map=key_map, transform_list=transform_list diff --git a/egomimic/scripts/tutorials/zarr_data_viz.ipynb b/egomimic/scripts/tutorials/zarr_data_viz.ipynb index 2f2e12d795..953fc1c802 100644 --- a/egomimic/scripts/tutorials/zarr_data_viz.ipynb +++ b/egomimic/scripts/tutorials/zarr_data_viz.ipynb @@ -291,7 +291,9 @@ ], "source": [ "key_map = Human.get_keymap(keymap_mode=\"keypoints\")\n", - "transform_list = Human.get_transform_list(mode=\"keypoints_headframe_ypr\", stride=3)\n", + "transform_list = Human.get_transform_list(\n", + " mode=\"keypoints_headframe_ypr\", stride=3, allow_legacy_rotation=True\n", + ")\n", "\n", "resolver = S3EpisodeResolver(\n", " TEMP_DIR,\n", diff --git a/egomimic/utils/action_utils.py b/egomimic/utils/action_utils.py index 7a80fbd970..fc6d9b098f 100644 --- a/egomimic/utils/action_utils.py +++ b/egomimic/utils/action_utils.py @@ -1,12 +1,16 @@ """pi0.5 action I/O. Actions arrive from the data pipeline already in their model-native layout -and already normalized by the standard MultiDataset pipeline: xyz+rot6d -(+gripper) per arm, the ypr->6D conversion being the ``CartesianYPRToRot6D`` -data transform, with norm stats computed on 6D. +and already normalized by the standard MultiDataset pipeline: -The forward pass only *packs* the normalized vector into openpi's 32-slot -action vector (``to32_norm_6d``) and the eval path unpacks it +- cartesian: xyz+rot6d(+gripper) per arm (the ypr->6D conversion is the + ``CartesianYPRToRot6D`` data transform, norm stats are computed on 6D); +- hand keypoints: the 144-D wrist-first ``[wrist xyz | wrist rot6d | + 21 keypoints] x 2`` vector (``keypoints_*_6d`` modes). + +The forward pass only *packs* the normalized vector into openpi's action +vector (``to32_norm_6d`` builds the canonical 32-slot block layout, and the +PI algo pads that to ``model.action_dim``) and the eval path unpacks it (``from32_norm_6d``). No rotation math and no normalization happen here. """ @@ -44,15 +48,28 @@ def _ensure_bsd(x: torch.Tensor) -> torch.Tensor: return x -def _pad32(x: torch.Tensor) -> torch.Tensor: +def pad_to_width(x: torch.Tensor, width: int) -> torch.Tensor: + """Zero-pad the last dim of a (B,S,D) tensor up to ``width``. + + Raises when ``D > width``: a wider native action than the model's + ``action_dim`` would otherwise be truncated silently. + """ x = _ensure_bsd(x) B, S, D = x.shape - if D == 32: + if D == width: return x - if D < 32: - pad = torch.zeros(B, S, 32 - D, dtype=x.dtype, device=x.device) - return torch.cat([x, pad], dim=-1) - return x[..., :32] + if D > width: + raise ValueError( + f"Action has {D} dims but the model action_dim is {width}; set " + "model.robomimic_model.config.model.action_dim >= the widest " + "embodiment action" + ) + pad = torch.zeros(B, S, width - D, dtype=x.dtype, device=x.device) + return torch.cat([x, pad], dim=-1) + + +def _pad32(x: torch.Tensor) -> torch.Tensor: + return pad_to_width(x, 32) def _stat_tensor(stats: dict[str, Any], key: str, ref: torch.Tensor) -> torch.Tensor: @@ -276,3 +293,33 @@ def from32_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: L = actions32[..., 0:9] # drop left gripper slot at idx 9 R = actions32[..., 10:19] # drop right gripper slot at idx 19 return torch.cat([L, R], dim=-1) # (B,S,18) + + +class HumanBimanualKeypoints(BaseActionConverter): + """144-D wrist-first MANO keypoint action, per hand + ``[wrist xyz+rot6d in its own wrist frame (9) | 21 keypoints in that + wrist frame (63)]`` x {left, right} (``keypoints_wristframe_6d``), packed + identity-first into the model's action vector. The PI algo resizes + openpi's action projections to ``model.action_dim`` (>= 144) and pads the + vector up to it; ``from32_norm_6d`` slices the native width back out. + """ + + native_dim = 144 + + def to32_norm_6d(self, actions: torch.Tensor) -> torch.Tensor: + actions = _ensure_bsd(actions) + if actions.shape[-1] != self.native_dim: + raise ValueError( + f"HumanBimanualKeypoints: expected {self.native_dim}-dim, got " + f"{actions.shape[-1]}" + ) + return actions + + def from32_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: + actions32 = _ensure_bsd(actions32) + if actions32.shape[-1] < self.native_dim: + raise ValueError( + f"HumanBimanualKeypoints: expected >={self.native_dim} dims, got " + f"{actions32.shape[-1]}" + ) + return actions32[..., : self.native_dim] diff --git a/egomimic/utils/pose_utils.py b/egomimic/utils/pose_utils.py index fa2095ae15..9f37e879dc 100644 --- a/egomimic/utils/pose_utils.py +++ b/egomimic/utils/pose_utils.py @@ -227,6 +227,49 @@ def bimanual_cartesian_layout(width: int) -> dict | None: return BIMANUAL_CARTESIAN_LAYOUTS.get(int(width)) +# Wrist-first bimanual hand-keypoint layouts, per hand +# [wrist xyz (3) | wrist rot (3 ypr or 6 rot6d) | 21 MANO keypoints (63)] +# x {left, right}. ``wrist_xyz`` / ``rot`` index the wrist pose channels, +# ``keypoints`` the flattened (21, 3) keypoint block of each hand. +def _keypoint_layout(rot_width: int) -> dict: + per_hand = 3 + rot_width + 63 + wrist_xyz, rot, kp = [], [], [] + for hand in range(2): + o = hand * per_hand + wrist_xyz += list(range(o, o + 3)) + rot += list(range(o + 3, o + 3 + rot_width)) + kp += list(range(o + 3 + rot_width, o + per_hand)) + return { + "wrist_xyz": tuple(wrist_xyz), + "rot": tuple(rot), + "keypoints": tuple(kp), + "per_hand": per_hand, + } + + +BIMANUAL_KEYPOINT_LAYOUTS = { + 138: _keypoint_layout(3), # ypr wrist rotation + 144: _keypoint_layout(6), # rot6d wrist rotation (the default) +} + + +def bimanual_keypoint_layout(width: int) -> dict | None: + """Index layout for a wrist-first bimanual keypoint action/proprio vector. + + Returns a dict with ``wrist_xyz`` / ``rot`` / ``keypoints`` index tuples + and the per-hand block width, or ``None`` for an unrecognized width + (138 = ypr wrist, 144 = rot6d wrist). + + MANO keypoint 0 IS the wrist, so per hand the wrist position appears twice: + once as ``wrist_xyz`` and once as the first keypoint. Deliberate -- 21 + keypoints is the MANO interface every consumer reshapes to -- and it is a + duplicate, not a constant: the only place it degenerates is the wrist-frame + PROPRIO, where kp0 sits at its own frame origin and normalization treats it + as a constant channel (NORM_MIN_RANGE). + """ + return BIMANUAL_KEYPOINT_LAYOUTS.get(int(width)) + + def _matrix_to_xyzwxyz(mats: np.ndarray) -> np.ndarray: """ args: @@ -327,10 +370,25 @@ def _split_action_pose(actions): return left_xyz, left_ypr, right_xyz, right_ypr -def _split_keypoints(keypoints, wrist_in_data: bool = False, is_quat: bool = True): +def _split_keypoints( + keypoints, + wrist_in_data: bool = False, + is_quat: bool = True, + is_rot6d: bool = False, +): + """Split a bimanual keypoint vector into its per-hand blocks. + + With ``wrist_in_data`` the layout is wrist-first per hand, + ``[wrist xyz | wrist rot | 21 keypoints]`` x 2, where the rotation block is + a quaternion (4, ``is_quat``), rot6d columns (6, ``is_rot6d``) or ypr (3); + returns the six blocks in order. Without it the vector is just + ``[left keypoints (63) | right keypoints (63)]`` and two blocks are returned. + """ if wrist_in_data: xyz_size = 3 - if is_quat: + if is_rot6d: + angle_size = 6 + elif is_quat: angle_size = 4 else: angle_size = 3 diff --git a/egomimic/utils/viz_utils.py b/egomimic/utils/viz_utils.py index 69b8d34fbf..0b2f31942e 100644 --- a/egomimic/utils/viz_utils.py +++ b/egomimic/utils/viz_utils.py @@ -311,7 +311,11 @@ def _viz_keypoints( vis = base.copy() h, w = vis.shape[:2] - if actions.shape[-1] == 140: + if actions.shape[-1] == 144: + _, _, left_keypoints, _, _, right_keypoints = _split_keypoints( + actions, wrist_in_data=True, is_rot6d=True + ) + elif actions.shape[-1] == 140: _, _, left_keypoints, _, _, right_keypoints = _split_keypoints( actions, wrist_in_data=True ) diff --git a/tests/fixtures/synthetic_episodes.py b/tests/fixtures/synthetic_episodes.py index b4317a057f..a63c8324ae 100644 --- a/tests/fixtures/synthetic_episodes.py +++ b/tests/fixtures/synthetic_episodes.py @@ -101,6 +101,16 @@ def write_episode( } # what the poses were composed through if v.has_head_pose: numeric["obs_head_pose"] = _pose(T, rng) + if v.embodiment == "human_bimanual": + # Hand keypoints (the default human action): a constant wrist pose per + # hand plus 21 MANO keypoints held at fixed offsets from it, so the + # wrist-frame keypoint action is constant in time like the poses. + for side in ("left", "right"): + wrist = _pose(T, rng) + offsets = rng.normal(0, 0.03, (1, 21, 3)) + kp = wrist[:, None, :3] + offsets # (T, 21, 3) world frame + numeric[f"{side}.obs_wrist_pose"] = wrist + numeric[f"{side}.obs_keypoints"] = kp.reshape(T, 63) return ZarrWriter.create_and_write( root / f"{vendor}_{seed:02d}.zarr", numeric_data=numeric, diff --git a/tests/fixtures/train_harness.py b/tests/fixtures/train_harness.py index 1dc9db0db4..6bdb415f8b 100644 --- a/tests/fixtures/train_harness.py +++ b/tests/fixtures/train_harness.py @@ -14,9 +14,11 @@ import egomimic.trainHydra as train_hydra from fixtures.synthetic_episodes import write_episode -# Width of the concatenated proprio the Pi wrapper hands openpi: the -# 6D-encoded ee_pose, grip-padded to the robot 20-D layout on both. -STATE_DIM = {"eva_bimanual": 20, "human_bimanual": 20} +# Width of the concatenated proprio the Pi wrapper hands openpi: eva = the +# 6D-encoded 20-D ee_pose; human = the 144-D keypoint proprio + the 20-D +# grip-padded ee_pose the keypoint data configs also emit for the prompt +# (include_ee_pose + pad_proprio_gripper). +STATE_DIM = {"eva_bimanual": 20, "human_bimanual": 144 + 20} PI_TOKENIZER = ( "google/paligemma-3b-mix-224" # tokenizer_model_name in model/pi0.5_base.yaml ) @@ -131,6 +133,8 @@ class StubPI0(nn.Module): def __init__(self, config): super().__init__() self.config = config + # Like openpi: 32-wide regardless of config.action_dim. PI swaps them + # for action_dim-wide ones when action_dim != 32. self.action_in_proj = nn.Linear(32, 64) self.action_out_proj = nn.Linear(64, 32) @@ -153,11 +157,20 @@ def forward(self, observation, actions): k, mask, ) - assert actions.shape == (B, self.config.action_horizon, 32), actions.shape + action_dim = self.config.action_dim + assert actions.shape == ( + B, + self.config.action_horizon, + action_dim, + ), actions.shape + assert self.action_in_proj.in_features == action_dim + assert self.action_out_proj.out_features == action_dim pred = self.action_out_proj(torch.tanh(self.action_in_proj(actions))) return ((pred - actions) ** 2).mean(dim=(-1, -2)) @torch.no_grad() def sample_actions(self, device, observation, noise=None, num_steps=10): B = observation.state.shape[0] - return torch.zeros(B, self.config.action_horizon, 32, device=device) + return torch.zeros( + B, self.config.action_horizon, self.config.action_dim, device=device + ) diff --git a/tests/unit/test_infra_port_wiring.py b/tests/unit/test_infra_port_wiring.py new file mode 100644 index 0000000000..381a975d7e --- /dev/null +++ b/tests/unit/test_infra_port_wiring.py @@ -0,0 +1,258 @@ +"""Wiring checks for the hand-keypoint default: recipes, stems, the mecka +left-wrist fix in keypoint modes, and the packed-width loss reduction.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import torch +from scipy.spatial.transform import Rotation as R + +import egomimic +from egomimic.rldb.embodiment.human import Human + +CONFIG_DIR = Path(egomimic.__file__).parent / "hydra_configs" +TRAIN_TOP_LEVEL = sorted(p.stem for p in CONFIG_DIR.glob("train_*.yaml")) +HUMAN = "human_bimanual" + +# The all-cartesian cam-frame cotrain recipe is cotrain_pi_lang plus these +# overrides (model/pi0.5_cotrain_eva_aria_6d.yaml); there is no data config +# that differs from its parent only in the transform mode. +CAM_FRAME_COTRAIN = [ + "data.train_datasets.eva_bimanual.resolver.transform_list.mode=cartesian_6d", + f"data.train_datasets.{HUMAN}.resolver.key_map.keymap_mode=cartesian", + f"data.train_datasets.{HUMAN}.resolver.key_map.include_ee_pose=false", + f"data.train_datasets.{HUMAN}.resolver.transform_list.mode=cartesian_6d", + f"data.train_datasets.{HUMAN}.resolver.transform_list.include_ee_pose=false", +] + + +def _human_dataset_nodes(cfg): + for split in ( + "train_datasets", + "valid_datasets", + "train_viz_datasets", + "unseen_op_valid_datasets", + ): + node = (cfg.data.get(split) or {}).get(HUMAN) + if node is not None and node.get("resolver") is not None: + yield split, node.resolver + + +@pytest.mark.parametrize("top", TRAIN_TOP_LEVEL) +def test_human_action_width_matches_data_mode(top, compose_resolve): + """A model declaring the 144-D human action must be fed the + keypoints_*_6d data (and vice versa); a stale ypr/cartesian mode would + only fail at the first forward pass.""" + cfg = compose_resolve(top, []) + rm = cfg.model.robomimic_model + if HUMAN not in (rm.get("domains") or []): + pytest.skip(f"{top}: no human domain") + dims = (rm.get("dims") or {}).get(HUMAN) + ac_key = rm.ac_keys[HUMAN] + for split, resolver in _human_dataset_nodes(cfg): + mode = resolver.transform_list.mode + keymap_mode = resolver.key_map.keymap_mode + if ac_key == "actions_keypoints": + assert keymap_mode == "keypoints", (top, split, keymap_mode) + assert mode.startswith("keypoints_") and mode.endswith("_6d"), ( + top, + split, + mode, + ) + if dims is not None: + assert dims.action == 144, (top, split, dims) + else: + assert keymap_mode == "cartesian", (top, split, keymap_mode) + assert mode.startswith("cartesian"), (top, split, mode) + + +@pytest.mark.parametrize( + "model,data,extra", + [ + ("hpt_bc_flow_aria", "aria", []), + ("hpt_bc_flow_mecka", "mecka", []), + ("hpt_bc_flow_scale", "scale", []), + ("hpt_bc_flow_human", "human", []), + ("pi0.5_bc_aria", "aria", []), + ("pi0.5_bc_mecka", "mecka", []), + ("pi0.5_bc_scale", "scale", []), + ("pi0.5_cotrain_eva_aria", "cotrain_pi_base", []), + ("pi0.5_cotrain_eva_aria_6d", "cotrain_pi_lang", CAM_FRAME_COTRAIN), + ], +) +def test_vendor_pairings_agree_on_the_human_action(model, data, extra, compose_resolve): + cfg = compose_resolve( + "train_zarr_cartesian", [f"model={model}", f"data={data}", *extra] + ) + rm = cfg.model.robomimic_model + ac_key = rm.ac_keys[HUMAN] + resolver = cfg.data.train_datasets[HUMAN].resolver + mode = resolver.transform_list.mode + if ac_key == "actions_keypoints": + assert mode.endswith("_6d") and mode.startswith("keypoints_"), (model, mode) + if rm.get("dims"): + assert rm.dims[HUMAN].action == 144 + if "pi0.5" in model: + assert rm.config.model.action_dim == 144 + assert "Keypoints" in rm.action_converters.rules.HUMAN_BIMANUAL._target_ + # the pi0.5 prompt state is the cartesian ee_pose + assert resolver.key_map.get("include_ee_pose") is True + assert resolver.transform_list.get("include_ee_pose") is True + else: + assert mode.startswith("cartesian"), (model, mode) + if "pi0.5" in model: + assert mode.endswith("_6d"), (model, mode) + assert rm.config.model.action_dim == 32 + + +def test_hpt_stems_follow_the_human_action(compose_resolve): + """HPT human stems: keypoint models expose exactly ``state_keypoints`` + (a base's ``state_ee_pose`` is nulled out, not inherited).""" + from egomimic.algo.hpt import HPTModel + + for model in ("hpt_bc_flow_human", "hpt_cotrain_flow_seperate_head"): + cfg = compose_resolve("train_zarr_cartesian", [f"model={model}"]) + stems = cfg.model.robomimic_model.stem_specs[HUMAN] + live = {k for k, v in stems.items() if v is not None} + assert live == {"state_keypoints"}, (model, live) + cfg = compose_resolve( + "train_zarr_cartesian", ["model=hpt_cotrain_flow_shared_head"] + ) + stems = cfg.model.robomimic_model.stem_specs[HUMAN] + assert {k for k, v in stems.items() if v is not None} == {"state_ee_pose"} + + # the model side drops null stems instead of registering None + m = HPTModel.__new__(HPTModel) + torch.nn.Module.__init__(m) + m.stem_spec, m.modalities, m.stems = {}, {}, {} + m.init_domain_stem("human_bimanual", {"state_ee_pose": None, "x": object()}) + assert m.modalities["human_bimanual"] == ["x"] + assert set(m.stems) == {"human_bimanual_x"} + + +# ------------------------------------------------------ mecka left-wrist fix +def _pose(rng): + q = R.random(random_state=int(rng.integers(1 << 31))).as_quat() + return np.concatenate([rng.uniform(-1, 1, 3), q[[3, 0, 1, 2]]]) + + +def _chunk(rng, start, n): + out = np.zeros((n, 7)) + p, r = start[:3].copy(), R.from_quat(start[[4, 5, 6, 3]]) + for t in range(n): + if t: + p = p + rng.normal(0, 0.01, 3) + r = R.from_rotvec(rng.normal(0, 0.05, 3)) * r + out[t] = np.concatenate([p, r.as_quat()[[3, 0, 1, 2]]]) + return out + + +def _rz180(pose7): + """What the fixed converter would have written: the same pose with its + local axes relabelled by Rz(180 deg).""" + out = np.array(pose7, dtype=np.float64, copy=True) + q = R.from_quat(out[..., [4, 5, 6, 3]]) * R.from_euler("z", np.pi) + out[..., 3:7] = q.as_quat()[..., [3, 0, 1, 2]] + return out + + +def _apply(tl, s): + s = {k: (v.copy() if isinstance(v, np.ndarray) else v) for k, v in s.items()} + for t in tl: + s = t.transform(s) + return s + + +def test_fix_left_wrist_convention_equals_reconverting_in_keypoint_mode(): + rng = np.random.default_rng(31) + H = Human.ACTION_HORIZON + raw = {"obs_head_pose": _pose(rng)} + for side in ("left", "right"): + wrist = _pose(rng) + raw[f"{side}.obs_wrist_pose"] = wrist + raw[f"{side}.action_wrist_pose"] = _chunk(rng, wrist, H) + kp = raw[f"{side}.action_wrist_pose"][:, None, :3] + rng.uniform( + -0.1, 0.1, (H, 21, 3) + ) + raw[f"{side}.action_keypoints"] = kp.reshape(H, 63) + raw[f"{side}.obs_keypoints"] = kp[0].reshape(63) + ee = _pose(rng) + raw[f"{side}.obs_ee_pose"] = ee + raw[f"{side}.action_ee_pose"] = _chunk(rng, ee, H) + # "reconverted" twin: the fixed converter relabels the LEFT hand's axes on + # every pose it writes (wrist_pose and ee_pose share the rotation). + fixed = dict(raw) + for k in ( + "left.obs_wrist_pose", + "left.action_wrist_pose", + "left.obs_ee_pose", + "left.action_ee_pose", + ): + fixed[k] = _rz180(raw[k]) + + kw = dict(stride=1, include_ee_pose=True, pad_proprio_gripper=True) + with_flag = _apply( + Human.get_transform_list( + "keypoints_wristframe_6d", fix_left_wrist_convention=True, **kw + ), + raw, + ) + reconverted = _apply( + Human.get_transform_list("keypoints_wristframe_6d", **kw), fixed + ) + assert set(with_flag) == set(reconverted) + for k in with_flag: + np.testing.assert_allclose(with_flag[k], reconverted[k], atol=1e-9, err_msg=k) + # and the flag really changes the left hand's frame + unfixed = _apply(Human.get_transform_list("keypoints_wristframe_6d", **kw), raw) + assert ( + np.abs(unfixed["actions_keypoints"] - with_flag["actions_keypoints"]).max() + > 1e-3 + ) + # the same equivalence without the ee_pose side (plain HPT keypoint data) + plain = {k: v for k, v in raw.items() if "ee_pose" not in k} + plain_fixed = {k: v for k, v in fixed.items() if "ee_pose" not in k} + a = _apply( + Human.get_transform_list( + "keypoints_wristframe_6d", stride=1, fix_left_wrist_convention=True + ), + plain, + ) + b = _apply( + Human.get_transform_list("keypoints_wristframe_6d", stride=1), plain_fixed + ) + for k in a: + np.testing.assert_allclose(a[k], b[k], atol=1e-9, err_msg=k) + + +def test_pi_loss_is_reduced_over_the_packed_width(): + from egomimic.algo.pi import PI + from egomimic.rldb.embodiment.embodiment import EMBODIMENT + from egomimic.utils.action_utils import ( + ConverterRegistry, + HumanBimanualKeypoints, + RobotBimanualCartesianEuler, + ) + + pi = PI.__new__(PI) + pi.action_registry = ConverterRegistry() + pi._packed_widths = {} + eva, human = EMBODIMENT.EVA_BIMANUAL.value, EMBODIMENT.HUMAN_BIMANUAL.value + pi.action_registry.register(eva, "actions_cartesian", RobotBimanualCartesianEuler()) + pi.action_registry.register(human, "actions_keypoints", HumanBimanualKeypoints()) + losses = torch.zeros(2, 5, 144) + losses[..., 32:] = 1.0 # error only in the zero-padded slots + eva_loss = pi._reduce_loss(losses, torch.zeros(2, 5, 20), eva, "actions_cartesian") + assert eva_loss.item() == 0.0 + kp_loss = pi._reduce_loss( + losses, torch.zeros(2, 5, 144), human, "actions_keypoints" + ) + assert abs(kp_loss.item() - 112 / 144) < 1e-6 + # already-reduced losses pass through unchanged + assert ( + pi._reduce_loss(torch.tensor([1.0, 3.0]), None, eva, "actions_cartesian").item() + == 2.0 + ) diff --git a/tests/unit/test_keypoints_6d.py b/tests/unit/test_keypoints_6d.py new file mode 100644 index 0000000000..c4a523006f --- /dev/null +++ b/tests/unit/test_keypoints_6d.py @@ -0,0 +1,276 @@ +"""Independent end-to-end checks of the 144-D wrist-frame hand-keypoint +pipeline (the default human action): + + raw world-frame keypoints + wrist / head poses + -> Human.get_transform_list("keypoints_wristframe_6d", include_ee_pose=True) + -> actions_keypoints (T, 144), observations.state.keypoints (144,), + observations.state.ee_pose (18 / 20,) + -> _build_human_keypoints_revert_6d_wristframe_transform_list (evaluator) + -> 126-D head-frame keypoints == independent inv(T_head) @ kp_world + +Ground truth is plain numpy/scipy SE(3) math sharing no code with the +pipeline. Interpolation endpoints are exact, so the chunk is compared at +t = 0 and t = -1 against raw frames 0 and horizon-1. +""" + +import numpy as np +import pytest +import torch +from scipy.spatial.transform import Rotation as R + +from egomimic.rldb.embodiment.embodiment import Embodiment +from egomimic.rldb.embodiment.human import ( + Human, + _build_human_keypoints_revert_6d_transform_list, + _build_human_keypoints_revert_6d_wristframe_transform_list, +) +from egomimic.rldb.zarr.action_chunk_transforms import ( + KeypointsRot6DToYPR, + KeypointsYPRToRot6D, +) +from egomimic.utils.pose_utils import bimanual_keypoint_layout + +HORIZON = Human.ACTION_HORIZON # raw frames per chunk (30) +CHUNK = 100 + + +def _rng(seed): + return np.random.default_rng(seed) + + +def _rand_pose(rng, scale=1.0): + q = R.random(random_state=int(rng.integers(1 << 31))).as_quat() # xyzw + return np.concatenate([rng.uniform(-scale, scale, 3), q[[3, 0, 1, 2]]]) + + +def _rand_chunk(rng, start, n): + out = np.zeros((n, 7)) + p = start[:3].copy() + r = R.from_quat(start[[4, 5, 6, 3]]) + for t in range(n): + if t > 0: + p = p + rng.normal(0, 0.01, 3) + r = R.from_rotvec(rng.normal(0, 0.05, 3)) * r + q = r.as_quat() + out[t] = np.concatenate([p, q[[3, 0, 1, 2]]]) + return out + + +def _T(p7): + M = np.eye(4) + M[:3, :3] = R.from_quat(p7[[4, 5, 6, 3]]).as_matrix() + M[:3, 3] = p7[:3] + return M + + +def _in_frame(T_frame, pts): + """Points (..., 3) world -> frame.""" + inv = np.linalg.inv(T_frame) + return pts @ inv[:3, :3].T + inv[:3, 3] + + +def _apply(transform_list, sample): + s = {k: (v.copy() if isinstance(v, np.ndarray) else v) for k, v in sample.items()} + for t in transform_list: + s = t.transform(s) + return s + + +def _raw_sample(rng): + head = _rand_pose(rng) + raw = {"obs_head_pose": head} + world = {} + for side in ("left", "right"): + wrist_obs = _rand_pose(rng) + wrist_act = _rand_chunk(rng, wrist_obs, HORIZON) # (30, 7) world + offsets = rng.uniform(-0.1, 0.1, size=(HORIZON, 21, 3)) + kp_act = wrist_act[:, None, :3] + offsets # (30, 21, 3) world + kp_obs = kp_act[0] + ee_obs = _rand_pose(rng) + raw[f"{side}.obs_wrist_pose"] = wrist_obs + raw[f"{side}.action_wrist_pose"] = wrist_act + raw[f"{side}.obs_keypoints"] = kp_obs.reshape(63) + raw[f"{side}.action_keypoints"] = kp_act.reshape(HORIZON, 63) + raw[f"{side}.obs_ee_pose"] = ee_obs + raw[f"{side}.action_ee_pose"] = _rand_chunk(rng, ee_obs, HORIZON) + world[side] = (wrist_obs, wrist_act, kp_obs, kp_act) + return raw, head, world + + +def test_keypoint_transform_round_trips_138_144(): + rng = _rng(0) + x = np.concatenate( + [ + rng.uniform(-1, 1, (5, 3)), + rng.uniform(-1, 1, (5, 3)), + rng.uniform(-1, 1, (5, 63)), + rng.uniform(-1, 1, (5, 3)), + rng.uniform(-1, 1, (5, 3)), + rng.uniform(-1, 1, (5, 63)), + ], + axis=-1, + ) + assert x.shape == (5, 138) + y = KeypointsYPRToRot6D(action_key="k").transform({"k": x.copy()})["k"] + assert y.shape == (5, 144) + layout = bimanual_keypoint_layout(144) + # keypoint blocks pass through untouched + np.testing.assert_allclose(y[:, list(layout["keypoints"])][:, :63], x[:, 6:69]) + back = KeypointsRot6DToYPR(action_key="k").transform({"k": y})["k"] + np.testing.assert_allclose(back, x, atol=1e-9) + # a single proprio vector works the same way + single = KeypointsYPRToRot6D(action_key="k").transform({"k": x[0].copy()})["k"] + assert single.shape == (144,) + t = KeypointsYPRToRot6D(action_key="k").transform({"k": torch.from_numpy(x)})["k"] + assert isinstance(t, torch.Tensor) and t.shape == (5, 144) + + +@pytest.mark.parametrize("pad", [False, True]) +def test_keypoints_wristframe_6d_pipeline_matches_independent_math(pad): + rng = _rng(3 + int(pad)) + raw, head, world = _raw_sample(rng) + fwd = Human.get_transform_list( + "keypoints_wristframe_6d", + stride=1, + include_ee_pose=True, + pad_proprio_gripper=pad, + ) + out = _apply(fwd, raw) + assert set(out) == { + "actions_keypoints", + "observations.state.keypoints", + "observations.state.ee_pose", + }, sorted(out) + act = np.asarray(out["actions_keypoints"]) + obs = np.asarray(out["observations.state.keypoints"]) + ee = np.asarray(out["observations.state.ee_pose"]) + assert act.shape == (CHUNK, 144) and obs.shape == (144,) + assert ee.shape == ((20,) if pad else (18,)) + if pad: + assert ee[9] == 0.0 and ee[19] == 0.0 + + layout = bimanual_keypoint_layout(144) + per_hand = layout["per_hand"] + Th = _T(head) + for si, side in enumerate(("left", "right")): + wrist_obs, wrist_act, kp_obs, kp_act = world[side] + Tw = _T(wrist_obs) + o = si * per_hand + # proprio: wrist pose in HEAD frame (xyz + rot6d), keypoints in the + # wrist's own frame + Twh = np.linalg.inv(Th) @ Tw + np.testing.assert_allclose(obs[o : o + 3], Twh[:3, 3], atol=1e-9) + np.testing.assert_allclose(obs[o + 3 : o + 6], Twh[:3, 0], atol=1e-9) + np.testing.assert_allclose(obs[o + 6 : o + 9], Twh[:3, 1], atol=1e-9) + np.testing.assert_allclose( + obs[o + 9 : o + per_hand].reshape(21, 3), _in_frame(Tw, kp_obs), atol=1e-9 + ) + # action: wrist pose relative to the obs wrist (identity at t=0), + # keypoints in the obs-wrist frame; endpoints are exact under + # interpolation. + for t_chunk, t_raw in ((0, 0), (CHUNK - 1, HORIZON - 1)): + Ta = _T(wrist_act[t_raw]) + Trel = np.linalg.inv(Tw) @ Ta + row = act[t_chunk] + np.testing.assert_allclose(row[o : o + 3], Trel[:3, 3], atol=1e-6) + np.testing.assert_allclose(row[o + 3 : o + 6], Trel[:3, 0], atol=1e-6) + np.testing.assert_allclose(row[o + 6 : o + 9], Trel[:3, 1], atol=1e-6) + np.testing.assert_allclose( + row[o + 9 : o + per_hand].reshape(21, 3), + _in_frame(Tw, kp_act[t_raw]), + atol=1e-6, + ) + np.testing.assert_allclose(act[0, list(layout["wrist_xyz"])], 0.0, atol=1e-12) + + # evaluator revert: back to 126-D head-frame keypoints (batched) + rev = _build_human_keypoints_revert_6d_wristframe_transform_list() + reverted = Embodiment.apply_transform( + { + "actions_keypoints": torch.from_numpy(act[None]).float(), + "observations.state.keypoints": torch.from_numpy(obs[None]).float(), + }, + rev, + ) + kp_head = np.asarray(reverted["actions_keypoints"]) + assert kp_head.shape == (1, CHUNK, 126) + for si, side in enumerate(("left", "right")): + _, _, _, kp_act = world[side] + for t_chunk, t_raw in ((0, 0), (CHUNK - 1, HORIZON - 1)): + np.testing.assert_allclose( + kp_head[0, t_chunk, 63 * si : 63 * si + 63].reshape(21, 3), + _in_frame(Th, kp_act[t_raw]), + atol=1e-4, + ) + + +def test_keypoints_headframe_6d_pipeline_and_revert(): + rng = _rng(9) + raw, head, world = _raw_sample(rng) + for k in list(raw): + if "ee_pose" in k: + del raw[k] + fwd = Human.get_transform_list("keypoints_headframe_6d", stride=1) + out = _apply(fwd, raw) + assert set(out) == {"actions_keypoints", "observations.state.keypoints"} + act = np.asarray(out["actions_keypoints"]) + assert act.shape == (CHUNK, 144) + Th = _T(head) + layout = bimanual_keypoint_layout(144) + per_hand = layout["per_hand"] + for si, side in enumerate(("left", "right")): + _, wrist_act, _, kp_act = world[side] + o = si * per_hand + Tah = np.linalg.inv(Th) @ _T(wrist_act[0]) + np.testing.assert_allclose(act[0, o : o + 3], Tah[:3, 3], atol=1e-6) + np.testing.assert_allclose(act[0, o + 3 : o + 6], Tah[:3, 0], atol=1e-6) + np.testing.assert_allclose( + act[0, o + 9 : o + per_hand].reshape(21, 3), + _in_frame(Th, kp_act[0]), + atol=1e-6, + ) + rev = _build_human_keypoints_revert_6d_transform_list() + reverted = Embodiment.apply_transform( + { + "actions_keypoints": torch.from_numpy(act[None]).float(), + "observations.state.keypoints": torch.from_numpy( + np.asarray(out["observations.state.keypoints"])[None] + ).float(), + }, + rev, + ) + assert np.asarray(reverted["actions_keypoints"]).shape == (1, CHUNK, 138) + assert np.asarray(reverted["observations.state.keypoints"]).shape == (1, 138) + + +def test_keypoint_modes_without_ee_pose_emit_no_ee_pose(): + rng = _rng(12) + raw, _, _ = _raw_sample(rng) + for k in list(raw): + if "ee_pose" in k: + del raw[k] + for mode in ("keypoints_wristframe_6d", "keypoints_wristframe_ypr"): + out = _apply( + Human.get_transform_list(mode, stride=1, allow_legacy_rotation=True), raw + ) + assert set(out) == {"actions_keypoints", "observations.state.keypoints"} + ypr = _apply( + Human.get_transform_list( + "keypoints_wristframe_ypr", stride=1, allow_legacy_rotation=True + ), + raw, + ) + assert np.asarray(ypr["actions_keypoints"]).shape == (CHUNK, 138) + + +def test_keymap_include_ee_pose_adds_the_ee_pose_keys(): + km = Human.get_keymap("keypoints") + assert "left.obs_ee_pose" not in km + km = Human.get_keymap("keypoints", include_ee_pose=True) + for side in ("left", "right"): + assert km[f"{side}.obs_ee_pose"]["key_type"] == "proprio_keys" + assert km[f"{side}.action_ee_pose"]["key_type"] == "action_keys" + assert km[f"{side}.action_ee_pose"]["horizon"] == HORIZON + with pytest.raises(ValueError, match="include_ee_pose"): + Human.get_transform_list("cartesian_6d", include_ee_pose=True) + with pytest.raises(ValueError, match="pad_proprio_gripper"): + Human.get_transform_list("keypoints_wristframe_6d", pad_proprio_gripper=True) diff --git a/tests/unit/test_model_dims.py b/tests/unit/test_model_dims.py index c68f993d2b..f8419f855e 100644 --- a/tests/unit/test_model_dims.py +++ b/tests/unit/test_model_dims.py @@ -37,21 +37,27 @@ def _target(name): if _target(p.stem) == "egomimic.algo.hpt.HPT" and p.stem != "egobridge" ) -COTRAIN = {"eva_bimanual": (14, 14), "human_bimanual": (12, 12)} +COTRAIN = {"eva_bimanual": (14, 14), "human_bimanual": (20, 18)} # One shared head: human actions are zero-padded to its width. -SHARED = (14, {"eva_bimanual": (14, 14), "human_bimanual": (12, 14)}) +SHARED = (20, {"eva_bimanual": (14, 20), "human_bimanual": (20, 20)}) # {config: (action_width or None, {emb: (proprio, action)})}; None = no shared # head, so no action_width key. EXPECTED = { "hpt_bc_flow_eva": (None, {"eva_bimanual": (20, 20)}), - "hpt_bc_flow_aria": (None, {"human_bimanual": (20, 18)}), - "hpt_bc_flow_human": (None, {"human_bimanual": (20, 18)}), - "hpt_bc_flow_mecka": (None, {"human_bimanual": (20, 18)}), - "hpt_bc_flow_scale": (None, {"human_bimanual": (20, 18)}), - "hpt_bc_keypoints_base": (None, {"human_bimanual": (138, 138)}), + # Human data defaults to the 144-D wrist-frame hand-keypoint action. + "hpt_bc_flow_aria": (None, {"human_bimanual": (144, 144)}), + "hpt_bc_flow_human": (None, {"human_bimanual": (144, 144)}), + "hpt_bc_flow_mecka": (None, {"human_bimanual": (144, 144)}), + "hpt_bc_flow_scale": (None, {"human_bimanual": (144, 144)}), + "hpt_bc_flow_human_cartesian": (None, {"human_bimanual": (20, 18)}), + "hpt_bc_keypoints_base": (None, {"human_bimanual": (144, 144)}), "hpt_cotrain_enc_dec_base": (None, COTRAIN), - "hpt_cotrain_flow_seperate_head": (None, COTRAIN), + # Separate heads: the human head follows the keypoint default. + "hpt_cotrain_flow_seperate_head": ( + None, + {"eva_bimanual": (14, 14), "human_bimanual": (144, 144)}, + ), "hpt_cotrain_flow_shared_head": SHARED, "hpt_cotrain_mecka_flow_shared_head": SHARED, "hpt_cotrain_scale_flow_shared_head": SHARED, @@ -71,6 +77,8 @@ def _width_leaves(rm): should read: ``.proprio``, ``.action`` or ``action_width``.""" for emb, stems in (rm.get("stem_specs") or {}).items(): for key, stem in (stems or {}).items(): + if stem is None: + continue # a child config dropping a base stem if key.startswith("state_"): path = f"stem_specs.{emb}.{key}.input_dim" yield path, stem["input_dim"], f"{emb}.proprio" diff --git a/tests/unit/test_pi05_norm_rot6d.py b/tests/unit/test_pi05_norm_rot6d.py index 89c40d9279..c9e13bc2ad 100644 --- a/tests/unit/test_pi05_norm_rot6d.py +++ b/tests/unit/test_pi05_norm_rot6d.py @@ -20,7 +20,9 @@ ) from egomimic.utils.action_utils import ( BaseActionConverter, + HumanBimanualKeypoints, RobotBimanualCartesianEuler, + pad_to_width, ) @@ -279,3 +281,57 @@ def test_rotate_local_frame_flips_left_wrist_convention(): single = t.transform({"k": chunk[0].copy()})["k"] np.testing.assert_allclose(single, out[0], atol=1e-12) + + +def test_keypoint_converter_is_identity_and_pads_to_model_width(): + # The 144-D keypoint action is packed identity-first; PI pads it (and the + # cartesian 32-slot layout) up to model.action_dim and the decode slices + # the native width back out of the wider vector. + conv = HumanBimanualKeypoints() + kp = torch.randn(2, 4, 144) + packed = conv.to32_norm_6d(kp) + torch.testing.assert_close(packed, kp) + torch.testing.assert_close(conv.from32_norm_6d(pad_to_width(packed, 160)), kp) + with pytest.raises(ValueError, match="expected 144-dim"): + conv.to32_norm_6d(torch.zeros(1, 1, 138)) + + eva32 = RobotBimanualCartesianEuler().to32_norm_6d(torch.ones(1, 3, 20)) + wide = pad_to_width(eva32, 144) + assert wide.shape[-1] == 144 and torch.all(wide[..., 32:] == 0) + torch.testing.assert_close( + RobotBimanualCartesianEuler().from32_norm_6d(wide)[..., :20], + torch.ones(1, 3, 20), + ) + with pytest.raises(ValueError, match="action_dim is 32"): + pad_to_width(kp, 32) + + +def test_fix_left_wrist_convention_flag_prepends_correction(): + from egomimic.rldb.embodiment.human import Human + from egomimic.rldb.zarr.action_chunk_transforms import RotateLocalFrame + + tl = Human.get_transform_list( + "cartesian_wristframe_6d", stride=1, fix_left_wrist_convention=True + ) + assert isinstance(tl[0], RotateLocalFrame) + assert set(tl[0].keys) == {"left.action_ee_pose", "left.obs_ee_pose"} + # default off: other vendors' data must be untouched + tl_off = Human.get_transform_list("cartesian_wristframe_6d", stride=1) + assert not isinstance(tl_off[0], RotateLocalFrame) + # keypoints modes correct the wrist_pose keys their frames are built on + # (the converter wrote the same double-mirrored rotation into both), plus + # the ee_pose keys when those are built too. + tl_kp = Human.get_transform_list( + "keypoints_wristframe_6d", fix_left_wrist_convention=True + ) + assert isinstance(tl_kp[0], RotateLocalFrame) + assert set(tl_kp[0].keys) == {"left.action_wrist_pose", "left.obs_wrist_pose"} + tl_kp_ee = Human.get_transform_list( + "keypoints_wristframe_6d", fix_left_wrist_convention=True, include_ee_pose=True + ) + assert set(tl_kp_ee[0].keys) == { + "left.action_wrist_pose", + "left.obs_wrist_pose", + "left.action_ee_pose", + "left.obs_ee_pose", + } diff --git a/tests/unit/test_train_step.py b/tests/unit/test_train_step.py index e5fbe91141..0094544868 100644 --- a/tests/unit/test_train_step.py +++ b/tests/unit/test_train_step.py @@ -43,6 +43,26 @@ def _run(cfg) -> float: return loss +def _expected_keys(embodiment: str) -> set[str]: + """Batch keys every recipe must hand the algo: eva trains on the cartesian + action, human on the wrist-frame hand-keypoint action (plus the cartesian + ee_pose proprio the keypoint data configs emit for Pi's prompt).""" + if embodiment == "eva_bimanual": + return { + "actions_cartesian", + "observations.state.ee_pose", + Eva.VIZ_IMAGE_KEY, + "observations.images.right_wrist_img", + "observations.images.left_wrist_img", + } + return { + "actions_keypoints", + "observations.state.keypoints", + "observations.state.ee_pose", + Human.VIZ_IMAGE_KEY, + } + + @pytest.mark.parametrize("vendor", VENDOR_NAMES) def test_hpt_train_step(tmp_path, monkeypatch, vendor): hermetic_env(monkeypatch) @@ -65,14 +85,7 @@ def test_hpt_train_step(tmp_path, monkeypatch, vendor): spy = BatchKeySpy(monkeypatch, HPT) _run(cfg) seen = spy.keys[recipe.embodiment] - is_eva = recipe.embodiment == "eva_bimanual" - front = Eva.VIZ_IMAGE_KEY if is_eva else Human.VIZ_IMAGE_KEY - expected = {"actions_cartesian", "observations.state.ee_pose", front} - if is_eva: - expected |= { - "observations.images.right_wrist_img", - "observations.images.left_wrist_img", - } + expected = _expected_keys(recipe.embodiment) assert ( expected <= seen ), f"{vendor}/hpt missing {expected - seen}; saw {sorted(seen)}" @@ -118,13 +131,7 @@ def test_pi_train_step(tmp_path, monkeypatch, vendor): # The wrapper sees dataset camera names; StubPI0 asserts it was handed # openpi's slot names, which is the remap working end to end. seen = spy.keys[recipe.embodiment] - front = Eva.VIZ_IMAGE_KEY if is_eva else Human.VIZ_IMAGE_KEY - expected = {"actions_cartesian", "observations.state.ee_pose", front, "annotations"} - if is_eva: - expected |= { - "observations.images.left_wrist_img", - "observations.images.right_wrist_img", - } + expected = _expected_keys(recipe.embodiment) | {"annotations"} assert ( expected <= seen ), f"{vendor}/pi missing {expected - seen}; saw {sorted(seen)}" From c1f9660739e8de4b2526ece2ae7a73fd9f04c13e Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Wed, 16 Sep 2026 17:24:24 -0400 Subject: [PATCH 02/79] norm: guard degenerate channels, and add openpi-style horizon pooling Wrist-relative actions made a latent division safe-looking: the t=0 cell of a wrist-frame chunk is exactly the identity pose, so its quantile range is 0 and the historical `(q99 - q1) + 1e-6` denominator scaled any off-convention value by 1e6. A channel whose range is below NORM_MIN_RANGE (1e-4) is now treated as constant -- it normalizes to 0 and unnormalizes to its centre rather than dividing by epsilon. Diffusion Policy / robomimic ignore ranges under 1e-4 and GR00T maps min == max to 0; this is the same rule, applied in the shared _apply_norm_one / _apply_unnorm_one so the dataset and the tests agree. norm_stats.pool_horizon (default false, so nothing changes unless asked) computes one set of stats per channel over every timestep of the chunk, the way openpi does, instead of one per (timestep, channel) cell. The pooled stats are tiled back to the (T, D) shape, so normalize, the bounds check and checkpoints see the same layout either way. It is part of the norm-stat cache key, so flipping it cannot silently reuse stats computed the other way. precompute_norm_stats.py computes and caches the stats for a recipe outside a training run, keyed by the same file + content hash. Which of the two is better is unsettled: the literature is split (LBM per-timestep with a stated rationale, openpi / RDT / GR00T / Diffusion Policy / UMI / lerobot / Octo pooled) and nothing here has compared them on a real recipe, so the default stays the inherited per-timestep one and the config comment says what would make it worth flipping. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- .../hydra_configs/train_zarr_cartesian.yaml | 7 + egomimic/rldb/zarr/norm_cache.py | 2 + egomimic/rldb/zarr/zarr_dataset_multi.py | 20 +- egomimic/scripts/precompute_norm_stats.py | 193 ++++++++++++++++++ egomimic/trainHydra.py | 5 + egomimic/utils/action_utils.py | 32 ++- tests/unit/test_eva_wristframe_norm_pool.py | 116 +++++++++++ 7 files changed, 370 insertions(+), 5 deletions(-) create mode 100644 egomimic/scripts/precompute_norm_stats.py create mode 100644 tests/unit/test_eva_wristframe_norm_pool.py diff --git a/egomimic/hydra_configs/train_zarr_cartesian.yaml b/egomimic/hydra_configs/train_zarr_cartesian.yaml index 333a2372de..bfb110d679 100644 --- a/egomimic/hydra_configs/train_zarr_cartesian.yaml +++ b/egomimic/hydra_configs/train_zarr_cartesian.yaml @@ -31,6 +31,13 @@ seed: 42 # Normalization and norm-stat cache settings (override per-run as needed). norm_stats: norm_mode: quantile + # true: one set of stats per channel over the whole action chunk (openpi, + # RDT, GR00T, Diffusion Policy, lerobot); false: per (timestep, channel) + # cell, which keeps resolution on the near-future steps of a relative-frame + # chunk (TRI's LBM makes the same choice for the same reason). Turn it on + # for checkpoint parity with an openpi-shaped recipe; the two have not been + # compared on a real run here. + pool_horizon: false sample_frac: 1.0 num_workers: 4 save_cache_dir: ${hydra:runtime.output_dir} diff --git a/egomimic/rldb/zarr/norm_cache.py b/egomimic/rldb/zarr/norm_cache.py index f706a62610..a03c25223c 100644 --- a/egomimic/rldb/zarr/norm_cache.py +++ b/egomimic/rldb/zarr/norm_cache.py @@ -153,6 +153,7 @@ def cache_inputs( episodes: Mapping[str, str | None], dataset_cfg, sample_frac: float, + pool_horizon: bool = False, ) -> dict: """Everything the key hashes. ``episodes`` maps episode hash to its ``episode_fingerprint``. norm_mode is not an input: every stat is computed @@ -164,6 +165,7 @@ def cache_inputs( "episodes": sorted([str(h), fp] for h, fp in episodes.items()), "recipe": recipe_inputs(cfg), "sample_frac": float(sample_frac), + "pool_horizon": bool(pool_horizon), "code": code_hash(cfg), } diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 833f11741f..8de84164bc 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -1379,7 +1379,17 @@ def infer_norm_from_dataset( batch_size: int = 512, num_workers: int = 4, precomputed_norm_path: str | None = None, + pool_horizon: bool = False, ): + """Fill ``self.norm_stats[embodiment]`` from a precomputed file or by + sampling ``dataset``. + + ``pool_horizon``: compute each chunked key's stats over every timestep + of the chunk (one set per channel, like openpi) instead of per + (timestep, channel) cell. The pooled stats are tiled back to the + ``(T, D)`` shape, so normalize / bounds check / checkpoints see the + same layout either way. + """ embodiment = dataset_name if isinstance(embodiment, str): embodiment = get_embodiment_id(embodiment) @@ -1464,7 +1474,15 @@ def infer_norm_from_dataset( collected[k] = self._drop_nonfinite_rows( np.concatenate(collected[k], axis=0), k ) - stats_np = self._compute_stats_for_array(collected[k]) + X = collected[k] + if pool_horizon and X.ndim > 2: + stats_np = self._compute_stats_for_array(X.reshape(-1, X.shape[-1])) + stats_np = { + name: np.broadcast_to(arr, X.shape[1:]).copy() + for name, arr in stats_np.items() + } + else: + stats_np = self._compute_stats_for_array(X) self.norm_stats[embodiment][k] = { name: np.asarray(arr, dtype=np.float32) for name, arr in stats_np.items() diff --git a/egomimic/scripts/precompute_norm_stats.py b/egomimic/scripts/precompute_norm_stats.py new file mode 100644 index 0000000000..482c613bcd --- /dev/null +++ b/egomimic/scripts/precompute_norm_stats.py @@ -0,0 +1,193 @@ +"""Precompute normalization stats on a CPU node, so GPU-node time isn't spent +computing them at training startup. + +Replicates trainHydra's norm loop EXACTLY (same hydra config + data config, +same keymap/transform, same sample_frac/seed) but builds NO model and NO +trainer: it only instantiates the train datasets (which s5cmd-syncs any +missing episodes from S3 as a side effect), infers shapes, computes norm +stats, and stores them twice: + + 1. an explicit file, ``/norm_stats/norm_stats.json``, for + ``norm_stats.precomputed_norm_path=/norm_stats``; + 2. the content-keyed cache under ``norm_stats.cache_dir`` (default + ``paths.cache_dir``), keyed exactly as trainHydra keys it, so a plain + training run on the same episodes + recipe hits the cache without any + override. + +The norm-mode keymap strips camera + annotation keys, so the stats pass reads +only the numeric proprio/action arrays: pure CPU work, no GPU, no JPEG decode +beyond one shape-inference sample. + +Usage (CPU node, repo root, emimic venv): + python egomimic/scripts/precompute_norm_stats.py \\ + --data mecka_all_6d --model pi0.5_bc_mecka_6d \\ + --sample-frac 0.1 --num-workers 30 --out /path/to/norm_stats/mecka_all_6d +""" + +from __future__ import annotations + +import argparse +import copy +import os +import time + +import hydra +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import OmegaConf + +import egomimic +from egomimic.rldb.embodiment.embodiment import get_embodiment_id +from egomimic.rldb.zarr import norm_cache +from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset +from egomimic.utils.env import load_env + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--config-name", default="train_zarr_cartesian_pi") + ap.add_argument( + "--data", + required=True, + help="data config group; MUST match the training run's data config", + ) + ap.add_argument( + "--model", + required=True, + help="only needed so the config composes; the model is never built", + ) + ap.add_argument("--sample-frac", type=float, default=0.1) + ap.add_argument( + "--max-samples", + type=int, + default=None, + help="optional hard cap on collected samples: the (N, 100, D) float32 " + "action stack plus np.percentile's sort copy is ~2 x N x 100 x D x 4 " + "bytes of RAM, so an uncapped 0.1 frac of the full mecka set (~8.5M " + "frames) needs >100GB at D=18; 2M samples still gives 2M draws per " + "(step, dim) cell. A cap is not part of the cache key (training never " + "caps), so the cache entry written here stands in for the uncapped " + "stats a training run would otherwise compute.", + ) + ap.add_argument("--num-workers", type=int, default=30) + ap.add_argument( + "--out", + required=True, + help="save_cache_dir; writes /norm_stats/norm_stats.json", + ) + ap.add_argument( + "--no-cache", + action="store_true", + help="only write the explicit file; skip the content-keyed cache entry", + ) + ap.add_argument( + "overrides", + nargs="*", + help="extra hydra overrides appended verbatim (e.g. " + "paths.dataset_dir=/path/to/zarr/mirror)", + ) + args = ap.parse_args() + + cfg_dir = os.path.join(os.path.dirname(egomimic.__file__), "hydra_configs") + GlobalHydra.instance().clear() + overrides = [ + f"data={args.data}", + f"model={args.model}", + f"norm_stats.sample_frac={args.sample_frac}", + f"norm_stats.num_workers={args.num_workers}", + f"norm_stats.save_cache_dir={args.out}", + "norm_stats.precomputed_norm_path=null", + "seed=42", + *args.overrides, + ] + with initialize_config_dir(version_base=None, config_dir=cfg_dir): + cfg = compose(config_name=args.config_name, overrides=overrides) + + import lightning as L + + L.seed_everything(cfg.seed, workers=True) + load_env() + + norm_mode = OmegaConf.select(cfg, "norm_stats.norm_mode", default="quantile") + pool_horizon = bool(OmegaConf.select(cfg, "norm_stats.pool_horizon", default=False)) + cache_dir = None if args.no_cache else OmegaConf.select(cfg, "norm_stats.cache_dir") + + # Mirrors trainHydra: instantiate train datasets (resolver syncs from S3 + # as needed), then a stats-only MultiDataset computes the norm stats from + # a norm-mode (numerics-only) copy of each dataset. + train_datasets = {} + for dataset_name in cfg.data.train_datasets: + if cfg.data.train_datasets[dataset_name] is None: + continue + print(f"[precompute] dataset={dataset_name}: instantiating (syncs S3) ...") + train_datasets[dataset_name] = hydra.utils.instantiate( + cfg.data.train_datasets[dataset_name], dataset_name=dataset_name + ) + + norm_stats = MultiDataset(state={}, norm_mode=norm_mode) + norm_stats.populate_from_datasets(train_datasets) + + for dataset_name, dataset in train_datasets.items(): + print(f"[precompute] dataset={dataset_name}: inferring shapes ...") + norm_stats.infer_shapes_from_batch(dataset[0]) + + inst = copy.deepcopy(cfg.data.train_datasets[dataset_name]) + km = OmegaConf.to_container(inst.resolver.key_map, resolve=False) + km["norm_mode"] = True # strips image + annotation keys + inst.resolver.key_map = km + norm_dataset = hydra.utils.instantiate(inst, dataset_name=dataset_name) + + t0 = time.perf_counter() + norm_stats.infer_norm_from_dataset( + norm_dataset, + dataset_name, + sample_frac=args.sample_frac, + max_samples=args.max_samples, + num_workers=args.num_workers, + precomputed_norm_path=None, # force compute + pool_horizon=pool_horizon, + ) + print( + f"[precompute] {dataset_name}: norm computed in " + f"{time.perf_counter() - t0:.1f}s" + ) + + if cache_dir: + # Same key trainHydra derives, so its next run is a cache hit. + episodes = { + h: norm_cache.episode_fingerprint(getattr(ds, "episode_path", None)) + for h, ds in dataset.datasets.items() + } + inputs = norm_cache.cache_inputs( + dataset_name, + episodes, + cfg.data.train_datasets[dataset_name], + args.sample_frac, + pool_horizon, + ) + key = norm_cache.norm_cache_key(inputs) + emb = get_embodiment_id(dataset_name) + written = norm_cache.write_cached( + cache_dir, + dataset_name, + key, + inputs, + emb, + norm_stats.norm_stats[emb], + norm_stats._norm_run_metadata, + ) + print(f"[precompute] {dataset_name}: cache entry {written}") + + norm_stats.cache_stats(save_cache_dir=args.out) + out_dir = os.path.join(args.out, "norm_stats") + print("\nDONE. Use this in training:") + print(f" norm_stats.precomputed_norm_path={out_dir}") + if cache_dir: + print( + f"(or nothing: the content-keyed cache under {cache_dir} now holds " + "these stats for the same episodes + recipe)" + ) + + +if __name__ == "__main__": + main() diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 4b2b5a17b0..f43f4630ed 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -239,6 +239,9 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: from egomimic.rldb.zarr import norm_cache sample_frac = OmegaConf.select(cfg, "norm_stats.sample_frac", default=1.0) + pool_horizon = bool( + OmegaConf.select(cfg, "norm_stats.pool_horizon", default=False) + ) explicit_path = OmegaConf.select( cfg, "norm_stats.precomputed_norm_path", default=None ) @@ -274,6 +277,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: episodes, cfg.data.train_datasets[dataset_name], sample_frac, + pool_horizon, ) key = norm_cache.norm_cache_key(inputs) cached = norm_cache.find_cached(cache_dir, dataset_name, key, emb) @@ -289,6 +293,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: precomputed_norm_path=explicit_path if explicit_path is not None else cached, + pool_horizon=pool_horizon, ) if key is not None and cached is None: if norm_stats.norm_stats.get(emb): diff --git a/egomimic/utils/action_utils.py b/egomimic/utils/action_utils.py index fc6d9b098f..0ebcc2653b 100644 --- a/egomimic/utils/action_utils.py +++ b/egomimic/utils/action_utils.py @@ -77,6 +77,21 @@ def _stat_tensor(stats: dict[str, Any], key: str, ref: torch.Tensor) -> torch.Te return value.to(dtype=ref.dtype if ref.is_floating_point() else torch.float32) +# A channel whose stat range (std, max - min, or q99 - q1) is below this is +# treated as constant: it normalizes to 0 and unnormalizes to its centre. +# Structural zeros, not rare accidents: the t=0 cell of a wrist-frame action +# chunk is exactly the identity pose, and in a wrist-frame keypoint proprio the +# wrist keypoint sits at its own frame origin while kp9 defines the forward +# axis, so kp0 and kp9's off-axis components are zero at EVERY timestep (28 +# action cells and 10 proprio channels on a real 144-D episode). A bare +# ``+ 1e-6`` denominator scaled any off-convention value in them by 1e6. +# One absolute threshold across channels with different units is deliberate: +# 1e-4 is 0.1 mm of translation but 1/20000 of a rot6d channel's range. Same +# rule and the same constant as Diffusion Policy / robomimic; GR00T instead +# clamps the range, which keeps dividing. +NORM_MIN_RANGE = 1e-4 + + def _norm_center_scale( stats: dict[str, Any], norm_mode: str, ref: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: @@ -99,14 +114,21 @@ def _eps(norm_mode: str) -> float: return 1e-6 if norm_mode == "zscore" else 0.5e-6 +def _degenerate(scale: torch.Tensor, norm_mode: str) -> torch.Tensor: + full_range = scale if norm_mode == "zscore" else 2.0 * scale + return full_range < NORM_MIN_RANGE + + def _apply_norm_one( tensor: torch.Tensor, stats: dict[str, Any], norm_mode: str, ) -> torch.Tensor: - """Normalize with per-key stats (zscore, or [-1, 1] for minmax/quantile).""" + """Normalize with per-key stats (zscore, or [-1, 1] for minmax/quantile). + Channels with a range below ``NORM_MIN_RANGE`` map to 0.""" center, scale = _norm_center_scale(stats, norm_mode, tensor) - return (tensor - center) / (scale + _eps(norm_mode)) + out = (tensor - center) / (scale + _eps(norm_mode)) + return torch.where(_degenerate(scale, norm_mode), torch.zeros_like(out), out) def _apply_unnorm_one( @@ -114,9 +136,11 @@ def _apply_unnorm_one( stats: dict[str, Any], norm_mode: str, ) -> torch.Tensor: - """Inverse of :func:`_apply_norm_one`.""" + """Inverse of :func:`_apply_norm_one`; degenerate channels return their + centre (mean / midpoint) whatever the model predicted.""" center, scale = _norm_center_scale(stats, norm_mode, tensor) - return tensor * (scale + _eps(norm_mode)) + center + out = tensor * (scale + _eps(norm_mode)) + center + return torch.where(_degenerate(scale, norm_mode), center.expand_as(out), out) # ---------- rotation helpers (shared with the eval metrics) ---------- diff --git a/tests/unit/test_eva_wristframe_norm_pool.py b/tests/unit/test_eva_wristframe_norm_pool.py new file mode 100644 index 0000000000..6041ca70f8 --- /dev/null +++ b/tests/unit/test_eva_wristframe_norm_pool.py @@ -0,0 +1,116 @@ +"""Horizon-pooled norm stats and the constant-channel normalization guard.""" + +import numpy as np +import pytest +import torch +from fixtures.synthetic_episodes import write_episode +from hydra import compose, initialize_config_module + +from egomimic.rldb.embodiment.eva import Eva +from egomimic.rldb.zarr import norm_cache +from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset + + +def _compose(name, overrides=()): + with initialize_config_module( + config_module="egomimic.hydra_configs", version_base=None + ): + return compose(config_name=name, overrides=list(overrides)) + + +def test_hpt_eva_dims_match_the_wrist_frame_6d_layout(): + cfg = _compose("train_zarr_cartesian") + dims = cfg.model.robomimic_model.dims.eva_bimanual + assert (dims.proprio, dims.action) == (20, 20) + + +def _eva_stats(tmp_path, pool_horizon): + for i in range(3): + write_episode(tmp_path, "eva", seed=i) + resolver = LocalEpisodeResolver( + tmp_path, + key_map=Eva.get_keymap(keymap_mode="cartesian", norm_mode=True), + transform_list=Eva.get_transform_list(mode="cartesian_wristframe_6d"), + ) + ds = MultiDataset._from_resolver(resolver, mode="total") + stats = MultiDataset(state={}, norm_mode="quantile") + stats.populate_from_datasets({"eva_bimanual": ds}) + stats.infer_shapes_from_batch(ds[0]) + stats.infer_norm_from_dataset( + ds, "eva_bimanual", sample_frac=1.0, num_workers=0, pool_horizon=pool_horizon + ) + return next(iter(stats.norm_stats.values())) + + +def test_pool_horizon_tiles_one_stat_per_channel(tmp_path): + per_t = _eva_stats(tmp_path / "a", pool_horizon=False) + pooled = _eva_stats(tmp_path / "b", pool_horizon=True) + act_t, act_p = per_t["actions_cartesian"], pooled["actions_cartesian"] + for name in act_t: + assert act_p[name].shape == act_t[name].shape + assert act_p[name].ndim == 2 + # every timestep carries the same per-channel value + np.testing.assert_array_equal( + act_p[name], act_p[name][:1].repeat(len(act_p[name]), 0) + ) + # every timestep has equal weight in the pool, so a pooled quantile lies + # between the smallest and largest per-timestep quantile of that channel + for name in ("quantile_1", "quantile_99", "median"): + lo = act_t[name].min(axis=0) - 1e-4 + hi = act_t[name].max(axis=0) + 1e-4 + assert np.all((act_p[name][0] >= lo) & (act_p[name][0] <= hi)), name + # proprio has no horizon axis: identical either way + for name, arr in per_t["observations.state.ee_pose"].items(): + np.testing.assert_allclose( + pooled["observations.state.ee_pose"][name], arr, rtol=1e-6 + ) + + +def test_pool_horizon_is_part_of_the_cache_key(): + args = ("eva_bimanual", {"h": "fp"}, {"resolver": {}}, 1.0) + assert norm_cache.norm_cache_key( + norm_cache.cache_inputs(*args, pool_horizon=True) + ) != norm_cache.norm_cache_key(norm_cache.cache_inputs(*args)) + + +# ------------------------------------------------- degenerate-range guard +@pytest.mark.parametrize("mode", ["zscore", "minmax", "quantile"]) +def test_norm_constant_channel_maps_to_zero_and_back_to_its_centre(mode): + from egomimic.utils.action_utils import ( + NORM_MIN_RANGE, + _apply_norm_one, + _apply_unnorm_one, + ) + + lo = np.array([0.0, -1.0, 2.0], dtype=np.float32) + hi = np.array([0.0, 1.0, 2.0 + NORM_MIN_RANGE / 4], dtype=np.float32) + stats = { + "mean": 0.5 * (lo + hi), + "std": hi - lo, + "min": lo, + "max": hi, + "quantile_1": lo, + "quantile_99": hi, + } + x = torch.tensor([[1e-3, 0.5, 2.0 + 1e-3]]) + n = _apply_norm_one(x, stats, mode) + # channels 0 and 2 are (near-)constant: 0 whatever the offset, not 1e3 + assert n[0, 0].item() == 0.0 and n[0, 2].item() == 0.0 + back = _apply_unnorm_one(torch.full_like(n, 0.7), stats, mode) + np.testing.assert_allclose(back[0, [0, 2]].numpy(), stats["mean"][[0, 2]]) + + # the regular channel keeps the historical formula exactly + if mode == "zscore": + ref = (0.5 - stats["mean"][1]) / (stats["std"][1] + 1e-6) + else: + ref = 2.0 * (0.5 - lo[1]) / (hi[1] - lo[1] + 1e-6) - 1.0 + assert n[0, 1].item() == pytest.approx(float(ref), rel=1e-6) + assert _apply_unnorm_one(n, stats, mode)[0, 1].item() == pytest.approx(0.5) + + +def test_multidataset_normalize_uses_the_guard(): + stats = {"quantile_1": np.zeros(2), "quantile_99": np.array([0.0, 1.0])} + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + n = md._apply_norm_one(torch.tensor([[5e-4, 0.5]]), stats) + assert n.tolist() == [[0.0, 0.0]] From fa31304f8479f49bb0bd5c1ab7c95cef93950fe3 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sun, 20 Sep 2026 16:48:27 -0400 Subject: [PATCH 03/79] norm: stop normalizing the rot6d rotation channels A 6D rotation is two columns of a rotation matrix: already bounded in [-1, 1], and per-channel centre-and-scale pulls the pair off the manifold without buying any resolution. The two projects that anchor an action chunk at near-identity the way this repo does -- TRI's LBM and UMI -- both exclude rotation from normalization for exactly that reason (LBM: "We avoid this normalization procedure for the 6D rotation in the poses to avoid corrupting the rotation matrix"; UMI sends its rotation channels through an identity normalizer). Per timestep AND normalized was the one combination with no precedent, and it is what manufactures the t=0 zero-range rotation cells NORM_MIN_RANGE exists to catch -- this removes them at the source. rot6d_channels() reads the indices off the existing layout tables, so the rule follows the layout rather than a second hardcoded list: the cartesian 18/20 widths and the 144-D keypoint layout, action and proprio alike. The ypr widths keep the old behaviour; an Euler angle is not a matrix column and those modes are legacy-only now. _apply_norm_one / _apply_unnorm_one take the channels and pass them through untouched, so the dataset path, the evaluator revert and the rollout revert all follow from one place, and the degenerate-channel guard is left to do what it is actually good at: the structural kp0 / kp9 position zeros. Norm-stat files stay valid -- what changes is how they are applied -- but a run before this commit and a run after it are not comparable, and the loss is now weighted differently between translation (unit scale) and rotation (raw). Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/rldb/zarr/zarr_dataset_multi.py | 35 +++++++++-- egomimic/utils/action_utils.py | 23 +++++-- egomimic/utils/pose_utils.py | 19 ++++++ tests/unit/test_eva_wristframe_norm_pool.py | 69 +++++++++++++++++++++ tests/unit/test_wrist6d_roundtrip.py | 30 ++++++--- 5 files changed, 156 insertions(+), 20 deletions(-) diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 8de84164bc..92ce38c15c 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -52,6 +52,7 @@ _apply_unnorm_one, ) from egomimic.utils.env import load_env +from egomimic.utils.pose_utils import rot6d_channels def create_default_engine(): @@ -1598,11 +1599,31 @@ def cache_stats(self, save_cache_dir: str): # ---- normalize / unnormalize ---- - def _apply_norm_one(self, tensor, stats): - return _apply_norm_one(tensor, stats, self.norm_mode) + def _apply_norm_one(self, tensor, stats, zarr_key=None): + return _apply_norm_one( + tensor, stats, self.norm_mode, self._identity_channels(zarr_key, tensor) + ) - def _apply_unnorm_one(self, tensor, stats): - return _apply_unnorm_one(tensor, stats, self.norm_mode) + def _apply_unnorm_one(self, tensor, stats, zarr_key=None): + return _apply_unnorm_one( + tensor, stats, self.norm_mode, self._identity_channels(zarr_key, tensor) + ) + + @staticmethod + def _identity_channels(zarr_key, tensor): + """Channels this key leaves unnormalized: the rot6d columns. + + A 6D rotation is two columns of a rotation matrix, already bounded in + [-1, 1], and per-channel scaling and shifting pulls the pair off the + manifold for no gain in resolution -- which is why the two projects + that anchor an action chunk at near-identity, TRI's LBM and UMI, both + exclude rotation from normalization outright. It also removes the t=0 + identity-rotation degeneracy at its source instead of catching it with + NORM_MIN_RANGE. + """ + if zarr_key is None: + return None + return rot6d_channels(zarr_key, tensor.shape[-1]) def normalize(self, data: dict, embodiment_id: int) -> dict: if not self.norm_stats.get(embodiment_id): @@ -1623,7 +1644,7 @@ def normalize(self, data: dict, embodiment_id: int) -> dict: tensor = torch.from_numpy(tensor).float() else: continue - out[zarr_key] = self._apply_norm_one(tensor, stats) + out[zarr_key] = self._apply_norm_one(tensor, stats, zarr_key) return out def unnormalize(self, data: dict, embodiment_id: int) -> dict: @@ -1647,7 +1668,9 @@ def unnormalize(self, data: dict, embodiment_id: int) -> dict: value = torch.from_numpy(value).float() else: continue - out[data_key] = self._apply_unnorm_one(value, stats) + out[data_key] = self._apply_unnorm_one( + value, stats, self.zarr_keys.get(embodiment_id, {}).get(key_name) + ) return out # ---- transform attachment ---- diff --git a/egomimic/utils/action_utils.py b/egomimic/utils/action_utils.py index 0ebcc2653b..399e0edb59 100644 --- a/egomimic/utils/action_utils.py +++ b/egomimic/utils/action_utils.py @@ -14,7 +14,7 @@ (``from32_norm_6d``). No rotation math and no normalization happen here. """ -from typing import Any, Dict, Tuple +from typing import Any, Dict, Sequence, Tuple import torch @@ -119,28 +119,43 @@ def _degenerate(scale: torch.Tensor, norm_mode: str) -> torch.Tensor: return full_range < NORM_MIN_RANGE +def _passthrough( + out: torch.Tensor, tensor: torch.Tensor, channels: Sequence[int] | None +) -> torch.Tensor: + if not channels: + return out + out = out.clone() + out[..., list(channels)] = tensor[..., list(channels)] + return out + + def _apply_norm_one( tensor: torch.Tensor, stats: dict[str, Any], norm_mode: str, + identity_channels: Sequence[int] | None = None, ) -> torch.Tensor: """Normalize with per-key stats (zscore, or [-1, 1] for minmax/quantile). - Channels with a range below ``NORM_MIN_RANGE`` map to 0.""" + Channels with a range below ``NORM_MIN_RANGE`` map to 0; + ``identity_channels`` pass through untouched.""" center, scale = _norm_center_scale(stats, norm_mode, tensor) out = (tensor - center) / (scale + _eps(norm_mode)) - return torch.where(_degenerate(scale, norm_mode), torch.zeros_like(out), out) + out = torch.where(_degenerate(scale, norm_mode), torch.zeros_like(out), out) + return _passthrough(out, tensor, identity_channels) def _apply_unnorm_one( tensor: torch.Tensor, stats: dict[str, Any], norm_mode: str, + identity_channels: Sequence[int] | None = None, ) -> torch.Tensor: """Inverse of :func:`_apply_norm_one`; degenerate channels return their centre (mean / midpoint) whatever the model predicted.""" center, scale = _norm_center_scale(stats, norm_mode, tensor) out = tensor * (scale + _eps(norm_mode)) + center - return torch.where(_degenerate(scale, norm_mode), center.expand_as(out), out) + out = torch.where(_degenerate(scale, norm_mode), center.expand_as(out), out) + return _passthrough(out, tensor, identity_channels) # ---------- rotation helpers (shared with the eval metrics) ---------- diff --git a/egomimic/utils/pose_utils.py b/egomimic/utils/pose_utils.py index 9f37e879dc..a87148dd79 100644 --- a/egomimic/utils/pose_utils.py +++ b/egomimic/utils/pose_utils.py @@ -270,6 +270,25 @@ def bimanual_keypoint_layout(width: int) -> dict | None: return BIMANUAL_KEYPOINT_LAYOUTS.get(int(width)) +# Keys whose last dim is one of the layouts above. +_CARTESIAN_KEYS = ("actions_cartesian", "observations.state.ee_pose") +_KEYPOINT_KEYS = ("actions_keypoints", "observations.state.keypoints") + + +def rot6d_channels(zarr_key: str, width: int) -> tuple[int, ...] | None: + """Indices of the rot6d channels of ``zarr_key`` at ``width``. + + ``None`` when the key or the width is not a known 6D layout, including the + ypr layouts: an Euler angle is not a rotation-matrix column and has none of + the properties that make these channels special. + """ + if zarr_key in _CARTESIAN_KEYS and int(width) in (18, 20): + return bimanual_cartesian_layout(width)["rot"] + if zarr_key in _KEYPOINT_KEYS and int(width) == 144: + return bimanual_keypoint_layout(width)["rot"] + return None + + def _matrix_to_xyzwxyz(mats: np.ndarray) -> np.ndarray: """ args: diff --git a/tests/unit/test_eva_wristframe_norm_pool.py b/tests/unit/test_eva_wristframe_norm_pool.py index 6041ca70f8..c4a91d0621 100644 --- a/tests/unit/test_eva_wristframe_norm_pool.py +++ b/tests/unit/test_eva_wristframe_norm_pool.py @@ -114,3 +114,72 @@ def test_multidataset_normalize_uses_the_guard(): md.norm_mode = "quantile" n = md._apply_norm_one(torch.tensor([[5e-4, 0.5]]), stats) assert n.tolist() == [[0.0, 0.0]] + + +# --- rot6d channels are never normalized ----------------------------------- + + +def _stats_for(width): + lo = -np.ones(width, dtype=np.float32) + hi = np.ones(width, dtype=np.float32) * 3.0 + return { + "mean": 0.5 * (lo + hi), + "std": hi - lo, + "min": lo, + "max": hi, + "quantile_1": lo, + "quantile_99": hi, + } + + +@pytest.mark.parametrize( + "zarr_key,width", + [ + ("actions_cartesian", 18), + ("actions_cartesian", 20), + ("observations.state.ee_pose", 20), + ("actions_keypoints", 144), + ("observations.state.keypoints", 144), + ], +) +def test_rot6d_channels_pass_through_normalization(zarr_key, width): + from egomimic.utils.action_utils import _apply_norm_one, _apply_unnorm_one + from egomimic.utils.pose_utils import rot6d_channels + + rot = list(rot6d_channels(zarr_key, width)) + other = [i for i in range(width) if i not in rot] + x = torch.arange(width, dtype=torch.float32).unsqueeze(0) / width + stats = _stats_for(width) + + n = _apply_norm_one(x, stats, "quantile", rot) + torch.testing.assert_close(n[:, rot], x[:, rot], atol=0, rtol=0) + assert not torch.allclose(n[:, other], x[:, other]) + back = _apply_unnorm_one(n, stats, "quantile", rot) + torch.testing.assert_close(back, x, atol=1e-5, rtol=0) + + +@pytest.mark.parametrize( + "zarr_key,width", + [ + ("actions_cartesian", 12), # ypr: not a rot6d layout + ("actions_keypoints", 138), + ("observations.state.joint_positions", 20), # not a pose key + ("actions_cartesian", 7), # unknown width + ], +) +def test_no_rot6d_channels_outside_the_6d_layouts(zarr_key, width): + from egomimic.utils.pose_utils import rot6d_channels + + assert rot6d_channels(zarr_key, width) is None + + +def test_multidataset_leaves_the_rot6d_block_alone(): + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + stats = _stats_for(18) + x = torch.arange(18, dtype=torch.float32).unsqueeze(0) / 18 + n = md._apply_norm_one(x, stats, "actions_cartesian") + rot = list(range(3, 9)) + list(range(12, 18)) + torch.testing.assert_close(n[:, rot], x[:, rot], atol=0, rtol=0) + # without the key there is no layout to consult, so nothing is skipped + assert not torch.allclose(md._apply_norm_one(x, stats)[:, rot], x[:, rot]) diff --git a/tests/unit/test_wrist6d_roundtrip.py b/tests/unit/test_wrist6d_roundtrip.py index ff98b4857b..4c8dcadfb0 100644 --- a/tests/unit/test_wrist6d_roundtrip.py +++ b/tests/unit/test_wrist6d_roundtrip.py @@ -35,7 +35,11 @@ _reconstruct_R_from_cols, _ypr_to_matrix, ) -from egomimic.utils.pose_utils import _rot6d_to_ypr, _ypr_to_rot6d +from egomimic.utils.pose_utils import ( + _rot6d_to_ypr, + _ypr_to_rot6d, + rot6d_channels, +) T = 100 @@ -245,15 +249,17 @@ def test_human_wristframe_6d_pipeline_round_trips_to_headframe(fix_left): st_act = _quantile_stats(act6, 18) st_obs = _quantile_stats(obs6, 20) a_t, o_t = torch.from_numpy(act6).float(), torch.from_numpy(obs6).float() - a_n = _apply_norm_one(a_t, st_act, "quantile") - o_n = _apply_norm_one(o_t, st_obs, "quantile") + act_rot = rot6d_channels("actions_cartesian", a_t.shape[-1]) + obs_rot = rot6d_channels("observations.state.ee_pose", o_t.shape[-1]) + a_n = _apply_norm_one(a_t, st_act, "quantile", act_rot) + o_n = _apply_norm_one(o_t, st_obs, "quantile", obs_rot) conv = HumanBimanualCartesianEuler() a32 = conv.to32_norm_6d(a_n) assert a32.shape == (B, T, 32) assert torch.all(a32[..., [9, 19]] == 0) and torch.all(a32[..., 20:] == 0) torch.testing.assert_close(conv.from32_norm_6d(a32), a_n, atol=0, rtol=0) - a_un = _apply_unnorm_one(conv.from32_norm_6d(a32), st_act, "quantile") - o_un = _apply_unnorm_one(o_n, st_obs, "quantile") + a_un = _apply_unnorm_one(conv.from32_norm_6d(a32), st_act, "quantile", act_rot) + o_un = _apply_unnorm_one(o_n, st_obs, "quantile", obs_rot) torch.testing.assert_close(a_un, a_t, atol=1e-5, rtol=0) torch.testing.assert_close(o_un, o_t, atol=1e-5, rtol=0) @@ -268,7 +274,9 @@ def test_human_wristframe_6d_pipeline_round_trips_to_headframe(fix_left): _assert_pose12_close(out["observations.state.ee_pose"], gt_obs, atol=1e-4) # (5) a noisy (non-orthonormal) prediction still reverts to finite poses - a_noisy = _apply_unnorm_one(a_n + 0.05 * torch.randn_like(a_n), st_act, "quantile") + a_noisy = _apply_unnorm_one( + a_n + 0.05 * torch.randn_like(a_n), st_act, "quantile", act_rot + ) out_n = Embodiment.apply_transform( {"actions_cartesian": a_noisy, "observations.state.ee_pose": o_un}, rev ) @@ -336,13 +344,15 @@ def test_eva_wristframe_6d_pipeline_round_trips_to_camframe(): st_act, st_obs = _quantile_stats(act6, 20), _quantile_stats(obs6, 20) a_t, o_t = torch.from_numpy(act6).float(), torch.from_numpy(obs6).float() - a_n = _apply_norm_one(a_t, st_act, "quantile") - o_n = _apply_norm_one(o_t, st_obs, "quantile") + act_rot = rot6d_channels("actions_cartesian", a_t.shape[-1]) + obs_rot = rot6d_channels("observations.state.ee_pose", o_t.shape[-1]) + a_n = _apply_norm_one(a_t, st_act, "quantile", act_rot) + o_n = _apply_norm_one(o_t, st_obs, "quantile", obs_rot) conv = RobotBimanualCartesianEuler() a32 = conv.to32_norm_6d(a_n) torch.testing.assert_close(conv.from32_norm_6d(a32), a_n, atol=0, rtol=0) - a_un = _apply_unnorm_one(conv.from32_norm_6d(a32), st_act, "quantile") - o_un = _apply_unnorm_one(o_n, st_obs, "quantile") + a_un = _apply_unnorm_one(conv.from32_norm_6d(a32), st_act, "quantile", act_rot) + o_un = _apply_unnorm_one(o_n, st_obs, "quantile", obs_rot) rev = _build_eva_cartesian_revert_6d_wristframe_transform_list() out = Embodiment.apply_transform( From e57b442c0446d8905de15e840785b552e078d0ea Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Mon, 12 Oct 2026 23:14:18 +0000 Subject: [PATCH 04/79] norm: keep normalizing rot6d for checkpoints saved before passthrough MultiDataset.to_state now records rot6d_passthrough. A checkpoint state that has stats but no flag predates the passthrough and normalized the rot6d channels, so from_state turns it off for those; fresh stats (empty state) keep it on. Without this, eval/rescore/requeue of such a checkpoint fed raw rot6d proprio and left predicted rotations in normalized space, silently. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GqAcWA2Lk9zubUhKX5DW1v --- egomimic/rldb/zarr/zarr_dataset_multi.py | 15 ++++++++-- tests/unit/test_eva_wristframe_norm_pool.py | 32 +++++++++++++++++++++ 2 files changed, 44 insertions(+), 3 deletions(-) diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 92ce38c15c..0859cbb40c 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -986,6 +986,9 @@ class MultiDataset(torch.utils.data.Dataset): """ NORMALIZE_KEY_TYPES = ("proprio_keys", "action_keys") + # Saved in the state: a model trained before rot6d passed through + # normalization must keep normalizing it (see _load_state). + rot6d_passthrough = True def __init__( self, @@ -1609,8 +1612,7 @@ def _apply_unnorm_one(self, tensor, stats, zarr_key=None): tensor, stats, self.norm_mode, self._identity_channels(zarr_key, tensor) ) - @staticmethod - def _identity_channels(zarr_key, tensor): + def _identity_channels(self, zarr_key, tensor): """Channels this key leaves unnormalized: the rot6d columns. A 6D rotation is two columns of a rotation matrix, already bounded in @@ -1621,7 +1623,7 @@ def _identity_channels(zarr_key, tensor): identity-rotation degeneracy at its source instead of catching it with NORM_MIN_RANGE. """ - if zarr_key is None: + if zarr_key is None or not self.rot6d_passthrough: return None return rot6d_channels(zarr_key, tensor.shape[-1]) @@ -1721,6 +1723,7 @@ def to_state(self) -> dict: "zarr_keys": copy.deepcopy(self.zarr_keys), "shapes": copy.deepcopy(self.shapes), "norm_stats": self._clone_norm_stats(self.norm_stats), + "rot6d_passthrough": self.rot6d_passthrough, } @classmethod @@ -1738,6 +1741,12 @@ def _load_state(self, state: dict) -> None: self.zarr_keys = copy.deepcopy(state.get("zarr_keys", {})) self.shapes = copy.deepcopy(state.get("shapes", {})) self.norm_stats = self._clone_norm_stats(state.get("norm_stats", {})) + # A checkpoint state saved before the flag existed normalized the rot6d + # channels; an empty state (fresh stats, e.g. trainHydra) keeps the default. + if "rot6d_passthrough" in state: + self.rot6d_passthrough = bool(state["rot6d_passthrough"]) + elif state.get("norm_stats"): + self.rot6d_passthrough = False for emb in self.embodiments: self.key_types.setdefault(emb, {}) self.zarr_keys.setdefault(emb, {}) diff --git a/tests/unit/test_eva_wristframe_norm_pool.py b/tests/unit/test_eva_wristframe_norm_pool.py index c4a91d0621..acbe8ae94a 100644 --- a/tests/unit/test_eva_wristframe_norm_pool.py +++ b/tests/unit/test_eva_wristframe_norm_pool.py @@ -183,3 +183,35 @@ def test_multidataset_leaves_the_rot6d_block_alone(): torch.testing.assert_close(n[:, rot], x[:, rot], atol=0, rtol=0) # without the key there is no layout to consult, so nothing is skipped assert not torch.allclose(md._apply_norm_one(x, stats)[:, rot], x[:, rot]) + + +def _state_with_stats(width, **extra): + return { + "norm_mode": "quantile", + "embodiments": [4], + "zarr_keys": {4: {"actions_cartesian": "actions_cartesian"}}, + "norm_stats": { + 4: {"actions_cartesian": {k: v for k, v in _stats_for(width).items()}} + }, + **extra, + } + + +def test_fresh_stats_pass_rot6d_through_and_save_the_flag(): + stats = MultiDataset(state={}, norm_mode="quantile") + assert stats.rot6d_passthrough + assert stats.to_state()["rot6d_passthrough"] is True + + +@pytest.mark.parametrize("saved,expect_identity", [(None, False), (True, True)]) +def test_state_round_trip_keeps_the_rot6d_convention(saved, expect_identity): + from egomimic.utils.pose_utils import rot6d_channels + + extra = {} if saved is None else {"rot6d_passthrough": saved} + restored = MultiDataset.from_state(_state_with_stats(20, **extra)) + x = torch.arange(20, dtype=torch.float32).unsqueeze(0) / 20 + rot = list(rot6d_channels("actions_cartesian", 20)) + n = restored._apply_norm_one(x, _stats_for(20), "actions_cartesian") + assert torch.equal(n[:, rot], x[:, rot]) == expect_identity + back = restored._apply_unnorm_one(n, _stats_for(20), "actions_cartesian") + torch.testing.assert_close(back, x, atol=1e-5, rtol=0) From 32b73de7f4273963938956404393f25350115016 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Wed, 16 Sep 2026 17:49:00 -0400 Subject: [PATCH 05/79] data: stop letting one bad episode or a rotation channel take a run down Three independent ways MultiDataset could fail a whole run on data that was fine, or on data it should have rejected up front. Bounds check. Quantile bounds on a rotation channel are meaningless: euler ypr wraps at +-pi and rot6d columns are matrix entries, so neither respects a quantile interval. _bounds_check_channels uses the layout tables to check only translation, gripper and keypoint channels; unrecognized widths still get the full-vector check, and the NaN/Inf check always covers everything. The bounds themselves get slack -- 1e-6 absolute plus half of each cell's own quantile range -- because the check is any-cell over ~13 200 cells per keypoint chunk. Measured on the 2815-frame flagship episode through the real keypoints_wristframe_6d pipeline (zscore, stats from sample_frac 0.10 over its own frames, 403 probe frames, any-key rejection): 39.5 % of frames rejected at slack 0, 17.9 % at 0.05, 11.4 % at 0.1, 5.2 % at 0.25, 3.2 % at 0.5, 1.0 % at 1.0. The 132-channel action chunk drives all of it. Half a range is the knee, not a cushion -- though the absolute numbers move with which frames the stats were drawn from, so the comment says how to reproduce them. Corrupt values sit orders of magnitude outside the range and are still caught. Cells the normalizer treats as constant are exempt, since a tiny offset there cannot blow up the loss -- and which stat decides that is norm_mode's business (std under zscore, the quantile range under quantile), so the exemption asks _degenerate instead of hardcoding q99 - q1; every shipped config is zscore, and a channel that is 0 in 99.6 % of samples and 1.0 in the rest has q99 - q1 = 0 but std = 0.063, i.e. it is still divided by and must still be checked. quantile_0_01 / quantile_99_99 now snap to observed samples ("lower" / "higher") instead of interpolating: under 10k samples the linear percentile lands strictly inside the sample range, so the extreme frames of the very episodes the stats came from were rejected (3-16 % of a single episode in the overfit runs). A stats/sample shape mismatch is logged once per key instead of silently disabling the check. Sampler. Retries stayed inside the failing episode until they exhausted it, so one wholly-bad episode killed the run and NCCL-timed-out every other rank. Retries now widen to the global index space after 25 attempts, and 100 consecutive failures raise. The cap is a systemic-failure detector, not a tolerance dial: attempts resets per __getitem__, so with p = the bad fraction of the whole index space it fires with probability p**(N - 25), a step function in p. At N = 100 the firing band sits at p >~ 0.8, so everything it can reach is systemic; at N = 1000 it needed ~99 % of the store to be bad and got there through ~975 sequential zarr/S3 reads per call per worker, which reads as a hang rather than a failure. Retry logging was one WARNING per attempt -- tens of thousands of lines per run -- now the first 20 per worker and then one in 500. Annotation cutoff. Many Scale-"completed" episodes have no usable annotation span because the injection step lagged, which makes the cutoff a silent no-op on exactly the runs whose point is to clamp chunks at annotation boundaries. S3AnnotationCutoffEpisodeResolver(require_annotations=True) drops them, logs the count it dropped and the count it kept, and raises if nothing is left. The per-episode span line drops to DEBUG: one line per episode per rank is thousands of lines on a real split, and the resolver's aggregate says the same thing. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/rldb/zarr/zarr_dataset_multi.py | 268 +++++++++++++++++--- tests/unit/test_eva_wristframe_norm_pool.py | 18 ++ tests/unit/test_infra_port_wiring.py | 104 ++++++++ tests/unit/test_pi05_norm_rot6d.py | 198 +++++++++++++++ 4 files changed, 552 insertions(+), 36 deletions(-) diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 0859cbb40c..65ee853157 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -50,9 +50,15 @@ from egomimic.utils.action_utils import ( _apply_norm_one, _apply_unnorm_one, + _degenerate, + _norm_center_scale, ) from egomimic.utils.env import load_env -from egomimic.utils.pose_utils import rot6d_channels +from egomimic.utils.pose_utils import ( + bimanual_cartesian_layout, + bimanual_keypoint_layout, + rot6d_channels, +) def create_default_engine(): @@ -205,6 +211,21 @@ def get_fallback_idx( return random.choice(valid_candidates), attempts +def _bounds_check_channels(zarr_key: str, width: int) -> list[int] | None: + """Channels of a known bimanual layout that quantile bounds apply to + (translation, gripper, keypoints), or ``None`` to check the full vector. + Rotation channels (Euler or rot6d columns) are excluded.""" + if zarr_key in ("actions_cartesian", "observations.state.ee_pose"): + layout = bimanual_cartesian_layout(width) + if layout is not None: + return list(layout["xyz"]) + list(layout["grip"]) + elif zarr_key in ("actions_keypoints", "observations.state.keypoints"): + layout = bimanual_keypoint_layout(width) + if layout is not None: + return sorted(list(layout["wrist_xyz"]) + list(layout["keypoints"])) + return None + + def _embodiment_matches(name: object, expected: str) -> bool: """Compare an episode's embodiment string with a dataset's expected name, honouring the legacy vendor aliases (MECKA_BIMANUAL == human_bimanual).""" @@ -990,6 +1011,10 @@ class MultiDataset(torch.utils.data.Dataset): # normalization must keep normalizing it (see _load_state). rot6d_passthrough = True + # Bounds-check slack per (timestep, channel) cell; see _check_bounds. + BOUNDS_ABS_TOL = 1e-6 + BOUNDS_REL_SLACK = 0.5 + def __init__( self, datasets: ( @@ -1142,6 +1167,19 @@ def _check_bounds( q_low = torch.broadcast_to(q_low, arr.shape) q_high = torch.broadcast_to(q_high, arr.shape) except RuntimeError: + # Stats were computed for a different layout than this sample + # (e.g. a stale precomputed norm_stats.json). Say so once + # instead of silently disabling the bounds check for the key; + # normalize() raises on the same mismatch anyway. + warn_key = f"bounds-shape:{zarr_key}" + if warn_key not in self._warned_violations: + self._warned_violations.add(warn_key) + logger.warning( + f"[MultiDataset] bounds check skipped for {zarr_key}: " + f"stats shape {tuple(q_low.shape)} does not broadcast to " + f"sample shape {tuple(arr.shape)} (norm stats computed " + "for a different layout?)" + ) continue if torch.any(torch.isnan(arr)) or torch.any(torch.isinf(arr)): @@ -1152,8 +1190,53 @@ def _check_bounds( logger.warning(prefix) return prefix - below = arr < q_low - above = arr > q_high + # Rotation channels are either Euler ypr (wraps at +-pi) or + # continuous 6D columns; quantile bounds on them are meaningless + # and reject otherwise-valid frames, so only the translation / + # gripper / keypoint channels are bounds-checked. Unrecognized + # widths fall through to a full-vector check; NaN/Inf above still + # covers the full vector. + check_idx = _bounds_check_channels(zarr_key, arr.shape[-1]) + if check_idx is not None: + arr_q = arr[..., check_idx] + q_low = q_low[..., check_idx] + q_high = q_high[..., check_idx] + else: + arr_q = arr + + # Slack on the quantile bounds: a relative part (a fraction of + # each cell's own quantile range) plus a 1e-6 absolute floor. + # The bounds are per (timestep, channel) cell and the check is + # any-cell, so with 13 200 cells per keypoint chunk a strict + # compare against a subsampled stats pass rejects valid frames in + # bulk. Measured on the 2815-frame flagship episode through the + # real keypoints_wristframe_6d pipeline, norm_mode=zscore, stats + # from infer_norm_from_dataset(sample_frac=0.10) over its own + # frames, 403 probe frames, a frame counted as rejected if ANY of + # actions_keypoints / state.keypoints / state.ee_pose fails: + # slack 0 0.05 0.1 0.25 0.5 1.0 + # any 39.5% 17.9% 11.4% 5.2% 3.2% 1.0% + # The 132-channel action chunk drives it (39.0 % -> 3.2 %); the 8 + # checked ee_pose channels never exceed 2.7 %. Which frames the + # stats came from moves the absolute numbers, so re-measure the + # same way before comparing. Corrupt values + # (fill constants, wrong-frame data) sit orders of magnitude + # outside the range and are still caught. + tol = self.BOUNDS_ABS_TOL + self.BOUNDS_REL_SLACK * (q_high - q_low) + below = arr_q < q_low - tol + above = arr_q > q_high + tol + # Cells normalize() treats as constant (wrist-frame t=0, the kp0 / + # kp9 structural zeros) cannot blow up the loss whatever offset + # they carry, so don't reject on them. Which stat decides that is + # norm_mode's business -- std under zscore, the quantile range + # under quantile -- and every shipped config is zscore, so reading + # q99 - q1 here would exempt cells the normalizer still divides by. + constant = self._degenerate_cells(stats, arr.shape) + if constant is not None: + if check_idx is not None: + constant = constant[..., check_idx] + below = below & ~constant + above = above & ~constant if torch.any(below) or torch.any(above): prefix = f"Bounds violation in {zarr_key} ep={episode_name} frame={idx}" warn_key = f"bounds:{episode_name}:{zarr_key}" @@ -1163,11 +1246,36 @@ def _check_bounds( n_above = int(above.sum().item()) logger.warning( f"{prefix} | n_below={n_below} n_above={n_above} " - f"arr_range=[{arr.min().item():.4f}, {arr.max().item():.4f}]" + f"arr_range=[{arr_q.min().item():.4f}, {arr_q.max().item():.4f}]" ) return prefix return None + def _degenerate_cells(self, stats: dict, shape) -> torch.Tensor | None: + """Mask of the cells ``normalize()`` treats as constant, or None when + this key's stats do not carry the range ``norm_mode`` reads.""" + try: + _, scale = _norm_center_scale( + stats, self.norm_mode, torch.empty(0, dtype=torch.float32) + ) + except (KeyError, ValueError): + return None + return torch.broadcast_to(_degenerate(scale, self.norm_mode), shape) + + # Retries stay within the failing episode this many times, then widen to + # the full index space: a single wholly-bad episode must not exhaust the + # sampler and take the run down (NCCL-timeouts the other ranks). + GLOBAL_FALLBACK_ATTEMPTS = 25 + # Cap on consecutive bad samples, after which the worker raises. It is a + # systemic-failure detector, not a tolerance dial: `attempts` resets per + # __getitem__, so with p = the bad fraction of the whole index space the + # cap fires with probability p**(N - GLOBAL_FALLBACK_ATTEMPTS), a step + # function in p. At N = 100 the firing band sits at p >~ 0.8, so anything + # it can detect is systemic; at N = 1000 it needed ~99 % of the store to be + # bad, and got there via 975 sequential zarr/S3 reads per call per worker, + # which looks like a hang rather than a failure. + MAX_FALLBACK_ATTEMPTS = 100 + def __getitem__(self, idx, _attempts: int | None = None): attempts = _attempts while True: @@ -1176,14 +1284,13 @@ def __getitem__(self, idx, _attempts: int | None = None): try: data = dataset[local_idx] except Exception as e: - next_idx, attempts = self._next_after_failure( - idx, - dataset_name, - attempts, - reason=f"Sample failed ({type(e).__name__}: {e}) at " - f"{dataset_name}[{local_idx}]", + reason = ( + f"Sample failed ({type(e).__name__}: {e}) at " + f"{dataset_name}[{local_idx}]" + ) + idx, attempts = self._next_after_failure( + idx, dataset_name, attempts, reason=reason ) - idx = next_idx continue # If this leaf is itself a MultiDataset, it already ran bounds + @@ -1193,38 +1300,50 @@ def __getitem__(self, idx, _attempts: int | None = None): violation = self._check_bounds(data, dataset, local_idx, dataset_name) if violation is not None: - next_idx, attempts = self._next_after_failure( - idx, - dataset_name, - attempts, - reason=violation, + idx, attempts = self._next_after_failure( + idx, dataset_name, attempts, reason=violation ) - idx = next_idx continue - # Bounds passed — normalize and return. + # Bounds passed: normalize and return. if self.norm_stats and data.get("embodiment") in self.norm_stats: data = self.normalize(data, data["embodiment"]) return data def _next_after_failure( self, idx: int, dataset_name: str, attempts: int | None, *, reason: str - ) -> tuple[int, int]: - global_candidates = self._global_indices_by_dataset[dataset_name] - next_idx, attempts = get_fallback_idx( - idx=idx, - candidates=global_candidates, - _attempts=attempts, - max_attempts=len(global_candidates), - exhausted_error=( - f"Entire dataset bad (no valid indices): dataset={dataset_name}" - ), - ) + ) -> tuple[int | None, int]: + """Pick the next index to try after a bad sample at ``idx``. + + Raises once ``MAX_FALLBACK_ATTEMPTS`` consecutive samples failed.""" + attempts = (attempts or 0) + 1 + if attempts >= self.MAX_FALLBACK_ATTEMPTS: + raise RuntimeError( + f"[MultiDataset] {attempts} consecutive bad samples, the last " + f"{attempts - self.GLOBAL_FALLBACK_ATTEMPTS} drawn uniformly " + "from the whole index space: the store, the keymap or the norm " + f"stats are broken, not one episode. Last: {reason}" + ) + if attempts <= self.GLOBAL_FALLBACK_ATTEMPTS: + candidates = [ + c for c in self._global_indices_by_dataset[dataset_name] if c != idx + ] + else: + candidates = None + if candidates: + next_idx = random.choice(candidates) + else: + next_idx = random.randrange(len(self.index_map)) next_dataset_name, next_local_idx = self.index_map[next_idx] - logger.warning( - f"{reason} | attempt {attempts}, " - f"trying {next_dataset_name}[{next_local_idx}]" - ) + # Every retry used to log at WARNING: tens of thousands of lines per + # run. Keep the first few per process, then one in every 500. + self._retry_log_count = getattr(self, "_retry_log_count", 0) + 1 + if self._retry_log_count <= 20 or self._retry_log_count % 500 == 0: + logger.warning( + f"{reason} | attempt {attempts}, " + f"trying {next_dataset_name}[{next_local_idx}] " + f"(retries so far in this worker: {self._retry_log_count})" + ) return next_idx, attempts @classmethod @@ -1571,8 +1690,14 @@ def _compute_stats_for_array(X): "median": np.median(X, axis=0), "quantile_1": np.percentile(X, 1, axis=0), "quantile_99": np.percentile(X, 99, axis=0), - "quantile_0_01": np.percentile(X, 0.01, axis=0), - "quantile_99_99": np.percentile(X, 99.99, axis=0), + # Bounds-check quantiles snap to observed samples ("lower" / + # "higher") instead of interpolating: with fewer than 10k samples + # the default linear percentile lands strictly inside the sample + # range, so the extreme frames of the very episodes the stats were + # computed on were rejected at train time (3-16 % of a single + # episode in the overfit runs). + "quantile_0_01": np.percentile(X, 0.01, axis=0, method="lower"), + "quantile_99_99": np.percentile(X, 99.99, axis=0, method="higher"), } def cache_stats(self, save_cache_dir: str): @@ -2114,13 +2239,26 @@ def _build_frame_to_ann_end(self) -> dict[int, int]: annotation span. Annotations use half-open ``[start_idx, end_idx)``. """ mapping: dict[int, int] = {} + n_spans = 0 for ann in self._load_annotations(): start_idx = int(ann.get("start_idx", -1)) end_idx = int(ann.get("end_idx", -1)) if start_idx < 0 or end_idx <= start_idx: continue + n_spans += 1 for idx in range(start_idx, end_idx): mapping[idx] = end_idx + # Visibility into annotation-cutoff usage: spans/frames_covered of 0 + # means the cutoff is a no-op for this episode. One line per episode + # per rank is thousands of lines on a real split, so this is DEBUG and + # the resolver's kept N/M line carries the aggregate. + logger.debug( + "[AnnotationCutoff] ep=%s spans=%d frames_covered=%d/%d", + Path(self.episode_path).name, + n_spans, + len(mapping), + self.total_frames, + ) return mapping def _chunk_end_idx(self, start_idx: int, horizon: int, key_type: str | None) -> int: @@ -2135,11 +2273,69 @@ def _chunk_end_idx(self, start_idx: int, horizon: int, key_type: str | None) -> return min(end_idx, ann_end) +def _episode_has_annotation_spans(ds: "ZarrDataset") -> bool: + """True if the episode has at least one usable ``[start_idx, end_idx)`` span. + + Many Scale-"completed" episodes have an empty (or span-less) zarr + ``annotations`` array because the annotation-injection step lagged; the + AnnotationCutoff is a no-op for those, so they should be dropped when the + point of the run is to clamp chunks at annotation boundaries. + """ + try: + anns = ds._load_annotations() + except Exception: + return False + return any( + isinstance(a, dict) + and 0 <= int(a.get("start_idx", -1)) < int(a.get("end_idx", -1)) + for a in anns + ) + + class S3AnnotationCutoffEpisodeResolver(S3EpisodeResolver): - """S3EpisodeResolver that loads ZarrAnnotationCutoffDataset instances.""" + """S3EpisodeResolver that loads ZarrAnnotationCutoffDataset instances. + + When ``require_annotations`` is set (default), episodes whose zarr + ``annotations`` array has no usable span are dropped; otherwise the + annotation cutoff would silently no-op on them. + """ _dataset_class = ZarrAnnotationCutoffDataset + def __init__(self, *args, require_annotations: bool = True, **kwargs): + super().__init__(*args, **kwargs) + self.require_annotations = require_annotations + + def resolve(self, filters=None, expected_embodiment=None): + datasets = super().resolve( + filters=filters, expected_embodiment=expected_embodiment + ) + if not self.require_annotations: + return datasets + kept = { + h: ds for h, ds in datasets.items() if _episode_has_annotation_spans(ds) + } + dropped = sorted(set(datasets) - set(kept)) + if dropped: + logger.warning( + "[AnnotationCutoff] dropped %d/%d episodes with no usable " + "annotation spans (e.g. %s)", + len(dropped), + len(datasets), + dropped[:5], + ) + logger.info( + "[AnnotationCutoff] kept %d/%d episodes with usable annotations", + len(kept), + len(datasets), + ) + if not kept: + raise ValueError( + "[AnnotationCutoff] no resolved episodes contain usable annotation " + "spans; check the filter / annotation injection for this dataset." + ) + return kept + class LocalAnnotationCutoffEpisodeResolver(LocalEpisodeResolver): """LocalEpisodeResolver that loads ZarrAnnotationCutoffDataset instances.""" diff --git a/tests/unit/test_eva_wristframe_norm_pool.py b/tests/unit/test_eva_wristframe_norm_pool.py index acbe8ae94a..31bf3a670c 100644 --- a/tests/unit/test_eva_wristframe_norm_pool.py +++ b/tests/unit/test_eva_wristframe_norm_pool.py @@ -215,3 +215,21 @@ def test_state_round_trip_keeps_the_rot6d_convention(saved, expect_identity): assert torch.equal(n[:, rot], x[:, rot]) == expect_identity back = restored._apply_unnorm_one(n, _stats_for(20), "actions_cartesian") torch.testing.assert_close(back, x, atol=1e-5, rtol=0) + + +def test_bounds_check_ignores_constant_cells(): + key = "actions_cartesian" + q1 = np.full((4, 14), -1.0, dtype=np.float32) + q99 = np.full((4, 14), 1.0, dtype=np.float32) + q1[0, :3] = q99[0, :3] = 0.0 # wrist-frame t=0: xyz exactly 0 + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + md.norm_stats = {0: {key: {"quantile_1": q1, "quantile_99": q99}}} + md.zarr_keys = {0: {key: key}} + md._warned_violations = set() + + arr = np.zeros((4, 14), dtype=np.float32) + arr[0, 0] = 1e-3 # off-convention offset at a constant cell: admitted + assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is None + arr[1, 0] = 50.0 # corrupt value at a regular cell: still rejected + assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is not None diff --git a/tests/unit/test_infra_port_wiring.py b/tests/unit/test_infra_port_wiring.py index 381a975d7e..549c03acb0 100644 --- a/tests/unit/test_infra_port_wiring.py +++ b/tests/unit/test_infra_port_wiring.py @@ -3,7 +3,9 @@ from __future__ import annotations +import inspect from pathlib import Path +from types import SimpleNamespace import numpy as np import pytest @@ -256,3 +258,105 @@ def test_pi_loss_is_reduced_over_the_packed_width(): pi._reduce_loss(torch.tensor([1.0, 3.0]), None, eva, "actions_cartesian").item() == 2.0 ) + + +# --------------------------------------------------- annotation-cutoff filter +def test_episode_has_annotation_spans(): + from egomimic.rldb.zarr.zarr_dataset_multi import _episode_has_annotation_spans + + def ds(anns): + return SimpleNamespace(_load_annotations=lambda: anns) + + assert _episode_has_annotation_spans( + ds([{"text": "a", "start_idx": 0, "end_idx": 5}]) + ) + assert not _episode_has_annotation_spans(ds([])) + assert not _episode_has_annotation_spans(ds([{"text": "a"}])) # span-less + assert not _episode_has_annotation_spans(ds([{"start_idx": 5, "end_idx": 5}])) + + def boom(): + raise OSError("corrupt") + + assert not _episode_has_annotation_spans(SimpleNamespace(_load_annotations=boom)) + + +def test_annotation_cutoff_resolver_defaults_to_requiring_spans(): + from egomimic.rldb.zarr.zarr_dataset_multi import S3AnnotationCutoffEpisodeResolver + + r = S3AnnotationCutoffEpisodeResolver.__new__(S3AnnotationCutoffEpisodeResolver) + sig = inspect.signature(S3AnnotationCutoffEpisodeResolver.__init__) + assert sig.parameters["require_annotations"].default is True + good = SimpleNamespace(_load_annotations=lambda: [{"start_idx": 0, "end_idx": 3}]) + bad = SimpleNamespace(_load_annotations=lambda: []) + r.require_annotations = True + # bypass the S3 base resolve: patch it on the instance's class chain + import egomimic.rldb.zarr.zarr_dataset_multi as m + + orig = m.S3EpisodeResolver.resolve + m.S3EpisodeResolver.resolve = lambda self, filters=None, expected_embodiment=None: { + "g": good, + "b": bad, + } + try: + kept = r.resolve(filters=None, expected_embodiment="human_bimanual") + assert set(kept) == {"g"} + r.require_annotations = False + assert set(r.resolve()) == {"g", "b"} + m.S3EpisodeResolver.resolve = ( + lambda self, filters=None, expected_embodiment=None: {"b": bad} + ) + r.require_annotations = True + with pytest.raises(ValueError, match="no resolved episodes"): + r.resolve() + finally: + m.S3EpisodeResolver.resolve = orig + + +# ------------------------------------------------ adversarial-review fixes +def test_bounds_quantiles_include_the_observed_extremes(): + """With fewer than 10k samples a linear 0.01 / 99.99 percentile lands + strictly inside the sample range, so the stats' own extreme frames were + rejected at train time; the bounds quantiles must snap to samples.""" + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + X = np.random.default_rng(0).normal(size=(500, 4)).astype(np.float32) + st = MultiDataset._compute_stats_for_array(X) + np.testing.assert_array_equal(st["quantile_0_01"], X.min(axis=0)) + np.testing.assert_array_equal(st["quantile_99_99"], X.max(axis=0)) + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + md.norm_stats = {0: {"actions_keypoints": st}} + md.zarr_keys = {0: {"actions_keypoints": "actions_keypoints"}} + md._warned_violations = set() + for row in (X.min(axis=0), X.max(axis=0)): + assert ( + md._check_bounds( + {"embodiment": 0, "actions_keypoints": row[None]}, None, 0, "ep" + ) + is None + ) + + +def test_bounds_check_has_relative_slack_but_catches_corrupt_values(): + """Per-cell bounds tolerate frames moderately beyond the stats sample's + range (valid extreme motion) and still reject values orders of magnitude + off (fill constants, wrong-frame data).""" + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + lo, hi = np.full(18, -1.0, np.float32), np.full(18, 1.0, np.float32) + md.norm_stats = {0: {"actions_cartesian": {"quantile_1": lo, "quantile_99": hi}}} + md.zarr_keys = {0: {"actions_cartesian": "actions_cartesian"}} + md._warned_violations = set() + + def check(v): + arr = np.zeros((3, 18), np.float32) + arr[1, 0] = v # an xyz channel + return md._check_bounds( + {"embodiment": 0, "actions_cartesian": arr}, None, 0, "e" + ) + + assert check(1.0 + 0.4 * 2.0) is None # within 50 % of the [-1, 1] range + assert check(1.0 + 0.6 * 2.0) is not None + assert check(1e9) is not None diff --git a/tests/unit/test_pi05_norm_rot6d.py b/tests/unit/test_pi05_norm_rot6d.py index c9e13bc2ad..6a1dce7ef5 100644 --- a/tests/unit/test_pi05_norm_rot6d.py +++ b/tests/unit/test_pi05_norm_rot6d.py @@ -335,3 +335,201 @@ def test_fix_left_wrist_convention_flag_prepends_correction(): "left.action_ee_pose", "left.obs_ee_pose", } + + +def _bounds_check_dataset(key: str, width: int): + """Minimal MultiDataset shell exposing _check_bounds with +-1 quantile + bounds on ``key`` for embodiment 0.""" + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + md.norm_stats = { + 0: { + key: { + "quantile_1": np.full(width, -1.0, dtype=np.float32), + "quantile_99": np.full(width, 1.0, dtype=np.float32), + } + } + } + md.zarr_keys = {0: {key: key}} + md._warned_violations = set() + return md + + +@pytest.mark.parametrize("key", ["actions_cartesian", "observations.state.ee_pose"]) +@pytest.mark.parametrize("width,rot_idx,xyz_idx", [(14, 3, 0), (20, 4, 0), (18, 5, 9)]) +def test_bounds_check_ignores_rotation_channels(key, width, rot_idx, xyz_idx): + # Rotation channels (Euler wraps at +-pi; 6D columns are ~[-1, 1]) must be + # excluded from quantile bounds checking, while translation/gripper + # channels are still checked and NaN/Inf still rejects the full vector. + md = _bounds_check_dataset(key, width) + arr = np.zeros((5, width), dtype=np.float32) + + arr[2, rot_idx] = 50.0 # far outside +-1, but a rotation channel + assert md._check_bounds({"embodiment": 0, key: arr.copy()}, None, 0, "ep") is None + + bad = arr.copy() + bad[2, xyz_idx] = 50.0 # translation channel out of bounds -> violation + assert md._check_bounds({"embodiment": 0, key: bad}, None, 0, "ep") is not None + + nan = arr.copy() + nan[2, rot_idx] = np.nan # NaN anywhere (even rotation) -> violation + assert md._check_bounds({"embodiment": 0, key: nan}, None, 0, "ep") is not None + + +@pytest.mark.parametrize("key", ["actions_keypoints", "observations.state.keypoints"]) +@pytest.mark.parametrize("width,rot_idx,kp_idx", [(138, 4, 10), (144, 78, 90)]) +def test_bounds_check_ignores_wrist_rotation_in_keypoint_layout( + key, width, rot_idx, kp_idx +): + md = _bounds_check_dataset(key, width) + arr = np.zeros((5, width), dtype=np.float32) + arr[1, rot_idx] = 50.0 # wrist rotation channel: not bounds-checked + assert md._check_bounds({"embodiment": 0, key: arr.copy()}, None, 0, "ep") is None + bad = arr.copy() + bad[1, kp_idx] = 50.0 # a keypoint coordinate: checked + assert md._check_bounds({"embodiment": 0, key: bad}, None, 0, "ep") is not None + + +def test_bounds_check_full_vector_for_other_keys(): + # Keys without a known layout (or unrecognized widths) keep the + # full-vector check. + md = _bounds_check_dataset("some_other_key", 20) + arr = np.zeros((5, 20), dtype=np.float32) + arr[2, 4] = 50.0 + assert ( + md._check_bounds({"embodiment": 0, "some_other_key": arr}, None, 0, "ep") + is not None + ) + + md16 = _bounds_check_dataset("actions_cartesian", 16) + arr16 = np.zeros((5, 16), dtype=np.float32) + arr16[2, 4] = 50.0 + assert ( + md16._check_bounds({"embodiment": 0, "actions_cartesian": arr16}, None, 0, "ep") + is not None + ) + + +def test_constant_cell_exemption_follows_norm_mode(): + """A channel the zscore normalizer still divides by must still be + bounds-checked, even though its quantile range is zero: 0 in 99.6 % of + samples and 1.0 in the rest gives q99 - q1 = 0 but std = 0.063.""" + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + stats = { + "quantile_1": np.zeros(18, dtype=np.float32), + "quantile_99": np.zeros(18, dtype=np.float32), + "mean": np.full(18, 0.004, dtype=np.float32), + "std": np.full(18, 0.063, dtype=np.float32), + } + arr = np.zeros((5, 18), dtype=np.float32) + arr[0, 0] = 50.0 # corrupt value in a quantile-collapsed channel + + for mode, rejected in (("quantile", False), ("zscore", True)): + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = mode + md.norm_stats = {0: {"actions_cartesian": stats}} + md.zarr_keys = {0: {"actions_cartesian": "actions_cartesian"}} + md._warned_violations = set() + got = md._check_bounds( + {"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep" + ) + assert (got is not None) is rejected, mode + + +def test_bounds_check_tolerates_roundoff_on_collapsed_bounds(): + """Wrist-frame t=0 cells have bounds [0, 0]; values there must not reject + (normalize() maps such constant cells to 0), while other cells are still + checked.""" + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + q99 = np.ones(18, dtype=np.float32) + q99[0] = 0.0 # channel 0 collapsed to [0, 0] + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + md.norm_stats = { + 0: { + "actions_cartesian": { + "quantile_1": np.zeros(18, dtype=np.float32), + "quantile_99": q99, + } + } + } + md.zarr_keys = {0: {"actions_cartesian": "actions_cartesian"}} + md._warned_violations = set() + arr = np.zeros((5, 18), dtype=np.float32) + arr[0, 0] = 1e-9 # a roundoff-scale xyz value at a [0, 0] bound + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") + is None + ) + arr[0, 0] = 1e-3 # off-convention offset at the constant cell: admitted + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") + is None + ) + arr[0, 1] = 50.0 # a real violation on a regular cell is still caught + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") + is not None + ) + + +def test_bounds_check_warns_once_on_stat_shape_mismatch(caplog): + md = _bounds_check_dataset("actions_cartesian", 18) + arr = np.zeros((5, 20), dtype=np.float32) # stats are 18-wide + with caplog.at_level("WARNING"): + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") + is None + ) + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 1, "ep") + is None + ) + msgs = [r.message for r in caplog.records if "bounds check skipped" in r.message] + assert len(msgs) == 1, msgs + + +def _fallback_dataset(): + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + md = MultiDataset.__new__(MultiDataset) + md.index_map = [("bad", i) for i in range(10)] + [("good", i) for i in range(1000)] + md._global_indices_by_dataset = { + "bad": list(range(10)), + "good": list(range(10, 1010)), + } + return md + + +def test_fallback_widens_to_global_after_local_attempts(): + # A wholly-bad episode must not exhaust the sampler: retries stay inside + # the failing episode for GLOBAL_FALLBACK_ATTEMPTS, then widen to the full + # index space. + md = _fallback_dataset() + attempts = None + seen_local, seen_global = set(), set() + idx = 0 + for _ in range(md.GLOBAL_FALLBACK_ATTEMPTS): + idx, attempts = md._next_after_failure(0, "bad", attempts, reason="r") + seen_local.add(md.index_map[idx][0]) + assert seen_local == {"bad"}, "early retries must stay within the episode" + + for _ in range(50): + idx, attempts = md._next_after_failure(idx, "bad", attempts, reason="r") + seen_global.add(md.index_map[idx][0]) + assert "good" in seen_global, "post-threshold retries must sample globally" + + +def test_fallback_cap_raises(): + # The cap is a systemic-failure detector: everything it can reach has + # already been drawn uniformly from the whole index space and failed. + md = _fallback_dataset() + attempts, idx, n = None, 0, 0 + with pytest.raises(RuntimeError, match="consecutive bad samples"): + for n in range(1, md.MAX_FALLBACK_ATTEMPTS + 2): + idx, attempts = md._next_after_failure(idx, "bad", attempts, reason="boom") + assert n == md.MAX_FALLBACK_ATTEMPTS From ce8666eae9f45b31fee618e2405383e9acd7fba8 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Wed, 16 Sep 2026 18:18:11 -0400 Subject: [PATCH 06/79] eval: layout-aware frame-invariant metrics, and a train-split viz head Metrics were paired MSE over whatever the action vector happened to be, which mixes metres and rotation-matrix entries into one number that means nothing in either unit. action_metrics dispatches on the layout instead: cartesian widths report xyz MSE (m^2), geodesic rotation error in degrees (after Gram-Schmidt, so it survives a non-orthonormal prediction), raw rot6d MSE as a training-tracking diagnostic, DTW and Frechet; keypoint widths report keypoint L2 in metres, wrist xyz / rotation error and keypoint DTW. Each also splits early / mid / late (0-10 / 10-50 / 50-100 % of the chunk), because a chunk model can be right at t=0 and drift, and one average hides that. DTW forgives temporal misalignment that paired MSE punishes: a correct trajectory executed slightly early scores near zero. These are all frame-invariant -- the cam-frame revert is a rigid transform applied to prediction and ground truth alike -- so the revert no longer runs for metrics at all, only when rendering the overlay. That also removes the cam-frame MSE, which was the same number as the native one. Rendering is CPU-bound at roughly a second a frame and dominated viz-epoch wall time (unbounded rendering added ~2h per viz epoch on the combined pretrain), so viz_every_n_epochs gates it on Lightning's (current_epoch + 1) convention, the one validation actually runs on, with the final epoch always rendering so short runs get a video; viz_max_batches caps frames per rank. Metrics still run on every batch of every pass. Video fps is divided by world_size, since distributed validation deals indices round-robin and rank 0's consecutive frames are world_size source frames apart -- they were being written as a world_size-times timelapse. train_viz adds a second val loader over the train split (metrics train_viz/, videos videos_train_viz/), on by default for training runs with an evaluator, so train and val curves come from the same evaluator on the same axes. A third unseen_op_valid loader and data.valid_prefix let the held-out-operator splits log seen_op_valid/ and unseen_op_valid/ separately, instead of the seen operators landing in the train_viz slot and the unseen ones in Valid/. ModelWrapper dispatches on val_loader_names by dataloader_idx. val_at_start runs a pre-fit baseline validation, skipped on resumes because trainer.validate there would score the un-resumed base model. Reverse KL is removed: M forward passes per val batch for a number nothing was read off. Co-Authored-By: Claude Opus 5 (1M context) --- egomimic/algo/hpt.py | 2 - egomimic/eval/action_metrics.py | 210 ++++++ egomimic/eval/eval_hpt.py | 117 +-- egomimic/eval/eval_pi.py | 88 +-- egomimic/eval/eval_train_viz.py | 65 ++ egomimic/eval/eval_video.py | 129 ++-- .../hydra_configs/evaluator/eval_hpt.yaml | 18 +- .../evaluator/eval_hpt_cartesian.yaml | 11 + .../evaluator/eval_hpt_keypoints_head.yaml | 17 + .../evaluator/eval_hpt_keypoints_wrist.yaml | 27 + .../evaluator/eval_hpt_wristframe_6d.yaml | 24 + egomimic/hydra_configs/evaluator/eval_pi.yaml | 32 +- .../evaluator/eval_pi_camframe_6d.yaml | 29 + .../evaluator/eval_pi_wristframe.yaml | 17 + .../evaluator/eval_pi_wristframe_6d.yaml | 29 + .../train_viz_hpt_keypoints_wrist.yaml | 23 + .../train_viz_hpt_wristframe_6d.yaml | 24 + .../hydra_configs/evaluator/train_viz_pi.yaml | 25 + .../evaluator/train_viz_pi_camframe_6d.yaml | 26 + .../evaluator/train_viz_pi_wristframe_6d.yaml | 25 + .../evaluator/viz/cartesian_wrist.yaml | 4 +- .../evaluator/viz/keypoints.yaml | 7 +- .../evaluator/viz/keypoints_wrist.yaml | 6 +- .../viz/pi_keypoints_lang_wrist.yaml | 10 + .../hydra_configs/model/hpt_bc_flow_eva.yaml | 1 - .../model/hpt_bc_flow_human.yaml | 3 +- .../model/hpt_bc_keypoints_base.yaml | 2 - .../model/hpt_cotrain_enc_dec_base.yaml | 1 - .../hydra_configs/train_zarr_cartesian.yaml | 11 + .../train_zarr_keypoint_wrist.yaml | 2 +- .../hydra_configs/train_zarr_keypoints.yaml | 2 +- egomimic/pl_utils/pl_data_utils.py | 77 +- egomimic/pl_utils/pl_model.py | 50 +- egomimic/trainHydra.py | 169 ++++- egomimic/utils/metrics.py | 64 +- tests/unit/test_eva_wristframe_norm_pool.py | 37 + tests/unit/test_hydra_configs.py | 11 +- tests/unit/test_infra_port_wiring.py | 188 +++++ tests/unit/test_pi05_norm_rot6d.py | 696 ++++++++++++++---- 39 files changed, 1853 insertions(+), 426 deletions(-) create mode 100644 egomimic/eval/action_metrics.py create mode 100644 egomimic/eval/eval_train_viz.py create mode 100644 egomimic/hydra_configs/evaluator/eval_hpt_cartesian.yaml create mode 100644 egomimic/hydra_configs/evaluator/eval_hpt_keypoints_head.yaml create mode 100644 egomimic/hydra_configs/evaluator/eval_hpt_keypoints_wrist.yaml create mode 100644 egomimic/hydra_configs/evaluator/eval_hpt_wristframe_6d.yaml create mode 100644 egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml create mode 100644 egomimic/hydra_configs/evaluator/eval_pi_wristframe.yaml create mode 100644 egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml create mode 100644 egomimic/hydra_configs/evaluator/train_viz_hpt_keypoints_wrist.yaml create mode 100644 egomimic/hydra_configs/evaluator/train_viz_hpt_wristframe_6d.yaml create mode 100644 egomimic/hydra_configs/evaluator/train_viz_pi.yaml create mode 100644 egomimic/hydra_configs/evaluator/train_viz_pi_camframe_6d.yaml create mode 100644 egomimic/hydra_configs/evaluator/train_viz_pi_wristframe_6d.yaml create mode 100644 egomimic/hydra_configs/evaluator/viz/pi_keypoints_lang_wrist.yaml diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index 8184f5ac74..979a536b79 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -912,8 +912,6 @@ def __init__( self.freeze_depth = kwargs.get("freeze_depth", 8) model.depth = self.depth - self.rkl_samples = kwargs.get("reverse_kl_samples", 4) - if self.ot: self.ot_warm_start_steps = kwargs.get("ot_warm_start_steps", 0) self.ot_6dof = kwargs.get("ot_6dof", False) diff --git a/egomimic/eval/action_metrics.py b/egomimic/eval/action_metrics.py new file mode 100644 index 0000000000..9c1d099562 --- /dev/null +++ b/egomimic/eval/action_metrics.py @@ -0,0 +1,210 @@ +"""Layout-aware action metrics shared by the PI and HPT evaluators. + +Every metric here is computed on the model's native (unnormalized) output and +is invariant to the rigid transform the evaluator's cam-frame revert applies +to prediction and ground truth alike, so it never needs the revert. + +Two layouts are recognised by width: + +- bimanual cartesian (12/14 ypr, 18/20 rot6d; see + ``pose_utils.bimanual_cartesian_layout``): position MSE, geodesic rotation + error, raw rotation-channel MSE, DTW + Frechet on xyz; +- wrist-first bimanual hand keypoints (138 ypr / 144 rot6d wrist; see + ``pose_utils.bimanual_keypoint_layout``): mean keypoint L2 (m), wrist + position MSE, wrist geodesic rotation error, DTW on the keypoint block. + +Unknown widths yield no metrics. +""" + +from __future__ import annotations + +import math + +import torch + +from egomimic.utils.action_utils import _reconstruct_R_from_cols, _ypr_to_matrix +from egomimic.utils.metrics import dtw_distance, frechet_gaussian_over_time +from egomimic.utils.pose_utils import ( + bimanual_cartesian_layout, + bimanual_keypoint_layout, +) + +_RAD2DEG = 180.0 / math.pi + + +def _paired_mse(pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor: + """Plain elementwise MSE, stateless (no torchmetrics accumulator).""" + return (pred.float() - gt.float()).pow(2).mean() + + +# Horizon segments for the per-segment L2 metrics, as fractions of the chunk: +# early = the first 10 % of timesteps, mid = 10-50 %, late = the second half. +# They separate near-term precision from long-range drift, which one +# chunk-mean error hides (and which per-timestep vs horizon-pooled norm stats +# trade against each other). +_HORIZON_SEGMENTS = (("early", 0.0, 0.1), ("mid", 0.1, 0.5), ("late", 0.5, 1.0)) + + +def _segment_l2(pred: torch.Tensor, gt: torch.Tensor, idx: list, prefix: str) -> dict: + """Mean euclidean error (m) of the xyz triples at ``idx`` per horizon + segment of a ``(B, T, D)`` chunk.""" + T = pred.shape[1] + diff = pred[..., idx].float() - gt[..., idx].float() + l2 = diff.reshape(*diff.shape[:-1], -1, 3).norm(dim=-1) # (B, T, P) + out = {} + for name, lo, hi in _HORIZON_SEGMENTS: + start = min(int(round(lo * T)), T - 1) + end = max(int(round(hi * T)), start + 1) + out[f"{prefix}_{name}"] = l2[:, start:end].mean() + return out + + +def _split_mse(pred_t: torch.Tensor, gt_t: torch.Tensor): + """(translation MSE, rotation-channel MSE) over a bimanual cartesian + vector, so a translation problem reads apart from a rotation one. The + rotation channels are whatever the layout holds: the continuous 6D + rotation-matrix columns for the 18D (human) / 20D (robot) widths, + yaw/pitch/roll for the 12/14D widths. Either way this is a plain + per-channel MSE (pre-Gram-Schmidt, axis-dependent, unitless for 6D; +-pi + wrap and gimbal lock for ypr), i.e. a training-tracking diagnostic, not a + rotation error; see :func:`_rot_geodesic_error` for that. Returns + (None, None) for an unknown width. + """ + layout = bimanual_cartesian_layout(pred_t.shape[-1]) + if layout is None: + return None, None + xyz_idx = list(layout["xyz"]) + rot_idx = list(layout["rot"]) + xyz = _paired_mse(pred_t[..., xyz_idx], gt_t[..., xyz_idx]) + rot = _paired_mse(pred_t[..., rot_idx], gt_t[..., rot_idx]) + return xyz, rot + + +def _geodesic_from_rot_channels(pred: torch.Tensor, gt: torch.Tensor, rot: list): + """Mean geodesic angle (rad) over both arms given the rotation channel + indices of a per-arm-concatenated vector (3 ypr or 6 rot6d per arm).""" + per_arm = len(rot) // 2 + # float64: arccos near 1 loses ~sqrt(eps), which in float32 puts a + # ~1e-4 rad floor under identical rotations; double makes it ~1e-8. + pred = pred.double() + gt = gt.double() + errs = [] + for arm in (rot[:per_arm], rot[per_arm:]): + p, g = pred[..., arm], gt[..., arm] + if per_arm == 3: + Rp, Rg = _ypr_to_matrix(p), _ypr_to_matrix(g) + else: + Rp = _reconstruct_R_from_cols(p[..., 0:3], p[..., 3:6]) + Rg = _reconstruct_R_from_cols(g[..., 0:3], g[..., 3:6]) + tr = (Rp.transpose(-1, -2) @ Rg).diagonal(dim1=-2, dim2=-1).sum(-1) + errs.append(torch.arccos(((tr - 1.0) / 2.0).clamp(-1.0, 1.0))) + return torch.stack(errs, dim=-1).mean().float() + + +def _rot_geodesic_error(pred: torch.Tensor, gt: torch.Tensor): + """Mean geodesic rotation error in radians over batch / time / both arms + of a bimanual cartesian vector. + + Euler-free: per arm, builds proper rotation matrices (``_ypr_to_matrix`` + for the 12/14-dim ypr widths; Gram-Schmidt on the two 6D columns for the + 18/20-dim widths, the same reconstruction the model decode uses) and + takes ``arccos((tr(R_pred^T R_gt) - 1) / 2)``. Unlike a ypr MSE this is + immune to the +-pi wrap AND to the yaw/roll degeneracy at pitch ~ +-pi/2. + Returns None for an unknown width. + """ + layout = bimanual_cartesian_layout(pred.shape[-1]) + if layout is None: + return None + return _geodesic_from_rot_channels(pred, gt, list(layout["rot"])) + + +def cartesian_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: + """Metrics for a bimanual cartesian chunk ``(B, T, D)``; ``{}`` if ``D`` + is not a cartesian width. ``pred``/``gt`` may live on any device.""" + pred_cpu, gt_cpu = pred.detach().cpu(), gt.detach().cpu() + layout = bimanual_cartesian_layout(pred_cpu.shape[-1]) + if layout is None: + return {} + metrics = {} + xyz_idx = list(layout["xyz"]) + xyz_p, rot_p = _split_mse(pred_cpu, gt_cpu) + metrics[f"{prefix}_xyz_paired_mse_avg"] = xyz_p + metrics[f"{prefix}_rot6d_paired_mse_avg"] = rot_p + metrics[f"{prefix}_xyz_final_mse_avg"] = _paired_mse( + pred_cpu[:, -1, xyz_idx], gt_cpu[:, -1, xyz_idx] + ) + metrics.update(_segment_l2(pred_cpu, gt_cpu, xyz_idx, f"{prefix}_xyz_l2")) + # Distributional / alignment metrics on the position channels only + # (metres), so a unitless 6D column never trades off against a metre. + # DTW forgives temporal misalignment (a correct motion executed early / + # late scores near zero where paired MSE penalizes it); Frechet compares + # the time-distribution shape of the motion. + pred_xyz = pred[..., xyz_idx] + gt_xyz = gt[..., xyz_idx].to(pred_xyz.device) + metrics[f"{prefix}_xyz_dtw_avg"] = dtw_distance(pred_xyz, gt_xyz).mean().item() + metrics[f"{prefix}_xyz_frechet_gauss_avg"] = ( + frechet_gaussian_over_time(pred_xyz, gt_xyz).mean().item() + ) + # Rotation error on the manifold, in degrees: what the decode actually + # produces after Gram-Schmidt, free of wrap / gimbal artefacts. + metrics[f"{prefix}_rot_err_deg_avg"] = ( + _rot_geodesic_error(pred_cpu, gt_cpu) * _RAD2DEG + ) + metrics[f"{prefix}_rot_err_deg_final"] = ( + _rot_geodesic_error(pred_cpu[:, -1], gt_cpu[:, -1]) * _RAD2DEG + ) + return metrics + + +def keypoint_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: + """Metrics for a wrist-first bimanual keypoint chunk ``(B, T, D)`` + (D = 138 ypr wrist or 144 rot6d wrist); ``{}`` for other widths. + + ``kp_l2`` is the mean euclidean keypoint error in metres, in whatever + frame the layout holds (the wrist frame for the ``*_wristframe_*`` modes, + head frame otherwise); ``wrist_xyz`` / ``wrist_rot_err_deg`` score the + wrist pose channels; ``kp_dtw`` is DTW over the flattened keypoint block. + """ + pred_cpu, gt_cpu = pred.detach().cpu(), gt.detach().cpu() + layout = bimanual_keypoint_layout(pred_cpu.shape[-1]) + if layout is None: + return {} + kp_idx = list(layout["keypoints"]) + wrist_idx = list(layout["wrist_xyz"]) + rot_idx = list(layout["rot"]) + + def _kp_l2(p, g): + diff = p[..., kp_idx].float() - g[..., kp_idx].float() + return diff.reshape(*diff.shape[:-1], -1, 3).norm(dim=-1).mean() + + metrics = { + f"{prefix}_kp_l2_avg": _kp_l2(pred_cpu, gt_cpu), + f"{prefix}_kp_l2_final": _kp_l2(pred_cpu[:, -1], gt_cpu[:, -1]), + f"{prefix}_wrist_xyz_paired_mse_avg": _paired_mse( + pred_cpu[..., wrist_idx], gt_cpu[..., wrist_idx] + ), + f"{prefix}_wrist_xyz_final_mse_avg": _paired_mse( + pred_cpu[:, -1, wrist_idx], gt_cpu[:, -1, wrist_idx] + ), + f"{prefix}_wrist_rot_err_deg_avg": ( + _geodesic_from_rot_channels(pred_cpu, gt_cpu, rot_idx) * _RAD2DEG + ), + f"{prefix}_wrist_rot_err_deg_final": ( + _geodesic_from_rot_channels(pred_cpu[:, -1], gt_cpu[:, -1], rot_idx) + * _RAD2DEG + ), + } + metrics.update(_segment_l2(pred_cpu, gt_cpu, kp_idx, f"{prefix}_kp_l2")) + metrics.update(_segment_l2(pred_cpu, gt_cpu, wrist_idx, f"{prefix}_wrist_xyz_l2")) + pred_kp = pred[..., kp_idx] + gt_kp = gt[..., kp_idx].to(pred_kp.device) + metrics[f"{prefix}_kp_dtw_avg"] = dtw_distance(pred_kp, gt_kp).mean().item() + return metrics + + +def layout_metrics(pred: torch.Tensor, gt: torch.Tensor, prefix: str) -> dict: + """Dispatch on the last-dim width: cartesian, keypoint, or ``{}``.""" + metrics = cartesian_metrics(pred, gt, prefix) + if metrics: + return metrics + return keypoint_metrics(pred, gt, prefix) diff --git a/egomimic/eval/eval_hpt.py b/egomimic/eval/eval_hpt.py index ee7af2a0e2..c6632e78a5 100644 --- a/egomimic/eval/eval_hpt.py +++ b/egomimic/eval/eval_hpt.py @@ -3,9 +3,10 @@ import torch from torchmetrics import MeanSquaredError +from egomimic.eval.action_metrics import layout_metrics from egomimic.eval.eval_video import EvalVideo from egomimic.rldb.embodiment.embodiment import Embodiment, get_embodiment -from egomimic.utils.metrics import frechet_gaussian_over_time, reverse_kl_from_samples +from egomimic.utils.metrics import frechet_gaussian_over_time class HPTEvalVideo(EvalVideo): @@ -13,14 +14,15 @@ class HPTEvalVideo(EvalVideo): Eval class for HPT models. Per embodiment, computes: - val loss (BC loss, same as training; also aggregated as ``Valid/action_loss``) - paired/final MSE + Frechet over time for the main / shared / auxiliary heads + - the layout metrics of :mod:`egomimic.eval.action_metrics` on the main + ``ac_key`` (position / rotation / keypoint errors, frame-invariant) - paired/final MSE in cam frame on the main ``ac_key``, when a - ``transform_lists`` entry is configured - - optional Reverse KL from samples + ``transform_lists`` entry is configured (viz batches only) The revert transform is applied once and reused by both the cam-frame MSE and the viz video. """ - def compute_metrics_and_viz(self, batch): + def compute_metrics_and_viz(self, batch, do_viz=True): algo = self.model preds = algo.forward_eval(batch) @@ -44,24 +46,20 @@ def compute_metrics_and_viz(self, batch): n_loss_embodiments += 1 if f"{embodiment_name}_{ac_key}" in preds and ac_key != algo.shared_ac_key: - metrics[f"Valid/{embodiment_name}_{ac_key}_paired_mse_avg"] = mse( - preds[f"{embodiment_name}_{ac_key}"].cpu(), _batch[ac_key].cpu() + pred_key = f"{embodiment_name}_{ac_key}" + metrics[f"Valid/{pred_key}_paired_mse_avg"] = mse( + preds[pred_key].cpu(), _batch[ac_key].cpu() ) - metrics[f"Valid/{embodiment_name}_{ac_key}_final_mse_avg"] = mse( - preds[f"{embodiment_name}_{ac_key}"][:, -1].cpu(), + metrics[f"Valid/{pred_key}_final_mse_avg"] = mse( + preds[pred_key][:, -1].cpu(), _batch[ac_key][:, -1].cpu(), ) - fd = frechet_gaussian_over_time( - preds[f"{embodiment_name}_{ac_key}"], _batch[ac_key] - ) - metrics[f"Valid/{embodiment_name}_{ac_key}_frechet_gauss_avg"] = ( - fd.mean().item() - ) - metrics[f"Valid/{embodiment_name}_{ac_key}_frechet_gauss_min"] = ( - fd.min().item() - ) - metrics[f"Valid/{embodiment_name}_{ac_key}_frechet_gauss_max"] = ( - fd.max().item() + fd = frechet_gaussian_over_time(preds[pred_key], _batch[ac_key]) + metrics[f"Valid/{pred_key}_frechet_gauss_avg"] = fd.mean().item() + metrics[f"Valid/{pred_key}_frechet_gauss_min"] = fd.min().item() + metrics[f"Valid/{pred_key}_frechet_gauss_max"] = fd.max().item() + metrics.update( + layout_metrics(preds[pred_key], _batch[ac_key], f"Valid/{pred_key}") ) if embodiment_name in algo.auxiliary_ac_keys: @@ -102,62 +100,8 @@ def compute_metrics_and_viz(self, batch): metrics[f"Valid/{pred_key}_frechet_gauss_min"] = fd.min().item() metrics[f"Valid/{pred_key}_frechet_gauss_max"] = fd.max().item() - if algo.rkl_samples and algo.rkl_samples > 1: - hpt_batch = { - "domain": embodiment_name, - "data": algo._robomimic_to_hpt_data( - batch[embodiment_id], - algo.camera_keys[embodiment_id], - algo.proprio_keys[embodiment_id], - algo.lang_keys[embodiment_id], - ac_key, - algo.auxiliary_ac_keys.get(embodiment_name, []), - ), - } - rkl_targets = [] - - if ( - f"{embodiment_name}_{ac_key}" in preds - and ac_key != algo.shared_ac_key - ): - rkl_targets.append( - ( - f"{embodiment_name}_{ac_key}", - _batch[ac_key].to(algo.device), - embodiment_name, - ) - ) - - if embodiment_name in algo.auxiliary_ac_keys: - for aux_key in algo.auxiliary_ac_keys[embodiment_name]: - aux_pred_key = f"{embodiment_name}_{aux_key}" - if aux_pred_key in preds: - rkl_targets.append( - ( - aux_pred_key, - _batch[aux_key].to(algo.device), - aux_key, - ) - ) - - if algo.shared_ac_key: - shared_pred_key = f"{embodiment_name}_{algo.shared_ac_key}" - if shared_pred_key in preds: - rkl_targets.append( - ( - shared_pred_key, - _batch[algo.shared_ac_key].to(algo.device), - "shared", - ) - ) - - M = int(algo.rkl_samples) - for pred_key_name, gt_tensor, head_key in rkl_targets: - samples = self._collect_policy_samples( - hpt_batch, ref=gt_tensor, key_name=head_key, M=M - ) - rkl = reverse_kl_from_samples(samples, gt_tensor) - metrics[f"Valid/{pred_key_name}_reverse_kl_M{M}"] = rkl.item() + if not do_viz: + continue transform_list = self.transform_lists.get(embodiment_name) main_pred_key = f"{embodiment_name}_{ac_key}" @@ -202,26 +146,3 @@ def _visualize_preds(self, predictions, batch): embodiment_id = batch["embodiment"][0].item() embodiment_name = get_embodiment(embodiment_id).lower() return self.viz_func[embodiment_name](predictions, batch) - - @torch.no_grad() - def _collect_policy_samples(self, hpt_batch, ref, key_name, M): - """Collect policy samples for Reverse KL.""" - algo = self.model - B, T, D = ref.shape - samples = [] - was_training = algo.nets.training - algo.nets.eval() - for _ in range(M): - out = algo.nets["policy"].forward( - hpt_batch["domain"], algo._clone_batch(hpt_batch["data"]) - ) - if key_name in out: - pred = out[key_name] - else: - pred = out[hpt_batch["domain"]] - - pred = pred[:, :T, :D] - samples.append(pred.unsqueeze(0)) - if was_training: - algo.nets.train() - return torch.cat(samples, dim=0) diff --git a/egomimic/eval/eval_pi.py b/egomimic/eval/eval_pi.py index 7206197701..4cf3fda014 100644 --- a/egomimic/eval/eval_pi.py +++ b/egomimic/eval/eval_pi.py @@ -1,29 +1,43 @@ import copy import torch -from torchmetrics import MeanSquaredError +from egomimic.eval.action_metrics import ( # noqa: F401 (re-exported for tests) + _paired_mse, + _rot_geodesic_error, + _split_mse, + layout_metrics, +) from egomimic.eval.eval_video import EvalVideo from egomimic.rldb.embodiment.embodiment import Embodiment, get_embodiment class PIEvalVideo(EvalVideo): """ - Eval class for PI models. Per embodiment, computes: - - val loss (flow-matching loss, same as training; also aggregated as ``Valid/action_loss``) - - paired/final MSE in the model's native wrist frame - - paired/final MSE in cam frame, when a ``transform_lists`` entry is configured - The revert transform is applied once and reused for both the cam-frame MSE - and the viz video. + Eval class for PI models. Per embodiment, on the model's native + (unnormalized) output, computes: + - ``action_loss``: flow-matching val loss, same as training (also logged + per embodiment as ``_loss``) + - the layout metrics of :mod:`egomimic.eval.action_metrics`: + cartesian widths get ``xyz_paired_mse_avg`` / ``xyz_final_mse_avg`` + (m^2), ``rot_err_deg_avg`` / ``rot_err_deg_final`` (geodesic, after + Gram-Schmidt), ``rot6d_paired_mse_avg`` (raw rotation channels, a + training-tracking diagnostic), ``xyz_dtw_avg`` and + ``xyz_frechet_gauss_avg``; keypoint widths get ``kp_l2_avg`` / + ``kp_l2_final`` (m), ``wrist_xyz_*_mse``, ``wrist_rot_err_deg_*`` and + ``kp_dtw_avg``. + + Metrics are frame-invariant (the cam-frame revert is a rigid transform + shared by prediction and gt), so the ``transform_lists`` revert runs only + when rendering the viz video (``do_viz``), never for metrics. """ - def compute_metrics_and_viz(self, batch): + def compute_metrics_and_viz(self, batch, do_viz=True): algo = self.model preds = algo.forward_eval(batch) metrics = {} images_dict = {} - mse = MeanSquaredError() total_loss = None n_loss_embodiments = 0 @@ -43,43 +57,29 @@ def compute_metrics_and_viz(self, batch): n_loss_embodiments += 1 if pred_key in preds: - metrics[f"Valid/{pred_key}_paired_mse_avg"] = mse( - preds[pred_key].cpu(), _batch[ac_key].cpu() + metrics.update( + layout_metrics(preds[pred_key], _batch[ac_key], f"Valid/{pred_key}") ) - metrics[f"Valid/{pred_key}_final_mse_avg"] = mse( - preds[pred_key][:, -1].cpu(), _batch[ac_key][:, -1].cpu() - ) - - transform_list = self.transform_lists.get(embodiment_name) - gt_batch_viz = _batch - preds_for_viz = preds - if transform_list is not None and pred_key in preds: - pred_batch = copy.deepcopy(_batch) - pred_batch[ac_key] = preds[pred_key] - gt_t = Embodiment.apply_transform(_batch, transform_list) - pred_t = Embodiment.apply_transform(pred_batch, transform_list) - # apply_transform drops keys whose shape[0] != batch_size - # (e.g. ``embodiment``, ``annotations``). Merge to preserve them. - gt_batch_viz = {**_batch, **gt_t} - pred_batch_viz = {**_batch, **pred_t} - - # ``.contiguous()`` because ``apply_transform`` returns CPU tensors, - # so ``.cpu()`` here is a no-op and ``[:, -1]`` leaves a non-contiguous - # view that torchmetrics' MSE doesn't accept. - metrics[f"Valid/{pred_key}_cam_paired_mse_avg"] = mse( - pred_batch_viz[ac_key].cpu().contiguous(), - gt_batch_viz[ac_key].cpu().contiguous(), - ) - metrics[f"Valid/{pred_key}_cam_final_mse_avg"] = mse( - pred_batch_viz[ac_key][:, -1].cpu().contiguous(), - gt_batch_viz[ac_key][:, -1].cpu().contiguous(), - ) - - preds_for_viz = dict(preds) - preds_for_viz[pred_key] = pred_batch_viz[ac_key] - ims = self._visualize_preds(preds_for_viz, gt_batch_viz) - images_dict[embodiment_id] = ims + if do_viz: + gt_batch_viz = _batch + preds_for_viz = preds + transform_list = self.transform_lists.get(embodiment_name) + if transform_list is not None and pred_key in preds: + # Revert to cam (head) frame + the ypr layout for the + # overlay only; the metrics above are frame-invariant. + pred_batch = copy.deepcopy(_batch) + pred_batch[ac_key] = preds[pred_key] + gt_t = Embodiment.apply_transform(_batch, transform_list) + pred_t = Embodiment.apply_transform(pred_batch, transform_list) + # apply_transform drops keys whose shape[0] != batch_size + # (e.g. ``embodiment``, ``annotations``). Merge to preserve them. + gt_batch_viz = {**_batch, **gt_t} + pred_batch_viz = {**_batch, **pred_t} + preds_for_viz = dict(preds) + preds_for_viz[pred_key] = pred_batch_viz[ac_key] + ims = self._visualize_preds(preds_for_viz, gt_batch_viz) + images_dict[embodiment_id] = ims if total_loss is not None and n_loss_embodiments > 0: metrics["Valid/action_loss"] = total_loss / n_loss_embodiments diff --git a/egomimic/eval/eval_train_viz.py b/egomimic/eval/eval_train_viz.py new file mode 100644 index 0000000000..3c9771eea3 --- /dev/null +++ b/egomimic/eval/eval_train_viz.py @@ -0,0 +1,65 @@ +"""Train-set visualization evaluator. + +Wraps a concrete EvalVideo (PIEvalVideo, HPTEvalVideo, ...) so the same +forward/metric/viz logic can run a second time against a separate +``train_viz`` dataloader. Videos go to ``/videos_train_viz/`` and +metric keys are prefixed with ``train_viz/`` so they don't collide with the +canonical ``Valid/...`` keys. + +Instantiated via Hydra from a config like +``hydra_configs/evaluator/train_viz_pi_wristframe_6d.yaml`` and enabled with +``+evaluator@train_viz_evaluator=`` on a data config that +defines ``train_viz_datasets``. +""" + +from __future__ import annotations + +import os + +from egomimic.eval.eval_video import EvalVideo + + +class TrainVizEvalVideo(EvalVideo): + def __init__( + self, base: EvalVideo, limit_val_batches: int = 50, prefix: str = "train_viz" + ): + """``prefix`` names the metric-key prefix and the ``videos_/`` + subdirectory; trainHydra reuses this wrapper for the opsplit heads + (``unseen_op_valid``, and ``seen_op_valid`` via ``data.valid_prefix``).""" + # `base` must be set before super().__init__: the trainer/model + # property setters fire on the base attribute during construction. + self.base = base + self.prefix = prefix + super().__init__( + limit_val_batches=limit_val_batches, + viz_func=base.viz_func, + transform_lists=base.transform_lists, + viz_every_n_epochs=base.viz_every_n_epochs, + viz_max_batches=base.viz_max_batches, + ) + + @property + def trainer(self): + return self._trainer + + @trainer.setter + def trainer(self, value): + self._trainer = value + self.base.trainer = value + + @property + def model(self): + return self._model + + @model.setter + def model(self, value): + self._model = value + self.base.model = value + + def video_dir(self): + return os.path.join(self.root_dir(), f"videos_{self.prefix}") + + def compute_metrics_and_viz(self, batch, do_viz=True): + metrics, images_dict = self.base.compute_metrics_and_viz(batch, do_viz=do_viz) + metrics = {f"{self.prefix}/{k}": v for k, v in metrics.items()} + return metrics, images_dict diff --git a/egomimic/eval/eval_video.py b/egomimic/eval/eval_video.py index 9cf44d121c..039e3cba11 100644 --- a/egomimic/eval/eval_video.py +++ b/egomimic/eval/eval_video.py @@ -20,14 +20,24 @@ def __init__( limit_val_batches: int = 400, viz_func: dict = None, transform_lists: dict | None = None, + viz_every_n_epochs: int = 1, + viz_max_batches: int | None = None, ): super().__init__() self.trainer = None self.model = None self.viz_func = viz_func + # Render the overlay video only on validation passes where + # (current_epoch + 1) is a multiple of this; metrics log every pass. + self.viz_every_n_epochs = viz_every_n_epochs + # On viz epochs, render overlay frames for only the first N val + # batches (None = all). Rendering is CPU-bound (~1s/frame) and + # dominates viz-epoch wall time; metrics are still computed on EVERY + # batch regardless. + self.viz_max_batches = viz_max_batches # Per-embodiment list[Transform] applied once during eval to project - # the model's wrist-frame actions back into cam (head) frame. Reused for - # both cam-frame MSE and the viz video so we don't transform twice. + # the model's wrist-frame actions back into cam (head) frame for the + # viz video. self.transform_lists = transform_lists or {} self.val_image_buffer = {} self.val_counter = {} @@ -44,52 +54,80 @@ def __init__( def video_dir(self): return os.path.join(self.root_dir(), "videos") + def _should_viz(self) -> bool: + if not self.viz_every_n_epochs or self.viz_every_n_epochs <= 0: + return False + # Lightning runs validation when (current_epoch + 1) is a multiple of + # check_val_every_n_epoch, with current_epoch still the pre-increment + # value during the val hooks (19, 39, ...). Gate on the same (+1) + # convention: a plain `current_epoch % n` never aligns with a + # validation epoch unless check_val_every_n_epoch == 1. + epoch = self.trainer.current_epoch + 1 + # The last epoch always renders, so short runs (trainer=debug, eval + # mode's single validate) get a video without retuning the interval. + max_epochs = getattr(self.trainer, "max_epochs", None) + if max_epochs is not None and max_epochs > 0 and epoch == max_epochs: + return True + return epoch % self.viz_every_n_epochs == 0 + @abstractmethod - def compute_metrics_and_viz(self, batch): + def compute_metrics_and_viz(self, batch, do_viz=True): """ Run the model's eval forward and compute metrics and visualization frames. Args: batch (dict): processed batch produced by the algo's `process_batch_for_training`. + do_viz (bool): render overlay frames for this batch. When False the + returned ``images_dict`` must be empty. Returns: metrics (dict[str, torch.Tensor | float]) images_dict (dict[embodiment_id, np.ndarray (B, H, W, 3)]) """ raise NotImplementedError + def _video_fps(self, source_fps: int = 30) -> int: + """Playback fps compensating for the DistributedSampler stride. + + Distributed validation deals the (unshuffled) val indices round-robin, + so rank 0's consecutive rendered frames are ``world_size`` source + frames apart. Writing them at the raw source fps produces a + world_size-times timelapse; scale playback down to restore wall-clock + pacing.""" + world = max(1, int(getattr(self.trainer, "world_size", 1) or 1)) + return max(1, round(source_fps / world)) + + def _write_video(self, key, frames) -> None: + path = os.path.join( + self.video_dir(), + f"epoch_{self.trainer.current_epoch}", + str(get_embodiment(key)), + f"validation_video_{self.val_counter[key]}.mp4", + ) + os.makedirs(os.path.dirname(path), exist_ok=True) + tvio.write_video(path, frames, fps=self._video_fps(), video_codec="h264") + def on_validation_start(self): - if self.trainer.is_global_zero: + if self.trainer.is_global_zero and self._should_viz(): os.makedirs( os.path.join(self.video_dir(), f"epoch_{self.trainer.current_epoch}"), exist_ok=True, ) def on_validation_end(self): + if not self._should_viz(): + return for key, buffer in self.val_image_buffer.items(): - os.makedirs( - os.path.join( - self.video_dir(), - f"epoch_{self.trainer.current_epoch}", - str(get_embodiment(key)), - ), - exist_ok=True, - ) if len(buffer) != 0: - frames = torch.stack(buffer) - path = os.path.join( - self.video_dir(), - f"epoch_{self.trainer.current_epoch}", - str(get_embodiment(key)), - f"validation_video_{self.val_counter[key]}.mp4", - ) - tvio.write_video(path, frames, fps=30, video_codec="h264") - + self._write_video(key, torch.stack(buffer)) self.val_counter[key] = 0 self.val_image_buffer[key] = [] def on_validation_step(self, batch, batch_idx, dataloader_idx=0): - metrics, images_dict = self.compute_metrics_and_viz(batch) + do_viz = self._should_viz() and ( + self.viz_max_batches is None or batch_idx < self.viz_max_batches + ) + metrics, images_dict = self.compute_metrics_and_viz(batch, do_viz=do_viz) device = self.trainer.lightning_module.device metrics = { @@ -97,30 +135,23 @@ def on_validation_step(self, batch, batch_idx, dataloader_idx=0): for k, v in metrics.items() } - ## images is now a dict - for key, images in images_dict.items(): - os.makedirs( - os.path.join( - self.video_dir(), - f"epoch_{self.trainer.current_epoch}", - str(get_embodiment(key)), - ), - exist_ok=True, - ) - if key not in self.val_image_buffer or self.val_image_buffer[key] is None: - self.val_image_buffer[key] = [] - self.val_counter[key] = 0 - self.val_image_buffer[key].extend(torch.from_numpy(images)) - if len(self.val_image_buffer[key]) >= 1000: - frames = torch.stack(self.val_image_buffer[key]) - path = os.path.join( - self.video_dir(), - f"epoch_{self.trainer.current_epoch}", - str(get_embodiment(key)), - f"validation_video_{self.val_counter[key]}.mp4", - ) - tvio.write_video(path, frames, fps=30, video_codec="h264") - self.val_image_buffer[key].clear() - self.val_counter[key] += 1 - - self.trainer.lightning_module.log_dict(metrics, sync_dist=True) + if do_viz: + for key, images in images_dict.items(): + if ( + key not in self.val_image_buffer + or self.val_image_buffer[key] is None + ): + self.val_image_buffer[key] = [] + self.val_counter[key] = 0 + self.val_image_buffer[key].extend(torch.from_numpy(images)) + if len(self.val_image_buffer[key]) >= 1000: + self._write_video(key, torch.stack(self.val_image_buffer[key])) + self.val_image_buffer[key].clear() + self.val_counter[key] += 1 + + # add_dataloader_idx=False: with the train_viz loader Lightning would + # otherwise suffix every key with "/dataloader_idx_N"; the train-viz + # wrapper prefixes its keys itself. + self.trainer.lightning_module.log_dict( + metrics, sync_dist=True, add_dataloader_idx=False + ) diff --git a/egomimic/hydra_configs/evaluator/eval_hpt.yaml b/egomimic/hydra_configs/evaluator/eval_hpt.yaml index 901be30749..55e0acf4ec 100644 --- a/egomimic/hydra_configs/evaluator/eval_hpt.yaml +++ b/egomimic/hydra_configs/evaluator/eval_hpt.yaml @@ -1,5 +1,21 @@ +# Default HPT evaluator, matching the per-vendor data defaults: eva on the +# wrist-frame cartesian 6D action (reverted to cam-frame ypr for the overlay), +# human on the 144-D wrist-frame hand-keypoint action (reverted to head-frame +# keypoints for the overlay). +# For cartesian human data use eval_hpt_cartesian / eval_hpt_wrist. defaults: - - viz@viz_func: cartesian + - viz@viz_func: keypoints_wrist - _self_ _target_: egomimic.eval.eval_hpt.HPTEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of +# this; metrics still log on every validation pass. Keep it a multiple of +# trainer.check_val_every_n_epoch or viz never coincides with a validation run. +viz_every_n_epochs: 200 + +transform_lists: + eva_bimanual: + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_cartesian.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_cartesian.yaml new file mode 100644 index 0000000000..17a3ef6d3d --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_cartesian.yaml @@ -0,0 +1,11 @@ +# HPT evaluator for cam/head-frame CARTESIAN ypr data on every embodiment +# (the pre-keypoint default; e.g. model=hpt_bc_flow_human_cartesian with a +# cartesian data override). No revert: the actions are already in the frame +# the overlay draws in. +defaults: + - viz@viz_func: cartesian + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +viz_every_n_epochs: 200 diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_head.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_head.yaml new file mode 100644 index 0000000000..abb99f2f8f --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_head.yaml @@ -0,0 +1,17 @@ +# HPT head-frame keypoint evaluator: use with a data config that applies the +# `keypoints_headframe_6d` transform mode (e.g. data=aria_keypoints). The +# keypoints are already in head frame; the revert only converts the wrist +# rotation rot6d -> ypr (144 -> 138) so the overlay's keypoint splitter sees +# the ypr layout. +defaults: + - viz@viz_func: keypoints + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +viz_every_n_epochs: 100 +viz_max_batches: 10 + +transform_lists: + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_wrist.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_wrist.yaml new file mode 100644 index 0000000000..1549a3d080 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_keypoints_wrist.yaml @@ -0,0 +1,27 @@ +# HPT wrist-frame keypoint evaluator: use with a data config that applies the +# `keypoints_wristframe_6d` transform mode (e.g. +# mecka_fold_freeform_opsplit_hpt_keypoints, or the per-vendor defaults). The +# model predicts the 144-D wrist-first action [wrist xyz+rot6d (9) | 21 +# keypoints (63)] x 2 in each hand's wrist frame; the revert converts the +# wrist rotation back to ypr and projects the two 63-D keypoint chunks back to +# head (camera) frame via the head-frame wrist pose in +# `observations.state.keypoints`, giving a 126-D head-frame +# `actions_keypoints` for the cam-frame keypoint MSE +# (Valid/human_bimanual_actions_keypoints_cam_{paired,final}_mse_avg) and the +# overlay video. Native-frame metrics (paired/final MSE, Frechet, the +# keypoint layout metrics) run on the full 144-D. +defaults: + - viz@viz_func: keypoints_wrist + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of this. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank; metrics still +# run on every batch. The train_viz wrapper inherits this cap. +viz_max_batches: 10 + +transform_lists: + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_hpt_wristframe_6d.yaml b/egomimic/hydra_configs/evaluator/eval_hpt_wristframe_6d.yaml new file mode 100644 index 0000000000..ad998e6498 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_hpt_wristframe_6d.yaml @@ -0,0 +1,24 @@ +# HPT wrist-frame 6D CARTESIAN evaluator: HPT twin of eval_pi_wristframe_6d. +# Use with a data config that applies the `cartesian_wristframe_6d` transform +# mode with `keymap_mode: cartesian` (e.g. mecka_fold_flagship_opsplit_hpt_6d). +# The model predicts the 18-D [xyz | 6D rot] x 2 action in each wrist's frame; +# the revert converts 6D -> ypr (Gram-Schmidt), unpads the 20-D proprio, and +# projects predictions + gt back to head (cam) frame, giving the cam-frame MSE +# (Valid/human_bimanual_actions_cartesian_cam_{paired,final}_mse_avg) and the +# overlay video. Native-frame metrics (paired/final MSE, Frechet, xyz / +# rotation errors) run on the 18-D output. +defaults: + - viz@viz_func: cartesian_wrist_lang + - _self_ + +_target_: egomimic.eval.eval_hpt.HPTEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of this. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank; metrics still +# run on every batch. The train_viz wrapper inherits this cap. +viz_max_batches: 10 + +transform_lists: + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_pi.yaml b/egomimic/hydra_configs/evaluator/eval_pi.yaml index 97b9c1a1f3..62fdd2c39b 100644 --- a/egomimic/hydra_configs/evaluator/eval_pi.yaml +++ b/egomimic/hydra_configs/evaluator/eval_pi.yaml @@ -1,18 +1,30 @@ +# Default Pi evaluator, matching the cotrain_pi_base data defaults: eva on the +# wrist-frame cartesian 6D action, human on the 144-D wrist-frame hand-keypoint +# action. Metrics are computed on the native output (frame-invariant; see +# egomimic/eval/action_metrics.py) and never go through the reverts below, +# which only feed the viz overlay: +# eva: xyz+6D -> xyz+ypr, then each arm's chunk wrist -> cam frame via +# the cam-frame proprio ee_pose +# human: wrist rot6d -> ypr, then each hand's keypoints wrist -> head frame +# via the head-frame wrist pose in observations.state.keypoints +# For the all-cartesian recipes use eval_pi_camframe_6d / eval_pi_wristframe_6d. defaults: - - viz@viz_func: pi_cartesian_lang + - viz@viz_func: pi_keypoints_lang_wrist - _self_ _target_: egomimic.eval.eval_pi.PIEvalVideo -# Per-embodiment revert transform. Applied once during validation to project -# the model's wrist-frame action chunks back to cam (head) frame, then reused -# for both the cam-frame MSE and the viz video. Each value resolves to a -# list[Transform] via its ``_target_``. Must match whichever viz config is -# mounted above. +# Render the val viz video only when (current_epoch + 1) is a multiple of +# this; metrics still log on every validation pass. Keep it a multiple of +# trainer.check_val_every_n_epoch or viz never coincides with a validation +# run. Without it the EvalVideo default of 1 renders video on every pass. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank (~1s/frame CPU +# draw); metrics still run on every batch. +viz_max_batches: 10 + transform_lists: eva_bimanual: - _target_: egomimic.rldb.embodiment.eva._build_eva_bimanual_revert_eef_frame_transform_list - is_quat: false + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list human_bimanual: - _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_eef_frame_transform_list - is_quat: false + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml b/egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml new file mode 100644 index 0000000000..eceaf649de --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml @@ -0,0 +1,29 @@ +defaults: + - viz@viz_func: pi_cartesian_lang + - _self_ + +_target_: egomimic.eval.eval_pi.PIEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of +# this; metrics still log on every validation pass. Keep it a multiple of +# trainer.check_val_every_n_epoch or viz never coincides with a validation +# run. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank; metrics still +# run on every batch. The train_viz wrapper inherits this cap. +viz_max_batches: 10 + +# All-cartesian cam-frame evaluator: use when every embodiment expresses its +# cartesian action in head/camera frame with the continuous 6D rotation +# (`transform_list: mode: cartesian_6d`, e.g. data=mecka_fold_freeform_6d +# with data.train_datasets.human_bimanual.resolver.transform_list.mode=cartesian_6d). +# Actions are already in cam (head) frame, so +# no frame change is needed; the revert only converts the rotation back from +# xyz+6D (9/arm) to xyz+ypr (6/arm) for the viz overlay, which draws in ypr. +# Metrics are computed on the native 6D output and never go through this +# revert. Each value resolves to a list[Transform] via its ``_target_``. +transform_lists: + eva_bimanual: + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_transform_list + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_6d_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_pi_wristframe.yaml b/egomimic/hydra_configs/evaluator/eval_pi_wristframe.yaml new file mode 100644 index 0000000000..62d092ddc1 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_pi_wristframe.yaml @@ -0,0 +1,17 @@ +# Wrist-frame twin of eval_pi (data=cotrain_pi_lang_wrist): eva on the +# wrist-frame cartesian 6D action (reverted 6D -> ypr, then wrist -> cam frame +# via the proprio ee_pose), human on the wrist-frame keypoint action. +defaults: + - viz@viz_func: pi_keypoints_lang_wrist + - _self_ + +_target_: egomimic.eval.eval_pi.PIEvalVideo + +viz_every_n_epochs: 100 +viz_max_batches: 10 + +transform_lists: + eva_bimanual: + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_keypoints_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml b/egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml new file mode 100644 index 0000000000..581d454fc0 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml @@ -0,0 +1,29 @@ +defaults: + - viz@viz_func: pi_cartesian_lang + - _self_ + +_target_: egomimic.eval.eval_pi.PIEvalVideo + +# Render the val viz video only when (current_epoch + 1) is a multiple of +# this (see eval_pi_camframe_6d.yaml); without it the EvalVideo default of 1 +# renders video on every validation pass. +viz_every_n_epochs: 100 +# On viz epochs, render only the first N val batches per rank (~1s/frame +# CPU draw; unbounded rendering added ~2h per viz epoch on the combined +# pretrain). Metrics still run on every batch. The train_viz wrapper +# inherits this cap. +viz_max_batches: 10 + +# All-cartesian wrist-frame evaluator: use with a data config that applies the +# `cartesian_wristframe_6d` transform mode (e.g. mecka_all_6d, +# mecka_fold_freeform*_6d). The model predicts in each wrist's local frame +# using the continuous 6D rotation; these revert transforms first convert the +# rotation xyz+6D -> xyz+ypr and then project predictions + gt back to cam +# (head) frame for the viz video only (metrics are frame-invariant and +# computed on the native output). Each value resolves to a list[Transform] +# via its ``_target_``. Must match the viz config mounted above. +transform_lists: + eva_bimanual: + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/evaluator/train_viz_hpt_keypoints_wrist.yaml b/egomimic/hydra_configs/evaluator/train_viz_hpt_keypoints_wrist.yaml new file mode 100644 index 0000000000..d9631fefd5 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_hpt_keypoints_wrist.yaml @@ -0,0 +1,23 @@ +# Keypoint HPT twin of train_viz_hpt_wristframe_6d: wraps +# eval_hpt_keypoints_wrist so the same forward/metric/viz logic runs a second +# time against the `train_viz` dataloader (metrics `train_viz/`, videos +# `videos_train_viz/`). +# +# Optional: the train_viz head is on by default (train_viz=true) and wraps +# the canonical `evaluator` over `train_viz_datasets`, or the train split +# when the data config declares none. Select this explicitly only to pin a +# different base: +# `+evaluator@train_viz_evaluator=train_viz_hpt_keypoints_wrist` + +# +# `base` is composed separately from the canonical `evaluator`, so CLI +# overrides of `evaluator.viz_every_n_epochs` / `viz_max_batches` do NOT +# reach this head; override `train_viz_evaluator.base.` as well. + +defaults: + - /evaluator@base: eval_hpt_keypoints_wrist + - _self_ + +_target_: egomimic.eval.eval_train_viz.TrainVizEvalVideo + +limit_val_batches: 50 diff --git a/egomimic/hydra_configs/evaluator/train_viz_hpt_wristframe_6d.yaml b/egomimic/hydra_configs/evaluator/train_viz_hpt_wristframe_6d.yaml new file mode 100644 index 0000000000..3cc9e4b13b --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_hpt_wristframe_6d.yaml @@ -0,0 +1,24 @@ +# Wraps eval_hpt_wristframe_6d so the same forward/metric/viz logic runs a +# second time against the `train_viz` dataloader (metrics `train_viz/`, videos +# `videos_train_viz/`). HPT twin of train_viz_pi_wristframe_6d. +# +# Optional: the train_viz head is on by default (train_viz=true) and wraps +# the canonical `evaluator` over `train_viz_datasets`, or the train split +# when the data config declares none. Select this explicitly only to pin a +# different base: +# `+evaluator@train_viz_evaluator=train_viz_hpt_wristframe_6d` +# +# `base` is composed separately from the canonical `evaluator`, so CLI +# overrides of `evaluator.viz_every_n_epochs` / `viz_max_batches` do NOT +# reach this head; override `train_viz_evaluator.base.` as well. + +defaults: + - /evaluator@base: eval_hpt_wristframe_6d + - _self_ + +_target_: egomimic.eval.eval_train_viz.TrainVizEvalVideo + +# Cap on batches consumed from the train_viz loader per validation epoch in +# eval-only mode. During fit-validation the trainer's limit_val_batches +# applies per-dataloader. +limit_val_batches: 50 diff --git a/egomimic/hydra_configs/evaluator/train_viz_pi.yaml b/egomimic/hydra_configs/evaluator/train_viz_pi.yaml new file mode 100644 index 0000000000..9f1bdbd797 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_pi.yaml @@ -0,0 +1,25 @@ +# Wraps the default Pi evaluator (eval_pi) so the same forward/metric/viz +# logic runs a second time against the `train_viz` dataloader. The wrapped +# `base` evaluator owns compute_metrics_and_viz; this wrapper writes videos to +# `videos_train_viz/` and prefixes metric keys with `train_viz/`. +# +# Optional: the train_viz head is on by default (train_viz=true) and wraps +# the canonical `evaluator` over `train_viz_datasets`, or the train split +# when the data config declares none. Select this explicitly only to pin a +# different base: +# `+evaluator@train_viz_evaluator=train_viz_pi` +# +# `base` is composed separately from the canonical `evaluator`, so CLI +# overrides of `evaluator.viz_every_n_epochs` / `viz_max_batches` do NOT +# reach this head; override `train_viz_evaluator.base.` as well. + +defaults: + - /evaluator@base: eval_pi + - _self_ + +_target_: egomimic.eval.eval_train_viz.TrainVizEvalVideo + +# Cap on batches consumed from the train_viz loader per validation epoch in +# eval-only mode. During fit-validation the trainer's limit_val_batches +# applies per-dataloader. +limit_val_batches: 50 diff --git a/egomimic/hydra_configs/evaluator/train_viz_pi_camframe_6d.yaml b/egomimic/hydra_configs/evaluator/train_viz_pi_camframe_6d.yaml new file mode 100644 index 0000000000..5eef381d55 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_pi_camframe_6d.yaml @@ -0,0 +1,26 @@ +# Cam-frame all-cartesian twin of train_viz_pi: wraps eval_pi_camframe_6d so +# the same forward/metric/viz logic runs a second time against the +# `train_viz` dataloader (metrics `train_viz/`, videos `videos_train_viz/`). +# +# Optional: the train_viz head is on by default (train_viz=true) and wraps +# the canonical `evaluator` over `train_viz_datasets`, or the train split +# when the data config declares none. Select this explicitly only to pin a +# different base: +# `+evaluator@train_viz_evaluator=train_viz_pi_camframe_6d` +# Needs the human transform mode overridden to `cartesian_6d`. + +# +# `base` is composed separately from the canonical `evaluator`, so CLI +# overrides of `evaluator.viz_every_n_epochs` / `viz_max_batches` do NOT +# reach this head; override `train_viz_evaluator.base.` as well. + +defaults: + - /evaluator@base: eval_pi_camframe_6d + - _self_ + +_target_: egomimic.eval.eval_train_viz.TrainVizEvalVideo + +# Cap on batches consumed from the train_viz loader per validation epoch in +# eval-only mode. During fit-validation the trainer's limit_val_batches +# applies per-dataloader. +limit_val_batches: 50 diff --git a/egomimic/hydra_configs/evaluator/train_viz_pi_wristframe_6d.yaml b/egomimic/hydra_configs/evaluator/train_viz_pi_wristframe_6d.yaml new file mode 100644 index 0000000000..15863c7cee --- /dev/null +++ b/egomimic/hydra_configs/evaluator/train_viz_pi_wristframe_6d.yaml @@ -0,0 +1,25 @@ +# Wraps the all-cartesian wrist-frame Pi evaluator (eval_pi_wristframe_6d) so +# the same forward/metric/viz logic runs a second time against the +# `train_viz` dataloader (metrics `train_viz/`, videos `videos_train_viz/`). +# +# Optional: the train_viz head is on by default (train_viz=true) and wraps +# the canonical `evaluator` over `train_viz_datasets`, or the train split +# when the data config declares none. Select this explicitly only to pin a +# different base: +# `+evaluator@train_viz_evaluator=train_viz_pi_wristframe_6d` + +# +# `base` is composed separately from the canonical `evaluator`, so CLI +# overrides of `evaluator.viz_every_n_epochs` / `viz_max_batches` do NOT +# reach this head; override `train_viz_evaluator.base.` as well. + +defaults: + - /evaluator@base: eval_pi_wristframe_6d + - _self_ + +_target_: egomimic.eval.eval_train_viz.TrainVizEvalVideo + +# Cap on batches consumed from the train_viz loader per validation epoch in +# eval-only mode. During fit-validation the trainer's limit_val_batches +# applies per-dataloader. +limit_val_batches: 50 diff --git a/egomimic/hydra_configs/evaluator/viz/cartesian_wrist.yaml b/egomimic/hydra_configs/evaluator/viz/cartesian_wrist.yaml index e4aae52a72..9863646876 100644 --- a/egomimic/hydra_configs/evaluator/viz/cartesian_wrist.yaml +++ b/egomimic/hydra_configs/evaluator/viz/cartesian_wrist.yaml @@ -1,12 +1,12 @@ eva_bimanual: _target_: egomimic.rldb.embodiment.eva.Eva.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_cartesian mode: traj human_bimanual: _target_: egomimic.rldb.embodiment.human.Human.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_cartesian mode: traj diff --git a/egomimic/hydra_configs/evaluator/viz/keypoints.yaml b/egomimic/hydra_configs/evaluator/viz/keypoints.yaml index cf5d1d218c..8f000e5ca7 100644 --- a/egomimic/hydra_configs/evaluator/viz/keypoints.yaml +++ b/egomimic/hydra_configs/evaluator/viz/keypoints.yaml @@ -1,12 +1,15 @@ +# Head-frame keypoint overlay: eva draws its cartesian trajectory, human its +# MANO keypoints (the 138-D ypr layout after the evaluator's 6d -> ypr revert, +# or 126-D head-frame keypoints after the wrist-frame revert). eva_bimanual: _target_: egomimic.rldb.embodiment.eva.Eva.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_cartesian mode: traj human_bimanual: _target_: egomimic.rldb.embodiment.human.Human.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_keypoints mode: keypoints diff --git a/egomimic/hydra_configs/evaluator/viz/keypoints_wrist.yaml b/egomimic/hydra_configs/evaluator/viz/keypoints_wrist.yaml index cf5d1d218c..2f92baf889 100644 --- a/egomimic/hydra_configs/evaluator/viz/keypoints_wrist.yaml +++ b/egomimic/hydra_configs/evaluator/viz/keypoints_wrist.yaml @@ -1,12 +1,14 @@ +# Wrist-frame keypoint overlay (same drawing as keypoints.yaml; the evaluator's +# transform_lists revert the wrist-frame prediction to head frame first). eva_bimanual: _target_: egomimic.rldb.embodiment.eva.Eva.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_cartesian mode: traj human_bimanual: _target_: egomimic.rldb.embodiment.human.Human.viz_gt_preds _partial_: true - image_key: front_img_1 + image_key: observations.images.front_img_1 action_key: actions_keypoints mode: keypoints diff --git a/egomimic/hydra_configs/evaluator/viz/pi_keypoints_lang_wrist.yaml b/egomimic/hydra_configs/evaluator/viz/pi_keypoints_lang_wrist.yaml new file mode 100644 index 0000000000..65d79d73a4 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/viz/pi_keypoints_lang_wrist.yaml @@ -0,0 +1,10 @@ +# Pi twin of keypoints_wrist: eva cartesian trajectory + human MANO keypoint +# overlay with the sampled prompt drawn on the frame. +defaults: + - keypoints_wrist + - _self_ + +eva_bimanual: + annotation_key: sampled_prompt +human_bimanual: + annotation_key: sampled_prompt diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml index 196dc7a716..6440e4dc08 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml @@ -32,7 +32,6 @@ robomimic_model: multitask: false pretrained: false pretrained_checkpoint: "" # TODO - reverse_kl_samples: 8 domains: ["eva_bimanual"] shared_obs_keys: ["front_img_1"] diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml index c8daab5660..227cc40e1d 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml @@ -36,8 +36,7 @@ robomimic_model: multitask: false pretrained: false pretrained_checkpoint: "" # TODO - reverse_kl_samples: 8 - + domains: ["human_bimanual"] shared_obs_keys: ["front_img_1"] diff --git a/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml b/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml index b4d1465fcc..9ed1ae6a03 100644 --- a/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_keypoints_base.yaml @@ -9,12 +9,10 @@ robomimic_model: dims: human_bimanual: {proprio: 144, action: 144} - ac_keys: human_bimanual: "actions_keypoints" diffusion: true - reverse_kl_samples: 8 trunk: embed_dim: 256 diff --git a/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml b/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml index 1c10c3480e..7edd9e4904 100644 --- a/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml +++ b/egomimic/hydra_configs/model/hpt_cotrain_enc_dec_base.yaml @@ -13,7 +13,6 @@ robomimic_model: human_bimanual: "actions_cartesian" eva_bimanual: "actions_joints" - reverse_kl_samples: 8 6dof: true trunk: diff --git a/egomimic/hydra_configs/train_zarr_cartesian.yaml b/egomimic/hydra_configs/train_zarr_cartesian.yaml index bfb110d679..0800f35631 100644 --- a/egomimic/hydra_configs/train_zarr_cartesian.yaml +++ b/egomimic/hydra_configs/train_zarr_cartesian.yaml @@ -28,6 +28,17 @@ launch_params: seed: 42 +# Optional pre-fit baseline validation (metrics at epoch 0; the overlay video +# still follows evaluator.viz_every_n_epochs). Off by default; enable with +# `val_at_start=true`. Skipped automatically on requeues/checkpoint resumes. +val_at_start: false + +# Second eval head on TRAIN data at every validation (metrics `train_viz/`, +# videos under `videos_train_viz/`). Wraps the canonical evaluator and reuses +# the train split unless the data config declares `train_viz_datasets` / +# the top-level config a `train_viz_evaluator`. `train_viz=false` turns it off. +train_viz: true + # Normalization and norm-stat cache settings (override per-run as needed). norm_stats: norm_mode: quantile diff --git a/egomimic/hydra_configs/train_zarr_keypoint_wrist.yaml b/egomimic/hydra_configs/train_zarr_keypoint_wrist.yaml index e45cfeea0a..b1571252c1 100644 --- a/egomimic/hydra_configs/train_zarr_keypoint_wrist.yaml +++ b/egomimic/hydra_configs/train_zarr_keypoint_wrist.yaml @@ -2,5 +2,5 @@ defaults: - train_zarr_cartesian - override model: hpt_bc_keypoints_base - override data: aria_keypoints_wrist - - override evaluator/viz@evaluator.viz_func: keypoints_wrist + - override evaluator: eval_hpt_keypoints_wrist - _self_ diff --git a/egomimic/hydra_configs/train_zarr_keypoints.yaml b/egomimic/hydra_configs/train_zarr_keypoints.yaml index 170bffd287..3cbe07a386 100644 --- a/egomimic/hydra_configs/train_zarr_keypoints.yaml +++ b/egomimic/hydra_configs/train_zarr_keypoints.yaml @@ -2,5 +2,5 @@ defaults: - train_zarr_cartesian - override model: hpt_bc_keypoints_base - override data: aria_keypoints - - override evaluator/viz@evaluator.viz_func: keypoints + - override evaluator: eval_hpt_keypoints_head - _self_ diff --git a/egomimic/pl_utils/pl_data_utils.py b/egomimic/pl_utils/pl_data_utils.py index e16a1a1b13..a414760529 100644 --- a/egomimic/pl_utils/pl_data_utils.py +++ b/egomimic/pl_utils/pl_data_utils.py @@ -20,6 +20,12 @@ def __init__( valid_datasets: dict, train_dataloader_params: dict, valid_dataloader_params: dict, + train_viz_datasets: dict | None = None, + train_viz_dataloader_params: dict | None = None, + unseen_op_valid_datasets: dict | None = None, + unseen_op_valid_dataloader_params: dict | None = None, + valid_prefix: str | None = None, + held_out_operators: list | None = None, ): """ Args: @@ -27,10 +33,31 @@ def __init__( valid_datasets: dictionary of valid datasets train_dataloader_params: dictionary of train dataloader parameters valid_dataloader_params: dictionary of valid dataloader parameters + train_viz_datasets: optional dict of datasets iterated like a + second val loader. Used by TrainVizEvalVideo to visualize the + policy on training data alongside the canonical validation. + train_viz_dataloader_params: dict of per-dataset DataLoader kwargs + for the train_viz loader. + unseen_op_valid_datasets: optional dict of datasets iterated as a + third val loader, evaluated by a prefixed copy of the canonical + evaluator (metrics ``unseen_op_valid/``). The held-out-operator split configs put the + unseen operators here, so ``valid_datasets`` can hold the seen + operators' held-out episodes and ``train_viz`` the train split. + unseen_op_valid_dataloader_params: dict of per-dataset DataLoader + kwargs for the unseen_op_valid loader. + valid_prefix: config-only. When set, trainHydra wraps the canonical + evaluator so the valid loader logs ``/...`` and writes + ``videos_/`` (``seen_op_valid`` in the opsplit configs). + held_out_operators: config-only. The held-out-operator split + configs (data/mecka_fold_*_opsplit_*.yaml) keep the operator + id list once at the data-config root and interpolate it from + the train/valid filter lambdas; hydra.instantiate(cfg.data) + forwards every root key here, so it must be accepted. Kept as + an attribute for provenance only. Tokenization (sampling a prompt from per-sample annotation lists, splicing in embodiment / control-mode / proprio blocks, and running - the HF tokenizer) lives on the algo side now — see + the HF tokenizer) lives on the algo side now; see ``PI.process_batch_for_training``. The collate here only stacks tensors and preserves variable-length list-valued keys (e.g. raw ``annotations``) so the algo can consume them downstream. @@ -43,6 +70,16 @@ def __init__( self.valid_datasets = {k: v for k, v in valid_datasets.items() if v is not None} self.train_dataloader_params = train_dataloader_params self.valid_dataloader_params = valid_dataloader_params + self.train_viz_datasets = { + k: v for k, v in (train_viz_datasets or {}).items() if v is not None + } + self.train_viz_dataloader_params = train_viz_dataloader_params or {} + self.unseen_op_valid_datasets = { + k: v for k, v in (unseen_op_valid_datasets or {}).items() if v is not None + } + self.unseen_op_valid_dataloader_params = unseen_op_valid_dataloader_params or {} + self.valid_prefix = valid_prefix + self.held_out_operators = list(held_out_operators or []) self.collate_fn = annotation_collate def train_dataloader(self): @@ -62,13 +99,13 @@ def train_dataloader(self): return CombinedLoader(iterables, "max_size_cycle") - def val_dataloader(self): + def _build_val_style_loader(self, datasets: dict, params: dict, kind: str): iterables = dict() - for dataset_name, dataset in self.valid_datasets.items(): - dataset_params = self.valid_dataloader_params.get(dataset_name) + for dataset_name, dataset in datasets.items(): + dataset_params = params.get(dataset_name) if dataset_params is None or len(dataset_params) == 0: raise ValueError( - f"No dataloader params found for dataset {dataset_name}. Please add {dataset_name} into your data config valid_dataloader_params." + f"No dataloader params found for dataset {dataset_name}. Please add {dataset_name} into your data config {kind}_dataloader_params." ) dataset_params = dict(dataset_params) shuffle = dataset_params.pop("shuffle", False) @@ -78,9 +115,37 @@ def val_dataloader(self): collate_fn=self.collate_fn, **dataset_params, ) - return CombinedLoader(iterables, "max_size_cycle") + def val_loader_names(self) -> list[str]: + """Names of the val loaders in ``val_dataloader()`` order; position i is + Lightning's ``dataloader_idx`` i. ModelWrapper dispatches on these.""" + names = ["valid"] + if self.train_viz_datasets: + names.append("train_viz") + if self.unseen_op_valid_datasets: + names.append("unseen_op_valid") + return names + + def val_dataloader(self): + sources = { + "valid": (self.valid_datasets, self.valid_dataloader_params), + "train_viz": (self.train_viz_datasets, self.train_viz_dataloader_params), + "unseen_op_valid": ( + self.unseen_op_valid_datasets, + self.unseen_op_valid_dataloader_params, + ), + } + names = self.val_loader_names() + loaders = [ + self._build_val_style_loader(*sources[name], kind=name) for name in names + ] + if len(loaders) == 1: + return loaders[0] + # Several heads: return a list so Lightning populates dataloader_idx + # (the position in val_loader_names()) and ModelWrapper can dispatch. + return loaders + def _extract_list_keys(batch): """Pop all list-valued keys from *batch* samples and return them separately. diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index b5657baf7c..3a2cdd9ae6 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/egomimic/pl_utils/pl_model.py @@ -75,6 +75,16 @@ def __init__( self.epoch_memory_stats = [] # Store memory stats per epoch self.evaluator = evaluator + # Optional second evaluator that runs against a train-sampled + # dataloader during the same validation pass (dataloader_idx=1). + # See egomimic/eval/eval_train_viz.py. + self.train_viz_evaluator = None + # Optional third evaluator for the unseen_op_valid loader (held-out + # operators in the opsplit data configs). + self.unseen_op_valid_evaluator = None + # Val loader names in dataloader_idx order, from + # MultiDataModuleWrapper.val_loader_names(); None = [valid, train_viz]. + self.val_loader_names = None @staticmethod def _as_config(cfg): @@ -205,30 +215,56 @@ def on_before_optimizer_step(self, optimizer): sync_dist=True, ) + def _val_heads(self) -> dict: + return { + "valid": self.evaluator, + "train_viz": self.train_viz_evaluator, + "unseen_op_valid": getattr(self, "unseen_op_valid_evaluator", None), + } + def on_validation_start(self): - if self.evaluator is None: + heads = [h for h in self._val_heads().values() if h is not None] + if not heads: return self.model.device = self.device - self.evaluator.on_validation_start() + for head in heads: + head.on_validation_start() def validation_step(self, batch, batch_idx, dataloader_idx=0): """ Run a validation step on the batch, and save that batch of images into the val_image_buffer. Once the buffer hits 1000 images, save that as a 30fps video using torchvision.io.write_video. + + ``dataloader_idx`` is populated by Lightning when ``val_dataloader()`` + returns a list; ``val_loader_names[dataloader_idx]`` picks the head + (valid -> evaluator, train_viz -> train_viz_evaluator, unseen_op_valid -> + unseen_op_valid_evaluator). """ - if self.evaluator is None: + names = getattr(self, "val_loader_names", None) or ["valid", "train_viz"] + name = names[dataloader_idx] if dataloader_idx < len(names) else None + active = self._val_heads().get(name) + if active is None: return + # When val_dataloader returns a list of CombinedLoaders (valid + + # train_viz), Lightning wraps it in an outer sequential CombinedLoader. + # The outer iterator calls next() on each inner CombinedLoader, which + # itself yields a (batch_dict, batch_idx, dataloader_idx) triple; that + # triple lands here as `batch`. Unwrap to recover the dict. + if isinstance(batch, tuple) and len(batch) == 3 and isinstance(batch[0], dict): + batch = batch[0] batch = self.model.process_batch_for_training(batch) print( - f"[VAL_STEP] rank={self.global_rank}, batch_idx={batch_idx}", + f"[VAL_STEP] rank={self.global_rank}, batch_idx={batch_idx}, " + f"dataloader_idx={dataloader_idx}", flush=True, ) - self.evaluator.on_validation_step(batch, batch_idx, dataloader_idx) + active.on_validation_step(batch, batch_idx, dataloader_idx) def on_validation_end(self): print(f"[ON_VALIDATION_END] rank={self.global_rank}", flush=True) - if self.evaluator is not None: - self.evaluator.on_validation_end() + for head in self._val_heads().values(): + if head is not None: + head.on_validation_end() print( f"Rank {self.global_rank} on validation end, waiting for all ranks to synchronize", diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index f43f4630ed..d5cf8787d0 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -108,7 +108,11 @@ def _weights_from_checkpoint(cfg: DictConfig) -> bool: return bool(ckpt_path) and os.path.isfile(ckpt_path) and not cfg.get("pretrained") -def _log_dataset_frame_counts(train_datasets: dict, valid_datasets: dict) -> None: +def _log_dataset_frame_counts( + train_datasets: dict, + valid_datasets: dict, + unseen_op_valid_datasets: dict | None = None, +) -> None: rows = [] for name, ds in train_datasets.items(): rows.append(("train", name, len(ds))) @@ -122,6 +126,8 @@ def _log_dataset_frame_counts(train_datasets: dict, valid_datasets: dict) -> Non rows.append( ("TOTAL", "(valid)", sum(len(ds) for ds in valid_datasets.values())) ) + for name, ds in (unseen_op_valid_datasets or {}).items(): + rows.append(("unseen_op_valid", name, len(ds))) table = tabulate( rows, headers=["Split", "Dataset", "Frames"], @@ -184,6 +190,91 @@ def load_checkpoint( return pl_load(path, map_location=map_location, weights_only=weights_only) +def _resolve_mode(cfg: DictConfig) -> str: + """``mode`` key, else the legacy ``train`` / ``eval`` booleans.""" + if cfg.get("mode") is not None: + return cfg.mode + if cfg.get("train", False): + return "train" + if cfg.get("eval", False): + return "eval" + raise ValueError("Config must specify either `mode` or `train`/`eval` booleans") + + +def _train_viz_enabled(cfg: DictConfig) -> bool: + """The train_viz head (metrics + video on TRAIN data each validation) is on + by default for training runs with an evaluator; ``train_viz=false`` turns + it off, including the heads a data/top-level config declares itself.""" + return ( + bool(cfg.get("train_viz", True)) + and _resolve_mode(cfg) == "train" + and cfg.get("evaluator") is not None + ) + + +def _train_viz_datasets(cfg: DictConfig, train_datasets: dict, instantiate): + """Datasets for the train_viz loader and, when they are derived here, its + loader params (None = leave ``data.train_viz_dataloader_params`` to Hydra). + + A data config's own ``train_viz_datasets`` win. Otherwise the train split + itself is reused (same dataset objects, so no second resolve) through an + unshuffled loader with the valid loader's params.""" + if not _train_viz_enabled(cfg): + return {}, None + explicit = cfg.data.get("train_viz_datasets") + if explicit is not None: + return { + name: None if node is None else instantiate(node, dataset_name=name) + for name, node in explicit.items() + }, None + valid_params = cfg.data.get("valid_dataloader_params") or {} + params = { + name: {**OmegaConf.to_container(valid_params[name]), "shuffle": False} + for name in train_datasets + if name in valid_params + } + return {name: train_datasets[name] for name in params}, params + + +def _build_train_viz_evaluator(cfg: DictConfig): + """``train_viz_evaluator`` if configured, else a TrainVizEvalVideo around a + second instance of the canonical evaluator (own frame buffers).""" + if not _train_viz_enabled(cfg): + return None + if cfg.get("train_viz_evaluator") is not None: + return hydra.utils.instantiate(cfg.train_viz_evaluator) + from egomimic.eval.eval_train_viz import TrainVizEvalVideo + + return TrainVizEvalVideo(hydra.utils.instantiate(cfg.evaluator)) + + +def _unseen_op_valid_datasets(cfg: DictConfig, instantiate) -> dict: + """Datasets for the third (unseen_op_valid) val loader: a data config's own + ``unseen_op_valid_datasets``, in training runs with an evaluator only (eval + mode validates ``valid_datasets`` alone).""" + explicit = cfg.data.get("unseen_op_valid_datasets") + if ( + explicit is None + or _resolve_mode(cfg) != "train" + or cfg.get("evaluator") is None + ): + return {} + return { + name: None if node is None else instantiate(node, dataset_name=name) + for name, node in explicit.items() + } + + +def _build_unseen_op_valid_evaluator(cfg: DictConfig): + """The canonical evaluator (fresh instance, own frame buffers) wrapped to + log ``unseen_op_valid/`` and write ``videos_unseen_op_valid/``.""" + from egomimic.eval.eval_train_viz import TrainVizEvalVideo + + return TrainVizEvalVideo( + hydra.utils.instantiate(cfg.evaluator), prefix="unseen_op_valid" + ) + + @task_wrapper def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Trains the model. Can additionally evaluate on a testset, using best weights obtained during @@ -205,8 +296,9 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: load_env() - # One SQL pull / path resolution per dataset spec across train, valid and - # the norm-stat copies; dropped on exit so nothing outlives this run. + # One SQL pull / path resolution per dataset spec across train, valid, + # train_viz and the norm-stat copies; dropped on exit so nothing outlives + # this run. with resolve_once(): train_datasets = {} for dataset_name in cfg.data.train_datasets: @@ -220,12 +312,30 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: cfg.data.valid_datasets[dataset_name], dataset_name=dataset_name ) + train_viz_datasets, train_viz_params = _train_viz_datasets( + cfg, train_datasets, instantiate=_instantiate_dataset + ) + unseen_op_valid_datasets = _unseen_op_valid_datasets( + cfg, instantiate=_instantiate_dataset + ) + log.info(f"Instantiating datamodule <{cfg.data._target_}>") assert ( "MultiDataModuleWrapper" in cfg.data._target_ ), "cfg.data._target_ must be 'MultiDataModuleWrapper'" + datamodule_kwargs = dict( + train_datasets=train_datasets, valid_datasets=valid_datasets + ) + # Always passed, so train_viz=false also drops a data config's own + # train_viz_datasets. + datamodule_kwargs["train_viz_datasets"] = train_viz_datasets + if train_viz_params is not None: + datamodule_kwargs["train_viz_dataloader_params"] = train_viz_params + # Always passed too: hydra would otherwise instantiate the raw + # config nodes itself (eval mode included). + datamodule_kwargs["unseen_op_valid_datasets"] = unseen_op_valid_datasets datamodule: LightningDataModule = hydra.utils.instantiate( - cfg.data, train_datasets=train_datasets, valid_datasets=valid_datasets + cfg.data, **datamodule_kwargs ) # Stats-only MultiDataset (no graph of its own; explicitly populated from @@ -321,6 +431,10 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: ds.set_norm_stats_from(norm_stats) for ds in datamodule.valid_datasets.values(): ds.set_norm_stats_from(norm_stats) + for ds in getattr(datamodule, "train_viz_datasets", {}).values(): + ds.set_norm_stats_from(norm_stats) + for ds in getattr(datamodule, "unseen_op_valid_datasets", {}).values(): + ds.set_norm_stats_from(norm_stats) _prepare_checkpoint_resume(cfg) @@ -331,20 +445,16 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: scheduler_interval=cfg.model.get("scheduler_interval", "step"), ) - _log_dataset_frame_counts(datamodule.train_datasets, datamodule.valid_datasets) + _log_dataset_frame_counts( + datamodule.train_datasets, + datamodule.valid_datasets, + datamodule.unseen_op_valid_datasets, + ) log.info("Instantiating callbacks...") callbacks: List[Callback] = instantiate_callbacks(cfg.get("callbacks")) - # Resolve mode: support both new `mode` key and legacy `train`/`eval` booleans - if cfg.get("mode") is not None: - mode = cfg.mode - elif cfg.get("train", False): - mode = "train" - elif cfg.get("eval", False): - mode = "eval" - else: - raise ValueError("Config must specify either `mode` or `train`/`eval` booleans") + mode = _resolve_mode(cfg) # In eval mode, apply trainer overrides from the eval object and disable logger if mode == "eval": @@ -394,9 +504,40 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: if mode == "train": if cfg.get("evaluator") is not None: eval_obj: Eval = hydra.utils.instantiate(cfg.evaluator) + # data.valid_prefix (e.g. seen_op_valid in the opsplit configs) + # renames the canonical head: `/Valid/...` metrics and + # `videos_/` instead of `Valid/...` and `videos/`. + valid_prefix = cfg.data.get("valid_prefix") + if valid_prefix: + from egomimic.eval.eval_train_viz import TrainVizEvalVideo + + eval_obj = TrainVizEvalVideo(eval_obj, prefix=valid_prefix) eval_obj.trainer = trainer eval_obj.model = model.model model.evaluator = eval_obj + train_viz_eval_obj = ( + _build_train_viz_evaluator(cfg) if datamodule.train_viz_datasets else None + ) + if train_viz_eval_obj is not None: + train_viz_eval_obj.trainer = trainer + train_viz_eval_obj.model = model.model + model.train_viz_evaluator = train_viz_eval_obj + if datamodule.unseen_op_valid_datasets: + unseen_eval_obj = _build_unseen_op_valid_evaluator(cfg) + unseen_eval_obj.trainer = trainer + unseen_eval_obj.model = model.model + model.unseen_op_valid_evaluator = unseen_eval_obj + model.val_loader_names = datamodule.val_loader_names() + # Pre-fit baseline val. Skipped on requeues AND checkpoint resumes: + # trainer.validate here runs BEFORE fit restores ckpt_path weights, so + # on a resume it would score the un-resumed base model. + # (_prepare_checkpoint_resume already folded a requeue into ckpt_path.) + if cfg.get("val_at_start", False) and not cfg.get("ckpt_path"): + log.info( + "val_at_start: running validation at epoch 0 (pre-fit baseline; " + "the overlay video follows evaluator.viz_every_n_epochs)" + ) + trainer.validate(model=model, datamodule=datamodule) log.info("Starting training!") trainer.fit( model=model, diff --git a/egomimic/utils/metrics.py b/egomimic/utils/metrics.py index aa6c04756c..8b1ace86c0 100644 --- a/egomimic/utils/metrics.py +++ b/egomimic/utils/metrics.py @@ -1,33 +1,7 @@ -"""Distribution-distance metrics used by evaluation. - -Moved out of egomimicUtils; code unchanged. -""" - -import math +"""Distribution-distance metrics used by evaluation.""" import torch -# ---- moved from egomimicUtils.py (code unchanged) ---- - - -def reverse_kl_from_samples(pred_samples, targets): - M, B, T, D = pred_samples.shape - - TD = T * D - const = -0.5 * TD * math.log(2.0 * math.pi) - - A = pred_samples.permute(1, 0, 2, 3).reshape(B, M, TD) # (B,M,TD) - MU = targets.reshape(B, 1, TD) # (B,1,TD) - - d2 = torch.cdist(A, A).pow(2) # (B,M,M) - log_q_each = torch.logsumexp(const - 0.5 * d2, dim=-1) - math.log(M) # (B,M) - - d2p = ((A - MU) ** 2).sum(dim=-1) # (B,M) - log_p_each = const - 0.5 * d2p # (B,M) - - rkl_each = (log_q_each - log_p_each).mean(dim=-1) # (B,) - return rkl_each.mean() - def frechet_gaussian_over_time( pred: torch.Tensor, @@ -112,3 +86,39 @@ def frechet_gaussian_over_time( "max": dist.max().item(), } return dist + + +def dtw_distance(pred: torch.Tensor, tgt: torch.Tensor, normalize: bool = True): + """Batched dynamic-time-warping distance between (B, T1, D) and (B, T2, D) + trajectories. + + Local cost is the euclidean distance between D-dim frames; the classic + DTW recurrence runs as a wavefront over the T1+T2-1 anti-diagonals of the + cost grid so the whole batch advances in vectorized steps (no per-cell + python loop). Unlike paired MSE, DTW forgives temporal misalignment: a + correct trajectory executed slightly early/late scores near zero. + + Args: + pred/tgt: (B, T, D) action chunks (any device). + normalize: divide the accumulated path cost by (T1 + T2), giving a + per-step average frame distance comparable across chunk lengths. + Returns: + (B,) distances. + """ + B, T1, _ = pred.shape + _, T2, _ = tgt.shape + C = torch.cdist(pred.float(), tgt.float()) # (B, T1, T2) + acc = torch.full((B, T1 + 1, T2 + 1), float("inf"), device=C.device, dtype=C.dtype) + acc[:, 0, 0] = 0.0 + for k in range(2, T1 + T2 + 1): + i = torch.arange(max(1, k - T2), min(T1, k - 1) + 1, device=C.device) + j = k - i + prev = torch.minimum( + torch.minimum(acc[:, i - 1, j], acc[:, i, j - 1]), + acc[:, i - 1, j - 1], + ) + acc[:, i, j] = C[:, i - 1, j - 1] + prev + out = acc[:, T1, T2] + if normalize: + out = out / (T1 + T2) + return out diff --git a/tests/unit/test_eva_wristframe_norm_pool.py b/tests/unit/test_eva_wristframe_norm_pool.py index 31bf3a670c..18b16a3f7b 100644 --- a/tests/unit/test_eva_wristframe_norm_pool.py +++ b/tests/unit/test_eva_wristframe_norm_pool.py @@ -6,6 +6,7 @@ from fixtures.synthetic_episodes import write_episode from hydra import compose, initialize_config_module +from egomimic.eval.action_metrics import cartesian_metrics, keypoint_metrics from egomimic.rldb.embodiment.eva import Eva from egomimic.rldb.zarr import norm_cache from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset @@ -233,3 +234,39 @@ def test_bounds_check_ignores_constant_cells(): assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is None arr[1, 0] = 50.0 # corrupt value at a regular cell: still rejected assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is not None + + +EVA_WRIST_REVERT = "egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list" + + +@pytest.mark.parametrize("top", ["train_zarr_cartesian", "train_zarr_cartesian_pi"]) +def test_eva_defaults_are_wrist_frame_6d(top): + cfg = _compose(top) + eva = cfg.data.train_datasets.eva_bimanual + assert eva.resolver.transform_list.mode == "cartesian_wristframe_6d" + assert cfg.evaluator.transform_lists.eva_bimanual._target_ == EVA_WRIST_REVERT + assert cfg.norm_stats.pool_horizon is False + + +def test_segment_l2_separates_early_from_late_error(): + B, T = 2, 100 + gt = torch.zeros(B, T, 20) + pred = gt.clone() + pred[:, :10, 0] = 0.03 # 3 cm on the left arm x, first 10 % only + m = cartesian_metrics(pred, gt, "p") + # mean over both arms: 3 cm on one of two + assert m["p_xyz_l2_early"].item() == pytest.approx(0.015) + assert m["p_xyz_l2_mid"].item() == pytest.approx(0.0) + assert m["p_xyz_l2_late"].item() == pytest.approx(0.0) + + +def test_segment_l2_on_keypoints(): + B, T = 1, 100 + gt = torch.zeros(B, T, 144) + pred = gt.clone() + pred[:, 50:, 9:72] = 0.02 # every left-hand keypoint off by (2,2,2) cm late + m = keypoint_metrics(pred, gt, "k") + expected = 0.5 * float(np.linalg.norm([0.02] * 3)) # one hand of two + assert m["k_kp_l2_late"].item() == pytest.approx(expected, rel=1e-5) + assert m["k_kp_l2_early"].item() == pytest.approx(0.0) + assert m["k_wrist_xyz_l2_late"].item() == pytest.approx(0.0) diff --git a/tests/unit/test_hydra_configs.py b/tests/unit/test_hydra_configs.py index 0cdf0de78b..66e34e3856 100644 --- a/tests/unit/test_hydra_configs.py +++ b/tests/unit/test_hydra_configs.py @@ -31,13 +31,20 @@ def _options(group: str) -> list[str]: TRAIN_TOP_LEVEL = [c for c in TOP_LEVEL if c.startswith("train_")] -@pytest.mark.parametrize("config_name", TOP_LEVEL) +# train_* configs are composed (with more checks) by +# test_infra_port_wiring.py::test_human_action_width_matches_data_mode. +@pytest.mark.parametrize( + "config_name", [c for c in TOP_LEVEL if c not in TRAIN_TOP_LEVEL] +) def test_top_level_config_resolves(config_name, compose_resolve): compose_resolve(config_name, []) @pytest.mark.parametrize("data", _options("data")) -@pytest.mark.parametrize("top", TRAIN_TOP_LEVEL) +# train_zarr_cartesian x data is composed by test_data_filters_instantiate_to_dataset_filter. +@pytest.mark.parametrize( + "top", [t for t in TRAIN_TOP_LEVEL if t != "train_zarr_cartesian"] +) def test_data_option_resolves_under_every_train_config(top, data, compose_resolve): compose_resolve(top, [f"data={data}"]) diff --git a/tests/unit/test_infra_port_wiring.py b/tests/unit/test_infra_port_wiring.py index 549c03acb0..d9559b41bb 100644 --- a/tests/unit/test_infra_port_wiring.py +++ b/tests/unit/test_infra_port_wiring.py @@ -10,6 +10,7 @@ import numpy as np import pytest import torch +from lightning import LightningModule from scipy.spatial.transform import Rotation as R import egomimic @@ -360,3 +361,190 @@ def check(v): assert check(1.0 + 0.4 * 2.0) is None # within 50 % of the [-1, 1] range assert check(1.0 + 0.6 * 2.0) is not None assert check(1e9) is not None + + +# ------------------------------------------------------ train_viz second loader +class _Dicts: + def __init__(self, n, tag): + self.n, self.tag = n, tag + + def __len__(self): + return self.n + + def __getitem__(self, i): + return {"x": np.float32(i), "tag": [self.tag]} + + +def test_val_dataloader_returns_two_loaders_with_train_viz(): + from lightning.pytorch.utilities.combined_loader import CombinedLoader + + from egomimic.pl_utils.pl_data_utils import MultiDataModuleWrapper + + params = {"human_bimanual": {"batch_size": 2, "num_workers": 0}} + dm = MultiDataModuleWrapper( + train_datasets={"human_bimanual": _Dicts(4, "train")}, + valid_datasets={"human_bimanual": _Dicts(4, "valid")}, + train_dataloader_params=params, + valid_dataloader_params=params, + held_out_operators=["op1"], + ) + assert isinstance(dm.val_dataloader(), CombinedLoader) + dm = MultiDataModuleWrapper( + train_datasets={"human_bimanual": _Dicts(4, "train")}, + valid_datasets={"human_bimanual": _Dicts(4, "valid")}, + train_dataloader_params=params, + valid_dataloader_params=params, + train_viz_datasets={"human_bimanual": _Dicts(4, "viz")}, + train_viz_dataloader_params={ + "human_bimanual": {**params["human_bimanual"], "shuffle": False} + }, + ) + loaders = dm.val_dataloader() + assert isinstance(loaders, list) and len(loaders) == 2 + batch, _, _ = next(iter(loaders[1])) + assert batch["human_bimanual"]["tag"] == [["viz"], ["viz"]] + + +def test_validation_step_routes_by_dataloader_idx(): + from egomimic.pl_utils.pl_model import ModelWrapper + + calls = [] + + class _Ev: + def __init__(self, name): + self.name = name + + def on_validation_step(self, batch, batch_idx, dataloader_idx=0): + calls.append((self.name, batch, batch_idx, dataloader_idx)) + + w = ModelWrapper.__new__(ModelWrapper) + LightningModule.__init__(w) # global_rank etc. without a trainer + w.model = SimpleNamespace(process_batch_for_training=lambda b: {"processed": b}) + w.evaluator = _Ev("valid") + w.train_viz_evaluator = _Ev("train_viz") + # Lightning hands the outer sequential loader's inner triple as `batch` + w.validation_step(({"k": 1}, 0, 0), 0, dataloader_idx=0) + w.validation_step(({"k": 2}, 0, 1), 3, dataloader_idx=1) + w.validation_step({"k": 3}, 1, dataloader_idx=0) + assert calls == [ + ("valid", {"processed": {"k": 1}}, 0, 0), + ("train_viz", {"processed": {"k": 2}}, 3, 1), + ("valid", {"processed": {"k": 3}}, 1, 0), + ] + w.train_viz_evaluator = None + w.validation_step({"k": 4}, 0, dataloader_idx=1) # no second head: ignored + assert len(calls) == 3 + + +def test_unseen_op_valid_third_loader_and_routing(): + """unseen_op_valid_datasets adds a loader after train_viz; without train_viz it + takes idx 1, and ModelWrapper routes by val_loader_names, not position.""" + from egomimic.pl_utils.pl_data_utils import MultiDataModuleWrapper + from egomimic.pl_utils.pl_model import ModelWrapper + + params = {"human_bimanual": {"batch_size": 2, "num_workers": 0}} + common = dict( + train_datasets={"human_bimanual": _Dicts(4, "train")}, + valid_datasets={"human_bimanual": _Dicts(4, "valid")}, + train_dataloader_params=params, + valid_dataloader_params=params, + unseen_op_valid_datasets={"human_bimanual": _Dicts(4, "unseen")}, + unseen_op_valid_dataloader_params=params, + ) + dm = MultiDataModuleWrapper( + **common, + train_viz_datasets={"human_bimanual": _Dicts(4, "viz")}, + train_viz_dataloader_params=params, + ) + assert dm.val_loader_names() == ["valid", "train_viz", "unseen_op_valid"] + loaders = dm.val_dataloader() + tags = [next(iter(ld))[0]["human_bimanual"]["tag"][0][0] for ld in loaders] + assert tags == ["valid", "viz", "unseen"] + + dm = MultiDataModuleWrapper(**common) + assert dm.val_loader_names() == ["valid", "unseen_op_valid"] + + calls = [] + + class _Ev: + def __init__(self, name): + self.name = name + + def on_validation_step(self, batch, batch_idx, dataloader_idx=0): + calls.append((self.name, dataloader_idx)) + + w = ModelWrapper.__new__(ModelWrapper) + LightningModule.__init__(w) + w.model = SimpleNamespace(process_batch_for_training=lambda b: b) + w.evaluator = _Ev("valid") + w.train_viz_evaluator = None + w.unseen_op_valid_evaluator = _Ev("unseen_op_valid") + w.val_loader_names = dm.val_loader_names() + w.validation_step({"k": 1}, 0, dataloader_idx=0) + w.validation_step({"k": 2}, 0, dataloader_idx=1) + assert calls == [("valid", 0), ("unseen_op_valid", 1)] + + +# ------------------------------------------------------ train_viz on by default +def _viz_cfg(**top): + from omegaconf import OmegaConf + + data = top.pop("data", {}) + cfg = { + "mode": "train", + "evaluator": { + "_target_": "egomimic.eval.eval_hpt.HPTEvalVideo", + "viz_every_n_epochs": 7, + }, + "data": { + "valid_dataloader_params": {HUMAN: {"batch_size": 4, "num_workers": 3}}, + **data, + }, + **top, + } + return OmegaConf.create(cfg) + + +def test_train_viz_defaults_to_the_train_datasets(): + import egomimic.trainHydra as th + + train = {HUMAN: object()} + viz, params = th._train_viz_datasets(_viz_cfg(), train, instantiate=None) + assert viz[HUMAN] is train[HUMAN], "reuses the train split, no second resolve" + # valid loader params, never shuffled (same leading slice every val) + assert params == {HUMAN: {"batch_size": 4, "num_workers": 3, "shuffle": False}} + + +def test_train_viz_explicit_datasets_and_opt_outs(): + import egomimic.trainHydra as th + + train = {HUMAN: object()} + explicit = _viz_cfg( + data={ + "train_viz_datasets": {HUMAN: {"tag": "explicit"}}, + "train_viz_dataloader_params": {HUMAN: {"batch_size": 32}}, + } + ) + viz, params = th._train_viz_datasets( + explicit, train, instantiate=lambda node, dataset_name: ("built", node.tag) + ) + assert viz == {HUMAN: ("built", "explicit")} + assert params is None, "the data config's own params are left to Hydra" + + for off in ( + _viz_cfg(train_viz=False), + _viz_cfg(mode="eval"), + _viz_cfg(evaluator=None), + ): + assert th._train_viz_datasets(off, train, instantiate=None) == ({}, None) + + +def test_train_viz_evaluator_wraps_the_canonical_evaluator(): + import egomimic.trainHydra as th + from egomimic.eval.eval_train_viz import TrainVizEvalVideo + + cfg = _viz_cfg() + ev = th._build_train_viz_evaluator(cfg) + assert isinstance(ev, TrainVizEvalVideo) + assert ev.viz_every_n_epochs == 7 + assert th._build_train_viz_evaluator(_viz_cfg(train_viz=False)) is None diff --git a/tests/unit/test_pi05_norm_rot6d.py b/tests/unit/test_pi05_norm_rot6d.py index 6a1dce7ef5..6c69eaae34 100644 --- a/tests/unit/test_pi05_norm_rot6d.py +++ b/tests/unit/test_pi05_norm_rot6d.py @@ -5,10 +5,14 @@ 20D 6D, with gripper) and human bimanual (12D ypr / 18D 6D, no gripper) layouts, plus the proprio ee_pose: the 6D transform modes convert the proprio too (a single pose vector, same per-arm layout as one action row), and the 6D revert -lists convert it back before the eef-frame revert reads it. +lists convert it back before the eef-frame revert reads it. Also the pieces +that ride along with the 6D pipeline: the rotation-free bounds check, the +mecka left-wrist fix, embodiment aliasing, the sampler fallback and the eval +helpers. """ import math +from types import SimpleNamespace import numpy as np import pytest @@ -24,6 +28,7 @@ RobotBimanualCartesianEuler, pad_to_width, ) +from egomimic.utils.pose_utils import _ypr_to_rot6d def _eva_ypr_chunk(T: int = 5) -> np.ndarray: @@ -103,6 +108,29 @@ def test_robot_bimanual_norm_6d_pack_round_trips(): torch.testing.assert_close(decoded, six6d, atol=1e-6, rtol=1e-6) +def test_keypoint_converter_is_identity_and_pads_to_model_width(): + # The 144-D keypoint action is packed identity-first; PI pads it (and the + # cartesian 32-slot layout) up to model.action_dim and the decode slices + # the native width back out of the wider vector. + conv = HumanBimanualKeypoints() + kp = torch.randn(2, 4, 144) + packed = conv.to32_norm_6d(kp) + torch.testing.assert_close(packed, kp) + torch.testing.assert_close(conv.from32_norm_6d(pad_to_width(packed, 160)), kp) + with pytest.raises(ValueError, match="expected 144-dim"): + conv.to32_norm_6d(torch.zeros(1, 1, 138)) + + eva32 = RobotBimanualCartesianEuler().to32_norm_6d(torch.ones(1, 3, 20)) + wide = pad_to_width(eva32, 144) + assert wide.shape[-1] == 144 and torch.all(wide[..., 32:] == 0) + torch.testing.assert_close( + RobotBimanualCartesianEuler().from32_norm_6d(wide)[..., :20], + torch.ones(1, 3, 20), + ) + with pytest.raises(ValueError, match="action_dim is 32"): + pad_to_width(kp, 32) + + @pytest.mark.parametrize( "chunk_fn,ypr_dim,six_dim", [(_eva_ypr_chunk, 14, 20), (_aria_ypr_chunk, 12, 18)], @@ -186,157 +214,6 @@ def test_6d_revert_lists_revert_proprio(): }, f"{build.__name__} must revert both action and proprio to ypr" -def test_vendor_embodiment_names_collapse_to_human(): - # Mirror episodes written by the vendor-split registry carry names like - # MECKA_BIMANUAL in their zarr metadata; locally all human demo data is - # one embodiment, so these must resolve to the HUMAN_* ids. - from egomimic.rldb.embodiment.embodiment import ( - EMBODIMENT, - get_embodiment_id, - is_legacy_vendor_embodiment, - ) - - for vendor in ("mecka", "scale", "aria", "lightwheel"): - assert ( - get_embodiment_id(f"{vendor}_bimanual") == EMBODIMENT.HUMAN_BIMANUAL.value - ) - assert ( - get_embodiment_id(f"{vendor}_right_arm") == EMBODIMENT.HUMAN_RIGHT_ARM.value - ) - assert ( - get_embodiment_id(f"{vendor}_left_arm") == EMBODIMENT.HUMAN_LEFT_ARM.value - ) - assert is_legacy_vendor_embodiment(f"{vendor.upper()}_BIMANUAL") - assert get_embodiment_id("human_bimanual") == EMBODIMENT.HUMAN_BIMANUAL.value - assert get_embodiment_id("eva_bimanual") == EMBODIMENT.EVA_BIMANUAL.value - assert not is_legacy_vendor_embodiment("human_bimanual") - with pytest.raises(KeyError): - get_embodiment_id("yam_bimanual") # robot names are never aliased - - -def test_base_converter_rejects_norm_6d_encoding(): - converter = BaseActionConverter() - with pytest.raises(NotImplementedError, match="normalized-rot6d"): - converter.to32_norm_6d(torch.zeros(1, 1, 20)) - with pytest.raises(NotImplementedError, match="normalized-rot6d"): - converter.from32_norm_6d(torch.zeros(1, 1, 32)) - - -def test_unpad_gripper_zeros_inverts_pad_and_noops_unpadded(): - from egomimic.rldb.zarr.action_chunk_transforms import ( - PadGripperZeros, - UnpadGripperZeros, - ) - - rng = np.random.default_rng(6) - for width in (12, 18): - v = rng.uniform(-1, 1, size=(width,)) - padded = PadGripperZeros(action_key="k").transform({"k": v.copy()})["k"] - assert padded.shape == (width + 2,) - back = UnpadGripperZeros(action_key="k").transform({"k": padded.copy()})["k"] - np.testing.assert_allclose(back, v) - # no-op on already-unpadded widths - same = UnpadGripperZeros(action_key="k").transform({"k": v.copy()})["k"] - np.testing.assert_allclose(same, v) - - -def test_human_6d_reverts_unpad_proprio(): - from egomimic.rldb.embodiment.human import ( - _build_human_cartesian_revert_6d_transform_list, - _build_human_cartesian_revert_6d_wristframe_transform_list, - ) - from egomimic.rldb.zarr.action_chunk_transforms import UnpadGripperZeros - - for build in ( - _build_human_cartesian_revert_6d_transform_list, - _build_human_cartesian_revert_6d_wristframe_transform_list, - ): - assert any(isinstance(t, UnpadGripperZeros) for t in build()), build.__name__ - - -def test_rotate_local_frame_flips_left_wrist_convention(): - # Right-multiplying by Rz(180 deg) must flip the pose's own x/y axes, keep - # z (knuckle-forward) and the position, skip zero-quat padding rows, and - # handle both (7,) poses and (T, 7) chunks. - from scipy.spatial.transform import Rotation as R - - from egomimic.rldb.zarr.action_chunk_transforms import RotateLocalFrame - - rng = np.random.default_rng(4) - q = R.random(3, random_state=5) - chunk = np.zeros((4, 7)) - chunk[:3, :3] = rng.uniform(-1, 1, size=(3, 3)) - chunk[:3, 3:] = q.as_quat()[:, [3, 0, 1, 2]] # wxyz; row 3 stays zero-padded - - t = RotateLocalFrame(keys=["k"]) - out = t.transform({"k": chunk.copy()})["k"] - - np.testing.assert_allclose(out[:, :3], chunk[:, :3]) # positions unchanged - np.testing.assert_allclose(out[3], np.zeros(7)) # padding untouched - R_old = q.as_matrix() - R_new = R.from_quat(out[:3, [4, 5, 6, 3]]).as_matrix() - np.testing.assert_allclose(R_new[:, :, 0], -R_old[:, :, 0], atol=1e-12) # x flip - np.testing.assert_allclose(R_new[:, :, 1], -R_old[:, :, 1], atol=1e-12) # y flip - np.testing.assert_allclose(R_new[:, :, 2], R_old[:, :, 2], atol=1e-12) # z kept - - single = t.transform({"k": chunk[0].copy()})["k"] - np.testing.assert_allclose(single, out[0], atol=1e-12) - - -def test_keypoint_converter_is_identity_and_pads_to_model_width(): - # The 144-D keypoint action is packed identity-first; PI pads it (and the - # cartesian 32-slot layout) up to model.action_dim and the decode slices - # the native width back out of the wider vector. - conv = HumanBimanualKeypoints() - kp = torch.randn(2, 4, 144) - packed = conv.to32_norm_6d(kp) - torch.testing.assert_close(packed, kp) - torch.testing.assert_close(conv.from32_norm_6d(pad_to_width(packed, 160)), kp) - with pytest.raises(ValueError, match="expected 144-dim"): - conv.to32_norm_6d(torch.zeros(1, 1, 138)) - - eva32 = RobotBimanualCartesianEuler().to32_norm_6d(torch.ones(1, 3, 20)) - wide = pad_to_width(eva32, 144) - assert wide.shape[-1] == 144 and torch.all(wide[..., 32:] == 0) - torch.testing.assert_close( - RobotBimanualCartesianEuler().from32_norm_6d(wide)[..., :20], - torch.ones(1, 3, 20), - ) - with pytest.raises(ValueError, match="action_dim is 32"): - pad_to_width(kp, 32) - - -def test_fix_left_wrist_convention_flag_prepends_correction(): - from egomimic.rldb.embodiment.human import Human - from egomimic.rldb.zarr.action_chunk_transforms import RotateLocalFrame - - tl = Human.get_transform_list( - "cartesian_wristframe_6d", stride=1, fix_left_wrist_convention=True - ) - assert isinstance(tl[0], RotateLocalFrame) - assert set(tl[0].keys) == {"left.action_ee_pose", "left.obs_ee_pose"} - # default off: other vendors' data must be untouched - tl_off = Human.get_transform_list("cartesian_wristframe_6d", stride=1) - assert not isinstance(tl_off[0], RotateLocalFrame) - # keypoints modes correct the wrist_pose keys their frames are built on - # (the converter wrote the same double-mirrored rotation into both), plus - # the ee_pose keys when those are built too. - tl_kp = Human.get_transform_list( - "keypoints_wristframe_6d", fix_left_wrist_convention=True - ) - assert isinstance(tl_kp[0], RotateLocalFrame) - assert set(tl_kp[0].keys) == {"left.action_wrist_pose", "left.obs_wrist_pose"} - tl_kp_ee = Human.get_transform_list( - "keypoints_wristframe_6d", fix_left_wrist_convention=True, include_ee_pose=True - ) - assert set(tl_kp_ee[0].keys) == { - "left.action_wrist_pose", - "left.obs_wrist_pose", - "left.action_ee_pose", - "left.obs_ee_pose", - } - - def _bounds_check_dataset(key: str, width: int): """Minimal MultiDataset shell exposing _check_bounds with +-1 quantile bounds on ``key`` for embodiment 0.""" @@ -493,6 +370,134 @@ def test_bounds_check_warns_once_on_stat_shape_mismatch(caplog): assert len(msgs) == 1, msgs +def test_rotate_local_frame_flips_left_wrist_convention(): + # Right-multiplying by Rz(180 deg) must flip the pose's own x/y axes, keep + # z (knuckle-forward) and the position, skip zero-quat padding rows, and + # handle both (7,) poses and (T, 7) chunks. + from scipy.spatial.transform import Rotation as R + + from egomimic.rldb.zarr.action_chunk_transforms import RotateLocalFrame + + rng = np.random.default_rng(4) + q = R.random(3, random_state=5) + chunk = np.zeros((4, 7)) + chunk[:3, :3] = rng.uniform(-1, 1, size=(3, 3)) + chunk[:3, 3:] = q.as_quat()[:, [3, 0, 1, 2]] # wxyz; row 3 stays zero-padded + + t = RotateLocalFrame(keys=["k"]) + out = t.transform({"k": chunk.copy()})["k"] + + np.testing.assert_allclose(out[:, :3], chunk[:, :3]) # positions unchanged + np.testing.assert_allclose(out[3], np.zeros(7)) # padding untouched + R_old = q.as_matrix() + R_new = R.from_quat(out[:3, [4, 5, 6, 3]]).as_matrix() + np.testing.assert_allclose(R_new[:, :, 0], -R_old[:, :, 0], atol=1e-12) # x flip + np.testing.assert_allclose(R_new[:, :, 1], -R_old[:, :, 1], atol=1e-12) # y flip + np.testing.assert_allclose(R_new[:, :, 2], R_old[:, :, 2], atol=1e-12) # z kept + + single = t.transform({"k": chunk[0].copy()})["k"] + np.testing.assert_allclose(single, out[0], atol=1e-12) + + +def test_fix_left_wrist_convention_flag_prepends_correction(): + from egomimic.rldb.embodiment.human import Human + from egomimic.rldb.zarr.action_chunk_transforms import RotateLocalFrame + + tl = Human.get_transform_list( + "cartesian_wristframe_6d", stride=1, fix_left_wrist_convention=True + ) + assert isinstance(tl[0], RotateLocalFrame) + assert set(tl[0].keys) == {"left.action_ee_pose", "left.obs_ee_pose"} + # default off: other vendors' data must be untouched + tl_off = Human.get_transform_list("cartesian_wristframe_6d", stride=1) + assert not isinstance(tl_off[0], RotateLocalFrame) + # keypoints modes correct the wrist_pose keys their frames are built on + # (the converter wrote the same double-mirrored rotation into both), plus + # the ee_pose keys when those are built too. + tl_kp = Human.get_transform_list( + "keypoints_wristframe_6d", fix_left_wrist_convention=True + ) + assert isinstance(tl_kp[0], RotateLocalFrame) + assert set(tl_kp[0].keys) == {"left.action_wrist_pose", "left.obs_wrist_pose"} + tl_kp_ee = Human.get_transform_list( + "keypoints_wristframe_6d", fix_left_wrist_convention=True, include_ee_pose=True + ) + assert set(tl_kp_ee[0].keys) == { + "left.action_wrist_pose", + "left.obs_wrist_pose", + "left.action_ee_pose", + "left.obs_ee_pose", + } + + +def test_vendor_embodiment_names_collapse_to_human(): + # Mirror episodes written by the vendor-split registry carry names like + # MECKA_BIMANUAL in their zarr metadata; locally all human demo data is + # one embodiment, so these must resolve to the HUMAN_* ids. + from egomimic.rldb.embodiment.embodiment import ( + EMBODIMENT, + get_embodiment_id, + is_legacy_vendor_embodiment, + ) + + for vendor in ("mecka", "scale", "aria", "lightwheel"): + assert ( + get_embodiment_id(f"{vendor}_bimanual") == EMBODIMENT.HUMAN_BIMANUAL.value + ) + assert ( + get_embodiment_id(f"{vendor}_right_arm") == EMBODIMENT.HUMAN_RIGHT_ARM.value + ) + assert ( + get_embodiment_id(f"{vendor}_left_arm") == EMBODIMENT.HUMAN_LEFT_ARM.value + ) + assert is_legacy_vendor_embodiment(f"{vendor.upper()}_BIMANUAL") + assert get_embodiment_id("human_bimanual") == EMBODIMENT.HUMAN_BIMANUAL.value + assert get_embodiment_id("eva_bimanual") == EMBODIMENT.EVA_BIMANUAL.value + assert not is_legacy_vendor_embodiment("human_bimanual") + with pytest.raises(KeyError): + get_embodiment_id("yam_bimanual") # robot names are never aliased + + +def test_base_converter_rejects_norm_6d_encoding(): + converter = BaseActionConverter() + with pytest.raises(NotImplementedError, match="normalized-rot6d"): + converter.to32_norm_6d(torch.zeros(1, 1, 20)) + with pytest.raises(NotImplementedError, match="normalized-rot6d"): + converter.from32_norm_6d(torch.zeros(1, 1, 32)) + + +def test_unpad_gripper_zeros_inverts_pad_and_noops_unpadded(): + from egomimic.rldb.zarr.action_chunk_transforms import ( + PadGripperZeros, + UnpadGripperZeros, + ) + + rng = np.random.default_rng(6) + for width in (12, 18): + v = rng.uniform(-1, 1, size=(width,)) + padded = PadGripperZeros(action_key="k").transform({"k": v.copy()})["k"] + assert padded.shape == (width + 2,) + back = UnpadGripperZeros(action_key="k").transform({"k": padded.copy()})["k"] + np.testing.assert_allclose(back, v) + # no-op on already-unpadded widths + same = UnpadGripperZeros(action_key="k").transform({"k": v.copy()})["k"] + np.testing.assert_allclose(same, v) + + +def test_human_6d_reverts_unpad_proprio(): + from egomimic.rldb.embodiment.human import ( + _build_human_cartesian_revert_6d_transform_list, + _build_human_cartesian_revert_6d_wristframe_transform_list, + ) + from egomimic.rldb.zarr.action_chunk_transforms import UnpadGripperZeros + + for build in ( + _build_human_cartesian_revert_6d_transform_list, + _build_human_cartesian_revert_6d_wristframe_transform_list, + ): + assert any(isinstance(t, UnpadGripperZeros) for t in build()), build.__name__ + + def _fallback_dataset(): from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset @@ -533,3 +538,390 @@ def test_fallback_cap_raises(): for n in range(1, md.MAX_FALLBACK_ATTEMPTS + 2): idx, attempts = md._next_after_failure(idx, "bad", attempts, reason="boom") assert n == md.MAX_FALLBACK_ATTEMPTS + + +def test_video_fps_compensates_for_world_size(): + # Distributed val strides an episode's frames by world_size on each rank; + # playback fps must scale down to keep videos wall-clock real-time. + from egomimic.eval.eval_video import EvalVideo + + class _Stub(EvalVideo): + def compute_metrics_and_viz(self, batch, do_viz=True): + raise NotImplementedError + + ev = _Stub.__new__(_Stub) + for world, expected in [(1, 30), (2, 15), (4, 8), (8, 4)]: + ev.trainer = SimpleNamespace(world_size=world) + assert ev._video_fps() == expected, (world, ev._video_fps()) + ev.trainer = SimpleNamespace() # no world_size attr -> assume 1 + assert ev._video_fps() == 30 + + +def test_viz_gate_follows_lightning_epoch_convention(): + from egomimic.eval.eval_video import EvalVideo + + class _Stub(EvalVideo): + def compute_metrics_and_viz(self, batch, do_viz=True): + raise NotImplementedError + + ev = _Stub(viz_every_n_epochs=20) + # Lightning validates at current_epoch 19, 39, ... for check_val_every_n_epoch=20 + for epoch, expected in [(0, False), (19, True), (20, False), (39, True)]: + ev.trainer = SimpleNamespace(current_epoch=epoch) + assert ev._should_viz() is expected, epoch + ev.viz_every_n_epochs = 0 + assert ev._should_viz() is False + + +def test_viz_gate_always_renders_the_final_epoch(): + from egomimic.eval.eval_video import EvalVideo + + class _Stub(EvalVideo): + def compute_metrics_and_viz(self, batch, do_viz=True): + raise NotImplementedError + + ev = _Stub(viz_every_n_epochs=100) + # trainer=debug: 4 epochs, val at current_epoch 1 and 3 -> video only at 3 + for epoch, expected in [(1, False), (3, True)]: + ev.trainer = SimpleNamespace(current_epoch=epoch, max_epochs=4) + assert ev._should_viz() is expected, epoch + # eval mode validates once at epoch 0 with max_epochs=1 + ev.trainer = SimpleNamespace(current_epoch=0, max_epochs=1) + assert ev._should_viz() is True + # unbounded runs (max_epochs None / -1) only follow the interval + for max_epochs in (None, -1): + ev.trainer = SimpleNamespace(current_epoch=0, max_epochs=max_epochs) + assert ev._should_viz() is False + # viz_every_n_epochs=0 still turns video off entirely + ev.viz_every_n_epochs = 0 + ev.trainer = SimpleNamespace(current_epoch=3, max_epochs=4) + assert ev._should_viz() is False + + +def test_frechet_helper(): + from egomimic.utils.metrics import frechet_gaussian_over_time + + B, T, D = 4, 10, 6 + torch.manual_seed(0) + pred = torch.randn(B, T, D) + fd = frechet_gaussian_over_time(pred, pred.clone()) + assert fd.shape == (B,) + assert fd.max() < 1e-2, "identical distributions must score ~0" + fd_shift = frechet_gaussian_over_time(pred + 5.0, pred) + assert (fd_shift > fd).all(), "mean shift must increase the distance" + + +def test_reverse_kl_is_gone(): + import egomimic.utils.metrics as metrics + + assert not hasattr(metrics, "reverse_kl_from_samples") + + +def test_train_viz_wrapper_prefixes_and_forwards(): + from egomimic.eval.eval_train_viz import TrainVizEvalVideo + from egomimic.eval.eval_video import EvalVideo + + class _Base(EvalVideo): + def __init__(self): + super().__init__(viz_func={}, transform_lists={}, viz_every_n_epochs=7) + self.seen = None + + def compute_metrics_and_viz(self, batch, do_viz=True): + self.seen = (self.model, do_viz) + return {"Valid/x": 1.0}, {} + + base = _Base() + tv = TrainVizEvalVideo(base) + algo = object() + tv.model = algo # property setter forwards to base too + metrics, _ = tv.compute_metrics_and_viz({}, do_viz=False) + assert set(metrics) == {"train_viz/Valid/x"}, metrics + assert base.seen == (algo, False) + assert tv.viz_every_n_epochs == 7, "wrapper inherits the base viz gate" + tv.trainer = SimpleNamespace(default_root_dir="/tmp/run") + assert base.trainer is tv.trainer + assert tv.video_dir().endswith("videos_train_viz") + + op = TrainVizEvalVideo(base, prefix="unseen_op_valid") + metrics, _ = op.compute_metrics_and_viz({}, do_viz=False) + assert set(metrics) == {"unseen_op_valid/Valid/x"}, metrics + op.trainer = SimpleNamespace(default_root_dir="/tmp/run") + assert op.video_dir().endswith("videos_unseen_op_valid") + + +def test_dtw_distance_matches_bruteforce_and_tolerates_shift(): + from egomimic.utils.metrics import dtw_distance + + def _dtw_ref(x, y): + t1, t2 = len(x), len(y) + cost = np.linalg.norm(x[:, None, :] - y[None, :, :], axis=-1) + acc = np.full((t1 + 1, t2 + 1), np.inf) + acc[0, 0] = 0.0 + for i in range(1, t1 + 1): + for j in range(1, t2 + 1): + acc[i, j] = cost[i - 1, j - 1] + min( + acc[i - 1, j], acc[i, j - 1], acc[i - 1, j - 1] + ) + return acc[t1, t2] + + torch.manual_seed(3) + pred = torch.randn(3, 9, 4) + tgt = torch.randn(3, 9, 4) + got = dtw_distance(pred, tgt, normalize=False) + for b in range(3): + ref = _dtw_ref(pred[b].numpy(), tgt[b].numpy()) + assert abs(got[b].item() - ref) < 1e-4, (b, got[b].item(), ref) + + # identical trajectories -> 0 + assert dtw_distance(pred, pred.clone()).max() < 1e-6 + + # a time-shifted copy of the same smooth trajectory: DTW must forgive the + # shift (score far below the paired per-step distance of the shifted pair) + t = torch.linspace(0, 6.28, 40) + traj = torch.stack([torch.sin(t), torch.cos(t)], dim=-1)[None] # (1, 40, 2) + shifted = torch.roll(traj, shifts=3, dims=1) + dtw_shift = dtw_distance(traj, shifted, normalize=False).item() + paired = (traj - shifted).norm(dim=-1).sum().item() + assert dtw_shift < 0.5 * paired, (dtw_shift, paired) + + +def test_split_mse_is_stateless_and_matches_manual(): + from egomimic.eval.action_metrics import _paired_mse, _split_mse + + rng = np.random.default_rng(21) + pred = torch.from_numpy(rng.normal(size=(4, 10, 18))).float() + gt = torch.from_numpy(rng.normal(size=(4, 10, 18))).float() + xyz_idx = [0, 1, 2, 9, 10, 11] + rot_idx = [i for i in range(18) if i not in xyz_idx] + xyz, rot = _split_mse(pred, gt) + torch.testing.assert_close( + xyz, (pred[..., xyz_idx] - gt[..., xyz_idx]).pow(2).mean() + ) + torch.testing.assert_close( + rot, (pred[..., rot_idx] - gt[..., rot_idx]).pow(2).mean() + ) + # stateless: a second, unrelated call is unaffected by the first + a, b = torch.zeros(2, 3, 18), torch.ones(2, 3, 18) + torch.testing.assert_close(_split_mse(a, b)[0], torch.tensor(1.0)) + torch.testing.assert_close(_paired_mse(a, b), torch.tensor(1.0)) + assert _split_mse(torch.zeros(2, 3, 7), torch.zeros(2, 3, 7)) == (None, None) + + +def test_rot_geodesic_error_matches_angle_and_survives_gimbal_lock(): + from scipy.spatial.transform import Rotation as R + + from egomimic.eval.action_metrics import _rot_geodesic_error + + rng = np.random.default_rng(22) + ypr = rng.uniform(-1.0, 1.0, size=(6, 5, 3)) + theta = 0.3 + # rotate every pose by theta about a random axis -> geodesic error == theta + axes = rng.normal(size=(6, 5, 3)) + axes /= np.linalg.norm(axes, axis=-1, keepdims=True) + Rg = R.from_euler("ZYX", ypr.reshape(-1, 3)) + Rp = R.from_rotvec(theta * axes.reshape(-1, 3)) * Rg + ypr_p = Rp.as_euler("ZYX").reshape(6, 5, 3) + + # 12-dim ypr layout (both arms the same pose) + gt12 = torch.from_numpy(np.concatenate([ypr, ypr], -1)).float() + gt12 = torch.cat( + [torch.zeros(6, 5, 3), gt12[..., :3], torch.zeros(6, 5, 3), gt12[..., 3:]], -1 + ) + pr12 = torch.from_numpy(np.concatenate([ypr_p, ypr_p], -1)).float() + pr12 = torch.cat( + [torch.zeros(6, 5, 3), pr12[..., :3], torch.zeros(6, 5, 3), pr12[..., 3:]], -1 + ) + assert abs(_rot_geodesic_error(pr12, gt12).item() - theta) < 1e-4 + assert _rot_geodesic_error(gt12, gt12).item() < 1e-5 + + # 18-dim 6D layout, same rotations -> same answer + def to18(y): + six = _ypr_to_rot6d(y) + arm = np.concatenate([np.zeros(y.shape[:-1] + (3,)), six], -1) + return torch.from_numpy(np.concatenate([arm, arm], -1)).float() + + assert abs(_rot_geodesic_error(to18(ypr_p), to18(ypr)).item() - theta) < 1e-4 + assert _rot_geodesic_error(torch.zeros(2, 3, 7), torch.zeros(2, 3, 7)) is None + + # Gimbal lock: pitch ~ pi/2, pred = gt rotated 0.004 rad about local y. + # A ypr MSE (even wrap-aware) explodes because yaw and roll trade off; + # the geodesic error reports the true 0.004. + gt = np.array([0.3, np.pi / 2 - 0.002, 0.2]) + Rgt = R.from_euler("ZYX", gt) + Rpr = Rgt * R.from_rotvec([0.0, 0.004, 0.0]) + pr = Rpr.as_euler("ZYX") + v_gt = torch.tensor([[0, 0, 0, *gt, 0, 0, 0, *gt]], dtype=torch.float32) + v_pr = torch.tensor([[0, 0, 0, *pr, 0, 0, 0, *pr]], dtype=torch.float32) + geo = _rot_geodesic_error(v_pr, v_gt).item() + assert abs(geo - 0.004) < 1e-3, geo + euler_mse = (v_pr - v_gt).pow(2).mean().item() + assert euler_mse > 0.1, euler_mse # the Euler metric is fooled here + + +# ------------------------------------------------ PIEvalVideo metric contract +def _rand_6d_chunks(seed, B=3, T=6): + """(pred, gt) native 18D human chunks with a small, known perturbation, + plus the per-(batch, arm) rotation matrices so a test can re-express both + in another frame. 6D columns are built directly from rotation matrices + (col 0 | col 1), matching ``_reconstruct_R_from_cols``.""" + from scipy.spatial.transform import Rotation as R + + rng = np.random.default_rng(seed) + xyz_g = rng.uniform(-1.0, 1.0, size=(B, T, 2, 3)) + R_g = R.random(B * T * 2, random_state=seed).as_matrix().reshape(B, T, 2, 3, 3) + xyz_p = xyz_g + 0.05 * rng.normal(size=xyz_g.shape) + R_p = ( + ( + R.from_rotvec(0.1 * rng.normal(size=(B * T * 2, 3))) + * R.from_matrix(R_g.reshape(-1, 3, 3)) + ) + .as_matrix() + .reshape(B, T, 2, 3, 3) + ) + return (xyz_p, R_p), (xyz_g, R_g) + + +def _chunk18(xyz, Rm): + six = np.concatenate([Rm[..., :, 0], Rm[..., :, 1]], -1) # (B,T,2,6) + arm = np.concatenate([xyz, six], -1) # (B,T,2,9) + return torch.from_numpy(arm.reshape(arm.shape[0], arm.shape[1], 18)).float() + + +def _run_pi_eval(pred, gt, ac_key="actions_cartesian"): + from egomimic.eval.eval_pi import PIEvalVideo + from egomimic.rldb.embodiment.embodiment import EMBODIMENT, get_embodiment + + eid = EMBODIMENT.HUMAN_BIMANUAL.value + name = get_embodiment(eid).lower() + + class _NormStats: + def unnormalize(self, batch, embodiment_id): + return batch + + class _Algo: + ac_keys = {eid: ac_key} + norm_stats = _NormStats() + + def forward_eval(self, batch): + return {f"{name}_loss": torch.tensor(0.5), f"{name}_{ac_key}": pred} + + ev = PIEvalVideo(viz_func={}, transform_lists={}, viz_every_n_epochs=1) + ev.model = _Algo() + batch = {eid: {ac_key: gt, "embodiment": torch.tensor([eid])}} + metrics, images = ev.compute_metrics_and_viz(batch, do_viz=False) + return f"Valid/{name}_{ac_key}_", f"Valid/{name}_loss", metrics, images + + +def test_pi_eval_metric_set_and_identity(): + (xyz_p, R_p), (xyz_g, R_g) = _rand_6d_chunks(40) + gt = _chunk18(xyz_g, R_g) + prefix, loss_key, metrics, images = _run_pi_eval(gt.clone(), gt) + expected = { + prefix + k + for k in ( + "xyz_paired_mse_avg", + "xyz_final_mse_avg", + "rot6d_paired_mse_avg", + "rot_err_deg_avg", + "rot_err_deg_final", + "xyz_dtw_avg", + "xyz_frechet_gauss_avg", + "xyz_l2_early", + "xyz_l2_mid", + "xyz_l2_late", + ) + } | {loss_key, "Valid/action_loss"} + assert set(metrics) == expected, sorted(metrics) + assert images == {}, "do_viz=False must not render" + assert float(metrics[prefix + "xyz_paired_mse_avg"]) < 1e-10 + assert float(metrics[prefix + "rot_err_deg_avg"]) < 1e-3 + assert float(metrics[prefix + "xyz_dtw_avg"]) < 1e-6 + assert float(metrics[prefix + "xyz_l2_late"]) < 1e-5 + assert float(metrics[prefix + "xyz_frechet_gauss_avg"]) < 1e-2 + + +def test_pi_eval_metrics_are_invariant_to_a_shared_rigid_transform(): + # The evaluator logs no cam-frame twins of the native metrics: the revert + # applies the SAME rigid transform (the observed wrist pose) to prediction + # and gt, and every kept metric is invariant to that. Re-express pred and + # gt in a random frame per (batch, arm) and check. + from scipy.spatial.transform import Rotation as R + + (xyz_p, R_p), (xyz_g, R_g) = _rand_6d_chunks(41) + B = xyz_g.shape[0] + R_w = R.random(B * 2, random_state=7).as_matrix().reshape(B, 1, 2, 3, 3) + t_w = np.random.default_rng(8).uniform(-2, 2, size=(B, 1, 2, 3)) + + def move(xyz, Rm): + return np.einsum("btaij,btaj->btai", R_w, xyz) + t_w, R_w @ Rm + + prefix, _, native, _ = _run_pi_eval(_chunk18(xyz_p, R_p), _chunk18(xyz_g, R_g)) + prefix, _, moved, _ = _run_pi_eval( + _chunk18(*move(xyz_p, R_p)), _chunk18(*move(xyz_g, R_g)) + ) + for key, atol in ( + ("xyz_paired_mse_avg", 1e-5), + ("xyz_final_mse_avg", 1e-5), + ("rot_err_deg_avg", 1e-3), + ("rot_err_deg_final", 1e-3), + ("xyz_dtw_avg", 1e-4), + ("xyz_frechet_gauss_avg", 1e-3), + ("xyz_l2_early", 1e-5), + ("xyz_l2_late", 1e-5), + ): + a, b = float(native[prefix + key]), float(moved[prefix + key]) + assert abs(a - b) < atol, (key, a, b) + # sanity: the perturbation is visible in the headline metrics + assert float(native[prefix + "rot_err_deg_avg"]) > 1.0 + assert float(native[prefix + "xyz_paired_mse_avg"]) > 1e-4 + + +def test_pi_eval_keypoint_metrics(): + from scipy.spatial.transform import Rotation as R + + rng = np.random.default_rng(50) + B, T = 2, 5 + wrist_xyz = rng.uniform(-1, 1, size=(B, T, 2, 3)) + ypr = rng.uniform(-1, 1, size=(B, T, 2, 3)) + kp = rng.uniform(-0.1, 0.1, size=(B, T, 2, 63)) + + def chunk(wx, y, k): + hand = np.concatenate([wx, _ypr_to_rot6d(y), k], -1) # (B,T,2,72) + return torch.from_numpy(hand.reshape(B, T, 144)).float() + + gt = chunk(wrist_xyz, ypr, kp) + theta = 0.2 + axes = rng.normal(size=(B * T * 2, 3)) + axes /= np.linalg.norm(axes, axis=-1, keepdims=True) + ypr_p = ( + (R.from_rotvec(theta * axes) * R.from_euler("ZYX", ypr.reshape(-1, 3))) + .as_euler("ZYX") + .reshape(B, T, 2, 3) + ) + pred = chunk(wrist_xyz + 0.1, ypr_p, kp + 0.02) + prefix, loss_key, metrics, _ = _run_pi_eval(pred, gt, ac_key="actions_keypoints") + expected = { + prefix + k + for k in ( + "kp_l2_avg", + "kp_l2_final", + "wrist_xyz_paired_mse_avg", + "wrist_xyz_final_mse_avg", + "wrist_rot_err_deg_avg", + "wrist_rot_err_deg_final", + "kp_dtw_avg", + *( + f"{m}_l2_{seg}" + for m in ("kp", "wrist_xyz") + for seg in ("early", "mid", "late") + ), + ) + } | {loss_key, "Valid/action_loss"} + assert set(metrics) == expected, sorted(metrics) + # every keypoint is off by 0.02 in each axis -> L2 = 0.02 * sqrt(3) + assert abs(float(metrics[prefix + "kp_l2_avg"]) - 0.02 * math.sqrt(3)) < 1e-5 + assert abs(float(metrics[prefix + "wrist_xyz_paired_mse_avg"]) - 0.01) < 1e-6 + assert ( + abs(float(metrics[prefix + "wrist_rot_err_deg_avg"]) - math.degrees(theta)) + < 1e-2 + ) From 3e94ca756dcc7087b8d8aaabd54e20612d64dda2 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Wed, 16 Sep 2026 18:30:21 -0400 Subject: [PATCH 07/79] recipes: mecka 6D / keypoint configs, 300M models, viz scripts, defaults The fold recipes this stack was built for: mecka_all_6d, mecka_fold_freeform_6d and its opsplit variant (wrist frame; cam frame is a transform_list.mode override rather than a near-duplicate config, and test_data_configs_compose rejects any data config that differs from its parent only in that mode), mecka_fold_flagship_opsplit_hpt_6d, mecka_fold_flagship_topop_hpt_6d and mecka_fold_freeform_opsplit_hpt_keypoints, plus the two 300M HPT models and pi0.5_bc_mecka_6d. Defaults, each for a failure seen in a run: - checkpoints every 50 epochs, saved on train-epoch end. With check_val_every_n_epoch > 1 the validation-end hook only fires on val epochs, so every_n_epochs silently skipped the others. - ddp_pi inference_mode: false. The compiled pi sampler does in-place setup on its static buffers during cudagraph capture, which is illegal on inference tensors ("Inplace update to inference tensor outside InferenceMode"); no_grad is equivalent for eval and creates none. - ddp_pi validates every 20 epochs instead of 200, and eval_hpt renders viz every 200. - pi0.5 gets a constant LR after warmup. The cosine schedule decayed to 0 over num_training_steps=60000, so runs past ~60k steps trained at ~0 lr and shorter runs never saw the configured lr for most of training. - The Qwen stem resolves a repo id to its cached snapshot directory when offline: transformers >= 4.57 probes the Hub while loading a tokenizer by repo id even under HF_HUB_OFFLINE, which raises on compute nodes with no internet. Loading from the directory skips the probe. Offline is read the way huggingface_hub reads it -- 1 / ON / YES / TRUE in any case, from the live environment -- rather than by matching the literal "1", which left HF_HUB_OFFLINE=true hitting the probe the helper exists to avoid. viz_mecka_episode.py and viz_eva_episode.py render an episode straight from its zarr, which is how the left-wrist convention bug was found. Both take the dataset dir from egomimic.utils.paths.default_dataset_dir(), which reads paths/default.yaml, instead of each re-stating the cluster path. HumanBimanualCartesianEuler / RobotBimanualCartesianEuler are renamed to HumanBimanualCartesian6D / RobotBimanualCartesian6D: both pack xyz + rot6d and have not touched an Euler angle since the 6D conversion, and this is the commit that makes the human one the converter for a new default recipe. Also prunes unit tests whose assertions another test already makes on the same code path, and parametrized cases that repeat a stricter test's compose call. Co-Authored-By: Claude Opus 5 (1M context) --- .gitignore | 10 + .../hydra_configs/callbacks/checkpoints.yaml | 6 +- egomimic/hydra_configs/data/mecka_all_6d.yaml | 91 ++++ .../mecka_fold_flagship_opsplit_hpt_6d.yaml | 138 ++++++ .../mecka_fold_flagship_topop_hpt_6d.yaml | 45 ++ .../data/mecka_fold_freeform_6d.yaml | 96 +++++ .../data/mecka_fold_freeform_opsplit_6d.yaml | 144 +++++++ ...a_fold_freeform_opsplit_hpt_keypoints.yaml | 51 +++ .../model/hpt_bc_keypoints_wrist_300M.yaml | 158 +++++++ .../model/hpt_bc_mecka_6d_300M.yaml | 161 +++++++ egomimic/hydra_configs/model/pi0.5_base.yaml | 8 +- .../hydra_configs/model/pi0.5_bc_eva.yaml | 2 +- .../model/pi0.5_bc_mecka_6d.yaml | 30 ++ .../model/pi0.5_cotrain_eva_aria.yaml | 2 +- .../model/pi0.5_cotrain_eva_aria_6d.yaml | 4 +- .../train_zarr_cartesian_pi.yaml | 1 + .../train_zarr_mecka_flagship_6d_hpt.yaml | 20 + .../train_zarr_mecka_kp_wrist_hpt.yaml | 17 + egomimic/hydra_configs/trainer/ddp_pi.yaml | 11 +- egomimic/models/hpt_nets.py | 33 +- egomimic/scripts/abc_process/__init__.py | 0 .../scripts/abc_process/viz_eva_episode.py | 392 ++++++++++++++++++ .../mecka_process/viz_mecka_episode.py | 264 ++++++++++++ egomimic/utils/action_utils.py | 4 +- egomimic/utils/paths.py | 21 + tests/fixtures/recipes.py | 11 +- tests/integration/test_train_step_gpu.py | 7 +- tests/unit/test_compile_cache.py | 6 - tests/unit/test_data_configs_compose.py | 30 ++ tests/unit/test_dataset_filter.py | 7 - tests/unit/test_eva_wristframe_norm_pool.py | 73 ++-- tests/unit/test_infra_port_wiring.py | 194 +++++---- tests/unit/test_latent_dataset_kwargs.py | 6 - tests/unit/test_model_dims.py | 2 + tests/unit/test_pi05_norm_rot6d.py | 8 +- tests/unit/test_wrist6d_roundtrip.py | 8 +- 36 files changed, 1893 insertions(+), 168 deletions(-) create mode 100644 egomimic/hydra_configs/data/mecka_all_6d.yaml create mode 100644 egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml create mode 100644 egomimic/hydra_configs/data/mecka_fold_flagship_topop_hpt_6d.yaml create mode 100644 egomimic/hydra_configs/data/mecka_fold_freeform_6d.yaml create mode 100644 egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml create mode 100644 egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_hpt_keypoints.yaml create mode 100644 egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml create mode 100644 egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml create mode 100644 egomimic/hydra_configs/model/pi0.5_bc_mecka_6d.yaml create mode 100644 egomimic/hydra_configs/train_zarr_mecka_flagship_6d_hpt.yaml create mode 100644 egomimic/hydra_configs/train_zarr_mecka_kp_wrist_hpt.yaml create mode 100644 egomimic/scripts/abc_process/__init__.py create mode 100644 egomimic/scripts/abc_process/viz_eva_episode.py create mode 100644 egomimic/scripts/mecka_process/viz_mecka_episode.py create mode 100644 egomimic/utils/paths.py diff --git a/.gitignore b/.gitignore index 8f3d361847..1380f8906a 100644 --- a/.gitignore +++ b/.gitignore @@ -45,10 +45,20 @@ annotations_test/** log_conversion/** debug_model_inputs/** temp_dir/** +/datasets/ +apikey.txt +slurm-*.out +slurmoutputs/ +*.log +.inductor_cache/ + scratch/ sample_data/ external_ckpts/ external/MANO/ +ann_backup/ +aria_ann_backup/ +sort_annotations/ # setuptools build output build/ diff --git a/egomimic/hydra_configs/callbacks/checkpoints.yaml b/egomimic/hydra_configs/callbacks/checkpoints.yaml index 86d21a792f..fbd6d13277 100644 --- a/egomimic/hydra_configs/callbacks/checkpoints.yaml +++ b/egomimic/hydra_configs/callbacks/checkpoints.yaml @@ -8,4 +8,8 @@ model_checkpoint: filename: "epoch_{epoch}" save_last: true save_top_k: -1 - every_n_epochs: 100 \ No newline at end of file + every_n_epochs: 50 + # Save at the end of the train epoch, not after validation: with + # check_val_every_n_epoch > 1 the validation-end hook only fires on val + # epochs and every_n_epochs would silently miss the others. + save_on_train_epoch_end: true diff --git a/egomimic/hydra_configs/data/mecka_all_6d.yaml b/egomimic/hydra_configs/data/mecka_all_6d.yaml new file mode 100644 index 0000000000..c7ba564068 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_all_6d.yaml @@ -0,0 +1,91 @@ +# MECKA_ALL_PI_6D: pi0.5 on ALL mecka episodes in the SQL table (filters +# lab=mecka with no task filter), stored under the human_bimanual embodiment. +# +# WRIST-frame CARTESIAN actions with the continuous 6D rotation representation +# (18D = xyz+6D per arm, no gripper): each arm's action chunk is expressed +# relative to that wrist's current pose, so both the arbitrary world frame AND +# the moving head frame cancel out of the action targets. The proprio ee_pose +# stays in the head/cam frame (it defines where the wrists are for the model +# and for the eval revert) and is 6D-encoded. stride=1: mecka actions are +# recorded at the target rate (Aria uses 3). Language prompts are sampled +# from the `annotations` track; prompt assembly (proprio-in-prompt on, +# embodiment label off) lives in model=pi0.5_bc_mecka_6d, which this config +# pairs with, alongside evaluator=eval_pi_wristframe_6d (reverts wrist -> +# head frame via the proprio, then 6D -> ypr, for the viz). +# +# This is the explicit cartesian recipe; the repo default for human data is +# the hand-keypoint action (data/mecka.yaml, data/cotrain_pi_base.yaml). + +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + # Per-cluster zarr mirror (paths/default.yaml; $EGOVERSE_DATASET_DIR + # overrides). The resolver s5cmd-syncs only the SQL-matched episodes + # missing here. + folder_path: ${paths.dataset_dir} + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: cartesian + annotation_key: annotations + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: cartesian_wristframe_6d + stride: 1 + # processed_v3 mecka zarrs predate the mecka_to_zarr rot_left fix: + # the LEFT wrist frame was double-mirrored onto the right hand's + # spatial convention. Rz(180)-corrects the raw left pose keys before + # any frame math (exactly equivalent to reconverting). + fix_left_wrist_convention: true + # 18 -> 20-dim proprio (zero grip slots at 9/19): State: prompt bins + # align positionally with the robot 20-dim layout. + pad_proprio_gripper: true + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + mode: train + valid_ratio: 0.05 + +valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: ${data.train_datasets.human_bimanual.filters} + mode: valid + valid_ratio: ${data.train_datasets.human_bimanual.valid_ratio} + +# Train-set spot check: a second val-style loader over the TRAIN split, +# consumed by the TrainVizEvalVideo wrapper on every validation epoch +# (metrics prefixed `train_viz/`, videos under `videos_train_viz/`). No +# shuffle -> the same leading slice every val, so the curves are comparable +# across time. Declared here only to set its loader params; without it the +# default train_viz head reuses the train split with the valid loader params +# (the head is on by default, train_viz=true; `train_viz=false` skips it). +train_viz_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: ${data.train_datasets.human_bimanual.filters} + mode: train + valid_ratio: ${data.train_datasets.human_bimanual.valid_ratio} + +train_viz_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 2 + shuffle: false + +train_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + +valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml new file mode 100644 index 0000000000..00d70ff479 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml @@ -0,0 +1,138 @@ +# MECKA_FOLD_FLAGSHIP_OPSPLIT_HPT_6D: HPT BC on the mecka FLAGSHIP fold-clothes +# episodes (SQL lab='mecka', task='fold_clothes', all 1607 eps / 316k frames +# under processed_v3/mecka/flagship/, 20 operators; a DIFFERENT set from the +# 462-ep FREEFORM 'folding_clothes' episodes of mecka_fold_freeform_opsplit_6d, +# zero operator overlap), with a SEEN- and an UNSEEN-operator validation set. +# Standalone (no defaults inheritance). +# +# THREE-WAY SPLIT (chosen 2026-09-11 from SQL app.episodes): +# unseen val = the 2 operators in `held_out_operators`, ALL their episodes +# (191 eps / 30,956 frames, 9.8 %). The two mid-sized operators +# with >= 70 episodes each; the three largest (65 % of frames) +# stay in train. +# seen val = the other 18 operators, seeded 5 % episode split +# (`mode: valid`, ~70 eps / 13.8k frames): known operators, +# unseen episodes. +# train = the other 18 operators, the remaining 95 % (~1346 eps / +# 271k frames). +# +# VAL LOADERS (three heads, one validation pass): +# valid_datasets -> seen_op_valid/Valid/human_bimanual_*, +# videos_seen_op_valid/: SEEN operators, unseen +# episodes (the canonical loader; `valid_prefix` +# renames its metrics and video dir) +# train_viz (default) -> train_viz/Valid/human_bimanual_*, +# videos_train_viz/: the TRAIN split itself (no +# train_viz_datasets here, so trainHydra reuses +# train_datasets, unshuffled); `train_viz=false` +# drops it +# unseen_op_valid_datasets -> unseen_op_valid/Valid/human_bimanual_*, +# videos_unseen_op_valid/: UNSEEN operators +# The train and seen-val entries share `filters` and `valid_ratio`, so the +# seeded split (split_dataset_names, SEED 42) makes them exact complements. +# +# PIPELINE: WRIST-frame CARTESIAN actions with the continuous 6D rotation +# (18D = xyz+6D per arm, no gripper): each arm's 100-step chunk is expressed +# relative to that wrist's current pose; the head-frame proprio ee_pose is +# 6D-encoded too and padded 18 -> 20 (zero grip slots at 9/19). +# fix_left_wrist_convention: processed_v3 mecka zarrs have the LEFT wrist frame +# double-mirrored; Rz(180)-corrected before any frame math. stride 1 (mecka +# is recorded at the target rate). Prompts from the `annotations` track. +# +# NORMALIZATION: quantile-normalize proprio + action keys from stats +# precomputed on THIS config (train side = the 18 operators): +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_flagship_6d_hpt \ +# --data mecka_fold_flagship_opsplit_hpt_6d --model hpt_bc_mecka_6d_300M --out +# The freeform stats share the key set but not the distribution; do not reuse. +# +# Pairs with model=hpt_bc_mecka_6d_300M and evaluator=eval_hpt_wristframe_6d +# (the train_viz head wraps it automatically); see +# train_zarr_mecka_flagship_6d_hpt.yaml for the composed recipe. + +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +# The 2 held-out (UNSEEN) operator ids (SQL app.episodes.operator). Referenced +# by the filter lambdas below via the mounted path `data.held_out_operators`. +# MultiDataModuleWrapper accepts this key (config-only). +held_out_operators: + - 6871d82a1aaf88556f0c6041 + - 6903a36c466be157ec64562d + +# Metric prefix / video dir for the canonical valid loader (config-only; read by +# trainHydra), so the seen-operator val is not logged under a bare Valid/. +valid_prefix: seen_op_valid + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: ${paths.dataset_dir} + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: cartesian + annotation_key: annotations + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: cartesian_wristframe_6d + stride: 1 + fix_left_wrist_convention: true + pad_proprio_gripper: true + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() == 'fold_clothes'" + - "lambda row: '/mecka/flagship/' in str(row.get('zarr_processed_path', ''))" + # OPERATOR SPLIT: train = every operator NOT in the held-out list. + - "lambda row: str(row.get('operator', '')) not in ${data.held_out_operators}" + # Seeded 95/5 episode split within the train operators; the 5 % is the + # seen-operator val (valid_datasets below). + mode: train + valid_ratio: 0.05 + +# SEEN operators, held-out episodes (the complement of train's 95 %) -> +# seen_op_valid/ metrics via the canonical evaluator. +valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: ${data.train_datasets.human_bimanual.filters} + mode: valid + valid_ratio: ${data.train_datasets.human_bimanual.valid_ratio} + +# UNSEEN operators: all episodes of the held-out operators -> unseen_op_valid/ +# metrics via a second copy of the canonical evaluator. +unseen_op_valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() == 'fold_clothes'" + - "lambda row: '/mecka/flagship/' in str(row.get('zarr_processed_path', ''))" + # OPERATOR SPLIT: unseen val = the held-out operators only. + - "lambda row: str(row.get('operator', '')) in ${data.held_out_operators}" + mode: total + +train_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + +# Unshuffled (both val loaders): the same leading slice every val, comparable +# across epochs with a coherent overlay video. Capped per pass by the +# trainer's limit_val_batches. The train_viz loader reuses these params. +valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false + +unseen_op_valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_topop_hpt_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_topop_hpt_6d.yaml new file mode 100644 index 0000000000..d4842d05cb --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_topop_hpt_6d.yaml @@ -0,0 +1,45 @@ +defaults: + - mecka_fold_flagship_opsplit_hpt_6d + - _self_ + +# MECKA_FOLD_FLAGSHIP_TOPOP_HPT_6D: single-operator twin of +# mecka_fold_flagship_opsplit_hpt_6d. TRAIN on the ONE largest flagship +# fold-clothes operator only; everything else (pipeline, unseen-operator val, +# loader params) is inherited from the parent. +# +# SPLIT (from SQL app.episodes, 2026-09-15): +# train = operator 6903686e0e94ce070afd1f24, the largest of the 20 +# (405 eps / 76,013 frames / ~0.70 h at 30 fps, 24 % of +# flagship), 95 % of its episodes (~385 eps / ~72k frames). +# seen val = the other 5 % of that operator's episodes (20 eps / 3,731 +# frames) -> seen_op_valid/. The parent's valid_datasets interpolates +# the train filters and valid_ratio, so it follows this +# override automatically. +# train_viz = the train split above (trainHydra's default) -> train_viz/. +# unseen val = UNCHANGED from the parent: all episodes of the 2 +# `held_out_operators` -> unseen_op_valid/, directly comparable +# with the 18-operator run. +# The other 17 operators are unused. +# +# Hydra replaces lists wholesale, so the train filter_lambdas are restated in +# full; only the last (operator) lambda differs from the parent. The operator id +# is inlined rather than put in a config key: MultiDataModuleWrapper rejects +# root keys it does not declare. +# +# NORMALIZATION: the train distribution differs from the parent's, so +# precompute stats on THIS config: +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_flagship_6d_hpt \ +# --data mecka_fold_flagship_topop_hpt_6d --model hpt_bc_mecka_6d_300M --out +# then train with +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_6d_hpt \ +# data=mecka_fold_flagship_topop_hpt_6d norm_stats.precomputed_norm_path=/norm_stats + +train_datasets: + human_bimanual: + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() == 'fold_clothes'" + - "lambda row: '/mecka/flagship/' in str(row.get('zarr_processed_path', ''))" + # SINGLE-OPERATOR TRAIN: the largest flagship operator only. + - "lambda row: str(row.get('operator', '')) == '6903686e0e94ce070afd1f24'" diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_6d.yaml new file mode 100644 index 0000000000..90fd9c6cac --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_6d.yaml @@ -0,0 +1,96 @@ +# MECKA_FOLD_FREEFORM_6D: pi0.5 BC on the mecka FREEFORM fold-clothes +# episodes only, stored under the human_bimanual embodiment: SQL +# task='folding_clothes', 462 eps / 1.13M frames under +# processed_v3/mecka/freeform/. "freeform" is only visible in +# zarr_processed_path, so the filter keys on '/mecka/freeform/' and accepts +# either fold-clothes task spelling (the flagship 'fold_clothes' rows fail the +# path check, so this resolves to the 462). +# +# RANDOM-EPISODE split (valid_ratio 0.05). For the held-out-OPERATOR split use +# mecka_fold_freeform_opsplit_6d instead. +# +# WRIST-frame CARTESIAN actions with the continuous 6D rotation (18D = +# xyz+6D per arm, no gripper): each arm's 100-step chunk is expressed +# relative to that wrist's current pose, so the world frame and the moving +# head frame cancel out of the targets; the head-frame proprio ee_pose is +# 6D-encoded too (it is the only head-frame signal the model sees). +# +# Same pipeline knobs as mecka_all_6d: fix_left_wrist_convention (processed_v3 +# mecka zarrs have the LEFT wrist frame double-mirrored; Rz(180)-corrected +# before any frame math), pad_proprio_gripper (18 -> 20-dim proprio so the +# State: prompt bins align with the robot layout), stride 1 (mecka is +# recorded at the target rate), prompts from the `annotations` track. +# Pairs with model=pi0.5_bc_mecka_6d and evaluator=eval_pi_wristframe_6d (the +# train-set spot-check head wraps it automatically). For CAM (head)-frame +# actions add +# data.train_datasets.human_bimanual.resolver.transform_list.mode=cartesian_6d +# with evaluator=eval_pi_camframe_6d. +# +# Norm stats: precompute on a CPU node with +# python egomimic/scripts/precompute_norm_stats.py --data mecka_fold_freeform_6d \ +# --model pi0.5_bc_mecka_6d --out +# and point training at norm_stats.precomputed_norm_path=/norm_stats +# (the script also fills the content-keyed cache, so a plain run hits it). + +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: ${paths.dataset_dir} + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: cartesian + annotation_key: annotations + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: cartesian_wristframe_6d + stride: 1 + fix_left_wrist_convention: true + pad_proprio_gripper: true + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + mode: train + valid_ratio: 0.05 + +valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: ${data.train_datasets.human_bimanual.filters} + mode: valid + valid_ratio: ${data.train_datasets.human_bimanual.valid_ratio} + +# Train-set spot check for the TrainVizEvalVideo head (train_viz/ metrics, +# videos_train_viz/). No shuffle -> same leading slice every val +# (the head is on by default, train_viz=true; `train_viz=false` skips it). +train_viz_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: ${data.train_datasets.human_bimanual.filters} + mode: train + valid_ratio: ${data.train_datasets.human_bimanual.valid_ratio} + +train_viz_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 2 + shuffle: false + +train_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + +valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml new file mode 100644 index 0000000000..0c878b6392 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml @@ -0,0 +1,144 @@ +# MECKA_FOLD_FREEFORM_OPSPLIT_6D: pi0.5 BC on the mecka FREEFORM fold-clothes +# episodes (SQL task='folding_clothes', 462 eps / 1.13M frames under +# processed_v3/mecka/freeform/; "freeform" is only visible in +# zarr_processed_path, so the filter keys on '/mecka/freeform/' and accepts +# either fold-clothes task spelling), stored under the human_bimanual +# embodiment, with the held-out-OPERATOR train/val split. +# +# OPERATOR SPLIT: validation = 31 operators held out ENTIRELY, so no val +# episode shares an operator with training. The list lives ONCE, in +# `held_out_operators` below; the train and valid filter lambdas both +# interpolate it, so the two sides cannot drift apart. The HPT keypoint twin +# (mecka_fold_freeform_opsplit_hpt_keypoints) inherit this file. +# +# Chosen deterministically on 2026-09-01 from SQL app.episodes (462 freeform +# fold eps / 178 operators): skip the 10 largest operators (their data stays in +# train), seeded shuffle (random.Random(42)) of the remaining operator ids, +# accumulate until >= 10 % of frames -> 31 operators, 50 episodes, +# 115,681 frames (10.2 %) held out; train = 147 operators, 412 episodes, +# 1,013,220 frames. Both sides use mode: total (no random episode split); the +# list IS the split. Do not edit it: checkpoints were trained against it. +# +# WRIST-frame CARTESIAN actions with the continuous 6D rotation (18D = +# xyz+6D per arm, no gripper): each arm's 100-step chunk is expressed +# relative to that wrist's current pose; the head-frame proprio ee_pose is +# 6D-encoded too. Pairs with model=pi0.5_bc_mecka_6d. +# +# Shared pipeline knobs (same as mecka_all_6d): fix_left_wrist_convention +# (processed_v3 mecka zarrs have the LEFT wrist frame double-mirrored; +# Rz(180)-corrected before any frame math), pad_proprio_gripper (18 -> 20-dim +# proprio so the State: prompt bins align with the robot layout), stride 1, +# prompts from the `annotations` track. Pairs with evaluator=eval_pi_wristframe_6d +# (the train_viz head wraps it automatically). For CAM (head)-frame actions add +# data.train_datasets.human_bimanual.resolver.transform_list.mode=cartesian_6d +# with evaluator=eval_pi_camframe_6d. +# +# Norm stats: precompute on a CPU node +# (egomimic/scripts/precompute_norm_stats.py) and point training at +# norm_stats.precomputed_norm_path=/norm_stats. + +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +# The 31 held-out operator ids (SQL app.episodes.operator), sorted. See the +# OPERATOR SPLIT note above. Referenced by the filter lambdas via the mounted +# path `data.held_out_operators`. MultiDataModuleWrapper accepts this key +# (config-only). +held_out_operators: + - 68a8e70d1ffcf0417e08f0e5 + - 68e2fc680b158ea94ab4f5bc + - 6930bf8a6e496b45b1e1122e + - 6931784c6e496b45b1e1b1cd + - 693a399ec67c9e4814b0c09f + - 693b9530c67c9e4814b20ab2 + - 693c78165cb6c2828a4c2087 + - 694604945c0479a0fb7f48a4 + - 6948d23b5c0479a0fb8302b3 + - 69529a85c0d9448524bc7880 + - 6954f0b720b100982d812ed8 + - 6957464720b100982d83b76c + - 695d07d783a9fdf2d84d9371 + - 695d084183a9fdf2d84d955a + - 6961c50f83a9fdf2d85ec95d + - 69624a8683a9fdf2d86013c0 + - 6973c3ae89fa4bc311d1bbe5 + - 69842643f8ca4464d9920799 + - 698937a5f8ca4464d99dce94 + - 6992b19ab9d55c9c2d48432c + - 699429e18ef0252970048624 + - 6996d5eded1708063b2fe419 + - 6999d17c8f173e14271f9890 + - 69a24256c4216c3e3ce33834 + - 69a39ec9c4216c3e3ce49734 + - 69a57a78c4216c3e3ce6323c + - 69a725a347e12c9b3390ba33 + - 69a7943547e12c9b33911c02 + - 69a84eed47e12c9b339208fd + - 69ab87207bf9cbce29a3cb36 + - 69ae2463d22d99655253f287 + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: ${paths.dataset_dir} + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: cartesian + annotation_key: annotations + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: cartesian_wristframe_6d + stride: 1 + fix_left_wrist_convention: true + pad_proprio_gripper: true + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # OPERATOR SPLIT: train = every operator NOT in the held-out list. + - "lambda row: str(row.get('operator', '')) not in ${data.held_out_operators}" + mode: total + +valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # OPERATOR SPLIT: valid = the held-out operators only. + - "lambda row: str(row.get('operator', '')) in ${data.held_out_operators}" + mode: total + +# Train-set spot check for the TrainVizEvalVideo head (train_viz/ metrics, +# videos_train_viz/): same operators as train, no shuffle -> same slice every val +# (the head is on by default, train_viz=true; `train_viz=false` skips it). +train_viz_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: ${data.train_datasets.human_bimanual.filters} + mode: total + +train_viz_dataloader_params: + human_bimanual: + batch_size: 32 + num_workers: 2 + shuffle: false + +train_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + +valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_hpt_keypoints.yaml new file mode 100644 index 0000000000..66ecee19bb --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_hpt_keypoints.yaml @@ -0,0 +1,51 @@ +defaults: + - mecka_fold_freeform_opsplit_6d + - _self_ + +# MECKA_FOLD_FREEFORM_OPSPLIT_HPT_KEYPOINTS: HPT keypoint-action twin of +# mecka_fold_freeform_opsplit_6d: the SAME mecka FREEFORM fold-clothes +# episodes and the SAME held-out-OPERATOR split (`held_out_operators`, the +# filters, the train_viz slice and the loader params are all inherited; edit +# the operator list only in the parent). Only the keymap and transform differ: +# +# keymap_mode: keypoints +# mode: keypoints_wristframe_6d +# +# Batch keys: +# actions_keypoints (100, 144) wrist-first wrist-frame keypoint +# action: per hand [wrist xyz+rot6d in +# its own current wrist frame (9) | +# 21 MANO keypoints in that wrist +# frame (63)] x {left, right} +# observations.state.keypoints (144,) per hand [wrist pose in HEAD frame +# (9) | keypoints in wrist frame (63)]; +# the HPT `state_keypoints` stem input +# and what the eval revert (wrist -> +# head frame) reads +# annotations task-description track -> the +# `sampled_prompt` for a Qwen text stem +# +# There is no `observations.state.ee_pose` (include_ee_pose off: HPT has no +# stem for it), so pad_proprio_gripper is off. fix_left_wrist_convention stays ON: +# in keypoint modes it corrects the left `obs/action_wrist_pose` keys, which +# carry the same double-mirrored rotation as the ee_pose on processed_v3. +# +# NORMALIZATION: proprio + action keys are quantile-normalized; precompute on +# a CPU node with +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_kp_wrist_hpt \ +# --data mecka_fold_freeform_opsplit_hpt_keypoints --model hpt_bc_keypoints_wrist_300M --out +# +# Pairs with model=hpt_bc_keypoints_wrist_300M and +# evaluator=eval_hpt_keypoints_wrist (the inherited train_viz slice gets the +# default head wrapping it); see train_zarr_mecka_kp_wrist_hpt.yaml for the +# composed recipe. + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + transform_list: + mode: keypoints_wristframe_6d + fix_left_wrist_convention: true + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml new file mode 100644 index 0000000000..d55bf1afc9 --- /dev/null +++ b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml @@ -0,0 +1,158 @@ +# HPT_BC_KEYPOINTS_WRIST_300M: ~300M-parameter HPT flow-matching BC on the +# 144-D wrist-frame MANO keypoint action, with a frozen Qwen3-Embedding text +# stem for the task prompt ("300M hpt qwen"). +# +# Data: mecka_fold_freeform_opsplit_hpt_keypoints; action `actions_keypoints`, +# per hand [wrist xyz+rot6d in its own wrist frame (9) | 21 keypoints in that +# wrist frame (63)] x {left, right}, quantile-normalized from precomputed +# stats. Inputs: front image (ResNet), the 144-D +# `observations.state.keypoints` (wrist pose in head frame + keypoints in +# wrist frame, per hand), and the sampled annotation through the Qwen +# per-token stem. Drop the text stem with +# ~model.robomimic_model.shared_stem_specs.annotation model.robomimic_model.shared_obs_keys=[front_img_1] +# +# Widths follow the lab's 300M reference recipe (hpt_bc_flow_mecka_300M.yaml on +# aniketh/hpt-run-configs; also the basis of the arc-E1 300M cotrain configs): +# embed_dim 840, 10 trunk heads, cross-attn 18 latents x 10 heads x 140 +# dim_head, head hidden_dim 320 / 5 heads / mlp 5x5. embed_dim is the shared +# bus width: every stem in/out, modality_embed_dim, encoder output and the +# head cond_dim move with it. num_blocks is set so the TRAINABLE count +# (frozen Qwen excluded) lands at ~300M (23 blocks ~= 298M measured on the +# 138-D layout; the 144-D stem/head add < 0.1M). +# +# Constant LR 5e-5, no scheduler (the setting of the mecka_fold_*_300M +# constlr_5e5 ablation whose keypoint arm fit best). +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.hpt.HPT + camera_transforms: {} + dims: + human_bimanual: {proprio: 144, action: 144} + ac_keys: + human_bimanual: "actions_keypoints" + + diffusion: true + + # Prompt for the Qwen stem: one annotation sampled per batch item from the + # `annotations` track (the data config's keymap exposes it). + annotation_key: "annotations" + annotation_sampling_mode: "random" + annotation_modality: "annotation" + default_prompt: "" + + trunk: + embed_dim: 840 + num_blocks: 23 + num_heads: 10 + token_postprocessing: "action_token" + observation_horizon: 1 + action_horizon: 64 + no_trunk: false + use_domain_embedding: true + drop_path: 0.1 + weight_init_style: "pytorch" + + multitask: false + pretrained: false + pretrained_checkpoint: null + domains: ["human_bimanual"] + shared_obs_keys: ["front_img_1", "annotation"] + + head_specs: + human_bimanual: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 50 + pooling: null + time_dist: "beta" + infer_ac_dims: + human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} + model: + _target_: egomimic.models.denoising_nets.CrossTransformer + nblocks: 6 + cond_dim: 840 + hidden_dim: 320 + act_dim: ${model.robomimic_model.dims.human_bimanual.action} + act_seq: 100 + n_heads: 5 + dropout: 0.1 + mlp_layers: 5 + mlp_ratio: 5 + + shared_stem_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: 840 + output_dim: 840 + widths: [840] + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 18 + crossattn_heads: 10 + crossattn_dim_head: 140 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 840 + annotation: + _target_: egomimic.models.hpt_nets.QwenPerTokenEncoder + model_name: "Qwen/Qwen3-Embedding-0.6B" + max_length: 128 + freeze: true + dtype: "float16" + output_dim: 840 + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 18 + crossattn_heads: 10 + crossattn_dim_head: 140 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 840 + + stem_specs: + human_bimanual: + state_keypoints: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} + output_dim: 840 + widths: [840] + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 18 + crossattn_heads: 10 + crossattn_dim_head: 140 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 840 + + encoder_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.ResNet + output_dim: 840 + num_of_copy: 1 + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.ColorJitter + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.05 + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + eval_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 5e-5 + weight_decay: 0.0001 + +scheduler: null diff --git a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml new file mode 100644 index 0000000000..8cf56b9a42 --- /dev/null +++ b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml @@ -0,0 +1,161 @@ +# HPT_BC_MECKA_6D_300M: ~300M-parameter HPT flow-matching BC on the 18-D +# WRIST-frame cartesian action with the continuous 6D rotation, plus a frozen +# Qwen3-Embedding text stem for the task prompt. HPT twin of pi0.5_bc_mecka_6d. +# +# Data: a `cartesian_wristframe_6d` data config with `keymap_mode: cartesian`, +# e.g. mecka_fold_flagship_opsplit_hpt_6d. Action `actions_cartesian` = per arm +# [xyz | 6D rotation] (18) relative to that wrist's current pose; proprio +# `observations.state.ee_pose` = head-frame wrist poses, 6D-encoded and padded +# 18 -> 20 (pad_proprio_gripper; zero grip slots at 9/19). Both are +# quantile-normalized from precomputed stats. Inputs: front image (ResNet), the +# 20-D ee_pose through the `state_ee_pose` stem, and one annotation sampled per +# batch item through the Qwen per-token stem. Drop the text stem with +# ~model.robomimic_model.shared_stem_specs.annotation model.robomimic_model.shared_obs_keys=[front_img_1] +# +# Widths follow the lab's 300M reference (hpt_bc_flow_mecka_300M on +# aniketh/hpt-run-configs): embed_dim 840, 10 trunk heads, cross-attn 18 +# latents x 10 heads x 140 dim_head, head hidden_dim 320 / 5 heads / mlp 5x5. +# num_blocks is set so the TRAINABLE count (frozen Qwen excluded) lands at +# ~300M: 23 blocks = 298.1M trainable (trunk 195.0M = 8.48M/block, head 74.2M, +# ResNet 11.6M, image stem 6.1M, Qwen stem projection 5.6M, ee_pose stem 5.4M; +# Qwen 595.8M frozen). 24 blocks would be 306.5M. +# +# Constant LR 5e-5, no scheduler (the setting of the mecka_fold_*_300M +# constlr_5e5 ablations; the CosineAnnealingLR of hpt_bc_flow_mecka sawtooths +# over a long run). +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.hpt.HPT + camera_transforms: {} + dims: + human_bimanual: {proprio: 20, action: 18} + ac_keys: + human_bimanual: "actions_cartesian" + + diffusion: true + 6dof: true + + # Prompt for the Qwen stem: one annotation sampled per batch item from the + # `annotations` track (the data config's keymap exposes it). + annotation_key: "annotations" + annotation_sampling_mode: "random" + annotation_modality: "annotation" + default_prompt: "" + + trunk: + embed_dim: 840 + num_blocks: 23 + num_heads: 10 + token_postprocessing: "action_token" + observation_horizon: 1 + action_horizon: 64 + no_trunk: false + use_domain_embedding: true + drop_path: 0.1 + weight_init_style: "pytorch" + + multitask: false + pretrained: false + pretrained_checkpoint: null + domains: ["human_bimanual"] + shared_obs_keys: ["front_img_1", "annotation"] + + head_specs: + human_bimanual: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 50 + pooling: null + time_dist: "beta" + infer_ac_dims: + human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} + model: + _target_: egomimic.models.denoising_nets.CrossTransformer + nblocks: 6 + cond_dim: 840 + hidden_dim: 320 + act_dim: ${model.robomimic_model.dims.human_bimanual.action} + act_seq: 100 + n_heads: 5 + dropout: 0.1 + mlp_layers: 5 + mlp_ratio: 5 + + shared_stem_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: 840 + output_dim: 840 + widths: [840] + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 18 + crossattn_heads: 10 + crossattn_dim_head: 140 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 840 + annotation: + _target_: egomimic.models.hpt_nets.QwenPerTokenEncoder + model_name: "Qwen/Qwen3-Embedding-0.6B" + max_length: 128 + freeze: true + dtype: "float16" + output_dim: 840 + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 18 + crossattn_heads: 10 + crossattn_dim_head: 140 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 840 + + stem_specs: + human_bimanual: + # 20-D: 6D-encoded head-frame wrist poses (18) + zero grip pads (2). + state_ee_pose: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} + output_dim: 840 + widths: [840] + specs: + random_horizon_masking: false + cross_attn: + crossattn_latent: 18 + crossattn_heads: 10 + crossattn_dim_head: 140 + crossattn_modality_dropout: 0.1 + modality_embed_dim: 840 + + encoder_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.ResNet + output_dim: 840 + num_of_copy: 1 + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.ColorJitter + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.05 + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + eval_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.Normalize + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +optimizer: + _target_: torch.optim.AdamW + _partial_: true + lr: 5e-5 + weight_decay: 0.0001 + +scheduler: null diff --git a/egomimic/hydra_configs/model/pi0.5_base.yaml b/egomimic/hydra_configs/model/pi0.5_base.yaml index 88f2a7947e..0f17969944 100644 --- a/egomimic/hydra_configs/model/pi0.5_base.yaml +++ b/egomimic/hydra_configs/model/pi0.5_base.yaml @@ -68,11 +68,13 @@ optimizer: eps: 1e-8 weight_decay: 0.0 +# Constant LR after warmup (matches the remote pi0.5 training stack). The +# previous cosine schedule decayed to 0 over num_training_steps=60000, so runs +# past ~60k steps trained at ~0 lr and shorter runs never saw the configured lr +# for most of training. scheduler: - _target_: transformers.get_cosine_schedule_with_warmup + _target_: transformers.optimization.get_constant_schedule_with_warmup _partial_: true num_warmup_steps: 2000 - num_training_steps: 60000 - num_cycles: 0.5 scheduler_interval: step \ No newline at end of file diff --git a/egomimic/hydra_configs/model/pi0.5_bc_eva.yaml b/egomimic/hydra_configs/model/pi0.5_bc_eva.yaml index c08fe331ce..f69fe9810a 100644 --- a/egomimic/hydra_configs/model/pi0.5_bc_eva.yaml +++ b/egomimic/hydra_configs/model/pi0.5_bc_eva.yaml @@ -11,7 +11,7 @@ robomimic_model: action_converters: rules: EVA_BIMANUAL: - _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler + _target_: egomimic.utils.action_utils.RobotBimanualCartesian6D # optional fallback if no match is found fallback: _target_: egomimic.utils.action_utils.BaseActionConverter diff --git a/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d.yaml b/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d.yaml new file mode 100644 index 0000000000..3a557c0b5b --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d.yaml @@ -0,0 +1,30 @@ +defaults: + - pi0.5_base + - _self_ + +# pi0.5 BC on mecka data alone (stored as human_bimanual) with the CARTESIAN +# action: actions arrive already in xyz+6D layout (18D per arm pair, no +# gripper, wrist-relative via the `cartesian_wristframe_6d` transform mode) +# and the head-frame proprio ee_pose arrives 6D-encoded too; the forward pass +# only packs them into openpi's 32D vector (pretrained projections kept). +# Pairs with data=mecka_all_6d / mecka_fold_freeform*_6d and +# evaluator=eval_pi_wristframe_6d (or, with the human transform mode +# overridden to `cartesian_6d`, eval_pi_camframe_6d). +robomimic_model: + camera_transforms: {} + ac_keys: + human_bimanual: "actions_cartesian" + domains: ["human_bimanual"] + + # Keep the discretized-proprio "State: " prompt block (pi0.5_base + # defaults proprio_in_prompt on) but do not splice "Embodiment: " into + # the prompt. + embodiment_label: false + + action_converters: + rules: + HUMAN_BIMANUAL: + _target_: egomimic.utils.action_utils.HumanBimanualCartesian6D + # optional fallback if no match is found + fallback: + _target_: egomimic.utils.action_utils.BaseActionConverter diff --git a/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml index c1ef9350b1..a11a9a24a0 100644 --- a/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml +++ b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml @@ -20,7 +20,7 @@ robomimic_model: action_converters: rules: EVA_BIMANUAL: - _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler + _target_: egomimic.utils.action_utils.RobotBimanualCartesian6D HUMAN_BIMANUAL: _target_: egomimic.utils.action_utils.HumanBimanualKeypoints # optional fallback if no match is found diff --git a/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria_6d.yaml b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria_6d.yaml index 301122433b..57159a8908 100644 --- a/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria_6d.yaml +++ b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria_6d.yaml @@ -25,9 +25,9 @@ robomimic_model: action_converters: rules: EVA_BIMANUAL: - _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler + _target_: egomimic.utils.action_utils.RobotBimanualCartesian6D HUMAN_BIMANUAL: - _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + _target_: egomimic.utils.action_utils.HumanBimanualCartesian6D # optional fallback if no match is found fallback: _target_: egomimic.utils.action_utils.BaseActionConverter diff --git a/egomimic/hydra_configs/train_zarr_cartesian_pi.yaml b/egomimic/hydra_configs/train_zarr_cartesian_pi.yaml index c625cbd71c..09ab5c5a56 100644 --- a/egomimic/hydra_configs/train_zarr_cartesian_pi.yaml +++ b/egomimic/hydra_configs/train_zarr_cartesian_pi.yaml @@ -5,3 +5,4 @@ defaults: - override trainer: ddp_pi - override evaluator: eval_pi - override hydra/launcher: submitit_pace + - _self_ diff --git a/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_hpt.yaml b/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_hpt.yaml new file mode 100644 index 0000000000..8a93235d67 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_hpt.yaml @@ -0,0 +1,20 @@ +# Composed recipe: ~300M HPT (+ Qwen text stem) BC on the 18-D wrist-frame +# 6D-rotation cartesian action, mecka FLAGSHIP fold-clothes, seen/unseen +# operator split (see data/mecka_fold_flagship_opsplit_hpt_6d.yaml). +# +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_6d_hpt \ +# norm_stats.precomputed_norm_path=/norm_stats +# +# The norm stats must come from data=mecka_fold_flagship_opsplit_hpt_6d (see +# that file for the precompute command). +defaults: + - train_zarr_cartesian + # Second eval head: metrics + videos on the TRAIN split each val + # (train_viz/). The data config renames the canonical head seen_op_valid/ + # and adds a third, unseen_op_valid/, for the held-out operators. (A plain + # entry, so it must precede the overrides.) + - evaluator@train_viz_evaluator: train_viz_hpt_wristframe_6d + - override model: hpt_bc_mecka_6d_300M + - override data: mecka_fold_flagship_opsplit_hpt_6d + - override evaluator: eval_hpt_wristframe_6d + - _self_ diff --git a/egomimic/hydra_configs/train_zarr_mecka_kp_wrist_hpt.yaml b/egomimic/hydra_configs/train_zarr_mecka_kp_wrist_hpt.yaml new file mode 100644 index 0000000000..59479cd4ed --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_kp_wrist_hpt.yaml @@ -0,0 +1,17 @@ +# Composed recipe: ~300M HPT (+ Qwen text stem) BC on the 144-D wrist-frame +# MANO keypoint action, mecka freeform fold-clothes, held-out-operator split. +# +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_kp_wrist_hpt \ +# norm_stats.precomputed_norm_path=/norm_stats +# +# The norm stats must come from data=mecka_fold_freeform_opsplit_hpt_keypoints +# (see that file for the precompute command). +defaults: + - train_zarr_cartesian + # Second eval head: metrics + videos on TRAIN data each val (train_viz/). + # (A plain entry, so it must precede the overrides.) + - evaluator@train_viz_evaluator: train_viz_hpt_keypoints_wrist + - override model: hpt_bc_keypoints_wrist_300M + - override data: mecka_fold_freeform_opsplit_hpt_keypoints + - override evaluator: eval_hpt_keypoints_wrist + - _self_ diff --git a/egomimic/hydra_configs/trainer/ddp_pi.yaml b/egomimic/hydra_configs/trainer/ddp_pi.yaml index 731316953a..fcfdb55c37 100644 --- a/egomimic/hydra_configs/trainer/ddp_pi.yaml +++ b/egomimic/hydra_configs/trainer/ddp_pi.yaml @@ -3,9 +3,16 @@ defaults: strategy: ddp_find_unused_parameters_true +# Validate under torch.no_grad instead of torch.inference_mode: the compiled +# pi sampler (max-autotune -> cudagraphs) does in-place setup on its static +# buffers during graph capture, which is illegal on inference tensors +# ("Inplace update to inference tensor outside InferenceMode"). no_grad is +# functionally equivalent for eval and creates no inference tensors. +inference_mode: false + accelerator: gpu devices: ${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'} num_nodes: ${launch_params.nodes} sync_batchnorm: True -check_val_every_n_epoch: 200 -num_sanity_val_steps: 0 \ No newline at end of file +check_val_every_n_epoch: 20 +num_sanity_val_steps: 0 diff --git a/egomimic/models/hpt_nets.py b/egomimic/models/hpt_nets.py index 3e76bcee43..3a6d4e0a3d 100644 --- a/egomimic/models/hpt_nets.py +++ b/egomimic/models/hpt_nets.py @@ -1,3 +1,4 @@ +import os from functools import partial from typing import Callable, List, Optional, Union @@ -676,6 +677,28 @@ def _qwen_last_token_pool( return last_hidden_states[batch_idx, seq_lens] +def _local_snapshot_if_offline(model_name: str) -> str: + """Resolve a HF repo id to its cached snapshot directory when the process + runs offline (``HF_HUB_OFFLINE`` / ``TRANSFORMERS_OFFLINE``); a local path + or an online process is returned unchanged.""" + if os.path.isdir(model_name): + return model_name + # huggingface_hub parses these with _is_true, i.e. 1 / ON / YES / TRUE in + # any case, and reads them once at import; match the values but read the + # live environment, since a launcher or a test may set them after import. + from huggingface_hub.constants import ENV_VARS_TRUE_VALUES + + offline = any( + (os.environ.get(var) or "").upper() in ENV_VARS_TRUE_VALUES + for var in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") + ) + if not offline: + return model_name + from huggingface_hub import snapshot_download + + return snapshot_download(model_name, local_files_only=True) + + class _Qwen3BaseEncoder(PolicyStem): """Shared base for Qwen3-Embedding stems used by HPT. @@ -713,8 +736,14 @@ def __init__( self.max_length = max_length self.freeze_encoder = freeze torch_dtype = getattr(torch, dtype) if isinstance(dtype, str) else dtype - self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left") - self.encoder = AutoModel.from_pretrained(model_name, torch_dtype=torch_dtype) + # transformers >= 4.57 probes the Hub while loading a tokenizer by repo + # id even under HF_HUB_OFFLINE=1 (tokenization_utils_base + # ._patch_mistral_regex -> model_info), which raises on compute nodes + # with no internet. Loading from the cached snapshot DIRECTORY skips + # that probe, so resolve the repo id to its local snapshot when offline. + load_path = _local_snapshot_if_offline(model_name) + self.tokenizer = AutoTokenizer.from_pretrained(load_path, padding_side="left") + self.encoder = AutoModel.from_pretrained(load_path, torch_dtype=torch_dtype) if freeze: for p in self.encoder.parameters(): p.requires_grad = False diff --git a/egomimic/scripts/abc_process/__init__.py b/egomimic/scripts/abc_process/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/egomimic/scripts/abc_process/viz_eva_episode.py b/egomimic/scripts/abc_process/viz_eva_episode.py new file mode 100644 index 0000000000..835742e52c --- /dev/null +++ b/egomimic/scripts/abc_process/viz_eva_episode.py @@ -0,0 +1,392 @@ +#!/usr/bin/env python +"""Visualize an EVA (ABC-130k) zarr episode -> mp4. + +Reads the camera frames straight out of the zarr store (no model / no GPU) and +writes a video. Optionally tiles multiple cameras side-by-side and burns in the +task description + frame index. + + ./emimic/bin/python egomimic/scripts/abc_process/viz_eva_episode.py \ + --episode 00022b62-0a07-4ac7-bb74-9b4da47f6c6e \ + --cams front_1 left_wrist right_wrist --out eva_ep.mp4 + +--episode accepts a bare hash (resolved under --folder, which defaults to the +dataset dir from hydra_configs/paths/default.yaml: $EGOVERSE_DATASET_DIR or +/coc/flash7/scratch/egoverseS3ZarrDataset) or a full .zarr path. + +--draw-axes overlays, on the front_1 tile, the per-arm EE pose triads AND the +world-frame origin triad (X=red, Y=green, Z=blue), projected with the SAME +math the training pipeline uses: world pose -> cam frame via +PoseCoordinateFrameTransform against EXTRINSICS[--extrinsics-key], then pixels +via INTRINSICS[--intrinsics-key]. The front_1 tile is stretched to 640x480 +first (the space that K and the abc_fold_viz solvePnP fit live in). Use +--still-frames to also dump overlaid PNGs of specific frames. +""" + +import argparse +import io +import os + +import numpy as np +import zarr + +from egomimic.utils.paths import default_dataset_dir + +DEFAULT_DATASET_DIR = default_dataset_dir() + + +def decode(elem): + """An image stored in the object array -> HxWx3 uint8. + + Frames are JPEG bytes wrapped in a 0-d object ndarray (sometimes nested), so + unwrap to the raw bytes and decode with PIL (imageio can't sniff a bare BytesIO). + """ + while isinstance(elem, np.ndarray) and elem.dtype == object and elem.ndim == 0: + elem = elem.item() + if isinstance(elem, np.ndarray) and elem.ndim >= 2: + return elem[..., :3].astype(np.uint8) + raw = bytes(elem) if not isinstance(elem, (bytes, bytearray)) else elem + from PIL import Image + + return np.asarray(Image.open(io.BytesIO(raw)).convert("RGB"), dtype=np.uint8) + + +# Viz space the overlay lives in: front_1 stretched to 640x480, matching both +# the pipeline's _resize_image_keys and the clicker/solvePnP fit. +VIZ_W, VIZ_H = 640, 480 +AXIS_COLORS = {"x": (255, 0, 0), "y": (0, 255, 0), "z": (0, 0, 255)} # RGB + +# Viz-only calibrations, kept local to this script (training reads intrinsics +# per-episode from zarr.json and bakes Eva.EXTRINSICS; nothing else in the repo +# needs these). Keys that exist on the Eva class are added lazily from it in +# AxisOverlay (see _registries) so this module stays import-light. +# eva K: ABC-130k RealSense top camera (640x480 space) +# yam K: rectified stereo front cam, calibration-derived, 640x480 space +# abc_fold_viz: solvePnP world->cam fit over ABC fold-clothes clicks (one +# shared world frame for both arms); yam: per-arm base->cam hand-eye calib +# eva (extrinsics, from Eva.EXTRINSICS): the per-arm base->cam bake training +# uses; aria (intrinsics, from Eva.INTRINSICS): the class fallback K, +# 640x480 space +INTRINSICS = { + "eva": np.array( + [[436.26, 0.0, 310.1, 0.0], [0.0, 435.13, 241.93, 0.0], [0.0, 0.0, 1.0, 0.0]] + ), + "yam": np.array( + [ + [250.15807185802416, 0.0, 267.5409836065574, 0.0], + [0.0, 260.01789323969166, 152.72727272727272, 0.0], + [0.0, 0.0, 1.0, 0.0], + ] + ), +} +_ABC_FOLD_VIZ = np.array( + [ + [-0.00540728648402, -0.91864395735192, 0.39504941573641, -0.07230219027034], + [-0.99924104589933, -0.01027605228132, -0.03757306135427, -0.23692835817344], + [0.03857581422213, -0.39495275966919, -0.91789118319482, 0.97783070025080], + [0.0, 0.0, 0.0, 1.0], + ] +) +EXTRINSICS = { + "abc_fold_viz": {"left": _ABC_FOLD_VIZ, "right": _ABC_FOLD_VIZ}, + "yam": { + "left": np.array( + [ + [0.0437291, -0.85821391, 0.5114261, 0.08690521], + [-0.99754573, -0.00948836, 0.06937217, -0.25327518], + [-0.05468356, -0.5132045, -0.85652253, 0.83357606], + [0.0, 0.0, 0.0, 1.0], + ] + ), + "right": np.array( + [ + [0.04251619, -0.84174054, 0.53820557, 0.04959916], + [-0.98836195, 0.04331659, 0.14582295, 0.27905066], + [-0.14605832, -0.53814174, -0.83010266, 0.87148985], + [0.0, 0.0, 0.0, 1.0], + ] + ), + }, +} +# Intrinsics keys whose K lives in the stretched 640x480 space (the pipeline's +# _resize_image_keys target); YAM's K is already in the stored image space. +_STRETCHED_INTRINSICS = {"eva", "aria"} + + +def _registries(): + """Local viz calibrations + the values that live on the Eva class.""" + from egomimic.rldb.embodiment.eva import Eva + + return ( + {**INTRINSICS, "aria": Eva.INTRINSICS}, + {**EXTRINSICS, "eva": Eva.EXTRINSICS}, + ) + + +class AxisOverlay: + """Project EE / world triads into front_1 pixels via the pipeline transform.""" + + def __init__(self, extrinsics_key: str, intrinsics_key: str = "eva"): + from egomimic.rldb.zarr.action_chunk_transforms import ( + PoseCoordinateFrameTransform, + ) + from egomimic.utils.pose_utils import _matrix_to_xyzwxyz, _xyzwxyz_to_matrix + + intrinsics, extrinsics = _registries() + self._to_mat = _xyzwxyz_to_matrix + self.K = intrinsics[intrinsics_key] + # EVA's K lives in the pipeline's stretched 640x480 space; YAM's K is + # already in the stored image space -> draw at native resolution there. + self.stretch = intrinsics_key in _STRETCHED_INTRINSICS + extr = extrinsics[extrinsics_key] + # cam<-world pose target per arm, exactly as the Eva transform list bakes it. + self.extr_pose = { + s: _matrix_to_xyzwxyz(np.asarray(extr[s], dtype=np.float64)[None])[0] + for s in ("left", "right") + } + # One shared world frame (e.g. abc_fold_viz) vs per-arm base frames (yam). + self.shared_base = np.allclose(extr["left"], extr["right"]) + self._t = PoseCoordinateFrameTransform( + target_world="extr", + pose_world="pose", + transformed_key_name="out", + mode="xyzwxyz", + ) + self.key = extrinsics_key + + def cam_pose_matrix(self, world_pose_xyzwxyz, side: str) -> np.ndarray: + """World-frame pose (7,) -> 4x4 cam<-ee matrix (pipeline-identical).""" + out = self._t.transform( + { + "extr": self.extr_pose[side], + "pose": np.asarray(world_pose_xyzwxyz, dtype=np.float64), + } + )["out"] + return self._to_mat(np.asarray(out, dtype=np.float64)[None])[0] + + def project(self, pts_cam: np.ndarray) -> np.ndarray: + """(N,3) cam-frame points -> (N,2) pixels; NaN where behind the camera.""" + px = np.full((len(pts_cam), 2), np.nan) + front = pts_cam[:, 2] > 1e-6 + if front.any(): + p = np.concatenate([pts_cam[front], np.ones((front.sum(), 1))], axis=1) + uv = self.K @ p.T + px[front] = (uv[:2] / uv[2]).T + return px + + def draw_triad(self, img, M_cam_pose, length: float, label: str, thickness=2): + """Draw an XYZ triad for a cam<-frame pose matrix onto img (RGB, 640x480).""" + import cv2 + + o = M_cam_pose[:3, 3] + pts = np.stack([o] + [o + M_cam_pose[:3, k] * length for k in range(3)]) + px = self.project(pts) + if np.isnan(px[0]).any(): + return + oi = tuple(np.round(px[0]).astype(int)) + for k, ax in enumerate("xyz"): + if np.isnan(px[k + 1]).any(): + continue + ei = tuple(np.round(px[k + 1]).astype(int)) + cv2.line(img, oi, ei, AXIS_COLORS[ax], thickness, cv2.LINE_AA) + cv2.putText( + img, + ax, + (ei[0] + 3, ei[1] - 3), + cv2.FONT_HERSHEY_SIMPLEX, + 0.45, + AXIS_COLORS[ax], + 1, + cv2.LINE_AA, + ) + cv2.circle(img, oi, 3, (255, 255, 255), -1) + cv2.putText( + img, + label, + (oi[0] + 5, oi[1] + 14), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + (255, 255, 0), + 1, + cv2.LINE_AA, + ) + + def overlay(self, front_native, lpose_w, rpose_w): + """front_1 native frame + world EE poses -> 640x480 RGB with triads. + + Returns (img, info) where info holds the projected EE origin pixels. + """ + import cv2 + + if self.stretch: + img = cv2.resize( + front_native, (VIZ_W, VIZ_H), interpolation=cv2.INTER_LINEAR + ) + else: + # PIL-decoded arrays are read-only; cv2 draws in place -> copy. + img = front_native.copy() + img = np.ascontiguousarray(img) + h = img.shape[0] + info = {} + # Base-frame triad(s): one shared world frame, or per-arm bases (yam). + ident = np.array([0, 0, 0, 1, 0, 0, 0], dtype=np.float64) + bases = ( + [("left", "world")] + if self.shared_base + else [("left", "L-base"), ("right", "R-base")] + ) + for side, lab in bases: + self.draw_triad( + img, self.cam_pose_matrix(ident, side), 0.15, lab, thickness=3 + ) + for side, pose_w, lab in ( + ("left", lpose_w, "L-ee"), + ("right", rpose_w, "R-ee"), + ): + # ABC episodes zero-pad the tail frames -> zero-norm quat; skip them. + if np.linalg.norm(np.asarray(pose_w, dtype=np.float64)[3:7]) < 1e-6: + info[side] = np.array([np.nan, np.nan]) + continue + M = self.cam_pose_matrix(pose_w, side) + self.draw_triad(img, M, 0.08, lab) + info[side] = self.project(M[:3, 3][None])[0] + cv2.putText( + img, + f"axes X=red Y=green Z=blue extr={self.key}", + (6, h - 8), + cv2.FONT_HERSHEY_SIMPLEX, + 0.45, + (255, 255, 255), + 1, + cv2.LINE_AA, + ) + return img, info + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--episode", required=True, help="episode hash or full .zarr path") + ap.add_argument( + "--folder", + default=DEFAULT_DATASET_DIR, + help="zarr dataset dir a bare --episode hash is resolved under " + "(default: $EGOVERSE_DATASET_DIR or the lab dataset dir)", + ) + ap.add_argument( + "--cams", nargs="+", default=["front_1"], help="front_1 left_wrist right_wrist" + ) + ap.add_argument( + "--out", default=None, help="output mp4 (default: ./eva_.mp4)" + ) + ap.add_argument("--max-frames", type=int, default=0, help="0 = all") + ap.add_argument( + "--draw-axes", + action="store_true", + help="overlay EE + world XYZ triads on front_1 (see module docstring)", + ) + ap.add_argument( + "--extrinsics-key", + default="abc_fold_viz", + help="EXTRINSICS key for the world->cam bake " + "(abc_fold_viz | yam | eva; default: abc_fold_viz)", + ) + ap.add_argument( + "--intrinsics-key", + default="eva", + help="INTRINSICS key for projection (eva | yam | aria)", + ) + ap.add_argument( + "--still-frames", + type=int, + nargs="*", + default=[], + help="also dump these frame indices as overlaid PNGs next to --out", + ) + args = ap.parse_args() + + path = ( + args.episode + if args.episode.endswith(".zarr") + else os.path.join(args.folder, f"{args.episode}.zarr") + ) + g = zarr.open_group(path, mode="r") + a = dict(g.attrs) + fps = int(a.get("fps", 30)) + n = int(a.get("total_frames", 0)) + print( + f"[viz] {path}\n task={a.get('task_description') or a.get('task_name')!r} " + f"frames={n} fps={fps} cams={args.cams}" + ) + + overlay = None + lobs = robs = None + if args.draw_axes: + if "front_1" not in args.cams: + args.cams = ["front_1"] + args.cams + overlay = AxisOverlay(args.extrinsics_key, args.intrinsics_key) + lobs, robs = g["left.obs_ee_pose"], g["right.obs_ee_pose"] + print( + f"[viz] axis overlay on front_1: extrinsics_key={args.extrinsics_key}, " + f"K=INTRINSICS['{args.intrinsics_key}'], stretch={overlay.stretch}" + ) + + cam_arrs = {c: g[f"images.{c}"] for c in args.cams} + # Arrays are zero-padded past attrs total_frames -> cap at the real length. + nf = min( + next(iter(cam_arrs.values())).shape[0], n or 10**9, args.max_frames or 10**9 + ) + out = args.out or f"eva_{os.path.basename(path).replace('.zarr', '')}.mp4" + stills = sorted(set(i for i in args.still_frames if 0 <= i < nf)) + + import imageio.v2 as imageio + + try: + from PIL import Image, ImageDraw + except Exception: + Image = None + writer = imageio.get_writer(out, fps=fps, macro_block_size=None) + for i in range(nf): + tiles = [] + for c in args.cams: + f = decode(cam_arrs[c][i]) + if overlay is not None and c == "front_1": + f, info = overlay.overlay(f, np.asarray(lobs[i]), np.asarray(robs[i])) + if i in stills: + print( + f"[viz] frame {i}: EE pixels L={np.round(info['left'], 1)} " + f"R={np.round(info['right'], 1)}" + ) + tiles.append(f) + h = min(t.shape[0] for t in tiles) + tiles = ( + [ + t[:h] + if t.shape[0] == h + else np.asarray( + Image.fromarray(t).resize((int(t.shape[1] * h / t.shape[0]), h)) + ) + for t in tiles + ] + if Image + else tiles + ) + frame = np.concatenate(tiles, axis=1) + if Image: + im = Image.fromarray(frame) + d = ImageDraw.Draw(im) + d.text( + (6, 6), + f"{i}/{nf} {(a.get('task_description') or '')[:60]}", + fill=(255, 255, 0), + ) + frame = np.asarray(im) + if i in stills: + still_path = f"{os.path.splitext(out)[0]}_f{i}.png" + Image.fromarray(frame).save(still_path) + print(f"[viz] still -> {still_path}") + writer.append_data(frame) + writer.close() + print(f"[viz] wrote {nf} frames -> {out}") + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/mecka_process/viz_mecka_episode.py b/egomimic/scripts/mecka_process/viz_mecka_episode.py new file mode 100644 index 0000000000..ef62e4f2bd --- /dev/null +++ b/egomimic/scripts/mecka_process/viz_mecka_episode.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python +"""Visualize a Mecka (egocentric human) zarr episode -> mp4 with frame overlays. + +Egocentric analog of abc_process/viz_eva_episode.py --draw-axes: projects into +the moving head camera (per-frame ``obs_head_pose`` is the world->cam target, +exactly the headframe transform training uses) with MECKA_INTRINSICS. + +Overlays, per frame: + - L/R wrist triads (X=red Y=green Z=blue) from ``obs_ee_pose`` — shows the + hand-tracking wrist-frame convention (diagonal to hand geometry, mirrored + between hands, unlike EVA/YAM tool frames). + - hand keypoints (gray dots) + a WHITE line wrist-root -> middle-MCP = + "knuckle-forward", the physical reference to compare the triads against. + - the WORLD origin triad (thick, length 0.3m): origin sits on the floor under + the operator's head at episode start, Z up (gravity-aligned). + + ./emimic/bin/python egomimic/scripts/mecka_process/viz_mecka_episode.py \ + --episode 696d0f031244464a4aac1b8f --still-frames 150 450 750 + +--episode accepts a bare hash (resolved under --folder, which defaults to the +dataset dir from hydra_configs/paths/default.yaml: $EGOVERSE_DATASET_DIR or +/coc/flash7/scratch/egoverseS3ZarrDataset) or a full .zarr path. +""" + +import argparse +import io +import os + +import numpy as np +import zarr + +from egomimic.utils.paths import default_dataset_dir + +DEFAULT_DATASET_DIR = default_dataset_dir() + + +def decode(elem): + while isinstance(elem, np.ndarray) and elem.dtype == object and elem.ndim == 0: + elem = elem.item() + if isinstance(elem, np.ndarray) and elem.ndim >= 2: + return elem[..., :3].astype(np.uint8) + raw = bytes(elem) if not isinstance(elem, (bytes, bytearray)) else elem + from PIL import Image + + return np.asarray(Image.open(io.BytesIO(raw)).convert("RGB"), dtype=np.uint8) + + +AXIS_COLORS = {"x": (255, 0, 0), "y": (0, 255, 0), "z": (0, 0, 255)} # RGB +# MediaPipe-style indices (verified on-data: kp0 roots all four finger chains, +# adjacent MCPs 5/9/13/17 are ~2.5cm apart). NOT Aria ordering. +WRIST_ROOT, MIDDLE_MCP = 0, 9 + + +class MeckaOverlay: + """World-frame geometry -> head-camera pixels via the training headframe math.""" + + def __init__(self): + # Intrinsics live on the embodiment module; imported lazily so --help + # stays cheap. + from egomimic.rldb.embodiment.human import MECKA_INTRINSICS + from egomimic.rldb.zarr.action_chunk_transforms import ( + PoseCoordinateFrameTransform, + ) + from egomimic.utils.pose_utils import _xyzwxyz_to_matrix + + self.K = MECKA_INTRINSICS + self._to_mat = _xyzwxyz_to_matrix + self._t_pose = PoseCoordinateFrameTransform( + target_world="head", + pose_world="pose", + transformed_key_name="out", + mode="xyzwxyz", + ) + self._t_xyz = PoseCoordinateFrameTransform( + target_world="head", + pose_world="pose", + transformed_key_name="out", + mode="xyz", + ) + + def cam_pose_matrix(self, head_pose, world_pose): + out = self._t_pose.transform( + { + "head": np.asarray(head_pose, np.float64), + "pose": np.asarray(world_pose, np.float64), + } + )["out"] + return self._to_mat(np.asarray(out, np.float64)[None])[0] + + def cam_xyz(self, head_pose, world_xyz): + """(N,3) world points -> (N,3) cam-frame points.""" + return np.stack( + [ + self._t_xyz.transform( + { + "head": np.asarray(head_pose, np.float64), + "pose": np.asarray(p, np.float64), + } + )["out"] + for p in np.asarray(world_xyz, np.float64).reshape(-1, 3) + ] + ) + + def project(self, pts_cam): + px = np.full((len(pts_cam), 2), np.nan) + front = pts_cam[:, 2] > 1e-6 + if front.any(): + p = np.concatenate([pts_cam[front], np.ones((front.sum(), 1))], axis=1) + uv = self.K @ p.T + px[front] = (uv[:2] / uv[2]).T + return px + + def draw_triad(self, img, M, length, label, thickness=2): + import cv2 + + o = M[:3, 3] + pts = np.stack([o] + [o + M[:3, k] * length for k in range(3)]) + px = self.project(pts) + if np.isnan(px[0]).any(): + return + oi = tuple(np.round(px[0]).astype(int)) + for k, ax in enumerate("xyz"): + if np.isnan(px[k + 1]).any(): + continue + ei = tuple(np.round(px[k + 1]).astype(int)) + cv2.line(img, oi, ei, AXIS_COLORS[ax], thickness, cv2.LINE_AA) + cv2.putText( + img, + ax, + (ei[0] + 3, ei[1] - 3), + cv2.FONT_HERSHEY_SIMPLEX, + 0.4, + AXIS_COLORS[ax], + 1, + cv2.LINE_AA, + ) + cv2.circle(img, oi, 3, (255, 255, 255), -1) + cv2.putText( + img, + label, + (oi[0] + 5, oi[1] + 14), + cv2.FONT_HERSHEY_SIMPLEX, + 0.45, + (255, 255, 0), + 1, + cv2.LINE_AA, + ) + + def overlay(self, img, head, lpose, rpose, lkp, rkp): + import cv2 + + # PIL-decoded arrays are read-only; cv2 draws in place -> copy. + img = np.ascontiguousarray(img.copy()) + # world origin triad (floor under head-at-start, Z up) + if np.linalg.norm(np.asarray(head)[3:7]) > 1e-6: + Mw = self.cam_pose_matrix(head, np.array([0, 0, 0, 1, 0, 0, 0.0])) + self.draw_triad(img, Mw, 0.3, "world", thickness=3) + for pose, kp, lab in ((lpose, lkp, "L-wrist"), (rpose, rkp, "R-wrist")): + if ( + np.linalg.norm(np.asarray(pose)[3:7]) < 1e-6 + or np.abs(kp).sum() < 1e-9 + ): + continue + # keypoints + knuckle-forward reference line + pc = self.cam_xyz(head, kp) + px = self.project(pc) + for p in px: + if not np.isnan(p).any(): + cv2.circle( + img, tuple(np.round(p).astype(int)), 2, (200, 200, 200), -1 + ) + a, b = px[WRIST_ROOT], px[MIDDLE_MCP] + if not (np.isnan(a).any() or np.isnan(b).any()): + cv2.line( + img, + tuple(np.round(a).astype(int)), + tuple(np.round(b).astype(int)), + (255, 255, 255), + 2, + cv2.LINE_AA, + ) + # the wrist-frame triad the pipeline actually uses + M = self.cam_pose_matrix(head, pose) + self.draw_triad(img, M, 0.07, lab) + cv2.putText( + img, + "axes X=red Y=green Z=blue | white line = knuckle-forward", + (6, img.shape[0] - 8), + cv2.FONT_HERSHEY_SIMPLEX, + 0.42, + (255, 255, 255), + 1, + cv2.LINE_AA, + ) + return img + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--episode", required=True, help="episode hash or full .zarr path") + ap.add_argument( + "--folder", + default=DEFAULT_DATASET_DIR, + help="zarr dataset dir a bare --episode hash is resolved under " + "(default: $EGOVERSE_DATASET_DIR or the lab dataset dir)", + ) + ap.add_argument( + "--out", default=None, help="output mp4 (default: ./mecka__axes.mp4)" + ) + ap.add_argument("--max-frames", type=int, default=0, help="0 = all") + ap.add_argument( + "--still-frames", + type=int, + nargs="*", + default=[], + help="also dump these frame indices as overlaid PNGs next to --out", + ) + args = ap.parse_args() + + path = ( + args.episode + if args.episode.endswith(".zarr") + else os.path.join(args.folder, f"{args.episode}.zarr") + ) + g = zarr.open_group(path, mode="r") + a = dict(g.attrs) + fps = int(a.get("fps", 30)) + print(f"[viz] {path}\n task={(a.get('task_description') or '')[:80]!r} fps={fps}") + + imgs = g["images.front_1"] + head = np.asarray(g["obs_head_pose"]) + lp, rp = np.asarray(g["left.obs_ee_pose"]), np.asarray(g["right.obs_ee_pose"]) + lk = np.asarray(g["left.obs_keypoints"]).reshape(len(head), 21, 3) + rk = np.asarray(g["right.obs_keypoints"]).reshape(len(head), 21, 3) + nf = min(imgs.shape[0], len(head), args.max_frames or 10**9) + out = args.out or f"mecka_{os.path.basename(path).replace('.zarr', '')}_axes.mp4" + stills = sorted(set(i for i in args.still_frames if 0 <= i < nf)) + + ov = MeckaOverlay() + import imageio.v2 as imageio + from PIL import Image, ImageDraw + + writer = imageio.get_writer(out, fps=fps, macro_block_size=None) + for i in range(nf): + frame = ov.overlay(decode(imgs[i]), head[i], lp[i], rp[i], lk[i], rk[i]) + im = Image.fromarray(frame) + d = ImageDraw.Draw(im) + d.text( + (6, 6), + f"{i}/{nf} {(a.get('task_description') or '')[:60]}", + fill=(255, 255, 0), + ) + frame = np.asarray(im) + if i in stills: + still = f"{os.path.splitext(out)[0]}_f{i}.png" + Image.fromarray(frame).save(still) + print(f"[viz] still -> {still}") + writer.append_data(frame) + writer.close() + print(f"[viz] wrote {nf} frames -> {out}") + + +if __name__ == "__main__": + main() diff --git a/egomimic/utils/action_utils.py b/egomimic/utils/action_utils.py index 399e0edb59..063410d292 100644 --- a/egomimic/utils/action_utils.py +++ b/egomimic/utils/action_utils.py @@ -270,7 +270,7 @@ def from32_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: # ============================================================ -class RobotBimanualCartesianEuler(BaseActionConverter): +class RobotBimanualCartesian6D(BaseActionConverter): """ Native: (B,S,20) = [L xyz(3) 6d(6) g(1) | R xyz(3) 6d(6) g(1)] 32-pack: left block 0..9, right block 10..19, zero pad 20..31 @@ -301,7 +301,7 @@ def from32_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: # ============================================================ -class HumanBimanualCartesianEuler(BaseActionConverter): +class HumanBimanualCartesian6D(BaseActionConverter): """ Native: (B,S,18) = [L xyz(3) 6d(6) | R xyz(3) 6d(6)] (no gripper) 32-pack: left block 0..9 (g=0), right block 10..19 (g=0), zero pad 20..31 diff --git a/egomimic/utils/paths.py b/egomimic/utils/paths.py new file mode 100644 index 0000000000..afed066cec --- /dev/null +++ b/egomimic/utils/paths.py @@ -0,0 +1,21 @@ +"""Path defaults for scripts that run outside hydra. + +A standalone script (a viz tool, a probe) still wants the cluster's dataset +mirror, and a second copy of the path silently drifts from the config the +training runs use. This reads the one hydra defines. +""" + +from __future__ import annotations + +from pathlib import Path + +from omegaconf import OmegaConf + +import egomimic.hydra_configs as _cfg_pkg + +_PATHS_YAML = Path(_cfg_pkg.__file__).parent / "paths" / "default.yaml" + + +def default_dataset_dir() -> str: + """``paths.dataset_dir`` from the hydra configs, env override included.""" + return str(OmegaConf.load(_PATHS_YAML).dataset_dir) diff --git a/tests/fixtures/recipes.py b/tests/fixtures/recipes.py index bce90a939f..78b90de88a 100644 --- a/tests/fixtures/recipes.py +++ b/tests/fixtures/recipes.py @@ -22,13 +22,10 @@ class Recipe: def _pi(data: str, model: str, emb: str) -> Recipe: - return Recipe( - "train_zarr_cartesian_pi", - data, - model, - emb, - (f"+data.train_datasets.{emb}.resolver.key_map.annotation_key=annotations",), - ) + # data/eva.yaml already emits the continuous-6D (wrist-frame) layout Pi + # requires, the same mode cotrain_pi_base / eva_pi set. + extra = [f"+data.train_datasets.{emb}.resolver.key_map.annotation_key=annotations"] + return Recipe("train_zarr_cartesian_pi", data, model, emb, tuple(extra)) RECIPES: dict[tuple[str, str], Recipe] = { diff --git a/tests/integration/test_train_step_gpu.py b/tests/integration/test_train_step_gpu.py index 561eabd857..9b348a66a2 100644 --- a/tests/integration/test_train_step_gpu.py +++ b/tests/integration/test_train_step_gpu.py @@ -132,7 +132,12 @@ def test_hpt_real_recipe(tmp_path, monkeypatch, vendor): _vendor_filter(cfg, recipe.embodiment, vendor) spy = BatchKeySpy(monkeypatch, HPT) _run(cfg, out) - assert "actions_cartesian" in spy.keys[recipe.embodiment] + action_key = ( + "actions_cartesian" + if recipe.embodiment == "eva_bimanual" + else "actions_keypoints" # human data defaults to hand keypoints + ) + assert action_key in spy.keys[recipe.embodiment] @pytest.mark.parametrize("vendor", VENDOR_NAMES) diff --git a/tests/unit/test_compile_cache.py b/tests/unit/test_compile_cache.py index b567423cd8..21b01b234c 100644 --- a/tests/unit/test_compile_cache.py +++ b/tests/unit/test_compile_cache.py @@ -32,12 +32,6 @@ def test_outside_slurm_keeps_torch_default(): assert env == {} -def test_explicit_value_is_respected(): - env = {KEY: "/scratch/ind", "SLURM_JOB_ID": "1"} - assert set_per_job_compile_cache_dir(env) == "/scratch/ind" - assert env[KEY] == "/scratch/ind" - - _LAUNCHER = """ import os, subprocess, sys from egomimic.utils.compile_cache import set_per_job_compile_cache_dir as s diff --git a/tests/unit/test_data_configs_compose.py b/tests/unit/test_data_configs_compose.py index 143bda48a8..88841b529c 100644 --- a/tests/unit/test_data_configs_compose.py +++ b/tests/unit/test_data_configs_compose.py @@ -9,6 +9,7 @@ import hydra import pytest +import yaml from hydra import compose, initialize_config_module from omegaconf import OmegaConf @@ -137,6 +138,35 @@ def test_no_data_config_uses_a_pi_keymap_mode(name): assert not str(mode).endswith("_pi"), (name, split, emb, mode) +# Aliases kept so configs saved by earlier runs still rebuild. +KNOWN_MODE_ONLY_CHILDREN = {"cotrain_pi_lang_wrist"} +MODE_KEYS = {"keymap_mode", "mode", "include_ee_pose"} + + +def _leaf_paths(node, prefix=()): + if isinstance(node, dict) and node: + for k, v in node.items(): + yield from _leaf_paths(v, (*prefix, k)) + else: + yield prefix + + +@pytest.mark.parametrize("name", DATA_CONFIGS) +def test_no_data_config_only_swaps_the_action_mode(name): + # A child that only changes keymap/transform mode (e.g. cam- vs wrist-frame) + # is a CLI override, not a data config; one data config per dataset slice. + raw = yaml.safe_load((DATA_DIR / f"{name}.yaml").read_text(encoding="utf-8")) + body = {k: v for k, v in (raw or {}).items() if k != "defaults"} + if not raw.get("defaults") or not body: + return + leaves = list(_leaf_paths(body)) + mode_only = all(path[-1] in MODE_KEYS for path in leaves) + if name in KNOWN_MODE_ONLY_CHILDREN: + assert mode_only, f"{name} is no longer mode-only; drop it from the allowlist" + else: + assert not mode_only, (name, leaves) + + def test_pi_train_config_defaults_to_all_eva_with_annotations(): # data=eva is a one-episode smoke config without annotations; Pi's default # must not silently train on that with an empty prompt. diff --git a/tests/unit/test_dataset_filter.py b/tests/unit/test_dataset_filter.py index fd0766ec45..48170797db 100644 --- a/tests/unit/test_dataset_filter.py +++ b/tests/unit/test_dataset_filter.py @@ -22,13 +22,6 @@ def test_dataset_filter_matches_rows_and_excludes_deleted_by_default() -> None: assert not filters.matches({"episode_hash": "episode-2"}) -def test_dataset_filter_empty_list_matches_all_non_deleted_rows() -> None: - filters = DatasetFilter() - - assert filters.matches({"episode_hash": "episode-1"}) - assert not filters.matches({"episode_hash": "episode-1", "is_deleted": True}) - - def test_dataset_filter_init_rejects_invalid_filter_and_prints_it(capsys) -> None: with pytest.raises(ValueError, match="Invalid filter"): DatasetFilter(filter_lambdas=["lambda row:"]) diff --git a/tests/unit/test_eva_wristframe_norm_pool.py b/tests/unit/test_eva_wristframe_norm_pool.py index 18b16a3f7b..cec5e3bd76 100644 --- a/tests/unit/test_eva_wristframe_norm_pool.py +++ b/tests/unit/test_eva_wristframe_norm_pool.py @@ -1,4 +1,4 @@ -"""Horizon-pooled norm stats and the constant-channel normalization guard.""" +"""Eva wrist-frame defaults, horizon-pooled norm stats, per-segment L2 metrics.""" import numpy as np import pytest @@ -11,6 +11,8 @@ from egomimic.rldb.zarr import norm_cache from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver, MultiDataset +EVA_WRIST_REVERT = "egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list" + def _compose(name, overrides=()): with initialize_config_module( @@ -19,6 +21,15 @@ def _compose(name, overrides=()): return compose(config_name=name, overrides=list(overrides)) +@pytest.mark.parametrize("top", ["train_zarr_cartesian", "train_zarr_cartesian_pi"]) +def test_eva_defaults_are_wrist_frame_6d(top): + cfg = _compose(top) + eva = cfg.data.train_datasets.eva_bimanual + assert eva.resolver.transform_list.mode == "cartesian_wristframe_6d" + assert cfg.evaluator.transform_lists.eva_bimanual._target_ == EVA_WRIST_REVERT + assert cfg.norm_stats.pool_horizon is False + + def test_hpt_eva_dims_match_the_wrist_frame_6d_layout(): cfg = _compose("train_zarr_cartesian") dims = cfg.model.robomimic_model.dims.eva_bimanual @@ -74,6 +85,30 @@ def test_pool_horizon_is_part_of_the_cache_key(): ) != norm_cache.norm_cache_key(norm_cache.cache_inputs(*args)) +def test_segment_l2_separates_early_from_late_error(): + B, T = 2, 100 + gt = torch.zeros(B, T, 20) + pred = gt.clone() + pred[:, :10, 0] = 0.03 # 3 cm on the left arm x, first 10 % only + m = cartesian_metrics(pred, gt, "p") + # mean over both arms: 3 cm on one of two + assert m["p_xyz_l2_early"].item() == pytest.approx(0.015) + assert m["p_xyz_l2_mid"].item() == pytest.approx(0.0) + assert m["p_xyz_l2_late"].item() == pytest.approx(0.0) + + +def test_segment_l2_on_keypoints(): + B, T = 1, 100 + gt = torch.zeros(B, T, 144) + pred = gt.clone() + pred[:, 50:, 9:72] = 0.02 # every left-hand keypoint off by (2,2,2) cm late + m = keypoint_metrics(pred, gt, "k") + expected = 0.5 * float(np.linalg.norm([0.02] * 3)) # one hand of two + assert m["k_kp_l2_late"].item() == pytest.approx(expected, rel=1e-5) + assert m["k_kp_l2_early"].item() == pytest.approx(0.0) + assert m["k_wrist_xyz_l2_late"].item() == pytest.approx(0.0) + + # ------------------------------------------------- degenerate-range guard @pytest.mark.parametrize("mode", ["zscore", "minmax", "quantile"]) def test_norm_constant_channel_maps_to_zero_and_back_to_its_centre(mode): @@ -234,39 +269,3 @@ def test_bounds_check_ignores_constant_cells(): assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is None arr[1, 0] = 50.0 # corrupt value at a regular cell: still rejected assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is not None - - -EVA_WRIST_REVERT = "egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list" - - -@pytest.mark.parametrize("top", ["train_zarr_cartesian", "train_zarr_cartesian_pi"]) -def test_eva_defaults_are_wrist_frame_6d(top): - cfg = _compose(top) - eva = cfg.data.train_datasets.eva_bimanual - assert eva.resolver.transform_list.mode == "cartesian_wristframe_6d" - assert cfg.evaluator.transform_lists.eva_bimanual._target_ == EVA_WRIST_REVERT - assert cfg.norm_stats.pool_horizon is False - - -def test_segment_l2_separates_early_from_late_error(): - B, T = 2, 100 - gt = torch.zeros(B, T, 20) - pred = gt.clone() - pred[:, :10, 0] = 0.03 # 3 cm on the left arm x, first 10 % only - m = cartesian_metrics(pred, gt, "p") - # mean over both arms: 3 cm on one of two - assert m["p_xyz_l2_early"].item() == pytest.approx(0.015) - assert m["p_xyz_l2_mid"].item() == pytest.approx(0.0) - assert m["p_xyz_l2_late"].item() == pytest.approx(0.0) - - -def test_segment_l2_on_keypoints(): - B, T = 1, 100 - gt = torch.zeros(B, T, 144) - pred = gt.clone() - pred[:, 50:, 9:72] = 0.02 # every left-hand keypoint off by (2,2,2) cm late - m = keypoint_metrics(pred, gt, "k") - expected = 0.5 * float(np.linalg.norm([0.02] * 3)) # one hand of two - assert m["k_kp_l2_late"].item() == pytest.approx(expected, rel=1e-5) - assert m["k_kp_l2_early"].item() == pytest.approx(0.0) - assert m["k_wrist_xyz_l2_late"].item() == pytest.approx(0.0) diff --git a/tests/unit/test_infra_port_wiring.py b/tests/unit/test_infra_port_wiring.py index d9559b41bb..58a50a0119 100644 --- a/tests/unit/test_infra_port_wiring.py +++ b/tests/unit/test_infra_port_wiring.py @@ -1,5 +1,6 @@ """Wiring checks for the hand-keypoint default: recipes, stems, the mecka -left-wrist fix in keypoint modes, and the packed-width loss reduction.""" +left-wrist fix in keypoint modes, the annotation-cutoff span filter and the +train_viz second val loader.""" from __future__ import annotations @@ -83,7 +84,14 @@ def test_human_action_width_matches_data_mode(top, compose_resolve): ("pi0.5_bc_mecka", "mecka", []), ("pi0.5_bc_scale", "scale", []), ("pi0.5_cotrain_eva_aria", "cotrain_pi_base", []), + ("pi0.5_bc_mecka_6d", "mecka_all_6d", []), ("pi0.5_cotrain_eva_aria_6d", "cotrain_pi_lang", CAM_FRAME_COTRAIN), + ("hpt_bc_mecka_6d_300M", "mecka_fold_flagship_opsplit_hpt_6d", []), + ( + "hpt_bc_keypoints_wrist_300M", + "mecka_fold_freeform_opsplit_hpt_keypoints", + [], + ), ], ) def test_vendor_pairings_agree_on_the_human_action(model, data, extra, compose_resolve): @@ -116,7 +124,7 @@ def test_hpt_stems_follow_the_human_action(compose_resolve): (a base's ``state_ee_pose`` is nulled out, not inherited).""" from egomimic.algo.hpt import HPTModel - for model in ("hpt_bc_flow_human", "hpt_cotrain_flow_seperate_head"): + for model in ("hpt_bc_flow_mecka", "hpt_cotrain_flow_seperate_head"): cfg = compose_resolve("train_zarr_cartesian", [f"model={model}"]) stems = cfg.model.robomimic_model.stem_specs[HUMAN] live = {k for k, v in stems.items() if v is not None} @@ -231,36 +239,6 @@ def test_fix_left_wrist_convention_equals_reconverting_in_keypoint_mode(): np.testing.assert_allclose(a[k], b[k], atol=1e-9, err_msg=k) -def test_pi_loss_is_reduced_over_the_packed_width(): - from egomimic.algo.pi import PI - from egomimic.rldb.embodiment.embodiment import EMBODIMENT - from egomimic.utils.action_utils import ( - ConverterRegistry, - HumanBimanualKeypoints, - RobotBimanualCartesianEuler, - ) - - pi = PI.__new__(PI) - pi.action_registry = ConverterRegistry() - pi._packed_widths = {} - eva, human = EMBODIMENT.EVA_BIMANUAL.value, EMBODIMENT.HUMAN_BIMANUAL.value - pi.action_registry.register(eva, "actions_cartesian", RobotBimanualCartesianEuler()) - pi.action_registry.register(human, "actions_keypoints", HumanBimanualKeypoints()) - losses = torch.zeros(2, 5, 144) - losses[..., 32:] = 1.0 # error only in the zero-padded slots - eva_loss = pi._reduce_loss(losses, torch.zeros(2, 5, 20), eva, "actions_cartesian") - assert eva_loss.item() == 0.0 - kp_loss = pi._reduce_loss( - losses, torch.zeros(2, 5, 144), human, "actions_keypoints" - ) - assert abs(kp_loss.item() - 112 / 144) < 1e-6 - # already-reduced losses pass through unchanged - assert ( - pi._reduce_loss(torch.tensor([1.0, 3.0]), None, eva, "actions_cartesian").item() - == 2.0 - ) - - # --------------------------------------------------- annotation-cutoff filter def test_episode_has_annotation_spans(): from egomimic.rldb.zarr.zarr_dataset_multi import _episode_has_annotation_spans @@ -313,56 +291,6 @@ def test_annotation_cutoff_resolver_defaults_to_requiring_spans(): m.S3EpisodeResolver.resolve = orig -# ------------------------------------------------ adversarial-review fixes -def test_bounds_quantiles_include_the_observed_extremes(): - """With fewer than 10k samples a linear 0.01 / 99.99 percentile lands - strictly inside the sample range, so the stats' own extreme frames were - rejected at train time; the bounds quantiles must snap to samples.""" - from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset - - X = np.random.default_rng(0).normal(size=(500, 4)).astype(np.float32) - st = MultiDataset._compute_stats_for_array(X) - np.testing.assert_array_equal(st["quantile_0_01"], X.min(axis=0)) - np.testing.assert_array_equal(st["quantile_99_99"], X.max(axis=0)) - md = MultiDataset.__new__(MultiDataset) - md.norm_mode = "quantile" - md.norm_stats = {0: {"actions_keypoints": st}} - md.zarr_keys = {0: {"actions_keypoints": "actions_keypoints"}} - md._warned_violations = set() - for row in (X.min(axis=0), X.max(axis=0)): - assert ( - md._check_bounds( - {"embodiment": 0, "actions_keypoints": row[None]}, None, 0, "ep" - ) - is None - ) - - -def test_bounds_check_has_relative_slack_but_catches_corrupt_values(): - """Per-cell bounds tolerate frames moderately beyond the stats sample's - range (valid extreme motion) and still reject values orders of magnitude - off (fill constants, wrong-frame data).""" - from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset - - md = MultiDataset.__new__(MultiDataset) - md.norm_mode = "quantile" - lo, hi = np.full(18, -1.0, np.float32), np.full(18, 1.0, np.float32) - md.norm_stats = {0: {"actions_cartesian": {"quantile_1": lo, "quantile_99": hi}}} - md.zarr_keys = {0: {"actions_cartesian": "actions_cartesian"}} - md._warned_violations = set() - - def check(v): - arr = np.zeros((3, 18), np.float32) - arr[1, 0] = v # an xyz channel - return md._check_bounds( - {"embodiment": 0, "actions_cartesian": arr}, None, 0, "e" - ) - - assert check(1.0 + 0.4 * 2.0) is None # within 50 % of the [-1, 1] range - assert check(1.0 + 0.6 * 2.0) is not None - assert check(1e9) is not None - - # ------------------------------------------------------ train_viz second loader class _Dicts: def __init__(self, n, tag): @@ -485,6 +413,108 @@ def on_validation_step(self, batch, batch_idx, dataloader_idx=0): assert calls == [("valid", 0), ("unseen_op_valid", 1)] +def test_flagship_opsplit_val_heads(compose_resolve): + """Flagship opsplit: valid (prefixed seen_op_valid) = seen operators' + held-out episodes (complement of train), train_viz = the train split (no + explicit datasets), unseen_op_valid = the held-out operators. The topop + twin inherits the same layout.""" + for data, train_op in ( + ("mecka_fold_flagship_opsplit_hpt_6d", "not in"), + ("mecka_fold_flagship_topop_hpt_6d", "== '6903686e0e94ce070afd1f24'"), + ): + cfg = compose_resolve("train_zarr_mecka_flagship_6d_hpt", [f"data={data}"]) + d = cfg.data + train, valid = d.train_datasets[HUMAN], d.valid_datasets[HUMAN] + unseen = d.unseen_op_valid_datasets[HUMAN] + assert d.get("train_viz_datasets") is None + assert d.valid_prefix == "seen_op_valid" + assert (train.mode, valid.mode, unseen.mode) == ("train", "valid", "total") + assert valid.valid_ratio == train.valid_ratio + assert list(valid.filters.filter_lambdas) == list(train.filters.filter_lambdas) + assert train_op in train.filters.filter_lambdas[-1] + ops = list(d.held_out_operators) + assert unseen.filters.filter_lambdas[-1].endswith(f" in {ops}") + + +def test_bounds_quantiles_include_the_observed_extremes(): + """With fewer than 10k samples a linear 0.01 / 99.99 percentile lands + strictly inside the sample range, so the stats' own extreme frames were + rejected at train time; the bounds quantiles must snap to samples.""" + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + X = np.random.default_rng(0).normal(size=(500, 4)).astype(np.float32) + st = MultiDataset._compute_stats_for_array(X) + np.testing.assert_array_equal(st["quantile_0_01"], X.min(axis=0)) + np.testing.assert_array_equal(st["quantile_99_99"], X.max(axis=0)) + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + md.norm_stats = {0: {"actions_keypoints": st}} + md.zarr_keys = {0: {"actions_keypoints": "actions_keypoints"}} + md._warned_violations = set() + for row in (X.min(axis=0), X.max(axis=0)): + assert ( + md._check_bounds( + {"embodiment": 0, "actions_keypoints": row[None]}, None, 0, "ep" + ) + is None + ) + + +def test_pi_loss_is_reduced_over_the_packed_width(): + from egomimic.algo.pi import PI + from egomimic.rldb.embodiment.embodiment import EMBODIMENT + from egomimic.utils.action_utils import ( + ConverterRegistry, + HumanBimanualKeypoints, + RobotBimanualCartesian6D, + ) + + pi = PI.__new__(PI) + pi.action_registry = ConverterRegistry() + pi._packed_widths = {} + eva, human = EMBODIMENT.EVA_BIMANUAL.value, EMBODIMENT.HUMAN_BIMANUAL.value + pi.action_registry.register(eva, "actions_cartesian", RobotBimanualCartesian6D()) + pi.action_registry.register(human, "actions_keypoints", HumanBimanualKeypoints()) + losses = torch.zeros(2, 5, 144) + losses[..., 32:] = 1.0 # error only in the zero-padded slots + eva_loss = pi._reduce_loss(losses, torch.zeros(2, 5, 20), eva, "actions_cartesian") + assert eva_loss.item() == 0.0 + kp_loss = pi._reduce_loss( + losses, torch.zeros(2, 5, 144), human, "actions_keypoints" + ) + assert abs(kp_loss.item() - 112 / 144) < 1e-6 + # already-reduced losses pass through unchanged + assert ( + pi._reduce_loss(torch.tensor([1.0, 3.0]), None, eva, "actions_cartesian").item() + == 2.0 + ) + + +def test_bounds_check_has_relative_slack_but_catches_corrupt_values(): + """Per-cell bounds tolerate frames moderately beyond the stats sample's + range (valid extreme motion) and still reject values orders of magnitude + off (fill constants, wrong-frame data).""" + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + lo, hi = np.full(18, -1.0, np.float32), np.full(18, 1.0, np.float32) + md.norm_stats = {0: {"actions_cartesian": {"quantile_1": lo, "quantile_99": hi}}} + md.zarr_keys = {0: {"actions_cartesian": "actions_cartesian"}} + md._warned_violations = set() + + def check(v): + arr = np.zeros((3, 18), np.float32) + arr[1, 0] = v # an xyz channel + return md._check_bounds( + {"embodiment": 0, "actions_cartesian": arr}, None, 0, "e" + ) + + assert check(1.0 + 0.4 * 2.0) is None # within 50 % of the [-1, 1] range + assert check(1.0 + 0.6 * 2.0) is not None + assert check(1e9) is not None + + # ------------------------------------------------------ train_viz on by default def _viz_cfg(**top): from omegaconf import OmegaConf diff --git a/tests/unit/test_latent_dataset_kwargs.py b/tests/unit/test_latent_dataset_kwargs.py index 237f9715ad..f9f2526bc4 100644 --- a/tests/unit/test_latent_dataset_kwargs.py +++ b/tests/unit/test_latent_dataset_kwargs.py @@ -5,8 +5,6 @@ from __future__ import annotations -import inspect - import hydra import pytest from fixtures.synthetic_episodes import write_episode @@ -19,10 +17,6 @@ LOCAL_RESOLVER = "egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver" -def test_build_dataset_accepts_dataset_name() -> None: - assert "dataset_name" in inspect.signature(build_dataset).parameters - - @pytest.mark.parametrize("mode", ["random", "custom"]) def test_build_dataset_forwards_dataset_name(monkeypatch, mode) -> None: seen: dict = {} diff --git a/tests/unit/test_model_dims.py b/tests/unit/test_model_dims.py index f8419f855e..259ed791be 100644 --- a/tests/unit/test_model_dims.py +++ b/tests/unit/test_model_dims.py @@ -52,6 +52,8 @@ def _target(name): "hpt_bc_flow_scale": (None, {"human_bimanual": (144, 144)}), "hpt_bc_flow_human_cartesian": (None, {"human_bimanual": (20, 18)}), "hpt_bc_keypoints_base": (None, {"human_bimanual": (144, 144)}), + "hpt_bc_keypoints_wrist_300M": (None, {"human_bimanual": (144, 144)}), + "hpt_bc_mecka_6d_300M": (None, {"human_bimanual": (20, 18)}), "hpt_cotrain_enc_dec_base": (None, COTRAIN), # Separate heads: the human head follows the keypoint default. "hpt_cotrain_flow_seperate_head": ( diff --git a/tests/unit/test_pi05_norm_rot6d.py b/tests/unit/test_pi05_norm_rot6d.py index 6c69eaae34..6a1ac744d3 100644 --- a/tests/unit/test_pi05_norm_rot6d.py +++ b/tests/unit/test_pi05_norm_rot6d.py @@ -25,7 +25,7 @@ from egomimic.utils.action_utils import ( BaseActionConverter, HumanBimanualKeypoints, - RobotBimanualCartesianEuler, + RobotBimanualCartesian6D, pad_to_width, ) from egomimic.utils.pose_utils import _ypr_to_rot6d @@ -95,7 +95,7 @@ def test_transform_preserves_tensor_type(): def test_robot_bimanual_norm_6d_pack_round_trips(): - converter = RobotBimanualCartesianEuler() + converter = RobotBimanualCartesian6D() six6d = torch.from_numpy( CartesianYPRToRot6D().transform({"actions_cartesian": _eva_ypr_chunk()})[ "actions_cartesian" @@ -120,11 +120,11 @@ def test_keypoint_converter_is_identity_and_pads_to_model_width(): with pytest.raises(ValueError, match="expected 144-dim"): conv.to32_norm_6d(torch.zeros(1, 1, 138)) - eva32 = RobotBimanualCartesianEuler().to32_norm_6d(torch.ones(1, 3, 20)) + eva32 = RobotBimanualCartesian6D().to32_norm_6d(torch.ones(1, 3, 20)) wide = pad_to_width(eva32, 144) assert wide.shape[-1] == 144 and torch.all(wide[..., 32:] == 0) torch.testing.assert_close( - RobotBimanualCartesianEuler().from32_norm_6d(wide)[..., :20], + RobotBimanualCartesian6D().from32_norm_6d(wide)[..., :20], torch.ones(1, 3, 20), ) with pytest.raises(ValueError, match="action_dim is 32"): diff --git a/tests/unit/test_wrist6d_roundtrip.py b/tests/unit/test_wrist6d_roundtrip.py index 4c8dcadfb0..882e86c2c1 100644 --- a/tests/unit/test_wrist6d_roundtrip.py +++ b/tests/unit/test_wrist6d_roundtrip.py @@ -27,8 +27,8 @@ SplitKeys, ) from egomimic.utils.action_utils import ( - HumanBimanualCartesianEuler, - RobotBimanualCartesianEuler, + HumanBimanualCartesian6D, + RobotBimanualCartesian6D, _apply_norm_one, _apply_unnorm_one, _matrix_to_ypr, @@ -253,7 +253,7 @@ def test_human_wristframe_6d_pipeline_round_trips_to_headframe(fix_left): obs_rot = rot6d_channels("observations.state.ee_pose", o_t.shape[-1]) a_n = _apply_norm_one(a_t, st_act, "quantile", act_rot) o_n = _apply_norm_one(o_t, st_obs, "quantile", obs_rot) - conv = HumanBimanualCartesianEuler() + conv = HumanBimanualCartesian6D() a32 = conv.to32_norm_6d(a_n) assert a32.shape == (B, T, 32) assert torch.all(a32[..., [9, 19]] == 0) and torch.all(a32[..., 20:] == 0) @@ -348,7 +348,7 @@ def test_eva_wristframe_6d_pipeline_round_trips_to_camframe(): obs_rot = rot6d_channels("observations.state.ee_pose", o_t.shape[-1]) a_n = _apply_norm_one(a_t, st_act, "quantile", act_rot) o_n = _apply_norm_one(o_t, st_obs, "quantile", obs_rot) - conv = RobotBimanualCartesianEuler() + conv = RobotBimanualCartesian6D() a32 = conv.to32_norm_6d(a_n) torch.testing.assert_close(conv.from32_norm_6d(a32), a_n, atol=0, rtol=0) a_un = _apply_unnorm_one(conv.from32_norm_6d(a32), st_act, "quantile", act_rot) From 49c9dbf79c7b1426f75e1d1c50d9bc57ea596e0f Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Wed, 16 Sep 2026 18:32:06 -0400 Subject: [PATCH 08/79] rollout: build inputs from the checkpoint's Eva transform mode, and revert its outputs NOT TESTED ON HARDWARE. The rollout hard-coded cartesian_wristframe_ypr inputs and used the policy's predictions as-is. Now that eva checkpoints can be trained in several modes (cartesian, cartesian_6d, cartesian_wristframe_6d), that silently feeds a model the wrong inputs and sends the wrong frame to the arms. It now reads data.train_datasets..resolver.transform_list.mode off the checkpoint's own config, builds the matching input pipeline, and maps predictions back to camera-frame xyz+ypr+gripper through Eva.get_revert_transform_list. The revert needs the un-normalized proprio, which process_batch_for_training normalizes in place, so the raw ee_pose is cloned before that call. A checkpoint whose config records no mode keeps exactly the old behaviour -- legacy inputs, no revert -- and says so loudly, so this cannot silently change what an existing checkpoint does. test_eva_revert_for_rollout_recovers_camframe_ypr covers the revert math by cross-checking the wrist-frame revert against the cam-frame cartesian_6d one, which never enters the wrist frame -- so the frame composition is genuinely validated, though not against independently written SE(3) math. The hardware path itself still needs a real rollout. Both get_transform_list calls here pass allow_legacy_rotation=True: this is inference, and a checkpoint that predates the 6D conversion records exactly the modes the training-time guard rejects. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/robot/rollout.py | 77 +++++++++++++++++++-- tests/unit/test_eva_wristframe_norm_pool.py | 76 ++++++++++++++++++++ 2 files changed, 148 insertions(+), 5 deletions(-) diff --git a/egomimic/robot/rollout.py b/egomimic/robot/rollout.py index fb252bb726..bad6744593 100644 --- a/egomimic/robot/rollout.py +++ b/egomimic/robot/rollout.py @@ -16,7 +16,7 @@ from egomimic.models.denoising_policy import DenoisingPolicy from egomimic.pl_utils.pl_model import ModelWrapper -from egomimic.rldb.embodiment.embodiment import get_embodiment +from egomimic.rldb.embodiment.embodiment import Embodiment, get_embodiment from egomimic.rldb.embodiment.eva import Eva from egomimic.rldb.embodiment.human import Human from egomimic.robot.robot_utils import RateLoop @@ -258,10 +258,25 @@ def __init__( self.debug_actions = None self.resampled_action_len = resampled_action_len self.debug = debug - # Inference on checkpoints trained before the 6D conversion. - self.transform_list = Eva.get_transform_list( - mode="cartesian_wristframe_ypr", allow_legacy_rotation=True - ) + self.transform_mode = self._checkpoint_transform_mode() + # Inference: a checkpoint recorded before the 6D conversion, or none at + # all, is exactly the case allow_legacy_rotation exists for. + if self.transform_mode is None: + # Legacy checkpoint: keep the historical input pipeline and use + # the predictions as-is. + self.transform_list = Eva.get_transform_list( + mode=self.LEGACY_TRANSFORM_MODE, allow_legacy_rotation=True + ) + self.revert_transform_list = None + else: + self.transform_list = Eva.get_transform_list( + mode=self.transform_mode, allow_legacy_rotation=True + ) + # Maps predictions (wrist frame and/or 6D rotation) back to the + # camera-frame xyz+ypr+gripper layout the command path expects. + self.revert_transform_list = Eva.get_revert_transform_list( + self.transform_mode + ) self.annotation = None self._tokenizer = None self.collate_fn = default_collate @@ -309,6 +324,50 @@ def _patch_checkpoint_paths(cls, ckpt_path): print(f"[rollout] Patched checkpoint saved to {patched_path}") return patched_path + # Input pipeline used for checkpoints whose config names no mode (the + # rollout hard-coded it before the mode was read from the checkpoint). + LEGACY_TRANSFORM_MODE = "cartesian_wristframe_ypr" + + def _checkpoint_transform_mode(self): + """Eva ``transform_list.mode`` the checkpoint was trained with, so the + rollout builds the same inputs and reverts the matching outputs; + ``None`` if the checkpoint config does not record it.""" + from omegaconf import OmegaConf + + name = { + "both": "eva_bimanual", + "right": "eva_right_arm", + "left": "eva_left_arm", + }[self.arm] + cfg = self.policy._as_config(getattr(self.policy.hparams, "config_tree", None)) + mode = None + if cfg is not None: + mode = OmegaConf.select( + cfg, f"data.train_datasets.{name}.resolver.transform_list.mode" + ) + if mode is None: + print( + f"[rollout] WARNING: checkpoint config has no {name} transform_list.mode;" + f" using legacy {self.LEGACY_TRANSFORM_MODE} inputs, no output revert" + ) + return None + print(f"[rollout] Eva transform_list mode from checkpoint: {mode}") + return mode + + def _revert_predictions(self, preds, obs_batch): + """(1, T, D) unnormalized predictions -> (1, T, 14) camera-frame + xyz+ypr+gripper per arm. ``obs_batch`` is the collated, not yet + normalized input batch (its proprio defines the wrist frame).""" + if self.revert_transform_list is None: + return preds + obs_key = "observations.state.ee_pose" + batch = { + "actions_cartesian": preds, + obs_key: torch.as_tensor(obs_batch[obs_key]).to(preds.device), + } + out = Embodiment.apply_transform(batch, self.revert_transform_list) + return torch.as_tensor(out["actions_cartesian"]) + def _load_policy(self): patched_path = self._patch_checkpoint_paths(self.policy_path) policy = ModelWrapper.load_from_checkpoint( @@ -372,6 +431,13 @@ def rollout_step(self, i, obs): for transform in self.transform_list: transform_list_batch = transform.transform(transform_list_batch) transform_list_batch = self.collate_fn([transform_list_batch]) + # process_batch_for_training normalizes in place; keep raw proprio + # for reverting the wrist-frame predictions. + raw_obs = { + "observations.state.ee_pose": transform_list_batch[ + "observations.state.ee_pose" + ].clone() + } if self.arm == "both": embodiment_name = "eva_bimanual" elif self.arm == "right": @@ -386,6 +452,7 @@ def rollout_step(self, i, obs): preds = self.policy.model.forward_eval(processed_batch)[ f"{embodiment_name}_actions_cartesian" ] + preds = self._revert_predictions(preds, raw_obs) self.actions = preds.detach().cpu().numpy().squeeze() self.debug_actions = self.actions.copy() if self.cartesian: diff --git a/tests/unit/test_eva_wristframe_norm_pool.py b/tests/unit/test_eva_wristframe_norm_pool.py index cec5e3bd76..9aa0b6a554 100644 --- a/tests/unit/test_eva_wristframe_norm_pool.py +++ b/tests/unit/test_eva_wristframe_norm_pool.py @@ -269,3 +269,79 @@ def test_bounds_check_ignores_constant_cells(): assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is None arr[1, 0] = 50.0 # corrupt value at a regular cell: still rejected assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is not None + + +# ------------------------------------------------------- rollout reverts +def _rand_pose7(rng): + from scipy.spatial.transform import Rotation as R + + q = R.random(random_state=int(rng.integers(1 << 31))).as_quat() + return np.concatenate([rng.uniform(-0.5, 0.5, 3), q[[3, 0, 1, 2]]]) + + +def _rand_chunk7(rng, start, n=45): + from scipy.spatial.transform import Rotation as R + + out = np.zeros((n, 7)) + p, r = start[:3].copy(), R.from_quat(start[[4, 5, 6, 3]]) + for t in range(n): + if t: + p = p + rng.normal(0, 0.01, 3) + r = R.from_rotvec(rng.normal(0, 0.05, 3)) * r + out[t] = np.concatenate([p, r.as_quat()[[3, 0, 1, 2]]]) + return out + + +# Compared against cartesian_6d, so that mode is the reference, not a case. +@pytest.mark.parametrize( + "mode", ["cartesian_wristframe_ypr", "cartesian_wristframe_6d"] +) +def test_eva_revert_for_rollout_recovers_camframe_ypr(mode): + """rollout.py feeds Eva.get_transform_list(mode) outputs to the model and + Eva.get_revert_transform_list(mode) must bring predictions back to the + same cam-frame xyz+ypr+gripper (14-D) whatever the training mode.""" + from scipy.spatial.transform import Rotation as R + + from egomimic.rldb.embodiment.embodiment import Embodiment + + rng = np.random.default_rng(3) + lobs, robs = _rand_pose7(rng), _rand_pose7(rng) + raw = { + "left.obs_ee_pose": lobs, + "right.obs_ee_pose": robs, + "left.cmd_ee_pose": _rand_chunk7(rng, lobs), + "right.cmd_ee_pose": _rand_chunk7(rng, robs), + "left.obs_gripper": np.array([0.2]), + "right.obs_gripper": np.array([0.8]), + "left.cmd_gripper": rng.uniform(0, 1, (45, 1)), + "right.cmd_gripper": rng.uniform(0, 1, (45, 1)), + } + + def run(m): + s = {k: v.copy() for k, v in raw.items()} + # the ypr half of this comparison is exactly what a legacy checkpoint + # feeds the rollout, which is the case allow_legacy_rotation covers + for t in Eva.get_transform_list(m, allow_legacy_rotation=True): + s = t.transform(s) + batch = { + "actions_cartesian": torch.as_tensor(s["actions_cartesian"])[None], + "observations.state.ee_pose": torch.as_tensor( + s["observations.state.ee_pose"] + )[None], + } + rev = Eva.get_revert_transform_list(m) + if rev is not None: + batch = Embodiment.apply_transform(batch, rev) + return np.asarray(batch["actions_cartesian"][0], dtype=np.float64) + + assert Eva.get_revert_transform_list("cartesian") is None + got, ref = run(mode), run("cartesian_6d") + assert got.shape == ref.shape == (100, 14) + for off in (0, 7): + np.testing.assert_allclose( + got[:, off : off + 3], ref[:, off : off + 3], atol=1e-5 + ) + np.testing.assert_allclose(got[:, off + 6], ref[:, off + 6], atol=1e-5) + Rg = R.from_euler("ZYX", got[:, off + 3 : off + 6]).as_matrix() + Rr = R.from_euler("ZYX", ref[:, off + 3 : off + 6]).as_matrix() + np.testing.assert_allclose(Rg, Rr, atol=1e-4) From f748613b28efa618c2cb396b7ce1de16ace4546c Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Mon, 12 Oct 2026 23:18:12 +0000 Subject: [PATCH 09/79] rollout: read the Eva transform mode trainHydra now saves with the checkpoint The rollout looked for data.train_datasets..resolver.transform_list.mode in the checkpoint's config tree, but _build_model_config_tree saves only model, so every checkpoint fell back to the legacy 14-D ypr pipeline and a 6D policy would have produced garbage commands. trainHydra now records {dataset: transform_list.mode} as data_transform_modes in that tree, and the rollout reads it from there. A checkpoint without it (every one trained before this commit) is taken as legacy only when its saved actions_cartesian width is 14; otherwise the rollout refuses to guess and asks for the new --transform-mode flag. EMBODIMENT_MAP still held the pre-collapse ids (8/7/6 for both/left/right); it now reads them from EMBODIMENT. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GqAcWA2Lk9zubUhKX5DW1v --- egomimic/robot/rollout.py | 75 +++++++++++++++-------- egomimic/trainHydra.py | 20 ++++++ tests/unit/test_rollout_transform_mode.py | 71 +++++++++++++++++++++ 3 files changed, 142 insertions(+), 24 deletions(-) create mode 100644 tests/unit/test_rollout_transform_mode.py diff --git a/egomimic/robot/rollout.py b/egomimic/robot/rollout.py index bad6744593..b6df442bd9 100644 --- a/egomimic/robot/rollout.py +++ b/egomimic/robot/rollout.py @@ -16,7 +16,11 @@ from egomimic.models.denoising_policy import DenoisingPolicy from egomimic.pl_utils.pl_model import ModelWrapper -from egomimic.rldb.embodiment.embodiment import Embodiment, get_embodiment +from egomimic.rldb.embodiment.embodiment import ( + EMBODIMENT, + Embodiment, + get_embodiment, +) from egomimic.rldb.embodiment.eva import Eva from egomimic.rldb.embodiment.human import Human from egomimic.robot.robot_utils import RateLoop @@ -154,9 +158,9 @@ def viz_rot_ee_pose(image, eepose, action_image_path, rot_image_path): LEFT_CAM_SERIAL = "" EMBODIMENT_MAP = { - "both": 8, - "left": 7, - "right": 6, + "both": EMBODIMENT.EVA_BIMANUAL.value, + "left": EMBODIMENT.EVA_LEFT_ARM.value, + "right": EMBODIMENT.EVA_RIGHT_ARM.value, } TEMP_DIR = "/home/robot/temp_dir" @@ -242,6 +246,7 @@ def __init__( resampled_action_len=None, debug=False, annotation_path=None, + transform_mode=None, ): super().__init__() self.arm = arm @@ -258,7 +263,7 @@ def __init__( self.debug_actions = None self.resampled_action_len = resampled_action_len self.debug = debug - self.transform_mode = self._checkpoint_transform_mode() + self.transform_mode = self._checkpoint_transform_mode(transform_mode) # Inference: a checkpoint recorded before the 6D conversion, or none at # all, is exactly the case allow_legacy_rotation exists for. if self.transform_mode is None: @@ -328,31 +333,44 @@ def _patch_checkpoint_paths(cls, ckpt_path): # rollout hard-coded it before the mode was read from the checkpoint). LEGACY_TRANSFORM_MODE = "cartesian_wristframe_ypr" - def _checkpoint_transform_mode(self): + def _checkpoint_transform_mode(self, override=None): """Eva ``transform_list.mode`` the checkpoint was trained with, so the rollout builds the same inputs and reverts the matching outputs; - ``None`` if the checkpoint config does not record it.""" + ``None`` for a legacy 14-D ypr checkpoint. ``override`` wins.""" from omegaconf import OmegaConf - name = { - "both": "eva_bimanual", - "right": "eva_right_arm", - "left": "eva_left_arm", - }[self.arm] - cfg = self.policy._as_config(getattr(self.policy.hparams, "config_tree", None)) - mode = None - if cfg is not None: - mode = OmegaConf.select( - cfg, f"data.train_datasets.{name}.resolver.transform_list.mode" - ) + name = get_embodiment(self.embodiment_id).lower() + mode, source = override, "--transform-mode" if mode is None: - print( - f"[rollout] WARNING: checkpoint config has no {name} transform_list.mode;" - f" using legacy {self.LEGACY_TRANSFORM_MODE} inputs, no output revert" + cfg = self.policy._as_config( + getattr(self.policy.hparams, "config_tree", None) ) - return None - print(f"[rollout] Eva transform_list mode from checkpoint: {mode}") - return mode + if cfg is not None: + mode = OmegaConf.select( + cfg, f"data_transform_modes.{name}", default=None + ) + source = "checkpoint" + if mode is None: + # Checkpoints from before trainHydra recorded the mode: only the + # legacy layout can be assumed; any other width is ambiguous. + width = self._checkpoint_action_width() + if width != 14: + raise ValueError( + f"[rollout] checkpoint records no {name} transform mode and its " + f"actions_cartesian width is {width}, not the legacy 14; pass " + "--transform-mode (e.g. cartesian_wristframe_6d)" + ) + mode, source = self.LEGACY_TRANSFORM_MODE, "14-D action width" + print(f"[rollout] Eva transform_list mode ({source}): {mode}") + return None if mode == self.LEGACY_TRANSFORM_MODE else mode + + def _checkpoint_action_width(self): + """Last dim of the checkpoint's Eva actions_cartesian, or None.""" + state = getattr(self.policy.hparams, "norm_stats_state", None) or {} + shapes = state.get("shapes", {}) + emb = shapes.get(self.embodiment_id, shapes.get(str(self.embodiment_id), {})) + shape = emb.get("actions_cartesian") + return None if shape is None else int(shape[-1]) def _revert_predictions(self, preds, obs_batch): """(1, T, D) unnormalized predictions -> (1, T, 14) camera-frame @@ -656,6 +674,7 @@ def main( offline_debug=False, offline_episode_path=None, annotation_path=None, + transform_mode=None, ): if arms == "both": arms_list = ["right", "left"] @@ -687,6 +706,7 @@ def main( resampled_action_len=resampled_action_len, debug=debug, annotation_path=annotation_path, + transform_mode=transform_mode, ) elif dataset_path is not None: rollout_type = "replay" @@ -883,6 +903,12 @@ def build_arg_parser(description="Rollout robot model."): type=str, help="path to the annotation file", ) + parser.add_argument( + "--transform-mode", + type=str, + help="Eva transform_list mode the policy was trained with; only needed for " + "checkpoints that do not record it (default: read from the checkpoint)", + ) return parser @@ -900,6 +926,7 @@ def run_from_args(args): offline_debug=args.offline_debug, offline_episode_path=args.offline_episode_path, annotation_path=args.annotation_path, + transform_mode=args.transform_mode, ) diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index d5cf8787d0..d4bdf00110 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -39,6 +39,21 @@ _PI_WEIGHT_KEY = "model.robomimic_model.config.pytorch_weight_path" +def _data_transform_modes(cfg: DictConfig) -> dict[str, str]: + """``{dataset name: transform_list.mode}`` for the train datasets that set one.""" + datasets = OmegaConf.select(cfg, "data.train_datasets", default=None) + if not isinstance(datasets, DictConfig): + return {} + modes = {} + for name in datasets: + mode = OmegaConf.select( + datasets, f"{name}.resolver.transform_list.mode", default=None + ) + if mode is not None: + modes[str(name)] = str(mode) + return modes + + def _build_model_config_tree(cfg: DictConfig) -> DictConfig: """The config tree ``ModelWrapper`` receives. Only this copy skips the PI base weights (see ``_weights_from_checkpoint``); ``cfg`` keeps the real path for @@ -51,6 +66,11 @@ def _build_model_config_tree(cfg: DictConfig) -> DictConfig: ): model_cfg.robomimic_model.norm_stats = None tree = OmegaConf.create({"model": model_cfg}) + # Saved with the checkpoint so the robot rollout rebuilds the same inputs + # and reverts the outputs; the data config itself is not saved. + modes = _data_transform_modes(cfg) + if modes: + tree.data_transform_modes = modes has_weights = OmegaConf.select(tree, _PI_WEIGHT_KEY, default=None) is not None if has_weights and _weights_from_checkpoint(cfg): log.info( diff --git a/tests/unit/test_rollout_transform_mode.py b/tests/unit/test_rollout_transform_mode.py new file mode 100644 index 0000000000..793612eed7 --- /dev/null +++ b/tests/unit/test_rollout_transform_mode.py @@ -0,0 +1,71 @@ +"""The robot rollout rebuilds a checkpoint's Eva input pipeline from the +transform mode trainHydra records in the model config tree, and refuses to +guess for an unrecorded non-legacy action width.""" + +from types import SimpleNamespace + +import pytest + +import egomimic.trainHydra as train_hydra +from egomimic.pl_utils.pl_model import ModelWrapper +from egomimic.robot.rollout import EMBODIMENT_MAP, PolicyRollout + + +@pytest.fixture(autouse=True) +def no_slurm(monkeypatch): + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + monkeypatch.delenv("SLURM_RESTART_COUNT", raising=False) + + +def test_config_tree_records_the_eva_transform_mode(compose_resolve, tmp_path): + cfg = compose_resolve( + "train_zarr_cartesian", ["data=eva", f"paths.output_dir={tmp_path}"] + ) + tree = train_hydra._build_model_config_tree(cfg) + assert tree.data_transform_modes.eva_bimanual == ( + cfg.data.train_datasets.eva_bimanual.resolver.transform_list.mode + ) + + +def _rollout(config_tree=None, action_width=None): + emb = EMBODIMENT_MAP["both"] + shapes = ( + {} + if action_width is None + else {emb: {"actions_cartesian": (100, action_width)}} + ) + policy = SimpleNamespace( + hparams=SimpleNamespace( + config_tree=config_tree, norm_stats_state={"shapes": shapes} + ), + _as_config=ModelWrapper._as_config, + ) + r = PolicyRollout.__new__(PolicyRollout) + r.arm, r.embodiment_id, r.policy = "both", emb, policy + return r + + +def test_recorded_mode_is_used(): + tree = { + "model": {}, + "data_transform_modes": {"eva_bimanual": "cartesian_wristframe_6d"}, + } + assert _rollout(tree, 20)._checkpoint_transform_mode() == "cartesian_wristframe_6d" + + +def test_override_wins_over_the_checkpoint(): + tree = {"model": {}, "data_transform_modes": {"eva_bimanual": "cartesian_6d"}} + r = _rollout(tree, 20) + assert r._checkpoint_transform_mode("cartesian_wristframe_6d") == ( + "cartesian_wristframe_6d" + ) + + +def test_unrecorded_legacy_width_uses_the_legacy_pipeline(): + assert _rollout({"model": {}}, 14)._checkpoint_transform_mode() is None + + +@pytest.mark.parametrize("width", [20, None]) +def test_unrecorded_other_width_refuses_to_guess(width): + with pytest.raises(ValueError, match="--transform-mode"): + _rollout({"model": {}}, width)._checkpoint_transform_mode() From e63242cac8a135c38e9e6ecdd6013fcbbbabdb86 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 15 Sep 2026 14:54:59 -0400 Subject: [PATCH 10/79] val: subsampled metric loaders + pinned contiguous video loaders The unshuffled val loaders were truncated by limit_val_batches (80 x 64 x W frames) to their leading episode-hash-sorted slice, so at W = 1 the seen-op head scored 25 of 70 episodes and 6 of 11 operators, and the unseen head 33 of 191 episodes with an 81/19 operator mix on a 52/48 split (lane-0 diagnostic, 2026-09-15, measured single-rank). The loaders are unshuffled only so the overlay video is coherent, and the video renders just the first viz_max_batches (10) batches anyway. Split each val head in two: * metric loader = EvenStrideDataset(split, frames_per_episode=K) with K from data.metric_frames_per_episode[head], so every episode -- and so every operator -- is scored, at the same 80 batches and with no rendering; * video loader = _video, the head's split restricted to the episode hashes pinned in data.video_episodes[head] (one per operator), contiguous and unsubsampled, logging no metrics. The pinned episodes are taken out of the already-resolved split, so resolve-once still holds. EvalVideo.on_validation_step gains a mode argument (metrics / video / both); ModelWrapper.validation_step routes _video to the same evaluator in video mode and skips it entirely on a non-viz epoch, before the forward pass. Heads with no pins are called exactly as before, so a data config that declares neither key behaves as today. K = floor(5120 / episodes-in-split), from the measured splits: opsplit valid 73 (70 eps), train_viz 3 (1346 eps), unseen_op_valid 26 (191 eps); topop valid 256 (20 eps), train_viz 13 (385 eps). Those are SINGLE-RANK numbers and trainHydra scales them by trainer.devices x num_nodes, because DistributedSampler strides rather than chunks: the W ranks between them walk the whole split, so the pre-split window was already 80 x 64 x W frames wide. Without the scaling the shipped 4-GPU ddp_pi run would score 5,110 frames of the seen-val head where it previously scored all 13,744 -- a regression on the one head small enough to score whole. EvenStrideDataset keeps an episode entire when K exceeds its length, so the scaling cannot overrun a short split. Expect a shift in the per-batch Frechet _min/_max curves: they now span many episodes instead of ~2 s of one, so they are not comparable with older runs. Eval.on_validation_step, the abstract base, takes the mode kwarg too: the metrics/video path passes it unconditionally, so an evaluator that predates the split would have raised TypeError at the first val step of any config that declares video_episodes. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/eval/eval.py | 5 +- egomimic/eval/eval_video.py | 29 +- .../mecka_fold_flagship_opsplit_hpt_6d.yaml | 76 +++ .../mecka_fold_flagship_topop_hpt_6d.yaml | 29 ++ egomimic/pl_utils/pl_data_utils.py | 90 +++- egomimic/pl_utils/pl_model.py | 31 +- egomimic/rldb/zarr/zarr_dataset_multi.py | 50 ++ egomimic/trainHydra.py | 158 +++++- tests/unit/test_eval_loaders.py | 462 ++++++++++++++++++ 9 files changed, 908 insertions(+), 22 deletions(-) create mode 100644 tests/unit/test_eval_loaders.py diff --git a/egomimic/eval/eval.py b/egomimic/eval/eval.py index aff032cf33..174d8947fd 100644 --- a/egomimic/eval/eval.py +++ b/egomimic/eval/eval.py @@ -18,5 +18,8 @@ def on_validation_end(self): pass @abstractmethod - def on_validation_step(self, batch, batch_idx, dataloader_idx=0): + def on_validation_step(self, batch, batch_idx, dataloader_idx=0, mode="both"): + """``mode``: "metrics" / "video" / "both". An evaluator that renders + nothing may ignore it, but it must accept it -- ModelWrapper passes it + for every loader of a config that declares ``video_episodes``.""" pass diff --git a/egomimic/eval/eval_video.py b/egomimic/eval/eval_video.py index 039e3cba11..9fbcc61723 100644 --- a/egomimic/eval/eval_video.py +++ b/egomimic/eval/eval_video.py @@ -123,10 +123,27 @@ def on_validation_end(self): self.val_counter[key] = 0 self.val_image_buffer[key] = [] - def on_validation_step(self, batch, batch_idx, dataloader_idx=0): - do_viz = self._should_viz() and ( - self.viz_max_batches is None or batch_idx < self.viz_max_batches + def on_validation_step(self, batch, batch_idx, dataloader_idx=0, mode="both"): + """``mode`` splits metrics from video so one head can serve two loaders: + + * ``"metrics"`` -- the per-episode subsampled metric loader. Never + renders or buffers a frame (rendering is ~1 s/frame of CPU, and a + subsampled video would be an incoherent time-lapse anyway). + * ``"video"`` -- the contiguous pinned-episode loader. Renders and + buffers, and logs NOTHING, so no ``Valid/...`` key is averaged over + it. On a non-viz epoch, or past ``viz_max_batches``, it returns before + the forward pass: there is nothing such a call could produce. + * ``"both"`` (default) -- today's single-loader behaviour, which is what + a data config without ``video_episodes`` still gets. + """ + if mode not in ("both", "metrics", "video"): + raise ValueError(f"unknown validation mode {mode!r}") + past_viz_cap = ( + self.viz_max_batches is not None and batch_idx >= self.viz_max_batches ) + if mode == "video" and (past_viz_cap or not self._should_viz()): + return + do_viz = mode != "metrics" and self._should_viz() and not past_viz_cap metrics, images_dict = self.compute_metrics_and_viz(batch, do_viz=do_viz) device = self.trainer.lightning_module.device @@ -149,6 +166,12 @@ def on_validation_step(self, batch, batch_idx, dataloader_idx=0): self.val_image_buffer[key].clear() self.val_counter[key] += 1 + if mode == "video": + # Video-only loader: the forward pass was for the overlay. Logging + # here would mix pinned-episode numbers into the head's metric and + # (with add_dataloader_idx=False) collide key-for-key with it. + return + # add_dataloader_idx=False: with the train_viz loader Lightning would # otherwise suffix every key with "/dataloader_idx_N"; the train-viz # wrapper prefixes its keys itself. diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml index 00d70ff479..b8172887b0 100644 --- a/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml @@ -31,6 +31,11 @@ # The train and seen-val entries share `filters` and `valid_ratio`, so the # seeded split (split_dataset_names, SEED 42) makes them exact complements. # +# METRIC vs VIDEO LOADERS (2026-09-15, lane A). Each head above is served by +# TWO loaders: a per-episode subsampled METRIC loader (no rendering) and a +# contiguous VIDEO-only loader over pinned episodes (no metrics). See +# `metric_frames_per_episode` / `video_episodes` at the bottom of this file. +# # PIPELINE: WRIST-frame CARTESIAN actions with the continuous 6D rotation # (18D = xyz+6D per arm, no gripper): each arm's 100-step chunk is expressed # relative to that wrist's current pose; the head-frame proprio ee_pose is @@ -136,3 +141,74 @@ unseen_op_valid_dataloader_params: batch_size: 64 num_workers: 8 shuffle: false + +# --------------------------------------------------------------------------- +# METRIC LOADERS: K evenly spaced frames per EPISODE (lane A, 2026-09-15). +# --------------------------------------------------------------------------- +# Config-only (trainHydra reads them; MultiDataModuleWrapper declares them +# because hydra.instantiate(cfg.data) forwards every root key). +# +# WHY: the val loaders are unshuffled so the overlay video is coherent, so the +# metric window was just the first limit_val_batches x batch_size x world_size +# = 80 x 64 x W frames in episode-hash order (DistributedSampler strides, so +# the ranks between them walk one contiguous head of the split). Measured +# 2026-09-15 at W = 1 (lane-0 diagnostic, docs .../reports/lane0-coverage.md): +# the seen-op head scored 25/70 episodes and 6/11 operators, and the unseen +# head 33/191 episodes with an 81/19 operator mix on a split that is 52/48. +# trainHydra now wraps each val head in EvenStrideDataset(base, +# frames_per_episode=K), which takes K evenly spaced frames from EVERY episode +# (all of them when the episode is shorter than K), so every episode and +# operator is scored at the SAME 80-batch cost. +# +# K below is the SINGLE-RANK value, floor(5120 / episodes-in-that-split); +# trainHydra multiplies it by trainer.devices x num_nodes, so a 4-GPU run +# scores 4x these frames in the same 80 batches per rank rather than a quarter +# of what it scored before subsampling existed. Episode counts measured from +# the real splits, 2026-09-15: +# valid 70 eps / 13,744 frames -> K = floor(5120/70) = 73 +# -> 70 x 73 = 5,110 frames = 79.8 batches +# train_viz 1,346 eps / 271,247 frames -> K = floor(5120/1346) = 3 +# -> 1346 x 3 = 4,038 frames = 63.1 batches +# unseen_op_valid 191 eps / 30,937 frames -> K = floor(5120/191) = 26 +# -> 191 x 26 = 4,966 frames = 77.6 batches +# NOTE: per-batch Frechet _min/_max now span many episodes instead of ~2 s of +# one, so those curves are NOT comparable with pre-2026-09-15 runs. +metric_frames_per_episode: + valid: 73 + train_viz: 3 + unseen_op_valid: 26 + +# --------------------------------------------------------------------------- +# VIDEO LOADERS: pinned, contiguous episodes (`_video` loaders). +# --------------------------------------------------------------------------- +# One episode_hash per operator per head. Pinning them here (rather than +# letting the loader order decide) means every run renders the SAME episodes, +# so videos are comparable across runs. The pinned episodes are taken out of +# that head's already-resolved split -- no second SQL pull -- and played in +# frame order; `evaluator.viz_max_batches` (10 batches = 640 frames) caps it. +# These loaders log NO metrics, and the metric loaders above render nothing. +# +# Chosen 2026-09-15 from the real splits; format is (operator, frames): +# valid 692e771c39719ab57395b408 (690364170e94ce070afd1da5, 498f) +# the longest seen-val episode: 498f = 7.8 of the 10 video +# batches from a single contiguous episode (the split mean +# is ~196f, which would only fill ~3). +# train_viz 692e9bda21fc595fbe6f941e (6903686e0e94ce070afd1f24, 496f) +# longest train episode of the largest operator; it is also +# in the topop train split, so that child inherits this pin. +# unseen_op_valid 692eadbaaec602a46af10686 (6871d82a1aaf88556f0c6041, 204f) +# 692ea49a727c13b350cb7c48 (6903a36c466be157ec64562d, 212f) +# one per held-out operator, both of typical length +# (416f = 6.5 of the 10 batches). Operator 6903a36c... does +# have a 661f episode (692fe8d889e6f7f45d175b78) that would +# fill the budget alone, but one-per-operator keeps both +# unseen operators on screen. Shared by every flagship +# config, since the unseen head is identical in all of them. +video_episodes: + valid: + - 692e771c39719ab57395b408 + train_viz: + - 692e9bda21fc595fbe6f941e + unseen_op_valid: + - 692eadbaaec602a46af10686 + - 692ea49a727c13b350cb7c48 diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_topop_hpt_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_topop_hpt_6d.yaml index d4842d05cb..063c9b5285 100644 --- a/egomimic/hydra_configs/data/mecka_fold_flagship_topop_hpt_6d.yaml +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_topop_hpt_6d.yaml @@ -43,3 +43,32 @@ train_datasets: - "lambda row: '/mecka/flagship/' in str(row.get('zarr_processed_path', ''))" # SINGLE-OPERATOR TRAIN: the largest flagship operator only. - "lambda row: str(row.get('operator', '')) == '6903686e0e94ce070afd1f24'" + +# --------------------------------------------------------------------------- +# METRIC / VIDEO LOADERS: only the SEEN-op and train_viz entries differ from the +# parent (this config trains on one operator, so those two splits are smaller). +# --------------------------------------------------------------------------- +# K is the SINGLE-RANK value, floor(5120 / episodes-in-that-split); trainHydra +# scales it by the world size. Episode counts measured 2026-09-15: +# valid 20 eps / 3,731 frames -> K = floor(5120/20) = 256. Every +# episode here is shorter than 256 (longest 239f), so this keeps +# the WHOLE seen-val split: 3,731 frames = 58.3 of the 80 batches. +# train_viz 385 eps / 72,282 frames -> K = floor(5120/385) = 13 +# -> 385 x 13 = 5,005 frames = 78.2 batches. +# `unseen_op_valid: 26` is inherited unchanged (same split as the parent), and +# hydra merges dicts, so only the two changed keys are restated. +metric_frames_per_episode: + valid: 256 + train_viz: 13 + +# The parent's seen-val pin belongs to operator 690364170e94ce070afd1da5, which +# is not in this config's single-operator split, so it is replaced (hydra +# replaces lists wholesale). The train_viz pin 692e9bda21fc595fbe6f941e IS an +# episode of operator 6903686e..., hence in this train split too, and the +# unseen pins are the same split as the parent's -- both inherited. +# valid 692e83ad567c97525f4011a5 (6903686e0e94ce070afd1f24, 239f) +# the longest episode in this 20-episode seen-val split: 3.7 of the 10 +# video batches. No longer contiguous episode exists here. +video_episodes: + valid: + - 692e83ad567c97525f4011a5 diff --git a/egomimic/pl_utils/pl_data_utils.py b/egomimic/pl_utils/pl_data_utils.py index a414760529..aa4fe2c7e8 100644 --- a/egomimic/pl_utils/pl_data_utils.py +++ b/egomimic/pl_utils/pl_data_utils.py @@ -6,6 +6,24 @@ logger = logging.getLogger(__name__) +# Val heads that can carry a metric loader, in val_dataloader() order. +VAL_HEADS = ("valid", "train_viz", "unseen_op_valid") +# Suffix of the video-only companion loader of a head (``valid`` -> +# ``valid_video``). ModelWrapper.validation_step splits on it. +VIDEO_SUFFIX = "_video" + + +def video_loader_name(head: str) -> str: + return f"{head}{VIDEO_SUFFIX}" + + +def head_of_loader(name: str) -> tuple[str, bool]: + """``("valid", False)`` for a metric loader, ``("valid", True)`` for its + video-only companion.""" + if name.endswith(VIDEO_SUFFIX): + return name[: -len(VIDEO_SUFFIX)], True + return name, False + class MultiDataModuleWrapper(LightningDataModule): """ @@ -26,6 +44,9 @@ def __init__( unseen_op_valid_dataloader_params: dict | None = None, valid_prefix: str | None = None, held_out_operators: list | None = None, + video_datasets: dict | None = None, + metric_frames_per_episode: dict | None = None, + video_episodes: dict | None = None, ): """ Args: @@ -54,6 +75,23 @@ def __init__( the train/valid filter lambdas; hydra.instantiate(cfg.data) forwards every root key here, so it must be accepted. Kept as an attribute for provenance only. + video_datasets: optional ``{head: {dataset_name: dataset}}`` for the + video-ONLY companion loaders (``_video``). Built by + trainHydra from ``data.video_episodes``: the head's own split + restricted to the pinned episode hashes, contiguous and + unsubsampled so the overlay video stays watchable while the + metric loader is per-episode subsampled. A head with no pins + gets no video loader (today's behaviour). The video loader + reuses that head's ``*_dataloader_params`` (unshuffled). + metric_frames_per_episode: config-only ``{head: K}``. trainHydra + wraps each val head's dataset in ``EvenStrideDataset(base, + frames_per_episode=K)`` so the ``limit_val_batches`` window + covers every episode of the split instead of the leading + hash-sorted slice. Forwarded here by hydra.instantiate (it + rejects undeclared root keys); kept for provenance only. + video_episodes: config-only ``{head: [episode_hash, ...]}`` -- the + pinned video episodes, one per operator per head. Provenance + only; trainHydra is what reads it. Tokenization (sampling a prompt from per-sample annotation lists, splicing in embodiment / control-mode / proprio blocks, and running @@ -80,6 +118,15 @@ def __init__( self.unseen_op_valid_dataloader_params = unseen_op_valid_dataloader_params or {} self.valid_prefix = valid_prefix self.held_out_operators = list(held_out_operators or []) + # {head: {dataset_name: dataset}}; heads with no pinned episodes are + # dropped so `val_loader_names()` only grows for configs that ask for it. + self.video_datasets = { + head: {k: v for k, v in (datasets or {}).items() if v is not None} + for head, datasets in (video_datasets or {}).items() + } + self.video_datasets = {h: d for h, d in self.video_datasets.items() if d} + self.metric_frames_per_episode = dict(metric_frames_per_episode or {}) + self.video_episodes = dict(video_episodes or {}) self.collate_fn = annotation_collate def train_dataloader(self): @@ -117,18 +164,8 @@ def _build_val_style_loader(self, datasets: dict, params: dict, kind: str): ) return CombinedLoader(iterables, "max_size_cycle") - def val_loader_names(self) -> list[str]: - """Names of the val loaders in ``val_dataloader()`` order; position i is - Lightning's ``dataloader_idx`` i. ModelWrapper dispatches on these.""" - names = ["valid"] - if self.train_viz_datasets: - names.append("train_viz") - if self.unseen_op_valid_datasets: - names.append("unseen_op_valid") - return names - - def val_dataloader(self): - sources = { + def _metric_sources(self) -> dict: + return { "valid": (self.valid_datasets, self.valid_dataloader_params), "train_viz": (self.train_viz_datasets, self.train_viz_dataloader_params), "unseen_op_valid": ( @@ -136,9 +173,36 @@ def val_dataloader(self): self.unseen_op_valid_dataloader_params, ), } + + def val_loader_names(self) -> list[str]: + """Names of the val loaders in ``val_dataloader()`` order; position i is + Lightning's ``dataloader_idx`` i. ModelWrapper dispatches on these. + + Metric loaders first, in the historical order, so a config with no + pinned video episodes gets exactly today's list and today's indices; + the video-only loaders are appended after them in the same head order. + """ + metric = [ + name + for name in VAL_HEADS + if name == "valid" or self._metric_sources()[name][0] + ] + return metric + [ + video_loader_name(name) for name in metric if name in self.video_datasets + ] + + def val_dataloader(self): + sources = self._metric_sources() + for head, datasets in self.video_datasets.items(): + # The video loader reuses its head's params (batch size / workers / + # unshuffled); only the dataset differs. + sources[video_loader_name(head)] = (datasets, sources[head][1]) names = self.val_loader_names() loaders = [ - self._build_val_style_loader(*sources[name], kind=name) for name in names + # kind names the *_dataloader_params block in the error message, and + # a video loader borrows its head's block. + self._build_val_style_loader(*sources[name], kind=head_of_loader(name)[0]) + for name in names ] if len(loaders) == 1: return loaders[0] diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index 3a2cdd9ae6..306f541ad8 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/egomimic/pl_utils/pl_model.py @@ -10,6 +10,7 @@ from omegaconf import DictConfig, OmegaConf import egomimic.utils.tensor_utils as TensorUtils +from egomimic.pl_utils.pl_data_utils import head_of_loader, video_loader_name from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset @@ -239,12 +240,31 @@ def validation_step(self, batch, batch_idx, dataloader_idx=0): returns a list; ``val_loader_names[dataloader_idx]`` picks the head (valid -> evaluator, train_viz -> train_viz_evaluator, unseen_op_valid -> unseen_op_valid_evaluator). + + A head with pinned video episodes also gets a ``_video`` loader, + which routes to the SAME evaluator in video-only mode while the metric + loader runs in metrics-only mode. Heads without one keep the single + combined call, so evaluators (and stub evaluators in tests) that take no + ``mode`` still work. """ names = getattr(self, "val_loader_names", None) or ["valid", "train_viz"] name = names[dataloader_idx] if dataloader_idx < len(names) else None - active = self._val_heads().get(name) + if name is None: + return + head, video_only = head_of_loader(name) + active = self._val_heads().get(head) if active is None: return + if video_only: + mode = "video" + # Non-viz epoch: skip before the forward pass, it would be wasted. + should_viz = getattr(active, "_should_viz", None) + if callable(should_viz) and not should_viz(): + return + elif video_loader_name(head) in names: + mode = "metrics" + else: + mode = "both" # When val_dataloader returns a list of CombinedLoaders (valid + # train_viz), Lightning wraps it in an outer sequential CombinedLoader. # The outer iterator calls next() on each inner CombinedLoader, which @@ -255,10 +275,15 @@ def validation_step(self, batch, batch_idx, dataloader_idx=0): batch = self.model.process_batch_for_training(batch) print( f"[VAL_STEP] rank={self.global_rank}, batch_idx={batch_idx}, " - f"dataloader_idx={dataloader_idx}", + f"dataloader_idx={dataloader_idx}, loader={name}, mode={mode}", flush=True, ) - active.on_validation_step(batch, batch_idx, dataloader_idx) + if mode == "both": + # The historical call: no `mode` kwarg, so evaluators that predate + # the split keep working unchanged. + active.on_validation_step(batch, batch_idx, dataloader_idx) + else: + active.on_validation_step(batch, batch_idx, dataloader_idx, mode=mode) def on_validation_end(self): print(f"[ON_VALIDATION_END] rank={self.global_rank}", flush=True) diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 65ee853157..70f5e7f9c6 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -1959,6 +1959,14 @@ def set_data_schematic(self, data_schematic) -> None: self.base.set_data_schematic(data_schematic) self.data_schematic = data_schematic + def set_norm_stats_from(self, source: "MultiDataset") -> None: + """Delegate to ``self.base`` as well: ``__getitem__`` here forwards to + ``self.base[...]``, so the bounds-check + normalize pass reads the + *base*'s stats. Setting them on the wrapper alone (the inherited + implementation) would leave the subsampled val loader unnormalized.""" + self.base.set_norm_stats_from(source) + MultiDataset.set_norm_stats_from(self, source) + def __getattr__(self, item): if item == "base": raise AttributeError(item) @@ -1968,6 +1976,48 @@ def __getattr__(self, item): return getattr(base, item) +def pinned_episode_subset( + base: MultiDataset, episode_hashes, dataset_name: str | None = None +) -> MultiDataset: + """A contiguous, unsubsampled ``MultiDataset`` over just ``episode_hashes`` + of an already-built ``base``. + + Used for the video-only val loaders: the overlay video must be watchable, so + it renders whole episodes in frame order rather than the metric loader's + per-episode subsample. + + The episodes are taken from ``base.datasets`` (the same leaf ``ZarrDataset`` + objects), NOT re-resolved: a second ``MultiDataset._from_resolver`` with + ``filters.episode_hashes`` set would be a different ``resolve_once`` memo key + and so a second path resolution / zarr open per head. A pin that is not in + ``base`` is a ``PinError``, matching ``_from_resolver``'s pin contract. + + ``episode_hashes`` order is preserved, so the video plays the pinned + episodes in config order. + """ + hashes = [str(h) for h in (episode_hashes or [])] + missing = [h for h in hashes if h not in base.datasets] + if missing: + raise PinError( + f"{len(missing)} pinned video episode(s) not in dataset " + f"'{dataset_name or ''}': {sorted(missing)}. The pin must " + "name an episode of THIS head's split (e.g. a seen-val episode for " + "`valid`, a held-out-operator episode for `unseen_op_valid`)." + ) + if not hashes: + raise PinError( + f"pinned_episode_subset called with no episode hashes for " + f"'{dataset_name or ''}'" + ) + subset = MultiDataset( + datasets={h: base.datasets[h] for h in hashes}, + mode="total", + norm_mode=base.norm_mode, + ) + subset.set_norm_stats_from(base) + return subset + + class ZarrDataset(torch.utils.data.Dataset): """ Base Zarr Dataset object, Just intializes as pass through to read from zarr episode diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index d4bdf00110..79f3291780 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -13,7 +13,7 @@ from lightning.fabric.utilities.cloud_io import get_filesystem from lightning.pytorch.loggers import Logger from lightning.pytorch.plugins.environments import SLURMEnvironment -from omegaconf import DictConfig, OmegaConf, open_dict +from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict from tabulate import tabulate import egomimic.utils.hydra_resolvers # noqa: F401 -- registers OmegaConf resolvers @@ -22,7 +22,12 @@ from egomimic.rldb.embodiment.embodiment import get_embodiment_id from egomimic.rldb.resolve_memo import resolve_once from egomimic.rldb.zarr.utils import set_global_seed -from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset, PinError +from egomimic.rldb.zarr.zarr_dataset_multi import ( + EvenStrideDataset, + MultiDataset, + PinError, + pinned_episode_subset, +) from egomimic.utils.checkpoint_utils import load_checkpoint_weights from egomimic.utils.compile_cache import set_per_job_compile_cache_dir from egomimic.utils.env import load_env @@ -285,6 +290,126 @@ def _unseen_op_valid_datasets(cfg: DictConfig, instantiate) -> dict: } +def _trainer_world_size(cfg: DictConfig) -> int: + """``devices * num_nodes`` as the config declares it, 1 if it cannot say. + + DistributedSampler strides rather than chunks, so W ranks between them walk + the WHOLE split; the per-rank ``limit_val_batches`` window is therefore W + times wider than it looks from one rank. + """ + devices = cfg.get("trainer", {}).get("devices", 1) + if isinstance(devices, (list, ListConfig)): + n = len(devices) + elif isinstance(devices, int) and devices > 0: + n = devices + else: # "auto", -1: resolved by lightning at runtime, unknown here + log.warning( + f"trainer.devices={devices!r} is not a count; treating the metric " + "loaders as single-rank, which under-scores the val split." + ) + n = 1 + nodes = cfg.get("trainer", {}).get("num_nodes", 1) + return n * (int(nodes) if isinstance(nodes, int) and nodes > 0 else 1) + + +def _metric_frames_per_episode(cfg: DictConfig, head: str) -> int | None: + """Frames per episode to keep on this val head, or None for no subsampling. + + ``data.metric_frames_per_episode[head]`` is written for ONE rank -- it is + ``floor(limit_val_batches * batch_size / n_episodes)`` -- so it is scaled by + the world size here: the ranks stride through the subsampled set together + and each still reads at most ``limit_val_batches`` batches. Without this a + 4-GPU run scores a quarter of the frames the same config scores on 1 GPU, + which on the seen-val head is fewer than it scored before subsampling + existed at all. EvenStrideDataset keeps a whole episode when K exceeds its + length, so a split that fits entirely is not subsampled. + """ + table = cfg.data.get("metric_frames_per_episode") + if table is None: + return None + k = table.get(head) + if k is None: + return None + k = int(k) + if k <= 0: + raise ValueError( + f"data.metric_frames_per_episode.{head} must be a positive int, got {k}" + ) + return k * _trainer_world_size(cfg) + + +def _subsample_val_datasets(cfg: DictConfig, head: str, datasets: dict) -> dict: + """Wrap each of this val head's datasets in ``EvenStrideDataset`` so the + ``limit_val_batches`` window spans EVERY episode of the split. + + The val loaders are unshuffled (so the overlay video is coherent), which + made the metric window the first ``limit_val_batches * batch_size * W`` + frames in episode-hash order: at W = 1 that is a third of the seen-val + split and effectively one operator of the unseen split (lane-0 coverage + report, 2026-09-15, measured single-rank). K evenly spaced frames per + episode gives every episode -- and so every operator -- equal weight at + the same batch count. + + No ``metric_frames_per_episode`` entry for this head => datasets pass + through untouched.""" + k = _metric_frames_per_episode(cfg, head) + if k is None or not datasets: + return datasets + wrapped = {} + for name, ds in datasets.items(): + if ds is None: + wrapped[name] = None + continue + sub = EvenStrideDataset(ds, frames_per_episode=k) + log.info( + f"val head '{head}' dataset '{name}': EvenStrideDataset " + f"frames_per_episode={k} -> {len(sub)} / {len(ds)} frames over " + f"{len(ds.datasets)} episodes" + ) + wrapped[name] = sub + return wrapped + + +def _video_datasets(cfg: DictConfig, heads: dict) -> dict: + """``{head: {dataset_name: MultiDataset}}`` for the video-only val loaders. + + ``data.video_episodes[head]`` pins one episode hash per operator; each + pinned episode is taken out of that head's ALREADY resolved split (no second + SQL pull or path resolution -- see ``pinned_episode_subset``), contiguous + and unsubsampled so the overlay video stays watchable. + + A head with no pins gets no video loader, so a config without the key keeps + today's behaviour (metrics AND video on the one metric loader).""" + table = cfg.data.get("video_episodes") + if table is None: + return {} + out: dict = {} + for head, datasets in heads.items(): + pins = [str(h) for h in (table.get(head) or [])] + if not pins or not datasets: + continue + live = {name: ds for name, ds in datasets.items() if ds is not None} + unplaced = sorted(set(pins) - {h for ds in live.values() for h in ds.datasets}) + if unplaced: + raise PinError( + f"{len(unplaced)} pinned video episode(s) for head '{head}' are in " + f"none of its datasets {sorted(live)}: {unplaced}. The pin must " + "name an episode of THIS head's split." + ) + built = {} + for name, ds in live.items(): + mine = [h for h in pins if h in ds.datasets] + if not mine: + continue + built[name] = pinned_episode_subset(ds, mine, dataset_name=name) + log.info( + f"val head '{head}' video loader from dataset '{name}': " + f"{len(mine)} pinned episode(s), {len(built[name])} frames" + ) + out[head] = built + return out + + def _build_unseen_op_valid_evaluator(cfg: DictConfig): """The canonical evaluator (fresh instance, own frame buffers) wrapped to log ``unseen_op_valid/`` and write ``videos_unseen_op_valid/``.""" @@ -339,6 +464,27 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: cfg, instantiate=_instantiate_dataset ) + # Split the val heads in two (lane A): a per-episode SUBSAMPLED metric + # loader that covers the whole split inside limit_val_batches, and a + # contiguous video-only loader over the pinned episodes. Build the video + # subsets from the full splits BEFORE wrapping, and let each head opt in + # independently (no keys in the data config => today's behaviour). + video_datasets = _video_datasets( + cfg, + { + "valid": valid_datasets, + "train_viz": train_viz_datasets, + "unseen_op_valid": unseen_op_valid_datasets, + }, + ) + valid_datasets = _subsample_val_datasets(cfg, "valid", valid_datasets) + train_viz_datasets = _subsample_val_datasets( + cfg, "train_viz", train_viz_datasets + ) + unseen_op_valid_datasets = _subsample_val_datasets( + cfg, "unseen_op_valid", unseen_op_valid_datasets + ) + log.info(f"Instantiating datamodule <{cfg.data._target_}>") assert ( "MultiDataModuleWrapper" in cfg.data._target_ @@ -354,6 +500,8 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: # Always passed too: hydra would otherwise instantiate the raw # config nodes itself (eval mode included). datamodule_kwargs["unseen_op_valid_datasets"] = unseen_op_valid_datasets + # Built above; hydra has nothing to instantiate for these. + datamodule_kwargs["video_datasets"] = video_datasets datamodule: LightningDataModule = hydra.utils.instantiate( cfg.data, **datamodule_kwargs ) @@ -455,6 +603,12 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: ds.set_norm_stats_from(norm_stats) for ds in getattr(datamodule, "unseen_op_valid_datasets", {}).values(): ds.set_norm_stats_from(norm_stats) + # The video subsets share leaves with their head's split but are separate + # MultiDatasets, so they need their own wiring (EvenStrideDataset forwards + # the call to the base it indexes into). + for head_datasets in getattr(datamodule, "video_datasets", {}).values(): + for ds in head_datasets.values(): + ds.set_norm_stats_from(norm_stats) _prepare_checkpoint_resume(cfg) diff --git a/tests/unit/test_eval_loaders.py b/tests/unit/test_eval_loaders.py new file mode 100644 index 0000000000..238319b6e3 --- /dev/null +++ b/tests/unit/test_eval_loaders.py @@ -0,0 +1,462 @@ +"""Lane A: split each val head into a subsampled METRIC loader and a +contiguous pinned-episode VIDEO loader. + +Before this change the unshuffled val loaders were truncated by +``limit_val_batches`` to their leading hash-sorted slice, so the metric saw a +third of the seen-op split and effectively one of the two held-out operators +(docs/hpt-experiments/supervisors/S1/reports/lane0-coverage.md). Now: + +* metric loader = ``EvenStrideDataset(split, frames_per_episode=K)``, K from + the data config, so every episode contributes and nothing renders; +* video loader = the head's split restricted to the pinned ``episode_hash``es, + contiguous, logging no metrics. + +A data config with neither key must behave exactly as before. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +from lightning import LightningModule +from omegaconf import OmegaConf + +import egomimic.trainHydra as th +from egomimic.pl_utils.pl_data_utils import MultiDataModuleWrapper +from egomimic.pl_utils.pl_model import ModelWrapper +from egomimic.rldb.zarr.zarr_dataset_multi import ( + EvenStrideDataset, + MultiDataset, + PinError, + pinned_episode_subset, +) + +HUMAN = "human_bimanual" +FLAGSHIP = "train_zarr_mecka_flagship_6d_hpt" +LIMIT_VAL_BATCHES = 80 +BATCH = 64 + +# (data config, {head: K}, {head: [pinned hashes]}) as set on this branch. +CONFIG_EXPECTATIONS = { + "mecka_fold_flagship_opsplit_hpt_6d": ( + {"valid": 73, "train_viz": 3, "unseen_op_valid": 26}, + { + "valid": ["692e771c39719ab57395b408"], + "train_viz": ["692e9bda21fc595fbe6f941e"], + "unseen_op_valid": [ + "692eadbaaec602a46af10686", + "692ea49a727c13b350cb7c48", + ], + }, + ), + "mecka_fold_flagship_topop_hpt_6d": ( + {"valid": 256, "train_viz": 13, "unseen_op_valid": 26}, + { + "valid": ["692e83ad567c97525f4011a5"], + "train_viz": ["692e9bda21fc595fbe6f941e"], + "unseen_op_valid": [ + "692eadbaaec602a46af10686", + "692ea49a727c13b350cb7c48", + ], + }, + ), +} + +# Episode counts these K values were computed from (measured 2026-09-15). +SPLIT_EPISODES = { + "mecka_fold_flagship_opsplit_hpt_6d": { + "valid": 70, + "train_viz": 1346, + "unseen_op_valid": 191, + }, + "mecka_fold_flagship_topop_hpt_6d": { + "valid": 20, + "train_viz": 385, + "unseen_op_valid": 191, + }, +} + + +# --------------------------------------------------------------- fixtures +class _Episode: + """Minimal leaf dataset: n frames, each sample tagged with its frame idx.""" + + def __init__(self, n: int): + self.n = n + + def __len__(self): + return self.n + + def __getitem__(self, i): + return {"frame": np.float32(i)} + + +def _split(sizes: dict[str, int]) -> MultiDataset: + return MultiDataset( + datasets={h: _Episode(n) for h, n in sizes.items()}, mode="total" + ) + + +def _cfg(metric=None, video=None, **data): + node = dict(data) + if metric is not None: + node["metric_frames_per_episode"] = metric + if video is not None: + node["video_episodes"] = video + return OmegaConf.create({"data": node}) + + +# ------------------------------------------------------------ 1. config keys +@pytest.mark.parametrize("data_config", sorted(CONFIG_EXPECTATIONS)) +def test_flagship_configs_declare_k_and_pins(data_config, compose_resolve): + """Both flagship configs compose with the new keys, K per head matches + floor(5120 / episodes-in-split), and each head pins its video episodes.""" + expected_k, expected_pins = CONFIG_EXPECTATIONS[data_config] + cfg = compose_resolve(FLAGSHIP, [f"data={data_config}"]) + assert int(cfg.trainer.limit_val_batches) == LIMIT_VAL_BATCHES + assert int(cfg.data.valid_dataloader_params[HUMAN].batch_size) == BATCH + + k_table = OmegaConf.to_container(cfg.data.metric_frames_per_episode) + assert k_table == expected_k + pins = OmegaConf.to_container(cfg.data.video_episodes) + assert pins == expected_pins + + for head, k in k_table.items(): + n_eps = SPLIT_EPISODES[data_config][head] + assert k == LIMIT_VAL_BATCHES * BATCH // n_eps, (head, k, n_eps) + # the subsampled set still fits the limit_val_batches window + assert k * n_eps <= LIMIT_VAL_BATCHES * BATCH, (head, k, n_eps) + # the unseen head is the same split in every flagship config, so its pins + # and K are shared + assert k_table["unseen_op_valid"] == 26 + assert len(pins["unseen_op_valid"]) == 2, "one pin per held-out operator" + + # and trainHydra reads exactly these + for head, k in expected_k.items(): + assert th._metric_frames_per_episode(cfg, head) == k + + +def test_config_without_the_keys_is_unchanged(compose_resolve): + """Backward compatibility: a data config that declares neither key gets the + old loader names and the unwrapped dataset objects.""" + cfg = compose_resolve("train_zarr_cartesian", ["data=mecka_all_6d"]) + assert cfg.data.get("metric_frames_per_episode") is None + assert cfg.data.get("video_episodes") is None + assert th._metric_frames_per_episode(cfg, "valid") is None + + split = _split({"ep0": 10, "ep1": 10}) + datasets = {HUMAN: split} + assert th._subsample_val_datasets(cfg, "valid", datasets) is datasets + assert th._video_datasets(cfg, {"valid": datasets}) == {} + + params = {HUMAN: {"batch_size": 2, "num_workers": 0}} + dm = MultiDataModuleWrapper( + train_datasets=datasets, + valid_datasets=datasets, + train_dataloader_params=params, + valid_dataloader_params=params, + train_viz_datasets=datasets, + train_viz_dataloader_params=params, + unseen_op_valid_datasets=datasets, + unseen_op_valid_dataloader_params=params, + ) + assert dm.val_loader_names() == ["valid", "train_viz", "unseen_op_valid"] + assert dm.valid_datasets[HUMAN] is split + + +# ------------------------------------------------------------ 2. metric loaders +def test_metric_loader_is_even_stride_over_every_episode(): + sizes = {"ep0": 500, "ep1": 300, "ep2": 40} + split = _split(sizes) + cfg = _cfg(metric={"valid": 73}) + wrapped = th._subsample_val_datasets(cfg, "valid", {HUMAN: split})[HUMAN] + + assert isinstance(wrapped, EvenStrideDataset) + assert isinstance(wrapped, MultiDataset), "trainHydra's isinstance checks" + assert wrapped.frames_per_episode == 73 + assert wrapped.base is split + + per_episode = {name: 0 for name in sizes} + starts = {} + offset = 0 + for name, n in sizes.items(): + starts[name] = offset + offset += n + for global_idx in wrapped.indices: + name, _ = split.index_map[global_idx] + per_episode[name] += 1 + # K frames from each episode, or the whole episode when it is shorter + assert per_episode == {"ep0": 73, "ep1": 73, "ep2": 40} + assert len(wrapped) == sum(per_episode.values()) + assert len(wrapped) <= LIMIT_VAL_BATCHES * BATCH + # ... and the frames really are spread across each episode, not its head + ep0 = [ + split.index_map[g][1] for g in wrapped.indices if split.index_map[g][0] == "ep0" + ] + assert ep0[0] == 0 and ep0[-1] == 499 + + # samples still come through (the wrapper indexes into the base) + assert wrapped[0]["frame"] == split[wrapped.indices[0]]["frame"] + + +def test_metric_wrapper_forwards_norm_stats_to_the_base(): + """__getitem__ delegates to base, so the stats must land on the base or the + subsampled val loader would be unnormalized.""" + split = _split({"ep0": 10}) + wrapped = EvenStrideDataset(split, frames_per_episode=5) + source = MultiDataset(state={}) + source.norm_stats = {0: {"actions_cartesian": {"mean": np.zeros(3)}}} + wrapped.set_norm_stats_from(source) + assert split.norm_stats is source.norm_stats + + +def test_every_head_can_be_subsampled_independently(): + cfg = _cfg(metric={"valid": 5, "unseen_op_valid": 2}) + heads = {h: {HUMAN: _split({"ep0": 100})} for h in ("valid", "train_viz")} + heads["unseen_op_valid"] = {HUMAN: _split({"ep0": 100})} + assert len(th._subsample_val_datasets(cfg, "valid", heads["valid"])[HUMAN]) == 5 + assert ( + len( + th._subsample_val_datasets( + cfg, "unseen_op_valid", heads["unseen_op_valid"] + )[HUMAN] + ) + == 2 + ) + # train_viz has no entry -> untouched + assert ( + th._subsample_val_datasets(cfg, "train_viz", heads["train_viz"]) + is heads["train_viz"] + ) + with pytest.raises(ValueError, match="positive int"): + th._metric_frames_per_episode(_cfg(metric={"valid": 0}), "valid") + + +def test_k_scales_with_the_world_size(): + """K in the config is the single-rank value. DistributedSampler strides, + so W ranks between them read W times the frames in the same number of + batches each; without the scaling a 4-GPU run scores a quarter of what the + same config scores on one GPU.""" + cfg = _cfg(metric={"valid": 7}) + for devices, nodes, expected in ( + (1, 1, 7), + (4, 1, 28), + (8, 2, 112), + ([0, 1], 1, 14), + ("auto", 1, 7), # unknown at config time: stay single-rank + (-1, 1, 7), + ): + cfg.trainer = {"devices": devices, "num_nodes": nodes} + assert th._metric_frames_per_episode(cfg, "valid") == expected, devices + # a split shorter than K is kept whole, so scaling cannot truncate + cfg.trainer = {"devices": 4, "num_nodes": 1} + split = {HUMAN: _split({"ep0": 20})} + assert len(th._subsample_val_datasets(cfg, "valid", split)[HUMAN]) == 20 + + +# ------------------------------------------------------------- 3. video loaders +def test_video_dataset_holds_only_the_pins_in_frame_order(): + split = _split({"ep0": 30, "ep1": 20, "ep2": 25}) + sub = pinned_episode_subset(split, ["ep2", "ep0"], dataset_name=HUMAN) + + assert list(sub.datasets) == ["ep2", "ep0"], "config order is playback order" + assert len(sub) == 55 + # contiguous: frame indices run 0..n-1 inside each episode, in order + by_episode: dict[str, list[int]] = {} + for name, local in sub.index_map: + by_episode.setdefault(name, []).append(local) + assert by_episode == {"ep2": list(range(25)), "ep0": list(range(30))} + + with pytest.raises(PinError, match="not in dataset"): + pinned_episode_subset(split, ["nope"], dataset_name=HUMAN) + + +def test_video_datasets_are_built_per_head_from_the_resolved_split(): + valid = _split({"a": 10, "b": 10}) + unseen = _split({"c": 10, "d": 10}) + cfg = _cfg(video={"valid": ["b"], "unseen_op_valid": ["c", "d"]}) + built = th._video_datasets( + cfg, + { + "valid": {HUMAN: valid}, + "train_viz": {HUMAN: _split({"t": 10})}, + "unseen_op_valid": {HUMAN: unseen}, + }, + ) + assert set(built) == {"valid", "unseen_op_valid"}, "no pins -> no video loader" + assert list(built["valid"][HUMAN].datasets) == ["b"] + assert list(built["unseen_op_valid"][HUMAN].datasets) == ["c", "d"] + # the video dataset reuses the SAME leaf objects: no second resolve + assert built["valid"][HUMAN].datasets["b"] is valid.datasets["b"] + + with pytest.raises(PinError, match="none of its datasets"): + th._video_datasets(_cfg(video={"valid": ["c"]}), {"valid": {HUMAN: valid}}) + + +def test_val_loader_names_append_the_video_loaders(): + params = {HUMAN: {"batch_size": 2, "num_workers": 0, "shuffle": False}} + datasets = {HUMAN: _split({"ep0": 4})} + dm = MultiDataModuleWrapper( + train_datasets=datasets, + valid_datasets=datasets, + train_dataloader_params=params, + valid_dataloader_params=params, + train_viz_datasets=datasets, + train_viz_dataloader_params=params, + unseen_op_valid_datasets=datasets, + unseen_op_valid_dataloader_params=params, + video_datasets={ + "valid": {HUMAN: _split({"pin": 4})}, + "unseen_op_valid": {HUMAN: _split({"pin2": 4})}, + }, + ) + # metric loaders keep their historical indices; video loaders follow + assert dm.val_loader_names() == [ + "valid", + "train_viz", + "unseen_op_valid", + "valid_video", + "unseen_op_valid_video", + ] + loaders = dm.val_dataloader() + assert len(loaders) == 5 + # the video loader really iterates the pinned dataset, unshuffled + batch, _, _ = next(iter(loaders[3])) + assert [float(f) for f in batch[HUMAN]["frame"]] == [0.0, 1.0] + + # a head whose video datasets were dropped gets no name + dm.video_datasets = {} + assert dm.val_loader_names() == ["valid", "train_viz", "unseen_op_valid"] + + +# ----------------------------------------------------------------- 4. dispatch +class _StubEval: + """Records how validation_step called it, with EvalVideo's viz gating.""" + + def __init__(self, should_viz=True): + self.calls: list[tuple] = [] + self._viz = should_viz + + def _should_viz(self): + return self._viz + + def on_validation_step(self, batch, batch_idx, dataloader_idx=0, mode="both"): + self.calls.append((batch_idx, dataloader_idx, mode)) + + +def _wrapper(names, heads): + w = ModelWrapper.__new__(ModelWrapper) + LightningModule.__init__(w) + w.model = SimpleNamespace(process_batch_for_training=lambda b: b) + w.evaluator = heads.get("valid") + w.train_viz_evaluator = heads.get("train_viz") + w.unseen_op_valid_evaluator = heads.get("unseen_op_valid") + w.val_loader_names = names + return w + + +def test_validation_step_splits_metrics_from_video(): + valid, unseen = _StubEval(), _StubEval() + names = ["valid", "unseen_op_valid", "valid_video", "unseen_op_valid_video"] + w = _wrapper(names, {"valid": valid, "unseen_op_valid": unseen}) + + for idx in range(4): + w.validation_step({"k": idx}, 0, dataloader_idx=idx) + + assert valid.calls == [(0, 0, "metrics"), (0, 2, "video")] + assert unseen.calls == [(0, 1, "metrics"), (0, 3, "video")] + + +def test_non_video_epoch_skips_the_video_loader_before_the_forward(): + valid = _StubEval(should_viz=False) + processed = [] + w = _wrapper(["valid", "valid_video"], {"valid": valid}) + w.model = SimpleNamespace( + process_batch_for_training=lambda b: processed.append(b) or b + ) + w.validation_step({"k": 0}, 0, dataloader_idx=0) + w.validation_step({"k": 1}, 0, dataloader_idx=1) + assert valid.calls == [(0, 0, "metrics")], "video loader skipped" + assert processed == [{"k": 0}], "no forward pass for the skipped loader" + + +def test_head_without_a_video_loader_keeps_the_old_combined_call(): + """Evaluators that predate the split (and the stubs in older tests) are + called with no ``mode`` kwarg at all.""" + seen = [] + + class _Old: + def on_validation_step(self, batch, batch_idx, dataloader_idx=0): + seen.append((batch_idx, dataloader_idx)) + + w = _wrapper(["valid"], {"valid": _Old()}) + w.validation_step({"k": 0}, 7, dataloader_idx=0) + assert seen == [(7, 0)] + + +# ----------------------------------------------- EvalVideo's two modes +class _Trainer: + def __init__(self, epoch=0, max_epochs=1): + self.current_epoch = epoch + self.max_epochs = max_epochs + self.world_size = 1 + self.is_global_zero = True + self.logged: list[dict] = [] + outer = self + + class _LM: + device = "cpu" + + def log_dict(self, metrics, **kw): + outer.logged.append(dict(metrics)) + + self.lightning_module = _LM() + + +def _eval_video(viz_max_batches=2): + from egomimic.eval.eval_video import EvalVideo + + class _Impl(EvalVideo): + def __init__(self, **kw): + super().__init__(**kw) + self.rendered: list[int] = [] + self.forwards = 0 + + def compute_metrics_and_viz(self, batch, do_viz=True): + self.forwards += 1 + if do_viz: + self.rendered.append(batch["i"]) + return {"Valid/mse": 1.0}, {HUMAN: np.zeros((1, 2, 2, 3), np.uint8)} + return {"Valid/mse": 1.0}, {} + + ev = _Impl(viz_every_n_epochs=1, viz_max_batches=viz_max_batches) + ev.trainer = _Trainer() + return ev + + +def test_metric_mode_never_renders_and_video_mode_never_logs(): + ev = _eval_video() + for i in range(3): + ev.on_validation_step({"i": i}, i, 0, mode="metrics") + assert ev.rendered == [], "no frame rendered on the metric loader" + assert ev.val_image_buffer == {}, "no frame buffered either" + assert len(ev.trainer.logged) == 3, "metrics logged for every metric batch" + + ev.trainer.logged.clear() + for i in range(4): + ev.on_validation_step({"i": i}, i, 1, mode="video") + assert ev.rendered == [0, 1], "rendered up to viz_max_batches" + assert ev.trainer.logged == [], "video loader logs no metric key" + # past viz_max_batches the video loader does not even run the forward + assert ev.forwards == 3 + 2 + + with pytest.raises(ValueError, match="unknown validation mode"): + ev.on_validation_step({"i": 0}, 0, 0, mode="bogus") + + +def test_default_mode_is_todays_behaviour(): + ev = _eval_video() + ev.on_validation_step({"i": 0}, 0, 0) + assert ev.rendered == [0] and len(ev.trainer.logged) == 1 From ce4c38393a473465ff4bcf4e82eaefc0ff7c5ddf Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 15 Sep 2026 15:13:28 -0400 Subject: [PATCH 11/79] data: mecka_fold_flagship_top3_hpt_6d, the 3-operator flagship train split Three-operator twin of mecka_fold_flagship_opsplit_hpt_6d (and the wider sibling of the single-operator topop config): train on the 3 largest flagship fold-clothes operators (1,079 eps / 206,276 frames, 65.3 % of flagship), none of them held out. Seen-op val is the seeded 5 % episode split within those operators (53 eps / 9,826 frames); the unseen-operator head is inherited unchanged, so the runs stay comparable. Lane A keys measured from the real split on a compute node 2026-09-15: metric_frames_per_episode valid=floor(5120/53)=96, train_viz=floor(5120/1026)=4 (single-rank values; trainHydra scales them by the world size) (unseen 26 inherited); video_episodes pins one seen-val episode per train operator (191+223+226 = 640 f = exactly the 10-batch video budget) and inherits the train_viz and unseen pins. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- .../data/mecka_fold_flagship_top3_hpt_6d.yaml | 93 +++++++++++++++++++ tests/unit/test_eval_loaders.py | 78 ++++++++++++++++ 2 files changed, 171 insertions(+) create mode 100644 egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_6d.yaml diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_6d.yaml new file mode 100644 index 0000000000..ba728fd656 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_6d.yaml @@ -0,0 +1,93 @@ +defaults: + - mecka_fold_flagship_opsplit_hpt_6d + - _self_ + +# MECKA_FOLD_FLAGSHIP_TOP3_HPT_6D: THREE-operator twin of +# mecka_fold_flagship_opsplit_hpt_6d (and the wider sibling of the +# single-operator mecka_fold_flagship_topop_hpt_6d). TRAIN on the 3 largest +# flagship fold-clothes operators only; everything else (pipeline, +# unseen-operator val, loader params) is inherited from the parent. +# +# SPLIT (operator totals from SQL app.episodes, 2026-09-15; the train/val +# numbers measured from the real seeded split on a compute node the same day): +# train ops = the 3 largest of the 20 flagship operators, by frames: +# 6903686e0e94ce070afd1f24 405 eps / 76,013 frames +# 690366b20e94ce070afd1e8a 354 eps / 67,085 frames +# 683785ac01ca734152093448 320 eps / 63,178 frames +# total 1,079 eps / 206,276 frames +# = 1.91 h at 30 fps, 65.3 % of the flagship set. +# None of the three is in the parent's `held_out_operators`. +# train = 95 % of those episodes: 1,026 eps / 196,450 frames +# (6903686e 382 eps/71,722 f, 690366b2 334/63,397, +# 683785ac 310/61,331). +# seen val = the other 5 % (53 eps / 9,826 frames) -> seen_op_valid/. The +# parent's valid_datasets interpolates the train filters and +# valid_ratio, so it follows this override automatically. +# train_viz = the train split above (trainHydra's default) -> train_viz/. +# unseen val = UNCHANGED from the parent: all episodes of the 2 +# `held_out_operators` (191 eps / 30,937 frames) -> +# unseen_op_valid/, directly comparable with the 18-operator +# and single-operator runs. +# The other 15 operators are unused. +# +# Hydra replaces lists wholesale, so the train filter_lambdas are restated in +# full; only the last (operator) lambda differs from the parent. The operator +# ids are inlined rather than put in a config key: MultiDataModuleWrapper +# rejects root keys it does not declare. +# +# NORMALIZATION: the train distribution differs from the parent's and from +# topop's, so precompute stats on THIS config: +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_flagship_6d_hpt \ +# --data mecka_fold_flagship_top3_hpt_6d --model hpt_bc_mecka_6d_300M --out +# then train with +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_6d_hpt \ +# data=mecka_fold_flagship_top3_hpt_6d norm_stats.precomputed_norm_path=/norm_stats + +train_datasets: + human_bimanual: + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() == 'fold_clothes'" + - "lambda row: '/mecka/flagship/' in str(row.get('zarr_processed_path', ''))" + # TOP-3 OPERATOR TRAIN: the three largest flagship operators only. + - "lambda row: str(row.get('operator', '')) in ['6903686e0e94ce070afd1f24', '690366b20e94ce070afd1e8a', '683785ac01ca734152093448']" + +# --------------------------------------------------------------------------- +# METRIC / VIDEO LOADERS: only the SEEN-op and train_viz entries differ from +# the parent (this config trains on 3 of the 18 operators, so those two splits +# are smaller); the unseen head is identical, so its K and pins are inherited. +# --------------------------------------------------------------------------- +# K is the SINGLE-RANK value, floor(5120 / episodes-in-that-split), 5120 = +# limit_val_batches (80) x batch_size (64); trainHydra scales it by the world +# size. Episode counts measured from the real splits 2026-09-15: +# valid 53 eps / 9,826 frames -> K = floor(5120/53) = 96. The +# shortest seen-val episode is 127f > 96, so every episode +# contributes exactly 96: 53 x 96 = 5,088 f = 79.5 batches. +# train_viz 1,026 eps / 196,450 frames -> K = floor(5120/1026) = 4 +# -> 1026 x 4 = 4,104 f = 64.1 batches. +# `unseen_op_valid: 26` is inherited unchanged (191 eps, same split as the +# parent), and hydra merges dicts, so only the two changed keys are restated. +metric_frames_per_episode: + valid: 96 + train_viz: 4 + +# The parent's seen-val pin belongs to operator 690364170e94ce070afd1da5, +# which is not in this config's 3-operator split, so `valid` is replaced +# (hydra replaces lists wholesale). Seen-val now spans 3 operators, so ONE +# episode is pinned per operator and the video shows all three. The train_viz +# pin 692e9bda21fc595fbe6f941e IS an episode of operator 6903686e..., hence in +# this train split too (confirmed on the measured split), and the unseen pins +# are the same split as the parent's -- both inherited. +# valid 692ea44dc621d7f4aac3aae6 (683785ac01ca734152093448, 191f) +# 692eb08997ffd8ec90340290 (690366b20e94ce070afd1e8a, 223f) +# 692fe79f34a99e18e25289ad (6903686e0e94ce070afd1f24, 226f) +# All three are at or above the split's 177f median length, and +# 191 + 223 + 226 = 640 f = EXACTLY the 10-batch +# (evaluator.viz_max_batches x 64) video budget, so each operator's +# episode is rendered whole whatever order the loader uses. +video_episodes: + valid: + - 692ea44dc621d7f4aac3aae6 + - 692eb08997ffd8ec90340290 + - 692fe79f34a99e18e25289ad diff --git a/tests/unit/test_eval_loaders.py b/tests/unit/test_eval_loaders.py index 238319b6e3..626e003715 100644 --- a/tests/unit/test_eval_loaders.py +++ b/tests/unit/test_eval_loaders.py @@ -16,6 +16,7 @@ from __future__ import annotations +import re from types import SimpleNamespace import numpy as np @@ -62,6 +63,23 @@ ], }, ), + "mecka_fold_flagship_top3_hpt_6d": ( + {"valid": 96, "train_viz": 4, "unseen_op_valid": 26}, + { + # one seen-val pin per train operator (683785ac / 690366b2 / + # 6903686e), 191 + 223 + 226 = 640f = the whole video budget + "valid": [ + "692ea44dc621d7f4aac3aae6", + "692eb08997ffd8ec90340290", + "692fe79f34a99e18e25289ad", + ], + "train_viz": ["692e9bda21fc595fbe6f941e"], + "unseen_op_valid": [ + "692eadbaaec602a46af10686", + "692ea49a727c13b350cb7c48", + ], + }, + ), } # Episode counts these K values were computed from (measured 2026-09-15). @@ -76,6 +94,11 @@ "train_viz": 385, "unseen_op_valid": 191, }, + "mecka_fold_flagship_top3_hpt_6d": { + "valid": 53, + "train_viz": 1026, + "unseen_op_valid": 191, + }, } @@ -138,6 +161,61 @@ def test_flagship_configs_declare_k_and_pins(data_config, compose_resolve): assert th._metric_frames_per_episode(cfg, head) == k +# The 3 largest flagship operators (SQL app.episodes, 2026-09-15), by frames. +TOP3_OPERATORS = [ + "6903686e0e94ce070afd1f24", + "690366b20e94ce070afd1e8a", + "683785ac01ca734152093448", +] + + +def test_top3_config_trains_on_exactly_the_three_largest_operators(compose_resolve): + """mecka_fold_flagship_top3_hpt_6d composes under the flagship recipe, its + train filter names exactly the 3 top operators (none of them held out), the + seen-op val follows that filter, and both lane A tables cover all 3 heads.""" + cfg = compose_resolve(FLAGSHIP, ["data=mecka_fold_flagship_top3_hpt_6d"]) + parent = compose_resolve(FLAGSHIP, ["data=mecka_fold_flagship_opsplit_hpt_6d"]) + d, pd = cfg.data, parent.data + train = d.train_datasets[HUMAN] + + # the operator lambda names exactly the three ids, largest first + operator_lambda = train.filters.filter_lambdas[-1] + assert "row.get('operator'" in operator_lambda + assert re.findall(r"[0-9a-f]{24}", operator_lambda) == TOP3_OPERATORS + held_out = [str(o) for o in d.held_out_operators] + assert not set(TOP3_OPERATORS) & set(held_out), "a train operator is held out" + + # the first three (lab / task / flagship-path) lambdas are the parent's + assert list(train.filters.filter_lambdas[:3]) == list( + pd.train_datasets[HUMAN].filters.filter_lambdas[:3] + ) + + # seen-op val interpolates the train filters, so it follows this override + valid = d.valid_datasets[HUMAN] + assert list(valid.filters.filter_lambdas) == list(train.filters.filter_lambdas) + assert valid.valid_ratio == train.valid_ratio == 0.05 + assert (train.mode, valid.mode) == ("train", "valid") + assert d.valid_prefix == "seen_op_valid" + assert d.get("train_viz_datasets") is None + + # the unseen-operator head is untouched + unseen = d.unseen_op_valid_datasets[HUMAN] + assert list(unseen.filters.filter_lambdas) == list( + pd.unseen_op_valid_datasets[HUMAN].filters.filter_lambdas + ) + assert unseen.mode == "total" + + # lane A keys cover all three heads + heads = {"valid", "train_viz", "unseen_op_valid"} + k_table = OmegaConf.to_container(d.metric_frames_per_episode) + pins = OmegaConf.to_container(d.video_episodes) + assert set(k_table) == heads + assert set(pins) == heads + assert all(isinstance(v, int) and v > 0 for v in k_table.values()) + assert len(pins["valid"]) == 3, "one seen-val video pin per train operator" + assert len(set(pins["valid"])) == 3 + + def test_config_without_the_keys_is_unchanged(compose_resolve): """Backward compatibility: a data config that declares neither key gets the old loader names and the unwrapped dataset objects.""" From 96bd8d718c743313aed5118ba70f6b30a2b33f7c Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sun, 13 Sep 2026 14:21:23 -0400 Subject: [PATCH 12/79] configs: declare per-embodiment dims once per HPT model config stem input_dim, infer_ac_dims, act_dim and enc-dec output_dim now interpolate robomimic_model.dims.{.proprio,.action,action_width}. Resolved values are unchanged (pinned by tests/unit/test_model_dims.py). Pi configs untouched. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01MWtR1x5tjREwoxQZAnUGrh Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01BbQNNaMAbTNMypegmJkv4s From 222faf85995ab4e666eeecf06e9c22debdfdb1f8 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 15 Sep 2026 14:41:09 -0400 Subject: [PATCH 13/79] eval: deterministic validation (seeded prompts + flow sampling); log-only grad-norm MAD Validating one checkpoint twice now gives identical prompts, actions and metrics. * HPT gets a per-pass RNG seed, base + rank * stride + pass counter, where the counter advances once per forward_eval call and is rewound by ModelWrapper.on_validation_start. forward_eval builds a torch.Generator from it and threads it down HPTModel.forward -> DenoisingPolicy.forward -> sample_action, and forks/seeds the global RNG around the BC val-loss call (DenoisingPolicy.predict draws its own noise and timestep, out of reach of any generator kwarg). Every generator kwarg defaults to None, so training and all other callers are untouched. The val-loss seed is offset by EVAL_LOSS_STRIDE rather than by 1: a torch.Generator seeded with S draws the same stream as a global manual_seed(S), so +1 would have handed batch i's val loss exactly the noise batch i+1 samples with. eval_base_seed is a real __init__ kwarg, so a second eval pass can be made to draw different noise on purpose. * _build_prompts picks annotations with that same seeded RNG whenever the algo is in eval mode, regardless of annotation_sampling_mode, so eval keeps prompt variety instead of collapsing to sample[0]. Train mode is unchanged. * The MAD grad-norm spike detector in on_after_backward is now log-only: it still logs the raw norm, the threshold, the flag and the [GRAD_NORM_SPIKE] line, but never rescales gradients, and every step (flagged or not) goes into the rolling median/MAD window. Fixed clipping moves to Lightning: trainer/default.yaml gains gradient_clip_val (null for now) and gradient_clip_algorithm: norm. The on_before_optimizer_step hook goes with it: Precision._after_closure calls that hook BEFORE _clip_gradients, so "policy_grad_norms_clipped" was the pre-clip norm under another name -- bit-identical to policy_grad_norms_raw, and no more informative once gradient_clip_val is set. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/algo/hpt.py | 132 +++++- egomimic/hydra_configs/trainer/default.yaml | 9 +- egomimic/models/denoising_policy.py | 11 +- egomimic/pl_utils/pl_model.py | 54 ++- tests/unit/test_eval_determinism.py | 435 ++++++++++++++++++++ 5 files changed, 597 insertions(+), 44 deletions(-) create mode 100644 tests/unit/test_eval_determinism.py diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index 979a536b79..cbed1282a9 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -23,6 +23,21 @@ # it has exactly one consumer, below). STD_SCALE = 0.02 +# Deterministic evaluation. Every RNG used while the algo is in eval mode is +# seeded from ``EVAL_BASE_SEED + rank * EVAL_RANK_STRIDE + pass_counter``, where +# the pass counter counts ``forward_eval`` calls since the last +# ``reset_eval_pass_counter()`` (called from ``ModelWrapper.on_validation_start``). +# Two validation passes over the same batches therefore replay the same prompts +# and the same sampling noise, while different ranks and different batches +# within a pass still see different draws. +EVAL_BASE_SEED = 0 +EVAL_RANK_STRIDE = 1_000_003 +# Offset of the val-loss RNG from the sampling generator's seed. It must not be +# a small number: the pass counter advances by 1 per forward_eval call, and a +# torch.Generator seeded with S draws the same stream as a global manual_seed +# of S, so an offset of 1 gave batch i's val loss the noise batch i+1 sampled. +EVAL_LOSS_STRIDE = 7_919_369 + class HPTModel(nn.Module): """ @@ -659,7 +674,7 @@ def compute_loss(self, batch): total_loss = action_loss + shared_action_loss + auxiliary_action_loss return total_loss - def forward(self, domain, data): + def forward(self, domain, data, generator=None): """ Forward pass of the HPTModel to compute actions. @@ -669,6 +684,10 @@ def forward(self, domain, data): The domain corresponding to the input data. data : dict Dictionary containing input data for various modalities. + generator : torch.Generator, optional + RNG handed to the denoising heads' sampler so that evaluation is + reproducible (see ``HPT.forward_eval``). ``None`` (the default) + keeps the global RNG, so every other caller is unchanged. Returns ------- @@ -678,19 +697,23 @@ def forward(self, domain, data): features, block_outputs = self.forward_features(domain, data) action = {} + # Only the denoising (diffusion / flow matching) heads sample noise and + # accept a generator; the plain MLP heads take the features alone. + head_kwargs = {"generator": generator} if self.diffusion else {} + if self.diffusion: features = (features, domain) if domain in self.heads: - action[domain] = self.heads[domain](features) + action[domain] = self.heads[domain](features, **head_kwargs) if self.shared_action: - action["shared"] = self.heads["shared"](features) + action["shared"] = self.heads["shared"](features, **head_kwargs) if domain in self.auxiliary_ac_keys: for key in self.auxiliary_ac_keys[domain]: if f"{domain}_{key}" in self.heads: - action[key] = self.heads[f"{domain}_{key}"](features) + action[key] = self.heads[f"{domain}_{key}"](features, **head_kwargs) return action @@ -855,6 +878,7 @@ def __init__( model = HPTModel(**trunk) model.auxiliary_ac_keys = self.auxiliary_ac_keys + self.eval_base_seed = int(kwargs.get("eval_base_seed", EVAL_BASE_SEED)) self.multitask = kwargs.get("multitask", False) self.device = kwargs.get( "device", torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -956,16 +980,82 @@ def __init__( self.training_step = 0 + # ------------------------------------------------------------------ + # Deterministic evaluation + # ------------------------------------------------------------------ + def reset_eval_pass_counter(self) -> None: + """Start a new validation pass: the next ``forward_eval`` call is pass 0. + + Called from ``ModelWrapper.on_validation_start``. Two passes over the + same loader (same order, same number of batches) then draw the same + prompts and the same sampling noise. + """ + self._eval_pass_counter = 0 + + @staticmethod + def _eval_rank() -> int: + """Distributed rank, or 0 outside a process group.""" + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_rank() + for var in ("RANK", "LOCAL_RANK", "SLURM_PROCID"): + value = os.environ.get(var) + if value is not None and value.lstrip("-").isdigit(): + return int(value) + return 0 + + def _eval_pass_seed(self) -> int: + """Seed for the current eval pass: base + rank stride + pass counter.""" + return ( + self.eval_base_seed + + self._eval_rank() * EVAL_RANK_STRIDE + + int(getattr(self, "_eval_pass_counter", 0)) + ) + + def _eval_device(self) -> torch.device: + return ( + torch.device(self.device) + if self.device is not None + else torch.device("cpu") + ) + + def _eval_generator(self) -> torch.Generator: + """Seeded generator on the device the sampling noise is drawn on.""" + device = self._eval_device() + generator = torch.Generator(device=device) + generator.manual_seed(self._eval_pass_seed()) + return generator + + def _eval_prompt_rng(self) -> random.Random: + """Seeded Python RNG for annotation choice during eval.""" + return random.Random(self._eval_pass_seed()) + + def _eval_fork_devices(self) -> list: + """Devices ``torch.random.fork_rng`` must save/restore (CPU is implicit).""" + device = self._eval_device() + if device.type != "cuda": + return [] + return [ + device.index if device.index is not None else torch.cuda.current_device() + ] + def _build_prompts(self, _batch, batch_size: int) -> list[str]: """Sample one annotation per batch item, falling back to default_prompt on empty / missing annotations. Mirrors the Pi algo flow. + + In eval the choice comes from a per-pass seeded RNG regardless of + ``annotation_sampling_mode``: that keeps prompt variety (unlike always + taking ``sample[0]``) while making two validation passes over the same + batches produce identical prompt lists. Training is unchanged. """ if self.annotation_key is None or self.annotation_key not in _batch: return [self.default_prompt] * batch_size + eval_rng = None if self.nets.training else self._eval_prompt_rng() prompts = [] for sample in _batch[self.annotation_key]: if not sample: prompts.append(self.default_prompt) + elif eval_rng is not None: + prompts.append(sample[eval_rng.randint(0, len(sample) - 1)]) elif self.annotation_sampling_mode == "random": prompts.append(sample[random.randint(0, len(sample) - 1)]) else: # "first" @@ -1101,6 +1191,12 @@ def forward_eval(self, batch): } """ unnorm_preds = {} + # One seeded generator per pass/batch: the flow-matching noise below and + # the prompts picked in ``process_batch_for_training`` share this pass's + # seed, so replaying the pass replays every draw. + generator = self._eval_generator() + pass_seed = self._eval_pass_seed() + fork_devices = self._eval_fork_devices() for embodiment_id, _batch in batch.items(): embodiment_name = get_embodiment(embodiment_id).lower() cam_keys = self.camera_keys[embodiment_id] @@ -1119,16 +1215,24 @@ def forward_eval(self, batch): # so keep a fresh copy for the forward() call below. forward_data = self._clone_batch(hpt_batch["data"]) - # BC val loss — same call as forward_training. - if self.freeze_repr: - val_loss = self.nets["policy"].compute_loss_depth( - hpt_batch, depth=self.freeze_depth - ) - else: - val_loss = self.nets["policy"].compute_loss(hpt_batch) + # BC val loss — same call as forward_training. The training loss + # draws its own noise / timestep from the global RNG deep inside the + # head (``DenoisingPolicy.predict``), which no generator kwarg + # reaches; fork and seed the global RNG around the call instead so + # the reported val loss is reproducible too. + with torch.random.fork_rng(devices=fork_devices): + torch.manual_seed(pass_seed + EVAL_LOSS_STRIDE) + if self.freeze_repr: + val_loss = self.nets["policy"].compute_loss_depth( + hpt_batch, depth=self.freeze_depth + ) + else: + val_loss = self.nets["policy"].compute_loss(hpt_batch) unnorm_preds[f"{embodiment_name}_loss"] = val_loss - actions = self.nets["policy"].forward(hpt_batch["domain"], forward_data) + actions = self.nets["policy"].forward( + hpt_batch["domain"], forward_data, generator=generator + ) predictions = OrderedDict() for key in actions: @@ -1153,6 +1257,10 @@ def forward_eval(self, batch): for key in unnorm_actions: unnorm_preds[f"{embodiment_name}_{key}"] = unnorm_actions[key] + # Advance the pass counter once per forward_eval call (after the batch, + # so the prompts built for this batch and the noise drawn for it share a + # seed). ``reset_eval_pass_counter`` rewinds it at on_validation_start. + self._eval_pass_counter = int(getattr(self, "_eval_pass_counter", 0)) + 1 return unnorm_preds @override diff --git a/egomimic/hydra_configs/trainer/default.yaml b/egomimic/hydra_configs/trainer/default.yaml index 62f147c27d..e75492a53f 100644 --- a/egomimic/hydra_configs/trainer/default.yaml +++ b/egomimic/hydra_configs/trainer/default.yaml @@ -17,4 +17,11 @@ check_val_every_n_epoch: 100 # set True to to ensure deterministic results # makes training slower but gives more reproducibility than just setting seeds -deterministic: False \ No newline at end of file +deterministic: False + +# Fixed gradient clipping, applied by Lightning in configure_gradient_clipping. +# The MAD spike detector in ModelWrapper.on_after_backward is log-only and never +# touches gradients; this is the only thing that clips. null = no clipping (the +# value is picked from the grad-norm curve of a log-only baseline run). +gradient_clip_val: null +gradient_clip_algorithm: norm diff --git a/egomimic/models/denoising_policy.py b/egomimic/models/denoising_policy.py index 645a8c44e4..b9180e0d37 100644 --- a/egomimic/models/denoising_policy.py +++ b/egomimic/models/denoising_policy.py @@ -80,9 +80,16 @@ def sample_action(self, global_cond, embodiment_name, generator=None): ) return self.inference(noise, global_cond, generator) - def forward(self, global_cond): + def forward(self, global_cond, generator=None): + """Sample actions from the conditioning. + + ``generator`` is threaded down to ``sample_action`` so the eval path can + make sampling reproducible (see ``HPT.forward_eval``). It defaults to + ``None``, i.e. the global RNG, so training and every other caller are + unchanged. + """ cond, embodiment = global_cond - return self.sample_action(cond, embodiment) + return self.sample_action(cond, embodiment, generator=generator) def predict(self, actions, global_cond) -> Tuple[torch.Tensor, torch.Tensor]: """ diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index 306f541ad8..dc60674c83 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/egomimic/pl_utils/pl_model.py @@ -184,38 +184,26 @@ def on_after_backward(self): info["policy_grad_norms_mad_threshold"] = threshold grad_norm_flagged = grad_norm_val > threshold info["policy_grad_norms_mad_flag"] = float(grad_norm_flagged) - if grad_norm_flagged: - torch.nn.utils.clip_grad_norm_(self.parameters(), max_norm=median) - if self.trainer.is_global_zero: - print( - "[GRAD_NORM_SPIKE] " - f"step={self.global_step} " - f"grad_norm={grad_norm_val:.4f} " - f"median={median:.4f} " - f"mad={mad:.4f} " - f"threshold={threshold:.4f}", - flush=True, - ) - - if not grad_norm_flagged: - self.grad_norm_history.append(grad_norm_val) + if grad_norm_flagged and self.trainer.is_global_zero: + print( + "[GRAD_NORM_SPIKE] " + f"step={self.global_step} " + f"grad_norm={grad_norm_val:.4f} " + f"median={median:.4f} " + f"mad={mad:.4f} " + f"threshold={threshold:.4f}", + flush=True, + ) + + # The MAD detector is log-only: it never rescales gradients (actual + # clipping is Lightning's fixed ``trainer.gradient_clip_val``), so every + # step goes into the rolling history. The old code excluded flagged + # steps only to stop the clip threshold from drifting upwards; as a pure + # diagnostic the plain rolling median/MAD is the honest statistic. + self.grad_norm_history.append(grad_norm_val) for k, v in info.items(): self.log("Train/" + k, v, on_step=False, on_epoch=True, sync_dist=True) - def on_before_optimizer_step(self, optimizer): - if not self.enable_grad_norm: - return - grad_norm = torch.nn.utils.clip_grad_norm_( - self.parameters(), max_norm=float("inf") - ) - self.log( - "Train/policy_grad_norms_clipped", - float(grad_norm), - on_step=False, - on_epoch=True, - sync_dist=True, - ) - def _val_heads(self) -> dict: return { "valid": self.evaluator, @@ -224,6 +212,14 @@ def _val_heads(self) -> dict: } def on_validation_start(self): + # Deterministic eval: rewind the algo's per-pass RNG counter so this + # pass replays the same prompts and sampling noise as any other pass + # over the same batches (egomimic/algo/hpt.py). Algos without the hook + # (e.g. PI) are unaffected. + reset_eval_rng = getattr(self.model, "reset_eval_pass_counter", None) + if callable(reset_eval_rng): + reset_eval_rng() + heads = [h for h in self._val_heads().values() if h is not None] if not heads: return diff --git a/tests/unit/test_eval_determinism.py b/tests/unit/test_eval_determinism.py new file mode 100644 index 0000000000..be5bea96af --- /dev/null +++ b/tests/unit/test_eval_determinism.py @@ -0,0 +1,435 @@ +"""Deterministic evaluation (lane B). + +Validating the same checkpoint twice must give the same prompts, the same +sampled actions and the same metrics. Three pieces make that true: + +* ``HPT._build_prompts`` picks annotations with a per-pass seeded RNG whenever + the algo is in eval mode (``annotation_sampling_mode`` only steers training); +* ``HPT.forward_eval`` builds a ``torch.Generator`` seeded from + ``EVAL_BASE_SEED + rank * EVAL_RANK_STRIDE + pass_counter`` and threads it + down to the denoising head's sampler, and forks/seeds the global RNG around + the BC val-loss call; ``ModelWrapper.on_validation_start`` rewinds the + counter, so pass N and pass N+1 replay the same draws. + +Plus the gradient-clipping change that rides along in this branch: the MAD +spike detector in ``ModelWrapper.on_after_backward`` is log-only (fixed +clipping is Lightning's ``trainer.gradient_clip_val``). + +Everything here is CPU-only and builds the smallest object that still runs the +real code path: the real ``FMPolicy`` Euler sampler, the real +``HPTModel.forward`` head dispatch and the real ``HPT.forward_eval``, with only +the stems/trunk (``forward_features``) and the norm stats stubbed out. +""" + +from __future__ import annotations + +import random +import types +from collections import deque + +import hydra +import pytest +import torch +import torch.nn as nn + +from egomimic.algo.hpt import EVAL_BASE_SEED, HPT, HPTModel +from egomimic.models.fm_policy import FMPolicy +from egomimic.pl_utils.pl_model import ModelWrapper + +DOMAIN = "eva_bimanual" +EMB_ID = 6 # EMBODIMENT.EVA_BIMANUAL; get_embodiment(6).lower() == DOMAIN +AC_KEY = "actions_cartesian" +ANNOTATION_KEY = "annotations" +BATCH = 4 +HORIZON = 3 +ACTION_DIM = 4 +FEAT_DIM = 5 + + +# --------------------------------------------------------------------------- +# Fixtures: a minimal but real HPT +# --------------------------------------------------------------------------- +class _TinyVelocity(nn.Module): + """Stands in for ConditionalUnet1D: v(x_t, t, cond), same call signature.""" + + def __init__(self, dim: int): + super().__init__() + self.lin = nn.Linear(dim, dim) + + def forward(self, x_t, t, global_cond): + return self.lin(x_t) * t.reshape(-1, 1, 1) + + +class _StubNormStats: + def unnormalize(self, predictions, embodiment_id): + return dict(predictions) + + +def _make_algo(annotation_key: str | None = ANNOTATION_KEY) -> HPT: + """A real HPT wired to a real FMPolicy head, without stems/trunk/dataset. + + ``HPT.__init__`` needs norm stats, configs and a dataset; every attribute + ``forward_eval`` / ``_build_prompts`` actually read is set here instead. + """ + torch.manual_seed(0) + head = FMPolicy( + model=_TinyVelocity(ACTION_DIM), + action_horizon=HORIZON, + infer_ac_dims={DOMAIN: ACTION_DIM}, + num_inference_steps=2, + ) + model = HPTModel.__new__(HPTModel) + nn.Module.__init__(model) + model.diffusion = True + model.shared_action = False + model.auxiliary_ac_keys = {} + model.heads = nn.ModuleDict({DOMAIN: head}) + model.device = torch.device("cpu") + features = ( + torch.arange(BATCH * FEAT_DIM, dtype=torch.float32).reshape(BATCH, FEAT_DIM) + / 10.0 + ) + model.forward_features = lambda domain, data: (features, None) + + algo = HPT.__new__(HPT) + algo.nets = nn.ModuleDict({"policy": model}) + algo.device = torch.device("cpu") + algo.norm_stats = _StubNormStats() + algo.camera_keys = {EMB_ID: []} + algo.proprio_keys = {EMB_ID: []} + algo.lang_keys = {EMB_ID: []} + algo.ac_keys = {EMB_ID: AC_KEY} + algo.auxiliary_ac_keys = {} + algo.shared_ac_key = None + algo.is_6dof = False + algo.freeze_repr = False + algo.freeze_depth = 8 + algo.eval_base_seed = EVAL_BASE_SEED + algo.annotation_key = annotation_key + algo.annotation_modality = "annotation" + algo.annotation_sampling_mode = "random" + algo.default_prompt = "do the task" + algo.train_image_augs = None + algo.eval_image_augs = None + algo.nets.eval() + return algo + + +def _make_batch() -> dict: + torch.manual_seed(1) + return { + EMB_ID: { + AC_KEY: torch.randn(BATCH, HORIZON, ACTION_DIM), + "pad_mask": torch.ones(BATCH, HORIZON, 1), + "embodiment": torch.tensor([EMB_ID]), + } + } + + +def _annotation_batch(n_options: int = 4) -> dict: + return { + ANNOTATION_KEY: [ + [f"sample{i}_option{j}" for j in range(n_options)] for i in range(BATCH) + ] + } + + +def _prompt_pass(algo: HPT, raw_batch: dict, n_batches: int = 3) -> list[list[str]]: + """One validation pass: N batches, counter advancing as forward_eval does.""" + algo.reset_eval_pass_counter() + out = [] + for _ in range(n_batches): + out.append(algo._build_prompts(raw_batch, BATCH)) + algo._eval_pass_counter += 1 # what forward_eval does after each batch + return out + + +# --------------------------------------------------------------------------- +# 1. Prompts +# --------------------------------------------------------------------------- +def test_eval_prompts_are_reproducible_across_passes(): + algo = _make_algo() + raw = _annotation_batch() + first = _prompt_pass(algo, raw) + second = _prompt_pass(algo, raw) + assert first == second + + +def test_eval_prompts_keep_variety_and_are_not_always_first(): + algo = _make_algo() + raw = _annotation_batch() + prompts = [p for batch in _prompt_pass(algo, raw, n_batches=5) for p in batch] + firsts = [options[0] for options in raw[ANNOTATION_KEY]] + assert any( + p not in firsts for p in prompts + ), f"eval prompts collapsed to sample[0]: {prompts}" + assert len(set(prompts)) > 1 + + +def test_eval_prompts_ignore_sampling_mode(): + """`annotation_sampling_mode` is an algo-level knob; eval overrides it.""" + algo = _make_algo() + raw = _annotation_batch() + algo.annotation_sampling_mode = "random" + as_random = _prompt_pass(algo, raw) + algo.annotation_sampling_mode = "first" + as_first = _prompt_pass(algo, raw) + assert as_random == as_first + assert as_first[0] != [options[0] for options in raw[ANNOTATION_KEY]] + + +def test_eval_prompts_fall_back_to_default_on_empty(): + algo = _make_algo() + raw = {ANNOTATION_KEY: [[], ["a", "b"], [], ["c"]]} + prompts = algo._build_prompts(raw, 4) + assert prompts[0] == algo.default_prompt + assert prompts[2] == algo.default_prompt + assert prompts[1] in {"a", "b"} and prompts[3] == "c" + + +def test_missing_annotation_key_uses_default_prompt(): + algo = _make_algo(annotation_key=None) + assert algo._build_prompts({}, 3) == [algo.default_prompt] * 3 + + +@pytest.mark.parametrize("mode", ["random", "first"]) +def test_train_prompt_behaviour_is_unchanged(mode): + """In train mode the parent's logic still applies verbatim.""" + algo = _make_algo() + algo.annotation_sampling_mode = mode + algo.nets.train() + raw = _annotation_batch() + + random.seed(12345) + got = algo._build_prompts(raw, BATCH) + + random.seed(12345) + expected = [] + for options in raw[ANNOTATION_KEY]: + if mode == "random": + expected.append(options[random.randint(0, len(options) - 1)]) + else: + expected.append(options[0]) + assert got == expected + + +# --------------------------------------------------------------------------- +# 2. Seeded flow sampling +# --------------------------------------------------------------------------- +def _eval_pass(algo: HPT, batch: dict, n_batches: int = 2) -> list[dict]: + algo.reset_eval_pass_counter() + return [algo.forward_eval(batch) for _ in range(n_batches)] + + +def _actions(preds: list[dict]) -> list[torch.Tensor]: + return [p[f"{DOMAIN}_{AC_KEY}"] for p in preds] + + +def test_two_validation_passes_give_bit_identical_actions(): + algo = _make_algo() + batch = _make_batch() + first = _eval_pass(algo, batch) + second = _eval_pass(algo, batch) + for a, b in zip(_actions(first), _actions(second)): + assert torch.equal(a, b) + + +def test_val_loss_is_identical_across_passes(): + algo = _make_algo() + batch = _make_batch() + first = _eval_pass(algo, batch) + second = _eval_pass(algo, batch) + for a, b in zip(first, second): + assert torch.equal(a[f"{DOMAIN}_loss"], b[f"{DOMAIN}_loss"]) + + +def test_batches_within_a_pass_use_different_noise(): + """Without the reset the counter has moved on, so the draws differ.""" + algo = _make_algo() + batch = _make_batch() + algo.reset_eval_pass_counter() + first = algo.forward_eval(batch) + second = algo.forward_eval(batch) # no reset: pass counter is now 1 + assert not torch.equal(first[f"{DOMAIN}_{AC_KEY}"], second[f"{DOMAIN}_{AC_KEY}"]) + + +def test_forward_eval_does_not_disturb_the_global_rng(): + """Eval must not shift the training RNG stream (fork_rng + generator).""" + algo = _make_algo() + batch = _make_batch() + torch.manual_seed(7) + before = torch.randn(3) + torch.manual_seed(7) + algo.reset_eval_pass_counter() + algo.forward_eval(batch) + after = torch.randn(3) + assert torch.equal(before, after) + + +def test_pass_seed_depends_on_rank_and_counter(monkeypatch): + algo = _make_algo() + algo.reset_eval_pass_counter() + monkeypatch.setenv("RANK", "0") + rank0 = algo._eval_pass_seed() + monkeypatch.setenv("RANK", "3") + rank3 = algo._eval_pass_seed() + assert rank0 != rank3 + algo._eval_pass_counter += 1 + assert algo._eval_pass_seed() == rank3 + 1 + + +def test_val_loss_noise_does_not_collide_with_the_next_batch(monkeypatch): + """A torch.Generator seeded with S draws the same stream as a global + manual_seed(S), so a +1 offset would give batch i's val loss the very noise + batch i+1 samples with.""" + from egomimic.algo.hpt import EVAL_LOSS_STRIDE + + assert EVAL_LOSS_STRIDE > 1000 + algo = _make_algo() + algo.reset_eval_pass_counter() + seeds = set() + for _ in range(4): + pass_seed = algo._eval_pass_seed() + seeds.add(pass_seed) + seeds.add(pass_seed + EVAL_LOSS_STRIDE) + algo._eval_pass_counter += 1 + assert len(seeds) == 8, "a val-loss seed repeats a sampling seed" + + +def test_eval_base_seed_is_a_real_knob(): + """A second eval pass can be made to draw DIFFERENT noise on purpose.""" + algo = _make_algo() + assert algo.eval_base_seed == EVAL_BASE_SEED + algo.reset_eval_pass_counter() + default = algo._eval_pass_seed() + algo.eval_base_seed = 12345 + assert algo._eval_pass_seed() == default + 12345 + + +def test_generator_is_threaded_to_the_head(): + """HPTModel.forward -> DenoisingPolicy.forward -> sample_action(generator).""" + algo = _make_algo() + model = algo.nets["policy"] + data = {"action": torch.zeros(BATCH, HORIZON, ACTION_DIM)} + gen = torch.Generator(device="cpu") + + gen.manual_seed(42) + a = model.forward(DOMAIN, data, generator=gen)[DOMAIN] + gen.manual_seed(42) + b = model.forward(DOMAIN, data, generator=gen)[DOMAIN] + gen.manual_seed(43) + c = model.forward(DOMAIN, data, generator=gen)[DOMAIN] + assert torch.equal(a, b) + assert not torch.equal(a, c) + + +def test_on_validation_start_resets_the_pass_counter(): + algo = _make_algo() + algo.reset_eval_pass_counter() + algo._eval_pass_counter = 17 + wrapper = types.SimpleNamespace( + model=algo, + device=torch.device("cpu"), + _val_heads=lambda: {"valid": None}, + ) + ModelWrapper.on_validation_start(wrapper) + assert algo._eval_pass_counter == 0 + + +# --------------------------------------------------------------------------- +# 3. Log-only MAD gradient detector +# --------------------------------------------------------------------------- +class _GradHookStub: + """Just enough LightningModule for ModelWrapper.on_after_backward.""" + + def __init__(self, param: nn.Parameter, history): + self.enable_grad_norm = True + self.grad_norm_mad_scale = ModelWrapper.grad_norm_mad_scale + self.grad_norm_mad_min_count = ModelWrapper.grad_norm_mad_min_count + self.grad_norm_mad_window = ModelWrapper.grad_norm_mad_window + self.grad_norm_history = deque(history, maxlen=self.grad_norm_mad_window) + self._param = param + self.global_step = 4242 + self.trainer = types.SimpleNamespace(is_global_zero=True) + self.logged = {} + + def parameters(self): + return [self._param] + + def log(self, name, value, **kwargs): + self.logged[name] = value + + +def _spiking_stub(spike_norm: float = 50.0): + param = nn.Parameter(torch.zeros(4)) + grad = torch.full((4,), spike_norm / 2.0) # ||grad||_2 == spike_norm + param.grad = grad + history = [ + 0.9 if i % 2 else 1.1 for i in range(ModelWrapper.grad_norm_mad_min_count) + ] + return _GradHookStub(param, history), param + + +def test_on_after_backward_never_modifies_gradients(): + stub, param = _spiking_stub() + before = param.grad.clone() + ModelWrapper.on_after_backward(stub) + assert torch.equal(param.grad, before), "MAD detector rescaled the gradients" + + +def test_on_after_backward_still_flags_and_logs_the_spike(): + stub, _ = _spiking_stub() + ModelWrapper.on_after_backward(stub) + logged = stub.logged + assert logged["Train/policy_grad_norms_mad_flag"] == 1.0 + assert logged["Train/policy_grad_norms_raw"] == pytest.approx(50.0, rel=1e-5) + assert "Train/policy_grad_norms_mad_threshold" in logged + + +def test_flagged_steps_are_appended_to_the_history(): + """Log-only detector: every step goes into the rolling median/MAD window.""" + stub, _ = _spiking_stub() + n_before = len(stub.grad_norm_history) + ModelWrapper.on_after_backward(stub) + assert len(stub.grad_norm_history) == n_before + 1 + assert stub.grad_norm_history[-1] == pytest.approx(50.0, rel=1e-5) + + +# --------------------------------------------------------------------------- +# 4. Trainer config knob +# --------------------------------------------------------------------------- +def test_trainer_default_exposes_gradient_clip_val(compose_resolve): + cfg = compose_resolve("train_zarr_cartesian", []) + assert "gradient_clip_val" in cfg.trainer + assert cfg.trainer.gradient_clip_val is None # no clipping until chosen + assert cfg.trainer.gradient_clip_algorithm == "norm" + + +@pytest.mark.parametrize("trainer_cfg", ["default", "ddp", "ddp_pi", "debug"]) +def test_every_trainer_config_keeps_the_knob(trainer_cfg, compose_resolve): + cfg = compose_resolve("train_zarr_cartesian", [f"trainer={trainer_cfg}"]) + assert cfg.trainer.gradient_clip_val is None + assert cfg.trainer.gradient_clip_algorithm == "norm" + + +def test_gradient_clip_override_reaches_the_trainer(tmp_path, compose_resolve): + cfg = compose_resolve( + "train_zarr_cartesian", + [ + "trainer=default", + "trainer.gradient_clip_val=1.0", + "trainer.accelerator=cpu", + "trainer.precision=32", + "trainer.max_epochs=1", + "trainer.min_epochs=1", + f"paths.output_dir={tmp_path}", + ], + ) + trainer = hydra.utils.instantiate(cfg.trainer, logger=False, callbacks=None) + assert trainer.gradient_clip_val == 1.0 + assert trainer.gradient_clip_algorithm == "norm" + + +def test_lightning_module_does_not_override_gradient_clipping(): + """Nothing in ModelWrapper may intercept Lightning's clipping hook.""" + assert "configure_gradient_clipping" not in vars(ModelWrapper) From feec00367a042bce5ba1b39179f5f00da19097a3 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sun, 13 Sep 2026 14:21:23 -0400 Subject: [PATCH 14/79] configs: declare per-embodiment dims once per HPT model config stem input_dim, infer_ac_dims, act_dim and enc-dec output_dim now interpolate robomimic_model.dims.{.proprio,.action,action_width}. Resolved values are unchanged (pinned by tests/unit/test_model_dims.py). Pi configs untouched. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01MWtR1x5tjREwoxQZAnUGrh Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01BbQNNaMAbTNMypegmJkv4s From 84bb0dbaf679b4e27d2f34d5d9241fdacb2e1f4b Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 15 Sep 2026 15:05:55 -0400 Subject: [PATCH 15/79] augs: per-sample image augmentation for HPT train-mode augs torchvision draws one random parameter set per call, so applying the configured train_image_augs to the stacked (B, 3, H, W) camera batch gave every image in the batch the same jitter. Wrap train_image_augs in a new PerSampleAugs module so each image draws its own parameters, and move the inline aug block of _robomimic_to_hpt_data into _apply_image_augs with the same gating. Deterministic eval augs stay batched. _apply_image_augs also checks train_image_augs before calling it: __init__ accepts None and train mode would have raised TypeError. The loop is not free: measured on an A40 at batch 64 with the shipped ColorJitter + Normalize, 13.8 -> 92.2 ms at 224x224 and 62.2 -> 105.5 ms at the camera's 360x640, per camera per training step. Correctness first here; the cost is removed further up the stack, where the jitter itself is vectorised over the batch. Co-Authored-By: Claude Fable 5.1 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/algo/hpt.py | 23 +++-- egomimic/models/image_augs.py | 30 +++++++ tests/unit/test_per_sample_augs.py | 133 +++++++++++++++++++++++++++++ 3 files changed, 181 insertions(+), 5 deletions(-) create mode 100644 egomimic/models/image_augs.py create mode 100644 tests/unit/test_per_sample_augs.py diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index cbed1282a9..ab9225f962 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -15,6 +15,7 @@ from egomimic.algo.algo import Algo from egomimic.models.hpt_nets import MultiheadAttention, SimpleTransformer +from egomimic.models.image_augs import PerSampleAugs from egomimic.rldb.embodiment.embodiment import get_embodiment, get_embodiment_id from egomimic.utils.hf_utils import download_from_huggingface from egomimic.utils.tensor_utils import EinOpsRearrange, get_sinusoid_encoding_table @@ -857,7 +858,11 @@ def __init__( self.annotation_modality = annotation_modality self.default_prompt = default_prompt - self.train_image_augs = train_image_augs + # Per-sample so every image draws its own jitter parameters; torchvision + # samples once per call, i.e. once for the whole stacked batch. + self.train_image_augs = ( + PerSampleAugs(train_image_augs) if train_image_augs is not None else None + ) self.eval_image_augs = eval_image_augs self.stem_specs = stem_specs self.head_specs = head_specs @@ -1327,6 +1332,17 @@ def _forward_ot(self, batch, embodiment1_id, embodiment2_id): supervised=self.supervised, ) + def _apply_image_augs(self, images, short): + """ + helper method that augments one camera's (B, 3, H, W) batch; train augs + are per-sample, the deterministic eval augs stay batched + """ + if self.nets.training and self.train_image_augs and short in self.encoders: + return self.train_image_augs(images) + elif self.eval_image_augs and short in self.encoders: + return self.eval_image_augs(images) + return images + def _robomimic_to_hpt_data( self, batch, cam_keys, proprio_keys, lang_keys, ac_key, aux_ac_keys=[] ): @@ -1348,10 +1364,7 @@ def _robomimic_to_hpt_data( short = key.rsplit(".", 1)[-1] _data = batch[key] if not torch.all(_data == 0): - if self.nets.training and short in self.encoders: - _data = self.train_image_augs(_data) - elif self.eval_image_augs and short in self.encoders: - _data = self.eval_image_augs(_data) + _data = self._apply_image_augs(_data, short) data[short] = _data.unsqueeze(1).unsqueeze(1) diff --git a/egomimic/models/image_augs.py b/egomimic/models/image_augs.py new file mode 100644 index 0000000000..a4350875a7 --- /dev/null +++ b/egomimic/models/image_augs.py @@ -0,0 +1,30 @@ +"""Per-sample image augmentation. + +torchvision's random transforms draw one parameter set per call, so applying a +``ColorJitter`` to a stacked (B, 3, H, W) batch jitters every image in the batch +by the same factors (and in the same op order). ``PerSampleAugs`` wraps the +configured transform and calls it once per image, so every sample draws its own +parameters. +""" + +import torch +from torch import nn + + +class PerSampleAugs(nn.Module): + """Apply ``augs`` to each image of a (B, C, H, W) batch independently. + + ``augs`` is the configured callable (typically a + ``torchvision.transforms.Compose``); a deterministic one is unaffected by + the wrapping, so this is safe for any existing aug list. + """ + + def __init__(self, augs): + super().__init__() + self.augs = augs + + def forward(self, images): + return torch.stack([self.augs(image) for image in images]) + + def extra_repr(self): + return repr(self.augs) diff --git a/tests/unit/test_per_sample_augs.py b/tests/unit/test_per_sample_augs.py new file mode 100644 index 0000000000..9aa678de15 --- /dev/null +++ b/tests/unit/test_per_sample_augs.py @@ -0,0 +1,133 @@ +"""Per-sample image augmentation (ideas 0.10). + +torchvision's random transforms draw one parameter set per call, so the old +inline ``self.train_image_augs(_data)`` on the stacked (B, 3, H, W) batch gave +every image in the batch the same jitter. These tests pin the per-sample +behaviour of ``PerSampleAugs`` and of ``HPT._apply_image_augs``, and pin that +the deterministic eval path is unchanged. +""" + +from __future__ import annotations + +import torch +from torch import nn +from torchvision.transforms import ColorJitter, Compose, Normalize + +from egomimic.models.image_augs import PerSampleAugs + +MEAN = [0.485, 0.456, 0.406] +STD = [0.229, 0.224, 0.225] + + +def _identical_batch(n=64, mean=0.5): + """``n`` copies of one image whose per-image mean is exactly ``mean``.""" + return torch.full((n, 3, 16, 16), float(mean)) + + +def _per_image_means(x): + return x.flatten(1).mean(dim=1) + + +def test_per_sample_augs_jitters_each_image_independently(): + """A batch of identical images must come out with per-image brightness + spread; the same ColorJitter called on the whole batch must not.""" + torch.manual_seed(0) + jitter = ColorJitter(brightness=0.5) + images = _identical_batch() + + per_sample_std = _per_image_means(PerSampleAugs(jitter)(images)).std().item() + batched_std = _per_image_means(jitter(images)).std().item() + print( + f"per-image brightness std: per-sample={per_sample_std:.6f} " + f"batched={batched_std:.6f}" + ) + + # brightness=0.5 scales each image by U(0.5, 1.5), so the per-image means + # of a mean-0.5 batch are U(0.25, 0.75): std = 0.5 / sqrt(12) ~ 0.144. + # 0.01 is ~14x below that and ~0 for the batched call. + assert per_sample_std > 0.01 + assert batched_std < 1e-6 + + +def test_per_sample_augs_matches_batched_for_deterministic_transform(): + """A deterministic Compose through PerSampleAugs is exactly the batched + result, so wrapping a Normalize-only aug list is a no-op.""" + augs = Compose([Normalize(mean=MEAN, std=STD)]) + images = torch.rand(8, 3, 16, 16) + torch.testing.assert_close( + PerSampleAugs(augs)(images), augs(images), rtol=0, atol=0 + ) + + +def test_per_sample_augs_preserves_shape_dtype_device_and_grad(): + images = torch.rand(4, 3, 16, 16, dtype=torch.float32, requires_grad=True) + out = PerSampleAugs(Compose([Normalize(mean=MEAN, std=STD)]))(images) + + assert out.shape == images.shape + assert out.dtype == images.dtype + assert out.device == images.device + out.sum().backward() + assert images.grad is not None and torch.isfinite(images.grad).all() + + +def _hpt_stub(train_image_augs=None, eval_image_augs=None, training=True): + """A bare HPT carrying only what ``_apply_image_augs`` reads; constructing + a real HPT would build the trunk and stems.""" + from egomimic.algo.hpt import HPT + + algo = HPT.__new__(HPT) + algo.nets = nn.ModuleDict() + algo.nets.train(training) + algo.encoders = {"front_img_1": {}} + algo.train_image_augs = ( + PerSampleAugs(train_image_augs) if train_image_augs is not None else None + ) + algo.eval_image_augs = eval_image_augs + return algo + + +def test_apply_image_augs_is_per_sample_in_train_mode(): + torch.manual_seed(0) + algo = _hpt_stub(train_image_augs=ColorJitter(brightness=0.5), training=True) + out = algo._apply_image_augs(_identical_batch(), "front_img_1") + assert _per_image_means(out).std().item() > 0.01 + + +def test_apply_image_augs_eval_mode_equals_batched_normalize(): + eval_augs = Compose([Normalize(mean=MEAN, std=STD)]) + algo = _hpt_stub( + train_image_augs=ColorJitter(brightness=0.5), + eval_image_augs=eval_augs, + training=False, + ) + images = torch.rand(8, 3, 16, 16) + torch.testing.assert_close( + algo._apply_image_augs(images, "front_img_1"), eval_augs(images), rtol=0, atol=0 + ) + + +def test_apply_image_augs_passes_through_uncoded_camera(): + """A camera with no encoder is left alone in either mode.""" + images = torch.rand(4, 3, 16, 16) + for training in (True, False): + algo = _hpt_stub( + train_image_augs=ColorJitter(brightness=0.5), + eval_image_augs=Compose([Normalize(mean=MEAN, std=STD)]), + training=training, + ) + out = algo._apply_image_augs(images, "wrist_img_1") + torch.testing.assert_close(out, images, rtol=0, atol=0) + + +def test_missing_train_augs_pass_through(): + """__init__ allows train_image_augs=None; the helper must not call it.""" + from egomimic.algo.hpt import HPT + + algo = HPT.__new__(HPT) + algo.nets = torch.nn.ModuleDict({"policy": torch.nn.Identity()}) + algo.nets.train() + algo.train_image_augs = None + algo.eval_image_augs = None + algo.encoders = {"front_img_1": object()} + images = torch.rand(2, 3, 8, 8) + assert algo._apply_image_augs(images, "front_img_1") is images From cb1db1daa3b070d5da90c3e1acd2a473a929c5ee Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 15 Sep 2026 19:09:19 -0400 Subject: [PATCH 16/79] resnet: keep a frozen backbone in eval mode; give backbones their own LR group Two defects that would confound any backbone comparison: * a frozen `ResNet` backbone kept its BatchNorm layers in train mode, so it normalized with batch statistics and kept updating its running stats. `requires_grad = False` does not touch either. `train()` now puts a frozen backbone in eval mode. Latent today -- no shipped config freezes the ResNet -- so no run to date was affected. * the pretrained backbone and the randomly initialised trunk shared one LR. The new `model.backbone_lr_scale` key gives the encoders' backbones (`backbone_parameters()`) their own param group at `lr * scale`; at the default 1.0 the optimizer is built exactly as before, as a single group over `trainer.model.parameters()`. Nothing shipped sets it to anything else, so the two-group path is exercised only by its tests for now. Co-Authored-By: Claude Fable 5.1 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- .../model/hpt_bc_mecka_6d_300M.yaml | 4 + egomimic/models/hpt_nets.py | 28 +++ egomimic/pl_utils/pl_model.py | 54 ++++- tests/unit/test_vision_backbone.py | 199 ++++++++++++++++++ 4 files changed, 279 insertions(+), 6 deletions(-) create mode 100644 tests/unit/test_vision_backbone.py diff --git a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml index 8cf56b9a42..46f6886cf1 100644 --- a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml @@ -152,6 +152,10 @@ robomimic_model: mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225] +# LR multiplier for the image encoders' pretrained backbones; 1.0 keeps the +# single optimizer param group (ModelWrapper.configure_optimizers). +backbone_lr_scale: 1.0 + optimizer: _target_: torch.optim.AdamW _partial_: true diff --git a/egomimic/models/hpt_nets.py b/egomimic/models/hpt_nets.py index 3a6d4e0a3d..fee4c7367c 100644 --- a/egomimic/models/hpt_nets.py +++ b/egomimic/models/hpt_nets.py @@ -586,6 +586,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ResNet(PolicyStem): + # ``self.net`` holds the ImageNet-pretrained torchvision backbone (``self.proj`` + # is built here and must keep its own init). Declared so a weight-protection + # pass can find it without knowing about this class. + _hpt_pretrained = True + _hpt_pretrained_attrs = ("net",) + def __init__( self, output_dim: int = 10, @@ -617,6 +623,7 @@ def __init__( self.avgpool = nn.AvgPool2d(7, stride=1) # Freeze the backbone if specified + self.freeze_backbone = freeze_backbone if freeze_backbone: self._freeze_backbone() @@ -630,6 +637,27 @@ def _freeze_backbone(self): for param in self.net.parameters(): param.requires_grad = False + def train(self, mode: bool = True): + """Keep a frozen backbone in eval mode regardless of the outer flag. + + ``requires_grad = False`` alone leaves the BatchNorm layers in train + mode: they would keep normalizing with batch statistics and keep + updating their running stats, so a "frozen" backbone would not be + frozen at all. + """ + super().train(mode) + if self.freeze_backbone: + if isinstance(self.net, nn.ModuleList): + for net in self.net: + net.eval() + else: + self.net.eval() + return self + + def backbone_parameters(self): + """Pretrained-backbone parameters (``proj`` stays in the main group).""" + return list(self.net.parameters()) + def forward(self, x: torch.Tensor) -> torch.Tensor: """ Performs a forward pass of the model. diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index dc60674c83..47644a6d6d 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/egomimic/pl_utils/pl_model.py @@ -297,6 +297,33 @@ def on_validation_end(self): flush=True, ) + def _backbone_param_groups(self, lr: float, scale: float): + """Two param groups: the encoders' pretrained backbones at ``lr * + scale``, every other trainable parameter at ``lr``. + + A pretrained image backbone and a randomly initialised trunk do not + want the same learning rate, and a frozen-vs-tuned comparison under one + LR is confounded. Encoders that do not declare ``backbone_parameters`` + (and frozen parameters) are left to the main group. + """ + encoders = getattr(self.model.nets["policy"], "encoders", None) or {} + backbone, seen = [], set() + for encoder in encoders.values(): + if not hasattr(encoder, "backbone_parameters"): + continue + for param in encoder.backbone_parameters(): + if param.requires_grad and id(param) not in seen: + seen.add(id(param)) + backbone.append(param) + rest = [ + param + for param in self.trainer.model.parameters() + if param.requires_grad and id(param) not in seen + ] + groups = [{"params": backbone, "lr": lr * scale}, {"params": rest}] + # a fully frozen (or absent) backbone leaves the main group alone + return [group for group in groups if group["params"]] + def configure_optimizers(self) -> Dict[str, Any]: """Choose what optimizers and learning-rate schedulers to use in your optimization. Normally you'd need one. But in the case of GANs or similar you might have multiple. @@ -309,12 +336,27 @@ def configure_optimizers(self) -> Dict[str, Any]: config_tree = getattr(self.hparams, "config_tree", None) if config_tree is not None: cfg = self._as_config(config_tree) - optimizer = hydra.utils.instantiate( - cfg.model.optimizer, - params=self.trainer.model.parameters(), - ) - if callable(optimizer): - optimizer = optimizer() + scale = float(cfg.model.get("backbone_lr_scale", 1.0)) + groups = None + if scale != 1.0: + lr = cfg.model.optimizer.get("lr") + if lr is None: + raise ValueError( + "model.backbone_lr_scale needs an explicit model.optimizer.lr" + ) + groups = self._backbone_param_groups(float(lr), scale) + if groups is None: + optimizer = hydra.utils.instantiate( + cfg.model.optimizer, + params=self.trainer.model.parameters(), + ) + if callable(optimizer): + optimizer = optimizer() + else: + # ``instantiate`` would turn the param-group dicts into config + # nodes, so bind them outside it. + factory = hydra.utils.instantiate(cfg.model.optimizer, _partial_=True) + optimizer = factory(params=groups) scheduler_cfg = cfg.model.get("scheduler") if scheduler_cfg is not None: scheduler = hydra.utils.instantiate( diff --git a/tests/unit/test_vision_backbone.py b/tests/unit/test_vision_backbone.py new file mode 100644 index 0000000000..e8d643d888 --- /dev/null +++ b/tests/unit/test_vision_backbone.py @@ -0,0 +1,199 @@ +"""The frozen-backbone eval-mode fix and the separate backbone LR group. + +Covers ``ResNet`` (unchanged forward, BatchNorm in eval when frozen, +``backbone_parameters``), ``ModelWrapper.configure_optimizers`` param groups, +and the flagship model config. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +import torchvision +from omegaconf import OmegaConf +from torch import nn + +from egomimic.models.hpt_nets import ResNet +from egomimic.pl_utils.pl_model import ModelWrapper + +# -------------------------------------------------------------------------- +# ResNet +# -------------------------------------------------------------------------- + + +def _resnet_stem(**kwargs): + torch.manual_seed(0) + return ResNet(output_dim=16, **kwargs) + + +def test_resnet_forward_matches_reference_implementation(): + """The baseline image path is bit-identical: the stem's output equals the + torchvision backbone followed by the stem's projection.""" + stem = _resnet_stem().eval() + x = torch.rand(2, 1, 1, 3, 96, 128) + with torch.no_grad(): + out = stem(x) + + backbone = nn.Sequential( + *list(torchvision.models.resnet18(weights="DEFAULT").children())[:-2] + ).eval() + with torch.no_grad(): + feat = backbone(x.view(-1, 3, 96, 128)) + expected = stem.proj(feat.view(2, feat.shape[1], -1).transpose(1, 2)) + assert out.shape == expected.shape + assert torch.equal(out, expected) + + +def test_resnet_frozen_backbone_batchnorm_stays_in_eval(): + """0.15: a frozen ResNet must not update / use batch statistics.""" + stem = _resnet_stem(freeze_backbone=True) + stem.train() + assert stem.training + bns = [ + m for m in stem.net.modules() if isinstance(m, nn.modules.batchnorm._BatchNorm) + ] + assert bns + assert all(not m.training for m in bns) + assert all(not p.requires_grad for p in stem.net.parameters()) + # an unfrozen stem keeps the default behaviour + trainable = _resnet_stem() + trainable.train() + assert all( + m.training + for m in trainable.net.modules() + if isinstance(m, nn.modules.batchnorm._BatchNorm) + ) + + +def test_resnet_frozen_backbone_running_stats_do_not_drift(): + stem = _resnet_stem(freeze_backbone=True) + stem.train() + bn = next( + m for m in stem.net.modules() if isinstance(m, nn.modules.batchnorm._BatchNorm) + ) + before = bn.running_mean.clone() + with torch.no_grad(): + stem(torch.rand(2, 1, 1, 3, 96, 128) * 5.0) + assert torch.equal(bn.running_mean, before) + + +def test_resnet_backbone_parameters_are_the_net(): + stem = _resnet_stem() + backbone = {id(p) for p in stem.backbone_parameters()} + assert backbone == {id(p) for p in stem.net.parameters()} + assert all(id(p) not in backbone for p in stem.proj.parameters()) + + +def test_resnet_backbone_parameters_with_multiple_copies(): + stem = ResNet(output_dim=16, num_of_copy=2) + backbone = {id(p) for p in stem.backbone_parameters()} + assert backbone == {id(p) for p in stem.net.parameters()} + + +# -------------------------------------------------------------------------- +# configure_optimizers +# -------------------------------------------------------------------------- + + +class _StubEncoder(nn.Module): + """Encoder with a pretrained-style backbone plus a trainable projection.""" + + def __init__(self, frozen: bool = False): + super().__init__() + self.net = nn.Linear(4, 4) + self.proj = nn.Linear(4, 4) + if frozen: + for p in self.net.parameters(): + p.requires_grad = False + + def backbone_parameters(self): + return list(self.net.parameters()) + + +def _stub_wrapper(backbone_lr_scale=None, frozen=False): + policy = nn.Module() + policy.encoders = nn.ModuleDict( + { + "front_img_1": _StubEncoder(frozen=frozen), + "other": nn.Linear(4, 4), # no backbone_parameters -> main group + } + ) + policy.trunk = nn.Linear(4, 4) + root = nn.Module() + root.nets = nn.ModuleDict({"policy": policy}) + + model_cfg = { + "optimizer": { + "_target_": "torch.optim.AdamW", + "_partial_": True, + "lr": 5e-5, + "weight_decay": 0.0001, + }, + "scheduler": None, + } + if backbone_lr_scale is not None: + model_cfg["backbone_lr_scale"] = backbone_lr_scale + config_tree = OmegaConf.create({"model": model_cfg}) + # A LightningModule cannot be built without a trainer / datamodule here, so + # call the real methods against the attributes they actually read. + stub = SimpleNamespace( + hparams=SimpleNamespace(config_tree=config_tree), + _as_config=ModelWrapper._as_config, + model=root, + trainer=SimpleNamespace(model=root), + ) + stub._backbone_param_groups = ModelWrapper._backbone_param_groups.__get__(stub) + return stub, root + + +@pytest.mark.parametrize("scale", [None, 1.0]) +def test_configure_optimizers_single_group_by_default(scale): + stub, _ = _stub_wrapper(backbone_lr_scale=scale) + out = ModelWrapper.configure_optimizers(stub) + groups = out["optimizer"].param_groups + assert len(groups) == 1 + assert groups[0]["lr"] == pytest.approx(5e-5) + + +def test_configure_optimizers_backbone_group(): + stub, root = _stub_wrapper(backbone_lr_scale=0.1) + out = ModelWrapper.configure_optimizers(stub) + groups = out["optimizer"].param_groups + assert len(groups) == 2 + by_lr = {round(g["lr"], 12): g for g in groups} + assert set(by_lr) == {round(5e-5, 12), round(5e-6, 12)} + + encoder = root.nets["policy"].encoders["front_img_1"] + backbone_ids = {id(p) for p in encoder.backbone_parameters()} + got_backbone = {id(p) for p in by_lr[round(5e-6, 12)]["params"]} + main = {id(p) for p in by_lr[round(5e-5, 12)]["params"]} + assert got_backbone == backbone_ids + assert not (got_backbone & main) + trainable = {id(p) for p in root.parameters() if p.requires_grad} + assert got_backbone | main == trainable + + +def test_configure_optimizers_skips_frozen_backbone_params(): + stub, root = _stub_wrapper(backbone_lr_scale=0.1, frozen=True) + out = ModelWrapper.configure_optimizers(stub) + every = [p for g in out["optimizer"].param_groups for p in g["params"]] + assert all(p.requires_grad for p in every) + assert {id(p) for p in every} == { + id(p) for p in root.parameters() if p.requires_grad + } + + +# -------------------------------------------------------------------------- +# Configs +# -------------------------------------------------------------------------- + +FLAGSHIP = "train_zarr_mecka_flagship_6d_hpt" + + +def test_base_model_config_declares_unit_backbone_lr_scale(compose_resolve): + cfg = compose_resolve(FLAGSHIP, []) + assert cfg.model.backbone_lr_scale == 1.0 + encoder = cfg.model.robomimic_model.encoder_specs.front_img_1 + assert encoder._target_ == "egomimic.models.hpt_nets.ResNet" From 133be3fb6f8540ca61ff76e79f11ccded21778c0 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 15 Sep 2026 14:50:07 -0400 Subject: [PATCH 17/79] loss: mask the repeat-last action padding out of the HPT loss ZarrDataset emits a real action_pad_mask (1.0 per real frame, 0.0 per repeat-last pad, honouring the annotation-cutoff EOS clamp) next to every horizoned action chunk. A new InterpolatePadMask transform carries it through the 30 -> 100 upsampling on the interpolators' own grid: an output step is real only if both source frames it interpolates between are real. HPT gains a use_pad_mask algo key (default false, declared in the two flagship model yamls; every other HPT config takes the __init__ default). With it on, process_batch_for_training builds pad_mask (B, S, 1) from the dataset mask instead of ones; DenoisingPolicy.loss_fn takes the mask and means the squared error over the valid steps only. With it off nothing changes: the mask is dropped from the batch and pad_mask stays all ones, which is exactly F.mse_loss. It ships off to match openpi, which asks LeRobotDataset for action_is_pad and then ignores it (Pi0.compute_loss means over the whole horizon). LeRobot's own policies do mask, but reduce with .mean() over all elements including the zeroed ones, so their loss shrinks as padding grows; this divides by the valid count instead, which keeps the gradient magnitude constant. Measured on the three pinned flagship episodes, 6.4-7.6 % of supervised action steps are repeat-last padding and 12.8-15.2 % of samples contain some, so turning it on changes the loss: a run before and a run after are not comparable, and the layout-aware eval metrics would have to be masked in the same change or train and val would be measuring different things. Padding is still repeat-last and no tail samples are dropped. Co-Authored-By: Claude Fable 5.1 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/algo/hpt.py | 25 +- .../model/hpt_bc_keypoints_wrist_300M.yaml | 4 + .../model/hpt_bc_mecka_6d_300M.yaml | 4 + egomimic/models/denoising_policy.py | 19 +- egomimic/rldb/embodiment/eva.py | 3 + egomimic/rldb/embodiment/human.py | 4 + egomimic/rldb/zarr/action_chunk_transforms.py | 47 +++ egomimic/rldb/zarr/zarr_dataset_multi.py | 33 ++ tests/unit/test_action_chunk_transforms.py | 20 +- tests/unit/test_loss_masking.py | 394 ++++++++++++++++++ 10 files changed, 539 insertions(+), 14 deletions(-) create mode 100644 tests/unit/test_loss_masking.py diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index ab9225f962..c152e5976d 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -847,6 +847,12 @@ def __init__( annotation_modality: str = "annotation", default_prompt: str = "", # --------------------------- + # Loss masking: supervise only the real action steps of a chunk + # (the dataset's ``action_pad_mask``) instead of the repeat-last + # padding at every episode / annotation tail. + # --------------------------- + use_pad_mask: bool = False, + # --------------------------- # Catch-all kwargs # --------------------------- **kwargs, @@ -857,6 +863,7 @@ def __init__( self.annotation_sampling_mode = annotation_sampling_mode self.annotation_modality = annotation_modality self.default_prompt = default_prompt + self.use_pad_mask = use_pad_mask # Per-sample so every image draws its own jitter parameters; torchvision # samples once per call, i.e. once for the whole stacked batch. @@ -1102,9 +1109,23 @@ def process_batch_for_training(self, batch): B, S, _ = processed_batch[embodiment_id][ac_key].shape device = processed_batch[embodiment_id][ac_key].device - processed_batch[embodiment_id]["pad_mask"] = torch.ones( - B, S, 1, device=device + # Dropped either way: downstream sees the same keys as before. + action_pad_mask = processed_batch[embodiment_id].pop( + "action_pad_mask", None ) + if self.use_pad_mask and action_pad_mask is not None: + if action_pad_mask.shape[-1] != S: + raise ValueError( + f"action_pad_mask length {action_pad_mask.shape[-1]} does " + f"not match the action horizon {S}" + ) + processed_batch[embodiment_id]["pad_mask"] = action_pad_mask[ + ..., None + ].to(device) + else: + processed_batch[embodiment_id]["pad_mask"] = torch.ones( + B, S, 1, device=device + ) # Sample one annotation per item (random/first, default fallback for # empty). Stays as list[str]; the Qwen stem owns tokenization. diff --git a/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml index d55bf1afc9..77f2f4f79c 100644 --- a/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml @@ -40,6 +40,10 @@ robomimic_model: annotation_modality: "annotation" default_prompt: "" + # Mask the repeat-last padding at every episode / annotation tail out of + # the action loss (the dataset's `action_pad_mask`). + use_pad_mask: false + trunk: embed_dim: 840 num_blocks: 23 diff --git a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml index 46f6886cf1..66ca79a60a 100644 --- a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml @@ -42,6 +42,10 @@ robomimic_model: annotation_modality: "annotation" default_prompt: "" + # Mask the repeat-last padding at every episode / annotation tail out of + # the action loss (the dataset's `action_pad_mask`). + use_pad_mask: false + trunk: embed_dim: 840 num_blocks: 23 diff --git a/egomimic/models/denoising_policy.py b/egomimic/models/denoising_policy.py index b9180e0d37..29f9a81d48 100644 --- a/egomimic/models/denoising_policy.py +++ b/egomimic/models/denoising_policy.py @@ -97,11 +97,24 @@ def predict(self, actions, global_cond) -> Tuple[torch.Tensor, torch.Tensor]: """ raise NotImplementedError - def loss_fn(self, pred, target): + def loss_fn(self, pred, target, pad_mask=None): """ Computes loss, function to override for stuff like adaptive loss weighting + + ``pad_mask`` (B, S, 1) or (B, S) marks the real action steps; padded + tail steps (0.0) drop out of the mean so the model is not supervised on + the dataset's repeat-last padding. ``None`` (or an all-ones mask) is + plain ``F.mse_loss``. """ - return F.mse_loss(pred, target) + if pad_mask is None: + return F.mse_loss(pred, target) + mask = pad_mask.to(dtype=pred.dtype, device=pred.device) + if mask.dim() == pred.dim() - 1: + mask = mask.unsqueeze(-1) + mask = mask.expand_as(pred) + # At least one element: an all-padding chunk contributes 0, not NaN. + denom = mask.sum().clamp(min=1.0) + return (((pred - target) ** 2) * mask).sum() / denom def preprocess_compute_loss(self, global_cond, data): if self.pooling == "mean": @@ -131,4 +144,4 @@ def preprocess_compute_loss(self, global_cond, data): def compute_loss(self, global_cond, data): actions, global_cond = self.preprocess_compute_loss(global_cond, data) pred, target = self.predict(actions, global_cond) - return self.loss_fn(pred, target) + return self.loss_fn(pred, target, pad_mask=data.get("pad_mask")) diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index 6067ad16fc..b3edfa4f48 100644 --- a/egomimic/rldb/embodiment/eva.py +++ b/egomimic/rldb/embodiment/eva.py @@ -18,6 +18,7 @@ ConcatKeys, DeleteKeys, InterpolateLinear, + InterpolatePadMask, InterpolatePose, NumpyToTensor, PoseCoordinateFrameTransform, @@ -396,6 +397,7 @@ def _build_eva_bimanual_eef_frame_transform_list( output_action_key=right_cmd_gripper, stride=stride, ), + InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), # Step 2: transform camera-frame actions into EEF-relative (wrist) frame ActionChunkCoordinateFrameTransform( target_world=left_obs_camframe, @@ -559,6 +561,7 @@ def _build_eva_bimanual_transform_list( output_action_key=right_cmd_gripper, stride=stride, ), + InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), ] if is_quat: diff --git a/egomimic/rldb/embodiment/human.py b/egomimic/rldb/embodiment/human.py index d30e701b43..d8e69ef3e6 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -16,6 +16,7 @@ CartesianYPRToRot6D, ConcatKeys, DeleteKeys, + InterpolatePadMask, InterpolatePose, KeypointsRot6DToYPR, KeypointsYPRToRot6D, @@ -996,6 +997,7 @@ def _build_human_keypoints_bimanual_transform_list( stride=stride, mode="xyzwxyz", ), + InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), ] if not is_quat: transform_list.extend( @@ -1269,6 +1271,7 @@ def _build_human_cartesian_eef_frame_transform_list( stride=stride, mode="xyzwxyz", ), + InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), ActionChunkCoordinateFrameTransform( target_world=left_obs_headframe, chunk_world=left_action_headframe, @@ -1382,6 +1385,7 @@ def _build_human_cartesian_bimanual_transform_list( stride=stride, mode="xyzwxyz", ), + InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), ] if target_world_is_quat: diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/zarr/action_chunk_transforms.py index 6a1352afbe..2be8127e66 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -140,6 +140,53 @@ def transform(self, batch: dict) -> dict: return batch +class InterpolatePadMask(Transform): + """Carry the dataset's action pad mask through the action interpolation. + + The interpolators resample a (T, D) chunk from ``np.linspace(0, 1, T)`` + onto ``np.linspace(0, 1, new_chunk_length)``; an output step is real only + if BOTH source frames it interpolates between are real (a step landing + exactly on a source frame inherits that frame). No-op when the mask key is + absent, and idempotent, so a mode that chains two builders + (``include_ee_pose``) still upsamples the mask once. + """ + + def __init__( + self, + new_chunk_length: int, + mask_key: str = "action_pad_mask", + stride: int = 1, + ): + if stride <= 0: + raise ValueError(f"stride must be positive, got {stride}") + self.new_chunk_length = new_chunk_length + self.mask_key = mask_key + self.stride = int(stride) + + def transform(self, batch: dict) -> dict: + mask = batch.get(self.mask_key) + if mask is None: + return batch + mask = np.asarray(mask, dtype=np.float32).reshape(-1) + if mask.shape[0] == self.new_chunk_length: + return batch # already at the model horizon + mask = mask[:: self.stride] + T = mask.shape[0] + if T < 2: + batch[self.mask_key] = np.full( + self.new_chunk_length, mask[0] if T else 0.0, dtype=np.float32 + ) + return batch + # Position of each output sample in source-frame index units. + pos = np.linspace(0, 1, self.new_chunk_length) * (T - 1) + nearest = np.rint(pos) + on_frame = np.abs(pos - nearest) < 1e-9 + lo = np.where(on_frame, nearest, np.floor(pos)).astype(int) + hi = np.where(on_frame, nearest, np.ceil(pos)).astype(int) + batch[self.mask_key] = ((mask[lo] > 0) & (mask[hi] > 0)).astype(np.float32) + return batch + + # --------------------------------------------------------------------------- # Coordinate Transforms # --------------------------------------------------------------------------- diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 70f5e7f9c6..89bd9733d2 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -2136,6 +2136,33 @@ def _chunk_end_idx(self, start_idx: int, horizon: int, key_type: str | None) -> """ return min(start_idx + horizon, self.total_frames) + def _action_horizon(self) -> int | None: + """Horizon of the keymap's ``action_keys`` entries, if it has any.""" + for spec in self.key_map.values(): + if ( + spec.get("key_type") == "action_keys" + and spec.get("horizon") is not None + ): + return int(spec["horizon"]) + return None + + def _action_pad_mask(self, idx: int) -> np.ndarray | None: + """1.0 per real action frame, 0.0 per repeat-last padded one, (horizon,). + + ``_pad_sequences`` pads a short action chunk by repeating the last + frame, which in a wrist-relative action space supervises "stop moving" + at every episode (or annotation) tail. This mask says which steps are + real so ``loss_fn`` can drop the rest; ``None`` when the keymap has no + horizoned action key. + """ + horizon = self._action_horizon() + if horizon is None: + return None + n_real = self._chunk_end_idx(idx, horizon, "action_keys") - idx + mask = np.zeros(horizon, dtype=np.float32) + mask[: max(n_real, 0)] = 1.0 + return mask + def _pad_sequences(self, data, horizon: int | None) -> dict: if horizon is None: return data @@ -2224,6 +2251,12 @@ def _next(reason: str, key: str = "") -> int: if retry: continue + pad_mask = self._action_pad_mask(idx) + if pad_mask is not None: + # Travels with the chunk through the interpolators + # (InterpolatePadMask) to the model horizon. + data["action_pad_mask"] = pad_mask + if self.transform: for transform in self.transform or []: data = transform.transform(data) diff --git a/tests/unit/test_action_chunk_transforms.py b/tests/unit/test_action_chunk_transforms.py index a9b9f9141c..bafe6b4357 100644 --- a/tests/unit/test_action_chunk_transforms.py +++ b/tests/unit/test_action_chunk_transforms.py @@ -344,6 +344,7 @@ def test_eva_transform_list_stepwise_keys_and_shapes() -> None: "InterpolatePose", "InterpolateLinear", "InterpolateLinear", + "InterpolatePadMask", "XYZWXYZ_to_XYZYPR", "ConcatKeys", "ConcatKeys", @@ -418,7 +419,7 @@ def test_eva_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 8, + 9, "XYZWXYZ_to_XYZYPR", base_keys | {"left.cmd_ee_pose_camframe", "right.cmd_ee_pose_camframe"}, { @@ -430,7 +431,7 @@ def test_eva_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 9, + 10, "ConcatKeys", { "actions_cartesian", @@ -451,7 +452,7 @@ def test_eva_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 10, + 11, "ConcatKeys", { "actions_cartesian", @@ -468,7 +469,7 @@ def test_eva_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 11, + 12, "DeleteKeys", {"actions_cartesian", "observations.state.ee_pose"}, { @@ -478,7 +479,7 @@ def test_eva_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 12, + 13, "NumpyToTensor", {"actions_cartesian", "observations.state.ee_pose"}, { @@ -516,6 +517,7 @@ def test_human_transform_list_stepwise_keys_and_shapes() -> None: "PoseCoordinateFrameTransform", "InterpolatePose", "InterpolatePose", + "InterpolatePadMask", "XYZWXYZ_to_XYZYPR", "ConcatKeys", "ConcatKeys", @@ -585,7 +587,7 @@ def test_human_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 6, + 7, "XYZWXYZ_to_XYZYPR", base_keys | { @@ -603,7 +605,7 @@ def test_human_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 7, + 8, "ConcatKeys", base_keys | { @@ -619,7 +621,7 @@ def test_human_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 8, + 9, "ConcatKeys", base_keys | {"actions_cartesian", "observations.state.ee_pose"}, { @@ -629,7 +631,7 @@ def test_human_transform_list_stepwise_keys_and_shapes() -> None: ) _assert_snapshot( snapshots, - 9, + 10, "DeleteKeys", {"actions_cartesian", "observations.state.ee_pose"}, { diff --git a/tests/unit/test_loss_masking.py b/tests/unit/test_loss_masking.py new file mode 100644 index 0000000000..89b0a63cda --- /dev/null +++ b/tests/unit/test_loss_masking.py @@ -0,0 +1,394 @@ +"""Episode-tail padding must not supervise "stop moving". + +The dataset emits a real ``action_pad_mask`` next to the action chunk, the +interpolators carry it to the model horizon, and ``DenoisingPolicy.loss_fn`` +drops the padded steps -- but only behind the ``use_pad_mask`` algo flag, so +the flag-off path is bit-for-bit today's behaviour. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +from fixtures.synthetic_episodes import write_episode +from omegaconf import OmegaConf + +import egomimic +from egomimic.algo.hpt import HPT +from egomimic.models.denoising_policy import DenoisingPolicy +from egomimic.rldb.embodiment.embodiment import get_embodiment_id +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.zarr.zarr_dataset_multi import ( + LocalAnnotationCutoffEpisodeResolver, + LocalEpisodeResolver, + MultiDataset, +) + +T_FRAMES = 48 # write_episode default +HORIZON = Human.ACTION_HORIZON # 30 raw action frames +CHUNK = 100 # model horizon after interpolation + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _pad_mask_transform_cls(): + """Imported inside the tests so the module still collects (and every test + fails on its own assertion) before the transform exists.""" + from egomimic.rldb.zarr.action_chunk_transforms import InterpolatePadMask + + return InterpolatePadMask + + +def _leaf(ds: MultiDataset): + """The single per-episode ZarrDataset behind a one-episode MultiDataset.""" + return next(iter(ds.datasets.values())) + + +def _mecka_dataset(tmp_path, *, transforms, cutoff: bool = False, n: int = 1): + for i in range(n): + write_episode(tmp_path, "mecka", seed=i) + # No annotation_key: the cutoff reads the zarr's spans directly, and the + # default collate cannot stack the per-sample annotation string lists. + key_map = Human.get_keymap(keymap_mode="cartesian") + resolver_cls = ( + LocalAnnotationCutoffEpisodeResolver if cutoff else LocalEpisodeResolver + ) + resolver = resolver_cls(tmp_path, key_map=key_map, transform_list=transforms) + return MultiDataset._from_resolver(resolver, mode="total") + + +def _expected_upsampled_mask(mask: np.ndarray, new_len: int) -> np.ndarray: + """Independent restatement of the rule: an output step is valid only if the + two source frames it interpolates between are both real; a step landing + exactly on a source frame inherits that frame.""" + T = mask.shape[0] + old = np.linspace(0, 1, T) + new = np.linspace(0, 1, new_len) + out = np.empty(new_len, dtype=np.float32) + for j, t in enumerate(new): + i = int(np.searchsorted(old, t)) + if i < T and np.isclose(old[i], t): + out[j] = mask[i] + else: + out[j] = mask[i - 1] * mask[i] + return out + + +# --------------------------------------------------------------------------- +# 1. dataset emits the real mask +# --------------------------------------------------------------------------- + + +def test_dataset_emits_action_pad_mask_for_a_tail_chunk(tmp_path): + ds = _mecka_dataset(tmp_path, transforms=[]) + leaf = _leaf(ds) + assert leaf.total_frames == T_FRAMES + + sample = leaf[T_FRAMES - 5] + mask = np.asarray(sample["action_pad_mask"]) + assert mask.shape == (HORIZON,) + expected = np.zeros(HORIZON, dtype=np.float32) + expected[:5] = 1.0 + np.testing.assert_array_equal(mask, expected) + + head = np.asarray(leaf[0]["action_pad_mask"]) + np.testing.assert_array_equal(head, np.ones(HORIZON, dtype=np.float32)) + + +def test_action_pad_mask_is_a_float32_tensor(tmp_path): + ds = _mecka_dataset(tmp_path, transforms=[]) + value = _leaf(ds)[0]["action_pad_mask"] + assert isinstance(value, torch.Tensor) + assert value.dtype == torch.float32 + + +def test_no_action_pad_mask_without_an_action_horizon(tmp_path): + """Keymaps with no horizoned ``action_keys`` entry must not grow a key.""" + write_episode(tmp_path, "mecka", seed=0) + key_map = Human.get_keymap(keymap_mode="cartesian") + for spec in key_map.values(): + if spec.get("key_type") == "action_keys": + spec["key_type"] = "proprio_keys" + spec.pop("horizon", None) + resolver = LocalEpisodeResolver(tmp_path, key_map=key_map, transform_list=[]) + ds = MultiDataset._from_resolver(resolver, mode="total") + assert "action_pad_mask" not in _leaf(ds)[0] + + +def test_mask_survives_bounds_check_and_normalization(tmp_path): + """``MultiDataset.__getitem__`` runs _check_bounds + normalize; the mask is + not in the norm stats and must come through untouched.""" + transforms = Human.get_transform_list(mode="cartesian_6d", stride=1) + ds = _mecka_dataset(tmp_path, transforms=transforms, n=3) + + stats = MultiDataset(state={}, norm_mode="quantile") + stats.populate_from_datasets({"human_bimanual": ds}) + stats.infer_shapes_from_batch(ds[0]) + stats.infer_norm_from_dataset(ds, "human_bimanual", sample_frac=1.0, num_workers=0) + emb = get_embodiment_id("human_bimanual") + assert "action_pad_mask" not in stats.norm_stats[emb] + + raw = _leaf(ds)[0]["action_pad_mask"].clone() + ds.set_norm_stats_from(stats) + normalized = ds[0] + assert "action_pad_mask" in normalized + assert normalized["action_pad_mask"].shape == (CHUNK,) + torch.testing.assert_close(normalized["action_pad_mask"], raw) + + +# --------------------------------------------------------------------------- +# 2. annotation cutoff +# --------------------------------------------------------------------------- + + +def test_annotation_cutoff_is_padding(tmp_path): + """Synthetic episodes carry spans [0, T//2) and [T//2, T-1); a chunk that + starts inside the first span is real only up to that span's end.""" + ds = _mecka_dataset(tmp_path, transforms=[], cutoff=True) + leaf = _leaf(ds) + ann_end = T_FRAMES // 2 + idx = ann_end - 4 + assert idx + HORIZON > ann_end # the cutoff, not the episode end, bites + + mask = np.asarray(leaf[idx]["action_pad_mask"]) + expected = np.zeros(HORIZON, dtype=np.float32) + expected[: ann_end - idx] = 1.0 + np.testing.assert_array_equal(mask, expected) + + +# --------------------------------------------------------------------------- +# 3. InterpolatePadMask +# --------------------------------------------------------------------------- + + +def test_interpolate_pad_mask_30_to_100_with_five_real_frames(): + mask = np.zeros(HORIZON, dtype=np.float32) + mask[:5] = 1.0 + batch = {"action_pad_mask": mask.copy()} + out = _pad_mask_transform_cls()(new_chunk_length=CHUNK).transform(batch)[ + "action_pad_mask" + ] + + expected = _expected_upsampled_mask(mask, CHUNK) + assert out.shape == (CHUNK,) + assert out.dtype == np.float32 + np.testing.assert_array_equal(out, expected) + assert out.sum() == expected.sum() + # The tail is padding and the head is real: the mask is a prefix of ones. + assert out[0] == 1.0 and out[-1] == 0.0 + assert np.all(np.diff(out) <= 0) + + +def test_interpolate_pad_mask_all_real_stays_all_ones(): + batch = {"action_pad_mask": np.ones(HORIZON, dtype=np.float32)} + out = _pad_mask_transform_cls()(new_chunk_length=CHUNK).transform(batch)[ + "action_pad_mask" + ] + np.testing.assert_array_equal(out, np.ones(CHUNK, dtype=np.float32)) + + +def test_interpolate_pad_mask_honours_stride(): + mask = np.zeros(HORIZON, dtype=np.float32) + mask[:16] = 1.0 + out = _pad_mask_transform_cls()(new_chunk_length=CHUNK, stride=3).transform( + {"action_pad_mask": mask.copy()} + )["action_pad_mask"] + np.testing.assert_array_equal(out, _expected_upsampled_mask(mask[::3], CHUNK)) + + +def test_interpolate_pad_mask_is_a_noop_without_the_key(): + batch = {"actions_cartesian": np.zeros((3, 4))} + assert _pad_mask_transform_cls()(new_chunk_length=CHUNK).transform(batch) == batch + + +def test_interpolate_pad_mask_is_idempotent(): + """``include_ee_pose`` chains two builders that both carry the mask.""" + mask = np.zeros(HORIZON, dtype=np.float32) + mask[:5] = 1.0 + t = _pad_mask_transform_cls()(new_chunk_length=CHUNK, stride=3) + once = t.transform({"action_pad_mask": mask.copy()}) + twice = t.transform(dict(once)) + np.testing.assert_array_equal(twice["action_pad_mask"], once["action_pad_mask"]) + + +def test_human_transform_lists_upsample_the_mask(): + modes = [ + "cartesian_6d", + "cartesian_wristframe_6d", + "keypoints_headframe_6d", + "keypoints_wristframe_6d", + ] + for mode in modes: + tl = Human.get_transform_list(mode=mode, stride=3) + assert any( + isinstance(t, _pad_mask_transform_cls()) for t in tl + ), f"{mode} does not carry the pad mask" + + +def test_keypoints_with_ee_pose_upsamples_the_mask_once(tmp_path): + """Two builders chain here; the mask must still come out (100,).""" + tl = Human.get_transform_list( + mode="keypoints_wristframe_6d", stride=1, include_ee_pose=True + ) + mask = np.zeros(HORIZON, dtype=np.float32) + mask[:5] = 1.0 + batch = {"action_pad_mask": mask.copy()} + for t in tl: + if isinstance(t, _pad_mask_transform_cls()): + batch = t.transform(batch) + assert batch["action_pad_mask"].shape == (CHUNK,) + np.testing.assert_array_equal( + batch["action_pad_mask"], _expected_upsampled_mask(mask, CHUNK) + ) + + +# --------------------------------------------------------------------------- +# 4. loss_fn +# --------------------------------------------------------------------------- + + +def _policy() -> DenoisingPolicy: + return DenoisingPolicy.__new__(DenoisingPolicy) + + +def test_loss_fn_without_mask_matches_mse_loss(): + g = torch.Generator().manual_seed(0) + pred = torch.randn(3, 7, 5, generator=g) + target = torch.randn(3, 7, 5, generator=g) + p = _policy() + torch.testing.assert_close(p.loss_fn(pred, target), F.mse_loss(pred, target)) + ones = torch.ones(3, 7, 1) + torch.testing.assert_close( + p.loss_fn(pred, target, pad_mask=ones), F.mse_loss(pred, target) + ) + + +def test_loss_fn_ignores_masked_steps(): + g = torch.Generator().manual_seed(1) + pred = torch.randn(2, 6, 4, generator=g) + target = torch.randn(2, 6, 4, generator=g) + mask = torch.ones(2, 6, 1) + mask[:, 4:] = 0.0 + p = _policy() + before = p.loss_fn(pred, target, pad_mask=mask) + + perturbed = pred.clone() + perturbed[:, 4:] += 100.0 + torch.testing.assert_close(p.loss_fn(perturbed, target, pad_mask=mask), before) + + # ... and it is exactly the mse over the kept steps. + torch.testing.assert_close(before, F.mse_loss(pred[:, :4], target[:, :4])) + + +def test_loss_fn_all_zero_mask_does_not_divide_by_zero(): + p = _policy() + loss = p.loss_fn( + torch.randn(2, 6, 4), torch.randn(2, 6, 4), pad_mask=torch.zeros(2, 6, 1) + ) + assert torch.isfinite(loss) and float(loss) == 0.0 + + +class _FixedPredPolicy(DenoisingPolicy): + """Minimal DenoisingPolicy: ``predict`` returns a canned prediction.""" + + def __init__(self, action_horizon: int, pred: torch.Tensor): + nn.Module.__init__(self) + self.action_horizon = action_horizon + self.pooling = None + self.padding = None + self._pred = pred + + def predict(self, actions, global_cond): + return self._pred, actions + + +def test_compute_loss_forwards_the_pad_mask(): + g = torch.Generator().manual_seed(2) + target = torch.randn(2, 6, 4, generator=g) + pred = torch.randn(2, 6, 4, generator=g) + policy = _FixedPredPolicy(6, pred) + mask = torch.ones(2, 6, 1) + mask[:, 4:] = 0.0 + global_cond = torch.zeros(2, 3) + loss = policy.compute_loss(global_cond, {"action": target, "pad_mask": mask}) + torch.testing.assert_close(loss, F.mse_loss(pred[:, :4], target[:, :4])) + + +# --------------------------------------------------------------------------- +# 5. process_batch_for_training + the config flag +# --------------------------------------------------------------------------- + + +class _NoRenameNormStats: + def zarr_key_to_keyname(self, zarr_key, embodiment_id): + return None + + +def _process(use_pad_mask: bool, action_pad_mask: torch.Tensor | None, S: int = 6): + algo = HPT.__new__(HPT) + algo.norm_stats = _NoRenameNormStats() + algo.device = "cpu" + algo.annotation_key = None + algo.use_pad_mask = use_pad_mask + emb = "human_bimanual" + emb_id = get_embodiment_id(emb) + algo.ac_keys = {emb_id: "actions_cartesian"} + inner = {"actions_cartesian": torch.zeros(2, S, 18)} + if action_pad_mask is not None: + inner["action_pad_mask"] = action_pad_mask + return HPT.process_batch_for_training(algo, {emb: inner})[emb_id] + + +def test_process_batch_keeps_ones_when_the_flag_is_off(): + mask = torch.ones(2, 6) + mask[:, 3:] = 0.0 + out = _process(False, mask) + torch.testing.assert_close(out["pad_mask"], torch.ones(2, 6, 1)) + assert "action_pad_mask" not in out + + +def test_process_batch_uses_the_dataset_mask_when_the_flag_is_on(): + mask = torch.ones(2, 6) + mask[:, 3:] = 0.0 + out = _process(True, mask) + assert out["pad_mask"].shape == (2, 6, 1) + torch.testing.assert_close(out["pad_mask"], mask[..., None]) + assert "action_pad_mask" not in out + + +def test_process_batch_falls_back_to_ones_without_the_dataset_key(): + out = _process(True, None) + torch.testing.assert_close(out["pad_mask"], torch.ones(2, 6, 1)) + + +def test_process_batch_raises_on_a_length_mismatch(): + with pytest.raises(ValueError, match="6"): + _process(True, torch.ones(2, 5)) + + +def test_use_pad_mask_defaults_to_false(): + import inspect + + assert inspect.signature(HPT.__init__).parameters["use_pad_mask"].default is False + + +# The two flagship recipes: the ones a run would flip the knob on. Every other +# HPT config takes the HPT.__init__ default, which the test above pins. +@pytest.mark.parametrize( + "yaml_name", + ["hpt_bc_mecka_6d_300M", "hpt_bc_keypoints_wrist_300M"], +) +def test_algo_yaml_declares_use_pad_mask_off(yaml_name): + from pathlib import Path + + path = ( + Path(egomimic.__file__).parent / "hydra_configs" / "model" / f"{yaml_name}.yaml" + ) + cfg = OmegaConf.load(path) + assert cfg.robomimic_model.use_pad_mask is False From dbdd171dd33720fea5524785950c5174ac0e2015 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 15 Sep 2026 15:44:48 -0400 Subject: [PATCH 18/79] obs: proprio history (K proprio steps, one token each, learned time embedding) The proprio keys the model consumes read K steps instead of one. Every change is gated on K > 1 (data) or history_len > 1 (model), and K = 1 / stride 1 is bit-for-bit today's pipeline. Data: Human.get_keymap gains proprio_history and history_stride; for K > 1 it marks the proprio entries with "history": K plus the stride, and ZarrDataset.__getitem__ reads the window idx - s*(K-1) ... idx in steps of s, clamped at the episode start and front-padded by repeating the first real frame, so the key is always (K, D) with the current frame last. A proprio_history_mask (1 = real, 0 = front padding) travels with it. The stride is why this ships on rather than at K = 1. Measured on a real flagship episode (30 fps, fold_clothes), the median wrist translation is 6.8 mm over 1 frame, 25.7 mm over 3, 42.9 mm over 5 and 80.8 mm over 10: an unstrided K = 3 window is 67 ms of past, close enough to a duplicate of the current pose to be mostly sensor noise, while the action chunk it conditions spans a full second. The flagship recipe therefore ships proprio_history: 3 / history_stride: 5 -- 0.33 s of past -- with the stem at history_len: 3. That changes the observation, so runs before and after this commit are not comparable, and history_dropout: 0.2 stops being inert: it is now a real hyperparameter, and the mitigation for the copycat/causal-confusion risk that proprio history classically brings. Transforms: the cartesian builders already treat a leading axis as batch; ActionChunkCoordinateFrameTransform / PoseCoordinateFrameTransform gain an opt-in target_history so a proprio target frame collapses to the CURRENT step (the action chunk must not move when history is added), and Reshape / QuaternionPoseToYPR tolerate the extra axis. observations.state.ee_pose is (K, 20). Eval: every human revert transform list starts with a SelectCurrentStep on its proprio key, so the val metrics, the cam-frame revert and the overlay video see the single current pose they always saw. No-op at K = 1 and idempotent. Norm stats are the baseline's: per-channel (D,) stats broadcast over (K, D), and both the training and the precompute stats pass force proprio_history=1 so K > 1 and K = 1 runs share one set. Model: MLPPolicyStem gains history_len / history_dropout. For history_len > 1 it owns a zero-init learned time embedding added after the projection, and HPT.stem_process skips the input-side sinusoid for such a stem (a constant [0, 1, 0, 1, ...] on normalized values is amplitude-1 noise on [-1, 1] data). history_dropout replaces the whole past by the current step per sample in training, exactly what an episode-start sample looks like after front padding. The K = 1 baseline stays one override away (proprio_history=1, history_stride=1, history_len=1) and the tests pin that it is bit-identical. Every other shipped config is untouched at K = 1, including the keypoint arm, whose wrist-frame proprio would need per-step framing before it can carry history. A compose test pins every train_*.yaml recipe's data K against its stem's history_len, because a recipe that inherits one model's data config and overrides the model is exactly how the two drift apart, and nothing else composes both halves. _stem_history_len takes the domain rather than matching a stem name by suffix: a cotrain run registers one state_ee_pose stem per domain, and a suffix match returns whichever comes first, so the human side was checked against eva's history_len. Latent while everything shipped K = 1; a blocker now. Co-Authored-By: Claude Fable 5.1 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/algo/hpt.py | 82 +- .../mecka_fold_flagship_opsplit_hpt_6d.yaml | 12 + .../model/hpt_bc_mecka_6d_300M.yaml | 10 + egomimic/models/hpt_nets.py | 49 +- egomimic/rldb/embodiment/human.py | 66 +- egomimic/rldb/zarr/action_chunk_transforms.py | 70 +- egomimic/rldb/zarr/zarr_dataset_multi.py | 65 ++ egomimic/scripts/precompute_norm_stats.py | 7 + egomimic/trainHydra.py | 7 + tests/unit/test_proprio_history.py | 881 ++++++++++++++++++ 10 files changed, 1226 insertions(+), 23 deletions(-) create mode 100644 tests/unit/test_proprio_history.py diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index c152e5976d..0d9fb1a60a 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -407,16 +407,24 @@ def stem_process(self, domain, data): horizon = np.random.randint(1, data_horizon + 1) data[modality] = data[modality][:, data_horizon - horizon :] - positional_embedding = get_sinusoid_encoding_table( - 0, horizon * int(np.prod(data_shape[2:-1])), data_shape[-1] - ).to(data[modality]) - positional_embedding = einops.repeat( - positional_embedding, "b h w -> (repeat b) h w", repeat=data_shape[0] - ) + # An observation-history stem owns its own LEARNED time-step + # embedding (added after the projection); the input-side sinusoid + # here is a constant [0, 1, 0, 1, ...] added to the NORMALIZED + # values, which would be amplitude-1 noise on [-1, 1] data once + # axis 1 carries more than the current step. + if getattr(stem, "history_len", 1) == 1: + positional_embedding = get_sinusoid_encoding_table( + 0, horizon * int(np.prod(data_shape[2:-1])), data_shape[-1] + ).to(data[modality]) + positional_embedding = einops.repeat( + positional_embedding, + "b h w -> (repeat b) h w", + repeat=data_shape[0], + ) - data[modality] = data[modality] + positional_embedding.view( - data[modality].shape - ) + data[modality] = data[modality] + positional_embedding.view( + data[modality].shape + ) stem_token = stem.compute_latent(data[modality]) feats.append(stem_token) feat_dict[modality] = stem_token @@ -1113,6 +1121,10 @@ def process_batch_for_training(self, batch): action_pad_mask = processed_batch[embodiment_id].pop( "action_pad_mask", None ) + # Front-padded history steps are copies of the first real frame -- + # exactly a history-dropout sample -- so the stem needs nothing + # extra and downstream keys stay as they were. + processed_batch[embodiment_id].pop("proprio_history_mask", None) if self.use_pad_mask and action_pad_mask is not None: if action_pad_mask.shape[-1] != S: raise ValueError( @@ -1173,7 +1185,13 @@ def forward_training(self, batch): ac_key = self.ac_keys[embodiment_id] aux_ac_keys = self.auxiliary_ac_keys.get(embodiment_name, []) data = self._robomimic_to_hpt_data( - _batch, cam_keys, proprio_keys, lang_keys, ac_key, aux_ac_keys + _batch, + cam_keys, + proprio_keys, + lang_keys, + ac_key, + aux_ac_keys, + domain=embodiment_name, ) hpt_batch = { "domain": embodiment_name, # readability on config side @@ -1231,7 +1249,13 @@ def forward_eval(self, batch): ac_key = self.ac_keys[embodiment_id] aux_ac_keys = self.auxiliary_ac_keys.get(embodiment_name, []) data = self._robomimic_to_hpt_data( - _batch, cam_keys, proprio_keys, lang_keys, ac_key, aux_ac_keys + _batch, + cam_keys, + proprio_keys, + lang_keys, + ac_key, + aux_ac_keys, + domain=embodiment_name, ) hpt_batch = { "domain": embodiment_name, # readability on config side @@ -1364,8 +1388,25 @@ def _apply_image_augs(self, images, short): return self.eval_image_augs(images) return images + def _stem_history_len(self, domain: str, modality: str) -> int: + """``history_len`` of ``domain``'s ``modality`` stem (1 when the stem + declares none). Stems are registered as ``f"{domain}_{modality}"``; a + cotrain run has one per domain, so the domain has to be named -- a + suffix match would return whichever came first in the dict. + """ + stems = getattr(self.nets["policy"], "stems", None) or {} + stem = stems.get(f"{domain}_{modality}", stems.get(modality)) + return int(getattr(stem, "history_len", 1)) if stem is not None else 1 + def _robomimic_to_hpt_data( - self, batch, cam_keys, proprio_keys, lang_keys, ac_key, aux_ac_keys=[] + self, + batch, + cam_keys, + proprio_keys, + lang_keys, + ac_key, + aux_ac_keys=[], + domain="", ): """ helper method that returns data in the format required for the HPT model @@ -1378,7 +1419,22 @@ def _robomimic_to_hpt_data( for key in proprio_keys: if key in batch: short = key.rsplit(".", 1)[-1] - data[f"state_{short}"] = batch[key].unsqueeze(1) + modality = f"state_{short}" + value = batch[key] + if value.ndim == 2: + data[modality] = value.unsqueeze(1) # (B, D) -> (B, 1, D) + else: + # Observation history: already (B, K, D), one token per step. + history_len = self._stem_history_len(domain, modality) + if value.shape[1] != history_len: + raise ValueError( + f"proprio '{key}' carries {value.shape[1]} history " + f"steps but the '{modality}' stem has " + f"history_len={history_len}; set " + "model...state_ee_pose.history_len and the keymap's " + "proprio_history to the same K" + ) + data[modality] = value for key in cam_keys: if key in batch: diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml index b8172887b0..36d76ec617 100644 --- a/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_6d.yaml @@ -77,6 +77,18 @@ train_datasets: _target_: egomimic.rldb.embodiment.human.Human.get_keymap keymap_mode: cartesian annotation_key: annotations + # Proprio observation history: the ee_pose keys read K steps spaced + # `history_stride` frames apart, (K, 20) with the current frame last. + # At 30 fps the median wrist translation is 6.8 mm over 1 frame but + # 42.9 mm over 5, so an unstrided window is mostly sensor noise while + # the action chunk it conditions spans a full second; K = 3, s = 5 is + # 0.33 s of past. Must match the stem's `history_len` (the run fails + # loudly if not), and the K = 1 baseline needs + # ...key_map.proprio_history=1 ...key_map.history_stride=1 + # model.robomimic_model.stem_specs.human_bimanual.state_ee_pose.history_len=1 + # The two val loaders interpolate this resolver, so they follow. + proprio_history: 3 + history_stride: 5 transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list mode: cartesian_wristframe_6d diff --git a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml index 66ca79a60a..84293e83ff 100644 --- a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml @@ -123,6 +123,16 @@ robomimic_model: input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} output_dim: 840 widths: [840] + # Proprio observation history: K tokens (one per step, current step + # last) plus a learned zero-init time-step embedding, instead of the + # single current-step token. Must equal the data config's + # `proprio_history` (which also sets how far apart the steps are). + # `history_dropout` replaces the whole past by the current step for a + # sample (train only) so the policy still works when history is + # missing at episode start -- it is live at history_len 3, and it is + # what keeps the policy from reading its own past as the answer. + history_len: 3 + history_dropout: 0.2 specs: random_horizon_masking: false cross_attn: diff --git a/egomimic/models/hpt_nets.py b/egomimic/models/hpt_nets.py index fee4c7367c..df97fb76a4 100644 --- a/egomimic/models/hpt_nets.py +++ b/egomimic/models/hpt_nets.py @@ -541,10 +541,35 @@ def __init__( tanh_end: bool = False, ln: bool = True, num_of_copy: int = 1, + history_len: int = 1, + history_dropout: float = 0.0, **kwargs, ) -> None: - """vanilla MLP class""" + """vanilla MLP class + + ``history_len`` (K) > 1 makes this an observation-history stem: the + input carries the last K steps on axis 1 (current step last), one token + per step, and a LEARNED per-step embedding (zero-init, so at init the K + tokens are exactly today's projection of each step) is added AFTER the + projection. ``HPT.stem_process`` skips its input-side sinusoid for such + a stem: that sinusoid is a constant [0, 1, 0, 1, ...] added to the + NORMALIZED values, i.e. amplitude-1 noise on [-1, 1] data. + + ``history_dropout`` is the per-sample probability, in training only, of + replacing every past step by the current one. That is exactly what an + episode-start sample looks like after the dataset's front padding, so + the policy trains on "history missing" the way it will see it at + deployment (the mitigation for copycat / causal confusion). Inert at + K = 1 and in eval. + """ super().__init__(**kwargs) + if history_len > 1 and num_of_copy > 1: + raise ValueError( + f"MLPPolicyStem: history_len={history_len} and " + f"num_of_copy={num_of_copy} both consume axis 1; use one." + ) + self.history_len = int(history_len) + self.history_dropout = float(history_dropout) modules = [nn.Linear(input_dim, widths[0]), nn.SiLU()] for i in range(len(widths) - 1): @@ -562,6 +587,15 @@ def __init__( self.net = nn.ModuleList( [nn.Sequential(*modules) for _ in range(num_of_copy)] ) + if self.history_len > 1: + self.time_embed = nn.Parameter(torch.zeros(self.history_len, output_dim)) + + def _drop_history(self, x: torch.Tensor) -> torch.Tensor: + """Per sample, with probability ``history_dropout``, replace every past + step by the current one (the last slot of axis 1).""" + drop = torch.rand(x.shape[0], device=x.device) < self.history_dropout + current = x[:, -1:].expand_as(x) + return torch.where(drop.view(-1, *([1] * (x.ndim - 1))), current, x) def forward(self, x: torch.Tensor) -> torch.Tensor: """ @@ -572,6 +606,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: Flatten tensor with shape [B, M, 512] """ + if self.history_len > 1: + if x.shape[1] != self.history_len: + raise ValueError( + f"MLPPolicyStem got {x.shape[1]} history steps on axis 1 but " + f"history_len={self.history_len} (input shape {tuple(x.shape)})" + ) + if self.training and self.history_dropout > 0: + x = self._drop_history(x) if self.num_of_copy > 1: out = [] iter_num = min(self.num_of_copy, x.shape[1]) @@ -582,6 +624,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y = torch.stack(out, dim=1) else: y = self.net(x) + if self.history_len > 1: + # Learned time-step embedding on axis 1, last slot = current step. + y = y + self.time_embed.view( + 1, self.history_len, *([1] * (y.ndim - 3)), y.shape[-1] + ) return y diff --git a/egomimic/rldb/embodiment/human.py b/egomimic/rldb/embodiment/human.py index d8e69ef3e6..57a12abf76 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -25,6 +25,7 @@ QuaternionPoseToYPR, Reshape, RotateLocalFrame, + SelectCurrentStep, SplitKeys, Transform, UnpadGripperZeros, @@ -232,6 +233,8 @@ def get_keymap( include_ee_pose: bool = False, norm_mode: bool = False, annotation_key: str = None, + proprio_history: int = 1, + history_stride: int = 1, ): """Build the keymap. Per-vendor knobs are explicit args from the data config: ``has_head_pose`` (False only for contributors that truly omit @@ -241,12 +244,23 @@ def get_keymap( the keypoint transform modes can build the cartesian ``observations.state.ee_pose`` proprio, the pi0.5 prompt state). ``norm_mode``/``annotation_key`` behave as in the base. + + ``proprio_history`` (K) makes the proprio keys the model consumes read + the last K frames instead of one: the dataset emits ``(K, D)`` with the + current frame last, plus a ``proprio_history_mask``. K = 1 (the + default) leaves the keymap exactly as it was. ``history_stride`` (s) + spaces that window out -- the steps are ``idx - s*(K-1) ... idx`` -- + because at 30 fps consecutive frames are nearly a duplicate of the + current pose (median wrist translation 6.8 mm at lag 1 against 42.9 mm + at lag 5), so s = 1 spends tokens on sensor noise. """ key_map = cls._get_keymap( keymap_mode, has_head_pose=has_head_pose, include_aria_keypoints=include_aria_keypoints, include_ee_pose=include_ee_pose, + proprio_history=proprio_history, + history_stride=history_stride, ) if annotation_key is not None and not norm_mode: key_map[annotation_key] = { @@ -270,12 +284,15 @@ def _get_keymap( has_head_pose: bool = True, include_aria_keypoints: bool = False, include_ee_pose: bool = False, + proprio_history: int = 1, + history_stride: int = 1, ): """Canonical MANO keymap. ``include_aria_keypoints`` additionally exposes the raw Aria-layout proprio keypoints alongside the MANO ones; ``include_ee_pose`` adds the palm-origin ee_pose keys to the keypoints keymap. The front image is always ``VIZ_IMAGE_KEY``; Pi renames it - onto its own slot. + onto its own slot. ``proprio_history`` > 1 marks the proprio entries + the model consumes with ``"history": K`` (see ``get_keymap``). """ base_mode = _strip_pi_keymap_mode(cls, keymap_mode) front_key = cls.VIZ_IMAGE_KEY @@ -386,6 +403,28 @@ def _get_keymap( "key_type": "proprio_keys", "zarr_key": "obs_head_pose", } + + if proprio_history > 1: + if history_stride < 1: + raise ValueError(f"history_stride must be >= 1, got {history_stride}") + # Backward window on the proprio the MODEL consumes. Not on + # ``obs_head_pose``: every history step is expressed in the CURRENT + # step's head frame (a per-step head frame would make the history + # unreadable), and not on the horizoned ``action_*`` entries, which + # are forward chunks. + if base_mode == "cartesian": + history_keys = ["left.obs_ee_pose", "right.obs_ee_pose"] + else: + history_keys = [ + f"{side}.{name}" + for side in ("left", "right") + for name in ("obs_keypoints", "obs_wrist_pose") + ] + if include_ee_pose: + history_keys += ["left.obs_ee_pose", "right.obs_ee_pose"] + for key in history_keys: + key_map[key]["history"] = int(proprio_history) + key_map[key]["history_stride"] = int(history_stride) return key_map @classmethod @@ -574,6 +613,9 @@ def _build_human_keypoints_revert_eef_frame_transform_list( else: pose_shape = 6 transform_list = [ + # proprio_history: the proprio defines the frame, so keep the CURRENT + # step (no-op at K = 1, idempotent after the 6D revert). + SelectCurrentStep(keys=[obs_key]), SplitKeys( input_key=obs_key, output_key_list=[ @@ -712,17 +754,22 @@ def _build_human_keypoints_eef_frame_transform_list( output_key=right_keypoints_action_headframe, shape=(chunk_length, 21, 3), ), + # target_history: the proprio wrist pose may carry a leading time + # axis (proprio_history); everything is expressed in the CURRENT + # step's wrist frame either way. ActionChunkCoordinateFrameTransform( target_world=left_wrist_obs_headframe, chunk_world=left_keypoints_action_headframe, transformed_key_name=left_keypoints_action_wristframe, mode="xyz", + target_history=True, ), ActionChunkCoordinateFrameTransform( target_world=right_wrist_obs_headframe, chunk_world=right_keypoints_action_headframe, transformed_key_name=right_keypoints_action_wristframe, mode="xyz", + target_history=True, ), Reshape( input_key=left_keypoints_action_wristframe, @@ -749,12 +796,14 @@ def _build_human_keypoints_eef_frame_transform_list( pose_world=left_keypoints_obs_headframe, transformed_key_name=left_keypoints_obs_wristframe, mode="xyz", + target_history=True, ), PoseCoordinateFrameTransform( target_world=right_wrist_obs_headframe, pose_world=right_keypoints_obs_headframe, transformed_key_name=right_keypoints_obs_wristframe, mode="xyz", + target_history=True, ), Reshape( input_key=left_keypoints_obs_wristframe, @@ -771,12 +820,14 @@ def _build_human_keypoints_eef_frame_transform_list( chunk_world=left_wrist_action_headframe, transformed_key_name=left_wrist_action_wristframe, mode="xyzwxyz", + target_history=True, ), ActionChunkCoordinateFrameTransform( target_world=right_wrist_obs_headframe, chunk_world=right_wrist_action_headframe, transformed_key_name=right_wrist_action_wristframe, mode="xyzwxyz", + target_history=True, ), ] ) @@ -1071,6 +1122,10 @@ def _build_human_cartesian_revert_eef_frame_transform_list( pose_shape = 7 if is_quat else 6 mode = "xyzwxyz" if is_quat else "xyzypr" transform_list = [ + # The proprio defines the frame the action chunk is projected out of; + # with proprio_history it arrives (K, D), so take the CURRENT step + # (no-op at K = 1, and idempotent when a 6D revert already did it). + SelectCurrentStep(keys=[obs_key]), SplitKeys( input_key=obs_key, output_key_list=[ @@ -1123,6 +1178,7 @@ def _build_human_cartesian_revert_6d_transform_list( grip-padded 18 -> 20) is reverted the same way. """ return [ + SelectCurrentStep(keys=[obs_key]), # proprio_history: keep the current step CartesianRot6DToYPR(action_key=action_key), CartesianRot6DToYPR(action_key=obs_key), UnpadGripperZeros(action_key=obs_key), @@ -1144,6 +1200,7 @@ def _build_human_cartesian_revert_6d_wristframe_transform_list( actions back into head frame using that ypr proprio to define the frame. """ return [ + SelectCurrentStep(keys=[obs_key]), # proprio_history: keep the current step CartesianRot6DToYPR(action_key=action_key), CartesianRot6DToYPR(action_key=obs_key), UnpadGripperZeros(action_key=obs_key), @@ -1162,6 +1219,7 @@ def _build_human_keypoints_revert_6d_transform_list( changes). The proprio is reverted the same way. """ return [ + SelectCurrentStep(keys=[obs_key]), # proprio_history: keep the current step KeypointsRot6DToYPR(action_key=action_key), KeypointsRot6DToYPR(action_key=obs_key), ] @@ -1180,6 +1238,7 @@ def _build_human_keypoints_revert_6d_wristframe_transform_list( ``actions_keypoints`` (the wrist-pose slices are consumed by the revert). """ return [ + SelectCurrentStep(keys=[obs_key]), # proprio_history: keep the current step KeypointsRot6DToYPR(action_key=action_key), KeypointsRot6DToYPR(action_key=obs_key), *_build_human_keypoints_revert_eef_frame_transform_list( @@ -1272,17 +1331,22 @@ def _build_human_cartesian_eef_frame_transform_list( mode="xyzwxyz", ), InterpolatePadMask(new_chunk_length=chunk_length, stride=stride), + # target_history: the proprio obs pose may carry a leading time axis + # (proprio_history); the action chunk is expressed in the CURRENT + # step's wrist frame either way. ActionChunkCoordinateFrameTransform( target_world=left_obs_headframe, chunk_world=left_action_headframe, transformed_key_name=left_action_wristframe, mode="xyzwxyz", + target_history=True, ), ActionChunkCoordinateFrameTransform( target_world=right_obs_headframe, chunk_world=right_action_headframe, transformed_key_name=right_action_wristframe, mode="xyzwxyz", + target_history=True, ), XYZWXYZ_to_XYZYPR( keys=[ diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/zarr/action_chunk_transforms.py index 2be8127e66..38117cad3e 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -201,6 +201,7 @@ def __init__( extra_batch_key: dict = None, mode: Literal["xyz", "xyzwxyz", "xyzypr"] = "xyzwxyz", inverse: bool = True, + target_history: bool = False, ): """ args: @@ -208,6 +209,10 @@ def __init__( chunk_world: transformed_key_name: is_quat: if True, inputs are xyz + quat(wxyz); otherwise xyz + ypr. + target_history: the target frame is a PROPRIO key that may be read + with history, ``(K, D)`` with the current frame last. Take that + current frame; without this flag a 2-D target is an error (a + target frame is a single pose). """ self.target_world = target_world self.chunk_world = chunk_world @@ -215,6 +220,7 @@ def __init__( self.extra_batch_key = extra_batch_key self.mode = mode self.inverse = inverse + self.target_history = target_history def transform(self, batch): """ @@ -237,6 +243,11 @@ def transform(self, batch): chunk_world = np.asarray(batch[self.chunk_world]) chunk_world_shape = None + if self.target_history and target_world.ndim > 1: + # Proprio read with history: (K, D), current frame last. The chunk + # is always expressed in the CURRENT step's frame. + target_world = target_world[-1] + if chunk_world.ndim > 2: chunk_world_shape = chunk_world.shape chunk_world = chunk_world.reshape(-1, chunk_world_shape[-1]) @@ -288,7 +299,11 @@ def transform(self, batch): class QuaternionPoseToYPR(Transform): - """Convert a single pose from xyz + quat(x,y,z,w) to xyz + ypr.""" + """Convert a single pose from xyz + quat(x,y,z,w) to xyz + ypr. + + A leading time axis (``(K, 7)``, a proprio key read with history) is + treated as a batch axis and comes back as ``(K, 6)``. + """ def __init__(self, pose_key: str, output_key: str): self.pose_key = pose_key @@ -296,15 +311,17 @@ def __init__(self, pose_key: str, output_key: str): def transform(self, batch: dict) -> dict: pose = np.asarray(batch[self.pose_key]) - if pose.shape != (7,): + if pose.ndim > 2 or pose.shape[-1] != 7: raise ValueError( - f"QuaternionPoseToYPR expects shape (7,), got {pose.shape} for key " - f"'{self.pose_key}'" + f"QuaternionPoseToYPR expects shape (7,) or (K, 7), got " + f"{pose.shape} for key '{self.pose_key}'" ) - xyz = pose[:3] - xyzw = wxyz_to_xyzw(pose[3:7]) + flat = pose.reshape(-1, 7) + xyz = flat[:, :3] + xyzw = wxyz_to_xyzw(flat[:, 3:7]) ypr = R.from_quat(xyzw).as_euler("ZYX", degrees=False) - batch[self.output_key] = np.concatenate([xyz, ypr], axis=0) + out = np.concatenate([xyz, ypr], axis=-1) + batch[self.output_key] = out.reshape(*pose.shape[:-1], 6) return batch @@ -380,6 +397,7 @@ def __init__( pose_world: str, transformed_key_name: str, mode: Literal["xyzwxyz", "xyzypr", "xyz"] = "xyzwxyz", + target_history: bool = False, ): self.target_world = target_world self.pose_world = pose_world @@ -390,6 +408,7 @@ def __init__( chunk_world=pose_world, transformed_key_name=transformed_key_name, mode=mode, + target_history=target_history, ) def transform(self, batch: dict) -> dict: @@ -769,6 +788,32 @@ def transform(self, batch: dict) -> dict: return batch +class SelectCurrentStep(Transform): + """Drop a proprio key's observation-history axis, keeping the CURRENT step. + + A proprio key read with ``proprio_history`` K > 1 is ``(K, D)`` per sample + with the current frame last. The revert / viz pipelines want the single + current pose the forward pipeline used as the frame, so they start by + collapsing it. ``ndim`` is the flat per-sample rank of the key (1 for a + ``(D,)`` pose vector); a value already at that rank is left alone, so the + K = 1 path is bit-for-bit unchanged and the transform is idempotent. + """ + + def __init__(self, keys: list[str], ndim: int = 1): + self.keys = list(keys) + self.ndim = ndim + + def transform(self, batch: dict) -> dict: + for key in self.keys: + value = batch.get(key) + if value is None or not hasattr(value, "ndim"): + continue + while value.ndim > self.ndim: + value = value[..., -1, :] + batch[key] = value + return batch + + class ConcatKeys(Transform): def __init__(self, key_list, new_key_name, delete_old_keys=False): self.key_list = list(key_list) @@ -850,13 +895,22 @@ def transform(self, batch: dict) -> dict: class Reshape(Transform): + """Reshape a key to ``shape``. A value holding a whole multiple of + ``shape`` (a proprio key read with history, ``(K, ...)``) keeps its leading + time axis and is reshaped per step. + """ + def __init__(self, input_key: str, output_key: str, shape: tuple): self.input_key = input_key self.output_key = output_key self.shape = shape def transform(self, batch: dict) -> dict: - batch[self.output_key] = batch[self.input_key].reshape(*self.shape) + value = batch[self.input_key] + n = value.numel() if isinstance(value, torch.Tensor) else np.asarray(value).size + target = int(np.prod(self.shape)) + shape = self.shape if n == target else (-1, *self.shape) + batch[self.output_key] = value.reshape(*shape) return batch diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 89bd9733d2..a5fff53a59 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -2163,6 +2163,50 @@ def _action_pad_mask(self, idx: int) -> np.ndarray | None: mask[: max(n_real, 0)] = 1.0 return mask + def _proprio_history(self) -> tuple[int, int] | None: + """``(K, stride)`` of the keymap's ``history`` entries, else None.""" + for spec in self.key_map.values(): + history = spec.get("history") + if history is not None and int(history) > 1: + return int(history), int(spec.get("history_stride", 1) or 1) + return None + + def _proprio_history_mask(self, idx: int) -> np.ndarray | None: + """1.0 per real history frame, 0.0 per front-padded one, ``(K,)``. + + The backward window is clamped at the episode start, so the first + stride * (K - 1) samples are short and ``_front_pad_history`` repeats + their first real frame. ``None`` when the keymap has no history key. + """ + spec = self._proprio_history() + if spec is None: + return None + history, stride = spec + n_real = min(history, idx // stride + 1) + mask = np.zeros(history, dtype=np.float32) + mask[history - n_real :] = 1.0 + return mask + + @staticmethod + def _stride_history(array: np.ndarray, history: int, stride: int) -> np.ndarray: + """Take the current frame and every ``stride``-th frame before it, in + order, from a contiguous backward window.""" + if stride > 1: + array = array[::-1][::stride][::-1] + return ZarrDataset._front_pad_history(array, history) + + @staticmethod + def _front_pad_history(array: np.ndarray, history: int) -> np.ndarray: + """Front-pad a backward window to ``history`` steps by repeating its + first real frame, so the current frame is always ``[-1]``. The mirror + of ``_pad_sequences``, which repeat-pads the TAIL of a forward chunk. + """ + seq_len = array.shape[0] + if seq_len >= history: + return array + padding = np.repeat(array[:1], history - seq_len, axis=0) + return np.concatenate([padding, array], axis=0) + def _pad_sequences(self, data, horizon: int | None) -> dict: if horizon is None: return data @@ -2219,6 +2263,7 @@ def _next(reason: str, key: str = "") -> int: zarr_key = self.key_map[k]["zarr_key"] key_type = self.key_map[k].get("key_type", None) horizon = self.key_map[k].get("horizon", None) + history = self.key_map[k].get("history", None) if key_type == "annotation_keys": data[k] = self._annotation_text_for_frame(idx) @@ -2227,12 +2272,26 @@ def _next(reason: str, key: str = "") -> int: if horizon is not None: end_idx = self._chunk_end_idx(idx, horizon, key_type) read_interval = (idx, end_idx) + elif history is not None and int(history) > 1: + # Backward window ending at (and including) the current + # frame, every `history_stride`-th step; clamped at the + # episode start and front-padded below, so the key is + # always (K, D) with `[..., -1]` the single frame a K = 1 + # keymap would have read. Read contiguously and thin after: + # the zarr chunks are contiguous anyway. + stride = int(self.key_map[k].get("history_stride", 1) or 1) + read_interval = ( + max(0, idx - stride * (int(history) - 1)), + idx + 1, + ) else: read_interval = (idx, None) read_dict = {zarr_key: read_interval} raw_data = self.episode_reader.read(read_dict) self._pad_sequences(raw_data, horizon) # should be able to pad images data[k] = raw_data[zarr_key] + if horizon is None and history is not None and int(history) > 1: + data[k] = self._stride_history(data[k], int(history), stride) if zarr_key in self._image_keys: jpeg_bytes = data[k] @@ -2257,6 +2316,12 @@ def _next(reason: str, key: str = "") -> int: # (InterpolatePadMask) to the model horizon. data["action_pad_mask"] = pad_mask + history_mask = self._proprio_history_mask(idx) + if history_mask is not None: + # Informational: the front padding already looks exactly like a + # history-dropout sample, so the stem needs nothing extra. + data["proprio_history_mask"] = history_mask + if self.transform: for transform in self.transform or []: data = transform.transform(data) diff --git a/egomimic/scripts/precompute_norm_stats.py b/egomimic/scripts/precompute_norm_stats.py index 482c613bcd..85d7c197ea 100644 --- a/egomimic/scripts/precompute_norm_stats.py +++ b/egomimic/scripts/precompute_norm_stats.py @@ -134,6 +134,13 @@ def main() -> None: inst = copy.deepcopy(cfg.data.train_datasets[dataset_name]) km = OmegaConf.to_container(inst.resolver.key_map, resolve=False) km["norm_mode"] = True # strips image + annotation keys + if "proprio_history" in km: + # Current-step proprio stats even on a history (K > 1) config: every + # history step is normalized with the CURRENT-step stats, so a K > 1 + # run and the K = 1 baseline share one norm-stat file. Reading one + # frame is exactly `[..., -1, :]` of the window, minus the K-fold + # read. + km["proprio_history"] = 1 inst.resolver.key_map = km norm_dataset = hydra.utils.instantiate(inst, dataset_name=dataset_name) diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 79f3291780..28a7de9aa4 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -537,6 +537,13 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: # this remove annotation and image keys from the keymap km["norm_mode"] = True + if "proprio_history" in km: + # Current-step proprio stats even on a history (K > 1) config: + # every history step is normalized with the CURRENT-step, + # per-channel (D,) stats, so a K > 1 run and the K = 1 baseline + # share exactly one set of stats. Same rule as + # scripts/precompute_norm_stats.py. + km["proprio_history"] = 1 instantiate_copy.resolver.key_map = km norm_dataset = _instantiate_dataset( diff --git a/tests/unit/test_proprio_history.py b/tests/unit/test_proprio_history.py new file mode 100644 index 0000000000..c05a24d029 --- /dev/null +++ b/tests/unit/test_proprio_history.py @@ -0,0 +1,881 @@ +"""Proprio observation history: the last K proprio vectors, one token per step. + +K = 1 must stay bit-for-bit today's pipeline -- same sample keys, same shapes, +same stem tokens, same sinusoidal PE -- so every check below has a K = 1 twin. +For K > 1 the dataset reads a backward window (front-padded with the first real +frame so the current frame is always ``[-1]``), the cartesian transforms treat +the leading axis as batch, the baseline's per-channel proprio stats broadcast +over it, and the stem replaces the input-side sinusoid with a learned +time-step embedding added after the projection. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import torch +from omegaconf import OmegaConf + +import egomimic +from egomimic.algo.hpt import HPT, HPTModel +from egomimic.models.hpt_nets import MLPPolicyStem +from egomimic.rldb.embodiment.embodiment import get_embodiment_id +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.zarr.zarr_dataset_multi import ( + LocalEpisodeResolver, + MultiDataset, +) +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +T_FRAMES = 24 +HORIZON = Human.ACTION_HORIZON +CHUNK = 100 +PROPRIO_KEYS = ("left.obs_ee_pose", "right.obs_ee_pose") + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _pose_series(T: int, offset: float) -> np.ndarray: + """(T, 7) xyz+quat(wxyz) whose translation is a distinct ramp per frame. + + The shared synthetic fixture holds every pose CONSTANT over time (so the + per-cell quantile bounds stay usable); history needs frames that differ, + so this writes its own episode with an identity rotation and a ramp. + """ + t = np.arange(T, dtype=np.float64)[:, None] + xyz = offset + t * np.array([[0.001, 0.002, 0.003]]) + quat = np.tile(np.array([[1.0, 0.0, 0.0, 0.0]]), (T, 1)) + return np.concatenate([xyz, quat], axis=1) + + +def _write_varying_episode(root, *, T: int = T_FRAMES, seed: int = 0): + rng = np.random.default_rng(seed) + H = W = 64 + K = np.array( + [[100.0, 0.0, W / 2, 0.0], [0.0, 100.0, H / 2, 0.0], [0.0, 0.0, 1.0, 0.0]] + ) + numeric = { + "left.obs_ee_pose": _pose_series(T, 0.10 + seed), + "right.obs_ee_pose": _pose_series(T, 0.20 + seed), + "obs_head_pose": _pose_series(T, 0.30 + seed), + } + images = {"images.front_1": rng.integers(0, 255, (T, H, W, 3), dtype=np.uint8)} + return ZarrWriter.create_and_write( + root / f"varying_{seed:02d}.zarr", + numeric_data=numeric, + image_data=images, + embodiment="human_bimanual", + fps=30, + task_name="synthetic", + intrinsics={"front_1": K}, + ) + + +def _leaf(ds: MultiDataset): + return next(iter(ds.datasets.values())) + + +def _dataset( + tmp_path, + *, + proprio_history: int = 1, + history_stride: int = 1, + transforms=None, + n: int = 1, +): + for i in range(n): + _write_varying_episode(tmp_path, seed=i) + key_map = Human.get_keymap( + keymap_mode="cartesian", + proprio_history=proprio_history, + history_stride=history_stride, + ) + resolver = LocalEpisodeResolver( + tmp_path, key_map=key_map, transform_list=transforms or [] + ) + return MultiDataset._from_resolver(resolver, mode="total") + + +def _raw_poses(tmp_path, key: str) -> np.ndarray: + """The written (T, 7) array for ``key``, straight off disk.""" + import zarr + + store = zarr.open_group(str(next(tmp_path.glob("*.zarr"))), mode="r") + return np.asarray(store[key][:]) + + +# --------------------------------------------------------------------------- +# 1. dataset: backward window, front padding, mask +# --------------------------------------------------------------------------- + + +def test_history_window_is_the_last_k_frames(tmp_path): + leaf = _leaf(_dataset(tmp_path, proprio_history=4)) + for key in PROPRIO_KEYS: + raw = _raw_poses(tmp_path, key) + got = np.asarray(leaf[10][key]) + assert got.shape == (4, 7) + np.testing.assert_allclose(got, raw[7:11], rtol=0, atol=1e-6) + # The current frame is always last. + np.testing.assert_allclose(got[-1], raw[10], rtol=0, atol=1e-6) + + +def test_history_stride_spaces_the_window_out(tmp_path): + """K = 3, s = 5 reads frames idx-10, idx-5, idx -- 0.33 s of past at 30 fps + instead of 0.07 s, with the current frame still last.""" + leaf = _leaf(_dataset(tmp_path, proprio_history=3, history_stride=5)) + for key in PROPRIO_KEYS: + raw = _raw_poses(tmp_path, key) + got = np.asarray(leaf[12][key]) + assert got.shape == (3, 7) + np.testing.assert_allclose(got, raw[[2, 7, 12]], rtol=0, atol=1e-6) + np.testing.assert_array_equal( + np.asarray(leaf[12]["proprio_history_mask"]), np.ones(3, dtype=np.float32) + ) + + +def test_strided_history_front_pads_at_the_episode_start(tmp_path): + leaf = _leaf(_dataset(tmp_path, proprio_history=3, history_stride=5)) + raw = _raw_poses(tmp_path, PROPRIO_KEYS[0]) + # idx 7: frames 7 and 2 are real, idx-10 would be negative + got = np.asarray(leaf[7][PROPRIO_KEYS[0]]) + np.testing.assert_allclose(got, raw[[2, 2, 7]], rtol=0, atol=1e-6) + np.testing.assert_array_equal( + np.asarray(leaf[7]["proprio_history_mask"]), + np.array([0, 1, 1], dtype=np.float32), + ) + # idx 3: only the current frame is real + np.testing.assert_array_equal( + np.asarray(leaf[3]["proprio_history_mask"]), + np.array([0, 0, 1], dtype=np.float32), + ) + + +def test_stride_one_is_the_contiguous_window(tmp_path): + a = _leaf(_dataset(tmp_path, proprio_history=4)) + b = _leaf(_dataset(tmp_path, proprio_history=4, history_stride=1)) + for key in PROPRIO_KEYS: + np.testing.assert_array_equal(np.asarray(a[10][key]), np.asarray(b[10][key])) + + +def test_history_mask_is_all_ones_away_from_the_episode_start(tmp_path): + leaf = _leaf(_dataset(tmp_path, proprio_history=4)) + mask = np.asarray(leaf[10]["proprio_history_mask"]) + assert mask.shape == (4,) + np.testing.assert_array_equal(mask, np.ones(4, dtype=np.float32)) + + +def test_history_front_pads_by_repeating_the_first_real_frame(tmp_path): + leaf = _leaf(_dataset(tmp_path, proprio_history=4)) + for key in PROPRIO_KEYS: + raw = _raw_poses(tmp_path, key) + got = np.asarray(leaf[1][key]) + assert got.shape == (4, 7) + expected = np.stack([raw[0], raw[0], raw[0], raw[1]]) + np.testing.assert_allclose(got, expected, rtol=0, atol=1e-6) + + # NOTE: the brief's example mask for idx=1 reads [0, 0, 0, 1]; the rule it + # states ("1 = real, 0 = front padding") gives [0, 0, 1, 1] -- frames 0 and + # 1 are both real reads, only the two leading copies of frame 0 are + # padding. The rule wins; the mask is informational only (it is dropped in + # process_batch_for_training and the padded steps already look exactly like + # a history-dropout sample). + mask = np.asarray(leaf[1]["proprio_history_mask"]) + np.testing.assert_array_equal(mask, np.array([0, 0, 1, 1], dtype=np.float32)) + + np.testing.assert_array_equal( + np.asarray(leaf[0]["proprio_history_mask"]), + np.array([0, 0, 0, 1], dtype=np.float32), + ) + + +def test_history_mask_is_float32(tmp_path): + value = _leaf(_dataset(tmp_path, proprio_history=4))[3]["proprio_history_mask"] + assert isinstance(value, torch.Tensor) + assert value.dtype == torch.float32 + + +def test_k1_sample_is_bit_for_bit_todays_sample(tmp_path): + """The K = 1 keymap must produce exactly what a keymap with no + ``proprio_history`` argument at all produces.""" + _write_varying_episode(tmp_path, seed=0) + baseline = Human.get_keymap(keymap_mode="cartesian") + with_arg = Human.get_keymap(keymap_mode="cartesian", proprio_history=1) + assert with_arg == baseline + assert not any("history" in spec for spec in with_arg.values()) + + def _sample(key_map): + resolver = LocalEpisodeResolver(tmp_path, key_map=key_map, transform_list=[]) + ds = MultiDataset._from_resolver(resolver, mode="total") + return _leaf(ds)[5] + + a, b = _sample(baseline), _sample(with_arg) + assert set(a) == set(b) + assert "proprio_history_mask" not in a + for key in PROPRIO_KEYS: + assert tuple(a[key].shape) == (7,) + for key, value in a.items(): + if isinstance(value, torch.Tensor): + torch.testing.assert_close(value, b[key]) + + +def test_keys_without_history_keep_the_single_frame_read(tmp_path): + """Only the proprio keys the model consumes get a time axis: the head pose + (the frame everything is expressed in) and the action chunks do not.""" + leaf = _leaf(_dataset(tmp_path, proprio_history=4)) + sample = leaf[10] + assert tuple(sample["obs_head_pose"].shape) == (7,) + assert tuple(sample["left.action_ee_pose"].shape) == (HORIZON, 7) + + +# --------------------------------------------------------------------------- +# 2. transforms: a leading time axis is a batch axis +# --------------------------------------------------------------------------- + + +def _wristframe_batch(poses: dict[str, np.ndarray]) -> dict: + rng = np.random.default_rng(0) + chunk = _pose_series(HORIZON, 0.4) + batch = { + "obs_head_pose": _pose_series(1, 0.3)[0], + "left.action_ee_pose": chunk.copy(), + "right.action_ee_pose": chunk.copy() + 0.01, + "action_pad_mask": np.ones(HORIZON, dtype=np.float32), + } + del rng + batch.update({k: v.copy() for k, v in poses.items()}) + return batch + + +def _run_wristframe(batch: dict) -> dict: + for t in Human.get_transform_list( + mode="cartesian_wristframe_6d", + stride=1, + fix_left_wrist_convention=True, + pad_proprio_gripper=True, + ): + batch = t.transform(batch) + return batch + + +def test_wristframe_chain_maps_k_by_7_proprio_to_k_by_20(): + K = 4 + hist = {key: _pose_series(K, 0.1 * (i + 1)) for i, key in enumerate(PROPRIO_KEYS)} + out = _run_wristframe(_wristframe_batch(hist)) + state = np.asarray(out["observations.state.ee_pose"]) + assert state.shape == (K, 20) + + current = {key: value[-1] for key, value in hist.items()} + out1 = _run_wristframe(_wristframe_batch(current)) + state1 = np.asarray(out1["observations.state.ee_pose"]) + assert state1.shape == (20,) + np.testing.assert_allclose(state[-1], state1, rtol=1e-6, atol=1e-8) + + +def test_wristframe_actions_are_unchanged_by_the_proprio_time_axis(): + """The action chunk is expressed in the CURRENT step's wrist frame, so it + must not move when history is added next to it.""" + K = 4 + hist = {key: _pose_series(K, 0.1 * (i + 1)) for i, key in enumerate(PROPRIO_KEYS)} + current = {key: value[-1] for key, value in hist.items()} + a = _run_wristframe(_wristframe_batch(hist))["actions_cartesian"] + b = _run_wristframe(_wristframe_batch(current))["actions_cartesian"] + np.testing.assert_allclose(np.asarray(a), np.asarray(b), rtol=1e-6, atol=1e-8) + + +def test_dataset_and_transforms_compose_to_a_k_by_20_state(tmp_path): + transforms = Human.get_transform_list( + mode="cartesian_wristframe_6d", + stride=1, + fix_left_wrist_convention=True, + pad_proprio_gripper=True, + ) + hist = _leaf(_dataset(tmp_path, proprio_history=4, transforms=transforms)) + base = _leaf(_dataset(tmp_path, proprio_history=1, transforms=transforms)) + for idx in (0, 3, 11): + h = hist[idx]["observations.state.ee_pose"] + b = base[idx]["observations.state.ee_pose"] + assert tuple(h.shape) == (4, 20) + assert tuple(b.shape) == (20,) + torch.testing.assert_close(h[-1], b, rtol=1e-5, atol=1e-6) + + +# --------------------------------------------------------------------------- +# 3. normalization / bounds over the time axis +# --------------------------------------------------------------------------- + + +def _norm_shell(key: str, width: int, *, lo=-1.0, hi=1.0): + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = "quantile" + md.norm_stats = { + 0: { + key: { + "quantile_1": np.full(width, lo, dtype=np.float32), + "quantile_99": np.full(width, hi, dtype=np.float32), + } + } + } + md.zarr_keys = {0: {key: key}} + md.key_types = {0: {key: "proprio_keys"}} + md._warned_violations = set() + return md + + +def test_per_channel_stats_broadcast_over_the_time_axis(): + key = "observations.state.ee_pose" + md = _norm_shell(key, 20, lo=-2.0, hi=4.0) + stats = md.norm_stats[0][key] + g = torch.Generator().manual_seed(0) + arr = torch.randn(4, 20, generator=g) + out = md._apply_norm_one(arr, stats) + assert out.shape == (4, 20) + for k in range(4): + torch.testing.assert_close(out[k], md._apply_norm_one(arr[k], stats)) + + +def test_bounds_check_catches_a_violation_in_a_past_step(): + key = "observations.state.ee_pose" + md = _norm_shell(key, 20) + arr = np.zeros((4, 20), dtype=np.float32) + assert md._check_bounds({"embodiment": 0, key: arr.copy()}, None, 0, "ep") is None + bad = arr.copy() + bad[0, 0] = 50.0 # oldest history step, translation channel + assert md._check_bounds({"embodiment": 0, key: bad}, None, 0, "ep") is not None + + +# --------------------------------------------------------------------------- +# 4. the stem +# --------------------------------------------------------------------------- + + +def _stems(D: int, out: int, **kwargs): + torch.manual_seed(0) + a = MLPPolicyStem(input_dim=D, output_dim=out, widths=[16], **kwargs) + torch.manual_seed(0) + b = MLPPolicyStem(input_dim=D, output_dim=out, widths=[16]) + return a, b + + +def test_history_stem_shape_and_zero_init_matches_the_baseline_stem(): + D, out, B, K = 5, 7, 3, 3 + hist, base = _stems(D, out, history_len=K) + x = torch.randn(B, K, 1, D) + y = hist(x) + assert y.shape == (B, K, 1, out) + # time_embed is zero at init, so every step is today's projection. + torch.testing.assert_close(y, base(x)) + for k in range(K): + torch.testing.assert_close(y[:, k], base(x[:, k])) + + +def test_time_embed_is_learned_zero_init_and_per_step(): + D, out, K = 5, 7, 3 + hist, _ = _stems(D, out, history_len=K) + assert hist.time_embed.shape == (K, out) + assert hist.time_embed.requires_grad + torch.testing.assert_close(hist.time_embed, torch.zeros(K, out)) + + with torch.no_grad(): + hist.time_embed.copy_( + torch.arange(K, dtype=torch.float32)[:, None].expand(K, out) + ) + x = torch.randn(2, K, 1, D) + y = hist(x) + assert not torch.allclose(y[:, 0], y[:, 1]) + # ... and the shift is exactly the per-step embedding. + hist.time_embed.data.zero_() + y0 = hist(x) + hist.time_embed.data.copy_( + torch.arange(K, dtype=torch.float32)[:, None].expand(K, out) + ) + torch.testing.assert_close( + hist(x) - y0, + torch.arange(K, dtype=torch.float32)[None, :, None, None].expand(2, K, 1, out), + ) + + +def test_history_len_one_stem_has_no_time_embed_and_is_todays_stem(): + D, out = 5, 7 + one, base = _stems(D, out, history_len=1) + assert not hasattr(one, "time_embed") + assert dict(one.named_parameters()).keys() == dict(base.named_parameters()).keys() + x = torch.randn(4, 1, 1, D) + torch.testing.assert_close(one(x), base(x)) + + +def test_history_dropout_replaces_the_past_with_the_current_step_in_train_mode(): + D, out, B, K = 5, 7, 6, 4 + hist, _ = _stems(D, out, history_len=K, history_dropout=1.0) + x = torch.randn(B, K, 1, D) + hist.train() + y = hist(x) + # Every step equals the current one: exactly an episode-start sample after + # front padding. + for k in range(K): + torch.testing.assert_close( + y[:, k], y[:, -1] - hist.time_embed[-1] + hist.time_embed[k] + ) + hist.eval() + torch.testing.assert_close(y, hist(_current_only(x))) + + +def _current_only(x: torch.Tensor) -> torch.Tensor: + out = x.clone() + out[:, :-1] = out[:, -1:] + return out + + +def test_history_dropout_is_a_noop_in_eval_mode(): + D, out, B, K = 5, 7, 6, 4 + drop, _ = _stems(D, out, history_len=K, history_dropout=1.0) + torch.manual_seed(0) + plain = MLPPolicyStem(input_dim=D, output_dim=out, widths=[16], history_len=K) + drop.eval() + plain.eval() + x = torch.randn(B, K, 1, D) + torch.testing.assert_close(drop(x), plain(x)) + assert not torch.allclose(drop(x), drop(_current_only(x))) + + +def test_history_dropout_zero_is_a_noop_in_train_mode(): + D, out, B, K = 5, 7, 6, 4 + a, _ = _stems(D, out, history_len=K, history_dropout=0.0) + torch.manual_seed(0) + b = MLPPolicyStem(input_dim=D, output_dim=out, widths=[16], history_len=K) + a.train() + b.eval() + x = torch.randn(B, K, 1, D) + torch.testing.assert_close(a(x), b(x)) + + +def test_history_stem_rejects_a_wrong_length_history(): + hist, _ = _stems(5, 7, history_len=3) + with pytest.raises(ValueError, match="3"): + hist(torch.randn(2, 4, 1, 5)) + + +# --------------------------------------------------------------------------- +# 5. stem_process: the input-side sinusoid goes once history_len > 1 +# --------------------------------------------------------------------------- + + +class _EchoStem: + """Stem stub whose ``compute_latent`` hands back what stem_process fed it.""" + + def __init__(self, history_len: int = 1): + self.history_len = history_len + self.seen = None + + def compute_latent(self, x): + self.seen = x.clone() + return x + + +def _stem_process(history_len: int, data_shape): + model = HPTModel.__new__(HPTModel) + model.modalities = {"human_bimanual": ["state_ee_pose"]} + model.shared_keys = [] + model.encoders = {} + stem = _EchoStem(history_len) + model.stems = {"human_bimanual_state_ee_pose": stem} + model.stem_spec = { + "human_bimanual": { + "state_ee_pose": OmegaConf.create( + {"specs": {"random_horizon_masking": False}} + ) + } + } + model.train_mode = False + torch.manual_seed(0) + x = torch.randn(*data_shape) + model.stem_process("human_bimanual", {"state_ee_pose": x.clone()}) + return x, stem.seen + + +def test_stem_process_skips_the_sinusoid_for_a_history_stem(): + x, seen = _stem_process(4, (2, 4, 20)) + torch.testing.assert_close(seen, x) + + +def test_stem_process_keeps_the_sinusoid_for_a_history_len_one_stem(): + x, seen = _stem_process(1, (2, 1, 20)) + from egomimic.utils.tensor_utils import get_sinusoid_encoding_table + + pe = get_sinusoid_encoding_table(0, 1, 20).to(x) + torch.testing.assert_close(seen, x + pe.view(1, 1, 20)) + assert not torch.allclose(seen, x) + + +# --------------------------------------------------------------------------- +# 6. HPT batch plumbing +# --------------------------------------------------------------------------- + + +class _NoRenameNormStats: + def zarr_key_to_keyname(self, zarr_key, embodiment_id): + return None + + +def _hpt_with_stem(history_len: int): + algo = HPT.__new__(HPT) + policy = HPTModel.__new__(HPTModel) + policy.stems = {"human_bimanual_state_ee_pose": _EchoStem(history_len)} + algo.nets = {"policy": policy} + algo.is_6dof = True + algo.shared_ac_key = None + algo.annotation_modality = "annotation" + algo.eval_image_augs = None + return algo + + +def _to_hpt(algo, proprio: torch.Tensor): + batch = { + "observations.state.ee_pose": proprio, + "pad_mask": torch.ones(proprio.shape[0], 5, 1), + "embodiment": torch.tensor([0]), + "actions_cartesian": torch.zeros(proprio.shape[0], 5, 18), + } + return HPT._robomimic_to_hpt_data( + algo, + batch, + cam_keys=[], + proprio_keys=["observations.state.ee_pose"], + lang_keys=[], + ac_key="actions_cartesian", + domain="human_bimanual", + ) + + +def test_two_dim_proprio_is_still_unsqueezed(): + algo = _hpt_with_stem(1) + data = _to_hpt(algo, torch.randn(3, 20)) + assert data["state_ee_pose"].shape == (3, 1, 20) + + +def test_three_dim_proprio_passes_through_without_unsqueeze(): + algo = _hpt_with_stem(4) + x = torch.randn(3, 4, 20) + data = _to_hpt(algo, x) + assert data["state_ee_pose"].shape == (3, 4, 20) + torch.testing.assert_close(data["state_ee_pose"], x) + + +def test_history_len_is_read_from_the_domain_not_by_suffix(): + """A cotrain run has one state_ee_pose stem per domain; a suffix match + returns whichever comes first in the dict, so the human side would be + checked against eva's history_len.""" + algo = _hpt_with_stem(3) + algo.nets["policy"].stems = { + "eva_bimanual_state_ee_pose": _EchoStem(1), + "human_bimanual_state_ee_pose": _EchoStem(3), + } + data = _to_hpt(algo, torch.randn(3, 3, 20)) + assert data["state_ee_pose"].shape == (3, 3, 20) + + +def test_history_length_mismatch_raises_naming_both(): + algo = _hpt_with_stem(2) + with pytest.raises(ValueError, match="state_ee_pose"): + _to_hpt(algo, torch.randn(3, 4, 20)) + + +def _process(action_pad_mask=None, history_mask=None, S: int = 6): + algo = HPT.__new__(HPT) + algo.norm_stats = _NoRenameNormStats() + algo.device = "cpu" + algo.annotation_key = None + algo.use_pad_mask = False + emb = "human_bimanual" + emb_id = get_embodiment_id(emb) + algo.ac_keys = {emb_id: "actions_cartesian"} + inner = { + "actions_cartesian": torch.zeros(2, S, 18), + "observations.state.ee_pose": torch.zeros(2, 4, 20), + } + if action_pad_mask is not None: + inner["action_pad_mask"] = action_pad_mask + if history_mask is not None: + inner["proprio_history_mask"] = history_mask + return HPT.process_batch_for_training(algo, {emb: inner})[emb_id] + + +def test_process_batch_drops_the_history_mask(): + out = _process(history_mask=torch.ones(2, 4)) + assert "proprio_history_mask" not in out + assert out["observations.state.ee_pose"].shape == (2, 4, 20) + + +def test_process_batch_is_unchanged_without_the_history_mask(): + out = _process() + assert "proprio_history_mask" not in out + torch.testing.assert_close(out["pad_mask"], torch.ones(2, 6, 1)) + + +# --------------------------------------------------------------------------- +# 7. configs +# --------------------------------------------------------------------------- + +RECIPE = "train_zarr_mecka_flagship_6d_hpt" +DATA = "data=mecka_fold_flagship_opsplit_hpt_6d" +KEYMAP = "data.train_datasets.human_bimanual.resolver.key_map" +STEM = "model.robomimic_model.stem_specs.human_bimanual.state_ee_pose" + + +def test_flagship_ships_k_three_stride_five(compose_resolve): + cfg = compose_resolve(RECIPE, [DATA]) + km = cfg.data.train_datasets.human_bimanual.resolver.key_map + assert km.proprio_history == 3 + assert km.history_stride == 5 + stem = cfg.model.robomimic_model.stem_specs.human_bimanual.state_ee_pose + assert stem.history_len == km.proprio_history, "data K and stem K must agree" + assert stem.history_dropout == 0.2 + + +def test_k_one_baseline_is_still_reachable(compose_resolve): + cfg = compose_resolve( + RECIPE, + [ + DATA, + f"{KEYMAP}.proprio_history=1", + f"{KEYMAP}.history_stride=1", + f"{STEM}.history_len=1", + ], + ) + km = cfg.data.train_datasets.human_bimanual.resolver.key_map + assert km.proprio_history == 1 and km.history_stride == 1 + assert ( + cfg.model.robomimic_model.stem_specs.human_bimanual.state_ee_pose.history_len + == 1 + ) + + +def test_k_two_override_reaches_every_loader(compose_resolve): + cfg = compose_resolve( + RECIPE, [DATA, f"{KEYMAP}.proprio_history=2", f"{STEM}.history_len=2"] + ) + data = cfg.data + for group in ("train_datasets", "valid_datasets", "unseen_op_valid_datasets"): + km = data[group]["human_bimanual"]["resolver"]["key_map"] + assert km["proprio_history"] == 2, group + assert ( + cfg.model.robomimic_model.stem_specs.human_bimanual.state_ee_pose.history_len + == 2 + ) + + +def test_k_two_keymap_marks_only_the_model_proprio_keys(): + km = Human.get_keymap(keymap_mode="cartesian", proprio_history=2) + assert {k for k, v in km.items() if "history" in v} == set(PROPRIO_KEYS) + assert all(km[k]["history"] == 2 for k in PROPRIO_KEYS) + + +# --------------------------------------------------------------------------- +# 8. the eval path: the revert transforms and the evaluator take the current +# step, so val metrics and the overlay video work at K > 1 +# --------------------------------------------------------------------------- + +REVERT_BUILDERS = [ + "_build_human_cartesian_revert_6d_wristframe_transform_list", + "_build_human_cartesian_revert_6d_transform_list", +] + + +def _revert_inputs(B: int = 2, K: int = 3): + """``(actions, current_proprio, history_proprio)`` whose last history step + is exactly the 2-D proprio.""" + g = torch.Generator().manual_seed(7) + actions = torch.randn(B, 100, 18, generator=g) + current = torch.randn(B, 20, generator=g) + history = torch.randn(B, K, 20, generator=g) + history[:, -1] = current + return actions, current, history + + +@pytest.mark.parametrize("builder_name", REVERT_BUILDERS) +def test_revert_transform_list_is_blind_to_proprio_history(builder_name): + from egomimic.rldb.embodiment import human as human_mod + from egomimic.rldb.embodiment.embodiment import Embodiment + + transform_list = getattr(human_mod, builder_name)() + actions, current, history = _revert_inputs() + + def _run(proprio): + return Embodiment.apply_transform( + { + "actions_cartesian": actions.clone(), + "observations.state.ee_pose": proprio.clone(), + }, + transform_list, + ) + + flat, hist = _run(current), _run(history) + assert flat["observations.state.ee_pose"].shape == (2, 12) + assert hist["observations.state.ee_pose"].shape == (2, 12) + for key in ("actions_cartesian", "observations.state.ee_pose"): + torch.testing.assert_close(hist[key], flat[key], rtol=1e-5, atol=1e-6) + + +def test_select_current_step_leaves_a_flat_value_alone_and_is_idempotent(): + from egomimic.rldb.zarr.action_chunk_transforms import SelectCurrentStep + + t = SelectCurrentStep(keys=["observations.state.ee_pose", "missing"]) + flat = np.arange(20, dtype=np.float32) + out = t.transform({"observations.state.ee_pose": flat}) + assert out["observations.state.ee_pose"] is flat # untouched, not a copy + + hist = np.stack([flat - 1, flat, flat + 1]) + once = t.transform({"observations.state.ee_pose": hist.copy()}) + np.testing.assert_array_equal(once["observations.state.ee_pose"], flat + 1) + twice = t.transform(dict(once)) + np.testing.assert_array_equal(twice["observations.state.ee_pose"], flat + 1) + + +def test_every_human_revert_list_selects_the_current_proprio_step(): + from egomimic.rldb.embodiment import human as human_mod + from egomimic.rldb.zarr.action_chunk_transforms import SelectCurrentStep + + builders = { + "_build_human_cartesian_revert_6d_transform_list": "observations.state.ee_pose", + "_build_human_cartesian_revert_6d_wristframe_transform_list": "observations.state.ee_pose", + "_build_human_cartesian_revert_eef_frame_transform_list": "observations.state.ee_pose", + "_build_human_keypoints_revert_6d_transform_list": "observations.state.keypoints", + "_build_human_keypoints_revert_6d_wristframe_transform_list": "observations.state.keypoints", + "_build_human_keypoints_revert_eef_frame_transform_list": "observations.state.keypoints", + } + for name, obs_key in builders.items(): + tl = getattr(human_mod, name)() + assert isinstance(tl[0], SelectCurrentStep), name + assert tl[0].keys == [obs_key], name + + +# --- the evaluator itself ----------------------------------------------- + + +class _IdentityNormStats: + def unnormalize(self, batch, embodiment_id): + return dict(batch) + + +class _StubAlgo: + """Just enough of HPT for HPTEvalVideo.compute_metrics_and_viz.""" + + def __init__(self, emb_id, preds): + self.norm_stats = _IdentityNormStats() + self.ac_keys = {emb_id: "actions_cartesian"} + self.shared_ac_key = None + self.auxiliary_ac_keys = {} + self._preds = preds + + def forward_eval(self, batch): + return self._preds + + +def _eval_once(proprio): + """Run the real HPTEvalVideo over one batch, returning its metrics and the + batch the viz function was handed.""" + from egomimic.eval.eval_hpt import HPTEvalVideo + from egomimic.rldb.embodiment.human import ( + _build_human_cartesian_revert_6d_wristframe_transform_list, + ) + + emb_id = get_embodiment_id("human_bimanual") + emb_name = "human_bimanual" + # Action horizon 1: the evaluator's paired/final MSE does + # ``pred[:, -1].cpu()``, which is a real (contiguous) copy for the CUDA + # tensors a run produces but a non-contiguous view for the CPU tensors a + # test hands it, and torchmetrics' ``view(-1)`` rejects that. S = 1 keeps + # the slice contiguous; the proprio path under test is unaffected. + g = torch.Generator().manual_seed(11) + actions = torch.randn(2, 1, 18, generator=g) + preds = { + f"{emb_name}_actions_cartesian": torch.randn(2, 1, 18, generator=g), + f"{emb_name}_loss": torch.tensor(0.5), + } + batch = { + emb_id: { + "actions_cartesian": actions.clone(), + "observations.state.ee_pose": proprio.clone(), + "embodiment": torch.tensor([emb_id, emb_id]), + } + } + + seen = {} + + def _viz(predictions, viz_batch): + seen["batch"] = viz_batch + seen["preds"] = predictions + return np.zeros((2, 4, 4, 3), dtype=np.uint8) + + ev = HPTEvalVideo.__new__(HPTEvalVideo) + ev.model = _StubAlgo(emb_id, preds) + ev.viz_func = {emb_name: _viz} + ev.transform_lists = { + emb_name: _build_human_cartesian_revert_6d_wristframe_transform_list() + } + metrics, images = ev.compute_metrics_and_viz(batch, do_viz=True) + return metrics, images, seen + + +def test_evaluator_metrics_and_viz_are_blind_to_proprio_history(): + _, current, history = _revert_inputs() + m_flat, im_flat, seen_flat = _eval_once(current) + m_hist, im_hist, seen_hist = _eval_once(history) + + assert set(m_flat) == set(m_hist) + cam_keys = [k for k in m_flat if "_cam_" in k] + assert cam_keys, "the cam-frame metrics must be produced" + for key in m_flat: + torch.testing.assert_close( + torch.as_tensor(m_flat[key]).float(), + torch.as_tensor(m_hist[key]).float(), + rtol=1e-5, + atol=1e-6, + ) + # The viz function sees the reverted, current-step proprio either way. + assert seen_flat["batch"]["observations.state.ee_pose"].shape == (2, 12) + torch.testing.assert_close( + seen_hist["batch"]["observations.state.ee_pose"], + seen_flat["batch"]["observations.state.ee_pose"], + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + seen_hist["preds"]["human_bimanual_actions_cartesian"], + seen_flat["preds"]["human_bimanual_actions_cartesian"], + rtol=1e-5, + atol=1e-6, + ) + assert set(im_flat) == set(im_hist) + + +TRAIN_CONFIGS = sorted( + p.stem + for p in (Path(egomimic.__file__).parent / "hydra_configs").glob("train_*.yaml") +) + + +@pytest.mark.parametrize("top", TRAIN_CONFIGS) +def test_every_recipe_agrees_with_its_stem_on_K(top, compose_resolve): + """The data's `proprio_history` and the stem's `history_len` must match, or + process_batch_for_training raises at the first training step. A recipe that + inherits a data config from one model and overrides the model is exactly + how they drift apart, and nothing else composes both halves.""" + cfg = compose_resolve(top, []) + stem_specs = (cfg.model.robomimic_model.get("stem_specs") or {}).get( + "human_bimanual" + ) or {} + for split in ("train_datasets", "valid_datasets"): + node = (cfg.data.get(split) or {}).get("human_bimanual") + if node is None or node.get("resolver") is None: + continue + k = int(node.resolver.key_map.get("proprio_history", 1) or 1) + for name, stem in stem_specs.items(): + if stem is None or "history_len" not in stem: + continue + assert int(stem.history_len) == k, (top, split, name, k, stem.history_len) From b84343b2cc2b641eae3427b447b7590522ea4647 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 15 Sep 2026 14:39:13 -0400 Subject: [PATCH 19/79] hpt: never re-initialise pretrained stems; hash-check them; mask text padding `HPTModel.finalize_modules` ran `self.apply(self._init_weights)` over the whole model, which xavier-randomised the pretrained HF encoders inside the text stems (the frozen Qwen3-Embedding weights were overwritten in place and their biases zeroed). Modules that own pretrained weights now declare them through the `PretrainedWeights` mixin, and the init pass walks the tree without ever entering such a subtree, so the ResNet backbone, `T5Encoder` and any future pretrained stem are covered by the same mechanism. `HPT.__init__` also called `finalize_modules` twice; the first call is gone. Once the model is built, moved and upcast, `verify_pretrained_weights` hashes every declared pretrained submodule and compares it with the checkpoint read straight off the local HF snapshot, failing the run on a mismatch. It costs a second full copy of the encoder (~1.1 GB fp16 at 0.6B) and ~3 s of hashing per rank, so `verify_pretrained: false` turns it off; it is on by default, because a silently re-initialised encoder trains to a plausible wrong result. Only stems that own an HF snapshot are checked -- `ResNet` returns no reference, so it is protected from the init pass but not hash-verified, which the mixin docstring now says. Finally, `QwenPerTokenEncoder.compute_latent` now passes the tokenizer attention mask into its cross-attention. Padded positions are zeroed before the biased `proj`, so they were the learned constant `proj.bias`, not zero, and got real attention weight: text-conditioned results from before this commit are not comparable with results after it. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/algo/hpt.py | 26 +- egomimic/models/hpt_nets.py | 231 ++++++++++++++- tests/unit/test_text_encoder_fixes.py | 388 ++++++++++++++++++++++++++ 3 files changed, 630 insertions(+), 15 deletions(-) create mode 100644 tests/unit/test_text_encoder_fixes.py diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index 0d9fb1a60a..6467fc1b05 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -14,7 +14,12 @@ from tslearn.metrics import SoftDTWLossPyTorch from egomimic.algo.algo import Algo -from egomimic.models.hpt_nets import MultiheadAttention, SimpleTransformer +from egomimic.models.hpt_nets import ( + MultiheadAttention, + SimpleTransformer, + apply_skipping_pretrained, + verify_pretrained_weights, +) from egomimic.models.image_augs import PerSampleAugs from egomimic.rldb.embodiment.embodiment import get_embodiment, get_embodiment_id from egomimic.utils.hf_utils import download_from_huggingface @@ -198,7 +203,12 @@ def finalize_modules(self): """ self.stems = nn.ModuleDict(self.stems) self.heads = nn.ModuleDict(self.heads) - self.apply(self._init_weights) + # Only initialize the modules HPT builds itself. ``self.apply`` would + # walk into pretrained submodules (the HF text encoders, the ImageNet + # ResNet backbone) and overwrite their weights with xavier noise; the + # traversal below skips anything a module declares via + # ``PretrainedWeights``, so future pretrained stems are covered too. + apply_skipping_pretrained(self, self._init_weights) # Shared action tokens if self.token_postprocessing == "action_token": @@ -943,8 +953,6 @@ def __init__( for modality, encoder_cfg in self.encoders.items(): model.init_encoder(modality, encoder_cfg) - model.finalize_modules() - self.ac_keys = {} self.camera_keys = {} self.proprio_keys = {} @@ -998,6 +1006,16 @@ def __init__( self.nets["policy"] = model self.nets = self.nets.float().to(self.device) + # Every pretrained submodule must still hold its checkpoint's weights + # once the model is fully built, moved and upcast. Raises on mismatch. + # It costs a second full copy of the encoder (~1.1 GB fp16 for 0.6B) + # and ~3 s of hashing per rank, so it is skippable -- but on by + # default: a silently re-initialised encoder trains to a plausible, + # wrong result. + self.verify_pretrained = bool(kwargs.get("verify_pretrained", True)) + if self.verify_pretrained: + verify_pretrained_weights(self.nets["policy"]) + self.training_step = 0 # ------------------------------------------------------------------ diff --git a/egomimic/models/hpt_nets.py b/egomimic/models/hpt_nets.py index df97fb76a4..30d47f107d 100644 --- a/egomimic/models/hpt_nets.py +++ b/egomimic/models/hpt_nets.py @@ -1,3 +1,5 @@ +import hashlib +import json import os from functools import partial from typing import Callable, List, Optional, Union @@ -8,6 +10,7 @@ import torch.utils.checkpoint as checkpoint import torchvision from einops import rearrange, repeat +from termcolor import cprint from timm.models.layers import DropPath, trunc_normal_ from timm.models.vision_transformer import VisionTransformer from torch import einsum @@ -16,6 +19,186 @@ from egomimic.utils.tensor_utils import get_sinusoid_encoding_table +# -------------------------------------------------------------------------- +# Pretrained-weight protection +# +# ``HPTModel.finalize_modules`` runs a xavier pass (``_init_weights``) over the +# whole model once the stems / heads are ModuleDicts. That pass used to visit +# every ``nn.Linear`` in every stem, including the HF text encoders loaded from +# a pretrained snapshot, silently randomising them. Modules that own pretrained +# weights declare them here; ``apply_skipping_pretrained`` then walks the tree +# without ever entering such a subtree, so any future pretrained stem (e.g. a +# Qwen VLM stem) is protected without touching ``finalize_modules`` again. +# -------------------------------------------------------------------------- + + +class PretrainedWeights: + """Mixin marking a module whose weights come from a pretrained checkpoint. + + Subclasses list the attribute names holding the pretrained submodules in + ``_hpt_pretrained_attrs``. Two things follow: + + * ``finalize_modules``'s init pass skips those subtrees entirely. + * if the subclass can produce a reference state dict for the checkpoint + (``pretrained_reference_state_dict``), the weights are hashed and + compared with that reference once the model is fully built. + """ + + _hpt_pretrained = True + _hpt_pretrained_attrs: tuple = () + + def pretrained_submodules(self) -> List[nn.Module]: + """The submodules whose weights must never be re-initialised.""" + mods = [] + for name in self._hpt_pretrained_attrs: + mod = getattr(self, name, None) + if isinstance(mod, nn.Module): + mods.append(mod) + return mods + + def pretrained_state_dict(self) -> dict: + """Flat ``{".": tensor}`` view of the pretrained weights.""" + state = {} + for name in self._hpt_pretrained_attrs: + mod = getattr(self, name, None) + if isinstance(mod, nn.Module): + for key, value in mod.state_dict().items(): + state[f"{name}.{key}"] = value + return state + + def pretrained_reference_state_dict(self) -> Optional[dict]: + """The checkpoint's own weights, read back independently of this model. + + ``None`` (the default) means "cannot be verified offline"; the checker + then skips this module instead of failing, so only the stems that own + an HF snapshot are hash-checked. ``ResNet`` is protected from the init + pass but not verified: its torchvision weights are not a file this can + re-read. + """ + return None + + def pretrained_hash_dtype(self) -> Optional[torch.dtype]: + """dtype both sides are cast to before hashing (``None`` = as stored).""" + return None + + +def pretrained_module_ids(root: nn.Module) -> set: + """ids of every module living under a declared pretrained submodule.""" + skip = set() + for module in root.modules(): + get = getattr(module, "pretrained_submodules", None) + if not callable(get): + continue + for sub in get(): + skip.update(id(m) for m in sub.modules()) + return skip + + +def apply_skipping_pretrained(root: nn.Module, fn: Callable) -> None: + """``root.apply(fn)`` that never enters a pretrained submodule. + + Same post-order traversal (and therefore the same RNG consumption) as + ``nn.Module.apply`` for a tree with no pretrained submodules. + """ + skip = pretrained_module_ids(root) + + def _walk(module: nn.Module) -> None: + if id(module) in skip: + return + for child in module.children(): + _walk(child) + fn(module) + + _walk(root) + + +def hash_state_dict(state: dict, dtype: Optional[torch.dtype] = None) -> str: + """Deterministic sha256 over parameter names + raw bytes, in name order.""" + digest = hashlib.sha256() + for name in sorted(state): + tensor = state[name].detach().to("cpu") + if dtype is not None: + tensor = tensor.to(dtype) + digest.update(name.encode("utf-8")) + digest.update(str(tuple(tensor.shape)).encode("utf-8")) + digest.update(str(tensor.dtype).encode("utf-8")) + digest.update(tensor.flatten().contiguous().view(torch.uint8).numpy().tobytes()) + return digest.hexdigest() + + +def verify_pretrained_weights(root: nn.Module, verbose: bool = True) -> dict: + """Hash every declared pretrained submodule and compare with its checkpoint. + + Raises ``RuntimeError`` naming the module and both hashes on a mismatch. + Returns ``{module_name: sha256}`` for the modules that could be checked. + """ + checked = {} + for name, module in root.named_modules(): + get_ref = getattr(module, "pretrained_reference_state_dict", None) + if not callable(get_ref): + continue + reference = get_ref() + if reference is None: + continue + dtype = module.pretrained_hash_dtype() + current = module.pretrained_state_dict() + if set(current) != set(reference): + raise RuntimeError( + f"pretrained weight check failed for '{name or root.__class__.__name__}'" + f" ({module.__class__.__name__}): the checkpoint and the live module" + f" disagree on which tensors exist; only in model:" + f" {sorted(set(current) - set(reference))[:5]}, only in checkpoint:" + f" {sorted(set(reference) - set(current))[:5]}" + ) + live_hash = hash_state_dict(current, dtype) + ref_hash = hash_state_dict(reference, dtype) + label = name or root.__class__.__name__ + if live_hash != ref_hash: + raise RuntimeError( + f"pretrained weights of '{label}' ({module.__class__.__name__}) do not" + f" match their checkpoint: model sha256={live_hash}," + f" checkpoint sha256={ref_hash}. Something modified them between" + " the load and here -- an init pass that descended into the" + " subtree, or an in-place edit. This runs during HPT.__init__," + " before any checkpoint restore, so a resumed checkpoint is not" + " the cause." + ) + checked[label] = live_hash + if verbose: + cprint( + f"[hpt] pretrained weights of '{label}' verified: sha256={live_hash}", + color="green", + ) + return checked + + +def _snapshot_state_dict( + snapshot_dir: str, prefix: str = "", dtype: Optional[torch.dtype] = None +) -> Optional[dict]: + """Read a local HF snapshot's safetensors shards into a flat state dict. + + ``None`` if the snapshot has no safetensors (the caller then skips the + check rather than hitting the network). + """ + from safetensors.torch import load_file + + index = os.path.join(snapshot_dir, "model.safetensors.index.json") + if os.path.isfile(index): + with open(index, encoding="utf-8") as handle: + shards = sorted(set(json.load(handle)["weight_map"].values())) + files = [os.path.join(snapshot_dir, shard) for shard in shards] + else: + single = os.path.join(snapshot_dir, "model.safetensors") + if not os.path.isfile(single): + return None + files = [single] + + state = {} + for path in files: + for key, tensor in load_file(path, device="cpu").items(): + state[f"{prefix}{key}"] = tensor if dtype is None else tensor.to(dtype) + return state + ## Taken directly from hpt/models/transformer with no modifications class CrossAttention(nn.Module): @@ -632,11 +815,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y -class ResNet(PolicyStem): - # ``self.net`` holds the ImageNet-pretrained torchvision backbone (``self.proj`` - # is built here and must keep its own init). Declared so a weight-protection - # pass can find it without knowing about this class. - _hpt_pretrained = True +class ResNet(PretrainedWeights, PolicyStem): + # ``self.net`` is the ImageNet-pretrained torchvision backbone; ``self.proj`` + # is built by HPT and must keep its xavier init. _hpt_pretrained_attrs = ("net",) def __init__( @@ -774,7 +955,7 @@ def _local_snapshot_if_offline(model_name: str) -> str: return snapshot_download(model_name, local_files_only=True) -class _Qwen3BaseEncoder(PolicyStem): +class _Qwen3BaseEncoder(PretrainedWeights, PolicyStem): """Shared base for Qwen3-Embedding stems used by HPT. Owns the tokenizer + transformer encoder so HPT.process_batch_for_training @@ -794,6 +975,7 @@ class _Qwen3BaseEncoder(PolicyStem): """ DEFAULT_MODEL = "Qwen/Qwen3-Embedding-0.6B" + _hpt_pretrained_attrs = ("encoder",) def __init__( self, @@ -819,6 +1001,10 @@ def __init__( load_path = _local_snapshot_if_offline(model_name) self.tokenizer = AutoTokenizer.from_pretrained(load_path, padding_side="left") self.encoder = AutoModel.from_pretrained(load_path, torch_dtype=torch_dtype) + # Remembered so the weight-hash check can read the same snapshot back + # and compare in the dtype the encoder was loaded in. + self._snapshot_dir = load_path + self._load_dtype = torch_dtype if freeze: for p in self.encoder.parameters(): p.requires_grad = False @@ -850,6 +1036,21 @@ def _encode(self, prompts): out = self.encoder(**tokens) return out.last_hidden_state, tokens["attention_mask"] + def pretrained_reference_state_dict(self) -> Optional[dict]: + """The snapshot's own safetensors, cast exactly the way + ``from_pretrained(torch_dtype=...)`` casts them.""" + if not os.path.isdir(self._snapshot_dir): + return None + return _snapshot_state_dict( + self._snapshot_dir, prefix="encoder.", dtype=self._load_dtype + ) + + def pretrained_hash_dtype(self) -> Optional[torch.dtype]: + """Hash in the load dtype. ``HPT.__init__`` upcasts the whole model with + ``nets.float()``; fp16 -> fp32 -> fp16 is exact, so casting the live + weights back down compares like for like without a second model load.""" + return self._load_dtype + def train(self, mode: bool = True): """Keep the frozen encoder in eval mode regardless of outer train flag.""" super().train(mode) @@ -881,18 +1082,26 @@ class QwenPerTokenEncoder(_Qwen3BaseEncoder): information from the learnable stem tokens. """ - def forward(self, prompts): + def forward_with_mask(self, prompts): + """Returns ``(features, attention_mask)``; the mask is what padded + positions must be excluded by, since zeroing them before the biased + ``proj`` turns them into the learned constant ``proj.bias``.""" hidden, mask = self._encode(prompts) feat = hidden.float() * mask.unsqueeze(-1).float() - return self.proj(feat) # (B, L, output_dim) + return self.proj(feat), mask.bool() # (B, L, output_dim), (B, L) + + def forward(self, prompts): + return self.forward_with_mask(prompts)[0] def compute_latent(self, prompts): - feat = self(prompts) # (B, L, hidden_size) + feat, mask = self.forward_with_mask(prompts) # (B, L, output_dim), (B, L) stem_tokens = self.tokens.repeat(feat.shape[0], 1, 1) - return self.cross_attention(stem_tokens, feat) + return self.cross_attention(stem_tokens, feat, mask=mask) + +class T5Encoder(PretrainedWeights, PolicyStem): + _hpt_pretrained_attrs = ("encoder",) -class T5Encoder(PolicyStem): def __init__(self, per_token=True, **kwargs) -> None: """T5 Encoder that expects pre-tokenized inputs diff --git a/tests/unit/test_text_encoder_fixes.py b/tests/unit/test_text_encoder_fixes.py new file mode 100644 index 0000000000..fda2736059 --- /dev/null +++ b/tests/unit/test_text_encoder_fixes.py @@ -0,0 +1,388 @@ +"""Lane C: the HPT text-encoder fixes. + +Three defects, one test file: + +1. ``HPTModel.finalize_modules`` used ``self.apply(self._init_weights)``, which + xavier-randomised the *pretrained* HF encoders inside the text stems (and + would do the same to any future pretrained stem). It now skips anything a + module declares through ``PretrainedWeights``. +2. ``HPT.__init__`` called ``finalize_modules`` twice. +3. ``QwenPerTokenEncoder.compute_latent`` never passed the tokenizer attention + mask to its cross-attention, so padded positions - which are the learned + constant ``proj.bias``, not zero, because the zeroing happens before the + biased ``proj`` - got real attention weight. + +Plus the new guard rail: the pretrained weights are hashed against their +on-disk snapshot once the model is built, and a mismatch fails the run. +""" + +from __future__ import annotations + +import glob +import inspect +import os + +import pytest +import torch +import torch.nn as nn +from omegaconf import OmegaConf + +from egomimic.models.hpt_nets import ( + PolicyStem, + PretrainedWeights, + QwenPerTokenEncoder, + apply_skipping_pretrained, + hash_state_dict, + verify_pretrained_weights, +) + +EMBED_DIM = 32 + + +def _cross_attn_specs(latent: int = 4) -> OmegaConf: + return OmegaConf.create( + { + "random_horizon_masking": False, + "cross_attn": { + "crossattn_latent": latent, + "crossattn_heads": 2, + "crossattn_dim_head": 8, + "crossattn_modality_dropout": 0.0, + "modality_embed_dim": EMBED_DIM, + }, + } + ) + + +# -------------------------------------------------------------------------- +# a tiny stand-in for a pretrained stem +# -------------------------------------------------------------------------- + + +class FakePretrainedStem(PretrainedWeights, nn.Module): + """Carries the marker, an "encoder" that must never be touched, and a + ``proj`` that HPT owns and therefore must still be initialised.""" + + _hpt_pretrained_attrs = ("encoder",) + + def __init__(self) -> None: + super().__init__() + self.encoder = nn.Sequential(nn.Linear(8, 8), nn.LayerNorm(8), nn.Linear(8, 8)) + self.proj = nn.Linear(8, EMBED_DIM) + with torch.no_grad(): # distinctive "checkpoint" values, no default zeros + for param in self.encoder.parameters(): + param.normal_(mean=1.0, std=1.0) + # pretend these came off disk + self._reference = { + k: v.detach().clone() for k, v in self.pretrained_state_dict().items() + } + + def pretrained_reference_state_dict(self): + return {k: v.clone() for k, v in self._reference.items()} + + +def _small_hpt_model(stems: dict): + from egomimic.algo.hpt import HPTModel + + model = HPTModel( + embed_dim=EMBED_DIM, + num_blocks=1, + num_heads=2, + observation_horizon=1, + action_horizon=2, + token_postprocessing="action_token", + ) + model.init_domain_stem("d", stems) + return model + + +# -------------------------------------------------------------------------- +# 1 / 2. the init pass skips pretrained weights and nothing else +# -------------------------------------------------------------------------- + + +def test_apply_skipping_pretrained_matches_apply_when_nothing_is_pretrained(): + """Modules HPT builds itself must be initialised exactly as before: with no + pretrained markers in the tree the new traversal is bit-identical to + ``nn.Module.apply`` (same post-order, same RNG draws).""" + from egomimic.algo.hpt import HPTModel + + def build(): + torch.manual_seed(0) + return nn.Sequential( + nn.Linear(8, 8), + nn.LayerNorm(8), + nn.Sequential(nn.Linear(8, 4), nn.Linear(4, 4)), + ) + + def init(module): # HPTModel._init_weights only ever looks at its argument + HPTModel._init_weights(None, module) + + reference, candidate = build(), build() + torch.manual_seed(1234) + reference.apply(init) + torch.manual_seed(1234) + apply_skipping_pretrained(candidate, init) + + ref_state, cand_state = reference.state_dict(), candidate.state_dict() + assert set(ref_state) == set(cand_state) + for key in ref_state: + assert torch.equal(ref_state[key], cand_state[key]), key + + +def test_finalize_modules_leaves_pretrained_weights_untouched(): + torch.manual_seed(0) + stem = FakePretrainedStem() + before_pretrained = { + k: v.detach().clone() for k, v in stem.pretrained_state_dict().items() + } + before_proj = stem.proj.weight.detach().clone() + + model = _small_hpt_model({"annotation": stem}) + model.finalize_modules() + + after = stem.pretrained_state_dict() + for key, value in before_pretrained.items(): + assert torch.equal(value, after[key]), f"pretrained tensor {key} was modified" + + # ... while the module HPT owns *was* re-initialised (xavier + zero bias). + assert not torch.equal(before_proj, stem.proj.weight) + assert torch.equal(stem.proj.bias, torch.zeros_like(stem.proj.bias)) + + +def test_finalize_modules_initialises_everything_hpt_builds(): + """Every Linear/LayerNorm outside a pretrained subtree gets the xavier + treatment (bias 0 / LayerNorm 1,0), trunk and head included.""" + torch.manual_seed(0) + stem = FakePretrainedStem() + model = _small_hpt_model({"annotation": stem}) + model.finalize_modules() + + skipped = {id(m) for m in stem.encoder.modules()} + seen = 0 + for module in model.modules(): + if id(module) in skipped: + continue + if isinstance(module, nn.Linear) and module.bias is not None: + assert torch.equal(module.bias, torch.zeros_like(module.bias)) + seen += 1 + elif isinstance(module, nn.LayerNorm): + assert torch.equal(module.weight, torch.ones_like(module.weight)) + assert torch.equal(module.bias, torch.zeros_like(module.bias)) + seen += 1 + assert seen > 5, "expected the trunk/stem modules to be visited" + + # the pretrained LayerNorm keeps its (non-default) loaded values + ln = stem.encoder[1] + assert not torch.equal(ln.bias, torch.zeros_like(ln.bias)) + + +# -------------------------------------------------------------------------- +# 3. finalize_modules is called once per HPT.__init__ +# -------------------------------------------------------------------------- + + +def test_finalize_modules_called_once_in_hpt_init(): + """Counted by walking HPT.__init__'s AST, not its source text: a comment + mentioning the call would satisfy a substring count.""" + import ast + import textwrap + + from egomimic.algo.hpt import HPT + + tree = ast.parse(textwrap.dedent(inspect.getsource(HPT.__init__))) + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "finalize_modules" + ] + assert len(calls) == 1, ( + "HPT.__init__ must call finalize_modules exactly once; a second call " + "re-runs the init pass and re-draws the shared action tokens" + ) + + +# -------------------------------------------------------------------------- +# 4. the weight hash check +# -------------------------------------------------------------------------- + + +def test_verify_pretrained_weights_passes_on_a_clean_model(): + torch.manual_seed(0) + stem = FakePretrainedStem() + model = _small_hpt_model({"annotation": stem}) + model.finalize_modules() + checked = verify_pretrained_weights(model, verbose=False) + assert checked, "the marked stem should have been hashed" + ((name, digest),) = checked.items() + assert "annotation" in name + assert digest == hash_state_dict(stem.pretrained_reference_state_dict()) + + +def test_verify_pretrained_weights_raises_when_a_weight_is_perturbed(): + torch.manual_seed(0) + stem = FakePretrainedStem() + model = _small_hpt_model({"annotation": stem}) + model.finalize_modules() + with torch.no_grad(): + stem.encoder[0].weight[0, 0] += 1e-3 + with pytest.raises(RuntimeError, match="do not match their checkpoint"): + verify_pretrained_weights(model, verbose=False) + + +# -------------------------------------------------------------------------- +# 5. the per-token text stem passes the attention mask +# -------------------------------------------------------------------------- + + +class StubPerTokenEncoder(QwenPerTokenEncoder): + """``QwenPerTokenEncoder`` with the HF model swapped for fixed features, so + the mask plumbing is tested without RoPE/position effects.""" + + def __init__(self, hidden: torch.Tensor, mask: torch.Tensor, specs) -> None: + PolicyStem.__init__(self, specs=specs) + self.hidden_size = hidden.shape[-1] + self.output_dim = EMBED_DIM + self.freeze_encoder = False # no HF encoder to keep in eval mode + self._snapshot_dir = "" + self._load_dtype = torch.float32 + self.proj = nn.Linear(self.hidden_size, EMBED_DIM) + self.stub = (hidden, mask) + + def _encode(self, prompts): + return self.stub + + +def _padded_stub(seq_short: int = 3, seq_long: int = 7, hidden_dim: int = 16): + torch.manual_seed(7) + hidden = torch.randn(2, seq_long, hidden_dim) + mask = torch.ones(2, seq_long, dtype=torch.long) + mask[0, : seq_long - seq_short] = 0 # left padding, as the tokenizer pads + return hidden, mask + + +def test_per_token_stem_masks_padding(): + hidden, mask = _padded_stub() + specs = _cross_attn_specs() + stem = StubPerTokenEncoder(hidden, mask, specs) + stem.init_cross_attn(specs.cross_attn) + stem.eval() + + with torch.no_grad(): + batched = stem.compute_latent(None) + + # the same (short) prompt encoded on its own: no padding at all + real = mask[0].bool() + stem.stub = (hidden[:1, real], mask[:1, real]) + alone = stem.compute_latent(None) + + # what the parent did: cross-attention over every position, mask ignored + stem.stub = (hidden, mask) + feat = stem(None) + tokens = stem.tokens.repeat(feat.shape[0], 1, 1) + unmasked = stem.cross_attention(tokens, feat) + + masked_diff = (batched[0] - alone[0]).abs().max().item() + unmasked_diff = (unmasked[0] - alone[0]).abs().max().item() + print( + f"\n[mask] max|short-in-padded-batch - short-alone|: " + f"with mask {masked_diff:.3e} / without mask (parent) {unmasked_diff:.3e}" + ) + assert masked_diff < 1e-5, masked_diff + assert ( + unmasked_diff > 1e-3 + ), "the stub is degenerate: padding made no difference even unmasked" + + +# -------------------------------------------------------------------------- +# the same, with the real Qwen encoder (skipped when the snapshot is absent) +# -------------------------------------------------------------------------- + + +def _qwen_snapshot() -> str | None: + cache = os.environ.get("HF_HUB_CACHE") or os.path.join( + os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface")), "hub" + ) + hits = sorted( + glob.glob( + os.path.join(cache, "models--Qwen--Qwen3-Embedding-0.6B", "snapshots", "*") + ) + ) + return hits[-1] if hits else None + + +QWEN_SNAPSHOT = _qwen_snapshot() +requires_qwen = pytest.mark.skipif( + QWEN_SNAPSHOT is None, reason="Qwen3-Embedding-0.6B snapshot not in the HF cache" +) + + +@pytest.fixture(scope="module") +def qwen_stem(): + if QWEN_SNAPSHOT is None: + pytest.skip("Qwen3-Embedding-0.6B snapshot not in the HF cache") + os.environ.setdefault("HF_HUB_OFFLINE", "1") + specs = _cross_attn_specs() + stem = QwenPerTokenEncoder( + model_name=QWEN_SNAPSHOT, max_length=32, output_dim=EMBED_DIM, specs=specs + ) + stem.init_cross_attn(specs.cross_attn) + return stem + + +@requires_qwen +def test_finalize_modules_keeps_the_real_qwen_encoder_bit_identical(qwen_stem): + before = hash_state_dict(qwen_stem.pretrained_state_dict()) + model = _small_hpt_model({"annotation": qwen_stem}) + model.finalize_modules() + after = hash_state_dict(qwen_stem.pretrained_state_dict()) + reference = hash_state_dict(qwen_stem.pretrained_reference_state_dict()) + print( + f"\n[qwen] sha256 after finalize_modules: {after}\n[qwen] snapshot: {reference}" + ) + assert after == before + assert after == reference + # and the check that runs at train start agrees, including after the upcast + model.float() + verify_pretrained_weights(model, verbose=False) + + +@requires_qwen +def test_verify_pretrained_weights_catches_a_perturbed_qwen_weight(qwen_stem): + param = next(iter(qwen_stem.encoder.parameters())) + original = param.detach().clone() + try: + with torch.no_grad(): + param[0, 0] += 1.0 + with pytest.raises(RuntimeError, match="do not match their checkpoint"): + verify_pretrained_weights(qwen_stem, verbose=False) + finally: + with torch.no_grad(): + param.copy_(original) + + +@requires_qwen +def test_real_qwen_per_token_stem_masks_padding(qwen_stem): + stem = qwen_stem + stem.eval() + prompts = ["pick up the cup", "pick up the red cup and place it on the blue plate"] + with torch.no_grad(): + feat, mask = stem.forward_with_mask(prompts) + latent = stem.compute_latent(prompts) + one_token = stem.tokens.repeat(1, 1, 1) + # reference: the short row's real positions only, same features + alone = stem.cross_attention(one_token, feat[0][mask[0]].unsqueeze(0)) + unmasked = stem.cross_attention(one_token, feat[:1]) + + masked_diff = (latent[0] - alone[0]).abs().max().item() + unmasked_diff = (unmasked[0] - alone[0]).abs().max().item() + n_pad = int((~mask[0]).sum()) + print( + f"\n[qwen-mask] {n_pad} padded positions; max|padded-row - real-tokens-only|: " + f"with mask {masked_diff:.3e} / without mask (parent) {unmasked_diff:.3e}" + ) + assert n_pad > 0, "prompts must differ in length for this test to mean anything" + assert masked_diff < 1e-5, masked_diff + assert unmasked_diff > 1e-3, unmasked_diff From 5d1054611e5fe16866e2d73578943b6b6e3659f9 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 15 Sep 2026 15:24:59 -0400 Subject: [PATCH 20/79] hpt: cache frozen Qwen text features per unique prompt The annotation vocabulary is tiny next to the number of training steps, so a frozen text encoder re-runs the same 600M-parameter forward every step. `_Qwen3BaseEncoder` now keeps a lazy in-memory cache of the per-token hidden states keyed by the exact prompt string: misses go through the tokenizer + encoder in ONE sub-batch of unique strings, their un-padded (L_i, D) rows are stored detached, and the batch is re-assembled with the same LEFT padding the tokenizer produces, so `forward_with_mask` / the pooled path are unchanged. Only legal when the encoder is frozen - a trainable encoder must stay online. Knobs on the stem: `cache_text_features` (default true) and `text_cache_max_entries` (default 8192, LRU); both added to the model yamls with a Qwen stem. The cap counts entries, and an entry is the prompt's un-padded (L, D) states -- 4-25 KB for real annotations, 262 KB at max_length 128 -- so the default is ~100-200 MB of device memory in practice and 2 GB in the worst case; the docstring says so next to the knob. A batch carrying more unique prompts than the cap reads its rows from the miss pass rather than from the cache, because the LRU trim runs first and would otherwise evict entries the very next line looks up. `_apply` (to/cuda/float) clears the cache, and `clear_text_feature_cache()` / `text_feature_cache_stats()` are exposed. Padded positions come back as zeros instead of the encoder's pad-token output; every consumer masks them out. Cached-vs-online on 50 real prompts (fp16, each warmed alone: worst-case padding mismatch): max |diff| 3.6e-3 on the latents and 9.4e-2 on raw hidden states, exactly equal to the online path's own batch-padding noise (9.4e-2); in fp32 the same comparison is 1.1e-6 / 1.6e-4. Co-Authored-By: Claude Fable 5.1 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- .../model/hpt_bc_keypoints_wrist_300M.yaml | 6 + .../model/hpt_bc_mecka_6d_300M.yaml | 6 + egomimic/models/hpt_nets.py | 127 +++++- tests/unit/test_text_feature_cache.py | 421 ++++++++++++++++++ 4 files changed, 559 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_text_feature_cache.py diff --git a/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml index 77f2f4f79c..1b7f4922bf 100644 --- a/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml @@ -103,6 +103,12 @@ robomimic_model: max_length: 128 freeze: true dtype: "float16" + # The frozen encoder's per-token features are cached per unique prompt + # string (LRU, in memory), so the 600M-parameter forward runs once per + # annotation instead of once per training step. Ignored when + # `freeze: false` - a trainable encoder must stay online. + cache_text_features: true + text_cache_max_entries: 8192 output_dim: 840 specs: random_horizon_masking: false diff --git a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml index 84293e83ff..48dea6f4fd 100644 --- a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml @@ -105,6 +105,12 @@ robomimic_model: max_length: 128 freeze: true dtype: "float16" + # The frozen encoder's per-token features are cached per unique prompt + # string (LRU, in memory), so the 600M-parameter forward runs once per + # annotation instead of once per training step. Ignored when + # `freeze: false` - a trainable encoder must stay online. + cache_text_features: true + text_cache_max_entries: 8192 output_dim: 840 specs: random_horizon_masking: false diff --git a/egomimic/models/hpt_nets.py b/egomimic/models/hpt_nets.py index 30d47f107d..746dcf650e 100644 --- a/egomimic/models/hpt_nets.py +++ b/egomimic/models/hpt_nets.py @@ -1,6 +1,7 @@ import hashlib import json import os +from collections import OrderedDict from functools import partial from typing import Callable, List, Optional, Union @@ -972,6 +973,15 @@ class _Qwen3BaseEncoder(PretrainedWeights, PolicyStem): cost down when frozen). normalize_pooled: only used by the pooled subclass; L2-normalizes the sentence embedding (Qwen3 official recipe). + cache_text_features: cache the frozen per-token hidden states per + unique prompt string (default True). Ignored when ``freeze`` is + False - a trainable encoder must stay online. + text_cache_max_entries: cache capacity in ENTRIES; LRU eviction beyond + it. An entry is the prompt's un-padded (L, D) hidden states, so + its size follows the prompt: 4-25 KB for a 2-12 token prompt at + 1024-dim fp16, up to max_length * D * 2 bytes (262 KB at + max_length 128). The 8192 default is therefore ~100-200 MB of + device memory on real annotations and 2 GB in the worst case. """ DEFAULT_MODEL = "Qwen/Qwen3-Embedding-0.6B" @@ -984,6 +994,8 @@ def __init__( freeze: bool = True, dtype: str = "float16", output_dim: int | None = None, + cache_text_features: bool = True, + text_cache_max_entries: int = 8192, **kwargs, ) -> None: super().__init__(**kwargs) @@ -1020,8 +1032,47 @@ def __init__( if self.output_dim != self.hidden_size else nn.Identity() ) + self._init_text_feature_cache(cache_text_features, text_cache_max_entries) - def _encode(self, prompts): + # ---------------------------------------------------------------- + # frozen text-feature cache + # ---------------------------------------------------------------- + + def _init_text_feature_cache( + self, cache_text_features: bool = True, text_cache_max_entries: int = 8192 + ) -> None: + """Set up the (empty) prompt -> per-token-features cache.""" + self.cache_text_features = bool(cache_text_features) + self.text_cache_max_entries = int(text_cache_max_entries) + self._text_feature_cache = OrderedDict() + self._text_cache_hits = 0 + self._text_cache_misses = 0 + + def clear_text_feature_cache(self) -> None: + """Drop every cached prompt (the hit/miss counters are cumulative).""" + self._text_feature_cache.clear() + + def text_feature_cache_stats(self) -> dict: + """``{"entries": n, "hits": h, "misses": m}``; hits/misses count prompt + occurrences, duplicates inside one batch included.""" + return { + "entries": len(self._text_feature_cache), + "hits": self._text_cache_hits, + "misses": self._text_cache_misses, + } + + def _text_cache_enabled(self) -> bool: + # A trainable encoder must stay online: its outputs change every step + # and the loss needs the graph. Only a frozen encoder is cacheable. + return bool( + self.cache_text_features + and self.freeze_encoder + and self.text_cache_max_entries > 0 + ) + + def _encode_online(self, prompts): + """Tokenize + run the encoder, returning ``(last_hidden_state, + attention_mask)`` left-padded (the tokenizer is ``padding_side="left"``).""" tokens = self.tokenizer( prompts, padding=True, @@ -1036,6 +1087,80 @@ def _encode(self, prompts): out = self.encoder(**tokens) return out.last_hidden_state, tokens["attention_mask"] + def _encode(self, prompts): + """``(last_hidden_state, attention_mask)`` for ``prompts``, served from + the per-prompt cache whenever the encoder is frozen. + + The annotation vocabulary is tiny next to the number of training steps, + so a frozen encoder recomputes the same 600M-parameter forward over and + over. Misses go through ``_encode_online`` in ONE sub-batch of unique + strings; each prompt's un-padded ``(L_i, D)`` rows are stored (the mask + is implied by ``L_i``) and the batch is re-assembled with the same LEFT + padding the tokenizer produces, so callers cannot tell the paths apart. + + One deliberate difference: padded positions come back as ZEROS here, + whereas the online path leaves the encoder's output for the pad tokens + there. Every consumer masks those positions out + (``forward_with_mask``; last-token pooling reads the final, always-real + position under left padding), so this only removes junk. + """ + if not self._text_cache_enabled(): + return self._encode_online(prompts) + + prompts = list(prompts) + device = self.device + misses = [] + for prompt in prompts: + entry = self._text_feature_cache.get(prompt) + # Entries computed on another device are stale (the module moved): + # drop them rather than silently copy across devices. + if entry is not None and entry.device != device: + del self._text_feature_cache[prompt] + entry = None + if entry is None: + self._text_cache_misses += 1 + if prompt not in misses: + misses.append(prompt) + else: + self._text_cache_hits += 1 + self._text_feature_cache.move_to_end(prompt) + + fresh = {} + if misses: + hidden, mask = self._encode_online(misses) + mask = mask.bool() + for i, prompt in enumerate(misses): + # un-padded per-token states; detached, no graph is kept + fresh[prompt] = hidden[i][mask[i]].detach() + self._text_feature_cache[prompt] = fresh[prompt] + self._text_feature_cache.move_to_end(prompt) + while len(self._text_feature_cache) > self.text_cache_max_entries: + self._text_feature_cache.popitem(last=False) # LRU victim + + # Read this batch's rows from what was just encoded first: a batch with + # more unique prompts than text_cache_max_entries evicts its own + # earliest entries in the trim above, and a KeyError here would be a + # crash rather than a slow path. + rows = [ + fresh[p] if p in fresh else self._text_feature_cache[p] for p in prompts + ] + longest = max(row.shape[0] for row in rows) + hidden = rows[0].new_zeros((len(rows), longest, rows[0].shape[-1])) + mask = torch.zeros((len(rows), longest), dtype=torch.long, device=device) + for i, row in enumerate(rows): # LEFT padding, as padding_side="left" + length = row.shape[0] + hidden[i, longest - length :] = row + mask[i, longest - length :] = 1 + return hidden, mask + + def _apply(self, *args, **kwargs): + """``.to()`` / ``.cuda()`` / ``.float()`` invalidate the cache: the + stored tensors belong to the old device / dtype.""" + out = super()._apply(*args, **kwargs) + if getattr(self, "_text_feature_cache", None): + self.clear_text_feature_cache() + return out + def pretrained_reference_state_dict(self) -> Optional[dict]: """The snapshot's own safetensors, cast exactly the way ``from_pretrained(torch_dtype=...)`` casts them.""" diff --git a/tests/unit/test_text_feature_cache.py b/tests/unit/test_text_feature_cache.py new file mode 100644 index 0000000000..db2678bf10 --- /dev/null +++ b/tests/unit/test_text_feature_cache.py @@ -0,0 +1,421 @@ +"""Lane C step 2: the frozen-text-feature cache on the Qwen stems. + +The annotation vocabulary is tiny compared with the number of training steps, +so a frozen text encoder recomputes the same ~600M-parameter forward over and +over. ``_Qwen3BaseEncoder`` now keeps a lazy in-memory cache of the per-token +hidden states, keyed by the exact prompt string, and re-assembles the batch +(left-padded, exactly like the tokenizer) from the cached rows. + +The cache is only legal when the encoder is frozen; a trainable encoder must +stay online (its outputs change every step and need a graph). +""" + +from __future__ import annotations + +import glob +import os +import time + +import pytest +import torch +import torch.nn as nn +from omegaconf import OmegaConf + +from egomimic.models.hpt_nets import PolicyStem, QwenPerTokenEncoder + +EMBED_DIM = 32 +HIDDEN_DIM = 16 + + +def _cross_attn_specs(latent: int = 4) -> OmegaConf: + return OmegaConf.create( + { + "random_horizon_masking": False, + "cross_attn": { + "crossattn_latent": latent, + "crossattn_heads": 2, + "crossattn_dim_head": 8, + "crossattn_modality_dropout": 0.0, + "modality_embed_dim": EMBED_DIM, + }, + } + ) + + +# -------------------------------------------------------------------------- +# a stub encoder: deterministic per-prompt features, counts its forwards +# -------------------------------------------------------------------------- + +PROMPT_LENGTHS = {"fold the towel": 2, "pick up the cup": 3, "wipe the table": 5} + + +class CountingStubEncoder(QwenPerTokenEncoder): + """``QwenPerTokenEncoder`` with the tokenizer + HF model replaced by a fixed + prompt -> per-token-features table, so the cache can be tested without + loading Qwen. ``_encode_online`` is what the cache is supposed to avoid; + every call is counted, along with how many prompts it saw.""" + + def __init__( + self, + specs, + freeze: bool = True, + cache_text_features: bool = True, + text_cache_max_entries: int = 8192, + lengths: dict | None = None, + ) -> None: + PolicyStem.__init__(self, specs=specs) + self.hidden_size = HIDDEN_DIM + self.output_dim = EMBED_DIM + self.freeze_encoder = freeze + self._snapshot_dir = "" + self._load_dtype = torch.float32 + self.proj = nn.Linear(HIDDEN_DIM, EMBED_DIM) + self.encoder = nn.Identity() # never called; `train()` keeps it in eval + self._init_text_feature_cache(cache_text_features, text_cache_max_entries) + + generator = torch.Generator().manual_seed(0) + self._table = { + prompt: torch.randn(length, HIDDEN_DIM, generator=generator) + for prompt, length in (lengths or PROMPT_LENGTHS).items() + } + self.calls = [] # one entry per _encode_online call: the prompt list + + def _encode_online(self, prompts): + self.calls.append(list(prompts)) + rows = [self._table[p] for p in prompts] + longest = max(row.shape[0] for row in rows) + hidden = torch.zeros(len(rows), longest, HIDDEN_DIM) + mask = torch.zeros(len(rows), longest, dtype=torch.long) + for i, row in enumerate(rows): # left padding, as padding_side="left" + hidden[i, longest - row.shape[0] :] = row + mask[i, longest - row.shape[0] :] = 1 + return hidden, mask + + +def _stub(**kwargs) -> CountingStubEncoder: + specs = _cross_attn_specs() + stem = CountingStubEncoder(specs, **kwargs) + stem.init_cross_attn(specs.cross_attn) + stem.eval() + return stem + + +A, B, C = list(PROMPT_LENGTHS) + + +# -------------------------------------------------------------------------- +# 1. misses are encoded once, in one sub-batch of unique strings +# -------------------------------------------------------------------------- + + +def test_cache_encodes_each_unique_prompt_exactly_once(): + stem = _stub() + with torch.no_grad(): + stem._encode([A, B]) + stem._encode([B, C, A]) + stem._encode([C]) + stem._encode([A, A, B, C]) + + # one call per batch that had misses, each carrying only the unique misses + assert stem.calls == [[A, B], [C]], stem.calls + stats = stem.text_feature_cache_stats() + assert stats["entries"] == 3 + assert stats["misses"] == 3 + assert stats["hits"] == 2 + 3 + 1 + 4 - 3 # every prompt seen, minus misses + + +def test_cached_output_equals_the_uncached_output(): + """Bit-identical, since the stub features do not depend on the batch.""" + batches = [[A, B], [B, C, A], [C], [A, A, C]] + cached, online = _stub(), _stub(cache_text_features=False) + online.load_state_dict(cached.state_dict()) # same proj / cross-attn weights + with torch.no_grad(): + for prompts in batches: + hid_c, mask_c = cached._encode(prompts) + hid_o, mask_o = online._encode(prompts) + assert hid_c.shape == hid_o.shape + assert torch.equal(mask_c.bool(), mask_o.bool()) + assert torch.equal(hid_c, hid_o), prompts + # and through the whole stem, cross-attention included + assert torch.equal( + cached.compute_latent(prompts), online.compute_latent(prompts) + ) + # the uncached stem re-encodes every batch (twice here: _encode + + # compute_latent), the cached one only ever saw the misses + assert online.calls == [b for b in batches for _ in range(2)] + assert cached.calls == [[A, B], [C]] + + +# -------------------------------------------------------------------------- +# 2. the re-assembled batch keeps the tokenizer's left padding +# -------------------------------------------------------------------------- + + +def test_reassembled_batch_is_left_padded(): + stem = _stub() + prompts = [C, A, B] # lengths 5, 2, 3 + with torch.no_grad(): + stem._encode([A]) # warm one entry so the batch mixes hits and misses + hidden, mask = stem._encode(prompts) + + lengths = [PROMPT_LENGTHS[p] for p in prompts] + longest = max(lengths) + assert hidden.shape == (len(prompts), longest, HIDDEN_DIM) + assert mask.shape == (len(prompts), longest) + for i, length in enumerate(lengths): + pad = longest - length + assert mask[i, :pad].sum() == 0, "padding must be on the LEFT" + assert mask[i, pad:].min() == 1 + assert torch.equal(hidden[i, pad:], stem._table[prompts[i]]) + assert torch.equal(hidden[i, :pad], torch.zeros(pad, HIDDEN_DIM)) + + +def test_cached_tensors_are_detached(): + stem = _stub(freeze=True) + with torch.no_grad(): + hidden, _ = stem._encode([A, B]) + assert not hidden.requires_grad + for entry in stem._text_feature_cache.values(): + assert not entry.requires_grad + assert entry.grad_fn is None + + +# -------------------------------------------------------------------------- +# 3. bypasses +# -------------------------------------------------------------------------- + + +def test_trainable_encoder_bypasses_the_cache(): + stem = _stub(freeze=False) + stem._encode([A, B]) + stem._encode([A, B]) + assert stem.calls == [[A, B], [A, B]] + assert stem.text_feature_cache_stats()["entries"] == 0 + + +def test_cache_text_features_false_bypasses_the_cache(): + stem = _stub(cache_text_features=False) + with torch.no_grad(): + stem._encode([A, B]) + stem._encode([A, B]) + assert stem.calls == [[A, B], [A, B]] + assert stem.text_feature_cache_stats()["entries"] == 0 + + +def test_clear_text_feature_cache(): + stem = _stub() + with torch.no_grad(): + stem._encode([A, B]) + stem.clear_text_feature_cache() + stem._encode([A, B]) + assert stem.calls == [[A, B], [A, B]] + assert stem.text_feature_cache_stats() == {"entries": 2, "hits": 0, "misses": 4} + + +# -------------------------------------------------------------------------- +# 4. LRU eviction +# -------------------------------------------------------------------------- + + +def test_lru_eviction_at_max_entries(): + stem = _stub(text_cache_max_entries=2) + with torch.no_grad(): + stem._encode([A, B]) # cache: A, B + stem._encode([A]) # hit; A becomes most-recent, B is the LRU victim + stem._encode([C]) # evicts B + assert set(stem._text_feature_cache) == {A, C} + stem._encode([A, C]) # both still cached: no new forward + assert stem.calls == [[A, B], [C]] + stem._encode([B]) # B was evicted -> re-encoded + assert stem.calls == [[A, B], [C], [B]] + assert stem.text_feature_cache_stats()["entries"] == 2 + + +def test_cache_never_exceeds_max_entries(): + lengths = {f"prompt {i}": 1 + (i % 4) for i in range(20)} + stem = _stub(text_cache_max_entries=5, lengths=lengths) + with torch.no_grad(): + for prompt in lengths: + stem._encode([prompt]) + assert len(stem._text_feature_cache) <= 5 + + +# -------------------------------------------------------------------------- +# 5. the real Qwen encoder (skipped when the snapshot is absent) +# -------------------------------------------------------------------------- + + +def _qwen_snapshot() -> str | None: + cache = os.environ.get("HF_HUB_CACHE") or os.path.join( + os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface")), "hub" + ) + hits = sorted( + glob.glob( + os.path.join(cache, "models--Qwen--Qwen3-Embedding-0.6B", "snapshots", "*") + ) + ) + return hits[-1] if hits else None + + +QWEN_SNAPSHOT = _qwen_snapshot() +requires_qwen = pytest.mark.skipif( + QWEN_SNAPSHOT is None, reason="Qwen3-Embedding-0.6B snapshot not in the HF cache" +) + +QWEN_PROMPTS = [ + "fold the towel", + "pick up the red cup and place it on the blue plate", + "wipe", + "open the drawer and take out the blue sponge, then close the drawer", +] + + +def _build_qwen_stems(dtype: str = "float16"): + """Two stems off the same snapshot: one cached, one always online.""" + if QWEN_SNAPSHOT is None: + pytest.skip("Qwen3-Embedding-0.6B snapshot not in the HF cache") + os.environ.setdefault("HF_HUB_OFFLINE", "1") + specs = _cross_attn_specs() + built = [] + for cache_text_features in (True, False): + stem = QwenPerTokenEncoder( + model_name=QWEN_SNAPSHOT, + max_length=32, + output_dim=EMBED_DIM, + dtype=dtype, + cache_text_features=cache_text_features, + specs=specs, + ) + stem.init_cross_attn(specs.cross_attn) + stem.eval() + built.append(stem) + # identical cross-attention / projection weights on both stems + built[1].load_state_dict(built[0].state_dict()) + return tuple(built) + + +@pytest.fixture(scope="module") +def qwen_stems(): + return _build_qwen_stems("float16") + + +@pytest.fixture(scope="module") +def qwen_stems_fp32(): + return _build_qwen_stems("float32") + + +@requires_qwen +def test_real_qwen_cached_matches_online(qwen_stems): + cached, online = qwen_stems + forwards = [] + original = type(cached)._encode_online + + def counting(self, prompts): + forwards.append(list(prompts)) + return original(self, prompts) + + with torch.no_grad(): + hid_on, mask_on = online._encode(QWEN_PROMPTS) + lat_on = online.compute_latent(QWEN_PROMPTS) + + # Warm one prompt on its own first, so the batch below is a mix of a + # hit (encoded with NO padding) and misses (encoded padded to each + # other) - the case where cached and online padding really differ. + cached.clear_text_feature_cache() + cached._encode([QWEN_PROMPTS[0]]) + cached._encode_online = counting.__get__(cached, type(cached)) + hid_c1, mask_c1 = cached._encode(QWEN_PROMPTS) + lat_c1 = cached.compute_latent(QWEN_PROMPTS) + first_calls = list(forwards) + n_first = len(forwards) + forwards.clear() + lat_c2 = cached.compute_latent(QWEN_PROMPTS) # fully warm + + assert n_first == 1, "the misses must go through the encoder in ONE sub-batch" + assert first_calls[0] == QWEN_PROMPTS[1:], "only the misses, unique, in order" + assert forwards == [], "a warm cache must perform zero encoder forwards" + assert torch.equal(mask_c1.bool(), mask_on.bool()) + assert hid_c1.shape == hid_on.shape + assert torch.equal(lat_c1, lat_c2), "the cache must be deterministic once warm" + + # Only the REAL token positions are comparable: the cache pads with zeros + # while the online path leaves the encoder's output for the pad tokens + # there. Both are masked out downstream (forward_with_mask), so this is the + # tensor the rest of the stem actually reads. + real = mask_on.bool() + hidden_diff = (hid_c1[real].float() - hid_on[real].float()).abs().max().item() + latent_diff = (lat_c1 - lat_on).abs().max().item() + print( + f"\n[qwen-cache] max|cached - online| hidden {hidden_diff:.3e} " + f"latent {latent_diff:.3e} (entries {cached.text_feature_cache_stats()})" + ) + # The cache encodes a miss with different left padding than the online + # batch would, and Qwen's RoPE positions are absolute, so fp16 rounding + # differs. The right yardstick is how much the ONLINE path itself moves + # under the same change of padding: encode one prompt alone and compare it + # with its row in the padded batch. + with torch.no_grad(): + solo = online._encode([QWEN_PROMPTS[1]])[0][0].float() + control = (solo - hid_on[1][mask_on[1].bool()].float()).abs().max().item() + print(f"[qwen-cache] online-vs-online padding noise (fp16): {control:.3e}") + assert hidden_diff <= max(2 * control, 1e-6), (hidden_diff, control) + # absolute ceilings: hidden states run to |x| ~ 27 here, so fp16's ~1e-3 + # relative precision is worth a few 1e-2 in absolute terms. + assert hidden_diff < 1e-1, hidden_diff + assert latent_diff < 5e-3, latent_diff + + +@requires_qwen +def test_real_qwen_cached_matches_online_in_fp32(qwen_stems_fp32): + """The same comparison without fp16 rounding: the cache is then exact to + within ordinary float noise (the brief's 1e-2 / 1e-3 bounds).""" + cached, online = qwen_stems_fp32 + with torch.no_grad(): + hid_on, mask_on = online._encode(QWEN_PROMPTS) + lat_on = online.compute_latent(QWEN_PROMPTS) + cached.clear_text_feature_cache() + cached._encode([QWEN_PROMPTS[0]]) # mixed hit / miss padding again + hid_c, _ = cached._encode(QWEN_PROMPTS) + lat_c = cached.compute_latent(QWEN_PROMPTS) + real = mask_on.bool() + hidden_diff = (hid_c[real] - hid_on[real]).abs().max().item() + latent_diff = (lat_c - lat_on).abs().max().item() + print( + f"\n[qwen-cache-fp32] max|cached - online| hidden {hidden_diff:.3e} " + f"latent {latent_diff:.3e}" + ) + assert hidden_diff < 1e-2, hidden_diff + assert latent_diff < 1e-3, latent_diff + + +@requires_qwen +def test_real_qwen_cache_is_faster_when_warm(qwen_stems): + cached, online = qwen_stems + cached.clear_text_feature_cache() + with torch.no_grad(): + cached.compute_latent(QWEN_PROMPTS) # warm + start = time.perf_counter() + cached.compute_latent(QWEN_PROMPTS) + warm = time.perf_counter() - start + start = time.perf_counter() + online.compute_latent(QWEN_PROMPTS) + live = time.perf_counter() - start + print(f"\n[qwen-cache] warm {warm * 1e3:.1f} ms vs online {live * 1e3:.1f} ms") + assert warm < live + + +def test_a_batch_wider_than_the_cache_still_serves_every_prompt(): + """The LRU trim runs after the miss pass, so a batch with more unique + prompts than the capacity evicts its own earliest entries; reading them + back must not KeyError.""" + stem = _stub(text_cache_max_entries=2) + with torch.no_grad(): + hidden, mask = stem._encode([A, B, C]) + online_hidden, online_mask = stem._encode_online([A, B, C]) + assert hidden.shape[0] == 3 and mask.shape[0] == 3 + assert len(stem._text_feature_cache) <= 2 + for i in range(3): + real, real_on = mask[i].bool(), online_mask[i].bool() + torch.testing.assert_close( + hidden[i][real], online_hidden[i][real_on], atol=0, rtol=0 + ) From e593181fa9d8d096f54cdaba59e1de028b80cf06 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Wed, 16 Sep 2026 12:25:58 -0400 Subject: [PATCH 21/79] hpt: build the policy on CPU and let Lightning place it HPT.__init__ defaulted self.device to torch.device('cuda') and moved the freshly built policy there. Under submitit every rank sees all 8 GPUs, so all 8 ranks resolved that to cuda:0 and built their policy on the same device: 8 x 3.34B, which OOMs during construction. The 1B arm fit and survived, which made the failure look config-specific. PI already builds on CPU and leaves self.device unset for Lightning to fill in through ModelWrapper.on_fit_start / on_validation_start. HPT now does the same: no device is chosen in __init__ and nothing leaves CPU there. Lightning's write then has to reach the HPTModel the algo holds. That model reads self.device to build the loss accumulators in compute_loss, while HPT.compute_losses accumulates on the algo's, so a write that stopped at the algo would leave the two on different devices -- 'Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!'. HPT.device is therefore a property whose setter also writes nets['policy'].device, so the two cannot diverge through the path the trainer uses. Verified on 8xH100 with the 3.34B hpt_3b arch-comparison arm: every rank reports CUDA_VISIBLE_DEVICES [0..7], all 8 reach fit start and train an epoch without an OOM. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01EAGuZKPVfr2J9L7W9gtxET --- egomimic/algo/hpt.py | 29 ++++++++++++---- tests/unit/test_device_propagation.py | 50 +++++++++++++++++++++++++++ tests/unit/test_train_step.py | 45 ++++++++++++++++++++++++ 3 files changed, 118 insertions(+), 6 deletions(-) create mode 100644 tests/unit/test_device_propagation.py diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index 6467fc1b05..9c143f37b8 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -2,7 +2,7 @@ import random from collections import OrderedDict from functools import partial -from typing import Literal +from typing import Literal, Union import einops import numpy as np @@ -132,7 +132,7 @@ def __init__( self.auxiliary_ac_keys = None self.shared_action = False - self.device = None + self.device: Union[torch.device, str, None] = None self.ot_6dof = False self.use_dtw = False @@ -829,6 +829,25 @@ def load_pretrained(self, checkpoint_path): class HPT(Algo): """ """ + @property + def device(self): + return self._device + + @device.setter + def device(self, value): + """Set the algo's device, and the policy's with it. + + The policy is built on CPU and moved by Lightning, which writes the + real device here (``ModelWrapper.on_fit_start`` / + ``on_validation_start``). ``HPTModel`` reads it to create the loss + accumulators in ``compute_loss``, so a write that stopped at the algo + would leave those on a different device than the parameters. + """ + self._device = value + nets = getattr(self, "nets", None) + if nets is not None and "policy" in nets: + nets["policy"].device = value + def __init__( self, norm_stats, @@ -910,9 +929,7 @@ def __init__( self.eval_base_seed = int(kwargs.get("eval_base_seed", EVAL_BASE_SEED)) self.multitask = kwargs.get("multitask", False) - self.device = kwargs.get( - "device", torch.device("cuda" if torch.cuda.is_available() else "cpu") - ) + self.device = None model.device = self.device self.diffusion = kwargs.get("diffusion", False) @@ -1004,7 +1021,7 @@ def __init__( model.finalize_modules() self.nets["policy"] = model - self.nets = self.nets.float().to(self.device) + self.nets = self.nets.float() # Every pretrained submodule must still hold its checkpoint's weights # once the model is fully built, moved and upcast. Raises on mismatch. diff --git a/tests/unit/test_device_propagation.py b/tests/unit/test_device_propagation.py new file mode 100644 index 0000000000..3103f70e76 --- /dev/null +++ b/tests/unit/test_device_propagation.py @@ -0,0 +1,50 @@ +"""The device the trainer actually runs on has to reach ``HPTModel``. + +``HPT`` builds its policy on CPU and leaves both devices unset; Lightning +learns the real one and ``ModelWrapper.on_fit_start`` / ``on_validation_start`` +write it to the algo. If that write stops at the algo, +``HPTModel.compute_loss`` builds its loss accumulators on the unset device +while ``HPT.compute_losses`` accumulates on the trainer's, and the run dies +with "Expected all tensors to be on the same device". +""" + +from __future__ import annotations + +import torch + +from egomimic.algo.hpt import HPT, HPTModel + + +def _algo_with_policy(construction_device: torch.device): + """An algo holding a policy, without running either ``__init__``. + + Building a real HPT means a full trunk, stems and heads; only the device + wiring is under test here, so both objects are created bare, exactly as + ``test_infra_port_wiring`` does for the stem registration. + """ + algo = HPT.__new__(HPT) + algo.nets = torch.nn.ModuleDict() + algo.device = construction_device + policy = HPTModel.__new__(HPTModel) + torch.nn.Module.__init__(policy) + policy.device = construction_device + algo.nets["policy"] = policy + return algo, policy + + +def test_assigning_the_algo_device_reaches_the_policy(): + algo, policy = _algo_with_policy(torch.device("cuda")) + + algo.device = torch.device("cpu") # what ModelWrapper.on_fit_start does + + assert policy.device == torch.device("cpu") + + +def test_algo_without_a_policy_still_takes_a_device(): + """The algo sets its device in ``__init__`` before the policy is built.""" + algo = HPT.__new__(HPT) + algo.nets = torch.nn.ModuleDict() + + algo.device = torch.device("cpu") + + assert algo.device == torch.device("cpu") diff --git a/tests/unit/test_train_step.py b/tests/unit/test_train_step.py index 0094544868..f4de8fde9a 100644 --- a/tests/unit/test_train_step.py +++ b/tests/unit/test_train_step.py @@ -9,6 +9,7 @@ import os import pytest +import torch from fixtures.recipes import ( RECIPES, common_overrides, @@ -91,6 +92,50 @@ def test_hpt_train_step(tmp_path, monkeypatch, vendor): ), f"{vendor}/hpt missing {expected - seen}; saw {sorted(seen)}" +def test_hpt_builds_the_policy_on_cpu_while_gpus_are_visible(tmp_path, monkeypatch): + """Under submitit every rank sees all 8 GPUs, so an algo that picks its own + construction device puts all 8 policies on cuda:0 and OOMs at 3B. Lightning + owns the move, so nothing may leave CPU in ``__init__``.""" + hermetic_env(monkeypatch) + recipe = RECIPES[("aria", "hpt")] + data, out, hashes = write_fixtures(tmp_path, "aria") + cfg = compose_recipe( + recipe, + common_overrides( + recipe.embodiment, + data, + out, + batch_size=2, + num_workers=0, + episode_hashes=hashes, + ) + + cpu_trainer_overrides(STEPS) + + hpt_small_overrides(recipe.embodiment), + out, + ) + + built = {} + init = HPT.__init__ + + def spy(self, *args, **kwargs): + # Only __init__ sees the GPUs: faking them for the whole run would send + # the rest of the stack down CUDA paths this CPU tier cannot take. + with monkeypatch.context() as gpus_visible: + gpus_visible.setattr(torch.cuda, "is_available", lambda: True) + gpus_visible.setattr(torch.cuda, "device_count", lambda: 8) + init(self, *args, **kwargs) + built["param_devices"] = {p.device.type for p in self.nets.parameters()} + built["algo_device"] = self.device + built["policy_device"] = self.nets["policy"].device + + monkeypatch.setattr(HPT, "__init__", spy) + _run(cfg) + + assert built["param_devices"] == {"cpu"} + assert built["algo_device"] is None + assert built["policy_device"] is None + + @pytest.mark.parametrize("vendor", VENDOR_NAMES) def test_pi_train_step(tmp_path, monkeypatch, vendor): reason = pi_unavailable() From 3c06f5f550a68b3ee933ad21d329844cbc218115 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 09:14:17 -0400 Subject: [PATCH 22/79] algo: DeterministicEvalMixin, the seeded eval helpers out of hpt.py The eval-time seeding (eval_base_seed + rank * EVAL_RANK_STRIDE + pass counter, the prompt sampling and the seeded generator) was 96 lines inside HPT, though nothing in it is HPT-specific: it reads only nets, device, eval_base_seed, annotation_key, annotation_sampling_mode and default_prompt. Moving it to egomimic/algo/eval_determinism.py as a mixin leaves hpt.py with the algo and gives the seeding contract one documented home. The attributes it reads are annotated, not defaulted: a class-level default would turn "the algo forgot to set this" from an AttributeError into a silent wrong value. The module docstring says plainly that this is HPT-only today -- PI mixes in nothing and has no eval seeding at all, so pi0.5 validation is still not reproducible -- and what wiring PI up would take. Pure code motion -- no behaviour change, the eight method bodies are byte-identical to the originals, and the existing determinism tests cover it unchanged. Co-Authored-By: Claude Fable 5.1 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/algo/eval_determinism.py | 124 ++++++++++++++++++++++++++++++ egomimic/algo/hpt.py | 105 ++----------------------- 2 files changed, 130 insertions(+), 99 deletions(-) create mode 100644 egomimic/algo/eval_determinism.py diff --git a/egomimic/algo/eval_determinism.py b/egomimic/algo/eval_determinism.py new file mode 100644 index 0000000000..46c9dc6437 --- /dev/null +++ b/egomimic/algo/eval_determinism.py @@ -0,0 +1,124 @@ +"""Deterministic evaluation. Mixed into HPT; nothing else uses it yet. + +Every RNG used while an algo is in eval mode is seeded from +``self.eval_base_seed + rank * EVAL_RANK_STRIDE + pass_counter``, where the +pass counter counts ``forward_eval`` calls since the last +``reset_eval_pass_counter()`` (called from ``ModelWrapper.on_validation_start``). +Two validation passes over the same batches therefore replay the same prompts +and the same sampling noise, while different ranks and different batches +within a pass still see different draws. + +PI does NOT mix this in and has no eval seeding of its own, so pi0.5 +validation still replays different prompts and different sampling noise every +pass. Seven of these eight methods are algo-agnostic and PI has everything +they read, so wiring it up is a small follow-up; only ``_build_prompts`` is +HPT-shaped (PI's takes an embodiment and reads ``self.sampling_mode``). + +The mixin reads ``self.nets`` (an ``nn.ModuleDict``, for ``.training``), +``self.device``, ``self.eval_base_seed``, ``self.annotation_key``, +``self.annotation_sampling_mode`` and ``self.default_prompt`` from the algo it +is mixed into; the algo is responsible for setting them. +""" + +from __future__ import annotations + +import os +import random + +import torch + +EVAL_BASE_SEED = 0 +EVAL_RANK_STRIDE = 1_000_003 +# Offset of the val-loss RNG from the sampling generator's seed. It must not be +# a small number: the pass counter advances by 1 per forward_eval call, and a +# torch.Generator seeded with S draws the same stream as a global manual_seed +# of S, so an offset of 1 gave batch i's val loss the noise batch i+1 sampled. +EVAL_LOSS_STRIDE = 7_919_369 + + +class DeterministicEvalMixin: + # Annotated, not defaulted: an algo that forgets to set one of these should + # raise rather than silently take an invented default. + eval_base_seed: int + annotation_key: str | None + annotation_sampling_mode: str + default_prompt: str + + def reset_eval_pass_counter(self) -> None: + """Start a new validation pass: the next ``forward_eval`` call is pass 0. + + Called from ``ModelWrapper.on_validation_start``. Two passes over the + same loader (same order, same number of batches) then draw the same + prompts and the same sampling noise. + """ + self._eval_pass_counter = 0 + + @staticmethod + def _eval_rank() -> int: + """Distributed rank, or 0 outside a process group.""" + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_rank() + for var in ("RANK", "LOCAL_RANK", "SLURM_PROCID"): + value = os.environ.get(var) + if value is not None and value.lstrip("-").isdigit(): + return int(value) + return 0 + + def _eval_pass_seed(self) -> int: + """Seed for the current eval pass: base + rank stride + pass counter.""" + return ( + self.eval_base_seed + + self._eval_rank() * EVAL_RANK_STRIDE + + int(getattr(self, "_eval_pass_counter", 0)) + ) + + def _eval_device(self) -> torch.device: + return ( + torch.device(self.device) + if self.device is not None + else torch.device("cpu") + ) + + def _eval_generator(self) -> torch.Generator: + """Seeded generator on the device the sampling noise is drawn on.""" + device = self._eval_device() + generator = torch.Generator(device=device) + generator.manual_seed(self._eval_pass_seed()) + return generator + + def _eval_prompt_rng(self) -> random.Random: + """Seeded Python RNG for annotation choice during eval.""" + return random.Random(self._eval_pass_seed()) + + def _eval_fork_devices(self) -> list: + """Devices ``torch.random.fork_rng`` must save/restore (CPU is implicit).""" + device = self._eval_device() + if device.type != "cuda": + return [] + return [ + device.index if device.index is not None else torch.cuda.current_device() + ] + + def _build_prompts(self, _batch, batch_size: int) -> list[str]: + """Sample one annotation per batch item, falling back to default_prompt + on empty / missing annotations. Mirrors the Pi algo flow. + + In eval the choice comes from a per-pass seeded RNG regardless of + ``annotation_sampling_mode``: that keeps prompt variety (unlike always + taking ``sample[0]``) while making two validation passes over the same + batches produce identical prompt lists. Training is unchanged. + """ + if self.annotation_key is None or self.annotation_key not in _batch: + return [self.default_prompt] * batch_size + eval_rng = None if self.nets.training else self._eval_prompt_rng() + prompts = [] + for sample in _batch[self.annotation_key]: + if not sample: + prompts.append(self.default_prompt) + elif eval_rng is not None: + prompts.append(sample[eval_rng.randint(0, len(sample) - 1)]) + elif self.annotation_sampling_mode == "random": + prompts.append(sample[random.randint(0, len(sample) - 1)]) + else: # "first" + prompts.append(sample[0]) + return prompts diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index 9c143f37b8..8d0ffcb428 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -1,5 +1,4 @@ import os -import random from collections import OrderedDict from functools import partial from typing import Literal, Union @@ -14,6 +13,11 @@ from tslearn.metrics import SoftDTWLossPyTorch from egomimic.algo.algo import Algo +from egomimic.algo.eval_determinism import ( + EVAL_BASE_SEED, + EVAL_LOSS_STRIDE, + DeterministicEvalMixin, +) from egomimic.models.hpt_nets import ( MultiheadAttention, SimpleTransformer, @@ -29,21 +33,6 @@ # it has exactly one consumer, below). STD_SCALE = 0.02 -# Deterministic evaluation. Every RNG used while the algo is in eval mode is -# seeded from ``EVAL_BASE_SEED + rank * EVAL_RANK_STRIDE + pass_counter``, where -# the pass counter counts ``forward_eval`` calls since the last -# ``reset_eval_pass_counter()`` (called from ``ModelWrapper.on_validation_start``). -# Two validation passes over the same batches therefore replay the same prompts -# and the same sampling noise, while different ranks and different batches -# within a pass still see different draws. -EVAL_BASE_SEED = 0 -EVAL_RANK_STRIDE = 1_000_003 -# Offset of the val-loss RNG from the sampling generator's seed. It must not be -# a small number: the pass counter advances by 1 per forward_eval call, and a -# torch.Generator seeded with S draws the same stream as a global manual_seed -# of S, so an offset of 1 gave batch i's val loss the noise batch i+1 sampled. -EVAL_LOSS_STRIDE = 7_919_369 - class HPTModel(nn.Module): """ @@ -826,7 +815,7 @@ def load_pretrained(self, checkpoint_path): self.load_trunk(os.path.join(checkpoint_path, "trunk.pth")) -class HPT(Algo): +class HPT(DeterministicEvalMixin, Algo): """ """ @property @@ -1035,88 +1024,6 @@ def __init__( self.training_step = 0 - # ------------------------------------------------------------------ - # Deterministic evaluation - # ------------------------------------------------------------------ - def reset_eval_pass_counter(self) -> None: - """Start a new validation pass: the next ``forward_eval`` call is pass 0. - - Called from ``ModelWrapper.on_validation_start``. Two passes over the - same loader (same order, same number of batches) then draw the same - prompts and the same sampling noise. - """ - self._eval_pass_counter = 0 - - @staticmethod - def _eval_rank() -> int: - """Distributed rank, or 0 outside a process group.""" - if torch.distributed.is_available() and torch.distributed.is_initialized(): - return torch.distributed.get_rank() - for var in ("RANK", "LOCAL_RANK", "SLURM_PROCID"): - value = os.environ.get(var) - if value is not None and value.lstrip("-").isdigit(): - return int(value) - return 0 - - def _eval_pass_seed(self) -> int: - """Seed for the current eval pass: base + rank stride + pass counter.""" - return ( - self.eval_base_seed - + self._eval_rank() * EVAL_RANK_STRIDE - + int(getattr(self, "_eval_pass_counter", 0)) - ) - - def _eval_device(self) -> torch.device: - return ( - torch.device(self.device) - if self.device is not None - else torch.device("cpu") - ) - - def _eval_generator(self) -> torch.Generator: - """Seeded generator on the device the sampling noise is drawn on.""" - device = self._eval_device() - generator = torch.Generator(device=device) - generator.manual_seed(self._eval_pass_seed()) - return generator - - def _eval_prompt_rng(self) -> random.Random: - """Seeded Python RNG for annotation choice during eval.""" - return random.Random(self._eval_pass_seed()) - - def _eval_fork_devices(self) -> list: - """Devices ``torch.random.fork_rng`` must save/restore (CPU is implicit).""" - device = self._eval_device() - if device.type != "cuda": - return [] - return [ - device.index if device.index is not None else torch.cuda.current_device() - ] - - def _build_prompts(self, _batch, batch_size: int) -> list[str]: - """Sample one annotation per batch item, falling back to default_prompt - on empty / missing annotations. Mirrors the Pi algo flow. - - In eval the choice comes from a per-pass seeded RNG regardless of - ``annotation_sampling_mode``: that keeps prompt variety (unlike always - taking ``sample[0]``) while making two validation passes over the same - batches produce identical prompt lists. Training is unchanged. - """ - if self.annotation_key is None or self.annotation_key not in _batch: - return [self.default_prompt] * batch_size - eval_rng = None if self.nets.training else self._eval_prompt_rng() - prompts = [] - for sample in _batch[self.annotation_key]: - if not sample: - prompts.append(self.default_prompt) - elif eval_rng is not None: - prompts.append(sample[eval_rng.randint(0, len(sample) - 1)]) - elif self.annotation_sampling_mode == "random": - prompts.append(sample[random.randint(0, len(sample) - 1)]) - else: # "first" - prompts.append(sample[0]) - return prompts - @override def process_batch_for_training(self, batch): """ From 9d8c5a2c61555f1cbbcf3079037902ed58bc3bf1 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 09:15:08 -0400 Subject: [PATCH 23/79] eval: overlay videos play at the source rate under DDP _video_fps divided 30 fps by world_size to undo a DistributedSampler stride that never applies: the val heads are CombinedLoaders Lightning does not shard (every rank runs every val batch; 1 rank = 223 val steps, 8 ranks = 8 x 223 in the 2026-09-16 DDP check), so rank 0 renders every frame and the 8-GPU videos came out as 8x slow motion (416 frames at 4 fps). Co-Authored-By: Claude Fable 5.1 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01THTuWdgk6M6pZ5Ya7ihZEF --- egomimic/eval/eval_video.py | 21 ++++++++++++--------- tests/unit/test_pi05_norm_rot6d.py | 14 ++++++++------ 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/egomimic/eval/eval_video.py b/egomimic/eval/eval_video.py index 9fbcc61723..49042f9a99 100644 --- a/egomimic/eval/eval_video.py +++ b/egomimic/eval/eval_video.py @@ -87,15 +87,18 @@ def compute_metrics_and_viz(self, batch, do_viz=True): raise NotImplementedError def _video_fps(self, source_fps: int = 30) -> int: - """Playback fps compensating for the DistributedSampler stride. - - Distributed validation deals the (unshuffled) val indices round-robin, - so rank 0's consecutive rendered frames are ``world_size`` source - frames apart. Writing them at the raw source fps produces a - world_size-times timelapse; scale playback down to restore wall-clock - pacing.""" - world = max(1, int(getattr(self.trainer, "world_size", 1) or 1)) - return max(1, round(source_fps / world)) + """Playback fps of the overlay video: the source rate. + + The val heads come back from ``val_dataloader()`` inside + ``CombinedLoader``s, which Lightning does NOT wrap in a + ``DistributedSampler`` (every rank runs every val batch; measured + 2026-09-16: 1 rank = 223 val steps, 8 ranks = 8 x 223), so rank 0 + renders every source frame and the video plays in real time at the + source rate. An earlier version divided by ``world_size`` to undo a + sampler stride that never applied, which made every multi-GPU video a + ``world_size``-times slow-motion (416 frames at 4 fps instead of 30). + """ + return int(source_fps) def _write_video(self, key, frames) -> None: path = os.path.join( diff --git a/tests/unit/test_pi05_norm_rot6d.py b/tests/unit/test_pi05_norm_rot6d.py index 6a1ac744d3..0b75298f2b 100644 --- a/tests/unit/test_pi05_norm_rot6d.py +++ b/tests/unit/test_pi05_norm_rot6d.py @@ -540,9 +540,10 @@ def test_fallback_cap_raises(): assert n == md.MAX_FALLBACK_ATTEMPTS -def test_video_fps_compensates_for_world_size(): - # Distributed val strides an episode's frames by world_size on each rank; - # playback fps must scale down to keep videos wall-clock real-time. +def test_video_fps_is_the_source_rate_regardless_of_world_size(): + """Val loaders are CombinedLoaders that Lightning never shards, so rank 0 + renders every frame: playback must stay at the source rate (an earlier + world_size division made 8-GPU videos 8x slow motion).""" from egomimic.eval.eval_video import EvalVideo class _Stub(EvalVideo): @@ -550,11 +551,12 @@ def compute_metrics_and_viz(self, batch, do_viz=True): raise NotImplementedError ev = _Stub.__new__(_Stub) - for world, expected in [(1, 30), (2, 15), (4, 8), (8, 4)]: + for world in (1, 2, 8): ev.trainer = SimpleNamespace(world_size=world) - assert ev._video_fps() == expected, (world, ev._video_fps()) - ev.trainer = SimpleNamespace() # no world_size attr -> assume 1 + assert ev._video_fps() == 30, (world, ev._video_fps()) + ev.trainer = SimpleNamespace(world_size=None) assert ev._video_fps() == 30 + assert ev._video_fps(source_fps=10) == 10 def test_viz_gate_follows_lightning_epoch_convention(): From 383f7989c811dcb31154052d1f7167c33a25a8dd Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 09:15:43 -0400 Subject: [PATCH 24/79] trainer: ddp_bf16_compress, DDP with a bf16 gradient all-reduce On the lab's 8 x A40 nodes -- no NVLink, two NUMA islands joined by cross-socket links -- the per-step gradient all-reduce, not the math, sets the step time once a model is large enough. Measured 2026-09-16 on a ~1.2 B -parameter policy: 6.4 s/step at 8 x batch 16 against 1.4 s/step on a single GPU, and 3.6 s/step with the reduction in bf16. Only the reduction is bf16; master weights and the optimizer stay fp32. It is plain DDP otherwise -- no unused-parameter detection -- so a model that leaves a trainable tensor without a gradient needs ddp_pi's `ddp_find_unused_parameters_true` instead and will crash under this one. The config says so, because nothing else would stop you selecting it for Pi. The hook is a zero-argument factory rather than a Hydra `_partial_` because Lightning logs the hook's __qualname__, which a functools.partial does not carry. `make_bf16_compress_hook` is named for what it is: torch's own hook is `bf16_compress_hook`, and sharing that name made the factory look like the hook. test_eval_determinism's trainer parametrization now enumerates the config group instead of a hardcoded list, so a new trainer config cannot land without the gradient-clip knob being checked. Co-Authored-By: Claude Fable 5.1 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01THTuWdgk6M6pZ5Ya7ihZEF --- .../trainer/ddp_bf16_compress.yaml | 17 +++++++++++++++++ egomimic/utils/ddp_hooks.py | 12 ++++++++++++ tests/unit/test_eval_determinism.py | 12 +++++++++++- 3 files changed, 40 insertions(+), 1 deletion(-) create mode 100644 egomimic/hydra_configs/trainer/ddp_bf16_compress.yaml create mode 100644 egomimic/utils/ddp_hooks.py diff --git a/egomimic/hydra_configs/trainer/ddp_bf16_compress.yaml b/egomimic/hydra_configs/trainer/ddp_bf16_compress.yaml new file mode 100644 index 0000000000..c266e66452 --- /dev/null +++ b/egomimic/hydra_configs/trainer/ddp_bf16_compress.yaml @@ -0,0 +1,17 @@ +# Plain DDP plus a bf16 gradient all-reduce, for a model big enough that the +# gradient reduction, not the math, sets the step time. On the lab's 8 x A40 +# nodes -- no NVLink, two NUMA islands joined by cross-socket links -- a ~1.2 B +# -parameter policy measured 6.4 s/step at 8 x batch 16 against 1.4 s/step on a +# single GPU, and 3.6 s/step once the reduction was halved (2026-09-16). Only +# the reduction is bf16; master weights and the optimizer stay fp32. +# +# This is plain DDP: no unused-parameter detection. A model that leaves any +# trainable tensor without a gradient needs `ddp_pi`-style +# `ddp_find_unused_parameters_true` instead, and will crash under this one. +defaults: + - ddp + +strategy: + _target_: lightning.pytorch.strategies.DDPStrategy + ddp_comm_hook: + _target_: egomimic.utils.ddp_hooks.make_bf16_compress_hook diff --git a/egomimic/utils/ddp_hooks.py b/egomimic/utils/ddp_hooks.py new file mode 100644 index 0000000000..1d86790a77 --- /dev/null +++ b/egomimic/utils/ddp_hooks.py @@ -0,0 +1,12 @@ +"""DDP gradient-communication hooks, exposed as zero-argument factories so a +Hydra config can select one (``ddp_comm_hook: {_target_: ...}``). Lightning +logs the hook's ``__qualname__``, so the config must hand it the real function, +not a ``functools.partial`` (Hydra's ``_partial_``).""" + +from torch.distributed.algorithms.ddp_comm_hooks import default_hooks + + +def make_bf16_compress_hook(): + """All-reduce gradients in bf16 (half the bytes of fp32), decompress after. + Master weights and the optimizer stay fp32.""" + return default_hooks.bf16_compress_hook diff --git a/tests/unit/test_eval_determinism.py b/tests/unit/test_eval_determinism.py index be5bea96af..d64a8a8fa8 100644 --- a/tests/unit/test_eval_determinism.py +++ b/tests/unit/test_eval_determinism.py @@ -26,12 +26,14 @@ import random import types from collections import deque +from pathlib import Path import hydra import pytest import torch import torch.nn as nn +import egomimic.hydra_configs from egomimic.algo.hpt import EVAL_BASE_SEED, HPT, HPTModel from egomimic.models.fm_policy import FMPolicy from egomimic.pl_utils.pl_model import ModelWrapper @@ -405,7 +407,15 @@ def test_trainer_default_exposes_gradient_clip_val(compose_resolve): assert cfg.trainer.gradient_clip_algorithm == "norm" -@pytest.mark.parametrize("trainer_cfg", ["default", "ddp", "ddp_pi", "debug"]) +# Derived from the config group, not listed: a new trainer config must not be +# able to join the repo without this knob being checked. +TRAINER_CONFIGS = sorted( + p.stem + for p in (Path(egomimic.hydra_configs.__file__).parent / "trainer").glob("*.yaml") +) + + +@pytest.mark.parametrize("trainer_cfg", TRAINER_CONFIGS) def test_every_trainer_config_keeps_the_knob(trainer_cfg, compose_resolve): cfg = compose_resolve("train_zarr_cartesian", [f"trainer={trainer_cfg}"]) assert cfg.trainer.gradient_clip_val is None From 0f56bf97be8ef76050cbb62008708093f22c186a Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 09:18:40 -0400 Subject: [PATCH 25/79] configs: the flagship keypoint HPT arm, on the seen/unseen operator split data/mecka_fold_flagship_{opsplit,top3}_hpt_keypoints: the 144-D wrist-frame MANO keypoint twins of the 6D flagship configs -- keymap_mode keypoints, mode keypoints_wristframe_6d, no observations.state.ee_pose (so pad_proprio_gripper is off); everything else, including the operator split, the val loaders and the video episodes, is inherited. train_zarr_mecka_flagship_kp_hpt pairs them with hpt_bc_keypoints_wrist_300M and eval_hpt_keypoints_wrist, for both the val evaluator and the train-time overlay. hpt_bc_keypoints_wrist_300M gains the history_len / history_dropout stem knobs and backbone_lr_scale, so its lanes match hpt_bc_mecka_6d_300M's; both are inert at their defaults. The tests check one thing composition cannot: `history_len` and the keymap's `proprio_history` are a single setting living in two files, and HPTModel raises when they disagree, so each data config this model pairs with is checked against the stem that consumes it. Co-Authored-By: Claude Fable 5.1 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01THTuWdgk6M6pZ5Ya7ihZEF --- ...a_fold_flagship_opsplit_hpt_keypoints.yaml | 32 ++++++++ ...ecka_fold_flagship_top3_hpt_keypoints.yaml | 23 ++++++ .../model/hpt_bc_keypoints_wrist_300M.yaml | 9 +++ .../train_zarr_mecka_flagship_kp_hpt.yaml | 23 ++++++ tests/unit/test_flagship_kp_configs.py | 79 +++++++++++++++++++ 5 files changed, 166 insertions(+) create mode 100644 egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_keypoints.yaml create mode 100644 egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_keypoints.yaml create mode 100644 egomimic/hydra_configs/train_zarr_mecka_flagship_kp_hpt.yaml create mode 100644 tests/unit/test_flagship_kp_configs.py diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_keypoints.yaml new file mode 100644 index 0000000000..11b8820795 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_opsplit_hpt_keypoints.yaml @@ -0,0 +1,32 @@ +# mecka_fold_flagship_opsplit_hpt_6d with the wrist-frame KEYPOINT pipeline +# (the keymap / transform of mecka_fold_freeform_opsplit_hpt_keypoints): +# actions_keypoints (100, 144) wrist-first wrist-frame keypoint +# action, per hand [wrist xyz+rot6d +# (9) | 21 keypoints (63)] +# observations.state.keypoints (144,) per hand [wrist pose in HEAD +# frame | keypoints in wrist frame] +# No `observations.state.ee_pose` (include_ee_pose off), so +# pad_proprio_gripper is off. The operator split, filters, val loaders and +# video episodes are inherited unchanged. Pairs with +# model=hpt_bc_keypoints_wrist_300M and +# evaluator=eval_hpt_keypoints_wrist. Norm stats are data-specific: +# precompute them for this config (egomimic/scripts/precompute_norm_stats.py). +defaults: + - mecka_fold_flagship_opsplit_hpt_6d + - _self_ + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + # The keypoint proprio is in the WRIST frame, where a history step + # would carry the wrist's own displacement into the articulation + # channels; stay at the parent's pre-history setting until the + # per-step framing lands. + proprio_history: 1 + history_stride: 1 + transform_list: + mode: keypoints_wristframe_6d + fix_left_wrist_convention: true + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_keypoints.yaml new file mode 100644 index 0000000000..226b603e2e --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_flagship_top3_hpt_keypoints.yaml @@ -0,0 +1,23 @@ +# mecka_fold_flagship_top3_hpt_6d (the 3-operator flagship train split) with +# the wrist-frame KEYPOINT pipeline; see +# mecka_fold_flagship_opsplit_hpt_keypoints for the key layout. Norm stats: +# logs/hpt_architecture/norm_stats/mecka_fold_flagship_top3_hpt_keypoints/norm_stats +defaults: + - mecka_fold_flagship_top3_hpt_6d + - _self_ + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + # The keypoint proprio is in the WRIST frame, where a history step + # would carry the wrist's own displacement into the articulation + # channels; stay at the parent's pre-history setting until the + # per-step framing lands. + proprio_history: 1 + history_stride: 1 + transform_list: + mode: keypoints_wristframe_6d + fix_left_wrist_convention: true + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml index 1b7f4922bf..c992e14b64 100644 --- a/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_keypoints_wrist_300M.yaml @@ -126,6 +126,11 @@ robomimic_model: input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} output_dim: 840 widths: [840] + # Proprio observation history (as hpt_bc_mecka_6d_300M's state_ee_pose + # stem): K tokens, learned zero-init time-step embedding, whole-past + # dropout. Must equal the data config's `proprio_history`; inert at 1. + history_len: 1 + history_dropout: 0.2 specs: random_horizon_masking: false cross_attn: @@ -159,6 +164,10 @@ robomimic_model: mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225] +# LR multiplier for the image encoders' pretrained backbones; 1.0 keeps the +# single optimizer param group (ModelWrapper.configure_optimizers). +backbone_lr_scale: 1.0 + optimizer: _target_: torch.optim.AdamW _partial_: true diff --git a/egomimic/hydra_configs/train_zarr_mecka_flagship_kp_hpt.yaml b/egomimic/hydra_configs/train_zarr_mecka_flagship_kp_hpt.yaml new file mode 100644 index 0000000000..791868c538 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_flagship_kp_hpt.yaml @@ -0,0 +1,23 @@ +# Composed recipe: HPT 300M flow-matching BC on the 144-D wrist-frame MANO +# keypoint action, mecka FLAGSHIP fold-clothes, seen/unseen operator split. +# The keypoint counterpart of train_zarr_mecka_flagship_6d_hpt, and the +# reference arm any other architecture on this split is compared against +# (the same recipe with `override model:` swapped). +# +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_kp_hpt \ +# data=mecka_fold_flagship_top3_hpt_keypoints \ +# norm_stats.precomputed_norm_path=/norm_stats +# +# The flagship annotations are empty, so every prompt is the default one +# ("fold clothes"), so a comparison arm has to set the same string. +defaults: + - train_zarr_cartesian + - evaluator@train_viz_evaluator: train_viz_hpt_keypoints_wrist + - override model: hpt_bc_keypoints_wrist_300M + - override data: mecka_fold_flagship_opsplit_hpt_keypoints + - override evaluator: eval_hpt_keypoints_wrist + - _self_ + +model: + robomimic_model: + default_prompt: "fold clothes" diff --git a/tests/unit/test_flagship_kp_configs.py b/tests/unit/test_flagship_kp_configs.py new file mode 100644 index 0000000000..05324f0cb8 --- /dev/null +++ b/tests/unit/test_flagship_kp_configs.py @@ -0,0 +1,79 @@ +"""The flagship keypoint HPT arm: recipe, data split and stem knobs compose.""" + +from __future__ import annotations + +import pytest + +RECIPE = "train_zarr_mecka_flagship_kp_hpt" + + +def test_kp_recipe_pairs_the_keypoint_data_evaluator_and_prompt(compose_resolve): + cfg = compose_resolve(RECIPE, []) + rm = cfg.model.robomimic_model + assert rm._target_ == "egomimic.algo.hpt.HPT" + assert rm.ac_keys.human_bimanual == "actions_keypoints" + assert rm.dims.human_bimanual.action == 144 + assert rm.dims.human_bimanual.proprio == 144 + # The flagship split carries no per-frame annotations, so every sample is + # meant to see one constant prompt; a comparison arm must set the same one. + assert rm.default_prompt == "fold clothes" + assert cfg.evaluator._target_ == "egomimic.eval.eval_hpt.HPTEvalVideo" + assert ( + "keypoints_revert_6d_wristframe" + in cfg.evaluator.transform_lists.human_bimanual._target_ + ) + # The train-time overlay wraps its own evaluator under `base`, and renders + # the keypoint action through Human.viz(mode="keypoints"). + train_viz = cfg.train_viz_evaluator.base + assert train_viz._target_ == "egomimic.eval.eval_hpt.HPTEvalVideo" + assert train_viz.viz_func.human_bimanual.mode == "keypoints" + assert train_viz.viz_func.human_bimanual.action_key == "actions_keypoints" + assert ( + train_viz.transform_lists.human_bimanual._target_ + == cfg.evaluator.transform_lists.human_bimanual._target_ + ) + resolver = cfg.data.train_datasets.human_bimanual.resolver + assert resolver.key_map.keymap_mode == "keypoints" + assert resolver.transform_list.mode == "keypoints_wristframe_6d" + assert resolver.transform_list.pad_proprio_gripper is False + + +def test_kp_top3_data_keeps_the_three_operators_and_the_keypoint_pipeline( + compose_resolve, +): + cfg = compose_resolve(RECIPE, ["data=mecka_fold_flagship_top3_hpt_keypoints"]) + lambdas = cfg.data.train_datasets.human_bimanual.filters.filter_lambdas + assert any("6903686e0e94ce070afd1f24" in fn for fn in lambdas) + assert ( + cfg.data.train_datasets.human_bimanual.resolver.transform_list.mode + == "keypoints_wristframe_6d" + ) + + +def test_hpt_kp_model_carries_the_history_and_backbone_lr_knobs(compose_resolve): + cfg = compose_resolve("train_zarr_cartesian", ["model=hpt_bc_keypoints_wrist_300M"]) + stem = cfg.model.robomimic_model.stem_specs.human_bimanual.state_keypoints + assert stem.history_len == 1 + assert stem.history_dropout == 0.2 + assert cfg.model.backbone_lr_scale == 1.0 + + +# The two knobs are one setting in two files: the keymap decides how many +# proprio steps a sample carries, the stem how many tokens consume them, and +# HPTModel.stem_process raises when they disagree. Composition alone cannot see +# it, so every keypoint data config this model is documented to pair with is +# checked here. +@pytest.mark.parametrize( + "data_cfg", + [ + "mecka_fold_flagship_opsplit_hpt_keypoints", + "mecka_fold_flagship_top3_hpt_keypoints", + ], +) +def test_keymap_proprio_history_matches_the_stem_that_consumes_it( + data_cfg, compose_resolve +): + cfg = compose_resolve(RECIPE, [f"data={data_cfg}"]) + key_map = cfg.data.train_datasets.human_bimanual.resolver.key_map + stem = cfg.model.robomimic_model.stem_specs.human_bimanual.state_keypoints + assert key_map.get("proprio_history", 1) == stem.history_len From c27cd6d2dd50b72c3ea166f60adc931e97d5baab Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 09:19:32 -0400 Subject: [PATCH 26/79] viz: one keypoint renderer, the horizon trace _viz_keypoints flattened a (T, D) chunk to (T * 21, 3) and then drew slots 0..20 only -- `for k in range(21)` for the dots, and edge_ranges over the same 21 -- so an action chunk collapsed to its first timestep and the overlay never moved, however long the horizon. Every projected point past frame 0 was computed and thrown away. That is what a whole training campaign's validation videos turned out to be (2026-09-18). This replaces it with _viz_keypoints_horizon_trace, which draws the chunk as a trajectory and degenerates to per-slot start markers for a lone timestep (it promotes (D,) to (1, D) itself). The old renderer is deleted rather than left reachable: Human.viz(mode="keypoints") has one path, and a renderer you can reach by accident is one you will reach. Consequences, all intended: the dataset inspector, the Aria tutorial and mecka_keypoints_viz render through the trace now, so the single-frame callers show per-slot start markers instead of a posed skeleton, and Aria's raw 21-keypoint edge tables go with the renderer that consumed them -- they had no other caller. tests/unit/test_chunk_renderer.py pins both halves: that the trace follows the chunk, and that no second renderer is reachable from Human. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01THTuWdgk6M6pZ5Ya7ihZEF --- egomimic/rldb/embodiment/human.py | 60 ++----- .../inspector_lib/dataset_view.py | 14 +- .../scripts/tutorials/aria_to_mano_convert.py | 12 +- egomimic/utils/viz_utils.py | 167 +++++++++++------- tests/unit/test_chunk_renderer.py | 114 ++++++++++++ 5 files changed, 245 insertions(+), 122 deletions(-) create mode 100644 tests/unit/test_chunk_renderer.py diff --git a/egomimic/rldb/embodiment/human.py b/egomimic/rldb/embodiment/human.py index 57a12abf76..189b3cee1a 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -34,7 +34,7 @@ from egomimic.utils.viz_utils import ( ColorPalette, _viz_gaze, - _viz_keypoints, + _viz_keypoints_horizon_trace, ) ARIA_INTRINSICS = np.array( @@ -87,38 +87,6 @@ ) -# Aria's raw 21-keypoint layout (0-4 fingertips, 5 palm root) — NOT MANO. Used -# only for the opt-in raw-Aria-keypoint viz; the canonical keypoints are MANO. -ARIA_FINGER_EDGES = [ - (5, 6), - (6, 7), - (7, 0), # thumb - (5, 8), - (8, 9), - (9, 10), - (10, 1), # index - (5, 11), - (11, 12), - (12, 13), - (13, 2), # middle - (5, 14), - (14, 15), - (15, 16), - (16, 3), # ring - (5, 17), - (17, 18), - (18, 19), - (19, 4), # pinky -] -ARIA_FINGER_EDGE_RANGES = [ - ("thumb", 0, 3), - ("index", 3, 7), - ("middle", 7, 11), - ("ring", 11, 15), - ("pinky", 15, 19), -] - - class Human(Embodiment): """Single data-driven human embodiment shared by every vendor's human data. @@ -183,8 +151,6 @@ def viz( "traj", "traj+rotation", "axes", "annotations", "keypoints", "gaze" ], intrinsics=None, - finger_edges=None, - finger_edge_ranges=None, **kwargs, ): K = intrinsics if intrinsics is not None else cls.INTRINSICS @@ -204,22 +170,24 @@ def viz( finger: ColorPalette.to_rgb(color, value=(i + 1) / (n + 1)) for i, finger in enumerate(cls.FINGER_COLORS) } - dot_color = ColorPalette.to_rgb(color, value=0.7) else: colors = cls.FINGER_COLORS - dot_color = cls.DOT_COLOR - return _viz_keypoints( + # The horizon trace is the ONLY keypoint renderer. A chunk draws + # as its trajectory; a lone timestep degenerates to its start + # markers. There is deliberately no static-skeleton path to fall + # back to: that fallback is what silently rendered motionless + # keypoint overlays for a whole 50k-step campaign, and a renderer + # you can reach by accident is one you will reach. + return _viz_keypoints_horizon_trace( image=image, - actions=viz_data, + actions=np.asarray(viz_data), intrinsics=K, - edges=finger_edges if finger_edges is not None else cls.FINGER_EDGES, - edge_ranges=( - finger_edge_ranges - if finger_edge_ranges is not None - else cls.FINGER_EDGE_RANGES - ), colors=colors, - dot_color=dot_color, + wrist_color=( + ColorPalette.to_rgb(color, value=0.85) + if color is not None and ColorPalette.is_valid(color) + else cls.DOT_COLOR + ), **kwargs, ) return super().viz(image, viz_data, mode=mode, intrinsics=intrinsics, **kwargs) diff --git a/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py b/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py index 8fbe2a8b31..759ca7a757 100644 --- a/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py +++ b/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py @@ -429,7 +429,7 @@ def _world_keypoints_to_cam(seq, T_world_head): _2D_DOT_COLORS_BGR = { "left": (0, 120, 255), "right": (255, 80, 0), -} # mirrors _viz_keypoints +} # the inspector's own per-arm dot colours; the 2D overlay is the horizon trace _ARM_COLORS = { arm: f"rgb({bgr[2]},{bgr[1]},{bgr[0]})" # BGR -> RGB for arm, bgr in _2D_DOT_COLORS_BGR.items() @@ -491,9 +491,10 @@ def build_3d_figure(grp, frame: int, overlay: str, traj_window: int = 60): emb_cls, "FINGER_EDGE_RANGES", Human.FINGER_EDGE_RANGES if Human else [] ) # Map each edge index -> finger color (via FINGER_EDGE_RANGES), fall back gray. - # FINGER_COLORS are BGR tuples (they're handed to cv2 in the 2D skeleton draw - # in `_viz_keypoints`), so reverse BGR -> RGB to make the 3D edges render the - # SAME color cv2 paints in 2D. + # FINGER_COLORS are BGR tuples because cv2 consumes them in the 2D overlay, so + # reverse BGR -> RGB here. The 2D overlay is the horizon trace, which paints + # only the fingertips, so the colours agree at the tips and the 3D view is the + # only place a whole finger takes its finger's colour. edge_color_css = ["rgb(180,180,180)"] * len(edges) for name, lo, hi in edge_ranges: col = finger_colors.get(name) @@ -805,8 +806,9 @@ def _chunk(seq, arm): if all(np.allclose(p, 0) for p in parts): return _badge(img_rgb.copy(), "keypoint: no pts"), False, "off" # canonical keypoints (wrist_in_data=False) layout: [L 63, R 63]; - # .viz(mode="keypoints") supplies the MANO edges/colors/ - # edge_ranges. Only Human supports keypoints; if the resolved class + # .viz(mode="keypoints") draws the horizon trace: for this + # single frame that is the per-slot start markers, not a posed + # skeleton. Only Human supports keypoints; if the resolved class # doesn't, report via badge rather than reintroducing a low-level draw call. kp = np.concatenate(parts).reshape(1, -1) try: diff --git a/egomimic/scripts/tutorials/aria_to_mano_convert.py b/egomimic/scripts/tutorials/aria_to_mano_convert.py index 305aa9fbe0..28c37ddb69 100644 --- a/egomimic/scripts/tutorials/aria_to_mano_convert.py +++ b/egomimic/scripts/tutorials/aria_to_mano_convert.py @@ -9,8 +9,8 @@ 4. Permute MANO's otaheri joint order (wrist, index*3, middle*3, pinky*3, ring*3, thumb*3, then 5 tips) to the canonical MANO ordering used by Human.FINGER_EDGES (wrist, thumb*4, index*4, middle*4, ring*4, pinky*4). -5. Render side-by-side: left half = Aria viz with Aria.FINGER_EDGES; right - half = fitted MANO viz with Human (=Mecka) FINGER_EDGES. +5. Render side-by-side: raw Aria keypoints on the left, the fitted MANO on + the right, both through Human.viz(mode="keypoints"). 6. Write mp4 to scratch/aria_to_mano.mp4. Prerequisites (not checked in): @@ -37,11 +37,7 @@ import numpy as np import torch -from egomimic.rldb.embodiment.human import ( - ARIA_FINGER_EDGE_RANGES, - ARIA_FINGER_EDGES, - Human, -) +from egomimic.rldb.embodiment.human import Human from egomimic.rldb.filters import DatasetFilter from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset, S3EpisodeResolver from egomimic.scripts.aria_process.aria_utils import fit_mano_to_aria_batched @@ -189,8 +185,6 @@ def render_side_by_side(rows, mano_left_canonical, mano_right_canonical): batch, mode="keypoints", viz_batch_key="actions_keypoints", - finger_edges=ARIA_FINGER_EDGES, - finger_edge_ranges=ARIA_FINGER_EDGE_RANGES, ) left_flat = mano_left_canonical[i].reshape(-1).numpy() right_flat = mano_right_canonical[i].reshape(-1).numpy() diff --git a/egomimic/utils/viz_utils.py b/egomimic/utils/viz_utils.py index 0b2f31942e..0e3bf8c591 100644 --- a/egomimic/utils/viz_utils.py +++ b/egomimic/utils/viz_utils.py @@ -293,83 +293,128 @@ def _viz_gaze( ) -def _viz_keypoints( +# MANO slots per hand. The keypoint action layouts are +# ``[wrist xyz | wrist rot | 21 keypoints] x 2`` (rot6d 144-D, quat 140-D, +# ypr 138-D) or bare ``[21 keypoints] x 2`` (126-D). +_N_KP = 21 +_KEYPOINT_WIDTHS = {2 * (3 * _N_KP + pose) for pose in (0, 6, 7, 9)} +# Each finger runs base -> tip, so the tips are the last slot of each group of +# four (slot 0 is the wrist). +FINGERTIP_SLOTS = (4, 8, 12, 16, 20) + + +def _viz_keypoints_horizon_trace( image, actions, intrinsics, - edges, colors, - edge_ranges, - dot_color=None, + wrist_color=None, + fingertip_indices=FINGERTIP_SLOTS, + finger_names=("thumb", "index", "middle", "ring", "pinky"), **kwargs, ): - """Visualize all 21 MANO keypoints per hand, projected onto the image.""" + """Draw the horizon TRACE of the 5 fingertips + the wrist, per hand. + + One polyline per tracked slot across the whole chunk, with a filled circle + at the first drawable step and an open square at the last. The wrist comes + from the action's wrist xyz block when the layout has one (144 / 140 / 138) + and from MANO slot 0 otherwise (126-D). + + Args: + actions: ``(T, D)`` chunk, or ``(D,)`` for a single step. + colors: ``{finger_name: RGB}``, as ``Human.FINGER_COLORS``. + wrist_color: RGB for the wrist trace; ``None`` repeats Human.DOT_COLOR. + """ alpha = kwargs.get("alpha", 1.0) image = _prepare_viz_image(image) - base = image.copy() vis = base.copy() h, w = vis.shape[:2] - if actions.shape[-1] == 144: - _, _, left_keypoints, _, _, right_keypoints = _split_keypoints( - actions, wrist_in_data=True, is_rot6d=True - ) - elif actions.shape[-1] == 140: - _, _, left_keypoints, _, _, right_keypoints = _split_keypoints( - actions, wrist_in_data=True - ) - elif actions.shape[-1] == 138: - _, _, left_keypoints, _, _, right_keypoints = _split_keypoints( - actions, wrist_in_data=True, is_quat=False + actions = np.asarray(actions) + if actions.ndim == 1: + actions = actions[None] + if actions.ndim != 2 or actions.shape[-1] not in _KEYPOINT_WIDTHS: + raise ValueError( + f"keypoint trace expects (T, D) with D in {sorted(_KEYPOINT_WIDTHS)}, " + f"got {actions.shape}" ) + + if actions.shape[-1] == 2 * (3 * _N_KP + 9): + blocks = _split_keypoints(actions, wrist_in_data=True, is_rot6d=True) + elif actions.shape[-1] == 2 * (3 * _N_KP + 7): + blocks = _split_keypoints(actions, wrist_in_data=True) + elif actions.shape[-1] == 2 * (3 * _N_KP + 6): + blocks = _split_keypoints(actions, wrist_in_data=True, is_quat=False) else: - left_keypoints, right_keypoints = _split_keypoints(actions, wrist_in_data=False) - keypoints = {} - keypoints["left"] = left_keypoints.reshape(-1, 3) - keypoints["right"] = right_keypoints.reshape(-1, 3) - _default_dot_colors = {"left": (0, 120, 255), "right": (255, 80, 0)} - for hand in ("left", "right"): - hand_dot_color = ( - dot_color if dot_color is not None else _default_dot_colors[hand] + left_kps, right_kps = _split_keypoints(actions, wrist_in_data=False) + blocks = (None, None, left_kps, None, None, right_kps) + left_wrist, _, left_kps, right_wrist, _, right_kps = blocks + + left_kps = left_kps.reshape(-1, _N_KP, 3) + right_kps = right_kps.reshape(-1, _N_KP, 3) + if left_wrist is None: + left_wrist, right_wrist = left_kps[:, 0, :], right_kps[:, 0, :] + + if wrist_color is None: + # Human.DOT_COLOR's value, duplicated: human.py imports this module, + # so importing Human here would be a cycle. + wrist_color = (255, 165, 0) + + def _project_and_draw(points_cam, color, thickness=2, marker_size=4): + """``(T, 3)`` camera-frame points -> polyline + start / end markers.""" + points_cam = np.asarray(points_cam) + if points_cam.shape[0] == 0: + return + finite = np.isfinite(points_cam).all(axis=-1) & (np.abs(points_cam) < 1e8).all( + axis=-1 ) - kps_cam = keypoints[hand] - # Camera frame -> pixels - kps_px = cam_frame_to_cam_pixels(kps_cam, intrinsics) # (42, 3+) 21 per arm - - # Identify valid keypoints (z > 0 and in image bounds) - valid = kps_cam[:, 2] > 0.01 - valid &= (kps_px[:, 0] >= 0) & (kps_px[:, 0] < w) - valid &= (kps_px[:, 1] >= 0) & (kps_px[:, 1] < h) - - # Draw skeleton edges (colored by finger) - for finger, start, end in edge_ranges: - color = colors[finger] - for edge_idx in range(start, end): - i, j = edges[edge_idx] - if valid[i] and valid[j]: - p1 = (int(kps_px[i, 0]), int(kps_px[i, 1])) - p2 = (int(kps_px[j, 0]), int(kps_px[j, 1])) - cv2.line(vis, p1, p2, color, 2) - - # Draw keypoint dots on top - for k in range(21): - if valid[k]: - center = (int(kps_px[k, 0]), int(kps_px[k, 1])) - cv2.circle(vis, center, 4, hand_dot_color, -1) - cv2.circle(vis, center, 4, (255, 255, 255), 1) # white border - - # Label wrist - if valid[0]: - wrist_px = (int(kps_px[0, 0]) + 6, int(kps_px[0, 1]) - 6) - cv2.putText( + safe = np.where(finite[:, None], points_cam, [0.0, 0.0, -1.0]) + with np.errstate(divide="ignore", invalid="ignore"): + px = cam_frame_to_cam_pixels(safe, intrinsics) # (T, 3+) + valid = finite & (points_cam[:, 2] > 0.01) + valid &= (px[:, 0] >= 0) & (px[:, 0] < w) + valid &= (px[:, 1] >= 0) & (px[:, 1] < h) + pts = np.round(np.nan_to_num(px[:, :2])).astype(np.int32) + + # Connect consecutive drawable steps; a gap breaks the line rather + # than jumping across the frame. + prev = None + for t in range(pts.shape[0]): + if not valid[t]: + prev = None + continue + cur = (int(pts[t, 0]), int(pts[t, 1])) + if prev is not None: + cv2.line(vis, prev, cur, color, thickness, cv2.LINE_AA) + prev = cur + + first_t = next((t for t in range(pts.shape[0]) if valid[t]), None) + last_t = next((t for t in range(pts.shape[0] - 1, -1, -1) if valid[t]), None) + if first_t is not None: + cv2.circle( vis, - f"{hand[0].upper()}", - wrist_px, - cv2.FONT_HERSHEY_SIMPLEX, - 0.5, - hand_dot_color, - 2, + (int(pts[first_t, 0]), int(pts[first_t, 1])), + marker_size, + color, + -1, + cv2.LINE_AA, + ) + if last_t is not None and last_t != first_t: + ex, ey = int(pts[last_t, 0]), int(pts[last_t, 1]) + r = marker_size + 1 + cv2.rectangle( + vis, (ex - r, ey - r), (ex + r, ey + r), color, 2, cv2.LINE_AA + ) + + for wrist_xyz, kps in ((left_wrist, left_kps), (right_wrist, right_kps)): + _project_and_draw(wrist_xyz, wrist_color, thickness=2, marker_size=5) + for tip_idx, finger_name in zip(fingertip_indices, finger_names): + _project_and_draw( + kps[:, tip_idx, :], + colors.get(finger_name, (200, 200, 200)), + thickness=2, + marker_size=4, ) if alpha < 1.0: diff --git a/tests/unit/test_chunk_renderer.py b/tests/unit/test_chunk_renderer.py new file mode 100644 index 0000000000..411cbb3965 --- /dev/null +++ b/tests/unit/test_chunk_renderer.py @@ -0,0 +1,114 @@ +"""The horizon trace is the only keypoint renderer, and it draws every frame. + +``Human.viz(mode="keypoints")`` has exactly one path: ``_viz_keypoints_horizon_trace``. +The static-skeleton renderer it used to fall back to is gone on purpose -- a +chunk that reached it rendered as a motionless overlay, which is what a +whole 50k-step campaign's validation videos turned out to be (2026-09-18). +""" + +import numpy as np +import pytest + +from egomimic.utils.viz_utils import _viz_keypoints_horizon_trace + +# px = 100 * x / z + 100, py = 100 * y / z + 100 on a 200x200 canvas. +K = np.array([[100, 0, 100, 0], [0, 100, 100, 0], [0, 0, 1, 0]]) + +BEHIND_CAMERA = [0.0, 0.0, -1.0] + +FINGERS = ("thumb", "index", "middle", "ring", "pinky") + + +def _trace(data, **kwargs): + return _viz_keypoints_horizon_trace( + np.zeros((200, 200, 3), np.uint8), + data, + K, + colors={name: (20, 255, 20) for name in FINGERS}, + **kwargs, + ) + + +def _sweep(length): + """A left hand whose slot 0 sweeps across the image, one x per frame. + + Slot 1 sits 0.4 above it; every other slot and the whole right hand stay + behind the camera. Returns a (T, 126) chunk. + """ + points = np.tile(np.array(BEHIND_CAMERA), (length, 2, 21, 1)) + xs = np.linspace(-0.7, 0.7, length) + points[:, 0, 0] = np.column_stack((xs, np.zeros(length), np.ones(length))) + points[:, 0, 1] = points[:, 0, 0] + [0.0, 0.4, 0.0] + return points.reshape(length, -1) + + +def test_horizon_trace_connects_every_step_of_the_chunk(): + """The whole sweep is a polyline, not just its endpoints.""" + image = _trace(_sweep(30)) + for x in np.linspace(-0.7, 0.7, 30): + assert image[100, round(100 + 100 * x)].any() + # the line fills the gaps between sampled steps too + assert image[100, 100].any() + + +def test_horizon_trace_accepts_the_126d_revert_output_and_144d_rot6d(): + """The evaluator hands it 126-D; raw wrist-frame actions are 144-D.""" + assert _trace(_sweep(8)).any() # 126-D, wrist falls back to MANO slot 0 + + hand = np.concatenate([np.zeros(9), np.tile(BEHIND_CAMERA, 21)]) + chunk = np.tile(np.concatenate([hand, hand]), (8, 1)) + chunk[:, 0:3] = np.column_stack( + (np.linspace(-0.5, 0.5, 8), np.zeros(8), np.ones(8)) + ) + assert chunk.shape[-1] == 144 + assert _trace(chunk).any() # wrist xyz block drives the trace + + +def test_a_lone_timestep_matches_its_own_flattened_form(): + """``(1, D)`` and ``(D,)`` are the same single-step trace, not two paths.""" + chunk = _sweep(1) + np.testing.assert_array_equal(_trace(chunk), _trace(chunk.reshape(-1))) + + +def test_non_finite_slots_are_skipped(): + """A NaN or absurd slot drops out of the trace; the finite ones still draw.""" + chunk = _sweep(3) + chunk[1, 0:3] = np.nan # frame 1's slot 0 + chunk[2, 0:3] = 1e12 + image = _trace(chunk) + # _sweep puts frame t's slot 0 at px = 100 * x + 100, i.e. 30 / 100 / 170. + assert image[100, 30].any() # frame 0 still drawn + assert not image[100, 100].any() # frame 1 was NaN + assert not image[100, 170].any() # frame 2 was 1e12 + + +def test_horizon_trace_rejects_a_bad_shape_and_an_unknown_width(): + with pytest.raises(ValueError, match="keypoint trace expects"): + _trace(np.zeros((2, 30, 126))) + with pytest.raises(ValueError, match="keypoint trace expects"): + _trace(np.zeros((4, 127))) + + +def test_human_viz_has_no_renderer_but_the_trace(monkeypatch): + """Every keypoint input shape routes to the trace -- there is no fallback.""" + import egomimic.rldb.embodiment.human as human_mod + from egomimic.rldb.embodiment.human import Human + + assert not hasattr(human_mod, "_viz_keypoints"), ( + "the static skeleton renderer is back; mode='keypoints' must have " + "exactly one path" + ) + + seen = [] + + def _record(**kw): + seen.append(kw["actions"].shape) + return np.zeros((2, 2, 3), np.uint8) + + monkeypatch.setattr(human_mod, "_viz_keypoints_horizon_trace", _record) + image = np.zeros((200, 200, 3), np.uint8) + Human.viz(image, _sweep(8), mode="keypoints", intrinsics=K) + Human.viz(image, _sweep(1), mode="keypoints", intrinsics=K) + Human.viz(image, _sweep(1).reshape(-1), mode="keypoints", intrinsics=K) + + assert seen == [(8, 126), (1, 126), (126,)] From 635b94a32f4a6b1875c6f03fb2bb861ccc4b6b0e Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 09:20:32 -0400 Subject: [PATCH 27/79] transforms: per-step wrist framing for proprio history PoseCoordinateFrameTransform gains per_step_target: with target_history K, step k's keypoints are expressed in step k's OWN wrist frame instead of the current step's. Re-expressing older steps in the current frame smears wrist motion into the articulation channels, and the palm-rigid MANO knuckles (1/5/13/17) have a very tight legitimate range, so they land far outside their per-cell bounds and worsen with age: measured 2026-09-17, 10.3 / 6.9 / 3.3 out-of-range cells per sample at history steps 0/1/2 against 0.0 at the current step. That cost 22,500 loader retries on one run and starved rank 0 into an NCCL watchdog kill. Per-step framing keeps every step a genuine state -- keypoint 0 is 0 at every step, so the norm stats and the bounds check describe all of them. The ACTION chunk keeps the current-step frame: its leading axis is the forward horizon, not history. It is the default, and it is inert unless the target itself was read with history and the pose carries the matching axis -- a keymap without proprio_history never emits that shape, so no existing call site changes. For anything that does, per-step is not a tuning choice: under the shipped K = 3 / stride 5 the oldest step is 10 frames back, where the measured wrist translation is 80.8 mm rather than the ~17 mm that already pushed 67 of 126 channels out of band. The flag remains as an escape hatch. Both halves are tested: the origin property at K=4, and bit-identical output at K=1 either way. Nothing shipped exercises this yet: no data config pairs proprio_history with a keypoint keymap, and the flagship keypoint configs are pinned at K = 1 because they sit below this commit. It is enablement for the keypoint history arm, not a fix to a live run. The per-step loop is also a Python loop over K (measured 59.5 -> 165.4 us per call at K=3, twice per sample); worth vectorising if that arm is ever configured, not before. CartesianYPRToRot6D also picks up the note about applying it to the proprio ee_pose, which the embodiment comment used to carry. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/rldb/embodiment/human.py | 17 +++- egomimic/rldb/zarr/action_chunk_transforms.py | 56 +++++++++++-- tests/unit/test_action_chunk_transforms.py | 84 +++++++++++++++++++ 3 files changed, 147 insertions(+), 10 deletions(-) diff --git a/egomimic/rldb/embodiment/human.py b/egomimic/rldb/embodiment/human.py index 189b3cee1a..d378dfbc8e 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -483,10 +483,8 @@ def get_transform_list( ) if mode == "cartesian_6d": # Head/camera-frame cartesian (12D xyz+ypr per arm pair) with the - # rotation re-expressed as the continuous 6D representation (18D) - # for pi0.5. The proprio ee_pose is 6D-encoded too: normalized YPR - # saturates yaw/roll at +-pi, so per-dim normalization needs the - # continuous representation. + # rotation as continuous 6D (18D) for pi0.5, on the proprio ee_pose + # as well as the actions -- see CartesianYPRToRot6D. return ( prefix + _build_human_cartesian_bimanual_transform_list(stride=stride) @@ -759,6 +757,17 @@ def _build_human_keypoints_eef_frame_transform_list( output_key=right_keypoints_obs_headframe, shape=(21, 3), ), + # per_step_target: with proprio_history K, step k's keypoints go + # into step k's OWN wrist frame, not the current step's. Otherwise + # a past hand is measured from today's wrist, which smears wrist + # motion into the articulation channels -- the palm-rigid knuckles + # (MANO 1/5/13/17) have a very tight legitimate range and land far + # outside it, worsening with age (measured 2026-09-17: 10.3/6.9/3.3 + # out-of-range cells per sample at steps 0/1/2, 0.0 at the current + # step). Per-step framing keeps every step a genuine state, so + # keypoint 0 is 0 at EVERY step and the norm stats / bounds check + # describe all of them. The ACTION chunk keeps the current-step + # frame: its leading axis is the forward horizon, not history. PoseCoordinateFrameTransform( target_world=left_wrist_obs_headframe, pose_world=left_keypoints_obs_headframe, diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/zarr/action_chunk_transforms.py index 38117cad3e..9150df551b 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -398,11 +398,31 @@ def __init__( transformed_key_name: str, mode: Literal["xyzwxyz", "xyzypr", "xyz"] = "xyzwxyz", target_history: bool = False, + per_step_target: bool = True, ): + """ + args: + per_step_target: put step k's pose in step k's OWN target frame + rather than every step in the current step's frame. On by + default, and inert unless the target itself was read with + history and the pose carries the matching axis -- which is the + only shape for which the other answer is wrong, not merely + different. + + For wrist-frame keypoints this is the difference between "where + the hand was, measured from today's wrist" and "what shape the + hand was in" -- only the latter is a state the norm stats + describe, and only it keeps keypoint 0 at the origin on every + step. Pass False to force the single-frame path; the ACTION + chunk does not come through here at all, since its leading axis + is the forward horizon and belongs in one current frame. + """ self.target_world = target_world self.pose_world = pose_world self.transformed_key_name = transformed_key_name self.mode = mode + self.target_history = target_history + self.per_step_target = per_step_target self._chunk_transform = ActionChunkCoordinateFrameTransform( target_world=target_world, chunk_world=pose_world, @@ -411,17 +431,37 @@ def __init__( target_history=target_history, ) - def transform(self, batch: dict) -> dict: - pose_world = np.asarray(batch[self.pose_world]) + def _one_step(self, target_pose, pose_world) -> np.ndarray: + """``pose_world`` in the single frame ``target_pose``.""" transformed = self._chunk_transform.transform( { - self.target_world: batch[self.target_world], + self.target_world: target_pose, self.pose_world: pose_world[None, :], } ) - batch[self.transformed_key_name] = np.asarray( - transformed[self.transformed_key_name] - )[0] + return np.asarray(transformed[self.transformed_key_name])[0] + + def transform(self, batch: dict) -> dict: + pose_world = np.asarray(batch[self.pose_world]) + target_world = np.asarray(batch[self.target_world]) + per_step = ( + self.per_step_target + and self.target_history + and target_world.ndim > 1 + and pose_world.ndim == target_world.ndim + 1 + and pose_world.shape[0] == target_world.shape[0] + ) + if per_step: + batch[self.transformed_key_name] = np.stack( + [ + self._one_step(target_world[k], pose_world[k]) + for k in range(target_world.shape[0]) + ] + ) + return batch + batch[self.transformed_key_name] = self._one_step( + batch[self.target_world], pose_world + ) return batch @@ -477,6 +517,10 @@ class CartesianYPRToRot6D(Transform): convention of the pi0.5 32D action blocks, so the resulting per-arm layout packs into them with no rotation math in the model. + Apply it to the proprio ee_pose as well as the actions: normalized YPR + saturates yaw/roll at +-pi, so per-dim normalization is only meaningful on + the continuous representation. + Input layouts (last dim), for action chunks ``(T, D)`` and single proprio poses ``(D,)`` alike: 12 -> [L xyz ypr, R xyz ypr] -> 18 [L xyz 6d, R xyz 6d] diff --git a/tests/unit/test_action_chunk_transforms.py b/tests/unit/test_action_chunk_transforms.py index bafe6b4357..5779212367 100644 --- a/tests/unit/test_action_chunk_transforms.py +++ b/tests/unit/test_action_chunk_transforms.py @@ -12,6 +12,7 @@ ActionChunkCoordinateFrameTransform, ConcatKeys, InterpolatePose, + PoseCoordinateFrameTransform, XYZWXYZ_to_XYZYPR, ) from egomimic.utils.pose_utils import _xyzwxyz_to_matrix @@ -639,3 +640,86 @@ def test_human_transform_list_stepwise_keys_and_shapes() -> None: "observations.state.ee_pose": (12,), }, ) + + +# --------------------------------------------------------------------------- +# PoseCoordinateFrameTransform.per_step_target +# --------------------------------------------------------------------------- +def _wrist_history(steps: int) -> np.ndarray: + """A wrist moving +0.1 in x per step, oldest first, identity rotation.""" + poses = np.zeros((steps, 7), dtype=np.float64) + poses[:, 0] = 0.1 * np.arange(steps) + poses[:, 3] = 1.0 # w of wxyz + return poses + + +def _keypoints_on_the_wrist(wrist: np.ndarray) -> np.ndarray: + """(K, 2, 3): slot 0 sits exactly on that step's wrist, slot 1 0.4 above.""" + kp = np.zeros((wrist.shape[0], 2, 3), dtype=np.float64) + kp[:, 0, :] = wrist[:, :3] + kp[:, 1, :] = wrist[:, :3] + np.array([0.0, 0.4, 0.0]) + return kp + + +def _wristframe(wrist, keypoints, per_step_target): + tf = PoseCoordinateFrameTransform( + target_world="wrist", + pose_world="kp_head", + transformed_key_name="kp_wrist", + mode="xyz", + target_history=True, + per_step_target=per_step_target, + ) + return np.asarray(tf.transform({"wrist": wrist, "kp_head": keypoints})["kp_wrist"]) + + +def test_per_step_target_keeps_every_history_step_in_its_own_frame(): + """The whole point: keypoint 0 is the wrist, so it must be 0 at EVERY step. + + Measured from the current step's wrist instead, an older step carries the + wrist's own displacement -- which is what pushed palm-rigid MANO knuckles + outside their per-cell quantile bounds and starved the loader on retries. + """ + wrist = _wrist_history(4) + keypoints = _keypoints_on_the_wrist(wrist) + + current_frame = _wristframe(wrist, keypoints, per_step_target=False) + per_step = _wristframe(wrist, keypoints, per_step_target=True) + + # Slot 0 == that step's wrist, so per-step framing sends it to the origin. + np.testing.assert_allclose(per_step[:, 0, :], 0.0, atol=1e-12) + # Without it, the three older steps keep the wrist's -0.3/-0.2/-0.1 offset + # and only the current step lands at the origin. + np.testing.assert_allclose( + current_frame[:, 0, 0], [-0.3, -0.2, -0.1, 0.0], atol=1e-12 + ) + # The articulation channel is unchanged either way: 0.4 above the wrist. + for framed in (current_frame, per_step): + np.testing.assert_allclose(framed[:, 1, 1], 0.4, atol=1e-12) + + +def test_per_step_target_is_the_default(): + """Anything read with history through this transform wants per-step + framing; the flag is the escape hatch, not the opt-in.""" + import inspect + + from egomimic.rldb.zarr.action_chunk_transforms import ( + PoseCoordinateFrameTransform, + ) + + default = ( + inspect.signature(PoseCoordinateFrameTransform.__init__) + .parameters["per_step_target"] + .default + ) + assert default is True + + +def test_per_step_target_is_inert_without_a_history_axis(): + """Existing single-step call sites cannot change behaviour: the per-step + path needs a leading axis on the target, which a K=1 keymap never emits.""" + wrist = _wrist_history(1)[0] # (7,), the shape a no-history keymap gives + keypoints = _keypoints_on_the_wrist(_wrist_history(1))[0] # (2, 3) + off = _wristframe(wrist, keypoints, per_step_target=False) + on = _wristframe(wrist, keypoints, per_step_target=True) + np.testing.assert_array_equal(off, on) From 4098a1c4daeb960b8b26c0e6434546ecbaece2fb Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 09:21:27 -0400 Subject: [PATCH 28/79] data: the mecka freeform folding operator-scaling ladder (k = 1/2/4) Three rungs, identical in everything but the train operator set, so an arm's only variable is how many operators its data came from: k=1 9 eps / 25633 f / 0.237 h 1 operator k=2 24 eps / 60374 f / 0.559 h 2 operators (nested) k=4 48 eps / 118388 f / 1.096 h 4 operators (nested) Operators are ranked by ANNOTATED hours (not raw hours) and nested, so k=2's set contains k=1's. Shared by all three: the 3-episode seen-op val (drawn from operator #1, which every rung trains on) and the 35-episode unseen-op val. The train split is an operator FILTER, not an episode pin, so the counts above are hand-measured documentation -- test_opscale_ladder.py drives the real DatasetFilter with synthetic rows and asserts the nesting and the exclusions, which are the parts that live in the config. The metric budget is 'auto': trainHydra derives K = floor(limit_val_batches * batch_size / episodes) per head from the resolved split, so K follows the live filters instead of a hand count that goes stale when SQL grows (or when limit_val_batches or the batch size change); limit_val_batches=0, which the bench scripts use to turn validation off, subsamples nothing. Explicit ints keep their meaning, so the flagship recipes that pin K are untouched. At the 2026-09-17 counts auto reproduces the old numbers (1706 / 146 / 106, 213, 568). Each rung has its own train distribution, so norm stats are per-arm; the 6D and keypoint variants of a rung share filters, loader params and budgets and differ only in keymap and transform. The ladder itself is edited in one place, mecka_fold_freeform_opscale4_6d, which the other five inherit from. No proprio_history: it is half a setting (the model's stem history_len is the other half, and HPTModel raises when they disagree), and every model in the repo is history_len 1, so a rung asking for more is a rung nothing can train. A test pins that, since composition cannot see it. The annotation guard is written once in the parent (a YAML anchor covers its three splits). mecka_fold_freeform_opsplit_6d, whose held-out operator list the ladder reuses, has no such guard and stays frozen; its header now says so, and HPT warns once when a sample falls back to an empty default_prompt, as Pi does. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01THTuWdgk6M6pZ5Ya7ihZEF --- egomimic/algo/eval_determinism.py | 16 + .../data/mecka_fold_freeform_opscale1_6d.yaml | 44 +++ ..._fold_freeform_opscale1_hpt_keypoints.yaml | 53 ++++ .../data/mecka_fold_freeform_opscale2_6d.yaml | 45 +++ ..._fold_freeform_opscale2_hpt_keypoints.yaml | 53 ++++ .../data/mecka_fold_freeform_opscale4_6d.yaml | 288 +++++++++++++++++ ..._fold_freeform_opscale4_hpt_keypoints.yaml | 53 ++++ .../data/mecka_fold_freeform_opsplit_6d.yaml | 7 + egomimic/trainHydra.py | 81 ++++- tests/unit/test_eval_loaders.py | 20 ++ tests/unit/test_opscale_ladder.py | 299 ++++++++++++++++++ 11 files changed, 945 insertions(+), 14 deletions(-) create mode 100644 egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_6d.yaml create mode 100644 egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_hpt_keypoints.yaml create mode 100644 egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_6d.yaml create mode 100644 egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_hpt_keypoints.yaml create mode 100644 egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_6d.yaml create mode 100644 egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_hpt_keypoints.yaml create mode 100644 tests/unit/test_opscale_ladder.py diff --git a/egomimic/algo/eval_determinism.py b/egomimic/algo/eval_determinism.py index 46c9dc6437..d25459d248 100644 --- a/egomimic/algo/eval_determinism.py +++ b/egomimic/algo/eval_determinism.py @@ -22,11 +22,14 @@ from __future__ import annotations +import logging import os import random import torch +logger = logging.getLogger(__name__) + EVAL_BASE_SEED = 0 EVAL_RANK_STRIDE = 1_000_003 # Offset of the val-loss RNG from the sampling generator's seed. It must not be @@ -114,6 +117,8 @@ def _build_prompts(self, _batch, batch_size: int) -> list[str]: prompts = [] for sample in _batch[self.annotation_key]: if not sample: + if not self.default_prompt: + self._warn_empty_prompt_once() prompts.append(self.default_prompt) elif eval_rng is not None: prompts.append(sample[eval_rng.randint(0, len(sample) - 1)]) @@ -122,3 +127,14 @@ def _build_prompts(self, _batch, batch_size: int) -> list[str]: else: # "first" prompts.append(sample[0]) return prompts + + def _warn_empty_prompt_once(self) -> None: + if getattr(self, "_empty_prompt_warned", False): + return + self._empty_prompt_warned = True + logger.warning( + "prompt fallback: a sample has no %r annotation and default_prompt " + "is empty, so the text stem sees ''. Filter unannotated episodes or " + "set model.robomimic_model.default_prompt.", + self.annotation_key, + ) diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_6d.yaml new file mode 100644 index 0000000000..e3b1396be3 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_6d.yaml @@ -0,0 +1,44 @@ +# MECKA_FOLD_FREEFORM_OPSCALE1_6D: the k=1-operator rung of the freeform +# fold-clothes OPERATOR-SCALING ladder. Everything -- pipeline, the annotation +# guard, the shared seen-op val, the 31-operator unseen val, the video pins, +# loader params -- is inherited from mecka_fold_freeform_opscale4_6d; only the +# train OPERATOR SET changes, so each rung only ADDS an operator. Edit the ladder +# only in the parent. +# +# TRAIN: ALL annotated episodes of the top 1 operator by annotated hours, +# minus the 3 shared seen-op val episodes -> +# 9 eps / 25633 f / 14.24 min = 0.237 h +# (1.00x the k=1 arm): +# #1 6964d38183a9fdf2d8663792 9 eps 25633 f = 14.24 min +# Counts measured 2026-09-17; the filter is live, so they grow if SQL gains more +# episodes for these operators. +# +# NORM STATS: this arm's train distribution differs from every other rung -- +# precompute on THIS config, do not reuse the parent's: +# python egomimic/scripts/precompute_norm_stats.py --data mecka_fold_freeform_opscale1_6d \ +# --model pi0.5_bc_mecka_6d --out + +defaults: + - mecka_fold_freeform_opscale4_6d + - _self_ + +# Hydra replaces lists wholesale, so the train lambdas are restated IN FULL; only +# the operator tuple differs from the parent. The resolver and `mode: total` are +# inherited untouched. +train_datasets: + human_bimanual: + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # Annotation guard: see mecka_fold_freeform_opscale4_6d. + - "lambda row: isinstance(row.get('segments'), list) and len(row.get('segments')) > 0" + # OPERATOR SET: the top 1 operator by ANNOTATED hours. Inlined rather + # than put in a config key -- MultiDataModuleWrapper rejects root keys it + # does not declare. + - "lambda row: str(row.get('operator', '')) in ('6964d38183a9fdf2d8663792',)" + # ...minus the 3 SEEN-OP VAL episodes (all from operator #1), which are + # the only episodes held back from training. Kept in sync with + # valid_datasets' `episode_hashes` pin by tests/unit/test_opscale_ladder.py. + - "lambda row: str(row.get('episode_hash', '')) not in ('696bcbe0e13d37dd4293f58b', '696bcbe86377d0871b6a9253', '696bcc71ed8ed158122ec794')" diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_hpt_keypoints.yaml new file mode 100644 index 0000000000..094efc9dc5 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale1_hpt_keypoints.yaml @@ -0,0 +1,53 @@ +# MECKA_FOLD_FREEFORM_OPSCALE1_HPT_KEYPOINTS: HPT keypoint-action twin of +# mecka_fold_freeform_opscale1_6d -- the SAME k=1-operator train set, the SAME +# annotation guard, the SAME shared seen-op val and the SAME unseen val (the +# filters, loader params and metric/video budgets are all inherited; edit the +# ladder only in mecka_fold_freeform_opscale4_6d). Only the keymap and transform +# differ: +# +# keymap_mode: keypoints +# mode: keypoints_wristframe_6d +# +# Batch keys: +# actions_keypoints (100, 144) wrist-first wrist-frame keypoint +# action: per hand [wrist xyz+rot6d in +# its own current wrist frame (9) | +# 21 MANO keypoints in that wrist +# frame (63)] x {left, right} +# observations.state.keypoints (144,) per hand [wrist pose in HEAD frame +# (9) | keypoints in wrist frame (63)]; +# the HPT `state_keypoints` stem input +# and what the eval revert (wrist -> +# head frame) reads +# annotations per-frame segment label -> the +# `sampled_prompt` for the Qwen text +# stem. Non-empty on every episode of +# every split (annotation guard). +# +# There is no `observations.state.ee_pose` (include_ee_pose off: HPT has no stem +# for it), so pad_proprio_gripper is off. fix_left_wrist_convention stays ON: in +# keypoint modes it corrects the left `obs/action_wrist_pose` keys, which carry +# the same double-mirrored rotation as the ee_pose on processed_v3. +# +# TRAIN: 9 eps / 25633 f / 14.24 min = 0.237 h across 1 operator +# (measured 2026-09-17; the filter is live). +# +# NORMALIZATION: proprio + action keys are quantile-normalized, and every rung of +# the ladder has its own train distribution -- precompute per arm: +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_kp_wrist_hpt \ +# --data mecka_fold_freeform_opscale1_hpt_keypoints --model hpt_bc_keypoints_wrist_300M --out +# +# Pairs with model=hpt_bc_keypoints_wrist_300M and +# evaluator=eval_hpt_keypoints_wrist. +defaults: + - mecka_fold_freeform_opscale1_6d + - _self_ + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + transform_list: + mode: keypoints_wristframe_6d + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_6d.yaml new file mode 100644 index 0000000000..ab7943eb4b --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_6d.yaml @@ -0,0 +1,45 @@ +# MECKA_FOLD_FREEFORM_OPSCALE2_6D: the k=2-operator rung of the freeform +# fold-clothes OPERATOR-SCALING ladder. Everything -- pipeline, the annotation +# guard, the shared seen-op val, the 31-operator unseen val, the video pins, +# loader params -- is inherited from mecka_fold_freeform_opscale4_6d; only the +# train OPERATOR SET changes, so each rung only ADDS an operator. Edit the ladder +# only in the parent. +# +# TRAIN: ALL annotated episodes of the top 2 operators by annotated hours, +# minus the 3 shared seen-op val episodes -> +# 24 eps / 60374 f / 33.54 min = 0.559 h +# (2.36x the k=1 arm): +# #1 6964d38183a9fdf2d8663792 9 eps 25633 f = 14.24 min +# #2 68e0b875c33a1abcb8fc55b9 15 eps 34741 f = 19.30 min +# Counts measured 2026-09-17; the filter is live, so they grow if SQL gains more +# episodes for these operators. +# +# NORM STATS: this arm's train distribution differs from every other rung -- +# precompute on THIS config, do not reuse the parent's: +# python egomimic/scripts/precompute_norm_stats.py --data mecka_fold_freeform_opscale2_6d \ +# --model pi0.5_bc_mecka_6d --out + +defaults: + - mecka_fold_freeform_opscale4_6d + - _self_ + +# Hydra replaces lists wholesale, so the train lambdas are restated IN FULL; only +# the operator tuple differs from the parent. The resolver and `mode: total` are +# inherited untouched. +train_datasets: + human_bimanual: + filters: + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # Annotation guard: see mecka_fold_freeform_opscale4_6d. + - "lambda row: isinstance(row.get('segments'), list) and len(row.get('segments')) > 0" + # OPERATOR SET: the top 2 operators by ANNOTATED hours. Inlined rather + # than put in a config key -- MultiDataModuleWrapper rejects root keys it + # does not declare. + - "lambda row: str(row.get('operator', '')) in ('6964d38183a9fdf2d8663792', '68e0b875c33a1abcb8fc55b9')" + # ...minus the 3 SEEN-OP VAL episodes (all from operator #1), which are + # the only episodes held back from training. Kept in sync with + # valid_datasets' `episode_hashes` pin by tests/unit/test_opscale_ladder.py. + - "lambda row: str(row.get('episode_hash', '')) not in ('696bcbe0e13d37dd4293f58b', '696bcbe86377d0871b6a9253', '696bcc71ed8ed158122ec794')" diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_hpt_keypoints.yaml new file mode 100644 index 0000000000..c6fd633ed4 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale2_hpt_keypoints.yaml @@ -0,0 +1,53 @@ +# MECKA_FOLD_FREEFORM_OPSCALE2_HPT_KEYPOINTS: HPT keypoint-action twin of +# mecka_fold_freeform_opscale2_6d -- the SAME k=2-operator train set, the SAME +# annotation guard, the SAME shared seen-op val and the SAME unseen val (the +# filters, loader params and metric/video budgets are all inherited; edit the +# ladder only in mecka_fold_freeform_opscale4_6d). Only the keymap and transform +# differ: +# +# keymap_mode: keypoints +# mode: keypoints_wristframe_6d +# +# Batch keys: +# actions_keypoints (100, 144) wrist-first wrist-frame keypoint +# action: per hand [wrist xyz+rot6d in +# its own current wrist frame (9) | +# 21 MANO keypoints in that wrist +# frame (63)] x {left, right} +# observations.state.keypoints (144,) per hand [wrist pose in HEAD frame +# (9) | keypoints in wrist frame (63)]; +# the HPT `state_keypoints` stem input +# and what the eval revert (wrist -> +# head frame) reads +# annotations per-frame segment label -> the +# `sampled_prompt` for the Qwen text +# stem. Non-empty on every episode of +# every split (annotation guard). +# +# There is no `observations.state.ee_pose` (include_ee_pose off: HPT has no stem +# for it), so pad_proprio_gripper is off. fix_left_wrist_convention stays ON: in +# keypoint modes it corrects the left `obs/action_wrist_pose` keys, which carry +# the same double-mirrored rotation as the ee_pose on processed_v3. +# +# TRAIN: 24 eps / 60374 f / 33.54 min = 0.559 h across 2 operators +# (measured 2026-09-17; the filter is live). +# +# NORMALIZATION: proprio + action keys are quantile-normalized, and every rung of +# the ladder has its own train distribution -- precompute per arm: +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_kp_wrist_hpt \ +# --data mecka_fold_freeform_opscale2_hpt_keypoints --model hpt_bc_keypoints_wrist_300M --out +# +# Pairs with model=hpt_bc_keypoints_wrist_300M and +# evaluator=eval_hpt_keypoints_wrist. +defaults: + - mecka_fold_freeform_opscale2_6d + - _self_ + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + transform_list: + mode: keypoints_wristframe_6d + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_6d.yaml new file mode 100644 index 0000000000..b6d6580247 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_6d.yaml @@ -0,0 +1,288 @@ +# MECKA_FOLD_FREEFORM_OPSCALE4_6D: pi0.5 BC on the mecka FREEFORM fold-clothes +# episodes (SQL lab='mecka', task='folding_clothes', 462 eps / 1.13 M frames / +# 178 operators under processed_v3/mecka/freeform/; "freeform" is only visible in +# zarr_processed_path, so the filter keys on '/mecka/freeform/' and accepts either +# fold-clothes task spelling). A DIFFERENT set from the FLAGSHIP 'fold_clothes' +# episodes -- zero operator overlap between the two. +# Standalone (no defaults inheritance); the k=1/2 arms inherit THIS file. +# +# OPERATOR-SCALING LADDER (chosen 2026-09-17 from SQL app.episodes). The question +# is how many DIFFERENT OPERATORS the training data should come from. Each arm +# trains on ALL the annotated data its operators have, minus the held-out +# validation episodes -- the per-operator amount is NOT capped, so the total +# grows both with the operator count and with how much those operators recorded: +# +# k=1 9 eps 25633 f 14.24 min = 0.237 h ( 1.00x the k=1 arm) +# k=2 24 eps 60374 f 33.54 min = 0.559 h ( 2.36x the k=1 arm) +# k=4 48 eps 118388 f 65.77 min = 1.096 h ( 4.62x the k=1 arm) +# +# Per-operator contributions (k=4; the smaller arms are prefixes of this list): +# #1 6964d38183a9fdf2d8663792 9 eps 25633 f = 14.24 min +# #2 68e0b875c33a1abcb8fc55b9 15 eps 34741 f = 19.30 min +# #3 695d09fc83a9fdf2d84d9c11 13 eps 30549 f = 16.97 min +# #4 69439ad5a2a8b5aee76602f5 11 eps 27465 f = 15.26 min +# +# So the arms are NOT a clean 1:2:4 in hours (1.00x, 2.36x, 4.62x) -- +# operator #1 is the smallest of the four after funding the val set. Operator +# count and data volume therefore move TOGETHER across this ladder; it measures +# "does adding another operator's data help", not operator diversity at a fixed +# budget. For the fixed-budget question, cap each operator instead. +# +# The three train operator sets are NESTED (k=1 subset k=2 subset k=4), so +# operator identity never changes underneath the ladder -- each rung only ADDS an +# operator. Because the split is an operator filter rather than an episode pin, +# that nesting cannot drift. +# +# LIVE FILTERS, NOT A SNAPSHOT: the train split is "every annotated episode of +# these operators", so it will GROW if more episodes for them land in SQL. The +# episode counts above were measured 2026-09-17 by hand and are documentation +# only -- nothing asserts them, because there is nothing to assert them against +# without the DB. tests/unit/test_opscale_ladder.py checks the parts that are +# in the config: the operator nesting, the exclusions and the metric budgets. +# +# ANNOTATED EPISODES ONLY. 112 of the 462 freeform episodes have segments = NULL, +# which becomes an empty `annotations` track and makes _build_prompts fall back to +# `default_prompt` ("" by default) for every one of their frames. Restricting to +# the 350 annotated episodes (7.91 h / 156 operators) makes the text stem see a +# real instruction on ~100 % of frames. The operator RANKING is by ANNOTATED +# hours, which is NOT the same order as raw hours -- do not re-derive it from +# unfiltered totals. +# +# THREE-WAY SPLIT, and both val sets are IDENTICAL in all three arms: +# train = all annotated episodes of the arm's operators, minus the +# 3 seen-op val episodes below. +# seen val = 3 episodes of operator #1 (6964d38183a9fdf2d8663792), +# 9745 f = 5.41 min -> seen_op_valid/. Operator #1 is in EVERY +# arm (the sets are nested), so this is the one seen-op val +# that is directly comparable across the whole ladder. These +# are the ONLY episodes withheld from training. +# unseen val = the ANNOTATED episodes of the 31 `held_out_operators` +# (35 of their 50 eps, 79850 f = 0.739 h, 26 of the 31 operators +# contribute) -> unseen_op_valid/. The operator LIST is reused +# verbatim from mecka_fold_freeform_opsplit_6d, so no unseen +# operator is ever trained on and the split is comparable with +# those checkpoints at the OPERATOR level; the annotation guard +# makes it a subset of that config's 50-episode val side, so it +# is NOT episode-identical. +# train_viz = the train split itself (no train_viz_datasets here, so +# trainHydra reuses train_datasets unshuffled). +# +# VAL LOADERS (three heads, one validation pass), as in the flagship opsplit +# configs: valid_datasets -> seen_op_valid/ (renamed by `valid_prefix`), +# unseen_op_valid_datasets -> unseen_op_valid/, train_viz -> train_viz/. Each is +# served by a subsampled METRIC loader and a pinned contiguous VIDEO loader; see +# `metric_frames_per_episode` / `video_episodes` at the bottom. +# +# PIPELINE: WRIST-frame CARTESIAN actions with the continuous 6D rotation (18D = +# xyz+6D per arm, no gripper): each arm's 100-step chunk is expressed relative to +# that wrist's current pose, so the world frame and the moving head frame cancel +# out of the targets; the head-frame proprio ee_pose is 6D-encoded too and padded +# 18 -> 20 (zero grip slots at 9/19) so the State: prompt bins align with the +# robot layout. fix_left_wrist_convention: processed_v3 mecka zarrs have the LEFT wrist +# frame double-mirrored; Rz(180)-corrected before any frame math. stride 1 (mecka +# is recorded at the target rate). Prompts come from the `annotations` track -- +# one segment label per frame, sampled by _build_prompts. +# +# Pairs with model=pi0.5_bc_mecka_6d and evaluator=eval_pi_wristframe_6d. For CAM +# (head)-frame actions add +# data.train_datasets.human_bimanual.resolver.transform_list.mode=cartesian_6d +# with evaluator=eval_pi_camframe_6d. +# +# NORMALIZATION: every arm has a DIFFERENT train distribution, so precompute +# stats per arm -- do NOT share them across the ladder or reuse the opsplit / +# flagship stats: +# python egomimic/scripts/precompute_norm_stats.py --data mecka_fold_freeform_opscale4_6d \ +# --model pi0.5_bc_mecka_6d --out +# then train with norm_stats.precomputed_norm_path=/norm_stats. + +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +# The 31 UNSEEN-operator ids, reused verbatim from mecka_fold_freeform_opsplit_6d +# so no unseen operator is ever trained on and the split stays operator-level +# comparable with the existing freeform opsplit checkpoints. Verified disjoint +# from the 4 train operators. Referenced by the unseen filter via +# `data.held_out_operators`; MultiDataModuleWrapper declares this key. +held_out_operators: + - 68a8e70d1ffcf0417e08f0e5 + - 68e2fc680b158ea94ab4f5bc + - 6930bf8a6e496b45b1e1122e + - 6931784c6e496b45b1e1b1cd + - 693a399ec67c9e4814b0c09f + - 693b9530c67c9e4814b20ab2 + - 693c78165cb6c2828a4c2087 + - 694604945c0479a0fb7f48a4 + - 6948d23b5c0479a0fb8302b3 + - 69529a85c0d9448524bc7880 + - 6954f0b720b100982d812ed8 + - 6957464720b100982d83b76c + - 695d07d783a9fdf2d84d9371 + - 695d084183a9fdf2d84d955a + - 6961c50f83a9fdf2d85ec95d + - 69624a8683a9fdf2d86013c0 + - 6973c3ae89fa4bc311d1bbe5 + - 69842643f8ca4464d9920799 + - 698937a5f8ca4464d99dce94 + - 6992b19ab9d55c9c2d48432c + - 699429e18ef0252970048624 + - 6996d5eded1708063b2fe419 + - 6999d17c8f173e14271f9890 + - 69a24256c4216c3e3ce33834 + - 69a39ec9c4216c3e3ce49734 + - 69a57a78c4216c3e3ce6323c + - 69a725a347e12c9b3390ba33 + - 69a7943547e12c9b33911c02 + - 69a84eed47e12c9b339208fd + - 69ab87207bf9cbce29a3cb36 + - 69ae2463d22d99655253f287 + +# Metric prefix / video dir for the canonical valid loader (config-only, read by +# trainHydra), so the SEEN-operator val is not logged under a bare Valid/. +valid_prefix: seen_op_valid + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + folder_path: ${paths.dataset_dir} + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + keymap_mode: cartesian + # No proprio_history: the keymap's K and the model's stem history_len are + # one setting in two files and HPTModel raises when they disagree, so a + # rung that asks for K > 1 only composes with a model configured for the + # same K. Raise both together (and see PoseCoordinateFrameTransform's + # per_step_target, which the keypoint twins rely on above K=1). + # Language: the per-frame segment label track. Every episode in every + # split of this config is annotation-guarded, so this is never empty. + annotation_key: annotations + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: cartesian_wristframe_6d + stride: 1 + fix_left_wrist_convention: true + pad_proprio_gripper: true + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # LANGUAGE ANNOTATIONS REQUIRED. 112 of the 462 freeform episodes have a + # NULL `segments` column, which becomes an EMPTY `annotations` track in the + # zarr; _annotation_text_for_frame then returns [] for every frame of them + # and _build_prompts silently substitutes `default_prompt` (an EMPTY string + # unless overridden), so they would train the text stem on nothing. This + # guard keeps the 350 annotated episodes, whose segments tile 99.7-100 % of + # their frames. `segments` is a list (or None) on the SQL row, so isinstance + # is the safe test -- do NOT use truthiness (a numpy round-trip makes + # `bool(...)` raise "truth value of an array is ambiguous"). + # Anchored so valid/unseen reuse this exact lambda. + - &annotated "lambda row: isinstance(row.get('segments'), list) and len(row.get('segments')) > 0" + # OPERATOR SET: the top 4 operators by ANNOTATED hours. Inlined rather + # than put in a config key -- MultiDataModuleWrapper rejects root keys it + # does not declare. + - "lambda row: str(row.get('operator', '')) in ('6964d38183a9fdf2d8663792', '68e0b875c33a1abcb8fc55b9', '695d09fc83a9fdf2d84d9c11', '69439ad5a2a8b5aee76602f5')" + # ...minus the 3 SEEN-OP VAL episodes (all from operator #1), which are + # the only episodes held back from training. Kept in sync with + # valid_datasets' `episode_hashes` pin by tests/unit/test_opscale_ladder.py. + - "lambda row: str(row.get('episode_hash', '')) not in ('696bcbe0e13d37dd4293f58b', '696bcbe86377d0871b6a9253', '696bcc71ed8ed158122ec794')" + # The filters ARE the split; no random episode split on top of them. + mode: total + +# SEEN operators, held-out episodes -> seen_op_valid/ metrics. 3 episodes of +# operator #1, which is in every arm of the ladder, so this val set is shared and +# comparable across all three rungs. Same resolver as train (so it follows any +# pipeline override); pinned by hash, and `mode: total` because the pin is the +# split. The train filter excludes exactly these 3 hashes. +valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # Annotation guard: see train_datasets. + - *annotated + # op #1 6964d38183a9fdf2d8663792: 3 eps, 9745 f = 5.41 min + episode_hashes: + - 696bcbe0e13d37dd4293f58b + - 696bcbe86377d0871b6a9253 + - 696bcc71ed8ed158122ec794 + mode: total + +# UNSEEN operators: the ANNOTATED episodes of the 31 held-out operators -> +# unseen_op_valid/ metrics via a second copy of the canonical evaluator. +# Identical in all three arms. +unseen_op_valid_datasets: + human_bimanual: + _target_: ${data.train_datasets.human_bimanual._target_} + resolver: ${data.train_datasets.human_bimanual.resolver} + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + - "lambda row: str(row.get('task', '')).lower() in ('folding_clothes', 'fold_clothes')" + - "lambda row: '/mecka/freeform/' in str(row.get('zarr_processed_path', ''))" + # Annotation guard: see train_datasets. + - *annotated + - "lambda row: str(row.get('operator', '')) in ${data.held_out_operators}" + mode: total + +train_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + +# Unshuffled (both val loaders): the same leading slice every val, comparable +# across epochs with a coherent overlay video. The train_viz loader reuses these. +valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false + +unseen_op_valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 8 + shuffle: false + +# --------------------------------------------------------------------------- +# METRIC LOADERS: K evenly spaced frames per EPISODE, so every episode and +# operator is scored inside the same limit_val_batches window (80 x 64 = 5120 +# frames per rank on the default trainer). `auto` sets +# K = floor(limit_val_batches * batch_size / episodes-in-split) from the RESOLVED +# split at startup, so K follows the live filters as SQL grows and needs no +# per-rung override. At the 2026-09-17 counts that is: +# valid 3 eps / 9745 f -> K = floor(5120/3) = 1706 +# unseen_op_valid 35 eps / 79850 f -> K = floor(5120/35) = 146 +# train_viz 48 eps / 118388 f -> K = floor(5120/48) = 106 (k=4) +# 24 eps / 60374 f -> K = 213 (k=2) +# 9 eps / 25633 f -> K = 568 (k=1) +metric_frames_per_episode: + valid: auto + unseen_op_valid: auto + train_viz: auto + +# --------------------------------------------------------------------------- +# VIDEO LOADERS: pinned contiguous episodes, so every run in the ladder renders +# the SAME episodes. Taken from that head's already-resolved split (no second SQL +# pull); `evaluator.viz_max_batches` (10 batches = 640 frames) caps them. Each +# pin is the LONGEST episode of its split, so one contiguous episode fills the +# whole video budget. Format (operator, frames, segments): +# valid 696bcbe86377d0871b6a9253 (6964d38183a9fdf2d8663792, 3688 f, 8 seg) +# unseen_op_valid 696bc7e021d8e8ab8ddc8380 (6957464720b100982d83b76c, 3358 f, 10 seg) +# train_viz 696dd090a56f03a27e59f69c (6964d38183a9fdf2d8663792, 3298 f, 8 seg) +# -- an operator #1 episode, so it is in EVERY arm's train +# split and the children inherit this pin unchanged. +video_episodes: + valid: + - 696bcbe86377d0871b6a9253 + unseen_op_valid: + - 696bc7e021d8e8ab8ddc8380 + train_viz: + - 696dd090a56f03a27e59f69c diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_hpt_keypoints.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_hpt_keypoints.yaml new file mode 100644 index 0000000000..cc67215aa5 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opscale4_hpt_keypoints.yaml @@ -0,0 +1,53 @@ +# MECKA_FOLD_FREEFORM_OPSCALE4_HPT_KEYPOINTS: HPT keypoint-action twin of +# mecka_fold_freeform_opscale4_6d -- the SAME k=4-operator train set, the SAME +# annotation guard, the SAME shared seen-op val and the SAME unseen val (the +# filters, loader params and metric/video budgets are all inherited; edit the +# ladder only in mecka_fold_freeform_opscale4_6d). Only the keymap and transform +# differ: +# +# keymap_mode: keypoints +# mode: keypoints_wristframe_6d +# +# Batch keys: +# actions_keypoints (100, 144) wrist-first wrist-frame keypoint +# action: per hand [wrist xyz+rot6d in +# its own current wrist frame (9) | +# 21 MANO keypoints in that wrist +# frame (63)] x {left, right} +# observations.state.keypoints (144,) per hand [wrist pose in HEAD frame +# (9) | keypoints in wrist frame (63)]; +# the HPT `state_keypoints` stem input +# and what the eval revert (wrist -> +# head frame) reads +# annotations per-frame segment label -> the +# `sampled_prompt` for the Qwen text +# stem. Non-empty on every episode of +# every split (annotation guard). +# +# There is no `observations.state.ee_pose` (include_ee_pose off: HPT has no stem +# for it), so pad_proprio_gripper is off. fix_left_wrist_convention stays ON: in +# keypoint modes it corrects the left `obs/action_wrist_pose` keys, which carry +# the same double-mirrored rotation as the ee_pose on processed_v3. +# +# TRAIN: 48 eps / 118388 f / 65.77 min = 1.096 h across 4 operators +# (measured 2026-09-17; the filter is live). +# +# NORMALIZATION: proprio + action keys are quantile-normalized, and every rung of +# the ladder has its own train distribution -- precompute per arm: +# python egomimic/scripts/precompute_norm_stats.py --config-name train_zarr_mecka_kp_wrist_hpt \ +# --data mecka_fold_freeform_opscale4_hpt_keypoints --model hpt_bc_keypoints_wrist_300M --out +# +# Pairs with model=hpt_bc_keypoints_wrist_300M and +# evaluator=eval_hpt_keypoints_wrist. +defaults: + - mecka_fold_freeform_opscale4_6d + - _self_ + +train_datasets: + human_bimanual: + resolver: + key_map: + keymap_mode: keypoints + transform_list: + mode: keypoints_wristframe_6d + pad_proprio_gripper: false diff --git a/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml b/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml index 0c878b6392..0eacf0e873 100644 --- a/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml +++ b/egomimic/hydra_configs/data/mecka_fold_freeform_opsplit_6d.yaml @@ -19,6 +19,13 @@ # 1,013,220 frames. Both sides use mode: total (no random episode split); the # list IS the split. Do not edit it: checkpoints were trained against it. # +# NO ANNOTATION GUARD: 112 of the 462 freeform episodes have segments = NULL, so +# their `annotations` track is empty and the prompt falls back to +# `default_prompt` ('' by default) on every frame: Pi sees "Task: , State: ...", +# HPT feeds '' to its text stem. Left as is because checkpoints were trained on +# this split; the opscale ladder (mecka_fold_freeform_opscale4_6d) reuses only +# its held-out OPERATOR list and filters to annotated episodes. +# # WRIST-frame CARTESIAN actions with the continuous 6D rotation (18D = # xyz+6D per arm, no gripper): each arm's 100-step chunk is expressed # relative to that wrist's current pose; the head-frame proprio ee_pose is diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 28a7de9aa4..43ad6134df 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -312,7 +312,12 @@ def _trainer_world_size(cfg: DictConfig) -> int: return n * (int(nodes) if isinstance(nodes, int) and nodes > 0 else 1) -def _metric_frames_per_episode(cfg: DictConfig, head: str) -> int | None: +def _metric_frames_per_episode( + cfg: DictConfig, + head: str, + n_episodes: int | None = None, + batch_size: int | None = None, +) -> int | None: """Frames per episode to keep on this val head, or None for no subsampling. ``data.metric_frames_per_episode[head]`` is written for ONE rank -- it is @@ -323,6 +328,11 @@ def _metric_frames_per_episode(cfg: DictConfig, head: str) -> int | None: which on the seen-val head is fewer than it scored before subsampling existed at all. EvenStrideDataset keeps a whole episode when K exceeds its length, so a split that fits entirely is not subsampled. + + ``auto`` computes that formula from the resolved split (``n_episodes``) and + the head's loader ``batch_size``, for splits defined by live SQL filters + whose episode count a hard-coded K would silently fall behind. With + ``limit_val_batches=0`` (validation off) it subsamples nothing. """ table = cfg.data.get("metric_frames_per_episode") if table is None: @@ -330,15 +340,37 @@ def _metric_frames_per_episode(cfg: DictConfig, head: str) -> int | None: k = table.get(head) if k is None: return None - k = int(k) - if k <= 0: - raise ValueError( - f"data.metric_frames_per_episode.{head} must be a positive int, got {k}" - ) + if k == "auto": + limit = cfg.get("trainer", {}).get("limit_val_batches") + if limit == 0: # validation is off; nothing to budget + return None + if ( + isinstance(limit, bool) + or not isinstance(limit, int) + or limit <= 0 + or not batch_size + or not n_episodes + ): + raise ValueError( + f"data.metric_frames_per_episode.{head}=auto needs an int " + "trainer.limit_val_batches, the head's loader batch_size and a " + f"resolved split; got limit_val_batches={limit!r}, " + f"batch_size={batch_size!r}, episodes={n_episodes!r}" + ) + k = max(1, limit * int(batch_size) // n_episodes) + else: + k = int(k) + if k <= 0: + raise ValueError( + f"data.metric_frames_per_episode.{head} must be a positive int " + f"or 'auto', got {k}" + ) return k * _trainer_world_size(cfg) -def _subsample_val_datasets(cfg: DictConfig, head: str, datasets: dict) -> dict: +def _subsample_val_datasets( + cfg: DictConfig, head: str, datasets: dict, loader_params=None +) -> dict: """Wrap each of this val head's datasets in ``EvenStrideDataset`` so the ``limit_val_batches`` window spans EVERY episode of the split. @@ -350,16 +382,24 @@ def _subsample_val_datasets(cfg: DictConfig, head: str, datasets: dict) -> dict: episode gives every episode -- and so every operator -- equal weight at the same batch count. - No ``metric_frames_per_episode`` entry for this head => datasets pass - through untouched.""" - k = _metric_frames_per_episode(cfg, head) - if k is None or not datasets: + ``loader_params`` is this head's ``{dataset_name: DataLoader kwargs}``, + read only for an ``auto`` K. No ``metric_frames_per_episode`` entry for this + head => datasets pass through untouched.""" + table = cfg.data.get("metric_frames_per_episode") + if table is None or table.get(head) is None or not datasets: return datasets wrapped = {} for name, ds in datasets.items(): if ds is None: wrapped[name] = None continue + params = (loader_params or {}).get(name) or {} + k = _metric_frames_per_episode( + cfg, head, n_episodes=len(ds.datasets), batch_size=params.get("batch_size") + ) + if k is None: + wrapped[name] = ds + continue sub = EvenStrideDataset(ds, frames_per_episode=k) log.info( f"val head '{head}' dataset '{name}': EvenStrideDataset " @@ -477,12 +517,25 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: "unseen_op_valid": unseen_op_valid_datasets, }, ) - valid_datasets = _subsample_val_datasets(cfg, "valid", valid_datasets) + valid_datasets = _subsample_val_datasets( + cfg, + "valid", + valid_datasets, + cfg.data.get("valid_dataloader_params"), + ) train_viz_datasets = _subsample_val_datasets( - cfg, "train_viz", train_viz_datasets + cfg, + "train_viz", + train_viz_datasets, + train_viz_params + if train_viz_params is not None + else cfg.data.get("train_viz_dataloader_params"), ) unseen_op_valid_datasets = _subsample_val_datasets( - cfg, "unseen_op_valid", unseen_op_valid_datasets + cfg, + "unseen_op_valid", + unseen_op_valid_datasets, + cfg.data.get("unseen_op_valid_dataloader_params"), ) log.info(f"Instantiating datamodule <{cfg.data._target_}>") diff --git a/tests/unit/test_eval_loaders.py b/tests/unit/test_eval_loaders.py index 626e003715..24e77fa740 100644 --- a/tests/unit/test_eval_loaders.py +++ b/tests/unit/test_eval_loaders.py @@ -334,6 +334,26 @@ def test_k_scales_with_the_world_size(): assert len(th._subsample_val_datasets(cfg, "valid", split)[HUMAN]) == 20 +def test_auto_k_is_derived_from_the_resolved_split(): + """`auto` = floor(limit_val_batches * batch_size / episodes), per dataset, + from the head's own loader batch size, then scaled by the world size.""" + cfg = _cfg(metric={"valid": "auto"}) + cfg.trainer = {"limit_val_batches": 4, "devices": 2} + split = {HUMAN: _split({"ep0": 100, "ep1": 100, "ep2": 100})} + params = {HUMAN: {"batch_size": 5}} + wrapped = th._subsample_val_datasets(cfg, "valid", split, params)[HUMAN] + assert wrapped.frames_per_episode == (4 * 5 // 3) * 2 + for limit in (1.0, None, True): + cfg.trainer.limit_val_batches = limit + with pytest.raises(ValueError, match="auto"): + th._subsample_val_datasets(cfg, "valid", split, params) + cfg.trainer.limit_val_batches = 0 # validation off, as the bench scripts run + assert th._subsample_val_datasets(cfg, "valid", split, params) == split + cfg.trainer.limit_val_batches = 4 + with pytest.raises(ValueError, match="batch_size"): + th._subsample_val_datasets(cfg, "valid", split, {}) + + # ------------------------------------------------------------- 3. video loaders def test_video_dataset_holds_only_the_pins_in_frame_order(): split = _split({"ep0": 30, "ep1": 20, "ep2": 25}) diff --git a/tests/unit/test_opscale_ladder.py b/tests/unit/test_opscale_ladder.py new file mode 100644 index 0000000000..f058274a06 --- /dev/null +++ b/tests/unit/test_opscale_ladder.py @@ -0,0 +1,299 @@ +"""The freeform fold-clothes OPERATOR-SCALING ladder holds its shape. + +mecka_fold_freeform_opscale{1,2,4} vary ONE thing: how many operators the +training data comes from. Each arm trains on ALL the annotated data its +operators have, minus the shared held-out val episodes. Everything that would +confound the comparison -- the seen-op val, the unseen-op val, the pipeline -- +must be identical across the three rungs, and the operator sets must be nested. + +The train split is an OPERATOR FILTER, not an episode pin, so these tests drive +the real DatasetFilter with synthetic rows rather than counting hashes. +""" + +from pathlib import Path + +import hydra +import pytest +import yaml +from hydra import compose, initialize_config_module +from omegaconf import OmegaConf + +import egomimic.hydra_configs as _cfg_pkg +from egomimic.rldb.filters import DatasetFilter + +DATA_DIR = Path(_cfg_pkg.__file__).parent / "data" +ARMS = (1, 2, 4) +VARIANTS = ("6d", "hpt_keypoints") + +# The 4 train operators, ranked by ANNOTATED hours (SQL app.episodes, 2026-09-17). +# Arm k trains on the first k of these. NOT the same order as raw hours. +OPERATORS = [ + "6964d38183a9fdf2d8663792", + "68e0b875c33a1abcb8fc55b9", + "695d09fc83a9fdf2d84d9c11", + "69439ad5a2a8b5aee76602f5", +] +# The only episodes withheld from training: 3 episodes of operator #1. +SEEN_VAL = [ + "696bcbe0e13d37dd4293f58b", + "696bcbe86377d0871b6a9253", + "696bcc71ed8ed158122ec794", +] +# Episodes per arm as measured 2026-09-17. The filter is LIVE, so these are +# documentation, not a guarantee; the metric budget is derived from the resolved +# split at runtime (`auto`), so nothing has to be kept in step with them. +TRAIN_EPISODES = {1: 9, 2: 24, 4: 48} +SEEN_VAL_EPISODES = 3 +UNSEEN_OPERATORS = 31 +SPLITS = ("train_datasets", "valid_datasets", "unseen_op_valid_datasets") + +# A row that clears every non-operator lambda: mecka, freeform, fold-clothes, +# annotated, and not one of the withheld episodes. +BASE_ROW = { + "lab": "mecka", + "task": "folding_clothes", + "zarr_processed_path": "s3://rldb/processed_v3/mecka/freeform/abc.zarr", + "episode_hash": "an_episode_not_in_the_val_set", + "segments": [{"label": "fold shirt", "start_seconds": 0, "end_seconds": 9}], +} + + +def _data(name: str): + with initialize_config_module( + config_module="egomimic.hydra_configs", version_base=None + ): + cfg = compose(config_name="train_zarr_cartesian", overrides=[f"data={name}"]) + cfg._set_flag("allow_objects", True) + OmegaConf.resolve(cfg.data) + return cfg.data + + +def _name(k: int, variant: str) -> str: + return f"mecka_fold_freeform_opscale{k}_{variant}" + + +def _filter(data, split: str) -> DatasetFilter: + node = data[split].human_bimanual + return DatasetFilter( + filter_lambdas=list(node.filters.filter_lambdas), + episode_hashes=list(node.filters.get("episode_hashes") or []), + ) + + +@pytest.fixture(scope="module") +def ladder(): + return {(k, v): _data(_name(k, v)) for k in ARMS for v in VARIANTS} + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_train_accepts_exactly_the_first_k_operators(ladder, k, variant): + """Nesting, stated as behaviour: arm k takes operators 1..k and no others, + so each rung only ADDS an operator to the one below it.""" + f = _filter(ladder[(k, variant)], "train_datasets") + for i, op in enumerate(OPERATORS): + row = {**BASE_ROW, "operator": op} + assert f.matches(row) is (i < k), f"operator #{i + 1} in the k={k} arm" + assert not f.matches({**BASE_ROW, "operator": "some_other_operator"}) + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_train_withholds_the_seen_val_episodes(ladder, k, variant): + """The seen-op val is the ONLY thing kept out of training, so the exclusion + must be present in every arm -- including k=1, whose single operator owns + all three of those episodes.""" + f = _filter(ladder[(k, variant)], "train_datasets") + for h in SEEN_VAL: + row = {**BASE_ROW, "operator": OPERATORS[0], "episode_hash": h} + assert not f.matches(row), f"k={k}: train accepted seen-val episode {h}" + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_seen_val_pin_matches_what_train_excludes(ladder, k, variant): + """The val pin and the train exclusion are written separately; if they drift, + an episode is either trained on AND validated, or dropped entirely.""" + data = ladder[(k, variant)] + pin = list(data.valid_datasets.human_bimanual.filters.episode_hashes) + assert sorted(pin) == sorted(SEEN_VAL) + assert len(pin) == SEEN_VAL_EPISODES + train = _filter(data, "train_datasets") + # every pinned episode is refused by train, and a non-pinned one is not + assert all( + not train.matches({**BASE_ROW, "operator": OPERATORS[0], "episode_hash": h}) + for h in pin + ) + assert train.matches({**BASE_ROW, "operator": OPERATORS[0]}) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_val_sets_are_identical_across_the_ladder(ladder, variant): + seen = { + k: list( + ladder[(k, variant)].valid_datasets.human_bimanual.filters.episode_hashes + ) + for k in ARMS + } + unseen = {k: list(ladder[(k, variant)].held_out_operators) for k in ARMS} + assert len({tuple(v) for v in seen.values()}) == 1, "seen-op val differs by rung" + assert len({tuple(v) for v in unseen.values()}) == 1, "unseen-op list differs" + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_train_operators_are_never_in_the_unseen_split(ladder, k, variant): + data = ladder[(k, variant)] + assert len(data.held_out_operators) == UNSEEN_OPERATORS + assert not (set(OPERATORS) & set(data.held_out_operators)) + # and the unseen filter refuses a train operator outright + unseen = _filter(data, "unseen_op_valid_datasets") + assert not unseen.matches({**BASE_ROW, "operator": OPERATORS[0]}) + assert unseen.matches({**BASE_ROW, "operator": data.held_out_operators[0]}) + + +def test_unseen_operators_match_the_opsplit_config(): + """Reusing mecka_fold_freeform_opsplit_6d's held-out list verbatim is what + keeps this ladder's unseen_op_valid/ numbers comparable with the existing + freeform opsplit checkpoints. Do not re-cut it.""" + raw = yaml.safe_load( + (DATA_DIR / "mecka_fold_freeform_opsplit_6d.yaml").read_text(encoding="utf-8") + ) + assert sorted(_data(_name(4, "6d")).held_out_operators) == sorted( + raw["held_out_operators"] + ) + + +# --- language annotations ------------------------------------------------- +# 112 of the 462 freeform episodes have segments = NULL -> an empty `annotations` +# track -> _annotation_text_for_frame returns [] for every frame -> _build_prompts +# substitutes `default_prompt` ("" unless overridden). Every split of every rung +# must filter those out, or a third of training runs on an empty instruction. + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +@pytest.mark.parametrize("split", SPLITS) +def test_every_split_requires_language_annotations(ladder, k, variant, split): + data = ladder[(k, variant)] + f = _filter(data, split) + operator = ( + data.held_out_operators[0] + if split == "unseen_op_valid_datasets" + else OPERATORS[0] + ) + good = {**BASE_ROW, "operator": operator} + if split == "valid_datasets": + good["episode_hash"] = SEEN_VAL[0] # that split is pinned by hash + for bad in (None, []): + assert not f.matches( + {**good, "segments": bad} + ), f"{split}: accepted an episode with segments={bad!r}" + assert f.matches(good), f"{split}: guard rejects an annotated row" + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_annotation_key_is_wired(ladder, k, variant): + """Without annotation_key the keymap emits no annotation entry at all, and + the prompt is `default_prompt` everywhere regardless of the guard.""" + km = ladder[(k, variant)].train_datasets.human_bimanual.resolver.key_map + assert km.annotation_key == "annotations" + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_no_rung_asks_for_proprio_history_alone(ladder, k, variant): + """proprio_history is half a setting: the keymap decides how many steps a + sample carries and the model's stem how many tokens consume them, and + HPTModel.stem_process raises when they disagree. Every model in the repo is + history_len 1, so a rung shipping K > 1 is a rung nothing can train -- which + is what happened when the ladder outlived the one model that wanted K=4.""" + km = ladder[(k, variant)].train_datasets.human_bimanual.resolver.key_map + assert km.get("proprio_history", 1) == 1 + + +# --- pipeline and loaders ------------------------------------------------- + + +@pytest.mark.parametrize("k", ARMS) +def test_keypoint_twin_shares_its_parents_split(ladder, k): + """The HPT twin may change only the keymap/transform; if its filters drift + from the 6d parent the two model families stop being comparable.""" + for split in SPLITS: + assert list( + ladder[(k, "hpt_keypoints")][split].human_bimanual.filters.filter_lambdas + ) == list(ladder[(k, "6d")][split].human_bimanual.filters.filter_lambdas), split + + +@pytest.mark.parametrize("k", ARMS) +def test_pipeline_per_variant(ladder, k): + for variant, keymap, mode, pad in ( + ("6d", "cartesian", "cartesian_wristframe_6d", True), + ("hpt_keypoints", "keypoints", "keypoints_wristframe_6d", False), + ): + res = ladder[(k, variant)].train_datasets.human_bimanual.resolver + assert res.key_map.keymap_mode == keymap + assert res.transform_list.mode == mode + assert res.transform_list.pad_proprio_gripper is pad + # processed_v3 mecka zarrs have the LEFT wrist frame double-mirrored. + # Instantiated, not just read off the node: a renamed kwarg would keep + # the assertion green while the config could no longer be built. + assert res.transform_list.fix_left_wrist_convention is True + assert hydra.utils.instantiate(res.transform_list) + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_filters_are_the_whole_split(ladder, k, variant): + """`mode: total` everywhere: the filters ARE the split, so no seeded episode + split may run on top of them and quietly drop a fraction of the data.""" + data = ladder[(k, variant)] + for split in SPLITS: + assert data[split].human_bimanual.mode == "total", split + assert data.valid_prefix == "seen_op_valid" + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_metric_budget_is_derived_from_the_resolved_split(ladder, k, variant): + """Every head is `auto`, so K follows the live filters; a hard-coded K would + push K x episodes past the limit_val_batches window as SQL grows. At the + documented counts `auto` reproduces floor(5120 / episodes).""" + from egomimic import trainHydra as th + + data = ladder[(k, variant)] + assert dict(data.metric_frames_per_episode) == { + "valid": "auto", + "unseen_op_valid": "auto", + "train_viz": "auto", + } + cfg = OmegaConf.create( + { + "data": {"metric_frames_per_episode": dict(data.metric_frames_per_episode)}, + "trainer": {"limit_val_batches": 80, "devices": 1}, + } + ) + bs = data.valid_dataloader_params.human_bimanual.batch_size + for head, n in (("train_viz", TRAIN_EPISODES[k]), ("valid", SEEN_VAL_EPISODES)): + kk = th._metric_frames_per_episode(cfg, head, n_episodes=n, batch_size=bs) + assert kk == 80 * bs // n + assert kk * n <= 80 * bs + + +@pytest.mark.parametrize("k", ARMS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_video_pins_belong_to_their_own_split(ladder, k, variant): + """A video pin is taken out of that head's already-resolved split; a pin from + elsewhere would make the loader re-resolve or render nothing.""" + data = ladder[(k, variant)] + vids = data.video_episodes + assert set(vids) == {"valid", "unseen_op_valid", "train_viz"} + assert set(vids.valid) <= set(SEEN_VAL) + assert len(vids.unseen_op_valid) == 1 and len(vids.train_viz) == 1 + # the train_viz pin is an operator #1 episode, so every arm -- including + # k=1 -- must actually accept it, which is why the children can inherit it + train = _filter(data, "train_datasets") + assert train.matches( + {**BASE_ROW, "operator": OPERATORS[0], "episode_hash": vids.train_viz[0]} + ), f"k={k}: the inherited train_viz video pin is not in this arm's split" From 7f6c8c32339d22f31e6173a6d7baca344deb08ea Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 11:09:09 -0400 Subject: [PATCH 29/79] checkpoints: last.ckpt as a symlink, not a second full copy ModelCheckpoint with `save_last: true` calls _save_checkpoint a second time for last.ckpt, so every checkpoint event writes the same tensors twice -- ~21 GB a time on Pi 0.5 (the size tests/integration/test_train_step_gpu.py records). It was measured at 9 GB and ~1.3 min on the now-retired ABC-DiT-L, inside an event that blocked the run for 2.7-3.4 min of a 42.3-min 1000-step cycle. "link" writes a relative symlink to the epoch checkpoint just saved instead. Safe here because save_top_k is -1, so the target is never collected, and _should_remove_checkpoint returns False when previous == current, which is always the case for last.ckpt. Resume is unaffected: _requeue_resume_path gates on os.path.isfile, which follows symlinks, and the opscale launcher already resumes through a hand-made symlink. A link can dangle where a file could not: the target deleted by hand, last.ckpt dropped by an rsync without -l, or a preemption between Lightning's unlink and relink. isfile is False in all three, and the requeue path used to warn and silently restart the run from scratch. It now raises when last.ckpt dangles or is missing next to other checkpoints; only an empty checkpoint dir falls back. One behaviour difference worth knowing: with "link" the callback relinks on EVERY train epoch once any checkpoint exists, not only on checkpoint epochs, so there is an extra os.remove + os.symlink + strategy.barrier per epoch. The test drives a real two-epoch fit and asserts both halves: "link" leaves a relative symlink to the newest epoch checkpoint, and `true` leaves a regular file the same size as it -- the copy this removes. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01M5dTGNydAZbRP1EsXgbp4W --- .../hydra_configs/callbacks/checkpoints.yaml | 6 +- egomimic/trainHydra.py | 27 +++- tests/unit/test_checkpoint_resume.py | 118 ++++++++++++++++++ 3 files changed, 149 insertions(+), 2 deletions(-) diff --git a/egomimic/hydra_configs/callbacks/checkpoints.yaml b/egomimic/hydra_configs/callbacks/checkpoints.yaml index fbd6d13277..6bb0f6f561 100644 --- a/egomimic/hydra_configs/callbacks/checkpoints.yaml +++ b/egomimic/hydra_configs/callbacks/checkpoints.yaml @@ -6,7 +6,11 @@ model_checkpoint: _target_: lightning.pytorch.callbacks.ModelCheckpoint dirpath: ${paths.output_dir}/checkpoints filename: "epoch_{epoch}" - save_last: true + # A symlink to the epoch checkpoint, not a second full copy of it (~21 GB on + # Pi 0.5). Requires save_top_k: -1 below, which is what keeps the target + # alive; relinks every epoch, not only checkpoint epochs. A requeue refuses to + # start over when last.ckpt dangles or is missing next to other checkpoints. + save_last: link save_top_k: -1 every_n_epochs: 50 # Save at the end of the train epoch, not after validation: with diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 43ad6134df..9326031106 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -104,7 +104,12 @@ def _requeue_resume_path(cfg: DictConfig) -> Optional[str]: def _prepare_checkpoint_resume(cfg: DictConfig) -> None: """Settle ``cfg.ckpt_path`` before the model config tree is built: a requeued job resumes from ``last.ckpt`` if it exists; one preempted before its first - checkpoint keeps the launch-time ``ckpt_path`` (warning).""" + checkpoint keeps the launch-time ``ckpt_path`` (warning). + + ``last.ckpt`` is a symlink (``save_last: link``), so it can dangle (target + deleted) or be missing while checkpoints exist (``rsync`` without ``-l``, a + preemption between Lightning's unlink and relink). Both raise: falling back + would silently restart a run that has checkpoints.""" requeue = _requeue_resume_path(cfg) if requeue is None: return @@ -112,6 +117,26 @@ def _prepare_checkpoint_resume(cfg: DictConfig) -> None: log.info(f"Detected SLURM requeue — resuming from {requeue}") cfg.ckpt_path = requeue return + if os.path.islink(requeue): + raise FileNotFoundError( + f"SLURM requeue: {requeue} links to {os.readlink(requeue)!r}, which " + "does not exist. Refusing to restart a run that has checkpoints; " + "restore the target or relink last.ckpt to an existing checkpoint." + ) + ckpt_dir = os.path.dirname(requeue) + existing = [ + os.path.join(ckpt_dir, f) + for f in (os.listdir(ckpt_dir) if os.path.isdir(ckpt_dir) else []) + if f.endswith(".ckpt") + ] + if existing: + newest = max(existing, key=os.path.getmtime) + raise FileNotFoundError( + f"SLURM requeue: {requeue} is missing but {ckpt_dir} holds " + f"{len(existing)} checkpoint(s), newest {newest}. Refusing to restart " + f"from scratch; link it with `ln -s {os.path.basename(newest)} " + f"{requeue}` and requeue." + ) fallback = cfg.get("ckpt_path") log.warning( f"SLURM requeue detected but {requeue} does not exist; falling back to " diff --git a/tests/unit/test_checkpoint_resume.py b/tests/unit/test_checkpoint_resume.py index 8c8e322e90..0e233e51ee 100644 --- a/tests/unit/test_checkpoint_resume.py +++ b/tests/unit/test_checkpoint_resume.py @@ -5,6 +5,7 @@ decision, and only when last.ckpt exists.""" import inspect +import os from pathlib import Path import pytest @@ -222,3 +223,120 @@ def test_load_checkpoint_weights_warns_on_unexpected_keys(tmp_path, caplog): load_checkpoint_weights(model, ckpt) assert "extra.buffer" in caplog.text assert torch.equal(model.weight, torch.ones(2, 2)) + + +# ------------------------------------------------ save_last writes a link, not 21 GB +def test_checkpoint_callback_links_last_instead_of_copying(compose_resolve, tmp_path): + """`save_last: true` wrote a second full copy of every checkpoint. "link" + symlinks it instead; save_top_k=-1 means the target is never collected.""" + import hydra + + cfg = _pi_cfg(compose_resolve, tmp_path) + assert cfg.callbacks.model_checkpoint.save_last == "link" + assert cfg.callbacks.model_checkpoint.save_top_k == -1, "the link's target" + + # lightning validates save_last at construction (it rejects "link" for a + # non-local dirpath), so instantiating is the real check + callback = hydra.utils.instantiate(cfg.callbacks.model_checkpoint) + assert callback.save_last == "link" + + +def test_requeue_resumes_through_a_symlinked_last_ckpt( + monkeypatch, compose_resolve, tmp_path +): + """What "link" leaves on disk: last.ckpt is a relative symlink to the real + epoch checkpoint, and the requeue path must still find it.""" + real = _touch(tmp_path / "checkpoints" / "epoch_epoch=49.ckpt") + link = tmp_path / "checkpoints" / "last.ckpt" + link.symlink_to(real.name) # relative, as ModelCheckpoint._link_checkpoint writes + + _requeue(monkeypatch) + cfg = _pi_cfg(compose_resolve, tmp_path) + train_hydra._prepare_checkpoint_resume(cfg) + assert Path(cfg.ckpt_path) == link + assert Path(cfg.ckpt_path).resolve() == real.resolve() + + +def test_requeue_with_a_dangling_last_ckpt_raises( + monkeypatch, compose_resolve, tmp_path +): + """os.path.isfile is False on a dangling link, which used to fall through to + "training starts over" on a run that has checkpoints.""" + link = tmp_path / "checkpoints" / "last.ckpt" + link.parent.mkdir(parents=True) + link.symlink_to("epoch_epoch=49.ckpt") + + _requeue(monkeypatch) + cfg = _pi_cfg(compose_resolve, tmp_path) + with pytest.raises(FileNotFoundError, match="does not exist"): + train_hydra._prepare_checkpoint_resume(cfg) + + +def test_requeue_with_checkpoints_but_no_last_ckpt_raises( + monkeypatch, compose_resolve, tmp_path +): + """rsync without -l, or a preemption between Lightning's unlink and relink, + leaves epoch checkpoints and no last.ckpt.""" + _touch(tmp_path / "checkpoints" / "epoch_epoch=49.ckpt") + + _requeue(monkeypatch) + cfg = _pi_cfg(compose_resolve, tmp_path) + with pytest.raises(FileNotFoundError, match="epoch_epoch=49.ckpt"): + train_hydra._prepare_checkpoint_resume(cfg) + + +@pytest.mark.parametrize("save_last", ["link", True]) +def test_save_last_link_writes_a_symlink_not_a_second_copy( + compose_resolve, tmp_path, save_last +): + """The point of the change, through a real fit: `true` writes last.ckpt as a + full second copy of the epoch checkpoint (~21 GB on Pi 0.5); "link" makes it + a symlink to the one already on disk.""" + import hydra + from lightning import LightningModule, Trainer + from torch.utils.data import DataLoader, TensorDataset + + class _M(LightningModule): + def __init__(self): + super().__init__() + self.layer = torch.nn.Linear(1, 1) + + def training_step(self, batch, batch_idx): + return self.layer(batch[0]).sum() + + def configure_optimizers(self): + return torch.optim.SGD(self.parameters(), lr=0.0) + + cfg = _pi_cfg(compose_resolve, tmp_path) + ckpt_dir = tmp_path / "checkpoints" + callback = hydra.utils.instantiate( + cfg.callbacks.model_checkpoint, + dirpath=str(ckpt_dir), + every_n_epochs=1, + save_last=save_last, + ) + Trainer( + max_epochs=2, + accelerator="cpu", + devices=1, + limit_train_batches=1, + callbacks=[callback], + logger=False, + enable_progress_bar=False, + enable_model_summary=False, + default_root_dir=str(tmp_path), + ).fit(_M(), DataLoader(TensorDataset(torch.zeros(2, 1)), batch_size=1)) + + last = ckpt_dir / "last.ckpt" + epochs = sorted(p.name for p in ckpt_dir.glob("epoch_*.ckpt")) + assert epochs == ["epoch_epoch=0.ckpt", "epoch_epoch=1.ckpt"], "save_top_k=-1" + assert last.is_file(), "either way the requeue path finds it" + + if save_last == "link": + assert last.is_symlink() + assert Path(os.readlink(last)) == Path("epoch_epoch=1.ckpt"), "relative" + assert last.resolve() == (ckpt_dir / "epoch_epoch=1.ckpt").resolve() + else: + # what we are getting rid of: a full second copy of the same checkpoint + assert not last.is_symlink() + assert last.stat().st_size == (ckpt_dir / "epoch_epoch=1.ckpt").stat().st_size From 64303ff9724d493aad7e222ee188920f6dce495b Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 11:09:55 -0400 Subject: [PATCH 30/79] eval: build the pinned-episode video batches only where they are used Two problems on the _video loaders, one of them a correctness bug. EvalVideo._write_video had no rank guard. The val heads come back from val_dataloader() inside CombinedLoaders, which Lightning does not wrap in a DistributedSampler, so every rank iterates every val batch and renders the same frames -- and every rank then wrote the same mp4 path. The file was whichever rank finished last, and could be torn. Separately, the video loaders were iterated in full on every validation pass. ModelWrapper.validation_step already returned before the forward on a non-viz pass, but Lightning had built the batch by then, so limit_val_batches x 64 pinned frames were decoded and discarded on every rank on the ~90% of passes that render nothing: ~1.7 min a pass, ~4% of a 1000-step cycle. _VizGatedLoader yields nothing when its gate is closed, which is on a non-viz pass and on every rank but 0. The val loaders are built once (_EvaluationLoop.setup_data returns early once _combined_loader is set), so the gate cannot be a length; it is an empty iterator, which _Sequential already reads as "move to the next loader". __len__ answers WITHOUT creating an iterator, because creating one forks that loader's workers and starts them prefetching -- doing the skipped work at setup instead, and never freeing it. It asks a fresh _MaxSizeCycle for the length, which stores its loaders and limits and builds no iterator until __iter__ (lightning 2.6.1). Ranks therefore iterate different numbers of val batches. That is safe because what they LOG is not: the metric loaders run everywhere and log the same keys, the video loader logs nothing, so the sync_dist reduction stays symmetric. It does mean the other ranks sit in that reduction while rank 0 renders -- 8.4 min for 76 batches back when all four ranks rendered at once, so well inside the 30-minute process-group timeout, but that timeout is now the ceiling on evaluator.viz_max_batches. trainHydra therefore refuses a head that has pinned video episodes and an evaluator with viz_max_batches=None. Every recipe that pins video episodes already caps it at 10; the evaluators that leave it unset (eval_hpt, eval_hpt_cartesian, eval_hpt_wrist) pair with configs that pin nothing, so they keep the ungated path where every rank renders together, as before. A desync here would hang a campaign rather than fail it, so it is tested against real DDP rather than argued: tests/fixtures/ddp_video_gate.py runs a 2-process gloo fit with the same collective pattern, at num_workers > 0 so the prefetch is visible. Both ranks finish; rank 0 steps the video loader on the viz epoch, rank 1 never touches the pinned episode. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01M5dTGNydAZbRP1EsXgbp4W --- egomimic/eval/eval_video.py | 12 +- egomimic/pl_utils/pl_data_utils.py | 97 +++++++++-- egomimic/trainHydra.py | 21 +++ tests/fixtures/ddp_video_gate.py | 130 +++++++++++++++ tests/unit/test_eval_loaders.py | 250 +++++++++++++++++++++++++++++ 5 files changed, 498 insertions(+), 12 deletions(-) create mode 100644 tests/fixtures/ddp_video_gate.py diff --git a/egomimic/eval/eval_video.py b/egomimic/eval/eval_video.py index 49042f9a99..3e0f188971 100644 --- a/egomimic/eval/eval_video.py +++ b/egomimic/eval/eval_video.py @@ -31,9 +31,10 @@ def __init__( # (current_epoch + 1) is a multiple of this; metrics log every pass. self.viz_every_n_epochs = viz_every_n_epochs # On viz epochs, render overlay frames for only the first N val - # batches (None = all). Rendering is CPU-bound (~1s/frame) and - # dominates viz-epoch wall time; metrics are still computed on EVERY - # batch regardless. + # batches (None = all). Rendering is CPU-bound -- ~0.1 s/frame, 4864 + # frames in 8.4 min with four ranks rendering at once -- and dominates + # viz-epoch wall time; metrics are still computed on EVERY batch + # regardless. Required on heads with pinned video loaders (trainHydra). self.viz_max_batches = viz_max_batches # Per-embodiment list[Transform] applied once during eval to project # the model's wrist-frame actions back into cam (head) frame for the @@ -101,6 +102,11 @@ def _video_fps(self, source_fps: int = 30) -> int: return int(source_fps) def _write_video(self, key, frames) -> None: + # Every rank that gets here renders the same frames to the same path + # (see _video_fps on why no DistributedSampler splits them), so without + # this the file is whatever the last rank to finish wrote. + if not self.trainer.is_global_zero: + return path = os.path.join( self.video_dir(), f"epoch_{self.trainer.current_epoch}", diff --git a/egomimic/pl_utils/pl_data_utils.py b/egomimic/pl_utils/pl_data_utils.py index aa4fe2c7e8..b41a6b64d6 100644 --- a/egomimic/pl_utils/pl_data_utils.py +++ b/egomimic/pl_utils/pl_data_utils.py @@ -1,7 +1,7 @@ import logging from lightning import LightningDataModule -from lightning.pytorch.utilities.combined_loader import CombinedLoader +from lightning.pytorch.utilities.combined_loader import CombinedLoader, _MaxSizeCycle from torch.utils.data import DataLoader, default_collate logger = logging.getLogger(__name__) @@ -25,6 +25,43 @@ def head_of_loader(name: str) -> tuple[str, bool]: return name, False +class _VizGatedLoader(CombinedLoader): + """A val loader that yields nothing when its gate is closed. + + ``ModelWrapper.validation_step`` already returns before the forward for a + ``_video`` loader it does not want, but Lightning has built the batch by + then, so the pinned frames are decoded and discarded anyway. + + The val loaders are built once (``_EvaluationLoop.setup_data`` returns early + once ``_combined_loader`` is set), so the gate cannot be a length; it is an + empty iterator. ``_Sequential`` already reads an early ``StopIteration`` from + one dataloader as "move to the next". + """ + + def __init__(self, iterables, mode, gate): + if mode != "max_size_cycle": + raise ValueError(f"only max_size_cycle is gated, not {mode!r}") + super().__init__(iterables, mode) + self._gate = gate + + def __iter__(self): + if not self._gate(): + return iter(()) + return super().__iter__() + + def __len__(self) -> int: + if self._iterator is not None: + return super().__len__() + # `CombinedLoader.__len__` reads a live iterator, and `__iter__` is what + # forks this loader's workers and starts them prefetching. Lightning asks + # for the length once per rank in `setup_data`, where the gate is already + # closed on every rank but 0, so answering WITHOUT an iterator is what + # keeps those ranks off the pinned episode. The length does not depend on + # the gate. `_MaxSizeCycle` (the mode is pinned in __init__) only stores + # its iterables until __iter__, so asking it forks nothing. + return len(_MaxSizeCycle(self.flattened, self.limits)) + + class MultiDataModuleWrapper(LightningDataModule): """ New functionality for dictionary based multi embodiment loading using CombinedLoader. @@ -146,7 +183,9 @@ def train_dataloader(self): return CombinedLoader(iterables, "max_size_cycle") - def _build_val_style_loader(self, datasets: dict, params: dict, kind: str): + def _build_val_style_loader( + self, datasets: dict, params: dict, kind: str, gate=None + ): iterables = dict() for dataset_name, dataset in datasets.items(): dataset_params = params.get(dataset_name) @@ -162,7 +201,41 @@ def _build_val_style_loader(self, datasets: dict, params: dict, kind: str): collate_fn=self.collate_fn, **dataset_params, ) - return CombinedLoader(iterables, "max_size_cycle") + if gate is None: + return CombinedLoader(iterables, "max_size_cycle") + return _VizGatedLoader(iterables, "max_size_cycle", gate) + + def _video_gate(self, head: str) -> bool: + """Does this validation pass render ``head``'s overlay video? + + Open only on rank 0 and only on a viz pass. Elsewhere the loader's whole + output -- decode, eval forward, render -- is discarded: the val loaders + carry no ``DistributedSampler`` (see ``EvalVideo._video_fps``), so every + rank was producing the same frames and only rank 0 ever wrote them. + + Ranks therefore iterate different numbers of val batches. That is safe + because what they log is not: the metric loaders run everywhere and log + the same keys, the video loader logs nothing, so the ``sync_dist`` + reduction in ``_evaluation_epoch_end`` stays symmetric. It does mean the + other ranks sit in that reduction for as long as rank 0 renders. Measured + 8.4 min for 76 batches back when all four ranks rendered at once, so rank + 0 alone is inside that and well under the 30-minute process-group timeout + -- but that timeout is now the ceiling on ``evaluator.viz_max_batches``, + which trainHydra therefore requires on every head with a video loader. + + Reads the same ``_should_viz()`` ``ModelWrapper.validation_step`` does, + back through the trainer so the datamodule needs no extra wiring, and + falls open whenever a link is missing: eval-only mode, tests, an + evaluator that predates the hook. + """ + trainer = getattr(self, "trainer", None) + if trainer is not None and not getattr(trainer, "is_global_zero", True): + return False + module = getattr(trainer, "lightning_module", None) + val_heads = getattr(module, "_val_heads", None) + evaluator = val_heads().get(head) if callable(val_heads) else None + should_viz = getattr(evaluator, "_should_viz", None) + return should_viz() if callable(should_viz) else True def _metric_sources(self) -> dict: return { @@ -198,12 +271,18 @@ def val_dataloader(self): # unshuffled); only the dataset differs. sources[video_loader_name(head)] = (datasets, sources[head][1]) names = self.val_loader_names() - loaders = [ - # kind names the *_dataloader_params block in the error message, and - # a video loader borrows its head's block. - self._build_val_style_loader(*sources[name], kind=head_of_loader(name)[0]) - for name in names - ] + loaders = [] + for name in names: + head, video_only = head_of_loader(name) + loaders.append( + # kind names the *_dataloader_params block in the error message, + # and a video loader borrows its head's block. + self._build_val_style_loader( + *sources[name], + kind=head, + gate=(lambda h=head: self._video_gate(h)) if video_only else None, + ) + ) if len(loaders) == 1: return loaders[0] # Several heads: return a list so Lightning populates dataloader_idx diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 9326031106..224a3fc13e 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -475,6 +475,26 @@ def _video_datasets(cfg: DictConfig, heads: dict) -> dict: return out +def _require_capped_video_heads(model, datamodule) -> None: + """A pinned video loader renders on rank 0 alone while the other ranks wait + in the val reduction, so an uncapped render (``viz_max_batches=None``) is + bounded only by the 30-minute process-group timeout.""" + heads = model._val_heads() + for head, datasets in (getattr(datamodule, "video_datasets", None) or {}).items(): + evaluator = heads.get(head) + if ( + datasets + and getattr(evaluator, "viz_func", None) is not None + and evaluator.viz_max_batches is None + ): + raise ValueError( + f"val head '{head}' has pinned video episodes " + f"(data.video_episodes.{head}) but its evaluator sets no " + "viz_max_batches: rank 0 would render every video batch while " + "the other ranks wait on it. Set evaluator.viz_max_batches." + ) + + def _build_unseen_op_valid_evaluator(cfg: DictConfig): """The canonical evaluator (fresh instance, own frame buffers) wrapped to log ``unseen_op_valid/`` and write ``videos_unseen_op_valid/``.""" @@ -786,6 +806,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: unseen_eval_obj.trainer = trainer unseen_eval_obj.model = model.model model.unseen_op_valid_evaluator = unseen_eval_obj + _require_capped_video_heads(model, datamodule) model.val_loader_names = datamodule.val_loader_names() # Pre-fit baseline val. Skipped on requeues AND checkpoint resumes: # trainer.validate here runs BEFORE fit restores ckpt_path weights, so diff --git a/tests/fixtures/ddp_video_gate.py b/tests/fixtures/ddp_video_gate.py new file mode 100644 index 0000000000..75cd906285 --- /dev/null +++ b/tests/fixtures/ddp_video_gate.py @@ -0,0 +1,130 @@ +"""Run a 2-rank DDP fit where only rank 0 iterates the video loader. + +Launched as a subprocess by tests/unit/test_eval_loaders.py. Each rank writes +what it saw to ``/rank.json``; the test reads them back. Module scope +(not a closure) so the ddp launcher can re-import it. +""" + +from __future__ import annotations + +import json +import multiprocessing as mp +import sys +from pathlib import Path + +import numpy as np +import torch +from lightning import LightningModule, Trainer +from lightning.fabric.plugins.environments import LightningEnvironment + +from egomimic.pl_utils.pl_data_utils import MultiDataModuleWrapper +from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + +HUMAN = "human_bimanual" + + +class _Episode: + """Counts reads in shared memory: the loader workers are other processes.""" + + def __init__(self, n: int): + self.n = n + self.reads = mp.Value("i", 0) + + def __len__(self): + return self.n + + def __getitem__(self, i): + with self.reads.get_lock(): + self.reads.value += 1 + return {"frame": np.float32(i)} + + +class _Module(LightningModule): + """Mimics ModelWrapper's collective pattern: sync_dist metrics off the + metric loader, nothing at all off the video loader.""" + + def __init__(self, pinned): + super().__init__() + self.layer = torch.nn.Linear(1, 1) + self.pinned = pinned + self.video_steps: list[int] = [] + + def _val_heads(self): + return {"valid": _Viz(self)} + + def training_step(self, batch, batch_idx): + return self.layer(batch[HUMAN]["frame"].reshape(-1, 1)).sum() + + def validation_step(self, batch, batch_idx, dataloader_idx=0): + if dataloader_idx == 1: + self.video_steps.append(self.current_epoch) + return # the video loader logs nothing, as EvalVideo does + self.log("val/mse", torch.tensor(1.0), sync_dist=True, add_dataloader_idx=False) + + def on_validation_end(self): + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.barrier() + + def configure_optimizers(self): + return torch.optim.SGD(self.parameters(), lr=0.0) + + +class _Viz: + def __init__(self, module): + self.module = module + + def _should_viz(self) -> bool: + return self.module.current_epoch == 1 + + +def main(out_dir: str) -> None: + pinned, metric = _Episode(8), _Episode(8) + dm = MultiDataModuleWrapper( + train_datasets={ + HUMAN: MultiDataset(datasets={"ep0": _Episode(8)}, mode="total") + }, + valid_datasets={HUMAN: MultiDataset(datasets={"ep0": metric}, mode="total")}, + train_dataloader_params={HUMAN: {"batch_size": 2, "num_workers": 0}}, + # num_workers > 0 on purpose: with 0 a DataLoader prefetches nothing, so + # a gate that still builds an iterator would read no frame and look clean. + valid_dataloader_params={ + HUMAN: {"batch_size": 2, "num_workers": 2, "shuffle": False} + }, + video_datasets={ + "valid": {HUMAN: MultiDataset(datasets={"pin": pinned}, mode="total")} + }, + ) + model = _Module(pinned) + trainer = Trainer( + strategy="ddp", + # force the subprocess launcher: run under `srun pytest` lightning would + # otherwise detect SLURM and wait for ranks srun was never asked for + plugins=[LightningEnvironment()], + accelerator="cpu", + devices=2, + max_epochs=2, + limit_train_batches=2, + limit_val_batches=2, + check_val_every_n_epoch=1, + num_sanity_val_steps=2, + logger=False, + enable_checkpointing=False, + enable_progress_bar=False, + enable_model_summary=False, + default_root_dir=out_dir, + ) + trainer.fit(model, datamodule=dm) + Path(out_dir, f"rank{trainer.global_rank}.json").write_text( + json.dumps( + { + "rank": trainer.global_rank, + "video_steps": model.video_steps, + "pinned_reads": pinned.reads.value, + "finished": True, + } + ) + ) + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/tests/unit/test_eval_loaders.py b/tests/unit/test_eval_loaders.py index 24e77fa740..3dc822589d 100644 --- a/tests/unit/test_eval_loaders.py +++ b/tests/unit/test_eval_loaders.py @@ -16,7 +16,10 @@ from __future__ import annotations +import os import re +import sys +from pathlib import Path from types import SimpleNamespace import numpy as np @@ -558,3 +561,250 @@ def test_default_mode_is_todays_behaviour(): ev = _eval_video() ev.on_validation_step({"i": 0}, 0, 0) assert ev.rendered == [0] and len(ev.trainer.logged) == 1 + + +# --------------------------------------------- video writes happen on rank 0 only +def test_only_rank_zero_writes_the_video(monkeypatch, tmp_path): + """Every rank iterates every val batch and renders the same frames, so + without a guard N ranks race to write one mp4 path.""" + import egomimic.eval.eval_video as ev_mod + + written: list[str] = [] + monkeypatch.setattr( + ev_mod.tvio, "write_video", lambda path, *a, **kw: written.append(path) + ) + + buffers = {} + for rank_zero in (True, False): + ev = _eval_video(viz_max_batches=None) + ev.trainer.is_global_zero = rank_zero + monkeypatch.setattr(type(ev), "video_dir", lambda self: str(tmp_path)) + for i in range(3): + ev.on_validation_step({"i": i}, i, 0, mode="video") + ev.on_validation_end() + # rendering, buffering and counters are identical on every rank... + buffers[rank_zero] = (ev.rendered, dict(ev.val_counter), ev.val_image_buffer) + + assert buffers[True] == buffers[False], "only the write differs between ranks" + assert len(written) == 1, "one file, from rank 0" + + +# ------------------------------------------- video loaders are gated on non-viz +class _GateTrainer: + """Just enough trainer for MultiDataModuleWrapper._video_gate.""" + + def __init__(self, evaluators): + self.lightning_module = SimpleNamespace(_val_heads=lambda: evaluators) + self.is_global_zero = True + + +def _dm(): + """A datamodule with one metric head and its pinned video companion.""" + params = {HUMAN: {"batch_size": 2, "num_workers": 0, "shuffle": False}} + datasets = {HUMAN: _split({"ep0": 4})} + return MultiDataModuleWrapper( + train_datasets=datasets, + valid_datasets=datasets, + train_dataloader_params=params, + valid_dataloader_params=params, + video_datasets={"valid": {HUMAN: _split({"pin": 4})}}, + ) + + +def test_video_loader_yields_nothing_when_the_gate_is_closed(): + from egomimic.pl_utils.pl_data_utils import _VizGatedLoader + + dm = _dm() + metric, video = dm.val_dataloader() + assert not isinstance(metric, _VizGatedLoader), "metric loaders run every pass" + assert isinstance(video, _VizGatedLoader) + + viz = _StubEval(should_viz=True) + dm.trainer = _GateTrainer({"valid": viz}) + assert len(list(iter(video))) == 2, "viz pass: the pinned episode is read" + + viz._viz = False + assert list(iter(video)) == [], "non-viz pass: not one batch is built" + # the length Lightning caches at setup does not move with the gate + assert len(video) == 2 + + +def test_gate_is_open_without_a_trainer_or_an_evaluator(): + dm = _dm() + _, video = dm.val_dataloader() + assert len(list(iter(video))) == 2, "no trainer attached -> gate open" + + dm.trainer = _GateTrainer({"valid": None}) + assert len(list(iter(video))) == 2, "head with no evaluator -> gate open" + + dm.trainer = _GateTrainer({"valid": object()}) + assert len(list(iter(video))) == 2, "evaluator without _should_viz -> gate open" + + +def test_gated_loader_length_survives_a_closed_gate_at_setup(): + """Lightning reads len() once, in setup_data, which can land on a non-viz + epoch: CombinedLoader.__len__ needs a live iterator the gate never made.""" + dm = _dm() + _, video = dm.val_dataloader() + dm.trainer = _GateTrainer({"valid": _StubEval(should_viz=False)}) + assert list(iter(video)) == [] + assert len(video) == 2 + + +def test_gated_loader_length_is_lightnings_own_and_opens_no_iterator(): + from lightning.pytorch.utilities.combined_loader import _MaxSizeCycle + + dm = _dm() + _, video = dm.val_dataloader() + video.limits = [1] + assert len(video) == len(_MaxSizeCycle(video.flattened, [1])) == 1 + assert video._iterator is None + + +def test_video_head_without_a_viz_cap_is_rejected(): + """Rank 0 renders a pinned video loader alone while the other ranks wait, so + an uncapped head is bounded only by the process-group timeout.""" + capped = SimpleNamespace(viz_func=print, viz_max_batches=10) + uncapped = SimpleNamespace(viz_func=print, viz_max_batches=None) + no_viz = SimpleNamespace(viz_func=None, viz_max_batches=None) + + def _run(heads, video): + model = SimpleNamespace(_val_heads=lambda: heads) + th._require_capped_video_heads(model, SimpleNamespace(video_datasets=video)) + + _run({"valid": capped}, {"valid": {HUMAN: object()}}) + _run({"valid": uncapped}, {"valid": {}}) + _run({"valid": uncapped}, {}) + _run({"valid": no_viz}, {"valid": {HUMAN: object()}}) + with pytest.raises(ValueError, match="viz_max_batches"): + _run({"valid": capped, "train_viz": uncapped}, {"train_viz": {HUMAN: object()}}) + + +# sanity steps on: the loaders (and their len()) are set up on a pass whose gate +# is already closed, which is what a real run does -- trainer/default.yaml sets no +# num_sanity_val_steps, so lightning's default of 2 applies. +@pytest.mark.parametrize("sanity_steps", [0, 2]) +def test_a_real_fit_builds_video_batches_only_on_viz_passes(tmp_path, sanity_steps): + """The gate against the real Lightning evaluation loop: a val dataloader + that yields nothing mid-sequence must not derail the pass, and the pinned + episode must not be touched at all on a non-viz epoch.""" + import torch + from lightning import Trainer + + class _Counting(_Episode): + def __init__(self, n): + super().__init__(n) + self.reads = 0 + + def __getitem__(self, i): + self.reads += 1 + return super().__getitem__(i) + + pinned, metric = _Counting(4), _Counting(4) + params = {HUMAN: {"batch_size": 2, "num_workers": 0, "shuffle": False}} + dm = MultiDataModuleWrapper( + train_datasets={HUMAN: _split({"ep0": 4})}, + valid_datasets={HUMAN: MultiDataset(datasets={"ep0": metric}, mode="total")}, + train_dataloader_params={HUMAN: {"batch_size": 2, "num_workers": 0}}, + valid_dataloader_params=params, + video_datasets={ + "valid": {HUMAN: MultiDataset(datasets={"pin": pinned}, mode="total")} + }, + ) + + class _Module(LightningModule): + def __init__(self): + super().__init__() + self.layer = torch.nn.Linear(1, 1) + self.val_loaders_seen: list[tuple[int, int]] = [] + self.pinned_reads: dict[int, int] = {} + + # viz only on the second validation pass + def _val_heads(self): + return { + "valid": SimpleNamespace(_should_viz=lambda: self.current_epoch == 1) + } + + def training_step(self, batch, batch_idx): + return self.layer(batch[HUMAN]["frame"].reshape(-1, 1)).sum() + + def validation_step(self, batch, batch_idx, dataloader_idx=0): + if not self.trainer.sanity_checking: + self.val_loaders_seen.append((self.current_epoch, dataloader_idx)) + + def on_validation_epoch_end(self): + self.pinned_reads[self.current_epoch] = pinned.reads + + def configure_optimizers(self): + return torch.optim.SGD(self.parameters(), lr=0.0) + + model = _Module() + Trainer( + max_epochs=2, + accelerator="cpu", + devices=1, + limit_train_batches=2, + limit_val_batches=2, + check_val_every_n_epoch=1, + num_sanity_val_steps=sanity_steps, + logger=False, + enable_checkpointing=False, + enable_progress_bar=False, + enable_model_summary=False, + default_root_dir=str(tmp_path), + ).fit(model, datamodule=dm) + + assert model.pinned_reads[0] == 0, "non-viz pass never touched the pinned episode" + assert model.pinned_reads[1] > 0, "viz pass still reads it" + # the metric loader is unaffected, and dataloader_idx stays stable + by_epoch = {e: [i for ep, i in model.val_loaders_seen if ep == e] for e in (0, 1)} + assert by_epoch[0] == [0, 0], "epoch 0: metric loader only" + assert by_epoch[1] == [0, 0, 1, 1], "epoch 1: metric loader then video loader" + + +def test_video_loader_is_closed_off_rank_zero(): + """Only rank 0 writes the overlay, and the val loaders carry no + DistributedSampler, so every other rank was decoding, forwarding and + rendering the same frames for nothing.""" + dm = _dm() + _, video = dm.val_dataloader() + viz = _StubEval(should_viz=True) + + dm.trainer = _GateTrainer({"valid": viz}) + dm.trainer.is_global_zero = True + assert len(list(iter(video))) == 2 + + dm.trainer.is_global_zero = False + assert list(iter(video)) == [], "rank 1 builds no video batch at all" + + +def test_two_rank_ddp_fit_survives_the_rank_divergence(tmp_path): + """The gate makes ranks iterate *different* numbers of val batches. Proven + against real DDP (gloo, 2 CPU processes) with ModelWrapper's collective + pattern -- sync_dist metrics off the metric loader, nothing off the video + one -- because a desync here would hang a campaign, not fail it.""" + import json + import subprocess + + script = Path(__file__).resolve().parents[1] / "fixtures" / "ddp_video_gate.py" + repo = Path(__file__).resolve().parents[2] + proc = subprocess.run( + [sys.executable, str(script), str(tmp_path)], + env={**os.environ, "PYTHONPATH": str(repo)}, + capture_output=True, + text=True, + timeout=600, + ) + assert proc.returncode == 0, proc.stderr[-3000:] + + ranks = { + int(p.stem.removeprefix("rank")): json.loads(p.read_text()) + for p in tmp_path.glob("rank*.json") + } + assert set(ranks) == {0, 1}, "both ranks finished; neither hung" + # epoch 0 renders nothing, epoch 1 does, and only on rank 0 + assert ranks[0]["video_steps"] == [1, 1] + assert ranks[0]["pinned_reads"] >= 4, "rank 0 read the frames it stepped" + # exactly 0, not "few": a gate that still builds an iterator makes the + # loader's workers prefetch, which is the cost this whole change is about + assert ranks[1]["video_steps"] == [] and ranks[1]["pinned_reads"] == 0 From 119657d7c9fc6b1fe4637a237eac2873893c9634 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 11:10:34 -0400 Subject: [PATCH 31/79] loader: one zarr read per key, not one per keymap entry ZarrDataset.__getitem__ issued a separate ZarrEpisode.read for every keymap entry, and several entries name the same zarr key with different windows: in the mecka keypoint keymap at proprio_history=4, left/right.obs_keypoints and the two wrist poses are each read once as a forward action chunk (i, i+H) and once as a backward proprio window (i-K+1, i+1). Ten reads over six distinct keys, each paying zarr's per-call metadata hop. It now makes one zarr call per key and slices each window out of what it read. Measured on a real mecka episode with that keymap, back to back in one process: 27.5 -> 15.6 ms a sample on the read path (1.76x), one read call instead of ten and six metadata hops instead of ten. read is ~70% of __getitem__, which is also what sets the epoch-boundary stall, since an epoch's first batch is one worker building the whole batch. Windows on one key merge into a single contiguous read only when every row between them lies in a chunk one of them already touches (ZarrEpisode.plan_reads); otherwise the groups are fetched together with one orthogonal index. Numeric keys are chunked 100 rows, so their windows always merge and the numbers above are this path. JPEG keys are one frame per chunk, where a plain union would read and ship every frame between two windows far apart -- the image-history window #678 adds on the front camera is exactly that case. Output is unchanged, except that an array shorter than a window now raises a ValueError naming the key; the old per-entry slice returned the short rows and padded them, and a lookup miss here would have escaped as StopIteration, which the DataLoader reads as the end of the epoch. The union is sliced first and padded second, so a chunk running off the end of the episode pads exactly as before, and each slice is copied because a slice of a shared span is a view where a per-entry read was not -- transforms are handed these arrays and may write into them. The test compares against a per-entry oracle over a RAMPED episode ([t, d] = t*1000 + d) rather than the shared synthetic fixture, which holds every numeric key constant in time and so cannot see a wrong frame offset at all. 5 keymaps x 9 indices x {ZarrDataset, ZarrAnnotationCutoffDataset}, plus the two paths nothing else in the repo covers: a key read as both a window and a single frame, and the JPEG-decode retry, which changes idx and so must recompute the windows and the read. The oracle compares values, so it would pass a loader that read the whole episode; separate tests count the ROWS read per key and pin plan_reads on both chunk sizes. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01M5dTGNydAZbRP1EsXgbp4W --- egomimic/rldb/zarr/zarr_dataset_multi.py | 174 ++++++++--- tests/unit/test_loader_read_coalescing.py | 342 ++++++++++++++++++++++ 2 files changed, 481 insertions(+), 35 deletions(-) create mode 100644 tests/unit/test_loader_read_coalescing.py diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index a5fff53a59..6eecbc61bb 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -2149,7 +2149,7 @@ def _action_horizon(self) -> int | None: def _action_pad_mask(self, idx: int) -> np.ndarray | None: """1.0 per real action frame, 0.0 per repeat-last padded one, (horizon,). - ``_pad_sequences`` pads a short action chunk by repeating the last + ``_pad_to_horizon`` pads a short action chunk by repeating the last frame, which in a wrist-relative action space supervises "stop moving" at every episode (or annotation) tail. This mask says which steps are real so ``loss_fn`` can drop the rest; ``None`` when the keymap has no @@ -2199,7 +2199,7 @@ def _stride_history(array: np.ndarray, history: int, stride: int) -> np.ndarray: def _front_pad_history(array: np.ndarray, history: int) -> np.ndarray: """Front-pad a backward window to ``history`` steps by repeating its first real frame, so the current frame is always ``[-1]``. The mirror - of ``_pad_sequences``, which repeat-pads the TAIL of a forward chunk. + of ``_pad_to_horizon``, which repeat-pads the TAIL of a forward chunk. """ seq_len = array.shape[0] if seq_len >= history: @@ -2207,23 +2207,56 @@ def _front_pad_history(array: np.ndarray, history: int) -> np.ndarray: padding = np.repeat(array[:1], history - seq_len, axis=0) return np.concatenate([padding, array], axis=0) + @staticmethod + def _pad_to_horizon(array, horizon: int | None): + """Repeat-last out to ``horizon`` a chunk the episode (or annotation) + ended before filling.""" + if horizon is None or not isinstance(array, np.ndarray): + return array + missing = horizon - array.shape[0] + if missing <= 0: + return array + return np.concatenate([array, np.repeat(array[-1:], missing, axis=0)], axis=0) + def _pad_sequences(self, data, horizon: int | None) -> dict: + """``_pad_to_horizon`` over a read result, in place. Keys are zarr keys.""" if horizon is None: return data - - # Note that k is zarr key for k in data: - if isinstance(data[k], np.ndarray): - seq_len = data[k].shape[0] - if seq_len < horizon: - # Pad by repeating the last frame - pad_len = horizon - seq_len - last_frame = data[k][-1:] # Keep dims: (1, action_dim) - padding = np.repeat(last_frame, pad_len, axis=0) - data[k] = np.concatenate([data[k], padding], axis=0) - + data[k] = self._pad_to_horizon(data[k], horizon) return data + def _read_windows(self, idx: int) -> dict[str, tuple[str, tuple[int, int | None]]]: + """``{keymap key: (zarr key, (start, end))}`` for frame *idx*. + + ``end`` None means the single frame at ``start``. Annotation keys read + nothing and are left out. + """ + windows = {} + for k, spec in self.key_map.items(): + if spec.get("key_type") == "annotation_keys": + continue + horizon = spec.get("horizon") + history = spec.get("history") + if horizon is not None: + interval = ( + idx, + self._chunk_end_idx(idx, horizon, spec.get("key_type")), + ) + elif history is not None and int(history) > 1: + # Backward window ending at (and including) the current frame, + # every `history_stride`-th step; clamped at the episode start, + # thinned and front-padded by the caller, so the key is always + # (K, D) with `[..., -1]` the single frame a K = 1 keymap would + # have read. Read contiguously: numeric keys are chunked 100 rows, + # so the skipped rows sit in chunks the read fetches anyway. + stride = int(spec.get("history_stride", 1) or 1) + interval = (max(0, idx - stride * (int(history) - 1)), idx + 1) + else: + interval = (idx, None) + windows[k] = (spec["zarr_key"], interval) + return windows + def __getitem__( self, idx: int, @@ -2257,6 +2290,21 @@ def _next(reason: str, key: str = "") -> int: return next_idx while True: + windows = self._read_windows(idx) + # One zarr call per key, not one per keymap entry: an action chunk + # and the proprio history of the same pose name the same key with + # different windows, and zarr's per-call metadata hop costs more + # than the handful of extra rows a merged read carries. Windows are + # merged only across rows in chunks the read fetches anyway (see + # ZarrEpisode.plan_reads), so two frames far apart on a JPEG key + # stay two rows. + intervals: dict[str, list[tuple[int, int]]] = {} + for zarr_key, (start, end) in windows.values(): + intervals.setdefault(zarr_key, []).append( + (start, start + 1 if end is None else end) + ) + groups = self.episode_reader.read_intervals(intervals) + data = {} retry = False for k in self.key_map: @@ -2269,29 +2317,42 @@ def _next(reason: str, key: str = "") -> int: data[k] = self._annotation_text_for_frame(idx) continue - if horizon is not None: - end_idx = self._chunk_end_idx(idx, horizon, key_type) - read_interval = (idx, end_idx) - elif history is not None and int(history) > 1: - # Backward window ending at (and including) the current - # frame, every `history_stride`-th step; clamped at the - # episode start and front-padded below, so the key is - # always (K, D) with `[..., -1]` the single frame a K = 1 - # keymap would have read. Read contiguously and thin after: - # the zarr chunks are contiguous anyway. - stride = int(self.key_map[k].get("history_stride", 1) or 1) - read_interval = ( - max(0, idx - stride * (int(history) - 1)), - idx + 1, + start, end = windows[k][1] + stop = start + 1 if end is None else end + found = next( + ( + (lo, rows) + for lo, rows in groups[zarr_key] + if lo <= start and stop <= lo + len(rows) + ), + None, + ) + if found is None: + # A StopIteration escaping __getitem__ reads as "epoch + # over" to the DataLoader; a short array must fail loudly. + raise ValueError( + f"ep={Path(self.episode_path).name} key={zarr_key}: " + f"window [{start}, {stop}) is past the rows read " + f"{[(lo, lo + len(rows)) for lo, rows in groups[zarr_key]]}" ) - else: - read_interval = (idx, None) - read_dict = {zarr_key: read_interval} - raw_data = self.episode_reader.read(read_dict) - self._pad_sequences(raw_data, horizon) # should be able to pad images - data[k] = raw_data[zarr_key] + base, block = found + chunk = ( + block[start - base] + if end is None + else block[start - base : end - base] + ) + if isinstance(chunk, np.ndarray): + # A slice of a shared read is a view; every keymap entry + # owned its own array before this coalescing and transforms + # are free to write into it. + chunk = chunk.copy() + data[k] = self._pad_to_horizon(chunk, horizon) if horizon is None and history is not None and int(history) > 1: - data[k] = self._stride_history(data[k], int(history), stride) + data[k] = self._stride_history( + data[k], + int(history), + int(self.key_map[k].get("history_stride", 1) or 1), + ) if zarr_key in self._image_keys: jpeg_bytes = data[k] @@ -2371,7 +2432,7 @@ class ZarrAnnotationCutoffDataset(ZarrDataset): Standard chunking from the start frame, but action reads stop at EOS+1 of the annotation span containing the start frame. The chunk is then padded out to - ``horizon`` via the base ``_pad_sequences`` (repeat-last), so frames beyond + ``horizon`` via the base ``_pad_to_horizon`` (repeat-last), so frames beyond EOS become the last action of the interval rather than crossing into the next annotation. @@ -2562,6 +2623,49 @@ def read( result[key] = data return result + @staticmethod + def plan_reads( + intervals: list[tuple[int, int]], chunk_len: int + ) -> list[tuple[int, int]]: + """Merge half-open ``[start, end)`` intervals on one key into read groups. + + Two merge only when every row between them lies in a chunk one of them + already touches, so the bridge costs no extra chunk fetch: the windows + on a numeric key (100-row chunks) collapse into one read, while two + frames far apart on a JPEG key (one frame per chunk) stay two rows + rather than every frame in between.""" + groups: list[list[int]] = [] + for start, end in sorted(intervals): + if groups and start // chunk_len <= (groups[-1][1] - 1) // chunk_len + 1: + groups[-1][1] = max(groups[-1][1], end) + else: + groups.append([start, end]) + return [(lo, hi) for lo, hi in groups] + + def read_intervals( + self, intervals: dict[str, list[tuple[int, int]]] + ) -> dict[str, list[tuple[int, np.ndarray]]]: + """One zarr call per key for all its ``[start, end)`` intervals. + + Returns ``{key: [(group_start, rows), ...]}``, one entry per group of + ``plan_reads``; several groups are fetched with a single orthogonal + index.""" + result = {} + for key, ivs in intervals.items(): + arr = self._store[key] + groups = self.plan_reads(ivs, arr.chunks[0]) + if len(groups) == 1: + lo, hi = groups[0] + result[key] = [(lo, arr[lo:hi])] + continue + rows = arr.oindex[np.concatenate([np.arange(lo, hi) for lo, hi in groups])] + out, offset = [], 0 + for lo, hi in groups: + out.append((lo, rows[offset : offset + hi - lo])) + offset += hi - lo + result[key] = out + return result + def _collect_keys(self) -> list[str]: """ Collect all array keys from the store. diff --git a/tests/unit/test_loader_read_coalescing.py b/tests/unit/test_loader_read_coalescing.py new file mode 100644 index 0000000000..3c8929e044 --- /dev/null +++ b/tests/unit/test_loader_read_coalescing.py @@ -0,0 +1,342 @@ +"""``ZarrDataset.__getitem__`` reads each zarr key once, not once per keymap entry. + +The mecka keypoint keymap with ``proprio_history=4`` names ``left/right.obs_keypoints`` +and ``left/right.obs_wrist_pose`` twice each -- once as a forward action chunk +``(i, i+H)``, once as a backward proprio window ``(i-K+1, i+1)``. Those were ten +separate ``ZarrEpisode.read`` calls over six distinct keys; they are now one zarr call +per key (``ZarrEpisode.read_intervals``), sliced twice. Nothing the dataset emits may change. + +The episode here is RAMPED -- every cell is ``frame * 1000 + column`` -- so a slice +taken at the wrong offset into the coalesced span produces different numbers rather +than the same ones. ``tests/fixtures/synthetic_episodes.py`` deliberately holds each +numeric key constant in time (for the norm-stat quantile bounds), which would make an +off-by-one here invisible: the single most likely bug in this change. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import simplejpeg +import zarr + +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.zarr.zarr_dataset_multi import ( + ZarrAnnotationCutoffDataset, + ZarrDataset, + ZarrEpisode, +) +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +T = 160 +HORIZON = Human.ACTION_HORIZON +# half-open [start, end); ZarrAnnotationCutoffDataset clamps action chunks at the +# end of the span holding the start frame, so idx 60 and 130 both hit a cutoff +ANN_SPANS = ((0, 80), (80, T - 1)) + +# (label, keymap kwargs). Covers both action spaces, history on and off, and the +# ee_pose variant that gives four entries over two shared keys. +KEYMAPS = { + "keypoints_k4": {"keymap_mode": "keypoints", "proprio_history": 4}, + "keypoints_k1": {"keymap_mode": "keypoints", "proprio_history": 1}, + "keypoints_k4_ee": { + "keymap_mode": "keypoints", + "proprio_history": 4, + "include_ee_pose": True, + }, + "cartesian_k4": {"keymap_mode": "cartesian", "proprio_history": 4}, + "cartesian_k1": {"keymap_mode": "cartesian", "proprio_history": 1}, +} + + +def _keymap(label): + return Human.get_keymap( + has_head_pose=True, annotation_key="annotations", **KEYMAPS[label] + ) + + +def _ramp(width: int) -> np.ndarray: + """``[t, d] = t * 1000 + d`` -- every cell identifies its own frame.""" + return ( + np.arange(T, dtype=np.float64)[:, None] * 1000.0 + + np.arange(width, dtype=np.float64)[None, :] + ) + + +@pytest.fixture(scope="module") +def episode(tmp_path_factory): + numeric = {"obs_head_pose": _ramp(7)} + for side in ("left", "right"): + # distinct per side so a left/right mix-up is caught too + off = 0.0 if side == "left" else 0.5 + numeric[f"{side}.obs_wrist_pose"] = _ramp(7) + off + numeric[f"{side}.obs_keypoints"] = _ramp(63) + off + numeric[f"{side}.obs_ee_pose"] = _ramp(7) + off + rng = np.random.default_rng(0) + return ZarrWriter.create_and_write( + tmp_path_factory.mktemp("ramped") / "ramped.zarr", + numeric_data=numeric, + image_data={"images.front_1": rng.integers(0, 255, (T, 32, 32, 3), np.uint8)}, + embodiment="human_bimanual", + task_name="ramp", + annotations=[(t, a, b) for t, (a, b) in zip(("fold", "put down"), ANN_SPANS)], + intrinsics={"front_1": np.eye(3, 4)}, + ) + + +def _ann_end(idx: int) -> int | None: + return next((e for s, e in ANN_SPANS if s <= idx < e), None) + + +def _expected_window(spec, idx, total, cutoff: bool): + """The read window for one keymap entry, written out independently of the + dataset: the spec this change must not move.""" + horizon, history = spec.get("horizon"), spec.get("history") + if horizon is not None: + end = min(idx + horizon, total) + if cutoff and spec.get("key_type") == "action_keys": + end = min(end, _ann_end(idx) or end) + return (idx, end) + if history is not None and int(history) > 1: + return (max(0, idx - int(history) + 1), idx + 1) + return (idx, None) + + +def _oracle(episode, key_map, idx, cutoff=False): + """One zarr slice per keymap entry -- the pre-coalescing read path.""" + store = zarr.open_group(str(episode), mode="r") + out = {} + for k, spec in key_map.items(): + if spec.get("key_type") == "annotation_keys": + continue + start, end = _expected_window(spec, idx, T, cutoff) + arr = store[spec["zarr_key"]] + chunk = arr[start:end] if end is not None else arr[start : start + 1][0] + horizon, history = spec.get("horizon"), spec.get("history") + if horizon is not None and chunk.shape[0] < horizon: + pad = np.repeat(chunk[-1:], horizon - chunk.shape[0], axis=0) + chunk = np.concatenate([chunk, pad], axis=0) + if horizon is None and history is not None and int(history) > 1: + if chunk.shape[0] < int(history): + pad = np.repeat(chunk[:1], int(history) - chunk.shape[0], axis=0) + chunk = np.concatenate([pad, chunk], axis=0) + out[k] = chunk + return out + + +# 0..3 clamp the backward window, 60 is interior, the rest truncate the forward chunk +@pytest.mark.parametrize("idx", [0, 1, 3, 60, 79, 130, T - HORIZON, T - 3, T - 1]) +@pytest.mark.parametrize("label", sorted(KEYMAPS)) +@pytest.mark.parametrize("cls", [ZarrDataset, ZarrAnnotationCutoffDataset]) +def test_getitem_matches_the_per_entry_read_path(episode, idx, label, cls): + key_map = _keymap(label) + sample = cls(episode, key_map=key_map, transform_list=None)[idx] + expected = _oracle(episode, key_map, idx, cutoff=cls is ZarrAnnotationCutoffDataset) + + # __getitem__ emits float32; compare bit-exactly in the dtype it emits + image_keys = {k for k, s in key_map.items() if s.get("key_type") == "camera_keys"} + for k, want in expected.items(): + if k in image_keys: + want = simplejpeg.decode_jpeg(want, colorspace="RGB") + want = np.transpose(want, (2, 0, 1)) / 255.0 + got = np.asarray(sample[k]) + assert got.shape == want.shape, k + np.testing.assert_array_equal(got, want.astype(np.float32), err_msg=k) + + +def test_shapes_are_the_keymap_contract(episode): + sample = ZarrDataset(episode, key_map=_keymap("keypoints_k4"))[60] + assert sample["left.action_keypoints"].shape == (HORIZON, 63) + assert sample["left.action_wrist_pose"].shape == (HORIZON, 7) + assert sample["left.obs_keypoints"].shape == (4, 63) + assert sample["left.obs_wrist_pose"].shape == (4, 7) + assert sample["obs_head_pose"].shape == (7,) + + +def _count_reads(monkeypatch) -> list[dict]: + """Per ``read_intervals`` call, ``{key: [(group_start, rows read), ...]}``.""" + calls: list[dict] = [] + inner = ZarrEpisode.read_intervals + + def spy(self, intervals): + out = inner(self, intervals) + calls.append({k: [(lo, len(rows)) for lo, rows in g] for k, g in out.items()}) + return out + + # ZarrEpisode has __slots__, so the counter goes on the class + monkeypatch.setattr(ZarrEpisode, "read_intervals", spy) + return calls + + +def test_one_read_per_distinct_zarr_key(episode, monkeypatch): + key_map = _keymap("keypoints_k4") + calls = _count_reads(monkeypatch) + + ZarrDataset(episode, key_map=key_map, transform_list=None)[60] + + assert len(calls) == 1, "one read call per __getitem__" + read_keys = set(calls[0]) + entries = [s for s in key_map.values() if s.get("key_type") != "annotation_keys"] + assert read_keys == {s["zarr_key"] for s in entries} + assert len(read_keys) == 6 < len(entries) == 10, "the four duplicates collapsed" + + +def test_rows_read_are_the_union_of_the_windows(episode, monkeypatch): + """Asserts rows READ, which the value oracle cannot see: a loader that read + the whole episode and sliced it would still match every value.""" + calls = _count_reads(monkeypatch) + ZarrDataset(episode, key_map=_keymap("keypoints_k4"), transform_list=None)[60] + rows = calls[0] + assert rows["left.obs_keypoints"] == [ + (57, 3 + HORIZON) + ], "history + chunk, one group" + assert rows["obs_head_pose"] == [(60, 1)], "a key read once is not widened" + assert rows["images.front_1"] == [(60, 1)] + + +@pytest.mark.parametrize( + "intervals, chunk, groups", + [ + # numeric keys: 100-row chunks + ([(57, 61), (60, 160)], 100, [(57, 160)]), + ([(0, 1), (150, 151)], 100, [(0, 151)]), # adjacent chunks: no new fetch + ([(0, 1), (250, 251)], 100, [(0, 1), (250, 251)]), # chunk 1 untouched + ([(99, 100), (100, 101)], 100, [(99, 101)]), # straddles a boundary + # JPEG keys: one frame per chunk + ([(60, 61), (57, 58)], 1, [(57, 58), (60, 61)]), + ([(59, 60), (60, 61)], 1, [(59, 61)]), + ([(10, 20), (15, 25)], 1, [(10, 25)]), + ([(0, 1), (0, 1)], 1, [(0, 1)]), + ], +) +def test_plan_reads_bridges_only_chunks_already_fetched(intervals, chunk, groups): + assert ZarrEpisode.plan_reads(intervals, chunk) == groups + + +def test_far_apart_jpeg_frames_are_two_rows_not_the_gap(episode): + reader = ZarrEpisode(episode) + store = zarr.open_group(str(episode), mode="r") + got = reader.read_intervals({"images.front_1": [(100, 101), (10, 11)]}) + assert [(lo, len(rows)) for lo, rows in got["images.front_1"]] == [ + (10, 1), + (100, 1), + ] + for lo, rows in got["images.front_1"]: + assert rows[0] == store["images.front_1"][lo : lo + 1][0] + + +def test_union_does_not_depend_on_keymap_order(episode): + """The real keymaps happen to list the forward chunk before the backward + window, which would hide a union that just took the last entry's bounds.""" + pose = "left.obs_wrist_pose" + order = { + "proprio": {"key_type": "proprio_keys", "zarr_key": pose, "history": 4}, + "action": {"key_type": "action_keys", "zarr_key": pose, "horizon": 10}, + } + for key_map in (order, dict(reversed(order.items()))): + ds = ZarrDataset(episode, key_map=key_map, transform_list=None) + starts = [w[1][0] for w in ds._read_windows(60).values()] + assert min(starts) == 57 and max(starts) == 60 + sample = ds[60] + np.testing.assert_array_equal( + np.asarray(sample["proprio"]), _ramp(7)[57:61].astype(np.float32) + ) + np.testing.assert_array_equal( + np.asarray(sample["action"]), _ramp(7)[60:70].astype(np.float32) + ) + + +def test_shared_span_slices_do_not_alias(episode): + """Two entries sliced out of one span must own their arrays: a transform is + handed the raw ndarrays and is free to write into them. Asserted at that + point, because __getitem__'s closing float64 -> float32 cast would sever the + aliasing before anything downstream could see it.""" + seen = {} + + class _Spy: + def transform(self, data): + seen["shares"] = np.shares_memory( + data["left.action_keypoints"], data["left.obs_keypoints"] + ) + data["left.action_keypoints"][0] += 1.0 # a transform writing in place + seen["proprio_last"] = data["left.obs_keypoints"][-1].copy() + return data + + ds = ZarrDataset(episode, key_map=_keymap("keypoints_k4"), transform_list=[_Spy()]) + ds[60] + assert seen["shares"] is False, "the two windows must not share a buffer" + # frame 60 of the ramp, unperturbed by the write to the action chunk + np.testing.assert_array_equal(seen["proprio_last"], _ramp(63)[60]) + + +def test_single_frame_entry_is_sliced_at_the_span_offset(episode): + """A key read as a single frame AND as a window starting earlier: the + single-frame slice needs the span offset like any other. No shipped keymap + mixes the two on one key (the history keys are exactly the obs entries), so + nothing else would catch a `block[0]` here.""" + pose = "left.obs_wrist_pose" + key_map = { + "history": {"key_type": "proprio_keys", "zarr_key": pose, "history": 4}, + "current": {"key_type": "proprio_keys", "zarr_key": pose}, + } + sample = ZarrDataset(episode, key_map=key_map, transform_list=None)[60] + np.testing.assert_array_equal( + np.asarray(sample["current"]), _ramp(7)[60].astype(np.float32) + ) + np.testing.assert_array_equal( + np.asarray(sample["history"]), _ramp(7)[57:61].astype(np.float32) + ) + + +def test_a_jpeg_failure_rereads_at_the_new_index(episode, monkeypatch): + """The decode retry picks a different frame and restarts the loop, so the + windows, spans and one read all have to be recomputed against the new idx -- + which the coalescing moved.""" + real = simplejpeg.decode_jpeg + calls: list[int] = [] + + def _fail_once(data, **kw): + calls.append(len(data)) + if len(calls) == 1: + raise ValueError("corrupt jpeg") + return real(data, **kw) + + monkeypatch.setattr( + "egomimic.rldb.zarr.zarr_dataset_multi.simplejpeg.decode_jpeg", _fail_once + ) + ds = ZarrDataset(episode, key_map=_keymap("keypoints_k4"), transform_list=None) + sample = ds[60] + + assert len(calls) == 2, "one failure, then the retry's decode" + # every numeric key comes from the RETRY's frame, not frame 60 + frame = int(round(float(np.asarray(sample["obs_head_pose"])[0]) / 1000.0)) + assert frame != 60, "the retry moved to another frame" + np.testing.assert_array_equal( + np.asarray(sample["obs_head_pose"]), _ramp(7)[frame].astype(np.float32) + ) + np.testing.assert_array_equal( + np.asarray(sample["left.obs_keypoints"]), + _ramp(63)[max(0, frame - 3) : frame + 1].astype(np.float32), + ) + + +def test_an_array_shorter_than_its_window_fails_loudly(tmp_path): + """A StopIteration out of __getitem__ would end the epoch silently.""" + numeric = {"obs_head_pose": _ramp(7)} + for side in ("left", "right"): + numeric[f"{side}.obs_wrist_pose"] = _ramp(7) + numeric[f"{side}.obs_keypoints"] = _ramp(63) + numeric[f"{side}.obs_ee_pose"] = _ramp(7) + path = ZarrWriter.create_and_write( + tmp_path / "short.zarr", + numeric_data=numeric, + image_data={"images.front_1": np.zeros((T, 32, 32, 3), np.uint8)}, + embodiment="human_bimanual", + task_name="ramp", + annotations=[("fold", 0, T - 1)], + intrinsics={"front_1": np.eye(3, 4)}, + ) + zarr.open_group(str(path), mode="r+")["left.obs_keypoints"].resize((130, 63)) + ds = ZarrDataset(path, key_map=_keymap("keypoints_k4"), transform_list=None) + with pytest.raises(ValueError, match="left.obs_keypoints"): + ds[115] From ed9e42dde258579b749f817e4acb1bf26002a139 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 19:02:55 +0000 Subject: [PATCH 32/79] perf: vectorized image augs and a torch.compile knob for HPT and Pi The 300M HPT train step was launch-bound, not compute-bound: ~14,600 CUDA kernel launches per step, 230 ms of CPU against 111 ms of GPU. Two things dominated, and both are fixed here. PerSampleAugs called torchvision's ColorJitter once per image so each sample drew its own factors -- 64 Python round trips a step, each launching a few dozen kernels and copying its factors from the host. BatchedColorJitter draws the factors as one on-device tensor and runs each adjustment once over the batch; a degenerate factor range reproduces torchvision's scalar adjustment exactly (tests/unit/test_image_augs.py). Anything it cannot batch still goes through the per-image loop. model.compile.enabled=true routes to .compile_for_training(), which uses nn.Module.compile rather than torch.compile(module) so state_dict keys survive -- the wrapper form renames every parameter _orig_mod.* and would break checkpoints written before the flag. That only routes __call__, so HPT compiles the trunk, the flow head's denoiser and the image encoders (the stems and heads are reached through named methods), and PI's call site moves to __call__. On one H100 at batch 64 and the 360x640 frames mecka decodes to: HPT 245 -> 173 ms from the augs, -> 109 ms compiled (2.24x). Pi at batch 16: 827 -> 588 ms compiled (1.41x). In fp32 the compiled loss is bit-identical to eager (bench/check_compile_parity.py). Also here: the flow head's attentions pass need_weights=False so they dispatch to scaled_dot_product_attention; frames decode to float32 instead of numpy's float64 default, halving the collate and H2D traffic (338 -> 169 MB a camera a batch, bit-identical); a per-camera per-step device->host sync is gone; and the batch is no longer deep-copied every step for an OT loss that is off. The all-zero (absent camera) guard moves into _apply_image_augs and becomes per-sample: the old test was one all() over the batch, so a batch mixing present and absent samples normalized the absent ones to -mean/std. The augs now run on every sample and the selector keeps the absent ones at zero, so RNG use no longer depends on the data. Jitter factors are drawn in float32 (uniform_ into bfloat16 gives ~40 distinct values; into uint8 it raises). enable_grad_norm was never passed to ModelWrapper, so pi0.5_base's false did nothing. It is wired now, and Pi follows openpi instead of opting out: it keeps logging the pre-clip grad norm every step and clips at openpi's 1.0, set as gradient_clip_val in the model config and handed to the Trainer unless trainer.gradient_clip_val is set. Pi recipes trained unclipped until now, so this is a behaviour change for every Pi run. One grad-norm pass remains (the second was removed with the log-only MAD detector), not the two bench_hpt.py models. bench/README.md records the harnesses, the numbers, what was tried and did not pay (persistent_workers, pin_memory, more workers, max-autotune), and which transformers each venv carries. The compile comments note that CUDA-graph modes are where a rank can wedge in cudagraph_trees and keep its GPUs after a cancel. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VQa3o51gPcrK8N1LLoqL95 --- .gitignore | 3 + bench/README.md | 190 ++++++ bench/__init__.py | 0 bench/bench_hpt.py | 544 ++++++++++++++++++ bench/bench_pi.py | 220 +++++++ bench/check_compile_parity.py | 176 ++++++ bench/e2e_hpt.sh | 25 + bench/e2e_loader.sh | 15 + bench/e2e_sweep.sh | 25 + bench/e2e_workers.sh | 17 + bench/step_timer.py | 86 +++ egomimic/algo/algo.py | 11 + egomimic/algo/hpt.py | 55 +- egomimic/algo/pi.py | 25 +- .../model/hpt_bc_mecka_6d_300M.yaml | 17 + egomimic/hydra_configs/model/pi0.5_base.yaml | 26 +- egomimic/models/denoising_nets.py | 10 +- egomimic/models/image_augs.py | 130 ++++- egomimic/pl_utils/pl_model.py | 24 + egomimic/rldb/zarr/zarr_dataset_multi.py | 13 +- egomimic/trainHydra.py | 17 +- egomimic/utils/compile_utils.py | 48 ++ tests/unit/test_image_augs.py | 109 ++++ tests/unit/test_image_decode_dtype.py | 28 + tests/unit/test_train_perf_wiring.py | 87 +++ 25 files changed, 1876 insertions(+), 25 deletions(-) create mode 100644 bench/README.md create mode 100644 bench/__init__.py create mode 100644 bench/bench_hpt.py create mode 100644 bench/bench_pi.py create mode 100644 bench/check_compile_parity.py create mode 100755 bench/e2e_hpt.sh create mode 100755 bench/e2e_loader.sh create mode 100755 bench/e2e_sweep.sh create mode 100755 bench/e2e_workers.sh create mode 100644 bench/step_timer.py create mode 100644 egomimic/utils/compile_utils.py create mode 100644 tests/unit/test_image_augs.py create mode 100644 tests/unit/test_image_decode_dtype.py create mode 100644 tests/unit/test_train_perf_wiring.py diff --git a/.gitignore b/.gitignore index 1380f8906a..9427ef3638 100644 --- a/.gitignore +++ b/.gitignore @@ -64,3 +64,6 @@ sort_annotations/ build/ dist/ docs/ + +# per-worktree venv for the Pi benchmarks (the shared emimic/ venv is off-pin) +.venv-pi/ diff --git a/bench/README.md b/bench/README.md new file mode 100644 index 0000000000..d239f19671 --- /dev/null +++ b/bench/README.md @@ -0,0 +1,190 @@ +# Training-speed experiments (H100, `ryanco/torch-compile-train`) + +Where the HPT and Pi train steps actually spend their time on the `loaner` +H100 nodes, what was changed, and what each change bought. Every number below +is one H100 (gpu11), bf16 autocast, measured with the harnesses in this folder. + +## The harnesses + +| file | what it does | +| --- | --- | +| `bench_hpt.py` | HPT train step (`forward_training` + `compute_losses` + backward + optimizer) on a synthetic batch, real model config. `--variants` A/Bs one change at a time; `--profile` dumps a torch-profiler table. | +| `bench_pi.py` | The same for Pi (openpi pi0.5), random init -- step time does not depend on the base weights. | +| `check_compile_parity.py` | Eager vs compiled, same weights, same batch: is `torch.compile` changing the numbers? | +| `e2e_hpt.sh`, `e2e_sweep.sh` | The real `trainHydra` path on locally cached mecka episodes, so the dataloader is in the measurement. | + +```bash +salloc -p loaner -A loaner --gres=gpu:h100:1 -c 28 --mem=200G +source ./wt-env.sh +python bench/bench_hpt.py --variants base,compile_hook --steps 25 --warmup 10 +``` + +`bench_pi.py` needs a venv on the pinned transformers (see "Loose ends", below). + +## What the baseline was doing + +The 300M HPT step at batch 64 was **not** limited by its own arithmetic. The +profiler put ~14,600 CUDA kernel launches in one step, 230 ms of CPU against +111 ms of GPU: the step was bound by the rate the host could issue work. The +model's matmuls accounted for ~26 ms of that. + +Two things dominated the launch count: + +1. **Per-sample image augmentation.** `PerSampleAugs` called torchvision's + `ColorJitter` once per image so each sample drew its own factors. At batch + 64 that is 64 Python round trips, each launching a few dozen kernels and + copying its factors from the host: ~470 `cudaMemcpyAsync` per step, and + **89 ms of a 196 ms step**. +2. **Single-tensor AdamW.** Not a bug in the repo -- `torch.optim.AdamW` with + neither `foreach` nor `fused` set picks `foreach` automatically, which the + shipped configs get. It is called out because passing `fused=False` + *explicitly* silently turns the foreach default off too, which is how the + first version of this benchmark ended up 17 ms/step slower than the real + training loop. + +## Changes + +| change | where | +| --- | --- | +| Vectorized per-sample `ColorJitter`: per-sample factors drawn on-device, each adjustment run once over the batch | `egomimic/models/image_augs.py` | +| `torch.compile` knob: `model.compile.enabled=true` -> `.compile_for_training()` | `pl_utils/pl_model.py`, `utils/compile_utils.py`, `algo/{algo,hpt,pi}.py` | +| `need_weights=False` on the flow head's attentions, so they dispatch to `scaled_dot_product_attention` instead of materialising per-head scores | `egomimic/models/denoising_nets.py` | +| Decode frames to float32 instead of numpy's float64 default, halving what the collate, the pin and the H2D copy move (338 -> 169 MB per camera per batch at 360x640, batch 64; bit-identical) | `egomimic/rldb/zarr/zarr_dataset_multi.py` | +| `non_blocking=True` on the batch's H2D copy (a no-op today, needed by anything that later pins) | `egomimic/algo/{hpt,pi}.py` | +| Drop a per-camera, per-step device->host sync (`if torch.all(img == 0)`) | `egomimic/algo/hpt.py` | +| Skip the per-step deep copy of the batch unless the OT loss is on | `egomimic/algo/hpt.py` | + +`compile_for_training` uses `nn.Module.compile`, not `torch.compile(module)`: +the former patches the call in place, the latter returns an `OptimizedModule` +whose parameters are all renamed `_orig_mod.*`, which would break every +checkpoint written before the flag was turned on. The flip side is that +`Module.compile` only routes `__call__`, so a submodule the algo invokes +through a named method (`stem.compute_latent`, `head.compute_loss`) stays +uncompiled no matter what -- that is why HPT compiles the trunk, the flow +head's denoiser and the image encoders and nothing else, and why `PI`'s call +site moved from `.forward(...)` to `(...)`. + +## Results + +### HPT model step, isolated (`bench_hpt.py`) + +One H100, `hpt_bc_mecka_6d_300M` (298 M trainable), batch 64, bf16 autocast, +synthetic batch, `AdamW` as the configs have it (no explicit `foreach`/`fused`), +including two grad-norm passes (`--grad-norm current`). That models a wrapper +older than this branch: `ModelWrapper` has run one pass since the MAD detector +went log-only, so every absolute time below carries one extra ~3 ms sweep; the +ratios are unaffected. + +360x640 -- the resolution mecka frames actually decode to: + +| | ms/step | vs before | peak GB | +| --- | ---: | ---: | ---: | +| before (per-image `ColorJitter`) | 245.2 | 1.00x | 16.8 | +| vectorized augs | 173.1 | 1.42x | 16.8 | +| + `model.compile.enabled=true` | 109.4 | **2.24x** | 16.3 | + +224x224, for comparison with anything already measured at that size: + +| | ms/step | vs before | peak GB | +| --- | ---: | ---: | ---: | +| before | 205.3 | 1.00x | 14.0 | +| vectorized augs | 120.2 | 1.71x | 13.9 | +| + compile | 82.1 | 2.50x | 13.8 | +| + compile, `mode=reduce-overhead` | 77.7 | 2.64x | 13.8 | + +Compile costs ~15-40 s at the first batch (it grows with image size) and +`reduce-overhead` roughly doubles that for another 5%. Inductor's cache is +keyed per Slurm job (`utils/compile_cache.py`), so every job -- including every +requeue -- pays it once. + +### Pi model step, isolated (`bench_pi.py`) + +One H100, `pi0.5_bc_mecka_6d` (3.6 B), batch 16, random init: + +| | ms/step | vs base | +| --- | ---: | ---: | +| base | 826.8 | 1.00x | +| `fused` AdamW | 794.8 | 1.04x | +| `model.compile.enabled=true` | 588.1 | 1.41x | +| compile + `fused` AdamW | 560.1 | **1.48x** | + +Compile costs ~4 min at the first batch. Forcing the gemma tower back to `sdpa` +attention did nothing for training -- openpi only pins it to `eager` inside +`sample_actions`, which the training pass never enters. + +### Is compile changing the numbers? (`check_compile_parity.py`) + +With every source of randomness pinned (dropout, drop-path, colour jitter, the +flow-matching noise and time) and both copies starting from identical weights: + +* **fp32: the loss is bit-identical.** The compiled graph is the same function. +* The remaining fp32 gradient differences are confined to the ResNet's conv + weights (cosine >= 0.9989), which is cudnn's non-deterministic + backward-filter, not compile. +* Under bf16 autocast the loss drifts ~0.3% on a random-init network, which is + reduction-order noise amplified through 23 trunk blocks. + +### End to end on real data (`e2e_hpt.sh`) + +The real `trainHydra` path, 20 locally cached mecka episodes, 5 x 40 batches, +median of the steps after the first epoch (`bench/step_timer.py` -- Lightning's +`simple` profiler averages the compile and the worker spin-up into every step, +which at 240 steps is a 70 ms/step artefact): + +| | step ms (median) | step ms (p10) | batch wait ms | +| --- | ---: | ---: | ---: | +| before | 298.6 | 266.2 | 41.1 | +| vectorized augs + float32 decode | 274.4 | 196.7 | 41.6 | +| + compile | 245.8 | 157.9 | 40.0 | + +Read the p10 column. The isolated model step drops 64 ms when compiled +(173 -> 109) and the real run's *fast* steps drop 59 ms (217 -> 158), which +matches; the median only drops 33 ms because the real step time is bimodal, +with roughly +/-70 ms of input-side jitter on top of every step. Compile is +doing what the microbenchmark says it does -- the run just spends a chunk of +each step somewhere else. + +That "somewhere else" is not worker count. 8 / 24 / 48 workers give 270.9 / +279.4 / 279.2 ms with a 40 ms batch wait throughout, so the pool is not the +limit and raising `num_workers` off the configs' 8 buys nothing here. The +constant ~40 ms wait is the main process receiving a 169 MB batch out of the +worker queue, which more workers cannot parallelise. + +## Things that were tried and did not pay + +Worth recording so nobody re-runs them: + +| | result | +| --- | --- | +| `persistent_workers=True` on the train loader | No effect. The epoch-start wait is ~1.6 s either way (measured per epoch), and the steady-state step is identical (272.9 vs 274.4 ms). | +| `pin_memory=True` (+ `non_blocking` H2D) | A wash. It moves the 41 ms batch wait into the step (wait 0.9 ms, step 351 vs 312) -- the pinning thread does the copy the main thread used to. It would only win with a stream prefetcher that issues the next batch's H2D during the current step. | +| `prefetch_factor=4` | No measurable effect. | +| `torch.set_float32_matmul_precision("high")`, `cudnn.benchmark` | Under bf16 autocast there is almost no fp32 matmul left to speed up: <1%. | +| `channels_last` for the ResNet | 2-3% in the isolated benchmark, inside the run-to-run noise end to end. | +| `mode="max-autotune-no-cudagraphs"` | 697 s of autotuning for 0.3%. Not worth it. | +| `mode="reduce-overhead"` (CUDA graphs) | A further 5% for roughly double the compile time. Left off by default. | +| Caching the sinusoid position tables | ~0.2%. The tables are rebuilt on the CPU and copied to the device every forward, which is ugly, but it is not where the time goes. | +| `fused=True` AdamW | 5% on HPT, 4% on Pi. Not wired into the configs: it is a numerics change to the optimizer and belongs in its own PR, not this one. Note that passing `fused=False` **explicitly** also disables the `foreach` default and costs 17 ms/step -- the configs correctly pass neither. | + +## Loose ends + +* **`enable_grad_norm` was never passed to `ModelWrapper`**, so + `model/pi0.5_base.yaml`'s `false` did nothing. It is wired now; Pi drops the + `false` and keeps logging its grad norm, as openpi does, and clips at openpi's + 1.0 (`gradient_clip_val` in the model config). HPT pays one pass, a + `clip_grad_norm_(inf)` sweep plus a host sync, about half the 6.5 ms/step the + two-pass bench measured; Pi pays that plus Lightning's clipping sweep, as + openpi does. +* **Which transformers you get depends on the venv.** `pyproject.toml` pins + 4.53.2 plus openpi's `transformers_replace` patch. The venv these numbers were + taken in had 5.17.0, where that patch breaks Pi (`PaliGemmaForConditionalGeneration` + has no `language_model`) and `tests/fixtures/train_harness.py` + (`transformers.utils.hub._is_offline_mode` was removed in 5.x); the Pi + benchmarks used a pinned, git-ignored `.venv-pi`, in which the whole unit suite + passes (1127 passed, 8 skipped). +* **A stream prefetcher** is the obvious next input-side move: issue the next + batch's H2D on a side stream during the current step, which is what would + make `pin_memory` pay. +* **uint8 images** would cut the batch another 4x (169 MB -> 42 MB) over the + float32 change, moving the `/255` and the transpose onto the GPU. Bigger + blast radius -- the transform list and norm stats would need checking. diff --git a/bench/__init__.py b/bench/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/bench/bench_hpt.py b/bench/bench_hpt.py new file mode 100644 index 0000000000..3d520e5dd8 --- /dev/null +++ b/bench/bench_hpt.py @@ -0,0 +1,544 @@ +"""Isolated HPT train-step benchmark: real model config, synthetic batches. + +No dataloader, no Lightning -- just forward_training + compute_losses + +backward + optimizer step, so a change to the model path shows up undiluted. + + python bench/bench_hpt.py --variants base,tf32,sdpa,compile_trunk + +Variants are monkeypatches applied before the model is built, so an A/B needs +no source edit; the winners get ported into egomimic/ afterwards. +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import sys +import time + +import hydra +import torch +from hydra import compose, initialize_config_module + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +EMB_HUMAN = 3 +EMB_EVA = 6 + + +class FakeNormStats: + """The slice of MultiDataset that HPT.__init__ reads: key discovery.""" + + norm_mode = "quantile" + + def __init__(self, keys: dict, embodiment_id: int): + self._keys = keys + self._emb = embodiment_id + + def keys_of_type(self, kind, embodiment_id): + return list(self._keys.get(kind, [])) + + def is_key_with_embodiment(self, key, embodiment_id): + return True + + def zarr_key_to_keyname(self, key, embodiment_id): + return key + + def unnormalize(self, preds, embodiment_id): + return preds + + +# ---------------------------------------------------------------- variants + + +def v_tf32(cfg): + torch.set_float32_matmul_precision("high") + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.backends.cudnn.benchmark = True + + +def v_sdpa(cfg): + """need_weights=False on the flow head's attentions. + + nn.MultiheadAttention defaults to need_weights=True, which forces the slow + path: it materialises B x H x L x L scores and averages them over heads, + instead of dispatching to scaled_dot_product_attention. + """ + from egomimic.models import denoising_nets as dn + + def forward_cross(self, x, cond): + res = x + x = self.ln1(x) + x, _ = self.mha(x, x, x, need_weights=False) + x = x + res + res = x + x = self.ln2(x) + x, _ = self.cmha(x, cond, cond, need_weights=False) + x = x + res + res = x + x = self.ln3(x) + x = self.mlp(x) + return x + res + + dn.CrossBlock.forward_cross = forward_cross + + +def v_poscache(cfg): + """Cache the sinusoid position tables instead of rebuilding them on the CPU + (and H2D-copying them) inside every forward.""" + import egomimic.algo.hpt as hpt_mod + import egomimic.utils.tensor_utils as tu + + real = tu.get_sinusoid_encoding_table + cache: dict = {} + + def cached(position_start, position_end, d_hid): + key = (position_start, position_end, d_hid) + if key not in cache: + cache[key] = real(position_start, position_end, d_hid) + return cache[key] + + tu.get_sinusoid_encoding_table = cached + hpt_mod.get_sinusoid_encoding_table = cached + + +def v_noclone(cfg): + """_clone_batch deep-copies every image tensor each step; only the optional + OT loss reads the copy.""" + import egomimic.algo.hpt as hpt_mod + + hpt_mod.HPT._clone_batch = lambda self, batch: batch + + +def v_nozerosync(cfg): + """`if not torch.all(img == 0)` in _robomimic_to_hpt_data is a blocking + device->host sync per camera per step. Always augment instead.""" + import egomimic.algo.hpt as hpt_mod + + def patched(self, batch, cam_keys, proprio_keys, lang_keys, ac_key, aux_ac_keys=[]): + data = {} + for key in proprio_keys: + if key in batch: + short = key.rsplit(".", 1)[-1] + modality = f"state_{short}" + value = batch[key] + data[modality] = value.unsqueeze(1) if value.ndim == 2 else value + for key in cam_keys: + if key in batch: + short = key.rsplit(".", 1)[-1] + _data = self._apply_image_augs(batch[key], short) + data[short] = _data.unsqueeze(1).unsqueeze(1) + for key in lang_keys: + if key in batch: + data[key] = batch[key] + if "sampled_prompt" in batch: + data[self.annotation_modality] = batch["sampled_prompt"] + data["is_6dof"] = self.is_6dof + data["pad_mask"] = batch["pad_mask"] + data["embodiment"] = batch["embodiment"] + for aux in aux_ac_keys: + data[aux] = batch[aux] + data["action"] = batch[self.shared_ac_key or ac_key] + return data + + hpt_mod.HPT._robomimic_to_hpt_data = patched + + +# Every attribute a variant may replace. Without restoring these between +# variants the second row of a sweep measures the first row plus itself, and +# the per-variant column becomes a running total. +def _patch_targets(): + import egomimic.algo.hpt as hpt_mod + import egomimic.utils.tensor_utils as tu + from egomimic.models import denoising_nets, image_augs + + return [ + (image_augs.PerSampleAugs, "forward"), + (hpt_mod.HPT, "_apply_image_augs"), + (hpt_mod.HPT, "_robomimic_to_hpt_data"), + (hpt_mod.HPT, "_clone_batch"), + (denoising_nets.CrossBlock, "forward_cross"), + (tu, "get_sinusoid_encoding_table"), + (hpt_mod, "get_sinusoid_encoding_table"), + ] + + +_PRISTINE = None +_PRISTINE_TF32 = None + + +def _restore_patched_globals(): + global _PRISTINE, _PRISTINE_TF32 + targets = _patch_targets() + if _PRISTINE is None: + _PRISTINE = [getattr(obj, name) for obj, name in targets] + _PRISTINE_TF32 = ( + torch.get_float32_matmul_precision(), + torch.backends.cudnn.benchmark, + torch.backends.cudnn.allow_tf32, + ) + return + for (obj, name), original in zip(targets, _PRISTINE): + setattr(obj, name, original) + precision, benchmark, cudnn_tf32 = _PRISTINE_TF32 + torch.set_float32_matmul_precision(precision) + torch.backends.cudnn.benchmark = benchmark + torch.backends.cudnn.allow_tf32 = cudnn_tf32 + + +def v_loopaugs(cfg): + """The pre-2026-09-19 image-aug path: torchvision's ColorJitter called once + per image. Kept so the table can show what the vectorized one replaced.""" + import torch + + from egomimic.models import image_augs + + def forward(self, images): + return torch.stack([self.augs(image) for image in images]) + + image_augs.PerSampleAugs.forward = forward + + +def v_noaugs(cfg): + """Upper bound for the image-aug path: PerSampleAugs loops ColorJitter over + the batch one image at a time.""" + import egomimic.algo.hpt as hpt_mod + + hpt_mod.HPT._apply_image_augs = lambda self, images, short: images + + +def v_channels_last(cfg): + os.environ["BENCH_CHANNELS_LAST"] = "1" + + +def v_compile_trunk(cfg): + os.environ["BENCH_COMPILE"] = "trunk" + + +def v_compile_head(cfg): + os.environ["BENCH_COMPILE"] = "head" + + +def v_compile_both(cfg): + os.environ["BENCH_COMPILE"] = "both" + + +def v_compile_block(cfg): + os.environ["BENCH_COMPILE"] = "block" + + +def v_compile_all(cfg): + os.environ["BENCH_COMPILE"] = "all" + + +def v_compile_hook(cfg): + """What the shipped knob does: HPT.compile_for_training via nn.Module.compile.""" + os.environ["BENCH_COMPILE"] = "hook" + + +def v_maxautotune(cfg): + os.environ["BENCH_COMPILE_MODE"] = "max-autotune-no-cudagraphs" + + +def v_cudagraphs(cfg): + os.environ["BENCH_COMPILE_MODE"] = "reduce-overhead" + + +VARIANTS = { + "base": lambda cfg: None, + "tf32": v_tf32, + "sdpa": v_sdpa, + "poscache": v_poscache, + "noclone": v_noclone, + "nozerosync": v_nozerosync, + "channels_last": v_channels_last, + "noaugs": v_noaugs, + "loopaugs": v_loopaugs, + "compile_trunk": v_compile_trunk, + "compile_head": v_compile_head, + "compile_both": v_compile_both, + "compile_block": v_compile_block, + "compile_all": v_compile_all, + "compile_hook": v_compile_hook, + "maxautotune": v_maxautotune, + "cudagraphs": v_cudagraphs, +} + + +# ---------------------------------------------------------------- build + + +def build_cfg(model_name: str, overrides: list[str]): + with initialize_config_module( + config_module="egomimic.hydra_configs", version_base=None + ): + return compose( + config_name="train_zarr_cartesian", + overrides=[f"model={model_name}", *overrides], + ) + + +def build_algo(cfg, embodiment: str, emb_id: int, keys: dict, device): + node = cfg.model.robomimic_model + algo = hydra.utils.instantiate(node, norm_stats=FakeNormStats(keys, emb_id)) + algo.device = device + algo.nets = algo.nets.to(device) + return algo + + +def synth_batch( + keys, + emb_id, + batch_size, + device, + action_dim, + proprio_dim, + horizon, + img_hw, + history_len, + prompt, +): + b = {} + for k in keys["camera_keys"]: + b[k] = torch.rand(batch_size, 3, *img_hw, device=device) + for k in keys["proprio_keys"]: + shape = ( + (batch_size, proprio_dim) + if history_len == 1 + else (batch_size, history_len, proprio_dim) + ) + b[k] = torch.randn(*shape, device=device) + for k in keys["action_keys"]: + b[k] = torch.randn(batch_size, horizon, action_dim, device=device) + b["pad_mask"] = torch.ones(batch_size, horizon, 1, device=device) + b["embodiment"] = torch.tensor([emb_id], device=device, dtype=torch.int64) + if prompt is not None: + b["sampled_prompt"] = [prompt] * batch_size + return {emb_id: b} + + +def maybe_compile(algo): + what = os.environ.get("BENCH_COMPILE") + if not what: + return + mode = os.environ.get("BENCH_COMPILE_MODE") or None + kw = {"dynamic": False} + if mode: + kw["mode"] = mode + policy = algo.nets["policy"] + if what == "hook": + algo.compile_for_training(mode=mode, dynamic=False) + return + if what == "all": + policy.trunk["trunk"] = torch.compile(policy.trunk["trunk"], **kw) + for name, head in policy.heads.items(): + head.model = torch.compile(head.model, **kw) + for name, enc in policy.encoders.items(): + policy.encoders[name] = torch.compile(enc, **kw) + for name, st in policy.stems.items(): + if name.endswith("_annotation"): + continue # frozen + text-feature cached; input is list[str] + policy.stems[name] = torch.compile(st, **kw) + return + if what in ("trunk", "both"): + policy.trunk["trunk"] = torch.compile(policy.trunk["trunk"], **kw) + if what in ("head", "both"): + for name, head in policy.heads.items(): + head.model = torch.compile(head.model, **kw) + if what == "block": + # per-block compile: one small graph reused 23x, so compile time stays + # flat in depth and the whole trunk still runs compiled kernels + blocks = policy.trunk["trunk"].blocks + for i, blk in enumerate(blocks): + blocks[i] = torch.compile(blk, **kw) + for name, head in policy.heads.items(): + for i, layer in enumerate(head.model.layers): + head.model.layers[i] = torch.compile(layer, **kw) + + +def maybe_channels_last(algo): + if os.environ.get("BENCH_CHANNELS_LAST"): + algo.nets["policy"].encoders.to(memory_format=torch.channels_last) + + +# ---------------------------------------------------------------- run + + +def run(args, variant: str): + for key in ("BENCH_COMPILE", "BENCH_COMPILE_MODE", "BENCH_CHANNELS_LAST"): + os.environ.pop(key, None) + _restore_patched_globals() + torch.manual_seed(0) + device = torch.device("cuda") + + emb_id = EMB_EVA if args.embodiment == "eva_bimanual" else EMB_HUMAN + keys = { + "camera_keys": [f"observations.images.{c}" for c in args.cameras.split(",")], + "proprio_keys": [f"observations.state.{args.proprio}"], + "action_keys": [args.action_key], + "lang_keys": [], + } + + for name in variant.split("+"): + VARIANTS[name](None) + + cfg = build_cfg(args.model, args.override or []) + algo = build_algo(cfg, args.embodiment, emb_id, keys, device) + maybe_channels_last(algo) + maybe_compile(algo) + + params = [p for p in algo.nets.parameters() if p.requires_grad] + # An explicit fused=False would also switch OFF the foreach auto-default, + # i.e. benchmark the single-tensor path against the real config's foreach. + adam_kw = {"foreach": None, "fused": None} + if args.adam == "fused": + adam_kw["fused"] = True + elif args.adam == "foreach": + adam_kw["foreach"] = True + elif args.adam == "single": + adam_kw["foreach"] = False + opt = torch.optim.AdamW(params, lr=1e-5, **adam_kw) + + stem = algo.nets["policy"].stems + hist = 1 + for n, s in stem.items(): + if n.endswith(f"_state_{args.proprio}"): + hist = int(getattr(s, "history_len", 1)) + prompt = "fold the towel" if algo.annotation_key else None + batch = synth_batch( + keys, + emb_id, + args.batch_size, + device, + args.action_dim, + args.proprio_dim, + args.horizon, + tuple(int(x) for x in args.img.split("x")), + hist, + prompt, + ) + + algo.nets.train() + torch.cuda.reset_peak_memory_stats() + amp = torch.autocast("cuda", dtype=torch.bfloat16, enabled=not args.fp32) + + def grad_norm_hooks(): + """What ModelWrapper does between backward and the optimizer step. + + ``once`` is today's wrapper: on_after_backward runs a full + clip_grad_norm_(inf) over every parameter and reads it on the host. + ``current`` adds the second pass on_before_optimizer_step used to run, + which is what the README's numbers were taken with. + """ + if args.grad_norm == "none": + return + params = list(algo.nets.parameters()) + n1 = torch.nn.utils.clip_grad_norm_(params, max_norm=float("inf")) + float(n1) + if args.grad_norm == "current": + n2 = torch.nn.utils.clip_grad_norm_(params, max_norm=float("inf")) + float(n2) + + def step(): + opt.zero_grad(set_to_none=True) + with amp: + preds = algo.forward_training({k: dict(v) for k, v in batch.items()}) + losses = algo.compute_losses(preds, batch) + losses["action_loss"].backward() + grad_norm_hooks() + opt.step() + return losses["action_loss"] + + t_compile = time.perf_counter() + for _ in range(args.warmup): + step() + torch.cuda.synchronize() + t_compile = time.perf_counter() - t_compile + + if args.profile: + from torch.profiler import ProfilerActivity, profile + + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + record_shapes=False, + ) as prof: + for _ in range(3): + step() + torch.cuda.synchronize() + print( + prof.key_averages().table( + sort_by="self_cuda_time_total", row_limit=args.profile_rows + ) + ) + print( + prof.key_averages().table( + sort_by="self_cpu_time_total", row_limit=args.profile_rows + ) + ) + + times = [] + for _ in range(args.steps): + torch.cuda.synchronize() + t0 = time.perf_counter() + step() + torch.cuda.synchronize() + times.append((time.perf_counter() - t0) * 1e3) + + return { + "variant": variant, + "median_ms": round(statistics.median(times), 2), + "p10_ms": round(sorted(times)[len(times) // 10], 2), + "warmup_s": round(t_compile, 1), + "peak_mem_gb": round(torch.cuda.max_memory_allocated() / 2**30, 2), + "params_M": round(sum(p.numel() for p in params) / 1e6, 1), + } + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="hpt_bc_mecka_6d_300M") + p.add_argument("--embodiment", default="human_bimanual") + p.add_argument("--cameras", default="front_img_1") + p.add_argument("--proprio", default="ee_pose") + p.add_argument("--action-key", default="actions_cartesian") + p.add_argument("--action-dim", type=int, default=18) + p.add_argument("--proprio-dim", type=int, default=20) + p.add_argument("--horizon", type=int, default=100) + p.add_argument("--img", default="224x224") + p.add_argument("--batch-size", type=int, default=64) + p.add_argument("--steps", type=int, default=30) + p.add_argument("--warmup", type=int, default=5) + p.add_argument("--fp32", action="store_true") + p.add_argument( + "--adam", default="default", choices=["default", "foreach", "fused", "single"] + ) + p.add_argument("--grad-norm", default="none", choices=["none", "once", "current"]) + p.add_argument("--override", action="append") + p.add_argument("--variants", default="base") + p.add_argument("--json", default=None) + p.add_argument("--profile", action="store_true") + p.add_argument("--profile-rows", type=int, default=25) + args = p.parse_args() + + rows = [] + for variant in args.variants.split(","): + rows.append(run(args, variant)) + print(json.dumps(rows[-1]), flush=True) + + base = rows[0]["median_ms"] + print(f"\n{'variant':<34}{'ms':>9}{'speedup':>9}{'mem GB':>9}{'warmup s':>10}") + for r in rows: + print( + f"{r['variant']:<34}{r['median_ms']:>9.2f}{base / r['median_ms']:>8.2f}x" + f"{r['peak_mem_gb']:>9.2f}{r['warmup_s']:>10.1f}" + ) + if args.json: + with open(args.json, "w") as f: + json.dump(rows, f, indent=2) + + +if __name__ == "__main__": + main() diff --git a/bench/bench_pi.py b/bench/bench_pi.py new file mode 100644 index 0000000000..51fd761e0c --- /dev/null +++ b/bench/bench_pi.py @@ -0,0 +1,220 @@ +"""Isolated Pi (openpi pi0.5) train-step benchmark: real model config, synthetic +batches, random init (step time does not depend on the base weights). + + python bench/bench_pi.py --variants base,compile_expert +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import sys +import time + +import hydra +import torch +from hydra import compose, initialize_config_module + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from bench.bench_hpt import FakeNormStats # noqa: E402 + + +def _patch_transformers_5_image_features() -> bool: + """The venv ships transformers 5.x while openpi pins 4.53.2, where + `get_image_features` returned a bare tensor. Unwrap it so the benchmark can + run; the numbers are indicative of the pinned stack, not a substitute.""" + from openpi.models_pytorch import gemma_pytorch as gp + + real = gp.PaliGemmaWithExpertModel.embed_image + + def embed_image(self, image): + out = real(self, image) + return getattr(out, "last_hidden_state", out) + + gp.PaliGemmaWithExpertModel.embed_image = embed_image + return True + + +EMB_HUMAN = 3 + + +def v_sdpa(cfg): + """openpi pins the gemma tower's attention to `eager` inside sample_actions + and never puts it back, so a train step that follows a val pass runs + materialised-score attention. Force sdpa before every forward.""" + os.environ["BENCH_PI_SDPA"] = "1" + + +def v_compile_expert(cfg): + os.environ["BENCH_PI_COMPILE"] = "expert" + + +def v_compile_forward(cfg): + os.environ["BENCH_PI_COMPILE"] = "forward" + + +def v_compile_hook(cfg): + """What the shipped knob does: PI.compile_for_training via Module.compile.""" + os.environ["BENCH_PI_COMPILE"] = "hook" + + +def v_grad_ckpt(cfg): + os.environ["BENCH_PI_GRAD_CKPT"] = "1" + + +def v_fused_adam(cfg): + os.environ["BENCH_PI_FUSED"] = "1" + + +VARIANTS = { + "base": lambda cfg: None, + "sdpa": v_sdpa, + "compile_expert": v_compile_expert, + "compile_forward": v_compile_forward, + "compile_hook": v_compile_hook, + "grad_ckpt": v_grad_ckpt, + "fused_adam": v_fused_adam, +} + + +def run(args, variant): + for key in ( + "BENCH_PI_COMPILE", + "BENCH_PI_GRAD_CKPT", + "BENCH_PI_FUSED", + "BENCH_PI_SDPA", + ): + os.environ.pop(key, None) + for name in variant.split("+"): + VARIANTS[name](None) + + if args.compat_transformers5: + _patch_transformers_5_image_features() + torch.manual_seed(0) + device = torch.device("cuda") + keys = { + "camera_keys": [f"observations.images.{c}" for c in args.cameras.split(",")], + "proprio_keys": [f"observations.state.{args.proprio}"], + "action_keys": [args.action_key], + "lang_keys": [], + } + + with initialize_config_module( + config_module="egomimic.hydra_configs", version_base=None + ): + cfg = compose( + config_name="train_zarr_cartesian_pi", + overrides=[ + f"model={args.model}", + "model.robomimic_model.config.pytorch_weight_path=null", + *(args.override or []), + ], + ) + algo = hydra.utils.instantiate( + cfg.model.robomimic_model, norm_stats=FakeNormStats(keys, EMB_HUMAN) + ) + algo.device = device + algo.nets = algo.nets.to(device) + + if os.environ.get("BENCH_PI_GRAD_CKPT"): + algo.nets["policy"].gradient_checkpointing_enable() + if os.environ.get("BENCH_PI_COMPILE") == "expert": + pwe = algo.nets["policy"].paligemma_with_expert + pwe.forward = torch.compile(pwe.forward, dynamic=False) + if os.environ.get("BENCH_PI_COMPILE") == "hook": + algo.compile_for_training() + if os.environ.get("BENCH_PI_COMPILE") == "forward": + policy = algo.nets["policy"] + policy.forward = torch.compile(policy.forward, dynamic=False) + + params = [p for p in algo.nets.parameters() if p.requires_grad] + opt = torch.optim.AdamW( + params, lr=1e-5, fused=True if os.environ.get("BENCH_PI_FUSED") else None + ) + + b, hz = args.batch_size, args.horizon + hw = tuple(int(x) for x in args.img.split("x")) + raw = { + keys["camera_keys"][0]: torch.rand(b, 3, *hw, device=device), + keys["proprio_keys"][0]: torch.randn(b, args.proprio_dim, device=device), + args.action_key: torch.randn(b, hz, args.action_dim, device=device), + "annotations": [["fold the towel"] for _ in range(b)], + } + batch = algo.process_batch_for_training({"human_bimanual": raw}) + + algo.nets.train() + torch.cuda.reset_peak_memory_stats() + + def force_sdpa(): + if not os.environ.get("BENCH_PI_SDPA"): + return + pwe = algo.nets["policy"].paligemma_with_expert + pwe.paligemma.language_model.config._attn_implementation = "sdpa" + pwe.gemma_expert.model.config._attn_implementation = "sdpa" + + def step(): + force_sdpa() + opt.zero_grad(set_to_none=True) + preds = algo.forward_training(batch) + loss = algo.compute_losses(preds, batch)["action_loss"] + loss.backward() + opt.step() + + t0 = time.perf_counter() + for _ in range(args.warmup): + step() + torch.cuda.synchronize() + warm = time.perf_counter() - t0 + + times = [] + for _ in range(args.steps): + torch.cuda.synchronize() + t = time.perf_counter() + step() + torch.cuda.synchronize() + times.append((time.perf_counter() - t) * 1e3) + + return { + "variant": variant, + "median_ms": round(statistics.median(times), 2), + "warmup_s": round(warm, 1), + "peak_mem_gb": round(torch.cuda.max_memory_allocated() / 2**30, 2), + "params_M": round(sum(p.numel() for p in params) / 1e6, 1), + } + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="pi0.5_bc_mecka_6d") + p.add_argument("--cameras", default="front_img_1") + p.add_argument("--proprio", default="ee_pose") + p.add_argument("--action-key", default="actions_cartesian") + p.add_argument("--action-dim", type=int, default=18) + p.add_argument("--proprio-dim", type=int, default=20) + p.add_argument("--horizon", type=int, default=100) + p.add_argument("--img", default="224x224") + p.add_argument("--batch-size", type=int, default=16) + p.add_argument("--steps", type=int, default=15) + p.add_argument("--warmup", type=int, default=5) + p.add_argument("--override", action="append") + p.add_argument("--variants", default="base") + p.add_argument("--compat-transformers5", action="store_true") + args = p.parse_args() + + rows = [run(args, v) for v in args.variants.split(",")] + for r in rows: + print(json.dumps(r), flush=True) + base = rows[0]["median_ms"] + print(f"\n{'variant':<28}{'ms':>9}{'speedup':>9}{'mem GB':>9}{'warmup s':>10}") + for r in rows: + print( + f"{r['variant']:<28}{r['median_ms']:>9.2f}{base / r['median_ms']:>8.2f}x" + f"{r['peak_mem_gb']:>9.2f}{r['warmup_s']:>10.1f}" + ) + + +if __name__ == "__main__": + main() diff --git a/bench/check_compile_parity.py b/bench/check_compile_parity.py new file mode 100644 index 0000000000..efb97c1878 --- /dev/null +++ b/bench/check_compile_parity.py @@ -0,0 +1,176 @@ +"""Does compile_for_training change the numbers? + +Same seed, same synthetic batch, eager vs compiled: the loss and every gradient +must agree to bf16-autocast tolerance. A mismatch bigger than that means the +compiled graph is not the model we think it is. +""" + +from __future__ import annotations + +import argparse +import os +import sys + +import hydra +import torch +from hydra import compose, initialize_config_module + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import torch.nn as nn # noqa: E402 + +from bench.bench_hpt import EMB_HUMAN, FakeNormStats, synth_batch # noqa: E402 +from egomimic.models.fm_policy import FMPolicy # noqa: E402 + + +def make_deterministic(algo, noise, time): + """Strip every source of randomness from the train step. + + Inductor does not reproduce eager's RNG stream, so a stochastic model gives + two different losses for reasons that have nothing to do with whether the + compiled graph is right. Dropout, drop-path, the colour jitter and the + flow-matching noise/time all get pinned, and what is left must match. + """ + for module in algo.nets.modules(): + if isinstance(module, nn.Dropout): + module.p = 0.0 + if isinstance(module, nn.MultiheadAttention): + module.dropout = 0.0 + if type(module).__name__ == "DropPath": + module.drop_prob = 0.0 + algo.train_image_augs = algo.eval_image_augs + + def predict(self, actions, global_cond): + t = time.to(actions.device).view(-1, 1, 1) + x_t = t * noise + (1 - t) * actions + return self.model(x_t, time.to(actions.device), global_cond), noise - actions + + FMPolicy.predict = predict + + +def build(keys, compile_it, mode): + torch.manual_seed(0) + with initialize_config_module( + config_module="egomimic.hydra_configs", version_base=None + ): + cfg = compose( + config_name="train_zarr_cartesian", overrides=["model=hpt_bc_mecka_6d_300M"] + ) + algo = hydra.utils.instantiate( + cfg.model.robomimic_model, norm_stats=FakeNormStats(keys, EMB_HUMAN) + ) + algo.device = torch.device("cuda") + algo.nets = algo.nets.to("cuda") + if compile_it: + algo.compile_for_training(mode=mode) + return algo + + +def run(algo, batch, seed, bf16=True): + algo.nets.train() + algo.nets.zero_grad(set_to_none=True) + torch.manual_seed(seed) + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=bf16): + preds = algo.forward_training({k: dict(v) for k, v in batch.items()}) + loss = algo.compute_losses(preds, batch)["action_loss"] + loss.backward() + grads = { + n: p.grad.detach().clone() + for n, p in algo.nets.named_parameters() + if p.grad is not None + } + return float(loss), grads + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--mode", default=None) + ap.add_argument("--batch-size", type=int, default=16) + ap.add_argument("--rtol", type=float, default=2e-2) + ap.add_argument("--atol", type=float, default=2e-3) + ap.add_argument("--min-cos", type=float, default=0.999) + ap.add_argument("--fp32", action="store_true") + args = ap.parse_args() + + keys = { + "camera_keys": ["observations.images.front_img_1"], + "proprio_keys": ["observations.state.ee_pose"], + "action_keys": ["actions_cartesian"], + "lang_keys": [], + } + torch.manual_seed(1234) + batch = synth_batch( + keys, + EMB_HUMAN, + args.batch_size, + torch.device("cuda"), + 18, + 20, + 100, + (224, 224), + 1, + "fold the towel", + ) + + noise = torch.randn(args.batch_size, 100, 18, device="cuda") + time = torch.rand(args.batch_size, device="cuda") * 0.9 + 0.05 + + eager = build(keys, False, None) + make_deterministic(eager, noise, time) + compiled = build(keys, True, args.mode) + make_deterministic(compiled, noise, time) + + # The comparison is meaningless unless both copies start from the same + # weights, and instantiate() draws from the global RNG. + a = dict(eager.nets.named_parameters()) + b = dict(compiled.nets.named_parameters()) + assert set(a) == set(b) + bad = [n for n in a if not torch.equal(a[n], b[n])] + assert not bad, f"{len(bad)} weights differ before the step, e.g. {bad[:3]}" + + eager_loss, eager_grads = run(eager, batch, seed=7, bf16=not args.fp32) + comp_loss, comp_grads = run(compiled, batch, seed=7, bf16=not args.fp32) + + print(f"loss eager={eager_loss:.6f} compiled={comp_loss:.6f}") + assert abs(eager_loss - comp_loss) <= args.atol + args.rtol * abs( + eager_loss + ), "loss differs beyond tolerance" + + assert set(eager_grads) == set(comp_grads), "gradient key sets differ" + + # Per tensor: max |delta| against that tensor's own scale, and the cosine of + # the two gradient vectors. The elementwise ratio is noisy under bf16 + # autocast (inductor fuses and reorders reductions); the direction is what + # the optimizer actually follows, so it is the honest check. + rows, dead = [], [] + for name, g in eager_grads.items(): + h = comp_grads[name] + if g.count_nonzero() == 0 and h.count_nonzero() == 0: + dead.append(name) # unused parameter in both; no direction to compare + continue + scale = g.abs().max().clamp(min=1e-8) + rel = ((g - h).abs().max() / scale).item() + cos = torch.nn.functional.cosine_similarity( + g.flatten().float(), h.flatten().float(), dim=0 + ).item() + rows.append((rel, cos, name)) + if dead: + print(f"{len(dead)} parameters have a zero gradient in both, e.g. {dead[:2]}") + rows.sort(reverse=True) + print(f"{len(rows)} gradient tensors") + print("worst by relative max-abs delta:") + for rel, cos, name in rows[:5]: + print(f" rel={rel:.2e} cos={cos:.6f} {name}") + min_cos = min(c for _, c, _ in rows) + worst_name = min(rows, key=lambda r: r[1])[2] + print( + f"min cosine similarity: {min_cos:.6f} ({worst_name}) " + f"|eager|={eager_grads[worst_name].norm():.3e} " + f"|compiled|={comp_grads[worst_name].norm():.3e}" + ) + assert min_cos >= args.min_cos, "a gradient points somewhere else" + print("PARITY OK") + + +if __name__ == "__main__": + main() diff --git a/bench/e2e_hpt.sh b/bench/e2e_hpt.sh new file mode 100755 index 0000000000..5ff44515dd --- /dev/null +++ b/bench/e2e_hpt.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# End-to-end HPT throughput: the real trainHydra path on locally cached mecka +# episodes, so the dataloader is in the measurement too. +# bench/e2e_hpt.sh [extra hydra overrides...] +set -euo pipefail +cd "$(dirname "$0")/.." +source ./wt-env.sh +export EGOVERSE_DATASET_DIR="${EGOVERSE_DATASET_DIR:-/workspace/users/copierceryan/datasets}" +export TOKENIZERS_PARALLELISM=false +HASHES=$(ls "$EGOVERSE_DATASET_DIR" | head -20 | paste -sd, -) +RUN="$1"; shift +python egomimic/trainHydra.py --config-name=train_zarr_cartesian \ + data=mecka_fold_freeform_opscale1_6d model=hpt_bc_mecka_6d_300M \ + "data.train_datasets.human_bimanual.resolver._target_=egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver" \ + "data.train_datasets.human_bimanual.filters=null" \ + "++data.train_datasets.human_bimanual.filters={_target_: egomimic.rldb.filters.DatasetFilter, episode_hashes: [${HASHES}]}" \ + "data.valid_datasets.human_bimanual.filters=null" \ + "++data.valid_datasets.human_bimanual.filters={_target_: egomimic.rldb.filters.DatasetFilter, episode_hashes: [${HASHES}]}" \ + "~data.unseen_op_valid_datasets" "~data.unseen_op_valid_dataloader_params" \ + "~data.video_episodes" \ + trainer=default trainer.accelerator=gpu trainer.devices=1 trainer.max_epochs=5 trainer.min_epochs=5 \ + trainer.limit_train_batches=40 +trainer.enable_checkpointing=false trainer.limit_val_batches=0 trainer.check_val_every_n_epoch=1000 \ + +trainer.enable_progress_bar=false logger=debug '~evaluator' '~callbacks' train_viz=false \ + '+callbacks={step_timer: {_target_: bench.step_timer.StepTimer, warmup_steps: 40}}' \ + norm_stats.sample_frac=0.01 +trainer.profiler=simple name=perf_e2e "description=$RUN" "$@" diff --git a/bench/e2e_loader.sh b/bench/e2e_loader.sh new file mode 100755 index 0000000000..faaf7d8d99 --- /dev/null +++ b/bench/e2e_loader.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +# persistent_workers, isolated: everything else identical. +set -uo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +OUT="$HERE/../logs/loader" +P="+data.train_dataloader_params.human_bimanual" +for v in false true; do + rm -f "$OUT.$v.log" + bash "$HERE/e2e_hpt.sh" "persist_$v" "$P.persistent_workers=$v" \ + "$P.pin_memory=false" "$P.prefetch_factor=2" > "$OUT.$v.log" 2>&1 +done +for v in false true; do + echo "--- persistent_workers=$v" + grep -h "\[step\]" "$OUT.$v.log" || echo FAILED +done diff --git a/bench/e2e_sweep.sh b/bench/e2e_sweep.sh new file mode 100755 index 0000000000..833ba21db8 --- /dev/null +++ b/bench/e2e_sweep.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# The DataLoader knobs, end to end. Each row is a full trainHydra run on the +# locally cached mecka episodes; read `run_training_batch` (per-step cost) and +# `run_training_epoch` minus 40 x that (per-epoch worker overhead). +set -uo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +OUT="${1:-$HERE/../logs/e2e}" +P="+data.train_dataloader_params.human_bimanual" + +declare -A ROWS=( + [off]="$P.persistent_workers=false $P.pin_memory=false $P.prefetch_factor=2" + [persist]="$P.persistent_workers=true $P.pin_memory=false $P.prefetch_factor=4" + [persist_pin]="$P.persistent_workers=true $P.pin_memory=true $P.prefetch_factor=4" + [persist_compile]="$P.persistent_workers=true $P.pin_memory=false $P.prefetch_factor=4 model.compile.enabled=true" +) +for row in off persist persist_pin persist_compile; do + rm -f "$OUT.$row.log" + # shellcheck disable=SC2086 + bash "$HERE/e2e_hpt.sh" "$row" ${ROWS[$row]} > "$OUT.$row.log" 2>&1 +done + +for row in off persist persist_pin persist_compile; do + printf '%-18s ' "$row" + grep -h "steady state over" "$OUT.$row.log" || echo FAILED +done diff --git a/bench/e2e_workers.sh b/bench/e2e_workers.sh new file mode 100755 index 0000000000..51d6eefea6 --- /dev/null +++ b/bench/e2e_workers.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +# num_workers, end to end. The data configs ship 8; a loaner node has 224 cores. +set -uo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +OUT="$HERE/../logs/workers" +for n in 8 24 48; do + rm -f "$OUT.$n.log" + bash "$HERE/e2e_hpt.sh" "w$n" "data.train_dataloader_params.human_bimanual.num_workers=$n" \ + > "$OUT.$n.log" 2>&1 +done +rm -f "$OUT.24c.log" +bash "$HERE/e2e_hpt.sh" w24c "data.train_dataloader_params.human_bimanual.num_workers=24" \ + model.compile.enabled=true > "$OUT.24c.log" 2>&1 +for n in 8 24 48 24c; do + printf '%-6s ' "$n" + grep -h "steady state over" "$OUT.$n.log" || echo FAILED +done diff --git a/bench/step_timer.py b/bench/step_timer.py new file mode 100644 index 0000000000..6c90754bbe --- /dev/null +++ b/bench/step_timer.py @@ -0,0 +1,86 @@ +"""Steady-state step timing for a real trainHydra run. + +Lightning's `simple` profiler reports the MEAN over every training batch, which +buries what we want to compare: the first batch carries worker spin-up, cudnn +autotuning and (with `model.compile.enabled`) a whole inductor compile, and at +240 steps a 17 s compile alone shifts the mean by 70 ms/step. This callback +reports the median of the steps after a warmup instead, split into the wait for +the batch and the step itself. +""" + +from __future__ import annotations + +import statistics +import time + +import torch +from lightning import Callback + + +class StepTimer(Callback): + def __init__(self, warmup_steps: int = 40, label: str = "step"): + self.warmup_steps = int(warmup_steps) + self.label = label + self._waits: list[float] = [] + self._steps: list[float] = [] + self._t_end = None + self._t_start = None + self._epoch_start = None + self._epoch_first_waits: list[float] = [] + + def _sync(self): + if torch.cuda.is_available(): + torch.cuda.synchronize() + + def on_train_epoch_start(self, trainer, pl_module): + self._epoch_start = time.perf_counter() + + def on_train_batch_start(self, trainer, pl_module, batch, batch_idx): + now = time.perf_counter() + if batch_idx == 0 and self._epoch_start is not None: + # Wait for the FIRST batch of an epoch. Without persistent_workers + # Lightning rebuilds the iterator here, which forks the pool and + # refills every prefetch queue before this returns. + self._epoch_first_waits.append(now - self._epoch_start) + elif self._t_end is not None: + self._waits.append(now - self._t_end) + self._t_start = now + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): + self._sync() + self._t_end = time.perf_counter() + self._steps.append(self._t_end - self._t_start) + + def on_fit_end(self, trainer, pl_module): + n = self.warmup_steps + steps, waits = self._steps[n:], self._waits[n:] + if not steps: + print(f"[{self.label}] only {len(self._steps)} steps, all warmup") + return + + def stats(xs): + xs = sorted(xs) + return ( + statistics.median(xs) * 1e3, + xs[len(xs) // 10] * 1e3, + xs[-max(1, len(xs) // 10)] * 1e3, + ) + + s_med, s_p10, s_p90 = stats(steps) + w_med, w_p10, w_p90 = stats(waits) if waits else (0.0, 0.0, 0.0) + print( + f"[{self.label}] steady state over {len(steps)} steps " + f"(first {n} dropped): " + f"step {s_med:.1f} ms (p10 {s_p10:.1f} / p90 {s_p90:.1f}), " + f"batch wait {w_med:.1f} ms (p10 {w_p10:.1f} / p90 {w_p90:.1f}), " + f"first step {self._steps[0] * 1e3:.0f} ms", + flush=True, + ) + epoch_waits = self._epoch_first_waits + if epoch_waits: + per_epoch = ", ".join(f"{w * 1e3:.0f}" for w in epoch_waits) + print( + f"[{self.label}] epoch-start batch wait (ms), one per epoch: " + f"{per_epoch} (mid-epoch median {w_med:.1f})", + flush=True, + ) diff --git a/egomimic/algo/algo.py b/egomimic/algo/algo.py index a01e4cc861..d5040ca26b 100644 --- a/egomimic/algo/algo.py +++ b/egomimic/algo/algo.py @@ -72,3 +72,14 @@ def log_info(self, info): loss_log (dict): name -> summary statistic """ raise NotImplementedError("Must implement log_info in subclass") + + def compile_for_training(self, mode=None, dynamic=False): + """Hand the hot submodules to ``torch.compile``. + + Called once by ``ModelWrapper`` when ``model.compile.enabled`` is set, + before DDP wraps anything -- inductor then sees one module at a time and + DDP's bucketing is unaffected. Algos that have nothing worth compiling + leave this alone; the default no-ops so a run that turns the flag on + against such an algo is slower to start, not broken. + """ + return [] diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index 8d0ffcb428..de7359cb67 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -26,6 +26,7 @@ ) from egomimic.models.image_augs import PerSampleAugs from egomimic.rldb.embodiment.embodiment import get_embodiment, get_embodiment_id +from egomimic.utils.compile_utils import compile_modules from egomimic.utils.hf_utils import download_from_huggingface from egomimic.utils.tensor_utils import EinOpsRearrange, get_sinusoid_encoding_table @@ -1024,6 +1025,23 @@ def __init__( self.training_step = 0 + @override + def compile_for_training(self, mode=None, dynamic=False): + """Compile the trunk, the flow head's denoiser and the image encoders. + + Only these three are invoked through ``__call__``; the stems run via + ``compute_latent`` and the heads via ``compute_loss``, which + ``Module.compile`` does not route (and which measured no gain anyway). + """ + policy = self.nets["policy"] + targets = [("trunk", policy.trunk["trunk"] if not policy.no_trunk else None)] + targets += [ + (f"head[{name}].model", getattr(head, "model", None)) + for name, head in policy.heads.items() + ] + targets += [(f"encoder[{name}]", enc) for name, enc in policy.encoders.items()] + return compile_modules(targets, mode=mode, dynamic=dynamic) + @override def process_batch_for_training(self, batch): """ @@ -1095,7 +1113,11 @@ def process_batch_for_training(self, batch): # TODO make this work with any fp type for key, value in processed_batch[embodiment_id].items(): if isinstance(value, torch.Tensor): - value = value.to(self.device) + # non_blocking: a no-op from pageable memory, and from a + # pin_memory loader it lets the copy overlap the previous + # step. Everything downstream is queued on the same stream, + # so ordering holds without an explicit wait. + value = value.to(self.device, non_blocking=True) if value.is_floating_point(): value = value.float() processed_batch[embodiment_id][key] = value @@ -1139,7 +1161,10 @@ def forward_training(self, batch): "domain": embodiment_name, # readability on config side "data": data, } - hpt_batches[embodiment_id] = self._clone_batch(hpt_batch) + # Only the OT loss reads this second copy; cloning every image + # tensor each step for a loss that is off is pure allocation. + if self.ot: + hpt_batches[embodiment_id] = self._clone_batch(hpt_batch) if self.freeze_repr: loss = self.nets["policy"].compute_loss_depth( @@ -1322,13 +1347,23 @@ def _forward_ot(self, batch, embodiment1_id, embodiment2_id): def _apply_image_augs(self, images, short): """ helper method that augments one camera's (B, 3, H, W) batch; train augs - are per-sample, the deterministic eval augs stay batched + are per-sample, the deterministic eval augs stay batched. + + An absent camera arrives as an all-zero image and must stay zero + (normalizing it would make it -mean/std). The test is per sample, so a + batch mixing present and absent samples is covered, and it is a + torch.where rather than an `if`, which would read the predicate on the + host and serialise the step against the next batch's copy. """ if self.nets.training and self.train_image_augs and short in self.encoders: - return self.train_image_augs(images) + augs = self.train_image_augs elif self.eval_image_augs and short in self.encoders: - return self.eval_image_augs(images) - return images + augs = self.eval_image_augs + else: + return images + absent = (images == 0).flatten(1).all(1) + out = augs(images) + return torch.where(absent.view(-1, *[1] * (out.dim() - 1)), 0.0, out) def _stem_history_len(self, domain: str, modality: str) -> int: """``history_len`` of ``domain``'s ``modality`` stem (1 when the stem @@ -1381,11 +1416,9 @@ def _robomimic_to_hpt_data( for key in cam_keys: if key in batch: short = key.rsplit(".", 1)[-1] - _data = batch[key] - if not torch.all(_data == 0): - _data = self._apply_image_augs(_data, short) - - data[short] = _data.unsqueeze(1).unsqueeze(1) + data[short] = ( + self._apply_image_augs(batch[key], short).unsqueeze(1).unsqueeze(1) + ) for key in lang_keys: if key in batch: diff --git a/egomimic/algo/pi.py b/egomimic/algo/pi.py index 7e26e352e5..72f64ac313 100644 --- a/egomimic/algo/pi.py +++ b/egomimic/algo/pi.py @@ -26,6 +26,7 @@ ) from egomimic.rldb.embodiment.embodiment import get_embodiment, get_embodiment_id from egomimic.utils.action_utils import ConverterRegistry, pad_to_width +from egomimic.utils.compile_utils import compile_modules logger = logging.getLogger(__name__) # Ensure logger propagates to root logger and has appropriate level @@ -419,7 +420,9 @@ def process_batch_for_training(self, batch): for key, value in processed_batch[embodiment_id].items(): if isinstance(value, torch.Tensor): - value = value.to(self.device) + # See HPT.process_batch_for_training: no-op from pageable + # memory, overlapped copy from a pin_memory loader. + value = value.to(self.device, non_blocking=True) if value.is_floating_point(): value = value.float() processed_batch[embodiment_id][key] = value @@ -432,6 +435,20 @@ def process_batch_for_training(self, batch): return processed_batch + @override + def compile_for_training(self, mode=None, dynamic=False): + """Compile openpi's ``PI0Pytorch.forward`` (the flow-matching training + pass: SigLIP tower, PaliGemma prefix and the action expert). + + ``sample_actions`` is already compiled by openpi itself, so only the + training pass is left. The observation preprocessing at the top of that + forward is Python-level dict work, so expect dynamo to graph-break there + and compile the transformer stack behind it. + """ + return compile_modules( + [("pi0", self.nets["policy"])], mode=mode, dynamic=dynamic + ) + @override def forward_training(self, batch): """ @@ -457,7 +474,9 @@ def forward_training(self, batch): embodiment_name, ) - losses = self.nets["policy"].forward(processed_obs, action) + # __call__, not .forward: nn.Module.compile only routes __call__, so + # a direct .forward() would silently run uncompiled. + losses = self.nets["policy"](processed_obs, action) loss = self._reduce_loss(losses, _batch[ac_key], embodiment_id, ac_key) predictions[f"{embodiment_name}_{ac_key}"] = _batch[ac_key] @@ -495,7 +514,7 @@ def forward_eval(self, batch): ) # Flow-matching val loss — same call as forward_training. - losses = self.nets["policy"].forward(processed_obs, action) + losses = self.nets["policy"](processed_obs, action) unnorm_preds[f"{embodiment_name}_loss"] = self._reduce_loss( losses, _batch[ac_key], embodiment_id, ac_key ) diff --git a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml index 48dea6f4fd..894b747092 100644 --- a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml @@ -183,3 +183,20 @@ optimizer: weight_decay: 0.0001 scheduler: null + +# torch.compile for the training step: ModelWrapper._maybe_compile hands the +# trunk, the flow head's denoiser and the image encoders to torch.compile +# (HPT.compile_for_training) before DDP wraps anything. Measured on one H100 at +# batch 64 and the 360x640 frames mecka decodes to (2026-09-19): 173 ms/step -> +# 109 ms (1.58x), for ~40 s of compile at the first batch. mode="reduce-overhead" +# is worth a further ~5% for roughly double that compile, but CUDA-graph modes +# (reduce-overhead, max-autotune) are where a rank can wedge in inductor's +# cudagraph_trees, ignore SIGTERM and keep its GPUs after a cancel. +# +# Off by default -- a run that never reaches steady state pays the compile for +# nothing. Per run: model.compile.enabled=true. `mode` goes straight to +# torch.compile; null is the default inductor mode the numbers were taken with. +compile: + enabled: false + mode: null + dynamic: false diff --git a/egomimic/hydra_configs/model/pi0.5_base.yaml b/egomimic/hydra_configs/model/pi0.5_base.yaml index 0f17969944..552db7a31b 100644 --- a/egomimic/hydra_configs/model/pi0.5_base.yaml +++ b/egomimic/hydra_configs/model/pi0.5_base.yaml @@ -58,7 +58,11 @@ robomimic_model: size: 224 interpolation: 3 -enable_grad_norm: false +# openpi's trainers clip the global grad norm at 1.0 (AdamW +# clip_gradient_norm) and log the pre-clip norm every step. trainHydra hands this +# to the Trainer unless trainer.gradient_clip_val is set; ModelWrapper logs the +# norm as Train/policy_grad_norms_raw (enable_grad_norm: false turns that off). +gradient_clip_val: 1.0 optimizer: _target_: torch.optim.AdamW @@ -77,4 +81,22 @@ scheduler: _partial_: true num_warmup_steps: 2000 -scheduler_interval: step \ No newline at end of file +scheduler_interval: step + +# torch.compile for the training step: ModelWrapper._maybe_compile compiles +# openpi's PI0Pytorch.forward (PI.compile_for_training). Measured on one H100 at +# batch 16 (2026-09-19): 827 ms/step -> 588 ms (1.41x), for ~4 min of compile at +# the first batch. `sample_actions` is compiled by openpi itself and is not +# affected by this flag. +# +# CUDA-graph modes (reduce-overhead, max-autotune) are where a rank can wedge in +# inductor's cudagraph_trees, ignore SIGTERM and keep its GPUs after a cancel; +# openpi already compiles sample_actions with max-autotune whatever this says. +# +# Off by default -- a run that never reaches steady state pays the compile for +# nothing. Per run: model.compile.enabled=true. `mode` goes straight to +# torch.compile; null is the default inductor mode the numbers were taken with. +compile: + enabled: false + mode: null + dynamic: false diff --git a/egomimic/models/denoising_nets.py b/egomimic/models/denoising_nets.py index 18471c59f5..6407340423 100644 --- a/egomimic/models/denoising_nets.py +++ b/egomimic/models/denoising_nets.py @@ -575,7 +575,7 @@ def forward(self, x, t, cond, timestep_table): res = x x = self.ln1(x) x = self.aln1(x, t, cond, timestep_table) - (x, _) = self.mha(x, x, x) + (x, _) = self.mha(x, x, x, need_weights=False) x = self.alns1(x, t, cond, timestep_table) x = x + res res = x @@ -617,13 +617,17 @@ def __init__( self.cond_proj = nn.Linear(cond_dim, hidden_dim) def forward_cross(self, x, cond): + # need_weights=False: the default True makes nn.MultiheadAttention take + # the slow path -- it materialises the B x H x L x L scores and averages + # them over heads to return weights nobody here reads -- instead of + # dispatching to scaled_dot_product_attention. The outputs are the same. res = x x = self.ln1(x) - x, _ = self.mha(x, x, x) + x, _ = self.mha(x, x, x, need_weights=False) x = x + res res = x x = self.ln2(x) - x, _ = self.cmha(x, cond, cond) + x, _ = self.cmha(x, cond, cond, need_weights=False) x = x + res res = x x = self.ln3(x) diff --git a/egomimic/models/image_augs.py b/egomimic/models/image_augs.py index a4350875a7..a2d6a23302 100644 --- a/egomimic/models/image_augs.py +++ b/egomimic/models/image_augs.py @@ -2,13 +2,129 @@ torchvision's random transforms draw one parameter set per call, so applying a ``ColorJitter`` to a stacked (B, 3, H, W) batch jitters every image in the batch -by the same factors (and in the same op order). ``PerSampleAugs`` wraps the -configured transform and calls it once per image, so every sample draws its own -parameters. +by the same factors (and in the same op order). ``PerSampleAugs`` gives every +sample its own parameters. + +The obvious way to do that -- call the transform once per image -- costs one +Python round trip and a few dozen tiny CUDA launches *per image*: measured at +89 ms of a 196 ms HPT train step at batch 64 (H100, 2026-09-19), i.e. 46% of +the step, because every factor is a CPU scalar that has to be copied to the +device. ``BatchedColorJitter`` draws the factors as one on-device (B,) tensor +and runs each adjustment once over the whole batch instead. """ +from typing import List, Optional, Tuple + import torch from torch import nn +from torchvision import transforms as T +from torchvision.transforms import _functional_tensor as _ft + +_GRAY_W = (0.2989, 0.587, 0.114) + + +def _rgb_to_gray(img: torch.Tensor) -> torch.Tensor: + w = torch.tensor(_GRAY_W, dtype=img.dtype, device=img.device).view(3, 1, 1) + return (img * w).sum(dim=-3, keepdim=True) + + +def _blend(img1: torch.Tensor, img2: torch.Tensor, ratio: torch.Tensor) -> torch.Tensor: + return (ratio * img1 + (1.0 - ratio) * img2).clamp(0.0, 1.0) + + +class BatchedColorJitter(nn.Module): + """``torchvision.transforms.ColorJitter`` with per-sample factors, batched. + + Same adjustments and same factor distributions as torchvision, drawn once + per image instead of once per call. The op ORDER is drawn once per call + (torchvision's own behaviour), not once per image as the per-image loop + did; each adjustment clamps to [0, 1], so the order is not quite inert, but + it is a permutation of the same four factors either way. + + Input is (B, 3, H, W) float in [0, 1]. + """ + + def __init__( + self, + brightness: Optional[Tuple[float, float]] = None, + contrast: Optional[Tuple[float, float]] = None, + saturation: Optional[Tuple[float, float]] = None, + hue: Optional[Tuple[float, float]] = None, + ): + super().__init__() + self.brightness = brightness + self.contrast = contrast + self.saturation = saturation + self.hue = hue + + @classmethod + def from_color_jitter(cls, jitter: T.ColorJitter) -> "BatchedColorJitter": + return cls(jitter.brightness, jitter.contrast, jitter.saturation, jitter.hue) + + @staticmethod + def _factors(bounds, n, img) -> torch.Tensor: + # Drawn in float32: uniform_ into bfloat16 yields ~40 distinct values, + # and into uint8 it raises. + lo, hi = bounds + f = torch.empty(n, 1, 1, 1, dtype=torch.float32, device=img.device) + return f.uniform_(lo, hi).to(img.dtype) + + def _apply_brightness(self, img, f): + return (img * f).clamp(0.0, 1.0) + + def _apply_contrast(self, img, f): + mean = _rgb_to_gray(img).mean(dim=(-3, -2, -1), keepdim=True) + return _blend(img, mean, f) + + def _apply_saturation(self, img, f): + return _blend(img, _rgb_to_gray(img), f) + + def _apply_hue(self, img, f): + hsv = _ft._rgb2hsv(img) + h, s, v = hsv.unbind(dim=-3) + h = (h + f.squeeze(-3)) % 1.0 + return _ft._hsv2rgb(torch.stack((h, s, v), dim=-3)).clamp(0.0, 1.0) + + def forward(self, img: torch.Tensor) -> torch.Tensor: + n = img.shape[0] + ops = [ + (self.brightness, self._apply_brightness), + (self.contrast, self._apply_contrast), + (self.saturation, self._apply_saturation), + (self.hue, self._apply_hue), + ] + for idx in torch.randperm(len(ops)).tolist(): + bounds, fn = ops[idx] + if bounds is None: + continue + img = fn(img, self._factors(bounds, n, img)) + return img + + def extra_repr(self) -> str: + return ( + f"brightness={self.brightness}, contrast={self.contrast}, " + f"saturation={self.saturation}, hue={self.hue}" + ) + + +# Transforms that are already batch-safe AND sample no parameters, so running +# them on the stacked batch is identical to running them per image. +_DETERMINISTIC = (T.Normalize, T.Resize, T.CenterCrop, T.ConvertImageDtype) + + +def _vectorize(augs) -> Optional[List[nn.Module]]: + """The batched equivalent of ``augs``, or None if any member has no batched + form (then the caller keeps the per-image loop).""" + members = augs.transforms if isinstance(augs, T.Compose) else [augs] + out: List[nn.Module] = [] + for t in members: + if isinstance(t, T.ColorJitter): + out.append(BatchedColorJitter.from_color_jitter(t)) + elif isinstance(t, _DETERMINISTIC): + out.append(t) + else: + return None + return out class PerSampleAugs(nn.Module): @@ -22,9 +138,15 @@ class PerSampleAugs(nn.Module): def __init__(self, augs): super().__init__() self.augs = augs + vectorized = _vectorize(augs) + self.vectorized = nn.ModuleList(vectorized) if vectorized is not None else None def forward(self, images): - return torch.stack([self.augs(image) for image in images]) + if self.vectorized is None: + return torch.stack([self.augs(image) for image in images]) + for t in self.vectorized: + images = t(images) + return images def extra_repr(self): return repr(self.augs) diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index 47644a6d6d..2d2441e711 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/egomimic/pl_utils/pl_model.py @@ -57,6 +57,7 @@ def __init__( if config_tree is not None: self.model = self._instantiate_model(config_tree, norm_stats_state) + self._maybe_compile(config_tree) elif robomimic_model is not None: # legacy support self.model = robomimic_model else: @@ -103,6 +104,29 @@ def _instantiate_model(self, config_tree, norm_stats_state): norm_stats=norm_stats, ) + def _maybe_compile(self, config_tree) -> None: + """Hand the algo's hot submodules to ``torch.compile`` when + ``model.compile.enabled`` is set. + + Here, not in the algo's ``__init__``: this runs on every DDP rank right + after the rank builds its own copy on CPU and before Lightning wraps it, + so each rank compiles for the GPU it is about to land on. Compilation + itself is lazy -- nothing is traced until the first batch. + """ + cfg = self._as_config(config_tree) + opts = cfg.model.get("compile") + if not opts or not opts.get("enabled", False): + return + compiled = self.model.compile_for_training( + mode=opts.get("mode"), dynamic=bool(opts.get("dynamic", False)) + ) + if not compiled: + print( + f"[compile] model.compile.enabled is set but " + f"{type(self.model).__name__} compiled nothing", + flush=True, + ) + # batch is now a dict, handle on model side def training_step(self, batch, batch_idx): self.train() diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 6eecbc61bb..d2e43ddc22 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -2362,7 +2362,18 @@ def _next(reason: str, key: str = "") -> int: idx = _next("JPEG decode failed", key=k) retry = True break - data[k] = np.transpose(decoded, (2, 0, 1)) / 255.0 + # float32, not numpy's float64 default for `uint8 / float`: + # process_batch_for_training casts to float32 on the GPU + # anyway, so the extra 4 bytes a pixel buy nothing and + # double the collate, the pin and the H2D copy. At 360x640 + # and batch 64 that is 338 MB a camera a batch instead of + # 169. The result is bit-identical: there are only 256 + # possible values and float32(u)/float32(255) equals + # float32(u / 255.0) for every one of them (see + # tests/unit/test_image_decode_dtype.py). + data[k] = np.transpose(decoded, (2, 0, 1)).astype( + np.float32 + ) / np.float32(255.0) elif zarr_key in self._json_keys: if isinstance(data[k], np.ndarray): data[k] = [self._decode_json_entry(v) for v in data[k]] diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 224a3fc13e..783f437dbe 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -86,6 +86,16 @@ def _build_model_config_tree(cfg: DictConfig) -> DictConfig: return tree +def _model_trainer_defaults(cfg: DictConfig) -> dict: + """Trainer kwargs a model config supplies: the gradient clip it was + published with (pi0.5: openpi's 1.0), unless ``trainer.gradient_clip_val`` + is set explicitly.""" + clip = cfg.model.get("gradient_clip_val") + if clip is None or cfg.trainer.get("gradient_clip_val") is not None: + return {} + return {"gradient_clip_val": clip} + + def _requeue_resume_path(cfg: DictConfig) -> Optional[str]: """``/last.ckpt`` when this process is a Slurm requeue (``SLURM_RESTART_COUNT`` > 0), else None. The dir is the ModelCheckpoint @@ -722,6 +732,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: config_tree=_build_model_config_tree(cfg), norm_stats_state=norm_stats.to_state(), scheduler_interval=cfg.model.get("scheduler_interval", "step"), + enable_grad_norm=cfg.model.get("enable_grad_norm", True), ) _log_dataset_frame_counts( @@ -762,7 +773,11 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: plugins.append(SLURMEnvironment(requeue_signal=signal.SIGUSR1)) print("SLURM REQUEUE ENABLED") trainer: Trainer = hydra.utils.instantiate( - cfg.trainer, callbacks=callbacks, logger=logger, plugins=plugins or None + cfg.trainer, + callbacks=callbacks, + logger=logger, + plugins=plugins or None, + **_model_trainer_defaults(cfg), ) object_dict = { diff --git a/egomimic/utils/compile_utils.py b/egomimic/utils/compile_utils.py new file mode 100644 index 0000000000..eb93e40e6f --- /dev/null +++ b/egomimic/utils/compile_utils.py @@ -0,0 +1,48 @@ +"""torch.compile wiring shared by the algos. + +``nn.Module.compile`` (not ``torch.compile(module)``) is what the algos call: +it swaps the module's call implementation IN PLACE, so the module keeps its +identity and -- the part that matters here -- its ``state_dict`` keys. Wrapping +a module in ``torch.compile`` instead returns an ``OptimizedModule`` whose +parameters are all prefixed ``_orig_mod.``, which silently breaks every +checkpoint written before the flag was turned on. + +The other half of the contract: ``Module.compile`` only routes ``__call__``. +A submodule the algo invokes through a named method (``stem.compute_latent``, +``head.compute_loss``) runs uncompiled no matter what, so only modules that are +actually *called* are worth handing over. +""" + +from __future__ import annotations + +import logging +from typing import Iterable, List, Optional + +import torch.nn as nn + +logger = logging.getLogger(__name__) + + +def compile_modules( + modules: Iterable[tuple[str, Optional[nn.Module]]], + mode: Optional[str] = None, + dynamic: bool = False, +) -> List[str]: + """``module.compile(...)`` for each named module; returns the names done.""" + kwargs = {"dynamic": dynamic} + if mode: + kwargs["mode"] = mode + done = [] + for name, module in modules: + if not isinstance(module, nn.Module): + continue # a head with no sub-network, or a None slot + module.compile(**kwargs) + done.append(name) + if done: + logger.info( + "torch.compile (mode=%s, dynamic=%s): %s", + mode or "default", + dynamic, + ", ".join(done), + ) + return done diff --git a/tests/unit/test_image_augs.py b/tests/unit/test_image_augs.py new file mode 100644 index 0000000000..4854370d5f --- /dev/null +++ b/tests/unit/test_image_augs.py @@ -0,0 +1,109 @@ +"""BatchedColorJitter must match torchvision's per-image adjustments, keep +per-sample randomness, and hand anything it cannot batch back to the loop.""" + +import pytest +import torch +import torchvision.transforms.functional as F +from torchvision import transforms as T + +from egomimic.models.image_augs import BatchedColorJitter, PerSampleAugs + +IMAGENET = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + + +@pytest.fixture +def images(): + torch.manual_seed(0) + return torch.rand(8, 3, 16, 16) + + +@pytest.mark.parametrize( + "kwargs, reference", + [ + ({"brightness": (0.7, 0.7)}, lambda x: F.adjust_brightness(x, 0.7)), + ({"contrast": (0.6, 0.6)}, lambda x: F.adjust_contrast(x, 0.6)), + ({"saturation": (1.3, 1.3)}, lambda x: F.adjust_saturation(x, 1.3)), + ({"hue": (0.04, 0.04)}, lambda x: F.adjust_hue(x, 0.04)), + ], +) +def test_matches_torchvision(images, kwargs, reference): + """A degenerate factor range pins every sample to one factor, so the batched + adjustment must equal torchvision's scalar one elementwise.""" + out = BatchedColorJitter(**kwargs)(images) + assert torch.allclose(out, reference(images), atol=1e-5) + + +def test_factors_are_per_sample(images): + jitter = BatchedColorJitter(brightness=(0.5, 1.5)) + ratios = (jitter(images) / images.clamp(min=1e-6)).flatten(1).median(dim=1).values + assert ratios.std() > 1e-3, "every sample got the same brightness factor" + + +def test_vectorizes_the_shipped_aug_list(): + augs = PerSampleAugs( + T.Compose( + [ + T.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05), + IMAGENET, + ] + ) + ) + assert augs.vectorized is not None + assert isinstance(augs.vectorized[0], BatchedColorJitter) + + +def test_falls_back_for_unbatchable_transforms(images): + augs = PerSampleAugs(T.Compose([T.RandomHorizontalFlip(), IMAGENET])) + assert augs.vectorized is None + assert augs(images).shape == images.shape + + +def test_deterministic_aug_list_is_unchanged(images): + """Eval augs have no random member, so the batched path must be exact.""" + augs = PerSampleAugs(T.Compose([IMAGENET])) + assert augs.vectorized is not None + assert torch.allclose(augs(images), IMAGENET(images), atol=1e-6) + + +def test_output_stays_in_range(images): + augs = PerSampleAugs( + T.Compose( + [T.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.2)] + ) + ) + out = augs(images) + assert out.min() >= 0.0 and out.max() <= 1.0 + + +def test_factors_are_drawn_in_float32(): + """uniform_ into bfloat16 collapses 4096 draws onto a few dozen values, and + into uint8 it raises.""" + img = torch.zeros(4096, 3, 1, 1, dtype=torch.bfloat16) + f = BatchedColorJitter._factors((0.9, 1.1), 4096, img) + assert f.dtype == torch.bfloat16 + assert f.float().unique().numel() > 30 + BatchedColorJitter._factors( + (0.9, 1.1), 4, torch.zeros(4, 3, 1, 1, dtype=torch.uint8) + ) + + +def test_absent_camera_stays_zero_per_sample(): + """An all-zero sample is an absent camera: it must not be normalized to + -mean/std, even when the rest of its batch is present.""" + from types import SimpleNamespace + + from egomimic.algo.hpt import HPT + + stub = SimpleNamespace( + nets=SimpleNamespace(training=True), + train_image_augs=PerSampleAugs(T.Compose([T.ColorJitter(0.4), IMAGENET])), + eval_image_augs=None, + encoders={"front_img_1": object()}, + ) + torch.manual_seed(0) + images = torch.rand(3, 3, 8, 8) + images[1] = 0.0 + out = HPT._apply_image_augs(stub, images, "front_img_1") + assert torch.equal(out[1], torch.zeros_like(out[1])) + assert not torch.equal(out[0], images[0]), "present samples are augmented" + assert torch.equal(HPT._apply_image_augs(stub, images, "not_an_encoder"), images) diff --git a/tests/unit/test_image_decode_dtype.py b/tests/unit/test_image_decode_dtype.py new file mode 100644 index 0000000000..31c405280a --- /dev/null +++ b/tests/unit/test_image_decode_dtype.py @@ -0,0 +1,28 @@ +"""Decoded frames leave the dataset as float32. + +`uint8 / 255.0` is float64 in numpy, which doubles what the collate, the +pinned-memory staging and the H2D copy move for every image in every batch -- +and the model casts to float32 on arrival regardless. +""" + +import numpy as np + + +def test_float32_division_is_bit_identical_to_float64(): + """The cheap dtype is also the same number: there are only 256 inputs, so + this is exhaustive, not a sample.""" + u = np.arange(256, dtype=np.uint8) + wide = (u / 255.0).astype(np.float32) + narrow = u.astype(np.float32) / np.float32(255.0) + assert np.array_equal(wide, narrow) + + +def test_decode_path_emits_float32(): + """Mirror of the expression in MultiDataset._read_keys' image branch.""" + decoded = np.random.randint(0, 256, size=(8, 8, 3), dtype=np.uint8) + out = np.transpose(decoded, (2, 0, 1)).astype(np.float32) / np.float32(255.0) + assert out.dtype == np.float32 + assert out.shape == (3, 8, 8) + assert np.array_equal( + out, (np.transpose(decoded, (2, 0, 1)) / 255.0).astype(np.float32) + ) diff --git a/tests/unit/test_train_perf_wiring.py b/tests/unit/test_train_perf_wiring.py new file mode 100644 index 0000000000..b55bcd4e7e --- /dev/null +++ b/tests/unit/test_train_perf_wiring.py @@ -0,0 +1,87 @@ +"""The torch.compile knob: off unless asked, in-place so checkpoints survive.""" + +import torch.nn as nn +from omegaconf import OmegaConf + +from egomimic.pl_utils.pl_model import ModelWrapper +from egomimic.utils.compile_utils import compile_modules + + +class _Tiny(nn.Module): + def __init__(self): + super().__init__() + self.lin = nn.Linear(4, 4) + + def forward(self, x): + return self.lin(x) + + +def test_compile_keeps_state_dict_keys(): + """nn.Module.compile is in-place; torch.compile(module) would prefix every + key with `_orig_mod.` and break checkpoints written before the flag.""" + module = _Tiny() + before = list(module.state_dict()) + compile_modules([("tiny", module)]) + assert list(module.state_dict()) == before + + +def test_compile_modules_skips_non_modules_and_reports_names(): + """A head without a sub-network hands back None (or something that is not a + Module); compiling it must be a skip, not a crash.""" + assert compile_modules([("a", _Tiny()), ("b", None), ("c", object())]) == ["a"] + + +class _SpyAlgo: + def __init__(self): + self.calls = [] + self.nets = nn.ModuleDict({"policy": _Tiny()}) + + def compile_for_training(self, mode=None, dynamic=False): + self.calls.append((mode, dynamic)) + return ["policy"] + + +def _wrapper_with(compile_cfg, monkeypatch): + algo = _SpyAlgo() + monkeypatch.setattr(ModelWrapper, "_instantiate_model", lambda *a, **k: algo) + model = {"robomimic_model": {}} + if compile_cfg is not None: + model["compile"] = compile_cfg + ModelWrapper(config_tree=OmegaConf.create({"model": model}), norm_stats_state=None) + return algo + + +def test_compile_is_off_unless_asked(monkeypatch): + assert _wrapper_with(None, monkeypatch).calls == [] + assert _wrapper_with({"enabled": False, "mode": None}, monkeypatch).calls == [] + + +def test_compile_forwards_mode_and_dynamic(monkeypatch): + cfg = {"enabled": True, "mode": "reduce-overhead", "dynamic": True} + assert _wrapper_with(cfg, monkeypatch).calls == [("reduce-overhead", True)] + + +def test_enable_grad_norm_reaches_the_wrapper(): + """It used to be dropped between the config and ModelWrapper, so a model's + `enable_grad_norm: false` did nothing.""" + import inspect + + from egomimic import trainHydra + + src = inspect.getsource(trainHydra) + assert 'enable_grad_norm=cfg.model.get("enable_grad_norm", True)' in src + + +def test_pi_clips_like_openpi_unless_the_trainer_says_otherwise(compose_resolve): + from egomimic.trainHydra import _model_trainer_defaults + + cfg = compose_resolve("train_zarr_cartesian", ["model=pi0.5_base"]) + assert cfg.model.gradient_clip_val == 1.0 + assert cfg.model.get("enable_grad_norm", True) is True + assert _model_trainer_defaults(cfg) == {"gradient_clip_val": 1.0} + + cfg.trainer.gradient_clip_val = 0.5 + assert _model_trainer_defaults(cfg) == {} + + hpt = compose_resolve("train_zarr_cartesian", []) + assert _model_trainer_defaults(hpt) == {} From d3e5b0ddacdd542647b8036fa7c6b611197eca3e Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 22:15:30 +0000 Subject: [PATCH 33/79] bench: verify the perf changes do not move accuracy The speed commit changed three things that are not bit-identical to what they replaced, and it only proved one of them. This adds the other two. need_weights=False moves the flow head's attentions onto scaled_dot_product_attention. check_attn_parity.py forces the old materialized-softmax path back on and compares a full fp32 step: loss 197.009293 vs 197.009033 (1.3e-6 relative), 581 gradient tensors at min cosine 0.999999. Reduction order, nothing else. The batched ColorJitter is the one real behaviour change -- it draws the op order once per call rather than once per image, so a batch shares a permutation. check_aug_distribution.py shows a single sample's marginal is untouched (6400 samples a path, two-sample KS over eight per-sample statistics, worst p = 0.61); only a training run can see the within-batch correlation. ab_augs.sh runs both paths at two seeds -- 2000 steps, 32 train episodes, held-out loss every 100 steps on 15 disjoint episodes, eval-mode augs so the held-out pass measures what was learned rather than what was shown. The batched arm is better at both seeds (-9.0%, -23.2%) and the 1.26x wall-clock gap confirms the control arm really took the per-image path. That rules out a regression; it is NOT evidence the rewrite helps. Two seeds cannot carry that, and the arms consume different amounts of RNG, so each arm/seed pair is effectively its own seed. bench/README.md says so where someone would otherwise quote the number. The control arm reuses the shipped fallback: PerSampleAugs keeps its per-image loop for aug lists _vectorize declines, so stubbing _vectorize out runs the old code rather than a reimplementation of it. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VQa3o51gPcrK8N1LLoqL95 --- bench/README.md | 70 ++++++++++++++++++ bench/ab_augs.sh | 40 +++++++++++ bench/ab_compare.py | 108 ++++++++++++++++++++++++++++ bench/ab_metrics.py | 84 ++++++++++++++++++++++ bench/check_attn_parity.py | 122 +++++++++++++++++++++++++++++++ bench/check_aug_distribution.py | 123 ++++++++++++++++++++++++++++++++ bench/train_loop_augs.py | 26 +++++++ 7 files changed, 573 insertions(+) create mode 100755 bench/ab_augs.sh create mode 100644 bench/ab_compare.py create mode 100644 bench/ab_metrics.py create mode 100644 bench/check_attn_parity.py create mode 100644 bench/check_aug_distribution.py create mode 100644 bench/train_loop_augs.py diff --git a/bench/README.md b/bench/README.md index d239f19671..aa009203cf 100644 --- a/bench/README.md +++ b/bench/README.md @@ -11,12 +11,21 @@ is one H100 (gpu11), bf16 autocast, measured with the harnesses in this folder. | `bench_hpt.py` | HPT train step (`forward_training` + `compute_losses` + backward + optimizer) on a synthetic batch, real model config. `--variants` A/Bs one change at a time; `--profile` dumps a torch-profiler table. | | `bench_pi.py` | The same for Pi (openpi pi0.5), random init -- step time does not depend on the base weights. | | `check_compile_parity.py` | Eager vs compiled, same weights, same batch: is `torch.compile` changing the numbers? | +| `check_attn_parity.py` | `need_weights=True` (the old path) vs `False`, same weights, same batch. | +| `check_aug_distribution.py` | Two-sample KS between the per-image and batched jitter, over per-sample image statistics. | +| `ab_augs.sh`, `ab_compare.py` | Accuracy A/B for the aug rewrite: both paths, two seeds, held-out loss on disjoint episodes. | | `e2e_hpt.sh`, `e2e_sweep.sh` | The real `trainHydra` path on locally cached mecka episodes, so the dataloader is in the measurement. | ```bash salloc -p loaner -A loaner --gres=gpu:h100:1 -c 28 --mem=200G source ./wt-env.sh python bench/bench_hpt.py --variants base,compile_hook --steps 25 --warmup 10 + +# accuracy A/B: ~12 min an arm on one H100 +for s in 42 43; do for arm in loop batched; do + AB_SEED=$s AB_EPOCHS=20 bench/ab_augs.sh $arm /${arm}_s${s}.json +done; done +python bench/ab_compare.py ``` `bench_pi.py` needs a venv on the pinned transformers (see "Loose ends", below). @@ -124,6 +133,67 @@ flow-matching noise and time) and both copies starting from identical weights: * Under bf16 autocast the loss drifts ~0.3% on a random-init network, which is reduction-order noise amplified through 23 trunk blocks. +### Does the aug rewrite change what the model learns? (`ab_augs.sh`, `ab_compare.py`) + +The vectorized `ColorJitter` is the only change here that is not exactly +equivalent to what it replaced. Everything else is either bit-identical +(float32 decode, the OT clone, the `torch.where` camera select, `non_blocking`) +or floating-point reassociation (compile, `need_weights`). The jitter is +different in two ways, both deliberate: + +* the **op order** is drawn once per call, not once per image, so samples in a + batch share a permutation. A single sample's marginal is unaffected -- the + order is still uniform over the 4! permutations and the factors are still + per-sample -- but within-batch independence is gone. +* the RNG stream differs, so no run reproduces an old seed. + +`check_aug_distribution.py` covers the marginal: 6400 augmented samples through +each path, two-sample KS on eight per-sample statistics (per-channel mean and +std, min, max). Worst p = 0.61. The marginals are indistinguishable. + +The within-batch correlation only a training run can see. `ab_augs.sh` runs both +arms at two seeds: same config, same 32 train episodes, compile off, 2000 steps, +held-out loss on 15 **disjoint** episodes every 100 steps. The held-out pass runs +in eval mode -- deterministic eval augs -- and reseeds, so it measures what the +model learned, not what it was shown. The `loop` arm stubs out `_vectorize` so +`PerSampleAugs` takes its own per-image fallback: the old code path, not a +reimplementation of it. + +| arm | seed | train (last 100) | held-out (mean of last 5) | wall s | +| --- | ---: | ---: | ---: | ---: | +| loop | 42 | 0.617 | 0.752 | 740 | +| loop | 43 | 0.647 | 0.822 | 734 | +| batched | 42 | 0.585 | 0.685 | 589 | +| batched | 43 | 0.598 | 0.631 | 566 | + +**No regression: the batched arm is better at both seeds** (-9.0%, -23.2%), and +the 1.26x wall-clock gap independently confirms the `loop` arm really ran the +per-image path. + +Do NOT read that as the rewrite improving accuracy. Two seeds cannot support +that claim, the two arms consume different amounts of RNG so each arm/seed pair +is effectively its own seed, and the held-out loss still swings 0.6-1.2 between +adjacent evals at step 2000. `ab_compare.py` prints `ARM EFFECT EXCEEDS SEED +NOISE` because the mean gap (0.129) is larger than the within-arm seed spread +(0.070); with n=2 that is weak evidence of a real difference and no evidence at +all of its direction. What it does rule out is the thing worth ruling out -- +sharing the op order across a batch is not costing accuracy. Confirming the +apparent gain would need ~6 seeds an arm. + +### Does `need_weights=False` change the numbers? (`check_attn_parity.py`) + +The flow head's attentions used to take `nn.MultiheadAttention`'s default +`need_weights=True`, which forces the materialized-softmax path; they now +dispatch to `scaled_dot_product_attention`. Forcing the old path back on and +comparing a full fp32 train step: + +* loss 197.009293 vs 197.009033, a relative delta of 1.3e-6 +* 581 gradient tensors, min cosine similarity 0.999999, worst relative max-abs + delta 1.1e-3 (all on ResNet conv weights, i.e. the same cudnn + backward-filter nondeterminism the compile parity check sees) + +Same function, different reduction order. + ### End to end on real data (`e2e_hpt.sh`) The real `trainHydra` path, 20 locally cached mecka episodes, 5 x 40 batches, diff --git a/bench/ab_augs.sh b/bench/ab_augs.sh new file mode 100755 index 0000000000..43305aae1a --- /dev/null +++ b/bench/ab_augs.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# Accuracy A/B for the vectorized augmentations: same seed, same data, same +# config, only the augmentation path differs. +# bench/ab_augs.sh [extra hydra overrides...] +# Train and held-out episodes are disjoint, so the held-out loss the ABMetrics +# callback records is a generalization measure, not a fit measure. +set -euo pipefail +cd "$(dirname "$0")/.." +source ./wt-env.sh +export EGOVERSE_DATASET_DIR="${EGOVERSE_DATASET_DIR:-/workspace/users/copierceryan/datasets}" +export TOKENIZERS_PARALLELISM=false +ALL=($(ls "$EGOVERSE_DATASET_DIR")) +TRAIN=$(printf '%s\n' "${ALL[@]:0:32}" | paste -sd, -) +VALID=$(printf '%s\n' "${ALL[@]:32}" | paste -sd, -) + +ARM="$1"; OUT="$2"; shift 2 +case "$ARM" in + batched) ENTRY=egomimic/trainHydra.py ;; + loop) ENTRY=bench/train_loop_augs.py ;; + *) echo "arm must be batched or loop" >&2; exit 2 ;; +esac + +EPOCHS="${AB_EPOCHS:-20}" +python "$ENTRY" --config-name=train_zarr_cartesian \ + data=mecka_fold_freeform_opscale1_6d model=hpt_bc_mecka_6d_300M \ + "seed=${AB_SEED:-42}" \ + "data.train_datasets.human_bimanual.resolver._target_=egomimic.rldb.zarr.zarr_dataset_multi.LocalEpisodeResolver" \ + "data.train_datasets.human_bimanual.filters=null" \ + "++data.train_datasets.human_bimanual.filters={_target_: egomimic.rldb.filters.DatasetFilter, episode_hashes: [${TRAIN}]}" \ + "data.valid_datasets.human_bimanual.filters=null" \ + "++data.valid_datasets.human_bimanual.filters={_target_: egomimic.rldb.filters.DatasetFilter, episode_hashes: [${VALID}]}" \ + "~data.unseen_op_valid_datasets" "~data.unseen_op_valid_dataloader_params" \ + "~data.video_episodes" \ + trainer=default trainer.accelerator=gpu trainer.devices=1 \ + "trainer.max_epochs=${EPOCHS}" "trainer.min_epochs=${EPOCHS}" \ + trainer.limit_train_batches=100 +trainer.enable_checkpointing=false \ + trainer.limit_val_batches=0 trainer.check_val_every_n_epoch=100000 \ + +trainer.enable_progress_bar=false logger=debug '~evaluator' '~callbacks' train_viz=false \ + "+callbacks={ab: {_target_: bench.ab_metrics.ABMetrics, out_path: '${OUT}', eval_batches: 12}}" \ + norm_stats.sample_frac=0.02 name=ab_augs "description=$ARM" "$@" diff --git a/bench/ab_compare.py b/bench/ab_compare.py new file mode 100644 index 0000000000..bd52945f2d --- /dev/null +++ b/bench/ab_compare.py @@ -0,0 +1,108 @@ +"""Compare ABMetrics runs: is the arm gap bigger than the seed gap? + +A single pair of runs cannot answer whether an augmentation change moved +accuracy, because two seeds of the SAME arm already differ. This reports the +between-arm difference next to the between-seed difference; the change is only +suspicious if the first is the larger of the two. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path + + +def load(path): + with open(path) as f: + return json.load(f) + + +def tail_mean(series, k): + return statistics.fmean(v for _, v in series[-k:]) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("results_dir") + ap.add_argument("--arms", nargs=2, default=["loop", "batched"]) + ap.add_argument("--seeds", nargs="+", type=int, default=[42, 43]) + ap.add_argument("--tail", type=int, default=5, help="eval points to average") + args = ap.parse_args() + + d = Path(args.results_dir) + runs = {} + for arm in args.arms: + for seed in args.seeds: + p = d / f"{arm}_s{seed}.json" + if p.exists(): + runs[(arm, seed)] = load(p) + else: + print(f"missing {p}") + + print( + f"{'arm':>9} {'seed':>5} {'steps':>6} {'train(last100)':>15} " + f"{'heldout(final)':>15} {f'heldout(last{args.tail})':>17} {'wall s':>8}" + ) + for (arm, seed), r in sorted(runs.items()): + tr = r["train_losses"] + ho = r["heldout"] + print( + f"{arm:>9} {seed:>5} {len(tr):>6} " + f"{statistics.fmean(v for _, v in tr[-100:]):>15.3f} " + f"{ho[-1][1]:>15.3f} {tail_mean(ho, args.tail):>17.3f} " + f"{r.get('wall_s') or 0:>8.0f}" + ) + + a0, a1 = args.arms + print(f"\nheld-out loss, mean of last {args.tail} eval points") + arm_gaps = [] + for seed in args.seeds: + if (a0, seed) in runs and (a1, seed) in runs: + x = tail_mean(runs[(a0, seed)]["heldout"], args.tail) + y = tail_mean(runs[(a1, seed)]["heldout"], args.tail) + arm_gaps.append(y - x) + print( + f" seed {seed}: {a0} {x:.3f} -> {a1} {y:.3f} " + f"delta {y - x:+.3f} ({100 * (y - x) / x:+.2f}%)" + ) + + seed_gaps = [] + for arm in args.arms: + vals = [ + tail_mean(runs[(arm, s)]["heldout"], args.tail) + for s in args.seeds + if (arm, s) in runs + ] + if len(vals) > 1: + seed_gaps.append(max(vals) - min(vals)) + print( + f" {arm}: seed spread {max(vals) - min(vals):.3f} " + f"over seeds {args.seeds}" + ) + + if arm_gaps and seed_gaps: + mean_arm = statistics.fmean(arm_gaps) + max_seed = max(seed_gaps) + print(f"\nmean arm effect {mean_arm:+.3f}, largest seed spread {max_seed:.3f}") + if abs(mean_arm) <= max_seed: + print("ARM EFFECT WITHIN SEED NOISE") + else: + print("ARM EFFECT EXCEEDS SEED NOISE -- investigate") + + print( + "\nheld-out curve (step: " + + ", ".join(f"{a}/{s}" for a in args.arms for s in args.seeds) + + ")" + ) + keys = [(a, s) for a in args.arms for s in args.seeds if (a, s) in runs] + n = min(len(runs[k]["heldout"]) for k in keys) + for i in range(n): + step = runs[keys[0]]["heldout"][i][0] + cells = " ".join(f"{runs[k]['heldout'][i][1]:9.3f}" for k in keys) + print(f" {step:>6} {cells}") + + +if __name__ == "__main__": + main() diff --git a/bench/ab_metrics.py b/bench/ab_metrics.py new file mode 100644 index 0000000000..50ee198ba1 --- /dev/null +++ b/bench/ab_metrics.py @@ -0,0 +1,84 @@ +"""Per-step train loss and a fixed held-out loss, written to JSONL. + +For comparing two training runs that differ only in the augmentation path. The +held-out pass runs in eval mode, so it uses the deterministic eval augs and is +unaffected by the change under test -- it measures what the model learned, not +what it was shown. The RNG is reseeded at the start of every pass so the +flow-matching noise and timesteps are identical across passes and across runs. +""" + +from __future__ import annotations + +import json +import statistics +import time + +import torch +from lightning import Callback + + +class ABMetrics(Callback): + def __init__(self, out_path: str, eval_batches: int = 12, eval_seed: int = 4242): + self.out_path = out_path + self.eval_batches = int(eval_batches) + self.eval_seed = int(eval_seed) + self.train_losses: list[tuple[int, float]] = [] + self.heldout: list[tuple[int, float]] = [] + self._loader = None + self._t0 = None + + def on_train_start(self, trainer, pl_module): + self._t0 = time.perf_counter() + dm = trainer.datamodule + self._loader = dm._build_val_style_loader( + dm.valid_datasets, dm.valid_dataloader_params, "valid" + ) + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): + loss = outputs["loss"] if isinstance(outputs, dict) else outputs + self.train_losses.append((int(trainer.global_step), float(loss))) + + @torch.no_grad() + def on_train_epoch_end(self, trainer, pl_module): + if self._loader is None: + return + model = pl_module.model + was_training = model.nets.training + model.nets.eval() + torch.manual_seed(self.eval_seed) + losses = [] + for i, batch in enumerate(self._loader): + if i >= self.eval_batches: + break + if ( + isinstance(batch, tuple) + and len(batch) == 3 + and isinstance(batch[0], dict) + ): + batch = batch[0] + processed = model.process_batch_for_training(batch) + with torch.autocast("cuda", dtype=torch.bfloat16): + preds = model.forward_training(processed) + losses.append( + float(model.compute_losses(preds, processed)["action_loss"]) + ) + if was_training: + model.nets.train() + if losses: + value = statistics.fmean(losses) + self.heldout.append((int(trainer.global_step), value)) + print( + f"[ab] step {trainer.global_step} heldout_loss {value:.6f} " + f"({len(losses)} batches, {time.perf_counter() - self._t0:.0f}s)", + flush=True, + ) + + def on_fit_end(self, trainer, pl_module): + payload = { + "train_losses": self.train_losses, + "heldout": self.heldout, + "wall_s": time.perf_counter() - self._t0 if self._t0 else None, + } + with open(self.out_path, "w") as f: + json.dump(payload, f) + print(f"[ab] wrote {self.out_path}", flush=True) diff --git a/bench/check_attn_parity.py b/bench/check_attn_parity.py new file mode 100644 index 0000000000..9ea765758d --- /dev/null +++ b/bench/check_attn_parity.py @@ -0,0 +1,122 @@ +"""Does need_weights=False change the numbers? + +The flow head's attentions used to call nn.MultiheadAttention with the default +need_weights=True, which forces the materialized-softmax path; they now pass +False and dispatch to scaled_dot_product_attention. That is the same function +with a different reduction order, so the loss and gradients should agree to +floating-point tolerance and nothing more. This forces the old path back on by +wrapping MultiheadAttention.forward and compares a full train step. +""" + +from __future__ import annotations + +import argparse +import os +import sys + +import torch +import torch.nn as nn + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from bench.bench_hpt import EMB_HUMAN, synth_batch # noqa: E402 +from bench.check_compile_parity import build, make_deterministic, run # noqa: E402 + +_MHA_FORWARD = nn.MultiheadAttention.forward + + +def force_need_weights(enabled: bool): + """Restore the pre-change call: need_weights defaulted to True.""" + if not enabled: + nn.MultiheadAttention.forward = _MHA_FORWARD + return + + def forward(self, *a, **kw): + kw["need_weights"] = True + kw.pop("average_attn_weights", None) + return _MHA_FORWARD(self, *a, **kw) + + nn.MultiheadAttention.forward = forward + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--batch-size", type=int, default=16) + ap.add_argument("--rtol", type=float, default=2e-2) + ap.add_argument("--atol", type=float, default=2e-3) + ap.add_argument("--min-cos", type=float, default=0.999) + ap.add_argument("--fp32", action="store_true") + args = ap.parse_args() + + keys = { + "camera_keys": ["observations.images.front_img_1"], + "proprio_keys": ["observations.state.ee_pose"], + "action_keys": ["actions_cartesian"], + "lang_keys": [], + } + torch.manual_seed(1234) + batch = synth_batch( + keys, + EMB_HUMAN, + args.batch_size, + torch.device("cuda"), + 18, + 20, + 100, + (224, 224), + 1, + "fold the towel", + ) + noise = torch.randn(args.batch_size, 100, 18, device="cuda") + time = torch.rand(args.batch_size, device="cuda") * 0.9 + 0.05 + + results = {} + for label, old_path in ( + ("sdpa (need_weights=False)", False), + ("materialized (need_weights=True)", True), + ): + force_need_weights(old_path) + algo = build(keys, compile_it=False, mode=None) + make_deterministic(algo, noise, time) + results[label] = run(algo, batch, seed=7, bf16=not args.fp32) + del algo + torch.cuda.empty_cache() + force_need_weights(False) + + (new_label, old_label) = list(results) + new_loss, new_grads = results[new_label] + old_loss, old_grads = results[old_label] + print( + f"loss {old_label}={old_loss:.6f} {new_label}={new_loss:.6f} " + f"delta={abs(new_loss - old_loss):.3e}" + ) + assert abs(new_loss - old_loss) <= args.atol + args.rtol * abs( + old_loss + ), "loss differs beyond floating-point tolerance" + + assert set(new_grads) == set(old_grads), "gradient key sets differ" + rows, dead = [], 0 + for name, g in old_grads.items(): + h = new_grads[name] + if g.count_nonzero() == 0 and h.count_nonzero() == 0: + dead += 1 + continue + scale = g.abs().max().clamp(min=1e-8) + rel = ((g - h).abs().max() / scale).item() + cos = torch.nn.functional.cosine_similarity( + g.flatten().float(), h.flatten().float(), dim=0 + ).item() + rows.append((rel, cos, name)) + rows.sort(reverse=True) + print(f"{len(rows)} gradient tensors compared ({dead} zero in both)") + print("worst by relative max-abs delta:") + for rel, cos, name in rows[:5]: + print(f" rel={rel:.2e} cos={cos:.6f} {name}") + min_cos = min(c for _, c, _ in rows) + print(f"min cosine similarity: {min_cos:.6f}") + assert min_cos >= args.min_cos, "a gradient points somewhere else" + print("ATTENTION PARITY OK") + + +if __name__ == "__main__": + main() diff --git a/bench/check_aug_distribution.py b/bench/check_aug_distribution.py new file mode 100644 index 0000000000..4d71286ba0 --- /dev/null +++ b/bench/check_aug_distribution.py @@ -0,0 +1,123 @@ +"""Does BatchedColorJitter draw from the same distribution as the per-image loop? + +The two cannot be compared pointwise: they consume different RNG streams, so +every sample differs by construction. What must match is the DISTRIBUTION of +augmented images. This draws the same images through both paths many times and +compares the empirical distributions of a set of per-sample statistics with a +two-sample Kolmogorov-Smirnov test. + +The one known deviation is op ORDER: the loop drew a permutation per image, the +batched version draws one per call. A single sample's marginal is unaffected +(the order is still uniform), so the KS test should pass; what it cannot see is +the within-batch correlation that introduces. --order-check measures that +separately. +""" + +from __future__ import annotations + +import argparse +import os +import sys + +import torch +from torchvision import transforms as T + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from egomimic.models.image_augs import PerSampleAugs # noqa: E402 + +JITTER = dict(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05) + + +def ks_2samp(a: torch.Tensor, b: torch.Tensor) -> tuple[float, float]: + """Two-sample KS statistic and its asymptotic p-value.""" + a, b = torch.sort(a)[0], torch.sort(b)[0] + n, m = a.numel(), b.numel() + allv = torch.cat([a, b]) + cdf_a = torch.searchsorted(a, allv, right=True).double() / n + cdf_b = torch.searchsorted(b, allv, right=True).double() / m + d = float((cdf_a - cdf_b).abs().max()) + en = (n * m / (n + m)) ** 0.5 + lam = (en + 0.12 + 0.11 / en) * d + p = 2.0 * sum( + (-1) ** (j - 1) * torch.exp(torch.tensor(-2.0 * j * j * lam * lam)).item() + for j in range(1, 101) + ) + return d, min(max(p, 0.0), 1.0) + + +def stats_of(img: torch.Tensor) -> torch.Tensor: + """Per-sample summary: per-channel mean and std, then overall min and max.""" + per_ch_mean = img.mean(dim=(-2, -1)) + per_ch_std = img.std(dim=(-2, -1)) + lo = img.amin(dim=(-3, -2, -1), keepdim=False).unsqueeze(-1) + hi = img.amax(dim=(-3, -2, -1), keepdim=False).unsqueeze(-1) + return torch.cat([per_ch_mean, per_ch_std, lo, hi], dim=-1) + + +STAT_NAMES = ["mean_R", "mean_G", "mean_B", "std_R", "std_G", "std_B", "min", "max"] + + +def collect(module, images, draws, device, seed): + torch.manual_seed(seed) + out = [] + for _ in range(draws): + out.append(stats_of(module(images.clone())).cpu()) + return torch.cat(out, dim=0) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--draws", type=int, default=200) + ap.add_argument("--batch-size", type=int, default=32) + ap.add_argument("--size", type=int, default=64) + ap.add_argument("--alpha", type=float, default=1e-3) + ap.add_argument("--device", default="cuda") + args = ap.parse_args() + + device = torch.device(args.device) + augs = T.Compose([T.ColorJitter(**JITTER)]) + + # A fixed, non-degenerate image set: flat colour would make several of the + # adjustments indistinguishable from each other. + torch.manual_seed(0) + images = torch.rand(args.batch_size, 3, args.size, args.size, device=device) + + batched = PerSampleAugs(augs).to(device) + assert batched.vectorized is not None, "the shipped aug list did not vectorize" + + loop = PerSampleAugs(augs).to(device) + loop.vectorized = None # the pre-change path + + n = args.draws * args.batch_size + print(f"{n} augmented samples per path ({args.draws} draws of {args.batch_size})") + a = collect(loop, images, args.draws, device, seed=1) + b = collect(batched, images, args.draws, device, seed=2) + + print( + f"{'statistic':>10} {'loop mean':>12} {'batched mean':>13} " + f"{'loop sd':>10} {'batched sd':>11} {'KS D':>8} {'p':>8}" + ) + worst_p = 1.0 + for i, name in enumerate(STAT_NAMES): + x, y = a[:, i].double(), b[:, i].double() + d, p = ks_2samp(x, y) + worst_p = min(worst_p, p) + print( + f"{name:>10} {x.mean():12.6f} {y.mean():13.6f} " + f"{x.std():10.6f} {y.std():11.6f} {d:8.4f} {p:8.4f}" + ) + + print( + f"\nworst p = {worst_p:.4f} over {len(STAT_NAMES)} statistics " + f"(alpha = {args.alpha})" + ) + if worst_p < args.alpha: + print("DISTRIBUTIONS DIFFER") + return 1 + print("DISTRIBUTION OK: no detectable difference in the per-sample marginal") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bench/train_loop_augs.py b/bench/train_loop_augs.py new file mode 100644 index 0000000000..c992b0afcf --- /dev/null +++ b/bench/train_loop_augs.py @@ -0,0 +1,26 @@ +"""trainHydra with the pre-change, per-image augmentation path. + +PerSampleAugs falls back to the per-image loop whenever _vectorize declines the +aug list, so stubbing it out reproduces the old behaviour exactly -- including +drawing the op order once per image -- without a second checkout. + +trainHydra is executed rather than imported: @hydra.main resolves its relative +config_path against the __file__ of the module that defines the task function, +so importing main from here makes Hydra look for the configs under bench/. +""" + +import os +import runpy + +import egomimic.models.image_augs as image_augs + +image_augs._vectorize = lambda augs: None + +from torchvision import transforms as _T # noqa: E402 + +_probe = image_augs.PerSampleAugs(_T.Compose([_T.ColorJitter(brightness=0.1)])) +assert _probe.vectorized is None, "loop-augs patch did not take" +print("[ab] per-image augmentation path active", flush=True) + +_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +runpy.run_path(os.path.join(_ROOT, "egomimic", "trainHydra.py"), run_name="__main__") From dcace1f5d84581b5fca546df9284b269d7328e85 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 21:42:20 +0000 Subject: [PATCH 34/79] tests: follow transformers 5's is_offline_mode() _hf_offline monkeypatched transformers.utils.hub._is_offline_mode, a module flag that 5.x replaced with is_offline_mode(). setattr on a name the module no longer has raises, so every test standing on the harness errored out before it reached the network: 12 of them, including the whole PI train-step matrix. Verified failing identically on an unmodified tree, so this is the venv's transformers 5.17 meeting a 4.x assumption, not a regression from this stack. Probe for whichever spelling the installed version has rather than pinning either, since the lockfile still says 4.53.2. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Xh3wMw9pXYPHUSGjU7qmix --- tests/fixtures/train_harness.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/fixtures/train_harness.py b/tests/fixtures/train_harness.py index 6bdb415f8b..025692925f 100644 --- a/tests/fixtures/train_harness.py +++ b/tests/fixtures/train_harness.py @@ -73,7 +73,12 @@ def _hf_offline(monkeypatch) -> None: monkeypatch.setenv("HF_HUB_OFFLINE", "1") monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) - monkeypatch.setattr(transformers.utils.hub, "_is_offline_mode", True) + # transformers <= 4 caches the flag in `_is_offline_mode`; 5.x reads it + # through `is_offline_mode()` instead. Whichever this version has. + if hasattr(transformers.utils.hub, "_is_offline_mode"): + monkeypatch.setattr(transformers.utils.hub, "_is_offline_mode", True) + else: + monkeypatch.setattr(transformers.utils.hub, "is_offline_mode", lambda: True) def write_fixtures( From a4b4fd00dd4d66cdb7a7a94dd04a1f5f8be55e11 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 21:42:51 +0000 Subject: [PATCH 35/79] pi: initialize from PaliGemma with a scratch action expert An arm for training pi0.5 from the released VLM instead of the pi0.5 base checkpoint, to separate what the VLM pretraining buys from what pi0.5's own robot pretraining buys. The prefix (SigLIP tower, Gemma-2B, the projector) comes from google/paligemma-3b-pt-224; everything pi0.5 adds on top -- the Gemma-300M action expert, action_in/out_proj, the adaRMS time MLPs -- stays at its init. pytorch_weight_path and paligemma_weight_path are mutually exclusive rather than ordered by precedence: silently preferring one would make a typo in the other look like it worked. Setting neither now warns instead of loading nothing quietly. The loader streams shard by shard and then fails if any parameter went untouched or any checkpoint tensor went unused, because a partial load of a 3.6B model is invisible in the loss for a long time. Two places where the release does not line up with what openpi builds, both handled in the key remap: - transformers 5 flattened SiglipVisionModel, so vision_tower.vision_model.* is vision_tower.* there; below 5 (the pin, 4.53.2) the hop is still in the module tree. The loader maps onto whichever the live model has. Mapping to the 5.x spelling unconditionally failed 21 SigLIP tensors on the pin, which the key-string tests could not see; a unit test now loads a tiny real PaliGemmaForConditionalGeneration under the installed transformers, and passes on both 4.53.2 and 5.17.0. - The checkpoint pads the vocabulary to 257216, 64 rows past the 257152 openpi allocates ( plus 63 unused). embed_tokens and lm_head are truncated; openpi reaches the image slots through the vision tower, never through an embed_tokens lookup, so the dropped rows are unreachable. google/paligemma-3b-pt-224 is license-gated. With paligemma_allow_mirror: true, resolve_paligemma_dir falls back on a 401/403 (not on a rate limit or an outage) to ungated mirrors whose shards are sha256-checked against the canonical repo's LFS metadata first -- a mirror that has drifted is refused, not used. It is off by default: integrity is checked, but a research arm's weights of record should come from the canonical repo unless a run opts in, and this cluster's HF token is authorized for it. Only *.safetensors are pulled. openpi's own loader source is not an option: its GCS bucket 403s anonymously. pi0.5_paligemma_init is an overlay, not a runnable config, so a keypoint or 6D arm is four lines. It changes the init source only; the constant LR stays, since pi0.5_base documents dropping the cosine schedule on purpose. openpi_compat carries three shims for a transformers 5 venv, which openpi's pin (4.53.2, also this repo's) does not expect; all should be deleted once the fork's transformers_replace patch is ported forward. The third restores openpi's fp32 keep-list, which spells the pre-flattening vision paths, so on 5.x the SigLIP patch and position embeddings ran in bf16. On 5.x only this arm loads: the pi0.5 base checkpoint stores the nested SigLIP names and its strict load_model stops at startup. pi05.md says which setup is supported and what differs. Verified with a full trainHydra run to exit 0 on an H100 (4 epochs, all four validation heads, checkpoint and videos written) and the real 12 GB checkpoint into the real model in tests/integration, both on transformers 5.17, plus the unit tests on 4.53.2 and 5.17.0. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Xh3wMw9pXYPHUSGjU7qmix --- README.md | 1 + egomimic/algo/pi.py | 37 +- egomimic/hydra_configs/model/pi0.5_base.yaml | 11 + .../pi0.5_bc_mecka_6d_paligemma_init.yaml | 15 + .../model/pi0.5_paligemma_init.yaml | 30 ++ egomimic/models/openpi_compat.py | 129 +++++ egomimic/models/paligemma_init.py | 296 ++++++++++++ pi05.md | 117 +++++ tests/integration/test_paligemma_init_real.py | 253 ++++++++++ tests/unit/test_openpi_compat.py | 115 +++++ tests/unit/test_paligemma_init.py | 457 ++++++++++++++++++ 11 files changed, 1453 insertions(+), 8 deletions(-) create mode 100644 egomimic/hydra_configs/model/pi0.5_bc_mecka_6d_paligemma_init.yaml create mode 100644 egomimic/hydra_configs/model/pi0.5_paligemma_init.yaml create mode 100644 egomimic/models/openpi_compat.py create mode 100644 egomimic/models/paligemma_init.py create mode 100644 tests/integration/test_paligemma_init_real.py create mode 100644 tests/unit/test_openpi_compat.py create mode 100644 tests/unit/test_paligemma_init.py diff --git a/README.md b/README.md index 6573e5419f..4d77d988be 100644 --- a/README.md +++ b/README.md @@ -101,6 +101,7 @@ Every cluster-specific default in the configs can be overridden without editing | `EGOVERSE_CACHE_DIR` | `paths.cache_dir`; norm stats cached by content under `/norm_stats/` | `$EGOVERSE_LOG_DIR/cache` | | `EGOVERSE_ROOT` | `paths.root_dir` | the directory you launch from | | `EGOVERSE_PI05_WEIGHTS` | pi0.5 base weights directory (`model.robomimic_model.config.pytorch_weight_path`) | lab PACE path | +| `EGOVERSE_PALIGEMMA_WEIGHTS` | PaliGemma checkpoint for the from-scratch pi arm (`...config.paligemma_weight_path`), a local dir or an HF repo id | `google/paligemma-3b-pt-224` | | `WANDB_ENTITY`, `WANDB_PROJECT` | `logger.wandb.entity` / `.project` | `rl2-group` / `zarr_test` | ### Other Settings diff --git a/egomimic/algo/pi.py b/egomimic/algo/pi.py index 72f64ac313..b0f5f6f90b 100644 --- a/egomimic/algo/pi.py +++ b/egomimic/algo/pi.py @@ -16,6 +16,8 @@ from transformers import AutoTokenizer from egomimic.algo.algo import Algo +from egomimic.models import openpi_compat +from egomimic.models.paligemma_init import load_paligemma_weights, select_init_source from egomimic.models.preprocess_pi_obs import ( PI_CAMERA_SLOTS, _concat_proprio, @@ -34,6 +36,10 @@ logger.setLevel(logging.INFO) logger.propagate = True # Explicitly enable propagation (default, but ensures it works) +# openpi targets transformers 4.53 (this repo's pin); these shims let it run on a +# 5.x venv too. See egomimic/models/openpi_compat.py. +openpi_compat.apply() + class PI(Algo): """ """ @@ -173,7 +179,20 @@ def __init__( self.model = openpi.models_pytorch.pi0_pytorch.PI0Pytorch(model_cfg) - if self.config.pytorch_weight_path is not None: + # Three mutually exclusive starting points, in decreasing order of what + # is pretrained: the full pi0.5 base checkpoint, PaliGemma only (the + # action expert and the action/time projections stay at their init), or + # nothing at all. + paligemma_weight_path = getattr(self.config, "paligemma_weight_path", None) + init_source = select_init_source( + self.config.pytorch_weight_path, paligemma_weight_path + ) + target = ( + self.model.module + if isinstance(self.model, torch.nn.parallel.DistributedDataParallel) + else self.model + ) + if init_source == "pi05_base": model_path = os.path.join( self.config.pytorch_weight_path, "model.safetensors" ) @@ -181,21 +200,23 @@ def __init__( raise FileNotFoundError( f"Pretrained weight file not found: {model_path}" ) - target = ( - self.model.module - if isinstance(self.model, torch.nn.parallel.DistributedDataParallel) - else self.model - ) safetensors.torch.load_model(target, model_path) logger.info( "Loaded pretrained weights from %s (%d parameters)", model_path, sum(p.numel() for p in target.parameters()), ) + elif init_source == "paligemma": + load_paligemma_weights( + target, + paligemma_weight_path, + allow_mirror=getattr(self.config, "paligemma_allow_mirror", False), + ) else: logger.warning( - "No pytorch_weight_path: no base weights loaded. The weights must " - "come from a checkpoint; otherwise this trains from scratch." + "Neither pytorch_weight_path nor paligemma_weight_path is set: no " + "base weights loaded. The weights must come from a checkpoint; " + "otherwise this trains from scratch, VLM included." ) # openpi's PyTorch port hard-codes 32-wide action projections # regardless of ``Pi0Config.action_dim``. Wider action spaces (the diff --git a/egomimic/hydra_configs/model/pi0.5_base.yaml b/egomimic/hydra_configs/model/pi0.5_base.yaml index 552db7a31b..0658feca5d 100644 --- a/egomimic/hydra_configs/model/pi0.5_base.yaml +++ b/egomimic/hydra_configs/model/pi0.5_base.yaml @@ -5,6 +5,17 @@ robomimic_model: pytorch_training_precision: bfloat16 # directory holding model.safetensors of the pi0.5 base ($EGOVERSE_PI05_WEIGHTS overrides) pytorch_weight_path: ${oc.env:EGOVERSE_PI05_WEIGHTS,/storage/home/hcoda1/4/paphiwetsa3/r-dxu345-0/projects/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch} + # The alternative to pytorch_weight_path: a local directory or an HF repo id + # holding the paligemma-3b-*-224 safetensors. Only the VLM prefix (SigLIP + # tower, multimodal projector, Gemma-2B) is loaded from it; everything pi0.5 + # adds -- the action expert, the action in/out projections, the adaRMS time + # MLPs -- trains from scratch. Setting both paths is an error. Flip it on + # with the model/pi0.5_paligemma_init.yaml overlay. + paligemma_weight_path: null + # google/paligemma-3b-pt-224 is license-gated. true = without an HF token + # that has accepted the Gemma terms, fall back to a third-party mirror whose + # shards sha256-match the canonical repo; false (default) = raise. + paligemma_allow_mirror: false model: pi05: true # Width of the model's action vector. 32 is openpi's pretrained width diff --git a/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d_paligemma_init.yaml b/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d_paligemma_init.yaml new file mode 100644 index 0000000000..91abc72867 --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d_paligemma_init.yaml @@ -0,0 +1,15 @@ +defaults: + - pi0.5_bc_mecka_6d + - pi0.5_paligemma_init + - _self_ + +# pi0.5_bc_mecka_6d with the pi0.5 base checkpoint swapped for a PaliGemma-only +# init -- the ablation that asks how much of the mecka 6D result comes from +# pi0.5's robot pretraining rather than from the VLM. Everything else (the +# cartesian 6D action layout, the action converter, the prompt blocks, the +# optimizer) is the parent's; see pi0.5_paligemma_init.yaml for the two keys +# this arm changes and for the lr knobs worth sweeping. +# +# Pairs with the same data configs as the parent (data=mecka_all_6d, +# mecka_fold_freeform*_6d) and evaluator=eval_pi_wristframe_6d. Norm stats are +# data-config-specific, not model-specific, so this arm reuses the parent's. diff --git a/egomimic/hydra_configs/model/pi0.5_paligemma_init.yaml b/egomimic/hydra_configs/model/pi0.5_paligemma_init.yaml new file mode 100644 index 0000000000..4e56103988 --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_paligemma_init.yaml @@ -0,0 +1,30 @@ +# OVERLAY, not a runnable model config. Compose it AFTER any pi0.5_* config to +# swap the pi0.5 base checkpoint for a PaliGemma-only init: the VLM prefix comes +# from the public paligemma-3b-pt-224 release and the action expert, the action +# in/out projections and the adaRMS time MLPs start from their PyTorch init. +# +# defaults: +# - pi0.5_bc_mecka_6d +# - pi0.5_paligemma_init +# - _self_ +# +# It sets only the two init keys, so it never clobbers the config it is layered +# on (action_dim, ac_keys, domains, the prompt flags all survive). Per-run, +# without a new file: +# model.robomimic_model.config.pytorch_weight_path=null \ +# model.robomimic_model.config.paligemma_weight_path=google/paligemma-3b-pt-224 +# +# The optimizer is left at pi0.5_base's fine-tuning recipe (constant 3e-5 after +# 2000 warmup steps). A randomly initialized action expert may well want a +# higher peak and a longer warmup; that is a per-run decision, e.g. +# model.optimizer.lr=1e-4 model.scheduler.num_warmup_steps=5000 +# +# Multi-rank runs should pre-fetch the checkpoint once, or every rank races on +# the same Hugging Face download lock: +# python -c "from egomimic.models.paligemma_init import resolve_paligemma_dir as r; print(r('google/paligemma-3b-pt-224'))" +robomimic_model: + config: + pytorch_weight_path: null + # $EGOVERSE_PALIGEMMA_WEIGHTS overrides: point it at a local snapshot dir to + # skip the Hub entirely. + paligemma_weight_path: ${oc.env:EGOVERSE_PALIGEMMA_WEIGHTS,google/paligemma-3b-pt-224} diff --git a/egomimic/models/openpi_compat.py b/egomimic/models/openpi_compat.py new file mode 100644 index 0000000000..ce6f5e15a3 --- /dev/null +++ b/egomimic/models/openpi_compat.py @@ -0,0 +1,129 @@ +"""Adapters that let openpi's PyTorch pi0 run against transformers 5. + +`external/openpi` pins `transformers==4.53.2` and ships a `transformers_replace` +patch written against it (its own `check.py` asserts that version), which is also +this repo's pin. A venv on transformers 5 works too, but there the stock modules +have moved on, so a few of openpi's call sites no longer line up. Nothing here +runs below transformers 5, and each shim is a no-op when it finds openpi's own +spelling already in place -- so porting the fork's patch forward makes this +module inert. Delete it then. + +Applied from `egomimic.algo.pi`, which is the only importer of openpi. +""" + +import logging + +import torch + +logger = logging.getLogger(__name__) + +_APPLIED = False + + +def _patch_embed_image() -> bool: + """`get_image_features` returns a `BaseModelOutputWithPooling`, not a tensor. + + transformers 5's `PaliGemmaModel.get_image_features` returns the vision + output object with the projected features stashed on `pooler_output`, while + openpi's `embed_prefix` reads `.shape` off the result. Unwrap it. + """ + from openpi.models_pytorch.gemma_pytorch import PaliGemmaWithExpertModel + + original = PaliGemmaWithExpertModel.embed_image + + def embed_image(self, image: torch.Tensor) -> torch.Tensor: + features = original(self, image) + if isinstance(features, torch.Tensor): + return features + return features.pooler_output + + PaliGemmaWithExpertModel.embed_image = embed_image + return True + + +def _patch_paligemma_bc_properties() -> bool: + """`paligemma.language_model` / `.vision_tower` / `.multi_modal_projector`. + + transformers 5 moved those submodules under `PaliGemmaModel` and dropped the + back-compat properties that used to forward to them from + `PaliGemmaForConditionalGeneration`. openpi still reads all three off the + outer model. Put the forwarders back. + """ + from transformers import PaliGemmaForConditionalGeneration + + added = [] + for name in ("language_model", "vision_tower", "multi_modal_projector"): + if hasattr(PaliGemmaForConditionalGeneration, name): + continue + setattr( + PaliGemmaForConditionalGeneration, + name, + property(lambda self, _name=name: getattr(self.model, _name)), + ) + added.append(name) + return bool(added) + + +# openpi keeps SigLIP's patch and position embeddings in float32 by substring on +# `vision_tower.vision_model.embeddings.*`; the flattened 5.x spelling of the same +# three tensors. +_FP32_KEEP_FLAT = ( + "vision_tower.embeddings.patch_embedding.", + "vision_tower.embeddings.position_embedding.", +) + + +def _patch_fp32_keep_list() -> bool: + """openpi's float32 keep-list misses the SigLIP embeddings. + + `to_bfloat16_for_selected_params` casts everything to bf16 and then restores + float32 on names containing `vision_tower.vision_model.embeddings.*`. + transformers 5 flattened `SiglipVisionModel`, so nothing matches and the + patch embedding and position embedding stay bf16 -- the tensors openpi + protects because small updates vanish in bf16. Re-apply the keep under the + flattened names; with the nested names present they match nothing. + """ + from openpi.models_pytorch.gemma_pytorch import PaliGemmaWithExpertModel + + original = PaliGemmaWithExpertModel.to_bfloat16_for_selected_params + + def to_bfloat16_for_selected_params(self, precision="bfloat16"): + original(self, precision) + if precision != "bfloat16": + return + for name, param in self.named_parameters(): + if any(selector in name for selector in _FP32_KEEP_FLAT): + param.data = param.data.to(dtype=torch.float32) + + PaliGemmaWithExpertModel.to_bfloat16_for_selected_params = ( + to_bfloat16_for_selected_params + ) + return True + + +def apply() -> None: + """Install every shim once. Safe to call repeatedly.""" + global _APPLIED + if _APPLIED: + return + _APPLIED = True + + import transformers + + major = int(transformers.__version__.split(".")[0]) + if major < 5: + return + + shims = ( + ("embed_image", _patch_embed_image), + ("paligemma_bc_properties", _patch_paligemma_bc_properties), + ("fp32_keep_list", _patch_fp32_keep_list), + ) + applied = [name for name, fn in shims if fn()] + if applied: + logger.info( + "openpi targets transformers 4.53 and this environment has %s; " + "applied compatibility shims: %s", + transformers.__version__, + ", ".join(applied), + ) diff --git a/egomimic/models/paligemma_init.py b/egomimic/models/paligemma_init.py new file mode 100644 index 0000000000..55799580ea --- /dev/null +++ b/egomimic/models/paligemma_init.py @@ -0,0 +1,296 @@ +"""PaliGemma-only initialization for the pi0.5 chassis. + +``egomimic.algo.pi.PI`` normally starts from the full pi0.5 base checkpoint. +This module is the other entry point: the VLM half (SigLIP tower, multimodal +projector, Gemma-2B prefix) comes from the public PaliGemma release, and +everything pi0.5 adds on top -- the action expert, the action in/out +projections, the adaRMS time MLPs -- keeps its PyTorch init. + +openpi builds the prefix as a stock ``PaliGemmaForConditionalGeneration`` whose +config is the ``gemma_2b`` variant (width 2048, depth 18, head_dim 256) next to +a SigLIP-so400m/14 tower, which is exactly the ``paligemma-3b-*-224`` +architecture, so the released checkpoint drops straight in. +""" + +import hashlib +import logging +import os +import re +from pathlib import Path + +import torch +from safetensors import safe_open + +logger = logging.getLogger(__name__) + +CANONICAL_REPO = "google/paligemma-3b-pt-224" + +# `CANONICAL_REPO` is license-gated: a machine without an HF token that has +# accepted the Gemma terms gets a 401 on the weight shards (the repo *metadata* +# stays anonymously readable, which is what makes the check below possible). +# These repos carry the same three shards. Every one is sha256-checked against +# the canonical repo's file metadata before it is used, so a mirror cannot +# substitute different weights -- it is only a transport. Opt-in +# (`paligemma_allow_mirror`): integrity is checked, but the weights of record +# should come from the canonical repo unless a run says otherwise. +MIRROR_REPOS = ( + "leo009/paligemma-3b-pt-224", + "hehe156/paligemma-3b-pt-224", +) + +# Only the shards: the loader reads nothing else and openpi builds the config in +# code, so an unverified json would only sit in the HF cache. +_ALLOW_PATTERNS = ["*.safetensors"] + +# Released PaliGemma checkpoint keys -> the module tree. These are what +# transformers applies via +# ``PaliGemmaForConditionalGeneration._checkpoint_conversion_mapping``; spelled +# out here because we stream the shards ourselves rather than going through +# ``from_pretrained`` (which would build a second full-size model alongside the +# one openpi already made). The vision rule is the transformers >= 5 spelling, +# which flattened ``SiglipVisionModel``; below 5 the ``vision_model.`` hop is +# still there, and ``_module_name`` picks whichever the live tree has. +# +# ``load_paligemma_weights`` fails loudly if these leave any parameter +# uncovered, so a future rename surfaces as an error, not a silent partial load. +_KEY_REMAP = ( + (re.compile(r"^language_model\.model\."), "model.language_model."), + (re.compile(r"^language_model\.lm_head\."), "lm_head."), + (re.compile(r"^vision_tower\.vision_model\."), "model.vision_tower."), + (re.compile(r"^vision_tower\."), "model.vision_tower."), + (re.compile(r"^multi_modal_projector\."), "model.multi_modal_projector."), +) + + +# The released checkpoint pads the token embedding from PaliGemma's 257152-token +# vocabulary up to 257216 rows. Row 257152 is -- which openpi never looks +# up, because ``embed_prefix`` concatenates SigLIP features directly instead of +# scattering them into an placeholder -- and the 63 rows after it are +# unused padding. openpi builds the table at exactly 257152, so these two +# tensors are truncated on their vocabulary axis; every real text token keeps its +# row, since the padding is at the end. +_VOCAB_TRUNCATABLE = frozenset( + {"model.language_model.embed_tokens.weight", "lm_head.weight"} +) + + +def _remap_key(key: str) -> str: + for pattern, replacement in _KEY_REMAP: + new, n = pattern.subn(replacement, key, count=1) + if n: + return new + return key + + +def _module_name(key: str, own: dict) -> str: + """Where checkpoint ``key`` lives in ``own``, the model's named tensors.""" + name = _remap_key(key) + if name not in own and key.startswith("vision_tower.vision_model."): + nested = "model." + key + if nested in own: + return nested + return name + + +def _fit(name: str, tensor: torch.Tensor, param: torch.Tensor) -> torch.Tensor: + if tensor.shape == param.shape: + return tensor + if ( + name in _VOCAB_TRUNCATABLE + and tensor.ndim == param.ndim == 2 + and tensor.shape[0] > param.shape[0] + and tensor.shape[1] == param.shape[1] + ): + logger.info( + "Truncating %s from %d to %d vocabulary rows (padding tokens).", + name, + tensor.shape[0], + param.shape[0], + ) + return tensor[: param.shape[0]] + raise ValueError( + f"Shape mismatch for {name}: checkpoint {tuple(tensor.shape)} " + f"vs model {tuple(param.shape)}" + ) + + +def select_init_source( + pytorch_weight_path: str | None, paligemma_weight_path: str | None +) -> str: + """Which set of starting weights a PI run uses: the full pi0.5 base + checkpoint (``"pi05_base"``), PaliGemma with a fresh action expert + (``"paligemma"``), or nothing at all (``"none"``).""" + if pytorch_weight_path is not None and paligemma_weight_path is not None: + raise ValueError( + "pytorch_weight_path and paligemma_weight_path are both set. The " + "pi0.5 base checkpoint already contains a trained PaliGemma; set " + "pytorch_weight_path=null to start from PaliGemma with a fresh " + "action expert." + ) + if pytorch_weight_path is not None: + return "pi05_base" + if paligemma_weight_path is not None: + return "paligemma" + return "none" + + +def _shard_digests(repo: str) -> dict[str, str]: + """``{shard filename: sha256}`` from a repo's file metadata.""" + from huggingface_hub import HfApi + + info = HfApi().model_info(repo, files_metadata=True) + digests = {} + for sibling in info.siblings: + if not sibling.rfilename.endswith(".safetensors"): + continue + lfs = getattr(sibling, "lfs", None) + sha = ( + lfs.get("sha256") if isinstance(lfs, dict) else getattr(lfs, "sha256", None) + ) + if sha: + digests[sibling.rfilename] = sha + return digests + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(16 << 20), b""): + digest.update(block) + return digest.hexdigest() + + +def _verify_shards(directory: str, expected: dict[str, str]) -> None: + root = Path(directory) + for name, sha in expected.items(): + shard = root / name + if not shard.is_file(): + raise FileNotFoundError(f"{shard} is missing from the mirror snapshot") + actual = _sha256(shard) + if actual != sha: + raise ValueError( + f"{shard} does not match {CANONICAL_REPO}: expected sha256 {sha}, " + f"got {actual}. Refusing to initialize from it." + ) + + +def resolve_paligemma_dir(source: str, *, allow_mirror: bool = False) -> str: + """Local directory holding the PaliGemma checkpoint named by ``source``. + + ``source`` is a local path (returned as-is) or a Hugging Face repo id. When + it is the gated canonical repo, this machine is not authorized for it (401 + or 403, not a rate limit or an outage) and ``allow_mirror`` is set, fall + back to a sha256-verified mirror. + """ + if os.path.isdir(source): + return source + + from huggingface_hub import snapshot_download + from huggingface_hub.errors import GatedRepoError, HfHubHTTPError + + try: + return snapshot_download(source, allow_patterns=_ALLOW_PATTERNS) + except HfHubHTTPError as gated: + status = getattr(getattr(gated, "response", None), "status_code", None) + unauthorized = isinstance(gated, GatedRepoError) or status in (401, 403) + if not (allow_mirror and unauthorized and source == CANONICAL_REPO): + raise + logger.warning( + "%s is gated for this machine (%s) and paligemma_allow_mirror is set. " + "Falling back to a mirror, sha256-verified against %s. Set HF_TOKEN " + "to a token that has accepted the Gemma terms to pull the canonical " + "repo instead.", + CANONICAL_REPO, + type(gated).__name__, + CANONICAL_REPO, + ) + + expected = _shard_digests(CANONICAL_REPO) + if not expected: + raise RuntimeError( + f"Cannot read shard checksums for {CANONICAL_REPO}; refusing to use " + "an unverified mirror." + ) + + failures = [] + for mirror in MIRROR_REPOS: + try: + if _shard_digests(mirror) != expected: + raise ValueError("shard checksums differ from the canonical repo") + directory = snapshot_download(mirror, allow_patterns=_ALLOW_PATTERNS) + _verify_shards(directory, expected) + except Exception as exc: # try the next mirror, report them all if none works + failures.append(f"{mirror}: {type(exc).__name__}: {exc}") + continue + logger.warning( + "Initializing PaliGemma from mirror %s (all %d shards sha256-match %s).", + mirror, + len(expected), + CANONICAL_REPO, + ) + return directory + + raise RuntimeError( + f"{CANONICAL_REPO} is gated and no verified mirror worked:\n " + + "\n ".join(failures) + ) + + +def load_paligemma_weights(model, source: str, *, allow_mirror: bool = False) -> str: + """Copy a PaliGemma checkpoint into ``model``'s VLM prefix, in place. + + ``model`` is an ``openpi.models_pytorch.pi0_pytorch.PI0Pytorch``. Only + ``paligemma_with_expert.paligemma`` is touched; the action expert and the + action/time projections are left at their init. Returns the directory the + weights came from. + """ + directory = resolve_paligemma_dir(source, allow_mirror=allow_mirror) + target = model.paligemma_with_expert.paligemma + + # remove_duplicate=False so a tied lm_head.weight is reachable under both + # names; the checkpoint only carries one of them. + params = dict(target.named_parameters(remove_duplicate=False)) + # Buffers are accepted if the checkpoint ships them, but are not required: + # rotary inv_freq, Gemma's embed_scale and the SigLIP position_ids are all + # derived from the config at construction time and are absent from the + # release. + own = dict(target.named_buffers(remove_duplicate=False)) + own.update(params) + + shards = sorted(Path(directory).glob("*.safetensors")) + if not shards: + raise FileNotFoundError(f"No .safetensors under {directory}") + + loaded, unexpected = set(), [] + for shard in shards: + with safe_open(shard, framework="pt") as handle: + for key in handle.keys(): # noqa: SIM118 - safe_open is not a Mapping + name = _module_name(key, own) + param = own.get(name) + if param is None: + unexpected.append(key) + continue + tensor = _fit(name, handle.get_tensor(key), param) + with torch.no_grad(): + param.copy_(tensor.to(param.dtype)) + loaded.add(name) + + # lm_head is tied to embed_tokens and, either way, dead weight here: no code + # path in openpi's PyTorch pi0 reads it. + tied = set(getattr(type(target), "_tied_weights_keys", None) or ()) + missing = {n for n in params if n not in loaded} - tied + if missing or unexpected: + raise RuntimeError( + f"PaliGemma load from {directory} did not cover the prefix exactly: " + f"{len(missing)} missing (e.g. {sorted(missing)[:5]}), " + f"{len(unexpected)} unexpected (e.g. {unexpected[:5]})." + ) + + logger.info( + "Initialized the VLM prefix from PaliGemma at %s (%d tensors, %d parameters). " + "The action expert and the action/time projections stay randomly initialized.", + directory, + len(loaded), + sum(p.numel() for p in target.parameters()), + ) + return directory diff --git a/pi05.md b/pi05.md index ddcd12b643..e380422434 100644 --- a/pi05.md +++ b/pi05.md @@ -25,3 +25,120 @@ uv run --active examples/convert_jax_model_to_pytorch.py \ --checkpoint_dir ../../egomimic/algo/pi_checkpoints/pi05_base \ --output_path ../../egomimic/algo/pi_checkpoints/pi05_base_pytorch ``` + +## Starting from PaliGemma instead of the pi0.5 base + +`model/pi0.5_paligemma_init.yaml` is an overlay that swaps the pi0.5 base +checkpoint for the public PaliGemma release: the VLM prefix (SigLIP tower, +multimodal projector, Gemma-2B) is loaded from `paligemma-3b-pt-224`, and +everything pi0.5 adds on top — the 300M action expert, the action in/out +projections, the adaRMS time MLPs — trains from scratch. That is the ablation +that separates "what the VLM knows" from "what pi0.5's robot pretraining added". + +A ready-made arm ships for the mecka 6D configs: + +```bash +python egomimic/trainHydra.py --config-name=train_zarr_cartesian_pi \ + data=mecka_fold_freeform_opscale1_6d \ + model=pi0.5_bc_mecka_6d_paligemma_init \ + evaluator=eval_pi_wristframe_6d +``` + +Layer the overlay onto any other `pi0.5_*` config the same way (`defaults: +[pi0.5_bc_aria, pi0.5_paligemma_init, _self_]`), or set the two keys per run: + +```bash +model.robomimic_model.config.pytorch_weight_path=null \ +model.robomimic_model.config.paligemma_weight_path=google/paligemma-3b-pt-224 +``` + +Setting both paths raises: the pi0.5 base checkpoint already contains a trained +PaliGemma, so letting one silently win would make a run's provenance +unreadable from its config. + +### Getting the weights + +`google/paligemma-3b-pt-224` is license-gated. Either authenticate + +```bash +huggingface-cli login # with a token that has accepted the Gemma terms +``` + +or set `paligemma_allow_mirror: true` to let `egomimic/models/paligemma_init.py` +fall back to a third-party mirror on a 401/403 — it compares every shard's +sha256 against the canonical repo's file metadata (readable without a token) +before using it, so a mirror is transport only and cannot substitute different +weights. It is off by default so the weights of record come from the canonical +repo. `$EGOVERSE_PALIGEMMA_WEIGHTS` points the arm at a local snapshot +directory and skips the Hub entirely. + +Pre-fetch once before a multi-rank run, or every rank races on the same +download lock: + +```bash +python -c "from egomimic.models.paligemma_init import resolve_paligemma_dir as r; print(r('google/paligemma-3b-pt-224'))" +``` + +### What the loader checks + +The released checkpoint does not drop in unchanged, so the load is strict about +the two places it differs and fails loudly anywhere else: + +- the checkpoint's `vision_tower.vision_model.*` keeps its `vision_model.` hop + below transformers 5 and loses it from 5 on (`SiglipVisionModel` was + flattened); the loader maps onto whichever the model has. +- the release pads the token embedding from PaliGemma's 257152-token vocabulary + to 257216 rows (`` plus 63 unused). openpi builds the table at exactly + 257152, so `embed_tokens` / `lm_head` are truncated on their vocabulary axis; + the padding is at the end, so every real text token keeps its row. + +Any parameter the checkpoint does not cover, and any checkpoint tensor with no +home in the module tree, raises — a partial load that quietly leaves half the +VLM random is the failure mode this exists to prevent. + +### Hyperparameters + +The overlay changes the init source and nothing else: the optimizer stays +`pi0.5_base`'s fine-tuning recipe (constant 3e-5 after 2000 warmup steps). A +randomly initialized action expert may want a higher peak and a longer warmup — +that is a per-run decision, e.g. + +```bash +model.optimizer.lr=1e-4 model.scheduler.num_warmup_steps=5000 +``` + +## transformers 4.53 vs 5.x + +`external/openpi` pins `transformers==4.53.2` and its `transformers_replace` +patch is written against that version (`models/siglip/check.py` asserts it). +So does `pyproject.toml`, and the install steps copy that patch into the venv: +4.53.2 plus the patch is the supported setup, what CI installs and what the +shared `emimic` venv carries. Which transformers a run gets is whatever its venv +has, so check before comparing numbers across machines. + +A venv on transformers 5.x (the PaliGemma-init arm was first run on 5.17 with a +partial hand-port of the patch: the Gemma modules carry openpi's adaRMS changes, +but `modeling_paligemma.py` and `modeling_siglip.py` are effectively stock 5.x) +also works, but three of openpi's assumptions do not survive it: + +- `PaliGemmaModel.get_image_features` now returns a `BaseModelOutputWithPooling` + with the projected features on `pooler_output`; openpi's `embed_prefix` reads + `.shape` off the result. +- transformers 5 dropped the back-compat properties that forwarded + `PaliGemmaForConditionalGeneration.language_model` / `.vision_tower` / + `.multi_modal_projector` to the inner `PaliGemmaModel`; openpi reads all three + off the outer model. +- openpi's float32 keep-list names the SigLIP patch and position embeddings + `vision_tower.vision_model.embeddings.*`; with `SiglipVisionModel` flattened + nothing matches and they would train in bf16. + +`egomimic/models/openpi_compat.py` bridges all three, applied from +`egomimic.algo.pi`. Each shim is guarded on what it finds and stands down +entirely below transformers 5, so porting the fork's patch forward makes the +module a no-op — delete it then. + +On 5.x only the PaliGemma-init arm loads. The pi0.5 base checkpoint +(`pi05_base_pytorch/model.safetensors`) stores the nested +`vision_tower.vision_model.*` names and `safetensors.torch.load_model` is +strict, so the default init path stops at startup on missing SigLIP keys. +Train from the pi0.5 base on the pinned 4.53.2. diff --git a/tests/integration/test_paligemma_init_real.py b/tests/integration/test_paligemma_init_real.py new file mode 100644 index 0000000000..e74135417f --- /dev/null +++ b/tests/integration/test_paligemma_init_real.py @@ -0,0 +1,253 @@ +"""The real PaliGemma release, loaded into the real openpi prefix. + +The unit tier proves the remapping and the coverage rules on a miniature module +tree; this is the part that can only be checked against the actual 3B +checkpoint -- that openpi's prefix and `paligemma-3b-pt-224` line up tensor for +tensor, that the values land unchanged, and that nothing outside the prefix +moves. It builds a 3.6B model and reads ~12 GB of weights, so it lives here. + + # the load itself: CPU, needs the checkpoint (or Hub access to fetch it) + python -m pytest --integration tests/integration/test_paligemma_init_real.py -q + # plus the flow-matching step, on a GPU node + python -m pytest --integration -m gpu tests/integration/test_paligemma_init_real.py -q +""" + +from __future__ import annotations + +import glob + +import pytest +import torch +from fixtures.train_harness import pi_unavailable +from safetensors import safe_open + +from egomimic.models.paligemma_init import ( + CANONICAL_REPO, + load_paligemma_weights, + resolve_paligemma_dir, +) + +pytestmark = pytest.mark.skipif( + pi_unavailable() is not None, reason=pi_unavailable() or "" +) + +PG = "paligemma_with_expert.paligemma." +EX = "paligemma_with_expert.gemma_expert." + +# These start random and must still be random afterwards: the whole point of +# this init is that pi0.5's own parameters are NOT pretrained. +FROM_SCRATCH = ( + EX + "model.layers.0.self_attn.q_proj.weight", + EX + "model.layers.17.mlp.up_proj.weight", + EX + "model.layers.0.input_layernorm.dense.weight", + EX + "model.norm.dense.weight", + "action_in_proj.weight", + "action_out_proj.weight", + "time_mlp_in.weight", + "time_mlp_out.weight", +) + +# module parameter -> the checkpoint key it must come from, one per remapping rule +PRETRAINED = { + PG + + "model.language_model.layers.0.self_attn.q_proj.weight": "language_model.model.layers.0.self_attn.q_proj.weight", + PG + + "model.language_model.layers.17.mlp.down_proj.weight": "language_model.model.layers.17.mlp.down_proj.weight", + PG + "model.language_model.norm.weight": "language_model.model.norm.weight", + PG + + "model.language_model.embed_tokens.weight": "language_model.model.embed_tokens.weight", + PG + + "model.vision_tower.encoder.layers.0.self_attn.q_proj.weight": "vision_tower.vision_model.encoder.layers.0.self_attn.q_proj.weight", + PG + + "model.vision_tower.encoder.layers.26.mlp.fc2.bias": "vision_tower.vision_model.encoder.layers.26.mlp.fc2.bias", + PG + + "model.vision_tower.embeddings.patch_embedding.weight": "vision_tower.vision_model.embeddings.patch_embedding.weight", + PG + + "model.multi_modal_projector.linear.weight": "multi_modal_projector.linear.weight", +} + + +@pytest.fixture(scope="module") +def checkpoint_dir() -> str: + try: + return resolve_paligemma_dir(CANONICAL_REPO) + except Exception as exc: # gated and offline, or no verified mirror reachable + pytest.skip(f"cannot resolve {CANONICAL_REPO}: {type(exc).__name__}: {exc}") + + +def _flat(named) -> dict: + """Parameters under their transformers-5 names, which the tables above use: + below 5 the SigLIP tree keeps a ``vision_model.`` hop.""" + return { + k.replace("vision_tower.vision_model.", "vision_tower."): v for k, v in named + } + + +@pytest.fixture(scope="module") +def loaded(checkpoint_dir): + """The real pi0.5 chassis with only its VLM prefix initialized.""" + import openpi.models.pi0_config + import openpi.models_pytorch.pi0_pytorch as pi0t + + # Production reaches openpi only through `egomimic.algo.pi`, which applies + # the transformers-5 shims at import. This builds PI0Pytorch directly, so + # take the same entry point rather than an unpatched openpi. + import egomimic.algo.pi # noqa: F401 + + config = openpi.models.pi0_config.Pi0Config( + dtype="bfloat16", + action_dim=32, + action_horizon=100, + max_token_len=180, + paligemma_variant="gemma_2b", + action_expert_variant="gemma_300m", + pi05=True, + ) + torch.manual_seed(0) + model = pi0t.PI0Pytorch(config) + before = _flat((k, v.detach().clone()) for k, v in model.named_parameters()) + load_paligemma_weights(model, checkpoint_dir) + return model, before, _flat(model.named_parameters()) + + +def test_pi05s_own_parameters_are_left_at_their_init(loaded): + _, before, after = loaded + for name in FROM_SCRATCH: + assert torch.equal(after[name], before[name]), f"{name} was overwritten" + + +def test_the_prefix_moved_off_its_init(loaded): + _, before, after = loaded + for name in PRETRAINED: + assert not torch.equal(after[name], before[name]), f"{name} was not loaded" + + +def test_the_prefix_matches_the_checkpoint_bit_for_bit(loaded, checkpoint_dir): + _, _, after = loaded + wanted = set(PRETRAINED.values()) + found = {} + for shard in sorted(glob.glob(checkpoint_dir + "/*.safetensors")): + with safe_open(shard, framework="pt") as handle: + for key in handle.keys(): # noqa: SIM118 - safe_open is not a Mapping + if key in wanted: + found[key] = handle.get_tensor(key) + assert set(found) == wanted, f"checkpoint is missing {wanted - set(found)}" + + for name, key in PRETRAINED.items(): + param = after[name] + reference = found[key] + if reference.shape != param.shape: # padded vocabulary + reference = reference[: param.shape[0]] + assert torch.equal( + param, reference.to(param.dtype) + ), f"value mismatch at {name}" + + +def test_openpis_mixed_precision_policy_survives_the_load(loaded): + # The norms openpi keeps in fp32 must not be dragged to bf16 by the copy, + # and the matmul weights must not be widened. + _, _, after = loaded + assert ( + after[PG + "model.language_model.layers.0.input_layernorm.weight"].dtype + is torch.float32 + ) + assert after[PG + "model.language_model.norm.weight"].dtype is torch.float32 + assert ( + after[PG + "model.language_model.layers.0.self_attn.q_proj.weight"].dtype + is torch.bfloat16 + ) + assert ( + after[PG + "model.vision_tower.encoder.layers.0.mlp.fc1.weight"].dtype + is torch.bfloat16 + ) + # openpi's keep-list, on either transformers (openpi_compat on 5.x) + for name in ( + "patch_embedding.weight", + "patch_embedding.bias", + "position_embedding.weight", + ): + assert ( + after[PG + "model.vision_tower.embeddings." + name].dtype is torch.float32 + ) + + +# ------------------------------------------------------------------ a real step + + +def _observation(batch: int, device, max_token_len: int): + """A minimal openpi observation: three camera slots, the wrists masked out + the way human data leaves them, and a prompt of real token ids.""" + from egomimic.models.preprocess_pi_obs import PI_CAMERA_SLOTS, _SimpleObservation + + image = torch.rand(batch, 3, 224, 224, device=device) * 2 - 1 + present = { + "base_0_rgb": True, + "left_wrist_0_rgb": False, + "right_wrist_0_rgb": False, + } + mask = torch.ones(batch, max_token_len, dtype=torch.bool, device=device) + return _SimpleObservation( + images={slot: image.clone() for slot in PI_CAMERA_SLOTS}, + image_masks={ + slot: torch.full((batch,), present[slot], dtype=torch.bool, device=device) + for slot in PI_CAMERA_SLOTS + }, + state=torch.zeros(batch, 32, device=device), + tokenized_prompt=torch.randint( + 0, 1000, (batch, max_token_len), device=device, dtype=torch.long + ), + tokenized_prompt_mask=mask, + token_ar_mask=mask.clone(), + token_loss_mask=mask.clone(), + ) + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_one_flow_matching_step_trains_both_halves(loaded): + """The point of the arm: a finite loss whose gradient reaches the frozen-in + PaliGemma prefix AND the action expert that started from noise. A load that + left half the model detached would still produce a loss here, so check the + gradients, not just the number.""" + model, _, _ = loaded + device = torch.device("cuda") + model = model.to(device) + model.train() + + observation = _observation(1, device, model.config.max_token_len) + actions = torch.randn( + 1, model.config.action_horizon, model.config.action_dim, device=device + ) + + torch.manual_seed(0) + losses = model.forward(observation, actions) + loss = ( + torch.stack(losses).mean() + if isinstance(losses, (list, tuple)) + else losses.mean() + ) + assert torch.isfinite(loss), f"non-finite flow-matching loss: {loss}" + + loss.backward() + probes = { + "prefix (pretrained)": model.paligemma_with_expert.paligemma.model.language_model.layers[ + 0 + ].self_attn.q_proj.weight, + "vision tower (pretrained)": model.paligemma_with_expert.paligemma.model.vision_tower.encoder.layers[ + 0 + ].self_attn.q_proj.weight, + "action expert (from scratch)": model.paligemma_with_expert.gemma_expert.model.layers[ + 0 + ].self_attn.q_proj.weight, + "action_out_proj (from scratch)": model.action_out_proj.weight, + } + for label, param in probes.items(): + assert param.grad is not None, f"no gradient reached {label}" + assert torch.isfinite(param.grad).all(), f"non-finite gradient at {label}" + assert param.grad.abs().sum() > 0, f"zero gradient at {label}" + + # and a step actually moves them + before = {k: v.detach().clone() for k, v in probes.items()} + torch.optim.SGD(model.parameters(), lr=1e-3).step() + for label, param in probes.items(): + assert not torch.equal(param, before[label]), f"{label} did not move" diff --git a/tests/unit/test_openpi_compat.py b/tests/unit/test_openpi_compat.py new file mode 100644 index 0000000000..951a90fa6a --- /dev/null +++ b/tests/unit/test_openpi_compat.py @@ -0,0 +1,115 @@ +"""openpi pins transformers 4.53 and ships a patch written against it; a venv +on transformers 5 needs these shims. The shims in egomimic/models/openpi_compat.py bridge the +gap and must disappear on their own once the fork's patch is ported forward.""" + +import pytest +import torch + +pytest.importorskip("openpi.models_pytorch.gemma_pytorch") + +from openpi.models_pytorch.gemma_pytorch import PaliGemmaWithExpertModel # noqa: E402 + +from egomimic.models import openpi_compat # noqa: E402 + + +@pytest.fixture(autouse=True) +def _restore_embed_image(): + original = PaliGemmaWithExpertModel.embed_image + cast = PaliGemmaWithExpertModel.to_bfloat16_for_selected_params + applied = openpi_compat._APPLIED + yield + PaliGemmaWithExpertModel.embed_image = original + PaliGemmaWithExpertModel.to_bfloat16_for_selected_params = cast + openpi_compat._APPLIED = applied + + +class _Outputs: + """What transformers 5's ``get_image_features`` hands back: the vision + output object with the projected features on ``pooler_output``.""" + + def __init__(self, features): + self.last_hidden_state = torch.zeros_like(features) + self.pooler_output = features + + +class _Prefix: + def __init__(self, result): + self._result = result + + class _Inner: + def __init__(self, result): + self._result = result + + def get_image_features(self, image): + return self._result + + @property + def model(self): + return self._Inner(self._result) + + +def _embed(result): + fake = object.__new__(PaliGemmaWithExpertModel) + fake.paligemma = _Prefix(result) + return PaliGemmaWithExpertModel.embed_image(fake, torch.zeros(1, 3, 224, 224)) + + +def test_a_wrapped_output_is_unwrapped_to_its_features(): + # Without this, openpi's embed_prefix reads `.shape` off the output object. + features = torch.randn(2, 256, 2048) + openpi_compat._patch_embed_image() + assert torch.equal(_embed(_Outputs(features)), features) + + +def test_a_plain_tensor_passes_straight_through(): + # openpi's own patched transformers already returns the tensor; the shim + # must not change that behaviour. + features = torch.randn(2, 256, 2048) + openpi_compat._patch_embed_image() + assert torch.equal(_embed(features), features) + + +def test_apply_is_a_no_op_the_second_time(caplog): + openpi_compat._APPLIED = False + openpi_compat.apply() + first = PaliGemmaWithExpertModel.embed_image + openpi_compat.apply() + assert PaliGemmaWithExpertModel.embed_image is first + + +def test_the_shim_stands_down_on_the_transformers_openpi_pins(monkeypatch): + import transformers + + original = PaliGemmaWithExpertModel.embed_image + monkeypatch.setattr(transformers, "__version__", "4.53.2") + openpi_compat._APPLIED = False + openpi_compat.apply() + assert PaliGemmaWithExpertModel.embed_image is original + + +def _siglip_embeddings(nested: bool) -> torch.nn.Module: + root = torch.nn.Module() + tower = torch.nn.Module() + parent = tower + if nested: + tower.vision_model = torch.nn.Module() + parent = tower.vision_model + parent.embeddings = torch.nn.Module() + parent.embeddings.patch_embedding = torch.nn.Conv2d(3, 4, 2) + parent.embeddings.position_embedding = torch.nn.Embedding(4, 4) + parent.encoder = torch.nn.Linear(4, 4) + root.vision_tower = tower + return root + + +@pytest.mark.parametrize("nested", [True, False]) +def test_the_siglip_embeddings_stay_float32(nested): + """openpi's keep-list matches the nested 4.x names only; with the shim the + flattened 5.x names are kept too, and everything else still goes bf16.""" + openpi_compat._patch_fp32_keep_list() + module = _siglip_embeddings(nested) + PaliGemmaWithExpertModel.to_bfloat16_for_selected_params(module, "bfloat16") + dtypes = {n: p.dtype for n, p in module.named_parameters()} + for name, dtype in dtypes.items(): + expected = torch.float32 if "embeddings" in name else torch.bfloat16 + assert dtype == expected, name diff --git a/tests/unit/test_paligemma_init.py b/tests/unit/test_paligemma_init.py new file mode 100644 index 0000000000..61e9f0d3ed --- /dev/null +++ b/tests/unit/test_paligemma_init.py @@ -0,0 +1,457 @@ +"""PaliGemma-only init: the VLM prefix comes from the public PaliGemma release +and everything pi0.5 adds on top starts from its PyTorch init. + +No network and no 3B model: the loader is exercised against a miniature module +tree, and against a tiny real ``PaliGemmaForConditionalGeneration`` built by +whichever transformers is installed (the vision tree differs across 4.x and +5.x), with safetensors written on the fly; the download resolution is driven +through stubs. The one thing that needs the real thing -- +that the released checkpoint covers openpi's prefix exactly -- is asserted by +``load_paligemma_weights`` itself (it raises on any missing or unexpected +tensor), so a rename upstream cannot pass silently. +""" + +import pytest +import torch +from safetensors.torch import save_file +from torch import nn + +from egomimic.models import paligemma_init as pgi + +CANONICAL = pgi.CANONICAL_REPO + + +# ---------------------------------------------------------------- key remapping + + +# Spelled exactly as google/paligemma-3b-pt-224 ships them. +@pytest.mark.parametrize( + ("checkpoint_key", "module_name"), + [ + ( + "language_model.model.layers.0.self_attn.q_proj.weight", + "model.language_model.layers.0.self_attn.q_proj.weight", + ), + ( + "language_model.model.embed_tokens.weight", + "model.language_model.embed_tokens.weight", + ), + ("language_model.model.norm.weight", "model.language_model.norm.weight"), + ("language_model.lm_head.weight", "lm_head.weight"), + # transformers >= 5 flattened SiglipVisionModel: the `.vision_model.` hop + # in the checkpoint has no counterpart in the module tree. + ( + "vision_tower.vision_model.encoder.layers.26.mlp.fc2.bias", + "model.vision_tower.encoder.layers.26.mlp.fc2.bias", + ), + ( + "vision_tower.vision_model.embeddings.patch_embedding.weight", + "model.vision_tower.embeddings.patch_embedding.weight", + ), + ( + "multi_modal_projector.linear.bias", + "model.multi_modal_projector.linear.bias", + ), + ], +) +def test_checkpoint_keys_map_onto_the_module_tree(checkpoint_key, module_name): + assert pgi._remap_key(checkpoint_key) == module_name + + +def test_unknown_prefix_is_left_alone(): + assert pgi._remap_key("action_in_proj.weight") == "action_in_proj.weight" + + +# ------------------------------------------------------------- vocab truncation + + +def test_padded_vocabulary_is_truncated_to_the_model_width(): + # The release pads 257152 -> 257216; the extra rows are + padding. + tensor, param = torch.arange(12.0).reshape(6, 2), torch.zeros(4, 2) + fitted = pgi._fit("model.language_model.embed_tokens.weight", tensor, param) + assert fitted.shape == param.shape + assert torch.equal(fitted, tensor[:4]) # padding is at the END + + +def test_truncation_is_not_applied_to_other_tensors(): + with pytest.raises(ValueError, match="Shape mismatch"): + pgi._fit( + "model.vision_tower.encoder.layers.0.mlp.fc1.weight", + torch.zeros(6, 2), + torch.zeros(4, 2), + ) + + +def test_a_narrower_checkpoint_is_never_padded_out(): + with pytest.raises(ValueError, match="Shape mismatch"): + pgi._fit("lm_head.weight", torch.zeros(3, 2), torch.zeros(4, 2)) + + +# ------------------------------------------------------------------ init source + + +def test_pi05_base_wins_when_only_it_is_set(): + assert pgi.select_init_source("/ckpt/pi05_base_pytorch", None) == "pi05_base" + + +def test_paligemma_when_the_base_checkpoint_is_off(): + assert pgi.select_init_source(None, CANONICAL) == "paligemma" + + +def test_neither_is_reported_as_no_weights(): + assert pgi.select_init_source(None, None) == "none" + + +def test_both_at_once_is_refused(): + # The pi0.5 base already contains a trained PaliGemma; silently letting one + # win would make the run's provenance unreadable from its config. + with pytest.raises(ValueError, match="both set"): + pgi.select_init_source("/ckpt/pi05_base_pytorch", CANONICAL) + + +# ------------------------------------------------------------ the loader itself + + +class _Prefix(nn.Module): + """Shaped like the part of ``PaliGemmaForConditionalGeneration`` the loader + walks: a ``model.`` subtree, a tied ``lm_head`` and a derived buffer.""" + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, vocab=4, width=3): + super().__init__() + self.model = nn.Module() + self.model.language_model = nn.Module() + self.model.language_model.embed_tokens = nn.Embedding(vocab, width) + self.model.vision_tower = nn.Module() + self.model.vision_tower.post_layernorm = nn.Linear(width, width) + self.lm_head = nn.Linear(width, vocab, bias=False) + self.lm_head.weight = self.model.language_model.embed_tokens.weight # tied + # config-derived, absent from every released checkpoint + self.model.language_model.register_buffer("inv_freq", torch.ones(width)) + + +class _Model(nn.Module): + def __init__(self): + super().__init__() + self.paligemma_with_expert = nn.Module() + self.paligemma_with_expert.paligemma = _Prefix() + self.action_in_proj = nn.Linear(3, 3) + + +def _write_checkpoint(tmp_path, tensors, shards=2): + items = list(tensors.items()) + size = -(-len(items) // shards) + for i in range(shards): + chunk = dict(items[i * size : (i + 1) * size]) + if chunk: + save_file(chunk, str(tmp_path / f"model-{i}.safetensors")) + return str(tmp_path) + + +def _full_checkpoint(vocab=4, width=3): + torch.manual_seed(1) + return { + "language_model.model.embed_tokens.weight": torch.randn(vocab, width), + "vision_tower.vision_model.post_layernorm.weight": torch.randn(width, width), + "vision_tower.vision_model.post_layernorm.bias": torch.randn(width), + } + + +def test_the_prefix_loads_and_nothing_else_moves(tmp_path): + model = _Model() + before = model.action_in_proj.weight.detach().clone() + tensors = _full_checkpoint() + directory = _write_checkpoint(tmp_path, tensors) + + assert pgi.load_paligemma_weights(model, directory) == directory + + prefix = model.paligemma_with_expert.paligemma + assert torch.equal( + prefix.model.language_model.embed_tokens.weight, + tensors["language_model.model.embed_tokens.weight"], + ) + assert torch.equal( + prefix.model.vision_tower.post_layernorm.bias, + tensors["vision_tower.vision_model.post_layernorm.bias"], + ) + # everything outside the prefix is untouched + assert torch.equal(model.action_in_proj.weight, before) + + +def test_the_tied_head_follows_the_embedding(tmp_path): + # The checkpoint carries no lm_head; tying must still leave it consistent. + model = _Model() + tensors = _full_checkpoint() + pgi.load_paligemma_weights(model, _write_checkpoint(tmp_path, tensors)) + prefix = model.paligemma_with_expert.paligemma + assert torch.equal( + prefix.lm_head.weight, prefix.model.language_model.embed_tokens.weight + ) + + +def test_a_derived_buffer_is_not_required(tmp_path): + model = _Model() + pgi.load_paligemma_weights(model, _write_checkpoint(tmp_path, _full_checkpoint())) + assert torch.equal( + model.paligemma_with_expert.paligemma.model.language_model.inv_freq, + torch.ones(3), + ) + + +def test_a_partial_checkpoint_is_refused(tmp_path): + # Silently leaving part of the VLM at its random init is the failure this + # whole module exists to avoid. + tensors = _full_checkpoint() + del tensors["vision_tower.vision_model.post_layernorm.weight"] + with pytest.raises(RuntimeError, match="missing"): + pgi.load_paligemma_weights(_Model(), _write_checkpoint(tmp_path, tensors)) + + +def test_an_unknown_tensor_is_refused(tmp_path): + tensors = _full_checkpoint() + tensors["vision_tower.vision_model.renamed_upstream.weight"] = torch.zeros(3) + with pytest.raises(RuntimeError, match="unexpected"): + pgi.load_paligemma_weights(_Model(), _write_checkpoint(tmp_path, tensors)) + + +def test_an_empty_directory_is_refused(tmp_path): + with pytest.raises(FileNotFoundError, match="safetensors"): + pgi.load_paligemma_weights(_Model(), str(tmp_path)) + + +def test_a_padded_vocabulary_checkpoint_loads(tmp_path): + # The real release is 64 rows wider than openpi's table. + tensors = _full_checkpoint() + wide = torch.randn(6, 3) + tensors["language_model.model.embed_tokens.weight"] = wide + pgi.load_paligemma_weights(_Model(), _write_checkpoint(tmp_path, tensors)) + + +# ------------------------------------------------------- download / mirror path + + +def _tiny_paligemma(): + from transformers import PaliGemmaConfig, PaliGemmaForConditionalGeneration + + config = PaliGemmaConfig( + text_config={ + "model_type": "gemma", + "num_hidden_layers": 1, + "hidden_size": 32, + "intermediate_size": 32, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "head_dim": 32, + "vocab_size": 300, + }, + vision_config={ + "num_hidden_layers": 1, + "hidden_size": 16, + "intermediate_size": 16, + "num_attention_heads": 1, + "image_size": 28, + "patch_size": 14, + "projection_dim": 32, + }, + image_token_index=299, + projection_dim=32, + ) + return PaliGemmaForConditionalGeneration(config) + + +def _release_key(name: str) -> str: + """The module-tree name as google/paligemma-3b-pt-224 spells it.""" + for old, new in ( + ("model.language_model.", "language_model.model."), + ("lm_head.", "language_model.lm_head."), + ("model.vision_tower.vision_model.", "vision_tower.vision_model."), + ("model.vision_tower.", "vision_tower.vision_model."), + ("model.multi_modal_projector.", "multi_modal_projector."), + ): + if name.startswith(old): + return new + name[len(old) :] + return name + + +def test_a_real_paligemma_loads_on_this_transformers(tmp_path): + """The miniature tree above is spelled the 5.x way; a real model is spelled + whichever way the installed transformers spells it, which is how a 4.53 + venv came to fail 21 SigLIP tensors while the key tests all passed.""" + source = _tiny_paligemma() + released = { + _release_key(n): p.detach().clone().contiguous() + for n, p in source.named_parameters() + } + target = _tiny_paligemma() + model = nn.Module() + model.paligemma_with_expert = nn.Module() + model.paligemma_with_expert.paligemma = target + pgi.load_paligemma_weights(model, _write_checkpoint(tmp_path, released)) + for name, param in target.named_parameters(): + assert torch.equal(param, dict(source.named_parameters())[name]), name + + +def test_a_nested_vision_tree_is_found(): + own = {"model.vision_tower.vision_model.post_layernorm.bias": None} + key = "vision_tower.vision_model.post_layernorm.bias" + assert ( + pgi._module_name(key, own) + == "model.vision_tower.vision_model.post_layernorm.bias" + ) + assert pgi._module_name(key, {}) == "model.vision_tower.post_layernorm.bias" + + +def test_only_the_shards_are_downloaded(): + assert pgi._ALLOW_PATTERNS == ["*.safetensors"] + + +def test_a_local_directory_is_used_as_is(tmp_path): + assert pgi.resolve_paligemma_dir(str(tmp_path)) == str(tmp_path) + + +def _gated_error(repo_id): + """The 401 huggingface_hub raises for a license-gated repo.""" + import requests + from huggingface_hub.errors import GatedRepoError + + response = requests.Response() + response.status_code = 401 + return GatedRepoError(f"401 for {repo_id}", response=response) + + +def _stub_hub(monkeypatch, *, gated, digests, downloads): + import huggingface_hub + + def snapshot_download(repo_id, **_): + if repo_id in gated: + raise _gated_error(repo_id) + return downloads[repo_id] + + monkeypatch.setattr(huggingface_hub, "snapshot_download", snapshot_download) + monkeypatch.setattr(pgi, "_shard_digests", lambda repo: digests[repo]) + monkeypatch.setattr(pgi, "_verify_shards", lambda directory, expected: None) + + +def test_the_canonical_repo_is_preferred(monkeypatch, tmp_path): + _stub_hub( + monkeypatch, gated=set(), digests={}, downloads={CANONICAL: str(tmp_path)} + ) + assert pgi.resolve_paligemma_dir(CANONICAL) == str(tmp_path) + + +def test_a_gated_canonical_repo_falls_back_to_a_matching_mirror(monkeypatch, tmp_path): + good = pgi.MIRROR_REPOS[0] + digests = { + CANONICAL: {"model.safetensors": "abc"}, + good: {"model.safetensors": "abc"}, + } + digests.update({m: {"model.safetensors": "abc"} for m in pgi.MIRROR_REPOS}) + _stub_hub( + monkeypatch, + gated={CANONICAL}, + digests=digests, + downloads={m: str(tmp_path) for m in pgi.MIRROR_REPOS}, + ) + assert pgi.resolve_paligemma_dir(CANONICAL, allow_mirror=True) == str(tmp_path) + + +def test_the_mirror_is_opt_in(monkeypatch, tmp_path): + from huggingface_hub.errors import GatedRepoError + + digests = {r: {"model.safetensors": "abc"} for r in (CANONICAL, *pgi.MIRROR_REPOS)} + _stub_hub( + monkeypatch, + gated={CANONICAL}, + digests=digests, + downloads={m: str(tmp_path) for m in pgi.MIRROR_REPOS}, + ) + with pytest.raises(GatedRepoError): + pgi.resolve_paligemma_dir(CANONICAL) + + +def test_a_rate_limit_is_not_a_reason_to_use_a_mirror(monkeypatch): + import huggingface_hub + import requests + from huggingface_hub.errors import HfHubHTTPError + + response = requests.Response() + response.status_code = 429 + + def snapshot_download(repo_id, **_): + raise HfHubHTTPError("429", response=response) + + monkeypatch.setattr(huggingface_hub, "snapshot_download", snapshot_download) + with pytest.raises(HfHubHTTPError): + pgi.resolve_paligemma_dir(CANONICAL, allow_mirror=True) + + +def test_a_mirror_whose_shards_differ_is_rejected(monkeypatch, tmp_path): + # A mirror is transport only: different bytes means it is not PaliGemma. + digests = {CANONICAL: {"model.safetensors": "abc"}} + digests.update({m: {"model.safetensors": "TAMPERED"} for m in pgi.MIRROR_REPOS}) + _stub_hub( + monkeypatch, + gated={CANONICAL}, + digests=digests, + downloads={m: str(tmp_path) for m in pgi.MIRROR_REPOS}, + ) + with pytest.raises(RuntimeError, match="no verified mirror"): + pgi.resolve_paligemma_dir(CANONICAL, allow_mirror=True) + + +def test_the_mirror_fallback_can_be_switched_off(monkeypatch, tmp_path): + from huggingface_hub.errors import GatedRepoError + + _stub_hub(monkeypatch, gated={CANONICAL}, digests={}, downloads={}) + with pytest.raises(GatedRepoError): + pgi.resolve_paligemma_dir(CANONICAL, allow_mirror=False) + + +def test_a_non_canonical_repo_never_falls_back(monkeypatch, tmp_path): + # Only the gated canonical repo has verified stand-ins; anything else the + # user names must fail loudly rather than silently load other weights. + from huggingface_hub.errors import GatedRepoError + + _stub_hub( + monkeypatch, gated={"someone/private-paligemma"}, digests={}, downloads={} + ) + with pytest.raises(GatedRepoError): + pgi.resolve_paligemma_dir("someone/private-paligemma", allow_mirror=True) + + +# -------------------------------------------------------------------- the config + + +def test_the_overlay_flips_the_init_source_and_keeps_its_parent(compose_resolve): + cfg = compose_resolve( + "train_zarr_cartesian_pi", ["model=pi0.5_bc_mecka_6d_paligemma_init"] + ) + model = cfg.model.robomimic_model + assert model.config.pytorch_weight_path is None + assert model.config.paligemma_weight_path == CANONICAL + assert ( + pgi.select_init_source( + model.config.pytorch_weight_path, model.config.paligemma_weight_path + ) + == "paligemma" + ) + # the two keys are all the overlay touches: pi0.5_bc_mecka_6d's own choices + # (cartesian 6D action, no Embodiment prompt block) must survive it + assert model.ac_keys.human_bimanual == "actions_cartesian" + assert model.embodiment_label is False + assert model.config.model.action_dim == 32 + assert ( + model.action_converters.rules.HUMAN_BIMANUAL._target_ + == "egomimic.utils.action_utils.HumanBimanualCartesian6D" + ) + + +def test_the_base_pi_config_still_starts_from_pi05_base(compose_resolve): + cfg = compose_resolve("train_zarr_cartesian_pi", ["model=pi0.5_bc_mecka_6d"]) + config = cfg.model.robomimic_model.config + assert config.paligemma_weight_path is None + assert ( + pgi.select_init_source(config.pytorch_weight_path, config.paligemma_weight_path) + == "pi05_base" + ) From 53036e37d491bd132ec493cb5a191d6e5dca3b11 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 22:18:17 +0000 Subject: [PATCH 36/79] loader: image history (a past front frame, gap in seconds) and the episode fps in every sample Human.get_keymap(image_history_gap_s=...) adds a _hist camera entry; ZarrDataset reads it round(gap x fps) frames back, clamped at the episode start. Every sample also carries the episode's fps (zarr attr). The past frame names the front camera's zarr key, so the loader's one-call-per- key read covers both windows. JPEG keys are one frame per chunk, and a plain union span would read every frame in between; the read planner keeps them two rows. Per sample on a synthetic 30 fps episode of 360x640 JPEGs (median of 7x40 samples, login node, noisy): no history 11-12 ms; gap 0.1 s 21-24 -> 19-23 ms (the second decode, which history needs); 0.5 s 30-34 -> 20-21 ms; 2 s 53-61 -> 19 ms, flat in the gap. A test counts the JPEG rows read at 0.1, 0.5 and 2 s. Nothing consumes _hist here; the keymap docstring says so. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_019DbPWJCKtVg5c2PNHs7egw --- egomimic/rldb/embodiment/embodiment.py | 3 + egomimic/rldb/embodiment/human.py | 15 ++++ egomimic/rldb/zarr/zarr_dataset_multi.py | 17 ++++ tests/unit/test_image_history.py | 102 +++++++++++++++++++++++ 4 files changed, 137 insertions(+) create mode 100644 tests/unit/test_image_history.py diff --git a/egomimic/rldb/embodiment/embodiment.py b/egomimic/rldb/embodiment/embodiment.py index 83a00db10b..1013d3c429 100644 --- a/egomimic/rldb/embodiment/embodiment.py +++ b/egomimic/rldb/embodiment/embodiment.py @@ -15,6 +15,9 @@ _viz_traj, ) +# A camera keymap entry "_hist" is the same camera, `lag_s` seconds earlier. +IMAGE_HISTORY_SUFFIX = "_hist" + class EMBODIMENT(Enum): # All human demonstration data is one embodiment (HUMAN_*); the robot Eva is diff --git a/egomimic/rldb/embodiment/human.py b/egomimic/rldb/embodiment/human.py index d378dfbc8e..111f9e5dee 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -5,6 +5,7 @@ import numpy as np from egomimic.rldb.embodiment.embodiment import ( + IMAGE_HISTORY_SUFFIX, Embodiment, _reject_legacy_rotation, _strip_pi_keymap_mode, @@ -203,6 +204,7 @@ def get_keymap( annotation_key: str = None, proprio_history: int = 1, history_stride: int = 1, + image_history_gap_s: float | None = None, ): """Build the keymap. Per-vendor knobs are explicit args from the data config: ``has_head_pose`` (False only for contributors that truly omit @@ -221,6 +223,13 @@ def get_keymap( because at 30 fps consecutive frames are nearly a duplicate of the current pose (median wrist translation 6.8 mm at lag 1 against 42.9 mm at lag 5), so s = 1 spends tokens on sensor noise. + + ``image_history_gap_s`` adds a second front-camera entry, + ``_hist``: the frame that many seconds before the current + one (frames = gap x the episode's fps; frame 0 at an episode start). + The algo must consume it explicitly: HPT has no stem for it and skips + augs and normalization on cameras outside its encoders, so under plain + HPT it is decoded and dropped. """ key_map = cls._get_keymap( keymap_mode, @@ -230,6 +239,12 @@ def get_keymap( proprio_history=proprio_history, history_stride=history_stride, ) + if image_history_gap_s is not None: + front = key_map[cls.VIZ_IMAGE_KEY] + key_map[f"{cls.VIZ_IMAGE_KEY}{IMAGE_HISTORY_SUFFIX}"] = { + **front, + "lag_s": float(image_history_gap_s), + } if annotation_key is not None and not norm_mode: key_map[annotation_key] = { "key_type": "annotation_keys", diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index d2e43ddc22..204db1918f 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -2187,6 +2187,16 @@ def _proprio_history_mask(self, idx: int) -> np.ndarray | None: mask[history - n_real :] = 1.0 return mask + def _fps(self) -> float: + """The episode's frame rate (zarr ``fps`` attr; 30 when absent).""" + return float(self.metadata.get("fps") or 30) + + def _lagged_idx(self, idx: int, lag_s: float) -> int: + """Frame ``lag_s`` seconds before ``idx`` (at least one frame back), + clamped at the episode start -- so the first samples read frame 0, + the same front padding ``_front_pad_history`` gives proprio.""" + return max(0, idx - max(1, round(float(lag_s) * self._fps()))) + @staticmethod def _stride_history(array: np.ndarray, history: int, stride: int) -> np.ndarray: """Take the current frame and every ``stride``-th frame before it, in @@ -2238,11 +2248,16 @@ def _read_windows(self, idx: int) -> dict[str, tuple[str, tuple[int, int | None] continue horizon = spec.get("horizon") history = spec.get("history") + lag_s = spec.get("lag_s") if horizon is not None: interval = ( idx, self._chunk_end_idx(idx, horizon, spec.get("key_type")), ) + elif lag_s is not None: + # A single PAST frame (image history): the same zarr array as + # the current-frame entry, read ``lag_s`` seconds back. + interval = (self._lagged_idx(idx, lag_s), None) elif history is not None and int(history) > 1: # Backward window ending at (and including) the current frame, # every `history_stride`-th step; clamped at the episode start, @@ -2403,6 +2418,8 @@ def _next(reason: str, key: str = "") -> int: data[k] = torch.from_numpy(v).to(torch.float32) data["embodiment"] = get_embodiment_id(self.embodiment) + # Control frequency of this episode (RDT's ctrl-freq token). + data["fps"] = torch.tensor(self._fps(), dtype=torch.float32) # Per-episode camera intrinsics travel with the batch so a single # data-driven embodiment can project without a hardcoded class const. # Single-camera -> (3,4) K matrix; no/multi-camera -> NaN sentinel, diff --git a/tests/unit/test_image_history.py b/tests/unit/test_image_history.py new file mode 100644 index 0000000000..3bdaa6da67 --- /dev/null +++ b/tests/unit/test_image_history.py @@ -0,0 +1,102 @@ +"""Image history in the data path: ``Human.get_keymap(image_history_gap_s=...)`` +adds a ``_hist`` camera entry that ``ZarrDataset`` reads that many +seconds back (in the episode's own fps), and every sample carries ``fps``.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from egomimic.rldb.embodiment.human import Human +from egomimic.rldb.zarr.zarr_dataset_multi import ( + LocalEpisodeResolver, + MultiDataset, + ZarrEpisode, +) +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +T = 90 +CAM = Human.VIZ_IMAGE_KEY +PAST = f"{CAM}_hist" + + +def _leaf(tmp_path, fps: int, **keymap_kwargs): + rng = np.random.default_rng(0) + t = np.arange(T, dtype=np.float64)[:, None] + pose = np.concatenate( + [t * 0.001 + np.zeros((T, 3)), np.tile([[1.0, 0.0, 0.0, 0.0]], (T, 1))], axis=1 + ) + ZarrWriter.create_and_write( + tmp_path / "ep.zarr", + numeric_data={ + "left.obs_ee_pose": pose, + "right.obs_ee_pose": pose, + "obs_head_pose": pose, + }, + image_data={"images.front_1": rng.integers(0, 255, (T, 32, 32, 3), np.uint8)}, + embodiment="human_bimanual", + fps=fps, + task_name="synthetic", + intrinsics={"front_1": np.eye(3, 4)}, + ) + key_map = Human.get_keymap(keymap_mode="cartesian", **keymap_kwargs) + resolver = LocalEpisodeResolver(tmp_path, key_map=key_map, transform_list=[]) + ds = MultiDataset._from_resolver(resolver, mode="total") + return next(iter(ds.datasets.values())) + + +@pytest.mark.parametrize("fps,gap_frames", [(30, 3), (10, 1), (60, 6)]) +def test_past_frame_is_gap_seconds_back_in_the_episodes_fps(tmp_path, fps, gap_frames): + leaf = _leaf(tmp_path, fps, image_history_gap_s=0.1) + idx = 8 + assert torch.equal(leaf[idx][PAST], leaf[idx - gap_frames][CAM]) + assert not torch.equal(leaf[idx][PAST], leaf[idx][CAM]) + assert leaf[idx]["fps"].item() == fps + + +def test_past_frame_clamps_to_the_episode_start(tmp_path): + leaf = _leaf(tmp_path, 30, image_history_gap_s=0.1) + for idx in (0, 1, 2): + assert torch.equal(leaf[idx][PAST], leaf[0][CAM]) + + +def test_a_gap_below_one_frame_still_reads_the_previous_frame(tmp_path): + leaf = _leaf(tmp_path, 30, image_history_gap_s=0.001) + assert torch.equal(leaf[5][PAST], leaf[4][CAM]) + + +def test_default_keymap_has_no_history_entry_but_samples_carry_fps(tmp_path): + leaf = _leaf(tmp_path, 30) + sample = leaf[3] + assert PAST not in sample and sample[CAM].shape == (3, 32, 32) + assert sample["fps"].dtype == torch.float32 + assert "proprio_history_mask" not in sample # a camera lag is not proprio history + + +def test_norm_mode_keymap_drops_the_history_camera_too(): + key_map = Human.get_keymap("cartesian", norm_mode=True, image_history_gap_s=0.1) + assert not [k for k in key_map if "img" in k] + + +@pytest.mark.parametrize("gap_s", [0.1, 0.5, 2.0]) +def test_the_past_frame_costs_one_jpeg_row_whatever_the_gap( + tmp_path, monkeypatch, gap_s +): + """The past frame shares the front camera's zarr key, so the read merges the + two windows. JPEGs are one frame per chunk: a plain union span would read + every frame between them (61 rows at 2 s and 30 fps) to use two.""" + leaf = _leaf(tmp_path, 30, image_history_gap_s=gap_s) + rows: list[dict] = [] + inner = ZarrEpisode.read_intervals + + def spy(self, intervals): + out = inner(self, intervals) + rows.append({k: sum(len(r) for _, r in g) for k, g in out.items()}) + return out + + monkeypatch.setattr(ZarrEpisode, "read_intervals", spy) + idx = 80 + sample = leaf[idx] + assert rows[-1]["images.front_1"] == 2 + assert torch.equal(sample[PAST], leaf[idx - round(gap_s * 30)][CAM]) From 07b119188549f26e5307b6151d4ec44ec7356968 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 22:18:20 +0000 Subject: [PATCH 37/79] vision: DINOv3Stem, a timm DINOv3 ViT image stem (patch tokens, any aspect ratio) Same call convention as ResNet. Loaded through timm: the locked transformers (4.53.2) has no DINOv3 class, and the timm hub repos are ungated. RoPE-only, so image_size may be [H, W]; exposes grid_size. pretrained_reference_state_dict reads the checkpoint back from the local HF cache (never the network), so the pretrained-weight hash check covers this tower as it does the Qwen stems; an uncached checkpoint or a random tower is skipped. Against the real timm/vit_base_patch16_dinov3.lvd1689m file the key sets match (162 tensors) and every value is equal. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_019DbPWJCKtVg5c2PNHs7egw --- egomimic/models/hpt_nets.py | 118 +++++++++++++++++++++++++++++++++ tests/unit/test_dinov3_stem.py | 101 ++++++++++++++++++++++++++++ 2 files changed, 219 insertions(+) create mode 100644 tests/unit/test_dinov3_stem.py diff --git a/egomimic/models/hpt_nets.py b/egomimic/models/hpt_nets.py index 746dcf650e..4fe4ac0cac 100644 --- a/egomimic/models/hpt_nets.py +++ b/egomimic/models/hpt_nets.py @@ -956,6 +956,124 @@ def _local_snapshot_if_offline(model_name: str) -> str: return snapshot_download(model_name, local_files_only=True) +class DINOv3Stem(PretrainedWeights, PolicyStem): + """Pretrained DINOv3 ViT (timm), a drop-in replacement for ``ResNet``: + ``(B, T, N, 3, H, W)`` images in [0, 1] in, ``(B, patches, output_dim)`` out. + The stem owns the resize and the checkpoint's mean / std, so the config + must NOT normalize these images in its augs. + + Loaded through timm rather than transformers: the ``timm/*_dinov3`` hub + repos are not gated (``facebook/dinov3-*`` is), and DINOv3 in transformers + needs >= 4.56, above our pin. + + Args: + model_name: timm model id, e.g. ``vit_base_patch16_dinov3.lvd1689m`` + (86M) or ``vit_large_patch16_dinov3.lvd1689m`` (300M). + output_dim: width of the projected tokens. + freeze_backbone: freeze the tower, keep it in eval mode and run it + under ``no_grad``. + image_size: ``S`` or ``[H, W]``, multiples of the patch size. DINOv3 is + RoPE-based, so non-square inputs need no position interpolation. + pretrained: ``False`` builds a randomly initialised tower (tests). + tower_kwargs: timm architecture overrides (tests). + """ + + DEFAULT_MODEL = "vit_base_patch16_dinov3.lvd1689m" + _hpt_pretrained_attrs = ("tower",) + + def __init__( + self, + model_name: str = DEFAULT_MODEL, + output_dim: int = 10, + freeze_backbone: bool = True, + image_size=256, + pretrained: bool = True, + tower_kwargs: Optional[dict] = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + import timm + + self.model_name = model_name + self.pretrained = pretrained + self.tower = timm.create_model( + model_name, pretrained=pretrained, num_classes=0, **dict(tower_kwargs or {}) + ) + self.hidden_size = int(self.tower.embed_dim) + self.num_prefix_tokens = int(self.tower.num_prefix_tokens) + size = [image_size] * 2 if isinstance(image_size, int) else list(image_size) + patch = self.tower.patch_embed.patch_size[0] + if any(s % patch for s in size): + raise ValueError( + f"DINOv3Stem image_size {size} is not a multiple of {patch}" + ) + self.image_size = tuple(int(s) for s in size) + self.grid_size = tuple(s // patch for s in self.image_size) + cfg = self.tower.pretrained_cfg + self.register_buffer( + "image_mean", torch.tensor(cfg["mean"]).view(1, 3, 1, 1), persistent=False + ) + self.register_buffer( + "image_std", torch.tensor(cfg["std"]).view(1, 3, 1, 1), persistent=False + ) + + self.out_dim = output_dim + self.proj = nn.Linear(self.hidden_size, output_dim) + + self.freeze_backbone = freeze_backbone + if freeze_backbone: + for param in self.tower.parameters(): + param.requires_grad = False + + def train(self, mode: bool = True): + """Keep a frozen tower in eval mode regardless of the outer flag.""" + super().train(mode) + if self.freeze_backbone: + self.tower.eval() + return self + + def backbone_parameters(self): + """Pretrained-tower parameters (``proj`` stays in the main group).""" + return list(self.tower.parameters()) + + def pretrained_reference_state_dict(self) -> Optional[dict]: + """The cached timm checkpoint, keyed ``tower.*``. ``None`` for a random + tower, or when the file is not in the local HF cache.""" + if not self.pretrained: + return None + from huggingface_hub import hf_hub_download + from safetensors.torch import load_file + + try: + path = hf_hub_download( + self.tower.pretrained_cfg["hf_hub_id"], + "model.safetensors", + local_files_only=True, + ) + except Exception: # not cached, or not an HF-hosted timm model + return None + return {f"tower.{k}": v for k, v in load_file(path).items()} + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, *_, H, W = x.shape + x = x.reshape(-1, 3, H, W) + x = F.interpolate( + x, + size=self.image_size, + mode="bilinear", + align_corners=False, + antialias=True, + ) + x = (x - self.image_mean) / self.image_std + with torch.set_grad_enabled( + not self.freeze_backbone and torch.is_grad_enabled() + ): + feat = self.tower.forward_features(x) + # drop CLS + register tokens; RDT conditions on patch tokens only + feat = feat[:, self.num_prefix_tokens :] + return self.proj(feat.reshape(B, -1, self.hidden_size)) + + class _Qwen3BaseEncoder(PretrainedWeights, PolicyStem): """Shared base for Qwen3-Embedding stems used by HPT. diff --git a/tests/unit/test_dinov3_stem.py b/tests/unit/test_dinov3_stem.py new file mode 100644 index 0000000000..dfc1608478 --- /dev/null +++ b/tests/unit/test_dinov3_stem.py @@ -0,0 +1,101 @@ +"""``DINOv3Stem`` on a random-init tiny tower (no network, no weights).""" + +from __future__ import annotations + +import pytest +import torch + +from egomimic.models.hpt_nets import DINOv3Stem + +D, B = 32, 3 +TINY_TOWER = {"embed_dim": 32, "depth": 1, "num_heads": 2} + + +def _dino(**kwargs): + torch.manual_seed(0) + kwargs.setdefault("freeze_backbone", True) + return DINOv3Stem( + model_name="vit_small_patch16_dinov3", + output_dim=D, + image_size=[32, 48], + pretrained=False, + tower_kwargs=TINY_TOWER, + **kwargs, + ) + + +def test_dinov3_stem_returns_patch_tokens_only(): + stem = _dino() + assert stem.grid_size == (2, 3) + assert stem(torch.rand(B, 1, 1, 3, 90, 160)).shape == (B, 2 * 3, D) + # 2 frames on the time axis: past frame's patches first + assert stem(torch.rand(B, 2, 1, 3, 90, 160)).shape == (B, 2 * 2 * 3, D) + + +def test_dinov3_stem_frozen_tower(): + stem = _dino().train() + assert not stem.tower.training + stem(torch.rand(B, 1, 1, 3, 32, 48)).sum().backward() + assert all(p.grad is None for p in stem.tower.parameters()) + assert stem.proj.weight.grad is not None + assert stem.pretrained_submodules() == [stem.tower] + + +def test_dinov3_stem_finetuned_tower_gets_gradients(): + stem = _dino(freeze_backbone=False).train() + stem(torch.rand(B, 1, 1, 3, 32, 48)).sum().backward() + assert any(p.grad is not None for p in stem.tower.parameters()) + assert {id(p) for p in stem.backbone_parameters()} == { + id(p) for p in stem.tower.parameters() + } + + +def test_dinov3_stem_rejects_off_grid_size(): + with pytest.raises(ValueError, match="multiple"): + DINOv3Stem( + model_name="vit_small_patch16_dinov3", + image_size=250, + pretrained=False, + tower_kwargs=TINY_TOWER, + ) + + +def test_a_random_tower_is_not_verified(): + assert _dino().pretrained_reference_state_dict() is None + + +def test_the_verifier_checks_the_tower_against_its_cached_checkpoint( + tmp_path, monkeypatch +): + """The stem reads its timm checkpoint back from the HF cache, so #645's + hash check covers it like the Qwen stems.""" + import huggingface_hub + from safetensors.torch import save_file + + from egomimic.models.hpt_nets import verify_pretrained_weights + + stem = _dino() + path = tmp_path / "model.safetensors" + save_file({k: v.contiguous() for k, v in stem.tower.state_dict().items()}, path) + monkeypatch.setattr(huggingface_hub, "hf_hub_download", lambda *a, **k: str(path)) + stem.pretrained = True + + assert list(verify_pretrained_weights(stem, verbose=False)) == ["DINOv3Stem"] + with torch.no_grad(): + next(stem.tower.parameters()).add_(1.0) + with pytest.raises(RuntimeError, match="do not"): + verify_pretrained_weights(stem, verbose=False) + + +def test_an_uncached_checkpoint_is_skipped_not_fetched(monkeypatch): + import huggingface_hub + from huggingface_hub.errors import LocalEntryNotFoundError + + def offline(*a, **k): + assert k.get("local_files_only") is True + raise LocalEntryNotFoundError("not cached") + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", offline) + stem = _dino() + stem.pretrained = True + assert stem.pretrained_reference_state_dict() is None From a7ad86d0a04ee1645302246456d0cc5a32849d70 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 22:18:24 +0000 Subject: [PATCH 38/79] algo: RDT (Robotics Diffusion Transformer) on the HPT chassis RDT(HPT) swaps the policy network: no Perceiver pooling; stems hand full token sequences to the head; the DiT blocks live in the trunk slot, shared across embodiments, and cross-attend to language / image tokens in alternating blocks; each FMPolicy head holds one embodiment's action adaptor, position table and output layer. DINOv3 ViT-B/L (fine-tuned at 0.1x) + frozen Qwen3-Embedding, 2-frame image history, ctrl-freq token, per-stream condition dropout. Configs: model/rdt_*; train_zarr_mecka_flagship_6d_rdt. The (past, current) pair goes through the train augs as (B * 2, 3, H, W) with frames=2: one jitter draw per sample covers both frames, and every op runs per frame. Laid side by side instead, ColorJitter's contrast took one mean over both frames (the past frame shifted by up to 0.20 with the current frame's content) and a Resize or CenterCrop would have mixed pixels across the seam. The pair also gets HPT's per-sample absent-camera guard, since both go through _apply_image_augs. compile.enabled=true leaves the trunk off HPT's compile list (no_trunk), but the DiT is still compiled: every head's denoiser calls it, so dynamo traces its blocks into the head's graph (asserted in test_rdt with a capturing backend). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_019DbPWJCKtVg5c2PNHs7egw --- egomimic/algo/hpt.py | 20 +- egomimic/algo/rdt.py | 236 +++++++++ .../model/rdt_1b_bc_mecka_6d_dinov3_vitb.yaml | 11 + .../model/rdt_bc_mecka_6d_dinov3_vitb.yaml | 173 +++++++ .../model/rdt_bc_mecka_6d_dinov3_vitl.yaml | 11 + .../train_zarr_mecka_flagship_6d_rdt.yaml | 32 ++ egomimic/models/image_augs.py | 28 +- egomimic/models/rdt_nets.py | 283 +++++++++++ tests/fixtures/recipes.py | 28 + tests/unit/test_rdt.py | 478 ++++++++++++++++++ tests/unit/test_train_step.py | 31 ++ 11 files changed, 1320 insertions(+), 11 deletions(-) create mode 100644 egomimic/algo/rdt.py create mode 100644 egomimic/hydra_configs/model/rdt_1b_bc_mecka_6d_dinov3_vitb.yaml create mode 100644 egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitb.yaml create mode 100644 egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitl.yaml create mode 100644 egomimic/hydra_configs/train_zarr_mecka_flagship_6d_rdt.yaml create mode 100644 egomimic/models/rdt_nets.py create mode 100644 tests/unit/test_rdt.py diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index de7359cb67..394a2ae452 100644 --- a/egomimic/algo/hpt.py +++ b/egomimic/algo/hpt.py @@ -914,7 +914,7 @@ def __init__( self.is_6dof = kwargs.get("6dof", False) self.kinematics_solver = kwargs.get("kinematics_solver", None) - model = HPTModel(**trunk) + model = self._build_policy(trunk) model.auxiliary_ac_keys = self.auxiliary_ac_keys self.eval_base_seed = int(kwargs.get("eval_base_seed", EVAL_BASE_SEED)) @@ -1025,6 +1025,10 @@ def __init__( self.training_step = 0 + def _build_policy(self, trunk: dict) -> HPTModel: + """The policy network; subclasses (``RDT``) swap in their own.""" + return HPTModel(**trunk) + @override def compile_for_training(self, mode=None, dynamic=False): """Compile the trunk, the flow head's denoiser and the image encoders. @@ -1034,6 +1038,8 @@ def compile_for_training(self, mode=None, dynamic=False): ``Module.compile`` does not route (and which measured no gain anyway). """ policy = self.nets["policy"] + # no_trunk: HPT's forward does not call it. RDT's heads do, so its DiT + # is traced into each head's graph without an entry of its own. targets = [("trunk", policy.trunk["trunk"] if not policy.no_trunk else None)] targets += [ (f"head[{name}].model", getattr(head, "model", None)) @@ -1344,10 +1350,12 @@ def _forward_ot(self, batch, embodiment1_id, embodiment2_id): supervised=self.supervised, ) - def _apply_image_augs(self, images, short): + def _apply_image_augs(self, images, short, frames: int = 1): """ helper method that augments one camera's (B, 3, H, W) batch; train augs - are per-sample, the deterministic eval augs stay batched. + are per-sample, the deterministic eval augs stay batched. ``frames`` > 1 + is (B * frames, 3, H, W), a sample's frames adjacent and sharing one + draw (see ``PerSampleAugs``). An absent camera arrives as an all-zero image and must stay zero (normalizing it would make it -mean/std). The test is per sample, so a @@ -1362,7 +1370,11 @@ def _apply_image_augs(self, images, short): else: return images absent = (images == 0).flatten(1).all(1) - out = augs(images) + out = ( + augs(images, frames=frames) + if isinstance(augs, PerSampleAugs) + else augs(images) + ) return torch.where(absent.view(-1, *[1] * (out.dim() - 1)), 0.0, out) def _stem_history_len(self, domain: str, modality: str) -> int: diff --git a/egomimic/algo/rdt.py b/egomimic/algo/rdt.py new file mode 100644 index 0000000000..7d2bd22505 --- /dev/null +++ b/egomimic/algo/rdt.py @@ -0,0 +1,236 @@ +"""RDT (Robotics Diffusion Transformer) on the HPT chassis. + +``RDT`` is ``HPT`` with a different policy network: the same batch handling, +stems, encoders, denoising heads, eval and config layout, but no Perceiver +pooling. Every stem hands its full token sequence to the head, whose denoiser +(``rdt_nets.RDTDenoiser``) runs the noisy action chunk through the trunk -- the +DiT blocks (``rdt_nets.RDTBackbone``), shared across embodiments, which +cross-attend to the language and image tokens in alternating blocks. +""" + +import torch +import torch.nn as nn + +from egomimic.algo.hpt import HPT, HPTModel +from egomimic.models.rdt_nets import RDTBackbone, RDTConditions, mlp_gelu, sincos_2d +from egomimic.rldb.embodiment.embodiment import IMAGE_HISTORY_SUFFIX +from egomimic.utils.tensor_utils import get_sinusoid_encoding_table + +INIT_STD = 0.02 + + +class RDTModel(HPTModel): + """Stems -> ``RDTConditions`` -> head -> shared ``RDTBackbone``. Streams: a + modality with an encoder is an image, a stem with ``forward_with_mask`` is + language, the rest is state. + + ``cond_drop`` (RDT's ``cond_mask_prob``, per stream, training only): the + state is replaced by a learned null token; the prompt by ``""``, i.e. the + no-prompt deployment condition; a camera by a learned null token, never + every camera of a sample, so ``img`` is inert with a single camera. + + ``image_history`` is the number of frames per camera (axis 1 of the image + input, current frame last); ``image_history_dropout`` replaces the past + frames by the current one, which is also what an episode start looks like. + """ + + def __init__( + self, + embed_dim: int = 1024, + depth: int = 28, + num_heads: int = 16, + mlp_ratio: float = 1.0, + cond_drop: dict | None = None, + image_history: int = 1, + image_history_dropout: float = 0.0, + **kwargs, + ): + super().__init__( + embed_dim=embed_dim, no_trunk=True, token_postprocessing="no-op" + ) + self.trunk = nn.ModuleDict( + {"trunk": RDTBackbone(embed_dim, depth, num_heads, mlp_ratio)} + ) + self.cond_drop = {"state": 0.0, "lang": 0.0, "img": 0.0, **(cond_drop or {})} + self.image_history = int(image_history) + self.image_history_dropout = float(image_history_dropout) + self.stem_modality = {} + + def _create_policy_trunk(self, *args, **kwargs): + return nn.ModuleDict() + + def init_domain_stem(self, domain_name, stem_spec): + # no init_cross_attn: the stems' Perceiver latents would be dead weights + stem_spec = {k: v for k, v in stem_spec.items() if v is not None} + self.stem_spec[domain_name] = stem_spec + self.modalities[domain_name] = list(stem_spec) + for modality, stem in stem_spec.items(): + self.stems[f"{domain_name}_{modality}"] = stem + self.stem_modality[f"{domain_name}_{modality}"] = modality + + def _stream(self, modality, stem) -> str: + if modality in self.encoders: + return "img" + return "lang" if hasattr(stem, "forward_with_mask") else "state" + + def finalize_modules(self): + D = self.embed_dim + streams = { + name: self._stream(self.stem_modality[name], stem) + for name, stem in self.stems.items() + } + cameras = [name for name, s in streams.items() if s == "img"] + # RDT's mlp2x_gelu language adaptor; image / state adaptors are the stems + self.lang_adaptors = nn.ModuleDict( + {name: mlp_gelu(D, D, 2) for name, s in streams.items() if s == "lang"} + ) + super().finalize_modules() + self.camera_embed = nn.ParameterDict( + {name: torch.zeros(1, 1, 1, D) for name in cameras} + ) + self.frame_embed = nn.Parameter(torch.zeros(1, self.image_history, 1, D)) + # Only what can be used: an unused parameter crashes plain DDP. + if self.cond_drop["state"] > 0 and "state" in streams.values(): + self.null_state = nn.Parameter(torch.randn(1, 1, D) * INIT_STD) + if self.cond_drop["img"] > 0 and len(cameras) > 1: + self.null_img = nn.Parameter(torch.randn(1, 1, D) * INIT_STD) + self.trunk["trunk"].initialize_weights() + for head in self.heads.values(): + init = getattr(getattr(head, "model", None), "initialize_weights", None) + if callable(init): + init() + + def _drop(self, prob: float, batch_size: int, device): + """(B,) bool: which samples lose a condition; None outside training.""" + if not (self.training and prob > 0): + return None + return torch.rand(batch_size, device=device) < prob + + def _image_tokens(self, name, modality, stem, x): + """(B, T, N, 3, H, W) frames, current last -> (B, T * N * patches, D).""" + B, T = x.shape[:2] + if T != self.image_history: + raise ValueError( + f"'{modality}' carries {T} frames but trunk.image_history=" + f"{self.image_history}; set it to match the data config's " + "image_history_gap_s (2 frames when set, else 1)" + ) + drop = self._drop(self.image_history_dropout, B, x.device) if T > 1 else None + if drop is not None: + x = torch.where(drop.view(B, 1, 1, 1, 1, 1), x[:, -1:].expand_as(x), x) + encoder = self.encoders[modality] + tokens = stem(encoder(x)) + tokens = tokens.reshape(B, T, -1, tokens.shape[-1]) + # DINOv3 is RoPE-only: its tokens carry no absolute position. + grid = getattr(encoder, "grid_size", None) + if grid is not None and grid[0] * grid[1] == tokens.shape[2]: + pos = sincos_2d(tokens.shape[-1], grid) + else: + pos = get_sinusoid_encoding_table(0, tokens.shape[2], tokens.shape[-1])[0] + tokens = tokens + pos.to(tokens) + self.frame_embed + self.camera_embed[name] + return tokens.reshape(B, -1, tokens.shape[-1]) + + def _drop_cameras(self, img: list) -> list: + """Per-camera dropout that never blinds a sample completely.""" + if len(img) < 2 or not (self.training and self.cond_drop["img"] > 0): + return img + B, device = len(img[0]), img[0].device + drop = torch.rand(B, len(img), device=device) < self.cond_drop["img"] + keep = torch.randint(len(img), (B,), device=device) + blind = drop.all(dim=1) + drop[blind, keep[blind]] = False + return [ + torch.where(drop[:, i, None, None], self.null_img.to(tokens), tokens) + for i, tokens in enumerate(img) + ] + + def forward_features(self, domain, data): + if "fps" not in data: + raise ValueError("RDT needs the per-sample 'fps' (ZarrDataset emits it)") + img, state, lang, lang_mask = [], [], [], [] + for modality in self.modalities.get(domain, []) + self.shared_keys: + if modality not in data: + continue + owner = "shared" if modality in self.shared_keys else domain + name = f"{owner}_{modality}" + stem = self.stems[name] + stream = self._stream(modality, stem) + x = data[modality] + + if stream == "lang": + drop = self._drop(self.cond_drop["lang"], len(x), self.device) + if drop is not None: + x = ["" if d else p for p, d in zip(x, drop.tolist())] + tokens, mask = stem.forward_with_mask(x) + if not mask.any(dim=1).all(): + raise ValueError(f"'{modality}': a prompt encoded to zero tokens") + lang.append(self.lang_adaptors[name](tokens)) + lang_mask.append(mask) + elif stream == "img": + img.append(self._image_tokens(name, modality, stem, x)) + else: + tokens = stem(x) + state.append(tokens.reshape(len(tokens), -1, tokens.shape[-1])) + + if not img: + raise ValueError(f"RDT needs at least one camera for domain '{domain}'") + img = torch.cat(self._drop_cameras(img), dim=1) + cond = RDTConditions( + img=img, + freq=data["fps"].reshape(-1).expand(len(img)), + backbone=self.trunk["trunk"], + ) + if state: + cond.state = torch.cat(state, dim=1) + drop = self._drop(self.cond_drop["state"], len(img), img.device) + if drop is not None: + null = self.null_state.to(cond.state) + cond.state = torch.where(drop[:, None, None], null, cond.state) + if lang: + cond.lang = torch.cat(lang, dim=1) + cond.lang_mask = torch.cat(lang_mask, dim=1) + return cond, None + + +class RDT(HPT): + def __init__(self, *args, **kwargs): + for flag in ("ot", "freeze_repr"): + if kwargs.get(flag): + raise ValueError(f"RDT has no HPT trunk; '{flag}' is HPT-only") + super().__init__(*args, **kwargs) + + def _build_policy(self, trunk: dict) -> RDTModel: + return RDTModel(**trunk) + + def _robomimic_to_hpt_data( + self, + batch, + cam_keys, + proprio_keys, + lang_keys, + ac_key, + aux_ac_keys=[], + domain="", + ): + """HPT's layout plus ``fps`` and, per camera with a ``*_hist`` twin, the + (past, current) frame pair on the image input's time axis.""" + past = { + key[: -len(IMAGE_HISTORY_SUFFIX)]: key + for key in cam_keys + if key.endswith(IMAGE_HISTORY_SUFFIX) and key in batch + } + single = [k for k in cam_keys if k not in past and k not in past.values()] + data = super()._robomimic_to_hpt_data( + batch, single, proprio_keys, lang_keys, ac_key, aux_ac_keys, domain=domain + ) + for key, past_key in past.items(): + short = key.rsplit(".", 1)[-1] + # (B * 2, 3, H, W), each sample's (past, current) adjacent: one + # jitter draw per sample covers both frames, and every op -- a + # contrast mean, a crop, a resize -- still sees one frame. + pair = torch.stack([batch[past_key], batch[key]], dim=1) + n = pair.shape[0] + pair = self._apply_image_augs(pair.flatten(0, 1), short, frames=2) + data[short] = pair.reshape(n, 2, *pair.shape[1:]).unsqueeze(2) + data["fps"] = batch["fps"] + return data diff --git a/egomimic/hydra_configs/model/rdt_1b_bc_mecka_6d_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_1b_bc_mecka_6d_dinov3_vitb.yaml new file mode 100644 index 0000000000..d60059a7b1 --- /dev/null +++ b/egomimic/hydra_configs/model/rdt_1b_bc_mecka_6d_dinov3_vitb.yaml @@ -0,0 +1,11 @@ +# RDT_1B_BC_MECKA_6D_DINOV3_VITB: rdt_bc_mecka_6d_dinov3_vitb at the paper's +# RDT-1B shape (width 2048 x 28 blocks x 32 heads, ~1.2B). From scratch on one +# task this is a scaling-curve point, not the default. +defaults: + - rdt_bc_mecka_6d_dinov3_vitb + - _self_ + +robomimic_model: + width: 2048 + trunk: + num_heads: 32 diff --git a/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitb.yaml b/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitb.yaml new file mode 100644 index 0000000000..f0539898f9 --- /dev/null +++ b/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitb.yaml @@ -0,0 +1,173 @@ +# RDT_BC_MECKA_6D_DINOV3_VITB: Robotics Diffusion Transformer +# (thu-ml/RoboticsDiffusionTransformer) on the HPT chassis, ~300M trainable +# DiT (param-matched to hpt_bc_mecka_6d_300M), same data / action space: 18-D +# wrist-frame 6D-rotation cartesian action, 20-D ee_pose proprio, front camera, +# one sampled annotation. +# +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_6d_rdt \ +# norm_stats.precomputed_norm_path=/norm_stats +# +# egomimic.algo.rdt.RDT is HPT with the policy network swapped: no Perceiver +# pooling. The stems are RDT's condition adaptors and hand their full token +# sequences to the head. `trunk` is the DiT, shared across embodiments: +# self-attention over [timestep; ctrl freq; state; noisy action chunk], +# cross-attention alternating between the language tokens (even blocks) and the +# image tokens (odd blocks). A head holds one embodiment's action adaptor, +# position table and output layer, so co-training shares the blocks at native +# action widths. +# +# Against the paper (decided with ABC-DiT, abc.bot, as the second reference): +# - flow matching (the repo's FMPolicy, 10 Euler steps like pi0.5 / ABC) for +# DDPM sample-prediction + DPM-Solver; +# - DINOv3 ViT-B/16 for SigLIP-so400m, fine-tuned at 0.1x LR (ABC's setting; +# encoder_specs.front_img_1.freeze_backbone=true is the frozen control), at +# the native 16:9 aspect, 224x400 = 350 patch tokens per frame, all of them +# cross-attended. The stem resizes and normalizes; the augs must NOT; +# - Qwen3-Embedding-0.6B per-token states for T5-XXL (frozen, cached); +# - native action width + per-embodiment adaptors for the 128-D unified vector; +# - 2 frames per camera, the past one 0.1 s back (RDT: the previous step); +# needs data...key_map.image_history_gap_s (the train config sets it); +# - DINOv3 is RoPE-only, so patches get a fixed normalized 2D sin-cos code on +# top of RDT's (learned, zero-init) frame / camera tags; +# - condition dropout per stream: state and language 0.1, no image dropout +# with a single camera. +# +# DiT: width 1024 x 28 blocks x 16 heads (head_dim 64, RDT-1B's), RDT's FFN +# ratio 1 -> 10.5M / block. +# Trainable 392.7M = DiT 296.6M + DINOv3 85.6M + adaptors; Qwen 596M frozen. +# H100, batch 32, bf16: 212 ms / step, 30 GiB (ViT-L: 365 ms, 49 GiB; the 1B +# DiT: 359 ms, 57 GiB); sampling 0.3 s per batch. +_target_: egomimic.pl_utils.pl_model.ModelWrapper +robomimic_model: + _target_: egomimic.algo.rdt.RDT + camera_transforms: {} + dims: + human_bimanual: {proprio: 20, action: 18} + # DiT width; every stem / encoder output_dim and the head's hidden_dim. + width: 1024 + ac_keys: + human_bimanual: "actions_cartesian" + + diffusion: true + 6dof: true + + annotation_key: "annotations" + annotation_sampling_mode: "random" + annotation_modality: "annotation" + default_prompt: "" + + use_pad_mask: false + + trunk: + embed_dim: ${model.robomimic_model.width} + depth: 28 + num_heads: 16 + mlp_ratio: 1.0 + # RDT's cond_mask_prob, per stream (train only). A dropped prompt becomes + # "", the no-prompt deployment condition; `img` is per camera and never + # blinds a sample, so it is inert with one camera. + cond_drop: {state: 0.1, lang: 0.1, img: 0.0} + # frames per camera (current last); 2 <-> key_map.image_history_gap_s set + image_history: 2 + # past frame := current frame, per sample (what an episode start looks like) + image_history_dropout: 0.2 + + multitask: false + pretrained: false + pretrained_checkpoint: null + domains: ["human_bimanual"] + shared_obs_keys: ["front_img_1", "annotation"] + + head_specs: + human_bimanual: + _target_: egomimic.models.fm_policy.FMPolicy + action_horizon: 100 + num_inference_steps: 10 + pooling: null + time_dist: "beta" + infer_ac_dims: + human_bimanual: ${model.robomimic_model.dims.human_bimanual.action} + model: + _target_: egomimic.models.rdt_nets.RDTDenoiser + act_dim: ${model.robomimic_model.dims.human_bimanual.action} + act_seq: 100 + hidden_dim: ${model.robomimic_model.width} + # proprio stems x history_len + n_state_tokens: 1 + # FMPolicy's t in [0, 1] -> RDT's [0, 1000) timestep embedding + time_scale: 1000.0 + + shared_stem_specs: + # image adaptor (DINOv3Stem.proj + this MLP ~ RDT's mlp2x_gelu) + front_img_1: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.width} + output_dim: ${model.robomimic_model.width} + widths: ["${model.robomimic_model.width}"] + # RDTModel adds the mlp2x_gelu language adaptor on top of these tokens. + annotation: + _target_: egomimic.models.hpt_nets.QwenPerTokenEncoder + model_name: "Qwen/Qwen3-Embedding-0.6B" + max_length: 128 + freeze: true + dtype: "float16" + cache_text_features: true + text_cache_max_entries: 8192 + output_dim: ${model.robomimic_model.width} + + stem_specs: + human_bimanual: + # state adaptor (RDT's mlp3x_gelu); one token per proprio step + state_ee_pose: + _target_: egomimic.models.hpt_nets.MLPPolicyStem + input_dim: ${model.robomimic_model.dims.human_bimanual.proprio} + output_dim: ${model.robomimic_model.width} + widths: + - ${model.robomimic_model.width} + - ${model.robomimic_model.width} + history_len: 1 + history_dropout: 0.2 + + encoder_specs: + front_img_1: + _target_: egomimic.models.hpt_nets.DINOv3Stem + # timm pulls model_name from the HF hub; prefetch it (the ViT-L config's + # too) on a login node before an offline or multi-rank run: + # python -c "import timm; timm.create_model('', pretrained=True)" + model_name: "vit_base_patch16_dinov3.lvd1689m" + output_dim: ${model.robomimic_model.width} + freeze_backbone: false + # [H, W]; 360x640 frames at their native aspect + image_size: [224, 400] + + train_image_augs: + _target_: torchvision.transforms.Compose + transforms: + - _target_: torchvision.transforms.ColorJitter + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.05 + eval_image_augs: null + +# DINOv3 tower at 1e-5 against 1e-4 everywhere else (ABC-DiT's 0.1). +backbone_lr_scale: 0.1 + +# RDT and ABC-DiT agree: AdamW 1e-4, wd 1e-2, constant LR. ABC's 1k-step +# warmup and (0.9, 0.95) betas; gradient clipping is in the train config. +optimizer: + _target_: torch.optim.AdamW + _partial_: true + _convert_: all + lr: 1e-4 + weight_decay: 0.01 + betas: [0.9, 0.95] + +# LinearLR holds end_factor after total_iters: warmup, then constant. +scheduler: + _target_: torch.optim.lr_scheduler.LinearLR + _partial_: true + start_factor: 1.0e-3 + end_factor: 1.0 + total_iters: 1000 +scheduler_interval: step diff --git a/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitl.yaml b/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitl.yaml new file mode 100644 index 0000000000..47096dd4d1 --- /dev/null +++ b/egomimic/hydra_configs/model/rdt_bc_mecka_6d_dinov3_vitl.yaml @@ -0,0 +1,11 @@ +# RDT_BC_MECKA_6D_DINOV3_VITL: rdt_bc_mecka_6d_dinov3_vitb with the DINOv3 +# ViT-L/16 tower (300M, 1024-wide tokens) -- the two differ only in the image +# encoder. +defaults: + - rdt_bc_mecka_6d_dinov3_vitb + - _self_ + +robomimic_model: + encoder_specs: + front_img_1: + model_name: "vit_large_patch16_dinov3.lvd1689m" diff --git a/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_rdt.yaml b/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_rdt.yaml new file mode 100644 index 0000000000..6ca6da36e4 --- /dev/null +++ b/egomimic/hydra_configs/train_zarr_mecka_flagship_6d_rdt.yaml @@ -0,0 +1,32 @@ +# Composed recipe: train_zarr_mecka_flagship_6d_hpt with the RDT policy +# (model/rdt_bc_mecka_6d_dinov3_vitb.yaml; model=rdt_bc_mecka_6d_dinov3_vitl for +# the ViT-L tower, model=rdt_1b_bc_mecka_6d_dinov3_vitb for the 1B DiT). Same +# data, action space, evaluators and norm stats as the HPT recipe. +# +# python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_6d_rdt \ +# norm_stats.precomputed_norm_path=/norm_stats +defaults: + - train_zarr_mecka_flagship_6d_hpt + - override model: rdt_bc_mecka_6d_dinov3_vitb + - _self_ + +data: + train_datasets: + human_bimanual: + resolver: + key_map: + # RDT's 2-frame image history: also load the front frame this many + # seconds back (3 frames at 30 fps). The val heads interpolate this + # resolver, so they follow. Pairs with model...trunk.image_history=2. + image_history_gap_s: 0.1 + # PROPRIO history stays off for this arm: the HPT recipe ships K = 3 + # / stride 5, but RDT's state adaptor is declared at history_len 1 + # and the two must agree or process_batch_for_training raises. Raise + # both together if the RDT arm should carry history too. + proprio_history: 1 + history_stride: 1 + +trainer: + # RDT clips at 1.0 (ABC at 10): protects the fine-tuned tower from the + # from-scratch DiT's early gradients. + gradient_clip_val: 1.0 diff --git a/egomimic/models/image_augs.py b/egomimic/models/image_augs.py index a2d6a23302..e95e12b39e 100644 --- a/egomimic/models/image_augs.py +++ b/egomimic/models/image_augs.py @@ -41,7 +41,9 @@ class BatchedColorJitter(nn.Module): did; each adjustment clamps to [0, 1], so the order is not quite inert, but it is a permutation of the same four factors either way. - Input is (B, 3, H, W) float in [0, 1]. + Input is (B, 3, H, W) float in [0, 1]. With ``frames`` > 1 it is + (B * frames, 3, H, W) with a sample's frames adjacent: they share one draw, + and every adjustment (contrast's mean included) still runs per frame. """ def __init__( @@ -85,8 +87,8 @@ def _apply_hue(self, img, f): h = (h + f.squeeze(-3)) % 1.0 return _ft._hsv2rgb(torch.stack((h, s, v), dim=-3)).clamp(0.0, 1.0) - def forward(self, img: torch.Tensor) -> torch.Tensor: - n = img.shape[0] + def forward(self, img: torch.Tensor, frames: int = 1) -> torch.Tensor: + n = img.shape[0] // frames ops = [ (self.brightness, self._apply_brightness), (self.contrast, self._apply_contrast), @@ -97,7 +99,7 @@ def forward(self, img: torch.Tensor) -> torch.Tensor: bounds, fn = ops[idx] if bounds is None: continue - img = fn(img, self._factors(bounds, n, img)) + img = fn(img, self._factors(bounds, n, img).repeat_interleave(frames, 0)) return img def extra_repr(self) -> str: @@ -133,6 +135,11 @@ class PerSampleAugs(nn.Module): ``augs`` is the configured callable (typically a ``torchvision.transforms.Compose``); a deterministic one is unaffected by the wrapping, so this is safe for any existing aug list. + + ``frames`` > 1 takes (B * frames, C, H, W), each sample's frames adjacent: + the frames of a sample share one parameter draw, and every transform still + sees one frame at a time, so a crop, resize or contrast mean never spans + two frames. """ def __init__(self, augs): @@ -141,11 +148,18 @@ def __init__(self, augs): vectorized = _vectorize(augs) self.vectorized = nn.ModuleList(vectorized) if vectorized is not None else None - def forward(self, images): + def forward(self, images, frames: int = 1): if self.vectorized is None: - return torch.stack([self.augs(image) for image in images]) + if frames == 1: + return torch.stack([self.augs(image) for image in images]) + # torchvision draws once per call and works per image of a stack + return torch.cat([self.augs(clip) for clip in images.split(frames)]) for t in self.vectorized: - images = t(images) + images = ( + t(images, frames=frames) + if isinstance(t, BatchedColorJitter) + else t(images) + ) return images def extra_repr(self): diff --git a/egomimic/models/rdt_nets.py b/egomimic/models/rdt_nets.py new file mode 100644 index 0000000000..51d92d1177 --- /dev/null +++ b/egomimic/models/rdt_nets.py @@ -0,0 +1,283 @@ +"""Robotics Diffusion Transformer (RDT, https://github.com/thu-ml/RoboticsDiffusionTransformer) +building blocks. ``RDTBackbone`` is the DiT shared by every embodiment (the +policy's trunk); ``RDTDenoiser`` is one embodiment's adaptor / output layer +around it, used as the denoiser of a ``DenoisingPolicy`` head.""" + +import math +from dataclasses import dataclass +from typing import Optional + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def sincos_1d(embed_dim: int, positions) -> np.ndarray: + """(M,) positions -> (M, embed_dim) sin-cos table.""" + assert embed_dim % 2 == 0 + omega = 1.0 / 10000 ** ( + np.arange(embed_dim // 2, dtype=np.float64) / (embed_dim / 2) + ) + out = np.einsum("m,d->md", np.asarray(positions, dtype=np.float64), omega) + return np.concatenate([np.sin(out), np.cos(out)], axis=1) + + +def multimodal_sincos(embed_dim: int, lens: list) -> torch.Tensor: + """RDT's input position table: the first half of each row encodes which + segment the token belongs to, the second half its index inside it.""" + half = embed_dim // 2 + segment = sincos_1d(half, np.arange(len(lens))) + rows = [] + for idx, length in enumerate(lens): + row = np.zeros((length, embed_dim)) + row[:, :half] = segment[idx] + row[:, half:] = sincos_1d(embed_dim - half, np.arange(length)) + rows.append(row) + return torch.from_numpy(np.concatenate(rows, axis=0)).float() + + +def sincos_2d(embed_dim: int, grid_hw, ref_grid: float = 16.0) -> torch.Tensor: + """(h * w, embed_dim) raster-order table, half the channels per axis. + + Coordinates are normalized to [0, ref_grid] on both axes, so an edge or the + centre gets the same code at any resolution or aspect ratio. + """ + h, w = grid_hw + half = embed_dim // 2 + rows = np.linspace(0.0, ref_grid, h) if h > 1 else np.zeros(1) + cols = np.linspace(0.0, ref_grid, w) if w > 1 else np.zeros(1) + table = np.concatenate( + [ + np.repeat(sincos_1d(half, rows), w, axis=0), + np.tile(sincos_1d(embed_dim - half, cols), (h, 1)), + ], + axis=1, + ) + return torch.from_numpy(table).float() + + +def mlp_gelu(in_dim: int, out_dim: int, depth: int) -> nn.Sequential: + """RDT's ``mlp{depth}x_gelu`` condition adaptor.""" + layers = [nn.Linear(in_dim, out_dim)] + for _ in range(1, depth): + layers += [nn.GELU(approximate="tanh"), nn.Linear(out_dim, out_dim)] + return nn.Sequential(*layers) + + +@dataclass +class RDTConditions: + """What the denoiser is conditioned on, already at the DiT width. + + Quacks like the single conditioning tensor ``DenoisingPolicy`` expects + (``len`` / ``dtype`` / ``device``), so the FM / diffusion heads take it + unchanged. + """ + + img: torch.Tensor # (B, N_img, D) + freq: torch.Tensor # (B,) control frequency, Hz + backbone: "RDTBackbone" # the shared trunk the denoiser runs through + state: Optional[torch.Tensor] = None # (B, K, D) + lang: Optional[torch.Tensor] = None # (B, L, D) + lang_mask: Optional[torch.Tensor] = None # (B, L) bool, True = attend + + def __len__(self): + return self.img.shape[0] + + @property + def dtype(self): + return self.img.dtype + + @property + def device(self): + return self.img.device + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size), + ) + self.frequency_embedding_size = frequency_embedding_size + + def forward(self, t: torch.Tensor) -> torch.Tensor: + half = self.frequency_embedding_size // 2 + freqs = torch.exp( + -math.log(10000) + * torch.arange(half, dtype=torch.float32, device=t.device) + / half + ) + args = t[:, None].float() * freqs[None] + emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + return self.mlp(emb.to(self.mlp[0].weight.dtype)) + + +class _Attention(nn.Module): + """Self- (``context=None``) or cross-attention with QK RMSNorm.""" + + def __init__(self, dim: int, num_heads: int): + super().__init__() + assert dim % num_heads == 0, "dim should be divisible by num_heads" + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.q = nn.Linear(dim, dim) + self.kv = nn.Linear(dim, dim * 2) + self.q_norm = nn.RMSNorm(self.head_dim, eps=1e-6) + self.k_norm = nn.RMSNorm(self.head_dim, eps=1e-6) + self.proj = nn.Linear(dim, dim) + + def forward(self, x, context=None, mask=None): + B, N, C = x.shape + context = x if context is None else context + L = context.shape[1] + q = self.q(x).view(B, N, self.num_heads, self.head_dim).transpose(1, 2) + k, v = ( + self.kv(context) + .view(B, L, 2, self.num_heads, self.head_dim) + .permute(2, 0, 3, 1, 4) + ) + q, k = self.q_norm(q), self.k_norm(k) + if mask is not None: + mask = mask.view(B, 1, 1, L) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + return self.proj(out.transpose(1, 2).reshape(B, N, C)) + + +class _Mlp(nn.Module): + def __init__(self, dim: int, hidden: int, out_dim: Optional[int] = None): + super().__init__() + self.fc1 = nn.Linear(dim, hidden) + self.act = nn.GELU(approximate="tanh") + self.fc2 = nn.Linear(hidden, out_dim or dim) + + def forward(self, x): + return self.fc2(self.act(self.fc1(x))) + + +class RDTBlock(nn.Module): + """Pre-RMSNorm self-attention, cross-attention to one condition, FFN.""" + + def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 1.0): + super().__init__() + self.norm1 = nn.RMSNorm(hidden_size, eps=1e-6) + self.attn = _Attention(hidden_size, num_heads) + self.norm2 = nn.RMSNorm(hidden_size, eps=1e-6) + self.cross_attn = _Attention(hidden_size, num_heads) + self.norm3 = nn.RMSNorm(hidden_size, eps=1e-6) + self.ffn = _Mlp(hidden_size, int(hidden_size * mlp_ratio)) + + def forward(self, x, c, mask=None): + x = x + self.attn(self.norm1(x)) + x = x + self.cross_attn(self.norm2(x), c, mask) + return x + self.ffn(self.norm3(x)) + + +def _xavier(module): + if isinstance(module, nn.Linear): + nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + +class RDTBackbone(nn.Module): + """The embodiment-agnostic part of RDT: timestep / control-frequency + embedders and the blocks, which alternate the condition they cross-attend + to (language, image, language, ...; image only when there is no language). + """ + + def __init__( + self, + hidden_dim: int = 1024, + depth: int = 28, + n_heads: int = 16, + mlp_ratio: float = 1.0, + ): + super().__init__() + self.hidden_dim = hidden_dim + self.t_embedder = TimestepEmbedder(hidden_dim) + self.freq_embedder = TimestepEmbedder(hidden_dim) + self.blocks = nn.ModuleList( + [RDTBlock(hidden_dim, n_heads, mlp_ratio) for _ in range(depth)] + ) + self.initialize_weights() + + def initialize_weights(self): + self.apply(_xavier) + for embedder in (self.t_embedder, self.freq_embedder): + nn.init.normal_(embedder.mlp[0].weight, std=0.02) + nn.init.normal_(embedder.mlp[2].weight, std=0.02) + + def prefix(self, t: torch.Tensor, freq: torch.Tensor, batch_size: int): + """(B, 2, D): the timestep and control-frequency tokens.""" + t = self.t_embedder(t.reshape(-1)).unsqueeze(1).expand(batch_size, -1, -1) + return torch.cat([t, self.freq_embedder(freq.reshape(-1)).unsqueeze(1)], dim=1) + + def forward(self, h: torch.Tensor, cond: "RDTConditions") -> torch.Tensor: + conds = [(cond.img, None)] + if cond.lang is not None: + conds.insert(0, (cond.lang, cond.lang_mask)) + for i, block in enumerate(self.blocks): + c, mask = conds[i % len(conds)] + h = block(h, c, mask) + return h + + +class RDTDenoiser(nn.Module): + """One embodiment's view of RDT: the action adaptor, the position table of + ``[timestep; ctrl freq; state tokens; noisy action chunk]`` and the output + layer. The blocks in between are ``cond.backbone``, shared across + embodiments. ``forward`` has the ``model(x_t, t, cond)`` signature of the + other denoising nets, with ``cond`` an ``RDTConditions``. + + ``time_scale`` maps the head's time onto RDT's [0, 1000) embedding range: + 1000 for flow matching (t in [0, 1]), 1 for a DDPM head (integer steps). + """ + + def __init__( + self, + act_dim: int, + act_seq: int, + hidden_dim: int = 1024, + n_state_tokens: int = 1, + time_scale: float = 1000.0, + ): + super().__init__() + self.act_seq = act_seq + self.n_state_tokens = n_state_tokens + self.time_scale = time_scale + self.action_adaptor = mlp_gelu(act_dim, hidden_dim, depth=3) + self.x_pos_embed = nn.Parameter( + torch.zeros(1, 2 + n_state_tokens + act_seq, hidden_dim) + ) + self.norm_final = nn.RMSNorm(hidden_dim, eps=1e-6) + self.ffn_final = _Mlp(hidden_dim, hidden_dim, act_dim) + self.initialize_weights() + + def initialize_weights(self): + """RDT's init. ``RDTModel.finalize_modules`` calls it again, after HPT's + xavier pass over every Linear.""" + self.apply(_xavier) + lens = [1, 1, self.n_state_tokens, self.act_seq] + self.x_pos_embed.data.copy_( + multimodal_sincos(self.x_pos_embed.shape[-1], [n for n in lens if n > 0]) + ) + nn.init.constant_(self.ffn_final.fc2.weight, 0) + nn.init.constant_(self.ffn_final.fc2.bias, 0) + + def forward(self, x, timesteps, cond: RDTConditions, *args, **kwargs): + n_state = 0 if cond.state is None else cond.state.shape[1] + if n_state != self.n_state_tokens: + raise ValueError( + f"RDTDenoiser got {n_state} state tokens but n_state_tokens=" + f"{self.n_state_tokens}; set it to (proprio stems x history_len)" + ) + tokens = [cond.backbone.prefix(timesteps * self.time_scale, cond.freq, len(x))] + if n_state: + tokens.append(cond.state) + tokens.append(self.action_adaptor(x)) + h = torch.cat(tokens, dim=1) + self.x_pos_embed + h = cond.backbone(h, cond) + return self.ffn_final(self.norm_final(h))[:, -self.act_seq :] diff --git a/tests/fixtures/recipes.py b/tests/fixtures/recipes.py index 78b90de88a..91de4221e0 100644 --- a/tests/fixtures/recipes.py +++ b/tests/fixtures/recipes.py @@ -41,6 +41,11 @@ def _pi(data: str, model: str, emb: str) -> Recipe: ("scale", "hpt"): Recipe( "train_zarr_cartesian", "scale", "hpt_bc_flow_scale", "human_bimanual" ), + # RDT's shipped config is the 6D cartesian one; rdt_small_overrides points + # it at the keypoint action this data config emits. + ("mecka", "rdt"): Recipe( + "train_zarr_cartesian", "mecka", "rdt_bc_mecka_6d_dinov3_vitb", "human_bimanual" + ), # Pi trains on the same per-vendor data configs as HPT: datasets emit one # camera naming and the Pi wrapper renames onto openpi's slots. The plain # vendor configs load no language field (that is a data choice the *_lang @@ -121,6 +126,29 @@ def hpt_small_overrides(emb: str) -> list[str]: ] +def rdt_small_overrides(emb: str) -> list[str]: + """Tiny DiT, random-init tiny DINOv3 tower, no Qwen stem (network-free); + the 2-frame image history the shipped train config turns on.""" + rm = "model.robomimic_model" + enc = f"{rm}.encoder_specs.front_img_1" + return [ + f"{rm}.dims.{emb}.action=144", + f"{rm}.ac_keys.{emb}=actions_keypoints", + f"{rm}.width=32", + f"{rm}.trunk.depth=2", + f"{rm}.trunk.num_heads=4", + f"{rm}.head_specs.{emb}.num_inference_steps=2", + f"~{rm}.shared_stem_specs.annotation", + f"{rm}.shared_obs_keys=[front_img_1]", + f"{rm}.annotation_key=null", + f"{enc}.model_name=vit_small_patch16_dinov3", + f"{enc}.image_size=[32,48]", + f"+{enc}.pretrained=false", + f"+{enc}.tower_kwargs={{embed_dim: 32, depth: 1, num_heads: 2}}", + f"+data.train_datasets.{emb}.resolver.key_map.image_history_gap_s=0.1", + ] + + def pi_cpu_overrides() -> list[str]: return [ "model.robomimic_model.config.pytorch_weight_path=null", diff --git a/tests/unit/test_rdt.py b/tests/unit/test_rdt.py new file mode 100644 index 0000000000..ee546e60e2 --- /dev/null +++ b/tests/unit/test_rdt.py @@ -0,0 +1,478 @@ +"""RDT on the HPT chassis: the shared DiT backbone and per-embodiment denoiser, +``RDTModel``'s condition streams (random-init tiny DINOv3 tower, no network) +and the algo's image-history batch layout.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn +from torchvision import transforms as T + +from egomimic.algo.rdt import RDT, RDTModel +from egomimic.models.fm_policy import FMPolicy +from egomimic.models.hpt_nets import DINOv3Stem, MLPPolicyStem, PolicyStem +from egomimic.models.rdt_nets import ( + RDTBackbone, + RDTConditions, + RDTDenoiser, + sincos_2d, +) + +D, ACT, SEQ, B = 32, 6, 5, 3 +DOMAIN = "human_bimanual" +TINY_TOWER = {"embed_dim": 32, "depth": 1, "num_heads": 2} + + +def _backbone(depth=4): + torch.manual_seed(0) + return RDTBackbone(hidden_dim=D, depth=depth, n_heads=4).eval() + + +def _denoiser(**kwargs): + torch.manual_seed(0) + net = RDTDenoiser(act_dim=ACT, act_seq=SEQ, hidden_dim=D, **kwargs) + # the zero-init output layer would hide every upstream effect + nn.init.normal_(net.ffn_final.fc2.weight, std=0.1) + return net.eval() + + +def _cond(backbone=None, lang_len=4, n_state=1): + torch.manual_seed(1) + return RDTConditions( + img=torch.randn(B, 7, D), + freq=torch.full((B,), 30.0), + backbone=backbone or _backbone(), + state=torch.randn(B, n_state, D) if n_state else None, + lang=torch.randn(B, lang_len, D), + lang_mask=torch.ones(B, lang_len, dtype=torch.bool), + ) + + +def test_denoiser_shape_and_zero_init(): + net = RDTDenoiser(act_dim=ACT, act_seq=SEQ, hidden_dim=D) + out = net(torch.randn(B, SEQ, ACT), torch.rand(B), _cond()) + assert out.shape == (B, SEQ, ACT) + assert torch.count_nonzero(out) == 0 + + +def test_denoiser_ignores_masked_language_tokens(): + net, x, t = _denoiser(), torch.randn(B, SEQ, ACT), torch.rand(B) + cond = _cond() + cond.lang_mask[:, -2:] = False + ref = net(x, t, cond) + cond.lang[:, -2:] += 10.0 + assert torch.allclose(net(x, t, cond), ref, atol=1e-6) + cond.lang[:, 0] += 10.0 + assert not torch.allclose(net(x, t, cond), ref, atol=1e-4) + + +def test_denoiser_uses_every_condition_time_and_frequency(): + net, x, t = _denoiser(), torch.randn(B, SEQ, ACT), torch.rand(B) + backbone = _backbone() + ref = net(x, t, _cond(backbone)) + for field in ("img", "state", "lang", "freq"): + cond = _cond(backbone) + setattr(cond, field, getattr(cond, field) + 1.0) + assert not torch.allclose(net(x, t, cond), ref, atol=1e-4), field + assert not torch.allclose(net(x, t * 0.5, _cond(backbone)), ref, atol=1e-4) + + +def test_denoisers_of_different_widths_share_one_backbone(): + """Co-training: per-embodiment adaptors, one set of DiT blocks.""" + backbone = _backbone() + wide = RDTDenoiser(act_dim=ACT + 3, act_seq=SEQ, hidden_dim=D) + cond = _cond(backbone) + assert _denoiser()(torch.randn(B, SEQ, ACT), torch.rand(B), cond).shape[-1] == ACT + assert wide(torch.randn(B, SEQ, ACT + 3), torch.rand(B), cond).shape[-1] == ACT + 3 + assert not any("blocks" in n for n, _ in wide.named_parameters()) + + +def test_denoiser_without_language_or_state(): + net = _denoiser(n_state_tokens=0) + cond = RDTConditions( + img=torch.randn(B, 7, D), freq=torch.full((B,), 30.0), backbone=_backbone() + ) + assert net(torch.randn(B, SEQ, ACT), torch.rand(B), cond).shape == (B, SEQ, ACT) + + +def test_denoiser_rejects_wrong_state_token_count(): + with pytest.raises(ValueError, match="n_state_tokens"): + _denoiser()(torch.randn(B, SEQ, ACT), torch.rand(B), _cond(n_state=2)) + + +def test_fm_head_samples_through_the_conditions_container(): + head = FMPolicy( + model=_denoiser(), + action_horizon=SEQ, + infer_ac_dims={DOMAIN: ACT}, + num_inference_steps=2, + ) + assert head((_cond(), DOMAIN)).shape == (B, SEQ, ACT) + + +def test_sincos_2d_is_resolution_agnostic_at_the_corners(): + small, large = sincos_2d(D, (2, 3)), sincos_2d(D, (14, 25)) + assert small.shape == (6, D) and large.shape == (350, D) + for a, b in ((0, 0), (2, 24), (3, 325), (5, 349)): # the four corners + assert torch.allclose(small[a], large[b], atol=1e-6) + # vertical neighbours share their column half + assert torch.equal(large[0, D // 2 :], large[25, D // 2 :]) + + +def _dino(**kwargs): + torch.manual_seed(0) + kwargs.setdefault("freeze_backbone", True) + return DINOv3Stem( + model_name="vit_small_patch16_dinov3", + output_dim=D, + image_size=[32, 48], + pretrained=False, + tower_kwargs=TINY_TOWER, + **kwargs, + ) + + +# -------------------------------------------------------------------------- +# RDTModel +# -------------------------------------------------------------------------- + + +class _StubText(PolicyStem): + """A per-token text stem: prompt i has i + 1 real tokens, left-padded.""" + + def __init__(self): + super().__init__() + self.embed = nn.Embedding(8, D) + + def forward_with_mask(self, prompts): + n = len(prompts) + ids = torch.arange(n)[None, :].expand(n, -1) + mask = ids >= (n - 1 - torch.arange(n))[:, None] + return self.embed(ids), mask + + +def _model(cond_drop=None, text=True, history_len=1, frames=2, frame_dropout=0.0): + torch.manual_seed(0) + model = RDTModel( + embed_dim=D, + depth=2, + num_heads=4, + cond_drop=cond_drop, + image_history=frames, + image_history_dropout=frame_dropout, + ) + shared = {"front_img_1": MLPPolicyStem(input_dim=D, output_dim=D, widths=[D])} + if text: + shared["annotation"] = _StubText() + model.init_domain_stem("shared", shared) + model.shared_keys = list(shared) + model.init_domain_stem( + DOMAIN, + { + "state_ee_pose": MLPPolicyStem( + input_dim=4, output_dim=D, widths=[D], history_len=history_len + ) + }, + ) + model.init_domain_head( + DOMAIN, + FMPolicy( + model=RDTDenoiser( + act_dim=ACT, act_seq=SEQ, hidden_dim=D, n_state_tokens=history_len + ), + action_horizon=SEQ, + infer_ac_dims={DOMAIN: ACT}, + num_inference_steps=2, + ), + ) + model.init_encoder("front_img_1", _dino()) + model.auxiliary_ac_keys = {} + model.diffusion = True + model.device = torch.device("cpu") + model.finalize_modules() + return model + + +def _data(history_len=1, text=True, frames=2): + torch.manual_seed(2) + data = { + "front_img_1": torch.rand(B, frames, 1, 3, 32, 48), + "state_ee_pose": torch.randn(B, history_len, 4), + "action": torch.randn(B, SEQ, ACT), + "pad_mask": torch.ones(B, SEQ, 1), + "fps": torch.full((B,), 30.0), + } + if text: + data["annotation"] = ["a", "b", "c"] + return data + + +ALL_DROPS = {"state": 1.0, "lang": 1.0, "img": 1.0} + + +def test_model_streams_and_shapes(): + model = _model(history_len=2).eval() + cond, _ = model.forward_features(DOMAIN, _data(history_len=2)) + assert cond.img.shape == (B, 2 * 6, D) # 2 frames x 6 patches + assert cond.state.shape == (B, 2, D) + assert cond.lang.shape == (B, B, D) + assert cond.lang_mask.sum(1).tolist() == [1, 2, 3] + assert cond.freq.tolist() == [30.0] * B + assert cond.backbone is model.trunk["trunk"] + + +def test_model_shares_the_dit_in_the_trunk_not_the_head(): + names = [n for n, _ in _model().named_parameters()] + assert any(n.startswith("trunk.trunk.blocks.") for n in names) + assert not any(n.startswith("heads.") and ".blocks." in n for n in names) + assert not any(".tokens" in n or "cross_attention" in n for n in names) + + +def test_model_keeps_rdt_init_after_hpt_xavier_pass(): + net = _model().heads[DOMAIN].model + assert torch.count_nonzero(net.ffn_final.fc2.weight) == 0 + assert net.x_pos_embed.abs().sum() > 0 + + +def test_model_loss_reaches_every_trainable_parameter(): + """Plain DDP crashes on a parameter that gets no gradient.""" + model = _model(cond_drop={"state": 0.5, "lang": 0.5}, frame_dropout=0.5).train() + nn.init.normal_(model.heads[DOMAIN].model.ffn_final.fc2.weight, std=0.1) + model.compute_loss({"domain": DOMAIN, "data": _data()}).backward() + missing = [ + n for n, p in model.named_parameters() if p.requires_grad and p.grad is None + ] + assert not missing, missing + + +def test_model_without_text_stem_and_single_frame(): + model = _model(cond_drop=ALL_DROPS, text=False, frames=1).train() + data = _data(text=False, frames=1) + cond, _ = model.forward_features(DOMAIN, data) + assert cond.lang is None and cond.img.shape == (B, 6, D) + model.compute_loss({"domain": DOMAIN, "data": data}).backward() + out = model.eval().forward(DOMAIN, _data(text=False, frames=1)) + assert out[DOMAIN].shape == (B, SEQ, ACT) + + +def test_model_rejects_a_frame_count_the_config_did_not_declare(): + with pytest.raises(ValueError, match="image_history"): + _model(frames=2).forward_features(DOMAIN, _data(frames=1)) + + +def test_condition_dropout_is_train_only_and_never_blinds_one_camera(): + model = _model(cond_drop=ALL_DROPS) + assert not hasattr(model, "null_img") # one camera: img dropout is inert + seen = [] + stem = model.stems["shared_annotation"] + real = stem.forward_with_mask + stem.forward_with_mask = lambda prompts: ( + seen.append(list(prompts)), + real(prompts), + )[1] + + ref, _ = model.eval().forward_features(DOMAIN, _data()) + assert seen[-1] == ["a", "b", "c"] + cond, _ = model.train().forward_features(DOMAIN, _data()) + assert seen[-1] == ["", "", ""] # a dropped prompt is the empty prompt + assert torch.allclose(cond.state, model.null_state.expand_as(cond.state)) + assert torch.allclose(cond.img, ref.img) # tower frozen + in eval: same tokens + + +def test_image_history_dropout_replaces_the_past_frame_by_the_current(): + model = _model(frame_dropout=1.0).train() + with torch.no_grad(): + model.frame_embed.zero_() + cond, _ = model.forward_features(DOMAIN, _data()) + past, current = cond.img[:, :6], cond.img[:, 6:] + assert torch.allclose(past, current, atol=1e-6) + cond, _ = model.eval().forward_features(DOMAIN, _data()) + assert not torch.allclose(cond.img[:, :6], cond.img[:, 6:], atol=1e-4) + + +def test_frame_and_camera_tags_start_at_zero(): + model = _model() + assert torch.count_nonzero(model.frame_embed) == 0 + assert all(torch.count_nonzero(p) == 0 for p in model.camera_embed.values()) + + +def test_model_needs_fps(): + data = _data() + del data["fps"] + with pytest.raises(ValueError, match="fps"): + _model().forward_features(DOMAIN, data) + + +# -------------------------------------------------------------------------- +# RDT algo +# -------------------------------------------------------------------------- + + +def test_rdt_rejects_hpt_trunk_features(): + with pytest.raises(ValueError, match="HPT-only"): + RDT(None, {}, None, None, ot=True) + + +def _pairing_algo(augs): + from egomimic.models.image_augs import PerSampleAugs + + algo = RDT.__new__(RDT) + algo.nets = nn.ModuleDict().train() + algo.encoders = {"front_img_1": None} + algo.train_image_augs = PerSampleAugs(augs) + algo.eval_image_augs = None + algo.annotation_modality, algo.is_6dof, algo.shared_ac_key = ( + "annotation", + True, + None, + ) + return algo + + +class _Shift(nn.Module): # a "random" aug with no batched form: one draw per call + def forward(self, image): + return image + torch.rand(()) + + +@pytest.mark.parametrize("vectorized", [True, False]) +def test_algo_pairs_history_frames_under_one_augmentation(vectorized): + """(past, current) reach the model as one (B, 2, 1, 3, H, W) input, jittered + with the same per-sample draw, on the batched and the per-image aug path.""" + from torchvision.transforms import ColorJitter + + algo = _pairing_algo(ColorJitter(0.4, 0.4, 0.4, 0.1) if vectorized else _Shift()) + assert (algo.train_image_augs.vectorized is not None) == vectorized + cam = "observations.images.front_img_1" + frame = torch.rand(1, 3, 8, 8).expand(B, -1, -1, -1).contiguous() + batch = { + cam: frame, + f"{cam}_hist": frame.clone(), + "actions": torch.zeros(B, SEQ, ACT), + "pad_mask": torch.ones(B, SEQ, 1), + "embodiment": torch.tensor([3]), + "fps": torch.full((B,), 30.0), + } + data = algo._robomimic_to_hpt_data(batch, [cam, f"{cam}_hist"], [], [], "actions") + frames = data["front_img_1"] + assert frames.shape == (B, 2, 1, 3, 8, 8) + assert "front_img_1_hist" not in data + # identical inputs: the pair stays identical, the samples do not + assert torch.allclose(frames[:, 0], frames[:, 1], atol=1e-6) + assert not torch.allclose(frames[0], frames[1], atol=1e-4) + assert data["fps"].tolist() == [30.0] * B + + +def test_algo_keeps_past_and_current_in_order(): + algo = _pairing_algo(nn.Identity()) + cam = "observations.images.front_img_1" + batch = { + cam: torch.ones(B, 3, 8, 8), + f"{cam}_hist": torch.zeros(B, 3, 8, 8), + "actions": torch.zeros(B, SEQ, ACT), + "pad_mask": torch.ones(B, SEQ, 1), + "embodiment": torch.tensor([3]), + "fps": torch.full((B,), 30.0), + } + frames = algo._robomimic_to_hpt_data( + batch, [cam, f"{cam}_hist"], [], [], "actions" + )["front_img_1"] + assert frames[:, 0].sum() == 0 and frames[:, 1].min() == 1 # current frame last + + +def _pair_batch(current, past): + cam = "observations.images.front_img_1" + n = current.shape[0] + return cam, { + cam: current, + f"{cam}_hist": past, + "actions": torch.zeros(n, SEQ, ACT), + "pad_mask": torch.ones(n, SEQ, 1), + "embodiment": torch.tensor([3]), + "fps": torch.full((n,), 30.0), + } + + +@pytest.mark.parametrize("vectorized", [True, False]) +def test_a_frames_contrast_mean_is_its_own(vectorized): + """Side by side, contrast took one mean over both frames, so the past + frame's pixels depended on the current frame's.""" + import torchvision.transforms.functional as F + from torchvision.transforms import ColorJitter + + jitter = ColorJitter(contrast=(0.5, 0.5)) + algo = _pairing_algo(jitter if vectorized else T.Compose([jitter, _Shift()])) + past = torch.rand(B, 3, 8, 8) + outs = [] + for current in (torch.zeros(B, 3, 8, 8) + 0.9, torch.rand(B, 3, 8, 8)): + cam, batch = _pair_batch(current, past.clone()) + outs.append( + algo._robomimic_to_hpt_data(batch, [cam, f"{cam}_hist"], [], [], "actions")[ + "front_img_1" + ][:, 0, 0] + ) + reference = F.adjust_contrast(past, 0.5) + for out in outs: + # _Shift adds one random offset per sample on the per-image path + offset = (out - reference).flatten(1) + assert torch.allclose(offset, offset[:, :1].expand_as(offset), atol=1e-5) + if vectorized: + assert torch.allclose(out, reference, atol=1e-5) + + +def test_a_resize_runs_per_frame(): + algo = _pairing_algo(T.Compose([T.ColorJitter(0.4), T.Resize(4)])) + frame = torch.rand(1, 3, 8, 8).expand(B, -1, -1, -1).contiguous() + cam, batch = _pair_batch(frame, frame.clone()) + frames = algo._robomimic_to_hpt_data( + batch, [cam, f"{cam}_hist"], [], [], "actions" + )["front_img_1"] + assert frames.shape == (B, 2, 1, 3, 4, 4) + assert torch.allclose(frames[:, 0], frames[:, 1], atol=1e-6) + + +def test_an_absent_past_frame_is_not_normalized(): + algo = _pairing_algo( + T.Compose([T.ColorJitter(0.4), T.Normalize([0.5] * 3, [0.25] * 3)]) + ) + cam, batch = _pair_batch(torch.rand(B, 3, 8, 8), torch.zeros(B, 3, 8, 8)) + frames = algo._robomimic_to_hpt_data( + batch, [cam, f"{cam}_hist"], [], [], "actions" + )["front_img_1"] + assert torch.equal(frames[:, 0], torch.zeros_like(frames[:, 0])) + assert frames[:, 1].min() < 0, "the present frame is normalized" + + +def test_compile_traces_the_dit_through_the_heads(): + """RDT's trunk is not a compile target of its own (no_trunk); each head's + denoiser calls it, so its blocks land in the head's graph.""" + import types + + graphs = [] + + def backend(gm, _inputs): + graphs.append(gm) + return gm.forward + + model = _model().train() + original = nn.Module.compile + + def compile_with(self, **kwargs): + kwargs.pop("mode", None) + return original(self, backend=backend, **kwargs) + + torch._dynamo.reset() + try: + nn.Module.compile = compile_with + RDT.compile_for_training(types.SimpleNamespace(nets={"policy": model})) + torch.manual_seed(0) + model.compute_loss({"domain": DOMAIN, "data": _data()}).backward() + finally: + nn.Module.compile = original + torch._dynamo.reset() + assert any("blocks" in g.code and "backbone" in g.code for g in graphs) + assert all( + p.grad is not None + for n, p in model.named_parameters() + if "trunk.trunk.blocks" in n + ) diff --git a/tests/unit/test_train_step.py b/tests/unit/test_train_step.py index f4de8fde9a..75a08e4ebe 100644 --- a/tests/unit/test_train_step.py +++ b/tests/unit/test_train_step.py @@ -16,6 +16,7 @@ cpu_trainer_overrides, hpt_small_overrides, pi_cpu_overrides, + rdt_small_overrides, ) from fixtures.train_harness import ( STATE_DIM, @@ -136,6 +137,36 @@ def spy(self, *args, **kwargs): assert built["policy_device"] is None +def test_rdt_train_step(tmp_path, monkeypatch): + hermetic_env(monkeypatch) + recipe = RECIPES[("mecka", "rdt")] + data, out, hashes = write_fixtures(tmp_path, "mecka") + cfg = compose_recipe( + recipe, + common_overrides( + recipe.embodiment, + data, + out, + batch_size=2, + num_workers=0, + episode_hashes=hashes, + ) + + cpu_trainer_overrides(STEPS) + + rdt_small_overrides(recipe.embodiment), + out, + ) + spy = BatchKeySpy(monkeypatch, HPT) + _run(cfg) + expected = { + "actions_keypoints", + "observations.state.ee_pose", + Human.VIZ_IMAGE_KEY, + f"{Human.VIZ_IMAGE_KEY}_hist", + "fps", + } + assert expected <= spy.keys[recipe.embodiment] + + @pytest.mark.parametrize("vendor", VENDOR_NAMES) def test_pi_train_step(tmp_path, monkeypatch, vendor): reason = pi_unavailable() From 76c173d779fe778847fc887334fe12bd72ba2b7e Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sun, 20 Sep 2026 16:15:13 +0000 Subject: [PATCH 39/79] slurm: stop cancelled jobs from leaking whole GPUs A cancelled DDP job left GPUs allocated with nvidia-smi blaming PIDs absent from /proc. The driver is fine: the memory is held by live processes reparented to init. gpu19 had 126 pt_data_worker survivors across three cancellations, each holding /dev/nvidia0-7, pinning 7 GPUs and ~1.7 TB of RAM. Workers fork after CUDA init, so they inherit the rank's device fds, and ranks open every visible GPU -- one survivor pins the node. A worker also inherits both ends of its result pipe, so a dead rank never sends EPIPE; it blocks in pipe_write and never reaches PyTorch's ppid check. Whole ranks leak the same way: one wedged in torch.cuda.synchronize() inside inductor cudagraph_trees ignores SIGTERM until Slurm gives up (UnkillableStepTimeout) and drops the job with its ranks still spinning. ProctrackType=proctrack/linuxproc walks the PPID tree, so it never sweeps anything reparented to init. egomimic/utils/gpu_orphans.py: * orphan_guarded/die_with_parent put PR_SET_PDEATHSIG on every worker, so the kernel kills it with its rank whatever it is blocked on. Wired into the train and val loaders and the norm-stats pass. * arm_rank_pdeathsig does the same for the rank process. A process whose parent is already gone when it arms exits 1, after one line on stderr saying why (os._exit skips logging's flush). SLURM_JOB_ID is also set inside a salloc shell, so a run started there with nohup is left unarmed (SIGHUP ignored): armed, it died with the shell it was meant to outlive. * reap_orphans runs from trainHydra.main before CUDA init and clears what an earlier job left, so a leak cannot outlive the job that made it. It only signals processes that are ours, reparented to init and holding an nvidia fd. Slurm-gated; EGOMIMIC_REAP_GPU_ORPHANS=0 (or false/no) disables it, any truthy spelling keeps it on, and anything else raises. reap_gpu_orphans.sh is the manual version for a node we are not submitting to. It reads PPid from /proc//status, as the module does: in /proc//stat, field 2 is comm, and a comm with spaces shifted field 4 far enough to read a live process's PPid as 1. The PDEATHSIG test waits for the unguarded worker to block on its pipe, which Linux >= 5.5 reports as pipe_write and 5.4 (sky2) as pipe_wait. It accepts either and skips on anything else, killing the worker first: failing there used to leave it running under init. The kill-switch test never scans the real /proc. The cluster-level fix is ProctrackType=proctrack/cgroup, which needs admin. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01DF1j2AzwzFztRsqvbTC35b --- AGENTS.md | 37 +++- egomimic/pl_utils/pl_data_utils.py | 6 +- egomimic/rldb/zarr/zarr_dataset_multi.py | 2 + egomimic/trainHydra.py | 6 + egomimic/utils/gpu_orphans.py | 175 ++++++++++++++++++ reap_gpu_orphans.sh | 78 ++++++++ tests/unit/test_worker_orphan_guard.py | 217 +++++++++++++++++++++++ 7 files changed, 518 insertions(+), 3 deletions(-) create mode 100644 egomimic/utils/gpu_orphans.py create mode 100755 reap_gpu_orphans.sh create mode 100644 tests/unit/test_worker_orphan_guard.py diff --git a/AGENTS.md b/AGENTS.md index c081b5e36f..9316b38e98 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -25,4 +25,39 @@ use plan mode for anything except extremely simple tasks ## Slurm rules If you're on a slurm cluster, request a GPU before running or testing training. -On sky1/sky2: salloc -p rl2-lab -A rl2-lab --gres=gpu:a40:1 -c 12 --mem=30G \ No newline at end of file +On sky1/sky2: salloc -p rl2-lab -A rl2-lab --gres=gpu:a40:1 -c 12 --mem=30G + +## Cancelled jobs and leaked GPU memory +A cancelled DDP job can leave whole GPUs allocated, with `nvidia-smi` blaming PIDs that are not in +`/proc`. Nothing is wrong with the driver: the memory is held by **live** `pt_data_worker` processes +reparented to init. + +Dataloader workers are forked after CUDA init, so each inherits the rank's `/dev/nvidia*` fds — and +ranks open *every* visible GPU, so one survivor pins the whole node. The worker cannot exit on its +own: it inherited both ends of its result pipe, so a dead rank never gives it an `EPIPE`, it blocks +in `pipe_write`, and PyTorch's parent check (which only runs between `index_queue` timeouts) is +never reached. `ProctrackType=proctrack/linuxproc` then walks the PPID tree and never sees it. + +Whole ranks leak the same way. A rank wedged in `torch.cuda.synchronize()` inside inductor's +`cudagraph_trees` does not answer SIGTERM, Slurm hits `UnkillableStepTimeout` (60s) and gives up +(`srun: error: Timed out waiting for job step to complete`), and the job leaves the queue with its +ranks still spinning on every GPU. Enabling `torch.compile` makes this more likely, not less. + +`egomimic/utils/gpu_orphans.py` owns the fix: + +- `orphan_guarded(params)` / `die_with_parent` set `PR_SET_PDEATHSIG` on every worker, so the kernel + kills it with its rank regardless of what it is blocked on. **Any new `DataLoader` on a training + path needs one of them** — they are already on the train/val loaders and the norm-stats pass. +- `arm_rank_pdeathsig()` does the same for the rank process itself, so a launcher that gives up + takes its ranks with it. Slurm-gated, and skipped when SIGHUP is ignored, so a `nohup` run in + a `salloc` shell is not shot when that shell exits. +- `reap_orphans()` runs from `trainHydra.main` before CUDA init and clears what an earlier job left + on the node, so a leak cannot outlive the job that made it. It only ever signals processes that + are yours, reparented to init and holding an nvidia fd; a live job's workers have a live parent, + so they are never touched. Off outside Slurm and on non-zero `SLURM_LOCALID`; disable with + `EGOMIMIC_REAP_GPU_ORPHANS=0`. +- `./reap_gpu_orphans.sh [--kill] ` is the manual version, for nodes you are not about to + submit to. It attaches to your own job on a busy node rather than queueing behind it. + +The cluster-level fix is `ProctrackType=proctrack/cgroup`, which sweeps by cgroup instead of by PPID +tree and would make all of this unnecessary. That needs admin; `/etc/slurm` is read-only to us. diff --git a/egomimic/pl_utils/pl_data_utils.py b/egomimic/pl_utils/pl_data_utils.py index b41a6b64d6..d7fd4ae114 100644 --- a/egomimic/pl_utils/pl_data_utils.py +++ b/egomimic/pl_utils/pl_data_utils.py @@ -4,6 +4,8 @@ from lightning.pytorch.utilities.combined_loader import CombinedLoader, _MaxSizeCycle from torch.utils.data import DataLoader, default_collate +from egomimic.utils.gpu_orphans import orphan_guarded + logger = logging.getLogger(__name__) # Val heads that can carry a metric loader, in val_dataloader() order. @@ -178,7 +180,7 @@ def train_dataloader(self): dataset, shuffle=True, collate_fn=self.collate_fn, - **dataset_params, + **orphan_guarded(dataset_params), ) return CombinedLoader(iterables, "max_size_cycle") @@ -193,7 +195,7 @@ def _build_val_style_loader( raise ValueError( f"No dataloader params found for dataset {dataset_name}. Please add {dataset_name} into your data config {kind}_dataloader_params." ) - dataset_params = dict(dataset_params) + dataset_params = orphan_guarded(dataset_params) shuffle = dataset_params.pop("shuffle", False) iterables[dataset_name] = DataLoader( dataset, diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 204db1918f..a1fb155a86 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -54,6 +54,7 @@ _norm_center_scale, ) from egomimic.utils.env import load_env +from egomimic.utils.gpu_orphans import die_with_parent from egomimic.utils.pose_utils import ( bimanual_cartesian_layout, bimanual_keypoint_layout, @@ -1569,6 +1570,7 @@ def infer_norm_from_dataset( num_workers=num_workers, shuffle=True, generator=torch.Generator().manual_seed(seed), + worker_init_fn=die_with_parent, ) N = len(dataset) if N <= 0: diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 783f437dbe..4b5905ed4f 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -31,6 +31,7 @@ from egomimic.utils.checkpoint_utils import load_checkpoint_weights from egomimic.utils.compile_cache import set_per_job_compile_cache_dir from egomimic.utils.env import load_env +from egomimic.utils.gpu_orphans import arm_rank_pdeathsig, reap_orphans from egomimic.utils.instantiators import instantiate_callbacks, instantiate_loggers from egomimic.utils.logging_utils import log_hyperparameters from egomimic.utils.pylogger import RankedLogger @@ -892,6 +893,11 @@ def main(cfg: DictConfig) -> Optional[float]: # the job runs main(), so each job keys the cache on its own SLURM_JOB_ID. set_per_job_compile_cache_dir() + # Before CUDA init: a cancelled job can leave dataloader workers -- or whole + # ranks -- holding GPU contexts, and nothing else on this cluster sweeps them. + reap_orphans() + arm_rank_pdeathsig() + # apply extra utilities # (e.g. ask for tags if none are provided in cfg, print cfg tree, etc.) extras(cfg) diff --git a/egomimic/utils/gpu_orphans.py b/egomimic/utils/gpu_orphans.py new file mode 100644 index 0000000000..92ec131ba3 --- /dev/null +++ b/egomimic/utils/gpu_orphans.py @@ -0,0 +1,175 @@ +"""Reclaim GPUs pinned by dataloader workers that outlived a cancelled job. + +A worker forked after CUDA init holds its rank's ``/dev/nvidia*`` fds, so the +rank's whole context stays resident after the rank is gone -- ``nvidia-smi`` then +reports the memory against a PID that is no longer in ``/proc``. The worker +cannot exit on its own (it blocks in ``pipe_write`` on a result pipe nobody +drains), and ``ProctrackType=proctrack/linuxproc`` walks the PPID tree, so +Slurm never sweeps it once it is reparented to init. + +``orphan_guarded`` / ``die_with_parent`` stop new leaks. This is the other half: a job +clears what an earlier job left on its node, before it touches CUDA itself. + +Nothing is imported from torch here -- this has to run before CUDA init. +""" + +import ctypes +import logging +import os +import signal + +from egomimic.utils.type_utils import str2bool + +logger = logging.getLogger(__name__) + +_DISABLE_ENV = "EGOMIMIC_REAP_GPU_ORPHANS" +_PR_SET_PDEATHSIG = 1 + + +def die_with_parent(_worker_id: int = 0) -> None: + """Ask the kernel to SIGKILL this dataloader worker when its rank dies. + + PyTorch's own parent check only runs between ``index_queue.get()`` timeouts, + which a worker blocked writing to a full result pipe never reaches. The + kernel does not care what the worker is blocked on. + """ + try: + ctypes.CDLL("libc.so.6", use_errno=True).prctl( + _PR_SET_PDEATHSIG, signal.SIGKILL + ) + except (OSError, AttributeError): # non-Linux, or no libc + logger.debug("PR_SET_PDEATHSIG unavailable; orphaned workers may leak GPUs") + return + # pdeathsig does not fire retroactively if the rank died during the fork. + if os.getppid() == 1: + # os._exit skips logging's flush; say why before a rank exits 1 silently. + os.write( + 2, + f"[gpu_orphans] pid {os.getpid()}: parent gone before " + "PR_SET_PDEATHSIG was armed; exiting\n".encode(), + ) + os._exit(1) + + +def orphan_guarded(params: dict) -> dict: + """``params`` with ``die_with_parent`` run before any caller ``worker_init_fn``.""" + params = dict(params) + if params.get("num_workers", 0) == 0: + return params + user_init = params.get("worker_init_fn") + if user_init is None: + params["worker_init_fn"] = die_with_parent + else: + + def _chained(worker_id: int, _user_init=user_init) -> None: + die_with_parent(worker_id) + _user_init(worker_id) + + params["worker_init_fn"] = _chained + return params + + +def _ppid(proc_root: str, pid: str) -> int | None: + try: + with open(f"{proc_root}/{pid}/status") as f: + for line in f: + if line.startswith("PPid:"): + return int(line.split()[1]) + except (OSError, ValueError): + pass + return None + + +def _holds_gpu(proc_root: str, pid: str) -> bool: + fd_dir = f"{proc_root}/{pid}/fd" + try: + fds = os.listdir(fd_dir) + except OSError: # gone, or not ours + return False + for fd in fds: + try: + if os.readlink(f"{fd_dir}/{fd}").startswith("/dev/nvidia"): + return True + except OSError: + continue + return False + + +def arm_rank_pdeathsig() -> bool: + """Have the kernel kill this process if its launcher dies first. + + Workers are not the only thing that leaks. A rank wedged in + ``torch.cuda.synchronize()`` inside inductor's cudagraph_trees does not + answer SIGTERM; Slurm hits ``UnkillableStepTimeout``, gives up, and the job + leaves the queue with its ranks still spinning on every GPU of the node. + + Only armed under Slurm, where the parent is slurmstepd or the launcher that + spawned this rank and never legitimately exits first. Off elsewhere, and off + for a ``nohup`` run (SIGHUP ignored): inside a ``salloc`` shell that is a run + meant to outlive the shell, which PR_SET_PDEATHSIG would kill. + """ + if "SLURM_JOB_ID" not in os.environ: + return False + if signal.getsignal(signal.SIGHUP) is signal.SIG_IGN: + logger.info("SIGHUP is ignored (nohup); not tying this rank to its parent") + return False + die_with_parent() + return True + + +def find_orphans(proc_root: str = "/proc", uid: int | None = None) -> list[int]: + """PIDs that are ours, reparented to init, and holding a GPU device open. + + A live job's workers always have a live parent, so the PPID-1 test alone + keeps this off anything still running -- ours or anyone else's. + """ + uid = os.getuid() if uid is None else uid + mine = {os.getpid(), os.getppid()} + orphans = [] + for pid in os.listdir(proc_root): + if not pid.isdigit() or int(pid) in mine: + continue + try: + if os.stat(f"{proc_root}/{pid}").st_uid != uid: + continue + except OSError: + continue + if _ppid(proc_root, pid) != 1: + continue + if _holds_gpu(proc_root, pid): + orphans.append(int(pid)) + return sorted(orphans) + + +def reap_orphans(dry_run: bool = False) -> list[int]: + """Kill this node's leaked GPU holders. Returns the PIDs signalled. + + Off outside Slurm: only there do we know every process of ours on the node + belongs to a job, so that a PPID-1 GPU holder can only be wreckage. + """ + if not str2bool(os.environ.get(_DISABLE_ENV, "1")): + return [] + if "SLURM_JOB_ID" not in os.environ: + return [] + if os.environ.get("SLURM_LOCALID", "0") != "0": # once per node + return [] + + orphans = find_orphans() + if not orphans: + return [] + logger.warning( + "Reaping %d orphaned GPU holder(s) left by an earlier job on %s: %s " + "(set %s=0 to disable)", + len(orphans), + os.uname().nodename, + orphans, + _DISABLE_ENV, + ) + if dry_run: + return orphans + for pid in orphans: + try: + os.kill(pid, signal.SIGKILL) + except OSError as e: + logger.warning("could not kill orphan %d: %s", pid, e) + return orphans diff --git a/reap_gpu_orphans.sh b/reap_gpu_orphans.sh new file mode 100755 index 0000000000..e277b6d006 --- /dev/null +++ b/reap_gpu_orphans.sh @@ -0,0 +1,78 @@ +#!/usr/bin/env bash +# Find (and optionally kill) your own dataloader workers that outlived a cancelled +# job and are still pinning GPU memory. +# +# A worker forked after CUDA init holds its rank's /dev/nvidia* fds, so the rank's +# whole context stays resident even after the rank is gone -- nvidia-smi then +# reports memory against a PID that is no longer in /proc. ProctrackType on this +# cluster is proctrack/linuxproc, which tracks the PPID tree, so workers that got +# reparented to init are never swept. +# +# ./reap_gpu_orphans.sh gpu19 # report only +# ./reap_gpu_orphans.sh --kill gpu19 # reclaim +# ./reap_gpu_orphans.sh --kill gpu11 gpu19 +# +# Only processes that are yours, orphaned (PPID 1) and holding an nvidia fd are +# ever signalled, so this is safe to run against a node with live jobs on it. +set -uo pipefail + +KILL=0 +[[ ${1:-} == --kill ]] && { KILL=1; shift; } +if [[ $# -eq 0 ]]; then + echo "usage: $0 [--kill] [node...]" >&2 + exit 2 +fi + +PROBE=' +ME=$(id -u); victims=() +for d in /proc/[0-9]*; do + p=${d#/proc/} + [[ $(stat -c %u "$d" 2>/dev/null) == "$ME" ]] || continue + # PPid from status, as gpu_orphans._ppid reads it: in stat, field 2 is comm, + # which may contain spaces and shift every field after it. + [[ $(awk "/^PPid:/{print \$2}" "$d/status" 2>/dev/null) == 1 ]] || continue + ls -l "$d/fd" 2>/dev/null | grep -q /dev/nvidia || continue + victims+=("$p") +done +echo "node=$(hostname -s) orphans=${#victims[@]}" +[[ ${#victims[@]} -eq 0 ]] && exit 0 +for p in "${victims[@]}"; do + echo " $p $(cat /proc/$p/comm 2>/dev/null) rss=$(awk "/VmRSS/{print \$2}" /proc/$p/status 2>/dev/null)kB age=$(ps -o etime= -p $p 2>/dev/null | tr -d " ")" +done +echo "before: $(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader | tr "\n" " ")" +if [[ $KILL == 1 ]]; then + kill -9 "${victims[@]}" 2>/dev/null + for i in 1 2 3 4 5 6; do + sleep 4 + n=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader | wc -l) + [[ $n -eq 0 ]] && break + done + echo "after: $(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader | tr "\n" " ")" + echo "compute apps left: $n" +fi +' + +probe() { # probe ; returns srun's status, not grep's + "$@" bash -c "KILL=$KILL; $PROBE" 2>&1 | grep -v '^srun: job' + return "${PIPESTATUS[0]}" +} + +rc=0 +for node in "$@"; do + echo "=== $node ===" + # --immediate everywhere: the partition is EXCLUSIVE, so a fresh allocation on a + # busy node queues forever, and a step inside a job whose training step already + # holds every CPU never starts. Better to say so than to hang. + job=$(squeue -h -u "$USER" -w "$node" -t R -o %i | head -1) + if [[ -n $job ]]; then + echo "(node held by your job $job; attaching to it)" + probe srun --overlap --immediate=20 --jobid="$job" -n1 -c 1 && continue + echo "(could not get a step in $job -- it is probably mid-training)" + fi + probe srun -p loaner -A loaner --nodelist="$node" -n1 -c 2 --mem=4G --overlap \ + --immediate=30 -t 5 && continue + echo "$node is busy; rerun when it is idle. A job starting there cleans it" >&2 + echo "automatically anyway (egomimic.utils.gpu_orphans.reap_orphans)." >&2 + rc=1 +done +exit "$rc" diff --git a/tests/unit/test_worker_orphan_guard.py b/tests/unit/test_worker_orphan_guard.py new file mode 100644 index 0000000000..27fce77db4 --- /dev/null +++ b/tests/unit/test_worker_orphan_guard.py @@ -0,0 +1,217 @@ +"""Dataloader workers must die with their rank. + +Workers are forked after CUDA init, so they inherit the rank's ``/dev/nvidia*`` +fds. One survivor keeps every GPU context on the node alive, and +``proctrack/linuxproc`` will not sweep it because it has been reparented to init. +""" + +import os +import subprocess +import sys +import textwrap +import time + +import pytest + +from egomimic.utils import gpu_orphans +from egomimic.utils.gpu_orphans import ( + arm_rank_pdeathsig, + die_with_parent, + find_orphans, + orphan_guarded, + reap_orphans, +) + +# A worker that inherits both ends of its result pipe never gets an EPIPE when +# the rank dies; it blocks in pipe_write, so PyTorch's own ppid check -- which +# only runs between index_queue timeouts -- is never reached. +RANK = textwrap.dedent( + """ + import os, sys, time + sys.path.insert(0, {root!r}) + from egomimic.utils.gpu_orphans import die_with_parent + guard = sys.argv[1] == "guard" + r, w = os.pipe() + if os.fork() == 0: + if guard: + die_with_parent(0) + with open(sys.argv[2], "w") as f: + f.write(str(os.getpid())) + while True: + os.write(w, b"x" * 65536) + time.sleep(600) + """ +) + + +def _alive(pid: int) -> bool: + return os.path.isdir(f"/proc/{pid}") + + +def _pid_written(pidfile) -> bool: + return pidfile.exists() and pidfile.read_text().strip().isdigit() + + +# The wait channel of a writer blocked on a full pipe: split out of pipe_wait in +# Linux 5.5, so a 5.4 kernel (sky2) still reports the old name. +_PIPE_BLOCKED = {"pipe_write", "pipe_wait"} + + +def _run_rank(tmp_path, mode: str) -> int: + """Start a rank, SIGKILL it once its worker is wedged, return the worker pid.""" + root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + pidfile = tmp_path / f"pid_{mode}" + rank = subprocess.Popen( + [sys.executable, "-c", RANK.format(root=root), mode, str(pidfile)] + ) + try: + deadline = time.time() + 180 # the rank imports torch first + while time.time() < deadline and not _pid_written(pidfile): + time.sleep(0.2) + assert _pid_written(pidfile), "worker never started" + worker = int(pidfile.read_text()) + deadline = time.time() + 30 + wchan = "" + while time.time() < deadline: + wchan = open(f"/proc/{worker}/wchan").read() + if wchan in _PIPE_BLOCKED: + break + time.sleep(0.2) + else: + os.kill(worker, 9) # an unguarded worker would outlive the test + pytest.skip(f"this kernel reports the blocked worker as {wchan!r}") + finally: + rank.kill() + rank.wait() + return worker + + +@pytest.mark.skipif(sys.platform != "linux", reason="PR_SET_PDEATHSIG is Linux-only") +def test_unguarded_worker_outlives_its_rank(tmp_path): + worker = _run_rank(tmp_path, "noguard") + time.sleep(3) + try: + assert _alive(worker), "baseline no longer reproduces the leak" + finally: + os.kill(worker, 9) + + +@pytest.mark.skipif(sys.platform != "linux", reason="PR_SET_PDEATHSIG is Linux-only") +def test_guarded_worker_dies_with_its_rank(tmp_path): + worker = _run_rank(tmp_path, "guard") + deadline = time.time() + 15 + while time.time() < deadline and _alive(worker): + time.sleep(0.2) + if _alive(worker): + os.kill(worker, 9) + pytest.fail("worker survived its rank despite PR_SET_PDEATHSIG") + + +def test_orphan_guard_runs_before_a_caller_init(): + calls = [] + params = orphan_guarded( + {"num_workers": 2, "worker_init_fn": lambda i: calls.append(("user", i))} + ) + params["worker_init_fn"](3) + assert calls == [("user", 3)] + + +def test_orphan_guard_is_a_noop_without_workers(): + assert orphan_guarded({"num_workers": 0}) == {"num_workers": 0} + + +def test_orphan_guard_installs_the_default(): + assert orphan_guarded({"num_workers": 4})["worker_init_fn"] is die_with_parent + + +def _fake_proc(root, pid: int, *, ppid: int, fds: list[str]) -> None: + d = root / str(pid) + (d / "fd").mkdir(parents=True) + (d / "status").write_text(f"Name:\tpt_data_worker\nPPid:\t{ppid}\n") + for i, target in enumerate(fds): + os.symlink(target, d / "fd" / str(i)) + + +def test_find_orphans_picks_only_ours_orphaned_and_on_a_gpu(tmp_path): + uid = os.getuid() + gpu = ["/dev/nvidia0", "/dev/nvidiactl"] + _fake_proc(tmp_path, 100, ppid=1, fds=gpu) # the leak + _fake_proc(tmp_path, 101, ppid=42, fds=gpu) # live job's worker + _fake_proc(tmp_path, 102, ppid=1, fds=["/dev/null"]) # no GPU + assert find_orphans(str(tmp_path), uid=uid) == [100] + + +def test_find_orphans_skips_other_users(tmp_path): + uid = os.getuid() + _fake_proc(tmp_path, 200, ppid=1, fds=["/dev/nvidia0"]) + assert find_orphans(str(tmp_path), uid=uid + 1) == [] + + +def test_find_orphans_tolerates_a_process_exiting_mid_scan(tmp_path): + uid = os.getuid() + (tmp_path / "300").mkdir() # no status/fd: raced away + _fake_proc(tmp_path, 301, ppid=1, fds=["/dev/nvidia0"]) + assert find_orphans(str(tmp_path), uid=uid) == [301] + + +def test_reap_is_off_outside_slurm(monkeypatch): + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + assert reap_orphans() == [] + + +def test_reap_runs_once_per_node(monkeypatch): + monkeypatch.setenv("SLURM_JOB_ID", "1") + monkeypatch.setenv("SLURM_LOCALID", "3") + assert reap_orphans() == [] + + +@pytest.mark.parametrize( + "value, on", + [("0", False), ("false", False), ("1", True), ("true", True), ("yes", True)], +) +def test_reap_honours_the_kill_switch(monkeypatch, value, on): + """Any truthy spelling keeps it on. Never the real /proc: a regression here + would SIGKILL the tester's own orphans.""" + monkeypatch.setenv("SLURM_JOB_ID", "1") + monkeypatch.setenv("SLURM_LOCALID", "0") + monkeypatch.setenv("EGOMIMIC_REAP_GPU_ORPHANS", value) + monkeypatch.setattr(gpu_orphans, "find_orphans", lambda: [4242]) + assert reap_orphans(dry_run=True) == ([4242] if on else []) + + +def test_reap_rejects_a_misspelled_switch(monkeypatch): + import argparse + + monkeypatch.setenv("SLURM_JOB_ID", "1") + monkeypatch.setenv("EGOMIMIC_REAP_GPU_ORPHANS", "of") + monkeypatch.setattr(gpu_orphans, "find_orphans", lambda: [4242]) + with pytest.raises(argparse.ArgumentTypeError): + reap_orphans(dry_run=True) + + +def test_rank_pdeathsig_is_off_outside_slurm(monkeypatch): + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + assert arm_rank_pdeathsig() is False + + +def test_rank_pdeathsig_arms_under_slurm(monkeypatch): + import ctypes + + monkeypatch.setenv("SLURM_JOB_ID", "1") + try: + assert arm_rank_pdeathsig() is True + finally: + # do not leave pytest itself armed + ctypes.CDLL("libc.so.6").prctl(1, 0) + + +def test_rank_pdeathsig_leaves_a_nohup_run_alone(monkeypatch): + """nohup inside salloc means "outlive this shell"; arming would kill it.""" + import signal + + monkeypatch.setenv("SLURM_JOB_ID", "1") + previous = signal.signal(signal.SIGHUP, signal.SIG_IGN) + try: + assert arm_rank_pdeathsig() is False + finally: + signal.signal(signal.SIGHUP, previous) From 9c51f3423aa50cb22478c1b1136b0a9549e5cdb6 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 22 Sep 2026 09:58:48 -0400 Subject: [PATCH 40/79] trainer: DDP devices are per node; eval-mode val subsampling is single-rank trainer/ddp.yaml and ddp_pi.yaml set devices to gpus_per_node * nodes next to num_nodes, but Lightning reads devices per node and multiplies by num_nodes itself, so a 2-node run asked each node for twice the GPUs it has. The submitit launchers made the same mistake in gres, which Slurm also reads per node. Both now use gpus_per_node; single-node runs are unchanged. trainHydra scales each val head's frames-per-episode K by devices * num_nodes, but eval mode forces devices=1 only after the loaders are built. An eval run of a multi-GPU recipe therefore kept W times the frames per episode that its single rank's limit_val_batches window holds, and the window ran out before the last episodes. _trainer_world_size now returns 1 in eval mode. Both predate this stack (the ddp devices expression is from 2025-01). Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UkAZwERXcFKjueZQKfRF6J --- .../hydra/launcher/submitit.yaml | 2 +- .../hydra/launcher/submitit_pace.yaml | 2 +- .../hydra/launcher/submitit_skynet.yaml | 2 +- egomimic/hydra_configs/trainer/ddp.yaml | 3 +- egomimic/hydra_configs/trainer/ddp_pi.yaml | 3 +- egomimic/trainHydra.py | 6 +++- tests/unit/test_eval_loaders.py | 28 +++++++++++++++++++ 7 files changed, 40 insertions(+), 6 deletions(-) diff --git a/egomimic/hydra_configs/hydra/launcher/submitit.yaml b/egomimic/hydra_configs/hydra/launcher/submitit.yaml index c56f2cd508..f154af112e 100644 --- a/egomimic/hydra_configs/hydra/launcher/submitit.yaml +++ b/egomimic/hydra_configs/hydra/launcher/submitit.yaml @@ -10,7 +10,7 @@ account: "rl2-lab" # Slurm account (e.g., 'my_accou cpus_per_task: 12 # Number of CPUs per task nodes: ${launch_params.nodes} # Number of nodes tasks_per_node: ${launch_params.gpus_per_node} # Use variable for tasks per node -gres: "gpu:a40:${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'}" # GPU type and count +gres: "gpu:a40:${launch_params.gpus_per_node}" # GPU type and count per node qos: "short" # Slurm QoS timeout_min: 2880 # Timeout in minutes (48 hours) exclude: "protocol, puma" # Nodes to exclude diff --git a/egomimic/hydra_configs/hydra/launcher/submitit_pace.yaml b/egomimic/hydra_configs/hydra/launcher/submitit_pace.yaml index ca724575bf..e61539521e 100644 --- a/egomimic/hydra_configs/hydra/launcher/submitit_pace.yaml +++ b/egomimic/hydra_configs/hydra/launcher/submitit_pace.yaml @@ -10,7 +10,7 @@ account: "gts-dxu345-rl2" # Slurm account cpus_per_task: 8 # Number of CPUs per task (max 4:1 CPU:GPU ratio) nodes: ${launch_params.nodes} # Number of nodes tasks_per_node: ${launch_params.gpus_per_node} # Use variable for tasks per node -gres: "gpu:h200:${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'}" # GPU type and count (h100 for H100 GPUs) +gres: "gpu:h200:${launch_params.gpus_per_node}" # GPU type and count per node (h100 for H100 GPUs) qos: "inferno" # Slurm QoS mem_per_gpu: 250G timeout_min: 2880 # Timeout in minutes (48 hours) diff --git a/egomimic/hydra_configs/hydra/launcher/submitit_skynet.yaml b/egomimic/hydra_configs/hydra/launcher/submitit_skynet.yaml index ca83317d8a..35bcdaae03 100644 --- a/egomimic/hydra_configs/hydra/launcher/submitit_skynet.yaml +++ b/egomimic/hydra_configs/hydra/launcher/submitit_skynet.yaml @@ -11,7 +11,7 @@ cpus_per_task: 12 nodes: ${launch_params.nodes} tasks_per_node: ${launch_params.gpus_per_node} -gres: "gpu:l40s:${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'}" +gres: "gpu:l40s:${launch_params.gpus_per_node}" qos: "short" timeout_min: 2880 # 80 hours diff --git a/egomimic/hydra_configs/trainer/ddp.yaml b/egomimic/hydra_configs/trainer/ddp.yaml index d3d90aca03..8b8c98f072 100644 --- a/egomimic/hydra_configs/trainer/ddp.yaml +++ b/egomimic/hydra_configs/trainer/ddp.yaml @@ -4,7 +4,8 @@ defaults: strategy: ddp accelerator: gpu -devices: ${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'} +# Per node: Lightning's world size is devices * num_nodes. +devices: ${launch_params.gpus_per_node} num_nodes: ${launch_params.nodes} sync_batchnorm: True check_val_every_n_epoch: 200 diff --git a/egomimic/hydra_configs/trainer/ddp_pi.yaml b/egomimic/hydra_configs/trainer/ddp_pi.yaml index fcfdb55c37..48ab3b76f5 100644 --- a/egomimic/hydra_configs/trainer/ddp_pi.yaml +++ b/egomimic/hydra_configs/trainer/ddp_pi.yaml @@ -11,7 +11,8 @@ strategy: ddp_find_unused_parameters_true inference_mode: false accelerator: gpu -devices: ${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'} +# Per node: Lightning's world size is devices * num_nodes. +devices: ${launch_params.gpus_per_node} num_nodes: ${launch_params.nodes} sync_batchnorm: True check_val_every_n_epoch: 20 diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 4b5905ed4f..173bfbe9a0 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -331,8 +331,12 @@ def _trainer_world_size(cfg: DictConfig) -> int: DistributedSampler strides rather than chunks, so W ranks between them walk the WHOLE split; the per-rank ``limit_val_batches`` window is therefore W - times wider than it looks from one rank. + times wider than it looks from one rank. Eval mode runs on one device + (``train`` forces ``devices=1`` after the loaders are built), so it is 1. """ + legacy = "train" if cfg.get("train") else "eval" if cfg.get("eval") else None + if (cfg.get("mode") or legacy) == "eval": # _resolve_mode, minus its raise + return 1 devices = cfg.get("trainer", {}).get("devices", 1) if isinstance(devices, (list, ListConfig)): n = len(devices) diff --git a/tests/unit/test_eval_loaders.py b/tests/unit/test_eval_loaders.py index 3dc822589d..5f9e30df4b 100644 --- a/tests/unit/test_eval_loaders.py +++ b/tests/unit/test_eval_loaders.py @@ -337,6 +337,34 @@ def test_k_scales_with_the_world_size(): assert len(th._subsample_val_datasets(cfg, "valid", split)[HUMAN]) == 20 +def test_eval_mode_is_single_rank(): + """Eval runs are forced onto one device after the loaders are built, so + their val split must be subsampled for one rank, not the training W.""" + cfg = _cfg(metric={"valid": 7}) + cfg.trainer = {"devices": 4, "num_nodes": 2} + assert th._metric_frames_per_episode(cfg, "valid") == 56 + cfg.mode = "eval" + assert th._metric_frames_per_episode(cfg, "valid") == 7 + del cfg["mode"] + cfg.eval = True + assert th._metric_frames_per_episode(cfg, "valid") == 7 + + +@pytest.mark.parametrize("trainer_cfg", ["ddp", "ddp_pi"]) +def test_ddp_devices_are_per_node(trainer_cfg, compose_resolve): + """Lightning reads devices per node; devices * num_nodes is the world.""" + cfg = compose_resolve( + "train_zarr_cartesian", + [ + f"trainer={trainer_cfg}", + "launch_params.gpus_per_node=8", + "launch_params.nodes=2", + ], + ) + assert cfg.trainer.devices == 8 and cfg.trainer.num_nodes == 2 + assert th._trainer_world_size(cfg) == 16 + + def test_auto_k_is_derived_from_the_resolved_split(): """`auto` = floor(limit_val_batches * batch_size / episodes), per dataset, from the head's own loader batch size, then scaled by the world size.""" From 015e4a7914d579f7b4b4234484c03a904a9d89a9 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 22 Sep 2026 13:14:03 +0000 Subject: [PATCH 41/79] eval: log validation videos to wandb Each mp4 EvalVideo writes is also logged to any WandbLogger under