tests/test_kinematics.py

"""Tests for fpgm.robot.kinematics.ArmKinematics.

Two fixtures are used:

* A synthetic 7-DOF URDF (data-free, written to ``tmp_path``, mirroring
  ``tests/test_robot_urdf.py``'s style) for the validation test, which needs
  a joint that deliberately violates the Rz-only assumption -- something the
  real URDF, by construction, never does.
* The real PointWorld Franka+Robotiq URDF (see
  ``scripts/fetch_robot_description.py --source pointworld``), skipped if not
  present on disk, for every test that needs the exact measured numbers this
  module's design was built against (the panda_link7 FK cross-check, the
  Jacobian finite-difference check, the TCP flange-table regression, and the
  full FK/IK round trip in ``test_ik.py``).
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest

from fpgm.geometry.transforms import matrix_to_rotvec
from fpgm.objects.interaction import InteractionConfig, InteractionModel
from fpgm.robot.kinematics import ArmKinematics, KinematicsError
from fpgm.robot.urdf import RobotModel

REPO_ROOT = Path(__file__).resolve().parents[1]
REAL_URDF = (
    REPO_ROOT
    / "third_party"
    / "robot_description"
    / "pointworld_franka_robotiq_2f85"
    / "franka_panda_robotiq_2f85.urdf"
)

requires_real_urdf = pytest.mark.skipif(
    not REAL_URDF.exists(),
    reason=f"real PointWorld URDF not fetched: {REAL_URDF}",
)


@pytest.fixture(scope="module")
def real_robot() -> RobotModel:
    return RobotModel(str(REAL_URDF), load_meshes=False)


@pytest.fixture(scope="module")
def real_arm(real_robot: RobotModel) -> ArmKinematics:
    return ArmKinematics.from_robot_model(real_robot)


def _random_in_limit_configs(arm: ArmKinematics, n: int, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    lo, hi = arm.limits[:, 0], arm.limits[:, 1]
    return rng.uniform(lo, hi, size=(n, arm.n_joints))


class TestFKMatchesLinkPoses:
    @requires_real_urdf
    def test_fk_matches_link_poses_panda_link7(
        self, real_robot: RobotModel, real_arm: ArmKinematics
    ):
        """Fast joint-chain FK reproduces link_poses()['panda_link7'] near machine epsilon.

        Measured: 3.3e-16 max abs error. This is the whole justification for
        the Rz-only fast path in the module docstring -- not an approximation.
        """
        configs = _random_in_limit_configs(real_arm, 20, seed=0)
        max_err = 0.0
        for q in configs:
            poses = real_robot.link_poses(q, gripper=0.0)
            transforms = real_arm.joint_transforms(q)
            err = np.max(np.abs(transforms[-1] - poses["panda_link7"]))
            max_err = max(max_err, err)
        assert max_err < 1e-12, f"max FK error {max_err} vs link_poses too large"


class TestJacobianMatchesFiniteDifferences:
    @requires_real_urdf
    def test_jacobian_matches_central_finite_differences(self, real_arm: ArmKinematics):
        """The single highest-value test here: a wrong Jacobian still lets IK
        converge (slowly, on the wrong redundancy branch) so a bug would not
        show up as an obvious failure anywhere else.

        Measured max abs deviation from central finite differences: 3.5e-7 at
        h=1e-6 -- FD truncation/roundoff noise, not a Jacobian bug.
        """
        configs = _random_in_limit_configs(real_arm, 15, seed=1)
        h = 1e-6
        max_dev = 0.0
        for q in configs:
            for g in (0.0, 0.5, 1.0):
                jac = real_arm.jacobian(q, g)
                jac_fd = np.zeros((6, real_arm.n_joints))
                for i in range(real_arm.n_joints):
                    qp, qm = q.copy(), q.copy()
                    qp[i] += h
                    qm[i] -= h
                    tp, tm = real_arm.fk(qp, g), real_arm.fk(qm, g)
                    jac_fd[:3, i] = (tp[:3, 3] - tm[:3, 3]) / (2 * h)
                    d_rot = tp[:3, :3] @ tm[:3, :3].T
                    jac_fd[3:, i] = matrix_to_rotvec(d_rot) / (2 * h)
                max_dev = max(max_dev, float(np.max(np.abs(jac - jac_fd))))
        assert max_dev < 1e-5, f"Jacobian vs FD max deviation {max_dev} too large"


class TestAxisValidation:
    def test_non_rz_axis_raises_naming_joint(self, tmp_path: Path):
        urdf_path = _write_synthetic_urdf(tmp_path, bad_axis=True)
        robot = RobotModel(
            urdf_path,
            base_link="base_link",
            arm_joints=[f"joint{i}" for i in range(1, 8)],
            gripper_drive_joint="finger_joint",
        )
        with pytest.raises(KinematicsError, match="joint3"):
            ArmKinematics.from_robot_model(robot)

    def test_non_revolute_joint_raises_naming_joint(self, tmp_path: Path):
        urdf_path = _write_synthetic_urdf(tmp_path, prismatic=True)
        robot = RobotModel(
            urdf_path,
            base_link="base_link",
            arm_joints=[f"joint{i}" for i in range(1, 8)],
            gripper_drive_joint="finger_joint",
        )
        with pytest.raises(KinematicsError, match="joint5"):
            ArmKinematics.from_robot_model(robot)


class TestTCPRegression:
    @requires_real_urdf
    def test_flange_swings_more_than_10mm_in_z(self, real_arm: ArmKinematics):
        """The subtle trap this module exists to avoid: the TCP is not a
        constant offset from panda_link8. Measured swing is 11.3 mm; a
        constant flange transform would be wrong by up to that much.
        """
        z_open = real_arm.flange(0.0)[2]
        z_closed = real_arm.flange(1.0)[2]
        assert abs(z_closed - z_open) > 0.010

    @requires_real_urdf
    def test_tcp_matches_interaction_grasp_point(
        self, real_robot: RobotModel, real_arm: ArmKinematics
    ):
        """ArmKinematics' TCP convention must agree with
        InteractionModel.grasp_point (fingertip midpoint) to well under
        InteractionConfig.grasp_distance_m (0.05 m) -- otherwise a "perfect"
        planned grasp would be silently rejected by the interaction stage.
        """
        interaction = InteractionModel(real_robot, InteractionConfig())
        configs = _random_in_limit_configs(real_arm, 5, seed=2)
        max_err = 0.0
        for g in (0.0, 0.5, 1.0):
            for q in configs:
                poses = real_robot.link_poses(q, g)
                grasp_point = interaction.grasp_point(poses)
                tcp = real_arm.fk(q, g)[:3, 3]
                max_err = max(max_err, float(np.max(np.abs(grasp_point - tcp))))
        assert max_err < 1e-9, f"TCP vs grasp_point max error {max_err} too large"


# --------------------------------------------------------------------------- #
# Synthetic 7-DOF fixture (data-free): base -> joint1..7 -> gripper, mirroring
# tests/test_robot_urdf.py's style. Only used for the axis-validation tests
# above, which need a deliberately-violating joint the real URDF never has.
# --------------------------------------------------------------------------- #


def _write_synthetic_urdf(
    tmp_path: Path, *, bad_axis: bool = False, prismatic: bool = False
) -> Path:
    links = ["base_link"] + [f"link{i}" for i in range(1, 8)] + ["gripper_drive_link"]
    link_xml = "\n".join(f'  <link name="{name}"/>' for name in links)

    joints = []
    parent = "base_link"
    for i in range(1, 8):
        child = f"link{i}"
        joint_type = "revolute"
        axis = "0 0 1"
        if bad_axis and i == 3:
            axis = "1 0 0"  # joint3 violates the Rz-only assumption
        if prismatic and i == 5:
            joint_type = "prismatic"
        joints.append(
            f'  <joint name="joint{i}" type="{joint_type}">\n'
            f'    <parent link="{parent}"/>\n'
            f'    <child link="{child}"/>\n'
            f'    <origin xyz="0 0 {0.1 * i}" rpy="0 0 0"/>\n'
            f'    <axis xyz="{axis}"/>\n'
            f'    <limit lower="-3.0" upper="3.0" effort="10" velocity="1"/>\n'
            f"  </joint>"
        )
        parent = child
    joints.append(
        '  <joint name="finger_joint" type="revolute">\n'
        f'    <parent link="{parent}"/>\n'
        '    <child link="gripper_drive_link"/>\n'
        '    <origin xyz="0 0 0" rpy="0 0 0"/>\n'
        '    <axis xyz="0 -1 0"/>\n'
        '    <limit lower="0.0" upper="0.8" effort="10" velocity="1"/>\n'
        "  </joint>"
    )

    urdf = f"""<?xml version="1.0"?>
<robot name="synthetic_7dof">
{link_xml}
{chr(10).join(joints)}
</robot>
"""
    urdf_path = tmp_path / "synthetic_7dof.urdf"
    urdf_path.write_text(urdf)
    return urdf_path