diff --git a/xr_rm_teleop/test/ik_method_comparison.py b/xr_rm_teleop/test/ik_method_comparison.py index bb39b10..97a4bf3 100644 --- a/xr_rm_teleop/test/ik_method_comparison.py +++ b/xr_rm_teleop/test/ik_method_comparison.py @@ -3,6 +3,7 @@ from __future__ import annotations from dataclasses import dataclass from pathlib import Path +import h5py import numpy as np @@ -88,3 +89,67 @@ def resample_trajectory( target_poses=target_poses, initial_joints=np.asarray(trajectory.initial_joints, dtype=float).copy(), ) + + +def _longest_true_run(mask: np.ndarray) -> slice: + values = np.asarray(mask, dtype=bool) + best_start = best_stop = start = 0 + for index, enabled in enumerate(np.r_[values, False]): + if enabled: + continue + if index - start > best_stop - best_start: + best_start, best_stop = start, index + start = index + 1 + if best_stop - best_start < 2: + raise ValueError("episode has no valid teleoperation run") + return slice(best_start, best_stop) + + +def load_episode(path: Path) -> EpisodeTrajectory: + source = Path(path).expanduser().resolve() + if not source.is_file() or source.suffix.lower() not in (".h5", ".hdf5"): + raise FileNotFoundError(f"episode not found: {source}") + with h5py.File(source, "r") as handle: + if str(handle.attrs.get("arm", "")) != "right_rm75": + raise ValueError("episode arm must be right_rm75") + if str(handle.attrs.get("pose_order", "")) != "x,y,z,qx,qy,qz,qw": + raise ValueError("episode pose_order is unsupported") + required = ( + "debug/tcp/final_target_pose", + "debug/timestamps/control_monotonic_ns", + "debug/control/teleop_active", + "debug/control/action_valid", + "debug/control/command_sent", + "observations/qpos", + ) + missing = [name for name in required if name not in handle] + if missing: + raise ValueError(f"episode datasets missing: {missing}") + poses = np.asarray(handle[required[0]], dtype=float) + timestamps = np.asarray(handle[required[1]], dtype=np.int64) + mask = np.logical_and.reduce( + [ + np.asarray(handle[required[2]], dtype=bool), + np.asarray(handle[required[3]], dtype=bool), + np.asarray(handle[required[4]], dtype=bool), + ] + ) + qpos = np.asarray(handle[required[5]], dtype=float) + if poses.shape != (timestamps.size, 7) or qpos.shape != (timestamps.size, 8): + raise ValueError("episode arrays have inconsistent shapes") + if not np.isfinite(poses).all() or not np.isfinite(qpos).all(): + raise ValueError("episode contains NaN/Inf") + selected = _longest_true_run(mask) + selected_times = timestamps[selected] + if np.any(np.diff(selected_times) <= 0): + raise ValueError("episode timestamps must be strictly increasing") + selected_poses = poses[selected].copy() + selected_poses[:, 3:] = np.asarray( + [_normalized_quaternion(value) for value in selected_poses[:, 3:]] + ) + return EpisodeTrajectory( + source_path=source, + times_s=(selected_times - selected_times[0]) * 1e-9, + target_poses=selected_poses, + initial_joints=qpos[selected.start, :7].copy(), + ) diff --git a/xr_rm_teleop/test/test_ik_method_comparison.py b/xr_rm_teleop/test/test_ik_method_comparison.py index 9fdc563..3073e08 100644 --- a/xr_rm_teleop/test/test_ik_method_comparison.py +++ b/xr_rm_teleop/test/test_ik_method_comparison.py @@ -4,6 +4,7 @@ import math import sys from pathlib import Path +import h5py import numpy as np import pytest @@ -46,3 +47,58 @@ def test_resample_trajectory_keeps_endpoints_and_uses_requested_rate() -> None: assert actual.target_poses[0] == pytest.approx(trajectory.target_poses[0]) assert actual.target_poses[-1] == pytest.approx(trajectory.target_poses[-1]) assert actual.target_poses[:, 0] == pytest.approx(actual.times_s) + + +def _write_episode(path: Path) -> None: + poses = np.asarray( + [ + [0.1, -0.2, 0.3, 0.0, 0.0, 0.0, 1.0], + [0.2, -0.2, 0.3, 0.0, 0.0, 0.0, 1.0], + [0.3, -0.2, 0.3, 0.0, 0.0, 0.0, 1.0], + [0.4, -0.2, 0.3, 0.0, 0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + with h5py.File(path, "w") as handle: + handle.attrs["arm"] = "right_rm75" + handle.attrs["pose_order"] = "x,y,z,qx,qy,qz,qw" + handle.create_dataset("debug/tcp/final_target_pose", data=poses) + handle.create_dataset( + "debug/timestamps/control_monotonic_ns", + data=np.asarray([0, 33_000_000, 66_000_000, 99_000_000]), + ) + handle.create_dataset( + "debug/control/teleop_active", data=[0, 1, 1, 0] + ) + handle.create_dataset( + "debug/control/action_valid", data=[1, 1, 1, 1] + ) + handle.create_dataset( + "debug/control/command_sent", data=[0, 1, 1, 0] + ) + qpos = np.zeros((4, 8), dtype=np.float32) + qpos[1, :7] = np.arange(7) * 0.1 + handle.create_dataset("observations/qpos", data=qpos) + + +def test_load_episode_uses_longest_valid_run_and_first_valid_qpos( + tmp_path: Path, +) -> None: + path = tmp_path / "episode.hdf5" + _write_episode(path) + + actual = comparison.load_episode(path) + + assert actual.times_s == pytest.approx([0.0, 0.033]) + assert actual.target_poses[:, 0] == pytest.approx([0.2, 0.3]) + assert actual.initial_joints == pytest.approx(np.arange(7) * 0.1) + + +def test_load_episode_rejects_wrong_arm(tmp_path: Path) -> None: + path = tmp_path / "episode.hdf5" + _write_episode(path) + with h5py.File(path, "r+") as handle: + handle.attrs.modify("arm", "left_rm75") + + with pytest.raises(ValueError, match="right_rm75"): + comparison.load_episode(path)