Merge pull request #79 from haosulab/nonconvex_mesh

Add nonconvex mesh support in SapienPlanningWorld. Add SapienPlanner

mplib/examples/collision_avoidance.py

@@ -1,6 +1,7 @@
 import sapien.core as sapien
 
 from mplib.examples.demo_setup import DemoSetup
+from mplib.sapien_utils import SapienPlanner, SapienPlanningWorld
 
 
 class PlanningDemo(DemoSetup):
@@ -20,7 +21,6 @@ def __init__(self):
         super().__init__()
         self.setup_scene()
         self.load_robot()
-        self.setup_planner()
 
         # Set initial joint positions
         init_qpos = [0, 0.19, 0.0, -2.61, 0.0, 2.94, 0.78, 0, 0]
@@ -54,6 +54,9 @@ def __init__(self):
         blue_cube = builder.build(name="blue_cube")
         blue_cube.set_pose(sapien.Pose([0.55, 0, 0.1]))
 
+        planning_world = SapienPlanningWorld(self.scene, [self.robot])
+        self.planner = SapienPlanner(planning_world, "panda_hand")
+
     def add_point_cloud(self):
         """We tell the planner about the obstacle through a point cloud"""
         import trimesh

mplib/planner.py

@@ -27,37 +27,40 @@ def __init__(
         self,
         urdf: str,
         move_group: str,
+        *,
         srdf: str = "",
         package_keyword_replacement: str = "",
+        use_convex: bool = False,
         user_link_names: Sequence[str] = [],
         user_joint_names: Sequence[str] = [],
         joint_vel_limits: Optional[Sequence[float] | np.ndarray] = None,
         joint_acc_limits: Optional[Sequence[float] | np.ndarray] = None,
+        verbose: bool = False,
         **kwargs,
     ):
         # constructor ankor end
-        """Motion planner for robots.
+        """
+        Motion planner for robots.
 
-        Args:
-            urdf: Unified Robot Description Format file.
-            user_link_names: names of links, the order matters.
-                If empty, all links will be used.
-            user_joint_names: names of the joints to plan.
-                If empty, all active joints will be used.
-            move_group: target link to move, usually the end-effector.
-            joint_vel_limits: maximum joint velocities for time parameterization,
-                which should have the same length as
-            joint_acc_limits: maximum joint accelerations for time parameterization,
-                which should have the same length as
-            srdf: Semantic Robot Description Format file.
         References:
             http://docs.ros.org/en/kinetic/api/moveit_tutorials/html/doc/urdf_srdf/urdf_srdf_tutorial.html
 
