[FEATURE] #2332 Implement Local Offset in Inverse Kinematics

Author: alexis779

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -1422,6 +1422,7 @@ def inverse_kinematics(
         link,
         pos=None,
         quat=None,
+        local_point=None,
         init_qpos=None,
         respect_joint_limit=True,
         max_samples=50,
@@ -1447,6 +1448,10 @@ def inverse_kinematics(
             The target position. If None, position error will not be considered. Defaults to None.
         quat : None | array_like, shape (4,), optional
             The target orientation. If None, orientation error will not be considered. Defaults to None.
+        local_point : None | array_like, shape (3,), optional
+            A point in the link's local frame to be positioned at `pos`. If None, the link origin is used.
+            This is useful for positioning a tool center point (TCP) or fingertip that is offset from the link origin.
+            Defaults to None (equivalent to [0, 0, 0]).
         init_qpos : None | array_like, shape (n_dofs,), optional
             Initial qpos used for solving IK. If None, the current qpos will be used. Defaults to None.
         respect_joint_limit : bool, optional
@@ -1495,6 +1500,7 @@ def inverse_kinematics(
             links=[link],
             poss=[pos] if pos is not None else [],
             quats=[quat] if quat is not None else [],
+            local_points=[local_point] if local_point is not None else [],
             init_qpos=init_qpos,
             respect_joint_limit=respect_joint_limit,
             max_samples=max_samples,
@@ -1521,6 +1527,7 @@ def inverse_kinematics_multilink(
         links,
         poss=None,
         quats=None,
+        local_points=None,
         init_qpos=None,
         respect_joint_limit=True,
         max_samples=50,
@@ -1546,6 +1553,11 @@ def inverse_kinematics_multilink(
             List of target positions. If empty, position error will not be considered. Defaults to None.
         quats : list, optional
             List of target orientations. If empty, orientation error will not be considered. Defaults to None.
+        local_points : list, optional
+            List of local points (one per link) in each link's local frame to be positioned at the corresponding target position.
+            If empty or None, link origins are used. Each element should be array_like of shape (3,) or None.
+            This is useful for positioning tool center points (TCP) or fingertips that are offset from the link origin.
+            Defaults to None.
         init_qpos : array_like, shape (n_dofs,), optional
             Initial qpos used for solving IK. If None, the current qpos will be used. Defaults to None.
         respect_joint_limit : bool, optional
@@ -1608,6 +1620,18 @@ def inverse_kinematics_multilink(
         elif len(quats) != n_links:
             gs.raise_exception("Accepting only `quats` with length equal to `links` or empty list.")
 
+        # Process local_points - default to origin [0, 0, 0] for each link
+        local_points = list(local_points) if local_points is not None else []
+        if not local_points:
+            local_points = [None for _ in range(n_links)]
+        elif len(local_points) != n_links:
+            gs.raise_exception("Accepting only `local_points` with length equal to `links` or empty list.")
+        for i, lp in enumerate(local_points):
+            if lp is None:
+                lp = [0.0, 0.0, 0.0]
+            local_points[i] = broadcast_tensor(lp, gs.tc_float, (3,), ("",)).contiguous()
+        local_points = torch.stack(local_points, dim=0)  # (n_links, 3)
+
         link_pos_mask, link_rot_mask = [], []
         for i, (pos, quat) in enumerate(zip(poss, quats)):
             if pos is None and quat is None:
@@ -1658,6 +1682,7 @@ def inverse_kinematics_multilink(
             links_idx,
             poss,
             quats,
+            local_points,
             n_links,
             dofs_idx,
             n_dofs,

genesis/engine/solvers/rigid/abd/inverse_kinematics.py

@@ -20,6 +20,7 @@ def kernel_rigid_entity_inverse_kinematics(
     links_idx: ti.types.ndarray(),
     poss: ti.types.ndarray(),
     quats: ti.types.ndarray(),
+    local_points: ti.types.ndarray(),
     n_links: ti.i32,
     dofs_idx: ti.types.ndarray(),
     n_dofs: ti.i32,
@@ -94,7 +95,11 @@ def kernel_rigid_entity_inverse_kinematics(
                     i_l_ee = links_idx[i_ee]
 
                     tgt_pos_i = ti.Vector([poss[i_ee, i_b_, 0], poss[i_ee, i_b_, 1], poss[i_ee, i_b_, 2]])
-                    err_pos_i = tgt_pos_i - links_state.pos[i_l_ee, i_b]
+                    local_point_i = ti.Vector([local_points[i_ee, 0], local_points[i_ee, 1], local_points[i_ee, 2]])
+                    pos_curr_i = links_state.pos[i_l_ee, i_b] + gu.ti_transform_by_quat(
+                        local_point_i, links_state.quat[i_l_ee, i_b]
+                    )
+                    err_pos_i = tgt_pos_i - pos_curr_i
                     for k in range(3):
                         err_pos_i[k] *= pos_mask[k] * link_pos_mask[i_ee]
                     if err_pos_i.norm() > pos_tol:
@@ -123,10 +128,11 @@ def kernel_rigid_entity_inverse_kinematics(
                 for i_ee in range(n_links):
                     # update jacobian for ee link
                     i_l_ee = links_idx[i_ee]
+                    local_point_i = ti.Vector([local_points[i_ee, 0], local_points[i_ee, 1], local_points[i_ee, 2]])
                     rigid_entity._func_get_jacobian(
                         tgt_link_idx=i_l_ee,
                         i_b=i_b,
-                        p_local=ti.Vector.zero(gs.ti_float, 3),
+                        p_local=local_point_i,
                         pos_mask=pos_mask,
                         rot_mask=rot_mask,
                         dofs_info=dofs_info,
@@ -222,7 +228,11 @@ def kernel_rigid_entity_inverse_kinematics(
                     i_l_ee = links_idx[i_ee]
 
                     tgt_pos_i = ti.Vector([poss[i_ee, i_b_, 0], poss[i_ee, i_b_, 1], poss[i_ee, i_b_, 2]])
-                    err_pos_i = tgt_pos_i - links_state.pos[i_l_ee, i_b]
+                    local_point_i = ti.Vector([local_points[i_ee, 0], local_points[i_ee, 1], local_points[i_ee, 2]])
+                    pos_curr_i = links_state.pos[i_l_ee, i_b] + gu.ti_transform_by_quat(
+                        local_point_i, links_state.quat[i_l_ee, i_b]
+                    )
+                    err_pos_i = tgt_pos_i - pos_curr_i
                     for k in range(3):
                         err_pos_i[k] *= pos_mask[k] * link_pos_mask[i_ee]
                     if err_pos_i.norm() > pos_tol:

tests/test_rigid_physics.py

@@ -1575,6 +1575,150 @@ def test_multilink_inverse_kinematics(show_viewer):
     assert_allclose(wrist.get_pos(envs_idx=(1,)), wrist_pos, tol=TOL)
 
 
+@pytest.mark.required
+@pytest.mark.parametrize("n_envs", [0, 2])
+@pytest.mark.parametrize("backend", [gs.cpu, gs.gpu])
+def test_inverse_kinematics_local_point(n_envs, show_viewer):
+    """Test IK with local_point parameter - positions an offset point at the target instead of link origin."""
+
+    TOL = 2e-3  # 2mm tolerance for final position check
+
+    scene = gs.Scene(
+        viewer_options=gs.options.ViewerOptions(
+            camera_pos=(2.5, 0.0, 1.5),
+            camera_lookat=(0.0, 0.0, 0.5),
+        ),
+        show_viewer=show_viewer,
+    )
+    robot = scene.add_entity(
+        morph=gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"),
+    )
+    scene.build(n_envs=n_envs)
+
+    end_effector = robot.get_link("hand")
+
+    # Define a local offset point in the end-effector frame (e.g., 10cm along Z-axis)
+    local_offset = torch.tensor([0.0, 0.0, 0.1], dtype=gs.tc_float, device=gs.device)
+
+    # Create different target positions and quaternions for each environment
+    num_envs = max(n_envs, 1)
+    target_pos_base = torch.tensor(
+        [[0.5, 0.2, 0.4], [0.45, 0.15, 0.35], [0.55, 0.25, 0.45]], dtype=gs.tc_float, device=gs.device
+    )[:num_envs]
+    target_quat_base = torch.tensor(
+        [[0.0, 1.0, 0.0, 0.0], [0.0, 0.9239, 0.3827, 0.0], [0.0, 0.9239, -0.3827, 0.0]],
+        dtype=gs.tc_float,
+        device=gs.device,
+    )[:num_envs]
+
+    # Handle different shapes based on n_envs
+    if n_envs > 0:
+        target_pos = target_pos_base
+        target_quat = target_quat_base
+    else:
+        target_pos = target_pos_base[0]
+        target_quat = target_quat_base[0]
+
+    # Solve IK with local_point (local_offset stays 1D - it gets broadcast internally)
+    qpos, err = robot.inverse_kinematics(
+        link=end_effector,
+        pos=target_pos,
+        quat=target_quat,
+        local_point=local_offset,
+        return_error=True,
+    )
+
+    # Apply the solution
+    robot.set_qpos(qpos)
+    scene.step()
+
+    # Verify the offset point is at the target position
+    link_pos = end_effector.get_pos()
+    link_quat = end_effector.get_quat()
+
+    # Transform local offset to world frame
+    world_offset = gu.transform_by_quat(local_offset, link_quat)
+    actual_point_pos = link_pos + world_offset
+
+    # Check that the offset point reached the target
+    assert_allclose(actual_point_pos, target_pos, tol=TOL)
+
+    # Also verify via forward kinematics
+    links_pos, links_quat = robot.forward_kinematics(qpos)
+
+    # Handle indexing based on n_envs
+    if n_envs > 0:
+        fk_link_pos = links_pos[:, end_effector.idx_local]
+        fk_link_quat = links_quat[:, end_effector.idx_local]
+    else:
+        fk_link_pos = links_pos[end_effector.idx_local]
+        fk_link_quat = links_quat[end_effector.idx_local]
+
+    fk_world_offset = gu.transform_by_quat(local_offset, fk_link_quat)
+    fk_actual_point_pos = fk_link_pos + fk_world_offset
+    assert_allclose(fk_actual_point_pos, target_pos, tol=TOL)
+
+    if show_viewer:
+        scene.visualizer.update()
+
+
+@pytest.mark.required
+@pytest.mark.parametrize("backend", [gs.cpu, gs.gpu])
+def test_inverse_kinematics_multilink_local_points(show_viewer):
+    """Test multi-link IK with local_points parameter."""
+
+    TOL = 2e-3  # 2mm tolerance for final position check
+
+    scene = gs.Scene(
+        viewer_options=gs.options.ViewerOptions(
+            camera_pos=(2.5, 0.0, 1.5),
+            camera_lookat=(0.0, 0.0, 0.5),
+        ),
+        show_viewer=show_viewer,
+    )
+    robot = scene.add_entity(
+        morph=gs.morphs.URDF(file="urdf/shadow_hand/shadow_hand.urdf"),
+    )
+    scene.build()
+
+    index_finger = robot.get_link("index_finger_distal")
+    middle_finger = robot.get_link("middle_finger_distal")
+
+    # Different local offsets for each finger (e.g., fingertip positions)
+    index_local_offset = torch.tensor([0.0, 0.0, 0.02], dtype=gs.tc_float, device=gs.device)
+    middle_local_offset = torch.tensor([0.0, 0.0, 0.02], dtype=gs.tc_float, device=gs.device)
+
+    # Target positions for the fingertips
+    index_target = torch.tensor([0.6, 0.5, 0.2], dtype=gs.tc_float, device=gs.device)
+    middle_target = torch.tensor([0.63, 0.5, 0.2], dtype=gs.tc_float, device=gs.device)
+
+    # Solve multi-link IK with local_points
+    qpos, err = robot.inverse_kinematics_multilink(
+        links=[index_finger, middle_finger],
+        poss=[index_target, middle_target],
+        local_points=[index_local_offset, middle_local_offset],
+        return_error=True,
+    )
+
+    # Apply solution
+    robot.set_qpos(qpos)
+    scene.step()
+
+    # Verify each offset point is at its target
+    for link, local_offset, target in [
+        (index_finger, index_local_offset, index_target),
+        (middle_finger, middle_local_offset, middle_target),
+    ]:
+        link_pos = link.get_pos()
+        link_quat = link.get_quat()
+        world_offset = gu.transform_by_quat(local_offset, link_quat)
+        actual_point_pos = link_pos + world_offset
+        assert_allclose(actual_point_pos, target, tol=TOL)
+
+    if show_viewer:
+        scene.visualizer.update()
+
+
 @pytest.mark.slow  # ~180s
 @pytest.mark.required
 @pytest.mark.parametrize("n_envs", [0, 2])