Support multi-shape Physx components in SapienPlanningWorld conversion

include/mplib/core/articulated_model.h

@@ -66,6 +66,7 @@ class ArticulatedModelTpl {
    * @param verbose: print debug information. Default: ``false``.
    * @return: a unique_ptr to ArticulatedModel
    */
+  // TODO(merge): remove pair (link name is not needed)
   static std::unique_ptr<ArticulatedModelTpl<S>> createFromURDFString(
       const std::string &urdf_string, const std::string &srdf_string,
       const std::vector<std::pair<std::string, collision_detection::FCLObjectPtr<S>>>

include/mplib/core/attached_body.h

@@ -51,10 +51,10 @@ class AttachedBodyTpl {
   /// @brief Gets the attached object (``FCLObjectPtr``)
   FCLObjectPtr getObject() const { return object_; }
 
-  /// @brief Gets the articulation this body is attached to
+  /// @brief Gets the articulation that this body is attached to
   ArticulatedModelPtr getAttachedArticulation() const { return attached_articulation_; }
 
-  /// @brief Gets the articulation this body is attached to
+  /// @brief Gets the articulation link id that this body is attached to
   int getAttachedLinkId() const { return attached_link_id_; }
 
   /// @brief Gets the attached pose (relative pose from attached link to object)
@@ -63,11 +63,14 @@ class AttachedBodyTpl {
   /// @brief Sets the attached pose (relative pose from attached link to object)
   void setPose(const Isometry3<S> &pose) { pose_ = pose; }
 
-  /// @brief Gets the global pose of the attached object
-  Isometry3<S> getGlobalPose() const {
-    return pinocchio_model_->getLinkPose(attached_link_id_) * pose_;
+  /// @brief Gets the global pose of the articulation link that this body is attached to
+  Isometry3<S> getAttachedLinkGlobalPose() const {
+    return pinocchio_model_->getLinkPose(attached_link_id_);
   }
 
+  /// @brief Gets the global pose of the attached object
+  Isometry3<S> getGlobalPose() const { return getAttachedLinkGlobalPose() * pose_; }
+
   /// @brief Updates the global pose of the attached object using current state
   void updatePose() const;
 

mplib/pymp/__init__.pyi

@@ -266,11 +266,15 @@ class AttachedBody:
         """
     def get_attached_articulation(self) -> ArticulatedModel:
         """
-        Gets the articulation this body is attached to
+        Gets the articulation that this body is attached to
+        """
+    def get_attached_link_global_pose(self) -> Pose:
+        """
+        Gets the global pose of the articulation link that this body is attached to
         """
     def get_attached_link_id(self) -> int:
         """
-        Gets the articulation this body is attached to
+        Gets the articulation link id that this body is attached to
         """
     def get_global_pose(self) -> Pose:
         """

mplib/sapien_utils/conversion.py

@@ -3,7 +3,7 @@
 from typing import Optional, Sequence
 
 import numpy as np
-from sapien import Entity, Pose, Scene
+from sapien import Entity, Scene
 from sapien.physx import (
     PhysxArticulation,
     PhysxArticulationLinkComponent,
@@ -19,14 +19,15 @@
 from transforms3d.euler import euler2quat
 
 from ..planner import Planner
-from ..pymp import ArticulatedModel, PlanningWorld
+from ..pymp import ArticulatedModel, PlanningWorld, Pose
 from ..pymp.collision_detection.fcl import (
     Box,
     BVHModel,
     Capsule,
     CollisionObject,
     Convex,
     Cylinder,
+    FCLObject,
     Halfspace,
     Sphere,
     collide,
@@ -39,7 +40,9 @@
 
 class SapienPlanningWorld(PlanningWorld):
     def __init__(
-        self, sim_scene: Scene, planned_articulations: list[PhysxArticulation] = []
+        self,
+        sim_scene: Scene,
+        planned_articulations: list[PhysxArticulation] = [],  # noqa: B006
     ):
         """
         Creates an mplib.pymp.planning_world.PlanningWorld from a sapien.Scene.
@@ -56,12 +59,12 @@ def __init__(
             urdf_str = export_kinematic_chain_urdf(articulation)
             srdf_str = export_srdf(articulation)
 
-            # Get all links with collision shapes at local_pose
-            collision_links = []  # [(link_name, [CollisionObject, ...]), ...]
-            for link in articulation.links:
-                col_objs = self.convert_sapien_col_shape(link)
-                if len(col_objs) > 0:
-                    collision_links.append((link.name, col_objs))
+            # Convert all links to FCLObject
+            collision_links = [
+                (link.name, fcl_obj)
+                for link in articulation.links
+                if (fcl_obj := self.convert_physx_component(link)) is not None
+            ]
 
             articulated_model = ArticulatedModel.create_from_urdf_string(
                 urdf_str,
@@ -81,76 +84,62 @@ def __init__(
 
         for entity in actors:
             component = entity.find_component_by_type(PhysxRigidBaseComponent)
-            assert (
-                component is not None
-            ), f"No PhysxRigidBaseComponent found in {entity.name}: {entity.components=}"
-            assert not isinstance(
-                component, PhysxArticulationLinkComponent
-            ), f"Component should not be PhysxArticulationLinkComponent: {component=}"
+            assert component is not None, (
+                f"No PhysxRigidBaseComponent found in {entity.name}: "
+                f"{entity.components=}"
+            )
 
             # Convert collision shapes at current global pose
-            col_objs = self.convert_sapien_col_shape(component)  # type: ignore
-            # TODO: multiple collision shapes
-            assert len(col_objs) == 1, (
-                f"Should only have 1 collision object, got {len(col_objs)} shapes for "
-                f"entity '{entity.name}'"
+            self.add_normal_object(
+                self.get_object_name(entity),
+                self.convert_physx_component(component),  # type: ignore
             )
 
-            self.add_normal_object(self.get_object_name(entity), col_objs[0])
-
     def update_from_simulation(self, *, update_attached_object: bool = True) -> None:
         """
-        Updates planning_world articulations/objects pose with current Scene state
+        Updates PlanningWorld's articulations/objects pose with current Scene state.
+        Note that shape's local_pose is not updated.
+        If those are changed, please recreate a new SapienPlanningWorld instance.
 
         :param update_attached_object: whether to update the attached pose of
             all attached objects
         """
         for articulation in self._sim_scene.get_all_articulations():
-            # set_qpos to update poses
-            self.get_articulation(self.get_object_name(articulation)).set_qpos(
-                articulation.qpos  # type: ignore
-            )
+            if art := self.get_articulation(self.get_object_name(articulation)):
+                # set_qpos to update poses
+                art.set_qpos(articulation.qpos)  # type: ignore
+            else:
+                raise RuntimeError(
+                    f"Articulation {articulation.name} not found in PlanningWorld! "
+                    "The scene might have changed since last update."
+                )
 
         for entity in self._sim_scene.get_all_actors():
-            component = entity.find_component_by_type(PhysxRigidBaseComponent)
-            assert (
-                component is not None
-            ), f"No PhysxRigidBaseComponent found in {entity.name}: {entity.components=}"
-            assert not isinstance(
-                component, PhysxArticulationLinkComponent
-            ), f"Component should not be PhysxArticulationLinkComponent: {component=}"
-
-            shapes = component.collision_shapes  # type: ignore
-            # TODO: multiple collision shapes
-            assert len(shapes) == 1, (
-                f"Should only have 1 collision shape, got {len(shapes)} shapes for "
-                f"entity '{entity.name}'"
-            )
-            shape = shapes[0]
-
-            pose: Pose = entity.pose * shape.local_pose
-            # NOTE: Convert poses for Capsule/Cylinder
-            if isinstance(
-                shape, (PhysxCollisionShapeCapsule, PhysxCollisionShapeCylinder)
-            ):
-                pose = pose * Pose(q=euler2quat(0, np.pi / 2, 0))  # type: ignore
-
-            # handle attached object
-            if self.is_normal_object_attached(self.get_object_name(entity)):
-                attached_body = self.get_attached_object(self.get_object_name(entity))
-                if update_attached_object:
-                    parent_posevec = (
-                        attached_body.get_attached_articulation()
-                        .get_pinocchio_model()
-                        .get_link_pose(attached_body.get_attached_link_id())
+            object_name = self.get_object_name(entity)
+
+            # If entity is an attached object
+            if attached_body := self.get_attached_object(object_name):
+                if update_attached_object:  # update attached pose
+                    attached_body.pose = (
+                        attached_body.get_attached_link_global_pose().inv()
+                        * entity.pose  # type: ignore
                     )
-                    parent_pose = Pose(parent_posevec[:3], parent_posevec[3:])  # type: ignore
-                    pose = parent_pose.inv() * pose  # new attached pose
-                    attached_body.set_pose(np.hstack((pose.p, pose.q)))
                 attached_body.update_pose()
+            elif fcl_obj := self.get_normal_object(object_name):
+                # Overwrite the object
+                self.add_normal_object(
+                    object_name,
+                    FCLObject(
+                        object_name,
+                        entity.pose,  # type: ignore
+                        fcl_obj.shapes,
+                        fcl_obj.shape_poses,
+                    ),
+                )
             else:
-                self.get_normal_object(self.get_object_name(entity)).set_transformation(
-                    np.hstack((pose.p, pose.q))
+                raise RuntimeError(
+                    f"Entity {entity.name} not found in PlanningWorld! "
+                    "The scene might have changed since last update."
                 )
 
     @staticmethod
@@ -170,65 +159,64 @@ def get_object_name(obj: PhysxArticulation | Entity) -> str:
             raise NotImplementedError(f"Unknown SAPIEN object type {type(obj)}")
 
     @staticmethod
-    def convert_sapien_col_shape(
-        component: PhysxRigidBaseComponent,
-    ) -> list[CollisionObject]:
-        """Converts a SAPIEN physx.PhysxRigidBaseComponent to an FCL CollisionObject
-        Returns a list of collision_obj at their current poses.
-
-        If the component is an articulation link, the returned collision_obj is at
-        the shape's local_pose.
-        Otherwise, the returned collision_obj is at the entity's global pose
+    def convert_physx_component(comp: PhysxRigidBaseComponent) -> FCLObject | None:
         """
-        shapes = component.collision_shapes
-        if len(shapes) == 0:
-            return []
+        Converts a SAPIEN physx.PhysxRigidBaseComponent to an FCLObject.
+        All shapes in the returned FCLObject are already set at their world poses.
 
-        # NOTE: MPlib currently only supports 1 collision shape per object
-        # TODO: multiple collision shapes
-        assert len(shapes) == 1, (
-            f"Should only have 1 collision shape, got {len(shapes)} shapes for "
-            f"entity '{component.entity.name}'"
-        )
+        :param comp: a SAPIEN physx.PhysxRigidBaseComponent.
+        :return: an FCLObject containing all collision shapes in the Physx component.
+            If the component has no collision shapes, return ``None``.
+        """
+        shapes: list[CollisionObject] = []
+        shape_poses: list[Pose] = []
+        for shape in comp.collision_shapes:
+            shape_poses.append(shape.local_pose)  # type: ignore
+
+            if isinstance(shape, PhysxCollisionShapeBox):
+                c_geom = Box(side=shape.half_size * 2)
+            elif isinstance(shape, PhysxCollisionShapeCapsule):
+                c_geom = Capsule(radius=shape.radius, lz=shape.half_length * 2)
+                # NOTE: physx Capsule has x-axis along capsule height
+                # FCL Capsule has z-axis along capsule height
+                shape_poses[-1] *= Pose(q=euler2quat(0, np.pi / 2, 0))
+            elif isinstance(shape, PhysxCollisionShapeConvexMesh):
+                assert np.allclose(
+                    shape.scale, 1.0
+                ), f"Not unit scale {shape.scale}, need to rescale vertices?"
+                c_geom = Convex(vertices=shape.vertices, faces=shape.triangles)
+            elif isinstance(shape, PhysxCollisionShapeCylinder):
+                c_geom = Cylinder(radius=shape.radius, lz=shape.half_length * 2)
+                # NOTE: physx Cylinder has x-axis along cylinder height
+                # FCL Cylinder has z-axis along cylinder height
+                shape_poses[-1] *= Pose(q=euler2quat(0, np.pi / 2, 0))
+            elif isinstance(shape, PhysxCollisionShapePlane):
+                raise NotImplementedError(
+                    "Support for Plane collision is not implemented yet in mplib, "
+                    "need fcl::Plane"
+                )
+            elif isinstance(shape, PhysxCollisionShapeSphere):
+                c_geom = Sphere(radius=shape.radius)
+            elif isinstance(shape, PhysxCollisionShapeTriangleMesh):
+                c_geom = BVHModel()
+                c_geom.begin_model()
+                c_geom.add_sub_model(vertices=shape.vertices, faces=shape.triangles)  # type: ignore
+                c_geom.end_model()
+            else:
+                raise TypeError(f"Unknown shape type: {type(shape)}")
+            shapes.append(CollisionObject(c_geom))
 
-        shape = shapes[0]
-        if isinstance(component, PhysxArticulationLinkComponent):  # articulation link
-            pose = shape.local_pose
-        else:
-            pose = component.entity.pose * shape.local_pose
-
-        if isinstance(shape, PhysxCollisionShapeBox):
-            collision_geom = Box(side=shape.half_size * 2)
-        elif isinstance(shape, PhysxCollisionShapeCapsule):
-            collision_geom = Capsule(radius=shape.radius, lz=shape.half_length * 2)
-            # NOTE: physx Capsule has x-axis along capsule height
-            # FCL Capsule has z-axis along capsule height
-            pose = pose * Pose(q=euler2quat(0, np.pi / 2, 0))  # type: ignore
-        elif isinstance(shape, PhysxCollisionShapeConvexMesh):
-            assert np.allclose(
-                shape.scale, 1.0
-            ), f"Not unit scale {shape.scale}, need to rescale vertices?"
-            collision_geom = Convex(vertices=shape.vertices, faces=shape.triangles)
-        elif isinstance(shape, PhysxCollisionShapeCylinder):
-            collision_geom = Cylinder(radius=shape.radius, lz=shape.half_length * 2)
-            # NOTE: physx Cylinder has x-axis along cylinder height
-            # FCL Cylinder has z-axis along cylinder height
-            pose = pose * Pose(q=euler2quat(0, np.pi / 2, 0))  # type: ignore
-        elif isinstance(shape, PhysxCollisionShapePlane):
-            raise NotImplementedError(
-                "Support for Plane collision is not implemented yet in mplib, "
-                "need fcl::Plane"
-            )
-        elif isinstance(shape, PhysxCollisionShapeSphere):
-            collision_geom = Sphere(radius=shape.radius)
-        elif isinstance(shape, PhysxCollisionShapeTriangleMesh):
-            collision_geom = BVHModel()
-            collision_geom.begin_model()
-            collision_geom.add_sub_model(vertices=shape.vertices, faces=shape.triangles)
-            collision_geom.end_model()
-        else:
-            raise TypeError(f"Unknown shape type: {type(shape)}")
-        return [CollisionObject(collision_geom, pose.p, pose.q)]  # type: ignore
+        if len(shapes) == 0:
+            return None
+
+        return FCLObject(
+            comp.name
+            if isinstance(comp, PhysxArticulationLinkComponent)
+            else SapienPlanningWorld.get_object_name(comp.entity),
+            comp.entity.pose,  # type: ignore
+            shapes,
+            shape_poses,
+        )
 
 
 class SapienPlanner(Planner):
@@ -303,7 +291,10 @@ def __init__(
 
     def update_from_simulation(self, *, update_attached_object: bool = True) -> None:
         """
-        Updates planning_world articulations/objects pose with current Scene state.
+        Updates PlanningWorld's articulations/objects pose with current Scene state.
+        Note that shape's local_pose is not updated.
+        If those are changed, please recreate a new SapienPlanningWorld instance.
+
         Directly calls ``SapienPlanningWorld.update_from_simulation()``
 
         :param update_attached_object: whether to update the attached pose of

pybind/core/pybind_attached_body.cpp

@@ -53,6 +53,12 @@ void build_pyattached_body(py::module &pymp) {
             body.setPose(pose.toIsometry());
           },
           py::arg("pose"), DOC(mplib, AttachedBodyTpl, setPose))
+      .def(
+          "get_attached_link_global_pose",
+          [](const AttachedBody &body) {
+            return Pose<S>(body.getAttachedLinkGlobalPose());
+          },
+          DOC(mplib, AttachedBodyTpl, getAttachedLinkGlobalPose))
       .def(
           "get_global_pose",
           [](const AttachedBody &body) { return Pose<S>(body.getGlobalPose()); },

pybind/docstring/core/attached_body.h

@@ -45,11 +45,15 @@ Construct an attached body for a specified link.
 
 static const char *__doc_mplib_AttachedBodyTpl_getAttachedArticulation =
 R"doc(
-Gets the articulation this body is attached to)doc";
+Gets the articulation that this body is attached to)doc";
+
+static const char *__doc_mplib_AttachedBodyTpl_getAttachedLinkGlobalPose =
+R"doc(
+Gets the global pose of the articulation link that this body is attached to)doc";
 
 static const char *__doc_mplib_AttachedBodyTpl_getAttachedLinkId =
 R"doc(
-Gets the articulation this body is attached to)doc";
+Gets the articulation link id that this body is attached to)doc";
 
 static const char *__doc_mplib_AttachedBodyTpl_getGlobalPose =
 R"doc(