+        :param urdf: Unified Robot Description Format file.
+        :param move_group: target link to move, usually the end-effector.
+        :param srdf: Semantic Robot Description Format file.
+        :param package_keyword_replacement: replace ``package://`` keyword in URDF
+        :param use_convex: if True, load collision mesh as convex mesh.
+            If mesh is not convex, a ``RuntimeError`` will be raised.
+        :param user_link_names: names of links, the order matters.
+            If empty, all links will be used.
+        :param user_joint_names: names of the joints to plan.
+            If empty, all active joints will be used.
+        :param joint_vel_limits: maximum joint velocities for time parameterization,
+            which should have the same length as ``self.move_group_joint_indices``
+        :param joint_acc_limits: maximum joint accelerations for time parameterization,
+            which should have the same length as ``self.move_group_joint_indices``
+        :param verbose: if True, print verbose logs for debugging
         """
-        if joint_vel_limits is None:
-            joint_vel_limits = []
-        if joint_acc_limits is None:
-            joint_acc_limits = []
         self.urdf = urdf
         self.srdf = srdf
         if self.srdf == "" and os.path.exists(urdf.replace(".urdf", ".srdf")):
@@ -70,10 +73,10 @@ def __init__(
         self.robot = ArticulatedModel(
             self.urdf,
             self.srdf,
-            link_names=user_link_names,
-            joint_names=user_joint_names,
-            convex=kwargs.get("convex", False),
-            verbose=False,
+            link_names=user_link_names,  # type: ignore
+            joint_names=user_joint_names,  # type: ignore
+            convex=use_convex,
+            verbose=verbose,
         )
         self.pinocchio_model = self.robot.get_pinocchio_model()
         self.user_link_names = self.pinocchio_model.get_link_names()
@@ -95,28 +98,24 @@ def __init__(
 
         self.joint_types = self.pinocchio_model.get_joint_types()
         self.joint_limits = np.concatenate(self.pinocchio_model.get_joint_limits())
-        self.joint_vel_limits = (
-            joint_vel_limits
-            if len(joint_vel_limits)
-            else np.ones(len(self.move_group_joint_indices))
-        )
-        self.joint_acc_limits = (
-            joint_acc_limits
-            if len(joint_acc_limits)
-            else np.ones(len(self.move_group_joint_indices))
-        )
+        if joint_vel_limits is None:
+            joint_vel_limits = np.ones(len(self.move_group_joint_indices))
+        if joint_acc_limits is None:
+            joint_acc_limits = np.ones(len(self.move_group_joint_indices))
+        self.joint_vel_limits = joint_vel_limits
+        self.joint_acc_limits = joint_acc_limits
         self.move_group_link_id = self.link_name_2_idx[self.move_group]
-        assert len(self.joint_vel_limits) == len(self.joint_acc_limits), (
-            f"length of joint_vel_limits ({len(self.joint_vel_limits)}) =/= "
-            f"length of joint_acc_limits ({len(self.joint_acc_limits)})"
-        )
-        assert len(self.joint_vel_limits) == len(self.move_group_joint_indices), (
-            f"length of joint_vel_limits ({len(self.joint_vel_limits)}) =/= "
-            f"length of move_group ({len(self.move_group_joint_indices)})"
-        )
-        assert len(self.joint_vel_limits) <= len(self.joint_limits), (
-            f"length of joint_vel_limits ({len(self.joint_vel_limits)}) > "
-            f"number of total joints ({len(self.joint_limits)})"
+
+        assert (
+            len(self.joint_vel_limits)
+            == len(self.joint_acc_limits)
+            == len(self.move_group_joint_indices)
+            <= len(self.joint_limits)
+        ), (
+            "length of joint_vel_limits, joint_acc_limits, and move_group_joint_indices"
+            " should equal and be <= number of total joints. "
+            f"{len(self.joint_vel_limits)} == {len(self.joint_acc_limits)} "
+            f"== {len(self.move_group_joint_indices)} <= {len(self.joint_limits)}"
         )
 
         # Mask for joints that have equivalent values (revolute joints with range > 2pi)
@@ -138,7 +137,9 @@ def __init__(
 
         self.planner = ompl.OMPLPlanner(world=self.planning_world)
 
-    def replace_package_keyword(self, package_keyword_replacement):
+    def replace_package_keyword(self, package_keyword_replacement: str):
+        # TODO(merge): fix file writing
+        # TODO(merge): convert to staticmethod
         """
         some ROS URDF files use package:// keyword to refer the package dir
         replace it with the given string (default is empty)
@@ -158,6 +159,7 @@ def replace_package_keyword(self, package_keyword_replacement):
         return rtn_urdf
 
     def generate_collision_pair(self, num_samples=100000):
