diff --git a/.gitignore b/.gitignore index 6a9d43c83..4fd335414 100644 --- a/.gitignore +++ b/.gitignore @@ -45,7 +45,18 @@ annotations_test/** log_conversion/** debug_model_inputs/** temp_dir/** +datasets/ +**/datasets/ +apikey.txt +slurm-*.out +slurmoutputs/ +*.log +.inductor_cache/ + scratch/ sample_data/ external_ckpts/ external/MANO/ +ann_backup/ +aria_ann_backup/ +sort_annotations/ \ No newline at end of file diff --git a/egomimic/algo/pi.py b/egomimic/algo/pi.py index a0d856415..198772a6e 100644 --- a/egomimic/algo/pi.py +++ b/egomimic/algo/pi.py @@ -26,9 +26,10 @@ ) from egomimic.rldb.embodiment.embodiment import get_embodiment, get_embodiment_id from egomimic.utils.action_utils import ( - ConverterRegistry, PI05_CARTESIAN_ACTION_ENCODING_LEGACY, + PI05_CARTESIAN_ACTION_ENCODING_NORM_ROT_6D, PI05_CARTESIAN_ACTION_ENCODING_RAW_ROT_6D, + ConverterRegistry, ) logger = logging.getLogger(__name__) @@ -308,7 +309,9 @@ def _action_stats(self, embodiment_id: int, ac_key: str) -> dict: f"and embodiment id {embodiment_id}" ) from exc - def _unnormalize_action(self, action: torch.Tensor, embodiment_id: int, ac_key: str): + def _unnormalize_action( + self, action: torch.Tensor, embodiment_id: int, ac_key: str + ): return self.norm_stats.unnormalize( {ac_key: action.clone(), "embodiment": embodiment_id}, embodiment_id, @@ -469,15 +472,14 @@ def forward_eval(self, batch): num_steps=self.num_steps, ) + pred_actions = pred_actions.clone() + predictions = OrderedDict() ref = _batch[ac_key] B, T, D = ref.shape converter = self.action_registry.get(embodiment_id, ac_key) - if ( - self.action_encoding - == PI05_CARTESIAN_ACTION_ENCODING_RAW_ROT_6D - ): + if self.action_encoding == PI05_CARTESIAN_ACTION_ENCODING_RAW_ROT_6D: pred_actions_orig = converter.from32_raw_rotation( pred_actions, stats=self._action_stats(embodiment_id, ac_key), @@ -485,6 +487,15 @@ def forward_eval(self, batch): unnormalize_non_rotation=True, ) unnorm_actions = {ac_key: pred_actions_orig[:, :T, :D]} + elif self.action_encoding == PI05_CARTESIAN_ACTION_ENCODING_NORM_ROT_6D: + # Extract the normalized xyz+6D(+gripper) action, then + # unnormalize via the standard pipeline (stats were computed + # over the 6D representation) to get raw 6D actions. + pred_6d = converter.from32_norm_6d(pred_actions) + predictions[ac_key] = pred_6d[:, :T, :D] + unnorm_actions = self.norm_stats.unnormalize( + predictions, embodiment_id + ) elif self.action_encoding == PI05_CARTESIAN_ACTION_ENCODING_LEGACY: pred_actions_orig = converter.from32(pred_actions) pred = pred_actions_orig[:, :T, :D] @@ -578,6 +589,11 @@ def _robomimic_to_pi_data( stats=self._action_stats(emb_id, ac_key), norm_mode=self.norm_stats.norm_mode, ) + elif self.action_encoding == PI05_CARTESIAN_ACTION_ENCODING_NORM_ROT_6D: + # Action is already a normalized xyz+6D(+gripper) chunk (the + # ypr->6D conversion happened in the CartesianYPRToRot6D data + # transform). Just pack it into the 32D vector. + action32 = converter.to32_norm_6d(action) elif self.action_encoding == PI05_CARTESIAN_ACTION_ENCODING_LEGACY: action32 = converter.to32(action) else: diff --git a/egomimic/hydra_configs/callbacks/checkpoints.yaml b/egomimic/hydra_configs/callbacks/checkpoints.yaml index 86d21a792..a73448099 100644 --- a/egomimic/hydra_configs/callbacks/checkpoints.yaml +++ b/egomimic/hydra_configs/callbacks/checkpoints.yaml @@ -8,4 +8,5 @@ model_checkpoint: filename: "epoch_{epoch}" save_last: true save_top_k: -1 - every_n_epochs: 100 \ No newline at end of file + every_n_epochs: 50 + save_on_train_epoch_end: true \ No newline at end of file diff --git a/egomimic/hydra_configs/data/cotrain_pi_base.yaml b/egomimic/hydra_configs/data/cotrain_pi_base.yaml index 6d682e2fd..e8ce9b86c 100644 --- a/egomimic/hydra_configs/data/cotrain_pi_base.yaml +++ b/egomimic/hydra_configs/data/cotrain_pi_base.yaml @@ -35,17 +35,17 @@ train_datasets: valid_datasets: eva_bimanual: - _target_: ${train_datasets.eva_bimanual._target_} - resolver: ${train_datasets.eva_bimanual.resolver} - filters: ${train_datasets.eva_bimanual.filters} + _target_: ${...train_datasets.eva_bimanual._target_} + resolver: ${...train_datasets.eva_bimanual.resolver} + filters: ${...train_datasets.eva_bimanual.filters} mode: valid - valid_ratio: ${train_datasets.eva_bimanual.valid_ratio} + valid_ratio: ${...train_datasets.eva_bimanual.valid_ratio} human_bimanual: - _target_: ${train_datasets.human_bimanual._target_} - resolver: ${train_datasets.human_bimanual.resolver} - filters: ${train_datasets.human_bimanual.filters} + _target_: ${...train_datasets.human_bimanual._target_} + resolver: ${...train_datasets.human_bimanual.resolver} + filters: ${...train_datasets.human_bimanual.filters} mode: valid - valid_ratio: ${train_datasets.human_bimanual.valid_ratio} + valid_ratio: ${...train_datasets.human_bimanual.valid_ratio} train_dataloader_params: eva_bimanual: diff --git a/egomimic/hydra_configs/data/cotrain_pi_lang.yaml b/egomimic/hydra_configs/data/cotrain_pi_lang.yaml index 430252b5d..26acecfb0 100644 --- a/egomimic/hydra_configs/data/cotrain_pi_lang.yaml +++ b/egomimic/hydra_configs/data/cotrain_pi_lang.yaml @@ -13,6 +13,9 @@ train_datasets: filter_lambdas: - "lambda row: (row['embodiment'] == 'eva_bimanual') & (row['task'] == 'pick_place') & (row['zarr_processed_path'] != '')" human_bimanual: + resolver: + key_map: + keymap_mode: cartesian_pi filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter project_name: "dense-language" diff --git a/egomimic/hydra_configs/data/cotrain_pi_lang_6d.yaml b/egomimic/hydra_configs/data/cotrain_pi_lang_6d.yaml new file mode 100644 index 000000000..960863461 --- /dev/null +++ b/egomimic/hydra_configs/data/cotrain_pi_lang_6d.yaml @@ -0,0 +1,19 @@ +defaults: + - cotrain_pi_lang + - _self_ + +# Same dense-language cotrain data as `cotrain_pi_lang`, but the cartesian action +# chunk is expressed with the continuous 6D rotation representation +# (xyz+6D per arm, +gripper for eva) instead of xyz+ypr. Pairs with the +# `cartesian_normalized_rot6d` action_encoding (model=pi0.5_cotrain_eva_aria_6d). +# Valid datasets inherit the train resolver via `${...}` interpolation in +# cotrain_pi_base, so overriding the train transform mode is sufficient. +train_datasets: + eva_bimanual: + resolver: + transform_list: + mode: cartesian_6d + human_bimanual: + resolver: + transform_list: + mode: cartesian_6d diff --git a/egomimic/hydra_configs/data/mecka_all_pi_6d.yaml b/egomimic/hydra_configs/data/mecka_all_pi_6d.yaml new file mode 100644 index 000000000..4a4d4dbe9 --- /dev/null +++ b/egomimic/hydra_configs/data/mecka_all_pi_6d.yaml @@ -0,0 +1,72 @@ +# MECKA_ALL_PI_6D — pi0.5 on ALL mecka episodes in the SQL table (filters +# lab=mecka with no task filter), stored under the human_bimanual embodiment. +# +# WRIST-frame cartesian actions with the continuous 6D rotation +# representation (18D = xyz+6D per arm, no gripper): each arm's action chunk +# is expressed relative to that wrist's current pose, so both the arbitrary +# world frame AND the moving head frame cancel out of the action targets. The +# proprio ee_pose stays in the head/cam frame (it defines where the wrists +# are for the model and for the eval revert) and is 6D-encoded. stride=1: +# mecka actions are recorded at the target rate (Aria uses 3). Language +# prompts are sampled from the `annotations` track; prompt assembly +# (proprio-in-prompt on, embodiment label off) lives in +# model=pi0.5_bc_mecka_6d, which this config pairs with, alongside +# evaluator=eval_pi_wristframe_6d (reverts wrist -> head frame via the +# proprio, then 6D -> ypr, for cam-frame MSE + viz). +# +# Validation/checkpoint cadence stays the epoch-based trainer default — no +# step-based val_check_interval / every_n_train_steps overrides here. + +_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper + +train_datasets: + human_bimanual: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver + # Shared PACE mirror of processed_v3 (55k episodes; 41,613/41,617 SQL + # mecka episodes already present). The resolver s5cmd-syncs only the + # SQL-matched episodes missing here (a handful), not the full set. + folder_path: /storage/project/r-dxu345-0/shared/egoverseS3ZarrDatasets + key_map: + _target_: egomimic.rldb.embodiment.human.Human.get_keymap + # "_pi" suffix -> PaliGemma-style camera names (base_0_rgb); PI's + # _fill_missing_images duplicates base into the absent wrist slots. + keymap_mode: cartesian_pi + annotation_key: annotations + transform_list: + _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + mode: cartesian_wristframe_6d + stride: 1 + # processed_v3 mecka zarrs predate the mecka_to_zarr rot_left fix: + # the LEFT wrist frame was double-mirrored onto the right hand's + # spatial convention. Rz(180°)-corrects the raw left pose keys before + # any frame math (exactly equivalent to reconverting). + fix_mecka_left_wrist: true + # 18 -> 20-dim proprio (zero grip slots at 9/19): State: prompt bins + # align positionally with the robot 20-dim layout. + pad_proprio_gripper: true + filters: + _target_: egomimic.rldb.filters.DatasetFilter + filter_lambdas: + - "lambda row: str(row.get('lab', '')).lower() == 'mecka'" + mode: train + valid_ratio: 0.05 + +valid_datasets: + human_bimanual: + _target_: ${...train_datasets.human_bimanual._target_} + resolver: ${...train_datasets.human_bimanual.resolver} + filters: ${...train_datasets.human_bimanual.filters} + mode: valid + valid_ratio: ${...train_datasets.human_bimanual.valid_ratio} + +train_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 6 + +valid_dataloader_params: + human_bimanual: + batch_size: 64 + num_workers: 6 diff --git a/egomimic/hydra_configs/data/obj_gen_pi_lang.yaml b/egomimic/hydra_configs/data/obj_gen_pi_lang.yaml new file mode 100644 index 000000000..08651b1dd --- /dev/null +++ b/egomimic/hydra_configs/data/obj_gen_pi_lang.yaml @@ -0,0 +1,33 @@ +defaults: + - cotrain_pi_base + - _self_ + +# Motion-generalization cotrain: annotated-only eva + annotated-only aria +# (base/object descriptions). PI-style camera keys for both embodiments. + +train_datasets: + eva_bimanual: + resolver: + folder_path: /storage/home/hcoda1/5/agao81/r-dxu345-0/pick_place + key_map: + keymap_mode: cartesian_pi + 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()) and (row.get('episode_hash') or '') != '2026-04-22-02-30-32-296000'" + aria_bimanual: + resolver: + folder_path: /storage/home/hcoda1/5/agao81/r-dxu345-0/pick_place + key_map: + keymap_mode: cartesian_pi + filters: + _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter + project_name: "dense-language" + filter_lambdas: + # Exclude these aria episodes: every base_0_rgb JPEG decode fails + # (entire episode is corrupted), exhausts random-retry budget at train time. + # 2026-04-26-00-17-53-000000 + # 2026-04-26-00-27-26-000000 + # 2026-05-01-02-52-58-000000 + - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'aria_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')" diff --git a/egomimic/hydra_configs/data/obj_gen_pi_lang_6d.yaml b/egomimic/hydra_configs/data/obj_gen_pi_lang_6d.yaml new file mode 100644 index 000000000..daf77f990 --- /dev/null +++ b/egomimic/hydra_configs/data/obj_gen_pi_lang_6d.yaml @@ -0,0 +1,25 @@ +defaults: + - obj_gen_pi_lang + - _self_ + +# 6D-rotation variant of the motion-generalization cotrain (obj_gen_pi_lang): +# - cartesian action chunk expressed with the continuous 6D rotation +# representation (xyz+6D per arm, +gripper for eva) via the `cartesian_6d` +# transform mode. Pairs with model=pi0.5_cotrain_eva_aria_6d +# (action_encoding=cartesian_normalized_rot6d) and evaluator=eval_pi_camframe_6d. +# - dataset switched from the plain Scale-annotation resolver to the +# AnnotationCutoff resolver, which clamps each action chunk at the end of the +# enclosing language-annotation span (ZarrAnnotationCutoffDataset). +# Valid datasets inherit the train resolver via `${...}` interpolation in +# cotrain_pi_base, so overriding the train resolver is sufficient. +train_datasets: + eva_bimanual: + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver + transform_list: + mode: cartesian_6d + aria_bimanual: + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver + transform_list: + mode: cartesian_6d diff --git a/egomimic/hydra_configs/data/obj_gen_pi_lang_wristframe.yaml b/egomimic/hydra_configs/data/obj_gen_pi_lang_wristframe.yaml new file mode 100644 index 000000000..c5d58a3e5 --- /dev/null +++ b/egomimic/hydra_configs/data/obj_gen_pi_lang_wristframe.yaml @@ -0,0 +1,40 @@ +defaults: + - cotrain_pi_base + - _self_ + +# Wristframe variant of motion_gen_pi_lang: annotated-only eva + annotated-only +# aria (base/object descriptions). Actions are expressed in each wrist's own +# frame (cartesian_wristframe_ypr) instead of the head/camera frame. Pair with +# evaluator=eval_pi so the revert transform projects predictions back to cam +# frame for the viz video. + +train_datasets: + eva_bimanual: + resolver: + folder_path: /storage/home/hcoda1/5/agao81/r-dxu345-0/pick_place + key_map: + keymap_mode: cartesian_pi + 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()) and (row.get('episode_hash') or '') != '2026-04-22-02-30-32-296000'" + aria_bimanual: + resolver: + folder_path: /storage/home/hcoda1/5/agao81/r-dxu345-0/pick_place + key_map: + keymap_mode: cartesian_pi + transform_list: + mode: cartesian_wristframe_ypr + filters: + _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter + project_name: "dense-language" + filter_lambdas: + # Exclude these aria episodes: every base_0_rgb JPEG decode fails + # (entire episode is corrupted), exhausts random-retry budget at train time. + # 2026-04-26-00-17-53-000000 + # 2026-04-26-00-27-26-000000 + # 2026-05-01-02-52-58-000000 + - "lambda row: row.get('task') == 'pick_place' and row.get('embodiment') == 'aria_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')" diff --git a/egomimic/hydra_configs/data/obj_gen_pi_lang_wristframe_6d.yaml b/egomimic/hydra_configs/data/obj_gen_pi_lang_wristframe_6d.yaml new file mode 100644 index 000000000..8d6b2121d --- /dev/null +++ b/egomimic/hydra_configs/data/obj_gen_pi_lang_wristframe_6d.yaml @@ -0,0 +1,25 @@ +defaults: + - obj_gen_pi_lang_wristframe + - _self_ + +# 6D-rotation + AnnotationCutoff variant of the wrist-frame motion-generalization +# cotrain (obj_gen_pi_lang_wristframe): +# - actions in each wrist's own frame, rotation expressed as the continuous 6D +# representation via the `cartesian_wristframe_6d` transform mode. Pairs with +# model=pi0.5_cotrain_eva_aria_6d and evaluator=eval_pi_wristframe_6d (which +# un-6Ds then projects wrist-frame preds back to cam frame for viz/MSE). +# - dataset switched to the AnnotationCutoff resolver (clamps each action chunk +# at the end of its enclosing language-annotation span). +# Valid datasets inherit the train resolver via `${...}` interpolation in +# cotrain_pi_base, so overriding the train resolver is sufficient. +train_datasets: + eva_bimanual: + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver + transform_list: + mode: cartesian_wristframe_6d + aria_bimanual: + resolver: + _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3AnnotationCutoffEpisodeResolver + transform_list: + mode: cartesian_wristframe_6d diff --git a/egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml b/egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml new file mode 100644 index 000000000..c5cf468dd --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_pi_camframe_6d.yaml @@ -0,0 +1,18 @@ +defaults: + - viz@viz_func: pi_cartesian_lang + - _self_ + +_target_: egomimic.eval.eval_pi.PIEvalVideo + +# Cam-frame 6D-rotation variant: use when the data config expresses cartesian +# actions in head/camera frame with the continuous 6D rotation representation +# (e.g. a data config with `transform_list: mode: cartesian_6d`). Actions are +# already in cam (head) frame, so no frame change is needed — the revert only +# converts the rotation back from xyz+6D (9/arm) to xyz+ypr (6/arm) so the viz +# video and cam-frame MSE see the same ypr layout as the plain cartesian mode. +# Each value resolves to a list[Transform] via its ``_target_``. +transform_lists: + eva_bimanual: + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_transform_list + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_6d_transform_list diff --git a/egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml b/egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml new file mode 100644 index 000000000..43eb169b6 --- /dev/null +++ b/egomimic/hydra_configs/evaluator/eval_pi_wristframe_6d.yaml @@ -0,0 +1,18 @@ +defaults: + - viz@viz_func: pi_cartesian_lang + - _self_ + +_target_: egomimic.eval.eval_pi.PIEvalVideo + +# Wrist-frame 6D-rotation variant: use with a data config that applies the +# `cartesian_wristframe_6d` transform mode (e.g. obj_gen_pi_lang_wristframe_6d). +# The model predicts in each wrist's local frame using the continuous 6D +# rotation; these revert transforms first convert the rotation xyz+6D -> xyz+ypr +# and then project predictions + gt back to cam (head) frame for the cam-frame +# MSE and the viz video. Each value resolves to a list[Transform] via its +# ``_target_``. Must match the viz config mounted above. +transform_lists: + eva_bimanual: + _target_: egomimic.rldb.embodiment.eva._build_eva_cartesian_revert_6d_wristframe_transform_list + human_bimanual: + _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_6d_wristframe_transform_list diff --git a/egomimic/hydra_configs/model/pi0.5_base.yaml b/egomimic/hydra_configs/model/pi0.5_base.yaml index b39d26e40..60e01e213 100644 --- a/egomimic/hydra_configs/model/pi0.5_base.yaml +++ b/egomimic/hydra_configs/model/pi0.5_base.yaml @@ -33,7 +33,7 @@ robomimic_model: tokenizer_max_length: 128 sampling_mode: "random" annotation_key: "annotations" - default_prompt: "This is a bad action." + default_prompt: "" proprio_in_prompt: true embodiment_label: true state_num_bins: 256 @@ -62,11 +62,13 @@ optimizer: eps: 1e-8 weight_decay: 0.0 +# Constant LR after warmup (matches the remote pi0.5 training stack). The +# previous cosine schedule decayed to 0 over num_training_steps=60000, so runs +# past ~60k steps trained at ~0 lr and shorter runs never saw the configured lr +# for most of training. scheduler: - _target_: transformers.get_cosine_schedule_with_warmup + _target_: transformers.optimization.get_constant_schedule_with_warmup _partial_: true num_warmup_steps: 2000 - num_training_steps: 60000 - num_cycles: 0.5 scheduler_interval: step \ No newline at end of file diff --git a/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d.yaml b/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d.yaml new file mode 100644 index 000000000..c92ba3a13 --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d.yaml @@ -0,0 +1,30 @@ +defaults: + - pi0.5_base + - _self_ + +# pi0.5 BC on mecka data alone (stored as human_bimanual — all human demo data +# is one embodiment locally). Normalized continuous-6D rotation encoding: +# actions arrive already in xyz+6D layout (18D per arm pair, no gripper, +# wrist-relative via the `cartesian_wristframe_6d` transform mode) and the +# head-frame proprio ee_pose arrives 6D-encoded too; the forward pass only +# packs them into the 32D vector. +# Pairs with data=mecka_all_pi_6d and evaluator=eval_pi_wristframe_6d. +robomimic_model: + camera_transforms: {} + ac_keys: + human_bimanual: "actions_cartesian" + domains: ["human_bimanual"] + + action_encoding: "cartesian_normalized_rot6d" + # Keep the discretized-proprio "State: " prompt block (pi0.5_base + # defaults proprio_in_prompt on) but do not splice "Embodiment: " into + # the prompt. + embodiment_label: false + + action_converters: + rules: + HUMAN_BIMANUAL: + _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler + # optional fallback if no match is found + fallback: + _target_: egomimic.utils.action_utils.BaseActionConverter diff --git a/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria_6d.yaml b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria_6d.yaml new file mode 100644 index 000000000..74c9ccba6 --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria_6d.yaml @@ -0,0 +1,14 @@ +defaults: + - pi0.5_cotrain_eva_aria + - _self_ + +# Normalized continuous-6D rotation encoding. Actions arrive already in +# xyz+6D(+gripper) layout (via the `cartesian_6d` transform mode) and normalized +# by the standard pipeline; the forward pass only packs them into the 32D vector. +# The proprio ee_pose arrives 6D-encoded as well (eva 20D, human 18D) — the 6D +# transform modes convert both, since normalized YPR proprio saturates yaw/roll +# at ±π just like the action rotations did. +robomimic_model: + action_encoding: "cartesian_normalized_rot6d" + # Do not splice "Embodiment: " into the prompt (pi0.5_base defaults it on). + embodiment_label: false diff --git a/egomimic/rldb/embodiment/embodiment.py b/egomimic/rldb/embodiment/embodiment.py index 4ae2f8da9..959e5f132 100644 --- a/egomimic/rldb/embodiment/embodiment.py +++ b/egomimic/rldb/embodiment/embodiment.py @@ -50,8 +50,20 @@ 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 get_embodiment_id(embodiment_name): - return EMBODIMENT[embodiment_name.upper()].value + name = embodiment_name.upper() + vendor, _, suffix = name.partition("_") + if vendor in _HUMAN_VENDOR_PREFIXES and suffix: + name = f"HUMAN_{suffix}" + return EMBODIMENT[name].value class Embodiment(ABC): @@ -199,12 +211,22 @@ def viz_gt_preds( pred_action = pred_actions[i] K_i = _intrinsics_from_batch(batch, i) ims = cls.viz( - image, action, mode=mode, color="Greens", alpha=gt_alpha, - intrinsics=K_i, **kwargs + image, + action, + mode=mode, + color="Greens", + alpha=gt_alpha, + intrinsics=K_i, + **kwargs, ) ims = cls.viz( - ims, pred_action, mode=mode, color="Reds", alpha=pred_alpha, - intrinsics=K_i, **kwargs + ims, + pred_action, + mode=mode, + color="Reds", + alpha=pred_alpha, + intrinsics=K_i, + **kwargs, ) if annotation_key is not None: ims = cls.viz(ims, [annotations[i]], mode="annotations", **kwargs) diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index 4c9cdbb87..d84733723 100644 --- a/egomimic/rldb/embodiment/eva.py +++ b/egomimic/rldb/embodiment/eva.py @@ -9,6 +9,8 @@ from egomimic.rldb.zarr.action_chunk_transforms import ( ActionChunkCoordinateFrameTransform, BatchQuaternionPoseToYPR, + CartesianRot6DToYPR, + CartesianYPRToRot6D, ConcatKeys, DeleteKeys, InterpolateLinear, @@ -49,13 +51,39 @@ class Eva(Embodiment): @staticmethod def get_transform_list( mode: Literal[ - "cartesian", "cartesian_wristframe_ypr", "cartesian_wristframe_quat" + "cartesian", + "cartesian_6d", + "cartesian_wristframe_ypr", + "cartesian_wristframe_6d", + "cartesian_wristframe_quat", ], ) -> list[Transform]: if mode == "cartesian": return _build_eva_bimanual_transform_list(is_quat=True) + elif mode == "cartesian_6d": + # Camera-frame cartesian (14D xyz+ypr+gripper per arm) with the + # rotation re-expressed as the continuous 6D representation + # (20D xyz+6d+gripper per arm) for pi0.5 normalized-rot6d encoding. + # The proprio ee_pose is 6D-encoded too: normalized YPR proprio + # saturates yaw/roll at ±π (wraparound), so per-dim normalization + # is only meaningful on the continuous rep — same fix as actions. + return _build_eva_bimanual_transform_list(is_quat=True) + [ + CartesianYPRToRot6D(action_key="actions_cartesian"), + CartesianYPRToRot6D(action_key="observations.state.ee_pose"), + ] elif mode == "cartesian_wristframe_ypr": return _build_eva_bimanual_eef_frame_transform_list(is_quat=False) + elif mode == "cartesian_wristframe_6d": + # Wrist-frame cartesian (14D xyz+ypr+gripper per arm) with the + # rotation re-expressed as the continuous 6D representation + # (20D) for pi0.5 normalized-rot6d encoding. The cam-frame proprio + # ee_pose is 6D-encoded too (see cartesian_6d) — extra important + # here since the proprio is the only cam-frame signal the model + # sees with wrist-relative action targets. + return _build_eva_bimanual_eef_frame_transform_list(is_quat=False) + [ + CartesianYPRToRot6D(action_key="actions_cartesian"), + CartesianYPRToRot6D(action_key="observations.state.ee_pose"), + ] elif mode == "cartesian_wristframe_quat": return _build_eva_bimanual_eef_frame_transform_list(is_quat=True) @@ -151,6 +179,49 @@ 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. + + Used by the cam-frame 6D evaluator: the action chunk is already in camera + frame (produced by the ``cartesian_6d`` transform mode), so only the + rotation representation is converted from xyz+6D (+gripper, 10/arm) back to + xyz+ypr (+gripper, 7/arm) so cam-frame MSE and the viz video see the same + ypr layout as the plain ``cartesian`` mode. The proprio ee_pose (also + 6D-encoded by the ``cartesian_6d`` mode) 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. + + Three stages for the cam-frame 6D wristframe evaluator: (1) convert the + action rotation from xyz+6D (+gripper) back to xyz+ypr (+gripper) via + ``CartesianRot6DToYPR`` (Gram-Schmidt re-orthonormalizes the possibly + non-orthonormal model prediction); (2) likewise revert the proprio + ``observations.state.ee_pose`` (6D-encoded by the ``cartesian_wristframe_6d`` + mode) back to ypr; (3) project the wrist-frame ypr actions back into camera + frame using the standard eef-frame revert, which reads 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 673ccc5e0..a3ef5786f 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -1,6 +1,5 @@ from __future__ import annotations -from abc import abstractmethod from typing import Literal import numpy as np @@ -9,6 +8,8 @@ from egomimic.rldb.zarr.action_chunk_transforms import ( ActionChunkCoordinateFrameTransform, BatchQuaternionPoseToYPR, + CartesianRot6DToYPR, + CartesianYPRToRot6D, ConcatKeys, DeleteKeys, InterpolatePose, @@ -16,8 +17,10 @@ PoseCoordinateFrameTransform, QuaternionPoseToYPR, Reshape, + RotateLocalFrame, SplitKeys, Transform, + UnpadGripperZeros, XYZWXYZ_to_XYZYPR, ) from egomimic.utils.viz_utils import ( @@ -26,7 +29,6 @@ _viz_keypoints, ) - ARIA_INTRINSICS = np.array( [ [133.25430222 * 2, 0.0, 320, 0], @@ -80,14 +82,32 @@ # Aria's raw 21-keypoint layout (0-4 fingertips, 5 palm root) — NOT MANO. Used # only for the opt-in raw-Aria-keypoint viz; the canonical keypoints are MANO. ARIA_FINGER_EDGES = [ - (5, 6), (6, 7), (7, 0), # thumb - (5, 8), (8, 9), (9, 10), (10, 1), # index - (5, 11), (11, 12), (12, 13), (13, 2), # middle - (5, 14), (14, 15), (15, 16), (16, 3), # ring - (5, 17), (17, 18), (18, 19), (19, 4), # pinky + (5, 6), + (6, 7), + (7, 0), # thumb + (5, 8), + (8, 9), + (9, 10), + (10, 1), # index + (5, 11), + (11, 12), + (12, 13), + (13, 2), # middle + (5, 14), + (14, 15), + (15, 16), + (16, 3), # ring + (5, 17), + (17, 18), + (18, 19), + (19, 4), # pinky ] ARIA_FINGER_EDGE_RANGES = [ - ("thumb", 0, 3), ("index", 3, 7), ("middle", 7, 11), ("ring", 11, 15), ("pinky", 15, 19), + ("thumb", 0, 3), + ("index", 3, 7), + ("middle", 7, 11), + ("ring", 11, 15), + ("pinky", 15, 19), ] @@ -103,6 +123,7 @@ class Human(Embodiment): zarr.json); ``cls.INTRINSICS`` is only a fallback for legacy episodes that lack them. The canonical keypoints are MANO for every vendor. """ + INTRINSICS = ARIA_INTRINSICS # fallback only — real value comes from the batch ACTION_HORIZON = 30 # Front-image key for Pi/PaliGemma-style naming (any "_pi"-suffixed mode); @@ -111,11 +132,26 @@ class Human(Embodiment): T_RGB_CPF = ARIA_T_RGB_CPF # for the opt-in aria gaze viz # Canonical MANO 21-keypoint topology: 0=wrist, 1-4 thumb, 5-8 index, ... FINGER_EDGES = [ - (0, 1), (1, 2), (2, 3), (3, 4), # thumb - (0, 5), (5, 6), (6, 7), (7, 8), # index - (0, 9), (9, 10), (10, 11), (11, 12), # middle - (0, 13), (13, 14), (14, 15), (15, 16), # ring - (0, 17), (17, 18), (18, 19), (19, 20), # pinky + (0, 1), + (1, 2), + (2, 3), + (3, 4), # thumb + (0, 5), + (5, 6), + (6, 7), + (7, 8), # index + (0, 9), + (9, 10), + (10, 11), + (11, 12), # middle + (0, 13), + (13, 14), + (14, 15), + (15, 16), # ring + (0, 17), + (17, 18), + (18, 19), + (19, 20), # pinky ] FINGER_COLORS = { "thumb": (255, 100, 100), @@ -191,10 +227,13 @@ def get_keymap( include_aria_keypoints: bool = False, norm_mode: bool = False, annotation_key: str = None, + high_annotation_key=None, ): """Build the keymap. Per-vendor knobs are explicit args from the data config: ``has_head_pose`` (Scale=False) and ``include_aria_keypoints`` - (Aria=True). ``norm_mode``/``annotation_key`` behave as in the base. + (Aria=True). ``norm_mode``/``annotation_key``/``high_annotation_key`` + behave as in the base (subtask mode splits the single annotation array + into a ``level == "low"`` target and a ``level == "high"`` prompt). """ key_map = cls._get_keymap( keymap_mode, @@ -206,6 +245,15 @@ def get_keymap( "key_type": "annotation_keys", "zarr_key": annotation_key, } + if high_annotation_key is not None: + # Subtask mode: split the single annotation array into a + # low-level (target) and high-level (prompt) view. + key_map[annotation_key]["level"] = "low" + key_map[high_annotation_key] = { + "key_type": "annotation_keys", + "zarr_key": annotation_key, + "level": "high", + } if norm_mode: to_delete = [ k @@ -327,26 +375,98 @@ 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, + fix_mecka_left_wrist: bool = False, + pad_proprio_gripper: bool = False, ) -> list[Transform]: """Transform pipeline. ``stride`` is the per-vendor action stride (Aria/LightWheel=3, Scale/Mecka=1), supplied by the data config. + + ``fix_mecka_left_wrist`` retroactively corrects the LEFT wrist-frame + convention of mecka zarrs converted before the ``rot_left`` fix in + ``mecka_to_zarr.compute_hand_pose_xyzquat`` (which double-mirrored the + left hand onto the right hand's spatial convention): the raw left pose + keys are right-multiplied by Rz(180°) before any frame math, exactly + equivalent to reconverting. Set it from mecka data configs only — + do NOT enable for aria/scale (different converters) or for mecka data + reconverted after the fix (it would double-flip). """ + prefix: list[Transform] = [] + if fix_mecka_left_wrist: + if mode.startswith("keypoints"): + raise ValueError( + "fix_mecka_left_wrist only applies to cartesian modes " + "(keypoints modes never read the constructed wrist pose)" + ) + prefix = [ + RotateLocalFrame(keys=["left.action_ee_pose", "left.obs_ee_pose"]) + ] if mode == "cartesian": - return _build_human_cartesian_bimanual_transform_list(stride=stride) - if mode == "cartesian_padded": - return _build_human_cartesian_bimanual_transform_list( + return prefix + _build_human_cartesian_bimanual_transform_list( stride=stride - ) + [PadGripperZeros(action_key="actions_cartesian")] + ) + if mode == "cartesian_6d": + # Head/camera-frame cartesian (12D xyz+ypr per arm) with rotation + # re-expressed as the continuous 6D representation (18D xyz+6d per + # arm) for pi0.5 normalized-rot6d encoding. The proprio ee_pose is + # 6D-encoded too: normalized YPR saturates yaw/roll at ±π + # (wraparound), so per-dim normalization needs the continuous rep. + return ( + prefix + + _build_human_cartesian_bimanual_transform_list(stride=stride) + + [ + CartesianYPRToRot6D(action_key="actions_cartesian"), + CartesianYPRToRot6D(action_key="observations.state.ee_pose"), + ] + # 18 -> 20: zero grip slots at 9/19 so the proprio State: bins + # align positionally with the robot 20-dim layout in the prompt. + + ( + [PadGripperZeros(action_key="observations.state.ee_pose")] + if pad_proprio_gripper + else [] + ) + ) + if mode == "cartesian_padded": + return ( + prefix + + _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) + return prefix + _build_human_cartesian_eef_frame_transform_list( + stride=stride + ) + if mode == "cartesian_wristframe_6d": + # Wrist-frame cartesian (12D xyz+ypr per arm) with rotation + # re-expressed as the continuous 6D representation (18D) for pi0.5 + # normalized-rot6d encoding. The headframe proprio ee_pose is + # 6D-encoded too (see cartesian_6d) — extra important here since + # the proprio is the only head-frame signal the model sees with + # wrist-relative action targets. + return ( + prefix + + _build_human_cartesian_eef_frame_transform_list(stride=stride) + + [ + CartesianYPRToRot6D(action_key="actions_cartesian"), + CartesianYPRToRot6D(action_key="observations.state.ee_pose"), + ] + # 18 -> 20: zero grip slots at 9/19 so the proprio State: bins + # align positionally with the robot 20-dim layout in the prompt. + + ( + [PadGripperZeros(action_key="observations.state.ee_pose")] + if pad_proprio_gripper + else [] + ) + ) if mode == "keypoints_headframe_ypr": return _build_human_keypoints_bimanual_transform_list( stride=stride, is_quat=False @@ -923,6 +1043,52 @@ 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. + + Used by the cam-frame 6D evaluator: the action chunk is already in + head/camera frame (produced by the ``cartesian_6d`` transform 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 cam-frame MSE and + the viz video see the same ypr layout as the plain ``cartesian`` mode. The + proprio ee_pose (also 6D-encoded by ``cartesian_6d``) 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 ARIA actions back to head/camera-frame ypr. + + (1) ``CartesianRot6DToYPR`` converts the action rotation xyz+6D -> xyz+ypr + (Gram-Schmidt re-orthonormalizes the possibly non-orthonormal model + prediction); (2) the proprio ``observations.state.ee_pose`` (6D-encoded by + the ``cartesian_wristframe_6d`` mode) is reverted to ypr the same way; + (3) the standard eef-frame revert projects 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), + # padded proprio arrives 20-dim -> 14 after 6D->ypr; the eef revert's + # SplitKeys expects the gripperless 12-dim layout (no-op if unpadded). + 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 0388d386a..778eae7df 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -28,9 +28,11 @@ _matrix_to_xyz, _matrix_to_xyzwxyz, _matrix_to_xyzypr, + _rot6d_to_ypr, _xyz_to_matrix, _xyzwxyz_to_matrix, _xyzypr_to_matrix, + _ypr_to_rot6d, wxyz_to_xyzw, xyzw_to_wxyz, ) @@ -387,6 +389,101 @@ def transform(self, batch: dict) -> dict: return batch +class CartesianYPRToRot6D(Transform): + """Convert a bimanual cartesian action chunk from per-arm xyz+ypr(+gripper) + to per-arm xyz+rot6d(+gripper). + + ``rot6d`` is the continuous 6D rotation representation = the first two + columns of the rotation matrix, packed as [col0(3), col1(3)] (see + :func:`egomimic.utils.pose_utils._ypr_to_rot6d`). This matches the column + convention of the ``to32``/``from32`` packers in + ``egomimic.utils.action_utils``, so the resulting per-arm layout maps + directly into the pi0.5 32D action blocks. + + Input layouts (last dim): + 12 -> [L xyz ypr, R xyz ypr] -> 18 [L xyz 6d, R xyz 6d] + 14 -> [L xyz ypr g, R xyz ypr g] -> 20 [L xyz 6d g, R xyz 6d g] + + Preserves the numpy/tensor type of the input (like ``PadGripperZeros``). + """ + + def __init__( + self, action_key: str = "actions_cartesian", output_key: str | None = None + ): + self.action_key = action_key + self.output_key = output_key or 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) + D = arr.shape[-1] + if D == 14: + l_xyz, l_ypr, l_g = arr[..., 0:3], arr[..., 3:6], arr[..., 6:7] + r_xyz, r_ypr, r_g = arr[..., 7:10], arr[..., 10:13], arr[..., 13:14] + out = np.concatenate( + [l_xyz, _ypr_to_rot6d(l_ypr), l_g, r_xyz, _ypr_to_rot6d(r_ypr), r_g], + axis=-1, + ) + elif D == 12: + l_xyz, l_ypr = arr[..., 0:3], arr[..., 3:6] + r_xyz, r_ypr = arr[..., 6:9], arr[..., 9:12] + out = np.concatenate( + [l_xyz, _ypr_to_rot6d(l_ypr), r_xyz, _ypr_to_rot6d(r_ypr)], + axis=-1, + ) + else: + raise ValueError( + f"CartesianYPRToRot6D expects last-dim 12 or 14, got {arr.shape} " + f"for '{self.action_key}'" + ) + batch[self.output_key] = torch.from_numpy(out) if is_tensor else out + return batch + + +class CartesianRot6DToYPR(Transform): + """Inverse of :class:`CartesianYPRToRot6D`: per-arm xyz+rot6d(+gripper) -> + xyz+ypr(+gripper). + + Input layouts (last dim): + 18 -> [L xyz 6d, R xyz 6d] -> 12 [L xyz ypr, R xyz ypr] + 20 -> [L xyz 6d g, R xyz 6d g] -> 14 [L xyz ypr g, R xyz ypr g] + """ + + def __init__( + self, action_key: str = "actions_cartesian", output_key: str | None = None + ): + self.action_key = action_key + self.output_key = output_key or 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) + D = arr.shape[-1] + if D == 20: + l_xyz, l_6d, l_g = arr[..., 0:3], arr[..., 3:9], arr[..., 9:10] + r_xyz, r_6d, r_g = arr[..., 10:13], arr[..., 13:19], arr[..., 19:20] + out = np.concatenate( + [l_xyz, _rot6d_to_ypr(l_6d), l_g, r_xyz, _rot6d_to_ypr(r_6d), r_g], + axis=-1, + ) + elif D == 18: + l_xyz, l_6d = arr[..., 0:3], arr[..., 3:9] + r_xyz, r_6d = arr[..., 9:12], arr[..., 12:18] + out = np.concatenate( + [l_xyz, _rot6d_to_ypr(l_6d), r_xyz, _rot6d_to_ypr(r_6d)], + axis=-1, + ) + else: + raise ValueError( + f"CartesianRot6DToYPR expects last-dim 18 or 20, got {arr.shape} " + f"for '{self.action_key}'" + ) + batch[self.output_key] = torch.from_numpy(out) if is_tensor else out + return batch + + class CartesianWithGripperCoordinateTransform(Transform): def __init__( self, @@ -482,9 +579,24 @@ def __init__(self, input_key: str, output_key_list: list[(str, int)]): 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 @@ -512,6 +624,77 @@ def transform(self, batch): return batch +class RotateLocalFrame(Transform): + """Right-multiply listed xyz+quat(wxyz) pose keys by a constant LOCAL + rotation: ``R_new = R_old @ R_fix``. Relabels the pose's own axes without + moving its origin. Handles ``(7,)`` poses and ``(T, 7)`` chunks. + + Used to retroactively fix the mecka LEFT wrist-frame convention without + reconverting the zarrs: ``compute_hand_pose_xyzquat`` built the palm + normal as ``cross(thumb_dir, pinky_dir)``, which already mirrors chirality + between hands, and then ``rot_left`` flipped x/y again — double-mirroring + the left hand onto the right hand's spatial convention. Since + ``rot_left == rot_right @ diag(-1, -1, 1)``, right-multiplying the stored + left pose by Rz(180°) (the default ``quat_wxyz``) is exactly equivalent to + reconverting with ``rot_right`` for both hands. + """ + + def __init__( + self, + keys: list[str], + quat_wxyz: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0), + ): + self.keys = list(keys) + self.quat_wxyz = tuple(float(v) for v in quat_wxyz) + w, x, y, z = self.quat_wxyz + self._fix = R.from_quat([x, y, z, w]) # scipy xyzw + + def transform(self, batch: dict) -> dict: + for key in self.keys: + pose = np.asarray(batch[key]) + if pose.shape[-1] != 7: + raise ValueError( + f"RotateLocalFrame expects xyz+quat(wxyz) with last dim 7, " + f"got {pose.shape} for '{key}'" + ) + flat = pose.reshape(-1, 7).astype(np.float64, copy=True) + # Zero-norm quats mark padded/invalid frames — leave them alone. + valid = np.linalg.norm(flat[:, 3:7], axis=-1) > 1e-6 + if valid.any(): + q_xyzw = flat[valid][:, [4, 5, 6, 3]] + rotated = (R.from_quat(q_xyzw) * self._fix).as_quat() # xyzw + flat[np.flatnonzero(valid), 3:7] = rotated[:, [3, 0, 1, 2]] + batch[key] = flat.reshape(pose.shape) + 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 (6D: drop 9, 19). Widths 12 + and 18 pass through unchanged, so revert pipelines work 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 = batch[self.action_key] + is_tensor = isinstance(arr, torch.Tensor) + a = arr.cpu().numpy() if is_tensor else np.asarray(arr) + D = a.shape[-1] + if D in (12, 18): + return batch + if D == 14: + keep = [0, 1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12] + elif D == 20: + keep = [i for i in range(20) if i not in (9, 19)] + else: + raise ValueError(f"UnpadGripperZeros: unexpected width {a.shape}") + out = a[..., keep] + batch[self.action_key] = torch.from_numpy(out) if is_tensor else out + return 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 @@ -528,6 +711,12 @@ 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] == 18: + # 6D layout: [L xyz 6d | R xyz 6d] -> insert grip zeros at 9, 19 + zero = np.zeros_like(arr[..., :1]) + out = np.concatenate([arr[..., 0:9], zero, arr[..., 9:18], zero], axis=-1) + batch[self.action_key] = torch.from_numpy(out) if is_tensor else out + return batch if arr.shape[-1] != 12: raise ValueError( f"PadGripperZeros expects last-dim 12, got {arr.shape} for " @@ -535,12 +724,8 @@ def transform(self, batch: dict) -> dict: ) 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 - ) + padded = np.concatenate((arr[..., :6], pad, arr[..., 6:], pad), axis=-1) + batch[self.action_key] = torch.from_numpy(padded) if is_tensor else padded return batch diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index a4d5d0a1b..9d560e0d7 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -48,6 +48,7 @@ create_default_engine, episode_table_to_df, ) +from egomimic.utils.pose_utils import bimanual_cartesian_layout if TYPE_CHECKING: # Annotation-only import — avoids a runtime circular import with @@ -876,6 +877,19 @@ def _check_bounds( q_low = torch.broadcast_to(q_low, arr.shape) q_high = torch.broadcast_to(q_high, arr.shape) except RuntimeError: + # Stats were computed for a different layout than this sample + # (e.g. a stale precomputed norm_stats.json). Say so once + # instead of silently disabling the bounds check for the key; + # normalize() will raise on the same mismatch anyway. + warn_key = f"bounds-shape:{zarr_key}" + if warn_key not in self._warned_violations: + self._warned_violations.add(warn_key) + logger.warning( + f"[MultiDataset] bounds check skipped for {zarr_key}: " + f"stats shape {tuple(q_low.shape)} does not broadcast to " + f"sample shape {tuple(arr.shape)} (norm stats computed " + "for a different layout?)" + ) continue if torch.any(torch.isnan(arr)) or torch.any(torch.isinf(arr)): @@ -886,8 +900,35 @@ def _check_bounds( logger.warning(prefix) return prefix - below = arr < q_low - above = arr > q_high + # The bimanual cartesian action chunk and the ee_pose proprio share + # a [L | R] layout whose rotation channels are either Euler ypr + # (wraps at ±π) or continuous 6D columns. In both cases quantile + # bounds on the rotation channels are meaningless and reject + # otherwise-valid frames, so only the translation (and gripper) + # channels are bounds-checked. Unrecognized widths fall through to + # a full-vector check; NaN/Inf above still covers the full vector. + cartesian_layout = None + if zarr_key in ("actions_cartesian", "observations.state.ee_pose"): + cartesian_layout = bimanual_cartesian_layout(arr.shape[-1]) + if cartesian_layout is not None: + check_idx = list(cartesian_layout["xyz"]) + list( + cartesian_layout["grip"] + ) + arr_q = arr[..., check_idx] + q_low = q_low[..., check_idx] + q_high = q_high[..., check_idx] + else: + arr_q = arr + + # Absolute slack on the quantile bounds. Wrist-frame action chunks + # are the identity pose at t=0 (the reference IS the obs pose), so + # those cells' bounds collapse to [0, 0] and a strict compare would + # reject every frame on any roundoff (today the cells are exactly + # 0.0, so this only guards against a different BLAS/dtype path). + # 1e-6 (m / normalized grip) is far below any real outlier. + tol = 1e-6 + below = arr_q < q_low - tol + above = arr_q > q_high + tol if torch.any(below) or torch.any(above): prefix = f"Bounds violation in {zarr_key} ep={episode_name} frame={idx}" warn_key = f"bounds:{episode_name}:{zarr_key}" @@ -897,7 +938,7 @@ def _check_bounds( n_above = int(above.sum().item()) logger.warning( f"{prefix} | n_below={n_below} n_above={n_above} " - f"arr_range=[{arr.min().item():.4f}, {arr.max().item():.4f}]" + f"arr_range=[{arr_q.min().item():.4f}, {arr_q.max().item():.4f}]" ) return prefix return None @@ -1116,19 +1157,7 @@ def infer_norm_from_dataset( ) return if os.path.isfile(precomputed_file): - with open(precomputed_file, "r") as f: - payload = json.load(f) - if str(embodiment) not in payload["stats"]: - raise ValueError( - f"norm_stats file {precomputed_file} has no entry for " - f"embodiment id {embodiment} (available: " - f"{sorted(payload['stats'])}). Stats are keyed by numeric " - "EMBODIMENT id, and ids were renumbered by the human/eva " - "embodiment collapse — recompute norm stats instead of " - "reusing a pre-collapse norm_stats.json." - ) - self.norm_stats[embodiment] = payload["stats"][str(embodiment)] - self._norm_run_metadata = payload.get("norm_run_metadata", None) + self._load_precomputed_stats(precomputed_file, embodiment, norm_keys) logger.info( f"[MultiDataset] Loaded precomputed stats for embodiment={embodiment}" ) @@ -1185,6 +1214,55 @@ def infer_norm_from_dataset( f"[MultiDataset] Finished norm inference, loading={loading_time:.2f}s, computing={computing_time:.2f}s" ) + def _load_precomputed_stats( + self, precomputed_file: str, embodiment: int, norm_keys: list[str] + ) -> None: + """Load ``norm_stats.json`` for one embodiment, refusing a file whose + provenance does not match this dataset. + + The stats are only meaningful for the exact (norm_mode, key set) + they were computed under; the payload's ``provenance`` block (written + by :meth:`cache_stats`) carries both. Files written before provenance + existed load as before, with a warning. + """ + with open(precomputed_file, "r") as f: + payload = json.load(f) + if str(embodiment) not in payload["stats"]: + raise ValueError( + f"norm_stats file {precomputed_file} has no entry for " + f"embodiment id {embodiment} (available: " + f"{sorted(payload['stats'])}). Stats are keyed by numeric " + "EMBODIMENT id, and ids were renumbered by the human/eva " + "embodiment collapse — recompute norm stats instead of " + "reusing a pre-collapse norm_stats.json." + ) + provenance = payload.get("provenance") + if provenance is None: + logger.warning( + f"[MultiDataset] {precomputed_file} carries no provenance block " + "(written by an older cache_stats); cannot verify it matches " + f"norm_mode={self.norm_mode!r} and this dataset's keys." + ) + else: + file_mode = provenance.get("norm_mode") + if file_mode != self.norm_mode: + raise ValueError( + f"norm_stats file {precomputed_file} was computed with " + f"norm_mode={file_mode!r} but this dataset uses " + f"norm_mode={self.norm_mode!r}; recompute the stats." + ) + file_keys = set(payload["stats"][str(embodiment)]) + want_keys = set(norm_keys) + if norm_keys and file_keys != want_keys: + raise ValueError( + f"norm_stats file {precomputed_file} keys for embodiment " + f"{embodiment} are {sorted(file_keys)} but this dataset " + f"normalizes {sorted(want_keys)}; the file was computed for a " + "different keymap/transform mode — recompute the stats." + ) + self.norm_stats[embodiment] = payload["stats"][str(embodiment)] + self._norm_run_metadata = payload.get("norm_run_metadata", None) + def _collect_norm_samples( self, loader, norm_keys, embodiment, n_samples, batch_size, num_workers ): @@ -1212,7 +1290,10 @@ def _collect_norm_samples( x = batch[zarr_key][:take] if hasattr(x, "detach"): x = x.detach().cpu().numpy() - collected[k].append(x) + # float32: stats are consumed as float32 anyway, and the + # float64 poses double the stacked-sample footprint (an + # (N, 100, 18) action stack at large N is tens of GB). + collected[k].append(np.asarray(x, dtype=np.float32)) cur += take pbar.update(take) return collected @@ -1244,6 +1325,19 @@ def cache_stats(self, save_cache_dir: str): } payload = { "stats": stats_out, + # What the stats are valid for — checked by _load_precomputed_stats + # so a cached file from another norm_mode / keymap / transform mode + # (same dims, different meaning) is refused instead of applied. + "provenance": { + "norm_mode": self.norm_mode, + "stat_shapes": { + str(emb): { + k: list(np.asarray(next(iter(sd.values()))).shape) + for k, sd in keys_dict.items() + } + for emb, keys_dict in self.norm_stats.items() + }, + }, "loading_time": None, "computing_time": None, "frames": None, @@ -1259,6 +1353,8 @@ def cache_stats(self, save_cache_dir: str): # ---- normalize / unnormalize ---- def _apply_norm_one(self, tensor, stats): + if self.norm_mode == "none": + return tensor if self.norm_mode == "zscore": mean = torch.as_tensor( stats["mean"], device=tensor.device, dtype=torch.float32 @@ -1286,6 +1382,8 @@ def _apply_norm_one(self, tensor, stats): raise ValueError(f"Invalid normalization mode: {self.norm_mode}") def _apply_unnorm_one(self, tensor, stats): + if self.norm_mode == "none": + return tensor if self.norm_mode == "zscore": mean = torch.as_tensor( stats["mean"], device=tensor.device, dtype=torch.float32 @@ -1753,13 +1851,18 @@ def _next(reason: str, key: str = "") -> int: # Normalize a 3x3 K to the canonical 3x4 (zeros last column); # some contributors store 3x3 (e.g. microagi). K = np.concatenate([K, np.zeros((3, 1), dtype=np.float32)], axis=1) - if K.shape != (3, 4): # unexpected -> sentinel (viz falls back to const) + if K.shape != ( + 3, + 4, + ): # unexpected -> sentinel (viz falls back to const) K = np.full((3, 4), np.nan, dtype=np.float32) else: K = np.full((3, 4), np.nan, dtype=np.float32) data["intrinsics"] = torch.from_numpy(np.ascontiguousarray(K)) ep_name = Path(self.episode_path).name - data["episode_hash"] = ep_name[:-5] if ep_name.endswith(".zarr") else ep_name + data["episode_hash"] = ( + ep_name[:-5] if ep_name.endswith(".zarr") else ep_name + ) _ = origin # preserved for symmetry with prior API return data @@ -1785,13 +1888,26 @@ def _build_frame_to_ann_end(self) -> dict[int, int]: annotation span. Annotations use half-open ``[start_idx, end_idx)``. """ mapping: dict[int, int] = {} + n_spans = 0 for ann in self._load_annotations(): start_idx = int(ann.get("start_idx", -1)) end_idx = int(ann.get("end_idx", -1)) if start_idx < 0 or end_idx <= start_idx: continue + n_spans += 1 for idx in range(start_idx, end_idx): mapping[idx] = end_idx + # One-time per-episode visibility into annotation-cutoff usage: if + # spans/frames_covered are 0 the cutoff is a no-op (episode has no usable + # annotations); >0 confirms action chunks are being clamped at EOS. + ep = Path(self.episode_path).name + logger.info( + "[AnnotationCutoff] ep=%s spans=%d frames_covered=%d/%d", + ep, + n_spans, + len(mapping), + self.total_frames, + ) return mapping def _chunk_end_idx(self, start_idx: int, horizon: int, key_type: str | None) -> int: @@ -1806,11 +1922,67 @@ def _chunk_end_idx(self, start_idx: int, horizon: int, key_type: str | None) -> return min(end_idx, ann_end) +def _episode_has_annotation_spans(ds: "ZarrDataset") -> bool: + """True if the episode has at least one usable ``[start_idx, end_idx)`` span. + + Many Scale-"completed" episodes have an empty (or span-less) zarr + ``annotations`` array because the annotation-injection step lagged; the + AnnotationCutoff is a no-op for those, so they should be dropped when the + point of the run is to clamp chunks at annotation boundaries. + """ + try: + anns = ds._load_annotations() + except Exception: + return False + return any( + isinstance(a, dict) + and 0 <= int(a.get("start_idx", -1)) < int(a.get("end_idx", -1)) + for a in anns + ) + + class S3AnnotationCutoffEpisodeResolver(S3EpisodeResolver): - """S3EpisodeResolver that loads ZarrAnnotationCutoffDataset instances.""" + """S3EpisodeResolver that loads ZarrAnnotationCutoffDataset instances. + + When ``require_annotations`` is set (default), episodes whose zarr + ``annotations`` array has no usable span are dropped — otherwise the + annotation cutoff would silently no-op on them. + """ _dataset_class = ZarrAnnotationCutoffDataset + def __init__(self, *args, require_annotations: bool = True, **kwargs): + super().__init__(*args, **kwargs) + self.require_annotations = require_annotations + + def resolve(self, filters=None): + datasets = super().resolve(filters=filters) + if not self.require_annotations: + return datasets + kept = { + h: ds for h, ds in datasets.items() if _episode_has_annotation_spans(ds) + } + dropped = sorted(set(datasets) - set(kept)) + if dropped: + logger.warning( + "[AnnotationCutoff] dropped %d/%d episodes with no usable " + "annotation spans (e.g. %s)", + len(dropped), + len(datasets), + dropped[:5], + ) + logger.info( + "[AnnotationCutoff] kept %d/%d episodes with usable annotations", + len(kept), + len(datasets), + ) + if not kept: + raise ValueError( + "[AnnotationCutoff] no resolved episodes contain usable annotation " + "spans — check the filter / annotation injection for this dataset." + ) + return kept + class LocalAnnotationCutoffEpisodeResolver(LocalEpisodeResolver): """LocalEpisodeResolver that loads ZarrAnnotationCutoffDataset instances.""" diff --git a/egomimic/scripts/abc_process/viz_eva_episode.py b/egomimic/scripts/abc_process/viz_eva_episode.py new file mode 100644 index 000000000..f9a90a688 --- /dev/null +++ b/egomimic/scripts/abc_process/viz_eva_episode.py @@ -0,0 +1,363 @@ +#!/usr/bin/env python +"""Visualize a local EVA (ABC-130k) zarr episode -> mp4. + +Reads the camera frames straight out of the zarr store (no model / no GPU) and +writes a video. Optionally tiles multiple cameras side-by-side and burns in the +task description + frame index. + + ./emimic/bin/python egomimic/scripts/abc_process/viz_eva_episode.py \ + --episode 00022b62-0a07-4ac7-bb74-9b4da47f6c6e \ + --cams front_1 left_wrist right_wrist --out /workspace/EgoVerse/eva_ep.mp4 + +--episode accepts a bare hash (resolved under --folder) or a full .zarr path. + +--draw-axes overlays, on the front_1 tile, the per-arm EE pose triads AND the +world-frame origin triad (X=red, Y=green, Z=blue), projected with the SAME +math the training pipeline uses: world pose -> cam frame via +PoseCoordinateFrameTransform against EXTRINSICS[--extrinsics-key], then pixels +via INTRINSICS["eva"]. The front_1 tile is stretched to 640x480 first (the +space that K and the abc_fold_viz solvePnP fit live in). Use --still-frames to +also dump overlaid PNGs of specific frames. +""" + +import argparse +import io +import os + +import numpy as np +import zarr + + +def decode(elem): + """An image stored in the object array -> HxWx3 uint8. + + Frames are JPEG bytes wrapped in a 0-d object ndarray (sometimes nested), so + unwrap to the raw bytes and decode with PIL (imageio can't sniff a bare BytesIO). + """ + while isinstance(elem, np.ndarray) and elem.dtype == object and elem.ndim == 0: + elem = elem.item() + if isinstance(elem, np.ndarray) and elem.ndim >= 2: + return elem[..., :3].astype(np.uint8) + raw = bytes(elem) if not isinstance(elem, (bytes, bytearray)) else elem + from PIL import Image + + return np.asarray(Image.open(io.BytesIO(raw)).convert("RGB"), dtype=np.uint8) + + +# Viz space the overlay lives in: front_1 stretched to 640x480, matching both +# the pipeline's _resize_image_keys and the clicker/solvePnP fit. +VIZ_W, VIZ_H = 640, 480 +AXIS_COLORS = {"x": (255, 0, 0), "y": (0, 255, 0), "z": (0, 0, 255)} # RGB + +# Inlined from the remote (wristframe-6d) repo's egomimicUtils registries — +# the local repo has no INTRINSICS/EXTRINSICS dicts (training reads intrinsics +# per-episode from zarr.json; these viz calibrations exist only on the remote +# branch). +# eva K: ABC-130k RealSense top camera (640x480 space) +# yam K: rectified stereo front cam, calibration-derived, 640x480 space +# abc_fold_viz: solvePnP world->cam fit over ABC fold-clothes clicks (one +# shared world frame for both arms); yam: per-arm base->cam hand-eye calib +INTRINSICS = { + "eva": np.array( + [[436.26, 0.0, 310.1, 0.0], [0.0, 435.13, 241.93, 0.0], [0.0, 0.0, 1.0, 0.0]] + ), + "yam": np.array( + [ + [250.15807185802416, 0.0, 267.5409836065574, 0.0], + [0.0, 260.01789323969166, 152.72727272727272, 0.0], + [0.0, 0.0, 1.0, 0.0], + ] + ), +} +_ABC_FOLD_VIZ = np.array( + [ + [-0.00540728648402, -0.91864395735192, 0.39504941573641, -0.07230219027034], + [-0.99924104589933, -0.01027605228132, -0.03757306135427, -0.23692835817344], + [0.03857581422213, -0.39495275966919, -0.91789118319482, 0.97783070025080], + [0.0, 0.0, 0.0, 1.0], + ] +) +EXTRINSICS = { + "abc_fold_viz": {"left": _ABC_FOLD_VIZ, "right": _ABC_FOLD_VIZ}, + "yam": { + "left": np.array( + [ + [0.0437291, -0.85821391, 0.5114261, 0.08690521], + [-0.99754573, -0.00948836, 0.06937217, -0.25327518], + [-0.05468356, -0.5132045, -0.85652253, 0.83357606], + [0.0, 0.0, 0.0, 1.0], + ] + ), + "right": np.array( + [ + [0.04251619, -0.84174054, 0.53820557, 0.04959916], + [-0.98836195, 0.04331659, 0.14582295, 0.27905066], + [-0.14605832, -0.53814174, -0.83010266, 0.87148985], + [0.0, 0.0, 0.0, 1.0], + ] + ), + }, +} + + +class AxisOverlay: + """Project EE / world triads into front_1 pixels via the pipeline transform.""" + + def __init__(self, extrinsics_key: str, intrinsics_key: str = "eva"): + from egomimic.rldb.zarr.action_chunk_transforms import ( + PoseCoordinateFrameTransform, + ) + from egomimic.utils.pose_utils import _matrix_to_xyzwxyz, _xyzwxyz_to_matrix + + self._to_mat = _xyzwxyz_to_matrix + self.K = INTRINSICS[intrinsics_key] + # EVA's K lives in the pipeline's stretched 640x480 space; YAM's K is + # already in the stored image space -> draw at native resolution there. + self.stretch = intrinsics_key == "eva" + extr = EXTRINSICS[extrinsics_key] + # cam<-world pose target per arm, exactly as Eva.get_transform_list bakes it. + self.extr_pose = { + s: _matrix_to_xyzwxyz(np.asarray(extr[s], dtype=np.float64)[None])[0] + for s in ("left", "right") + } + # One shared world frame (e.g. abc_fold_viz) vs per-arm base frames (yam). + self.shared_base = np.allclose(extr["left"], extr["right"]) + self._t = PoseCoordinateFrameTransform( + target_world="extr", + pose_world="pose", + transformed_key_name="out", + mode="xyzwxyz", + ) + self.key = extrinsics_key + + def cam_pose_matrix(self, world_pose_xyzwxyz, side: str) -> np.ndarray: + """World-frame pose (7,) -> 4x4 cam<-ee matrix (pipeline-identical).""" + out = self._t.transform( + { + "extr": self.extr_pose[side], + "pose": np.asarray(world_pose_xyzwxyz, dtype=np.float64), + } + )["out"] + return self._to_mat(np.asarray(out, dtype=np.float64)[None])[0] + + def project(self, pts_cam: np.ndarray) -> np.ndarray: + """(N,3) cam-frame points -> (N,2) pixels; NaN where behind the camera.""" + px = np.full((len(pts_cam), 2), np.nan) + front = pts_cam[:, 2] > 1e-6 + if front.any(): + p = np.concatenate([pts_cam[front], np.ones((front.sum(), 1))], axis=1) + uv = self.K @ p.T + px[front] = (uv[:2] / uv[2]).T + return px + + def draw_triad(self, img, M_cam_pose, length: float, label: str, thickness=2): + """Draw an XYZ triad for a cam<-frame pose matrix onto img (RGB, 640x480).""" + import cv2 + + o = M_cam_pose[:3, 3] + pts = np.stack([o] + [o + M_cam_pose[:3, k] * length for k in range(3)]) + px = self.project(pts) + if np.isnan(px[0]).any(): + return + oi = tuple(np.round(px[0]).astype(int)) + for k, ax in enumerate("xyz"): + if np.isnan(px[k + 1]).any(): + continue + ei = tuple(np.round(px[k + 1]).astype(int)) + cv2.line(img, oi, ei, AXIS_COLORS[ax], thickness, cv2.LINE_AA) + cv2.putText( + img, + ax, + (ei[0] + 3, ei[1] - 3), + cv2.FONT_HERSHEY_SIMPLEX, + 0.45, + AXIS_COLORS[ax], + 1, + cv2.LINE_AA, + ) + cv2.circle(img, oi, 3, (255, 255, 255), -1) + cv2.putText( + img, + label, + (oi[0] + 5, oi[1] + 14), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + (255, 255, 0), + 1, + cv2.LINE_AA, + ) + + def overlay(self, front_native, lpose_w, rpose_w): + """front_1 native frame + world EE poses -> 640x480 RGB with triads. + + Returns (img, info) where info holds the projected EE origin pixels. + """ + import cv2 + + if self.stretch: + img = cv2.resize( + front_native, (VIZ_W, VIZ_H), interpolation=cv2.INTER_LINEAR + ) + else: + # PIL-decoded arrays are read-only; cv2 draws in place -> copy. + img = front_native.copy() + img = np.ascontiguousarray(img) + h = img.shape[0] + info = {} + # Base-frame triad(s): one shared world frame, or per-arm bases (yam). + ident = np.array([0, 0, 0, 1, 0, 0, 0], dtype=np.float64) + bases = ( + [("left", "world")] + if self.shared_base + else [("left", "L-base"), ("right", "R-base")] + ) + for side, lab in bases: + self.draw_triad( + img, self.cam_pose_matrix(ident, side), 0.15, lab, thickness=3 + ) + for side, pose_w, lab in ( + ("left", lpose_w, "L-ee"), + ("right", rpose_w, "R-ee"), + ): + # ABC episodes zero-pad the tail frames -> zero-norm quat; skip them. + if np.linalg.norm(np.asarray(pose_w, dtype=np.float64)[3:7]) < 1e-6: + info[side] = np.array([np.nan, np.nan]) + continue + M = self.cam_pose_matrix(pose_w, side) + self.draw_triad(img, M, 0.08, lab) + info[side] = self.project(M[:3, 3][None])[0] + cv2.putText( + img, + f"axes X=red Y=green Z=blue extr={self.key}", + (6, h - 8), + cv2.FONT_HERSHEY_SIMPLEX, + 0.45, + (255, 255, 255), + 1, + cv2.LINE_AA, + ) + return img, info + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--episode", required=True, help="episode hash or full .zarr path") + ap.add_argument("--folder", default="/workspace/eva/abc130k_zarr") + ap.add_argument( + "--cams", nargs="+", default=["front_1"], help="front_1 left_wrist right_wrist" + ) + ap.add_argument("--out", default=None) + ap.add_argument("--max-frames", type=int, default=0, help="0 = all") + ap.add_argument( + "--draw-axes", + action="store_true", + help="overlay EE + world XYZ triads on front_1 (see module docstring)", + ) + ap.add_argument( + "--extrinsics-key", + default="abc_fold_viz", + help="EXTRINSICS key for the world->cam bake (default: abc_fold_viz)", + ) + ap.add_argument( + "--intrinsics-key", + default="eva", + help="INTRINSICS key for projection (eva | yam)", + ) + ap.add_argument( + "--still-frames", + type=int, + nargs="*", + default=[], + help="also dump these frame indices as overlaid PNGs next to --out", + ) + args = ap.parse_args() + + path = ( + args.episode + if args.episode.endswith(".zarr") + else os.path.join(args.folder, f"{args.episode}.zarr") + ) + g = zarr.open_group(path, mode="r") + a = dict(g.attrs) + fps = int(a.get("fps", 30)) + n = int(a.get("total_frames", 0)) + print( + f"[viz] {path}\n task={a.get('task_description') or a.get('task_name')!r} frames={n} fps={fps} cams={args.cams}" + ) + + overlay = None + lobs = robs = None + if args.draw_axes: + if "front_1" not in args.cams: + args.cams = ["front_1"] + args.cams + overlay = AxisOverlay(args.extrinsics_key, args.intrinsics_key) + lobs, robs = g["left.obs_ee_pose"], g["right.obs_ee_pose"] + print( + f"[viz] axis overlay on front_1: extrinsics_key={args.extrinsics_key}, " + f"K=INTRINSICS['{args.intrinsics_key}'], stretch={overlay.stretch}" + ) + + cam_arrs = {c: g[f"images.{c}"] for c in args.cams} + # Arrays are zero-padded past attrs total_frames -> cap at the real length. + nf = min( + next(iter(cam_arrs.values())).shape[0], n or 10**9, args.max_frames or 10**9 + ) + out = ( + args.out + or f"/workspace/EgoVerse/eva_{os.path.basename(path).replace('.zarr','')}.mp4" + ) + stills = sorted(set(i for i in args.still_frames if 0 <= i < nf)) + + import imageio.v2 as imageio + + try: + from PIL import Image, ImageDraw + except Exception: + Image = None + writer = imageio.get_writer(out, fps=fps, macro_block_size=None) + for i in range(nf): + tiles = [] + for c in args.cams: + f = decode(cam_arrs[c][i]) + if overlay is not None and c == "front_1": + f, info = overlay.overlay(f, np.asarray(lobs[i]), np.asarray(robs[i])) + if i in stills: + print( + f"[viz] frame {i}: EE pixels L={np.round(info['left'],1)} " + f"R={np.round(info['right'],1)}" + ) + tiles.append(f) + h = min(t.shape[0] for t in tiles) + tiles = ( + [ + t[:h] + if t.shape[0] == h + else np.asarray( + Image.fromarray(t).resize((int(t.shape[1] * h / t.shape[0]), h)) + ) + for t in tiles + ] + if Image + else tiles + ) + frame = np.concatenate(tiles, axis=1) + if Image: + im = Image.fromarray(frame) + d = ImageDraw.Draw(im) + d.text( + (6, 6), + f"{i}/{nf} {(a.get('task_description') or '')[:60]}", + fill=(255, 255, 0), + ) + frame = np.asarray(im) + if i in stills: + still_path = f"{os.path.splitext(out)[0]}_f{i}.png" + Image.fromarray(frame).save(still_path) + print(f"[viz] still -> {still_path}") + writer.append_data(frame) + writer.close() + print(f"[viz] wrote {nf} frames -> {out}") + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/mecka_process/mecka_to_zarr.py b/egomimic/scripts/mecka_process/mecka_to_zarr.py index fd0f605b5..139987c92 100644 --- a/egomimic/scripts/mecka_process/mecka_to_zarr.py +++ b/egomimic/scripts/mecka_process/mecka_to_zarr.py @@ -241,10 +241,16 @@ def compute_hand_pose_xyzquat(keypoints: np.ndarray, hand_index: int) -> np.ndar rot_matrix = np.column_stack([forward, right, up]) - if hand_index == 0: - rot_matrix = rot_matrix @ rot_left - else: - rot_matrix = rot_matrix @ rot_right + # Same axis relabel for BOTH hands: ``up = cross(thumb_dir, pinky_dir)`` + # already mirrors chirality between left and right (thumb/pinky are + # anatomically swapped), so the frames come out anatomically mirrored for + # free. The old ``rot_left`` (= rot_right @ diag(-1, -1, 1)) flipped x/y a + # second time, double-mirroring the left hand onto the right hand's + # spatial convention. Data converted with the old code is corrected at + # train time by RotateLocalFrame (Rz(180°) right-multiply on the left + # pose) — see Human.get_transform_list(fix_mecka_left_wrist=True). + del rot_left, hand_index # convention no longer differs per hand + rot_matrix = rot_matrix @ rot_right quat_xyzw = Rotation.from_matrix(rot_matrix).as_quat() # SciPy returns (x, y, z, w) quat_wxyz = np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]]) diff --git a/egomimic/scripts/mecka_process/viz_mecka_episode.py b/egomimic/scripts/mecka_process/viz_mecka_episode.py new file mode 100644 index 000000000..35c060461 --- /dev/null +++ b/egomimic/scripts/mecka_process/viz_mecka_episode.py @@ -0,0 +1,248 @@ +#!/usr/bin/env python +"""Visualize a Mecka (egocentric human) zarr episode -> mp4 with frame overlays. + +Egocentric analog of abc_process/viz_eva_episode.py --draw-axes: projects into +the moving head camera (per-frame ``obs_head_pose`` is the world->cam target, +exactly the headframe transform training uses) with INTRINSICS["mecka"]. + +Overlays, per frame: + - L/R wrist triads (X=red Y=green Z=blue) from ``obs_ee_pose`` — shows the + hand-tracking wrist-frame convention (diagonal to hand geometry, mirrored + between hands, unlike EVA/YAM tool frames). + - hand keypoints (gray dots) + a WHITE line wrist-root -> middle-MCP = + "knuckle-forward", the physical reference to compare the triads against. + - the WORLD origin triad (thick, length 0.3m): origin sits on the floor under + the operator's head at episode start, Z up (gravity-aligned). + + ./emimic/bin/python egomimic/scripts/mecka_process/viz_mecka_episode.py \ + --episode 696d0f031244464a4aac1b8f --still-frames 150 450 750 +""" + +import argparse +import io +import os + +import numpy as np +import zarr + + +def decode(elem): + while isinstance(elem, np.ndarray) and elem.dtype == object and elem.ndim == 0: + elem = elem.item() + if isinstance(elem, np.ndarray) and elem.ndim >= 2: + return elem[..., :3].astype(np.uint8) + raw = bytes(elem) if not isinstance(elem, (bytes, bytearray)) else elem + from PIL import Image + + return np.asarray(Image.open(io.BytesIO(raw)).convert("RGB"), dtype=np.uint8) + + +AXIS_COLORS = {"x": (255, 0, 0), "y": (0, 255, 0), "z": (0, 0, 255)} # RGB +# MediaPipe-style indices (verified on-data: kp0 roots all four finger chains, +# adjacent MCPs 5/9/13/17 are ~2.5cm apart). NOT Aria ordering. +WRIST_ROOT, MIDDLE_MCP = 0, 9 + + +class MeckaOverlay: + """World-frame geometry -> head-camera pixels via the training headframe math.""" + + def __init__(self): + # Local repo keeps intrinsics on the embodiment classes (there is no + # egomimicUtils.INTRINSICS registry like the remote's). + from egomimic.rldb.embodiment.human import MECKA_INTRINSICS + from egomimic.rldb.zarr.action_chunk_transforms import ( + PoseCoordinateFrameTransform, + ) + from egomimic.utils.pose_utils import _xyzwxyz_to_matrix + + self.K = MECKA_INTRINSICS + self._to_mat = _xyzwxyz_to_matrix + self._t_pose = PoseCoordinateFrameTransform( + target_world="head", + pose_world="pose", + transformed_key_name="out", + mode="xyzwxyz", + ) + self._t_xyz = PoseCoordinateFrameTransform( + target_world="head", + pose_world="pose", + transformed_key_name="out", + mode="xyz", + ) + + def cam_pose_matrix(self, head_pose, world_pose): + out = self._t_pose.transform( + { + "head": np.asarray(head_pose, np.float64), + "pose": np.asarray(world_pose, np.float64), + } + )["out"] + return self._to_mat(np.asarray(out, np.float64)[None])[0] + + def cam_xyz(self, head_pose, world_xyz): + """(N,3) world points -> (N,3) cam-frame points.""" + return np.stack( + [ + self._t_xyz.transform( + { + "head": np.asarray(head_pose, np.float64), + "pose": np.asarray(p, np.float64), + } + )["out"] + for p in np.asarray(world_xyz, np.float64).reshape(-1, 3) + ] + ) + + def project(self, pts_cam): + px = np.full((len(pts_cam), 2), np.nan) + front = pts_cam[:, 2] > 1e-6 + if front.any(): + p = np.concatenate([pts_cam[front], np.ones((front.sum(), 1))], axis=1) + uv = self.K @ p.T + px[front] = (uv[:2] / uv[2]).T + return px + + def draw_triad(self, img, M, length, label, thickness=2): + import cv2 + + o = M[:3, 3] + pts = np.stack([o] + [o + M[:3, k] * length for k in range(3)]) + px = self.project(pts) + if np.isnan(px[0]).any(): + return + oi = tuple(np.round(px[0]).astype(int)) + for k, ax in enumerate("xyz"): + if np.isnan(px[k + 1]).any(): + continue + ei = tuple(np.round(px[k + 1]).astype(int)) + cv2.line(img, oi, ei, AXIS_COLORS[ax], thickness, cv2.LINE_AA) + cv2.putText( + img, + ax, + (ei[0] + 3, ei[1] - 3), + cv2.FONT_HERSHEY_SIMPLEX, + 0.4, + AXIS_COLORS[ax], + 1, + cv2.LINE_AA, + ) + cv2.circle(img, oi, 3, (255, 255, 255), -1) + cv2.putText( + img, + label, + (oi[0] + 5, oi[1] + 14), + cv2.FONT_HERSHEY_SIMPLEX, + 0.45, + (255, 255, 0), + 1, + cv2.LINE_AA, + ) + + def overlay(self, img, head, lpose, rpose, lkp, rkp): + import cv2 + + # PIL-decoded arrays are read-only; cv2 draws in place -> copy. + img = np.ascontiguousarray(img.copy()) + # world origin triad (floor under head-at-start, Z up) + if np.linalg.norm(np.asarray(head)[3:7]) > 1e-6: + Mw = self.cam_pose_matrix(head, np.array([0, 0, 0, 1, 0, 0, 0.0])) + self.draw_triad(img, Mw, 0.3, "world", thickness=3) + for pose, kp, lab in ((lpose, lkp, "L-wrist"), (rpose, rkp, "R-wrist")): + if ( + np.linalg.norm(np.asarray(pose)[3:7]) < 1e-6 + or np.abs(kp).sum() < 1e-9 + ): + continue + # keypoints + knuckle-forward reference line + pc = self.cam_xyz(head, kp) + px = self.project(pc) + for p in px: + if not np.isnan(p).any(): + cv2.circle( + img, tuple(np.round(p).astype(int)), 2, (200, 200, 200), -1 + ) + a, b = px[WRIST_ROOT], px[MIDDLE_MCP] + if not (np.isnan(a).any() or np.isnan(b).any()): + cv2.line( + img, + tuple(np.round(a).astype(int)), + tuple(np.round(b).astype(int)), + (255, 255, 255), + 2, + cv2.LINE_AA, + ) + # the wrist-frame triad the pipeline actually uses + M = self.cam_pose_matrix(head, pose) + self.draw_triad(img, M, 0.07, lab) + cv2.putText( + img, + "axes X=red Y=green Z=blue | white line = knuckle-forward", + (6, img.shape[0] - 8), + cv2.FONT_HERSHEY_SIMPLEX, + 0.42, + (255, 255, 255), + 1, + cv2.LINE_AA, + ) + return img + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--episode", required=True) + ap.add_argument( + "--folder", default="/storage/project/r-dxu345-0/shared/egoverseS3ZarrDatasets" + ) + ap.add_argument("--out", default=None) + ap.add_argument("--max-frames", type=int, default=0) + ap.add_argument("--still-frames", type=int, nargs="*", default=[]) + args = ap.parse_args() + + path = ( + args.episode + if args.episode.endswith(".zarr") + else os.path.join(args.folder, f"{args.episode}.zarr") + ) + g = zarr.open_group(path, mode="r") + a = dict(g.attrs) + fps = int(a.get("fps", 30)) + print(f"[viz] {path}\n task={(a.get('task_description') or '')[:80]!r} fps={fps}") + + imgs = g["images.front_1"] + head = np.asarray(g["obs_head_pose"]) + lp, rp = np.asarray(g["left.obs_ee_pose"]), np.asarray(g["right.obs_ee_pose"]) + lk = np.asarray(g["left.obs_keypoints"]).reshape(len(head), 21, 3) + rk = np.asarray(g["right.obs_keypoints"]).reshape(len(head), 21, 3) + nf = min(imgs.shape[0], len(head), args.max_frames or 10**9) + out = ( + args.out + or f"/workspace/EgoVerse/mecka_{os.path.basename(path).replace('.zarr','')}_axes.mp4" + ) + stills = sorted(set(i for i in args.still_frames if 0 <= i < nf)) + + ov = MeckaOverlay() + import imageio.v2 as imageio + from PIL import Image, ImageDraw + + writer = imageio.get_writer(out, fps=fps, macro_block_size=None) + for i in range(nf): + frame = ov.overlay(decode(imgs[i]), head[i], lp[i], rp[i], lk[i], rk[i]) + im = Image.fromarray(frame) + d = ImageDraw.Draw(im) + d.text( + (6, 6), + f"{i}/{nf} {(a.get('task_description') or '')[:60]}", + fill=(255, 255, 0), + ) + frame = np.asarray(im) + if i in stills: + still = f"{os.path.splitext(out)[0]}_f{i}.png" + Image.fromarray(frame).save(still) + print(f"[viz] still -> {still}") + writer.append_data(frame) + writer.close() + print(f"[viz] wrote {nf} frames -> {out}") + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/norm_mecka_all_6d.sbatch b/egomimic/scripts/norm_mecka_all_6d.sbatch new file mode 100644 index 000000000..81515fbb7 --- /dev/null +++ b/egomimic/scripts/norm_mecka_all_6d.sbatch @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH -A gts-dxu345-rl2 +#SBATCH -q inferno +#SBATCH -N1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=24 +#SBATCH --mem=96G +#SBATCH --time=12:00:00 +#SBATCH --job-name=norm_mecka_all_6d +#SBATCH --output=/storage/project/r-dxu345-0/agao81/EgoVerse/logs/norm_mecka_all_6d_%j.log + +# CPU-only: compute 0.1-frac quantile norm stats over all SQL lab=mecka +# episodes (shared PACE mirror already holds the data; the resolver syncs +# only the handful of missing episodes). cartesian_6d: 18D action + proprio. +# +# Output: /storage/project/r-dxu345-0/agao81/norm_stats/mecka_all_6d/norm_stats/norm_stats.json +# Training uses: +# norm_stats.precomputed_norm_path=/storage/project/r-dxu345-0/agao81/norm_stats/mecka_all_6d/norm_stats + +set -euo pipefail +cd /storage/project/r-dxu345-0/agao81/EgoVerse +PY=/storage/home/hcoda1/5/agao81/r-dxu345-0/EgoVerse/emimic/bin/python +export PATH="/storage/home/hcoda1/5/agao81/r-dxu345-0/EgoVerse/emimic/bin:$PATH" +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 + +"$PY" egomimic/scripts/precompute_norm_stats.py \ + --data mecka_all_pi_6d \ + --model pi0.5_bc_mecka_6d \ + --sample-frac 0.1 \ + --num-workers 22 \ + --out /storage/project/r-dxu345-0/agao81/norm_stats/mecka_all_6d diff --git a/egomimic/scripts/precompute_norm_stats.py b/egomimic/scripts/precompute_norm_stats.py new file mode 100644 index 000000000..2b98890a0 --- /dev/null +++ b/egomimic/scripts/precompute_norm_stats.py @@ -0,0 +1,132 @@ +"""Precompute normalization stats on a CPU node, so GPU-node time isn't spent +computing them at training startup. + +Replicates trainHydra's norm loop EXACTLY (same hydra config + data config, +same keymap/transform, same sample_frac/seed) but builds NO model and NO +trainer — it only instantiates the train datasets (which s5cmd-syncs any +missing episodes from S3 as a side effect), infers shapes, computes norm +stats, and caches them. The norm-mode keymap strips camera + annotation keys, +so the stats pass reads only the numeric proprio/action arrays — pure CPU +work, no GPU, no JPEG decode beyond one shape-inference sample. + +Output: /norm_stats/norm_stats.json — point training at it with: + norm_stats.precomputed_norm_path=/norm_stats + +Usage (CPU node, repo root, emimic venv): + python egomimic/scripts/precompute_norm_stats.py \ + --data mecka_all_pi_6d --model pi0.5_bc_mecka_6d \ + --sample-frac 0.1 --num-workers 30 \ + --out /storage/project/r-dxu345-0/agao81/norm_stats/mecka_all_6d +""" + +import argparse +import copy +import os +import time + +import hydra +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import OmegaConf + +import egomimic +from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--config-name", default="train_zarr_cartesian_pi") + ap.add_argument( + "--data", + default="mecka_all_pi_6d", + help="data config group — MUST match the training run's data config", + ) + ap.add_argument( + "--model", + default="pi0.5_bc_mecka_6d", + help="only needed so the config composes; the model is never built", + ) + ap.add_argument("--sample-frac", type=float, default=0.1) + ap.add_argument( + "--max-samples", + type=int, + default=2_000_000, + help="hard cap on collected samples: the (N, 100, 18) float32 action " + "stack plus np.percentile's sort copy is ~2 x N x 7.2KB of RAM, so an " + "uncapped 0.1 frac of the full mecka set (~8.5M) needs ~125GB. 2M " + "samples (~15GB stack) still gives 2M draws per (step, dim) cell.", + ) + ap.add_argument("--num-workers", type=int, default=30) + ap.add_argument( + "--out", + required=True, + help="save_cache_dir; writes /norm_stats/norm_stats.json", + ) + args = ap.parse_args() + + cfg_dir = os.path.join(os.path.dirname(egomimic.__file__), "hydra_configs") + GlobalHydra.instance().clear() + overrides = [ + f"data={args.data}", + f"model={args.model}", + f"norm_stats.sample_frac={args.sample_frac}", + f"norm_stats.num_workers={args.num_workers}", + f"norm_stats.save_cache_dir={args.out}", + "norm_stats.precomputed_norm_path=null", + "seed=42", + ] + with initialize_config_dir(version_base=None, config_dir=cfg_dir): + cfg = compose(config_name=args.config_name, overrides=overrides) + + import lightning as L + + L.seed_everything(cfg.seed, workers=True) + + # Mirrors trainHydra: instantiate train datasets (resolver syncs from S3 + # as needed), then a stats-only MultiDataset computes the norm stats from + # a norm-mode (numerics-only) copy of each dataset. + train_datasets = {} + for dataset_name in cfg.data.train_datasets: + print(f"[precompute] dataset={dataset_name}: instantiating (syncs S3) ...") + train_datasets[dataset_name] = hydra.utils.instantiate( + cfg.data.train_datasets[dataset_name] + ) + + norm_stats = MultiDataset( + state={}, + norm_mode=OmegaConf.select(cfg, "norm_stats.norm_mode", default="quantile"), + ) + norm_stats.populate_from_datasets(train_datasets) + + for dataset_name, dataset in train_datasets.items(): + print(f"[precompute] dataset={dataset_name}: inferring shapes ...") + norm_stats.infer_shapes_from_batch(dataset[0]) + + inst = copy.deepcopy(cfg.data.train_datasets[dataset_name]) + km = OmegaConf.to_container(inst.resolver.key_map, resolve=False) + km["norm_mode"] = True # strips image + annotation keys + inst.resolver.key_map = km + norm_dataset = hydra.utils.instantiate(inst) + + t0 = time.perf_counter() + norm_stats.infer_norm_from_dataset( + norm_dataset, + dataset_name, + sample_frac=args.sample_frac, + max_samples=args.max_samples, + num_workers=args.num_workers, + precomputed_norm_path=None, # force compute + ) + print( + f"[precompute] {dataset_name}: norm computed in " + f"{time.perf_counter() - t0:.1f}s" + ) + + norm_stats.cache_stats(save_cache_dir=args.out) + out_dir = os.path.join(args.out, "norm_stats") + print("\nDONE. Use this in training:") + print(f" norm_stats.precomputed_norm_path={out_dir}") + + +if __name__ == "__main__": + main() diff --git a/egomimic/scripts/smoke_mecka_all_6d.sbatch b/egomimic/scripts/smoke_mecka_all_6d.sbatch new file mode 100644 index 000000000..f9d40cbf4 --- /dev/null +++ b/egomimic/scripts/smoke_mecka_all_6d.sbatch @@ -0,0 +1,62 @@ +#!/bin/bash +#SBATCH -A gts-dxu345-rl2 +#SBATCH -q inferno +#SBATCH -N1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=128G +#SBATCH --gres=gpu:h200:1 +#SBATCH --time=03:00:00 +#SBATCH --job-name=smoke_mecka_6d +#SBATCH --output=/storage/project/r-dxu345-0/agao81/EgoVerse/logs/smoke_mecka_6d_%j.log + +# 1xH200 smoke of the mecka cartesian-6d pi0.5 pipeline with trainer=debug +# (4 epochs, val every 2, 2 train / 3 val batches). +# +# viz gate check: evaluator.viz_every_n_epochs=4 with val at epochs 1 and 3 +# means epoch 1's validation must NOT render video and epoch 3's must +# ((3+1) % 4 == 0) -> exactly one videos/epoch_3 dir proves the gate. +# +# Env knobs: +# DEBUG_EPS - limit resolver to N episodes (default 12; empty = all) +# NORM_PATH - precomputed norm_stats dir; empty = compute fresh at startup + +set -euo pipefail +cd /storage/project/r-dxu345-0/agao81/EgoVerse +PY=/storage/home/hcoda1/5/agao81/r-dxu345-0/EgoVerse/emimic/bin/python +export PATH="/storage/home/hcoda1/5/agao81/r-dxu345-0/EgoVerse/emimic/bin:$PATH" +export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 WANDB_MODE=offline +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 + +DEBUG_EPS="${DEBUG_EPS-12}" +NORM_PATH="${NORM_PATH-}" + +ARGS=( + --config-name=train_zarr_cartesian_pi + data=mecka_all_pi_6d + model=pi0.5_bc_mecka_6d + evaluator=eval_pi_wristframe_6d + trainer=debug + launch_params.gpus_per_node=1 + evaluator.viz_every_n_epochs=4 + model.robomimic_model.config.pytorch_weight_path=/storage/project/r-dxu345-0/agao81/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch + name=smoke_mecka_6d + description=smoke_mecka_all_pi_6d +) +if [[ -n "$DEBUG_EPS" ]]; then + ARGS+=( + "++data.train_datasets.human_bimanual.resolver.debug=$DEBUG_EPS" + # 12-episode subset: default valid_ratio 0.05 floors to 0 valid episodes. + "data.train_datasets.human_bimanual.valid_ratio=0.25" + ) +fi +if [[ -n "$NORM_PATH" ]]; then + ARGS+=("norm_stats.precomputed_norm_path=$NORM_PATH") +else + # inline stats on a tiny fraction: the smoke only needs plausible + # normalizers, not accurate ones — avoids the startup stats cost. + ARGS+=("norm_stats.sample_frac=0.02") +fi + +echo "[smoke] python -u egomimic/trainHydra.py ${ARGS[*]}" +"$PY" -u egomimic/trainHydra.py "${ARGS[@]}" diff --git a/egomimic/scripts/train_mecka_all_6d.sbatch b/egomimic/scripts/train_mecka_all_6d.sbatch new file mode 100644 index 000000000..3a651fcc7 --- /dev/null +++ b/egomimic/scripts/train_mecka_all_6d.sbatch @@ -0,0 +1,77 @@ +#!/bin/bash +#SBATCH -A gts-dxu345-rl2 +#SBATCH -q inferno +#SBATCH -N1 +#SBATCH --ntasks=8 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=515750M +#SBATCH --gres=gpu:h200:8 +#SBATCH --time=72:00:00 +#SBATCH --job-name=pi05_mecka_all_6d +#SBATCH --output=/storage/project/r-dxu345-0/agao81/EgoVerse/logs/pi05_mecka_all_6d_%j.log + +# Full pi0.5 BC on ALL SQL mecka episodes (cartesian_6d, 6D proprio), 8xH200. +# +# Cadence (epoch = limit_train_batches = 100 steps): +# val every 50 epochs (5k steps, 80 batches, no shuffle) +# viz every 100 epochs (10k steps) ckpt every 100 epochs, top-3 by step +# +# Requires the CPU norm job's output (norm_mecka_all_6d.sbatch): +# /storage/project/r-dxu345-0/agao81/norm_stats/mecka_all_6d/norm_stats +# +# Env knobs: +# NORM_PATH - precomputed norm_stats dir (default above) +# CKPT_PATH - resume from a Lightning checkpoint +# LR - optimizer lr (default 5e-5, matching the remote mecka runs) + +set -euo pipefail +cd /storage/project/r-dxu345-0/agao81/EgoVerse +PY=/storage/home/hcoda1/5/agao81/r-dxu345-0/EgoVerse/emimic/bin/python +export PATH="/storage/home/hcoda1/5/agao81/r-dxu345-0/EgoVerse/emimic/bin:$PATH" +export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 +# wandb online (key from apikey.txt); falls back to offline if absent. +if grep -q WANDB_API_KEY apikey.txt 2>/dev/null; then + export WANDB_API_KEY="$(grep -oP 'WANDB_API_KEY=\K.*' apikey.txt)" + export WANDB_MODE=online +elif grep -q api.wandb.ai ~/.netrc 2>/dev/null; then + export WANDB_MODE=online # netrc-authenticated +else + export WANDB_MODE=offline +fi +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 + +NORM_PATH="${NORM_PATH-/storage/project/r-dxu345-0/agao81/norm_stats/mecka_all_6d/norm_stats}" +CKPT_PATH="${CKPT_PATH-}" +LR="${LR-5e-5}" + +if [[ ! -f "$NORM_PATH/norm_stats.json" ]]; then + echo "ERROR: precomputed norm stats not found: $NORM_PATH/norm_stats.json" >&2 + echo "Run egomimic/scripts/norm_mecka_all_6d.sbatch first." >&2 + exit 1 +fi + +ARGS=( + --config-name=train_zarr_cartesian_pi + data=mecka_all_pi_6d + model=pi0.5_bc_mecka_6d + evaluator=eval_pi_wristframe_6d + launch_params.gpus_per_node=8 + "model.optimizer.lr=$LR" + "norm_stats.precomputed_norm_path=$NORM_PATH" + model.robomimic_model.config.pytorch_weight_path=/storage/project/r-dxu345-0/agao81/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch + # val every 50 epochs (5k steps); viz stays 100 (fires every 2nd val); + # ckpt every 100 epochs (10k steps), keep 3 most recent (top-k by step). + trainer.check_val_every_n_epoch=50 + callbacks.model_checkpoint.every_n_epochs=100 + callbacks.model_checkpoint.save_top_k=3 + +callbacks.model_checkpoint.monitor=step + +callbacks.model_checkpoint.mode=max + name=pi05_mecka_all_6d + description=pi05_mecka_all_6d_bs64x8_constlr_wrist6d +) +# Lightning names checkpoints "epoch_epoch=..." — quote so hydra's grammar +# doesn't choke on '='. +[[ -n "$CKPT_PATH" ]] && ARGS+=("ckpt_path='$CKPT_PATH'") + +echo "[train] srun python -u egomimic/trainHydra.py ${ARGS[*]}" +srun "$PY" -u egomimic/trainHydra.py "${ARGS[@]}" diff --git a/egomimic/utils/action_utils.py b/egomimic/utils/action_utils.py index 57602f5a9..a755f4983 100644 --- a/egomimic/utils/action_utils.py +++ b/egomimic/utils/action_utils.py @@ -4,6 +4,11 @@ PI05_CARTESIAN_ACTION_ENCODING_RAW_ROT_6D = "cartesian_ypr_raw_rot6d" PI05_CARTESIAN_ACTION_ENCODING_LEGACY = "legacy_normalized_ypr_rot6d" +# Actions arrive already in xyz+6D(+gripper) layout (the ypr->6D conversion is +# done by the ``CartesianYPRToRot6D`` data transform) and already normalized by +# the standard MultiDataset pipeline. The forward pass only *packs* the +# normalized 6D action into the 32D vector (see ``to32_norm_6d`` below). +PI05_CARTESIAN_ACTION_ENCODING_NORM_ROT_6D = "cartesian_normalized_rot6d" # Bimanual robot Cartesian layout: [x, y, z, yaw, pitch, roll, gripper] x 2. ROBOT_BIMANUAL_CARTESIAN_ROT_DIMS = (3, 4, 5, 10, 11, 12) @@ -243,6 +248,25 @@ def from32_raw_rotation( f"{type(self).__name__} does not support raw-rotation action decoding" ) + def to32_norm_6d(self, actions: torch.Tensor) -> torch.Tensor: + """Pack an already-normalized xyz+6D(+gripper) action into the 32D vector. + + The ypr->6D conversion happens upstream in the ``CartesianYPRToRot6D`` + data transform and the result is normalized by the standard data + pipeline, so this is a pure rearrange (no rotation math, no + normalization). + """ + raise NotImplementedError( + f"{type(self).__name__} does not support normalized-rot6d encoding" + ) + + def from32_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: + """Inverse of :meth:`to32_norm_6d`: extract the normalized xyz+6D(+gripper) + action from the 32D vector (pure rearrange).""" + raise NotImplementedError( + f"{type(self).__name__} does not support normalized-rot6d decoding" + ) + # ============================================================ # ROBOT CONVERTERS @@ -380,7 +404,9 @@ def to20_raw_rotation( ) if normalized_actions is None: if stats is None: - raise ValueError("stats are required when normalized_actions is omitted") + raise ValueError( + "stats are required when normalized_actions is omitted" + ) model_actions = _normalize_robot_bimanual_non_rot( raw_actions, stats, norm_mode ) @@ -449,6 +475,25 @@ def from32_raw_rotation( unnormalize_non_rotation=unnormalize_non_rotation, ) + def to32_norm_6d(self, actions: torch.Tensor) -> torch.Tensor: + # actions: (B,S,20) = [L xyz(3) 6d(6) g(1), R xyz(3) 6d(6) g(1)] — already + # the canonical 32D block layout (left 0..9, right 10..19), just pad. + actions = _ensure_bsd(actions) + if actions.shape[-1] != 20: + raise ValueError( + f"RobotBimanual.to32_norm_6d expected 20-dim, got {actions.shape[-1]}" + ) + return _pad32(actions) + + def from32_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: + actions32 = _ensure_bsd(actions32) + if actions32.shape[-1] < 20: + raise ValueError( + f"RobotBimanual.from32_norm_6d expected >=20 dims, got " + f"{actions32.shape[-1]}" + ) + return actions32[..., 0:20] + # ============================================================ # HUMAN CONVERTERS @@ -545,3 +590,30 @@ def from32(self, actions32: torch.Tensor) -> torch.Tensor: 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) + + def to32_norm_6d(self, actions: torch.Tensor) -> torch.Tensor: + # actions: (B,S,18) = [L xyz(3) 6d(6), R xyz(3) 6d(6)]. 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] != 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_norm_6d(self, actions32: torch.Tensor) -> torch.Tensor: + actions32 = _ensure_bsd(actions32) + 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/egomimic/utils/pose_utils.py b/egomimic/utils/pose_utils.py index f6b6bb7b7..19f6eb472 100644 --- a/egomimic/utils/pose_utils.py +++ b/egomimic/utils/pose_utils.py @@ -129,6 +129,102 @@ def _xyzypr_to_matrix(xyzypr: np.ndarray) -> np.ndarray: return mats +def _ypr_to_rot6d(ypr: np.ndarray) -> np.ndarray: + """Convert euler ypr to the continuous 6D rotation representation. + + args: + ypr: (..., 3) array of [yaw, pitch, roll] (radians, ZYX convention) + returns: + (..., 6) array = first two columns of the rotation matrix, + concatenated as [col0(3), col1(3)]. + + Matches the column convention used by the torch packers in + ``egomimic.utils.action_utils`` (``_ypr_to_matrix`` = Rz@Ry@Rx, and + ``to32`` taking ``R[..., 0]`` / ``R[..., 1]``). + """ + ypr = np.asarray(ypr) + if ypr.shape[-1] != 3: + raise ValueError(f"Expected (..., 3) ypr, got shape {ypr.shape}") + dtype = ypr.dtype if np.issubdtype(ypr.dtype, np.floating) else np.float64 + shape = ypr.shape[:-1] + flat = ypr.reshape(-1, 3).astype(np.float64) + mats = R.from_euler("ZYX", flat, degrees=False).as_matrix() # (N, 3, 3) + six = np.concatenate([mats[:, :, 0], mats[:, :, 1]], axis=-1) # cols 0,1 + return six.reshape(*shape, 6).astype(dtype, copy=False) + + +def _rot6d_to_ypr(six: np.ndarray) -> np.ndarray: + """Inverse of :func:`_ypr_to_rot6d`. + + args: + six: (..., 6) array = [col0(3), col1(3)] of a rotation matrix. + returns: + (..., 3) array of [yaw, pitch, roll] (radians, ZYX convention). + + Reconstructs a proper rotation via Gram-Schmidt (mirroring + ``_reconstruct_R_from_cols`` in ``action_utils``) before extracting euler + angles, so ``_rot6d_to_ypr(_ypr_to_rot6d(ypr)) == ypr``. + """ + six = np.asarray(six) + if six.shape[-1] != 6: + raise ValueError(f"Expected (..., 6) rot6d, got shape {six.shape}") + dtype = six.dtype if np.issubdtype(six.dtype, np.floating) else np.float64 + shape = six.shape[:-1] + flat = six.reshape(-1, 6).astype(np.float64) + c1 = flat[:, 0:3] + c2 = flat[:, 3:6] + eps = 1e-8 + c1n = c1 / np.clip(np.linalg.norm(c1, axis=-1, keepdims=True), eps, None) + proj = np.sum(c2 * c1n, axis=-1, keepdims=True) * c1n + c2o = c2 - proj + c2n = c2o / np.clip(np.linalg.norm(c2o, axis=-1, keepdims=True), eps, None) + c3n = np.cross(c1n, c2n) + mats = np.stack([c1n, c2n, c3n], axis=-1) # columns + ypr = R.from_matrix(mats).as_euler("ZYX", degrees=False) + return ypr.reshape(*shape, 3).astype(dtype, copy=False) + + +# [left arm | right arm]; per-arm blocks are one of: +# ypr: xyz(3) + ypr(3) [+ gripper(1)] +# 6d: xyz(3) + col1(3) + col2(3) [+ gripper(1)] +# Mapping width -> {xyz, rot, grip} channel indices. xyz/grip are bounded, +# linearly-interpolated channels; the rot channels are either Euler (wrap at +# +-pi) or continuous 6D columns (bounded in ~[-1, 1]), so quantile bounds on +# them are meaningless — norm-stat bounds checking consumes this to know which +# channel is which. +BIMANUAL_CARTESIAN_LAYOUTS = { + 12: { # human ypr: [L xyz ypr | R xyz ypr] + "xyz": (0, 1, 2, 6, 7, 8), + "rot": (3, 4, 5, 9, 10, 11), + "grip": (), + }, + 14: { # robot ypr: [L xyz ypr g | R xyz ypr g] + "xyz": (0, 1, 2, 7, 8, 9), + "rot": (3, 4, 5, 10, 11, 12), + "grip": (6, 13), + }, + 18: { # human 6d: [L xyz c1 c2 | R xyz c1 c2] + "xyz": (0, 1, 2, 9, 10, 11), + "rot": (3, 4, 5, 6, 7, 8, 12, 13, 14, 15, 16, 17), + "grip": (), + }, + 20: { # robot 6d: [L xyz c1 c2 g | R xyz c1 c2 g] + "xyz": (0, 1, 2, 10, 11, 12), + "rot": (3, 4, 5, 6, 7, 8, 13, 14, 15, 16, 17, 18), + "grip": (9, 19), + }, +} + + +def bimanual_cartesian_layout(width: int) -> dict | None: + """Index layout for a bimanual cartesian action/proprio vector. + + Returns a dict with ``xyz`` / ``rot`` / ``grip`` index tuples, or ``None`` + if ``width`` is not a recognized native width (12/14 ypr, 18/20 6D). + """ + return BIMANUAL_CARTESIAN_LAYOUTS.get(int(width)) + + def _matrix_to_xyzwxyz(mats: np.ndarray) -> np.ndarray: """ args: diff --git a/egomimic/utils/test_pi05_norm_rot6d.py b/egomimic/utils/test_pi05_norm_rot6d.py new file mode 100644 index 000000000..6fe4b9165 --- /dev/null +++ b/egomimic/utils/test_pi05_norm_rot6d.py @@ -0,0 +1,307 @@ +"""Round-trip tests for the normalized continuous-6D rotation encoding. + +Covers the data transform (ypr <-> 6D) and the converter 32D 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 numpy as np +import pytest +import torch + +from egomimic.rldb.zarr.action_chunk_transforms import ( + CartesianRot6DToYPR, + CartesianYPRToRot6D, +) +from egomimic.utils.action_utils import ( + BaseActionConverter, + HumanBimanualCartesianEuler, + RobotBimanualCartesianEuler, +) +from egomimic.utils.pose_utils import _rot6d_to_ypr, _ypr_to_rot6d + + +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 + + +def test_ypr_rot6d_helpers_round_trip(): + ypr = np.random.default_rng(2).uniform(-1.0, 1.0, size=(7, 3)) + six = _ypr_to_rot6d(ypr) + assert six.shape == (7, 6) + np.testing.assert_allclose(_rot6d_to_ypr(six), ypr, atol=1e-6) + + +@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_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() + six = torch.from_numpy(_eva_ypr_chunk()).float() + six6d = torch.from_numpy( + CartesianYPRToRot6D().transform({"actions_cartesian": six.numpy()})[ + "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) + + +def test_human_bimanual_norm_6d_pack_round_trips_and_zeros_gripper(): + converter = HumanBimanualCartesianEuler() + six6d = torch.from_numpy( + CartesianYPRToRot6D().transform({"actions_cartesian": _aria_ypr_chunk()})[ + "actions_cartesian" + ] + ).float()[None] # (1, T, 18) + + packed = converter.to32_norm_6d(six6d) + assert packed.shape[-1] == 32 + # gripper slots (9, 19) must be zero for human (no gripper signal). + torch.testing.assert_close(packed[..., 9], torch.zeros_like(packed[..., 9])) + torch.testing.assert_close(packed[..., 19], torch.zeros_like(packed[..., 19])) + + 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)} + + +@pytest.mark.parametrize("mode", ["cartesian_6d", "cartesian_wristframe_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" + + +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 _bounds_check_dataset(key: str, width: int): + """Minimal MultiDataset shell exposing _check_bounds with ±1 quantile + bounds on ``key`` for embodiment 0.""" + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + md = MultiDataset.__new__(MultiDataset) + md.norm_stats = { + 0: { + key: { + "quantile_1": np.full(width, -1.0, dtype=np.float32), + "quantile_99": np.full(width, 1.0, dtype=np.float32), + } + } + } + md.zarr_keys = {0: {key: key}} + md._warned_violations = set() + return md + + +@pytest.mark.parametrize("key", ["actions_cartesian", "observations.state.ee_pose"]) +@pytest.mark.parametrize("width,rot_idx,xyz_idx", [(14, 3, 0), (20, 4, 0), (18, 5, 9)]) +def test_bounds_check_ignores_rotation_channels(key, width, rot_idx, xyz_idx): + # Rotation channels (Euler wraps at ±π; 6D columns are ~[-1, 1]) must be + # excluded from quantile bounds checking — matching the remote pipeline — + # while translation/gripper channels are still checked and NaN/Inf still + # rejects the full vector. + md = _bounds_check_dataset(key, width) + arr = np.zeros((5, width), dtype=np.float32) + + arr[2, rot_idx] = 50.0 # far outside ±1, but a rotation channel + assert md._check_bounds({"embodiment": 0, key: arr.copy()}, None, 0, "ep") is None + + bad = arr.copy() + bad[2, xyz_idx] = 50.0 # translation channel out of bounds -> violation + assert md._check_bounds({"embodiment": 0, key: bad}, None, 0, "ep") is not None + + nan = arr.copy() + nan[2, rot_idx] = np.nan # NaN anywhere (even rotation) -> violation + assert md._check_bounds({"embodiment": 0, key: nan}, None, 0, "ep") is not None + + +def test_bounds_check_full_vector_for_other_keys(): + # Keys without the bimanual cartesian layout (or unrecognized widths) keep + # the full-vector check. + md = _bounds_check_dataset("some_other_key", 20) + arr = np.zeros((5, 20), dtype=np.float32) + arr[2, 4] = 50.0 + assert ( + md._check_bounds({"embodiment": 0, "some_other_key": arr}, None, 0, "ep") + is not None + ) + + md16 = _bounds_check_dataset("actions_cartesian", 16) + arr16 = np.zeros((5, 16), dtype=np.float32) + arr16[2, 4] = 50.0 + assert ( + md16._check_bounds({"embodiment": 0, "actions_cartesian": arr16}, None, 0, "ep") + is not None + ) + + +def test_rotate_local_frame_flips_left_wrist_convention(): + # Right-multiplying by Rz(180°) must flip the pose's own x/y axes, keep z + # (knuckle-forward) and the position, skip zero-quat padding rows, and + # handle both (7,) poses and (T, 7) chunks. + from scipy.spatial.transform import Rotation as R + + from egomimic.rldb.zarr.action_chunk_transforms import RotateLocalFrame + + rng = np.random.default_rng(4) + q = R.random(3, random_state=5) + chunk = np.zeros((4, 7)) + chunk[:3, :3] = rng.uniform(-1, 1, size=(3, 3)) + chunk[:3, 3:] = q.as_quat()[:, [3, 0, 1, 2]] # wxyz; row 3 stays zero-padded + + t = RotateLocalFrame(keys=["k"]) + out = t.transform({"k": chunk.copy()})["k"] + + np.testing.assert_allclose(out[:, :3], chunk[:, :3]) # positions unchanged + np.testing.assert_allclose(out[3], np.zeros(7)) # padding untouched + R_old = q.as_matrix() + R_new = R.from_quat(out[:3, [4, 5, 6, 3]]).as_matrix() + np.testing.assert_allclose(R_new[:, :, 0], -R_old[:, :, 0], atol=1e-12) # x flip + np.testing.assert_allclose(R_new[:, :, 1], -R_old[:, :, 1], atol=1e-12) # y flip + np.testing.assert_allclose(R_new[:, :, 2], R_old[:, :, 2], atol=1e-12) # z kept + + single = t.transform({"k": chunk[0].copy()})["k"] + np.testing.assert_allclose(single, out[0], atol=1e-12) + + +def test_fix_mecka_left_wrist_flag_prepends_correction(): + from egomimic.rldb.embodiment.human import Human + from egomimic.rldb.zarr.action_chunk_transforms import RotateLocalFrame + + tl = Human.get_transform_list( + "cartesian_wristframe_6d", stride=1, fix_mecka_left_wrist=True + ) + assert isinstance(tl[0], RotateLocalFrame) + assert set(tl[0].keys) == {"left.action_ee_pose", "left.obs_ee_pose"} + # default off — other vendors' data must be untouched + tl_off = Human.get_transform_list("cartesian_wristframe_6d", stride=1) + assert not isinstance(tl_off[0], RotateLocalFrame) + with pytest.raises(ValueError, match="keypoints"): + Human.get_transform_list("keypoints_headframe_ypr", fix_mecka_left_wrist=True) + + +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 + + 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 get_embodiment_id("human_bimanual") == EMBODIMENT.HUMAN_BIMANUAL.value + assert get_embodiment_id("eva_bimanual") == EMBODIMENT.EVA_BIMANUAL.value + 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)) diff --git a/egomimic/utils/test_wrist6d_roundtrip.py b/egomimic/utils/test_wrist6d_roundtrip.py new file mode 100644 index 000000000..3aa2da013 --- /dev/null +++ b/egomimic/utils/test_wrist6d_roundtrip.py @@ -0,0 +1,482 @@ +"""Independent end-to-end checks of the wrist-frame 6D pipeline. + +The round-trip tests in ``test_pi05_norm_rot6d.py`` compare the pipeline +against itself (``_rot6d_to_ypr(_ypr_to_rot6d(x)) == x``), which a +consistently-wrong convention would pass, and the transform-list tests there +only assert which Transform classes are present. These tests instead run the +REAL transform lists on synthetic world-frame poses and compare every stage +against plain numpy/scipy SE(3) math that shares no code with the pipeline: + + raw world-frame poses + -> Human/Eva.get_transform_list("cartesian_wristframe_6d") (data) + -> quantile normalize (the dataset formula) + -> to32_norm_6d / from32_norm_6d (model I/O) + -> unnormalize + -> _build_*_cartesian_revert_6d_wristframe_transform_list (evaluator) + -> head/cam-frame xyz+ypr == independent inv(T_head) @ T_action +""" + +import json + +import numpy as np +import pytest +import torch +from scipy.spatial.transform import Rotation as R + +from egomimic.rldb.embodiment.embodiment import Embodiment +from egomimic.rldb.zarr.action_chunk_transforms import ( + CartesianYPRToRot6D, + SplitKeys, +) +from egomimic.utils.action_utils import ( + HumanBimanualCartesianEuler, + RobotBimanualCartesianEuler, + _apply_norm_one, + _apply_unnorm_one, + _matrix_to_ypr, + _reconstruct_R_from_cols, + _ypr_to_matrix, +) +from egomimic.utils.pose_utils import _rot6d_to_ypr, _ypr_to_rot6d + +T = 100 + + +# ---------------------------------------------------------------- helpers +def _rng(seed): + return np.random.default_rng(seed) + + +def _rand_pose(rng, scale=1.0): + q = R.random(random_state=int(rng.integers(1 << 31))).as_quat() # xyzw + return np.concatenate([rng.uniform(-scale, scale, 3), q[[3, 0, 1, 2]]]) + + +def _rand_chunk(rng, start): + """Smooth random walk of xyz+quat(wxyz) poses starting AT ``start``.""" + out = np.zeros((T, 7)) + p = start[:3].copy() + r = R.from_quat(start[[4, 5, 6, 3]]) + for t in range(T): + if t > 0: + p = p + rng.normal(0, 0.01, 3) + r = R.from_rotvec(rng.normal(0, 0.05, 3)) * r + q = r.as_quat() + out[t] = np.concatenate([p, q[[3, 0, 1, 2]]]) + return out + + +def _T(p7): + M = np.eye(4) + M[:3, :3] = R.from_quat(p7[[4, 5, 6, 3]]).as_matrix() + M[:3, 3] = p7[:3] + return M + + +def _T_chunk(c): + return np.stack([_T(row) for row in c]) + + +def _xyzypr(M): + return np.concatenate( + [M[..., :3, 3], R.from_matrix(M[..., :3, :3]).as_euler("ZYX")], axis=-1 + ) + + +def _R_of_ypr(ypr): + return R.from_euler("ZYX", ypr).as_matrix() + + +def _apply(transform_list, sample): + s = {k: (v.copy() if isinstance(v, np.ndarray) else v) for k, v in sample.items()} + for t in transform_list: + s = t.transform(s) + return s + + +def _assert_pose12_close(got, ref, atol): + """Compare (..., 12) xyz+ypr vectors: xyz directly, rotation via matrices + (immune to ±π wrap and gimbal ambiguity).""" + got = np.asarray(got, dtype=np.float64) + ref = np.asarray(ref, dtype=np.float64) + for off in (0, 6): + np.testing.assert_allclose(got[..., off : off + 3], ref[..., off : off + 3], atol=atol) + Rg = _R_of_ypr(got[..., off + 3 : off + 6].reshape(-1, 3)) + Rr = _R_of_ypr(ref[..., off + 3 : off + 6].reshape(-1, 3)) + np.testing.assert_allclose(Rg, Rr, atol=atol) + + +def _quantile_stats(x, width): + flat = x.reshape(-1, width) + return { + "quantile_1": np.percentile(flat, 1, axis=0).astype(np.float32), + "quantile_99": np.percentile(flat, 99, axis=0).astype(np.float32), + } + + +def _wrist6d_ref(obs_pose12, act_pose12, side): + """Independent wrist-frame 6D block (T, 9) for one arm from head-frame + xyz+ypr proprio (12,) and actions (T, 12).""" + o = 6 * side + To = np.eye(4) + To[:3, :3] = _R_of_ypr(obs_pose12[o + 3 : o + 6]) + To[:3, 3] = obs_pose12[o : o + 3] + Ta = np.stack([np.eye(4)] * act_pose12.shape[0]) + Ta[:, :3, :3] = _R_of_ypr(act_pose12[:, o + 3 : o + 6]) + Ta[:, :3, 3] = act_pose12[:, o : o + 3] + Tw = np.linalg.inv(To)[None] @ Ta + return np.concatenate([Tw[:, :3, 3], Tw[:, :3, 0], Tw[:, :3, 1]], axis=-1) + + +# ------------------------------------------------ convention cross-checks +def test_torch_ypr_matrix_matches_scipy_zyx(): + """The torch packers and the numpy transforms must agree on what 'ypr' + means: intrinsic Z-Y-X (yaw about z, then pitch about the new y, then + roll about the new x), radians.""" + ypr = _rng(0).uniform(-3.0, 3.0, size=(64, 3)) + R_torch = _ypr_to_matrix(torch.from_numpy(ypr)).numpy() + R_scipy = R.from_euler("ZYX", ypr).as_matrix() + np.testing.assert_allclose(R_torch, R_scipy, atol=1e-12) + # ...and _matrix_to_ypr inverts it on the principal branch. + back = _matrix_to_ypr(torch.from_numpy(R_scipy)).numpy() + np.testing.assert_allclose(R.from_euler("ZYX", back).as_matrix(), R_scipy, atol=1e-12) + # numpy 6D helpers use the same columns as the torch packers. + six = _ypr_to_rot6d(ypr) + np.testing.assert_allclose(six[:, :3], R_scipy[:, :, 0], atol=1e-12) + np.testing.assert_allclose(six[:, 3:], R_scipy[:, :, 1], atol=1e-12) + np.testing.assert_allclose( + R.from_euler("ZYX", _rot6d_to_ypr(six)).as_matrix(), R_scipy, atol=1e-12 + ) + + +def test_gram_schmidt_matches_independent_and_is_proper(): + rng = _rng(1) + c1 = rng.normal(size=(32, 3)) + c2 = rng.normal(size=(32, 3)) + Rt = _reconstruct_R_from_cols(torch.from_numpy(c1), torch.from_numpy(c2)).numpy() + # independent GS + a = c1 / np.linalg.norm(c1, axis=-1, keepdims=True) + b = c2 - (c2 * a).sum(-1, keepdims=True) * a + b = b / np.linalg.norm(b, axis=-1, keepdims=True) + Rn = np.stack([a, b, np.cross(a, b)], axis=-1) + np.testing.assert_allclose(Rt, Rn, atol=1e-12) + np.testing.assert_allclose(np.linalg.det(Rt), 1.0, atol=1e-12) + np.testing.assert_allclose( + Rt @ np.transpose(Rt, (0, 2, 1)), np.broadcast_to(np.eye(3), Rt.shape), atol=1e-12 + ) + + +# ------------------------------------------------------------ human path +@pytest.mark.parametrize("fix_left", [False, True]) +def test_human_wristframe_6d_pipeline_round_trips_to_headframe(fix_left): + from egomimic.rldb.embodiment.human import ( + Human, + _build_human_cartesian_revert_6d_wristframe_transform_list, + ) + + rng = _rng(2 + int(fix_left)) + B = 3 + raws = [] + for _ in range(B): + head, lobs, robs = _rand_pose(rng), _rand_pose(rng), _rand_pose(rng) + raws.append( + { + "obs_head_pose": head, + "left.obs_ee_pose": lobs, + "right.obs_ee_pose": robs, + "left.action_ee_pose": _rand_chunk(rng, lobs), + "right.action_ee_pose": _rand_chunk(rng, robs), + } + ) + fwd = Human.get_transform_list( + "cartesian_wristframe_6d", + stride=1, + fix_mecka_left_wrist=fix_left, + pad_proprio_gripper=True, + ) + outs = [_apply(fwd, r) for r in raws] + act6 = np.stack([o["actions_cartesian"] for o in outs]) + obs6 = np.stack([o["observations.state.ee_pose"] for o in outs]) + assert act6.shape == (B, T, 18) and obs6.shape == (B, 20) + assert set(outs[0]) == {"actions_cartesian", "observations.state.ee_pose"} + assert np.all(obs6[:, [9, 19]] == 0.0) + + # Independent head-frame ground truth. The Rz(180°) fix relabels the + # LEFT hand's local axes (right-multiply), before any frame math. + Rfix = np.eye(4) + Rfix[:3, :3] = R.from_euler("z", np.pi).as_matrix() + gt_act = np.zeros((B, T, 12)) + gt_obs = np.zeros((B, 12)) + for b, r in enumerate(raws): + Th = _T(r["obs_head_pose"]) + for si, side in enumerate(("left", "right")): + Ta = _T_chunk(r[f"{side}.action_ee_pose"]) + To = _T(r[f"{side}.obs_ee_pose"]) + if fix_left and side == "left": + Ta = Ta @ Rfix + To = To @ Rfix + gt_act[b, :, 6 * si : 6 * si + 6] = _xyzypr(np.linalg.inv(Th)[None] @ Ta) + gt_obs[b, 6 * si : 6 * si + 6] = _xyzypr(np.linalg.inv(Th) @ To) + + # (1) proprio = head-frame obs pose, 6D-encoded, grip slots zero + obs_ref = np.stack( + [CartesianYPRToRot6D(action_key="k").transform({"k": g})["k"] for g in gt_obs] + ) + keep = [i for i in range(20) if i not in (9, 19)] + np.testing.assert_allclose(obs6[:, keep], obs_ref, atol=1e-9) + + # (2) actions = each arm's pose in that arm's obs-wrist frame, 6D-encoded + for b in range(B): + for si in range(2): + np.testing.assert_allclose( + act6[b, :, 9 * si : 9 * si + 9], + _wrist6d_ref(gt_obs[b], gt_act[b], si), + atol=1e-9, + ) + # t = 0 is the identity pose exactly (reference IS the obs pose) — the + # bounds-check tolerance in MultiDataset._check_bounds relies on this. + np.testing.assert_allclose(act6[:, 0, [0, 1, 2, 9, 10, 11]], 0.0, atol=1e-12) + + # (3) normalize -> 32D pack -> unpack -> unnormalize is exact + st_act = _quantile_stats(act6, 18) + st_obs = _quantile_stats(obs6, 20) + a_t, o_t = torch.from_numpy(act6).float(), torch.from_numpy(obs6).float() + a_n = _apply_norm_one(a_t, st_act, "quantile") + o_n = _apply_norm_one(o_t, st_obs, "quantile") + conv = HumanBimanualCartesianEuler() + a32 = conv.to32_norm_6d(a_n) + assert a32.shape == (B, T, 32) + assert torch.all(a32[..., [9, 19]] == 0) and torch.all(a32[..., 20:] == 0) + torch.testing.assert_close(conv.from32_norm_6d(a32), a_n, atol=0, rtol=0) + a_un = _apply_unnorm_one(conv.from32_norm_6d(a32), st_act, "quantile") + o_un = _apply_unnorm_one(o_n, st_obs, "quantile") + torch.testing.assert_close(a_un, a_t, atol=1e-5, rtol=0) + torch.testing.assert_close(o_un, o_t, atol=1e-5, rtol=0) + + # (4) evaluator revert (batched, like eval_pi) lands back in head frame + rev = _build_human_cartesian_revert_6d_wristframe_transform_list() + out = Embodiment.apply_transform( + {"actions_cartesian": a_un, "observations.state.ee_pose": o_un}, rev + ) + assert np.asarray(out["actions_cartesian"]).shape == (B, T, 12) + assert np.asarray(out["observations.state.ee_pose"]).shape == (B, 12) + _assert_pose12_close(out["actions_cartesian"], gt_act, atol=1e-4) + _assert_pose12_close(out["observations.state.ee_pose"], gt_obs, atol=1e-4) + + # (5) a noisy (non-orthonormal) prediction still reverts to finite poses + a_noisy = _apply_unnorm_one(a_n + 0.05 * torch.randn_like(a_n), st_act, "quantile") + out_n = Embodiment.apply_transform( + {"actions_cartesian": a_noisy, "observations.state.ee_pose": o_un}, rev + ) + assert np.isfinite(np.asarray(out_n["actions_cartesian"])).all() + + +def test_human_wristframe_actions_are_headframe_invariant(): + """Wrist-relative targets must not depend on the head pose at all.""" + from egomimic.rldb.embodiment.human import Human + + rng = _rng(7) + lobs, robs = _rand_pose(rng), _rand_pose(rng) + raw = { + "left.obs_ee_pose": lobs, + "right.obs_ee_pose": robs, + "left.action_ee_pose": _rand_chunk(rng, lobs), + "right.action_ee_pose": _rand_chunk(rng, robs), + } + fwd = Human.get_transform_list("cartesian_wristframe_6d", stride=1) + a1 = _apply(fwd, {**raw, "obs_head_pose": _rand_pose(rng)})["actions_cartesian"] + a2 = _apply(fwd, {**raw, "obs_head_pose": _rand_pose(rng)})["actions_cartesian"] + np.testing.assert_allclose(a1, a2, atol=1e-12) + + +# -------------------------------------------------------------- eva path +def _eva_extrinsics_variants(): + """Every calibration the checkout knows about. Older trees expose one + ``Eva.EXTRINSICS`` dict; newer ones a keyed ``EVA_EXTRINSICS`` registry + selected via ``get_transform_list(..., extrinsics_key=...)``.""" + import egomimic.rldb.embodiment.eva as eva_mod + + registry = getattr(eva_mod, "EVA_EXTRINSICS", None) + if registry is None: + return [pytest.param(None, eva_mod.Eva.EXTRINSICS, id="default")] + return [pytest.param(k, v, id=k) for k, v in registry.items()] + + +@pytest.mark.parametrize("extrinsics_key,extrinsics", _eva_extrinsics_variants()) +def test_eva_wristframe_6d_pipeline_round_trips_to_camframe(extrinsics_key, extrinsics): + from egomimic.rldb.embodiment.eva import ( + Eva, + _build_eva_cartesian_revert_6d_wristframe_transform_list, + ) + + rng = _rng(11) + B = 3 + raws = [] + for _ in range(B): + lobs, robs = _rand_pose(rng), _rand_pose(rng) + raws.append( + { + "left.obs_ee_pose": lobs, + "right.obs_ee_pose": robs, + "left.cmd_ee_pose": _rand_chunk(rng, lobs), + "right.cmd_ee_pose": _rand_chunk(rng, robs), + "left.obs_gripper": rng.uniform(0, 1, (1,)), + "right.obs_gripper": rng.uniform(0, 1, (1,)), + "left.cmd_gripper": rng.uniform(0, 1, (T, 1)), + "right.cmd_gripper": rng.uniform(0, 1, (T, 1)), + } + ) + kwargs = {} if extrinsics_key is None else {"extrinsics_key": extrinsics_key} + fwd = Eva.get_transform_list("cartesian_wristframe_6d", **kwargs) + outs = [_apply(fwd, r) for r in raws] + act6 = np.stack([o["actions_cartesian"] for o in outs]) + obs6 = np.stack([o["observations.state.ee_pose"] for o in outs]) + assert act6.shape == (B, T, 20) and obs6.shape == (B, 20) + + # independent cam-frame GT: T_cam = inv(E_side) @ T_base + gt_act = np.zeros((B, T, 14)) + gt_obs = np.zeros((B, 14)) + for b, r in enumerate(raws): + for si, side in enumerate(("left", "right")): + Einv = np.linalg.inv(np.asarray(extrinsics[side])) + gt_act[b, :, 7 * si : 7 * si + 6] = _xyzypr(Einv[None] @ _T_chunk(r[f"{side}.cmd_ee_pose"])) + gt_act[b, :, 7 * si + 6] = r[f"{side}.cmd_gripper"][:, 0] + gt_obs[b, 7 * si : 7 * si + 6] = _xyzypr(Einv @ _T(r[f"{side}.obs_ee_pose"])) + gt_obs[b, 7 * si + 6] = r[f"{side}.obs_gripper"][0] + + # the x5Dec13_2 rig calibration is ~6e-9 off orthonormal, hence 1e-7 + obs_ref = np.stack( + [CartesianYPRToRot6D(action_key="k").transform({"k": g})["k"] for g in gt_obs] + ) + np.testing.assert_allclose(obs6, obs_ref, atol=1e-7) + pose_idx = [0, 1, 2, 3, 4, 5, 7, 8, 9, 10, 11, 12] + for b in range(B): + for si in range(2): + ref = _wrist6d_ref(gt_obs[b][pose_idx], gt_act[b][:, pose_idx], si) + got = act6[b, :, 10 * si : 10 * si + 10] + np.testing.assert_allclose(got[:, :9], ref, atol=1e-7) + np.testing.assert_allclose(got[:, 9], gt_act[b, :, 7 * si + 6], atol=1e-12) + + st_act, st_obs = _quantile_stats(act6, 20), _quantile_stats(obs6, 20) + a_t, o_t = torch.from_numpy(act6).float(), torch.from_numpy(obs6).float() + a_n = _apply_norm_one(a_t, st_act, "quantile") + o_n = _apply_norm_one(o_t, st_obs, "quantile") + conv = RobotBimanualCartesianEuler() + a32 = conv.to32_norm_6d(a_n) + torch.testing.assert_close(conv.from32_norm_6d(a32), a_n, atol=0, rtol=0) + a_un = _apply_unnorm_one(conv.from32_norm_6d(a32), st_act, "quantile") + o_un = _apply_unnorm_one(o_n, st_obs, "quantile") + + rev = _build_eva_cartesian_revert_6d_wristframe_transform_list() + out = Embodiment.apply_transform( + {"actions_cartesian": a_un, "observations.state.ee_pose": o_un}, rev + ) + ra = np.asarray(out["actions_cartesian"]) + ro = np.asarray(out["observations.state.ee_pose"]) + assert ra.shape == (B, T, 14) and ro.shape == (B, 14) + _assert_pose12_close(ra[..., pose_idx], gt_act[..., pose_idx], atol=1e-4) + np.testing.assert_allclose(ra[..., [6, 13]], gt_act[..., [6, 13]], atol=1e-5) + _assert_pose12_close(ro[..., pose_idx], gt_obs[..., pose_idx], atol=1e-4) + + +# ------------------------------------------------------ silent-failure guards +def test_split_keys_rejects_width_mismatch(): + """A ypr revert list fed a 6D batch used to slice 'xyz + col0' as + 'xyz + ypr' without complaint; now it must fail loudly.""" + sk = SplitKeys(input_key="k", output_key_list=[("a", 6), ("b", 6)]) + ok = sk.transform({"k": np.zeros((4, 12))}) + assert ok["a"].shape == (4, 6) and ok["b"].shape == (4, 6) + with pytest.raises(ValueError, match="last dim 18"): + sk.transform({"k": np.zeros((4, 18))}) + with pytest.raises(ValueError, match="sums to 12"): + sk.transform({"k": torch.zeros(4, 20)}) + + +def test_ypr_revert_on_6d_batch_fails_loudly(): + """The evaluator/data-config mismatch the guard is for: eval_pi.yaml's + ypr revert applied to a cartesian_wristframe_6d batch.""" + from egomimic.rldb.embodiment.human import ( + _build_human_cartesian_revert_eef_frame_transform_list, + ) + + rev = _build_human_cartesian_revert_eef_frame_transform_list(is_quat=False) + batch = { + "actions_cartesian": torch.zeros(2, T, 18), + "observations.state.ee_pose": torch.zeros(2, 20), + } + with pytest.raises(ValueError, match="SplitKeys"): + Embodiment.apply_transform(batch, rev) + + +def _bounds_dataset(key, width, q_low, q_high): + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + md = MultiDataset.__new__(MultiDataset) + md.norm_stats = { + 0: { + key: { + "quantile_1": np.full(width, q_low, dtype=np.float32), + "quantile_99": np.full(width, q_high, dtype=np.float32), + } + } + } + md.zarr_keys = {0: {key: key}} + md._warned_violations = set() + return md + + +def test_bounds_check_tolerates_roundoff_on_collapsed_bounds(): + """Wrist-frame t=0 cells have bounds [0, 0]; roundoff must not reject.""" + md = _bounds_dataset("actions_cartesian", 18, 0.0, 0.0) + arr = np.zeros((5, 18), dtype=np.float32) + arr[0, 0] = 1e-9 # a roundoff-scale xyz value at a [0, 0] bound + assert md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") is None + arr[0, 0] = 1e-3 # a real violation is still caught + assert md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") is not None + + +def test_bounds_check_warns_once_on_stat_shape_mismatch(caplog): + md = _bounds_dataset("actions_cartesian", 18, -1.0, 1.0) + arr = np.zeros((5, 20), dtype=np.float32) # stats are 18-wide + with caplog.at_level("WARNING"): + assert md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") is None + assert md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 1, "ep") is None + msgs = [r.message for r in caplog.records if "bounds check skipped" in r.message] + assert len(msgs) == 1, msgs + + +def test_precomputed_norm_stats_provenance_is_checked(tmp_path): + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + writer = MultiDataset.__new__(MultiDataset) + writer.norm_mode = "quantile" + writer._norm_run_metadata = None + writer.norm_stats = { + 1: { + "actions_cartesian": {"quantile_1": np.zeros((T, 18)), "quantile_99": np.ones((T, 18))}, + "observations.state.ee_pose": {"quantile_1": np.zeros(20), "quantile_99": np.ones(20)}, + } + } + writer.cache_stats(str(tmp_path)) + path = tmp_path / "norm_stats" / "norm_stats.json" + payload = json.loads(path.read_text()) + assert payload["provenance"]["norm_mode"] == "quantile" + assert payload["provenance"]["stat_shapes"]["1"]["actions_cartesian"] == [T, 18] + + def reader(norm_mode): + md = MultiDataset.__new__(MultiDataset) + md.norm_mode = norm_mode + md.norm_stats = {1: {}} + md._norm_run_metadata = None + return md + + keys = ["actions_cartesian", "observations.state.ee_pose"] + good = reader("quantile") + good._load_precomputed_stats(str(path), 1, keys) + assert set(good.norm_stats[1]) == set(keys) + with pytest.raises(ValueError, match="norm_mode='zscore'"): + reader("zscore")._load_precomputed_stats(str(path), 1, keys) + with pytest.raises(ValueError, match="different keymap/transform mode"): + reader("quantile")._load_precomputed_stats(str(path), 1, ["actions_cartesian"]) + with pytest.raises(ValueError, match="no entry for embodiment id 2"): + reader("quantile")._load_precomputed_stats(str(path), 2, keys)