[FEATURE] Add support of USD import for Rigid Body.

examples/usd/import_stage.py

@@ -0,0 +1,33 @@
+import genesis as gs
+
+gs.init(backend=gs.cpu)
+
+scene = gs.Scene(
+    viewer_options=gs.options.ViewerOptions(
+        camera_pos=(3.5, 0.0, 2.5),
+        camera_lookat=(0.0, 0.0, 0.5),
+        camera_fov=40,
+        enable_interaction=True,
+    ),
+    rigid_options=gs.options.RigidOptions(
+        # constraint_solver=gs.constraint_solver.Newton,
+        gravity=(0, 0, -9.8),
+        enable_collision=True,
+        enable_joint_limit=True,
+        max_collision_pairs=1000,
+    ),
+    show_viewer=True,
+)
+
+AssetRoot = "D:/Assets"
+
+# load a stage from USD file
+entities = scene.add_stage(f"{AssetRoot}/Lightwheel_Kitchen001/Kitchen001/Kitchen001.usd")
+
+# Build the scene
+scene.build()
+
+# Run the simulation for visualization
+while scene.viewer.is_alive():
+    scene.step()
+    pass

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -138,9 +138,20 @@ def _load_model(self):
         self._joints = gs.List()
         self._equalities = gs.List()
 
-        if isinstance(self._morph, gs.morphs.Mesh):
-            self._load_mesh(self._morph, self._surface)
-        elif isinstance(self._morph, (gs.morphs.MJCF, gs.morphs.URDF, gs.morphs.Drone)):
+        if isinstance(self._morph, (gs.morphs.Mesh, gs.morphs.USDRigidBody)):
+            if isinstance(self._morph, gs.morphs.Mesh):
+                self._load_mesh(self._morph, self._surface)
+            else:
+                self._load_usd_rigid_body(self._morph, self._surface)
+        elif isinstance(
+            self._morph,
+            (
+                gs.morphs.MJCF,
+                gs.morphs.URDF,
+                gs.morphs.Drone,
+                gs.morphs.USDArticulation,
+            ),
+        ):
             self._load_scene(self._morph, self._surface)
         elif isinstance(self._morph, gs.morphs.Primitive):
             self._load_primitive(self._morph, self._surface)
@@ -322,6 +333,24 @@ def _load_mesh(self, morph, surface):
             surface=surface,
         )
 
+    def _load_usd_rigid_body(self, morph, surface):
+        """
+        Load a USD rigid body, similar to _load_mesh but parsing from USD file.
+        """
+        from genesis.utils.usd import parse_usd_rigid_body
+
+        # Parse USD rigid body to get l_info, j_infos, and g_infos
+        l_info, j_infos, g_infos = parse_usd_rigid_body(morph, surface)
+
+        # Create link and joint using _add_by_info
+        link, (joint,) = self._add_by_info(
+            l_info=l_info,
+            j_infos=j_infos,
+            g_infos=g_infos,
+            morph=morph,
+            surface=surface,
+        )
+
     def _load_terrain(self, morph, surface):
         vmesh, mesh, self.terrain_hf = tu.parse_terrain(morph, surface)
         self.terrain_scale = np.array((morph.horizontal_scale, morph.vertical_scale), dtype=gs.np_float)
@@ -372,60 +401,54 @@ def _load_terrain(self, morph, surface):
             surface=surface,
         )
 
