src/fpgm/motion/types.py

"""Data contracts for the Cartesian action-spec -> joint-trajectory planner.

Every pose in this module -- every :class:`CartesianWaypoint`, every
``tcp_poses`` entry on :class:`JointTrajectory` -- is expressed in the
``panda_link0`` world frame, exactly like every other frame in this repo (the
frame DROID's recorded camera extrinsics and :class:`fpgm.robot.urdf.RobotModel`
poses are already in). No stage-local frame is ever introduced here.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum

import numpy as np

from fpgm.geometry.transforms import rpy_to_matrix
from fpgm.robot.kinematics import ArmKinematics
from fpgm.types import FpgmError

__all__ = [
    "ActionSpec",
    "ActionStep",
    "CartesianWaypoint",
    "Dwell",
    "GripperAction",
    "IssueKind",
    "JointTrajectory",
    "PlanningError",
    "TrajectoryIssue",
]


class PlanningError(FpgmError):
    """Raised by :class:`fpgm.motion.trajectory.TrajectoryPlanner` when
    ``on_failure="raise"`` (the default, matching this repo's eager-validation
    posture) and a segment cannot be reached or a joint's velocity limit
    cannot be satisfied even after retiming. Also raised by an
    :class:`ActionSpec` step that is malformed (e.g. a :class:`GripperAction`
    given both or neither of ``value``/``width_m``).
    """


@dataclass
class CartesianWaypoint:
    """A target TCP pose, in the ``panda_link0`` world frame.

    Attributes:
        xyz: ``(3,)`` target position, metres.
        rpy: ``(3,)`` target orientation as URDF ``roll pitch yaw`` (see
            :func:`fpgm.geometry.transforms.rpy_to_matrix` for the exact
            axis-order convention this pins down).
        speed_mps: Desired peak Cartesian linear speed for the move into this
            waypoint -- the ``v_max`` fed to
            :func:`fpgm.motion.profile.trapezoidal_profile`. Not necessarily
            achieved: a short move never reaches cruise speed (see the
            triangular-profile branch), and a fast move that would demand more
            joint-rate than the arm has is re-timed slower (see
            :class:`fpgm.motion.trajectory.TrajectoryPlanner`).
    """

    xyz: np.ndarray
    rpy: np.ndarray
    speed_mps: float = 0.15

    def matrix(self) -> np.ndarray:
        """``(4, 4)`` SE3 pose in the ``panda_link0`` frame."""
        mat = np.eye(4, dtype=np.float64)
        mat[:3, :3] = rpy_to_matrix(np.asarray(self.rpy, dtype=np.float64))
        mat[:3, 3] = np.asarray(self.xyz, dtype=np.float64).reshape(3)
        return mat


@dataclass
class GripperAction:
    """Command the gripper to a value -- exactly one of ``value``/``width_m``.

    Attributes:
        value: Normalised DROID gripper command in ``[0, 1]`` (0 = open),
            same convention as :meth:`fpgm.robot.kinematics.ArmKinematics.flange`.
        width_m: Desired fingertip-to-fingertip separation, metres; resolved
            to a normalised value via
            :meth:`~fpgm.robot.kinematics.ArmKinematics.gripper_value_for_width`
            (the measured, monotone finger-separation table -- not a guessed
            linear mapping).
    """

    value: float | None = None
    width_m: float | None = None

    def __post_init__(self) -> None:
        if (self.value is None) == (self.width_m is None):
            raise PlanningError(
                "GripperAction requires exactly one of value or width_m, got "
                f"value={self.value!r} width_m={self.width_m!r}"
            )

    def resolve(self, arm: ArmKinematics) -> float:
        """Normalised gripper value in ``[0, 1]``, resolving ``width_m`` if needed."""
        if self.value is not None:
            return float(self.value)
        assert self.width_m is not None  # guaranteed by __post_init__
        return arm.gripper_value_for_width(self.width_m)


@dataclass
class Dwell:
    """Hold the current joint config and gripper value for ``seconds``."""

    seconds: float


ActionStep = CartesianWaypoint | GripperAction | Dwell


@dataclass
class ActionSpec:
    """An ordered list of :class:`CartesianWaypoint`/:class:`GripperAction`/:class:`Dwell`."""

    steps: list[ActionStep] = field(default_factory=list)

    def __iter__(self):
        return iter(self.steps)

    def __len__(self) -> int:
        return len(self.steps)

    def append(self, step: ActionStep) -> None:
        self.steps.append(step)


class IssueKind(Enum):
    """Why a :class:`TrajectoryIssue` was raised."""

    #: IK did not converge to the commanded target within tolerance.
    UNREACHABLE = "unreachable"
    #: The returned q sits at (or within ``joint_limit_atol`` of) a joint limit.
    JOINT_LIMIT = "joint_limit"
    #: A per-joint velocity demand could not be brought under the joint's
    #: limit by re-timing -- i.e. the Cartesian motion demanded is fine, but
    #: the *direction* in joint space needs unboundedly fast joint motion for
    #: bounded Cartesian motion. That is the signature of a true kinematic
    #: singularity, not an ordinary "go slower" fix (see
    #: :class:`fpgm.motion.trajectory.TrajectoryPlanner`).
    SINGULARITY = "singularity"
    #: A segment was re-timed (slowed down, more grid ticks) to bring its
    #: peak per-joint velocity ratio back under 1.0.
    SPEED_CLAMPED = "speed_clamped"


@dataclass
class TrajectoryIssue:
    """One flagged problem, attributable to a specific spec step and sample.

    Attributes:
        kind: See :class:`IssueKind`.
        segment_index: Index into the originating :class:`ActionSpec`'s steps
            (matches :attr:`JointTrajectory.segment_index`).
        sample_index: Index into :attr:`JointTrajectory.timestamps`, or ``-1``
            when the issue describes the whole segment rather than one sample
            (e.g. a re-timing decision).
        message: Human-readable detail, useful in logs/debuggers.
    """

    kind: IssueKind
    segment_index: int
    sample_index: int
    message: str


@dataclass
class JointTrajectory:
    """A planned, IK-realised joint trajectory at the DROID 15 Hz grid.

    **Critical interface contract**: ``(joint_positions, gripper, timestamps)``
    is *exactly* the positional signature
    :meth:`fpgm.objects.interaction.InteractionModel.solve` already accepts
    (``solve(self, alignment, joint_positions, gripper, timestamps)``). That
    is deliberate, not a coincidence: a synthesized plan from this module can
    drive the existing interaction stage and every downstream renderer with
    zero glue code, just ``model.solve(alignment, traj.joint_positions,
    traj.gripper, traj.timestamps)``. See
    ``tests/test_trajectory.py::test_interface_contract_matches_interaction_solve``.

    Attributes:
        timestamps: ``(T,)`` seconds, ``k / rate_hz`` for integer ``k`` --
            constructed directly from the grid index, never by accumulating
            ``+= dt`` in a loop, so ``np.diff(timestamps)`` is exactly
            ``1/rate_hz`` everywhere with zero float drift across segment
            boundaries (see :mod:`fpgm.motion.profile`'s grid-snapping
            docstring for why every segment lands on a tick in the first
            place). ``timestamps[0] == 0.0`` is the ``(q_start, g_start)``
            state handed to :meth:`~fpgm.motion.trajectory.TrajectoryPlanner.plan`,
            before any spec step has run.
        joint_positions: ``(T, 7)`` arm joint angles.
        gripper: ``(T,)`` normalised gripper value, 0 = open.
        tcp_poses: ``(T, 4, 4)`` FK of ``(joint_positions[t], gripper[t])`` --
            **not** necessarily the commanded Cartesian target: during a
            gripper ramp ``joint_positions`` holds constant while the TCP
            still physically translates up to 11.3 mm (the mimic-driven
            finger travel documented in
            :mod:`fpgm.robot.kinematics`'s module docstring). That is real
            motion, not tracking error -- see :attr:`position_error_m` below.
        position_error_m: ``(T,)`` IK position residual at each sample.
            Identically ``0.0`` on samples that were not produced by an IK
            solve (gripper ramps, dwells): those samples are defined as
            ``fk(q, g)`` at the held ``q``, not as a Cartesian target that IK
            was asked to hit, so there is no residual to report. Do **not**
            read a gripper-ramp sample's TCP drift (see ``tcp_poses`` above)
            as if it belonged in this column.
        orientation_error_rad: ``(T,)`` IK orientation residual, same
            convention as :attr:`position_error_m`.
        min_singular_value: ``(T,)`` smallest singular value of the Jacobian
            at each sample (see ``fpgm.robot.ik.NEAR_SINGULAR_THRESHOLD``).
        joints_at_limit: ``(T, 7)`` bool, True where a joint sits within
            ``joint_limit_atol`` of its lower or upper bound.
        segment_index: ``(T,)`` int32, the originating :class:`ActionSpec`
            step index for each sample; ``-1`` for the single initial
            ``(q_start, g_start)`` row, which precedes any step.
        issues: Flagged problems (see :class:`TrajectoryIssue`); empty when
            everything converged clean.
        success: False iff planning hit an issue that ``on_failure="best_effort"``
            tolerated instead of raising (unreachable target, or an
            unresolved-after-retiming velocity/singularity issue). Always
            True when ``on_failure="raise"`` was used, since that mode raises
            :class:`PlanningError` instead of ever returning such a result.
        joint_speed_ratio: ``(T,)`` ``max_j |dq_j| / (dt * vel_lim_j)`` between
            each sample and the previous one (0.0 for the first). This is the
            same quantity :class:`~fpgm.motion.trajectory.TrajectoryPlanner`'s
            per-joint velocity clamp checks against 1.0; kept on the result so
            :func:`~fpgm.motion.trajectory.summarize_plan` (and any other
            diagnostic) does not need to recompute it from ``vel_lim``, which
            is otherwise only known to the planner's :class:`ArmKinematics`.
        segment_labels: ``{segment_index: "cartesian" | "gripper" | "dwell"}``,
            one entry per :class:`ActionSpec` step that produced samples.
            Auxiliary bookkeeping for :func:`~fpgm.motion.trajectory.summarize_plan`;
            not part of the interface contract above.
    """

    timestamps: np.ndarray
    joint_positions: np.ndarray
    gripper: np.ndarray
    tcp_poses: np.ndarray
    position_error_m: np.ndarray
    orientation_error_rad: np.ndarray
    min_singular_value: np.ndarray
    joints_at_limit: np.ndarray
    segment_index: np.ndarray
    issues: list[TrajectoryIssue] = field(default_factory=list)
    success: bool = True
    joint_speed_ratio: np.ndarray = field(default_factory=lambda: np.zeros(0))
    segment_labels: dict[int, str] = field(default_factory=dict)

    def __len__(self) -> int:
        return int(self.timestamps.shape[0])