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

.github/workflows/examples.yml

@@ -51,7 +51,7 @@ jobs:
         run: |
           pip install --upgrade pip setuptools wheel
           pip install torch --index-url https://download.pytorch.org/whl/cpu
-          pip install -e '.[dev]' pynput
+          pip install -e '.[dev,usd]' pynput
 
       - name: Get gstaichi version
         id: gstaichi_version

.github/workflows/generic.yml

@@ -131,6 +131,8 @@ jobs:
         shell: bash
         run: |
           PYTHON_DEPS="dev"
+          # Install USD for all platforms except ARM (usd-core doesn't support ARM)
+          # This is required for test_mesh.py which tests USD parsing functionality
           if [[ "${{ matrix.OS }}" != 'ubuntu-24.04-arm' ]] ; then
             PYTHON_DEPS="${PYTHON_DEPS},usd"
           fi

examples/usd/import_stage.py

@@ -0,0 +1,102 @@
+import argparse
+import os
+
+import numpy as np
+from huggingface_hub import snapshot_download
+
+import genesis as gs
+from genesis.utils.misc import ti_to_numpy
+
+
+class JointAnimator:
+    """
+    A simple JointAnimator to animate the joints' positions of the scene.
+
+    It uses the sin function to interpolate between the lower and upper limits of the joints.
+    """
+
+    def __init__(self, scene: gs.Scene):
+        self.rigid = scene.sim.rigid_solver
+        joint_limits = ti_to_numpy(self.rigid.dofs_info.limit)
+        joint_limits = np.clip(joint_limits, -np.pi, np.pi)
+
+        init_positions = self.rigid.get_dofs_position().numpy()
+
+        self.joint_lower = joint_limits[:, 0]
+        self.joint_upper = joint_limits[:, 1]
+
+        valid_range_mask = (self.joint_upper - self.joint_lower) > gs.EPS
+
+        normalized_init_pos = np.where(
+            valid_range_mask,
+            2.0 * (init_positions - self.joint_lower) / (self.joint_upper - self.joint_lower) - 1.0,
+            0.0,
+        )
+        self.init_phase = np.arcsin(normalized_init_pos)
+
+    def animate(self, scene: gs.Scene):
+        """Calculate target positions using sin function to interpolate between lower and upper limits"""
+        t = scene.t * scene.dt
+        theta = np.pi * t + self.init_phase
+        theta = theta % (2 * np.pi)
+        sin_values = np.sin(theta)
+        normalized = (sin_values + 1.0) / 2.0
+        target = self.joint_lower + (self.joint_upper - self.joint_lower) * normalized
+        self.rigid.control_dofs_position(target)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-n", "--num_steps", type=int, default=1000)
+    parser.add_argument("-v", "--vis", action="store_true", default=False)
+    args = parser.parse_args()
+
+    args.num_steps = 1 if "PYTEST_VERSION" in os.environ else args.num_steps
+    args.vis = False if "PYTEST_VERSION" in os.environ else args.vis
+
+    gs.init(backend=gs.cpu)
+
+    dt = 0.002
+    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,
+        ),
+        rigid_options=gs.options.RigidOptions(
+            dt=dt,
+            gravity=(0, 0, -9.8),
+            enable_collision=True,
+            enable_joint_limit=True,
+            max_collision_pairs=1000,
+        ),
+        show_viewer=args.vis,
+    )
+
+    asset_path = snapshot_download(
+        repo_type="dataset",
+        repo_id="Genesis-Intelligence/assets",
+        revision="main",
+        allow_patterns="usd/Lightwheel_Kitchen001/Kitchen001/*",
+        max_workers=1,
+    )
+
+    entities = scene.add_stage(
+        morph=gs.morphs.USD(
+            file=f"{asset_path}/usd/Lightwheel_Kitchen001/Kitchen001/Kitchen001.usd",
+        ),
+        # vis_mode="collision",
+        # visualize_contact=True,
+    )
+
+    scene.build()
+
+    joint_animator = JointAnimator(scene)
+
+    for _ in range(args.num_steps):
+        joint_animator.animate(scene)
+        scene.step()
+
+
+if __name__ == "__main__":
+    main()

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -140,7 +140,7 @@ def _load_model(self):
 
         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)):
+        elif isinstance(self._morph, (gs.morphs.MJCF, gs.morphs.URDF, gs.morphs.Drone, gs.morphs.USD)):
             self._load_scene(self._morph, self._surface)
         elif isinstance(self._morph, gs.morphs.Primitive):
             self._load_primitive(self._morph, self._surface)
@@ -371,11 +371,10 @@ def _load_scene(self, morph, surface):
         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:
+        elif isinstance(morph, (gs.morphs.URDF, gs.morphs.Drone)):
             # 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
@@ -416,7 +415,11 @@ def _load_scene(self, morph, surface):
                             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.")
+        elif isinstance(morph, gs.morphs.USD):
+            from genesis.utils.usd import parse_usd_rigid_entity
 
+            # Unified parser handles both articulations and rigid bodies
+            l_infos, links_j_infos, links_g_infos, eqs_info = parse_usd_rigid_entity(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))

