From 08586107dd034e36e3bc37426637ceea56851a8c Mon Sep 17 00:00:00 2001 From: YikaiFu-cart Date: Mon, 24 Aug 2026 17:43:47 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=E9=80=86=E8=BF=90?= =?UTF-8?q?=E5=8A=A8=E5=AD=A6=E8=BD=A8=E8=BF=B9=E9=87=8D=E9=87=87=E6=A0=B7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- xr_rm_teleop/test/ik_method_comparison.py | 90 +++++++++++++++++++ .../test/test_ik_method_comparison.py | 48 ++++++++++ 2 files changed, 138 insertions(+) create mode 100644 xr_rm_teleop/test/ik_method_comparison.py create mode 100644 xr_rm_teleop/test/test_ik_method_comparison.py diff --git a/xr_rm_teleop/test/ik_method_comparison.py b/xr_rm_teleop/test/ik_method_comparison.py new file mode 100644 index 0000000..bb39b10 --- /dev/null +++ b/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(), + ) diff --git a/xr_rm_teleop/test/test_ik_method_comparison.py b/xr_rm_teleop/test/test_ik_method_comparison.py new file mode 100644 index 0000000..9fdc563 --- /dev/null +++ b/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)