diff --git a/egomimic/hydra_configs/data/aria.yaml b/egomimic/hydra_configs/data/aria.yaml index 47428fcd4..34fe9b800 100644 --- a/egomimic/hydra_configs/data/aria.yaml +++ b/egomimic/hydra_configs/data/aria.yaml @@ -11,7 +11,9 @@ train_datasets: keymap_mode: cartesian transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler stride: 3 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/aria_keypoints.yaml b/egomimic/hydra_configs/data/aria_keypoints.yaml index f291b0990..8be34b20f 100644 --- a/egomimic/hydra_configs/data/aria_keypoints.yaml +++ b/egomimic/hydra_configs/data/aria_keypoints.yaml @@ -11,7 +11,9 @@ train_datasets: keymap_mode: keypoints transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: keypoints_headframe_ypr + action_mode: keypoints + coord_frame: camframe + rotation_mode: euler stride: 3 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml b/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml index d0076c101..9caab3c96 100644 --- a/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml +++ b/egomimic/hydra_configs/data/aria_keypoints_wrist.yaml @@ -11,7 +11,9 @@ train_datasets: keymap_mode: keypoints transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: keypoints_wristframe_ypr + action_mode: keypoints + coord_frame: eef_frame + rotation_mode: euler stride: 3 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/aria_pi.yaml b/egomimic/hydra_configs/data/aria_pi.yaml index ff9e2af07..b9944f7a8 100644 --- a/egomimic/hydra_configs/data/aria_pi.yaml +++ b/egomimic/hydra_configs/data/aria_pi.yaml @@ -10,6 +10,9 @@ train_datasets: _target_: egomimic.rldb.embodiment.human.Human.get_keymap transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler stride: 3 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/aria_qwen.yaml b/egomimic/hydra_configs/data/aria_qwen.yaml index 8815fcdff..1d4822662 100644 --- a/egomimic/hydra_configs/data/aria_qwen.yaml +++ b/egomimic/hydra_configs/data/aria_qwen.yaml @@ -16,7 +16,9 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler stride: 3 filters: null mode: train diff --git a/egomimic/hydra_configs/data/bc_pickplace_eva_qwen.yaml b/egomimic/hydra_configs/data/bc_pickplace_eva_qwen.yaml index bfeba77eb..3401d05be 100644 --- a/egomimic/hydra_configs/data/bc_pickplace_eva_qwen.yaml +++ b/egomimic/hydra_configs/data/bc_pickplace_eva_qwen.yaml @@ -17,7 +17,9 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter project_name: dense-language @@ -37,7 +39,9 @@ valid_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter project_name: dense-language diff --git a/egomimic/hydra_configs/data/cotrain_pi_base.yaml b/egomimic/hydra_configs/data/cotrain_pi_base.yaml index 6d682e2fd..f14e1ff12 100644 --- a/egomimic/hydra_configs/data/cotrain_pi_base.yaml +++ b/egomimic/hydra_configs/data/cotrain_pi_base.yaml @@ -12,7 +12,9 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler filters: null mode: train valid_ratio: 0.05 @@ -27,7 +29,9 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler stride: 3 filters: null mode: train diff --git a/egomimic/hydra_configs/data/cotrain_pi_lang_wrist.yaml b/egomimic/hydra_configs/data/cotrain_pi_lang_wrist.yaml index 641e483fa..447de7eb3 100644 --- a/egomimic/hydra_configs/data/cotrain_pi_lang_wrist.yaml +++ b/egomimic/hydra_configs/data/cotrain_pi_lang_wrist.yaml @@ -8,10 +8,14 @@ train_datasets: key_map: keymap_mode: cartesian_wristframe_ypr transform_list: - mode: cartesian_wristframe_ypr + action_mode: cartesian + coord_frame: eef_frame + rotation_mode: euler human_bimanual: resolver: key_map: keymap_mode: cartesian transform_list: - mode: cartesian_wristframe_ypr + action_mode: cartesian + coord_frame: eef_frame + rotation_mode: euler diff --git a/egomimic/hydra_configs/data/cotrain_pi_pickplace_qwen.yaml b/egomimic/hydra_configs/data/cotrain_pi_pickplace_qwen.yaml index 0dcccb548..f8297c29d 100644 --- a/egomimic/hydra_configs/data/cotrain_pi_pickplace_qwen.yaml +++ b/egomimic/hydra_configs/data/cotrain_pi_pickplace_qwen.yaml @@ -16,7 +16,9 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter project_name: dense-language @@ -34,7 +36,9 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler stride: 3 filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter @@ -55,7 +59,9 @@ valid_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter project_name: dense-language diff --git a/egomimic/hydra_configs/data/cotrain_pickplace_qwen.yaml b/egomimic/hydra_configs/data/cotrain_pickplace_qwen.yaml index c05902d4e..ebfd6ba82 100644 --- a/egomimic/hydra_configs/data/cotrain_pickplace_qwen.yaml +++ b/egomimic/hydra_configs/data/cotrain_pickplace_qwen.yaml @@ -25,7 +25,9 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter project_name: dense-language @@ -43,7 +45,9 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler stride: 3 filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter @@ -64,7 +68,9 @@ valid_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter project_name: dense-language diff --git a/egomimic/hydra_configs/data/cotrain_pickplace_qwen_objgen.yaml b/egomimic/hydra_configs/data/cotrain_pickplace_qwen_objgen.yaml index abb5958ea..8d860e5ad 100644 --- a/egomimic/hydra_configs/data/cotrain_pickplace_qwen_objgen.yaml +++ b/egomimic/hydra_configs/data/cotrain_pickplace_qwen_objgen.yaml @@ -20,7 +20,9 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter project_name: dense-language @@ -38,7 +40,9 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_padded + action_mode: cartesian_gripper_padded + coord_frame: camframe + rotation_mode: euler stride: 3 filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter @@ -59,7 +63,9 @@ valid_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter project_name: dense-language @@ -77,7 +83,9 @@ valid_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_padded + action_mode: cartesian_gripper_padded + coord_frame: camframe + rotation_mode: euler stride: 3 filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter diff --git a/egomimic/hydra_configs/data/cotrain_pickplace_qwen_wrist.yaml b/egomimic/hydra_configs/data/cotrain_pickplace_qwen_wrist.yaml index 9ef6424d5..489e7dfe8 100644 --- a/egomimic/hydra_configs/data/cotrain_pickplace_qwen_wrist.yaml +++ b/egomimic/hydra_configs/data/cotrain_pickplace_qwen_wrist.yaml @@ -1,5 +1,5 @@ # Wristframe variant of cotrain_pickplace_qwen.yaml. -# Transform list uses cartesian_wristframe_ypr — actions are in each EEF wrist +# Transform list uses cartesian + eef_frame + euler — actions are in each EEF # frame instead of head/base frame. Action chunk key + dim unchanged # (actions_cartesian, 14D for eva). Pair with viz=cartesian_wrist_lang. @@ -17,7 +17,9 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian_wristframe_ypr + action_mode: cartesian + coord_frame: eef_frame + rotation_mode: euler filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter project_name: dense-language @@ -35,7 +37,9 @@ train_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_wristframe_ypr + action_mode: cartesian + coord_frame: eef_frame + rotation_mode: euler stride: 3 filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter @@ -56,7 +60,9 @@ valid_datasets: annotation_key: annotations transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian_wristframe_ypr + action_mode: cartesian + coord_frame: eef_frame + rotation_mode: euler filters: _target_: egomimic.rldb.filters.ScaleAnnotationDatasetFilter project_name: dense-language diff --git a/egomimic/hydra_configs/data/eva.yaml b/egomimic/hydra_configs/data/eva.yaml index dfe59a883..65dd9060a 100644 --- a/egomimic/hydra_configs/data/eva.yaml +++ b/egomimic/hydra_configs/data/eva.yaml @@ -11,7 +11,9 @@ train_datasets: keymap_mode: cartesian transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler filters: _target_: egomimic.rldb.filters.DatasetFilter filter_lambdas: diff --git a/egomimic/hydra_configs/data/human.yaml b/egomimic/hydra_configs/data/human.yaml index 13cab563e..956bbf517 100644 --- a/egomimic/hydra_configs/data/human.yaml +++ b/egomimic/hydra_configs/data/human.yaml @@ -17,7 +17,9 @@ train_datasets: keymap_mode: cartesian transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler stride: 3 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/industry_eva_pi.yaml b/egomimic/hydra_configs/data/industry_eva_pi.yaml index 5b3728ce3..c7b07a235 100644 --- a/egomimic/hydra_configs/data/industry_eva_pi.yaml +++ b/egomimic/hydra_configs/data/industry_eva_pi.yaml @@ -15,6 +15,9 @@ train_datasets: _target_: egomimic.rldb.embodiment.eva.Eva.get_keymap transform_list: _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler filters: _target_: egomimic.rldb.filters.DatasetFilter filter_lambdas: @@ -30,6 +33,9 @@ train_datasets: has_head_pose: false transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler stride: 1 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/mecka.yaml b/egomimic/hydra_configs/data/mecka.yaml index 12a22ad45..ed911094a 100644 --- a/egomimic/hydra_configs/data/mecka.yaml +++ b/egomimic/hydra_configs/data/mecka.yaml @@ -11,7 +11,9 @@ train_datasets: keymap_mode: cartesian transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler stride: 1 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/mecka_pi.yaml b/egomimic/hydra_configs/data/mecka_pi.yaml index ee6d183b6..3b3f9d1b8 100644 --- a/egomimic/hydra_configs/data/mecka_pi.yaml +++ b/egomimic/hydra_configs/data/mecka_pi.yaml @@ -10,6 +10,9 @@ train_datasets: _target_: egomimic.rldb.embodiment.human.Human.get_keymap transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler stride: 1 filters: lab: "mecka" diff --git a/egomimic/hydra_configs/data/mecka_scale_cotrain_pi.yaml b/egomimic/hydra_configs/data/mecka_scale_cotrain_pi.yaml index fac02bc39..af00732ff 100644 --- a/egomimic/hydra_configs/data/mecka_scale_cotrain_pi.yaml +++ b/egomimic/hydra_configs/data/mecka_scale_cotrain_pi.yaml @@ -16,6 +16,9 @@ train_datasets: has_head_pose: false transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler stride: 1 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/scale.yaml b/egomimic/hydra_configs/data/scale.yaml index 63748b8e6..feb39e158 100644 --- a/egomimic/hydra_configs/data/scale.yaml +++ b/egomimic/hydra_configs/data/scale.yaml @@ -12,7 +12,9 @@ train_datasets: has_head_pose: false transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler stride: 1 filters: _target_: egomimic.rldb.filters.DatasetFilter diff --git a/egomimic/hydra_configs/data/scale_pi.yaml b/egomimic/hydra_configs/data/scale_pi.yaml index d2c5f8417..0bfeb57ea 100644 --- a/egomimic/hydra_configs/data/scale_pi.yaml +++ b/egomimic/hydra_configs/data/scale_pi.yaml @@ -11,6 +11,9 @@ train_datasets: has_head_pose: false transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list + action_mode: cartesian + coord_frame: camframe + rotation_mode: euler stride: 1 filters: filter_lambdas: diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index 4c9cdbb87..128e5c3e5 100644 --- a/egomimic/rldb/embodiment/eva.py +++ b/egomimic/rldb/embodiment/eva.py @@ -8,17 +8,15 @@ from egomimic.rldb.embodiment.human import ARIA_INTRINSICS from egomimic.rldb.zarr.action_chunk_transforms import ( ActionChunkCoordinateFrameTransform, - BatchQuaternionPoseToYPR, ConcatKeys, DeleteKeys, InterpolateLinear, InterpolatePose, NumpyToTensor, PoseCoordinateFrameTransform, - QuaternionPoseToYPR, SplitKeys, Transform, - XYZWXYZ_to_XYZYPR, + transforms_for_rotation_mode, ) from egomimic.utils.pose_utils import ( _matrix_to_xyzwxyz, @@ -48,16 +46,34 @@ class Eva(Embodiment): @staticmethod def get_transform_list( - mode: Literal[ - "cartesian", "cartesian_wristframe_ypr", "cartesian_wristframe_quat" - ], + action_mode: Literal[ + "cartesian", + ] = "cartesian", + coord_frame: Literal[ + "camframe", + "eef_frame", + ] = "camframe", + rotation_mode: Literal[ + "euler", + "quat", + "6D", + ] = "euler", ) -> list[Transform]: - if mode == "cartesian": - return _build_eva_bimanual_transform_list(is_quat=True) - elif mode == "cartesian_wristframe_ypr": - return _build_eva_bimanual_eef_frame_transform_list(is_quat=False) - elif mode == "cartesian_wristframe_quat": - return _build_eva_bimanual_eef_frame_transform_list(is_quat=True) + """``action_mode`` is the action layout; ``coord_frame`` is where poses + live; ``rotation_mode`` is how rotation is stored. + + Cam-frame actions are expressed in the wrist cameras via :attr:`EXTRINSICS`. + EEF-frame actions are a delta from the current EEF pose. In both cases + the geometric hops run in xyz+quat, then ``rotation_mode`` converts rotation + to euler (xyz+ypr, 14D), quat (16D), or Zhou 6D (20D). + """ + if action_mode != "cartesian": + raise ValueError(f"unknown action_mode {action_mode!r}") + if coord_frame == "camframe": + return _build_eva_bimanual_transform_list(rotation_mode=rotation_mode) + if coord_frame == "eef_frame": + return _build_eva_bimanual_eef_frame_transform_list(rotation_mode=rotation_mode) + raise ValueError(f"unknown coord_frame {coord_frame!r}") @classmethod def _get_keymap(cls, keymap_mode: str): @@ -244,7 +260,7 @@ def _build_eva_bimanual_eef_frame_transform_list( obs_key: str = "observations.state.ee_pose", chunk_length: int = 100, stride: int = 1, - is_quat: bool = True, + rotation_mode: Literal["euler", "quat", "6D"] = "euler", ) -> list[Transform]: """EVA bimanual transform pipeline with actions expressed relative to the current EEF pose (wrist frame), analogous to keypoints relative to wrist pose.""" @@ -323,27 +339,17 @@ def _build_eva_bimanual_eef_frame_transform_list( ), ] - if not is_quat: - transform_list.extend( - [ - BatchQuaternionPoseToYPR( - pose_key=left_cmd_wristframe, - output_key=left_cmd_wristframe, - ), - BatchQuaternionPoseToYPR( - pose_key=right_cmd_wristframe, - output_key=right_cmd_wristframe, - ), - QuaternionPoseToYPR( - pose_key=left_obs_camframe, - output_key=left_obs_camframe, - ), - QuaternionPoseToYPR( - pose_key=right_obs_camframe, - output_key=right_obs_camframe, - ), - ] + transform_list.extend( + transforms_for_rotation_mode( + keys=[ + left_cmd_wristframe, + right_cmd_wristframe, + left_obs_camframe, + right_obs_camframe, + ], + rotation_mode=rotation_mode, ) + ) transform_list.extend( [ @@ -408,7 +414,7 @@ def _build_eva_bimanual_transform_list( obs_key: str = "observations.state.ee_pose", chunk_length: int = 100, stride: int = 1, - is_quat: bool = True, + rotation_mode: Literal["euler", "quat", "6D"] = "euler", ) -> list[Transform]: """Canonical EVA bimanual transform pipeline used by tests and notebooks.""" extrinsics = Eva.EXTRINSICS @@ -417,47 +423,46 @@ def _build_eva_bimanual_transform_list( left_extra_batch_key = {"left_extrinsics_pose": left_extrinsics_pose} right_extra_batch_key = {"right_extrinsics_pose": right_extrinsics_pose} - mode = "xyzwxyz" if is_quat else "xyzypr" transform_list = [ ActionChunkCoordinateFrameTransform( target_world=left_target_world, chunk_world=left_cmd_world, transformed_key_name=left_cmd_camframe, extra_batch_key=left_extra_batch_key, - mode=mode, + mode="xyzwxyz", ), ActionChunkCoordinateFrameTransform( target_world=right_target_world, chunk_world=right_cmd_world, transformed_key_name=right_cmd_camframe, extra_batch_key=right_extra_batch_key, - mode=mode, + mode="xyzwxyz", ), PoseCoordinateFrameTransform( target_world=left_target_world, pose_world=left_obs_pose, transformed_key_name=left_obs_pose, - mode=mode, + mode="xyzwxyz", ), PoseCoordinateFrameTransform( target_world=right_target_world, pose_world=right_obs_pose, transformed_key_name=right_obs_pose, - mode=mode, + mode="xyzwxyz", ), InterpolatePose( new_chunk_length=chunk_length, action_key=left_cmd_camframe, output_action_key=left_cmd_camframe, stride=stride, - mode=mode, + mode="xyzwxyz", ), InterpolatePose( new_chunk_length=chunk_length, action_key=right_cmd_camframe, output_action_key=right_cmd_camframe, stride=stride, - mode=mode, + mode="xyzwxyz", ), InterpolateLinear( new_chunk_length=chunk_length, @@ -473,17 +478,17 @@ def _build_eva_bimanual_transform_list( ), ] - if is_quat: - transform_list.append( - XYZWXYZ_to_XYZYPR( - keys=[ - left_cmd_camframe, - right_cmd_camframe, - left_obs_pose, - right_obs_pose, - ] - ) + transform_list.extend( + transforms_for_rotation_mode( + keys=[ + left_cmd_camframe, + right_cmd_camframe, + left_obs_pose, + right_obs_pose, + ], + rotation_mode=rotation_mode, ) + ) transform_list.extend( [ diff --git a/egomimic/rldb/embodiment/human.py b/egomimic/rldb/embodiment/human.py index 673ccc5e0..d817162b9 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -8,17 +8,15 @@ from egomimic.rldb.embodiment.embodiment import Embodiment from egomimic.rldb.zarr.action_chunk_transforms import ( ActionChunkCoordinateFrameTransform, - BatchQuaternionPoseToYPR, ConcatKeys, DeleteKeys, InterpolatePose, PadGripperZeros, PoseCoordinateFrameTransform, - QuaternionPoseToYPR, Reshape, SplitKeys, Transform, - XYZWXYZ_to_XYZYPR, + transforms_for_rotation_mode, ) from egomimic.utils.viz_utils import ( ColorPalette, @@ -98,7 +96,7 @@ class Human(Embodiment): by the data config, because get_keymap / get_transform_list resolve at hydra config time (before any episode is read): - get_keymap(keymap_mode, has_head_pose=True, include_aria_keypoints=False) - - get_transform_list(mode, stride=3) + - get_transform_list(action_mode, coord_frame, rotation_mode, stride=3) Per-episode camera intrinsics travel in ``batch["intrinsics"]`` (from zarr.json); ``cls.INTRINSICS`` is only a fallback for legacy episodes that lack them. The canonical keypoints are MANO for every vendor. @@ -325,45 +323,74 @@ def _get_keymap( @classmethod def get_transform_list( cls, - mode: Literal[ + action_mode: Literal[ "cartesian", - "cartesian_padded", - "cartesian_wristframe_ypr", - "keypoints_headframe_ypr", - "keypoints_headframe_quat", - "keypoints_wristframe_ypr", - "keypoints_wristframe_quat", - ], + "cartesian_gripper_padded", + "keypoints", + ] = "cartesian", + coord_frame: Literal[ + "camframe", + "eef_frame", + ] = "camframe", + rotation_mode: Literal[ + "euler", + "quat", + "6D", + ] = "euler", stride: int = 3, ) -> list[Transform]: - """Transform pipeline. ``stride`` is the per-vendor action stride - (Aria/LightWheel=3, Scale/Mecka=1), supplied by the data config. + """``action_mode`` is the action layout; ``coord_frame`` is where poses + live; ``rotation_mode`` is how rotation is stored. + + ``stride`` is the per-vendor action stride (Aria/LightWheel=3, + Scale/Mecka=1), supplied by the data config. + + Human cartesian has no gripper. ``cartesian_gripper_padded`` inserts a + zero gripper per arm so the layout matches Eva/Yam (14D euler, 16D quat, + 20D Zhou 6D). """ - if mode == "cartesian": - return _build_human_cartesian_bimanual_transform_list(stride=stride) - if mode == "cartesian_padded": - return _build_human_cartesian_bimanual_transform_list( - stride=stride - ) + [PadGripperZeros(action_key="actions_cartesian")] - if mode == "cartesian_wristframe_ypr": - return _build_human_cartesian_eef_frame_transform_list(stride=stride) - if mode == "keypoints_headframe_ypr": - return _build_human_keypoints_bimanual_transform_list( - stride=stride, is_quat=False - ) - if mode == "keypoints_headframe_quat": - return _build_human_keypoints_bimanual_transform_list( - stride=stride, is_quat=True + if action_mode in ("cartesian", "cartesian_gripper_padded"): + builders = { + "camframe": _build_human_cartesian_bimanual_transform_list, + "eef_frame": _build_human_cartesian_eef_frame_transform_list, + } + elif action_mode == "keypoints": + builders = { + "camframe": _build_human_keypoints_bimanual_transform_list, + "eef_frame": _build_human_keypoints_eef_frame_transform_list, + } + else: + raise ValueError( + f"Unsupported action_mode '{action_mode}' for {cls.__name__}" ) - if mode == "keypoints_wristframe_ypr": - return _build_human_keypoints_eef_frame_transform_list( - stride=stride, is_quat=False + if coord_frame not in builders: + raise ValueError( + f"Unsupported coord_frame '{coord_frame}' for {cls.__name__} " + f"action_mode '{action_mode}'" ) - if mode == "keypoints_wristframe_quat": - return _build_human_keypoints_eef_frame_transform_list( - stride=stride, is_quat=True + transform_list = builders[coord_frame]( + stride=stride, rotation_mode=rotation_mode + ) + if action_mode == "cartesian_gripper_padded": + return _pad_human_cartesian_gripper( + transform_list, rotation_mode=rotation_mode ) - raise ValueError(f"Unsupported transform_list mode '{mode}' for {cls.__name__}") + return transform_list + + +def _pad_human_cartesian_gripper( + transform_list: list[Transform], + *, + rotation_mode: Literal["euler", "quat", "6D"], + actions_key: str = "actions_cartesian", + obs_key: str = "observations.state.ee_pose", +) -> list[Transform]: + """Append zero-gripper slots so human cartesian matches Eva/Yam layout.""" + pose_dim = {"euler": 6, "quat": 7, "6D": 9}[rotation_mode] + return transform_list + [ + PadGripperZeros(action_key=actions_key, pose_dim=pose_dim), + PadGripperZeros(action_key=obs_key, pose_dim=pose_dim), + ] # this works for quat and ypr since actionChunkCoordinateFrameTransform works for both @@ -484,7 +511,7 @@ def _build_human_keypoints_eef_frame_transform_list( delete_target_world: bool = True, chunk_length: int = 100, stride: int = 3, - is_quat: bool = True, + rotation_mode: Literal["euler", "quat", "6D"] = "euler", ) -> list[Transform]: transform_list = _build_human_keypoints_bimanual_transform_list( target_world=target_world, @@ -494,7 +521,7 @@ def _build_human_keypoints_eef_frame_transform_list( chunk_length=chunk_length, stride=stride, concat_keys=False, - is_quat=True, + rotation_mode="quat", ) delete_keys = [ left_keypoints_action_world, @@ -596,27 +623,17 @@ def _build_human_keypoints_eef_frame_transform_list( ), ] ) - if not is_quat: - transform_list.extend( - [ - BatchQuaternionPoseToYPR( - pose_key=left_wrist_action_wristframe, - output_key=left_wrist_action_wristframe, - ), - BatchQuaternionPoseToYPR( - pose_key=right_wrist_action_wristframe, - output_key=right_wrist_action_wristframe, - ), - QuaternionPoseToYPR( - pose_key=left_wrist_obs_headframe, - output_key=left_wrist_obs_headframe, - ), - QuaternionPoseToYPR( - pose_key=right_wrist_obs_headframe, - output_key=right_wrist_obs_headframe, - ), - ] + transform_list.extend( + transforms_for_rotation_mode( + keys=[ + left_wrist_action_wristframe, + right_wrist_action_wristframe, + left_wrist_obs_headframe, + right_wrist_obs_headframe, + ], + rotation_mode=rotation_mode, ) + ) transform_list.extend( [ ConcatKeys( @@ -670,7 +687,7 @@ def _build_human_keypoints_bimanual_transform_list( chunk_length: int = 100, stride: int = 3, concat_keys: bool = True, - is_quat: bool = True, + rotation_mode: Literal["euler", "quat", "6D"] = "euler", ) -> list[Transform]: keys_to_delete = list( { @@ -814,28 +831,18 @@ def _build_human_keypoints_bimanual_transform_list( mode="xyzwxyz", ), ] - if not is_quat: + if concat_keys: transform_list.extend( - [ - BatchQuaternionPoseToYPR( - pose_key=left_wrist_action_headframe, - output_key=left_wrist_action_headframe, - ), - BatchQuaternionPoseToYPR( - pose_key=right_wrist_action_headframe, - output_key=right_wrist_action_headframe, - ), - QuaternionPoseToYPR( - pose_key=left_wrist_obs_headframe, - output_key=left_wrist_obs_headframe, - ), - QuaternionPoseToYPR( - pose_key=right_wrist_obs_headframe, - output_key=right_wrist_obs_headframe, - ), - ] + transforms_for_rotation_mode( + keys=[ + left_wrist_action_headframe, + right_wrist_action_headframe, + left_wrist_obs_headframe, + right_wrist_obs_headframe, + ], + rotation_mode=rotation_mode, + ) ) - if concat_keys: transform_list.extend( [ ConcatKeys( @@ -943,14 +950,15 @@ def _build_human_cartesian_eef_frame_transform_list( chunk_length: int = 100, stride: int = 3, delete_target_world: bool = True, + rotation_mode: Literal["euler", "quat", "6D"] = "euler", ) -> list[Transform]: """ARIA bimanual cartesian pipeline expressed in the current wrist frame. Action ee-pose chunks are first transformed world → headframe (via ``obs_head_pose``), then headframe → wristframe (via the proprio ``*.obs_ee_pose_headframe`` for each side). Proprio ee-poses remain in - headframe (wristframe of the wrist itself is identity). All retained poses - are converted to xyz-ypr. + headframe (wristframe of the wrist itself is identity). Rotation is + converted per ``rotation_mode``. """ keys_to_delete = list( { @@ -1018,26 +1026,33 @@ def _build_human_cartesian_eef_frame_transform_list( transformed_key_name=right_action_wristframe, mode="xyzwxyz", ), - XYZWXYZ_to_XYZYPR( + ] + transform_list.extend( + transforms_for_rotation_mode( keys=[ left_action_wristframe, right_action_wristframe, left_obs_headframe, right_obs_headframe, - ] - ), - ConcatKeys( - key_list=[left_action_wristframe, right_action_wristframe], - new_key_name=actions_key, - delete_old_keys=True, - ), - ConcatKeys( - key_list=[left_obs_headframe, right_obs_headframe], - new_key_name=obs_key, - delete_old_keys=True, - ), - DeleteKeys(keys_to_delete=keys_to_delete), - ] + ], + rotation_mode=rotation_mode, + ) + ) + transform_list.extend( + [ + ConcatKeys( + key_list=[left_action_wristframe, right_action_wristframe], + new_key_name=actions_key, + delete_old_keys=True, + ), + ConcatKeys( + key_list=[left_obs_headframe, right_obs_headframe], + new_key_name=obs_key, + delete_old_keys=True, + ), + DeleteKeys(keys_to_delete=keys_to_delete), + ] + ) return transform_list @@ -1059,6 +1074,7 @@ def _build_human_cartesian_bimanual_transform_list( chunk_length: int = 100, stride: int = 3, delete_target_world: bool = True, + rotation_mode: Literal["euler", "quat", "6D"] = "euler", ) -> list[Transform]: """Canonical ARIA bimanual transform pipeline used by tests and notebooks. @@ -1121,17 +1137,17 @@ def _build_human_cartesian_bimanual_transform_list( ), ] - if target_world_is_quat: - transform_list.append( - XYZWXYZ_to_XYZYPR( - keys=[ - left_action_headframe, - right_action_headframe, - left_obs_headframe, - right_obs_headframe, - ] - ) + transform_list.extend( + transforms_for_rotation_mode( + keys=[ + left_action_headframe, + right_action_headframe, + left_obs_headframe, + right_obs_headframe, + ], + rotation_mode=rotation_mode, ) + ) transform_list.extend( [ diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/zarr/action_chunk_transforms.py index 0388d386a..0e2015a1b 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -26,6 +26,7 @@ _interpolate_quat_wxyz, _interpolate_xyz, _matrix_to_xyz, + _matrix_to_xyzrot6d, _matrix_to_xyzwxyz, _matrix_to_xyzypr, _xyz_to_matrix, @@ -366,6 +367,27 @@ def transform(self, batch): return batch +class XYZWXYZ_to_XYZRot6D(Transform): + """Convert listed keys from xyz+quat(wxyz) to xyz+rot6d in-place.""" + + def __init__(self, keys: list[str]): + self.keys = list(keys) + + def transform(self, batch: dict) -> dict: + for key in self.keys: + value = np.asarray(batch[key]) + if value.ndim == 1 and value.shape[0] == 7: + batch[key] = _matrix_to_xyzrot6d(_xyzwxyz_to_matrix(value[None, :]))[0] + elif value.ndim == 2 and value.shape[1] == 7: + batch[key] = _matrix_to_xyzrot6d(_xyzwxyz_to_matrix(value)) + else: + raise ValueError( + f"XYZWXYZ_to_XYZRot6D expects key '{key}' to have shape (7,) " + f"or (T, 7), got {value.shape}" + ) + return batch + + class XYZWXYZ_to_XYZYPR(Transform): """Convert listed keys from xyz+quat(wxyz) to xyz+ypr in-place.""" @@ -387,6 +409,25 @@ def transform(self, batch: dict) -> dict: return batch +def transforms_for_rotation_mode( + keys: list[str], + rotation_mode: Literal["euler", "quat", "6D"], +) -> list[Transform]: + """Convert xyz+quat(wxyz) poses to the requested rotation representation. + + Geometric frame hops always run in quaternion form. This is the last step + that turns those 7D poses into what the policy sees: ``euler`` -> xyz+ypr + (6), ``quat`` -> leave 7D, ``6D`` -> Zhou 6D (xyz + first two columns of R). + """ + if rotation_mode == "quat": + return [] + if rotation_mode == "euler": + return [XYZWXYZ_to_XYZYPR(keys=keys)] + if rotation_mode == "6D": + return [XYZWXYZ_to_XYZRot6D(keys=keys)] + raise ValueError(f"unknown rotation_mode {rotation_mode!r}") + + class CartesianWithGripperCoordinateTransform(Transform): def __init__( self, @@ -513,30 +554,31 @@ def transform(self, batch): class PadGripperZeros(Transform): - """Pad a 12D bimanual cartesian action chunk to 14D by inserting a zero - gripper slot at position 6 (end of left arm) and position 13 (end of right - arm), matching the canonical [L xyz ypr g, R xyz ypr g] layout used by Eva. + """Insert a zero gripper slot after each arm's pose. - Used so aria (which has no gripper signal) can share an FM denoiser head - sized for 14D actions without needing in-model padding branches. + Default ``pose_dim=6`` (xyz+ypr) pads 12D -> 14D. ``pose_dim=9`` + (xyz+Zhou 6D) pads 18D -> 20D so human can share a 20D head with Eva/Yam. """ - def __init__(self, action_key: str = "actions_cartesian"): + def __init__(self, action_key: str = "actions_cartesian", pose_dim: int = 6): self.action_key = action_key + self.pose_dim = int(pose_dim) def transform(self, batch: dict) -> dict: actions = batch[self.action_key] is_tensor = isinstance(actions, torch.Tensor) arr = actions.cpu().numpy() if is_tensor else np.asarray(actions) - if arr.shape[-1] != 12: + expected = 2 * self.pose_dim + if arr.shape[-1] != expected: raise ValueError( - f"PadGripperZeros expects last-dim 12, got {arr.shape} for " - f"'{self.action_key}'" + f"PadGripperZeros expects last-dim {expected} (2 x pose_dim=" + f"{self.pose_dim}), got {arr.shape} for '{self.action_key}'" ) pad_shape = (*arr.shape[:-1], 1) pad = np.zeros(pad_shape, dtype=arr.dtype) + half = self.pose_dim padded = np.concatenate( - (arr[..., :6], pad, arr[..., 6:], pad), axis=-1 + (arr[..., :half], pad, arr[..., half:], pad), axis=-1 ) batch[self.action_key] = ( torch.from_numpy(padded) if is_tensor else padded diff --git a/egomimic/robot/rollout.py b/egomimic/robot/rollout.py index a9d93e485..1001dd3ef 100644 --- a/egomimic/robot/rollout.py +++ b/egomimic/robot/rollout.py @@ -254,7 +254,11 @@ def __init__( self.debug_actions = None self.resampled_action_len = resampled_action_len self.debug = debug - self.transform_list = Eva.get_transform_list(mode="cartesian_wristframe_ypr") + self.transform_list = Eva.get_transform_list( + action_mode="cartesian", + coord_frame="eef_frame", + rotation_mode="euler", + ) self.annotation = None self._tokenizer = None self.collate_fn = default_collate diff --git a/egomimic/scripts/tutorials/aria_to_mano_convert.py b/egomimic/scripts/tutorials/aria_to_mano_convert.py index 0ddfb9833..e96901110 100644 --- a/egomimic/scripts/tutorials/aria_to_mano_convert.py +++ b/egomimic/scripts/tutorials/aria_to_mano_convert.py @@ -3,7 +3,7 @@ Pipeline: 1. Pull first aria episode via SQL. 2. Build MultiDataset with Aria.get_keymap("keypoints") + transform_list - "keypoints_headframe_ypr" -> "actions_keypoints" (left wrist+kp, right wrist+kp). + keypoints + camframe + euler -> "actions_keypoints" (left wrist+kp, right wrist+kp). 3. For the first N frames, stack aria keypoints into a batched tensor and fit MANO_RIGHT / MANO_LEFT to them in parallel (batched Adam). 4. Permute MANO's otaheri joint order (wrist, index*3, middle*3, pinky*3, @@ -78,7 +78,12 @@ def fetch_episode_loader(): print(f"Aria episode: {episode_hash} (of {len(aria_df)})") key_map = Human.get_keymap(keymap_mode="keypoints") - transform_list = Human.get_transform_list(mode="keypoints_headframe_ypr", stride=3) + transform_list = Human.get_transform_list( + action_mode="keypoints", + coord_frame="camframe", + rotation_mode="euler", + stride=3, + ) resolver = S3EpisodeResolver( str(CACHE_DIR), key_map=key_map, transform_list=transform_list ) @@ -94,7 +99,7 @@ def fetch_episode_loader(): def gather_aria_keypoints(loader, n_frames): """Collect first n_frames worth of (image, left_kp, right_kp, batch) tuples. - Aria's batched actions_keypoints layout (after keypoints_headframe_ypr): + Aria's batched actions_keypoints layout (after keypoints + camframe + euler): [left_wrist_xyz(3), left_wrist_ypr(3), left_kp(63), right_wrist_xyz(3), right_wrist_ypr(3), right_kp(63)] = 138 per timestep, shape (T_chunk, 138). We take chunk[0] (current timestep). diff --git a/egomimic/scripts/tutorials/mecka_keypoints_viz.py b/egomimic/scripts/tutorials/mecka_keypoints_viz.py index 65786380c..e62a07c52 100644 --- a/egomimic/scripts/tutorials/mecka_keypoints_viz.py +++ b/egomimic/scripts/tutorials/mecka_keypoints_viz.py @@ -34,7 +34,12 @@ def main() -> None: print(f"Using mecka episode: {episode_hash} (of {len(mecka_df)} candidates)") key_map = Human.get_keymap(keymap_mode="keypoints") - transform_list = Human.get_transform_list(mode="keypoints_headframe_ypr", stride=1) + transform_list = Human.get_transform_list( + action_mode="keypoints", + coord_frame="camframe", + rotation_mode="euler", + stride=1, + ) resolver = S3EpisodeResolver( str(CACHE_DIR), key_map=key_map, transform_list=transform_list diff --git a/egomimic/scripts/tutorials/zarr_data_viz.ipynb b/egomimic/scripts/tutorials/zarr_data_viz.ipynb index a9eb11bc0..a7bc559f0 100644 --- a/egomimic/scripts/tutorials/zarr_data_viz.ipynb +++ b/egomimic/scripts/tutorials/zarr_data_viz.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "7cd73ea6", "metadata": {}, "outputs": [], @@ -23,7 +23,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "id": "32d9110f", "metadata": {}, "outputs": [], @@ -49,7 +49,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "id": "cc9edba1", "metadata": {}, "outputs": [], @@ -66,7 +66,7 @@ "outputs": [], "source": [ "key_map = Eva.get_keymap(keymap_mode=\"cartesian\")\n", - "transform_list = Eva.get_transform_list(mode=\"cartesian\")\n", + "transform_list = Eva.get_transform_list(action_mode=\"cartesian\", coord_frame=\"camframe\", rotation_mode=\"euler\")\n", "\n", "resolver = S3EpisodeResolver(\n", " TEMP_DIR, key_map=key_map, transform_list=transform_list\n", @@ -120,7 +120,7 @@ "metadata": {}, "source": [ "## Human Datasets\n", - "All human data uses the single `Human` embodiment. Per-source differences (head pose, action stride) are passed as args: `Human.get_keymap(keymap_mode=..., has_head_pose=...)`, `Human.get_transform_list(mode=..., stride=...)`." + "All human data uses the single `Human` embodiment. Per-source differences (head pose, action stride) are passed as args: `Human.get_keymap(keymap_mode=..., has_head_pose=...)`, `Human.get_transform_list(action_mode=..., coord_frame=..., rotation_mode=..., stride=...)`." ] }, { @@ -133,7 +133,7 @@ "from egomimic.rldb.zarr.zarr_dataset_multi import LocalEpisodeResolver\n", "\n", "key_map = Human.get_keymap(keymap_mode=\"cartesian\")\n", - "transform_list = Human.get_transform_list(mode=\"cartesian\", stride=3)\n", + "transform_list = Human.get_transform_list(action_mode=\"cartesian\", coord_frame=\"camframe\", rotation_mode=\"euler\", stride=3)\n", "\n", "TEMP_DIR = '/coc/flash7/scratch/egoverseDebugDatasets/egoverseS3DatasetTest/'\n", "resolver = LocalEpisodeResolver(\n", @@ -157,7 +157,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "id": "3cea6d5a", "metadata": {}, "outputs": [], @@ -167,53 +167,20 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "id": "c77c89d6", "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "dict_keys(['observations.images.front_img_1', 'actions_cartesian', 'observations.state.ee_pose', 'embodiment', 'intrinsics', 'episode_hash'])" - ] - }, - "execution_count": 6, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "batch.keys()" ] }, - { - "cell_type": "markdown", - "id": "08fb2784", - "metadata": {}, - "source": [] - }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "id": "af65095a", "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], + "outputs": [], "source": [ "ims = []\n", "for i, batch in enumerate(loader):\n", @@ -227,26 +194,10 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": null, "id": "e6d8d872", "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], + "outputs": [], "source": [ "# Human YPR video (same data loop, YPR overlay)\n", "ims_ypr = []\n", @@ -261,41 +212,13 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": null, "id": "db296b65", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Tables in schema 'app': ['episodes']\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "No episodes matched filters.\n" - ] - }, - { - "ename": "ValueError", - "evalue": "No valid collection names from _get_filtered_paths: filters matched no episodes in the SQL table.", - "output_type": "error", - "traceback": [ - "\u001b[31m---------------------------------------------------------------------------\u001b[39m", - "\u001b[31mValueError\u001b[39m Traceback (most recent call last)", - "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[9]\u001b[39m\u001b[32m, line 10\u001b[39m\n\u001b[32m 2\u001b[39m transform_list = Human.get_transform_list(mode=\u001b[33m\"\u001b[39m\u001b[33mkeypoints_headframe_ypr\u001b[39m\u001b[33m\"\u001b[39m, stride=\u001b[32m3\u001b[39m)\n\u001b[32m 4\u001b[39m resolver = S3EpisodeResolver(\n\u001b[32m 5\u001b[39m TEMP_DIR,\n\u001b[32m 6\u001b[39m key_map=key_map,\n\u001b[32m 7\u001b[39m transform_list=transform_list,\n\u001b[32m 8\u001b[39m )\n\u001b[32m---> \u001b[39m\u001b[32m10\u001b[39m cloudflare_ds = \u001b[43mMultiDataset\u001b[49m\u001b[43m.\u001b[49m\u001b[43m_from_resolver\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m 11\u001b[39m \u001b[43m \u001b[49m\u001b[43mresolver\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfilters\u001b[49m\u001b[43m=\u001b[49m\u001b[43mfilters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msync_from_s3\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmode\u001b[49m\u001b[43m=\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mtotal\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\n\u001b[32m 12\u001b[39m \u001b[43m)\u001b[49m\n\u001b[32m 14\u001b[39m loader = torch.utils.data.DataLoader(cloudflare_ds, batch_size=\u001b[32m1\u001b[39m, shuffle=\u001b[38;5;28;01mFalse\u001b[39;00m)\n", - "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/rl2/projects/EgoVerse/egomimic/rldb/zarr/zarr_dataset_multi.py:973\u001b[39m, in \u001b[36mMultiDataset._from_resolver\u001b[39m\u001b[34m(cls, resolver, **kwargs)\u001b[39m\n\u001b[32m 971\u001b[39m resolved = resolver.resolve(sync_from_s3=sync_from_s3, filters=filters)\n\u001b[32m 972\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m973\u001b[39m resolved = \u001b[43mresolver\u001b[49m\u001b[43m.\u001b[49m\u001b[43mresolve\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfilters\u001b[49m\u001b[43m=\u001b[49m\u001b[43mfilters\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 975\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mcls\u001b[39m(datasets=resolved, **kwargs)\n", - "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/rl2/projects/EgoVerse/egomimic/rldb/zarr/zarr_dataset_multi.py:298\u001b[39m, in \u001b[36mS3EpisodeResolver.resolve\u001b[39m\u001b[34m(self, filters)\u001b[39m\n\u001b[32m 296\u001b[39m valid_hashes = {hashes \u001b[38;5;28;01mfor\u001b[39;00m _, hashes \u001b[38;5;129;01min\u001b[39;00m filtered_paths}\n\u001b[32m 297\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m valid_hashes:\n\u001b[32m--> \u001b[39m\u001b[32m298\u001b[39m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[32m 299\u001b[39m \u001b[33m\"\u001b[39m\u001b[33mNo valid collection names from _get_filtered_paths: \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m 300\u001b[39m \u001b[33m\"\u001b[39m\u001b[33mfilters matched no episodes in the SQL table.\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m 301\u001b[39m )\n\u001b[32m 303\u001b[39m datasets = \u001b[38;5;28mself\u001b[39m._load_zarr_datasets(\n\u001b[32m 304\u001b[39m search_path=\u001b[38;5;28mself\u001b[39m.folder_path,\n\u001b[32m 305\u001b[39m valid_folder_names=valid_hashes,\n\u001b[32m 306\u001b[39m )\n\u001b[32m 308\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m datasets\n", - "\u001b[31mValueError\u001b[39m: No valid collection names from _get_filtered_paths: filters matched no episodes in the SQL table." - ] - } - ], + "outputs": [], "source": [ "key_map = Human.get_keymap(keymap_mode=\"keypoints\")\n", - "transform_list = Human.get_transform_list(mode=\"keypoints_headframe_ypr\", stride=3)\n", + "transform_list = Human.get_transform_list(action_mode=\"keypoints\", coord_frame=\"camframe\", rotation_mode=\"euler\", stride=3)\n", "\n", "resolver = S3EpisodeResolver(\n", " TEMP_DIR,\n", @@ -479,11 +402,191 @@ " break\n", "mpy.show_video(ims_annotations, fps=30)" ] + }, + { + "cell_type": "markdown", + "id": "d54205a1", + "metadata": {}, + "source": [ + "# ABC Data\n", + "\n", + "An ABC-130k episode on the two-arm YAM station, pulled from R2 with\n", + "`S3EpisodeResolver` and loaded with the `Yam` embodiment (bimanual, parallel-jaw,\n", + "like `Eva`).\n", + "\n", + "`Yam.EXTRINSICS` is the top-camera transform, so `coord_frame=\"camframe\"` gives\n", + "camera-frame poses and the overlays work as they do for Eva. ABC's MCAP records\n", + "no extrinsics; this one is composed from the published station model\n", + "(`i2rt-robotics/i2rt`, `robot_models/station/yam_station_*_4310_d405`) and refined\n", + "to the real mount. It covers the **RealSense D405** stations; episodes on the\n", + "wider camera carry intrinsics only, so use an `eef_frame` coord_frame there.\n", + "\n", + "`S3EpisodeResolver` filters against the **SQL episode table** (not local zarr\n", + "attrs), and syncs anything missing into `ABC_DATA_DIR`. Episodes already there\n", + "are skipped, so this is cheap on a warm cache.\n", + "\n", + "This section is self-contained: it does not depend on the Eva/Human cells above." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dae97782", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from pathlib import Path\n", + "\n", + "import imageio_ffmpeg\n", + "import mediapy as mpy\n", + "import numpy as np\n", + "import torch\n", + "import zarr\n", + "\n", + "from egomimic.rldb.embodiment.yam import Yam\n", + "from egomimic.rldb.filters import DatasetFilter\n", + "from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset, S3EpisodeResolver\n", + "from egomimic.utils.aws.aws_data_utils import load_env\n", + "\n", + "mpy.set_ffmpeg(imageio_ffmpeg.get_ffmpeg_exe())\n", + "load_env() # R2 + DB credentials from ~/.egoverse_env\n", + "\n", + "# Downloads land here. Same folder the local conversions live in, so episodes\n", + "# already present are reused rather than re-fetched: the resolver looks for a\n", + "# directory named exactly , and the converter's .zarr\n", + "# stores carry a symlink alongside them.\n", + "ABC_DATA_DIR = \"/coc/flash7/scratch/acheluva3/abc_data/zarr\"\n", + "\n", + "# The only thing to change to look at a different episode. It is the filter the\n", + "# resolver runs against the SQL episode table, and the resolver syncs it down if\n", + "# it is not already in ABC_DATA_DIR.\n", + "EPISODE_HASH = \"c7b01f60-b438-41a2-be41-a2bbb5c0bd0a\"\n", + "print(\"episode:\", EPISODE_HASH)" + ] + }, + { + "cell_type": "markdown", + "id": "9b4de8a6", + "metadata": {}, + "source": [ + "## Subtask annotations\n", + "\n", + "ABC ships subtask labels in a sibling `annotation.mcap`; the converter maps them to\n", + "frame ranges in the `annotations` array." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3d812703", + "metadata": {}, + "outputs": [], + "source": [ + "key_map = Yam.get_keymap(keymap_mode=\"cartesian\", annotation_key=\"annotations\")\n", + "# cartesian + camframe + euler = camera frame, 14D [L xyz ypr g, R xyz ypr g]\n", + "# — the layout the viz helpers expect, matching Eva.\n", + "transform_list = Yam.get_transform_list(action_mode=\"cartesian\", coord_frame=\"camframe\", rotation_mode=\"euler\")\n", + "\n", + "# The overlays only draw on the front camera, so drop the wrist streams: one JPEG\n", + "# decode per sample instead of three, which matters across a multi-thousand-frame\n", + "# episode.\n", + "viz_key_map = {\n", + " k: v\n", + " for k, v in key_map.items()\n", + " if v[\"zarr_key\"] not in (\"images.left_wrist\", \"images.right_wrist\")\n", + "}\n", + "\n", + "# Filters run against the SQL episode table here, so the fields are app.episodes\n", + "# columns rather than the zarr attrs LocalEpisodeResolver matches on. To pull a\n", + "# whole task instead:\n", + "# \"lambda row: row['lab']=='abc' and row['task']=='fold and stack the skirts'\"\n", + "filters = DatasetFilter(\n", + " filter_lambdas=[f\"lambda row: row['episode_hash'] == '{EPISODE_HASH}'\"]\n", + ")\n", + "\n", + "# sync_from_s3=True downloads the episode if ABC_DATA_DIR does not already have it.\n", + "abc_ds = MultiDataset._from_resolver(\n", + " S3EpisodeResolver(\n", + " ABC_DATA_DIR, key_map=viz_key_map, transform_list=transform_list\n", + " ),\n", + " filters=filters,\n", + " sync_from_s3=True,\n", + " mode=\"total\",\n", + ")\n", + "\n", + "# Only now is the store guaranteed on disk. resolves either to a directory\n", + "# the resolver just synced or to the symlink beside a local .zarr.\n", + "store = zarr.open(str(Path(ABC_DATA_DIR) / EPISODE_HASH), mode=\"r\", zarr_format=3)\n", + "T, FPS = store.attrs[\"total_frames\"], store.attrs[\"fps\"]\n", + "print(f\"{store.attrs['task_name']!r} | {T} frames @ {FPS}fps | {store.attrs['embodiment']}\")\n", + "\n", + "# The videos below render only the first PREVIEW_SECONDS at the capture rate.\n", + "PREVIEW_SECONDS = 30\n", + "PREVIEW_FRAMES = min(len(abc_ds), int(PREVIEW_SECONDS * FPS))\n", + "\n", + "loader = torch.utils.data.DataLoader(abc_ds, batch_size=1, shuffle=False)\n", + "print(f\"{len(abc_ds)} samples | preview = first {PREVIEW_FRAMES} frames ({PREVIEW_SECONDS}s)\")" + ] + }, + { + "cell_type": "markdown", + "id": "42c80c12", + "metadata": {}, + "source": [ + "## Subtask annotations\n", + "\n", + "ABC ships subtask labels in a sibling `annotation.mcap`; the converter maps them to\n", + "frame ranges in the `annotations` array." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7b37cf57", + "metadata": {}, + "outputs": [], + "source": [ + "for a in (json.loads(b) for b in store[\"annotations\"][:]):\n", + " span = (a[\"end_idx\"] - a[\"start_idx\"]) / FPS\n", + " print(f\" {a['start_idx']:>5} -> {a['end_idx']:<5} ({span:5.1f}s) {a['text']}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7e511011", + "metadata": {}, + "outputs": [], + "source": [ + "# Separate YPR visualization preview\n", + "for batch in loader:\n", + " vis_ypr = Yam.viz_transformed_batch(batch, mode=\"axes\")\n", + " mpy.show_image(vis_ypr)\n", + " break" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a5f67af4", + "metadata": {}, + "outputs": [], + "source": [ + "images = []\n", + "for i, batch in enumerate(loader):\n", + " vis = Yam.viz_transformed_batch(batch, mode=\"traj+rotation\")\n", + " images.append(vis)\n", + " if i + 1 >= PREVIEW_FRAMES:\n", + " break\n", + "\n", + "mpy.show_video(images, fps=int(FPS))\n" + ] } ], "metadata": { "kernelspec": { - "display_name": "EgoVerse (3.11.15.final.0)", + "display_name": ".venv", "language": "python", "name": "python3" }, @@ -497,7 +600,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.15" + "version": "3.11.13" } }, "nbformat": 4, diff --git a/egomimic/scripts/tutorials/zarr_data_viz.py b/egomimic/scripts/tutorials/zarr_data_viz.py index 23adfaa21..23dfe781c 100644 --- a/egomimic/scripts/tutorials/zarr_data_viz.py +++ b/egomimic/scripts/tutorials/zarr_data_viz.py @@ -19,7 +19,12 @@ load_env() key_map = Human.get_keymap(keymap_mode="cartesian") -transform_list = Human.get_transform_list(mode="cartesian", stride=3) +transform_list = Human.get_transform_list( + action_mode="cartesian", + coord_frame="camframe", + rotation_mode="euler", + stride=3, +) resolver = S3EpisodeResolver(TEMP_DIR, key_map=key_map, transform_list=transform_list) filters = DatasetFilter( diff --git a/egomimic/utils/pose_utils.py b/egomimic/utils/pose_utils.py index f6b6bb7b7..736c6ce49 100644 --- a/egomimic/utils/pose_utils.py +++ b/egomimic/utils/pose_utils.py @@ -171,6 +171,45 @@ def _xyzwxyz_to_matrix(xyzwxyz: np.ndarray) -> np.ndarray: return mats +def _matrix_to_xyzrot6d(mats: np.ndarray) -> np.ndarray: + """SE3 matrices -> xyz + Zhou 6D rotation (first two columns of R). + + Dropping the last column of the 3x3 leaves two 3-vectors, stacked + column-major as ``[r00, r10, r20, r01, r11, r21]``. + """ + if mats.ndim != 3 or mats.shape[-2:] != (4, 4): + raise ValueError(f"Expected (B, 4, 4) array, got shape {mats.shape}") + mats = np.asarray(mats) + dtype = mats.dtype if np.issubdtype(mats.dtype, np.floating) else np.float64 + xyz = mats[:, :3, 3] + # First two columns, each 3D, concatenated — not a C-order reshape of (3, 2), + # which would interleave row-wise. + rot6d = np.concatenate([mats[:, :3, 0], mats[:, :3, 1]], axis=-1) + return np.concatenate([xyz, rot6d], axis=-1).astype(dtype, copy=False) + + +def _xyzrot6d_to_matrix(xyzrot6d: np.ndarray) -> np.ndarray: + """Inverse of ``_matrix_to_xyzrot6d``: Gram-Schmidt the two columns, third = cross.""" + if xyzrot6d.ndim != 2 or xyzrot6d.shape[-1] != 9: + raise ValueError(f"Expected (B, 9) array, got shape {xyzrot6d.shape}") + B = xyzrot6d.shape[0] + dtype = xyzrot6d.dtype if np.issubdtype(xyzrot6d.dtype, np.floating) else np.float64 + a1 = xyzrot6d[:, 3:6] + a2 = xyzrot6d[:, 6:9] + n1 = np.linalg.norm(a1, axis=-1, keepdims=True) + b1 = a1 / np.clip(n1, 1e-12, None) + a2_proj = a2 - np.sum(b1 * a2, axis=-1, keepdims=True) * b1 + n2 = np.linalg.norm(a2_proj, axis=-1, keepdims=True) + b2 = a2_proj / np.clip(n2, 1e-12, None) + b3 = np.cross(b1, b2) + mats = np.broadcast_to(np.eye(4, dtype=dtype), (B, 4, 4)).copy() + mats[:, :3, 0] = b1 + mats[:, :3, 1] = b2 + mats[:, :3, 2] = b3 + mats[:, :3, 3] = xyzrot6d[:, :3] + return mats + + def T_rot_orientation(T: np.ndarray, rot_orientation: np.ndarray) -> np.ndarray: """ Permute the rotation matrix of a SE(3) transformation.