from __future__ import annotations import math import sys from pathlib import Path import h5py 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) def _write_episode(path: Path) -> None: poses = np.asarray( [ [0.1, -0.2, 0.3, 0.0, 0.0, 0.0, 1.0], [0.2, -0.2, 0.3, 0.0, 0.0, 0.0, 1.0], [0.3, -0.2, 0.3, 0.0, 0.0, 0.0, 1.0], [0.4, -0.2, 0.3, 0.0, 0.0, 0.0, 1.0], ], dtype=np.float32, ) with h5py.File(path, "w") as handle: handle.attrs["arm"] = "right_rm75" handle.attrs["pose_order"] = "x,y,z,qx,qy,qz,qw" handle.create_dataset("debug/tcp/final_target_pose", data=poses) handle.create_dataset( "debug/timestamps/control_monotonic_ns", data=np.asarray([0, 33_000_000, 66_000_000, 99_000_000]), ) handle.create_dataset( "debug/control/teleop_active", data=[0, 1, 1, 0] ) handle.create_dataset( "debug/control/action_valid", data=[1, 1, 1, 1] ) handle.create_dataset( "debug/control/command_sent", data=[0, 1, 1, 0] ) qpos = np.zeros((4, 8), dtype=np.float32) qpos[1, :7] = np.arange(7) * 0.1 handle.create_dataset("observations/qpos", data=qpos) def test_load_episode_uses_longest_valid_run_and_first_valid_qpos( tmp_path: Path, ) -> None: path = tmp_path / "episode.hdf5" _write_episode(path) actual = comparison.load_episode(path) assert actual.times_s == pytest.approx([0.0, 0.033]) assert actual.target_poses[:, 0] == pytest.approx([0.2, 0.3]) assert actual.initial_joints == pytest.approx(np.arange(7) * 0.1) def test_load_episode_rejects_wrong_arm(tmp_path: Path) -> None: path = tmp_path / "episode.hdf5" _write_episode(path) with h5py.File(path, "r+") as handle: handle.attrs.modify("arm", "left_rm75") with pytest.raises(ValueError, match="right_rm75"): comparison.load_episode(path) def test_orientation_error_and_joint_margin_match_definitions() -> None: identity = np.eye(3) quarter_turn = comparison._rotation_z(math.pi / 2.0) joints = np.asarray([0.0, -0.5]) lower = np.asarray([-1.0, -1.0]) upper = np.asarray([1.0, 3.0]) assert comparison.orientation_error_rad(identity, quarter_turn) \ == pytest.approx(math.pi / 2.0) assert comparison.normalized_joint_margin(joints, lower, upper) \ == pytest.approx(0.125) def test_choose_dls_damping_is_lexicographic() -> None: candidates = [ comparison.MethodSummary("dls", 0.01, 0.90, 0.004, 0.01, 50.0), comparison.MethodSummary("dls", 0.03, 0.95, 0.006, 0.02, 30.0), comparison.MethodSummary("dls", 0.10, 0.95, 0.004, 0.01, 40.0), ] assert comparison.choose_dls_damping(candidates) == pytest.approx(0.10) def test_limit_joint_command_reuses_production_limiter() -> None: target, velocity, limited = comparison.limit_joint_command( target=np.full(7, 1.0), previous_target=np.zeros(7), previous_velocity=np.zeros(7), max_speed=1.0, max_acceleration=10.0, dt=0.1, ) 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)