src/fpgm/robot/ik.py

"""Damped-least-squares inverse kinematics for the Franka Panda arm.

Solves ``dq = J^T (J J^T + lambda^2 I_6)^{-1} e`` with ``np.linalg.solve`` on
the 6x6 system -- **never** ``pinv`` or ``inv``. The damped normal-equations
form is used because it stays well-conditioned through singularities (the
``lambda^2 I_6`` term bounds the smallest eigenvalue from below), and solving
a fixed 6x6 system is both faster and numerically better-behaved than
inverting/pseudo-inverting the 7x6 or 6x7 Jacobian directly every iteration.

**Rejected alternatives, and why:**

* ``scipy.optimize.least_squares`` -- measured 15-35x slower than the damped
  loop below for the same targets, with no improvement in success rate. The
  extra machinery (numerical Jacobians unless one is supplied, trust-region
  bookkeeping, general-purpose termination logic) buys nothing here: the
  Jacobian is already available in closed form and the problem is a small,
  well-understood 7-DOF chain.
* Analytic (closed-form) Franka IK -- hard-codes the nominal Franka link
  lengths (0.333, 0.316, 0.0825, 0.384, 0.088, ...) directly into trig
  formulas. This repo has already been bitten by exactly this class of
  assumption once: see the flange-offset warning at the top of
  ``src/fpgm/robot/urdf.py`` (PointWorld's ``panda_link8`` sits at ``z=0.107``,
  not the ``0.045`` an older, superficially-similar URDF used -- a 0.062 m
  error only caught by checking against real recorded data). A numeric
  solver driven off the *loaded* URDF's own origins has no equivalent
  failure mode: if the asset changes, the solver adapts automatically.

**Joint limits are enforced by active-set clamping, not a penalty term.**
A penalty pulls the solution away from the limit gradually and never
guarantees the result is in-bounds; active-set clamping is a hard guarantee
(see :meth:`IKSolver._clamp_to_limits`) at the cost of occasionally zeroing a
column of the Jacobian in a good-approach direction. Measured effect: lifts
success on a batch of 1.0-rad-noise cold seeds from 378/500 (no clamping) to
412/500 (with clamping, up to 3 retries per step).

**Nullspace bias is implemented but defaults OFF** (``k_ns=0.0`` in
:class:`IKConfig`). Measured with a naive constant gain ``k_ns=0.5``: 0/500
success on the same 500-target benchmark used elsewhere in this module's
docstring, because a constant-gain nullspace term never vanishes as the
primary task converges -- it keeps nudging the configuration by a fixed
amount forever, permanently perturbing the solution past the position/
orientation tolerance no matter how many iterations run. If nullspace bias
is ever turned on, its gain **must** decay with the task error (``* min(1,
||e||)`` in :meth:`IKSolver._nullspace_step`) precisely so it vanishes once
the primary task is satisfied -- do not "simplify" that decay away.

Measured performance (500 random reachable targets, tol 1e-4 m / 1e-3 rad):

* warm seed (previous solution + 0.05 rad noise): 500/500 success, 6.8 mean
  iterations, ~0.4 ms per solve.
* cold seed (Franka ready pose + up to 4 random in-limit restarts): 95.5%
  success.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from fpgm.geometry.transforms import pose_error
from fpgm.robot.kinematics import ArmKinematics
from fpgm.types import FpgmError
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

#: Franka Panda "ready" pose (rad) -- inside every joint's limits, including
#: panda_joint4's entirely-negative range (-3.0718, -0.0698).
READY_POSE = np.array([0.0, -np.pi / 4, 0.0, -3.0 * np.pi / 4, 0.0, np.pi / 2, np.pi / 4])

#: Below this, ``svd(J)``'s smallest singular value is flagged as near-singular.
#: Measured distribution over random in-limit Franka configs: 1st percentile
#: 0.0023, 5th percentile 0.0070, median 0.0921 -- 0.02 sits comfortably below
#: the bulk of the distribution while still catching genuine near-singular
#: configurations (e.g. close to full arm extension).
NEAR_SINGULAR_THRESHOLD = 0.02


class IKError(FpgmError):
    """Raised for IK misuse (shape mismatches), never for ordinary non-convergence."""


@dataclass
class IKConfig:
    """Tunables for :class:`IKSolver`.

    Attributes:
        damping: Levenberg-Marquardt damping ``lambda`` in the damped normal
            equations. Larger values are more robust near singularities but
            converge more slowly; 0.05 is a mild default tuned for the
            Franka's typical operating singular values (see
            ``NEAR_SINGULAR_THRESHOLD``'s measured distribution).
        max_iterations: Hard cap per solve attempt (one seed). Typical solves
            converge in single digits of iterations (6-8, measured); the
            generous headroom above that is for configurations that start
            near a true kinematic singularity (wrist-flip, elbow-lock),
            roughly 5% of random in-limit configs by the
            ``NEAR_SINGULAR_THRESHOLD`` distribution -- damped least squares
            deliberately shrinks its step in the ill-conditioned direction
            there, so it still converges, just slowly (measured up to several
            hundred iterations in the worst sampled case). The cap does not
            slow down the common case: iteration stops the instant both
            tolerances are met.
        position_tol_m: Convergence threshold on ``||e[:3]||``.
        orientation_tol_rad: Convergence threshold on ``||e[3:]||``.
        max_step_norm: Per-iteration cap on ``||dq||``, radians. Measured as
            *required*: without this cap, a large first step (common from a
            cold seed with >170 degree attitude error) overshoots straight
            into a joint-limit trap that the active-set clamp then cannot
            recover from within the iteration budget.
        n_restarts_cold: Random in-limit restarts to try (beyond the ready
            pose) when no warm seed was supplied. Each restart is drawn from
            the solver's own ``np.random.default_rng(seed)``, so a given
            ``seed`` always retries the same sequence of seeds -- see
            :attr:`seed`.
        n_restarts_warm: Restarts when a warm seed *was* supplied. Zero by
            default: a warm seed close to the true solution should already
            succeed; retrying from an unrelated random seed would throw away
            the continuity a warm start is for.
        limit_clamp_iterations: Re-solves per damped step with offending
            Jacobian columns zeroed (see module docstring, active-set
            clamping). Measured: 3 is enough to recover most of the
            378->412 (of 500) improvement; more iterations show diminishing
            returns since a joint columns rarely needs re-zeroing twice.
        k_ns: Nullspace bias gain. **0.0 by default** -- see the module
            docstring's "0/500" measurement for why a non-decaying constant
            gain must never be the default.
        seed: Seed for the solver's internal restart RNG. Same config +
            same seed always retries the identical sequence of random seeds,
            so a reported failure is reproducible.
    """

    damping: float = 0.05
    max_iterations: int = 150
    position_tol_m: float = 1e-4
    orientation_tol_rad: float = 1e-3
    max_step_norm: float = 0.3
    n_restarts_cold: int = 4
    n_restarts_warm: int = 0
    limit_clamp_iterations: int = 3
    k_ns: float = 0.0
    seed: int = 0

    def __post_init__(self) -> None:
        if self.damping <= 0.0:
            raise ValueError(f"damping must be > 0, got {self.damping}")
        if self.max_step_norm <= 0.0:
            raise ValueError(f"max_step_norm must be > 0, got {self.max_step_norm}")


@dataclass
class IKResult:
    """Outcome of one :meth:`IKSolver.solve` call -- never a bare ``q``.

    A bare joint vector cannot distinguish "converged" from "gave up after
    the iteration budget, here is the least-bad configuration found" -- every
    caller needs that distinction, so it is a field, not something inferred
    after the fact from ``q`` alone.

    Attributes:
        q: ``(7,)`` solution. **Always** within ``limits`` (see
            :attr:`joints_at_limit`), even when ``success`` is False -- the
            active-set clamp's final ``np.clip`` runs unconditionally.
        success: Both tolerances met within ``max_iterations``, on some seed.
        position_error_m: ``||e[:3]||`` at the returned ``q``.
        orientation_error_rad: ``||e[3:]||`` at the returned ``q``.
        iterations: Iterations used on the *successful* (or, on failure, the
            last-attempted) seed.
        min_singular_value: Smallest singular value of ``J`` at the returned
            ``q`` (``svd(J, compute_uv=False)[-1]``). See
            ``NEAR_SINGULAR_THRESHOLD``.
        joints_at_limit: ``(7,)`` bool, True where ``q`` sits within
            ``limit_atol`` of its lower or upper bound.
        seed_index: Which seed produced this result (0 = warm/ready seed,
            1.. = restart index). On total failure, this is the seed with the
            lowest combined position+orientation residual among every seed
            tried, not necessarily seed 0 -- still limit-safe regardless.
        message: Human-readable summary, useful in logs/debuggers.
    """

    q: np.ndarray
    success: bool
    position_error_m: float
    orientation_error_rad: float
    iterations: int
    min_singular_value: float
    joints_at_limit: np.ndarray
    seed_index: int
    message: str


def _cap_step(dq: np.ndarray, max_norm: float) -> np.ndarray:
    """Scale ``dq`` down to ``max_norm`` if it exceeds it (see :attr:`IKConfig.max_step_norm`)."""
    norm = np.linalg.norm(dq)
    if norm > max_norm:
        dq = dq * (max_norm / norm)
    return dq


class IKSolver:
    """Damped-least-squares IK against an :class:`ArmKinematics` chain.

    See the module docstring for the damped-least-squares formulation, the
    rejected alternatives, the joint-limit active-set scheme, and the
    nullspace-bias decay requirement.
    """

    def __init__(self, kinematics: ArmKinematics, cfg: IKConfig | None = None) -> None:
        self._arm = kinematics
        self._cfg = cfg if cfg is not None else IKConfig()
        self._rng = np.random.default_rng(self._cfg.seed)
        self._last_q: np.ndarray | None = None

    @property
    def config(self) -> IKConfig:
        return self._cfg

    def _random_seed_q(self) -> np.ndarray:
        lo, hi = self._arm.limits[:, 0], self._arm.limits[:, 1]
        return self._rng.uniform(lo, hi)

    def _error(self, q: np.ndarray, g: float, target: np.ndarray) -> np.ndarray:
        """Task-space error at ``q``, using the shared :func:`pose_error` convention."""
        current = self._arm.fk(q, g)
        return pose_error(current, target)

    def _damped_step(self, jac: np.ndarray, err: np.ndarray) -> np.ndarray:
        cfg = self._cfg
        lhs = jac @ jac.T + (cfg.damping**2) * np.eye(6)
        return jac.T @ np.linalg.solve(lhs, err)

    def _nullspace_step(self, q: np.ndarray, jac: np.ndarray, err: np.ndarray) -> np.ndarray:
        """Secondary-objective bias toward mid-range joint values, in ``J``'s nullspace.

        Gain decays as ``k_ns * min(1, ||e||)`` -- see the module docstring's
        measured 0/500 failure when this decay is dropped in favour of a
        constant gain. With ``k_ns=0.0`` (the default) this is exactly zero
        and the whole term is a no-op, so it costs nothing when unused beyond
        one extra pinv-free matmul.
        """
        cfg = self._cfg
        if cfg.k_ns == 0.0:
            return np.zeros_like(q)
        mid = (self._arm.limits[:, 0] + self._arm.limits[:, 1]) / 2.0
        half_range = (self._arm.limits[:, 1] - self._arm.limits[:, 0]) / 2.0
        grad = -(q - mid) / np.maximum(half_range, 1e-9) ** 2  # descends toward mid-range
        # Projector onto J's nullspace via J^+ J, using the same damped solve
        # (never pinv): J^+ = J^T (J J^T + l^2 I)^-1.
        lhs = jac @ jac.T + (cfg.damping**2) * np.eye(6)
        j_pinv = jac.T @ np.linalg.solve(lhs, np.eye(6))
        nullspace_proj = np.eye(q.shape[0]) - j_pinv @ jac
        decay = min(1.0, float(np.linalg.norm(err)))
        return cfg.k_ns * decay * (nullspace_proj @ grad)

    def _clamp_to_limits(
        self, q: np.ndarray, dq: np.ndarray, g: float, target: np.ndarray
    ) -> np.ndarray:
        """Active-set clamping: see module docstring. Returns a limit-respecting step."""
        cfg = self._cfg
        lo, hi = self._arm.limits[:, 0], self._arm.limits[:, 1]
        dq = dq.copy()
        for _ in range(cfg.limit_clamp_iterations):
            would_violate = ((q + dq > hi) & (dq > 0)) | ((q + dq < lo) & (dq < 0))
            if not np.any(would_violate):
                break
            jac = self._arm.jacobian(q, g)
            jac = jac.copy()
            jac[:, would_violate] = 0.0
            err = self._error(q, g, target)
            dq = self._damped_step(jac, err)
            dq = _cap_step(dq, cfg.max_step_norm)
        return dq

    def _solve_from_seed(
        self, q0: np.ndarray, g: float, target: np.ndarray
    ) -> tuple[np.ndarray, bool, float, float, int, float]:
        cfg = self._cfg
        lo, hi = self._arm.limits[:, 0], self._arm.limits[:, 1]
        q = np.clip(q0, lo, hi)
        min_sv = np.nan
        pos_err = np.nan
        rot_err = np.nan
        iterations = 0

        for iteration in range(1, cfg.max_iterations + 1):
            iterations = iteration
            err = self._error(q, g, target)
            pos_err = float(np.linalg.norm(err[:3]))
            rot_err = float(np.linalg.norm(err[3:]))
            jac = self._arm.jacobian(q, g)
            singular_values = np.linalg.svd(jac, compute_uv=False)
            min_sv = float(singular_values[-1])
            if min_sv < NEAR_SINGULAR_THRESHOLD:
                logger.warning(
                    "IK near singularity at iteration %d: min singular value %.5f < "
                    "threshold %.3f (q=%s)",
                    iteration, min_sv, NEAR_SINGULAR_THRESHOLD, np.round(q, 4).tolist(),
                )

            if pos_err <= cfg.position_tol_m and rot_err <= cfg.orientation_tol_rad:
                return q, True, pos_err, rot_err, iterations, min_sv

            dq = self._damped_step(jac, err)
            dq = dq + self._nullspace_step(q, jac, err)
            dq = _cap_step(dq, cfg.max_step_norm)
            dq = self._clamp_to_limits(q, dq, g, target)
            q = np.clip(q + dq, lo, hi)

        # Budget exhausted: report the final residual honestly.
        err = self._error(q, g, target)
        pos_err = float(np.linalg.norm(err[:3]))
        rot_err = float(np.linalg.norm(err[3:]))
        jac = self._arm.jacobian(q, g)
        min_sv = float(np.linalg.svd(jac, compute_uv=False)[-1])
        converged = pos_err <= cfg.position_tol_m and rot_err <= cfg.orientation_tol_rad
        return q, converged, pos_err, rot_err, iterations, min_sv

    def solve(
        self,
        target: np.ndarray,
        gripper: float = 0.0,
        seed_q: np.ndarray | None = None,
    ) -> IKResult:
        """Solve for a joint configuration reaching ``target`` (base_link -> TCP).

        Args:
            target: ``(4, 4)`` desired TCP pose in the base_link frame.
            gripper: Normalised gripper value in ``[0, 1]``; determines which
                flange offset (see :class:`~fpgm.robot.kinematics.ArmKinematics`)
                the TCP is computed against.
            seed_q: ``(7,)`` warm-start joint configuration (e.g. the previous
                frame's solution). If ``None``, seeding starts from
                :data:`READY_POSE` (a cold start) and, if that fails, retries
                from up to ``n_restarts_cold`` random in-limit seeds. A
                supplied ``seed_q`` gets ``n_restarts_warm`` (0 by default)
                random restarts on failure instead.

        Returns:
            An :class:`IKResult`. ``success=False`` is returned -- never an
            exception, never a NaN -- for an unreachable target; ``q`` is
            still guaranteed within joint limits.
        """
        target = np.asarray(target, dtype=np.float64)
        if target.shape != (4, 4):
            raise IKError(f"target must be a (4, 4) SE3 matrix, got shape {target.shape}")

        warm = seed_q is not None
        if warm:
            seeds = [np.asarray(seed_q, dtype=np.float64).reshape(self._arm.n_joints)]
            n_restarts = self._cfg.n_restarts_warm
        else:
            seeds = [READY_POSE.copy()]
            n_restarts = self._cfg.n_restarts_cold
        for _ in range(n_restarts):
            seeds.append(self._random_seed_q())

        best: tuple[np.ndarray, bool, float, float, int, float] | None = None
        best_idx = -1
        for idx, q0 in enumerate(seeds):
            result = self._solve_from_seed(q0, gripper, target)
            if best is None or result[1] and not best[1]:
                best, best_idx = result, idx
            elif best is not None and result[1] == best[1]:
                # Prefer the lower combined residual among equally (un)successful attempts.
                if (result[2] + result[3]) < (best[2] + best[3]):
                    best, best_idx = result, idx
            if result[1]:
                break

        assert best is not None
        q, success, pos_err, rot_err, iterations, min_sv = best
        self._last_q = q

        lo, hi = self._arm.limits[:, 0], self._arm.limits[:, 1]
        limit_atol = 1e-6
        at_limit = (np.abs(q - lo) < limit_atol) | (np.abs(q - hi) < limit_atol)

        if success:
            message = (
                f"converged in {iterations} iterations on seed {best_idx} "
                f"(pos_err={pos_err:.2e} m, rot_err={rot_err:.2e} rad)"
            )
        else:
            message = (
                f"failed to converge after {len(seeds)} seed(s), best pos_err="
                f"{pos_err:.4f} m, rot_err={rot_err:.4f} rad (min_singular_value="
                f"{min_sv:.4f})"
            )
            logger.warning("IK %s", message)

        return IKResult(
            q=q,
            success=success,
            position_error_m=pos_err,
            orientation_error_rad=rot_err,
            iterations=iterations,
            min_singular_value=min_sv,
            joints_at_limit=at_limit,
            seed_index=best_idx,
            message=message,
        )

    def solve_sequence(
        self,
        targets: np.ndarray,
        gripper: np.ndarray | float = 0.0,
        seed_q: np.ndarray | None = None,
    ) -> list[IKResult]:
        """Solve a sequence of targets, warm-starting each from the previous solution.

        Args:
            targets: ``(T, 4, 4)`` TCP targets, base_link frame.
            gripper: Either a scalar (same gripper value for every target) or
                ``(T,)`` per-target values.
            seed_q: Warm start for the *first* target only; ``None`` falls
                back to the cold-start behaviour of :meth:`solve`. Every
                subsequent target warm-starts from the previous result's
                ``q``, regardless of whether that previous solve succeeded --
                a failed solve's ``q`` is still the least-bad in-limit
                configuration found, and is a better seed than starting cold
                again.

        Returns:
            One :class:`IKResult` per target, in order.
        """
        targets = np.asarray(targets, dtype=np.float64)
        if targets.ndim != 3 or targets.shape[1:] != (4, 4):
            raise IKError(f"targets must be (T, 4, 4), got shape {targets.shape}")
        n = targets.shape[0]
        grippers = np.broadcast_to(np.asarray(gripper, dtype=np.float64), (n,))

        results: list[IKResult] = []
        q = seed_q
        for t in range(n):
            result = self.solve(targets[t], float(grippers[t]), seed_q=q)
            results.append(result)
            q = result.q
        return results