Adds support for keyword arguments into ManagerBase (#198)

# Description

This MR adds support for keyword arguments into the `ManagerBase` class.
This helps in defining some default arguments into the term function
call, which makes the configuration for those terms easier. Earlier, we
had a lot of `SceneEntity("robot")` references, which is redundant for
most cases that we are interested in.

## Type of change

- New feature (non-breaking change which adds functionality)
- Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- This change requires a documentation update

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

---------

Signed-off-by: Mayank Mittal <12863862+Mayankm96@users.noreply.github.com>
Co-authored-by: Farbod Farshidian <ffarshidian@theaiinstitute.com>

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

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.9.15"
+version = "0.9.16"
 
 # Description
 title = "ORBIT framework for Robot Learning"

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

@@ -1,6 +1,21 @@
 Changelog
 ---------
 
+0.9.16 (2023-10-22)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added support for keyword arguments for terms in the :class:`omni.isaac.orbit.managers.ManagerBase`.
+
+Fixed
+^^^^^
+
+* Fixed resetting of buffers in the :class:`TerminationManager` class. Earlier, the values were being set
+  to ``0.0`` instead of ``False``.
+
+
 0.9.15 (2023-10-22)
 ~~~~~~~~~~~~~~~~~~~
 

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

@@ -22,7 +22,9 @@
     from omni.isaac.orbit.envs.rl_env import RLEnv
 
 
-def terrain_levels_vel(env: RLEnv, env_ids: Sequence[int], asset_cfg: SceneEntityCfg) -> torch.Tensor:
+def terrain_levels_vel(
+    env: RLEnv, env_ids: Sequence[int], asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
+) -> torch.Tensor:
     """Curriculum based on the distance the robot walked when commanded to move at a desired velocity.
 
     This term is used to increase the difficulty of the terrain when the robot walks far enough and decrease the

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

@@ -26,21 +26,21 @@
 """
 
 
-def base_lin_vel(env: BaseEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+def base_lin_vel(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
     """Root linear velocity in the asset's root frame."""
     # extract the used quantities (to enable type-hinting)
     asset: RigidObject = env.scene[asset_cfg.name]
     return asset.data.root_lin_vel_b
 
 
-def base_ang_vel(env: BaseEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+def base_ang_vel(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
     """Root angular velocity in the asset's root frame."""
     # extract the used quantities (to enable type-hinting)
     asset: RigidObject = env.scene[asset_cfg.name]
     return asset.data.root_ang_vel_b
 
 
-def projected_gravity(env: BaseEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+def projected_gravity(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
     """Gravity projection on the asset's root frame."""
     # extract the used quantities (to enable type-hinting)
     asset: RigidObject = env.scene[asset_cfg.name]
@@ -52,14 +52,14 @@ def projected_gravity(env: BaseEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
 """
 
 
-def joint_pos_rel(env: BaseEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+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."""
     # 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
 
 
-def joint_vel_rel(env: BaseEnv, asset_cfg: SceneEntityCfg):
+def joint_vel_rel(env: BaseEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")):
     """The joint velocities of the asset w.r.t. the default joint velocities."""
     # extract the used quantities (to enable type-hinting)
     asset: Articulation = env.scene[asset_cfg.name]
@@ -71,16 +71,12 @@ def joint_vel_rel(env: BaseEnv, asset_cfg: SceneEntityCfg):
 """
 
 
-def height_scan(env: BaseEnv, asset_cfg: SceneEntityCfg, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
-    """Height scan from the given sensor w.r.t. the asset's root frame."""
+def height_scan(env: BaseEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
+    """Height scan from the given sensor w.r.t. the sensor's frame."""
     # extract the used quantities (to enable type-hinting)
-    asset: RigidObject = env.scene[asset_cfg.name]
     sensor: RayCaster = env.scene.sensors[sensor_cfg.name]
-    # TODO (@dhoeller): is this sensor specific or we can generalize it?
-    hit_points_z = torch.nan_to_num(sensor.data.ray_hits_w[..., 2], posinf=-1.0)
-    # compute the height scan: robot_z - ground_z - offset
-    heights = asset.data.root_state_w[:, 2].unsqueeze(1) - hit_points_z - 0.5
     # return the height scan
+    heights = sensor.data.pos_w[:, 2].unsqueeze(1) - sensor.data.ray_hits_w[..., 2] - 0.5
     return heights
 
 
@@ -89,7 +85,7 @@ def height_scan(env: BaseEnv, asset_cfg: SceneEntityCfg, sensor_cfg: SceneEntity
 """
 
 
-def action(env: BaseEnv) -> torch.Tensor:
+def last_action(env: BaseEnv) -> torch.Tensor:
     """The last input action to the environment."""
     return env.action_manager.action
 

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

@@ -28,11 +28,11 @@
 def randomize_rigid_body_material(
     env: RLEnv,
     env_ids: torch.Tensor | None,
-    asset_cfg: SceneEntityCfg,
     static_friction_range: tuple[float, float],
     dynamic_friction_range: tuple[float, float],
     restitution_range: tuple[float, float],
     num_buckets: int,
+    asset_cfg: SceneEntityCfg,
 ):
     """Randomize the physics materials on all geometries of the asset.
 
@@ -79,7 +79,7 @@ def randomize_rigid_body_material(
     asset.body_physx_view.set_material_properties(materials, indices)
 
 
-def add_body_mass(env: RLEnv, env_ids: torch.Tensor | None, asset_cfg: SceneEntityCfg, mass_range: tuple[float, float]):
+def add_body_mass(env: RLEnv, env_ids: torch.Tensor | None, mass_range: tuple[float, float], asset_cfg: SceneEntityCfg):
     """Randomize the mass of the bodies by adding a random value sampled from the given range.
 
     .. tip::
@@ -109,9 +109,9 @@ def add_body_mass(env: RLEnv, env_ids: torch.Tensor | None, asset_cfg: SceneEnti
 def apply_external_force_torque(
     env: RLEnv,
     env_ids: torch.Tensor,
-    asset_cfg: SceneEntityCfg,
     force_range: tuple[float, float],
     torque_range: tuple[float, float],
+    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
 ):
     """Randomize the external forces and torques applied to the bodies.
 
@@ -137,7 +137,10 @@ def apply_external_force_torque(
 
 
 def push_by_setting_velocity(
-    env: RLEnv, env_ids: torch.Tensor, asset_cfg: SceneEntityCfg, velocity_range: dict[str, tuple[float, float]]
+    env: RLEnv,
+    env_ids: torch.Tensor,
+    velocity_range: dict[str, tuple[float, float]],
+    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
 ):
     """Push the asset by setting the root velocity to a random value within the given ranges.
 
@@ -167,9 +170,9 @@ def push_by_setting_velocity(
 def reset_root_state(
     env: RLEnv,
     env_ids: torch.Tensor,
-    asset_cfg: SceneEntityCfg,
     pose_range: dict[str, tuple[float, float]],
     velocity_range: dict[str, tuple[float, float]],
+    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
 ):
     """Reset the asset root state to a random position and velocity within the given ranges.
 
@@ -218,9 +221,9 @@ def reset_root_state(
 def reset_joints_by_scale(
     env: RLEnv,
     env_ids: torch.Tensor,
-    asset_cfg: SceneEntityCfg,
     position_range: tuple[float, float],
     velocity_range: tuple[float, float],
+    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
 ):
     """Reset the robot joints by scaling the default position and velocity by the given ranges.
 

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

@@ -21,27 +21,36 @@
 if TYPE_CHECKING:
     from omni.isaac.orbit.envs.rl_env import RLEnv
 
+"""
+General.
+"""
+
+
+def termination_penalty(env: RLEnv) -> torch.Tensor:
+    """Penalize terminated episodes that don't correspond to episodic timeouts."""
+    return env.reset_buf * (~env.termination_manager.time_outs)
+
 
 """
 Root penalties.
 """
 
 
-def lin_vel_z_l2(env: RLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+def lin_vel_z_l2(env: RLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
     """Penalize z-axis base linear velocity using L2-kernel."""
     # extract the used quantities (to enable type-hinting)
     asset: RigidObject = env.scene[asset_cfg.name]
     return torch.square(asset.data.root_lin_vel_b[:, 2])
 
 
-def ang_vel_xy_l2(env: RLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+def ang_vel_xy_l2(env: RLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
     """Penalize xy-axis base angular velocity using L2-kernel."""
     # extract the used quantities (to enable type-hinting)
     asset: RigidObject = env.scene[asset_cfg.name]
     return torch.sum(torch.square(asset.data.root_ang_vel_b[:, :2]), dim=1)
 
 
-def flat_orientation_l2(env: RLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+def flat_orientation_l2(env: RLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
     """Penalize non-flat base orientation using L2-kernel.
 
     This is computed by penalizing the xy-components of the projected gravity vector.
@@ -51,7 +60,9 @@ def flat_orientation_l2(env: RLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
     return torch.sum(torch.square(asset.data.projected_gravity_b[:, :2]), dim=1)
 
 
-def base_height_l2(env: RLEnv, asset_cfg: SceneEntityCfg, target_height: float) -> torch.Tensor:
+def base_height_l2(
+    env: RLEnv, target_height: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
+) -> torch.Tensor:
     """Penalize asset height from its target using L2-kernel.
 
     Note:
@@ -63,33 +74,39 @@ def base_height_l2(env: RLEnv, asset_cfg: SceneEntityCfg, target_height: float)
     return torch.square(asset.data.root_pos_w[:, 2] - target_height)
 
 
+def body_lin_acc_l2(env: RLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+    """Penalize the linear acceleration of bodies using L2-kernel."""
+    asset: Articulation = env.scene[asset_cfg.name]
+    return torch.sum(torch.norm(asset.data.body_lin_acc_w[:, asset_cfg.body_ids, :], dim=-1), dim=1)
+
+
 """
 Joint penalties.
 """
 
 
-def joint_torques_l2(env: RLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+def joint_torques_l2(env: RLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
     """Penalize torques applied on the articulation using L2-kernel."""
     # 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)
 
 
-def joint_vel_l2(env: RLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+def joint_vel_l2(env: RLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
     """Penalize joint velocities on the articulation."""
     # 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)
 
 
-def joint_acc_l2(env: RLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+def joint_acc_l2(env: RLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
     """Penalize joint accelerations on the articulation using L2-kernel."""
     # 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)
 
 
-def joint_pos_limits(env: RLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+def joint_pos_limits(env: RLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
     """Penalize joint positions if they cross the soft limits.
 
     This is computed as a sum of the absolute value of the difference between the joint position and the soft limits.
@@ -102,7 +119,9 @@ def joint_pos_limits(env: RLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
     return torch.sum(out_of_limits, dim=1)
 
 
-def joint_vel_limits(env: RLEnv, asset_cfg: SceneEntityCfg, soft_ratio: float) -> torch.Tensor:
+def joint_vel_limits(
+    env: RLEnv, soft_ratio: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
+) -> torch.Tensor:
     """Penalize joint velocities if they cross the soft limits.
 
     This is computed as a sum of the absolute value of the difference between the joint velocity and the soft limits.
@@ -124,7 +143,7 @@ def joint_vel_limits(env: RLEnv, asset_cfg: SceneEntityCfg, soft_ratio: float) -
 """
 
 
-def applied_torque_limits(env: RLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor:
+def applied_torque_limits(env: RLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
     """Penalize applied torques if they cross the limits.
 
     This is computed as a sum of the absolute value of the difference between the applied torques and the limits.
@@ -151,7 +170,7 @@ def action_rate_l2(env: RLEnv) -> torch.Tensor:
 """
 
 
-def undesired_contacts(env: RLEnv, sensor_cfg: SceneEntityCfg, threshold: float) -> torch.Tensor:
+def undesired_contacts(env: RLEnv, threshold: float, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
     """Penalize undesired contacts as the number of violations that are above a threshold."""
     # extract the used quantities (to enable type-hinting)
     contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
@@ -162,7 +181,7 @@ def undesired_contacts(env: RLEnv, sensor_cfg: SceneEntityCfg, threshold: float)
     return torch.sum(is_contact, dim=1)
 
 
-def contact_forces(env: RLEnv, sensor_cfg: SceneEntityCfg, threshold: float) -> torch.Tensor:
+def contact_forces(env: RLEnv, threshold: float, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
     """Penalize contact forces as the amount of violations of the net contact force."""
     # extract the used quantities (to enable type-hinting)
     contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
@@ -178,7 +197,7 @@ def contact_forces(env: RLEnv, sensor_cfg: SceneEntityCfg, threshold: float) ->
 """
 
 
-def track_lin_vel_xy_exp(env: RLEnv, asset_cfg: SceneEntityCfg, std: float) -> torch.Tensor:
+def track_lin_vel_xy_exp(env: RLEnv, std: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
     """Reward tracking of linear velocity commands (xy axes) using exponential kernel."""
     # extract the used quantities (to enable type-hinting)
     asset: RigidObject = env.scene[asset_cfg.name]
@@ -189,7 +208,7 @@ def track_lin_vel_xy_exp(env: RLEnv, asset_cfg: SceneEntityCfg, std: float) -> t
     return torch.exp(-lin_vel_error / std**2)
 
 
-def track_ang_vel_z_exp(env: RLEnv, asset_cfg: SceneEntityCfg, std: float) -> torch.Tensor:
+def track_ang_vel_z_exp(env: RLEnv, std: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
     """Reward tracking of angular velocity commands (yaw) using exponential kernel."""
     # extract the used quantities (to enable type-hinting)
     asset: RigidObject = env.scene[asset_cfg.name]