tests/test_trajectory.py

"""Tests for fpgm.motion.trajectory.TrajectoryPlanner.

Uses the real PointWorld Franka+Robotiq URDF (skipped if not present -- see
``scripts/fetch_robot_description.py --source pointworld``) throughout: the
path-fidelity, joint-velocity-clamp, and gripper-timing numbers this module's
design was validated against are only meaningful against the real 7-DOF chain
with its real joint/velocity limits, not a toy fixture.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest

from fpgm.motion.profile import ProfileConfig
from fpgm.motion.trajectory import PlannerConfig, TrajectoryPlanner, summarize_plan
from fpgm.motion.types import (
    ActionSpec,
    CartesianWaypoint,
    Dwell,
    GripperAction,
    IssueKind,
    JointTrajectory,
    PlanningError,
)
from fpgm.objects.interaction import InteractionConfig, InteractionModel
from fpgm.objects.types import ObjectAlignment
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}",
)

_DOWN = [np.pi, 0.0, 0.0]


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


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


@pytest.fixture()
def ik(arm: ArmKinematics) -> IKSolver:
    return IKSolver(arm, IKConfig(seed=0))


@pytest.fixture()
def planner(arm: ArmKinematics, ik: IKSolver, robot: RobotModel) -> TrajectoryPlanner:
    return TrajectoryPlanner(arm, ik, robot, PlannerConfig())


def _tcp_positions(traj: JointTrajectory) -> np.ndarray:
    return traj.tcp_poses[:, :3, 3]


def _cartesian_segment_track_error_m(traj: JointTrajectory, seg: int) -> float:
    """Worst perpendicular distance from the planned TCP path to the commanded
    straight line for one Cartesian segment (the "path fidelity" this module's
    design was validated against -- see the reference 94 um worst case).
    """
    idx = np.flatnonzero(traj.segment_index == seg)
    start_idx = idx[0] - 1
    positions = _tcp_positions(traj)
    p0, p1 = positions[start_idx], positions[idx[-1]]
    line_vec = p1 - p0
    line_len = float(np.linalg.norm(line_vec))
    if line_len < 1e-9:
        return 0.0
    line_dir = line_vec / line_len
    worst = 0.0
    for i in idx:
        p = positions[i]
        proj = float(np.clip(np.dot(p - p0, line_dir), 0.0, line_len))
        closest = p0 + proj * line_dir
        worst = max(worst, float(np.linalg.norm(p - closest)))
    return worst


@requires_real_urdf
class TestPathFidelity:
    def test_tcp_tracks_straight_line_within_2mm(self, planner: TrajectoryPlanner):
        """A validated end-to-end reference run measured a 94 um worst case;
        this asserts the much looser 2 mm bound the design was built against.
        """
        spec = ActionSpec()
        spec.append(CartesianWaypoint(xyz=[0.45, -0.15, 0.30], rpy=_DOWN, speed_mps=0.15))
        spec.append(CartesianWaypoint(xyz=[0.45, 0.15, 0.30], rpy=_DOWN, speed_mps=0.20))
        traj = planner.plan(spec, READY_POSE.copy(), 0.0)

        assert traj.success
        for seg in (0, 1):
            err = _cartesian_segment_track_error_m(traj, seg)
            assert err < 2e-3, f"segment {seg} worst tracking error {err * 1000:.3f} mm >= 2 mm"


@requires_real_urdf
class TestJointVelocityClamp:
    def test_normal_plan_respects_velocity_limits(
        self, planner: TrajectoryPlanner, arm: ArmKinematics
    ):
        spec = ActionSpec()
        spec.append(CartesianWaypoint(xyz=[0.45, -0.15, 0.30], rpy=_DOWN, speed_mps=0.20))
        spec.append(CartesianWaypoint(xyz=[0.45, 0.15, 0.20], rpy=_DOWN, speed_mps=0.25))
        traj = planner.plan(spec, READY_POSE.copy(), 0.0)

        dt = 1.0 / PlannerConfig().profile.rate_hz
        dq = np.diff(traj.joint_positions, axis=0)
        ratio = np.abs(dq) / (dt * arm.vel_lim)
        assert np.all(ratio <= 1.0 + 1e-6)
        assert not any(i.kind is IssueKind.SPEED_CLAMPED for i in traj.issues)

    def test_deliberately_over_fast_command_triggers_retiming(
        self, arm: ArmKinematics, ik: IKSolver, robot: RobotModel
    ):
        """A segment whose raw (unsnapped) profile would demand far more than
        any joint's velocity limit must be re-timed (slowed down, more grid
        ticks) until it fits -- and the re-timing must be reported.
        """
        extreme_cfg = PlannerConfig(
            profile=ProfileConfig(a_max=1000.0, omega_max=100.0, alpha_max=1000.0)
        )
        planner = TrajectoryPlanner(arm, ik, robot, extreme_cfg)
        spec = ActionSpec()
        spec.append(CartesianWaypoint(xyz=[0.5, 0.2, 0.3], rpy=_DOWN, speed_mps=50.0))

        traj = planner.plan(spec, READY_POSE.copy(), 0.0)

        assert any(i.kind is IssueKind.SPEED_CLAMPED for i in traj.issues), traj.issues
        dt = 1.0 / extreme_cfg.profile.rate_hz
        dq = np.diff(traj.joint_positions, axis=0)
        ratio = np.abs(dq) / (dt * arm.vel_lim)
        assert np.all(ratio <= 1.0 + 1e-6)

    def test_unfixable_velocity_demand_reported_as_singularity_not_speed(
        self, arm: ArmKinematics, ik: IKSolver, robot: RobotModel
    ):
        """With the retry budget forced to zero, the same over-fast command
        cannot be re-timed at all -- it must be reported as SINGULARITY
        (unbounded joint rate for bounded Cartesian rate), never as a mere
        speed problem, per the module's design.
        """
        extreme_cfg = PlannerConfig(
            profile=ProfileConfig(a_max=1000.0, omega_max=100.0, alpha_max=1000.0),
            max_retime_iterations=0,
            on_failure="best_effort",
        )
        planner = TrajectoryPlanner(arm, ik, robot, extreme_cfg)
        spec = ActionSpec()
        spec.append(CartesianWaypoint(xyz=[0.5, 0.2, 0.3], rpy=_DOWN, speed_mps=50.0))

        traj = planner.plan(spec, READY_POSE.copy(), 0.0)

        assert not traj.success
        assert any(i.kind is IssueKind.SINGULARITY for i in traj.issues)
        assert not any(i.kind is IssueKind.SPEED_CLAMPED for i in traj.issues)


@requires_real_urdf
class TestGripperStep:
    def test_arm_holds_q_gripper_ramps_monotone_min_duration_no_drift_issue(
        self, planner: TrajectoryPlanner, robot: RobotModel
    ):
        spec = ActionSpec()
        spec.append(GripperAction(value=1.0))
        traj = planner.plan(spec, READY_POSE.copy(), 0.0)

        assert np.allclose(traj.joint_positions, traj.joint_positions[0])
        assert np.all(np.diff(traj.gripper) >= -1e-12)
        assert traj.gripper[0] == pytest.approx(0.0)
        assert traj.gripper[-1] == pytest.approx(1.0)

        lower, upper = robot.gripper_limits
        min_duration = (upper - lower) / robot.gripper_drive_velocity_limit
        assert traj.timestamps[-1] >= min_duration - 1e-9

        # The TCP genuinely moves (mimic-driven finger travel), but that must
        # never be flagged as a tracking-error issue.
        tcp_drift_m = float(
            np.linalg.norm(traj.tcp_poses[-1, :3, 3] - traj.tcp_poses[0, :3, 3])
        )
        assert tcp_drift_m > 0.005  # ~11.3 mm measured in ArmKinematics's docstring
        assert traj.issues == []
        assert np.all(traj.position_error_m == 0.0)
        assert np.all(traj.orientation_error_rad == 0.0)

    def test_gripper_value_resolved_from_width(
        self, planner: TrajectoryPlanner, arm: ArmKinematics
    ):
        target_width = 0.05
        spec = ActionSpec()
        spec.append(GripperAction(width_m=target_width))
        traj = planner.plan(spec, READY_POSE.copy(), 0.0)
        assert traj.gripper[-1] == pytest.approx(arm.gripper_value_for_width(target_width))


@requires_real_urdf
class TestDwell:
    def test_dwell_holds_state_for_at_least_requested_seconds(self, planner: TrajectoryPlanner):
        spec = ActionSpec()
        spec.append(Dwell(seconds=0.5))
        traj = planner.plan(spec, READY_POSE.copy(), 0.3)

        assert np.allclose(traj.joint_positions, traj.joint_positions[0])
        assert np.all(traj.gripper == pytest.approx(0.3))
        assert traj.timestamps[-1] >= 0.5 - 1e-9
        assert traj.issues == []


@requires_real_urdf
class TestInterfaceContract:
    def test_produced_triple_accepted_by_interaction_model_solve(
        self, planner: TrajectoryPlanner, robot: RobotModel
    ):
        """(joint_positions, gripper, timestamps) must be exactly the
        signature InteractionModel.solve already accepts -- a synthesized
        plan should drive the interaction stage with zero glue code.
        """
        spec = ActionSpec()
        spec.append(CartesianWaypoint(xyz=[0.45, -0.1, 0.30], rpy=_DOWN, speed_mps=0.15))
        spec.append(GripperAction(value=1.0))
        traj = planner.plan(spec, READY_POSE.copy(), 0.0)

        alignment = ObjectAlignment(
            transform=np.eye(4, dtype=np.float64),
            scale=1.0,
            points_world=np.zeros((1, 3), dtype=np.float64),
        )
        model = InteractionModel(robot, InteractionConfig())

        # No adaptation, no reshaping, no glue -- straight from the planner's
        # output fields into solve()'s positional arguments.
        result = model.solve(alignment, traj.joint_positions, traj.gripper, traj.timestamps)

        assert len(result) == len(traj)


@requires_real_urdf
class TestOnFailure:
    def test_raise_mode_raises_planning_error_on_unreachable_target(
        self, arm: ArmKinematics, robot: RobotModel
    ):
        # A tight iteration budget and an aggressive profile keep this test
        # fast: the target is genuinely unreachable (2 m away), so every IK
        # call along the segment will exhaust its budget either way.
        fast_ik = IKSolver(arm, IKConfig(seed=0, max_iterations=15, n_restarts_cold=1))
        cfg = PlannerConfig(profile=ProfileConfig(a_max=5.0), on_failure="raise")
        planner = TrajectoryPlanner(arm, fast_ik, robot, cfg)
        spec = ActionSpec()
        spec.append(CartesianWaypoint(xyz=[2.0, 0.0, 0.5], rpy=_DOWN, speed_mps=5.0))

        with pytest.raises(PlanningError):
            planner.plan(spec, READY_POSE.copy(), 0.0)

    def test_best_effort_mode_returns_failure_with_issues(
        self, arm: ArmKinematics, robot: RobotModel
    ):
        fast_ik = IKSolver(arm, IKConfig(seed=0, max_iterations=15, n_restarts_cold=1))
        cfg = PlannerConfig(profile=ProfileConfig(a_max=5.0), on_failure="best_effort")
        planner = TrajectoryPlanner(arm, fast_ik, robot, cfg)
        spec = ActionSpec()
        spec.append(CartesianWaypoint(xyz=[2.0, 0.0, 0.5], rpy=_DOWN, speed_mps=5.0))

        traj = planner.plan(spec, READY_POSE.copy(), 0.0)

        assert traj.success is False
        assert len(traj.issues) > 0
        # q must still be a real (finite, in-limit) trajectory -- best_effort
        # keeps the closest achieved q, never NaN.
        assert not np.any(np.isnan(traj.joint_positions))
        lo, hi = arm.limits[:, 0], arm.limits[:, 1]
        assert np.all(traj.joint_positions >= lo - 1e-9)
        assert np.all(traj.joint_positions <= hi + 1e-9)


@requires_real_urdf
class TestSummarizePlan:
    def test_summary_has_one_row_per_segment(self, planner: TrajectoryPlanner):
        spec = ActionSpec()
        spec.append(CartesianWaypoint(xyz=[0.45, -0.1, 0.30], rpy=_DOWN, speed_mps=0.15))
        spec.append(GripperAction(value=1.0))
        spec.append(Dwell(seconds=0.2))
        traj = planner.plan(spec, READY_POSE.copy(), 0.0)

        table = summarize_plan(traj)
        lines = table.strip().splitlines()
        assert len(lines) == 1 + 3  # header + one row per step
        assert "cartesian" in table
        assert "gripper" in table
        assert "dwell" in table