diff --git a/egomimic/hydra_configs/data/industry_eva_pi.yaml b/egomimic/hydra_configs/data/industry_eva_pi.yaml deleted file mode 100644 index 34d2e421ef..0000000000 --- a/egomimic/hydra_configs/data/industry_eva_pi.yaml +++ /dev/null @@ -1,70 +0,0 @@ -_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper - -# Embodiment collapse: mecka + scale are both `human_bimanual` now, merged into ONE -# human_bimanual domain alongside eva_bimanual. Scale exports carry obs_head_pose -# (sfs_to_egoverse_zarr.py writes it from the per-frame camera pose), so the merged -# human keymap keeps the default has_head_pose=true. eva is unchanged. -train_datasets: - eva_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.eva.Eva.get_keymap - keymap_mode: cartesian - transform_list: - _target_: egomimic.rldb.embodiment.eva.Eva.get_transform_list - mode: cartesian_wristframe_6d - filters: - _target_: egomimic.rldb.filters.DatasetFilter - filter_lambdas: - - "lambda row: row['embodiment'] == 'eva_bimanual' and row['task'] == 'fold_clothes'" - mode: train - human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver - folder_path: ${paths.dataset_dir} - key_map: - _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian - transform_list: - _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d - # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 - # State: prompt bins align with the robot 20-D layout on every config. - pad_proprio_gripper: true - stride: 1 - filters: - _target_: egomimic.rldb.filters.DatasetFilter - filter_lambdas: - - "lambda row: (row['lab'] == 'mecka' and row['task'] == 'fold_clothes') or (row['lab'] == 'scale' and row['task'] == 'flagship_fold_clothes')" - mode: train - -valid_datasets: - eva_bimanual: - _target_: ${data.train_datasets.eva_bimanual._target_} - resolver: ${data.train_datasets.eva_bimanual.resolver} - filters: ${data.train_datasets.eva_bimanual.filters} - mode: valid - human_bimanual: - _target_: ${data.train_datasets.human_bimanual._target_} - resolver: ${data.train_datasets.human_bimanual.resolver} - filters: ${data.train_datasets.human_bimanual.filters} - mode: valid - -train_dataloader_params: - eva_bimanual: - batch_size: 21 - num_workers: 10 - human_bimanual: - batch_size: 21 - num_workers: 10 -valid_dataloader_params: - eva_bimanual: - batch_size: 21 - num_workers: 10 - human_bimanual: - batch_size: 21 - num_workers: 10 diff --git a/egomimic/hydra_configs/data/mecka.yaml b/egomimic/hydra_configs/data/mecka.yaml index 65befd7318..3fe7f6c6ca 100644 --- a/egomimic/hydra_configs/data/mecka.yaml +++ b/egomimic/hydra_configs/data/mecka.yaml @@ -12,6 +12,7 @@ train_datasets: transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list mode: cartesian_6d + fix_left_wrist_convention: true # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 # State: prompt bins align with the robot 20-D layout on every config. pad_proprio_gripper: true diff --git a/egomimic/hydra_configs/data/mecka_scale_cotrain_pi.yaml b/egomimic/hydra_configs/data/mecka_scale_cotrain_pi.yaml deleted file mode 100644 index 1571b3a9ca..0000000000 --- a/egomimic/hydra_configs/data/mecka_scale_cotrain_pi.yaml +++ /dev/null @@ -1,44 +0,0 @@ -_target_: egomimic.pl_utils.pl_data_utils.MultiDataModuleWrapper - -# Embodiment collapse: mecka + scale are both `human_bimanual` now, so they can no -# longer be separate cotrain domains. Merged into ONE human_bimanual dataset that -# loads both labs. Scale exports carry obs_head_pose (sfs_to_egoverse_zarr.py -# writes it from the per-frame camera pose), so the merged human keymap keeps -# the default has_head_pose=true; the cartesian transform pivots on it. -train_datasets: - human_bimanual: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.MultiDataset._from_resolver - resolver: - _target_: egomimic.rldb.zarr.zarr_dataset_multi.S3EpisodeResolver - folder_path: ${paths.dataset_dir} - key_map: - _target_: egomimic.rldb.embodiment.human.Human.get_keymap - keymap_mode: cartesian - transform_list: - _target_: egomimic.rldb.embodiment.human.Human.get_transform_list - mode: cartesian_6d - # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 - # State: prompt bins align with the robot 20-D layout on every config. - pad_proprio_gripper: true - stride: 1 - filters: - _target_: egomimic.rldb.filters.DatasetFilter - filter_lambdas: - - "lambda row: (row['lab'] == 'mecka' and row['task'] == 'fold_clothes') or (row['lab'] == 'scale' and row['task'] == 'flagship_fold_clothes')" - mode: train - -valid_datasets: - human_bimanual: - _target_: ${data.train_datasets.human_bimanual._target_} - resolver: ${data.train_datasets.human_bimanual.resolver} - filters: ${data.train_datasets.human_bimanual.filters} - mode: valid - -train_dataloader_params: - human_bimanual: - batch_size: 32 - num_workers: 10 -valid_dataloader_params: - human_bimanual: - batch_size: 32 - num_workers: 10 diff --git a/egomimic/hydra_configs/data/scale.yaml b/egomimic/hydra_configs/data/scale.yaml index 97ce8f0eae..81a486e8cc 100644 --- a/egomimic/hydra_configs/data/scale.yaml +++ b/egomimic/hydra_configs/data/scale.yaml @@ -12,6 +12,10 @@ train_datasets: transform_list: _target_: egomimic.rldb.embodiment.human.Human.get_transform_list mode: cartesian_6d + # Scale exports were written by the pre-fix hand-pose converter too: + # measured on all 20 production episodes, the left wrist matches + # rot_left to <= 2.06 deg with the right-hand control clean. + fix_left_wrist_convention: true # 18 -> 20-D ee_pose proprio (zero grip slots at 9/19) so the pi0.5 # State: prompt bins align with the robot 20-D layout on every config. pad_proprio_gripper: true diff --git a/egomimic/hydra_configs/model/pi0.5_cotrain_mecka_scale.yaml b/egomimic/hydra_configs/model/pi0.5_cotrain_mecka_scale.yaml deleted file mode 100644 index 79ed54e8f2..0000000000 --- a/egomimic/hydra_configs/model/pi0.5_cotrain_mecka_scale.yaml +++ /dev/null @@ -1,17 +0,0 @@ -defaults: - - pi0.5_base - - _self_ - -robomimic_model: - camera_transforms: {} - ac_keys: - human_bimanual: "actions_cartesian" - domains: ["human_bimanual"] - - 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/rldb/embodiment/human.py b/egomimic/rldb/embodiment/human.py index f18dd0e59a..cadef181c6 100644 --- a/egomimic/rldb/embodiment/human.py +++ b/egomimic/rldb/embodiment/human.py @@ -21,6 +21,7 @@ PoseCoordinateFrameTransform, QuaternionPoseToYPR, Reshape, + RotateLocalFrame, SplitKeys, Transform, UnpadGripperZeros, @@ -373,6 +374,7 @@ def get_transform_list( "keypoints_wristframe_quat", ], stride: int = 3, + fix_left_wrist_convention: bool = False, pad_proprio_gripper: bool = False, allow_legacy_rotation: bool = False, ) -> list[Transform]: @@ -383,12 +385,37 @@ def get_transform_list( the modes in ``LEGACY_ROTATION_MODES``; only the data/rollout boundary (viz, a checkpoint that predates the 6D conversion) may set it. + ``fix_left_wrist_convention`` retroactively corrects the LEFT + wrist-frame convention of zarrs written before the ``rot_left`` fix in + ``mecka_to_zarr.compute_hand_pose_xyzquat`` (which double-mirrored the + left hand onto the right hand's spatial convention): the raw left + pose keys (ee_pose in cartesian modes, wrist_pose in keypoints modes; + the converter writes the same rotation into both) are right-multiplied + by Rz(180 deg) before any frame math, exactly equivalent to + reconverting. It is a property of the DATA, not of the vendor: enable + it for any zarr written by the pre-fix hand-pose converter (measured: + every mecka and scale episode in the production set), and disable it + for data reconverted after the fix, which would double-flip. Aria and + LightWheel are not covered by that measurement -- no such episodes are + in the production set -- so leave them unset. + ``left_wrist_convention`` re-derives the verdict from a zarr's own + keypoints when a setting has to be checked. + ``pad_proprio_gripper`` pads the cartesian proprio ``observations.state.ee_pose`` 18 -> 20 (zero grip slots at 9/19) so the pi0.5 ``State:`` prompt bins align positionally with the robot 20D layout. 6D modes only. """ _reject_legacy_rotation(cls, mode, allow_legacy_rotation) + is_keypoints = mode.startswith("keypoints") + prefix: list[Transform] = [] + if fix_left_wrist_convention: + fix_keys = ( + ["left.action_wrist_pose", "left.obs_wrist_pose"] + if is_keypoints + else ["left.action_ee_pose", "left.obs_ee_pose"] + ) + prefix = [RotateLocalFrame(keys=fix_keys)] if pad_proprio_gripper and not mode.endswith("_6d"): raise ValueError( "pad_proprio_gripper needs a 6D-encoded ee_pose proprio: a " @@ -402,7 +429,9 @@ def get_transform_list( proprio_6d.append(PadGripperZeros(action_key="observations.state.ee_pose")) if mode == "cartesian": - return _build_human_cartesian_bimanual_transform_list(stride=stride) + return prefix + _build_human_cartesian_bimanual_transform_list( + stride=stride + ) if mode == "cartesian_6d": # Head/camera-frame cartesian (12D xyz+ypr per arm pair) with the # rotation re-expressed as the continuous 6D representation (18D) @@ -410,39 +439,45 @@ def get_transform_list( # saturates yaw/roll at +-pi, so per-dim normalization needs the # continuous representation. return ( - _build_human_cartesian_bimanual_transform_list(stride=stride) + prefix + + _build_human_cartesian_bimanual_transform_list(stride=stride) + [CartesianYPRToRot6D(action_key="actions_cartesian")] + proprio_6d ) if mode == "cartesian_padded": - return _build_human_cartesian_bimanual_transform_list(stride=stride) + [ - PadGripperZeros(action_key="actions_cartesian") - ] + 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 with the 6D rotation (18D). The head-frame # proprio ee_pose is 6D-encoded too (see cartesian_6d); it is the # only head-frame signal the model sees with wrist-relative targets. return ( - _build_human_cartesian_eef_frame_transform_list(stride=stride) + prefix + + _build_human_cartesian_eef_frame_transform_list(stride=stride) + [CartesianYPRToRot6D(action_key="actions_cartesian")] + proprio_6d ) if mode == "keypoints_headframe_ypr": - return _build_human_keypoints_bimanual_transform_list( + return prefix + _build_human_keypoints_bimanual_transform_list( stride=stride, is_quat=False ) if mode == "keypoints_headframe_quat": - return _build_human_keypoints_bimanual_transform_list( + return prefix + _build_human_keypoints_bimanual_transform_list( stride=stride, is_quat=True ) if mode == "keypoints_wristframe_ypr": - return _build_human_keypoints_eef_frame_transform_list( + return prefix + _build_human_keypoints_eef_frame_transform_list( stride=stride, is_quat=False ) if mode == "keypoints_wristframe_quat": - return _build_human_keypoints_eef_frame_transform_list( + return prefix + _build_human_keypoints_eef_frame_transform_list( stride=stride, is_quat=True ) raise ValueError(f"Unsupported transform_list mode '{mode}' for {cls.__name__}") diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/zarr/action_chunk_transforms.py index 158f7efd2c..62a0f91452 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -492,6 +492,50 @@ def transform(self, batch: dict) -> dict: 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 deg) (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 CartesianWithGripperCoordinateTransform(Transform): def __init__( self, diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index bef29cf5d4..833f11741f 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -407,10 +407,12 @@ def __init__( transform_list: list | None = None, debug: int | bool | None = None, norm_stats: dict | None = None, + validate_wrist_convention: bool = True, ): self.bucket_name = bucket_name self.main_prefix = main_prefix self.debug = debug + self.validate_wrist_convention = validate_wrist_convention super().__init__( folder_path, key_map=key_map, @@ -450,7 +452,72 @@ def resolve( "No valid collection names from _get_filtered_paths: " "filters matched no episodes in the SQL table." ) - return self.load(paths) + datasets = self.load(paths) + if self.validate_wrist_convention: + self._check_left_wrist_convention(datasets) + return datasets + + WRIST_CONVENTION_SAMPLE = 3 + _WRIST_KEYS = ( + "left.obs_keypoints", + "left.obs_wrist_pose", + "right.obs_keypoints", + "right.obs_wrist_pose", + ) + + def _check_left_wrist_convention(self, datasets: dict) -> None: + """Cross-check ``fix_left_wrist_convention`` against the episodes. + + The flag is a property of the data, not of the code: it is right today + only because nobody has reconverted, and a wrong setting produces a + plausible mirrored frame rather than an error. A few sampled episodes + re-derive the convention from their own keypoints. + + Only the double-flip raises. Enabled-but-post-fix means the correction + is being applied to data that was already corrected, which is a fresh + error introduced by a stale setting. Disabled-but-pre-fix is the state + every run was in before this commit -- wrong, but not newly wrong, and + the generic human configs (aria, human, the cotrains) have never set + the flag; turning their startup into a crash is not this correction's + business. It warns, naming the knob. An episode that matches neither + convention (a handful of 2025 ones do) warns too. + """ + from egomimic.rldb.zarr.action_chunk_transforms import RotateLocalFrame + from egomimic.scripts.mecka_process.mecka_to_zarr import left_wrist_convention + + fix_enabled = any( + isinstance(t, RotateLocalFrame) for t in (self.transform_list or []) + ) + for name, ds in list(datasets.items())[: self.WRIST_CONVENTION_SAMPLE]: + try: + store = zarr.open_group(str(ds.episode_path), mode="r") + if any(k not in store for k in self._WRIST_KEYS): + continue + verdict = left_wrist_convention(store) + except Exception as e: # a probe must never take a run down + logger.warning(f"[wrist-convention] {name}: probe failed ({e})") + continue + if verdict == "unknown": + logger.warning( + f"[wrist-convention] {name}: neither convention matches its own " + "keypoints (right-hand control failed); cannot check " + f"fix_left_wrist_convention={fix_enabled} against it" + ) + continue + if verdict == "post_fix" and fix_enabled: + raise ValueError( + f"[wrist-convention] {name} was written post_fix, but this " + "dataset runs with fix_left_wrist_convention=true: the " + "Rz(180) correction would double-flip data the converter " + "already fixed. Drop the flag for reconverted data." + ) + if verdict == "pre_fix" and not fix_enabled: + logger.warning( + f"[wrist-convention] {name} was written pre_fix (its LEFT " + "wrist is double-mirrored) and this dataset does not set " + "fix_left_wrist_convention: it trains on the uncorrected " + "frame. Set the flag on this data config unless you mean to." + ) @staticmethod def _get_filtered_paths( diff --git a/egomimic/scripts/mecka_process/mecka_to_zarr.py b/egomimic/scripts/mecka_process/mecka_to_zarr.py index fd0f605b52..f00cebe38e 100644 --- a/egomimic/scripts/mecka_process/mecka_to_zarr.py +++ b/egomimic/scripts/mecka_process/mecka_to_zarr.py @@ -15,7 +15,7 @@ import time import traceback from pathlib import Path -from typing import List, Optional, Tuple, Union +from typing import List, Literal, Optional, Tuple, Union import cv2 import numpy as np @@ -176,6 +176,35 @@ def transform_to_pose(T: np.ndarray) -> np.ndarray: return np.concatenate([pos, quat_wxyz]) +ROT_RELABEL_POST_FIX = np.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]]) # "rot_right" +ROT_RELABEL_PRE_FIX_LEFT = np.array([[0, 0, 1], [-1, 0, 0], [0, -1, 0]]) # "rot_left" + + +def _unrelabeled_hand_frame(keypoints: np.ndarray) -> Optional[np.ndarray]: + """Hand frame from 21 keypoints, before the axis relabel; None if degenerate. + + Forward: wrist -> middle-finger base. Up: cross(thumb_dir, pinky_dir), which + already mirrors chirality between hands. Right: cross(forward, up), flipped + to make the frame right-handed. + """ + wrist = keypoints[0] + forward = keypoints[9] - wrist + if np.linalg.norm(forward) < 1e-6: + return None + forward = forward / np.linalg.norm(forward) + + up = np.cross(keypoints[5] - wrist, keypoints[17] - wrist) + if np.linalg.norm(up) < 1e-6: + return None + up = up / np.linalg.norm(up) + + right = np.cross(forward, up) + right = right / np.linalg.norm(right) + up = np.cross(right, forward) + right = right * -1 + return np.column_stack([forward, right, up]) + + def compute_hand_pose_xyzquat(keypoints: np.ndarray, hand_index: int) -> np.ndarray: """ Compute 7DOF pose from 21 hand keypoints in camera frame. @@ -195,8 +224,8 @@ def compute_hand_pose_xyzquat(keypoints: np.ndarray, hand_index: int) -> np.ndar # Identity quaternion in WXYZ for degenerate cases quat_wxyz_identity = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64) - rot_left = np.array([[0, 0, 1], [-1, 0, 0], [0, -1, 0]]) - rot_right = np.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]]) + rot_left = ROT_RELABEL_PRE_FIX_LEFT + rot_right = ROT_RELABEL_POST_FIX if np.allclose(keypoints, 0): fallback = np.concatenate([np.zeros(3), quat_wxyz_identity]) @@ -208,43 +237,23 @@ def compute_hand_pose_xyzquat(keypoints: np.ndarray, hand_index: int) -> np.ndar ) position = centroid wrist = keypoints[0] - middle_base = keypoints[9] # Base of middle finger - - # Forward axis: wrist → middle finger base - forward = middle_base - wrist - if np.linalg.norm(forward) < 1e-6: - logger.warning( - "Forward direction norm too small, returning position + identity quat" - ) - fallback = np.concatenate([position, quat_wxyz_identity]) - return fallback, fallback.copy() - forward = forward / np.linalg.norm(forward) - - # Up axis: cross product of thumb→wrist and pinky→wrist - thumb_dir = keypoints[5] - wrist - pinky_dir = keypoints[17] - wrist - up = np.cross(thumb_dir, pinky_dir) - if np.linalg.norm(up) < 1e-6: - logger.warning( - "Up direction norm too small, returning position + identity quat" - ) + rot_matrix = _unrelabeled_hand_frame(keypoints) + if rot_matrix is None: + logger.warning("Degenerate hand frame, returning position + identity quat") fallback = np.concatenate([position, quat_wxyz_identity]) return fallback, fallback.copy() - up = up / np.linalg.norm(up) - - # Right axis: cross(forward, up), then re-orthonormalize - right = np.cross(forward, up) - right = right / np.linalg.norm(right) - up = np.cross(right, forward) - right = right * -1 # Flip for right-handed frame - 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 deg) right-multiply on the left + # pose); see Human.get_transform_list(fix_left_wrist_convention=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]]) @@ -254,6 +263,62 @@ def compute_hand_pose_xyzquat(keypoints: np.ndarray, hand_index: int) -> np.ndar ) +def _median_geodesic_deg( + keypoints: np.ndarray, poses: np.ndarray, relabel: np.ndarray +) -> Optional[float]: + """Median angle between the stored rotations and the ones ``relabel`` + reproduces from the keypoints; None if no frame is usable.""" + errors = [] + for kp, pose in zip(keypoints, poses): + frame = _unrelabeled_hand_frame(kp) if not np.allclose(kp, 0) else None + if frame is None or np.allclose(pose[3:], 0): + continue + qw, qx, qy, qz = pose[3:7] + stored = Rotation.from_quat([qx, qy, qz, qw]).as_matrix() + delta = (frame @ relabel).T @ stored + errors.append(np.degrees(Rotation.from_matrix(delta).magnitude())) + return float(np.median(errors)) if errors else None + + +def left_wrist_convention( + group, frames: int = 25, tol_deg: float = 10.0 +) -> Literal["pre_fix", "post_fix", "unknown"]: + """Which hand-pose convention an episode zarr was written with. + + Both wrist rotations are recomputed from the episode's own stored keypoints + under each relabel and compared with the stored quaternions. The RIGHT hand + is the control: it used the same relabel before and after the fix, so if it + matches neither, the left-hand comparison is confounded and the verdict is + ``unknown`` -- which is what the handful of episodes written by a third, + earlier converter return. + """ + try: + left_kp = np.asarray(group["left.obs_keypoints"][:frames], dtype=np.float64) + left_pose = np.asarray(group["left.obs_wrist_pose"][:frames], dtype=np.float64) + right_kp = np.asarray(group["right.obs_keypoints"][:frames], dtype=np.float64) + right_pose = np.asarray( + group["right.obs_wrist_pose"][:frames], dtype=np.float64 + ) + except (KeyError, IndexError, TypeError): + return "unknown" + if min(len(left_kp), len(left_pose), len(right_kp), len(right_pose)) == 0: + return "unknown" + + left_kp = left_kp.reshape(len(left_kp), -1, 3) + right_kp = right_kp.reshape(len(right_kp), -1, 3) + control = _median_geodesic_deg(right_kp, right_pose, ROT_RELABEL_POST_FIX) + if control is None or control > tol_deg: + return "unknown" + + post_fix = _median_geodesic_deg(left_kp, left_pose, ROT_RELABEL_POST_FIX) + pre_fix = _median_geodesic_deg(left_kp, left_pose, ROT_RELABEL_PRE_FIX_LEFT) + if pre_fix is not None and pre_fix <= tol_deg: + return "pre_fix" + if post_fix is not None and post_fix <= tol_deg: + return "post_fix" + return "unknown" + + class MeckaExtractor: """ Extract observation features from Mecka RL2 episode data. diff --git a/tests/unit/test_pi05_norm_rot6d.py b/tests/unit/test_pi05_norm_rot6d.py index 58598728ca..89c40d9279 100644 --- a/tests/unit/test_pi05_norm_rot6d.py +++ b/tests/unit/test_pi05_norm_rot6d.py @@ -250,3 +250,32 @@ def test_human_6d_reverts_unpad_proprio(): _build_human_cartesian_revert_6d_wristframe_transform_list, ): assert any(isinstance(t, UnpadGripperZeros) for t in build()), build.__name__ + + +def test_rotate_local_frame_flips_left_wrist_convention(): + # Right-multiplying by Rz(180 deg) must flip the pose's own x/y axes, keep + # z (knuckle-forward) and the position, skip zero-quat padding rows, and + # handle both (7,) poses and (T, 7) chunks. + from scipy.spatial.transform import Rotation as R + + from egomimic.rldb.zarr.action_chunk_transforms import RotateLocalFrame + + rng = np.random.default_rng(4) + q = R.random(3, random_state=5) + chunk = np.zeros((4, 7)) + chunk[:3, :3] = rng.uniform(-1, 1, size=(3, 3)) + chunk[:3, 3:] = q.as_quat()[:, [3, 0, 1, 2]] # wxyz; row 3 stays zero-padded + + t = RotateLocalFrame(keys=["k"]) + out = t.transform({"k": chunk.copy()})["k"] + + np.testing.assert_allclose(out[:, :3], chunk[:, :3]) # positions unchanged + np.testing.assert_allclose(out[3], np.zeros(7)) # padding untouched + R_old = q.as_matrix() + R_new = R.from_quat(out[:3, [4, 5, 6, 3]]).as_matrix() + np.testing.assert_allclose(R_new[:, :, 0], -R_old[:, :, 0], atol=1e-12) # x flip + np.testing.assert_allclose(R_new[:, :, 1], -R_old[:, :, 1], atol=1e-12) # y flip + np.testing.assert_allclose(R_new[:, :, 2], R_old[:, :, 2], atol=1e-12) # z kept + + single = t.transform({"k": chunk[0].copy()})["k"] + np.testing.assert_allclose(single, out[0], atol=1e-12) diff --git a/tests/unit/test_wrist6d_roundtrip.py b/tests/unit/test_wrist6d_roundtrip.py new file mode 100644 index 0000000000..ff98b4857b --- /dev/null +++ b/tests/unit/test_wrist6d_roundtrip.py @@ -0,0 +1,518 @@ +"""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 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, n=T): + """Smooth random walk of xyz+quat(wxyz) poses starting AT ``start``.""" + out = np.zeros((n, 7)) + p = start[:3].copy() + r = R.from_quat(start[[4, 5, 6, 3]]) + for t in range(n): + if t > 0: + p = p + rng.normal(0, 0.01, 3) + r = R.from_rotvec(rng.normal(0, 0.05, 3)) * r + q = r.as_quat() + out[t] = np.concatenate([p, q[[3, 0, 1, 2]]]) + return out + + +def _T(p7): + M = np.eye(4) + M[:3, :3] = R.from_quat(p7[[4, 5, 6, 3]]).as_matrix() + M[:3, 3] = p7[:3] + return M + + +def _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 +-pi 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 helpers 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_left_wrist_convention=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 deg) 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 -> 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() + + +# -------------------------------------------------------------- eva path +def test_eva_wristframe_6d_pipeline_round_trips_to_camframe(): + from egomimic.rldb.embodiment.eva import ( + Eva, + _build_eva_cartesian_revert_6d_wristframe_transform_list, + ) + + extrinsics = Eva.EXTRINSICS + 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)), + } + ) + fwd = Eva.get_transform_list("cartesian_wristframe_6d") + 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 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: a plain ypr + wrist-frame 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) + + +# --------------------------------------------------------------------------- +# The convention checker: does a zarr's stored wrist rotation match what its +# own keypoints produce under the pre-fix or the post-fix relabel? +# --------------------------------------------------------------------------- + + +def _synthetic_hand(rng, n_frames=5): + """(n, 21, 3) keypoints whose wrist frame is well-conditioned.""" + kp = rng.normal(scale=0.02, size=(n_frames, 21, 3)) + kp[:, 0] = [0.0, 0.0, 0.0] # wrist + kp[:, 9] += [0.10, 0.0, 0.0] # middle-finger base: forward axis + kp[:, 5] += [0.03, 0.05, 0.0] # thumb + kp[:, 17] += [0.03, -0.05, 0.0] # pinky + return kp + + +def _poses_from(kp, relabel): + from egomimic.scripts.mecka_process.mecka_to_zarr import _unrelabeled_hand_frame + + poses = np.zeros((len(kp), 7)) + for i, frame in enumerate(kp): + q = R.from_matrix(_unrelabeled_hand_frame(frame) @ relabel).as_quat() + poses[i, :3] = frame[0] + poses[i, 3:] = [q[3], q[0], q[1], q[2]] # stored WXYZ + return poses + + +def _episode(tmp_path, left_relabel, right_relabel, name="ep"): + import zarr + + from egomimic.scripts.mecka_process.mecka_to_zarr import ( + ROT_RELABEL_POST_FIX, + ROT_RELABEL_PRE_FIX_LEFT, + ) + + relabels = {"pre_fix": ROT_RELABEL_PRE_FIX_LEFT, "post_fix": ROT_RELABEL_POST_FIX} + rng = np.random.default_rng(0) + g = zarr.open_group(str(tmp_path / f"{name}.zarr"), mode="w") + for side, which in (("left", left_relabel), ("right", right_relabel)): + kp = _synthetic_hand(rng) + relabel = relabels[which] if isinstance(which, str) else which + g[f"{side}.obs_keypoints"] = kp.reshape(len(kp), -1) + g[f"{side}.obs_wrist_pose"] = _poses_from(kp, relabel) + return g + + +@pytest.mark.parametrize("written", ["pre_fix", "post_fix"]) +def test_left_wrist_convention_reads_the_data(tmp_path, written): + from egomimic.scripts.mecka_process.mecka_to_zarr import left_wrist_convention + + g = _episode(tmp_path, left_relabel=written, right_relabel="post_fix") + assert left_wrist_convention(g) == written + + +def test_left_wrist_convention_is_unknown_when_the_control_fails(tmp_path): + from egomimic.scripts.mecka_process.mecka_to_zarr import left_wrist_convention + + # A third converter: the right hand matches neither relabel, so the + # left-hand comparison is confounded (3 fold_clothes + 2 pick_place + # episodes in the production set look like this). + third = np.array([[0, 0, 1], [0, -1, 0], [1, 0, 0]]) + g = _episode(tmp_path, left_relabel=third, right_relabel=third) + assert left_wrist_convention(g) == "unknown" + + +def test_left_wrist_convention_is_unknown_without_keypoints(tmp_path): + import zarr + + from egomimic.scripts.mecka_process.mecka_to_zarr import left_wrist_convention + + g = zarr.open_group(str(tmp_path / "bare.zarr"), mode="w") + assert left_wrist_convention(g) == "unknown" + + +class _FakeDataset: + def __init__(self, path): + self.episode_path = path + + +def _resolver(tmp_path, fix_enabled): + from egomimic.rldb.embodiment.human import Human + from egomimic.rldb.zarr.zarr_dataset_multi import S3EpisodeResolver + + return S3EpisodeResolver( + tmp_path, + key_map=Human.get_keymap(keymap_mode="cartesian"), + transform_list=Human.get_transform_list( + mode="cartesian_6d", stride=1, fix_left_wrist_convention=fix_enabled + ), + ) + + +def test_resolver_rejects_a_double_flip(tmp_path): + """Correcting data the converter already corrected is a fresh error.""" + _episode(tmp_path, left_relabel="post_fix", right_relabel="post_fix") + datasets = {"ep": _FakeDataset(tmp_path / "ep.zarr")} + with pytest.raises(ValueError, match="double-flip"): + _resolver(tmp_path, True)._check_left_wrist_convention(datasets) + + +def test_resolver_only_warns_on_uncorrected_data(tmp_path, caplog): + """Pre-fix data with the flag off is what every run did before this + commit; it is wrong, but killing the generic human configs at startup is + not this correction's business.""" + import logging + + _episode(tmp_path, left_relabel="pre_fix", right_relabel="post_fix") + datasets = {"ep": _FakeDataset(tmp_path / "ep.zarr")} + with caplog.at_level(logging.WARNING): + _resolver(tmp_path, False)._check_left_wrist_convention(datasets) + assert any("uncorrected frame" in r.message for r in caplog.records) + + +@pytest.mark.parametrize( + "written,fix_enabled", [("pre_fix", True), ("post_fix", False)] +) +def test_resolver_accepts_a_flag_the_data_agrees_with(tmp_path, written, fix_enabled): + _episode(tmp_path, left_relabel=written, right_relabel="post_fix") + datasets = {"ep": _FakeDataset(tmp_path / "ep.zarr")} + _resolver(tmp_path, fix_enabled)._check_left_wrist_convention(datasets) + + +def test_resolver_only_warns_on_an_unrecognised_convention(tmp_path, caplog): + import logging + + third = np.array([[0, 0, 1], [0, -1, 0], [1, 0, 0]]) + _episode(tmp_path, left_relabel=third, right_relabel=third) + datasets = {"ep": _FakeDataset(tmp_path / "ep.zarr")} + with caplog.at_level(logging.WARNING): + _resolver(tmp_path, True)._check_left_wrist_convention(datasets) + assert any("wrist-convention" in r.message for r in caplog.records)