Spaces:
Running
Running
| """Pose comparison utilities for hardware integration tests. | |
| Provides trajectory comparison using distance_between_poses() from | |
| the reachy_mini SDK. Returns per-frame and aggregate metrics (RMSE, | |
| max, mean) for translation, rotation, and "magic distance". | |
| """ | |
| import numpy as np | |
| from dataclasses import dataclass | |
| from reachy_mini.utils.interpolation import distance_between_poses | |
| class TrajectoryMetrics: | |
| """Aggregate metrics from comparing two trajectories.""" | |
| n_frames: int | |
| # Magic distance (mm + degrees combined) | |
| magic_mean: float | |
| magic_rmse: float | |
| magic_max: float | |
| magic_p95: float | |
| # Translation (meters) | |
| trans_mean: float | |
| trans_rmse: float | |
| trans_max: float | |
| # Rotation (radians) | |
| rot_mean: float | |
| rot_rmse: float | |
| rot_max: float | |
| # Per-frame arrays | |
| magic_distances: np.ndarray | |
| trans_distances: np.ndarray | |
| rot_distances: np.ndarray | |
| def summary(self) -> str: | |
| return ( | |
| f"Trajectory comparison ({self.n_frames} frames):\n" | |
| f" Magic: mean={self.magic_mean:.1f} rmse={self.magic_rmse:.1f} " | |
| f"max={self.magic_max:.1f} p95={self.magic_p95:.1f}\n" | |
| f" Trans: mean={self.trans_mean*1000:.1f}mm " | |
| f"rmse={self.trans_rmse*1000:.1f}mm max={self.trans_max*1000:.1f}mm\n" | |
| f" Rot: mean={np.degrees(self.rot_mean):.1f}deg " | |
| f"rmse={np.degrees(self.rot_rmse):.1f}deg max={np.degrees(self.rot_max):.1f}deg" | |
| ) | |
| def frames_to_poses(frames: list[dict]) -> list[np.ndarray]: | |
| """Extract head poses (4x4 matrices) from recorded frames.""" | |
| poses = [] | |
| for f in frames: | |
| head = f["head"] | |
| poses.append(np.array(head, dtype=np.float64)) | |
| return poses | |
| def compare_trajectories( | |
| ref_times: list[float], | |
| ref_frames: list[dict], | |
| rec_times: list[float], | |
| rec_frames: list[dict], | |
| ) -> TrajectoryMetrics: | |
| """Compare two recorded trajectories frame-by-frame. | |
| Interpolates the recorded trajectory timestamps to match the reference | |
| timestamps using nearest-neighbor lookup, then computes per-frame | |
| distances using distance_between_poses(). | |
| Args: | |
| ref_times: timestamps from the reference recording | |
| ref_frames: frame dicts from the reference recording | |
| rec_times: timestamps from the new recording | |
| rec_frames: frame dicts from the new recording | |
| Returns: | |
| TrajectoryMetrics with per-frame and aggregate distances. | |
| """ | |
| ref_poses = frames_to_poses(ref_frames) | |
| rec_poses = frames_to_poses(rec_frames) | |
| rec_times_arr = np.array(rec_times) | |
| trans_dists = [] | |
| rot_dists = [] | |
| magic_dists = [] | |
| for i, (t, ref_pose) in enumerate(zip(ref_times, ref_poses)): | |
| # Find nearest recorded frame by timestamp | |
| idx = int(np.argmin(np.abs(rec_times_arr - t))) | |
| rec_pose = rec_poses[idx] | |
| d_trans, d_rot, d_magic = distance_between_poses(ref_pose, rec_pose) | |
| trans_dists.append(d_trans) | |
| rot_dists.append(d_rot) | |
| magic_dists.append(d_magic) | |
| trans_arr = np.array(trans_dists) | |
| rot_arr = np.array(rot_dists) | |
| magic_arr = np.array(magic_dists) | |
| return TrajectoryMetrics( | |
| n_frames=len(ref_times), | |
| magic_mean=float(np.mean(magic_arr)), | |
| magic_rmse=float(np.sqrt(np.mean(magic_arr**2))), | |
| magic_max=float(np.max(magic_arr)), | |
| magic_p95=float(np.percentile(magic_arr, 95)), | |
| trans_mean=float(np.mean(trans_arr)), | |
| trans_rmse=float(np.sqrt(np.mean(trans_arr**2))), | |
| trans_max=float(np.max(trans_arr)), | |
| rot_mean=float(np.mean(rot_arr)), | |
| rot_rmse=float(np.sqrt(np.mean(rot_arr**2))), | |
| rot_max=float(np.max(rot_arr)), | |
| magic_distances=magic_arr, | |
| trans_distances=trans_arr, | |
| rot_distances=rot_arr, | |
| ) | |