Share one MJCF kinematics between the drivers that solve for themselves

positronic/drivers/roboarm/mjcf_kinematics.py

@@ -0,0 +1,106 @@
+"""FK and IK on an MJCF model, for a driver that solves for its own arm.
+
+An arm whose controller solves Cartesian goals itself does not need this. One whose controller does not,
+or whose firmware solves them in a way the driver cannot accept, carries its own model and solves here.
+
+Its own module rather than beside the placo ``Kinematics``: that one imports ``placo``, which only the
+``hardware`` extra carries, and a driver that solves against an MJCF must not need it.
+"""
+
+from collections.abc import Callable, Iterable, Sequence
+
+import mujoco as mj
+import numpy as np
+
+from positronic import geom
+from positronic.utils import package_assets_path
+
+from .ik import qpos_from_site_pose
+
+
+class MjcfKinematics:
+    """FK and IK on an MJCF model at one site, in the model's own base frame.
+
+    ``mujoco`` exports every symbol below from a compiled extension, so a type checker cannot see them.
+
+    :param mjcf_path: the model, relative to the package assets.
+    :param site: the site FK reports and IK aims at.
+    :param joint_names: the joints, in the order a caller states them.
+    :param reach_postures: warm starts for a target, tried after the live posture. The arms this drives
+        have a 6-DoF wrist, which gives LM no null space to escape a bad basin, so a seed near the goal is
+        what makes limit-clamped IK reliable.
+    """
+
+    # FK-verify acceptance for a solution, after wrapping and clamping it into joint range
+    _POS_TOL = 1e-3  # meters
+    _ROT_TOL = 1e-2  # radians
+
+    def __init__(
+        self,
+        mjcf_path: str,
+        site: str,
+        joint_names: Sequence[str],
+        reach_postures: Callable[[geom.Transform3D], Iterable[np.ndarray]],
+    ):
+        # rules-allow: primitive-type — `package_assets_path(relative_path: str) -> str` owns the join, and
+        # every caller spells the model this way
+        self._model = mj.MjModel.from_xml_path(package_assets_path(mjcf_path))
+        self._data = mj.MjData(self._model)
+        self._site_id = mj.mj_name2id(self._model, mj.mjtObj.mjOBJ_SITE, site)
+        if self._site_id < 0:
+            # Indexing with -1 reads the model's last site, so the arm would solve for another frame.
+            raise ValueError(f'{mjcf_path} names no site {site!r}')
+        self._qpos_ids = np.array([self._model.joint(name).qposadr.item() for name in joint_names])
+        self._dof_ids = np.array([self._model.joint(name).dofadr.item() for name in joint_names])
+        ranges = np.array([self._model.joint(name).range for name in joint_names])
+        self.lower, self.upper = ranges[:, 0], ranges[:, 1]
+        self._reach_postures = reach_postures
+
+    def fk(self, q: np.ndarray) -> geom.Transform3D:
+        self._data.qpos[self._qpos_ids] = q
+        mj.mj_kinematics(self._model, self._data)
+        quat = np.empty(4)
+        mj.mju_mat2Quat(quat, self._data.site_xmat[self._site_id].copy())
+        return geom.Transform3D(self._data.site_xpos[self._site_id].copy(), geom.Rotation.from_quat(quat))
+
+    def ik(
+        self, target: geom.Transform3D, current_q: np.ndarray, max_jump: float | np.ndarray | None = None
+    ) -> np.ndarray | None:
+        """The joints that reach ``target``, warm-started from where the arm stands, or ``None``.
+
+        ``max_jump`` bounds how far the solution may sit from ``current_q``, per joint or over all of them,
+        and the search stops at the live posture: the arm keeps the shape it has, and a pose it can reach
+        only in another one comes back as nothing. Without it the reach postures are tried too, so the arm
+        may change shape to get there.
+
+        A solution is wrapped and clamped into joint range and then FK-verified, so a target the arm cannot
+        reach comes back as nothing rather than as the nearest thing the solver stopped at.
+        """
+        seeds = (current_q,) if max_jump is not None else (current_q, *self._reach_postures(target))
+        for start in seeds:
+            self._data.qpos[:] = 0.0
+            self._data.qpos[self._qpos_ids] = start
+            qpos, _, success = qpos_from_site_pose(
+                self._model,
+                self._data,
+                self._site_id,
+                self._dof_ids,
+                target.translation,
+                target.rotation.as_quat,
+                rot_weight=0.5,
+            )
+            if not success:
+                continue
+            q = qpos[self._qpos_ids].copy()
+            # A revolute joint at q ± 2π is the same pose; wrap out-of-range entries back in when they fit.
+            q = np.where(q > self.upper, q - 2 * np.pi, q)
+            q = np.where(q < self.lower, q + 2 * np.pi, q)
+            q = np.clip(q, self.lower, self.upper)
+            if max_jump is not None and np.any(np.abs(q - current_q) > max_jump):
+                continue
+            reached = self.fk(q)
+            rot_err = (reached.rotation.inv * target.rotation).angle
+            rot_err = min(rot_err, 2 * np.pi - rot_err)
+            if np.linalg.norm(reached.translation - target.translation) < self._POS_TOL and rot_err < self._ROT_TOL:
+                return q
+        return None

