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

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)