src/fpgm/robot/urdf.py

"""Franka Panda + Robotiq 2F-85 forward kinematics from a URDF.

:class:`RobotModel` wraps ``yourdfpy`` with the two things this repo's URDF needs
that ``yourdfpy`` doesn't give for free: resolving the Robotiq's ``<mimic>`` joints
from a single drive value, and mapping DROID's normalised gripper signal onto that
drive joint's actual limits. All poses are returned in ``base_link`` (``panda_link0``
by default), which is DROID's "world" frame -- the frame the recorded camera
extrinsics are expressed against.

The default asset is PointWorld's ``franka_panda_robotiq_2f85.urdf`` (see
``scripts/fetch_robot_description.py --source pointworld``), whose ``panda_link8``
sits at the real Franka flange offset (``panda_joint8`` origin ``xyz="0 0 0.107"``).
The older fairo/polymetis ``panda_robotiq_85.urdf`` placed it at ``0.045`` -- a 0.062 m
error confirmed against DROID's recorded ``cartesian_position``. The PointWorld asset
also layers a ``*_sc`` self-collision proxy link (coarse collision-only primitives)
next to each real arm link; those carry no ``<visual>`` of their own, so
:meth:`RobotModel.visual_meshes` already skips them for free and no special-casing is
needed here.
"""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import trimesh
import yourdfpy

from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

#: Franka arm joints, in DROID's ``joint_positions`` (T, 7) column order.
DEFAULT_ARM_JOINTS: tuple[str, ...] = tuple(f"panda_joint{i}" for i in range(1, 8))

#: Robotiq 2F-85 mimic wiring as documented for PointWorld's URDF (all offsets 0),
#: driven by ``finger_joint``. Used only as a fallback when the URDF itself carries
#: no ``<mimic>`` tags at all.
_FALLBACK_MIMIC_MULTIPLIERS: dict[str, float] = {
    "left_inner_knuckle_joint": 1.0,
    "left_inner_finger_joint": -1.0,
    "right_inner_knuckle_joint": -1.0,
    "right_inner_finger_joint": 1.0,
    "right_outer_knuckle_joint": -1.0,
}


def _link_visual_rgba(link: Any) -> np.ndarray | None:
    """The ``<visual><material><color rgba=...>`` a link declares, as ``(4,)`` uint8.

    Returns ``None`` when the link declares no material, which is the signal to leave
    the loaded mesh's own appearance untouched.

    **Why this is needed at all.** In this URDF the Franka's own links are COLLADA
    (``link*.dae``) and carry real materials -- trimesh loads them as ``TextureVisuals``
    with ``baseColorFactor`` white, which is correct and must not be overwritten. The
    Robotiq 2F-85's nine links are ``.STL``, a format that stores **no** colour at all;
    trimesh therefore hands back ``ColorVisuals`` filled with its own default grey,
    ``(102, 102, 102, 255)``. The URDF says those links are ``rgba="0.1 0.1 0.1 1"``
    (i.e. near-black, 25/255), but nothing was reading it, so every render drew a
    mid-grey gripper where the real Robotiq is black.

    That was not cosmetic. ``export_vace``'s ``control_vis`` channel composites this
    render over the real plate and feeds it to the video model as the appearance
    control, so a grey gripper is a grey gripper in the training signal -- and, as
    measured, in the generation.

    Only the *first* declared material is used when a link has several visuals with
    different ones; that does not occur in this URDF, and a warning fires if it ever
    does rather than silently picking one.
    """
    visuals = getattr(link, "visuals", None) or []
    rgbas = []
    for visual in visuals:
        material = getattr(visual, "material", None)
        color = getattr(material, "color", None) if material is not None else None
        rgba = getattr(color, "rgba", None) if color is not None else None
        if rgba is not None:
            rgbas.append(np.asarray(rgba, dtype=np.float64))
    if not rgbas:
        return None
    if len(rgbas) > 1 and not all(np.allclose(r, rgbas[0]) for r in rgbas[1:]):
        logger.warning(
            "link %r declares %d visuals with differing materials; using the first",
            getattr(link, "name", "<unnamed>"), len(rgbas),
        )
    return np.clip(np.round(rgbas[0] * 255.0), 0, 255).astype(np.uint8)


@dataclass(frozen=True)
class MimicJoint:
    """A joint whose position is a fixed affine function of another joint's.

    Attributes:
        name: The mimic (driven) joint's name.
        source: The joint it mimics.
        multiplier: Applied to the source joint's position.
        offset: Added after the multiplier.
    """

    name: str
    source: str
    multiplier: float
    offset: float