-    def _load_scene(self, morph, surface):
-        # Mujoco's unified MJCF+URDF parser is not good enough for now to be used for loading both MJCF and URDF files.
-        # First, it would happen when loading visual meshes having supported format (i.e. Collada files '.dae').
-        # Second, it does not take into account URDF 'mimic' joint constraints. However, it does a better job at
-        # initialized undetermined physics parameters.
-        if isinstance(morph, gs.morphs.MJCF):
-            # Mujoco's unified MJCF+URDF parser systematically for MJCF files
-            l_infos, links_j_infos, links_g_infos, eqs_info = mju.parse_xml(morph, surface)
-        else:
-            # Custom "legacy" URDF parser for loading geometries (visual and collision) and equality constraints.
-            # This is necessary because Mujoco cannot parse visual geometries (meshes) reliably for URDF.
-            l_infos, links_j_infos, links_g_infos, eqs_info = uu.parse_urdf(morph, surface)
-
-            # Mujoco's unified MJCF+URDF parser for only link, joints, and collision geometries properties.
-            morph_ = copy(morph)
-            morph_.visualization = False
-            try:
-                # Mujoco's unified MJCF+URDF parser for URDF files.
-                # Note that Mujoco URDF parser completely ignores equality constraints.
-                l_infos, links_j_infos_mj, links_g_infos_mj, _ = mju.parse_xml(morph_, surface)
-
-                # Mujoco is not parsing actuators properties
-                for j_info_gs in chain.from_iterable(links_j_infos):
-                    for j_info_mj in chain.from_iterable(links_j_infos_mj):
-                        if j_info_mj["name"] == j_info_gs["name"]:
-                            for name in ("dofs_force_range", "dofs_armature", "dofs_kp", "dofs_kv"):
-                                j_info_mj[name] = j_info_gs[name]
-                links_j_infos = links_j_infos_mj
-
-                # Take into account 'world' body if it was added automatically for our legacy URDF parser
-                if len(links_g_infos_mj) == len(links_g_infos) + 1:
-                    assert not links_g_infos_mj[0]
-                    links_g_infos.insert(0, [])
-                assert len(links_g_infos_mj) == len(links_g_infos)
-
-                # Update collision geometries, ignoring fake" visual geometries returned by Mujoco, (which is using
-                # collision as visual to avoid loading mesh files), and keeping the true visual geometries provided
-                # by our custom legacy URDF parser.
-                # Note that the Kinematic tree ordering is stable between Mujoco and Genesis (Hopefully!).
-                for link_g_infos, link_g_infos_mj in zip(links_g_infos, links_g_infos_mj):
-                    # Remove collision geometries from our legacy URDF parser
-                    for i_g, g_info in tuple(enumerate(link_g_infos))[::-1]:
-                        is_col = g_info["contype"] or g_info["conaffinity"]
-                        if is_col:
-                            del link_g_infos[i_g]
-
-                    # Add visual geometries from Mujoco's unified MJCF+URDF parser
-                    for g_info in link_g_infos_mj:
-                        is_col = g_info["contype"] or g_info["conaffinity"]
-                        if is_col:
-                            link_g_infos.append(g_info)
-            except (ValueError, AssertionError):
-                gs.logger.info("Falling back to legacy URDF parser. Default values of physics properties may be off.")
-
+    def _collect_urdf_articulation_info(self, morph, surface):
+        # Custom "legacy" URDF parser for loading geometries (visual and collision) and equality constraints.
+        # This is necessary because Mujoco cannot parse visual geometries (meshes) reliably for URDF.
+        l_infos, links_j_infos, links_g_infos, eqs_info = uu.parse_urdf(morph, surface)
+        # Mujoco's unified MJCF+URDF parser for only link, joints, and collision geometries properties.
+        morph_ = copy(morph)
+        morph_.visualization = False
+        try:
+            # Mujoco's unified MJCF+URDF parser for URDF files.
+            # Note that Mujoco URDF parser completely ignores equality constraints.
+            l_infos, links_j_infos_mj, links_g_infos_mj, _ = mju.parse_xml(morph_, surface)
+
+            # Mujoco is not parsing actuators properties
+            for j_info_gs in chain.from_iterable(links_j_infos):
+                for j_info_mj in chain.from_iterable(links_j_infos_mj):
+                    if j_info_mj["name"] == j_info_gs["name"]:
+                        for name in ("dofs_force_range", "dofs_armature", "dofs_kp", "dofs_kv"):
+                            j_info_mj[name] = j_info_gs[name]
+            links_j_infos = links_j_infos_mj
+
+            # Take into account 'world' body if it was added automatically for our legacy URDF parser
+            if len(links_g_infos_mj) == len(links_g_infos) + 1:
+                assert not links_g_infos_mj[0]
+                links_g_infos.insert(0, [])
+            assert len(links_g_infos_mj) == len(links_g_infos)
+
+            # Update collision geometries, ignoring fake" visual geometries returned by Mujoco, (which is using
+            # collision as visual to avoid loading mesh files), and keeping the true visual geometries provided
+            # by our custom legacy URDF parser.
+            # Note that the Kinematic tree ordering is stable between Mujoco and Genesis (Hopefully!).
+            for link_g_infos, link_g_infos_mj in zip(links_g_infos, links_g_infos_mj):
+                # Remove collision geometries from our legacy URDF parser
+                for i_g, g_info in tuple(enumerate(link_g_infos))[::-1]:
+                    is_col = g_info["contype"] or g_info["conaffinity"]
+                    if is_col:
+                        del link_g_infos[i_g]
+
+                # Add visual geometries from Mujoco's unified MJCF+URDF parser
+                for g_info in link_g_infos_mj:
+                    is_col = g_info["contype"] or g_info["conaffinity"]
+                    if is_col:
+                        link_g_infos.append(g_info)
+        except (ValueError, AssertionError):
+            gs.logger.info("Falling back to legacy URDF parser. Default values of physics properties may be off.")
+
+        return l_infos, links_j_infos, links_g_infos, eqs_info
+
+    def _build_up_articulation(self, l_infos, links_j_infos, links_g_infos, eqs_info, morph, surface):
         # Add free floating joint at root if necessary
         if (
             (isinstance(morph, gs.morphs.Drone) or (isinstance(morph, gs.morphs.URDF) and not morph.fixed))
@@ -591,6 +614,26 @@ def _load_scene(self, morph, surface):
                 sol_params=eq_info["sol_params"],
             )
 
