Adapts terms to deal with slice values for SceneEntityCfg.body_ids (#447)

# Description

With commit d682c8dd, the code stopped working. Somehow, the tests did not
catch the issue, but it shows up when you run the workflows. Anyway, this
MR makes the fix necessary to deal with slices.

Fixes https://github.com/NVIDIA-Omniverse/orbit/issues/277 ,
https://github.com/NVIDIA-Omniverse/orbit/issues/276

## Type of change

- Bug fix (non-breaking change which fixes an issue)

## 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
- [ ] 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

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

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

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

 Changelog
 ---------
 
+0.12.4 (2024-03-11)
+~~~~~~~~~~~~~~~~~~~
+
+Fixed
+^^^^^
+
+
+* Adapted randomization terms to deal with ``slice`` for the body indices. Earlier, the terms were not
+  able to handle the slice object and were throwing an error.
+* Added ``slice`` type-hinting to all body and joint related methods in the rigid body and articulation
+  classes. This is to make it clear that the methods can handle both list of indices and slices.
+
+
 0.12.3 (2024-03-11)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/articulation/articulation.py

@@ -235,8 +235,8 @@ class Articulation(RigidObject):
         self,
         position: torch.Tensor,
         velocity: torch.Tensor,
-        joint_ids: Sequence[int] | None = None,
-        env_ids: Sequence[int] | None = None,
+        joint_ids: Sequence[int] | slice | None = None,
+        env_ids: Sequence[int] | slice | None = None,
     ):
         """Write joint positions and velocities to the simulation.
 
@@ -265,7 +265,7 @@ class Articulation(RigidObject):
     def write_joint_stiffness_to_sim(
         self,
         stiffness: torch.Tensor | float,
-        joint_ids: Sequence[int] | None = None,
+        joint_ids: Sequence[int] | slice | None = None,
         env_ids: Sequence[int] | None = None,
     ):
         """Write joint stiffness into the simulation.
@@ -291,7 +291,7 @@ class Articulation(RigidObject):
     def write_joint_damping_to_sim(
         self,
         damping: torch.Tensor | float,
-        joint_ids: Sequence[int] | None = None,
+        joint_ids: Sequence[int] | slice | None = None,
         env_ids: Sequence[int] | None = None,
     ):
         """Write joint damping into the simulation.
@@ -319,7 +319,7 @@ class Articulation(RigidObject):
     def write_joint_effort_limit_to_sim(
         self,
         limits: torch.Tensor | float,
-        joint_ids: Sequence[int] | None = None,
+        joint_ids: Sequence[int] | slice | None = None,
         env_ids: Sequence[int] | None = None,
     ):
         """Write joint effort limits into the simulation.
@@ -349,7 +349,7 @@ class Articulation(RigidObject):
     def write_joint_armature_to_sim(
         self,
         armature: torch.Tensor | float,
-        joint_ids: Sequence[int] | None = None,
+        joint_ids: Sequence[int] | slice | None = None,
         env_ids: Sequence[int] | None = None,
     ):
         """Write joint armature into the simulation.
@@ -374,7 +374,7 @@ class Articulation(RigidObject):
     def write_joint_friction_to_sim(
         self,
         joint_friction: torch.Tensor | float,
-        joint_ids: Sequence[int] | None = None,
+        joint_ids: Sequence[int] | slice | None = None,
         env_ids: Sequence[int] | None = None,
     ):
         """Write joint friction into the simulation.
@@ -401,7 +401,7 @@ class Articulation(RigidObject):
     """
 
     def set_joint_position_target(
-        self, target: torch.Tensor, joint_ids: Sequence[int] | None = None, env_ids: Sequence[int] | None = None
+        self, target: torch.Tensor, joint_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | None = None
     ):
         """Set joint position targets into internal buffers.
 
@@ -423,7 +423,7 @@ class Articulation(RigidObject):
         self._data.joint_pos_target[env_ids, joint_ids] = target
 
     def set_joint_velocity_target(
-        self, target: torch.Tensor, joint_ids: Sequence[int] | None = None, env_ids: Sequence[int] | None = None
+        self, target: torch.Tensor, joint_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | None = None
     ):
         """Set joint velocity targets into internal buffers.
 
@@ -445,7 +445,7 @@ class Articulation(RigidObject):
         self._data.joint_vel_target[env_ids, joint_ids] = target
 
     def set_joint_effort_target(
-        self, target: torch.Tensor, joint_ids: Sequence[int] | None = None, env_ids: Sequence[int] | None = None
+        self, target: torch.Tensor, joint_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | None = None
     ):
         """Set joint efforts into internal buffers.
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/rigid_object/rigid_object.py

@@ -222,7 +222,7 @@ class RigidObject(AssetBase):
         self,
         forces: torch.Tensor,
         torques: torch.Tensor,
