6D rotation normalization to prevent large changes from euler

egomimic/eval/eval_pi.py

@@ -5,6 +5,7 @@
 
 from egomimic.eval.eval_video import EvalVideo
 from egomimic.rldb.embodiment.embodiment import Embodiment, get_embodiment
+from egomimic.utils.pose_utils import bimanual_cartesian_layout
 
 
 class PIEvalVideo(EvalVideo):
@@ -27,24 +28,31 @@ def compute_metrics_and_viz(self, batch):
         total_loss = None
         n_loss_embodiments = 0
 
-        # Bimanual 12-D layout from `HumanBimanualCartesianEuler.from32`:
-        # [L_xyz(3), L_ypr(3), R_xyz(3), R_ypr(3)]. Split lets us tell a
-        # translation problem apart from a rotation-reconstruction artifact
-        # (6D-cols → matrix → YPR can blow up near gimbal lock / ±π wrap).
+        # Split the bimanual cartesian vector into a translation MSE and a
+        # rotation MSE so a translation problem reads apart from a rotation
+        # one. Handles all four native widths via ``bimanual_cartesian_layout``:
+        #   - native model output: 18D (human) / 20D (robot) continuous 6D cols
+        #     — this metric is clean (6D has no ±π wrap).
+        #   - reverted cam-frame output: 12D (human) / 14D (robot) ypr — the
+        #     rotation MSE here keeps the old gimbal/±π caveat, but it's a
+        #     secondary viz metric.
+        # The metric keys keep the historical ``_ypr_`` name for dashboard
+        # continuity; it denotes "rotation channels" regardless of encoding.
         def _split_mse(pred_t, gt_t):
-            if pred_t.shape[-1] != 12:
+            layout = bimanual_cartesian_layout(pred_t.shape[-1])
+            if layout is None:
                 return None, None
-            xyz_idx = [0, 1, 2, 6, 7, 8]
-            ypr_idx = [3, 4, 5, 9, 10, 11]
+            xyz_idx = list(layout["xyz"])
+            rot_idx = list(layout["rot"])
             xyz = MeanSquaredError()(
                 pred_t[..., xyz_idx].cpu().contiguous(),
                 gt_t[..., xyz_idx].cpu().contiguous(),
             )
