# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
from __future__ import annotations
import re
import torch
import weakref
from collections.abc import Sequence
from typing import TYPE_CHECKING
import omni.kit.app
import omni.log
import omni.physics.tensors.impl.api as physx
import omni.timeline
from pxr import UsdPhysics
import omni.isaac.lab.sim as sim_utils
import omni.isaac.lab.utils.math as math_utils
import omni.isaac.lab.utils.string as string_utils
from ..asset_base import AssetBase
from .rigid_object_collection_data import RigidObjectCollectionData
if TYPE_CHECKING:
from .rigid_object_collection_cfg import RigidObjectCollectionCfg
[docs]class RigidObjectCollection(AssetBase):
"""A rigid object collection class.
This class represents a collection of rigid objects in the simulation, where the state of the rigid objects can be
accessed and modified using a batched ``(env_ids, object_ids)`` API.
For each rigid body in the collection, the root prim of the asset must have the `USD RigidBodyAPI`_
applied to it. This API is used to define the simulation properties of the rigid bodies. On playing the
simulation, the physics engine will automatically register the rigid bodies and create a corresponding
rigid body handle. This handle can be accessed using the :attr:`root_physx_view` attribute.
.. note::
Rigid objects in the collection are uniquely identified via the key of the dictionary
:attr:`~omni.isaac.lab.assets.RigidObjectCollectionCfg.rigid_objects` in :class:`~omni.isaac.lab.assets.RigidObjectCollectionCfg`.
This differs from the class :class:`~omni.isaac.lab.assets.RigidObject`, where a rigid object is identified by
the name of the Xform where the `USD RigidBodyAPI`_ is applied. This would not be possible for the rigid object
collection since the :attr:`~omni.isaac.lab.assets.RigidObjectCollectionCfg.rigid_objects` dictionary could
contain the same rigid object multiple times, leading to ambiguity.
.. _`USD RigidBodyAPI`: https://openusd.org/dev/api/class_usd_physics_rigid_body_a_p_i.html
"""
cfg: RigidObjectCollectionCfg
"""Configuration instance for the rigid object collection."""
[docs] def __init__(self, cfg: RigidObjectCollectionCfg):
"""Initialize the rigid object collection.
Args:
cfg: A configuration instance.
"""
# check that the config is valid
cfg.validate()
# store inputs
self.cfg = cfg
# flag for whether the asset is initialized
self._is_initialized = False
for rigid_object_cfg in self.cfg.rigid_objects.values():
# check if the rigid object path is valid
# note: currently the spawner does not work if there is a regex pattern in the leaf
# For example, if the prim path is "/World/Object_[1,2]" since the spawner will not
# know which prim to spawn. This is a limitation of the spawner and not the asset.
asset_path = rigid_object_cfg.prim_path.split("/")[-1]
asset_path_is_regex = re.match(r"^[a-zA-Z0-9/_]+$", asset_path) is None
# spawn the asset
if rigid_object_cfg.spawn is not None and not asset_path_is_regex:
rigid_object_cfg.spawn.func(
rigid_object_cfg.prim_path,
rigid_object_cfg.spawn,
translation=rigid_object_cfg.init_state.pos,
orientation=rigid_object_cfg.init_state.rot,
)
# check that spawn was successful
matching_prims = sim_utils.find_matching_prims(rigid_object_cfg.prim_path)
if len(matching_prims) == 0:
raise RuntimeError(f"Could not find prim with path {rigid_object_cfg.prim_path}.")
# stores object names
self._object_names_list = []
# note: Use weakref on all callbacks to ensure that this object can be deleted when its destructor is called.
# add callbacks for stage play/stop
# The order is set to 10 which is arbitrary but should be lower priority than the default order of 0
timeline_event_stream = omni.timeline.get_timeline_interface().get_timeline_event_stream()
self._initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
int(omni.timeline.TimelineEventType.PLAY),
lambda event, obj=weakref.proxy(self): obj._initialize_callback(event),
order=10,
)
self._invalidate_initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
int(omni.timeline.TimelineEventType.STOP),
lambda event, obj=weakref.proxy(self): obj._invalidate_initialize_callback(event),
order=10,
)
self._debug_vis_handle = None
"""
Properties
"""
@property
def data(self) -> RigidObjectCollectionData:
return self._data
@property
def num_instances(self) -> int:
"""Number of instances of the collection."""
return self.root_physx_view.count // self.num_objects
@property
def num_objects(self) -> int:
"""Number of objects in the collection.
This corresponds to the distinct number of rigid bodies in the collection.
"""
return len(self.object_names)
@property
def object_names(self) -> list[str]:
"""Ordered names of objects in the rigid object collection."""
return self._object_names_list
@property
def root_physx_view(self) -> physx.RigidBodyView:
"""Rigid body view for the rigid body collection (PhysX).
Note:
Use this view with caution. It requires handling of tensors in a specific way.
"""
return self._root_physx_view
"""
Operations.
"""
[docs] def reset(self, env_ids: torch.Tensor | None = None, object_ids: slice | torch.Tensor | None = None):
"""Resets all internal buffers of selected environments and objects.
Args:
env_ids: The indices of the object to reset. Defaults to None (all instances).
object_ids: The indices of the object to reset. Defaults to None (all objects).
"""
# resolve all indices
if env_ids is None:
env_ids = self._ALL_ENV_INDICES
if object_ids is None:
object_ids = self._ALL_OBJ_INDICES
# reset external wrench
self._external_force_b[env_ids[:, None], object_ids] = 0.0
self._external_torque_b[env_ids[:, None], object_ids] = 0.0
[docs] def write_data_to_sim(self):
"""Write external wrench to the simulation.
Note:
We write external wrench to the simulation here since this function is called before the simulation step.
This ensures that the external wrench is applied at every simulation step.
"""
# write external wrench
if self.has_external_wrench:
self.root_physx_view.apply_forces_and_torques_at_position(
force_data=self.reshape_data_to_view(self._external_force_b),
torque_data=self.reshape_data_to_view(self._external_torque_b),
position_data=None,
indices=self._env_obj_ids_to_view_ids(self._ALL_ENV_INDICES, self._ALL_OBJ_INDICES),
is_global=False,
)
[docs] def update(self, dt: float):
self._data.update(dt)
"""
Operations - Finders.
"""
[docs] def find_objects(
self, name_keys: str | Sequence[str], preserve_order: bool = False
) -> tuple[torch.Tensor, list[str]]:
"""Find objects in the collection based on the name keys.
Please check the :meth:`omni.isaac.lab.utils.string_utils.resolve_matching_names` function for more
information on the name matching.
Args:
name_keys: A regular expression or a list of regular expressions to match the object names.
preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False.
Returns:
A tuple containing the object indices and names.
"""
obj_ids, obj_names = string_utils.resolve_matching_names(name_keys, self.object_names, preserve_order)
return torch.tensor(obj_ids, device=self.device), obj_names
"""
Operations - Write to simulation.
"""
[docs] def write_object_state_to_sim(
self,
object_state: torch.Tensor,
env_ids: torch.Tensor | None = None,
object_ids: slice | torch.Tensor | None = None,
):
"""Set the object state over selected environment and object indices into the simulation.
The object state comprises of the cartesian position, quaternion orientation in (w, x, y, z), and linear
and angular velocity. All the quantities are in the simulation frame.
Args:
object_state: Object state in simulation frame. Shape is (len(env_ids), len(object_ids), 13).
env_ids: Environment indices. If None, then all indices are used.
object_ids: Object indices. If None, then all indices are used.
"""
# set into simulation
self.write_object_pose_to_sim(object_state[..., :7], env_ids=env_ids, object_ids=object_ids)
self.write_object_velocity_to_sim(object_state[..., 7:], env_ids=env_ids, object_ids=object_ids)
[docs] def write_object_pose_to_sim(
self,
object_pose: torch.Tensor,
env_ids: torch.Tensor | None = None,
object_ids: slice | torch.Tensor | None = None,
):
"""Set the object pose over selected environment and object indices into the simulation.
The object pose comprises of the cartesian position and quaternion orientation in (w, x, y, z).
Args:
object_pose: Object poses in simulation frame. Shape is (len(env_ids), len(object_ids), 7).
env_ids: Environment indices. If None, then all indices are used.
object_ids: Object indices. If None, then all indices are used.
"""
# resolve all indices
# -- env_ids
if env_ids is None:
env_ids = self._ALL_ENV_INDICES
# -- object_ids
if object_ids is None:
object_ids = self._ALL_OBJ_INDICES
# note: we need to do this here since tensors are not set into simulation until step.
# set into internal buffers
self._data.object_state_w[env_ids[:, None], object_ids, :7] = object_pose.clone()
# convert the quaternion from wxyz to xyzw
poses_xyzw = self._data.object_state_w[..., :7].clone()
poses_xyzw[..., 3:] = math_utils.convert_quat(poses_xyzw[..., 3:], to="xyzw")
# set into simulation
view_ids = self._env_obj_ids_to_view_ids(env_ids, object_ids)
self.root_physx_view.set_transforms(self.reshape_data_to_view(poses_xyzw), indices=view_ids)
[docs] def write_object_velocity_to_sim(
self,
object_velocity: torch.Tensor,
env_ids: torch.Tensor | None = None,
object_ids: slice | torch.Tensor | None = None,
):
"""Set the object velocity over selected environment and object indices into the simulation.
Args:
object_velocity: Object velocities in simulation frame. Shape is (len(env_ids), len(object_ids), 6).
env_ids: Environment indices. If None, then all indices are used.
object_ids: Object indices. If None, then all indices are used.
"""
# resolve all indices
# -- env_ids
if env_ids is None:
env_ids = self._ALL_ENV_INDICES
# -- object_ids
if object_ids is None:
object_ids = self._ALL_OBJ_INDICES
self._data.object_state_w[env_ids[:, None], object_ids, 7:] = object_velocity.clone()
self._data.object_acc_w[env_ids[:, None], object_ids] = 0.0
# set into simulation
view_ids = self._env_obj_ids_to_view_ids(env_ids, object_ids)
self.root_physx_view.set_velocities(
self.reshape_data_to_view(self._data.object_state_w[..., 7:]), indices=view_ids
)
"""
Operations - Setters.
"""
[docs] def set_external_force_and_torque(
self,
forces: torch.Tensor,
torques: torch.Tensor,
object_ids: slice | torch.Tensor | None = None,
env_ids: torch.Tensor | None = None,
):
"""Set external force and torque to apply on the objects' bodies in their local frame.
For many applications, we want to keep the applied external force on rigid bodies constant over a period of
time (for instance, during the policy control). This function allows us to store the external force and torque
into buffers which are then applied to the simulation at every step.
.. caution::
If the function is called with empty forces and torques, then this function disables the application
of external wrench to the simulation.
.. code-block:: python
# example of disabling external wrench
asset.set_external_force_and_torque(forces=torch.zeros(0, 0, 3), torques=torch.zeros(0, 0, 3))
.. note::
This function does not apply the external wrench to the simulation. It only fills the buffers with
the desired values. To apply the external wrench, call the :meth:`write_data_to_sim` function
right before the simulation step.
Args:
forces: External forces in bodies' local frame. Shape is (len(env_ids), len(object_ids), 3).
torques: External torques in bodies' local frame. Shape is (len(env_ids), len(object_ids), 3).
object_ids: Object indices to apply external wrench to. Defaults to None (all objects).
env_ids: Environment indices to apply external wrench to. Defaults to None (all instances).
"""
if forces.any() or torques.any():
self.has_external_wrench = True
# resolve all indices
# -- env_ids
if env_ids is None:
env_ids = self._ALL_ENV_INDICES
# -- object_ids
if object_ids is None:
object_ids = self._ALL_OBJ_INDICES
# set into internal buffers
self._external_force_b[env_ids[:, None], object_ids] = forces
self._external_torque_b[env_ids[:, None], object_ids] = torques
else:
self.has_external_wrench = False
"""
Internal helper.
"""
def _initialize_impl(self):
# create simulation view
self._physics_sim_view = physx.create_simulation_view(self._backend)
self._physics_sim_view.set_subspace_roots("/")
root_prim_path_exprs = []
for name, rigid_object_cfg in self.cfg.rigid_objects.items():
# obtain the first prim in the regex expression (all others are assumed to be a copy of this)
template_prim = sim_utils.find_first_matching_prim(rigid_object_cfg.prim_path)
if template_prim is None:
raise RuntimeError(f"Failed to find prim for expression: '{rigid_object_cfg.prim_path}'.")
template_prim_path = template_prim.GetPath().pathString
# find rigid root prims
root_prims = sim_utils.get_all_matching_child_prims(
template_prim_path, predicate=lambda prim: prim.HasAPI(UsdPhysics.RigidBodyAPI)
)
if len(root_prims) == 0:
raise RuntimeError(
f"Failed to find a rigid body when resolving '{rigid_object_cfg.prim_path}'."
" Please ensure that the prim has 'USD RigidBodyAPI' applied."
)
if len(root_prims) > 1:
raise RuntimeError(
f"Failed to find a single rigid body when resolving '{rigid_object_cfg.prim_path}'."
f" Found multiple '{root_prims}' under '{template_prim_path}'."
" Please ensure that there is only one rigid body in the prim path tree."
)
# check that no rigid object has an articulation root API, which decreases simulation performance
articulation_prims = sim_utils.get_all_matching_child_prims(
template_prim_path, predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI)
)
if len(articulation_prims) != 0:
if articulation_prims[0].GetAttribute("physxArticulation:articulationEnabled").Get():
raise RuntimeError(
f"Found an articulation root when resolving '{rigid_object_cfg.prim_path}' in the rigid object"
f" collection. These are located at: '{articulation_prims}' under '{template_prim_path}'."
" Please disable the articulation root in the USD or from code by setting the parameter"
" 'ArticulationRootPropertiesCfg.articulation_enabled' to False in the spawn configuration."
)
# resolve root prim back into regex expression
root_prim_path = root_prims[0].GetPath().pathString
root_prim_path_expr = rigid_object_cfg.prim_path + root_prim_path[len(template_prim_path) :]
root_prim_path_exprs.append(root_prim_path_expr.replace(".*", "*"))
self._object_names_list.append(name)
# -- object view
self._root_physx_view = self._physics_sim_view.create_rigid_body_view(root_prim_path_exprs)
# check if the rigid body was created
if self._root_physx_view._backend is None:
raise RuntimeError("Failed to create rigid body collection. Please check PhysX logs.")
# log information about the rigid body
omni.log.info(f"Number of instances: {self.num_instances}")
omni.log.info(f"Number of distinct objects: {self.num_objects}")
omni.log.info(f"Object names: {self.object_names}")
# container for data access
self._data = RigidObjectCollectionData(self.root_physx_view, self.num_objects, self.device)
# create buffers
self._create_buffers()
# process configuration
self._process_cfg()
# update the rigid body data
self.update(0.0)
def _create_buffers(self):
"""Create buffers for storing data."""
# constants
self._ALL_ENV_INDICES = torch.arange(self.num_instances, dtype=torch.long, device=self.device)
self._ALL_OBJ_INDICES = torch.arange(self.num_objects, dtype=torch.long, device=self.device)
# external forces and torques
self.has_external_wrench = False
self._external_force_b = torch.zeros((self.num_instances, self.num_objects, 3), device=self.device)
self._external_torque_b = torch.zeros_like(self._external_force_b)
# set information about rigid body into data
self._data.object_names = self.object_names
self._data.default_mass = self.reshape_view_to_data(self.root_physx_view.get_masses().clone())
self._data.default_inertia = self.reshape_view_to_data(self.root_physx_view.get_inertias().clone())
def _process_cfg(self):
"""Post processing of configuration parameters."""
# default state
# -- object state
default_object_states = []
for rigid_object_cfg in self.cfg.rigid_objects.values():
default_object_state = (
tuple(rigid_object_cfg.init_state.pos)
+ tuple(rigid_object_cfg.init_state.rot)
+ tuple(rigid_object_cfg.init_state.lin_vel)
+ tuple(rigid_object_cfg.init_state.ang_vel)
)
default_object_state = (
torch.tensor(default_object_state, dtype=torch.float, device=self.device)
.repeat(self.num_instances, 1)
.unsqueeze(1)
)
default_object_states.append(default_object_state)
# concatenate the default state for each object
default_object_states = torch.cat(default_object_states, dim=1)
self._data.default_object_state = default_object_states
[docs] def reshape_view_to_data(self, data: torch.Tensor) -> torch.Tensor:
"""Reshapes and arranges the data coming from the :attr:`root_physx_view` to (num_instances, num_objects, data_size).
Args:
data: The data coming from the :attr:`root_physx_view`. Shape is (num_instances*num_objects, data_size).
Returns:
The reshaped data. Shape is (num_instances, num_objects, data_size).
"""
return torch.einsum("ijk -> jik", data.reshape(self.num_objects, self.num_instances, -1))
[docs] def reshape_data_to_view(self, data: torch.Tensor) -> torch.Tensor:
"""Reshapes and arranges the data to the be consistent with data from the :attr:`root_physx_view`.
Args:
data: The data to be reshaped. Shape is (num_instances, num_objects, data_size).
Returns:
The reshaped data. Shape is (num_instances*num_objects, data_size).
"""
return torch.einsum("ijk -> jik", data).reshape(self.num_objects * self.num_instances, *data.shape[2:])
def _env_obj_ids_to_view_ids(
self, env_ids: torch.Tensor, object_ids: Sequence[int] | slice | torch.Tensor
) -> torch.Tensor:
"""Converts environment and object indices to indices consistent with data from :attr:`root_physx_view`.
Args:
env_ids: Environment indices.
object_ids: Object indices.
Returns:
The view indices.
"""
# the order is env_0/object_0, env_0/object_1, env_0/object_..., env_1/object_0, env_1/object_1, ...
# return a flat tensor of indices
if isinstance(object_ids, slice):
object_ids = self._ALL_OBJ_INDICES
elif isinstance(object_ids, Sequence):
object_ids = torch.tensor(object_ids, device=self.device)
return (object_ids.unsqueeze(1) * self.num_instances + env_ids).flatten()
"""
Internal simulation callbacks.
"""
def _invalidate_initialize_callback(self, event):
"""Invalidates the scene elements."""
# call parent
super()._invalidate_initialize_callback(event)
# set all existing views to None to invalidate them
self._physics_sim_view = None
self._root_physx_view = None