-        body_ids: Sequence[int] | None = None,
+        body_ids: Sequence[int] | slice | None = None,
         env_ids: Sequence[int] | None = None,
     ):
         """Set external force and torque to apply on the asset's bodies in their local frame.
@@ -261,6 +261,8 @@ class RigidObject(AssetBase):
             # -- body_ids
             if body_ids is None:
                 body_ids = torch.arange(self.num_bodies, dtype=torch.long, device=self.device)
+            elif isinstance(body_ids, slice):
+                body_ids = torch.arange(self.num_bodies, dtype=torch.long, device=self.device)[body_ids]
             elif not isinstance(body_ids, torch.Tensor):
                 body_ids = torch.tensor(body_ids, dtype=torch.long, device=self.device)
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/randomizations.py

@@ -62,7 +62,12 @@ def randomize_rigid_body_material(
     num_envs = env.scene.num_envs
     # resolve environment ids
     if env_ids is None:
-        env_ids = torch.arange(num_envs)
+        env_ids = torch.arange(num_envs, device="cpu")
+    # resolve body indices
+    if isinstance(asset_cfg.body_ids, slice):
+        body_ids = torch.arange(asset.num_bodies, dtype=torch.int, device="cpu")[asset_cfg.body_ids]
+    else:
+        body_ids = torch.tensor(asset_cfg.body_ids, dtype=torch.int, device="cpu")
 
     # sample material properties from the given ranges
     material_buckets = torch.zeros(num_buckets, 3)
@@ -74,11 +79,10 @@ def randomize_rigid_body_material(
     materials = material_buckets[material_ids]
     # resolve the global body indices from the env_ids and the env_body_ids
     bodies_per_env = asset.body_physx_view.count // num_envs  # - number of bodies per spawned asset
-    indices = torch.tensor(asset_cfg.body_ids, dtype=torch.int).repeat(len(env_ids), 1)
+    indices = body_ids.repeat(len(env_ids), 1)
     indices += env_ids.unsqueeze(1) * bodies_per_env
 
     # set the material properties into the physics simulation
-    # TODO: Need to use CPU tensors for now. Check if this changes in the new release
     asset.body_physx_view.set_material_properties(materials, indices)
 
 
@@ -96,18 +100,22 @@ def add_body_mass(
     num_envs = env.scene.num_envs
     # resolve environment ids
     if env_ids is None:
-        env_ids = torch.arange(num_envs)
+        env_ids = torch.arange(num_envs, device="cpu")
+    # resolve body indices
+    if isinstance(asset_cfg.body_ids, slice):
+        body_ids = torch.arange(asset.num_bodies, dtype=torch.int, device="cpu")[asset_cfg.body_ids]
+    else:
+        body_ids = torch.tensor(asset_cfg.body_ids, dtype=torch.int, device="cpu")
 
     # get the current masses of the bodies (num_assets x num_bodies)
     masses = asset.body_physx_view.get_masses()
     masses += sample_uniform(*mass_range, masses.shape, device=masses.device)
     # resolve the global body indices from the env_ids and the env_body_ids
     bodies_per_env = asset.body_physx_view.count // env.num_envs
-    indices = torch.tensor(asset_cfg.body_ids, dtype=torch.int).repeat(len(env_ids), 1)
+    indices = body_ids.repeat(len(env_ids), 1)
     indices += env_ids.unsqueeze(1) * bodies_per_env
 
     # set the mass into the physics simulation
-    # TODO: Need to use CPU tensors for now. Check if this changes in the new release
     asset.body_physx_view.set_masses(masses, indices)
 
 
@@ -123,7 +131,7 @@ def apply_external_force_torque(
     This function creates a set of random forces and torques sampled from the given ranges. The number of forces
     and torques is equal to the number of bodies times the number of environments. The forces and torques are
     applied to the bodies by calling ``asset.set_external_force_and_torque``. The forces and torques are only
-    applied when ``asset.write_data_to_sim()`` is called.
+    applied when ``asset.write_data_to_sim()`` is called in the environment.
     """
     # extract the used quantities (to enable type-hinting)
     asset: RigidObject = env.scene[asset_cfg.name]
@@ -131,9 +139,11 @@ def apply_external_force_torque(
     # resolve environment ids
     if env_ids is None:
         env_ids = torch.arange(num_envs)
+    # resolve number of bodies
+    num_bodies = len(asset_cfg.body_ids) if isinstance(asset_cfg.body_ids, list) else asset.num_bodies
 
     # sample random forces and torques
-    size = (len(env_ids), len(asset_cfg.body_ids), 3)
+    size = (len(env_ids), num_bodies, 3)
     forces = sample_uniform(*force_range, size, asset.device)
     torques = sample_uniform(*torque_range, size, asset.device)
     # set the forces and torques into the buffers