[FEATURE] Add support of 'capsule' primitive in URDF file. (#2045)

Author: duburcqa

genesis/ext/urdfpy/__init__.py

@@ -2,6 +2,7 @@
     URDFType,
     Box,
     Cylinder,
+    Capsule,
     Sphere,
     Mesh,
     Geometry,
@@ -29,6 +30,7 @@
     "URDFType",
     "Box",
     "Cylinder",
+    "Capsule",
     "Sphere",
     "Mesh",
     "Geometry",

genesis/ext/urdfpy/urdf.py

@@ -441,6 +441,87 @@ def copy(self, prefix="", scale=None):
         return c
 
 
+class Capsule(URDFType):
+    """A capsule whose center is at the local origin.
+
+    Parameters
+    ----------
+    radius : float
+        The radius of the capsule in meters.
+    length : float
+        The length of the capsule in meters.
+    """
+
+    _ATTRIBS = {
+        "radius": (float, True),
+        "length": (float, True),
+    }
+    _TAG = "capsule"
+
+    def __init__(self, radius, length):
+        self.radius = radius
+        self.length = length
+        self._meshes = []
+
+    @property
+    def radius(self):
+        """float : The radius of the capsule in meters."""
+        return self._radius
+
+    @radius.setter
+    def radius(self, value):
+        self._radius = float(value)
+        self._meshes = []
+
+    @property
+    def length(self):
+        """float : The length of the capsule in meters."""
+        return self._length
+
+    @length.setter
+    def length(self, value):
+        self._length = float(value)
+        self._meshes = []
+
+    @property
+    def meshes(self):
+        """list of :class:`~trimesh.base.Trimesh` : The triangular meshes
+        that represent this object.
+        """
+        if len(self._meshes) == 0:
+            self._meshes = [trimesh.creation.capsule(radius=self.radius, height=self.length)]
+        return self._meshes
+
+    def copy(self, prefix="", scale=None):
+        """Create a deep copy with the prefix applied to all names.
+
+        Parameters
+        ----------
+        prefix : str
+            A prefix to apply to all names.
+
+        Returns
+        -------
+        :class:`.Capsule`
+            A deep copy.
+        """
+        if scale is None:
+            scale = 1.0
+        if isinstance(scale, (list, np.ndarray)):
+            if scale[0] != scale[1]:
+                raise ValueError("Cannot rescale capsule geometry with asymmetry in x/y")
+            c = Capsule(
+                radius=self.radius * scale[0],
+                length=self.length * scale[2],
+            )
+        else:
+            c = Capsule(
+                radius=self.radius * scale,
+                length=self.length * scale,
+            )
+        return c
+
+
 class Sphere(URDFType):
     """A sphere whose center is at the local origin.
 
@@ -657,6 +738,8 @@ class Geometry(URDFType):
         Box geometry.
     cylinder : :class:`.Cylinder`
         Cylindrical geometry.
+    capsule : :class:`.Capsule`
+        Capsule geometry.
     sphere : :class:`.Sphere`
         Spherical geometry.
     mesh : :class:`.Mesh`
@@ -666,16 +749,18 @@ class Geometry(URDFType):
     _ELEMENTS = {
         "box": (Box, False, False),
         "cylinder": (Cylinder, False, False),
+        "capsule": (Capsule, False, False),
         "sphere": (Sphere, False, False),
         "mesh": (Mesh, False, False),
     }
     _TAG = "geometry"
 
-    def __init__(self, box=None, cylinder=None, sphere=None, mesh=None):
-        if box is None and cylinder is None and sphere is None and mesh is None:
+    def __init__(self, box=None, cylinder=None, capsule=None, sphere=None, mesh=None):
+        if box is None and cylinder is None and capsule is None and sphere is None and mesh is None:
             raise ValueError("At least one geometry element must be set")
         self.box = box
         self.cylinder = cylinder
+        self.capsule = capsule
         self.sphere = sphere
         self.mesh = mesh
 
@@ -701,6 +786,17 @@ def cylinder(self, value):
             raise TypeError("Expected Cylinder type")
         self._cylinder = value
 
+    @property
+    def capsule(self):
+        """:class:`.Capsule` : Capsule geometry."""
+        return self._capsule
+
+    @capsule.setter
+    def capsule(self, value):
+        if value is not None and not isinstance(value, Capsule):
+            raise TypeError("Expected Capsule type")
+        self._capsule = value
+
     @property
     def sphere(self):
         """:class:`.Sphere` : Spherical geometry."""
@@ -732,6 +828,8 @@ def geometry(self):
             return self.box
         if self.cylinder is not None:
             return self.cylinder
+        if self.capsule is not None:
+            return self.capsule
         if self.sphere is not None:
             return self.sphere
         if self.mesh is not None:
@@ -761,6 +859,7 @@ def copy(self, prefix="", scale=None):
         v = Geometry(
             box=(self.box.copy(prefix=prefix, scale=scale) if self.box else None),
             cylinder=(self.cylinder.copy(prefix=prefix, scale=scale) if self.cylinder else None),
+            capsule=(self.capsule.copy(prefix=prefix, scale=scale) if self.capsule else None),
             sphere=(self.sphere.copy(prefix=prefix, scale=scale) if self.sphere else None),
             mesh=(self.mesh.copy(prefix=prefix, scale=scale) if self.mesh else None),
         )

genesis/utils/urdf.py

@@ -138,17 +138,23 @@ def parse_urdf(morph, surface):
                     g_info = {"mesh" if geom_is_col else "vmesh": mesh}
                     link_g_infos_.append(g_info)
             else:
-                # Each geometry primitive is one RigidGeom in genesis.
+                # Each geometry primitive is one RigidGeom in genesis
                 if isinstance(geom.geometry.geometry, urdfpy.Box):
                     tmesh = trimesh.creation.box(extents=geom.geometry.geometry.size)
                     geom_type = gs.GEOM_TYPE.BOX
                     geom_data = np.array(geom.geometry.geometry.size)
+                elif isinstance(geom.geometry.geometry, urdfpy.Capsule):
+                    tmesh = trimesh.creation.capsule(
+                        radius=geom.geometry.geometry.radius, height=geom.geometry.geometry.length
+                    )
+                    geom_type = gs.GEOM_TYPE.CAPSULE
+                    geom_data = np.array([geom.geometry.geometry.radius, geom.geometry.geometry.length])
                 elif isinstance(geom.geometry.geometry, urdfpy.Cylinder):
                     tmesh = trimesh.creation.cylinder(
                         radius=geom.geometry.geometry.radius, height=geom.geometry.geometry.length
                     )
                     geom_type = gs.GEOM_TYPE.CYLINDER
-                    geom_data = None
+                    geom_data = np.array([geom.geometry.geometry.radius, geom.geometry.geometry.length])
                 elif isinstance(geom.geometry.geometry, urdfpy.Sphere):
                     if geom_is_col:
                         tmesh = trimesh.creation.icosphere(radius=geom.geometry.geometry.radius, subdivisions=2)

tests/test_rigid_physics.py

@@ -2466,6 +2466,52 @@ def _check_entity_positions(relative, tol):
     _check_entity_positions(relative=True, tol=2e-3)
 
 
+@pytest.mark.required
+def test_urdf_capsule(tmp_path, show_viewer, tol):
+    urdf_path = tmp_path / "capsule.urdf"
+    with open(urdf_path, "w") as f:
+        f.write(
+            """
+            <robot name="urdf_robot">
+                <link name="base_link">
+                    <inertial>
+                        <origin rpy="0 0 0" xyz="0 0 0"/>
+                        <mass value=".1"/>
+                        <inertia ixx="1" ixy="0" ixz="0" iyy="1" iyz="0" izz="1"/>
+                    </inertial>
+                    <collision>
+                        <origin rpy="0 0 0" xyz="0 0 0"/>
+                        <geometry>
+                            <capsule length="0.1" radius="0.02"/>
+                        </geometry>
+                    </collision>
+                </link>
+            </robot>
+            """
+        )
+
+    scene = gs.Scene(show_viewer=show_viewer)
+    scene.add_entity(gs.morphs.Plane())
+    robot = scene.add_entity(
+        gs.morphs.URDF(
+            file=urdf_path,
+            pos=(0.0, 0.0, 0.3),
+        ),
+        vis_mode="collision",
+    )
+    scene.build()
+
+    (geom,) = robot.geoms
+    assert geom.type == gs.GEOM_TYPE.CAPSULE
+    assert_allclose(geom.data[:2], (0.02, 0.1), tol=gs.EPS)
+
+    for _ in range(40):
+        scene.step()
+    geom_verts = tensor_to_array(geom.get_verts())
+    assert np.linalg.norm(geom_verts - np.zeros(3), axis=-1, ord=np.inf).min() < 1e-3
+    assert np.linalg.norm(geom_verts - np.array((0.0, 0.0, 0.14)), axis=-1, ord=np.inf).min() < 1e-3
+
+
 @pytest.mark.required
 def test_mjcf_parsing_with_include():
     scene = gs.Scene()