from __future__ import annotations import math import time from dataclasses import dataclass, replace from pathlib import Path import h5py import numpy as np import yaml 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) class EpisodeTrajectory: source_path: Path times_s: np.ndarray target_poses: np.ndarray initial_joints: np.ndarray @dataclass(frozen=True) class MethodSummary: method: str damping: float | None success_rate: float position_rmse_m: float orientation_rmse_rad: float max_joint_speed_deg_s: float min_joint_margin: float = 0.0 mean_solve_ms: float = 0.0 max_solve_ms: float = 0.0 failure_count: int = 0 longest_failure_streak: int = 0 command_limited_count: int = 0 @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(): 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(), ) 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(), ) def load_right_config(path: Path) -> dict[str, float]: with Path(path).open("r", encoding="utf-8") as stream: document = yaml.safe_load(stream) parameters = document["single_arm_velocity_teleop"]["ros__parameters"] names = ( "qp_j3_reference_deg", "qp_j3_weight", "qp_j4_min_deg", "qp_j4_warn_deg", "qp_j4_weight", "qp_manipulability_sigma_stop", "qp_manipulability_sigma_warn", "qp_manipulability_weight", "joint_max_speed", "joint_max_acc", ) return {name: float(parameters[name]) for name in names} def _rotation_z(angle: float) -> np.ndarray: cosine, sine = math.cos(angle), math.sin(angle) return np.asarray( [ [cosine, -sine, 0.0], [sine, cosine, 0.0], [0.0, 0.0, 1.0], ] ) def orientation_error_rad(actual: np.ndarray, target: np.ndarray) -> float: delta = np.asarray(target) @ np.asarray(actual).T cosine = float(np.clip((np.trace(delta) - 1.0) * 0.5, -1.0, 1.0)) return float(math.acos(cosine)) def normalized_joint_margin( joints: np.ndarray, lower: np.ndarray, upper: np.ndarray, ) -> float: values = np.asarray(joints, dtype=float) lower_values = np.asarray(lower, dtype=float) upper_values = np.asarray(upper, dtype=float) span = upper_values - lower_values if np.any(span <= 0.0): raise ValueError("joint limits must have positive spans") margins = np.minimum( values - lower_values, upper_values - values, ) / span return float(np.min(margins)) def choose_dls_damping(candidates: list[MethodSummary]) -> float: if not candidates or any(value.damping is None for value in candidates): raise ValueError("DLS candidates must contain damping values") selected = min( candidates, key=lambda value: ( -value.success_rate, value.position_rmse_m / 0.002 + value.orientation_rmse_rad / 0.005, value.max_joint_speed_deg_s, ), ) return float(selected.damping) def _longest_failure_streak(success: np.ndarray) -> int: longest = current = 0 for value in np.asarray(success, dtype=bool): current = 0 if value else current + 1 longest = max(longest, current) return longest def summarize_result( result: ReplayResult, damping: float | None = None, ) -> MethodSummary: return MethodSummary( method=result.method, damping=damping, success_rate=float(np.mean(result.success)), position_rmse_m=float( np.sqrt(np.mean(result.position_errors_m**2)) ), orientation_rmse_rad=float( np.sqrt(np.mean(result.orientation_errors_rad**2)) ), max_joint_speed_deg_s=float( np.max(np.abs(np.degrees(result.velocities))) ), min_joint_margin=float(np.min(result.joint_margins)), mean_solve_ms=float(np.mean(result.solve_durations_ms)), max_solve_ms=float(np.max(result.solve_durations_ms)), failure_count=int(np.count_nonzero(~result.success)), longest_failure_streak=_longest_failure_streak(result.success), command_limited_count=int(np.count_nonzero(result.command_limited)), ) def limit_joint_command( *, target: np.ndarray, previous_target: np.ndarray, previous_velocity: np.ndarray, max_speed: float, max_acceleration: float, dt: float, ) -> tuple[np.ndarray, np.ndarray, bool]: limited_target, limited_velocity = ( SingleArmVelocityTeleop._limit_joint_command_step( target=np.asarray(target, dtype=float).tolist(), previous_target=np.asarray(previous_target, dtype=float).tolist(), previous_velocity=np.asarray(previous_velocity, dtype=float).tolist(), max_speed=max_speed, max_acceleration=max_acceleration, dt=dt, ) ) target_array = np.asarray(limited_target, dtype=float) velocity_array = np.asarray(limited_velocity, dtype=float) return ( target_array, 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, config: dict[str, float], ) -> QpSolverAdapter: return QpSolverAdapter( PlacoIkSolver( str(urdf_path), dt, "right", j3_reference_deg=config["qp_j3_reference_deg"], j3_weight=config["qp_j3_weight"], j4_min_deg=config["qp_j4_min_deg"], j4_warn_deg=config["qp_j4_warn_deg"], j4_weight=config["qp_j4_weight"], manipulability_sigma_stop=config[ "qp_manipulability_sigma_stop" ], manipulability_sigma_warn=config[ "qp_manipulability_sigma_warn" ], manipulability_weight=config["qp_manipulability_weight"], ) ) 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, ) DLS_DAMPING_CANDIDATES = (0.001, 0.003, 0.01, 0.03, 0.1, 0.3) def evaluate_methods( trajectory: EpisodeTrajectory, urdf_path: Path, config: dict[str, float], timing_repeats: int, ) -> tuple[dict[str, ReplayResult], dict[str, MethodSummary], float]: if timing_repeats < 1: raise ValueError("timing_repeats must be positive") dt = float(np.median(np.diff(trajectory.times_s))) speed = math.radians(config["joint_max_speed"]) acceleration = math.radians(config["joint_max_acc"]) damping_results = [] for damping in DLS_DAMPING_CANDIDATES: replay = run_replay( "dls", DifferentialIkSolver(urdf_path, dt, "dls", damping), trajectory, max_speed=speed, max_acceleration=acceleration, measure_time=False, ) damping_results.append(summarize_result(replay, damping)) selected_damping = choose_dls_damping(damping_results) factories = { "pinv": lambda: DifferentialIkSolver(urdf_path, dt, "pinv"), "dls": lambda: DifferentialIkSolver( urdf_path, dt, "dls", selected_damping ), "qp": lambda: make_qp_solver(urdf_path, dt, config), } results = {} summaries = {} for name, factory in factories.items(): run_replay( name, factory(), trajectory, max_speed=speed, max_acceleration=acceleration, ) result = run_replay( name, factory(), trajectory, max_speed=speed, max_acceleration=acceleration, ) timed = [] for _ in range(timing_repeats): repeated = run_replay( name, factory(), trajectory, max_speed=speed, max_acceleration=acceleration, ) if not np.allclose( repeated.joints, result.joints, atol=1e-10, rtol=0.0, ): raise RuntimeError(f"{name} replay is not deterministic") timed.append(repeated.solve_durations_ms) timing_matrix = np.stack(timed) result = replace( result, solve_durations_ms=np.mean(timing_matrix, axis=0), ) results[name] = result summaries[name] = replace( summarize_result( result, selected_damping if name == "dls" else None, ), mean_solve_ms=float(np.mean(timing_matrix)), max_solve_ms=float(np.max(timing_matrix)), ) return results, summaries, selected_damping