+    def _load_scene(self, morph, surface):
+        # Mujoco's unified MJCF+URDF parser is not good enough for now to be used for loading both MJCF and URDF files.
+        # First, it would happen when loading visual meshes having supported format (i.e. Collada files '.dae').
+        # Second, it does not take into account URDF 'mimic' joint constraints. However, it does a better job at
+        # initialized undetermined physics parameters.
+        if isinstance(morph, gs.morphs.MJCF):
+            # Mujoco's unified MJCF+URDF parser systematically for MJCF files
+            l_infos, links_j_infos, links_g_infos, eqs_info = mju.parse_xml(morph, surface)
+        elif isinstance(morph, (gs.morphs.URDF, gs.morphs.Drone)):
+            l_infos, links_j_infos, links_g_infos, eqs_info = self._collect_urdf_articulation_info(morph, surface)
+        elif isinstance(morph, gs.morphs.USDArticulation):
+            from genesis.utils.usd import parse_usd_articulation
+
+            l_infos, links_j_infos, links_g_infos, eqs_info = parse_usd_articulation(morph, surface)
+        else:
+            gs.raise_exception(f"Unsupported morph type: {type(morph)}")
+
+        if len(l_infos) > 0:
+            self._build_up_articulation(l_infos, links_j_infos, links_g_infos, eqs_info, morph, surface)
+
     def _build(self):
         for link in self._links:
             link._build()

genesis/engine/scene.py

