feat: 添加三种逆运动学离线复放

xr_rm_teleop/test/ik_method_comparison.py

@@ -1,13 +1,26 @@
 from __future__ import annotations
 
 import math
+import time
 from dataclasses import dataclass
 from pathlib import Path
 
 import h5py
 import numpy as np
 
-from xr_rm_teleop.single_arm_velocity_teleop import SingleArmVelocityTeleop
+from xr_rm_teleop.placo_ik_solver import (
+    QP_MAX_ITERATIONS,
+    QP_ORIENTATION_TOLERANCE_RAD,
+    QP_POSITION_TOLERANCE_M,
+    PlacoIkSolver,
+    _validated_transform,
+)
+from xr_rm_teleop.single_arm_velocity_teleop import (
+    SingleArmVelocityTeleop,
+    _matrix_to_quaternion,
+    _quaternion_to_matrix,
+    _so3_log,
+)
 
 
 @dataclass(frozen=True)
@@ -28,6 +41,22 @@ class MethodSummary:
     max_joint_speed_deg_s: float
 
 
+@dataclass(frozen=True)
+class ReplayResult:
+    method: str
+    times_s: np.ndarray
+    target_poses: np.ndarray
+    actual_poses: np.ndarray
+    joints: np.ndarray
+    velocities: np.ndarray
+    position_errors_m: np.ndarray
+    orientation_errors_rad: np.ndarray
+    joint_margins: np.ndarray
+    solve_durations_ms: np.ndarray
+    success: np.ndarray
+    command_limited: 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():
@@ -244,3 +273,208 @@ def limit_joint_command(
         velocity_array,
         not np.allclose(target_array, target, atol=1e-12, rtol=0.0),
     )