-            ypr = MeanSquaredError()(
-                pred_t[..., ypr_idx].cpu().contiguous(),
-                gt_t[..., ypr_idx].cpu().contiguous(),
+            rot = MeanSquaredError()(
+                pred_t[..., rot_idx].cpu().contiguous(),
+                gt_t[..., rot_idx].cpu().contiguous(),
             )
-            return xyz, ypr
+            return xyz, rot
 
         for embodiment_id, _batch in batch.items():
             _batch = algo.norm_stats.unnormalize(_batch, embodiment_id)

egomimic/hydra_configs/data/cotrain_pi_base.yaml

@@ -12,7 +12,7 @@ train_datasets:
         annotation_key: annotations
       transform_list:
         _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list
-        mode: cartesian
+        mode: cartesian_6d
     filters: null
     mode: train
     valid_ratio: 0.05
@@ -27,7 +27,7 @@ train_datasets:
         annotation_key: annotations
       transform_list:
         _target_: egomimic.rldb.embodiment.human.Aria.get_transform_list
-        mode: cartesian
+        mode: cartesian_6d
     filters: null
     mode: train
     valid_ratio: 0.05

egomimic/hydra_configs/data/cotrain_pi_lang_wrist.yaml

@@ -8,10 +8,10 @@ train_datasets:
       key_map:
         keymap_mode: cartesian_wristframe_ypr
       transform_list:
-        mode: cartesian_wristframe_ypr
+        mode: cartesian_wristframe_6d
   aria_bimanual:
     resolver:
       key_map:
         keymap_mode: cartesian
       transform_list:
-        mode: cartesian_wristframe_ypr
+        mode: cartesian_wristframe_6d

egomimic/hydra_configs/data/mecka_pi.yaml

@@ -12,7 +12,7 @@ train_datasets:
         annotation_key: annotations
       transform_list:
         _target_: egomimic.rldb.embodiment.human.Mecka.get_transform_list
-        mode: cartesian
+        mode: cartesian_6d
     filters:
       _target_: egomimic.rldb.filters.DatasetFilter
       filter_lambdas:

egomimic/hydra_configs/data/mecka_pi_10_hrs.yaml

@@ -1,5 +1,8 @@
 _target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper
 
+# Visualization task allowlist consumed by evaluator.tasks by default.
+viz_tasks: [cutting_fabric, hanging_clothes, brushing_shoes]
+
 train_datasets:
   mecka_bimanual:
     _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
@@ -12,11 +15,11 @@ train_datasets:
         annotation_key: annotations
       transform_list:
         _target_: egomimic.rldb.embodiment.human.Mecka.get_transform_list
-        mode: cartesian
+        mode: cartesian_6d
     filters:
       _target_: egomimic.rldb.filters.DatasetFilter
       filter_lambdas:
-        - "lambda row: row['lab'] == 'mecka' and row['task'] in {'packaging_coffee', 'wrapping_gifts', 'cleaning_tools', 'folding_napkins', 'repairing_electronics', 'polishing_jewelry', 'disassembling_phone', 'assembling_flowers', 'making_dumplings', 'peeling_vegetables'}"
+        - "lambda row: row['lab'] == 'mecka' and row['task'] in {'cutting_fabric', 'hanging_clothes', 'brushing_shoes'}"
     mode: train
 
 valid_datasets:

egomimic/hydra_configs/data/mecka_pi_50_hrs.yaml

@@ -12,7 +12,7 @@ train_datasets:
         annotation_key: annotations
       transform_list:
         _target_: egomimic.rldb.embodiment.human.Mecka.get_transform_list
-        mode: cartesian
+        mode: cartesian_6d
     filters:
       _target_: egomimic.rldb.filters.DatasetFilter
       filter_lambdas:

egomimic/hydra_configs/data/mecka_pi_brushing_shoes.yaml

@@ -0,0 +1,51 @@
+_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper
+
+# Visualization task allowlist consumed by evaluator.tasks by default.
+viz_tasks: [brushing_shoes]
+
+train_datasets:
+  mecka_bimanual:
+    _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver
+    resolver:
+      _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver
+      folder_path: ${paths.dataset_dir}
+      key_map:
+        _target_: egomimic.rldb.embodiment.human.Mecka.get_keymap
+        mode: cartesian_pi
+        annotation_key: annotations
+      transform_list:
+        _target_: egomimic.rldb.embodiment.human.Mecka.get_transform_list
+        mode: cartesian_6d
+    filters:
+      _target_: egomimic.rldb.filters.DatasetFilter
+      filter_lambdas:
+        - "lambda row: row['lab'] == 'mecka' and row['task'] == 'brushing_shoes'"
+    mode: train
+
+valid_datasets:
+  mecka_bimanual:
+    _target_: ${data.train_datasets.mecka_bimanual._target_}
+    resolver: ${data.train_datasets.mecka_bimanual.resolver}
+    filters: ${data.train_datasets.mecka_bimanual.filters}
+    mode: valid
+
+train_dataloader_params:
+  mecka_bimanual:
+    batch_size: 64
+    num_workers: 10
+valid_dataloader_params:
+  mecka_bimanual:
+    batch_size: 64
+    num_workers: 10
+
+# `+evaluator@train_viz_evaluator=train_viz_pi`.
+train_viz_datasets:
+  mecka_bimanual:
+    _target_: ${data.train_datasets.mecka_bimanual._target_}
+    resolver: ${data.train_datasets.mecka_bimanual.resolver}
+    filters: ${data.train_datasets.mecka_bimanual.filters}
+    mode: train
+train_viz_dataloader_params:
+  mecka_bimanual:
+    batch_size: 64
+    num_workers: 10

egomimic/hydra_configs/data/mecka_pi_eval.yaml

@@ -16,7 +16,7 @@ train_datasets:
         annotation_key: annotations
       transform_list:
         _target_: egomimic.rldb.embodiment.human.Mecka.get_transform_list
-        mode: cartesian
+        mode: cartesian_6d
     filters:
       _target_: egomimic.rldb.filters.DatasetFilter
       filter_lambdas:

egomimic/hydra_configs/data/mecka_pi_fold_clothes_freeform.yaml

@@ -12,7 +12,7 @@ train_datasets:
         annotation_key: annotations
       transform_list:
         _target_: egomimic.rldb.embodiment.human.Mecka.get_transform_list
-        mode: cartesian
+        mode: cartesian_6d
     filters:
       _target_: egomimic.rldb.filters.DatasetFilter
       filter_lambdas:

egomimic/hydra_configs/evaluator/eval_hpt_wrist.yaml

@@ -12,5 +12,5 @@ transform_lists:
     _target_: egomimic.rldb.embodiment.eva._build_eva_bimanual_revert_eef_frame_transform_list
     is_quat: false
   aria_bimanual:
-    _target_: egomimic.rldb.embodiment.human._build_aria_cartesian_revert_eef_frame_transform_list
+    _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_eef_frame_transform_list
     is_quat: false

egomimic/hydra_configs/evaluator/eval_pi.yaml

@@ -20,7 +20,7 @@ max_frames_per_task: 1000
 transform_lists:
   eva_bimanual:
     _target_: egomimic.rldb.embodiment.eva._build_eva_bimanual_revert_eef_frame_transform_list
-    is_quat: false
+    rot_repr: "6d"
   aria_bimanual:
-    _target_: egomimic.rldb.embodiment.human._build_aria_cartesian_revert_eef_frame_transform_list
-    is_quat: false
+    _target_: egomimic.rldb.embodiment.human._build_human_cartesian_revert_eef_frame_transform_list
+    rot_repr: "6d"

egomimic/hydra_configs/hydra/launcher/submitit_cpu_pace.yaml

@@ -7,7 +7,7 @@ _target_: hydra_plugins.hydra_submitit_launcher.submitit_launcher.SlurmLauncher
 name: ${hydra.job.name}
 partition: "cpu-small"                              # PACE CPU partition — confirm before first submit
 account: "gts-dxu345-rl2"
-cpus_per_task: 24
+cpus_per_task: 12
 nodes: 1
 tasks_per_node: 1
 qos: "inferno"

egomimic/hydra_configs/model/pi0.5_bc_aria.yaml

@@ -14,7 +14,7 @@ robomimic_model:
   action_converters:
     rules:
       ARIA_BIMANUAL:
-        _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler
+        _target_: egomimic.utils.action_utils.HumanBimanualCartesian6D
     # optional fallback if no match is found
     fallback:
       _target_: egomimic.utils.action_utils.BaseActionConverter

egomimic/hydra_configs/model/pi0.5_bc_eva.yaml

@@ -14,7 +14,7 @@ robomimic_model:
   action_converters:
     rules:
       EVA_BIMANUAL:
-        _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler
+        _target_: egomimic.utils.action_utils.RobotBimanualCartesian6D
     # optional fallback if no match is found
     fallback:
       _target_: egomimic.utils.action_utils.BaseActionConverter

egomimic/hydra_configs/model/pi0.5_bc_mecka.yaml

@@ -14,7 +14,7 @@ robomimic_model:
   action_converters:
     rules:
       MECKA_BIMANUAL:
-        _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler
+        _target_: egomimic.utils.action_utils.HumanBimanualCartesian6D
     # optional fallback if no match is found
     fallback:
       _target_: egomimic.utils.action_utils.BaseActionConverter

egomimic/hydra_configs/model/pi0.5_bc_scale.yaml

@@ -14,7 +14,7 @@ robomimic_model:
   action_converters:
     rules:
       SCALE_BIMANUAL:
-        _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler
+        _target_: egomimic.utils.action_utils.HumanBimanualCartesian6D
     # optional fallback if no match is found
     fallback:
       _target_: egomimic.utils.action_utils.BaseActionConverter

egomimic/hydra_configs/model/pi0.5_cotrain_eva_aria.yaml

@@ -20,9 +20,9 @@ robomimic_model:
   action_converters:
     rules:
       EVA_BIMANUAL:
-        _target_: egomimic.utils.action_utils.RobotBimanualCartesianEuler
+        _target_: egomimic.utils.action_utils.RobotBimanualCartesian6D
       ARIA_BIMANUAL:
-        _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler
+        _target_: egomimic.utils.action_utils.HumanBimanualCartesian6D
     # optional fallback if no match is found
     fallback:
       _target_: egomimic.utils.action_utils.BaseActionConverter

egomimic/hydra_configs/model/pi0.5_cotrain_mecka_scale.yaml

@@ -20,9 +20,9 @@ robomimic_model:
   action_converters:
     rules:
       MECKA_BIMANUAL:
-        _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler
+        _target_: egomimic.utils.action_utils.HumanBimanualCartesian6D
       SCALE_BIMANUAL:
-        _target_: egomimic.utils.action_utils.HumanBimanualCartesianEuler
+        _target_: egomimic.utils.action_utils.HumanBimanualCartesian6D
     # optional fallback if no match is found
     fallback:
       _target_: egomimic.utils.action_utils.BaseActionConverter

egomimic/hydra_configs/trainer/ddp_pi.yaml

@@ -7,5 +7,5 @@ accelerator: gpu
 devices: ${eval:'${launch_params.gpus_per_node} * ${launch_params.nodes}'}
 num_nodes: ${launch_params.nodes}
 sync_batchnorm: True
-check_val_every_n_epoch: 150
+check_val_every_n_epoch: 50
 num_sanity_val_steps: 0

egomimic/hydra_configs/trainer/default.yaml

@@ -13,8 +13,8 @@ precision: bf16
 limit_train_batches: 100
 limit_val_batches: 80
 # perform a validation loop every N training epochs
-check_val_every_n_epoch: 100
+check_val_every_n_epoch: 50
 
 # set True to to ensure deterministic results
 # makes training slower but gives more reproducibility than just setting seeds
-deterministic: False
+deterministic: False