+        # TODO(merge): convert to staticmethod
         """
         We read the srdf file to get the link pairs that should not collide.
         If not provided, we need to randomly sample configurations
@@ -367,7 +369,7 @@ def IK(
         self,
         goal_pose: np.ndarray,
         start_qpos: np.ndarray,
-        mask: Optional[list[bool] | np.ndarray] = None,
+        mask: Optional[Sequence[bool] | np.ndarray] = None,
         *,
         n_init_qpos: int = 20,
         threshold: float = 1e-3,
@@ -404,8 +406,11 @@ def IK(
         q_goals = []
         qpos = start_qpos
         for _ in range(n_init_qpos):
-            ik_qpos, ik_success, ik_error = self.pinocchio_model.compute_IK_CLIK(
-                move_link_idx, goal_pose, qpos, mask
+            ik_qpos, ik_success, _ = self.pinocchio_model.compute_IK_CLIK(
+                move_link_idx,
+                goal_pose,
+                qpos,
+                mask,  # type: ignore
             )
             success = ik_success and self.wrap_joint_limit(ik_qpos)
 
@@ -497,11 +502,11 @@ def TOPP(self, path, step=0.1, verbose=False):
         jnt_traj = instance.compute_trajectory()
         if jnt_traj is None:
             raise RuntimeError("Fail to parameterize path")
-        ts_sample = np.linspace(0, jnt_traj.duration, int(jnt_traj.duration / step))
+        ts_sample = np.linspace(0, jnt_traj.duration, int(jnt_traj.duration / step))  # type: ignore
         qs_sample = jnt_traj(ts_sample)
         qds_sample = jnt_traj(ts_sample, 1)
         qdds_sample = jnt_traj(ts_sample, 2)
-        return ts_sample, qs_sample, qds_sample, qdds_sample, jnt_traj.duration
+        return ts_sample, qs_sample, qds_sample, qdds_sample, jnt_traj.duration  # type: ignore
 
     # TODO: change method name to align with PlanningWorld API?
     def update_point_cloud(self, points, resolution=1e-3, name="scene_pcd"):
@@ -646,8 +651,8 @@ def plan_qpos_to_qpos(
         fix_joint_limits: bool = True,
         fixed_joint_indices: Optional[list[int]] = None,
         simplify: bool = True,
-        constraint_function: Optional[Callable] = None,
-        constraint_jacobian: Optional[Callable] = None,
+        constraint_function: Optional[Callable[[np.ndarray, np.ndarray], None]] = None,
+        constraint_jacobian: Optional[Callable[[np.ndarray, np.ndarray], None]] = None,
         constraint_tolerance: float = 1e-3,
         verbose: bool = False,
     ) -> dict[str, str | np.ndarray | np.float64]:
@@ -705,8 +710,8 @@ def plan_qpos_to_qpos(
             range=rrt_range,
             fixed_joints=fixed_joints,
             simplify=simplify,
-            constraint_function=constraint_function,
-            constraint_jacobian=constraint_jacobian,
+            constraint_function=constraint_function,  # type: ignore
+            constraint_jacobian=constraint_jacobian,  # type: ignore
             constraint_tolerance=constraint_tolerance,
             verbose=verbose,
         )
@@ -791,14 +796,15 @@ def plan_qpos_to_pose(
         # we need to take only the move_group joints when planning
         # idx = self.move_group_joint_indices
 
+        # TODO(merge): verify this
         ik_status, goal_qpos = self.IK(goal_pose, current_qpos, mask)
         if ik_status != "Success":
             return {"status": ik_status}
 
         if verbose:
             print("IK results:")
-            for i in range(len(goal_qpos)):
-                print(goal_qpos[i])
+            for i in range(len(goal_qpos)):  # type: ignore
+                print(goal_qpos[i])  # type: ignore
 
         # goal_qpos_ = [goal_qpos[i][move_joint_idx] for i in range(len(goal_qpos))]
         self.robot.set_qpos(current_qpos, True)
@@ -809,11 +815,11 @@ def plan_qpos_to_pose(
 
         if verbose:
             print("IK results:")
-            for i in range(len(goal_qpos)):
-                print(goal_qpos[i])
+            for i in range(len(goal_qpos)):  # type: ignore
+                print(goal_qpos[i])  # type: ignore
 
         return self.plan_qpos_to_qpos(
-            goal_qpos,
+            goal_qpos,  # type: ignore
             current_qpos,
             time_step=time_step,
             rrt_range=rrt_range,

mplib/pymp/collision_detection/fcl.pyi

@@ -50,7 +50,7 @@ class BVHModel(CollisionGeometry):
     """
     def __init__(self) -> None: ...
     @typing.overload
