Fixes joint and body sub-indexing for observations and rewards (#434)

# Description

Replaces the default values for `joint_ids` and `body_ids` from `None`
to `slice(None)` in the `SceneEntityCfg` and adapted rewards and
observations so that the user can query a sub-set of joints and bodies.

## 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
- [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.11.0"
+version = "0.11.1"
 
 # Description
 title = "ORBIT framework for Robot Learning"

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

 Changelog
 ---------
 
+0.11.1 (2024-02-29)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Replaced the default values for ``joint_ids`` and ``body_ids`` from ``None`` to ``slice(None)``
+  in the :class:`omni.isaac.orbit.managers.SceneEntityCfg`.
+* Adapted rewards and observations terms so that the users can query a subset of joints and bodies.
+
+
 0.11.0 (2024-02-27)
 ~~~~~~~~~~~~~~~~~~~
 

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

@@ -61,26 +61,37 @@ Joint state.
 
 
 def joint_pos_rel(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
-    """The joint positions of the asset w.r.t. the default joint positions."""
+    """The joint positions of the asset w.r.t. the default joint positions.
+
+    NOTE: Only the joints configured in :attr:`asset_cfg.joint_ids` will have their positions returned.
+    """
     # extract the used quantities (to enable type-hinting)
     asset: Articulation = env.scene[asset_cfg.name]
-    return asset.data.joint_pos - asset.data.default_joint_pos
+    return asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]
 
 
 def joint_pos_norm(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
-    """The joint positions of the asset normalized with the asset's joint limits."""
+    """The joint positions of the asset normalized with the asset's joint limits.
+
+    NOTE: Only the joints configured in :attr:`asset_cfg.joint_ids` will have their normalized positions returned.
+    """
     # extract the used quantities (to enable type-hinting)
     asset: Articulation = env.scene[asset_cfg.name]
     return math_utils.scale_transform(
-        asset.data.joint_pos, asset.data.soft_joint_pos_limits[..., 0], asset.data.soft_joint_pos_limits[..., 1]
+        asset.data.joint_pos[:, asset_cfg.joint_ids],
+        asset.data.soft_joint_pos_limits[:, asset_cfg.joint_ids, 0],
+        asset.data.soft_joint_pos_limits[:, asset_cfg.joint_ids, 1],
     )
 
 
 def joint_vel_rel(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")):
-    """The joint velocities of the asset w.r.t. the default joint velocities."""
+    """The joint velocities of the asset w.r.t. the default joint velocities.
+
+    NOTE: Only the joints configured in :attr:`asset_cfg.joint_ids` will have their velocities returned.
+    """
     # extract the used quantities (to enable type-hinting)
     asset: Articulation = env.scene[asset_cfg.name]
-    return asset.data.joint_vel - asset.data.default_joint_vel
+    return asset.data.joint_vel[:, asset_cfg.joint_ids] - asset.data.default_joint_vel[:, asset_cfg.joint_ids]
 
 
 """

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

@@ -91,10 +91,13 @@ Joint penalties.
 
 
 def joint_torques_l2(env: RLTaskEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
-    """Penalize joint torques applied on the articulation using L2-kernel."""
+    """Penalize joint torques applied on the articulation using L2-kernel.
+
+    NOTE: Only the joints configured in :attr:`asset_cfg.joint_ids` will have their joint torques contribute to the L2 norm.
+    """
     # extract the used quantities (to enable type-hinting)
     asset: Articulation = env.scene[asset_cfg.name]
-    return torch.sum(torch.square(asset.data.applied_torque), dim=1)
+    return torch.sum(torch.square(asset.data.applied_torque[:, asset_cfg.joint_ids]), dim=1)
 
 
 def joint_vel_l1(env: RLTaskEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
@@ -105,17 +108,23 @@ def joint_vel_l1(env: RLTaskEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
 
 
 def joint_vel_l2(env: RLTaskEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
-    """Penalize joint velocities on the articulation using L2-kernel."""
+    """Penalize joint velocities on the articulation using L1-kernel.
+
+    NOTE: Only the joints configured in :attr:`asset_cfg.joint_ids` will have their joint velocities contribute to the L1 norm.
+    """
     # extract the used quantities (to enable type-hinting)
     asset: Articulation = env.scene[asset_cfg.name]
-    return torch.sum(torch.square(asset.data.joint_vel), dim=1)
+    return torch.sum(torch.square(asset.data.joint_vel[:, asset_cfg.joint_ids]), dim=1)
 
 
 def joint_acc_l2(env: RLTaskEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
-    """Penalize joint accelerations on the articulation using L2-kernel."""
+    """Penalize joint accelerations on the articulation using L2-kernel.
+
+    NOTE: Only the joints configured in :attr:`asset_cfg.joint_ids` will have their joint accelerations contribute to the L2 norm.
+    """
     # extract the used quantities (to enable type-hinting)
     asset: Articulation = env.scene[asset_cfg.name]
-    return torch.sum(torch.square(asset.data.joint_acc), dim=1)
+    return torch.sum(torch.square(asset.data.joint_acc[:, asset_cfg.joint_ids]), dim=1)
 
 
 def joint_deviation_l1(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
@@ -157,7 +166,10 @@ def joint_vel_limits(
     # extract the used quantities (to enable type-hinting)
     asset: Articulation = env.scene[asset_cfg.name]
     # compute out of limits constraints
-    out_of_limits = torch.abs(asset.data.joint_vel) - asset.data.soft_joint_vel_limits * soft_ratio
+    out_of_limits = (
+        torch.abs(asset.data.joint_vel[:, asset_cfg.joint_ids])
+        - asset.data.soft_joint_vel_limits[:, asset_cfg.joint_ids] * soft_ratio
+    )
     # clip to max error = 1 rad/s per joint to avoid huge penalties
     out_of_limits = out_of_limits.clip_(min=0.0, max=1.0)
     return torch.sum(out_of_limits, dim=1)
@@ -181,7 +193,9 @@ def applied_torque_limits(env: RLTaskEnv, asset_cfg: SceneEntityCfg = SceneEntit
     asset: Articulation = env.scene[asset_cfg.name]
     # compute out of limits constraints
     # TODO: We need to fix this to support implicit joints.
-    out_of_limits = torch.abs(asset.data.applied_torque - asset.data.computed_torque)
+    out_of_limits = torch.abs(
+        asset.data.applied_torque[:, asset_cfg.joint_ids] - asset.data.computed_torque[:, asset_cfg.joint_ids]
+    )
     return torch.sum(out_of_limits, dim=1)
 
 

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

@@ -38,8 +38,9 @@ class SceneEntityCfg:
     function as a list of joint indices under :attr:`joint_ids`.
     """
 
-    joint_ids: list[int] | None = None
-    """The indices of the joints from the asset required by the term. Defaults to None.
+    joint_ids: list[int] | slice = slice(None)
+    """The indices of the joints from the asset required by the term. Defaults to slice(None), which means
+    all the joints in the asset (if present).
 
     If :attr:`joint_names` is specified, this is filled in automatically on initialization of the
     manager.
@@ -54,8 +55,9 @@ class SceneEntityCfg:
     function as a list of body indices under :attr:`body_ids`.
     """
 
-    body_ids: list[int] | None = None
-    """The indices of the bodies from the asset required by the term. Defaults to None.
+    body_ids: list[int] | slice = slice(None)
+    """The indices of the bodies from the asset required by the term. Defaults to slice(None), which means
+    all the bodies in the asset.
 
     If :attr:`body_names` is specified, this is filled in automatically on initialization of the
     manager.
@@ -79,11 +81,12 @@ class SceneEntityCfg:
         # check if the entity is valid
         if self.name not in scene.keys():
             raise ValueError(f"The scene entity '{self.name}' does not exist. Available entities: {scene.keys()}.")
+
         # convert joint names to indices based on regex
-        if self.joint_names is not None or self.joint_ids is not None:
+        if self.joint_names is not None or self.joint_ids != slice(None):
             entity: Articulation = scene[self.name]
-            # -- if both are not None, check if they are valid
-            if self.joint_names is not None and self.joint_ids is not None:
+            # -- if both are not their default values, check if they are valid
+            if self.joint_names is not None and self.joint_ids != slice(None):
                 if isinstance(self.joint_names, str):
                     self.joint_names = [self.joint_names]
                 if isinstance(self.joint_ids, int):
@@ -102,16 +105,20 @@ class SceneEntityCfg:
                 if isinstance(self.joint_names, str):
                     self.joint_names = [self.joint_names]
                 self.joint_ids, _ = entity.find_joints(self.joint_names)
+                # performance optimization (slice offers faster indexing than list of indices)
+                if len(self.joint_ids) == entity.num_joints:
+                    self.joint_ids = slice(None)
             # -- from joint indices to joint names
-            elif self.joint_ids is not None:
+            elif self.joint_ids != slice(None):
                 if isinstance(self.joint_ids, int):
                     self.joint_ids = [self.joint_ids]
                 self.joint_names = [entity.joint_names[i] for i in self.joint_ids]
+
         # convert body names to indices based on regex
-        if self.body_names is not None or self.body_ids is not None:
+        if self.body_names is not None or self.body_ids != slice(None):
             entity: RigidObject = scene[self.name]
-            # -- if both are not None, check if they are valid
-            if self.body_names is not None and self.body_ids is not None:
+            # -- if both are not their default values, check if they are valid
+            if self.body_names is not None and self.body_ids != slice(None):
                 if isinstance(self.body_names, str):
                     self.body_names = [self.body_names]
                 if isinstance(self.body_ids, int):
@@ -130,8 +137,11 @@ class SceneEntityCfg:
                 if isinstance(self.body_names, str):
                     self.body_names = [self.body_names]
                 self.body_ids, _ = entity.find_bodies(self.body_names)
+                # performance optimization (slice offers faster indexing than list of indices)
+                if len(self.body_ids) == entity.num_bodies:
+                    self.body_ids = slice(None)
             # -- from body indices to body names
-            elif self.body_ids is not None:
+            elif self.body_ids != slice(None):
                 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]