acRealman_xr/xr_rm_teleop/test/test_placo_transforms.py

import math
from pathlib import Path
from types import SimpleNamespace
from xml.etree import ElementTree

import numpy as np
import pytest

from xr_rm_teleop.placo_ik_solver import (
    QP_POSITION_TOLERANCE_M,
    PlacoIkSolver,
    _validated_transform,
)


def test_fixed_urdf_has_seven_moving_joints_and_omnipicker_tcp() -> None:
    urdf_path = (
        Path(__file__).resolve().parents[1]
        / "models"
        / "rm75_omnipicker"
        / "urdf"
        / "RM75-B_OmniPicker_fixed.urdf"
    )
    root = ElementTree.parse(urdf_path).getroot()
    moving_joint_names = [
        joint.attrib["name"]
        for joint in root.findall("joint")
        if joint.attrib["type"] != "fixed"
    ]
    tcp_joint = root.find("joint[@name='omnipicker_tcp_joint']")
    mesh_filenames = [
        mesh.attrib["filename"]
        for mesh in root.findall(".//mesh")
    ]

    assert moving_joint_names == [f"joint_{index}" for index in range(1, 8)]
    assert all(
        filename.startswith(
            "package://xr_rm_teleop/models/rm75_omnipicker/meshes/"
        )
        for filename in mesh_filenames
    )
    assert tcp_joint is not None
    assert tcp_joint.attrib["type"] == "fixed"
    assert tcp_joint.find("parent").attrib["link"] == "omnipicker_base_link"
    assert tcp_joint.find("child").attrib["link"] == "omnipicker_tcp"
    assert tcp_joint.find("origin").attrib["xyz"] == "0 0 0.16"
    assert tcp_joint.find("origin").attrib["rpy"] == "0 0 0"


def _rm75_placo_solver() -> tuple[PlacoIkSolver, list[float]]:
    pytest.importorskip("placo")
    urdf_path = (
        Path(__file__).resolve().parents[1]
        / "models"
        / "rm75_omnipicker"
        / "urdf"
        / "RM75-B_OmniPicker_fixed.urdf"
    )
    joints = [
        math.radians(value)
        for value in [
            -90.14,
            3.76,
            -86.89,
            87.89,
            -96.53,
            -79.62,
            -90.04,
        ]
    ]
    return PlacoIkSolver(str(urdf_path), 1.0 / 90.0), joints


def test_qp_solve_converges_to_reachable_tcp_target() -> None:
    solver, joints = _rm75_placo_solver()
    start_pose = solver.update_joint_state(joints)
    target_pose = start_pose.copy()
    target_pose[0, 3] += 0.07

    result = solver.solve(target_pose)
    reached_pose = solver.update_joint_state(result)
    position_error = np.linalg.norm(
        target_pose[:3, 3] - reached_pose[:3, 3]
    )
    rotation_delta = (
        target_pose[:3, :3] @ reached_pose[:3, :3].T
    )
    orientation_error = math.acos(
        float(
            np.clip(
                (np.trace(rotation_delta) - 1.0) * 0.5,
                -1.0,
                1.0,
            )
        )
    )

    assert position_error <= QP_POSITION_TOLERANCE_M
    assert orientation_error <= 5e-3


def test_qp_solve_accepts_position_error_within_two_millimeters() -> None:
    solver = object.__new__(PlacoIkSolver)
    solver._actual_joints = np.zeros(7)
    solver._robot = SimpleNamespace(
        state=SimpleNamespace(q=np.zeros(14))
    )
    solver._frame_task = SimpleNamespace(T_world_frame=None)
    solver._target_errors = lambda: (1.5e-3, 0.0)

    result = solver.solve(np.eye(4))

    assert result == pytest.approx([0.0] * 7)


def test_qp_solve_rejects_position_error_above_two_millimeters() -> None:
    solver = object.__new__(PlacoIkSolver)
    solver._actual_joints = np.zeros(7)
    solver._robot = SimpleNamespace(
        state=SimpleNamespace(q=np.zeros(14)),
        update_kinematics=lambda: None,
    )
    solver._frame_task = SimpleNamespace(T_world_frame=None)
    solver._solver = SimpleNamespace(solve=lambda update: None)
    solver._validate_result = lambda result, previous: None
    solver._target_errors = lambda: (2.1e-3, 0.0)

    with pytest.raises(RuntimeError, match="QP did not converge after 30"):
        solver.solve(np.eye(4))


def test_validated_transform_accepts_finite_se3_and_returns_a_copy() -> None:
    transform = np.eye(4)
    transform[:3, 3] = [0.3, -0.1, 0.2]

    actual = _validated_transform(transform)

    assert actual == pytest.approx(transform)
    assert actual is not transform


@pytest.mark.parametrize(
    "transform",
    [
        np.eye(3),
        np.full((4, 4), np.nan),
        np.vstack([np.eye(3, 4), [0.0, 0.0, 0.0, 2.0]]),
        np.diag([2.0, 1.0, 1.0, 1.0]),
    ],
)
def test_validated_transform_rejects_invalid_se3(transform: np.ndarray) -> None:
    with pytest.raises(ValueError):
        _validated_transform(transform)


def test_qp_result_rejects_nan_position_and_velocity_violations() -> None:
    solver = object.__new__(PlacoIkSolver)
    solver._joint_limits = np.asarray([[-1.0, 1.0]] * 7)
    solver._velocity_limits = np.ones(7)
    solver._dt = 0.1
    solver._actual_joints = np.zeros(7)

    with pytest.raises(ValueError, match="finite"):
        solver._validate_result(np.full(7, np.nan))
    with pytest.raises(ValueError, match="position"):
        solver._validate_result(np.full(7, 2.0))
    with pytest.raises(ValueError, match="velocity"):
        solver._validate_result(np.full(7, 0.2))