diff --git a/CONTRIBUTING_DATA.md b/CONTRIBUTING_DATA.md index 57fc23ca85..b182a33638 100644 --- a/CONTRIBUTING_DATA.md +++ b/CONTRIBUTING_DATA.md @@ -644,7 +644,7 @@ If you are contributing egocentric human data, you use `human_bimanual` (or the There is a **single concrete `Human` embodiment class** ([`egomimic/rldb/embodiment/human.py`](egomimic/rldb/embodiment/human.py)) shared by all human data. You do **not** write a per-vendor subclass, and there is **no per-vendor embodiment identifier** — every human contributor uses `human_*` and records their source in the `lab` field. Camera intrinsics travel with the data (`zarr.attrs`, §6.3 / §6.4); per-vendor structural choices are passed as explicit arguments from the data config: - `Human.get_keymap(keymap_mode="cartesian"|"keypoints", has_head_pose=, include_aria_keypoints=)` -- `Human.get_transform_list(mode="cartesian"|"keypoints_headframe_ypr"|..., stride=)` +- `Human.get_transform_list(mode="cartesian_6d"|"keypoints_headframe_ypr"|..., stride=)` Onboarding human data is just two steps: @@ -947,7 +947,7 @@ import torch HAS_HEAD_POSE = True key_map = Human.get_keymap(keymap_mode="cartesian", has_head_pose=HAS_HEAD_POSE) transform_list = ( - Human.get_transform_list(mode="cartesian", stride=3) if HAS_HEAD_POSE else None + Human.get_transform_list(mode="cartesian_6d", stride=3) if HAS_HEAD_POSE else None ) resolver = LocalEpisodeResolver( @@ -987,7 +987,7 @@ from egomimic.rldb.filters import DatasetFilter resolver = LocalEpisodeResolver( folder_path = "/local/processed", key_map = Human.get_keymap(keymap_mode="cartesian"), - transform_list = Human.get_transform_list(mode="cartesian", stride=3), + transform_list = Human.get_transform_list(mode="cartesian_6d", stride=3), ) filters = DatasetFilter(filter_lambdas=[ "lambda row: row['episode_hash'] == '2026-03-15-14-22-10-000000'" diff --git a/egomimic/algo/pi.py b/egomimic/algo/pi.py index dc137c737d..d055eab546 100644 --- a/egomimic/algo/pi.py +++ b/egomimic/algo/pi.py @@ -421,13 +421,7 @@ def forward_training(self, batch): ) losses = self.nets["policy"].forward(processed_obs, action) - - 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) - - loss = losses.mean() + loss = self._reduce_loss(losses, _batch[ac_key], embodiment_id, ac_key) predictions[f"{embodiment_name}_{ac_key}"] = _batch[ac_key] predictions[f"{embodiment_name}_loss"] = loss @@ -465,13 +459,9 @@ def forward_eval(self, batch): # Flow-matching val loss — same call as forward_training. losses = self.nets["policy"].forward(processed_obs, action) - 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 - ) - unnorm_preds[f"{embodiment_name}_loss"] = losses.mean() + unnorm_preds[f"{embodiment_name}_loss"] = self._reduce_loss( + losses, _batch[ac_key], embodiment_id, ac_key + ) pred_actions = self.nets["policy"].sample_actions( device=self.device, @@ -480,22 +470,36 @@ def forward_eval(self, batch): num_steps=self.num_steps, ) - predictions = OrderedDict() - ref = _batch[ac_key] - B, T, D = ref.shape - - converter = self.action_registry.get(embodiment_id, ac_key) - pred_actions_orig = converter.from32(pred_actions) - - pred = pred_actions_orig[:, :T, :D] - predictions[ac_key] = pred - - unnorm_actions = self.norm_stats.unnormalize(predictions, embodiment_id) + unnorm_actions = self._postprocess_sampled_actions( + pred_actions, _batch, embodiment_id, ac_key + ) for key in unnorm_actions: unnorm_preds[f"{embodiment_name}_{key}"] = unnorm_actions[key] return unnorm_preds + def _reduce_loss(self, losses, action, embodiment_id, ac_key) -> torch.Tensor: + """Mean flow-matching loss for the embodiment.""" + 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) + return losses.mean() + + 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).""" + ref = _batch[ac_key] + _, T, D = ref.shape + converter = self.action_registry.get(embodiment_id, ac_key) + predictions = OrderedDict() + pred_native = converter.from32_norm_6d(pred_actions) + predictions[ac_key] = pred_native[:, :T, :D] + return self.norm_stats.unnormalize(predictions, embodiment_id) + @override def compute_losses(self, predictions, batch): """ @@ -555,7 +559,10 @@ def _robomimic_to_pi_data(self, batch, proprio_keys, lang_keys, ac_key, embodime emb_id = get_embodiment_id(embodiment) # embodiment is a name string converter = self.action_registry.get(emb_id, ac_key) - action32 = converter.to32(action) + # 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) # 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 @@ -630,40 +637,3 @@ def _clone_batch(self, batch): return batch.clone() else: return batch # Return as is for non-tensor types - - def _extract_xyz(self, x): - """ - Extract xyz (3D position) and rotation from 6DoF or 6DoF+gripper actions. - - Supports: - - 6: 6DoF (single arm) - - 7: 6DoF + gripper (single arm) - - 12: 2 arms × 6DoF - - 14: 2 arms × (6DoF + gripper) - - Returns: - xyz: Tensor with only xyz per arm (shape: ..., 3) or (..., 6) for dual-arm. - rot: Tensor with only rotation per arm (shape: ..., 3) or (..., 6) for dual-arm. - """ - if x.shape[-1] == 6: - return x[..., :3], x[..., 3:6] - elif x.shape[-1] == 7: - return x[..., :3], x[..., 3:6] - elif x.shape[-1] == 12: - xyz_right = x[..., :3] - rot_right = x[..., 3:6] - xyz_left = x[..., 6:9] - rot_left = x[..., 9:12] - return torch.cat([xyz_right, xyz_left], dim=-1), torch.cat( - [rot_right, rot_left], dim=-1 - ) - elif x.shape[-1] == 14: - xyz_right = x[..., :3] - rot_right = x[..., 3:6] - xyz_left = x[..., 7:10] - rot_left = x[..., 10:13] - return torch.cat([xyz_right, xyz_left], dim=-1), torch.cat( - [rot_right, rot_left], dim=-1 - ) - else: - raise ValueError(f"Unexpected shape for 6DoF input: {x.shape}") diff --git a/egomimic/eval/eval_latent.py b/egomimic/eval/eval_latent.py index 972c5e12ea..604900ce63 100644 --- a/egomimic/eval/eval_latent.py +++ b/egomimic/eval/eval_latent.py @@ -221,7 +221,7 @@ def on_validation_start(self): exist_ok=True, ) - def compute_metrics_and_viz(self, batch): + def compute_metrics_and_viz(self, batch, do_viz=True): algo = self.model metrics = {} images_dict = {} @@ -260,7 +260,7 @@ def compute_metrics_and_viz(self, batch): ref = _batch[ac_key] B, T, D = ref.shape converter = algo.action_registry.get(embodiment_id, ac_key) - pred_actions_orig = converter.from32(pred_actions) + pred_actions_orig = converter.from32_norm_6d(pred_actions) pred = pred_actions_orig[:, :T, :D] predictions = OrderedDict() @@ -280,7 +280,7 @@ def compute_metrics_and_viz(self, batch): unnorm_batch[ac_key][:, -1].cpu(), ) - if self.viz_func is not None: + if do_viz and self.viz_func is not None: images_dict[embodiment_id] = self._visualize_preds( unnorm_preds, unnorm_batch ) diff --git a/egomimic/hydra_configs/data/aria.yaml b/egomimic/hydra_configs/data/aria.yaml index f96c9750c5..ef1550dd1f 100644 --- a/egomimic/hydra_configs/data/aria.yaml +++ b/egomimic/hydra_configs/data/aria.yaml @@ -11,7 +11,10 @@ train_datasets: keymap_mode: cartesian transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + mode: cartesian_6d + # 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: 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 b1cf807ca5..9612b56990 100644 --- a/egomimic/hydra_configs/data/aria_qwen.yaml +++ b/egomimic/hydra_configs/data/aria_qwen.yaml @@ -16,7 +16,10 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + mode: cartesian_6d + # 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: 3 filters: null mode: train diff --git a/egomimic/hydra_configs/data/bc_pickplace_eva_qwen.yaml b/egomimic/hydra_configs/data/bc_pickplace_eva_qwen.yaml deleted file mode 100644 index bfeba77ebc..0000000000 --- a/egomimic/hydra_configs/data/bc_pickplace_eva_qwen.yaml +++ /dev/null @@ -1,55 +0,0 @@ -# BC eva-only pick_place dataset for the HPT + Qwen-annotation experiment. -# Same filters as cotrain_pickplace_qwen but human_bimanual stripped out. -# train = 216 eva_bimanual -# valid = 12 eva_bimanual alignment -# Persistent zarr cache is shared with cotrain (no re-sync). -_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper - -train_datasets: - eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - key_map: - _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap - keymap_mode: cartesian - annotation_key: annotations - transform_list: - _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'eva_bimanual' and (row.get('zarr_processed_path') or '') != '' and 'alignment' not in ((row.get('task_description') or '').lower())" - mode: total - -valid_datasets: - eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - key_map: - _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap - keymap_mode: cartesian - annotation_key: annotations - transform_list: - _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'eva_bimanual' and (row.get('zarr_processed_path') or '') != '' and 'alignment' in ((row.get('task_description') or '').lower())" - mode: total - -train_dataloader_params: - eva_bimanual: - batch_size: 32 - num_workers: 6 -valid_dataloader_params: - eva_bimanual: - batch_size: 32 - num_workers: 6 diff --git a/egomimic/hydra_configs/data/cotrain_pi_base.yaml b/egomimic/hydra_configs/data/cotrain_pi_base.yaml index e6ce59979c..bfa45b5bba 100644 --- a/egomimic/hydra_configs/data/cotrain_pi_base.yaml +++ b/egomimic/hydra_configs/data/cotrain_pi_base.yaml @@ -12,7 +12,7 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian + mode: cartesian_wristframe_6d filters: null mode: train valid_ratio: 0.05 @@ -27,7 +27,10 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + mode: cartesian_6d + # 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: 3 filters: null mode: train diff --git a/egomimic/hydra_configs/data/cotrain_pi_lang_wrist.yaml b/egomimic/hydra_configs/data/cotrain_pi_lang_wrist.yaml index 641e483fa4..19146085ed 100644 --- a/egomimic/hydra_configs/data/cotrain_pi_lang_wrist.yaml +++ b/egomimic/hydra_configs/data/cotrain_pi_lang_wrist.yaml @@ -2,16 +2,11 @@ defaults: - cotrain_pi_lang - _self_ +# Kept for configs saved by earlier runs: cotrain_pi_lang with eva on the +# wrist-frame 6D action, which is now the cotrain_pi_base default, so this +# composes the same as cotrain_pi_lang. Pair with the default evaluator=eval_pi. train_datasets: eva_bimanual: resolver: - key_map: - keymap_mode: cartesian_wristframe_ypr transform_list: - mode: cartesian_wristframe_ypr - human_bimanual: - resolver: - key_map: - keymap_mode: cartesian - transform_list: - mode: cartesian_wristframe_ypr + mode: cartesian_wristframe_6d diff --git a/egomimic/hydra_configs/data/cotrain_pi_pickplace_qwen.yaml b/egomimic/hydra_configs/data/cotrain_pi_pickplace_qwen.yaml deleted file mode 100644 index e225b9557b..0000000000 --- a/egomimic/hydra_configs/data/cotrain_pi_pickplace_qwen.yaml +++ /dev/null @@ -1,75 +0,0 @@ -# Pi0.5 cotrain pick_place data — same filters as cotrain_pickplace_qwen. -# Camera keys are the dataset names; the Pi wrapper renames them onto openpi's -# slots (base_0_rgb, *_wrist_0_rgb) itself. -_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper - -train_datasets: - eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - key_map: - _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap - keymap_mode: cartesian - annotation_key: annotations - transform_list: - _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'eva_bimanual' and (row.get('zarr_processed_path') or '') != '' and 'alignment' not in ((row.get('task_description') or '').lower())" - mode: total - human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - 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 - stride: 3 - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'human_bimanual' and (row.get('zarr_processed_path') or '') != '' and any(s in ((row.get('task_description') or '').lower()) for s in ('base', 'object')) and (row.get('episode_hash') or '') not in ('2026-04-26-00-17-53-000000', '2026-04-26-00-27-26-000000', '2026-05-01-02-52-58-000000')" - mode: total - -valid_datasets: - eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - key_map: - _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap - keymap_mode: cartesian - annotation_key: annotations - transform_list: - _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'eva_bimanual' and (row.get('zarr_processed_path') or '') != '' and 'alignment' in ((row.get('task_description') or '').lower())" - mode: total - -train_dataloader_params: - eva_bimanual: - batch_size: 32 - num_workers: 6 - human_bimanual: - batch_size: 32 - num_workers: 6 -valid_dataloader_params: - eva_bimanual: - batch_size: 32 - num_workers: 6 diff --git a/egomimic/hydra_configs/data/cotrain_pickplace_qwen.yaml b/egomimic/hydra_configs/data/cotrain_pickplace_qwen.yaml deleted file mode 100644 index c05902d4e0..0000000000 --- a/egomimic/hydra_configs/data/cotrain_pickplace_qwen.yaml +++ /dev/null @@ -1,86 +0,0 @@ -# defaults removed: hardcoding folder_path below for standalone load by the annotation pipeline - - - -# Cotrain pick_place dataset for the HPT + Qwen-annotation experiment. -# -# train = 216 eva_bimanual + 103 human_bimanual = 319 episodes (mode: total). -# valid = 12 eva_bimanual alignment episodes (aria has 0 alignment episodes in -# the dense-language Scale project today; left out intentionally). -# -# All four filters are gated by ScaleAnnotationDatasetFilter(dense-language) -# so the scale_to_bucket -> bucket_to_zarr pipeline can process exactly the -# episodes referenced here. -_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper - -train_datasets: - eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - key_map: - _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap - keymap_mode: cartesian - annotation_key: annotations - transform_list: - _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'eva_bimanual' and (row.get('zarr_processed_path') or '') != '' and 'alignment' not in ((row.get('task_description') or '').lower())" - mode: total - human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - 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 - stride: 3 - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'human_bimanual' and (row.get('zarr_processed_path') or '') != '' and any(s in ((row.get('task_description') or '').lower()) for s in ('base', 'object')) and (row.get('episode_hash') or '') not in ('2026-04-26-00-17-53-000000', '2026-04-26-00-27-26-000000', '2026-05-01-02-52-58-000000')" - mode: total - -valid_datasets: - eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - key_map: - _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap - keymap_mode: cartesian - annotation_key: annotations - transform_list: - _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'eva_bimanual' and (row.get('zarr_processed_path') or '') != '' and 'alignment' in ((row.get('task_description') or '').lower())" - mode: total - -train_dataloader_params: - eva_bimanual: - batch_size: 32 - num_workers: 6 - human_bimanual: - batch_size: 32 - num_workers: 6 - -valid_dataloader_params: - eva_bimanual: - batch_size: 32 - num_workers: 6 diff --git a/egomimic/hydra_configs/data/cotrain_pickplace_qwen_objgen.yaml b/egomimic/hydra_configs/data/cotrain_pickplace_qwen_objgen.yaml deleted file mode 100644 index abb5958ea5..0000000000 --- a/egomimic/hydra_configs/data/cotrain_pickplace_qwen_objgen.yaml +++ /dev/null @@ -1,102 +0,0 @@ -# Cotrain pick_place dataset with object_gen monitoring. -# -# Same train set as cotrain_pickplace_qwen.yaml. Valid additionally includes -# human_bimanual object_gen episodes (36 ep) — a subset of train, used for -# monitoring object_gen-specific train loss (NOT held-out generalization). -# -# train = 216 eva_bimanual + 105 human_bimanual = 321 episodes -# valid = 12 eva_bimanual alignment + ~36 human_bimanual object_gen = ~48 ep -_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper - -train_datasets: - eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - key_map: - _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap - keymap_mode: cartesian - annotation_key: annotations - transform_list: - _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'eva_bimanual' and (row.get('zarr_processed_path') or '') != '' and 'alignment' not in ((row.get('task_description') or '').lower())" - mode: total - human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - 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_padded - stride: 3 - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'human_bimanual' and (row.get('zarr_processed_path') or '') != '' and any(s in ((row.get('task_description') or '').lower()) for s in ('base', 'object')) and (row.get('episode_hash') or '') not in ('2026-04-26-00-17-53-000000', '2026-04-26-00-27-26-000000', '2026-05-01-02-52-58-000000')" - mode: total - -valid_datasets: - eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - key_map: - _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap - keymap_mode: cartesian - annotation_key: annotations - transform_list: - _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'eva_bimanual' and (row.get('zarr_processed_path') or '') != '' and 'alignment' in ((row.get('task_description') or '').lower())" - mode: total - human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - 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_padded - stride: 3 - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'human_bimanual' and (row.get('zarr_processed_path') or '') != '' and 'object' in ((row.get('task_description') or '').lower())" - mode: total - -train_dataloader_params: - eva_bimanual: - batch_size: 32 - num_workers: 6 - human_bimanual: - batch_size: 32 - num_workers: 6 -valid_dataloader_params: - eva_bimanual: - batch_size: 32 - num_workers: 6 - human_bimanual: - batch_size: 32 - num_workers: 6 diff --git a/egomimic/hydra_configs/data/cotrain_pickplace_qwen_wrist.yaml b/egomimic/hydra_configs/data/cotrain_pickplace_qwen_wrist.yaml deleted file mode 100644 index 9ef6424d5a..0000000000 --- a/egomimic/hydra_configs/data/cotrain_pickplace_qwen_wrist.yaml +++ /dev/null @@ -1,78 +0,0 @@ -# Wristframe variant of cotrain_pickplace_qwen.yaml. -# Transform list uses cartesian_wristframe_ypr — actions are in each EEF wrist -# frame instead of head/base frame. Action chunk key + dim unchanged -# (actions_cartesian, 14D for eva). Pair with viz=cartesian_wrist_lang. - -_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper - -train_datasets: - eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - key_map: - _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap - keymap_mode: cartesian - annotation_key: annotations - transform_list: - _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian_wristframe_ypr - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'eva_bimanual' and (row.get('zarr_processed_path') or '') != '' and 'alignment' not in ((row.get('task_description') or '').lower())" - mode: total - human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - 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_ypr - stride: 3 - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'human_bimanual' and (row.get('zarr_processed_path') or '') != '' and any(s in ((row.get('task_description') or '').lower()) for s in ('base', 'object')) and (row.get('episode_hash') or '') not in ('2026-04-26-00-17-53-000000', '2026-04-26-00-27-26-000000', '2026-05-01-02-52-58-000000')" - mode: total - -valid_datasets: - eva_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver - folder_path: /storage/project/r-dxu345-0/paphiwetsa3/datasets/pickplace_qwen - key_map: - _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap - keymap_mode: cartesian - annotation_key: annotations - transform_list: - _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian_wristframe_ypr - filters: - _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter - project_name: dense-language - filter_lambdas: - - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'eva_bimanual' and (row.get('zarr_processed_path') or '') != '' and 'alignment' in ((row.get('task_description') or '').lower())" - mode: total - -train_dataloader_params: - eva_bimanual: - batch_size: 32 - num_workers: 6 - human_bimanual: - batch_size: 32 - num_workers: 6 - -valid_dataloader_params: - eva_bimanual: - batch_size: 32 - num_workers: 6 diff --git a/egomimic/hydra_configs/data/eva.yaml b/egomimic/hydra_configs/data/eva.yaml index c1b366a3ca..c56f51c4e9 100644 --- a/egomimic/hydra_configs/data/eva.yaml +++ b/egomimic/hydra_configs/data/eva.yaml @@ -11,7 +11,9 @@ train_datasets: keymap_mode: cartesian transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian + # Wrist-relative actions: each arm's chunk in its current EEF frame, + # continuous 6D rotation (20-D); the proprio ee_pose stays cam frame. + mode: cartesian_wristframe_6d filters: _target_: egomimic.rldb.filters.DatasetFilter episode_hashes: ["2025-12-26-18-07-46-296000"] diff --git a/egomimic/hydra_configs/data/human.yaml b/egomimic/hydra_configs/data/human.yaml index 13cab563ee..59a472a346 100644 --- a/egomimic/hydra_configs/data/human.yaml +++ b/egomimic/hydra_configs/data/human.yaml @@ -17,7 +17,10 @@ train_datasets: keymap_mode: cartesian transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + mode: cartesian_6d + # 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: 3 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/industry_eva_pi.yaml b/egomimic/hydra_configs/data/industry_eva_pi.yaml index 39c88da22d..34d2e421ef 100644 --- a/egomimic/hydra_configs/data/industry_eva_pi.yaml +++ b/egomimic/hydra_configs/data/industry_eva_pi.yaml @@ -15,7 +15,7 @@ train_datasets: keymap_mode: cartesian transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian + mode: cartesian_wristframe_6d filters: _target_: egomimic.rldb.filters.DatasetFilter filter_lambdas: @@ -31,7 +31,10 @@ train_datasets: keymap_mode: cartesian transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + mode: cartesian_6d + # 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: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/mecka.yaml b/egomimic/hydra_configs/data/mecka.yaml index e038820b8c..65befd7318 100644 --- a/egomimic/hydra_configs/data/mecka.yaml +++ b/egomimic/hydra_configs/data/mecka.yaml @@ -11,7 +11,10 @@ train_datasets: keymap_mode: cartesian transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + mode: cartesian_6d + # 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: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/mecka_scale_cotrain_pi.yaml b/egomimic/hydra_configs/data/mecka_scale_cotrain_pi.yaml index a2cd727256..1571b3a9ca 100644 --- a/egomimic/hydra_configs/data/mecka_scale_cotrain_pi.yaml +++ b/egomimic/hydra_configs/data/mecka_scale_cotrain_pi.yaml @@ -16,7 +16,10 @@ train_datasets: keymap_mode: cartesian transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + mode: cartesian_6d + # 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: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/scale.yaml b/egomimic/hydra_configs/data/scale.yaml index 593685ac72..97ce8f0eae 100644 --- a/egomimic/hydra_configs/data/scale.yaml +++ b/egomimic/hydra_configs/data/scale.yaml @@ -11,7 +11,10 @@ train_datasets: keymap_mode: cartesian transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + mode: cartesian_6d + # 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: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_aria.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_aria.yaml index ef8a7fa67a..eb345dd1e9 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_aria.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_aria.yaml @@ -4,7 +4,9 @@ robomimic_model: camera_transforms: {} dims: - human_bimanual: {proprio: 12, action: 12} + # 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 diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml index 21108686ea..196dc7a716 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_eva.yaml @@ -4,7 +4,8 @@ robomimic_model: camera_transforms: {} dims: - eva_bimanual: {proprio: 14, action: 14} + # cartesian_wristframe_6d (data/eva.yaml): [xyz | rot6d | grip] x 2. + eva_bimanual: {proprio: 20, action: 20} diffusion: true 6dof: true diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml index 159ce6e69d..21c3f7fb01 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_human.yaml @@ -9,7 +9,9 @@ robomimic_model: camera_transforms: {} dims: - human_bimanual: {proprio: 12, action: 12} + # 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 diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_mecka.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_mecka.yaml index ef8a7fa67a..eb345dd1e9 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_mecka.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_mecka.yaml @@ -4,7 +4,9 @@ robomimic_model: camera_transforms: {} dims: - human_bimanual: {proprio: 12, action: 12} + # 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 diff --git a/egomimic/hydra_configs/model/hpt_bc_flow_scale.yaml b/egomimic/hydra_configs/model/hpt_bc_flow_scale.yaml index 4d446b7b72..6ae2b6d1d4 100644 --- a/egomimic/hydra_configs/model/hpt_bc_flow_scale.yaml +++ b/egomimic/hydra_configs/model/hpt_bc_flow_scale.yaml @@ -4,7 +4,9 @@ robomimic_model: camera_transforms: {} dims: - human_bimanual: {proprio: 12, action: 12} + # 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 diff --git a/egomimic/hydra_configs/model/hpt_bc_pickplace_qwen_pertoken.yaml b/egomimic/hydra_configs/model/hpt_bc_pickplace_qwen_pertoken.yaml deleted file mode 100644 index bd110d503a..0000000000 --- a/egomimic/hydra_configs/model/hpt_bc_pickplace_qwen_pertoken.yaml +++ /dev/null @@ -1,149 +0,0 @@ -# BC HPT (eva-only, robot data) + Qwen3-Embedding (per-token hidden states) text stem. -# Pair with: data=bc_pickplace_eva_qwen -# train = 216 eva_bimanual (pick_place, not alignment, dense-language gated) -# valid = 12 eva_bimanual alignment -# Action format = actions_cartesian (cartesian_wristframe_ypr), act_dim=14. -_target_: egomimic.pl_utils.pl_model.ModelWrapper -robomimic_model: - _target_: egomimic.algo.hpt.HPT - - camera_transforms: {} - - dims: - eva_bimanual: {proprio: 14, action: 14} - - diffusion: true - 6dof: true - - ac_keys: - eva_bimanual: "actions_cartesian" - - annotation_key: "annotations" - annotation_sampling_mode: "random" - annotation_modality: "annotation" - default_prompt: "" - - 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: "" - reverse_kl_samples: 8 - - domains: ["eva_bimanual"] - shared_obs_keys: ["front_img_1", "annotation"] - - 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 - annotation: - _target_: egomimic.models.hpt_nets.QwenPerTokenEncoder - model_name: "Qwen/Qwen3-Embedding-0.6B" - max_length: 128 - freeze: true - dtype: "float16" - output_dim: 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: - eva_bimanual: - state_ee_pose: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: ${model.robomimic_model.dims.eva_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: - eva_bimanual: - _target_: egomimic.models.fm_policy.FMPolicy - action_horizon: 100 - num_inference_steps: 50 - pooling: null - time_dist: "beta" - infer_ac_dims: - eva_bimanual: ${model.robomimic_model.dims.eva_bimanual.action} - model: - _target_: egomimic.models.denoising_nets.CrossTransformer - nblocks: 6 - cond_dim: 256 - hidden_dim: 128 - act_dim: ${model.robomimic_model.dims.eva_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 - 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: 2e-4 - weight_decay: 0.0001 - -scheduler: - _target_: egomimic.utils.scheduler_utils.warmup_then_cosine - _partial_: true - warmup_epochs: 100 - total_epochs: 2000 - eta_min: 1.0e-5 -scheduler_interval: epoch diff --git a/egomimic/hydra_configs/model/hpt_bc_pickplace_qwen_pooled.yaml b/egomimic/hydra_configs/model/hpt_bc_pickplace_qwen_pooled.yaml deleted file mode 100644 index a7beadb0b7..0000000000 --- a/egomimic/hydra_configs/model/hpt_bc_pickplace_qwen_pooled.yaml +++ /dev/null @@ -1,149 +0,0 @@ -# BC HPT (eva-only, robot data) + Qwen3-Embedding (last-token pooled) text stem. -# Pair with: data=bc_pickplace_eva_qwen -# train = 216 eva_bimanual (pick_place, not alignment, dense-language gated) -# valid = 12 eva_bimanual alignment -# Action format = actions_cartesian (cartesian_wristframe_ypr), act_dim=14. -_target_: egomimic.pl_utils.pl_model.ModelWrapper -robomimic_model: - _target_: egomimic.algo.hpt.HPT - - camera_transforms: {} - - dims: - eva_bimanual: {proprio: 14, action: 14} - - diffusion: true - 6dof: true - - ac_keys: - eva_bimanual: "actions_cartesian" - - annotation_key: "annotations" - annotation_sampling_mode: "random" - annotation_modality: "annotation" - default_prompt: "" - - 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: "" - reverse_kl_samples: 8 - - domains: ["eva_bimanual"] - shared_obs_keys: ["front_img_1", "annotation"] - - 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 - annotation: - _target_: egomimic.models.hpt_nets.QwenPooledEncoder - model_name: "Qwen/Qwen3-Embedding-0.6B" - max_length: 128 - freeze: true - dtype: "float16" - output_dim: 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: - eva_bimanual: - state_ee_pose: - _target_: egomimic.models.hpt_nets.MLPPolicyStem - input_dim: ${model.robomimic_model.dims.eva_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: - eva_bimanual: - _target_: egomimic.models.fm_policy.FMPolicy - action_horizon: 100 - num_inference_steps: 50 - pooling: null - time_dist: "beta" - infer_ac_dims: - eva_bimanual: ${model.robomimic_model.dims.eva_bimanual.action} - model: - _target_: egomimic.models.denoising_nets.CrossTransformer - nblocks: 6 - cond_dim: 256 - hidden_dim: 128 - act_dim: ${model.robomimic_model.dims.eva_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 - 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: 2e-4 - weight_decay: 0.0001 - -scheduler: - _target_: egomimic.utils.scheduler_utils.warmup_then_cosine - _partial_: true - warmup_epochs: 100 - total_epochs: 2000 - eta_min: 1.0e-5 -scheduler_interval: epoch 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 new file mode 100644 index 0000000000..301122433b --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria_6d.yaml @@ -0,0 +1,33 @@ +defaults: + - pi0.5_base + - _self_ + +# All-cartesian pi0.5 cotrain: data=cotrain_pi_lang plus +# data.train_datasets.eva_bimanual.resolver.transform_list.mode=cartesian_6d +# 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 +# with evaluator=eval_pi_camframe_6d. Both embodiments on +# the head/camera-frame cartesian action in the continuous 6D rotation layout +# (eva 20D, human 18D), normalized by the standard pipeline; the forward pass +# only packs them into openpi's 32D vector (pretrained projections kept). The +# proprio ee_pose arrives 6D-encoded as well. No "Embodiment: " prompt +# block (pi0.5_base defaults it on). +robomimic_model: + camera_transforms: {} + embodiment_label: false + ac_keys: + eva_bimanual: "actions_cartesian" + human_bimanual: "actions_cartesian" + domains: ["eva_bimanual", "human_bimanual"] + + action_converters: + rules: + EVA_BIMANUAL: + _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler + HUMAN_BIMANUAL: + _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + # optional fallback if no match is found + fallback: + _target_: egomimic.utils.action_utils.BaseActionConverter diff --git a/egomimic/rldb/embodiment/embodiment.py b/egomimic/rldb/embodiment/embodiment.py index 568b638a93..c6fb9d1870 100644 --- a/egomimic/rldb/embodiment/embodiment.py +++ b/egomimic/rldb/embodiment/embodiment.py @@ -51,8 +51,31 @@ def get_embodiment(index): return EMBODIMENT_ID_TO_KEY.get(index, None) +# Human demo data written by the vendor-split registry carries vendor-tagged +# embodiment metadata (e.g. MECKA_BIMANUAL, SCALE_LEFT_ARM). Locally all human +# demonstration data is ONE embodiment (see the EMBODIMENT docstring; the +# source lives only in the SQL `lab` field), so those names collapse to +# HUMAN_*. Robot names (EVA_*) are never aliased. +HUMAN_VENDOR_PREFIXES = ("MECKA", "SCALE", "ARIA", "LIGHTWHEEL") + + +def canonical_embodiment_name(embodiment_name: str) -> str: + """Upper-case EMBODIMENT member name, with legacy vendor prefixes + (``MECKA_BIMANUAL`` ...) collapsed onto ``HUMAN_*``. Does not validate.""" + name = embodiment_name.upper() + vendor, _, suffix = name.partition("_") + if vendor in HUMAN_VENDOR_PREFIXES and suffix: + return f"HUMAN_{suffix}" + return name + + +def is_legacy_vendor_embodiment(embodiment_name: str) -> bool: + """True for vendor-tagged human names that only resolve through aliasing.""" + return canonical_embodiment_name(embodiment_name) != embodiment_name.upper() + + def get_embodiment_id(embodiment_name): - return EMBODIMENT[embodiment_name.upper()].value + return EMBODIMENT[canonical_embodiment_name(embodiment_name)].value def _strip_pi_keymap_mode(cls, keymap_mode: str) -> str: @@ -73,6 +96,31 @@ def _strip_pi_keymap_mode(cls, keymap_mode: str) -> str: return base_mode +# Rotation representations that are discontinuous on SO(3) -- Euler angles wrap +# at +-pi and gimbal-lock, quaternions double-cover -- so per-dim normalization +# and regression losses on them are ill-defined. Training targets the continuous +# 6D representation everywhere; these modes survive only at the data/rollout +# boundary and for checkpoints that predate it. +LEGACY_ROTATION_MODES = { + "cartesian": "cartesian_6d", + "cartesian_padded": "cartesian_6d", + "cartesian_wristframe_ypr": "cartesian_wristframe_6d", + "cartesian_wristframe_quat": "cartesian_wristframe_6d", +} + + +def _reject_legacy_rotation(cls, mode: str, allow_legacy_rotation: bool) -> None: + replacement = LEGACY_ROTATION_MODES.get(mode) + if replacement is None or allow_legacy_rotation: + return + raise ValueError( + f"{cls.__name__} transform_list mode '{mode}' trains on a discontinuous " + f"rotation; use '{replacement}' instead. Reading data written before the " + "6D conversion (viz, rollout, a legacy checkpoint) needs an explicit " + "allow_legacy_rotation=True." + ) + + class Embodiment(ABC): """Base embodiment class. An embodiment is responsible for defining the transform pipeline that converts between the raw data in the dataset and the canonical representation used by the model.""" diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index 3304e80f4a..6067ad16fc 100644 --- a/egomimic/rldb/embodiment/eva.py +++ b/egomimic/rldb/embodiment/eva.py @@ -4,11 +4,17 @@ import numpy as np -from egomimic.rldb.embodiment.embodiment import Embodiment, _strip_pi_keymap_mode +from egomimic.rldb.embodiment.embodiment import ( + Embodiment, + _reject_legacy_rotation, + _strip_pi_keymap_mode, +) from egomimic.rldb.embodiment.human import ARIA_INTRINSICS from egomimic.rldb.zarr.action_chunk_transforms import ( ActionChunkCoordinateFrameTransform, BatchQuaternionPoseToYPR, + CartesianRot6DToYPR, + CartesianYPRToRot6D, ConcatKeys, DeleteKeys, InterpolateLinear, @@ -46,18 +52,67 @@ class Eva(Embodiment): ), } - @staticmethod + @classmethod def get_transform_list( + cls, mode: Literal[ - "cartesian", "cartesian_wristframe_ypr", "cartesian_wristframe_quat" + "cartesian", + "cartesian_6d", + "cartesian_wristframe_ypr", + "cartesian_wristframe_6d", + "cartesian_wristframe_quat", ], + allow_legacy_rotation: bool = False, ) -> list[Transform]: + """Transform pipeline for the requested mode. + + ``allow_legacy_rotation`` opts out of the continuous-rotation rule for + the modes in ``LEGACY_ROTATION_MODES``; only the data/rollout boundary + (viz, a checkpoint that predates the 6D conversion) may set it. + """ + _reject_legacy_rotation(cls, mode, allow_legacy_rotation) if mode == "cartesian": return _build_eva_bimanual_transform_list(is_quat=True) - elif mode == "cartesian_wristframe_ypr": + if mode == "cartesian_6d": + # Camera-frame cartesian (14D xyz+ypr+gripper per arm pair) with + # the rotation re-expressed as the continuous 6D representation + # (20D) for pi0.5. The proprio ee_pose is 6D-encoded too: + # normalized YPR proprio saturates yaw/roll at +-pi (wraparound), + # so per-dim normalization is only meaningful on the continuous + # representation, same as for the actions. + return _build_eva_bimanual_transform_list(is_quat=True) + [ + CartesianYPRToRot6D(action_key="actions_cartesian"), + CartesianYPRToRot6D(action_key="observations.state.ee_pose"), + ] + if mode == "cartesian_wristframe_ypr": return _build_eva_bimanual_eef_frame_transform_list(is_quat=False) - elif mode == "cartesian_wristframe_quat": + if mode == "cartesian_wristframe_6d": + # Wrist-frame cartesian with the 6D rotation (20D). The cam-frame + # proprio ee_pose is 6D-encoded too (see cartesian_6d); it is the + # only cam-frame signal the model sees with wrist-relative targets. + return _build_eva_bimanual_eef_frame_transform_list(is_quat=False) + [ + CartesianYPRToRot6D(action_key="actions_cartesian"), + CartesianYPRToRot6D(action_key="observations.state.ee_pose"), + ] + if mode == "cartesian_wristframe_quat": return _build_eva_bimanual_eef_frame_transform_list(is_quat=True) + raise ValueError(f"Unsupported transform_list mode '{mode}' for Eva") + + @staticmethod + def get_revert_transform_list(mode: str) -> list[Transform] | None: + """Transforms that map a (unnormalized) prediction batch of a model + trained with ``get_transform_list(mode)`` back to camera-frame + xyz+ypr+gripper (7/arm), given the batch's proprio + ``observations.state.ee_pose``. ``None``: already in that layout.""" + if mode == "cartesian": + return None + if mode == "cartesian_6d": + return _build_eva_cartesian_revert_6d_transform_list() + if mode == "cartesian_wristframe_ypr": + return _build_eva_bimanual_revert_eef_frame_transform_list(is_quat=False) + if mode == "cartesian_wristframe_6d": + return _build_eva_cartesian_revert_6d_wristframe_transform_list() + raise ValueError(f"No camera-frame revert for Eva transform_list mode '{mode}'") @classmethod def _get_keymap(cls, keymap_mode: str): @@ -144,6 +199,46 @@ def dinov3_keymap(cls): } +def _build_eva_cartesian_revert_6d_transform_list( + *, + action_key: str = "actions_cartesian", + obs_key: str = "observations.state.ee_pose", +) -> list[Transform]: + """Revert camera-frame 6D-rotation EVA cartesian actions back to ypr. + + For the cam-frame 6D evaluator: the action chunk is already in camera + frame (``cartesian_6d`` mode), so only the rotation representation is + converted from xyz+6D(+gripper, 10/arm) back to xyz+ypr(+gripper, 7/arm) + so the viz overlay sees the same layout as the plain ``cartesian`` mode. + The proprio ee_pose is reverted the same way. + """ + return [ + CartesianRot6DToYPR(action_key=action_key), + CartesianRot6DToYPR(action_key=obs_key), + ] + + +def _build_eva_cartesian_revert_6d_wristframe_transform_list( + *, + action_key: str = "actions_cartesian", + obs_key: str = "observations.state.ee_pose", +) -> list[Transform]: + """Revert wrist-frame 6D-rotation EVA actions back to camera-frame ypr. + + (1) ``CartesianRot6DToYPR`` converts the action rotation xyz+6D -> xyz+ypr + (Gram-Schmidt re-orthonormalizes the model prediction); (2) the proprio + ``observations.state.ee_pose`` (6D-encoded by ``cartesian_wristframe_6d``) + is reverted to ypr the same way; (3) the standard eef-frame revert + projects the wrist-frame ypr actions back into camera frame using that + ypr proprio to define the frame. + """ + return [ + CartesianRot6DToYPR(action_key=action_key), + CartesianRot6DToYPR(action_key=obs_key), + *_build_eva_bimanual_revert_eef_frame_transform_list(is_quat=False), + ] + + def _build_eva_bimanual_revert_eef_frame_transform_list( *, action_key: str = "actions_cartesian", diff --git a/egomimic/rldb/embodiment/human.py b/egomimic/rldb/embodiment/human.py index d0356e7f59..f18dd0e59a 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -4,10 +4,16 @@ import numpy as np -from egomimic.rldb.embodiment.embodiment import Embodiment, _strip_pi_keymap_mode +from egomimic.rldb.embodiment.embodiment import ( + Embodiment, + _reject_legacy_rotation, + _strip_pi_keymap_mode, +) from egomimic.rldb.zarr.action_chunk_transforms import ( ActionChunkCoordinateFrameTransform, BatchQuaternionPoseToYPR, + CartesianRot6DToYPR, + CartesianYPRToRot6D, ConcatKeys, DeleteKeys, InterpolatePose, @@ -17,6 +23,7 @@ Reshape, SplitKeys, Transform, + UnpadGripperZeros, XYZWXYZ_to_XYZYPR, ) from egomimic.utils.viz_utils import ( @@ -356,26 +363,72 @@ def get_transform_list( cls, mode: Literal[ "cartesian", + "cartesian_6d", "cartesian_padded", "cartesian_wristframe_ypr", + "cartesian_wristframe_6d", "keypoints_headframe_ypr", "keypoints_headframe_quat", "keypoints_wristframe_ypr", "keypoints_wristframe_quat", ], stride: int = 3, + pad_proprio_gripper: bool = False, + allow_legacy_rotation: bool = False, ) -> list[Transform]: """Transform pipeline. ``stride`` is the per-vendor action stride (Aria/LightWheel=3, Scale/Mecka=1), supplied by the data config. + + ``allow_legacy_rotation`` opts out of the continuous-rotation rule for + the modes in ``LEGACY_ROTATION_MODES``; only the data/rollout boundary + (viz, a checkpoint that predates the 6D conversion) may set it. + + ``pad_proprio_gripper`` pads the cartesian proprio + ``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. """ + _reject_legacy_rotation(cls, mode, allow_legacy_rotation) + if pad_proprio_gripper and not mode.endswith("_6d"): + raise ValueError( + "pad_proprio_gripper needs a 6D-encoded ee_pose proprio: a " + "cartesian *_6d mode" + ) + + proprio_6d: list[Transform] = [ + CartesianYPRToRot6D(action_key="observations.state.ee_pose") + ] + if pad_proprio_gripper: + proprio_6d.append(PadGripperZeros(action_key="observations.state.ee_pose")) + if mode == "cartesian": return _build_human_cartesian_bimanual_transform_list(stride=stride) + 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. + return ( + _build_human_cartesian_bimanual_transform_list(stride=stride) + + [CartesianYPRToRot6D(action_key="actions_cartesian")] + + proprio_6d + ) if mode == "cartesian_padded": return _build_human_cartesian_bimanual_transform_list(stride=stride) + [ PadGripperZeros(action_key="actions_cartesian") ] if mode == "cartesian_wristframe_ypr": return _build_human_cartesian_eef_frame_transform_list(stride=stride) + if mode == "cartesian_wristframe_6d": + # Wrist-frame cartesian with the 6D rotation (18D). The head-frame + # proprio ee_pose is 6D-encoded too (see cartesian_6d); it is the + # only head-frame signal the model sees with wrist-relative targets. + return ( + _build_human_cartesian_eef_frame_transform_list(stride=stride) + + [CartesianYPRToRot6D(action_key="actions_cartesian")] + + proprio_6d + ) if mode == "keypoints_headframe_ypr": return _build_human_keypoints_bimanual_transform_list( stride=stride, is_quat=False @@ -952,6 +1005,49 @@ def _build_human_cartesian_revert_eef_frame_transform_list( return transform_list +def _build_human_cartesian_revert_6d_transform_list( + *, + action_key: str = "actions_cartesian", + obs_key: str = "observations.state.ee_pose", +) -> list[Transform]: + """Revert head/camera-frame 6D-rotation cartesian actions back to ypr. + + For the cam-frame 6D evaluator: the action chunk is already in head frame + (``cartesian_6d`` mode), so no coordinate-frame change is needed; only the + rotation representation is converted from xyz+6D (9/arm) back to xyz+ypr + (6/arm) so the viz overlay sees the same layout as the plain + ``cartesian`` mode. The proprio ee_pose (also 6D-encoded, possibly + grip-padded 18 -> 20) is reverted the same way. + """ + return [ + CartesianRot6DToYPR(action_key=action_key), + CartesianRot6DToYPR(action_key=obs_key), + UnpadGripperZeros(action_key=obs_key), + ] + + +def _build_human_cartesian_revert_6d_wristframe_transform_list( + *, + action_key: str = "actions_cartesian", + obs_key: str = "observations.state.ee_pose", +) -> list[Transform]: + """Revert wrist-frame 6D-rotation human actions back to head-frame ypr. + + (1) ``CartesianRot6DToYPR`` converts the action rotation xyz+6D -> xyz+ypr + (Gram-Schmidt re-orthonormalizes the model prediction); (2) the proprio + ``observations.state.ee_pose`` (6D-encoded by ``cartesian_wristframe_6d``) + is reverted to ypr the same way and unpadded (20 -> 14 -> 12; no-op when + unpadded); (3) the standard eef-frame revert projects the wrist-frame ypr + actions back into head frame using that ypr proprio to define the frame. + """ + return [ + CartesianRot6DToYPR(action_key=action_key), + CartesianRot6DToYPR(action_key=obs_key), + UnpadGripperZeros(action_key=obs_key), + *_build_human_cartesian_revert_eef_frame_transform_list(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 c45d058d1e..158f7efd2c 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -582,14 +582,29 @@ def transform(self, batch): # Shape Transforms # --------------------------------------------------------------------------- class SplitKeys(Transform): - def __init__(self, input_key: str, output_key_list: list[(str, int)]): + def __init__(self, input_key: str, output_key_list: list[tuple[str, int]]): self.input_key = input_key self.output_key_list = list(output_key_list) def transform(self, batch: dict) -> dict: + value = batch[self.input_key] + expected = sum(size for _, size in self.output_key_list) + width = int(value.shape[-1]) + if width != expected: + # Every caller lays out the WHOLE vector, so a mismatch means the + # split layout and the data disagree: typically an evaluator + # revert list built for one transform mode (e.g. 12-dim ypr) fed a + # batch from another (18/20-dim 6D). Slicing silently would hand + # the frame math xyz + rot6d columns as "xyz + ypr". + raise ValueError( + f"SplitKeys: '{self.input_key}' has last dim {width} but the " + f"output layout {self.output_key_list} sums to {expected}. Check " + "that the evaluator transform_lists match the data config's " + "transform mode." + ) prev_end = 0 for key, size in self.output_key_list: - batch[key] = batch[self.input_key][..., prev_end : prev_end + size] + batch[key] = value[..., prev_end : prev_end + size] prev_end += size return batch @@ -618,30 +633,59 @@ def transform(self, batch): class PadGripperZeros(Transform): - """Pad a 12D bimanual cartesian action chunk to 14D by inserting a zero - gripper slot at position 6 (end of left arm) and position 13 (end of right - arm), matching the canonical [L xyz ypr g, R xyz ypr g] layout used by Eva. - - Used so aria (which has no gripper signal) can share an FM denoiser head - sized for 14D actions without needing in-model padding branches. + """Pad a gripperless bimanual cartesian vector to the robot layout by + inserting a zero gripper slot at the end of each arm block: + 12D ypr -> 14D (slots 6, 13) or 18D rot6d -> 20D (slots 9, 19), matching + the canonical [L xyz rot g, R xyz rot g] layout used by Eva. + + Used so human data (which has no gripper signal) can share an FM denoiser + head sized for robot actions, and so the human proprio ``ee_pose`` bins in + the pi0.5 prompt align positionally with the robot 20D layout. """ def __init__(self, action_key: str = "actions_cartesian"): self.action_key = action_key def transform(self, batch: dict) -> dict: - actions = batch[self.action_key] - is_tensor = isinstance(actions, torch.Tensor) - arr = actions.cpu().numpy() if is_tensor else np.asarray(actions) - if arr.shape[-1] != 12: + arr, is_tensor = _as_array(batch[self.action_key]) + pad = np.zeros((*arr.shape[:-1], 1), dtype=arr.dtype) + if arr.shape[-1] == 12: + out = np.concatenate((arr[..., :6], pad, arr[..., 6:], pad), axis=-1) + elif arr.shape[-1] == 18: + out = np.concatenate((arr[..., :9], pad, arr[..., 9:], pad), axis=-1) + else: raise ValueError( - f"PadGripperZeros expects last-dim 12, got {arr.shape} for " + f"PadGripperZeros expects last-dim 12 or 18, got {arr.shape} for " f"'{self.action_key}'" ) - pad_shape = (*arr.shape[:-1], 1) - pad = np.zeros(pad_shape, dtype=arr.dtype) - padded = np.concatenate((arr[..., :6], pad, arr[..., 6:], pad), axis=-1) - batch[self.action_key] = torch.from_numpy(padded) if is_tensor else padded + batch[self.action_key] = _like_input(out, is_tensor) + return batch + + +class UnpadGripperZeros(Transform): + """Inverse of :class:`PadGripperZeros`: drop the per-arm zero gripper + slots, 14 -> 12 (ypr: drop 6, 13) or 20 -> 18 (rot6d: drop 9, 19). Widths + 12 and 18 pass through unchanged, so a revert pipeline works whether or + not the forward pipeline padded (``pad_proprio_gripper``).""" + + def __init__(self, action_key: str = "observations.state.ee_pose"): + self.action_key = action_key + + def transform(self, batch: dict) -> dict: + arr, is_tensor = _as_array(batch[self.action_key]) + D = arr.shape[-1] + if D in (12, 18): + return batch + if D == 14: + keep = [i for i in range(14) if i not in (6, 13)] + elif D == 20: + keep = [i for i in range(20) if i not in (9, 19)] + else: + raise ValueError( + f"UnpadGripperZeros expects last-dim 12/14/18/20, got {arr.shape} " + f"for '{self.action_key}'" + ) + batch[self.action_key] = _like_input(arr[..., keep], is_tensor) return batch diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 6b63aba25e..bef29cf5d4 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -39,7 +39,10 @@ import zarr from tqdm import tqdm -from egomimic.rldb.embodiment.embodiment import get_embodiment_id +from egomimic.rldb.embodiment.embodiment import ( + canonical_embodiment_name, + get_embodiment_id, +) # from action_chunk_transforms import Transform from egomimic.rldb.filters import DatasetFilter @@ -201,6 +204,17 @@ def get_fallback_idx( return random.choice(valid_candidates), attempts +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).""" + try: + return canonical_embodiment_name(str(name)) == canonical_embodiment_name( + expected + ) + except AttributeError: + return False + + class PinError(ValueError): """A pinned episode_hash cannot be used: missing, deleted, no processed path, or its embodiment does not match the dataset it was pinned under.""" @@ -236,7 +250,7 @@ def _check_pins( problems.append(f"{h}: empty zarr_processed_path") if expected_embodiment is not None: emb = str(row.get("embodiment", "")) - if emb != expected_embodiment: + if not _embodiment_matches(emb, expected_embodiment): problems.append( f"{h}: embodiment '{emb}' != dataset '{expected_embodiment}'" ) @@ -253,7 +267,9 @@ def _warn_embodiment_mismatch( ) -> None: if expected_embodiment is None or "embodiment" not in rows: return - mismatched = rows.loc[rows["embodiment"].astype(str) != expected_embodiment] + mismatched = rows.loc[ + ~rows["embodiment"].map(lambda e: _embodiment_matches(e, expected_embodiment)) + ] if len(mismatched): details = sorted( f"{h} ({e})" @@ -802,7 +818,7 @@ def _get_local_filtered_paths( problems.append(f"{h}: not in local directory {search_path}") elif expected_embodiment is not None: emb = str(present[h].get("embodiment", "")) - if emb != expected_embodiment: + if not _embodiment_matches(emb, expected_embodiment): problems.append( f"{h}: embodiment '{emb}' != dataset '{expected_embodiment}'" ) @@ -823,7 +839,7 @@ def _get_local_filtered_paths( bad = sorted( (r["episode_hash"], str(r.get("embodiment", ""))) for r in matched - if str(r.get("embodiment", "")) != expected_embodiment + if not _embodiment_matches(r.get("embodiment", ""), expected_embodiment) ) if bad: details = [f"{h} ({e})" for h, e in bad] diff --git a/egomimic/robot/rollout.py b/egomimic/robot/rollout.py index e74f2900c9..fb252bb726 100644 --- a/egomimic/robot/rollout.py +++ b/egomimic/robot/rollout.py @@ -258,7 +258,10 @@ def __init__( self.debug_actions = None self.resampled_action_len = resampled_action_len self.debug = debug - self.transform_list = Eva.get_transform_list(mode="cartesian_wristframe_ypr") + # Inference on checkpoints trained before the 6D conversion. + self.transform_list = Eva.get_transform_list( + mode="cartesian_wristframe_ypr", allow_legacy_rotation=True + ) self.annotation = None self._tokenizer = None self.collate_fn = default_collate diff --git a/egomimic/scripts/tutorials/zarr_data_viz.ipynb b/egomimic/scripts/tutorials/zarr_data_viz.ipynb index fe6b436c71..2f2e12d795 100644 --- a/egomimic/scripts/tutorials/zarr_data_viz.ipynb +++ b/egomimic/scripts/tutorials/zarr_data_viz.ipynb @@ -66,7 +66,7 @@ "outputs": [], "source": [ "key_map = Eva.get_keymap(keymap_mode=\"cartesian\")\n", - "transform_list = Eva.get_transform_list(mode=\"cartesian\")\n", + "transform_list = Eva.get_transform_list(mode=\"cartesian\", allow_legacy_rotation=True)\n", "\n", "resolver = S3EpisodeResolver(TEMP_DIR, key_map=key_map, transform_list=transform_list)\n", "filters = DatasetFilter(\n", @@ -129,7 +129,9 @@ "from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver\n", "\n", "key_map = Human.get_keymap(keymap_mode=\"cartesian\")\n", - "transform_list = Human.get_transform_list(mode=\"cartesian\", stride=3)\n", + "transform_list = Human.get_transform_list(\n", + " mode=\"cartesian\", stride=3, allow_legacy_rotation=True\n", + ")\n", "\n", "TEMP_DIR = \"/coc/flash7/scratch/egoverseDebugDatasets/egoverseS3DatasetTest/\"\n", "resolver = LocalEpisodeResolver(\n", diff --git a/egomimic/scripts/tutorials/zarr_data_viz.py b/egomimic/scripts/tutorials/zarr_data_viz.py index 23adfaa211..9f55b1a3be 100644 --- a/egomimic/scripts/tutorials/zarr_data_viz.py +++ b/egomimic/scripts/tutorials/zarr_data_viz.py @@ -19,7 +19,9 @@ load_env() key_map = Human.get_keymap(keymap_mode="cartesian") -transform_list = Human.get_transform_list(mode="cartesian", stride=3) +transform_list = Human.get_transform_list( + mode="cartesian", stride=3, allow_legacy_rotation=True +) resolver = S3EpisodeResolver(TEMP_DIR, key_map=key_map, transform_list=transform_list) filters = DatasetFilter( diff --git a/egomimic/utils/action_utils.py b/egomimic/utils/action_utils.py index 2b5268e305..7a80fbd970 100644 --- a/egomimic/utils/action_utils.py +++ b/egomimic/utils/action_utils.py @@ -1,3 +1,15 @@ +"""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. + +The forward pass only *packs* the normalized vector into openpi's 32-slot +action vector (``to32_norm_6d``) and the eval path unpacks it +(``from32_norm_6d``). No rotation math and no normalization happen here. +""" + from typing import Any, Dict, Tuple import torch @@ -23,7 +35,7 @@ def get(self, embodiment_id: int, ac_key: str) -> "BaseActionConverter": ) -# ---------- shared helpers (same conventions you used) ---------- +# ---------- shared helpers ---------- def _ensure_bsd(x: torch.Tensor) -> torch.Tensor: if x.ndim == 2: return x.unsqueeze(1) @@ -174,16 +186,27 @@ def _reconstruct_R_from_cols(c1: torch.Tensor, c2: torch.Tensor) -> torch.Tensor # ---------- base interface ---------- class BaseActionConverter: """ - Implement both directions: - - to32(actions_orig) -> (B,S,32) - - from32(actions_32) -> original shape/dim + Pack / unpack between an embodiment's native normalized action and the + pi0.5 action vector: + - to32_norm_6d(actions) -> (B,S,>=32) canonical block layout + - from32_norm_6d(actions32) -> native shape/dim + The base class implements neither; it is the ``fallback`` converter in the + model yamls and raises so an unregistered embodiment fails loudly. """ - def to32(self, actions: torch.Tensor) -> torch.Tensor: - raise NotImplementedError + def to32_norm_6d(self, actions: torch.Tensor) -> torch.Tensor: + """Pack an already-normalized native action into the canonical + pi0.5 layout (pure rearrange: no rotation math, no normalization).""" + raise NotImplementedError( + f"{type(self).__name__} does not support normalized-rot6d encoding" + ) - def from32(self, actions32: torch.Tensor) -> torch.Tensor: - raise NotImplementedError + def from32_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: + """Inverse of :meth:`to32_norm_6d`: extract the normalized native + action from the model's action vector (pure rearrange).""" + raise NotImplementedError( + f"{type(self).__name__} does not support normalized-rot6d decoding" + ) # ============================================================ @@ -191,111 +214,30 @@ def from32(self, actions32: torch.Tensor) -> torch.Tensor: # ============================================================ -class RobotLeftCartesianEuler(BaseActionConverter): - """ - Input orig: (B,S,7) = [x,y,z,yaw,pitch,roll,grip] - 32-pack: indices 0..9 = [xyz, R[:,0], R[:,1], grip] - """ - - def to32(self, actions: torch.Tensor) -> torch.Tensor: - actions = _ensure_bsd(actions) - if actions.shape[-1] != 7: - raise ValueError(f"RobotLeft: expected 7-dim, got {actions.shape[-1]}") - xyz = actions[..., 0:3] - ypr = actions[..., 3:6] - g = actions[..., 6:7] - R = _ypr_to_matrix(ypr) - c1, c2 = R[..., 0], R[..., 1] - block = torch.cat([xyz, c1, c2, g], dim=-1) # (B,S,10) - return _pad32(block) - - def from32(self, actions32: torch.Tensor) -> torch.Tensor: - actions32 = _ensure_bsd(actions32) - block = actions32[..., 0:10] # left block - xyz = block[..., 0:3] - c1 = block[..., 3:6] - c2 = block[..., 6:9] - g = block[..., 9:10] - R = _reconstruct_R_from_cols(c1, c2) - ypr = _matrix_to_ypr(R) - return torch.cat([xyz, ypr, g], dim=-1) # (B,S,7) - - -class RobotRightCartesianEuler(BaseActionConverter): - """ - Input orig: (B,S,7) - 32-pack: indices 10..19 = right block; indices 0..9 are zeros - """ - - def to32(self, actions: torch.Tensor) -> torch.Tensor: - actions = _ensure_bsd(actions) - if actions.shape[-1] != 7: - raise ValueError(f"RobotRight: expected 7-dim, got {actions.shape[-1]}") - xyz = actions[..., 0:3] - ypr = actions[..., 3:6] - g = actions[..., 6:7] - R = _ypr_to_matrix(ypr) - c1, c2 = R[..., 0], R[..., 1] - right = torch.cat([xyz, c1, c2, g], dim=-1) # (B,S,10) - zeros = torch.zeros_like(right) - return _pad32(torch.cat([zeros, right], dim=-1)) # (B,S,20) -> pad 32 - - def from32(self, actions32: torch.Tensor) -> torch.Tensor: - actions32 = _ensure_bsd(actions32) - block = actions32[..., 10:20] # right block - xyz = block[..., 0:3] - c1 = block[..., 3:6] - c2 = block[..., 6:9] - g = block[..., 9:10] - R = _reconstruct_R_from_cols(c1, c2) - ypr = _matrix_to_ypr(R) - return torch.cat([xyz, ypr, g], dim=-1) # (B,S,7) - - class RobotBimanualCartesianEuler(BaseActionConverter): """ - Input orig: (B,S,14) = L7 | R7 - 32-pack: left block 0..9, right block 10..19 + 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 """ - def to32(self, actions: torch.Tensor) -> torch.Tensor: + def to32_norm_6d(self, actions: torch.Tensor) -> torch.Tensor: + # The native 20D layout already IS the canonical 32D block layout + # (left 0..9, right 10..19), just pad. actions = _ensure_bsd(actions) - if actions.shape[-1] != 14: - raise ValueError(f"RobotBimanual: expected 14-dim, got {actions.shape[-1]}") - L, R = actions[..., :7], actions[..., 7:14] - - # left - L_xyz, L_ypr, L_g = L[..., 0:3], L[..., 3:6], L[..., 6:7] - L_R = _ypr_to_matrix(L_ypr) - L_c1, L_c2 = L_R[..., 0], L_R[..., 1] - left_block = torch.cat([L_xyz, L_c1, L_c2, L_g], dim=-1) # (B,S,10) + if actions.shape[-1] != 20: + raise ValueError( + f"RobotBimanual.to32_norm_6d expected 20-dim, got {actions.shape[-1]}" + ) + return _pad32(actions) - # right - R_xyz, R_ypr, R_g = R[..., 0:3], R[..., 3:6], R[..., 6:7] - R_R = _ypr_to_matrix(R_ypr) - R_c1, R_c2 = R_R[..., 0], R_R[..., 1] - right_block = torch.cat([R_xyz, R_c1, R_c2, R_g], dim=-1) # (B,S,10) - - return _pad32(torch.cat([left_block, right_block], dim=-1)) # (B,S,20+) -> 32 - - def from32(self, actions32: torch.Tensor) -> torch.Tensor: + def from32_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: actions32 = _ensure_bsd(actions32) - Lb = actions32[..., 0:10] - Rb = actions32[..., 10:20] - - # left - L_xyz, L_c1, L_c2, L_g = Lb[..., 0:3], Lb[..., 3:6], Lb[..., 6:9], Lb[..., 9:10] - L_R = _reconstruct_R_from_cols(L_c1, L_c2) - L_ypr = _matrix_to_ypr(L_R) - - # right - R_xyz, R_c1, R_c2, R_g = Rb[..., 0:3], Rb[..., 3:6], Rb[..., 6:9], Rb[..., 9:10] - R_R = _reconstruct_R_from_cols(R_c1, R_c2) - R_ypr = _matrix_to_ypr(R_R) - - L7 = torch.cat([L_xyz, L_ypr, L_g], dim=-1) - R7 = torch.cat([R_xyz, R_ypr, R_g], dim=-1) - return torch.cat([L7, R7], dim=-1) # (B,S,14) + if actions32.shape[-1] < 20: + raise ValueError( + f"RobotBimanual.from32_norm_6d expected >=20 dims, got " + f"{actions32.shape[-1]}" + ) + return actions32[..., 0:20] # ============================================================ @@ -303,93 +245,34 @@ def from32(self, actions32: torch.Tensor) -> torch.Tensor: # ============================================================ -class HumanLeftCartesianEuler(BaseActionConverter): - """ - Input orig: (B,S,6) = [x,y,z,yaw,pitch,roll] - 32-pack: left block 0..9 with g=0 - """ - - def to32(self, actions: torch.Tensor) -> torch.Tensor: - actions = _ensure_bsd(actions) - if actions.shape[-1] != 6: - raise ValueError(f"HumanLeft: expected 6-dim, got {actions.shape[-1]}") - xyz, ypr = actions[..., 0:3], actions[..., 3:6] - R = _ypr_to_matrix(ypr) - c1, c2 = R[..., 0], R[..., 1] - g0 = torch.zeros_like(xyz[..., :1]) - block = torch.cat([xyz, c1, c2, g0], dim=-1) # (B,S,10) - return _pad32(block) - - def from32(self, actions32: torch.Tensor) -> torch.Tensor: - actions32 = _ensure_bsd(actions32) - block = actions32[..., 0:10] - xyz, c1, c2 = block[..., 0:3], block[..., 3:6], block[..., 6:9] - R = _reconstruct_R_from_cols(c1, c2) - ypr = _matrix_to_ypr(R) - return torch.cat([xyz, ypr], dim=-1) # (B,S,6) - - -class HumanRightCartesianEuler(BaseActionConverter): - """ - Input orig: (B,S,6) - 32-pack: zeros 0..9, right block 10..19 with g=0 in block - """ - - def to32(self, actions: torch.Tensor) -> torch.Tensor: - actions = _ensure_bsd(actions) - if actions.shape[-1] != 6: - raise ValueError(f"HumanRight: expected 6-dim, got {actions.shape[-1]}") - xyz, ypr = actions[..., 0:3], actions[..., 3:6] - R = _ypr_to_matrix(ypr) - c1, c2 = R[..., 0], R[..., 1] - g0 = torch.zeros_like(xyz[..., :1]) - block = torch.cat([xyz, c1, c2, g0], dim=-1) # (B,S,10) - zeros = torch.zeros_like(block) - return _pad32(torch.cat([zeros, block], dim=-1)) - - def from32(self, actions32: torch.Tensor) -> torch.Tensor: - actions32 = _ensure_bsd(actions32) - block = actions32[..., 10:20] - xyz, c1, c2 = block[..., 0:3], block[..., 3:6], block[..., 6:9] - R = _reconstruct_R_from_cols(c1, c2) - ypr = _matrix_to_ypr(R) - return torch.cat([xyz, ypr], dim=-1) # (B,S,6) - - class HumanBimanualCartesianEuler(BaseActionConverter): """ - Input orig: (B,S,12) = L6 | R6 - 32-pack: left 0..9 (g=0), right 10..19 (g=0) + 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 """ - def to32(self, actions: torch.Tensor) -> torch.Tensor: + def to32_norm_6d(self, actions: torch.Tensor) -> torch.Tensor: + # Human has no gripper, so insert a zero gripper slot at the end of + # each arm block to match the 32D block layout [xyz(3) c1(3) c2(3) g(1)] x 2. actions = _ensure_bsd(actions) - if actions.shape[-1] != 12: - raise ValueError(f"HumanBimanual: expected 12-dim, got {actions.shape[-1]}") - L, R = actions[..., :6], actions[..., 6:12] - - L_xyz, L_ypr = L[..., 0:3], L[..., 3:6] - L_R = _ypr_to_matrix(L_ypr) - L_c1, L_c2 = L_R[..., 0], L_R[..., 1] - g0L = torch.zeros_like(L_xyz[..., :1]) - Lblock = torch.cat([L_xyz, L_c1, L_c2, g0L], dim=-1) # (B,S,10) - - R_xyz, R_ypr = R[..., 0:3], R[..., 3:6] - R_R = _ypr_to_matrix(R_ypr) - R_c1, R_c2 = R_R[..., 0], R_R[..., 1] - g0R = torch.zeros_like(R_xyz[..., :1]) - Rblock = torch.cat([R_xyz, R_c1, R_c2, g0R], dim=-1) # (B,S,10) - + if actions.shape[-1] != 18: + raise ValueError( + f"HumanBimanual.to32_norm_6d expected 18-dim, got {actions.shape[-1]}" + ) + L = actions[..., 0:9] + R = actions[..., 9:18] + g0 = torch.zeros_like(actions[..., :1]) + Lblock = torch.cat([L, g0], dim=-1) # (B,S,10) + Rblock = torch.cat([R, g0], dim=-1) # (B,S,10) return _pad32(torch.cat([Lblock, Rblock], dim=-1)) - def from32(self, actions32: torch.Tensor) -> torch.Tensor: + def from32_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: actions32 = _ensure_bsd(actions32) - Lb = actions32[..., 0:10] - Rb = actions32[..., 10:20] - L_xyz, L_c1, L_c2 = Lb[..., 0:3], Lb[..., 3:6], Lb[..., 6:9] - R_xyz, R_c1, R_c2 = Rb[..., 0:3], Rb[..., 3:6], Rb[..., 6:9] - L_R = _reconstruct_R_from_cols(L_c1, L_c2) - L_ypr = _matrix_to_ypr(L_R) - R_R = _reconstruct_R_from_cols(R_c1, R_c2) - R_ypr = _matrix_to_ypr(R_R) - return torch.cat([L_xyz, L_ypr, R_xyz, R_ypr], dim=-1) # (B,S,12) + if actions32.shape[-1] < 20: + raise ValueError( + f"HumanBimanual.from32_norm_6d expected >=20 dims, got " + f"{actions32.shape[-1]}" + ) + 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) diff --git a/tests/fixtures/train_harness.py b/tests/fixtures/train_harness.py index 9d7e3828a9..1dc9db0db4 100644 --- a/tests/fixtures/train_harness.py +++ b/tests/fixtures/train_harness.py @@ -14,7 +14,9 @@ import egomimic.trainHydra as train_hydra from fixtures.synthetic_episodes import write_episode -STATE_DIM = {"eva_bimanual": 14, "human_bimanual": 12} +# 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} PI_TOKENIZER = ( "google/paligemma-3b-mix-224" # tokenizer_model_name in model/pi0.5_base.yaml ) diff --git a/tests/unit/test_data_configs_compose.py b/tests/unit/test_data_configs_compose.py index 8fcc8fead0..143bda48a8 100644 --- a/tests/unit/test_data_configs_compose.py +++ b/tests/unit/test_data_configs_compose.py @@ -13,6 +13,7 @@ from omegaconf import OmegaConf import egomimic.hydra_configs as _cfg_pkg +from egomimic.rldb.embodiment.embodiment import LEGACY_ROTATION_MODES from egomimic.rldb.embodiment.human import Human DATA_DIR = Path(_cfg_pkg.__file__).parent / "data" @@ -103,6 +104,26 @@ def test_scale_loads_head_pose(): assert "obs_head_pose" in Human.get_keymap(keymap_mode="cartesian") +@pytest.mark.parametrize("name", _params(KNOWN_BROKEN_COMPOSE)) +def test_no_data_config_trains_on_a_discontinuous_rotation(name): + # The 6D representation is the training target everywhere: Euler angles wrap + # at +-pi and gimbal-lock, quaternions double-cover, so per-dim + # normalization and regression losses on them are ill-defined. The guard in + # get_transform_list is what enforces it; this stops a config from opting + # out of the guard. + cfg = compose_data(name) + for split in ("train_datasets", "valid_datasets"): + for emb, node in (cfg.data.get(split) or {}).items(): + if node is None: + continue + transform_list = node.resolver.get("transform_list") + if not hasattr(transform_list, "get"): + continue + mode = transform_list.get("mode") + assert mode not in LEGACY_ROTATION_MODES, (name, split, emb, mode) + assert not transform_list.get("allow_legacy_rotation"), (name, split, emb) + + @pytest.mark.parametrize("name", _params(KNOWN_BROKEN_COMPOSE)) def test_no_data_config_uses_a_pi_keymap_mode(name): # Camera naming is the Pi wrapper's job (PI_CAMERA_SLOTS); a `_pi` keymap diff --git a/tests/unit/test_lazy_imports.py b/tests/unit/test_lazy_imports.py index 089230f075..b6ee19300d 100644 --- a/tests/unit/test_lazy_imports.py +++ b/tests/unit/test_lazy_imports.py @@ -17,7 +17,7 @@ def test_local_path_imports_without_cloud_packages(tmp_path) -> None: from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver from egomimic.rldb.embodiment.human import Human r = LocalEpisodeResolver({str(tmp_path)!r}, key_map=Human.get_keymap(keymap_mode="cartesian"), - transform_list=Human.get_transform_list(mode="cartesian")) + transform_list=Human.get_transform_list(mode="cartesian", allow_legacy_rotation=True)) ds = r.resolve(filters=DatasetFilter(episode_hashes=["aria_00"])) assert set(ds) == {{"aria_00"}}, ds print("ok") diff --git a/tests/unit/test_model_dims.py b/tests/unit/test_model_dims.py index 21b2e0c0a0..c68f993d2b 100644 --- a/tests/unit/test_model_dims.py +++ b/tests/unit/test_model_dims.py @@ -44,14 +44,12 @@ def _target(name): # {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": (14, 14)}), - "hpt_bc_flow_aria": (None, {"human_bimanual": (12, 12)}), - "hpt_bc_flow_human": (None, {"human_bimanual": (12, 12)}), - "hpt_bc_flow_mecka": (None, {"human_bimanual": (12, 12)}), - "hpt_bc_flow_scale": (None, {"human_bimanual": (12, 12)}), + "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)}), - "hpt_bc_pickplace_qwen_pertoken": (None, {"eva_bimanual": (14, 14)}), - "hpt_bc_pickplace_qwen_pooled": (None, {"eva_bimanual": (14, 14)}), "hpt_cotrain_enc_dec_base": (None, COTRAIN), "hpt_cotrain_flow_seperate_head": (None, COTRAIN), "hpt_cotrain_flow_shared_head": SHARED, diff --git a/tests/unit/test_norm_stats_finite.py b/tests/unit/test_norm_stats_finite.py index 9a5565ff60..ac4dacc9af 100644 --- a/tests/unit/test_norm_stats_finite.py +++ b/tests/unit/test_norm_stats_finite.py @@ -49,7 +49,9 @@ def test_infer_norm_from_dataset_ignores_episode_with_nan_gripper(tmp_path, capl resolver = LocalEpisodeResolver( tmp_path, key_map=Eva.get_keymap(keymap_mode="cartesian", norm_mode=True), - transform_list=Eva.get_transform_list(mode="cartesian"), + transform_list=Eva.get_transform_list( + mode="cartesian", allow_legacy_rotation=True + ), ) ds = MultiDataset._from_resolver(resolver, mode="total") stats = MultiDataset(state={}, norm_mode="quantile") diff --git a/tests/unit/test_pi05_norm_rot6d.py b/tests/unit/test_pi05_norm_rot6d.py new file mode 100644 index 0000000000..58598728ca --- /dev/null +++ b/tests/unit/test_pi05_norm_rot6d.py @@ -0,0 +1,252 @@ +"""Round-trip tests for the normalized continuous-6D rotation encoding. + +Covers the data transform (ypr <-> 6D) and the converter packers +(``to32_norm_6d`` / ``from32_norm_6d``) for both the robot bimanual (14D ypr / +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. +""" + +import math + +import numpy as np +import pytest +import torch + +from egomimic.rldb.zarr.action_chunk_transforms import ( + CartesianRot6DToYPR, + CartesianYPRToRot6D, +) +from egomimic.utils.action_utils import ( + BaseActionConverter, + RobotBimanualCartesianEuler, +) + + +def _eva_ypr_chunk(T: int = 5) -> np.ndarray: + # [L xyz ypr g, R xyz ypr g]; moderate angles to avoid gimbal/wrap ambiguity. + rng = np.random.default_rng(0) + xyz = rng.uniform(-1.0, 1.0, size=(T, 3)) + ypr = rng.uniform(-1.0, 1.0, size=(T, 3)) # radians, well inside (-pi, pi) + g = rng.uniform(0.0, 1.0, size=(T, 1)) + arm = np.concatenate([xyz, ypr, g], axis=-1) + return np.concatenate([arm, arm], axis=-1) # 14D + + +def _aria_ypr_chunk(T: int = 5) -> np.ndarray: + rng = np.random.default_rng(1) + xyz = rng.uniform(-1.0, 1.0, size=(T, 3)) + ypr = rng.uniform(-1.0, 1.0, size=(T, 3)) + arm = np.concatenate([xyz, ypr], axis=-1) + return np.concatenate([arm, arm], axis=-1) # 12D + + +@pytest.mark.parametrize( + "chunk_fn,ypr_dim,six_dim", + [(_eva_ypr_chunk, 14, 20), (_aria_ypr_chunk, 12, 18)], +) +def test_cartesian_ypr_rot6d_transform_round_trips(chunk_fn, ypr_dim, six_dim): + ypr = chunk_fn() + assert ypr.shape[-1] == ypr_dim + + fwd = CartesianYPRToRot6D(action_key="actions_cartesian") + rev = CartesianRot6DToYPR(action_key="actions_cartesian") + + batch = {"actions_cartesian": ypr.copy()} + batch = fwd.transform(batch) + assert batch["actions_cartesian"].shape[-1] == six_dim + + batch = rev.transform(batch) + np.testing.assert_allclose(batch["actions_cartesian"], ypr, atol=1e-6) + + +def test_ypr_to_rot6d_transform_is_continuous_at_yaw_wrap(): + # The whole point of encoding rotation as 6D BEFORE normalization: a yaw + # of +pi-eps and -pi+eps are the same orientation. As ypr they differ by + # ~2pi (and a normalized-ypr target would jump); as 6D columns they are + # within O(eps) of each other, so the regression target is continuous. + eps = 1e-3 + a = np.zeros((1, 12)) + b = np.zeros((1, 12)) + a[0, 3], b[0, 3] = math.pi - eps, -math.pi + eps # L yaw + a[0, 9], b[0, 9] = math.pi - eps, -math.pi + eps # R yaw + fwd = CartesianYPRToRot6D(action_key="actions_cartesian") + six_a = fwd.transform({"actions_cartesian": a.copy()})["actions_cartesian"] + six_b = fwd.transform({"actions_cartesian": b.copy()})["actions_cartesian"] + assert np.abs(a - b).max() > 6.0, "ypr inputs straddle the wrap" + assert np.abs(six_a - six_b).max() < 5 * eps, "6D columns must not jump" + + +def test_transform_preserves_tensor_type(): + ypr = torch.from_numpy(_eva_ypr_chunk()) + out = CartesianYPRToRot6D().transform({"actions_cartesian": ypr})[ + "actions_cartesian" + ] + assert isinstance(out, torch.Tensor) + assert out.shape[-1] == 20 + + +def test_robot_bimanual_norm_6d_pack_round_trips(): + converter = RobotBimanualCartesianEuler() + six6d = torch.from_numpy( + CartesianYPRToRot6D().transform({"actions_cartesian": _eva_ypr_chunk()})[ + "actions_cartesian" + ] + ).float()[None] # (1, T, 20) + + packed = converter.to32_norm_6d(six6d) + assert packed.shape[-1] == 32 + decoded = converter.from32_norm_6d(packed) + torch.testing.assert_close(decoded, six6d, atol=1e-6, rtol=1e-6) + + +@pytest.mark.parametrize( + "chunk_fn,ypr_dim,six_dim", + [(_eva_ypr_chunk, 14, 20), (_aria_ypr_chunk, 12, 18)], +) +def test_proprio_pose_vector_round_trips(chunk_fn, ypr_dim, six_dim): + # The proprio ee_pose is a single pose vector (D,) with the same per-arm + # layout as one action row; the same transforms must handle it. + pose = chunk_fn(T=1)[0] + assert pose.shape == (ypr_dim,) + + fwd = CartesianYPRToRot6D(action_key="observations.state.ee_pose") + rev = CartesianRot6DToYPR(action_key="observations.state.ee_pose") + + batch = {"observations.state.ee_pose": pose.copy()} + batch = fwd.transform(batch) + assert batch["observations.state.ee_pose"].shape == (six_dim,) + + batch = rev.transform(batch) + np.testing.assert_allclose(batch["observations.state.ee_pose"], pose, atol=1e-6) + + +def _keys_of(transforms, cls): + return {t.action_key for t in transforms if isinstance(t, cls)} + + +# cartesian_wristframe_6d is covered by the test_wrist6d_roundtrip pipeline tests. +@pytest.mark.parametrize("mode", ["cartesian_6d"]) +def test_6d_modes_convert_action_and_proprio(mode): + from egomimic.rldb.embodiment.eva import Eva + from egomimic.rldb.embodiment.human import Human + + for cls in (Eva, Human): + transform_list = cls.get_transform_list(mode) + assert _keys_of(transform_list, CartesianYPRToRot6D) == { + "actions_cartesian", + "observations.state.ee_pose", + }, f"{cls.__name__} {mode} must 6D-encode both action and proprio" + + +@pytest.mark.parametrize( + "cls_name,mode", + [ + ("Eva", "cartesian"), + ("Eva", "cartesian_wristframe_ypr"), + ("Eva", "cartesian_wristframe_quat"), + ("Human", "cartesian"), + ("Human", "cartesian_padded"), + ("Human", "cartesian_wristframe_ypr"), + ], +) +def test_legacy_rotation_modes_need_an_explicit_opt_in(cls_name, mode): + from egomimic.rldb.embodiment.eva import Eva + from egomimic.rldb.embodiment.human import Human + + cls = {"Eva": Eva, "Human": Human}[cls_name] + with pytest.raises(ValueError, match="discontinuous rotation"): + cls.get_transform_list(mode) + assert cls.get_transform_list(mode, allow_legacy_rotation=True) + + +def test_6d_revert_lists_revert_proprio(): + from egomimic.rldb.embodiment.eva import ( + _build_eva_cartesian_revert_6d_transform_list, + _build_eva_cartesian_revert_6d_wristframe_transform_list, + ) + from egomimic.rldb.embodiment.human import ( + _build_human_cartesian_revert_6d_transform_list, + _build_human_cartesian_revert_6d_wristframe_transform_list, + ) + + for build in ( + _build_eva_cartesian_revert_6d_transform_list, + _build_eva_cartesian_revert_6d_wristframe_transform_list, + _build_human_cartesian_revert_6d_transform_list, + _build_human_cartesian_revert_6d_wristframe_transform_list, + ): + transform_list = build() + assert _keys_of(transform_list, CartesianRot6DToYPR) == { + "actions_cartesian", + "observations.state.ee_pose", + }, 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__ diff --git a/tests/unit/test_pins.py b/tests/unit/test_pins.py index db0f4090d7..2e65cf4ec3 100644 --- a/tests/unit/test_pins.py +++ b/tests/unit/test_pins.py @@ -126,6 +126,16 @@ def test_local_pin_embodiment_mismatch_is_error(tmp_path) -> None: ) +def test_local_pin_accepts_legacy_vendor_embodiment(tmp_path) -> None: + _attrs_only_episode(tmp_path, "l-mecka", embodiment="MECKA_BIMANUAL") + paths = LocalEpisodeResolver._get_local_filtered_paths( + tmp_path, + DatasetFilter(episode_hashes=["l-mecka"]), + expected_embodiment="human_bimanual", + ) + assert [h for _, h in paths] == ["l-mecka"] + + def test_local_unpinned_mismatch_warns(tmp_path, caplog) -> None: _attrs_only_episode(tmp_path, "l-eva", embodiment="eva_bimanual") _attrs_only_episode(tmp_path, "l-hum", embodiment="human_bimanual") diff --git a/tests/unit/test_synthetic_episodes.py b/tests/unit/test_synthetic_episodes.py index c92a0d26b1..116c6531c5 100644 --- a/tests/unit/test_synthetic_episodes.py +++ b/tests/unit/test_synthetic_episodes.py @@ -18,12 +18,16 @@ def test_episode_loads_with_vendor_recipe(tmp_path, vendor): write_episode(tmp_path, vendor, seed=i) if v.embodiment == "eva_bimanual": key_map = Eva.get_keymap(keymap_mode="cartesian", annotation_key="annotations") - transforms = Eva.get_transform_list(mode="cartesian") + transforms = Eva.get_transform_list( + mode="cartesian", allow_legacy_rotation=True + ) else: key_map = Human.get_keymap( keymap_mode="cartesian", annotation_key="annotations" ) - transforms = Human.get_transform_list(mode="cartesian", stride=v.stride) + transforms = Human.get_transform_list( + mode="cartesian", stride=v.stride, allow_legacy_rotation=True + ) resolver = LocalEpisodeResolver( tmp_path, key_map=key_map, transform_list=transforms )