src/fpgm/robot/kinematics.py

"""Fast forward kinematics + analytic Jacobian for the Franka Panda arm.

:class:`RobotModel.link_poses` is the general-purpose FK path: it walks the
full ``yourdfpy`` scene graph (~40 links, including the PointWorld URDF's
``*_sc`` self-collision proxies) through ``update_cfg`` + ``get_transform``,
which costs ~0.694 ms per call (measured). That is too slow to call at every
iteration of an IK solve, and it does not expose a Jacobian at all.

:class:`ArmKinematics` instead exploits a fact specific to this URDF: **all
seven** ``panda_joint1..7`` are revolute about their own local +Z axis
(``<axis xyz="0 0 1"/>``, verified in the URDF at construction, not assumed --
see :func:`_validate_axes`). That means each joint's contribution to the
chain is just its fixed origin followed by a Z-rotation by ``cos``/``sin`` of
the joint angle, with no need to build a general rotation matrix. One loop
over 7 joints, each iteration multiplying two ``(4, 4)`` matrices, computes
FK in 31 microseconds (measured) -- 22x faster than ``link_poses()`` -- and
reproduces ``link_poses()['panda_link7']`` to 3.3e-16 max abs error (see
``tests/test_kinematics.py::test_fk_matches_link_poses``), i.e. this is not
an approximation, just a shortcut that is only valid because the axis
assumption holds.

The Jacobian is the standard analytic geometric Jacobian for a serial
revolute chain, vectorised over all 7 columns at once (see :meth:`jacobian`):
for joint i with base-frame axis ``z_i`` and origin ``p_i``, and end-effector
position ``p_e``, column i is ``[z_i x (p_e - p_i); z_i]``. Verified against
central finite differences at max abs deviation 3.5e-7 (h=1e-6) -- that
residual is FD truncation/roundoff noise, not a Jacobian bug (see
``tests/test_kinematics.py::test_jacobian_matches_finite_differences``).

The TCP (tool centre point, fingertip midpoint) is **not** a constant offset
from ``panda_link8``: because the Robotiq 2F-85's inner fingers are
mimic-driven off ``finger_joint``, the fingertip midpoint physically
translates in the ``panda_link8`` frame as the gripper opens and closes.
Measured over ``g in {0.0, 0.5, 1.0}``::

    g=0.00  offset=(0, 0.00024, 0.11498)  finger separation=0.0941 m
    g=0.50  offset=(0, 0.00029, 0.12420)  finger separation=0.0572 m
    g=1.00  offset=(0, 0.00030, 0.12627)  finger separation=0.0162 m

That is 11.3 mm of TCP travel along z between fully open and fully closed --
100x the IK solver's default 1e-4 m position tolerance. A single constant
flange transform (the usual shortcut for a fixed-jaw gripper) would silently
plan grasps up to 1.1 cm off target, which is enough to miss the
``grasp_distance_m = 0.05`` registration window in
:class:`fpgm.objects.interaction.InteractionConfig` on a marginal approach.
:meth:`ArmKinematics.flange` therefore looks the offset up from an 11-point
table built once at construction (:meth:`ArmKinematics.from_robot_model`,
11 x 0.694 ms = ~8 ms, paid once), interpolated with ``np.interp`` (measured
interpolation error < 0.05 mm).
"""

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np

from fpgm.robot.urdf import RobotModel
from fpgm.types import FpgmError
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

#: Number of gripper samples used to build the flange-offset table (0..1 inclusive).
_FLANGE_TABLE_SAMPLES = 11

#: Axis every arm joint's <axis xyz=...> must match, to float roundoff.
_Z_AXIS = np.array([0.0, 0.0, 1.0], dtype=np.float64)
_AXIS_TOL = 1e-9


class KinematicsError(FpgmError):
    """Raised when the URDF does not satisfy the Rz-only fast-path assumption."""


