Fixes resolution of indices inside the asset class (#251)

# Description

To resolve all environment indices, the previous implementation used
`self._ALL_INDICES`. However, this is not broadcastable when tensors are
indexed with two lists, for example:

```python
import torch

x = torch.zeros(100, 25)
x[[0, 2, 3, 4], [1, 2, 3]]
```

This MR fixes this bug (introduced in #248).

## 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`
- [ ] 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
- [ ] 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/omni/isaac/orbit/assets/articulation/articulation.py

@@ -161,8 +161,10 @@ class Articulation(RigidObject):
 
     def write_root_pose_to_sim(self, root_pose: torch.Tensor, env_ids: Sequence[int] | None = None):
         # resolve all indices
+        physx_env_ids = env_ids
         if env_ids is None:
-            env_ids = self._ALL_INDICES
+            env_ids = slice(None)
+            physx_env_ids = self._ALL_INDICES
         # note: we need to do this here since tensors are not set into simulation until step.
         # set into internal buffers
         self._data.root_state_w[env_ids, :7] = root_pose.clone()
@@ -170,17 +172,19 @@ class Articulation(RigidObject):
         root_poses_xyzw = self._data.root_state_w[:, :7].clone()
         root_poses_xyzw[:, 3:] = math_utils.convert_quat(root_poses_xyzw[:, 3:], to="xyzw")
         # set into simulation
-        self.root_physx_view.set_root_transforms(root_poses_xyzw, indices=env_ids)
+        self.root_physx_view.set_root_transforms(root_poses_xyzw, indices=physx_env_ids)
 
     def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None):
         # resolve all indices
+        physx_env_ids = env_ids
         if env_ids is None:
-            env_ids = self._ALL_INDICES
+            env_ids = slice(None)
+            physx_env_ids = self._ALL_INDICES
         # note: we need to do this here since tensors are not set into simulation until step.
         # set into internal buffers
         self._data.root_state_w[env_ids, 7:] = root_velocity.clone()
         # set into simulation
-        self.root_physx_view.set_root_velocities(self._data.root_state_w[:, 7:], indices=env_ids)
+        self.root_physx_view.set_root_velocities(self._data.root_state_w[:, 7:], indices=physx_env_ids)
 
     def write_joint_state_to_sim(
         self,
@@ -198,8 +202,10 @@ class Articulation(RigidObject):
             env_ids: The environment indices to set the targets for. Defaults to None (all environments).
         """
         # resolve indices
+        physx_env_ids = env_ids
         if env_ids is None:
-            env_ids = self._ALL_INDICES
+            env_ids = slice(None)
+            physx_env_ids = self._ALL_INDICES
         if joint_ids is None:
             joint_ids = slice(None)
         # set into internal buffers
@@ -208,8 +214,8 @@ class Articulation(RigidObject):
         self._previous_joint_vel[env_ids, joint_ids] = velocity
         self._data.joint_acc[env_ids, joint_ids] = 0.0
         # set into simulation
-        self.root_physx_view.set_dof_positions(self._data.joint_pos, indices=env_ids)
-        self.root_physx_view.set_dof_velocities(self._data.joint_vel, indices=env_ids)
+        self.root_physx_view.set_dof_positions(self._data.joint_pos, indices=physx_env_ids)
+        self.root_physx_view.set_dof_velocities(self._data.joint_vel, indices=physx_env_ids)
 
     def write_joint_stiffness_to_sim(
         self,
@@ -226,14 +232,16 @@ class Articulation(RigidObject):
         """
         # note: This function isn't setting the values for actuator models. (#128)
         # resolve indices
+        physx_env_ids = env_ids
         if env_ids is None:
-            env_ids = self._ALL_INDICES
+            env_ids = slice(None)
+            physx_env_ids = self._ALL_INDICES
         if joint_ids is None:
             joint_ids = slice(None)
         # set into internal buffers
         self._data.joint_stiffness[env_ids, joint_ids] = stiffness
         # set into simulation
-        self.root_physx_view.set_dof_stiffnesses(self._data.joint_stiffness.cpu(), indices=env_ids.cpu())
+        self.root_physx_view.set_dof_stiffnesses(self._data.joint_stiffness.cpu(), indices=physx_env_ids.cpu())
 
     def write_joint_damping_to_sim(
         self, damping: torch.Tensor, joint_ids: Sequence[int] | None = None, env_ids: Sequence[int] | None = None
@@ -249,14 +257,16 @@ class Articulation(RigidObject):
         """
         # note: This function isn't setting the values for actuator models. (#128)
         # resolve indices
+        physx_env_ids = env_ids
         if env_ids is None:
-            env_ids = self._ALL_INDICES
+            env_ids = slice(None)
+            physx_env_ids = self._ALL_INDICES
         if joint_ids is None:
             joint_ids = slice(None)
         # set into internal buffers
         self._data.joint_damping[env_ids, joint_ids] = damping
         # set into simulation
-        self.root_physx_view.set_dof_dampings(self._data.joint_damping.cpu(), indices=env_ids.cpu())
+        self.root_physx_view.set_dof_dampings(self._data.joint_damping.cpu(), indices=physx_env_ids.cpu())
 
     def write_joint_torque_limit_to_sim(
         self,
@@ -273,15 +283,17 @@ class Articulation(RigidObject):
         """
         # note: This function isn't setting the values for actuator models. (#128)
         # resolve indices
+        physx_env_ids = env_ids
         if env_ids is None:
-            env_ids = self._ALL_INDICES
+            env_ids = slice(None)
+            physx_env_ids = self._ALL_INDICES
         if joint_ids is None:
             joint_ids = slice(None)
         # set into internal buffers
         torque_limit_all = self.root_physx_view.get_dof_max_forces()
         torque_limit_all[env_ids, joint_ids] = limits
         # set into simulation
-        self.root_physx_view.set_dof_max_forces(torque_limit_all.cpu(), indices=env_ids.cpu())
+        self.root_physx_view.set_dof_max_forces(torque_limit_all.cpu(), indices=physx_env_ids.cpu())
 
     """
     Operations - State.

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

@@ -174,8 +174,10 @@ class RigidObject(AssetBase):
             env_ids: Environment indices. If :obj:`None`, then all indices are used.
         """
         # resolve all indices
+        physx_env_ids = env_ids
         if env_ids is None:
-            env_ids = self._ALL_INDICES
+            env_ids = slice(None)
+            physx_env_ids = self._ALL_INDICES
         # note: we need to do this here since tensors are not set into simulation until step.
         # set into internal buffers
         self._data.root_state_w[env_ids, :7] = root_pose.clone()
@@ -183,7 +185,7 @@ class RigidObject(AssetBase):
         root_poses_xyzw = self._data.root_state_w[:, :7].clone()
         root_poses_xyzw[:, 3:] = math_utils.convert_quat(root_poses_xyzw[:, 3:], to="xyzw")
         # set into simulation
-        self.root_physx_view.set_transforms(root_poses_xyzw, indices=env_ids)
+        self.root_physx_view.set_transforms(root_poses_xyzw, indices=physx_env_ids)
 
     def write_root_velocity_to_sim(self, root_velocity: torch.Tensor, env_ids: Sequence[int] | None = None):
         """Set the root velocity over selected environment indices into the simulation.
@@ -193,13 +195,15 @@ class RigidObject(AssetBase):
             env_ids: Environment indices. If :obj:`None`, then all indices are used.
         """
         # resolve all indices
+        physx_env_ids = env_ids
         if env_ids is None:
-            env_ids = self._ALL_INDICES
+            env_ids = slice(None)
+            physx_env_ids = self._ALL_INDICES
         # note: we need to do this here since tensors are not set into simulation until step.
         # set into internal buffers
         self._data.root_state_w[env_ids, 7:] = root_velocity.clone()
         # set into simulation
-        self.root_physx_view.set_velocities(self._data.root_state_w[:, 7:], indices=env_ids)
+        self.root_physx_view.set_velocities(self._data.root_state_w[:, 7:], indices=physx_env_ids)
 
     """
     Operations - Setters.