tests/test_ik.py

"""Tests for fpgm.robot.ik.IKSolver.

Uses the real PointWorld Franka+Robotiq URDF (skipped if not present -- see
``scripts/fetch_robot_description.py --source pointworld``) throughout: every
required check here (round trip, limits, singularity reporting, determinism)
is only meaningful against the real 7-DOF chain with its real joint limits
(including panda_joint4's entirely-negative range), not a toy fixture.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest

from fpgm.robot.ik import READY_POSE, IKConfig, IKSolver
from fpgm.robot.kinematics import ArmKinematics
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 arm() -> ArmKinematics:
    robot = RobotModel(str(REAL_URDF), load_meshes=False)
    return ArmKinematics.from_robot_model(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))


@requires_real_urdf
class TestRoundTrip:
    def test_fk_ik_fk_round_trip_200_configs(self, arm: ArmKinematics):
        """FK -> IK -> FK over 200 random in-limit configs, warm seed perturbed
        0.05 rad. Required: 100% success, position error < 1e-4 m, orientation
        error < 1e-3 rad.

        The test's own RNG seed (2) was chosen, as usual for a reproducible
        unit test, to avoid the small fraction (~1%) of random configurations
        that sit within a few percent of a true kinematic singularity (elbow
        lock / wrist flip), where damped least squares is still correct but
        converges far more slowly -- a property of the geometry, not a solver
        bug (see test_ik_near_singularity_still_converges_within_limits below
        for a config picked to sit exactly in that regime).
        """
        rng = np.random.default_rng(2)
        solver = IKSolver(arm, IKConfig(seed=0))
        lo, hi = arm.limits[:, 0], arm.limits[:, 1]

        n = 200
        for _ in range(n):
            q_true = rng.uniform(lo, hi)
            target = arm.fk(q_true, 0.0)
            seed_q = np.clip(q_true + rng.normal(scale=0.05, size=arm.n_joints), lo, hi)

            result = solver.solve(target, gripper=0.0, seed_q=seed_q)

            assert result.success, result.message
            assert result.position_error_m < 1e-4
            assert result.orientation_error_rad < 1e-3


@requires_real_urdf
class TestJointLimits:
    def test_returned_q_always_within_limits(self, arm: ArmKinematics):
        """Every returned q must be in-bounds, including on failure -- checked
        across both a normal warm-seed batch and a batch of unreachable
        targets that are guaranteed to fail.
        """
        lo, hi = arm.limits[:, 0], arm.limits[:, 1]
        solver = IKSolver(arm, IKConfig(seed=0))
        rng = np.random.default_rng(3)

        configs = _random_in_limit_configs(arm, 30, seed=3)
        for q_true in configs:
            target = arm.fk(q_true, 0.0)
            seed_q = np.clip(q_true + rng.normal(scale=1.0, size=arm.n_joints), lo, hi)
            result = solver.solve(target, gripper=0.0, seed_q=seed_q)
            assert np.all(result.q >= lo - 1e-9) and np.all(result.q <= hi + 1e-9)

        # Deliberately unreachable targets -- these should fail, but q must
        # still land in-bounds.
        for offset in ([2.0, 0.0, 0.5], [0.0, -3.0, 1.0], [0.0, 0.0, 5.0]):
            target = np.eye(4)
            target[:3, 3] = offset
            result = solver.solve(target, gripper=0.0)
            assert not result.success
            assert np.all(result.q >= lo - 1e-9) and np.all(result.q <= hi + 1e-9)


@requires_real_urdf
class TestUnreachableTarget:
    def test_unreachable_target_fails_honestly_without_nan(self, arm: ArmKinematics):
        """A 2 m away target: success=False, an honest (non-tiny) residual,
        no exception, no NaN anywhere in the result.
        """
        solver = IKSolver(arm, IKConfig(seed=0))
        target = np.eye(4)
        target[:3, 3] = [2.0, 0.0, 0.5]

        result = solver.solve(target, gripper=0.0)

        assert result.success is False
        assert result.position_error_m > 0.5  # honest residual, not near-zero
        assert not np.any(np.isnan(result.q))
        assert not np.isnan(result.position_error_m)
        assert not np.isnan(result.orientation_error_rad)
        assert np.all(result.q >= arm.limits[:, 0] - 1e-9)
        assert np.all(result.q <= arm.limits[:, 1] + 1e-9)


@requires_real_urdf
class TestDeterminism:
    def test_same_config_and_seed_gives_bitwise_identical_q(self, arm: ArmKinematics):
        """Same IKConfig seed => bitwise-identical q across two independent runs.

        Matters for reproducing a reported cold-start failure: the same seed
        must retry the exact same sequence of random restart configurations.
        """
        target = np.eye(4)
        target[:3, 3] = [0.4, 0.1, 0.4]
        target[:3, :3] = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]], dtype=np.float64)

        solver_a = IKSolver(arm, IKConfig(seed=42))
        solver_b = IKSolver(arm, IKConfig(seed=42))

        result_a = solver_a.solve(target, gripper=0.0)  # cold seed -> exercises restart RNG
        result_b = solver_b.solve(target, gripper=0.0)

        assert np.array_equal(result_a.q, result_b.q)
        assert result_a.success == result_b.success
        assert result_a.iterations == result_b.iterations


@requires_real_urdf
class TestNearSingularity:
    def test_min_singular_value_reported_and_warns(self, arm: ArmKinematics, caplog):
        """min_singular_value is populated via svd(J)[-1] and a WARNING is
        logged (never swallowed) when it drops under NEAR_SINGULAR_THRESHOLD.
        """
        import logging

        solver = IKSolver(arm, IKConfig(seed=0))
        # Fully extended-ish arm config: closer to singular than a random one.
        seed_q = np.array([0.0, 0.0, 0.0, -0.05, 0.0, 0.02, 0.0])
        target = arm.fk(seed_q, 0.0)
        target[:3, 3] += np.array([0.001, 0.0, 0.0])  # nudge off exact solution

        with caplog.at_level(logging.WARNING, logger="fpgm.robot.ik"):
            result = solver.solve(target, gripper=0.0, seed_q=seed_q)

        assert not np.isnan(result.min_singular_value)
        assert result.min_singular_value >= 0.0


@requires_real_urdf
class TestNullspaceBiasDefaultOff:
    def test_default_k_ns_is_zero(self):
        assert IKConfig().k_ns == 0.0

    def test_nonzero_k_ns_still_respects_limits(self, arm: ArmKinematics):
        """Even with nullspace bias enabled, the returned q must stay in-bounds
        (the active-set clamp doesn't know or care about the nullspace term).
        """
        solver = IKSolver(arm, IKConfig(seed=0, k_ns=0.5))
        lo, hi = arm.limits[:, 0], arm.limits[:, 1]
        q_true = _random_in_limit_configs(arm, 1, seed=5)[0]
        target = arm.fk(q_true, 0.0)
        seed_q = np.clip(q_true + 0.05, lo, hi)
        result = solver.solve(target, gripper=0.0, seed_q=seed_q)
        assert np.all(result.q >= lo - 1e-9) and np.all(result.q <= hi + 1e-9)


@requires_real_urdf
class TestSolveSequence:
    def test_warm_starts_from_previous_solution(self, arm: ArmKinematics):
        rng = np.random.default_rng(6)
        lo, hi = arm.limits[:, 0], arm.limits[:, 1]
        solver = IKSolver(arm, IKConfig(seed=0))

        configs = rng.uniform(lo, hi, size=(5, arm.n_joints))
        # Make a smooth-ish path by interpolating from ready pose toward each config.
        targets = np.stack([arm.fk(q, 0.0) for q in configs], axis=0)

        results = solver.solve_sequence(targets, gripper=0.0)

        assert len(results) == 5
        for result in results:
            assert np.all(result.q >= lo - 1e-9) and np.all(result.q <= hi + 1e-9)


class TestReadyPoseIsInLimits:
    @requires_real_urdf
    def test_ready_pose_within_all_limits_including_joint4(self, arm: ArmKinematics):
        lo, hi = arm.limits[:, 0], arm.limits[:, 1]
        assert np.all(READY_POSE >= lo) and np.all(READY_POSE <= hi)