IK joint limits and other various fixes

.github/workflows/python-package.yml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [ "3.8", "3.9", "3.10" ]
+        python-version: [ "3.9", "3.10", "3.11", "3.12" ]
 
     steps:
       - uses: actions/checkout@v3
@@ -30,7 +30,7 @@ jobs:
           python -m pip install .[test]
           python -m pip install flake8 pytest
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-          pip install mujoco
+          pip install "mujoco<=3.3.7"
       - name: Clone mujoco_menagerie repository into the tests/ folder
         run: |
           git clone https://github.com/google-deepmind/mujoco_menagerie

.gitignore

@@ -12,3 +12,4 @@ dist
 tests/MUJOCO_LOG.TXT
 tests/mujoco_menagerie/
 .ipynb_checkpoints
+CLAUDE.md

pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md" # Optional
 # 'Programming Language' classifiers above, 'pip install' will check this
 # and refuse to install the project if the version does not match. See
 # https://packaging.python.org/guides/distributing-packages-using-setuptools/#python-requires
-requires-python = ">=3.6"
+requires-python = ">=3.9"
 
 # This is either text indicating the license for the distribution, or a file
 # that contains the license

src/pytorch_kinematics/chain.py

@@ -463,6 +463,12 @@ def __init__(self, chain, end_frame_name, root_frame_name="", **kwargs):
         for i in range(len(ancestors) - 1):
             frames[i].children = [frames[i + 1]]
 
+        # The root frame's joint and link offset describe the transform from its parent
+        # to itself, which is outside this serial chain. Reset them so the root acts as
+        # the origin of the new coordinate system.
+        frames[0].joint = Joint()
+        frames[0].link.offset = None
+
         self._serial_frames = frames
         super().__init__(frames[0], **kwargs)
 

src/pytorch_kinematics/ik.py

@@ -3,6 +3,7 @@
 from pytorch_kinematics.transforms import rotation_conversions
 from typing import NamedTuple, Union, Optional, Callable
 import typing
+import math
 import torch
 import inspect
 from matplotlib import pyplot as plt, cm as cm