@dataclass
class ArmKinematics:
    """Fast FK + analytic Jacobian for the 7-DOF Franka arm, base_link -> TCP.

    Construct via :meth:`from_robot_model`, never directly -- the constructor
    assumes ``origins``/``axes``/``limits``/``vel_lim``/the flange table have
    already been validated and measured by the classmethod.

    All arrays are cached once at construction (never recomputed per call):

    Attributes:
        origins: ``(7, 4, 4)``, joint i's parent-link -> child-link origin
            transform (the URDF's fixed ``<origin rpy= xyz=>``).
        axes: ``(7, 3)``, joint i's rotation axis in its own parent-relative
            frame. Only consulted by :func:`_validate_axes` at construction;
            the FK/Jacobian hot paths assume every row equals ``[0, 0, 1]``.
        limits: ``(7, 2)`` ``[lower, upper]`` position limits, radians.
        vel_lim: ``(7,)`` velocity limits, rad/s.
        link7_to_link8: ``(4, 4)`` fixed transform from the last arm joint's
            child link (``panda_link7``) to ``panda_link8`` -- the URDF's
            ``panda_joint8`` (fixed, ``xyz="0 0 0.107"``). ``joint_transforms``
            only walks the 7 *actuated* arm joints, so this fixed hop has to
            be applied separately before the (``panda_link8``-frame) flange
            offset below means anything in the base frame.
        gripper_rotation: ``(3, 3)``, the TCP orientation relative to
            ``panda_link8`` -- constant across gripper value (only translation
            swings with the mimic-driven fingers).
        flange_g_samples: ``(11,)`` gripper values the offset table was sampled
            at, ``linspace(0, 1, 11)``.
        flange_offsets: ``(11, 3)`` TCP translation in the ``panda_link8``
            frame at each sampled gripper value. See the module docstring for
            the measured 11.3 mm swing this exists to capture.
        flange_separations: ``(11,)`` fingertip-to-fingertip separation
            (metres) at each sampled gripper value; monotone decreasing.
            Backs :meth:`gripper_value_for_width`.
        arm_joints: The 7 arm joint names, in column order.
    """

    origins: np.ndarray
    axes: np.ndarray
    limits: np.ndarray
    vel_lim: np.ndarray
    link7_to_link8: np.ndarray
    gripper_rotation: np.ndarray
    flange_g_samples: np.ndarray
    flange_offsets: np.ndarray
    flange_separations: np.ndarray
    arm_joints: tuple[str, ...] = field(default=())

    @classmethod
    def from_robot_model(cls, robot: RobotModel) -> ArmKinematics:
        """Build an :class:`ArmKinematics` from a loaded :class:`RobotModel`.

        Does 11 full :meth:`RobotModel.link_poses` calls (~0.694 ms each, so
        ~8 ms total) to build the flange-offset table -- paid once here, never
        again at solve time.

        Raises:
            KinematicsError: If any arm joint is not ``revolute`` or its axis
                is not (to 1e-9) ``[0, 0, 1]`` -- naming the offending joint.
                The Rz-only fast FK/Jacobian below is only correct under this
                assumption, so it is checked eagerly rather than assumed.
        """
        arm_joints = robot.arm_joints
        n = len(arm_joints)

        origins = np.empty((n, 4, 4), dtype=np.float64)
        axes = np.empty((n, 3), dtype=np.float64)
        for i, name in enumerate(arm_joints):
            origins[i] = robot.joint_origin(name)
            axes[i] = robot.joint_axis(name)

        _validate_axes(robot, arm_joints, axes)

        limits = robot.arm_joint_limits()
        vel_lim = robot.arm_joint_velocity_limits()

        link7_to_link8, gripper_rotation, g_samples, offsets, separations = _build_flange_table(
            robot
        )

        return cls(
            origins=origins,
            axes=axes,
            limits=limits,
            vel_lim=vel_lim,
            link7_to_link8=link7_to_link8,
            gripper_rotation=gripper_rotation,
            flange_g_samples=g_samples,
            flange_offsets=offsets,
            flange_separations=separations,
            arm_joints=arm_joints,
        )

    @property
    def n_joints(self) -> int:
        return self.origins.shape[0]

    def joint_transforms(self, q: np.ndarray) -> np.ndarray:
        """Base-frame pose of every joint's child link: ``(7, 4, 4)``.

        ``result[i]`` is ``base_link -> child_link_of_joint_i``, i.e. the same
        frame :meth:`RobotModel.link_poses` returns for ``panda_link{i+1}``.
        Built as ``T_i = T_{i-1} @ O_i @ Rz(q_i)`` in a single Python loop over
        the 7 joints, with ``Rz`` assembled by hand from ``cos``/``sin``
        (never via a general-purpose ``rotvec_to_matrix`` call) -- this is the
        31-microsecond fast path described in the module docstring.
        """
        q = np.asarray(q, dtype=np.float64).reshape(self.n_joints)
        transforms = np.empty((self.n_joints, 4, 4), dtype=np.float64)
        prev = np.eye(4, dtype=np.float64)
        for i in range(self.n_joints):
            c, s = np.cos(q[i]), np.sin(q[i])
            rz = np.array(
                [
                    [c, -s, 0.0, 0.0],
                    [s, c, 0.0, 0.0],
                    [0.0, 0.0, 1.0, 0.0],
                    [0.0, 0.0, 0.0, 1.0],
                ]
            )
            prev = prev @ self.origins[i] @ rz
            transforms[i] = prev
        return transforms

    def flange(self, g: float) -> np.ndarray:
        """TCP offset ``(3,)`` in the ``panda_link8`` frame, at gripper value ``g``.

        A *table* lookup, not a constant -- see the module docstring for why a
        constant flange offset would be wrong by up to 11.3 mm. ``g`` outside
        ``[0, 1]`` is extrapolated by ``np.interp``'s clamp-to-endpoint
        behaviour rather than raising, matching :meth:`RobotModel.joint_config`'s
        own clipping of an out-of-range gripper signal.
        """
        g = float(g)
        return np.array(
            [
                np.interp(g, self.flange_g_samples, self.flange_offsets[:, 0]),
                np.interp(g, self.flange_g_samples, self.flange_offsets[:, 1]),
                np.interp(g, self.flange_g_samples, self.flange_offsets[:, 2]),
            ],
            dtype=np.float64,
        )

    def flange_transform(self, g: float) -> np.ndarray:
        """``(4, 4)`` panda_link8 -> TCP transform at gripper value ``g``."""
        mat = np.eye(4, dtype=np.float64)
        mat[:3, :3] = self.gripper_rotation
        mat[:3, 3] = self.flange(g)
        return mat

    def link7_to_tcp_transform(self, g: float) -> np.ndarray:
        """``(4, 4)`` panda_link7 -> TCP transform at gripper value ``g``.

        Composes the fixed ``panda_joint8`` hop (:attr:`link7_to_link8`) with
        the gripper-value-dependent flange offset (:meth:`flange_transform`,
        measured in the ``panda_link8`` frame) -- both are needed to get from
        the last *actuated* joint's link to the TCP.
        """
        return self.link7_to_link8 @ self.flange_transform(g)

    def gripper_value_for_width(self, width_m: float) -> float:
        """Invert the measured (monotone decreasing) finger-separation table.

        ``width_m`` is fingertip-to-fingertip separation, metres. Extrapolates
        by clamping to the nearest sampled endpoint, same as :meth:`flange`.
        Preferred over a magic ``close=1.0`` constant because it is grounded
        in the same measured table the TCP offset itself comes from.
        """
        # flange_separations is monotone decreasing in g (0.0941 m @ g=0 down
        # to 0.0162 m @ g=1); np.interp requires its xp strictly increasing,
        # so both arrays are reversed for the lookup.
        widths_increasing = self.flange_separations[::-1]
        g_increasing = self.flange_g_samples[::-1]
        return float(np.interp(width_m, widths_increasing, g_increasing))

    def fk(self, q: np.ndarray, g: float) -> np.ndarray:
        """``(4, 4)`` base_link -> TCP pose at joint config ``q`` and gripper ``g``."""
        transforms = self.joint_transforms(q)
        return transforms[-1] @ self.link7_to_tcp_transform(g)

    def jacobian(self, q: np.ndarray, g: float) -> np.ndarray:
        """``(6, 7)`` analytic geometric Jacobian, base_link -> TCP, at ``(q, g)``.

        Rows 0:3 are linear velocity, rows 3:6 angular velocity, both expressed
        in the base frame. Column i (revolute joint i) is::

            J[:3, i] = z_i x (p_e - p_i)
            J[3:, i] = z_i

        with ``z_i = T_i[:3, 2]`` (joint i's base-frame axis) and
        ``p_i = T_i[:3, 3]`` (joint i's base-frame origin), vectorised over all
        7 columns via ``np.cross`` rather than looped. See the module
        docstring for the finite-difference cross-check (max abs deviation
        3.5e-7 at h=1e-6).
        """
        transforms = self.joint_transforms(q)
        p_e = (transforms[-1] @ self.link7_to_tcp_transform(g))[:3, 3]
        z = transforms[:, :3, 2]  # (7, 3)
        p = transforms[:, :3, 3]  # (7, 3)
        jac = np.empty((6, self.n_joints), dtype=np.float64)
        jac[:3] = np.cross(z, p_e - p).T
        jac[3:] = z.T
        return jac


