src/fpgm/objects/interaction.py

"""Stage 4: robot motion -> per-frame object pose.

Given the object's pose at frame 0 (stage 3's :class:`~fpgm.objects.types.
ObjectAlignment`) and the robot's recorded joint/gripper trajectory, this module
decides -- frame by frame -- whether the arm is doing nothing to the object,
pushing it, or has grasped it, and produces the resulting
:class:`~fpgm.objects.types.ObjectTrajectory`.

Two things are deliberately *not* physics:

* Pushed motion is amplified by ``push_gain`` for visibility (a DROID push is
  often a couple of centimetres -- invisible at typical render scale). This is
  an explicit, named, recorded parameter, never a silently baked-in fudge.
* The grasp/release decision uses hysteresis (separate thresholds) rather than
  a single threshold, because the raw normalised gripper signal is noisy enough
  that a bare threshold chatters the state every few frames.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from fpgm.geometry.transforms import invert_se3
from fpgm.objects.types import InteractionState, ObjectAlignment, ObjectTrajectory
from fpgm.robot.urdf import RobotModel
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)


@dataclass
class InteractionConfig:
    """Tunables for :class:`InteractionModel`.

    ``push_gain`` is a *visualisation* choice, not physics: it is always carried
    onto the output :class:`~fpgm.objects.types.ObjectTrajectory` (``push_gain``)
    so a later reader can never mistake the result for a physical simulation.
    """

    #: Multiplier on pushed object displacement. 4.0 means "5 cm of arm motion
    #: becomes 20 cm of object motion" -- exaggerated on purpose, for visibility.
    push_gain: float = 4.0
    #: Normalised DROID gripper command in [0, 1] (0 = open) above which the
    #: jaws are considered to be closing on the object.
    grasp_gripper_threshold: float = 0.15
    #: Hysteresis: once grasped, stay grasped until the gripper opens past this
    #: (lower) threshold. Prevents a noisy signal chattering grasped/free.
    release_gripper_threshold: float = 0.10
    #: A grasp is only established if the object is within this distance of the
    #: grasp point at the moment the gripper closes -- otherwise closing on
    #: empty air would incorrectly attach a distant object.
    grasp_distance_m: float = 0.05
    #: Below this distance, with jaws not yet closed enough to grasp, contact
    #: counts as pushing.
    contact_distance_m: float = 0.04
    #: URDF link names used to locate the grasp point and gripper orientation.
    left_finger_link: str = "left_inner_finger"
    right_finger_link: str = "right_inner_finger"
    gripper_base_link: str = "robotiq_85_base_link"
    #: Fallback offset (metres) along the gripper base's +Z when the finger
    #: links are absent, approximating the Robotiq 2F-85's base-to-fingertip
    #: reach.
    fallback_grasp_offset_m: float = 0.06

    def __post_init__(self) -> None:
        if self.release_gripper_threshold >= self.grasp_gripper_threshold:
            raise ValueError(
                "release_gripper_threshold must be < grasp_gripper_threshold for "
                f"hysteresis to have any effect, got release="
                f"{self.release_gripper_threshold} >= grasp="
                f"{self.grasp_gripper_threshold}"
            )


class InteractionModel:
    """Turns a robot trajectory into a per-frame object pose (pipeline stage 4).

    Works against any object exposing a ``RobotModel``-shaped ``link_poses(
    joint_positions, gripper) -> dict[str, (4, 4)]`` method -- real usage passes a
    :class:`~fpgm.robot.urdf.RobotModel`, tests pass a scripted stub.
    """

    def __init__(self, robot: RobotModel, cfg: InteractionConfig) -> None:
        """Args:
        robot: Forward-kinematics source; only ``link_poses`` is called.
        cfg: Thresholds and gain governing the grasp/push/free decision.
        """
        self._robot = robot
        self._cfg = cfg

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

    def grasp_point(self, link_poses: dict[str, np.ndarray]) -> np.ndarray:
        """Midpoint between the gripper's inner fingers, in the world frame.

        Falls back to an offset along the gripper base link's +Z axis when the
        finger links are not present in ``link_poses`` (e.g. a coarser URDF) --
        logged, since it means the grasp geometry is only approximate from that
        point on.

        Args:
            link_poses: ``{link_name: (4, 4) pose}``, as returned by
                :meth:`~fpgm.robot.urdf.RobotModel.link_poses`.

        Returns:
            ``(3,)`` world-frame position.
        """
        cfg = self._cfg
        if cfg.left_finger_link in link_poses and cfg.right_finger_link in link_poses:
            left = link_poses[cfg.left_finger_link][:3, 3]
            right = link_poses[cfg.right_finger_link][:3, 3]
            return (left + right) / 2.0

        logger.warning(
            "finger links %r/%r not found in link_poses; falling back to an "
            "offset of %.3f m along %r's +Z axis for the grasp point",
            cfg.left_finger_link,
            cfg.right_finger_link,
            cfg.fallback_grasp_offset_m,
            cfg.gripper_base_link,
        )
        base = link_poses[cfg.gripper_base_link]
        return base[:3, 3] + base[:3, :3] @ np.array([0.0, 0.0, cfg.fallback_grasp_offset_m])

    def _grip_pose(self, link_poses: dict[str, np.ndarray], grip_point: np.ndarray) -> np.ndarray:
        """Full SE3 grip frame: position at the grasp point, gripper-base rotation.

        :meth:`grasp_point` only returns a position because that's all the
        distance/contact checks need; this adds the orientation required to
        carry rotation into the grasp attachment in :meth:`solve`.
        """
        cfg = self._cfg
        if cfg.gripper_base_link in link_poses:
            orientation = link_poses[cfg.gripper_base_link][:3, :3]
        else:
            logger.warning(
                "gripper base link %r not found in link_poses; grasp attachment "
                "will not track gripper rotation",
                cfg.gripper_base_link,
            )
            orientation = np.eye(3)
        pose = np.eye(4, dtype=np.float64)
        pose[:3, :3] = orientation
        pose[:3, 3] = grip_point
        return pose

    def solve(
        self,
        alignment: ObjectAlignment,
        joint_positions: np.ndarray,
        gripper: np.ndarray,
        timestamps: np.ndarray,
    ) -> ObjectTrajectory:
        """Compute the object's per-frame world pose from the robot's motion.

        Args:
            alignment: The object's pose (and observed cloud) at frame 0 of the
                arrays below -- stage 3's output. Only ``alignment.transform`` is
                advanced; everything else on it is provenance, untouched.
            joint_positions: ``(T, 7)`` arm joint angles, one row per frame.
            gripper: ``(T,)`` normalised DROID gripper signal, 0 = open.
            timestamps: ``(T,)`` seconds, one per frame; copied onto the result.

        Returns:
            An :class:`~fpgm.objects.types.ObjectTrajectory` with one pose,
            state, and gripper distance per frame.

        Raises:
            ValueError: If ``joint_positions``, ``gripper``, and ``timestamps``
                disagree on the number of frames.
        """
        joint_positions = np.asarray(joint_positions, dtype=np.float64)
        gripper = np.asarray(gripper, dtype=np.float64)
        timestamps = np.asarray(timestamps, dtype=np.float64)
        n_frames = joint_positions.shape[0]
        if gripper.shape != (n_frames,) or timestamps.shape != (n_frames,):
            raise ValueError(
                f"frame count mismatch: joint_positions has {n_frames} frames, "
                f"gripper has shape {gripper.shape}, timestamps has shape "
                f"{timestamps.shape}"
            )

        cfg = self._cfg
        obj_pose = np.array(alignment.transform, dtype=np.float64, copy=True)
        transforms = np.empty((n_frames, 4, 4), dtype=np.float64)
        states: list[InteractionState] = []
        gripper_distance = np.empty(n_frames, dtype=np.float64)

        grasped = False
        grasp_relative: np.ndarray | None = None
        prev_grip_point: np.ndarray | None = None

        for t in range(n_frames):
            link_poses = self._robot.link_poses(joint_positions[t], float(gripper[t]))
            grip_point = self.grasp_point(link_poses)
            grip_pose = self._grip_pose(link_poses, grip_point)
            obj_position = obj_pose[:3, 3]
            dist = float(np.linalg.norm(grip_point - obj_position))
            gripper_distance[t] = dist

            # Hysteresis: only release below the lower threshold, so a signal
            # chattering around grasp_gripper_threshold does not toggle state.
            if grasped and gripper[t] < cfg.release_gripper_threshold:
                grasped = False
                grasp_relative = None

            if (
                not grasped
                and gripper[t] >= cfg.grasp_gripper_threshold
                and dist <= cfg.grasp_distance_m
            ):
                grasped = True
                # Attach by recording the object's pose relative to the grip
                # frame once, at the moment of grasping. Replaying
                # T_grip(t) @ T_grip_obj thereafter is what makes rotation and
                # lifting follow the gripper automatically -- no separate
                # rotation/translation bookkeeping needed while grasped.
                grasp_relative = invert_se3(grip_pose) @ obj_pose

            if grasped:
                state = InteractionState.GRASPED
                assert grasp_relative is not None  # set on the transition above
                obj_pose = grip_pose @ grasp_relative
            elif prev_grip_point is not None and dist <= cfg.contact_distance_m:
                push_vec = obj_position - grip_point
                push_norm = float(np.linalg.norm(push_vec))
                gripper_disp = grip_point - prev_grip_point
                advance = 0.0
                push_dir = None
                if push_norm > 1e-9:
                    push_dir = push_vec / push_norm
                    advance = float(np.dot(gripper_disp, push_dir))
                if advance > 0.0 and push_dir is not None:
                    state = InteractionState.PUSHED
                    # Only the component of gripper motion toward the object
                    # counts, so sideways or retreating motion never drags the
                    # object along. push_gain amplifies this for visibility --
                    # see InteractionConfig.push_gain; it is not physical.
                    new_pose = obj_pose.copy()
                    new_pose[:3, 3] = obj_position + push_dir * advance * cfg.push_gain
                    obj_pose = new_pose
                else:
                    state = InteractionState.FREE
            else:
                state = InteractionState.FREE

            states.append(state)
            transforms[t] = obj_pose
            prev_grip_point = grip_point

        trajectory = ObjectTrajectory(
            timestamps=timestamps,
            transforms=transforms,
            states=states,
            gripper_distance_m=gripper_distance,
            push_gain=cfg.push_gain,
        )
        logger.info("interaction timeline (%d frames):\n%s", n_frames, summarize(trajectory))
        return trajectory


def summarize(trajectory: ObjectTrajectory) -> str:
    """Render the trajectory's state spans as a compact debug table.

    This is the primary debugging artefact for this stage: eyeballing per-frame
    4x4 poses is impractical, but a handful of (state, frame range, seconds,
    displacement) rows is not.

    Args:
        trajectory: Output of :meth:`InteractionModel.solve`.

    Returns:
        A multi-line human-readable table, one row per contiguous state span.
    """
    spans = trajectory.state_spans()
    if not spans:
        return "(empty trajectory)"
    positions = trajectory.positions
    lines = [f"{'state':<8} {'frames':>13}  {'seconds':>15}  {'disp_m':>8}"]
    for state, start, end in spans:
        disp = float(np.linalg.norm(positions[end - 1] - positions[start]))
        frame_range = f"{start}-{end - 1}"
        t0, t1 = trajectory.timestamps[start], trajectory.timestamps[end - 1]
        time_range = f"{t0:.2f}-{t1:.2f}"
        lines.append(f"{state.value:<8} {frame_range:>13}  {time_range:>15}  {disp:8.4f}")
    return "\n".join(lines)