positronic/drivers/roboarm/tests/test_kinematics.py

@@ -0,0 +1,76 @@
+"""What ``MjcfKinematics`` answers on the YAM model, held to a fixed corpus.
+
+``yam_kinematics_goldens.npz`` holds ten joint vectors inside the model's range, the pose each puts
+``DEFAULT_FRAME`` at, the joints IK finds for that pose from a seed 0.15 rad away, and two targets outside
+the arm's reach. The numbers are the answers, not a sample of them: a run that disagrees with the file has
+changed what the class solves, so regenerating it is a change of behaviour rather than of test data.
+"""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from positronic import geom
+from positronic.drivers.roboarm.mjcf_kinematics import MjcfKinematics
+from positronic.drivers.roboarm.models import DEFAULT_FRAME
+
+_YAM_MJCF = 'assets/mujoco/i2rt_yam/yam.xml'
+_YAM_JOINTS = ('joint1', 'joint2', 'joint3', 'joint4', 'joint5', 'joint6')
+_GOLDENS = Path(__file__).with_name('yam_kinematics_goldens.npz')
+
+
+def _yam_reach_postures(target: geom.Transform3D) -> list[np.ndarray]:
+    """The YAM's warm starts: joint1 at the target's azimuth, elbow folded down at two heights. The corpus
+    is answers from these seeds, so it is read with them."""
+    az = np.arctan2(target.translation[1], target.translation[0])
+    return [np.array([az, 1.8, 2.2, 0.0, -0.9, 0.0]), np.array([az, 1.2, 1.2, 0.0, 0.6, 0.0])]
+
+
+@pytest.fixture(scope='module')
+def yam() -> MjcfKinematics:
+    return MjcfKinematics(_YAM_MJCF, DEFAULT_FRAME, _YAM_JOINTS, _yam_reach_postures)
+
+
+@pytest.fixture(scope='module')
+def goldens() -> dict[str, np.ndarray]:
+    return dict(np.load(_GOLDENS))
+
+
+def _pose(row: np.ndarray) -> geom.Transform3D:
+    return geom.Transform3D(row[:3], geom.Rotation.from_quat(row[3:]))
+
+
+def test_forward_kinematics_answer_what_the_yam_driver_answered(yam, goldens):
+    for q, want in zip(goldens['q'], goldens['fk'], strict=True):
+        got = yam.fk(q)
+        np.testing.assert_allclose(got.translation, want[:3], atol=1e-9)
+        assert geom.quat_closest(got.rotation, geom.Rotation.from_quat(want[3:])) == geom.Rotation.from_quat(want[3:])
+
+
+def test_inverse_kinematics_answer_what_the_yam_driver_answered(yam, goldens):
+    for target, seed, want, solved in zip(
+        goldens['ik_target'], goldens['ik_seed'], goldens['ik_q'], goldens['ik_solved'], strict=True
+    ):
+        got = yam.ik(_pose(target), seed)
+        if not solved:
+            assert got is None
+            continue
+        assert got is not None
+        np.testing.assert_allclose(got, want, atol=1e-9)
+
+
+def test_a_capped_search_keeps_the_shape_the_arm_stands_in(yam, goldens):
+    """``max_jump`` is what a streamed setpoint is solved under: no seed but the live posture, and no
+    solution further from it than the cap."""
+    target, seed = _pose(goldens['ik_target'][0]), goldens['ik_seed'][0]
+
+    assert yam.ik(target, seed, max_jump=1e-6) is None
+
+    got = yam.ik(target, seed, max_jump=1.0)
+    assert got is not None and np.all(np.abs(got - seed) <= 1.0)
+
+
+def test_a_target_beyond_the_arm_is_refused_rather_than_approached(yam):
+    """The solver stops wherever it got to, so an unverified answer would be a pose the arm cannot hold."""
+    assert yam.ik(geom.Transform3D(np.array([2.0, 0.0, 0.3])), np.zeros(6)) is None

