Allows setting USD variants when loading prim from USD file (#402)

# Description

This PR introduces a `variants` attribute in the `UsdFileCfg` which can
be used to set different variants when loading an asset from a USD file.

- New function `set_usd_variants` which applies variant sets to a prim
- New, optional `variants` attribute in the `UsdFileCfg` to specify the
attributes to be set
- Add variant setting in the `_spawn_from_usd_file` function

Fixes #401

## Type of change

- New feature (non-breaking change which adds functionality)

## 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 run all the tests with `./orbit.sh --test` and they pass
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

CONTRIBUTORS.md

@@ -38,6 +38,7 @@ Guidelines for modifications:
 * Chenyu Yang
 * Jia Lin Yuan
 * Jingzhou Liu
+* Lorenz Wellhausen
 * Muhong Guo
 * Kourosh Darvish
 * Özhan Özen

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

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.16.2"
+version = "0.16.3"
 
 # Description
 title = "ORBIT framework for Robot Learning"

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

 Changelog
 ---------
 
+0.16.3 (2024-05-13)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added ``variants`` attribute to the :class:`omni.isaac.orbit.sim.from_files.UsdFileCfg` class to select USD
+  variants when loading assets from USD files.
+
+
 0.16.2 (2024-04-26)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/spawners/from_files/from_files.py

@@ -14,7 +14,7 @@ import omni.kit.commands
 from pxr import Gf, Sdf, Usd
 
 from omni.isaac.orbit.sim import converters, schemas
-from omni.isaac.orbit.sim.utils import bind_physics_material, bind_visual_material, clone
+from omni.isaac.orbit.sim.utils import bind_physics_material, bind_visual_material, clone, select_usd_variants
 
 if TYPE_CHECKING:
     from . import from_files_cfg
@@ -233,6 +233,10 @@ def _spawn_from_usd_file(
     else:
         carb.log_warn(f"A prim already exists at prim path: '{prim_path}'.")
 
+    # modify variants
+    if hasattr(cfg, "variants") and cfg.variants is not None:
+        select_usd_variants(prim_path, cfg.variants)
+
     # modify rigid body properties
     if cfg.rigid_props is not None:
         schemas.modify_rigid_body_properties(prim_path, cfg.rigid_props)

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/spawners/from_files/from_files_cfg.py

@@ -69,6 +69,14 @@ class UsdFileCfg(FileCfg):
     usd_path: str = MISSING
     """Path to the USD file to spawn asset from."""
 
+    variants: object | dict[str, str] | None = None
+    """Variants to select from in the input USD file. Defaults to None, in which case no variants are applied.
+
+    This can either be a configclass object, in which case each attribute is used as a variant set name and its specified value,
+    or a dictionary mapping between the two. Please check the :meth:`~omni.isaac.orbit.sim.utils.select_usd_variants` function
+    for more information.
+    """
+
 
 @configclass
 class UrdfFileCfg(FileCfg, converters.UrdfConverterCfg):

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

@@ -777,3 +777,82 @@ def find_global_fixed_joint_prim(
                 break
 
     return fixed_joint_prim
+
+
+"""
+USD Variants.
+"""
+
+
+def select_usd_variants(prim_path: str, variants: object | dict[str, str], stage: Usd.Stage | None = None):
+    """Sets the variant selections from the specified variant sets on a USD prim.
+
+    `USD Variants`_ are a very powerful tool in USD composition that allows prims to have different options on
+    a single asset. This can be done by modifying variations of the same prim parameters per variant option in a set.
+    This function acts as a script-based utility to set the variant selections for the specified variant sets on a
+    USD prim.
+
+    The function takes a dictionary or a config class mapping variant set names to variant selections. For instance,
+    if we have a prim at ``"/World/Table"`` with two variant sets: "color" and "size", we can set the variant
+    selections as follows:
+
+    .. code-block:: python
+
+        select_usd_variants(
+            prim_path="/World/Table",
+            variants={
+                "color": "red",
+                "size": "large",
+            },
+        )
+
+    Alternatively, we can use a config class to define the variant selections:
+
+    .. code-block:: python
+
+        @configclass
+        class TableVariants:
+            color: Literal["blue", "red"] = "red"
+            size: Literal["small", "large"] = "large"
+
+        select_usd_variants(
+            prim_path="/World/Table",
+            variants=TableVariants(),
+        )
+
+    Args:
+        prim_path: The path of the USD prim.
+        variants: A dictionary or config class mapping variant set names to variant selections.
+        stage: The USD stage. Defaults to None, in which case, the current stage is used.
+
+    Raises:
+        ValueError: If the prim at the specified path is not valid.
+
+    .. _USD Variants: https://graphics.pixar.com/usd/docs/USD-Glossary.html#USDGlossary-Variant
+    """
+    # Resolve stage
+    if stage is None:
+        stage = stage_utils.get_current_stage()
+    # Obtain prim
+    prim = stage.GetPrimAtPath(prim_path)
+    if not prim.IsValid():
+        raise ValueError(f"Prim at path '{prim_path}' is not valid.")
+    # Convert to dict if we have a configclass object.
+    if not isinstance(variants, dict):
+        variants = variants.to_dict()
+
+    existing_variant_sets = prim.GetVariantSets()
+    for variant_set_name, variant_selection in variants.items():
+        # Check if the variant set exists on the prim.
+        if not existing_variant_sets.HasVariantSet(variant_set_name):
+            carb.log_warn(f"Variant set '{variant_set_name}' does not exist on prim '{prim_path}'.")
+            continue
+
+        variant_set = existing_variant_sets.GetVariantSet(variant_set_name)
+        # Only set the variant selection if it is different from the current selection.
+        if variant_set.GetVariantSelection() != variant_selection:
+            variant_set.SetVariantSelection(variant_selection)
+            carb.log_info(
+                f"Setting variant selection '{variant_selection}' for variant set '{variant_set_name}' on"
+                f" prim '{prim_path}'."
+            )

source/extensions/omni.isaac.orbit/test/sim/test_utils.py

@@ -18,6 +18,7 @@ import unittest
 
 import omni.isaac.core.utils.prims as prim_utils
 import omni.isaac.core.utils.stage as stage_utils
+from pxr import Sdf, Usd, UsdGeom
 
 import omni.isaac.orbit.sim as sim_utils
 from omni.isaac.orbit.utils.assets import ISAAC_NUCLEUS_DIR, ISAAC_ORBIT_NUCLEUS_DIR
@@ -106,6 +107,24 @@ class TestUtilities(unittest.TestCase):
         self.assertIsNotNone(sim_utils.find_global_fixed_joint_prim("/World/Franka"))
         self.assertIsNone(sim_utils.find_global_fixed_joint_prim("/World/Franka", check_enabled_only=True))
 
+    def test_select_usd_variants(self):
+        """Test select_usd_variants() function."""
+        stage = stage_utils.get_current_stage()
+        prim: Usd.Prim = UsdGeom.Xform.Define(stage, Sdf.Path("/World")).GetPrim()
+        stage.SetDefaultPrim(prim)
+
+        # Create the variant set and add your variants to it.
+        variants = ["red", "blue", "green"]
+        variant_set = prim.GetVariantSets().AddVariantSet("colors")
+        for variant in variants:
+            variant_set.AddVariant(variant)
+
+        # Set the variant selection
+        sim_utils.utils.select_usd_variants("/World", {"colors": "red"}, stage)
+
+        # Check if the variant selection is correct
+        self.assertEqual(variant_set.GetVariantSelection(), "red")
+
 
 if __name__ == "__main__":
     run_tests()