scripts/arm_kinematics.py

"""Robot-only forward/inverse kinematics for the SO-101 block-sorting scene.

The solver uses the robot description and joint encoders; it never reads object
positions. Coordinates are metres and joint targets are radians. The requested
grasp is vertical, with its closing axis at ``yaw`` around the world Z axis.
For this official SO-101 calibration, vertical grasps have approximately
``wrist_roll = shoulder_pan + 0.04868 + yaw``.
"""

from pathlib import Path

import mujoco
import numpy as np


class SO101Kinematics:
    joint_names = (
        "shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex",
        "wrist_roll", "gripper",
    )

    def __init__(self, scene_path):
        self.model = mujoco.MjModel.from_xml_path(str(Path(scene_path).resolve()))
        self.data = mujoco.MjData(self.model)
        self.joint_ids = np.array([self.model.joint(n).id for n in self.joint_names])
        self.qpos_indices = self.model.jnt_qposadr[self.joint_ids].copy()
        self.limits = self.model.jnt_range[self.joint_ids].copy()
        self.body_id = self.model.body("gripper").id
        self.grasp_site_id = self.model.site("grasp_center").id
        self.marker_geom_id = self.model.geom("vision_marker").id
        self.last_error = None

    def _forward(self, q):
        q = np.asarray(q, dtype=float)
        if q.shape != (6,) or not np.isfinite(q).all():
            raise ValueError("Provide six finite joint angles in radians.")
        self.data.qpos[self.qpos_indices] = q
        mujoco.mj_kinematics(self.model, self.data)
        return (
            self.data.site_xpos[self.grasp_site_id].copy(),
            self.data.xmat[self.body_id].reshape(3, 3).copy(),
        )

    def forward_grasp(self, q):
        """Return the grasp center XYZ for six robot joint angles."""
        return self._forward(q)[0]

    def forward_marker(self, q):
        """Return the visible magenta calibration marker's center XYZ."""
        self._forward(q)
        return self.data.geom_xpos[self.marker_geom_id].copy()

    def solve(self, xyz, seed=None, grip=0.6, yaw=0.0):
        """Return six motor targets for a vertical grasp.

        ``seed`` can contain five arm angles or all six angles. Pass the previous
        targets while interpolating a Cartesian path. ``grip`` is the requested
        gripper hinge angle: 0.6 opens it and -0.15 closes it under the configured
        actuator force limit. No joint or object state in the running simulator
        is modified by this method. Unreachable targets raise ValueError.
        """
        target = np.asarray(xyz, dtype=float)
        if target.shape != (3,) or not np.isfinite(target).all():
            raise ValueError("xyz must contain three finite coordinates in metres.")
        if not np.isfinite([grip, yaw]).all():
            raise ValueError("grip and yaw must be finite angles in radians.")
        if not self.limits[5, 0] <= grip <= self.limits[5, 1]:
            raise ValueError("Gripper target is outside the physical joint range.")

        desired_x = np.array([np.cos(yaw), np.sin(yaw), 0.0])
        desired_z = np.array([0.0, 0.0, 1.0])
        weight = 0.10

        def residual(q5):
            p, rotation = self._forward(np.r_[q5, grip])
            error = np.r_[target - p,
                          weight * (desired_x - rotation[:, 0]),
                          weight * (desired_z - rotation[:, 2])]
            return error, p, rotation

        pan = -np.arctan2(target[1], target[0] - 0.0388353)
        roll = pan + 0.04868 + yaw
        starts = []
        if seed is not None:
            seed = np.asarray(seed, dtype=float)
            if seed.shape not in ((5,), (6,)) or not np.isfinite(seed).all():
                raise ValueError("seed must contain five or six finite joint angles.")
            starts.append(seed[:5].copy())
        starts.extend([
            np.array([pan, 0.0, 0.3, 1.27, roll]),
            np.array([pan, 0.6, -0.4, 1.37, roll]),
            np.array([pan, -0.5, 0.7, 1.37, roll]),
        ])

        best = None
        for start in starts:
            q = np.clip(start, self.limits[:5, 0], self.limits[:5, 1])
            for _ in range(160):
                error, p, rotation = residual(q)
                score = float(error @ error)
                if best is None or score < best[0]:
                    best = (score, q.copy(), p.copy(), rotation.copy())
                if np.linalg.norm(error[:3]) < 2e-6 and np.linalg.norm(error[3:]) < 2e-6:
                    self.last_error = {
                        "position_m": float(np.linalg.norm(error[:3])),
                        "orientation": float(np.linalg.norm(error[3:]) / weight),
                    }
                    return np.r_[q, grip]

                jacobian = np.empty((9, 5))
                for j in range(5):
                    trial = q.copy()
                    trial[j] += 1e-5
                    shifted_error, _, _ = residual(trial)
                    jacobian[:, j] = (error - shifted_error) / 1e-5
                step = np.linalg.lstsq(
                    np.vstack([jacobian, np.eye(5) * 0.001]),
                    np.r_[error, np.zeros(5)], rcond=None,
                )[0]
                step *= min(1.0, 0.18 / max(np.linalg.norm(step), 1e-12))
                accepted = False
                for fraction in (1.0, 0.5, 0.25, 0.10, 0.025):
                    candidate = np.clip(q + fraction * step,
                                        self.limits[:5, 0], self.limits[:5, 1])
                    candidate_error, _, _ = residual(candidate)
                    if candidate_error @ candidate_error < score - 1e-15:
                        q = candidate
                        accepted = True
                        break
                if not accepted:
                    break

        _, q, p, rotation = best
        position_error = float(np.linalg.norm(target - p))
        orientation_error = float(np.linalg.norm(np.r_[
            desired_x - rotation[:, 0], desired_z - rotation[:, 2]]))
        self.last_error = {"position_m": position_error,
                           "orientation": orientation_error}
        if position_error <= 0.0008 and orientation_error <= 0.006:
            return np.r_[q, grip]
        raise ValueError(
            f"Unreachable vertical grasp {target.tolist()}: "
            f"position error {position_error * 1000:.2f} mm, "
            f"orientation error {orientation_error:.4f}. "
            "Reduce transfer height or bring the target closer to the arm."
        )