Support newer Mink and optimize finite-difference IK

.github/workflows/test.yaml

@@ -26,7 +26,10 @@ concurrency:
   cancel-in-progress: true
 jobs:
   pytest:
-    name: pytest
+    name: pytest (${{ matrix.mink }})
+    strategy:
+      matrix:
+        mink: ["mink==1.2.0", "mink"]
     runs-on: ubuntu-latest
     timeout-minutes: 10
     permissions:
@@ -37,6 +40,7 @@ jobs:
       - name: Sync environment
         run: |
           uv sync
+          uv pip install "${{ matrix.mink }}"
       - name: Run pytest
         run: |
-          uv run pytest
+          uv run --no-sync pytest

pyproject.toml

@@ -26,9 +26,9 @@ readme = "README.md"
 requires-python = ">=3.10"
 
 dependencies = [
-    "daqp==0.7.2",
-    "mink>=1.1.0",
-    "mujoco>=3.6.0",
+    "daqp>=0.8.2",
+    "mink>=1.2.0",
+    "mujoco>=3.8.1",
     "numpy",
     "openarm-mujoco>=2.0.1",
     "pyyaml",

src/openarm_control/geometry/jacobian.py

@@ -22,6 +22,32 @@
 import numpy.typing as npt
 
 
+class CachedConfiguration(mink.Configuration):
+    """Cache read-only geometric quantities until the next configuration update."""
+
+    def __init__(self, model: mujoco.MjModel, q: np.ndarray | None = None) -> None:
+        """Initialize caches before Mink's constructor calls update."""
+        self._frame_jacobians: dict[tuple[str | int, str], np.ndarray] = {}
+        self._arm_svds: dict[mink.Task, tuple[np.ndarray, np.ndarray]] = {}
+        super().__init__(model, q=q)
+
+    def update(self, q: np.ndarray | None = None) -> None:
+        """Invalidate geometry whenever forward kinematics is updated."""
+        self._frame_jacobians.clear()
+        self._arm_svds.clear()
+        super().update(q=q)
+
+    def get_frame_jacobian(self, frame_name: str | int, frame_type: str) -> np.ndarray:
+        """Return a read-only body Jacobian shared by tasks at this configuration."""
+        key = (frame_name, frame_type)
+        jacobian = self._frame_jacobians.get(key)
+        if jacobian is None:
+            jacobian = super().get_frame_jacobian(frame_name, frame_type)
+            jacobian.setflags(write=False)
+            self._frame_jacobians[key] = jacobian
+        return jacobian
+
+
 def relative_root_is_independent_of_dofs(
     frame_task: mink.Task,
     model: mujoco.MjModel,
@@ -34,25 +60,39 @@ def relative_root_is_independent_of_dofs(
     if not isinstance(native_task, mink.RelativeFrameTask):
         raise TypeError("Expected a Mink frame task or a frame-task wrapper.")
 
-    root_id = mujoco.mj_name2id(
-        model,
-        {
-            "body": mujoco.mjtObj.mjOBJ_BODY,
-            "site": mujoco.mjtObj.mjOBJ_SITE,
-            "geom": mujoco.mjtObj.mjOBJ_GEOM,
-        }[native_task.root_type],
-        native_task.root_name,
+    return frame_is_independent_of_dofs(
+        model, native_task.root_name, native_task.root_type, dof_indices
     )
-    if root_id < 0:
-        raise ValueError(
-            f"Unknown {native_task.root_type} frame {native_task.root_name!r}."
+
+
+def frame_is_independent_of_dofs(
+    model: mujoco.MjModel,
+    frame_name: str | int,
+    frame_type: str,
+    dof_indices: npt.ArrayLike,
+) -> bool:
+    """Return whether none of the selected DoFs is an ancestor of the frame."""
+    frame_id = (
+        mujoco.mj_name2id(
+            model,
+            {
+                "body": mujoco.mjtObj.mjOBJ_BODY,
+                "site": mujoco.mjtObj.mjOBJ_SITE,
+                "geom": mujoco.mjtObj.mjOBJ_GEOM,
+            }[frame_type],
+            frame_name,
         )
-    if native_task.root_type == "body":
-        body_id = root_id
-    elif native_task.root_type == "site":
-        body_id = int(model.site_bodyid[root_id])
+        if isinstance(frame_name, str)
+        else int(frame_name)
+    )
+    if frame_id < 0:
+        raise ValueError(f"Unknown {frame_type} frame {frame_name!r}.")
+    if frame_type == "body":
+        body_id = frame_id
+    elif frame_type == "site":
+        body_id = int(model.site_bodyid[frame_id])
     else:
-        body_id = int(model.geom_bodyid[root_id])
+        body_id = int(model.geom_bodyid[frame_id])
 
     selected_dofs = set(np.asarray(dof_indices, dtype=int).tolist())
     while body_id > 0:

src/openarm_control/kinematics.py

@@ -40,6 +40,7 @@
 import numpy as np
 
 from openarm_control.config import ArmSetup, frame_name
+from openarm_control.geometry.jacobian import CachedConfiguration
 from openarm_control.geometry.poses import pose_to_se3
 from openarm_control.ik_params import (
     IKParams,
@@ -155,7 +156,7 @@ def __init__(self, setup: ArmSetup, params: IKParams) -> None:
         self._substep_dt = params.dt / params.max_iters
         self._max_iters = params.max_iters
 
-        self._config = mink.Configuration(setup.model)
+        self._config = CachedConfiguration(setup.model)
         self._config.update(q=setup.data.qpos.copy())
         mid_qpos = self._config.data.qpos.copy()
 
@@ -271,6 +272,13 @@ def __init__(self, setup: ArmSetup, params: IKParams) -> None:
                 self._singularity_limits[side] = limit
                 self._limits.append(limit)
 
+        self._paired_singularity_limits: (
+            tuple[SingularityApproachLimit, SingularityApproachLimit] | None
+        ) = None
+        limits = tuple(self._singularity_limits.values())
+        if len(limits) == 2 and limits[0].can_share_perturbations(limits[1]):
+            self._paired_singularity_limits = (limits[0], limits[1])
+
         self._posture_task = mink.PostureTask(setup.model, cost=params.posture_cost)
         self._posture_task.set_target(mid_qpos)
 
@@ -291,6 +299,10 @@ def __init__(self, setup: ArmSetup, params: IKParams) -> None:
                     characteristic_length=params.jacobian_characteristic_length,
                 )
 
+        nullspace_tasks = tuple(self._nullspace_tasks.values())
+        if len(nullspace_tasks) == 2:
+            nullspace_tasks[0]._batch_svd_with(nullspace_tasks[1])
+
         self._kinetic_energy_task: KineticEnergyRegularizationTask | None = None
         if params.kinetic_energy_cost > 0.0:
             self._kinetic_energy_task = KineticEnergyRegularizationTask(
@@ -357,8 +369,22 @@ def solve(self) -> np.ndarray | None:
         constraints = [self._freeze_task] if self._freeze_task else []
 
         q_before = self._config.data.qpos.copy()
-        for limit in self._singularity_limits.values():
-            limit.prepare(self._config)
+        gradients: tuple[np.ndarray | None, ...] = (None,) * len(
+            self._singularity_limits
+        )
+        if self._paired_singularity_limits is not None:
+            first, second = self._paired_singularity_limits
+            gradients = first._paired_gradients(self._config, second)
+        for (side, limit), gradient in zip(
+            self._singularity_limits.items(), gradients, strict=True
+        ):
+            nullspace = self._nullspace_tasks.get(side)
+            command_ratio = (
+                nullspace.compute_singularity_ratio(self._config)
+                if nullspace is not None
+                else None
+            )
+            limit.prepare(self._config, command_ratio=command_ratio, gradient=gradient)
         self._pending = set(self._sides)
 
         for _ in range(self._max_iters):

src/openarm_control/qp/arm_joint_limit.py

@@ -24,8 +24,8 @@
 import numpy.typing as npt
 
 
-class ArmConfigurationLimit(mink.ConfigurationLimit):
-    """Apply Mink's configuration limits only to selected scalar arm joints."""
+class ArmConfigurationLimit(mink.Limit):
+    """Apply configuration limits only to selected scalar arm joints."""
 
     def __init__(
         self,
@@ -35,8 +35,14 @@ def __init__(
         gain: float,
     ) -> None:
         """Build a configuration limit projected onto selected arm DoFs."""
-        super().__init__(model, gain=gain)
+        if not 0.0 < gain <= 1.0:
+            raise mink.exceptions.LimitDefinitionError(
+                f"{self.__class__.__name__} gain must be in the range (0, 1]"
+            )
+        self.model = model
+        self.gain = gain
         selected_qpos = {int(index) for index in qpos_indices}
+        active_joints: list[int] = []
         active_dofs: list[int] = []
         for joint_id in range(model.njnt):
             if (
@@ -49,13 +55,34 @@ def __init__(
                 int(mujoco.mjtJoint.mjJNT_SLIDE),
             ):
                 raise ValueError("ArmConfigurationLimit only supports scalar joints.")
+            active_joints.append(joint_id)
             active_dofs.append(int(model.jnt_dofadr[joint_id]))
         self.indices = _readonly(active_dofs, dtype=int)
+        self.qpos_indices = _readonly(model.jnt_qposadr[active_joints], dtype=int)
+        self.lower = _readonly(model.jnt_range[active_joints, 0])
+        self.upper = _readonly(model.jnt_range[active_joints, 1])
         self.projection_matrix = (
             np.eye(model.nv, dtype=np.float64)[self.indices]
             if self.indices.size
             else None
         )
+        self._G = (
+            np.vstack([self.projection_matrix, -self.projection_matrix])
+            if self.projection_matrix is not None
+            else None
+        )
+
+    def compute_qp_inequalities(
+        self, configuration: mink.Configuration, dt: float
+    ) -> mink.Constraint:
+        """Bound selected joint displacements without constraining scene joints."""
+        if self.projection_matrix is None:
+            return mink.Constraint()
+        q = configuration.data.qpos[self.qpos_indices]
+        return mink.Constraint(
+            G=self._G,
+            h=self.gain * np.hstack([self.upper - q, q - self.lower]),
+        )
 
 
 class ArmJointLimit(mink.Limit):
@@ -143,6 +170,11 @@ def __init__(
             if self.dof_indices.size
             else None
         )
+        self._G = (
+            np.vstack([self.projection_matrix, -self.projection_matrix])
+            if self.projection_matrix is not None
+            else None
+        )
         self._measured_qpos: np.ndarray | None = None
 
     def update_measured_state(self, qpos: npt.ArrayLike) -> None:
@@ -169,7 +201,7 @@ def compute_qp_inequalities(
         if dt <= 0.0:
             raise ValueError("dt must be positive.")
 
-        command_q = configuration.q[self.qpos_indices]
+        command_q = configuration.data.qpos[self.qpos_indices]
         max_step = dt * self.max_velocity
         position_lower = self.position_gain * (self.lower - command_q)
         position_upper = self.position_gain * (self.upper - command_q)
@@ -205,7 +237,7 @@ def compute_qp_inequalities(
             upper_step = np.minimum(upper_step, dt * upper_velocity)
 
         return mink.Constraint(
-            G=np.vstack([self.projection_matrix, -self.projection_matrix]),
+            G=self._G,
             h=np.hstack([upper_step, -lower_step]),
         )
 

src/openarm_control/qp/bounded_frame_task.py

@@ -143,16 +143,24 @@ def compute_limited_error(self, configuration: mink.Configuration) -> np.ndarray
         """Return one substep of the total position and orientation budgets."""
         return self._clip_error(self.compute_full_error(configuration))
 
-    def _clip_error(self, full_error: np.ndarray) -> np.ndarray:
+    def _clip_error(
+        self,
+        full_error: np.ndarray,
+        *,
+        position: bool = True,
+        orientation: bool = True,
+    ) -> np.ndarray:
         """Limit a previously computed frame error."""
         error = full_error.copy()
         for part, total_limit in (
-            (slice(0, 3), self.position_error_limit),
-            (slice(3, 6), self.orientation_error_limit),
+            (slice(0, 3), self.position_error_limit if position else 0.0),
+            (slice(3, 6), self.orientation_error_limit if orientation else 0.0),
         ):
             limit = total_limit / self._substeps
+            if limit <= 0.0:
+                continue
             norm = float(np.linalg.norm(error[part]))
-            if limit > 0.0 and norm > limit:
+            if norm > limit:
                 error[part] *= limit / norm
         return error
 
@@ -167,39 +175,45 @@ def compute_jacobian(self, configuration: mink.Configuration) -> np.ndarray:
     def compute_qp_residual(
         self,
         configuration: mink.Configuration,
-    ) -> tuple[np.ndarray, np.ndarray, float] | None:
-        """Opt out of Mink's fused residual path so bounding is not bypassed.
+    ) -> tuple[np.ndarray, np.ndarray, float]:
+        """Return the bounded request for Mink's fused objective assembly."""
+        terms = self._bounded_terms(configuration)
+        if terms is None:
+            return self.frame_task.compute_qp_residual(configuration)
+        return self._weighted_residual(*terms)
+
+    def _bounded_terms(
+        self, configuration: mink.Configuration
+    ) -> tuple[np.ndarray, np.ndarray] | None:
+        bound_position = self.position_error_limit > 0.0 and self.limit_activation > 0.0
+        bound_orientation = self.orientation_error_limit > 0.0
+        if not bound_position and not bound_orientation:
+            return None
 
-        Mink 1.2's solver prefers a task's residual and only falls back to
-        :meth:`compute_qp_objective` when it is ``None``. The inherited
-        ``mink.Task`` residual is built from :meth:`compute_error`, which
-        deliberately reports the *full* error, so accepting it would silently
-        drop the bounding this class exists to apply.
-        """
-        return None
+        full_error, jacobian = self.frame_task._error_and_jacobian(configuration)
+        error = self._clip_error(
+            full_error, position=bound_position, orientation=bound_orientation
+        )
+        if bound_position:
+            error[:3] = full_error[:3] + self.limit_activation * (
+                error[:3] - full_error[:3]
+            )
+        return error, jacobian
 
     def compute_qp_objective(
         self,
         configuration: mink.Configuration,
     ) -> mink.Objective:
         """Modulate position by its schedule and always bound orientation."""
-        bound_position = self.position_error_limit > 0.0 and self.limit_activation > 0.0
-        bound_orientation = self.orientation_error_limit > 0.0
-        if not bound_position and not bound_orientation:
+        terms = self._bounded_terms(configuration)
+        if terms is None:
             return self.frame_task.compute_qp_objective(configuration)
-
-        full_error = self.compute_full_error(configuration)
-        limited_error = self._clip_error(full_error)
-        error = full_error.copy()
-        if bound_position:
-            error[:3] += self.limit_activation * (limited_error[:3] - full_error[:3])
-        if bound_orientation:
-            error[3:] = limited_error[3:]
+        error, jacobian = terms
         nv = configuration.model.nv
         if self._identity is None or self._identity.shape != (nv, nv):
             self._identity = np.eye(nv, dtype=np.float64)
         return self._assemble_qp(
             error,
-            self.compute_jacobian(configuration),
+            jacobian,
             self._identity,
         )