feat: 添加逆运动学轨迹重采样

xr_rm_teleop/test/ik_method_comparison.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class EpisodeTrajectory:
+    source_path: Path
+    times_s: np.ndarray
+    target_poses: np.ndarray
+    initial_joints: np.ndarray
+
+
+def _normalized_quaternion(values: np.ndarray) -> np.ndarray:
+    quaternion = np.asarray(values, dtype=float)
+    if quaternion.shape != (4,) or not np.isfinite(quaternion).all():
+        raise ValueError("quaternion must contain 4 finite values")
+    norm = float(np.linalg.norm(quaternion))
+    if norm <= 1e-12:
+        raise ValueError("quaternion norm must be positive")
+    return quaternion / norm
+
+
+def _slerp_quaternion(
+    start: np.ndarray,
+    end: np.ndarray,
+    fraction: float,
+) -> np.ndarray:
+    first = _normalized_quaternion(start)
+    second = _normalized_quaternion(end)
+    dot = float(np.dot(first, second))
+    if dot < 0.0:
+        second = -second
+        dot = -dot
+    dot = float(np.clip(dot, -1.0, 1.0))
+    if dot > 1.0 - 1e-8:
+        return _normalized_quaternion(
+            first + float(fraction) * (second - first)
+        )
+    angle = float(np.arccos(dot))
+    sine = float(np.sin(angle))
+    return _normalized_quaternion(
+        np.sin((1.0 - fraction) * angle) / sine * first
+        + np.sin(fraction * angle) / sine * second
+    )
+
+
+def resample_trajectory(
+    trajectory: EpisodeTrajectory,
+    sample_rate_hz: float,
+) -> EpisodeTrajectory:
+    if not np.isfinite(sample_rate_hz) or sample_rate_hz <= 0.0:
+        raise ValueError("sample_rate_hz must be finite and positive")
+    source_times = np.asarray(trajectory.times_s, dtype=float)
+    poses = np.asarray(trajectory.target_poses, dtype=float)
+    if source_times.ndim != 1 or poses.shape != (source_times.size, 7):
+        raise ValueError("trajectory must contain N timestamps and N x 7 poses")
+    if source_times.size < 2 or np.any(np.diff(source_times) <= 0.0):
+        raise ValueError("trajectory timestamps must be strictly increasing")
+
+    duration = float(source_times[-1] - source_times[0])
+    count = int(round(duration * sample_rate_hz)) + 1
+    target_times = np.linspace(source_times[0], source_times[-1], count)
+    target_poses = np.empty((count, 7), dtype=float)
+    for axis in range(3):
+        target_poses[:, axis] = np.interp(
+            target_times,
+            source_times,
+            poses[:, axis],
+        )
+    for index, timestamp in enumerate(target_times):
+        right = int(np.searchsorted(source_times, timestamp, side="right"))
+        right = min(max(right, 1), source_times.size - 1)
+        left = right - 1
+        interval = source_times[right] - source_times[left]
+        fraction = float((timestamp - source_times[left]) / interval)
+        target_poses[index, 3:] = _slerp_quaternion(
+            poses[left, 3:], poses[right, 3:], fraction
+        )
+    target_poses[0] = poses[0]
+    target_poses[-1] = poses[-1]
+    return EpisodeTrajectory(
+        source_path=trajectory.source_path,
+        times_s=target_times - target_times[0],
+        target_poses=target_poses,
+        initial_joints=np.asarray(trajectory.initial_joints, dtype=float).copy(),
+    )

xr_rm_teleop/test/test_ik_method_comparison.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import math
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+TEST_DIR = Path(__file__).resolve().parent
+if str(TEST_DIR) not in sys.path:
+    sys.path.insert(0, str(TEST_DIR))
+
+import ik_method_comparison as comparison
+
+
+def test_slerp_uses_shortest_arc_and_returns_unit_quaternion() -> None:
+    start = np.asarray([0.0, 0.0, 0.0, 1.0])
+    end = -np.asarray([0.0, 0.0, math.sin(0.1), math.cos(0.1)])
+
+    actual = comparison._slerp_quaternion(start, end, 0.5)
+
+    assert np.linalg.norm(actual) == pytest.approx(1.0)
+    assert actual == pytest.approx(
+        [0.0, 0.0, math.sin(0.05), math.cos(0.05)]
+    )
+
+
+def test_resample_trajectory_keeps_endpoints_and_uses_requested_rate() -> None:
+    trajectory = comparison.EpisodeTrajectory(
+        source_path=Path("episode.hdf5"),
+        times_s=np.asarray([0.0, 0.5, 1.0]),
+        target_poses=np.asarray(
+            [
+                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+                [0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+                [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+            ]
+        ),
+        initial_joints=np.zeros(7),
+    )
+
+    actual = comparison.resample_trajectory(trajectory, 4.0)
+
+    assert actual.times_s == pytest.approx([0.0, 0.25, 0.5, 0.75, 1.0])
+    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)