positronic/drivers/roboarm/yam.py

@@ -18,18 +18,16 @@
 from collections.abc import Callable, Generator, Iterator
 from typing import Any
 
-import mujoco as mj
 import numpy as np
 
 import pimm
 from positronic import geom
 from positronic.drivers import vendor_import
 from positronic.drivers.roboarm import keys as roboarm_keys
 from positronic.drivers.utils import DriverRun, MoveAbandoned, MoveStatus, log_failure
-from positronic.utils import package_assets_path
 
 from . import RobotStatus, State, command
-from .ik import qpos_from_site_pose
+from .mjcf_kinematics import MjcfKinematics
 from .models import DEFAULT_FRAME
 
 # i2rt lives in the `yam` extra, which the type-check environment does not install.
@@ -43,19 +41,17 @@
 # TODO(#517): centralise driver kinematics so driver and sim share one module.
 _JOINT_NAMES = ('joint1', 'joint2', 'joint3', 'joint4', 'joint5', 'joint6')
 _MJCF_PATH = 'assets/mujoco/i2rt_yam/yam.xml'
-_IK_POS_TOL = 1e-3  # meters; FK-verify acceptance for an IK solution after limit clamping
-_IK_ROT_TOL = 1e-2  # radians
 # Where the driver leaves the chain when it takes control: the menagerie "home" keyframe, folded up and back.
 _PARK_JOINTS = np.array([0.0, 1.047, 1.047, 0.0, 0.0, 0.0])
 # The vendor's observation contract
 _JOINT_POS, _JOINT_VEL, _GRIPPER_POS = 'joint_pos', 'joint_vel', 'gripper_pos'
 
 
-def _reach_postures(x: float, y: float) -> list[np.ndarray]:
-    """IK warm-start candidates for reaching toward arm-base-frame point (x, y): joint1 swung to the target's
-    azimuth, elbow folded down at two heights. The 6-DoF wrist gives LM no null space to escape bad basins,
-    so seeding near the goal is what makes limit-clamped IK reliable."""
-    az = np.arctan2(y, x)
+def _reach_postures(target: geom.Transform3D) -> list[np.ndarray]:
+    """IK warm-start candidates for reaching toward ``target``: joint1 swung to its azimuth, elbow folded
+    down at two heights. The 6-DoF wrist gives LM no null space to escape bad basins, so seeding near the
+    goal is what makes limit-clamped IK reliable."""
+    az = np.arctan2(target.translation[1], target.translation[0])
     return [np.array([az, 1.8, 2.2, 0.0, -0.9, 0.0]), np.array([az, 1.2, 1.2, 0.0, 0.6, 0.0])]
 
 