genesis/engine/mesh.py

@@ -355,7 +355,7 @@ def from_morph_surface(cls, morph, surface=None):
                         " and rotate glb mesh by default later and gradually enforce this option."
                     )
             elif morph.is_format(gs.options.morphs.USD_FORMATS):
-                import genesis.utils.usda as usda_utils
+                import genesis.utils.usd.usda as usda_utils
 
                 meshes = usda_utils.parse_mesh_usd(morph.file, morph.group_by_material, morph.scale, surface)
             elif isinstance(morph, gs.options.morphs.MeshSet):

genesis/engine/scene.py

@@ -3,7 +3,7 @@
 import sys
 import time
 import weakref
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Callable, Literal
 
 import numpy as np
 import torch
@@ -443,6 +443,17 @@ def add_entity(
 
         return entity
 
+    @gs.assert_unbuilt
+    def add_stage(
+        self,
+        morph: gs.morphs.USD,
+        vis_mode: Literal["visual", "collision"] = "visual",
+        visualize_contact: bool = False,
+    ):
+        from genesis.utils.usd import import_from_stage
+
+        return import_from_stage(self, morph.file, vis_mode, morph, visualize_contact)
+
     @gs.assert_unbuilt
     def add_mesh_light(
         self,

genesis/options/morphs.py

@@ -1284,3 +1284,73 @@ def default_params(self):
     @property
     def subterrain_params(self):
         return self._subterrain_parameters
+
+
+class USD(FileMorph):
+    """
+    Configuration class for USD file loading with advanced processing options.
+
+    This class encapsulates the file path and processing parameters for USD loading,
+    allowing users to control convexification, decimation, and decomposition behavior
+    when loading USD scenes via add_stage().
+
+    Parameters
+    ----------
+    file : str
+        The path to the USD file.
+    convexify : bool, optional
+        Whether to convexify the entity. When convexify is True, all the meshes in the entity will each be converted
+        to a set of convex hulls. The mesh will be decomposed into multiple convex components if the convex hull is not
+        sufficient to meet the desired accuracy. The module 'coacd' is used for this decomposition process.
+        If not given, it defaults to `True` for `RigidEntity` and `False` for other deformable entities.
+    decompose_object_error_threshold : float, optional
+        For basic rigid objects (mug, table...), skip convex decomposition if the relative difference between the
+        volume of original mesh and its convex hull is lower than this threshold.
+        0.0 to enforce decomposition, float("inf") to disable it completely. Defaults to 0.15 (15%).
+    decompose_robot_error_threshold : float, optional
+        For poly-articulated robots, skip convex decomposition if the relative difference between the volume of
+        original mesh and its convex hull is lower than this threshold.
+        0.0 to enforce decomposition, float("inf") to disable it completely. Defaults to float("inf").
+    coacd_options : CoacdOptions, optional
+        Options for configuring coacd convex decomposition. Needs to be a `gs.options.CoacdOptions` object.
+    decimate : bool, optional
+        Whether to decimate (simplify) the mesh. Defaults to True. **This is only used for RigidEntity.**
+    decimate_face_num : int, optional
+        The number of faces to decimate to. Defaults to 500. **This is only used for RigidEntity.**
+    decimate_aggressiveness : int, optional
+        How hard the decimation process will try to match the target number of faces, as an integer ranging from 0 to 8.
+        0 is lossless. 2 preserves all features of the original geometry. 5 may significantly alter the original
+        geometry if necessary. 8 does what needs to be done at all costs. Defaults to 2.
+        **This is only used for RigidEntity.**
+    _prim_path : str, optional
+        The parsing target prim path. Defaults to None.
+    _parsing_type : str, optional
+        The parsing type.
+        'articulation' for articulated body parsing, ArticulationRootAPI is required.
+        'rigid_body' for rigid body parsing, CollisionAPI|RigidBodyAPI is required.
+        Defaults to None, no parsing will be performed.
+    """
+
+    file: str
+    convexify: Optional[bool] = None
+    decompose_object_error_threshold: float = 0.15
+    decompose_robot_error_threshold: float = float("inf")
+    coacd_options: Optional[CoacdOptions] = None
+    decimate: bool = True
+    decimate_face_num: int = 500
+    decimate_aggressiveness: int = 2
+    prim_path: Optional[str] = None
+    parsing_type: Optional[Literal["articulation", "rigid_body"]] = None
+    parser_ctx: Any = None
+
+    def __init__(self, **data):
+        super().__init__(**data)
+
+        if not isinstance(self.file, str):
+            gs.raise_exception("`file` should be a string.")
+
+        if not self.file.lower().endswith(USD_FORMATS):
+            gs.raise_exception(f"USDMorph requires a USD file with extension {USD_FORMATS}, got: {self.file}")
+
+        if self.coacd_options is None:
+            self.coacd_options = CoacdOptions()

genesis/utils/geom.py

@@ -1,4 +1,5 @@
 import math
+from typing import Literal
 
 import numpy as np
 import numba as nb
@@ -1173,6 +1174,113 @@ def inv_transform_by_T(pos, T):
     return transform_by_R(pos - trans, R_inv)
 
 
+def polar(A, pure_rotation: bool = True, side: Literal["right", "left"] = "right"):
+    """
+    Compute the polar decomposition of a matrix.
+
+    Parameters
+    ----------
+    A : np.ndarray | torch.Tensor
+        The matrix to decompose. Must be a 2D matrix (M, N), not batched.
+    pure_rotation : bool, optional
+        If True, ensure the unitary matrix U has det(U) = 1 (pure rotation).
+        If False, U may have det(U) = -1 (contains reflection). Default is True.
+    side : Literal['right', 'left'], optional
+        The side of the decomposition. Either 'right' or 'left'. Default is 'right'.
+
+    Returns
+    -------
+    tuple[np.ndarray | torch.Tensor, np.ndarray | torch.Tensor]
+        A tuple of (U, P) where:
+        - U : The unitary matrix (rotation part), same shape as A (M, N).
+        - P : The positive semi-definite matrix (scaling part). For 'right' decomposition,
+          P has shape (N, N). For 'left' decomposition, P has shape (M, M).
+    """
+    if isinstance(A, torch.Tensor):
+        if A.ndim != 2:
+            gs.raise_exception(f"Input must be a 2D matrix. got: {A.ndim=} dimensions")
+
+        U_svd, Sigma, Vt = torch.linalg.svd(A, full_matrices=False)
+
+        # Handle pure_rotation: if det(U) < 0, flip signs to make it a pure rotation
+        if pure_rotation and A.shape[0] == A.shape[1]:  # Only for square matrices
+            # Compute U first to check its determinant
+            U_temp = U_svd @ Vt
+            det_U = torch.linalg.det(U_temp)
+            if det_U < 0:
+                U_svd = U_svd.clone()
+                Vt = Vt.clone()
+                Sigma = Sigma.clone()
+
+                # Flip the smallest singular value (last one) and corresponding column/row
+                # This maintains A = U_svd @ diag(Sigma) @ Vt while changing det(U)
+                if side == "right":
+                    # Flip last singular value and last row of Vt
+                    Sigma[-1] *= -1
+                    Vt[-1, :] *= -1
+                else:
+                    # For left polar, flip last singular value and last column of U_svd
+                    Sigma[-1] *= -1
+                    U_svd[:, -1] *= -1
+
+        U = U_svd @ Vt
+
+        # Construct diagonal matrix from singular values
+        # Use absolute value to ensure P is positive semi-definite
+        # (even if we flipped Sigma's sign to fix rotation)
+        Sigma_abs = torch.abs(Sigma)
+        Sigma_diag = torch.diag_embed(Sigma_abs)
+        if side == "right":
+            # P = Vt.T @ diag(|Sigma|) @ Vt
+            P = Vt.transpose(-2, -1) @ Sigma_diag @ Vt
+        else:
+            # P = U_svd @ diag(|Sigma|) @ U_svd.T (left polar: A = P @ U)
+            P = U_svd @ Sigma_diag @ U_svd.transpose(-2, -1)
+        return U, P
+    elif isinstance(A, np.ndarray):
+        if A.ndim != 2:
+            gs.raise_exception(f"Input must be a 2D matrix. got: {A.ndim=} dimensions")
+
+        U_svd, Sigma, Vt = np.linalg.svd(A, full_matrices=False)
+
+        # Handle pure_rotation: if det(U) < 0, flip signs to make it a pure rotation
+        if pure_rotation and A.shape[0] == A.shape[1]:  # Only for square matrices
+            # Compute U first to check its determinant
+            U_temp = U_svd @ Vt
+            det_U = np.linalg.det(U_temp)
+            if det_U < 0:
+                U_svd = U_svd.copy()
+                Vt = Vt.copy()
+                Sigma = Sigma.copy()
+
+                # Flip the smallest singular value (last one) and corresponding column/row
+                # This maintains A = U_svd @ diag(Sigma) @ Vt while changing det(U)
+                if side == "right":
+                    # Flip last singular value and last row of Vt
+                    Sigma[-1] *= -1
+                    Vt[-1, :] *= -1
+                else:
+                    # For left polar, flip last singular value and last column of U_svd
+                    Sigma[-1] *= -1
+                    U_svd[:, -1] *= -1
+
+        U = U_svd @ Vt
+
+        # Use absolute value to ensure P is positive semi-definite
+        # (even if we flipped Sigma's sign to fix rotation)
+        Sigma_abs = np.abs(Sigma)
+
+        if side == "right":
+            # P = Vt.T @ diag(|Sigma|) @ Vt
+            P = Vt.T @ np.diag(Sigma_abs) @ Vt
+        else:
+            # P = U_svd @ diag(|Sigma|) @ U_svd.T (left polar: A = P @ U)
+            P = U_svd @ np.diag(Sigma_abs) @ U_svd.T
+        return U, P
+    else:
+        gs.raise_exception(f"the input must be either torch.Tensor or np.ndarray. got: {type(A)=}")
+
+
 # ------------------------------------------------------------------------------------
 # ------------------------------------- numpy ----------------------------------------
 # ------------------------------------------------------------------------------------