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