@@ -102,60 +98,6 @@ def encode(self, q: np.ndarray, dq: np.ndarray, ee_pose: geom.Transform3D, statu
         self.array[YamState.STATUS_OFFSET] = status.value
 
 
-class _Kinematics:
-    """FK/IK on the vendored YAM MJCF at ``DEFAULT_FRAME``, in the arm-base frame.
-
-    ``mujoco`` exports every symbol below from a compiled extension, so a type checker cannot see them.
-    """
-
-    def __init__(self):
-        model_path = package_assets_path(_MJCF_PATH)
-        self._model = mj.MjModel.from_xml_path(model_path)
-        self._data = mj.MjData(self._model)
-        site = mj.mjtObj.mjOBJ_SITE
-        self._site_id = mj.mj_name2id(self._model, site, DEFAULT_FRAME)
-        self._qpos_ids = np.array([self._model.joint(name).qposadr.item() for name in _JOINT_NAMES])
-        self._dof_ids = np.array([self._model.joint(name).dofadr.item() for name in _JOINT_NAMES])
-        ranges = np.array([self._model.joint(name).range for name in _JOINT_NAMES])
-        self._lower, self._upper = ranges[:, 0], ranges[:, 1]
-
-    def fk(self, q: np.ndarray) -> geom.Transform3D:
-        self._data.qpos[self._qpos_ids] = q
-        mj.mj_kinematics(self._model, self._data)
-        quat = np.empty(4)
-        mj.mju_mat2Quat(quat, self._data.site_xmat[self._site_id].copy())
-        return geom.Transform3D(self._data.site_xpos[self._site_id].copy(), geom.Rotation.from_quat(quat))
-
-    def ik(self, target: geom.Transform3D, current_q: np.ndarray) -> np.ndarray | None:
-        """Multi-start LM IK: the live posture first, then the reach postures toward the target's azimuth.
-        Solutions are wrapped and clamped into joint range, then FK-verified before acceptance."""
-        for start in (current_q, *_reach_postures(*target.translation[:2])):
-            self._data.qpos[:] = 0.0
-            self._data.qpos[self._qpos_ids] = start
-            qpos, _, success = qpos_from_site_pose(
-                self._model,
-                self._data,
-                self._site_id,
-                self._dof_ids,
-                target.translation,
-                target.rotation.as_quat,
-                rot_weight=0.5,
-            )
-            if not success:
-                continue
-            q = qpos[self._qpos_ids].copy()
-            # A revolute joint at q ± 2π is the same pose; wrap out-of-range entries back in when they fit.
-            q = np.where(q > self._upper, q - 2 * np.pi, q)
-            q = np.where(q < self._lower, q + 2 * np.pi, q)
-            q = np.clip(q, self._lower, self._upper)
-            reached = self.fk(q)
-            rot_err = (reached.rotation.inv * target.rotation).angle
-            rot_err = min(rot_err, 2 * np.pi - rot_err)
-            if np.linalg.norm(reached.translation - target.translation) < _IK_POS_TOL and rot_err < _IK_ROT_TOL:
-                return q
-        return None
-
-
 class _Chain(DriverRun[command.CommandType]):
     """The chain the driver drives: the vendor handle, and the state and moves that go with it."""
 
@@ -181,7 +123,7 @@ def __init__(
         self.grip_out = grip_out
         self.state = YamState()
         self._base_pose = base_pose
-        self._kin = _Kinematics()
+        self._kin = MjcfKinematics(_MJCF_PATH, DEFAULT_FRAME, _JOINT_NAMES, _reach_postures)
 
     def __enter__(self) -> '_Chain':
         return self
@@ -461,7 +403,7 @@ def pump(seconds: float):
             pump(0.1)  # the opening move ramps the chain to the park pose over a couple of seconds
         assert state.value.status == RobotStatus.AVAILABLE, state.value.status
 
-        kin = _Kinematics()
+        kin = MjcfKinematics(_MJCF_PATH, DEFAULT_FRAME, _JOINT_NAMES, _reach_postures)
 
         if fake is not None:
             # State round-trip: the parked chain comes back through the driver's FK.