Adds a rigid body collection class (#1288)

# Description

Adds a rigid body collection class, which allows to spawn multiple
objects in each environment and access/modify the quantities with a
unified (env_ids, object_ids) API.

## 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
`./isaaclab.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
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

docs/source/api/lab/omni.isaac.lab.assets.rst

@@ -12,6 +12,9 @@
     RigidObject
     RigidObjectData
     RigidObjectCfg
+    RigidObjectCollection
+    RigidObjectCollectionData
+    RigidObjectCollectionCfg
     Articulation
     ArticulationData
     ArticulationCfg
@@ -51,6 +54,26 @@ Rigid Object
     :show-inheritance:
     :exclude-members: __init__, class_type
 
+Rigid Object Collection
+-----------------------
+
+.. autoclass:: RigidObjectCollection
+    :members:
+    :inherited-members:
+    :show-inheritance:
+
+.. autoclass:: RigidObjectCollectionData
+    :members:
+    :inherited-members:
+    :show-inheritance:
+    :exclude-members: __init__
+
+.. autoclass:: RigidObjectCollectionCfg
+    :members:
+    :inherited-members:
+    :show-inheritance:
+    :exclude-members: __init__, class_type
+
 Articulation
 ------------
 

docs/source/how-to/multi_asset_spawning.rst

+
 Spawning Multiple Assets
 ========================
 
 .. currentmodule:: omni.isaac.lab
 
-Typical, spawning configurations (introduced in the :ref:`tutorial-spawn-prims` tutorial) copy the same
+Typical spawning configurations (introduced in the :ref:`tutorial-spawn-prims` tutorial) copy the same
 asset (or USD primitive) across the different resolved prim paths from the expressions.
 For instance, if the user specifies to spawn the asset at "/World/Table\_.*/Object", the same
 asset is created at the paths "/World/Table_0/Object", "/World/Table_1/Object" and so on.
 
-However, at times, it might be desirable to spawn different assets under the prim paths to
-ensure a diversity in the simulation. This guide describes how to create different assets under
-each prim path using the spawning functionality.
+However, we also support multi-asset spawning with two mechanisms:
+
+1. Rigid object collections. This allows the user to spawn multiple rigid objects in each environment and access/modify
+   them with a unified API, improving performance.
+
+2. Spawning different assets under the same prim path. This allows the user to create diverse simulations, where each
+   environment has a different asset.
+
+This guide describes how to use these two mechanisms.
 
 The sample script ``multi_asset.py`` is used as a reference, located in the
 ``IsaacLab/source/standalone/demos`` directory.
@@ -20,20 +27,41 @@ The sample script ``multi_asset.py`` is used as a reference, located in the
 
    .. literalinclude:: ../../../source/standalone/demos/multi_asset.py
       :language: python
-      :emphasize-lines: 101-123, 130-149
+      :emphasize-lines: 109-131, 135-179, 184-203
       :linenos:
 
-This script creates multiple environments, where each environment has a rigid object that is either a cone,
-a cube, or a sphere, and an articulation that is either the ANYmal-C or ANYmal-D robot.
+This script creates multiple environments, where each environment has:
+
+* a rigid object collection containing a cone, a cube, and a sphere
+* a rigid object that is either a cone, a cube, or a sphere, chosen at random
+* an articulation that is either the ANYmal-C or ANYmal-D robot, chosen at random
 
 .. image:: ../_static/demos/multi_asset.jpg
   :width: 100%
   :alt: result of multi_asset.py
 
-Using Multi-Asset Spawning Functions
-------------------------------------
 
-It is possible to spawn different assets and USDs in each environment using the spawners
+Rigid Object Collections
+------------------------
+
+Multiple rigid objects can be spawned in each environment and accessed/modified with a unified ``(env_ids, obj_ids)`` API.
+While the user could also create multiple rigid objects by spawning them individually, the API is more user-friendly and
+more efficient since it uses a single physics view under the hood to handle all the objects.
+
+.. literalinclude:: ../../../source/standalone/demos/multi_asset.py
+   :language: python
+   :lines: 135-179
+   :dedent:
+
+The configuration :class:`~assets.RigidObjectCollectionCfg` is used to create the collection. It's attribute :attr:`~assets.RigidObjectCollectionCfg.rigid_objects`
+is a dictionary containing :class:`~assets.RigidObjectCfg` objects. The keys serve as unique identifiers for each
+rigid object in the collection.
+
+
+Spawning different assets under the same prim path
+--------------------------------------------------
+
+It is possible to spawn different assets and USDs under the same prim path in each environment using the spawners
 :class:`~sim.spawners.wrappers.MultiAssetSpawnerCfg` and :class:`~sim.spawners.wrappers.MultiUsdFileCfg`:
 
 * We set the spawn configuration in :class:`~assets.RigidObjectCfg` to be
@@ -41,7 +69,7 @@ It is possible to spawn different assets and USDs in each environment using the
 
   .. literalinclude:: ../../../source/standalone/demos/multi_asset.py
      :language: python
-     :lines: 99-125
+     :lines: 107-133
      :dedent:
 
   This function allows you to define a list of different assets that can be spawned as rigid objects.
@@ -53,14 +81,14 @@ It is possible to spawn different assets and USDs in each environment using the
 
   .. literalinclude:: ../../../source/standalone/demos/multi_asset.py
      :language: python
-     :lines: 128-161
+     :lines: 182-215
      :dedent:
 
   Similar to before, this configuration allows the selection of different USD files representing articulated assets.
 
 
 Things to Note
---------------
+~~~~~~~~~~~~~~
 
 Similar asset structuring
 ~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -85,7 +113,7 @@ anymore. Hence the flag :attr:`scene.InteractiveScene.replicate_physics` must be
 
 .. literalinclude:: ../../../source/standalone/demos/multi_asset.py
    :language: python
-   :lines: 221-224
+   :lines: 280-283
    :dedent:
 
 The Code Execution

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

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.27.13"
+version = "0.27.14"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

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

 Changelog
 ---------
 
+0.27.14 (2024-10-23)
+~~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added the class :class:`~omni.isaac.lab.assets.RigidObjectCollection` which allows to spawn
+  multiple objects in each environment and access/modify the quantities with a unified (env_ids, object_ids) API.
+
+
 0.27.13 (2024-10-30)
 ~~~~~~~~~~~~~~~~~~~~
 
@@ -13,7 +23,7 @@ Added
 
 
 0.27.12 (2024-01-04)
-~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~
 
 Removed
 ^^^^^^^

source/extensions/omni.isaac.lab/omni/isaac/lab/assets/__init__.py

@@ -43,3 +43,4 @@ from .asset_base import AssetBase
 from .asset_base_cfg import AssetBaseCfg
 from .deformable_object import DeformableObject, DeformableObjectCfg, DeformableObjectData
 from .rigid_object import RigidObject, RigidObjectCfg, RigidObjectData
+from .rigid_object_collection import RigidObjectCollection, RigidObjectCollectionCfg, RigidObjectCollectionData

source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/__init__.py

+# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Sub-module for rigid object collection."""
+
+from .rigid_object_collection import RigidObjectCollection
+from .rigid_object_collection_cfg import RigidObjectCollectionCfg
+from .rigid_object_collection_data import RigidObjectCollectionData

source/extensions/omni.isaac.lab/omni/isaac/lab/assets/rigid_object_collection/rigid_object_collection_cfg.py

+# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from dataclasses import MISSING
+
+from omni.isaac.lab.assets.rigid_object import RigidObjectCfg
+from omni.isaac.lab.utils import configclass
+
+from .rigid_object_collection import RigidObjectCollection
+
+
+@configclass
+class RigidObjectCollectionCfg:
+    """Configuration parameters for a rigid object collection."""
+
+    class_type: type = RigidObjectCollection
+    """The associated asset class.
+
+    The class should inherit from :class:`omni.isaac.lab.assets.asset_base.AssetBase`.
+    """
+
+    rigid_objects: dict[str, RigidObjectCfg] = MISSING
+    """Dictionary of rigid object configurations to spawn.
+
+    The keys are the names for the objects, which are used as unique identifiers throughout the code.
+    """

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/scene_entity_cfg.py

@@ -7,7 +7,7 @@
 
 from dataclasses import MISSING
 
-from omni.isaac.lab.assets import Articulation, RigidObject
+from omni.isaac.lab.assets import Articulation, RigidObject, RigidObjectCollection
 from omni.isaac.lab.scene import InteractiveScene
 from omni.isaac.lab.utils import configclass
 
@@ -78,16 +78,34 @@ class SceneEntityCfg:
     manager.
     """
 
+    object_collection_names: str | list[str] | None = None
+    """The names of the objects in the rigid object collection required by the term. Defaults to None.
+
+    The names can be either names or a regular expression matching the object names in the collection.
+
+    These are converted to object indices on initialization of the manager and passed to the term
+    function as a list of object indices under :attr:`object_collection_ids`.
+    """
+
+    object_collection_ids: list[int] | slice = slice(None)
+    """The indices of the objects from the rigid object collection required by the term. Defaults to slice(None),
+    which means all the objects in the collection.
+
+    If :attr:`object_collection_names` is specified, this is filled in automatically on initialization of the manager.
+    """
+
     preserve_order: bool = False
-    """Whether to preserve indices ordering to match with that in the specified joint or body names. Defaults to False.
+    """Whether to preserve indices ordering to match with that in the specified joint, body, or object collection names.
+    Defaults to False.
 
-    If False, the ordering of the indices are sorted in ascending order (i.e. the ordering in the entity's joints
-    or bodies). Otherwise, the indices are preserved in the order of the specified joint and body names.
+    If False, the ordering of the indices are sorted in ascending order (i.e. the ordering in the entity's joints,
+    bodies, or object in the object collection). Otherwise, the indices are preserved in the order of the specified
+    joint, body, or object collection names.
 
     For more details, see the :meth:`omni.isaac.lab.utils.string.resolve_matching_names` function.
 
     .. note::
-        This attribute is only used when :attr:`joint_names` or :attr:`body_names` are specified.
+        This attribute is only used when :attr:`joint_names`, :attr:`body_names`, or :attr:`object_collection_names` are specified.
 
     """
 
@@ -106,6 +124,7 @@ class SceneEntityCfg:
             ValueError: If both ``joint_names`` and ``joint_ids`` are specified and are not consistent.
             ValueError: If both ``fixed_tendon_names`` and ``fixed_tendon_ids`` are specified and are not consistent.
             ValueError: If both ``body_names`` and ``body_ids`` are specified and are not consistent.
+            ValueError: If both ``object_collection_names`` and ``object_collection_ids`` are specified and are not consistent.
         """
         # check if the entity is valid
         if self.name not in scene.keys():
@@ -120,6 +139,9 @@ class SceneEntityCfg:
         # convert body names to indices based on regex
         self._resolve_body_names(scene)
 
+        # convert object collection names to indices based on regex
+        self._resolve_object_collection_names(scene)
+
     def _resolve_joint_names(self, scene: InteractiveScene):
         # convert joint names to indices based on regex
         if self.joint_names is not None or self.joint_ids != slice(None):
@@ -228,3 +250,36 @@ class SceneEntityCfg:
                 if isinstance(self.body_ids, int):
                     self.body_ids = [self.body_ids]
                 self.body_names = [entity.body_names[i] for i in self.body_ids]
+
+    def _resolve_object_collection_names(self, scene: InteractiveScene):
+        # convert object names to indices based on regex
+        if self.object_collection_names is not None or self.object_collection_ids != slice(None):
+            entity: RigidObjectCollection = scene[self.name]
+            # -- if both are not their default values, check if they are valid
+            if self.object_collection_names is not None and self.object_collection_ids != slice(None):
+                if isinstance(self.object_collection_names, str):
+                    self.object_collection_names = [self.object_collection_names]
+                if isinstance(self.object_collection_ids, int):
+                    self.object_collection_ids = [self.object_collection_ids]
+                object_ids, _ = entity.find_objects(self.object_collection_names, preserve_order=self.preserve_order)
+                object_names = [entity.object_names[i] for i in self.object_collection_ids]
+                if object_ids != self.object_collection_ids or object_names != self.object_collection_names:
+                    raise ValueError(
+                        "Both 'object_collection_names' and 'object_collection_ids' are specified, and are not"
+                        " consistent.\n\tfrom object collection names:"
+                        f" {self.object_collection_names} [{object_ids}]\n\tfrom object collection ids:"
+                        f" {object_names} [{self.object_collection_ids}]\nHint: Use either 'object_collection_names' or"
+                        " 'object_collection_ids' to avoid confusion."
+                    )
+            # -- from object names to object indices
+            elif self.object_collection_names is not None:
+                if isinstance(self.object_collection_names, str):
+                    self.object_collection_names = [self.object_collection_names]
+                self.object_collection_ids, _ = entity.find_objects(
+                    self.object_collection_names, preserve_order=self.preserve_order
+                )
+            # -- from object indices to object names
+            elif self.object_collection_ids != slice(None):
+                if isinstance(self.object_collection_ids, int):
+                    self.object_collection_ids = [self.object_collection_ids]
+                self.object_collection_names = [entity.object_names[i] for i in self.object_collection_ids]

source/extensions/omni.isaac.lab/omni/isaac/lab/scene/interactive_scene.py

@@ -22,6 +22,8 @@ from omni.isaac.lab.assets import (
     DeformableObjectCfg,
     RigidObject,
     RigidObjectCfg,
+    RigidObjectCollection,
+    RigidObjectCollectionCfg,
 )
 from omni.isaac.lab.sensors import ContactSensorCfg, FrameTransformerCfg, SensorBase, SensorBaseCfg
 from omni.isaac.lab.terrains import TerrainImporter, TerrainImporterCfg
@@ -113,6 +115,7 @@ class InteractiveScene:
         self._articulations = dict()
         self._deformable_objects = dict()
         self._rigid_objects = dict()
+        self._rigid_object_collections = dict()
         self._sensors = dict()
         self._extras = dict()
         # obtain the current stage
@@ -309,6 +312,11 @@ class InteractiveScene:
         """A dictionary of rigid objects in the scene."""
         return self._rigid_objects
 
+    @property
+    def rigid_object_collections(self) -> dict[str, RigidObjectCollection]:
+        """A dictionary of rigid object collections in the scene."""
+        return self._rigid_object_collections
+
     @property
     def sensors(self) -> dict[str, SensorBase]:
         """A dictionary of the sensors in the scene, such as cameras and contact reporters."""
@@ -351,6 +359,8 @@ class InteractiveScene:
             deformable_object.reset(env_ids)
         for rigid_object in self._rigid_objects.values():
             rigid_object.reset(env_ids)
+        for rigid_object_collection in self._rigid_object_collections.values():
+            rigid_object_collection.reset(env_ids)
         # -- sensors
         for sensor in self._sensors.values():
             sensor.reset(env_ids)
@@ -364,6 +374,8 @@ class InteractiveScene:
             deformable_object.write_data_to_sim()
         for rigid_object in self._rigid_objects.values():
             rigid_object.write_data_to_sim()
+        for rigid_object_collection in self._rigid_object_collections.values():
+            rigid_object_collection.write_data_to_sim()
 
     def update(self, dt: float) -> None:
         """Update the scene entities.
@@ -378,6 +390,8 @@ class InteractiveScene:
             deformable_object.update(dt)
         for rigid_object in self._rigid_objects.values():
             rigid_object.update(dt)
+        for rigid_object_collection in self._rigid_object_collections.values():
+            rigid_object_collection.update(dt)
         # -- sensors
         for sensor in self._sensors.values():
             sensor.update(dt, force_recompute=not self.cfg.lazy_sensor_update)
@@ -397,6 +411,7 @@ class InteractiveScene:
             self._articulations,
             self._deformable_objects,
             self._rigid_objects,
+            self._rigid_object_collections,
             self._sensors,
             self._extras,
         ]:
@@ -422,6 +437,7 @@ class InteractiveScene:
             self._articulations,
             self._deformable_objects,
             self._rigid_objects,
+            self._rigid_object_collections,
             self._sensors,
             self._extras,
         ]:
@@ -454,6 +470,7 @@ class InteractiveScene:
             if asset_name in InteractiveSceneCfg.__dataclass_fields__ or asset_cfg is None:
                 continue
             # resolve regex
+            if hasattr(asset_cfg, "prim_path"):
                 asset_cfg.prim_path = asset_cfg.prim_path.format(ENV_REGEX_NS=self.env_regex_ns)
             # create asset
             if isinstance(asset_cfg, TerrainImporterCfg):
@@ -467,6 +484,14 @@ class InteractiveScene:
                 self._deformable_objects[asset_name] = asset_cfg.class_type(asset_cfg)
             elif isinstance(asset_cfg, RigidObjectCfg):
                 self._rigid_objects[asset_name] = asset_cfg.class_type(asset_cfg)
+            elif isinstance(asset_cfg, RigidObjectCollectionCfg):
+                for rigid_object_cfg in asset_cfg.rigid_objects.values():
+                    rigid_object_cfg.prim_path = rigid_object_cfg.prim_path.format(ENV_REGEX_NS=self.env_regex_ns)
+                self._rigid_object_collections[asset_name] = asset_cfg.class_type(asset_cfg)
+                for rigid_object_cfg in asset_cfg.rigid_objects.values():
+                    if hasattr(rigid_object_cfg, "collision_group") and rigid_object_cfg.collision_group == -1:
+                        asset_paths = sim_utils.find_matching_prim_paths(rigid_object_cfg.prim_path)
+                        self._global_prim_paths += asset_paths
             elif isinstance(asset_cfg, SensorBaseCfg):
                 # Update target frame path(s)' regex name space for FrameTransformer
                 if isinstance(asset_cfg, FrameTransformerCfg):

source/standalone/demos/multi_asset.py

@@ -23,7 +23,7 @@ from omni.isaac.lab.app import AppLauncher
 
 # add argparse arguments
 parser = argparse.ArgumentParser(description="Demo on spawning different objects in multiple environments.")
-parser.add_argument("--num_envs", type=int, default=1024, help="Number of environments to spawn.")
+parser.add_argument("--num_envs", type=int, default=512, help="Number of environments to spawn.")
 # append AppLauncher cli args
 AppLauncher.add_app_launcher_args(parser)
 # parse the arguments
@@ -41,7 +41,15 @@ import omni.usd
 from pxr import Gf, Sdf
 
 import omni.isaac.lab.sim as sim_utils
-from omni.isaac.lab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg
+from omni.isaac.lab.assets import (
+    Articulation,
+    ArticulationCfg,
+    AssetBaseCfg,
+    RigidObject,
+    RigidObjectCfg,
+    RigidObjectCollection,
+    RigidObjectCollectionCfg,
+)
 from omni.isaac.lab.scene import InteractiveScene, InteractiveSceneCfg
 from omni.isaac.lab.sim import SimulationContext
 from omni.isaac.lab.utils import Timer, configclass
@@ -124,6 +132,52 @@ class MultiObjectSceneCfg(InteractiveSceneCfg):
         init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 2.0)),
     )
 
+    # object collection
+    object_collection: RigidObjectCollectionCfg = RigidObjectCollectionCfg(
+        rigid_objects={
+            "object_A": RigidObjectCfg(
+                prim_path="/World/envs/env_.*/Object_A",
+                spawn=sim_utils.SphereCfg(
+                    radius=0.1,
+                    visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), metallic=0.2),
+                    rigid_props=sim_utils.RigidBodyPropertiesCfg(
+                        solver_position_iteration_count=4, solver_velocity_iteration_count=0
+                    ),
+                    mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
+                    collision_props=sim_utils.CollisionPropertiesCfg(),
+                ),
+                init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, -0.5, 2.0)),
+            ),
+            "object_B": RigidObjectCfg(
+                prim_path="/World/envs/env_.*/Object_B",
+                spawn=sim_utils.CuboidCfg(
+                    size=(0.1, 0.1, 0.1),
+                    visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), metallic=0.2),
+                    rigid_props=sim_utils.RigidBodyPropertiesCfg(
+                        solver_position_iteration_count=4, solver_velocity_iteration_count=0
+                    ),
+                    mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
+                    collision_props=sim_utils.CollisionPropertiesCfg(),
+                ),
+                init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.5, 2.0)),
+            ),
+            "object_C": RigidObjectCfg(
+                prim_path="/World/envs/env_.*/Object_C",
+                spawn=sim_utils.ConeCfg(
+                    radius=0.1,
+                    height=0.3,
+                    visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), metallic=0.2),
+                    rigid_props=sim_utils.RigidBodyPropertiesCfg(
+                        solver_position_iteration_count=4, solver_velocity_iteration_count=0
+                    ),
+                    mass_props=sim_utils.MassPropertiesCfg(mass=1.0),
+                    collision_props=sim_utils.CollisionPropertiesCfg(),
+                ),
+                init_state=RigidObjectCfg.InitialStateCfg(pos=(0.5, 0.0, 2.0)),
+            ),
+        }
+    )
+
     # articulation
     robot: ArticulationCfg = ArticulationCfg(
         prim_path="/World/envs/env_.*/Robot",
@@ -170,15 +224,16 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene):
     """Runs the simulation loop."""
     # Extract scene entities
     # note: we only do this here for readability.
-    rigid_object = scene["object"]
-    robot = scene["robot"]
+    rigid_object: RigidObject = scene["object"]
+    rigid_object_collection: RigidObjectCollection = scene["object_collection"]
+    robot: Articulation = scene["robot"]
     # Define simulation stepping
     sim_dt = sim.get_physics_dt()
     count = 0
     # Simulation loop
     while simulation_app.is_running():
         # Reset
-        if count % 500 == 0:
+        if count % 250 == 0:
             # reset counter
             count = 0
             # reset the scene entities
@@ -186,6 +241,10 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene):
             root_state = rigid_object.data.default_root_state.clone()
             root_state[:, :3] += scene.env_origins
             rigid_object.write_root_state_to_sim(root_state)
+            # object collection
+            object_state = rigid_object_collection.data.default_object_state.clone()
+            object_state[..., :3] += scene.env_origins.unsqueeze(1)
+            rigid_object_collection.write_object_state_to_sim(object_state)
             # robot
             # -- root state
             root_state = robot.data.default_root_state.clone()