def _validate_axes(robot: RobotModel, arm_joints: tuple[str, ...], axes: np.ndarray) -> None:
    """Assert every arm joint is revolute about its local +Z, naming the offender.

    The fast FK/Jacobian above hard-code ``Rz(q)`` per joint; if even one arm
    joint were prismatic, fixed, or revolute about a different axis, that
    shortcut would silently produce a wrong pose. Checked once, eagerly, at
    construction -- this repo's convention (see
    ``RobotModel._validate_actuated``) is to fail loudly at load time rather
    than compute a plausible-looking wrong answer.
    """
    bad: list[str] = []
    for name, axis in zip(arm_joints, axes, strict=True):
        joint_type = robot.joint_type(name)
        if joint_type != "revolute":
            bad.append(f"{name} (type={joint_type!r}, expected 'revolute')")
            continue
        if np.linalg.norm(axis - _Z_AXIS) >= _AXIS_TOL:
            bad.append(f"{name} (axis={axis.tolist()}, expected [0, 0, 1])")
    if bad:
        raise KinematicsError(
            "ArmKinematics requires every arm joint to be revolute about local +Z "
            "(the Rz-only fast path is only valid under that assumption); "
            f"offending joint(s): {bad}"
        )


def _build_flange_table(
    robot: RobotModel,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Sample :meth:`RobotModel.link_poses` at zero arm config across 11 gripper values.

    Returns ``(link7_to_link8, gripper_rotation, g_samples, offsets,
    separations)`` where ``offsets[k]`` is the TCP (fingertip midpoint)
    position in the ``panda_link8`` frame at ``g_samples[k]``,
    ``separations[k]`` is the fingertip-to-fingertip distance at the same
    sample, and ``link7_to_link8`` is the fixed ``panda_joint8`` hop (constant
    across both ``q`` and ``g``, so it is read once at ``k=0``).

    TCP position/orientation convention matches
    :meth:`fpgm.objects.interaction.InteractionModel.grasp_point` exactly
    (midpoint of the two inner-finger link origins; orientation from
    ``robotiq_85_base_link``) -- see the TCP regression test in
    ``tests/test_kinematics.py``.
    """
    zeros = np.zeros(len(robot.arm_joints), dtype=np.float64)
    g_samples = np.linspace(0.0, 1.0, _FLANGE_TABLE_SAMPLES)
    offsets = np.empty((_FLANGE_TABLE_SAMPLES, 3), dtype=np.float64)
    separations = np.empty(_FLANGE_TABLE_SAMPLES, dtype=np.float64)
    gripper_rotation: np.ndarray | None = None
    link7_to_link8: np.ndarray | None = None

    for k, g in enumerate(g_samples):
        poses = robot.link_poses(zeros, float(g))
        link7 = poses["panda_link7"]
        link8 = poses["panda_link8"]
        left = poses["left_inner_finger"][:3, 3]
        right = poses["right_inner_finger"][:3, 3]
        tcp_world = (left + right) / 2.0
        # panda_link8 -> TCP offset: rotate the world-frame difference into
        # panda_link8's own frame (its rotation block, transposed).
        offsets[k] = link8[:3, :3].T @ (tcp_world - link8[:3, 3])
        separations[k] = float(np.linalg.norm(left - right))
        if gripper_rotation is None:
            gripper_rotation = link8[:3, :3].T @ poses["robotiq_85_base_link"][:3, :3]
        if link7_to_link8 is None:
            link7_to_link8 = np.eye(4, dtype=np.float64)
            link7_to_link8[:3, :3] = link7[:3, :3].T @ link8[:3, :3]
            link7_to_link8[:3, 3] = link7[:3, :3].T @ (link8[:3, 3] - link7[:3, 3])

    logger.info(
        "flange table built: g=0.0 offset=%s sep=%.4f; g=1.0 offset=%s sep=%.4f",
        np.round(offsets[0], 5).tolist(),
        separations[0],
        np.round(offsets[-1], 5).tolist(),
        separations[-1],
    )
    assert gripper_rotation is not None
    assert link7_to_link8 is not None
    return link7_to_link8, gripper_rotation, g_samples, offsets, separations