-    def addSubModel(
+    def add_sub_model(
         self,
         vertices: list[
             numpy.ndarray[
@@ -64,7 +64,7 @@ class BVHModel(CollisionGeometry):
         :param vertices: vertices of the sub-model
         """
     @typing.overload
-    def addSubModel(
+    def add_sub_model(
         self,
         vertices: list[
             numpy.ndarray[
@@ -80,7 +80,7 @@ class BVHModel(CollisionGeometry):
         :param faces: faces of the sub-model represented by a list of vertex indices
         """
     @typing.overload
-    def addSubModel(
+    def add_sub_model(
         self,
         vertices: list[
             numpy.ndarray[
@@ -99,14 +99,14 @@ class BVHModel(CollisionGeometry):
         :param vertices: vertices of the sub-model
         :param faces: faces of the sub-model represented by a list of vertex indices
         """
-    def beginModel(self, num_faces: int = 0, num_vertices: int = 0) -> int:
+    def begin_model(self, num_faces: int = 0, num_vertices: int = 0) -> int:
         """
         Begin to construct a BVHModel.
 
         :param num_faces: number of faces of the mesh
         :param num_vertices: number of vertices of the mesh
         """
-    def endModel(self) -> int:
+    def end_model(self) -> int:
         """
         End the construction of a BVHModel.
         """
@@ -195,26 +195,26 @@ class CollisionGeometry:
     ]
     aabb_radius: float
     cost_density: float
-    def computeCOM(
+    def compute_com(
         self,
     ) -> numpy.ndarray[
         tuple[typing.Literal[3], typing.Literal[1]], numpy.dtype[numpy.float64]
     ]: ...
-    def computeLocalAABB(self) -> None: ...
-    def computeMomentofInertia(
+    def compute_local_aabb(self) -> None: ...
+    def compute_moment_of_inertia(
         self,
     ) -> numpy.ndarray[
         tuple[typing.Literal[3], typing.Literal[3]], numpy.dtype[numpy.float64]
     ]: ...
-    def computeMomentofInertiaRelatedToCOM(
+    def compute_moment_of_inertia_related_to_com(
         self,
     ) -> numpy.ndarray[
         tuple[typing.Literal[3], typing.Literal[3]], numpy.dtype[numpy.float64]
     ]: ...
-    def computeVolume(self) -> float: ...
-    def isFree(self) -> bool: ...
-    def isOccupied(self) -> bool: ...
-    def isUncertain(self) -> bool: ...
+    def compute_volume(self) -> float: ...
+    def is_free(self) -> bool: ...
+    def is_occupied(self) -> bool: ...
+    def is_uncertain(self) -> bool: ...
 
 class CollisionObject:
     """
@@ -269,7 +269,7 @@ class CollisionRequest:
         gjk_solver_type: GJKSolverType = ...,
         gjk_tolerance: float = 1e-06,
     ) -> None: ...
-    def isSatisfied(self, result: ...) -> bool: ...
+    def is_satisfied(self, result: ...) -> bool: ...
 
 class CollisionResult:
     def __init__(self) -> None: ...
@@ -503,7 +503,7 @@ class DistanceRequest:
         distance_tolerance: float = 1e-06,
         gjk_solver_type: GJKSolverType = ...,
     ) -> None: ...
-    def isSatisfied(self, result: ...) -> bool: ...
+    def is_satisfied(self, result: ...) -> bool: ...
 
 class DistanceResult:
     def __init__(self, min_distance: float = 1.7976931348623157e308) -> None: ...

mplib/sapien_utils/__init__.py

@@ -1 +1 @@
-from .conversion import SapienPlanningWorld
+from .conversion import SapienPlanner, SapienPlanningWorld