Adds methods to set physics-based schemas (#110)

# Description

This MR introduces wrappers around different [USD
Physics](https://openusd.org/dev/api/usd_physics_page_front.html) and
[PhysX
solver](https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/104.2/index.html)
schemas. The functions allow modifying the properties on an asset prim
using configuration objects.

The schemas supersede the current `omni.isaac.orbit.utils.kit.py` which
did the same job but had duplication in the implementations.

## Type of change

- New feature (non-breaking change which adds functionality)
- This change requires a documentation update

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./orbit.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file

---------

Signed-off-by: Mayank Mittal <[email protected]>

source/extensions/omni.isaac.orbit/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.8.7"
+version = "0.8.8"
 
 # Description
 title = "ORBIT framework for Robot Learning"

source/extensions/omni.isaac.orbit/docs/CHANGELOG.rst

@@ -1,6 +1,17 @@
 Changelog
 ---------
 
+0.8.8 (2023-08-09)
+~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added configuration classes and functions for setting different physics-based schemas in the
+  :mod:`omni.isaac.orbit.sim.schemas` module. These allow modifying properties of the physics solver
+  on the asset using configuration objects.
+
+
 0.8.7 (2023-08-03)
 ~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/__init__.py

@@ -3,13 +3,17 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-"""This sub-module provides the ``SimulationContext`` class.
+"""This sub-module contains utilities for simulation.
+
+To make it convenient to use the module, we recommend importing the module as follows:
+
+.. code-block:: python
+
+    import omni.isaac.orbit.sim as sim_utils
 
-The :class:`SimulationContext` inherits from the :class:`omni.isaac.core.simulation_context.SimulationContext` class
-to provide additional functionality for the Orbit extension. This includes configuring the simulation through
-the configuration class :class:`SimulationCfg` and providing a context manager for the simulation.
 """
 
+from .schemas import *  # noqa: F401, F403
 from .simulation_cfg import PhysicsMaterialCfg, PhysxCfg, SimulationCfg
 from .simulation_context import SimulationContext
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/schemas/__init__.py

@@ -0,0 +1,62 @@
+# Copyright [2023] Boston Dynamics AI Institute, Inc.
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Sub-module containing utilities for schemas used in Omniverse.
+
+We wrap the USD schemas for PhysX and USD Physics in a more convenient API for setting the parameters from
+Python. This is done so that configuration objects can define the schema properties to set and make it easier
+to tune the physics parameters without requiring to open Omniverse Kit and manually set the parameters into
+the respective USD attributes.
+
+.. caution::
+
+    Schema properties cannot be applied on prims that are prototypes as they are read-only prims. This
+    particularly affects instanced assets where some of the prims (usually the visual and collision meshes)
+    are prototypes so that the instancing can be done efficiently.
+
+    In such cases, it is assumed that the prototypes have sim-ready properties on them that don't need to be modified.
+    Trying to set properties into prototypes will throw a warning saying that the prim is a prototype and the
+    properties cannot be set.
+
+The schemas are defined in the following links:
+
+* `UsdPhysics schema <https://openusd.org/dev/api/usd_physics_page_front.html>`_
+* `PhysxSchema schema <https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/104.2/index.html>`_
+
+Locally, the schemas are defined in the following files:
+
+* ``_isaac_sim/kit/extsPhysics/omni.usd.schema.physics/plugins/UsdPhysics/resources/UsdPhysics/schema.usda``
+* ``_isaac_sim/kit/extsPhysics/omni.usd.schema.physx/plugins/PhysxSchema/resources/PhysxSchema/schema.usda``
+
+"""
+
+from .schemas import (
+    set_articulation_root_properties,
+    set_collision_properties,
+    set_mass_properties,
+    set_rigid_body_properties,
+)
+from .schemas_cfg import (
+    ArticulationRootPropertiesCfg,
+    CollisionPropertiesCfg,
+    MassPropertiesCfg,
+    RigidBodyPropertiesCfg,
+)
+
+__all__ = [
+    # articulation root
+    "ArticulationRootPropertiesCfg",
+    "set_articulation_root_properties",
+    # rigid bodies
+    "RigidBodyPropertiesCfg",
+    "set_rigid_body_properties",
+    # colliders
+    "CollisionPropertiesCfg",
+    "set_collision_properties",
+    # mass
+    "MassPropertiesCfg",
+    "set_mass_properties",
+]

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/schemas/schemas.py

@@ -0,0 +1,219 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import omni.isaac.core.utils.stage as stage_utils
+from pxr import PhysxSchema, Usd, UsdPhysics
+
+from ..utils import apply_nested, safe_set_attribute_on_usd_schema
+from . import schemas_cfg
+
+
+@apply_nested
+def set_articulation_root_properties(
+    prim_path: str, cfg: schemas_cfg.ArticulationRootPropertiesCfg, stage: Usd.Stage = None
+):
+    """Set PhysX parameters for an articulation root prim.
+
+    The `articulation root`_ marks the root of an articulation tree. For floating articulations, this should be on
+    the root body. For fixed articulations, this API can be on a direct or indirect parent of the root joint
+    which is fixed to the world.
+
+    The schema comprises of attributes that belong to the `ArticulationRootAPI`_ and `PhysxArticulationAPI`_.
+    schemas. The latter contains the PhysX parameters for the articulation root.
+
+    The properties are applied to the articulation root prim. The common properties (such as solver position
+    and velocity iteration counts, sleep threshold, stabilization threshold) take precedence over those specified
+    in the rigid body schemas for all the rigid bodies in the articulation.
+
+    .. note::
+        This function is decorated with :func:`apply_nested` that set the properties to all the prims
+        (that have the schema applied on them) under the input prim path.
+
+    .. _articulation root: https://nvidia-omniverse.github.io/PhysX/physx/5.2.1/docs/Articulations.html
+    .. _ArticulationRootAPI: https://openusd.org/dev/api/class_usd_physics_articulation_root_a_p_i.html
+    .. _PhysxArticulationAPI: https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/104.2/class_physx_schema_physx_articulation_a_p_i.html
+
+    Args:
+        prim_path (str): The prim path to the articulation root.
+        cfg (schemas_cfg.ArticulationRootPropertiesCfg): The configuration for the articulation root.
+        stage (Usd.Stage, optional): The stage where to find the prim. Defaults to None, in which case the
+            current stage is used.
+    """
+    # obtain stage
+    if stage is None:
+        stage = stage_utils.get_current_stage()
+    # get articulation USD prim
+    articulation_prim = stage.GetPrimAtPath(prim_path)
+    # check if prim has articulation applied on it
+    if not UsdPhysics.ArticulationRootAPI(articulation_prim):
+        return False
+    # retrieve the articulation api
+    physx_articulation_api = PhysxSchema.PhysxArticulationAPI(articulation_prim)
+    if not physx_articulation_api:
+        physx_articulation_api = PhysxSchema.PhysxArticulationAPI.Apply(articulation_prim)
+
+    # convert to dict
+    cfg = cfg.to_dict()
+    # set into physx api
+    for attr_name, value in cfg.items():
+        safe_set_attribute_on_usd_schema(physx_articulation_api, attr_name, value)
+    # success
+    return True
+
+
+@apply_nested
+def set_rigid_body_properties(prim_path: str, cfg: schemas_cfg.RigidBodyPropertiesCfg, stage: Usd.Stage = None):
+    """Set PhysX parameters for a rigid body prim.
+
+    A `rigid body`_ is a single body that can be simulated by PhysX. It can be either dynamic or kinematic.
+    A dynamic body responds to forces and collisions. A `kinematic body`_ can be moved by the user, but does not
+    respond to forces. They are similar to having static bodies that can be moved around.
+
+    The schema comprises of attributes that belong to the `RigidBodyAPI`_ and `PhysxRigidBodyAPI`_.
+    schemas. The latter contains the PhysX parameters for the rigid body.
+
+    .. note::
+        This function is decorated with :func:`apply_nested` that sets the properties to all the prims
+        (that have the schema applied on them) under the input prim path.
+
+    .. _rigid body: https://nvidia-omniverse.github.io/PhysX/physx/5.2.1/docs/RigidBodyOverview.html
+    .. _kinematic body: https://openusd.org/release/wp_rigid_body_physics.html#kinematic-bodies
+    .. _RigidBodyAPI: https://openusd.org/dev/api/class_usd_physics_rigid_body_a_p_i.html
+    .. _PhysxRigidBodyAPI: https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/104.2/class_physx_schema_physx_rigid_body_a_p_i.html
+
+    Args:
+        prim_path (str): The prim path to the rigid body.
+        cfg (schemas_cfg.RigidBodyPropertiesCfg): The configuration for the rigid body.
+        stage (Usd.Stage, optional): The stage where to find the prim. Defaults to None, in which case the
+            current stage is used.
+    """
+    # obtain stage
+    if stage is None:
+        stage = stage_utils.get_current_stage()
+    # get rigid-body USD prim
+    rigid_body_prim = stage.GetPrimAtPath(prim_path)
+    # check if prim has rigid-body applied on it
+    if not UsdPhysics.RigidBodyAPI(rigid_body_prim):
+        return False
+    # retrieve the USD rigid-body api
+    usd_rigid_body_api = UsdPhysics.RigidBodyAPI(rigid_body_prim)
+    # retrieve the physx rigid-body api
+    physx_rigid_body_api = PhysxSchema.PhysxRigidBodyAPI(rigid_body_prim)
+    if not physx_rigid_body_api:
+        physx_rigid_body_api = PhysxSchema.PhysxRigidBodyAPI.Apply(rigid_body_prim)
+
+    # convert to dict
+    cfg = cfg.to_dict()
+    # set into USD API
+    for attr_name in ["rigid_body_enabled", "kinematic_enabled"]:
+        value = cfg.pop(attr_name, None)
+        safe_set_attribute_on_usd_schema(usd_rigid_body_api, attr_name, value)
+    # set into PhysX API
+    for attr_name, value in cfg.items():
+        safe_set_attribute_on_usd_schema(physx_rigid_body_api, attr_name, value)
+    # success
+    return True
+
+
+@apply_nested
+def set_collision_properties(prim_path: str, cfg: schemas_cfg.CollisionPropertiesCfg, stage: Usd.Stage = None):
+    """Set PhysX properties of collider prim.
+
+    These properties are based on the `UsdPhysics.CollisionAPI` and `PhysxSchema.PhysxCollisionAPI`_ schemas.
+    For more information on the properties, please refer to the official documentation.
+
+    Tuning these parameters influence the contact behavior of the rigid body. For more information on
+    tune them and their effect on the simulation, please refer to the
+    `PhysX documentation <https://nvidia-omniverse.github.io/PhysX/physx/5.2.1/docs/AdvancedCollisionDetection.html>`_.
+
+    .. note::
+        This function is decorated with :func:`apply_nested` that sets the properties to all the prims
+        (that have the schema applied on them) under the input prim path.
+
+    .. UsdPhysics.CollisionAPI: https://openusd.org/dev/api/class_usd_physics_collision_a_p_i.html
+    .. PhysxSchema.PhysxCollisionAPI: https://docs.omniverse.nvidia.com/kit/docs/omni_usd_schema_physics/104.2/class_physx_schema_physx_collision_a_p_i.html
+
+    Args:
+        prim_path (str): The prim path of parent.
+        cfg (schemas_cfg.CollisionPropertiesCfg): The configuration for the collider.
+        stage (Usd.Stage, optional): The stage where to find the prim. Defaults to None, in which case the
+            current stage is used.
+    """
+    # obtain stage
+    if stage is None:
+        stage = stage_utils.get_current_stage()
+    # get USD prim
+    collider_prim = stage.GetPrimAtPath(prim_path)
+    # check if prim has collision applied on it
+    if not UsdPhysics.CollisionAPI(collider_prim):
+        return False
+    # retrieve the USD rigid-body api
+    usd_collision_api = UsdPhysics.CollisionAPI(collider_prim)
+    # retrieve the collision api
+    physx_collision_api = PhysxSchema.PhysxCollisionAPI(collider_prim)
+    if not physx_collision_api:
+        physx_collision_api = PhysxSchema.PhysxCollisionAPI.Apply(collider_prim)
+
+    # convert to dict
+    cfg = cfg.to_dict()
+    # set into USD API
+    for attr_name in ["collision_enabled"]:
+        value = cfg.pop(attr_name, None)
+        safe_set_attribute_on_usd_schema(usd_collision_api, attr_name, value)
+    # set into PhysX API
+    for attr_name, value in cfg.items():
+        safe_set_attribute_on_usd_schema(physx_collision_api, attr_name, value)
+    # success
+    return True
+
+
+@apply_nested
+def set_mass_properties(prim_path: str, cfg: schemas_cfg.MassPropertiesCfg, stage: Usd.Stage = None):
+    """Set properties for the mass of a rigid body prim.
+
+    These properties are based on the `UsdPhysics.MassAPI` schema. If the mass is not defined, the density is used
+    to compute the mass. However, in that case, a collision approximation of the rigid body is used to
+    compute the density. For more information on the properties, please refer to the
+    `documentation <https://openusd.org/release/wp_rigid_body_physics.html#body-mass-properties>`_.
+
+    .. caution::
+
+        The mass of an object can be specified in multiple ways and have several conflicting settings
+        that are resolved based on precedence. Please make sure to understand the precedence rules
+        before using this property.
+
+    .. note::
+        This function is decorated with :func:`apply_nested` that sets the properties to all the prims
+        (that have the schema applied on them) under the input prim path.
+
+    .. UsdPhysics.MassAPI: https://openusd.org/dev/api/class_usd_physics_mass_a_p_i.html
+
+    Args:
+        prim_path (str): The prim path of the rigid body.
+        cfg (schemas_cfg.MassPropertiesCfg): The configuration for the mass properties.
+        stage (Usd.Stage, optional): The stage where to find the prim. Defaults to None, in which case the
+            current stage is used.
+    """
+    # obtain stage
+    if stage is None:
+        stage = stage_utils.get_current_stage()
+    # get USD prim
+    rigid_prim = stage.GetPrimAtPath(prim_path)
+    # check if prim has collision applied on it
+    if not UsdPhysics.MassAPI(rigid_prim):
+        return False
+    # retrieve the USD rigid-body api
+    usd_physics_mass_api = UsdPhysics.MassAPI(rigid_prim)
+
+    # convert to dict
+    cfg = cfg.to_dict()
+    # set into USD API
+    for attr_name in ["mass", "density"]:
+        value = cfg.pop(attr_name, None)
+        safe_set_attribute_on_usd_schema(usd_physics_mass_api, attr_name, value)
+    # success
+    return True