class RobotModel:
    """Franka Panda + Robotiq 2F-85 kinematics from a URDF.

    Wraps a ``yourdfpy.URDF`` to add: mimic-joint resolution (``yourdfpy`` parses
    ``<mimic>`` tags but does not apply them automatically -- see
    :meth:`_resolve_mimic_definitions`), a DROID-gripper-signal to drive-joint-angle
    mapping, and link poses/meshes expressed in a chosen base frame.
    """

    def __init__(
        self,
        urdf_path: str | Path,
        *,
        base_link: str = "panda_link0",
        arm_joints: Sequence[str] | None = None,
        gripper_drive_joint: str = "finger_joint",
        load_meshes: bool = True,
    ) -> None:
        """Load and validate a URDF.

        Args:
            urdf_path: Path to the URDF file.
            base_link: Frame that :meth:`link_poses` and :meth:`link_origins` return
                poses relative to. Defaults to DROID's "world" frame.
            arm_joints: Names of the arm's actuated joints, in the order they map onto
                DROID's ``joint_positions`` columns. Defaults to ``panda_joint1..7``.
            gripper_drive_joint: The Robotiq's single actuated (non-mimic) joint.
            load_meshes: Whether to load referenced mesh files. Set False to skip disk
                I/O when only kinematics (no :meth:`visual_meshes`) are needed.

        Raises:
            ValueError: If ``base_link``, any ``arm_joints`` entry, or
                ``gripper_drive_joint`` is not found among the URDF's links/actuated
                joints, or if the drive joint has no usable ``<limit>``.
        """
        # force_mesh=False keeps each <visual> as its own scene-graph node (rather
        # than merging all of a link's geometry into one mesh), which is what
        # visual_meshes() needs to return real per-visual geometry.
        self._urdf = yourdfpy.URDF.load(
            str(urdf_path),
            build_scene_graph=True,
            load_meshes=load_meshes,
            force_mesh=False,
        )

        if base_link not in self._urdf.link_map:
            raise ValueError(
                f"base_link {base_link!r} not found. Available links: "
                f"{sorted(self._urdf.link_map)}"
            )
        self._base_link = base_link

        arm_joints = tuple(arm_joints) if arm_joints is not None else DEFAULT_ARM_JOINTS
        self._validate_actuated(arm_joints, "arm_joints")
        self._validate_actuated((gripper_drive_joint,), "gripper_drive_joint")
        self._arm_joints = arm_joints
        self._gripper_drive_joint = gripper_drive_joint

        self._mimic_joints = self._resolve_mimic_definitions()

        limit = self._urdf.joint_map[gripper_drive_joint].limit
        if limit is None or limit.lower is None or limit.upper is None:
            raise ValueError(
                f"drive joint {gripper_drive_joint!r} has no usable <limit lower= upper=> "
                "in the URDF"
            )
        self._gripper_limits = (float(limit.lower), float(limit.upper))

    def _validate_actuated(self, names: Sequence[str], label: str) -> None:
        """Raise with the full list of what IS available if any ``names`` is missing.

        A silent name mismatch here would still produce a pose -- just a wrong one --
        so every name is checked eagerly at construction time.
        """
        available = set(self._urdf.actuated_joint_names)
        missing = [n for n in names if n not in available]
        if missing:
            raise ValueError(
                f"{label} not found among actuated joints: {missing}. "
                f"Available actuated joints: {sorted(available)}"
            )

    def _resolve_mimic_definitions(self) -> list[MimicJoint]:
        """Read ``<mimic>`` tags from the parsed URDF rather than hardcoding them.

        Falls back to the documented Robotiq 2F-85 wiring (multiplier only, offset 0)
        if the URDF has no ``<mimic>`` tags at all, so the model still degrades
        gracefully on a hand-edited or partial file -- but this is logged loudly
        since it means the URDF and the code have silently diverged.
        """
        found = [
            MimicJoint(
                name=joint.name,
                source=joint.mimic.joint,
                multiplier=float(joint.mimic.multiplier),
                offset=float(joint.mimic.offset),
            )
            for joint in self._urdf.robot.joints
            if joint.mimic is not None
        ]
        if found:
            return found

        logger.warning(
            "no <mimic> tags found in URDF; falling back to the documented Robotiq "
            "2F-85 wiring driven by %r",
            self._gripper_drive_joint,
        )
        return [
            MimicJoint(name=name, source=self._gripper_drive_joint, multiplier=mult, offset=0.0)
            for name, mult in _FALLBACK_MIMIC_MULTIPLIERS.items()
            if name in self._urdf.joint_map
        ]

    @property
    def link_names(self) -> list[str]:
        """All link names in the URDF, in document order."""
        return [link.name for link in self._urdf.robot.links]

    @property
    def actuated_joint_names(self) -> list[str]:
        """Names of every actuated (non-mimic, non-fixed) joint in the URDF."""
        return list(self._urdf.actuated_joint_names)

    @property
    def mimic_joints(self) -> list[MimicJoint]:
        """The resolved mimic-joint wiring (from ``<mimic>`` tags, or the fallback)."""
        return list(self._mimic_joints)

    @property
    def gripper_limits(self) -> tuple[float, float]:
        """``(lower, upper)`` limit, in radians, of the gripper drive joint."""
        return self._gripper_limits

    @property
    def arm_joints(self) -> tuple[str, ...]:
        """Arm joint names, in the order they map onto ``joint_positions`` columns."""
        return self._arm_joints

    @property
    def base_link(self) -> str:
        """The frame every :meth:`link_poses` pose is expressed relative to."""
        return self._base_link

    @property
    def gripper_drive_velocity_limit(self) -> float:
        """``<limit velocity=...>`` of the gripper drive joint, in rad/s."""
        limit = self._urdf.joint_map[self._gripper_drive_joint].limit
        return float(limit.velocity)

    def arm_joint_limits(self) -> np.ndarray:
        """``(len(arm_joints), 2)`` array of ``[lower, upper]`` position limits, radians.

        Reads straight from the URDF's ``<limit>`` tags rather than being pinned
        as a literal in caller code (see e.g. the ``panda_joint4`` entirely-negative
        range, ``-3.0718..-0.0698``) -- one source of truth so a URDF edit cannot
        silently diverge from a copy-pasted constant elsewhere in the repo.
        """
        limits = np.empty((len(self._arm_joints), 2), dtype=np.float64)
        for i, name in enumerate(self._arm_joints):
            limit = self._urdf.joint_map[name].limit
            if limit is None or limit.lower is None or limit.upper is None:
                raise ValueError(f"arm joint {name!r} has no usable <limit lower= upper=>")
            limits[i] = (float(limit.lower), float(limit.upper))
        return limits

    def arm_joint_velocity_limits(self) -> np.ndarray:
        """``(len(arm_joints),)`` array of ``<limit velocity=...>``, rad/s."""
        vels = np.empty(len(self._arm_joints), dtype=np.float64)
        for i, name in enumerate(self._arm_joints):
            limit = self._urdf.joint_map[name].limit
            if limit is None or limit.velocity is None:
                raise ValueError(f"arm joint {name!r} has no usable <limit velocity=>")
            vels[i] = float(limit.velocity)
        return vels

    def joint_origin(self, name: str) -> np.ndarray:
        """``(4, 4)`` parent-link -> child-link transform from a joint's ``<origin>``.

        This is the *fixed* joint origin (rpy, xyz), already resolved to a matrix
        by ``yourdfpy`` -- not a configuration-dependent pose. Callers rotate it
        by the joint's own axis themselves (see
        :class:`fpgm.robot.kinematics.ArmKinematics`).
        """
        joint = self._urdf.joint_map[name]
        origin = joint.origin
        if origin is None:
            return np.eye(4, dtype=np.float64)
        return np.array(origin, dtype=np.float64)

    def joint_type(self, name: str) -> str:
        """URDF ``<joint type=...>`` string (e.g. ``"revolute"``, ``"fixed"``)."""
        return str(self._urdf.joint_map[name].type)

    def joint_axis(self, name: str) -> np.ndarray:
        """``(3,)`` joint axis, in the joint's own (parent-relative) frame."""
        joint = self._urdf.joint_map[name]
        axis = joint.axis
        if axis is None:
            axis = np.array([1.0, 0.0, 0.0])
        return np.asarray(axis, dtype=np.float64).reshape(3)

    def joint_config(self, joint_positions: np.ndarray, gripper: float) -> dict[str, float]:
        """Build the full joint-angle dict, including resolved mimic joints.

        Args:
            joint_positions: ``(len(arm_joints),)`` array, e.g. DROID's 7-vector.
            gripper: Normalised DROID gripper signal in ``[0, 1]``; 0 is open. Values
                outside ``[0, 1]`` are clipped rather than rejected, since DROID's
                recorded signal occasionally drifts slightly past its nominal range.

        Returns:
            Mapping from every arm joint, the gripper drive joint, and every mimic
            joint to its resolved position, in radians.

        Raises:
            ValueError: If ``joint_positions`` is not a 1D array of the expected
                length. Never silently truncated or padded -- a length mismatch means
                the caller mixed up which array they passed in.
        """
        q = np.asarray(joint_positions, dtype=np.float64)
        if q.ndim != 1 or q.shape[0] != len(self._arm_joints):
            raise ValueError(
                f"joint_positions has shape {q.shape}, expected a 1D array of length "
                f"{len(self._arm_joints)} (one per arm joint: {list(self._arm_joints)})"
            )

        config: dict[str, float] = dict(
            zip(self._arm_joints, (float(v) for v in q), strict=True)
        )

        lower, upper = self._gripper_limits
        g = float(np.clip(gripper, 0.0, 1.0))
        config[self._gripper_drive_joint] = lower + g * (upper - lower)

        # yourdfpy parses <mimic> tags onto Joint.mimic but does not apply them when
        # building the scene graph -- each mimic joint has to be resolved here.
        # Resolved iteratively (rather than assuming source == drive joint) so a
        # mimic-of-a-mimic chain would still work, though the current URDF has none.
        pending = list(self._mimic_joints)
        while pending:
            still_pending = []
            for m in pending:
                if m.source in config:
                    config[m.name] = m.multiplier * config[m.source] + m.offset
                else:
                    still_pending.append(m)
            if len(still_pending) == len(pending):
                raise ValueError(
                    f"could not resolve mimic joints with unknown source joint(s): "
                    f"{[m.name for m in still_pending]}"
                )
            pending = still_pending

        return config

    def link_poses(self, joint_positions: np.ndarray, gripper: float) -> dict[str, np.ndarray]:
        """Forward kinematics for every link, in the ``base_link`` frame.

        Args:
            joint_positions: See :meth:`joint_config`.
            gripper: See :meth:`joint_config`.

        Returns:
            ``{link_name: (4, 4) pose}``, base_link -> link, for every link in the URDF.
        """
        self._urdf.update_cfg(self.joint_config(joint_positions, gripper))
        return {
            name: np.array(self._urdf.get_transform(frame_to=name, frame_from=self._base_link))
            for name in self.link_names
        }

    def link_origins(
        self, joint_positions: np.ndarray, gripper: float
    ) -> tuple[list[str], np.ndarray]:
        """Convenience wrapper around :meth:`link_poses` for quick reprojection checks.

        Args:
            joint_positions: See :meth:`joint_config`.
            gripper: See :meth:`joint_config`.

        Returns:
            ``(names, origins)`` where ``origins`` is ``(N, 3)``, ``names[i]``'s
            position in the ``base_link`` frame.
        """
        poses = self.link_poses(joint_positions, gripper)
        names = list(poses)
        origins = np.stack([poses[name][:3, 3] for name in names], axis=0)
        return names, origins

    def visual_meshes(self) -> dict[str, list[tuple[trimesh.Trimesh, np.ndarray]]]:
        """Per-link visual geometry, with URDF ``<visual><origin>`` and mesh ``scale``
        already applied.

        Returns:
            ``{link_name: [(mesh, mesh_to_link_transform), ...]}``. Only links with at
            least one ``<visual>`` element *and* at least one successfully loaded
            geometry are included -- e.g. with ``load_meshes=False``, links whose
            visuals are external mesh files (as opposed to primitives) are omitted.
            This also means the PointWorld URDF's ``*_sc`` self-collision proxy links
            are excluded automatically, with no special-casing needed: they carry only
            ``<collision>`` primitives, never a ``<visual>``.
            ``mesh_to_link_transform`` is a ``(4, 4)`` matrix, constant across
            configurations (it is the visual's placement within its own link, not
            affected by joint angles).
        """
        scene = self._urdf.scene
        children = scene.graph.transforms.children
        node_data = scene.graph.transforms.node_data

        result: dict[str, list[tuple[trimesh.Trimesh, np.ndarray]]] = {}
        for link in self._urdf.robot.links:
            if not link.visuals:
                continue

            link_rgba = _link_visual_rgba(link)

            meshes: list[tuple[trimesh.Trimesh, np.ndarray]] = []
            for node in children.get(link.name, []):
                geom_name = node_data[node].get("geometry")
                if geom_name is None:
                    continue  # a link-to-link joint child, not a geometry node
                # yourdfpy already folded <visual><origin> and <mesh scale=...> into
                # this transform when it built the scene graph.
                transform, _ = scene.graph.get(frame_to=node, frame_from=link.name)
                mesh = scene.geometry[geom_name]
                if link_rgba is not None:
                    # Copy: `scene.geometry` entries are shared, and recolouring in
                    # place would leak this link's colour into any other consumer of
                    # the same loaded URDF.
                    mesh = mesh.copy()
                    mesh.visual = trimesh.visual.ColorVisuals(
                        mesh=mesh, face_colors=link_rgba
                    )
                meshes.append((mesh, np.array(transform)))

            if not meshes:
                logger.warning(
                    "link %r has <visual> elements but no geometry was loaded "
                    "(missing mesh file, or load_meshes=False?)",
                    link.name,
                )
            else:
                result[link.name] = meshes

        return result