Fix the MuJoCo IK fallback in the UR5e examples; add restarts to unseeded solves

CHANGELOG.md

@@ -23,6 +23,17 @@ follows [Semantic Versioning](https://semver.org/).
 
 ### Fixed
 
+- The UR5e examples' MuJoCo differential-IK fallback crashed on the first
+  IK update because the collision checker was passed in the `joint_limits`
+  positional slot. Both examples now select the backend through
+  `build_ik_solver(..., backend)` with a `--ik {auto,ssik,mujoco}` flag, and
+  the fallback is exercised by tests even when SSIK is installed (#65).
+- `MuJoCoIKSolver.solve` without a seed tries the current state and then
+  `restarts` (default 3) random initial configurations within limits and
+  returns every distinct converged solution. Previously it depended on
+  whatever the shared MuJoCo state was last left in and converged on only
+  about half of reachable poses, which made TSR goal sampling through the
+  fallback unreliable.
 - `SSIKSolver` copies `T_base` and `T_ee` at construction and stores them
   read-only, so mutating the caller's array can no longer desynchronize the
   stored transform from its cached inverse (#66).

examples/tsr_union_demo.py

@@ -50,6 +50,25 @@
     SSIK_AVAILABLE = False
 
 
+def build_ik_solver(model, data, joint_names, collision, menagerie_path: Path, backend: str = "auto"):
+    """Choose the IK backend: SSIK (enumerative analytical) or MuJoCo differential IK.
+
+    ``backend`` is "auto" (SSIK if installed, else MuJoCo), "ssik", or "mujoco".
+    Returns ``(solver, name)`` where ``name`` is "ssik" or "mujoco".
+    """
+    if backend not in ("auto", "ssik", "mujoco"):
+        raise ValueError(f"backend must be 'auto', 'ssik', or 'mujoco', got {backend!r}")
+    if backend == "ssik" and not SSIK_AVAILABLE:
+        raise RuntimeError('SSIK backend requested but ssik is not installed: pip install "pycbirrt[ssik]"')
+    if backend in ("auto", "ssik") and SSIK_AVAILABLE:
+        arm = ssik.Manipulator.from_mjcf(
+            menagerie_path / "universal_robots_ur5e" / "ur5e.xml", base="world", ee="wrist_3_link"
+        )
+        return SSIKSolver(arm, T_ee=site_offset_in_body(model, "attachment_site")), "ssik"
+    # Joint limits come from the MuJoCo model; the collision checker must be passed by keyword (#65).
+    return MuJoCoIKSolver(model, data, "attachment_site", joint_names, collision_checker=collision), "mujoco"
+
+
 def get_menagerie_path() -> Path:
     """Get path to MuJoCo Menagerie."""
     path = os.environ.get("MUJOCO_MENAGERIE_PATH")
@@ -411,6 +430,9 @@ def main():
     parser.add_argument("--grasps", type=int, default=6, help="Number of grasp cycles to show")
     parser.add_argument("--seed", type=int, default=None, help="Random seed (None = random each run)")
     parser.add_argument("--interactive", "-i", action="store_true", help="Use interactive viewer instead of video")
+    parser.add_argument(
+        "--ik", choices=["auto", "ssik", "mujoco"], default="auto", help="IK backend (default: SSIK if installed)"
+    )
     args = parser.parse_args()
 
     menagerie = get_menagerie_path()
@@ -436,25 +458,14 @@ def main():
     robot = MuJoCoRobotModel(model, data, "attachment_site", joints)
     collision = MuJoCoCollisionChecker(model, data, joints)
 
-    if SSIK_AVAILABLE:
+    ik_solver, ik_name = build_ik_solver(model, data, joints, collision, menagerie, args.ik)
+    if ik_name == "ssik":
         print("Using SSIK analytical IK solver")
-        arm = ssik.Manipulator.from_mjcf(
-            menagerie / "universal_robots_ur5e" / "ur5e.xml", base="world", ee="wrist_3_link"
-        )
-        ik_solver = SSIKSolver(arm, T_ee=site_offset_in_body(model, "attachment_site"))
-        config = CBiRRTConfig(
-            timeout=30.0,
-            goal_bias=0.15,
-        )
+        config = CBiRRTConfig(timeout=30.0, goal_bias=0.15)
     else:
         print("Using MuJoCo differential IK solver")
-        ik_solver = MuJoCoIKSolver(model, data, "attachment_site", joints, collision)
         # Differential solver may need more pose samples
-        config = CBiRRTConfig(
-            timeout=30.0,
-            goal_bias=0.15,
-            tsr_samples=100,
-        )
+        config = CBiRRTConfig(timeout=30.0, goal_bias=0.15, tsr_samples=100)
 
     planner = CBiRRT(robot, ik_solver, collision, config)
 

examples/ur5e_mujoco.py

@@ -72,6 +72,28 @@ def get_menagerie_path() -> Path:
 # =============================================================================
 
 
+def build_ik_solver(model, data, joint_names, collision_checker, menagerie_path: Path, backend: str = "auto"):
+    """Choose the IK backend: SSIK (enumerative analytical) or MuJoCo differential IK.
+
+    ``backend`` is "auto" (SSIK if installed, else MuJoCo), "ssik", or "mujoco".
+    SSIK is built from the same MJCF as the scene, with the world as base and the
+    attachment site's offset as T_ee, so its frames match MuJoCoRobotModel exactly.
+    Returns ``(solver, name)`` where ``name`` is "ssik" or "mujoco".
+    """
+    if backend not in ("auto", "ssik", "mujoco"):
+        raise ValueError(f"backend must be 'auto', 'ssik', or 'mujoco', got {backend!r}")
+    if backend == "ssik" and not SSIK_AVAILABLE:
+        raise RuntimeError('SSIK backend requested but ssik is not installed: pip install "pycbirrt[ssik]"')
+    if backend in ("auto", "ssik") and SSIK_AVAILABLE:
+        arm = ssik.Manipulator.from_mjcf(
+            menagerie_path / "universal_robots_ur5e" / "ur5e.xml", base="world", ee="wrist_3_link"
+        )
+        return SSIKSolver(arm, T_ee=site_offset_in_body(model, "attachment_site")), "ssik"
+    # Joint limits come from the MuJoCo model; the collision checker must be passed by keyword (#65).
+    solver = MuJoCoIKSolver(model, data, "attachment_site", joint_names, collision_checker=collision_checker)
+    return solver, "mujoco"
+
+
 def create_grasp_tsr(target_pos: np.ndarray) -> TSR:
     """Create a TSR for top-down grasp of a cylinder.
 
@@ -315,6 +337,9 @@ def main():
     parser.add_argument("--render", type=str, help="Render to video file (e.g., output.mp4)")
     parser.add_argument("--no-viz", action="store_true", help="Skip visualization")
     parser.add_argument("--seed", type=int, default=42, help="Random seed (default: 42)")
+    parser.add_argument(
+        "--ik", choices=["auto", "ssik", "mujoco"], default="auto", help="IK backend (default: SSIK if installed)"
+    )
     args = parser.parse_args()
 
     menagerie_path = get_menagerie_path()
@@ -338,18 +363,12 @@ def main():
     robot = MuJoCoRobotModel(model, data, "attachment_site", ur5e_joints)
     collision_checker = MuJoCoCollisionChecker(model, data, ur5e_joints)
 
-    # IK solver: prefer SSIK (enumerative analytical), fall back to MuJoCo (differential).
-    # SSIK is built from the same MJCF as the scene, with the world as base and the
-    # attachment site's offset as T_ee, so its frames match MuJoCoRobotModel exactly.
-    if SSIK_AVAILABLE:
+    # IK solver: prefer SSIK (enumerative analytical), fall back to MuJoCo (differential)
+    ik_solver, ik_name = build_ik_solver(model, data, ur5e_joints, collision_checker, menagerie_path, args.ik)
+    if ik_name == "ssik":
         print("Using SSIK (analytical) IK solver")
-        arm = ssik.Manipulator.from_mjcf(
-            menagerie_path / "universal_robots_ur5e" / "ur5e.xml", base="world", ee="wrist_3_link"
-        )
-        ik_solver = SSIKSolver(arm, T_ee=site_offset_in_body(model, "attachment_site"))
     else:
         print('Using MuJoCo (differential) IK solver (pip install "pycbirrt[ssik]" for faster planning)')
-        ik_solver = MuJoCoIKSolver(model, data, "attachment_site", ur5e_joints, collision_checker)
 
     config = CBiRRTConfig(
         max_iterations=5000,

src/pycbirrt/backends/mujoco.py

@@ -181,6 +181,8 @@ def __init__(
         damping: float = 0.1,
         max_iterations: int = 200,
         tolerance: float = 1e-3,
+        restarts: int = 3,
+        seed: int | None = None,
     ):
         """Initialize MuJoCo IK solver.
 
@@ -193,6 +195,13 @@ def __init__(
             collision_checker: Optional collision checker for validation
             damping: Damping factor for damped least squares
             max_iterations: Maximum iterations for IK convergence
+            restarts: When ``solve`` is called without ``q_init``, number of
+                random initial configurations tried after the current state.
+                Differential IK from a single start converges on only part of
+                the workspace and depends on whatever the shared MjData was
+                last left in; restarts make unseeded solves (e.g. TSR goal
+                sampling) reliable and surface several distinct solutions.
+            seed: Seed for the restart generator.
             tolerance: Position/orientation error tolerance for convergence
         """
         self.model = model
@@ -201,6 +210,10 @@ def __init__(
         self.max_iterations = max_iterations
         self.tolerance = tolerance
         self.joint_limits = joint_limits
+        if restarts < 0:
+            raise ValueError("restarts must be nonnegative")
+        self.restarts = restarts
+        self._rng = np.random.default_rng(seed)
         self.collision_checker = collision_checker
 
         # Get site ID
@@ -307,14 +320,34 @@ def solve(self, pose: np.ndarray, q_init: np.ndarray | None = None) -> list[np.n
 
         Args:
             pose: 4x4 homogeneous transform of desired end-effector pose
-            q_init: Initial configuration (if None, uses current model state)
+            q_init: Initial configuration. If given, one solve runs from it
+                (the seed-nearest solution, as projection wants). If None, one
+                solve runs from the current model state and then from
+                ``restarts`` random configurations within the joint limits,
+                and every distinct converged solution is returned.
 
         Returns:
-            List containing one solution if found, empty list otherwise
+            Converged solutions (possibly empty). One when seeded.
         """
-        # Set initial configuration
         if q_init is not None:
-            self._set_config(q_init)
+            return self._solve_from(pose, np.asarray(q_init, dtype=float))
+
+        solutions: list[np.ndarray] = []
+        inits = [self._get_config()]
+        if self.restarts and self.joint_limits is not None:
+            lower, upper = self.joint_limits
+            # Start within one turn: a differential solver has no use for far windings
+            lo, hi = np.maximum(lower, -np.pi), np.minimum(upper, np.pi)
+            inits += [self._rng.uniform(lo, hi) for _ in range(self.restarts)]
+        for q0 in inits:
+            for q in self._solve_from(pose, q0):
+                if not any(np.linalg.norm(q - s) < 1e-3 for s in solutions):
+                    solutions.append(q)
+        return solutions
+
+    def _solve_from(self, pose: np.ndarray, q_init: np.ndarray) -> list[np.ndarray]:
+        """Damped least squares from one initial configuration; one solution or none."""
+        self._set_config(q_init)
 
         for _ in range(self.max_iterations):
             # Forward kinematics

tests/test_examples_mujoco.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2025 Siddhartha Srinivasa
+
+"""The UR5e examples' MuJoCo differential-IK fallback works when forced (#65).
+
+Needs mujoco and MUJOCO_MENAGERIE_PATH, not ssik: the point is to exercise the
+non-SSIK path even in an environment where SSIK is installed.
+"""
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+mujoco = pytest.importorskip("mujoco")
+MENAGERIE = os.environ.get("MUJOCO_MENAGERIE_PATH")
+if not MENAGERIE or not (Path(MENAGERIE) / "universal_robots_ur5e").exists():
+    pytest.skip("MUJOCO_MENAGERIE_PATH not set to a mujoco_menagerie clone", allow_module_level=True)
+
+EXAMPLES = Path(__file__).resolve().parent.parent / "examples"
+sys.path.insert(0, str(EXAMPLES))
+
+import tsr_union_demo  # noqa: E402
+import ur5e_mujoco  # noqa: E402
+
+from pycbirrt import CBiRRT, CBiRRTConfig  # noqa: E402
+from pycbirrt.backends.mujoco import MuJoCoCollisionChecker, MuJoCoIKSolver, MuJoCoRobotModel  # noqa: E402
+
+JOINTS = [
+    "shoulder_pan_joint",
+    "shoulder_lift_joint",
+    "elbow_joint",
+    "wrist_1_joint",
+    "wrist_2_joint",
+    "wrist_3_joint",
+]
+HOME = np.array([0, -np.pi / 2, np.pi / 2, -np.pi / 2, -np.pi / 2, 0])
+
+
+class TestUr5eMujocoExample:
+    def test_forced_mujoco_fallback_solves_a_nontrivial_pose(self):
+        model = ur5e_mujoco.create_scene(Path(MENAGERIE))
+        data = mujoco.MjData(model)
+        robot = MuJoCoRobotModel(model, data, "attachment_site", JOINTS)
+        collision = MuJoCoCollisionChecker(model, data, JOINTS)
+        ik, name = ur5e_mujoco.build_ik_solver(model, data, JOINTS, collision, Path(MENAGERIE), backend="mujoco")
+        assert name == "mujoco" and isinstance(ik, MuJoCoIKSolver)
+        assert ik.collision_checker is collision
+        target = robot.forward_kinematics(HOME + np.array([0.3, 0.2, -0.2, 0.1, 0.1, 0.4]))
+        sols = ik.solve(target, q_init=HOME)  # requires iterative updates from the seed
+        assert sols
+        assert np.linalg.norm(robot.forward_kinematics(sols[0])[:3, 3] - target[:3, 3]) < 5e-3
+
+    def test_auto_prefers_ssik_when_installed(self):
+        pytest.importorskip("ssik")
+        from pycbirrt.backends.ssik import SSIKSolver
+
+        model = ur5e_mujoco.create_scene(Path(MENAGERIE))
+        data = mujoco.MjData(model)
+        collision = MuJoCoCollisionChecker(model, data, JOINTS)
+        ik, name = ur5e_mujoco.build_ik_solver(model, data, JOINTS, collision, Path(MENAGERIE), backend="auto")
+        assert name == "ssik" and isinstance(ik, SSIKSolver)
+
+    def test_bad_backend_rejected(self):
+        with pytest.raises(ValueError):
+            ur5e_mujoco.build_ik_solver(None, None, JOINTS, None, Path(MENAGERIE), backend="eaik")
+
+    def test_script_runs_through_the_mujoco_fallback(self, tmp_path):
+        env = dict(os.environ, MPLBACKEND="Agg", MUJOCO_MENAGERIE_PATH=MENAGERIE)
+        proc = subprocess.run(
+            [sys.executable, str(EXAMPLES / "ur5e_mujoco.py"), "--no-viz", "--ik", "mujoco"],
+            cwd=tmp_path,
+            env=env,
+            capture_output=True,
+            text=True,
+            timeout=600,
+        )
+        assert proc.returncode == 0, proc.stderr[-2000:]
+        assert "Using MuJoCo (differential) IK solver" in proc.stdout
+        assert "Found path" in proc.stdout
+
+
+class TestTsrUnionDemo:
+    def test_one_planning_query_through_the_mujoco_fallback(self):
+        target = np.array([0.45, 0.0, 0.38])
+        model = tsr_union_demo.create_scene(Path(MENAGERIE), target)
+        data = mujoco.MjData(model)
+        robot = MuJoCoRobotModel(model, data, "attachment_site", JOINTS)
+        collision = MuJoCoCollisionChecker(model, data, JOINTS)
+        ik, name = tsr_union_demo.build_ik_solver(model, data, JOINTS, collision, Path(MENAGERIE), backend="mujoco")
+        assert name == "mujoco" and isinstance(ik, MuJoCoIKSolver)
+        planner = CBiRRT(robot, ik, collision, CBiRRTConfig(timeout=60.0, goal_bias=0.15, tsr_samples=100))
+        top, side = tsr_union_demo.create_grasp_tsrs(target)
+        result = planner.plan(start=HOME, goal_tsrs=[top, side], seed=0, return_details=True)
+        assert result.success, result.failure_reason
+        assert all(collision.is_valid(q) for q in result.path)