+
+
+def _pose_to_transform(pose: np.ndarray) -> np.ndarray:
+    values = np.asarray(pose, dtype=float)
+    transform = np.eye(4)
+    transform[:3, 3] = values[:3]
+    transform[:3, :3] = _quaternion_to_matrix(tuple(values[3:]))
+    return transform
+
+
+def _transform_to_pose(transform: np.ndarray) -> np.ndarray:
+    quaternion = _matrix_to_quaternion(transform[:3, :3])
+    return np.asarray([*transform[:3, 3], *quaternion], dtype=float)
+
+
+class DifferentialIkSolver:
+    def __init__(
+        self,
+        urdf_path: Path,
+        dt: float,
+        method: str,
+        damping: float = 0.0,
+    ) -> None:
+        if method not in ("pinv", "dls"):
+            raise ValueError("method must be pinv or dls")
+        if method == "dls" and damping <= 0.0:
+            raise ValueError("DLS damping must be positive")
+        self._kinematics = PlacoIkSolver(str(urdf_path), dt, "right")
+        self._dt = dt
+        self._method = method
+        self._damping = float(damping)
+        self._actual_joints: np.ndarray | None = None
+
+    @property
+    def joint_limits(self) -> np.ndarray:
+        return self._kinematics._joint_limits.copy()
+
+    def update_joint_state(self, joints: list[float]) -> np.ndarray:
+        self._actual_joints = np.asarray(joints, dtype=float).copy()
+        return self._kinematics.update_joint_state(joints)
+
+    def _set_internal_joints(self, joints: np.ndarray) -> None:
+        robot = self._kinematics._robot
+        robot.state.q[self._kinematics._q_offsets] = joints
+        robot.update_kinematics()
+
+    def _errors(self, target: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        robot = self._kinematics._robot
+        world_base = robot.get_T_world_frame(self._kinematics._base_frame)
+        world_tcp = robot.get_T_world_frame(self._kinematics._tcp_frame)
+        world_target = world_base @ target
+        position = world_target[:3, 3] - world_tcp[:3, 3]
+        orientation = _so3_log(
+            world_target[:3, :3] @ world_tcp[:3, :3].T
+        )
+        return position, orientation
+
+    def solve(self, target: np.ndarray) -> list[float]:
+        if self._actual_joints is None:
+            raise RuntimeError("joint state must be initialized before IK solve")
+        target = _validated_transform(target)
+        actual = self._actual_joints.copy()
+        result = actual.copy()
+        try:
+            for _ in range(QP_MAX_ITERATIONS):
+                self._set_internal_joints(result)
+                position, orientation = self._errors(target)
+                if (
+                    np.linalg.norm(position) <= QP_POSITION_TOLERANCE_M
+                    and np.linalg.norm(orientation)
+                    <= QP_ORIENTATION_TOLERANCE_RAD
+                ):
+                    return result.tolist()
+                jacobian = self._kinematics._active_tcp_jacobian()
+                desired_twist = np.r_[position, orientation] / self._dt
+                if self._method == "pinv":
+                    joint_velocity = np.linalg.pinv(jacobian) @ desired_twist
+                else:
+                    system = (
+                        jacobian @ jacobian.T
+                        + self._damping**2 * np.eye(6)
+                    )
+                    joint_velocity = jacobian.T @ np.linalg.solve(
+                        system, desired_twist
+                    )
+                candidate = result + joint_velocity * self._dt
+                self._kinematics._validate_result(candidate, result)
+                result = candidate
+            raise RuntimeError(
+                f"{self._method} did not converge after "
+                f"{QP_MAX_ITERATIONS} iterations"
+            )
+        except Exception:
+            self._set_internal_joints(actual)
+            raise
+
+
+class QpSolverAdapter:
+    def __init__(self, solver: PlacoIkSolver) -> None:
+        self._solver = solver
+
+    @property
+    def joint_limits(self) -> np.ndarray:
+        return self._solver._joint_limits.copy()
+
+    def update_joint_state(self, joints: list[float]) -> np.ndarray:
+        return self._solver.update_joint_state(joints)
+
+    def solve(self, target: np.ndarray) -> list[float]:
+        return self._solver.solve(target)
+
+
+def make_qp_solver(urdf_path: Path, dt: float) -> QpSolverAdapter:
+    return QpSolverAdapter(
+        PlacoIkSolver(
+            str(urdf_path),
+            dt,
+            "right",
+            j3_reference_deg=-89.57,
+            j3_weight=1e-4,
+            j4_min_deg=10.0,
+            j4_warn_deg=25.0,
+            j4_weight=1e-4,
+            manipulability_sigma_stop=0.01,
+            manipulability_sigma_warn=0.04,
+            manipulability_weight=1e-4,
+        )
+    )
+
+
+def run_replay(
+    method: str,
+    solver,
+    trajectory: EpisodeTrajectory,
+    *,
+    max_speed: float,
+    max_acceleration: float,
+    measure_time: bool = True,
+) -> ReplayResult:
+    count = trajectory.times_s.size
+    dt = float(np.median(np.diff(trajectory.times_s)))
+    joints = np.empty((count, 7))
+    velocities = np.zeros((count, 7))
+    actual_poses = np.empty((count, 7))
+    position_errors = np.empty(count)
+    orientation_errors = np.empty(count)
+    margins = np.empty(count)
+    durations = np.zeros(count)
+    success = np.zeros(count, dtype=bool)
+    command_limited = np.zeros(count, dtype=bool)
+    current = trajectory.initial_joints.copy()
+    previous_velocity = np.zeros(7)
+    lower, upper = solver.joint_limits.T
+
+    for index, pose in enumerate(trajectory.target_poses):
+        solver.update_joint_state(current.tolist())
+        started = time.perf_counter_ns()
+        try:
+            candidate = np.asarray(
+                solver.solve(_pose_to_transform(pose)), dtype=float
+            )
+            success[index] = True
+        except Exception:
+            candidate = current.copy()
+            previous_velocity = np.zeros(7)
+        durations[index] = (
+            (time.perf_counter_ns() - started) * 1e-6 if measure_time else 0.0
+        )
+        if success[index]:
+            current, previous_velocity, command_limited[index] = (
+                limit_joint_command(
+                    target=candidate,
+                    previous_target=current,
+                    previous_velocity=previous_velocity,
+                    max_speed=max_speed,
+                    max_acceleration=max_acceleration,
+                    dt=dt,
+                )
+            )
+        actual_transform = solver.update_joint_state(current.tolist())
+        actual_pose = _transform_to_pose(actual_transform)
+        joints[index] = current
+        velocities[index] = previous_velocity
+        actual_poses[index] = actual_pose
+        position_errors[index] = np.linalg.norm(pose[:3] - actual_pose[:3])
+        orientation_errors[index] = orientation_error_rad(
+            _quaternion_to_matrix(tuple(actual_pose[3:])),
+            _quaternion_to_matrix(tuple(pose[3:])),
+        )
+        margins[index] = normalized_joint_margin(current, lower, upper)
+
+    return ReplayResult(
+        method=method,
+        times_s=trajectory.times_s.copy(),
+        target_poses=trajectory.target_poses.copy(),
+        actual_poses=actual_poses,
+        joints=joints,
+        velocities=velocities,
+        position_errors_m=position_errors,
+        orientation_errors_rad=orientation_errors,
+        joint_margins=margins,
+        solve_durations_ms=durations,
+        success=success,
+        command_limited=command_limited,
+    )

xr_rm_teleop/test/test_ik_method_comparison.py

@@ -140,3 +140,68 @@ def test_limit_joint_command_reuses_production_limiter() -> None:
     assert target == pytest.approx([0.1] * 7)
     assert velocity == pytest.approx([1.0] * 7)
     assert limited
+
+
+class _FakeSolver:
+    def __init__(self, fail: bool) -> None:
+        self.fail = fail
+        self.joint_limits = np.asarray([[-2.0, 2.0]] * 7)
+
+    def update_joint_state(self, joints: list[float]) -> np.ndarray:
+        pose = np.eye(4)
+        pose[0, 3] = joints[0]
+        return pose
+
+    def solve(self, target: np.ndarray) -> list[float]:
+        if self.fail:
+            raise RuntimeError("not converged")
+        return [float(target[0, 3])] + [0.0] * 6
+
+
+def test_replay_holds_previous_state_on_solver_failure() -> None:
+    trajectory = comparison.EpisodeTrajectory(
+        source_path=Path("episode.hdf5"),
+        times_s=np.asarray([0.0, 0.1]),
+        target_poses=np.asarray(
+            [
+                [0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+                [0.2, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+            ]
+        ),
+        initial_joints=np.zeros(7),
+    )
+
+    result = comparison.run_replay(
+        "fake",
+        _FakeSolver(fail=True),
+        trajectory,
+        max_speed=1.0,
+        max_acceleration=10.0,
+        measure_time=False,
+    )
+
+    assert result.joints == pytest.approx(np.zeros((2, 7)))
+    assert not result.success.any()
+    assert result.velocities == pytest.approx(np.zeros((2, 7)))
+
+
+def test_real_urdf_solvers_return_finite_safe_outputs() -> None:
+    pytest.importorskip("placo")
+    urdf = TEST_DIR.parent / "models" / "dual_rm75" / "Dual_arm.urdf"
+    joints = np.radians(
+        [-86.10, 22.80, -89.57, 93.98, -91.82, -87.32, -89.35]
+    )
+    solvers = [
+        comparison.DifferentialIkSolver(urdf, 1.0 / 90.0, "pinv"),
+        comparison.DifferentialIkSolver(urdf, 1.0 / 90.0, "dls", 0.03),
+        comparison.make_qp_solver(urdf, 1.0 / 90.0),
+    ]
+    for solver in solvers:
+        target = solver.update_joint_state(joints.tolist())
+        target = target.copy()
+        target[0, 3] += 0.003
+        result = np.asarray(solver.solve(target), dtype=float)
+        assert result.shape == (7,)
+        assert np.isfinite(result).all()
+        assert np.all(result >= solver.joint_limits[:, 0] - 1e-9)
+        assert np.all(result <= solver.joint_limits[:, 1] + 1e-9)