import math import time from types import SimpleNamespace import numpy as np import pytest from xr_rm_teleop.realman_adapter import JointStateSnapshot from xr_rm_teleop.single_arm_velocity_teleop import ( SingleArmVelocityTeleop, _make_transform, _matrix_to_quaternion, _normalize_quaternion, _project_rotation, _quaternion_to_matrix, _so3_exp, _so3_log, ) def _make_teleop_for_orientation() -> SingleArmVelocityTeleop: teleop = object.__new__(SingleArmVelocityTeleop) teleop._enable_orientation_control = True teleop._enable_orientation_axes = [True, True, True] teleop._controller_orientation_start = (0.0, 0.0, 0.0, 1.0) teleop._robot_start_transform = _make_transform( [0.3, 0.0, 0.2], _so3_exp(np.asarray([0.1, -0.2, 0.3])), ) teleop._xr_to_robot_matrix = [ 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, ] return teleop def test_identity_controller_orientation_keeps_tcp_orientation() -> None: teleop = _make_teleop_for_orientation() target = teleop._raw_orientation_from_controller((0.0, 0.0, 0.0, 1.0)) assert target == pytest.approx(teleop._robot_start_transform[:3, :3]) def test_xr_relative_rotation_maps_through_xr_to_robot_matrix() -> None: teleop = _make_teleop_for_orientation() teleop._robot_start_transform = np.eye(4) xr_roll = _matrix_to_quaternion(_so3_exp(np.asarray([0.2, 0.0, 0.0]))) target = teleop._raw_orientation_from_controller(xr_roll) assert _so3_log(target) == pytest.approx([0.0, 0.0, 0.2]) def test_quaternion_sign_does_not_change_rotation() -> None: quaternion = _normalize_quaternion((0.2, -0.3, 0.1, 0.9)) assert _quaternion_to_matrix(quaternion) == pytest.approx( _quaternion_to_matrix(tuple(-value for value in quaternion)) ) @pytest.mark.parametrize("pitch", [math.pi / 2.0 - 1e-5, -math.pi / 2.0 + 1e-5]) def test_small_rotation_near_gimbal_lock_stays_small(pitch: float) -> None: teleop = _make_teleop_for_orientation() start_rotation = _so3_exp(np.asarray([0.0, pitch, 0.0])) teleop._robot_start_transform = _make_transform([0.3, 0.0, 0.2], start_rotation) teleop._xr_to_robot_matrix = np.eye(3).reshape(-1).tolist() controller = _matrix_to_quaternion(_so3_exp(np.asarray([0.01, 0.0, 0.0]))) target = teleop._raw_orientation_from_controller(controller) error = _so3_log(target @ start_rotation.T) assert np.linalg.norm(error) == pytest.approx(0.01) def test_crossing_old_rpy_branch_uses_shortest_rotation() -> None: teleop = _make_teleop_for_orientation() start_rotation = _so3_exp(np.asarray([0.0, math.pi / 2.0 - 0.001, 0.0])) teleop._robot_start_transform = _make_transform([0.3, 0.0, 0.2], start_rotation) teleop._xr_to_robot_matrix = np.eye(3).reshape(-1).tolist() controller = _matrix_to_quaternion(_so3_exp(np.asarray([0.0, 0.002, 0.0]))) target = teleop._raw_orientation_from_controller(controller) assert _so3_log(target @ start_rotation.T) == pytest.approx( [0.0, 0.002, 0.0], abs=1e-9, ) def test_disabled_orientation_axis_zeros_robot_rotation_vector_component() -> None: teleop = _make_teleop_for_orientation() teleop._robot_start_transform = np.eye(4) teleop._xr_to_robot_matrix = np.eye(3).reshape(-1).tolist() teleop._enable_orientation_axes = [True, False, True] controller = _matrix_to_quaternion(_so3_exp(np.asarray([0.1, 0.2, 0.3]))) target = teleop._raw_orientation_from_controller(controller) assert _so3_log(target) == pytest.approx([0.1, 0.0, 0.3]) def test_orientation_deadband_filter_and_speed_limit_use_so3_angle() -> None: teleop = object.__new__(SingleArmVelocityTeleop) teleop._orientation_deadband_rad = 0.01 teleop._orientation_filter_alpha = 0.5 teleop._max_orientation_speed = 0.5 teleop._dt = 1.0 / 125.0 teleop._last_sent_orientation = np.eye(3) teleop._filtered_orientation_target = np.eye(3) inside_deadband = _so3_exp(np.asarray([0.006, 0.006, 0.0])) assert teleop._apply_orientation_deadband(inside_deadband) == pytest.approx(np.eye(3)) target = _so3_exp(np.asarray([0.2, 0.0, 0.0])) filtered = teleop._filter_orientation_target(target) assert _so3_log(filtered) == pytest.approx([0.1, 0.0, 0.0]) limited, was_limited = teleop._limit_orientation_step(target) assert was_limited assert np.linalg.norm(_so3_log(limited)) == pytest.approx(0.5 / 125.0) def test_rotation_matrix_to_debug_quaternion_is_normalized() -> None: quaternion = _matrix_to_quaternion(_so3_exp(np.asarray([0.2, -0.1, 0.3]))) assert np.isfinite(quaternion).all() assert np.linalg.norm(quaternion) == pytest.approx(1.0) def test_rotation_projection_accepts_small_error_and_rejects_invalid_matrix() -> None: near_rotation = np.eye(3) near_rotation[0, 1] = 1e-5 projected = _project_rotation(near_rotation) assert projected.T @ projected == pytest.approx(np.eye(3)) assert np.linalg.det(projected) == pytest.approx(1.0) with pytest.raises(ValueError): _project_rotation(np.diag([2.0, 1.0, 1.0])) def test_invalid_controller_quaternion_stops_current_tick() -> None: class FakeTime: def __sub__(self, other): del other return SimpleNamespace(nanoseconds=0) class FakeClock: def now(self): return FakeTime() class FakeLogger: def warn(self, *args, **kwargs): del args, kwargs teleop = object.__new__(SingleArmVelocityTeleop) teleop._last_msg = SimpleNamespace( grip=True, pose=SimpleNamespace( position=SimpleNamespace(x=0.0, y=0.0, z=0.0), orientation=SimpleNamespace(x=0.0, y=0.0, z=0.0, w=0.0), ), ) teleop._last_msg_time = FakeTime() teleop._arm_name = "test_rm75" teleop._command_timeout_sec = 0.12 teleop._enable_orientation_control = True teleop._adapter = SimpleNamespace( get_latest_joint_state=lambda: JointStateSnapshot( [0.1] * 7, time.monotonic(), ) ) teleop._ik_solver = SimpleNamespace(update_joint_state=lambda joints: np.eye(4)) teleop._active = False teleop._last_valid_joint_target = None teleop._last_current_pose = None teleop._joint_feedback_ready = True stopped = [] teleop.get_clock = lambda: FakeClock() teleop.get_logger = lambda: FakeLogger() teleop._safe_stop = lambda reset_active: stopped.append(reset_active) teleop._control_tick() assert stopped == [True] def test_zero_quaternion_is_invalid() -> None: with pytest.raises(ValueError): _normalize_quaternion([0.0, 0.0, 0.0, 0.0])