[BugFix]: Fetch IK bug / mobile manipulation IK bug fix (#996)

* Fix: fetch ik fix

* Fix: fetch cpu ik fix

* Fix: allow set custom root name for pd ee

* demo_manual_control_continuous play script

* Fix: remove comments

* Fix: fetch pd_ee_target_delta_pose by using torso link as root for ik

Author: arth-shukla

mani_skill/agents/controllers/pd_ee_pose.py

@@ -1,12 +1,12 @@
 from dataclasses import dataclass
-from typing import Literal, Sequence, Union
+from typing import Literal, Optional, Sequence, Union
 
 import numpy as np
 import torch
 from gymnasium import spaces
 
 from mani_skill.agents.controllers.utils.kinematics import Kinematics
-from mani_skill.utils import gym_utils
+from mani_skill.utils import gym_utils, sapien_utils
 from mani_skill.utils.geometry.rotation_conversions import (
     euler_angles_to_matrix,
     matrix_to_quaternion,
@@ -45,6 +45,13 @@ def _initialize_joints(self):
 
         self.ee_link = self.kinematics.end_link
 
+        if self.config.root_link_name is not None:
+            self.root_link = sapien_utils.get_obj_by_name(
+                self.articulation.get_links(), self.config.root_link_name
+            )
+        else:
+            self.root_link = self.articulation.root
+
     def _initialize_action_space(self):
         low = np.float32(np.broadcast_to(self.config.pos_lower, 3))
         high = np.float32(np.broadcast_to(self.config.pos_upper, 3))
@@ -60,7 +67,7 @@ def ee_pose(self):
 
     @property
     def ee_pose_at_base(self):
-        to_base = self.articulation.pose.inv()
+        to_base = self.root_link.pose.inv()
         return to_base * (self.ee_pose)
 
     def reset(self):
@@ -69,9 +76,9 @@ def reset(self):
             self._target_pose = self.ee_pose_at_base
         else:
             # TODO (stao): this is a strange way to mask setting individual batched pose parts
-            self._target_pose.raw_pose[
-                self.scene._reset_mask
-            ] = self.ee_pose_at_base.raw_pose[self.scene._reset_mask]
+            self._target_pose.raw_pose[self.scene._reset_mask] = (
+                self.ee_pose_at_base.raw_pose[self.scene._reset_mask]
+            )
 
     def compute_target_pose(self, prev_ee_pose_at_base, action):
         # Keep the current rotation and change the position
@@ -156,6 +163,8 @@ class PDEEPosControllerConfig(ControllerConfig):
     ] = "root_translation"
     """Choice of frame to use for translational and rotational control of the end-effector. To learn how these work explicitly
     with videos of each one's behavior see https://maniskill.readthedocs.io/en/latest/user_guide/concepts/controllers.html#pd-ee-end-effector-pose"""
+    root_link_name: Optional[str] = None
+    """Optionally set different root link for root translation control (e.g. if root is different than base)"""
     use_delta: bool = True
     """Whether to use delta-action control. If true then actions indicate the delta/change in position via translation and orientation via
     rotation. If false, then actions indicate in the base frame (typically wherever the root link of the robot is) what pose the end effector

mani_skill/agents/controllers/utils/kinematics.py

@@ -1,6 +1,7 @@
 """
 Code for kinematics utilities on CPU/GPU
 """
+
 from contextlib import contextmanager, redirect_stderr, redirect_stdout
 from os import devnull
 from typing import List
@@ -12,6 +13,7 @@
         "pytorch_kinematics_ms not installed. Install with pip install pytorch_kinematics_ms"
     )
 import torch
+from lxml import etree as ET
 from sapien.wrapper.pinocchio_model import PinocchioModel
 
 from mani_skill.utils import common
@@ -44,12 +46,18 @@ def __init__(
             articulation (Articulation): the articulation object
             active_joint_indices (torch.Tensor): indices of the active joints that can be controlled
         """
+
+        # NOTE (arth): urdf path with feasible kinematic chain. may not be same urdf used to
+        #   build the sapien articulation (e.g. sapien articulation may have added joints for
+        #   mobile base which should not be used in IK)
         self.urdf_path = urdf_path
         self.end_link = articulation.links_map[end_link_name]
-        self.end_link_idx = articulation.links.index(self.end_link)
-        self.active_joint_indices = active_joint_indices
+
         self.articulation = articulation
         self.device = articulation.device
+
+        self.active_joint_indices = active_joint_indices
+
         # note that everything past the end-link is ignored. Any joint whose ancestor is self.end_link is ignored
         cur_link = self.end_link
         active_ancestor_joints: List[ArticulationJoint] = []
@@ -58,17 +66,25 @@ def __init__(
                 active_ancestor_joints.append(cur_link.joint)
             cur_link = cur_link.joint.parent_link
         active_ancestor_joints = active_ancestor_joints[::-1]
-        self.active_ancestor_joints = active_ancestor_joints
 
-        # initially self.active_joint_indices references active joints that are controlled.
-        # we also make the assumption that the active index is the same across all parallel managed joints
-        self.active_ancestor_joint_idxs = [
-            (x.active_index[0]).cpu().item() for x in self.active_ancestor_joints
-        ]
-        self.controlled_joints_idx_in_qmask = [
-            self.active_ancestor_joint_idxs.index(idx)
-            for idx in self.active_joint_indices
+        # NOTE (arth): some robots, like Fetch, have dummy joints that can mess with IK solver.
+        #   we assume that the urdf_path provides a valid kinematic chain, and prune joints
+        #   which are in the ManiSkill articulation but not in the kinematic chain
+        with open(self.urdf_path, "r") as f:
+            urdf_string = f.read()
+        xml = ET.fromstring(urdf_string.encode("utf-8"))
+        self._kinematic_chain_joint_names = set(
+            node.get("name") for node in xml if node.tag == "joint"
+        )
+        self._kinematic_chain_link_names = set(
+            node.get("name") for node in xml if node.tag == "link"
+        )
+        self.active_ancestor_joints = [
+            x
+            for x in active_ancestor_joints
+            if x.name in self._kinematic_chain_joint_names
         ]
+
         if self.device.type == "cuda":
             self.use_gpu_ik = True
             self._setup_gpu()
@@ -79,14 +95,47 @@ def __init__(
     def _setup_cpu(self):
         """setup the kinematics solvers on the CPU"""
         self.use_gpu_ik = False
-        # NOTE (stao): currently using the pinnochio that comes packaged with SAPIEN
-        self.qmask = torch.zeros(
-            self.articulation.max_dof, dtype=bool, device=self.device
+
+        with open(self.urdf_path, "r") as f:
+            xml = f.read()
+
+        joint_order = [
+            j.name
+            for j in self.articulation.active_joints
+            if j.name in self._kinematic_chain_joint_names
+        ]
+        link_order = [
+            l.name
+            for l in self.articulation.links
+            if l.name in self._kinematic_chain_link_names
+        ]
+
+        self.pmodel = PinocchioModel(xml, [0, 0, -9.81])
+        self.pmodel.set_joint_order(joint_order)
+        self.pmodel.set_link_order(link_order)
+
+        controlled_joint_names = [
+            self.articulation.active_joints[i].name for i in self.active_joint_indices
+        ]
+        self.pmodel_controlled_joint_indices = torch.tensor(
+            [joint_order.index(cj) for cj in controlled_joint_names],
+            dtype=torch.int,
+            device=self.device,
         )
-        self.pmodel: PinocchioModel = self.articulation._objs[
-            0
-        ].create_pinocchio_model()
-        self.qmask[self.active_joint_indices] = 1
+
+        articulation_active_joint_names_to_idx = dict(
+            (j.name, i) for i, j in enumerate(self.articulation.active_joints)
+        )
+        self.pmodel_active_joint_indices = torch.tensor(
+            [articulation_active_joint_names_to_idx[jn] for jn in joint_order],
+            dtype=torch.int,
+            device=self.device,
+        )
+
+        # NOTE (arth): pmodel will use urdf_path, set values based on this xml
+        self.end_link_idx = link_order.index(self.end_link.name)
+        self.qmask = torch.zeros(len(joint_order), dtype=bool, device=self.device)
+        self.qmask[self.pmodel_controlled_joint_indices] = 1
 
     def _setup_gpu(self):
         """setup the kinematics solvers on the GPU"""
@@ -116,6 +165,16 @@ def suppress_stdout_stderr():
             num_retries=1,
         )
 
+        # initially self.active_joint_indices references active joints that are controlled.
+        # we also make the assumption that the active index is the same across all parallel managed joints
+        self.active_ancestor_joint_idxs = [
+            (x.active_index[0]).cpu().item() for x in self.active_ancestor_joints
+        ]
+        self.controlled_joints_idx_in_qmask = [
+            self.active_ancestor_joint_idxs.index(idx)
+            for idx in self.active_joint_indices
+        ]
+
         self.qmask = torch.zeros(
             len(self.active_ancestor_joints), dtype=bool, device=self.device
         )
@@ -167,10 +226,15 @@ def compute_ik(
                 if pos_only:
                     jacobian = jacobian[:, 0:3]
 
+                # NOTE (arth): use only the parts of the jacobian that correspond to the active joints
+                jacobian = jacobian[:, :, self.qmask]
+
                 # NOTE (stao): this method of IK is from https://mathweb.ucsd.edu/~sbuss/ResearchWeb/ikmethods/iksurvey.pdf by Samuel R. Buss
                 delta_joint_pos = torch.linalg.pinv(jacobian) @ action.unsqueeze(-1)
-                return q0 + delta_joint_pos.squeeze(-1)
+
+                return q0[:, self.qmask] + delta_joint_pos.squeeze(-1)
         else:
+            q0 = q0[:, self.pmodel_active_joint_indices]
             result, success, error = self.pmodel.compute_inverse_kinematics(
                 self.end_link_idx,
                 target_pose.sp,
@@ -180,7 +244,7 @@ def compute_ik(
             )
             if success:
                 return common.to_tensor(
-                    [result[self.active_ancestor_joint_idxs]], device=self.device
+                    [result[self.pmodel_controlled_joint_indices]], device=self.device
                 )
             else:
                 return None

mani_skill/agents/robots/fetch/fetch.py

@@ -27,6 +27,7 @@
 class Fetch(BaseAgent):
     uid = "fetch"
     urdf_path = f"{PACKAGE_ASSET_DIR}/robots/fetch/fetch.urdf"
+    urdf_arm_ik_path = f"{PACKAGE_ASSET_DIR}/robots/fetch/fetch_torso_up.urdf"
     urdf_config = dict(
         _materials=dict(
             gripper=dict(static_friction=2.0, dynamic_friction=2.0, restitution=0.0)
@@ -149,7 +150,8 @@ def _controller_configs(self):
             damping=self.arm_damping,
             force_limit=self.arm_force_limit,
             ee_link=self.ee_link_name,
-            urdf_path=self.urdf_path,
+            urdf_path=self.urdf_arm_ik_path,
+            root_link_name="torso_lift_link",
         )
         arm_pd_ee_delta_pose = PDEEPoseControllerConfig(
             joint_names=self.arm_joint_names,
@@ -161,7 +163,8 @@ def _controller_configs(self):
             damping=self.arm_damping,
             force_limit=self.arm_force_limit,
             ee_link=self.ee_link_name,
-            urdf_path=self.urdf_path,
+            urdf_path=self.urdf_arm_ik_path,
+            root_link_name="torso_lift_link",
         )
 
         arm_pd_ee_target_delta_pos = deepcopy(arm_pd_ee_delta_pos)