diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index 0768a8f..7ea4760 100644 --- a/.github/workflows/python-package.yml +++ b/.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 diff --git a/.gitignore b/.gitignore index 86eaeb1..1ab8fb4 100644 --- a/.gitignore +++ b/.gitignore @@ -12,3 +12,4 @@ dist tests/MUJOCO_LOG.TXT tests/mujoco_menagerie/ .ipynb_checkpoints +CLAUDE.md diff --git a/pyproject.toml b/pyproject.toml index f9b05eb..1f9035d 100644 --- a/pyproject.toml +++ b/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 diff --git a/src/pytorch_kinematics/chain.py b/src/pytorch_kinematics/chain.py index 0b21114..04b092a 100644 --- a/src/pytorch_kinematics/chain.py +++ b/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) diff --git a/src/pytorch_kinematics/ik.py b/src/pytorch_kinematics/ik.py index 356cffc..3e448ab 100644 --- a/src/pytorch_kinematics/ik.py +++ b/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,6 +173,10 @@ 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 @@ -173,6 +184,12 @@ def __init__(self, serial_chain: SerialChain, 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 diff --git a/src/pytorch_kinematics/jacobian.py b/src/pytorch_kinematics/jacobian.py index 4ff3b32..94c0796 100644 --- a/src/pytorch_kinematics/jacobian.py +++ b/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() diff --git a/src/pytorch_kinematics/transforms/transform3d.py b/src/pytorch_kinematics/transforms/transform3d.py index bad7949..000ddd4 100644 --- a/src/pytorch_kinematics/transforms/transform3d.py +++ b/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) diff --git a/tests/benchmark_ik.py b/tests/benchmark_ik.py new file mode 100644 index 0000000..9a454c9 --- /dev/null +++ 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) diff --git a/tests/test_ik_joint_limits.py b/tests/test_ik_joint_limits.py new file mode 100644 index 0000000..77d473c --- /dev/null +++ b/tests/test_ik_joint_limits.py @@ -0,0 +1,79 @@ +import os + +import torch +import pytorch_kinematics as pk +import pytorch_seed + +try: + import pybullet_data + URDF_PATH = os.path.join(pybullet_data.getDataPath(), "kuka_iiwa/model.urdf") +except ImportError: + URDF_PATH = os.path.join(os.path.dirname(__file__), "kuka_iiwa.urdf") + + +def test_ik_solutions_within_joint_limits(): + pytorch_seed.seed(2) + device = "cuda" if torch.cuda.is_available() else "cpu" + chain = pk.build_serial_chain_from_urdf(open(URDF_PATH).read(), "lbr_iiwa_link_7") + chain = chain.to(device=device) + + lim = torch.tensor(chain.get_joint_limits(), device=device) + M = 100 + 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=50, num_retries=num_retries, + joint_limits=lim.T, + early_stopping_any_converged=True, + early_stopping_no_improvement="all", + debug=False, lr=0.2, + enforce_joint_limits=True) + + sol = ik.solve(goal_tf) + print("IK solved %d / %d goals" % (sol.converged_any.sum(), M)) + + # All solutions must respect joint limits + q_all = sol.solutions + eps = 1e-6 + low = chain.low + high = chain.high + assert (q_all >= low - eps).all(), \ + f"Solutions below lower limits: max violation = {(low - q_all).clamp(min=0).max()}" + assert (q_all <= high + eps).all(), \ + f"Solutions above upper limits: max violation = {(q_all - high).clamp(min=0).max()}" + + # Convergence rate should be high since these goals are reachable + converged_rate = sol.converged_any.sum().item() / M + print(f"Convergence rate: {converged_rate:.1%}") + assert converged_rate > 0.8, f"Convergence rate too low: {converged_rate:.1%}" + + +def test_enforce_joint_limits_flag(): + """Test that enforce_joint_limits=False disables limit enforcement.""" + pytorch_seed.seed(2) + chain = pk.build_serial_chain_from_urdf(open(URDF_PATH).read(), "lbr_iiwa_link_7") + + lim = torch.tensor(chain.get_joint_limits()) + M = 50 + goal_q = torch.rand(M, lim.shape[1]) * (lim[1] - lim[0]) + lim[0] + goal_tf = chain.forward_kinematics(goal_q) + + ik = pk.PseudoInverseIK(chain, max_iterations=30, num_retries=10, + joint_limits=lim.T, + early_stopping_any_converged=True, + early_stopping_no_improvement="all", + debug=False, lr=0.2, + enforce_joint_limits=False) + + sol = ik.solve(goal_tf) + # With enforce_joint_limits=False, some solutions may exceed limits + q_all = sol.solutions + has_violation = ((q_all < chain.low).any() | (q_all > chain.high).any()).item() + # We don't assert violations must exist (they might not), but just check it ran without error + print(f"enforce_joint_limits=False: violations present = {has_violation}") + + +if __name__ == "__main__": + test_ik_solutions_within_joint_limits() + test_enforce_joint_limits_flag() diff --git a/tests/test_inverse_kinematics.py b/tests/test_inverse_kinematics.py index 2659f16..c26c041 100644 --- a/tests/test_inverse_kinematics.py +++ b/tests/test_inverse_kinematics.py @@ -207,7 +207,7 @@ def test_ik_in_place_no_err(robot="kuka_iiwa"): # do IK sol = ik.solve(goal_in_rob_frame_tf) assert sol.converged.sum() == M - assert torch.allclose(sol.solutions[0][0], cur_q) + assert torch.allclose(sol.solutions[0][0], cur_q, atol=1e-6) assert torch.allclose(sol.err_pos[0], torch.zeros(1, device=device), atol=1e-6) assert torch.allclose(sol.err_rot[0], torch.zeros(1, device=device), atol=1e-6) diff --git a/tests/test_jacobian.py b/tests/test_jacobian.py index 960c722..c244f73 100644 --- a/tests/test_jacobian.py +++ b/tests/test_jacobian.py @@ -192,6 +192,31 @@ def get_pt(th): pass +def test_mjcf_jacobian(): + """Test that the Jacobian is correct for MJCF models where link offsets are non-identity (issue #57).""" + chain = pk.build_serial_chain_from_mjcf(open(os.path.join(TEST_DIR, "ant.xml")).read(), 'front_left_foot') + chain = chain.to(dtype=torch.float64) + + def get_pt(th): + return chain.forward_kinematics(th).transform_points( + torch.zeros((1, 3), device=th.device, dtype=th.dtype)).squeeze(1) + + N = 20 + th = torch.rand(N, 2, dtype=torch.float64) + + j_autograd = torch.autograd.functional.jacobian(get_pt, inputs=th, vectorize=True) + j_autograd = j_autograd[range(N), :, range(N)] + + j_analytical = chain.jacobian(th) + + # the position Jacobian (first 3 rows) must match autograd + assert torch.allclose(j_autograd, j_analytical[:, :3], atol=1e-6), \ + f"MJCF Jacobian mismatch:\nanalytical:\n{j_analytical[0, :3]}\nautograd:\n{j_autograd[0]}" + + # sanity check: the linear velocity rows should NOT be all zeros + assert j_analytical[:, :3].abs().sum() > 0, "Linear velocity Jacobian should be non-zero" + + if __name__ == "__main__": test_correctness() test_parallel() @@ -200,3 +225,4 @@ def get_pt(th): test_jacobian_prismatic() test_jacobian_at_different_loc_than_ee() test_comparison_to_autograd() + test_mjcf_jacobian() diff --git a/tests/test_serial_chain_creation.py b/tests/test_serial_chain_creation.py index 057d54c..0f2dada 100644 --- a/tests/test_serial_chain_creation.py +++ b/tests/test_serial_chain_creation.py @@ -70,7 +70,7 @@ def test_extract_serial_chain_from_tree(): serial_chain = pk.SerialChain(chain, "ee_gripper_link", "gripper_link") serial_frame_expected = """ - gripper_link +gripper_link └── ee_arm_link └── gripper_bar_link └── fingers_link @@ -78,9 +78,59 @@ def test_extract_serial_chain_from_tree(): """ serial_frame = serial_chain.print_tree() assert serial_frame_expected.strip() == serial_frame.strip() - # only gripper_link is the parent frame of a joint in this serial chain - assert serial_chain.n_joints == 1 + # gripper_link's own joint connects wrist_link -> gripper_link, which is outside + # this sub-chain, so it should not be counted. All remaining joints are fixed. + assert serial_chain.n_joints == 0 + + +def test_serial_chain_non_root_start(): + """Test that SerialChain with a non-root root_frame_name excludes the root's own joint (issue #62).""" + data = ('' + '' + '' + '' + '' + '' + '' + '' + '' + '' + '' + '' + '' + '' + '' + '' + '' + '' + '' + '' + '') + full_chain = pk.build_chain_from_urdf(data) + + # SerialChain from link2 to link4 should only include joint2 and joint3, not joint1 + serial = pk.SerialChain(full_chain, "link4", "link2") + assert "joint1" not in serial.get_joint_parameter_names() + assert serial.get_joint_parameter_names() == ["joint2", "joint3"] + + # FK from a frame to itself should be identity + identity_chain = pk.SerialChain(full_chain, "link2", "link2") + tg = identity_chain.forward_kinematics(torch.zeros(0)) + assert torch.allclose(tg.get_matrix(), torch.eye(4).unsqueeze(0), atol=1e-7) + + # FK from link2 to link4 with zeros should give translation of (2, 0, 0) — two joints each offset by (1, 0, 0) + tg = serial.forward_kinematics(torch.zeros(2)) + pos = tg.get_matrix()[0, :3, 3] + assert torch.allclose(pos, torch.tensor([2.0, 0.0, 0.0]), atol=1e-7) + + # Compare: full chain FK at zeros, link4 relative to link2 + full_ret = full_chain.forward_kinematics(torch.zeros(3)) + full_link2 = full_ret['link2'].get_matrix() + full_link4 = full_ret['link4'].get_matrix() + relative = torch.linalg.inv(full_link2) @ full_link4 + assert torch.allclose(relative, tg.get_matrix(), atol=1e-7) if __name__ == "__main__": test_extract_serial_chain_from_tree() + test_serial_chain_non_root_start()