diff --git a/positronic/drivers/roboarm/mjcf_kinematics.py b/positronic/drivers/roboarm/mjcf_kinematics.py new file mode 100644 index 000000000..8aef4018c --- /dev/null +++ b/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 diff --git a/positronic/drivers/roboarm/tests/test_kinematics.py b/positronic/drivers/roboarm/tests/test_kinematics.py new file mode 100644 index 000000000..4249ca00c --- /dev/null +++ b/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 diff --git a/positronic/drivers/roboarm/tests/yam_kinematics_goldens.npz b/positronic/drivers/roboarm/tests/yam_kinematics_goldens.npz new file mode 100644 index 000000000..f2c516728 Binary files /dev/null and b/positronic/drivers/roboarm/tests/yam_kinematics_goldens.npz differ diff --git a/positronic/drivers/roboarm/yam.py b/positronic/drivers/roboarm/yam.py index 6611068ee..94c57a18d 100644 --- a/positronic/drivers/roboarm/yam.py +++ b/positronic/drivers/roboarm/yam.py @@ -18,7 +18,6 @@ from collections.abc import Callable, Generator, Iterator from typing import Any -import mujoco as mj import numpy as np import pimm @@ -26,10 +25,9 @@ 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.