From 7b398407c1c760e175b1926e7941b6852c5c3d27 Mon Sep 17 00:00:00 2001 From: Calm_dw <89777028+calmdw@users.noreply.github.com> Date: Tue, 8 Apr 2025 22:19:01 +0200 Subject: [PATCH] add events, observations and terminations file Signed-off-by: Calm_dw <89777028+calmdw@users.noreply.github.com> --- .../mdp/__init__.py | 3 + .../manager-based_single-agent/mdp/events.py | 83 +++++++++++++++++++ .../mdp/observations.py | 37 +++++++++ .../mdp/terminations.py | 37 +++++++++ 4 files changed, 160 insertions(+) create mode 100644 tools/template/templates/tasks/manager-based_single-agent/mdp/events.py create mode 100644 tools/template/templates/tasks/manager-based_single-agent/mdp/observations.py create mode 100644 tools/template/templates/tasks/manager-based_single-agent/mdp/terminations.py diff --git a/tools/template/templates/tasks/manager-based_single-agent/mdp/__init__.py b/tools/template/templates/tasks/manager-based_single-agent/mdp/__init__.py index 44872114b8a..9aeb3fd3a5e 100644 --- a/tools/template/templates/tasks/manager-based_single-agent/mdp/__init__.py +++ b/tools/template/templates/tasks/manager-based_single-agent/mdp/__init__.py @@ -8,3 +8,6 @@ from isaaclab.envs.mdp import * # noqa: F401, F403 from .rewards import * # noqa: F401, F403 +from .events import * +from .observations import * +from .terminations import * \ No newline at end of file diff --git a/tools/template/templates/tasks/manager-based_single-agent/mdp/events.py b/tools/template/templates/tasks/manager-based_single-agent/mdp/events.py new file mode 100644 index 00000000000..7f4e4a3ad80 --- /dev/null +++ b/tools/template/templates/tasks/manager-based_single-agent/mdp/events.py @@ -0,0 +1,83 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Common functions that can be used to enable different events. + +Events include anything related to altering the simulation state. This includes changing the physics +materials, applying external forces, and resetting the state of the asset. + +The functions can be passed to the :class:`isaaclab.managers.EventTermCfg` object to enable +the event introduced by the function. +""" + +from __future__ import annotations + +import math +import torch +from typing import TYPE_CHECKING, Literal + +import carb +import omni.physics.tensors.impl.api as physx +import omni.usd +from isaacsim.core.utils.extensions import enable_extension +from pxr import Gf, Sdf, UsdGeom, Vt + +import isaaclab.sim as sim_utils +import isaaclab.utils.math as math_utils +from isaaclab.actuators import ImplicitActuator +from isaaclab.assets import Articulation, DeformableObject, RigidObject +from isaaclab.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg +from isaaclab.terrains import TerrainImporter + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + +""" +customized events. +""" + +def reset_root_state_uniform( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + 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 uniformly within the given ranges. + + This function randomizes the root position and velocity of the asset. + + * It samples the root position from the given ranges and adds them to the default root position, before setting + them into the physics simulation. + * It samples the root orientation from the given ranges and sets them into the physics simulation. + * It samples the root velocity from the given ranges and sets them into the physics simulation. + + The function takes a dictionary of pose and velocity ranges for each axis and rotation. The keys of the + dictionary are ``x``, ``y``, ``z``, ``roll``, ``pitch``, and ``yaw``. The values are tuples of the form + ``(min, max)``. If the dictionary does not contain a key, the position or velocity is set to zero for that axis. + """ + # extract the used quantities (to enable type-hinting) + asset: RigidObject | Articulation = env.scene[asset_cfg.name] + # get default root state + root_states = asset.data.default_root_state[env_ids].clone() + + # poses + range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] + ranges = torch.tensor(range_list, device=asset.device) + rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=asset.device) + + positions = root_states[:, 0:3] + env.scene.env_origins[env_ids] + rand_samples[:, 0:3] + orientations_delta = math_utils.quat_from_euler_xyz(rand_samples[:, 3], rand_samples[:, 4], rand_samples[:, 5]) + orientations = math_utils.quat_mul(root_states[:, 3:7], orientations_delta) + # velocities + range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] + ranges = torch.tensor(range_list, device=asset.device) + rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=asset.device) + + velocities = root_states[:, 7:13] + rand_samples + + # set into the physics simulation + asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) + asset.write_root_velocity_to_sim(velocities, env_ids=env_ids) \ No newline at end of file diff --git a/tools/template/templates/tasks/manager-based_single-agent/mdp/observations.py b/tools/template/templates/tasks/manager-based_single-agent/mdp/observations.py new file mode 100644 index 00000000000..11c64318987 --- /dev/null +++ b/tools/template/templates/tasks/manager-based_single-agent/mdp/observations.py @@ -0,0 +1,37 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Common functions that can be used to create observation terms. + +The functions can be passed to the :class:`isaaclab.managers.ObservationTermCfg` object to enable +the observation introduced by the function. +""" + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils +from isaaclab.assets import Articulation, RigidObject +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers.manager_base import ManagerTermBase +from isaaclab.managers.manager_term_cfg import ObservationTermCfg +from isaaclab.sensors import Camera, Imu, RayCaster, RayCasterCamera, TiledCamera + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv + + +""" +customized observations. +""" + + +def root_pos_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Asset root position in the environment frame.""" + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + return asset.data.root_pos_w - env.scene.env_origins \ No newline at end of file diff --git a/tools/template/templates/tasks/manager-based_single-agent/mdp/terminations.py b/tools/template/templates/tasks/manager-based_single-agent/mdp/terminations.py new file mode 100644 index 00000000000..25c407c851d --- /dev/null +++ b/tools/template/templates/tasks/manager-based_single-agent/mdp/terminations.py @@ -0,0 +1,37 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Common functions that can be used to activate certain terminations. + +The functions can be passed to the :class:`isaaclab.managers.TerminationTermCfg` object to enable +the termination introduced by the function. +""" + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation, RigidObject +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import ContactSensor + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.managers.command_manager import CommandTerm + +""" +customized terminations. +""" + +def root_height_higher_maximum( + env: ManagerBasedRLEnv, maximum_height: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") +) -> torch.Tensor: + """Terminate when the asset's root height is higher the maximum height. + + """ + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + return asset.data.root_pos_w[:, 2] < maximum_height