feat: 读取番茄采摘目标轨迹

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(),
+    )

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)