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

2026-08-24 17:50:23 +08:00
parent fbd170c0be
commit 398a50b0b3
2 changed files with 300 additions and 1 deletions
@@ -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,
+    )