@@ -452,13 +452,12 @@ def add_entity(
         return entity
 
     @gs.assert_unbuilt
-    def link_entities(
-        self,
-        parent_entity: "Entity",
-        child_entity: "Entity",
-        parent_link_name="",
-        child_link_name="",
-    ):
+    def add_stage(self, file: str):
+        from ..utils.usd import import_from_usd
+        return import_from_usd(self, file)
+
+    @gs.assert_unbuilt
+    def link_entities(self, parent_entity: "Entity", child_entity: "Entity", parent_link_name="", child_link_name=""):
         """
         links two entities to act as single entity.
 

genesis/engine/solvers/rigid/constraint_solver_decomp.py

@@ -1048,7 +1048,6 @@ def add_joint_limit_constraints(
                     pos_delta_min = rigid_global_info.qpos[i_q, i_b] - dofs_info.limit[I_d][0]
                     pos_delta_max = dofs_info.limit[I_d][1] - rigid_global_info.qpos[i_q, i_b]
                     pos_delta = min(pos_delta_min, pos_delta_max)
-
                     if pos_delta < 0:
                         jac = (pos_delta_min < pos_delta_max) * 2 - 1
                         jac_qvel = jac * dofs_state.vel[i_d, i_b]

genesis/options/morphs.py

@@ -6,7 +6,7 @@
 """
 
 import os
-from typing import Any, List, Optional, Sequence, Tuple, Union
+from typing import Any, List, Optional, Sequence, Tuple, Union, ClassVar
 
 import numpy as np
 
@@ -1283,3 +1283,61 @@ def default_params(self):
     @property
     def subterrain_params(self):
         return self._subterrain_parameters
+
+class USDArticulation(FileMorph):
+    """
+    Morph loaded from a USD Prim inside a USD file. 
+
+    This morph only supports Prim with APISchema:
+
+    - PhysicsRigidBodyAPI
+    - PhysicsArticulationRootAPI
+
+    Parameters
+    ----------
+    file : str
+        The path to the file.
+
+    prim_path : str, optional
+        The path to the USD prim inside the USD file. If not specified, the default prim will be used.
+
+    parser_ctx : UsdParserContext, optional, for better performance when parsing large USD files
+    """
+    file: str
+    prim_path: str = None
+    parser_ctx: Any = None
+    def __init__(self, **data):
+        super().__init__(**data)
+        for USD_FORMAT in USD_FORMATS:
+            if self.is_format(USD_FORMAT):
+                break
+        else:
+            gs.raise_exception(f"Expected one of `{USD_FORMATS}` extensions for USD file: {self.file}")
+
+
+class USDRigidBody(FileMorph):
+    """
+    Morph loaded from a USD Prim with RigidBodyAPI (but not ArticulationRootAPI) inside a USD file.
+
+    This morph supports single rigid bodies (not articulated structures).
+
+    Parameters
+    ----------
+    file : str
+        The path to the file.
+
+    prim_path : str, optional
+        The path to the USD prim inside the USD file. If not specified, the first prim with RigidBodyAPI will be used.
+
+    parser_ctx : UsdParserContext, optional, for better performance when parsing large USD files
+    """
+    file: str
+    prim_path: str = None
+    parser_ctx: Any  = None
+    def __init__(self, **data):
+        super().__init__(**data)
+        for USD_FORMAT in USD_FORMATS:
+            if self.is_format(USD_FORMAT):
+                break
+        else:
+            gs.raise_exception(f"Expected one of `{USD_FORMATS}` extensions for USD file: {self.file}")

genesis/utils/urdf.py

@@ -59,7 +59,7 @@ def _order_links(l_infos, j_infos, links_g_infos=None):
 def parse_urdf(morph, surface):
     if isinstance(morph.file, (str, Path)):
         path = os.path.join(get_assets_dir(), morph.file)
-        robot = urdfpy.URDF.load(path)
+        robot:urdfpy.URDF = urdfpy.URDF.load(path)
     else:
         robot = morph.file
 

genesis/utils/usd/__init__.py

@@ -0,0 +1,16 @@
+"""
+USD Parser System for Genesis
+
+This package provides an extendable USD parser system with:
+- UsdParserUtils: Utility functions for transforms, mesh conversion, etc.
+- UsdParserContext: Context for tracking materials, articulations, and rigid bodies
+- UsdRenderingMaterialParser: Parser for rendering materials
+- UsdArticulationParser: Parser for articulations
+- UsdRigidBodyParser: Parser for rigid bodies
+- UsdParser: Main parser entrance with import_from_stage function
+"""
+
+from .usd_parser import UsdParser
+from .usd_parser import import_from_usd
+from .usd_articulation_parser import parse_usd_articulation
+from .usd_rigid_body_parser import parse_usd_rigid_body