@@ -61,13 +62,17 @@ def update(self, q: torch.tensor, err: torch.tensor, use_keep_mask=True, keep_ma
             # those that have converged are no longer remaining
             self.update_remaining_with_keep_mask(keep_mask)
 
-        self.solutions = qq
-        self.err_pos = err_pos
-        self.err_rot = err_rot
-        self.converged_pos = converged_pos
-        self.converged_rot = converged_rot
-        self.converged = converged
-        self.converged_any = converged_any
+        # sticky convergence: only overwrite retries that haven't already converged
+        # this prevents overwriting a good solution if the solver overshoots on the next step
+        already_converged = self.converged
+        update_mask = ~already_converged
+        self.solutions[update_mask] = qq[update_mask]
+        self.err_pos[update_mask] = err_pos[update_mask]
+        self.err_rot[update_mask] = err_rot[update_mask]
+        self.converged_pos = self.converged_pos | converged_pos
+        self.converged_rot = self.converged_rot | converged_rot
+        self.converged = self.converged | converged
+        self.converged_any = self.converged_any | converged_any
 
         return converged_any
 
@@ -141,7 +146,9 @@ def __init__(self, serial_chain: SerialChain,
                  debug=False,
                  early_stopping_any_converged=False,
                  early_stopping_no_improvement="any", early_stopping_no_improvement_patience=2,
-                 optimizer_method: Union[str, typing.Type[torch.optim.Optimizer]] = "sgd"
+                 optimizer_method: Union[str, typing.Type[torch.optim.Optimizer]] = "sgd",
+                 enforce_joint_limits: bool = True,
+                 num_limit_refinement_iterations: int = 10
                  ):
         """
         :param serial_chain:
@@ -166,13 +173,23 @@ def __init__(self, serial_chain: SerialChain,
         :param early_stopping_no_improvement_patience: number of consecutive iterations with no improvement before
         considering it no improvement
         :param optimizer_method: either a string or a torch.optim.Optimizer class
+        :param enforce_joint_limits: whether to enforce joint limits on the solution. Uses the chain's joint limits
+        (chain.low/chain.high). After solving, revolute joints are wrapped by multiples of 2*pi, then any remaining
+        violations are resolved via clamped refinement iterations. Set to False to disable.
+        :param num_limit_refinement_iterations: number of clamped IK refinement iterations for enforcing joint limits
         """
         self.chain = serial_chain
         self.dtype = serial_chain.dtype
         self.device = serial_chain.device
         joint_names = self.chain.get_joint_parameter_names(exclude_fixed=True)
         self.dof = len(joint_names)
         self.debug = debug
+        self.enforce_joint_limits = enforce_joint_limits
+        self.num_limit_refinement_iterations = num_limit_refinement_iterations
+        # precompute which joints are revolute (for 2*pi wrapping)
+        joints = self.chain.get_joints(exclude_fixed=True)
+        self._revolute_mask = torch.tensor([j.joint_type == 'revolute' for j in joints],
+                                           dtype=torch.bool, device=self.device)
         self.early_stopping_any_converged = early_stopping_any_converged
         self.early_stopping_no_improvement = early_stopping_no_improvement
         self.early_stopping_no_improvement_patience = early_stopping_no_improvement_patience
@@ -347,35 +364,37 @@ def solve(self, target_poses: Transform3d) -> IKSolution:
                     q = q + lr * dq
 
             with torch.no_grad():
-                self.err_all = dx.squeeze()
+                # recompute error at the new q so convergence check matches the stored solution
+                m_new = self.chain.forward_kinematics(q).get_matrix()
+                m_new = m_new.view(-1, self.num_retries, 4, 4)
+                dx_new, pos_diff, rot_diff = delta_pose(m_new, target_pos, target_wxyz)
+                self.err_all = dx_new.squeeze()
                 self.err = self.err_all.norm(dim=-1)
                 sol.update(q, self.err_all, use_keep_mask=self.early_stopping_any_converged)
 
                 if self.early_stopping_no_improvement is not None:
                     if self.no_improve_counter is None:
                         self.no_improve_counter = torch.zeros_like(self.err)
+                        self.err_min = self.err.clone()
                     else:
-                        if self.err_min is None:
-                            self.err_min = self.err.clone()
-                        else:
-                            improved = self.err < self.err_min
-                            self.err_min[improved] = self.err[improved]
-
-                            self.no_improve_counter[improved] = 0
-                            self.no_improve_counter[~improved] += 1
-
-                            # those that haven't improved
-                            could_improve = self.no_improve_counter <= self.early_stopping_no_improvement_patience
-                            # consider problems, and only throw out those whose all retries cannot be improved
-                            could_improve = could_improve.reshape(-1, self.num_retries)
-                            if self.early_stopping_no_improvement == "all":
-                                could_improve = could_improve.all(dim=1)
-                            elif self.early_stopping_no_improvement == "any":
-                                could_improve = could_improve.any(dim=1)
-                            elif isinstance(self.early_stopping_no_improvement, float):
-                                ratio_improved = could_improve.sum(dim=1) / self.num_retries
-                                could_improve = ratio_improved > self.early_stopping_no_improvement
-                            sol.update_remaining_with_keep_mask(could_improve)
+                        improved = self.err < self.err_min
+                        self.err_min[improved] = self.err[improved]
+
+                        self.no_improve_counter[improved] = 0
+                        self.no_improve_counter[~improved] += 1
+
+                        # those that haven't improved
+                        could_improve = self.no_improve_counter <= self.early_stopping_no_improvement_patience
+                        # consider problems, and only throw out those whose all retries cannot be improved
+                        could_improve = could_improve.reshape(-1, self.num_retries)
+                        if self.early_stopping_no_improvement == "all":
+                            could_improve = could_improve.all(dim=1)
+                        elif self.early_stopping_no_improvement == "any":
+                            could_improve = could_improve.any(dim=1)
+                        elif isinstance(self.early_stopping_no_improvement, float):
+                            ratio_improved = could_improve.sum(dim=1) / self.num_retries
+                            could_improve = ratio_improved > self.early_stopping_no_improvement
+                        sol.update_remaining_with_keep_mask(could_improve)
 
                 if self.debug:
                     pos_errors.append(pos_diff.reshape(-1, 3).norm(dim=1))
@@ -405,10 +424,68 @@ def solve(self, target_poses: Transform3d) -> IKSolution:
             ax[1].set_ylabel("rotation error")
             plt.show()
 
-        if i == self.max_iterations - 1:
-            sol.update(q, self.err_all, use_keep_mask=False)
+        if self.enforce_joint_limits and self._has_finite_limits():
+            self._enforce_limits(sol, target_pos, target_wxyz)
+
         return sol
 
+    def _has_finite_limits(self):
+        """Check if the chain has any finite joint limits."""
+        return torch.isfinite(self.chain.low).any() or torch.isfinite(self.chain.high).any()
+
+    def _wrap_revolute_joints(self, q):
+        """Wrap revolute joint values by multiples of 2*pi to bring them closer to the valid range.
+
+        For revolute joints, q and q + 2*n*pi produce identical FK, so we can freely shift
+        by multiples of 2*pi without affecting the solution. Prismatic joints are left unchanged.
+        """
+        low = self.chain.low
+        high = self.chain.high
+        center = (low + high) / 2
+        # compute the number of 2*pi shifts needed to bring q closest to the center of the limits
+        n = torch.round((center - q) / (2 * math.pi))
+        # only apply to revolute joints
+        q_wrapped = q.clone()
+        q_wrapped[..., self._revolute_mask] = q[..., self._revolute_mask] + n[..., self._revolute_mask] * 2 * math.pi
+        return q_wrapped
+
+    def _enforce_limits(self, sol, target_pos, target_wxyz):
+        """Enforce joint limits on IK solutions via wrapping and clamped refinement.
+
+        1. Wrap revolute joints by 2*pi to bring them into the valid range (preserves FK exactly).
+        2. Clamp any remaining violations and run refinement iterations to recover convergence.
+        3. Update the solution with the refined joint values.
+        """
+        M = target_pos.shape[0]
+        q = sol.solutions.reshape(-1, self.dof)
+
+        # Step 1: wrap revolute joints by 2*pi
+        q = self._wrap_revolute_joints(q)
+
+        # Step 2: clamp to limits
+        q = torch.clamp(q, self.chain.low, self.chain.high)
+
+        # Step 3: clamped refinement iterations to recover from any error introduced by clamping
+        for _ in range(self.num_limit_refinement_iterations):
+            J, m = self.chain.jacobian(q, ret_eef_pose=True)
+            m = m.view(-1, self.num_retries, 4, 4)
+            dx, _, _ = delta_pose(m, target_pos, target_wxyz)
+            dq = self.compute_dq(J, dx).squeeze(2)
+            q = q + self.lr * dq
+            q = torch.clamp(q, self.chain.low, self.chain.high)
+
+        # Reset convergence so all retries are re-evaluated with the limit-enforced solutions
+        sol.converged[:] = False
+        sol.converged_any[:] = False
+        sol.converged_pos[:] = False
+        sol.converged_rot[:] = False
+
+        # Recompute error and update solution
+        m_new = self.chain.forward_kinematics(q).get_matrix()
+        m_new = m_new.view(-1, self.num_retries, 4, 4)
+        dx_new, _, _ = delta_pose(m_new, target_pos, target_wxyz)
+        sol.update(q, dx_new.squeeze(), use_keep_mask=False)
+
 
 class PseudoInverseIKWithSVD(PseudoInverseIK):
     # generally slower, but allows for selective damping if needed

src/pytorch_kinematics/jacobian.py

@@ -12,8 +12,7 @@ def calc_jacobian(serial_chain, th, tool=None, ret_eef_pose=False):
     tool is the transformation wrt the end effector; default is identity. If specified, will have to
     specify for each of the N inputs
 
-    FIXME: this code assumes the joint frame and the child link frame are the same
-    """
+"""
     if not torch.is_tensor(th):
         th = torch.tensor(th, dtype=serial_chain.dtype, device=serial_chain.device)
     if len(th.shape) <= 1:
@@ -49,6 +48,8 @@ def calc_jacobian(serial_chain, th, tool=None, ret_eef_pose=False):
             j_eef[:, :3, -cnt] = (f.joint.axis.repeat(N, 1, 1) @ cur_transform[:, :3, :3])[:, 0, :]
         cur_frame_transform = f.get_transform(th[:, -cnt]).get_matrix()
         cur_transform = cur_frame_transform @ cur_transform
+        if f.link.offset is not None:
+            cur_transform = f.link.offset.get_matrix() @ cur_transform
 
     # currently j_eef is Jacobian in end-effector frame, convert to base/world frame
     pose = serial_chain.forward_kinematics(th).get_matrix()

src/pytorch_kinematics/transforms/transform3d.py

@@ -51,7 +51,7 @@ class Transform3d:
 
     .. code-block:: python
 
-        y1 = t3.transform_points(t2.transform_points(t1.transform_points(x)))
+        y1 = t1.transform_points(t2.transform_points(t3.transform_points(x)))
         y2 = t1.compose(t2).compose(t3).transform_points(x)
         y3 = t1.compose(t2, t3).transform_points(x)
 
@@ -746,6 +746,8 @@ def _broadcast_bmm(a, b):
             a = a.expand(len(b), -1, -1)
         if len(b) == 1:
             b = b.expand(len(a), -1, -1)
+    if a.dtype != b.dtype:
+        b = b.to(dtype=a.dtype)
     return a.bmm(b)
 
 

tests/benchmark_ik.py

@@ -0,0 +1,72 @@
+"""Benchmark IK performance: convergence rate, timing, and joint limit compliance."""
+import os
+import sys
+from timeit import default_timer as timer
+
+import torch
+import pytorch_kinematics as pk
+import pytorch_seed
+
+
+def benchmark(device="cpu"):
+    pytorch_seed.seed(2)
+
+    # Try pybullet_data kuka first, fall back to local test URDF
+    try:
+        import pybullet_data
+        full_urdf = os.path.join(pybullet_data.getDataPath(), "kuka_iiwa/model.urdf")
+    except ImportError:
+        full_urdf = os.path.join(os.path.dirname(__file__), "kuka_iiwa.urdf")
+    chain = pk.build_serial_chain_from_urdf(open(full_urdf).read(), "lbr_iiwa_link_7")
+    chain = chain.to(device=device)
+
+    lim = torch.tensor(chain.get_joint_limits(), device=device)
+    M = 500
+    goal_q = torch.rand(M, lim.shape[1], device=device) * (lim[1] - lim[0]) + lim[0]
+    goal_tf = chain.forward_kinematics(goal_q)
+
+    num_retries = 10
+    ik = pk.PseudoInverseIK(chain, max_iterations=30, num_retries=num_retries,
+                            joint_limits=lim.T,
+                            early_stopping_any_converged=True,
+                            early_stopping_no_improvement="all",
+                            debug=False,
+                            lr=0.2)
+
+    # Warmup
+    _ = ik.solve(goal_tf)
+
+    # Timed run
+    N_RUNS = 3
+    times = []
+    for _ in range(N_RUNS):
+        pytorch_seed.seed(2)
+        ik.initial_config = ik.sample_configs(num_retries)
+        t0 = timer()
+        sol = ik.solve(goal_tf)
+        t1 = timer()
+        times.append(t1 - t0)
+
+    avg_time = sum(times) / len(times)
+
+    q_all = sol.solutions
+    low, high = chain.low, chain.high
+    below = (low - q_all).clamp(min=0).max().item()
+    above = (q_all - high).clamp(min=0).max().item()
+
+    print(f"Goals solved:       {sol.converged_any.sum().item()} / {M}")
+    print(f"Convergence rate:   {sol.converged_any.sum().item() / M:.1%}")
+    print(f"Total converged:    {sol.converged.sum().item()} / {sol.converged.numel()}")
+    print(f"Iterations:         {sol.iterations}")
+    print(f"Avg time ({N_RUNS} runs): {avg_time:.4f}s")
+    print(f"Max violation below: {below:.6f}")
+    print(f"Max violation above: {above:.6f}")
+
+
+if __name__ == "__main__":
+    device = sys.argv[1] if len(sys.argv) > 1 else "cpu"
+    if device == "cuda" and not torch.cuda.is_available():
+        print("CUDA not available, falling back to cpu")
+        device = "cpu"
+    print(f"=== Device: {device} ===")
+    benchmark(device=device)