From fbd4ff04783d278ea26ef988ea79780023199fbc Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Tue, 20 May 2025 20:10:33 +0200 Subject: [PATCH 01/23] very quick prototyping of the Obs outputs --- scripts/environments/random_agent.py | 8 ++ .../isaaclab/envs/manager_based_env.py | 9 ++ .../isaaclab/envs/mdp/observations.py | 116 +++++++++++++++++- 3 files changed, 129 insertions(+), 4 deletions(-) diff --git a/scripts/environments/random_agent.py b/scripts/environments/random_agent.py index b3187c3b372..9422680e992 100644 --- a/scripts/environments/random_agent.py +++ b/scripts/environments/random_agent.py @@ -47,11 +47,19 @@ def main(): # create environment env = gym.make(args_cli.task, cfg=env_cfg) + + # print info (this is vectorized environment) print(f"[INFO]: Gym observation space: {env.observation_space}") print(f"[INFO]: Gym action space: {env.action_space}") # reset environment env.reset() + + out = env.unwrapped.get_IO_descriptors + for o in out: + print(f"--- Obs term: {o.name} ---") + for k, v in o.__dict__.items(): + print(f"{k}: {v}") # simulate environment while simulation_app.is_running(): # run everything in inference mode diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 0dc88b85e2e..20245446a33 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -209,6 +209,15 @@ def step_dt(self) -> float: def device(self): """The device on which the environment is running.""" return self.sim.device + + @property + def get_IO_descriptors(self): + """Get the IO descriptors for the environment. + + Returns: + A dictionary with keys as the group names and values as the IO descriptors. + """ + return self.observation_manager.get_IO_descriptors """ Operations - Setup. diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index 5d2cd7f96f7..b2d0a9a3458 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -24,12 +24,41 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv +from isaaclab.utils import configclass +from dataclasses import MISSING + +@configclass +class IODescriptor: + mdp_type: str = "Observation" + name: str = None + description: str = None + shape: tuple[int, ...] = None + dtype: torch.dtype = None + observation_type: str = None + """ Root state. """ +@configclass +class RootStateIODescriptor(IODescriptor): + observation_type: str = "RootState" + axes: list[str] = None + units: str = None + + +def root_state_io_descriptor(descriptor: RootStateIODescriptor): + def decorator(func): + def wrapper(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), inspect: bool = False): + return func(env, asset_cfg) + descriptor.name = func.__name__ + wrapper._has_descriptor = True + wrapper._descriptor = descriptor + return wrapper + return decorator +@root_state_io_descriptor(RootStateIODescriptor(description="Root height in the world frame.", units="m", axes=["Z"], shape=(1,), dtype=torch.float32)) def base_pos_z(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Root height in the simulation world frame.""" # extract the used quantities (to enable type-hinting) @@ -37,6 +66,7 @@ def base_pos_z(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg( return asset.data.root_pos_w[:, 2].unsqueeze(-1) +@root_state_io_descriptor(RootStateIODescriptor(description="Root linear velocity in the robot's frame.", units="m/s", axes=["X", "Y", "Z"], shape=(3,), dtype=torch.float32)) def base_lin_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Root linear velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) @@ -44,6 +74,7 @@ def base_lin_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCf return asset.data.root_lin_vel_b +@root_state_io_descriptor(RootStateIODescriptor(description="Root angular velocity in the robot's frame.", units="rad/s", axes=["X", "Y", "Z"], shape=(3,), dtype=torch.float32)) def base_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Root angular velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) @@ -51,6 +82,7 @@ def base_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCf return asset.data.root_ang_vel_b +@root_state_io_descriptor(RootStateIODescriptor(description="Projection of gravity in the robot's root frame.", units="m/s^2", axes=["X", "Y", "Z"] , shape=(3,), dtype=torch.float32)) def projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Gravity projection on the asset's root frame.""" # extract the used quantities (to enable type-hinting) @@ -58,6 +90,7 @@ def projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEnt return asset.data.projected_gravity_b +@root_state_io_descriptor(RootStateIODescriptor(description="Root body position in the world frame.", units="m", axes=["X", "Y", "Z"], shape=(3,), dtype=torch.float32)) 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) @@ -65,6 +98,7 @@ def root_pos_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg( return asset.data.root_pos_w - env.scene.env_origins +@root_state_io_descriptor(RootStateIODescriptor(description="Root body orientation in the world frame.", units="unit", axes=["W", "X", "Y", "Z"], shape=(4,), dtype=torch.float32)) def root_quat_w( env: ManagerBasedEnv, make_quat_unique: bool = False, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> torch.Tensor: @@ -82,6 +116,7 @@ def root_quat_w( return math_utils.quat_unique(quat) if make_quat_unique else quat +@root_state_io_descriptor(RootStateIODescriptor(description="Root body linear velocity in the world frame.", units="m/s", axes=["X", "Y", "Z"], shape=(3,), dtype=torch.float32)) def root_lin_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Asset root linear velocity in the environment frame.""" # extract the used quantities (to enable type-hinting) @@ -89,6 +124,7 @@ def root_lin_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntity return asset.data.root_lin_vel_w +@root_state_io_descriptor(RootStateIODescriptor(description="Root body angular velocity in the world frame.", units="rad/s", axes=["X", "Y", "Z"], shape=(3,), dtype=torch.float32)) def root_ang_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Asset root angular velocity in the environment frame.""" # extract the used quantities (to enable type-hinting) @@ -100,7 +136,30 @@ def root_ang_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntity Body state """ +@configclass +class BodyStateIODescriptor(IODescriptor): + observation_type: str = "BodyState" + body_ids: list[int] | int = [] + body_names: list[str] | str = [] + +def body_state_io_descriptor(descriptor: BodyStateIODescriptor, inspect: bool = False): + def decorator(func): + descriptor.name = func.__name__ + def wrapper(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): + if inspect: + out = func(env, asset_cfg) + descriptor.shape = (out.shape[-1],) + descriptor.body_ids = asset_cfg.body_ids + descriptor.body_names = asset_cfg.body_names + return out + else: + return func(env, asset_cfg) + wrapper._has_descriptor = True + wrapper._descriptor = descriptor + return wrapper + return decorator +@body_state_io_descriptor(BodyStateIODescriptor(description="The flattened body poses of the robot in the world frame. The output shape is 7 * num_bodies", dtype=torch.float32)) def body_pose_w( env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), @@ -124,6 +183,7 @@ def body_pose_w( return pose.reshape(env.num_envs, -1) +@body_state_io_descriptor(BodyStateIODescriptor(description="The direction of gravity projected on to bodies own frames. The output shape is 3 * num_bodies", dtype=torch.float32)) def body_projected_gravity_b( env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), @@ -151,8 +211,35 @@ def body_projected_gravity_b( """ Joint state. """ +@configclass +class JointStateIODescriptor(IODescriptor): + observation_type: str = "JointState" + joint_ids: list[int] | int = [] + joint_names: list[str] | str = [] + +def joint_state_io_descriptor(descriptor: JointStateIODescriptor): + def decorator(func): + descriptor.name = func.__name__ + def wrapper(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), inspect: bool = False): + if inspect: + out = func(env, asset_cfg) + descriptor.shape = (out.shape[-1],) + descriptor.joint_ids = asset_cfg.joint_ids + descriptor.joint_names = asset_cfg.joint_names + + if descriptor.joint_names is None: + asset: Articulation = env.scene[asset_cfg.name] + descriptor.joint_names = asset.joint_names + descriptor.joint_ids = list(range(len(asset.joint_names))) + return out + else: + return func(env, asset_cfg) + wrapper._has_descriptor = True + wrapper._descriptor = descriptor + return wrapper + return decorator - +@joint_state_io_descriptor(JointStateIODescriptor(description="The joint positions of the asset. The output shape is num_joints.", dtype=torch.float32)) def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset. @@ -163,6 +250,7 @@ def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" return asset.data.joint_pos[:, asset_cfg.joint_ids] +@joint_state_io_descriptor(JointStateIODescriptor(description="The joint positions of the asset w.r.t. the default joint positions. The output shape is num_joints", dtype=torch.float32)) def joint_pos_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset w.r.t. the default joint positions. @@ -173,6 +261,7 @@ def joint_pos_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC return asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids] +@joint_state_io_descriptor(JointStateIODescriptor(description="The joint positions of the asset normalized with the asset's joint limits. The output shape is num_joints", dtype=torch.float32)) def joint_pos_limit_normalized( env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> torch.Tensor: @@ -188,7 +277,7 @@ def joint_pos_limit_normalized( asset.data.soft_joint_pos_limits[:, asset_cfg.joint_ids, 1], ) - +@joint_state_io_descriptor(JointStateIODescriptor(description="The joint velocities of the asset. The output shape is num_joints", dtype=torch.float32)) def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset. @@ -198,7 +287,7 @@ def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" asset: Articulation = env.scene[asset_cfg.name] return asset.data.joint_vel[:, asset_cfg.joint_ids] - +@joint_state_io_descriptor(JointStateIODescriptor(description="The joint velocities of the asset w.r.t. the default joint velocities. The output shape is num_joints", dtype=torch.float32)) def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset w.r.t. the default joint velocities. @@ -208,7 +297,7 @@ def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC asset: Articulation = env.scene[asset_cfg.name] return asset.data.joint_vel[:, asset_cfg.joint_ids] - asset.data.default_joint_vel[:, asset_cfg.joint_ids] - +@joint_state_io_descriptor(JointStateIODescriptor(description="The joint applied effort of the robot. The output shape is num_joints", dtype=torch.float32)) def joint_effort(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint applied effort of the robot. @@ -593,6 +682,25 @@ def _inference(model, images: torch.Tensor) -> torch.Tensor: """ +@configclass +class ActionIODescriptor(IODescriptor): + observation_type: str = "Action" + +def root_state_io_descriptor(descriptor: RootStateIODescriptor): + def decorator(func): + def wrapper(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), inspect: bool = False): + if inspect: + out = func(env, asset_cfg) + descriptor.shape = (out.shape[-1],) + return out + else: + return func(env, asset_cfg) + descriptor.name = func.__name__ + wrapper._has_descriptor = True + wrapper._descriptor = descriptor + return wrapper + return decorator + def last_action(env: ManagerBasedEnv, action_name: str | None = None) -> torch.Tensor: """The last input action to the environment. From 7336b89f898132f112b375e29c9739b6c586bf28 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Wed, 21 May 2025 12:01:29 +0200 Subject: [PATCH 02/23] It sort of works, but it's messy. Maybe we should use latch on to the resolve function to get more data out of the term. --- .../isaaclab/envs/mdp/observations.py | 18 ++++++---- .../isaaclab/managers/manager_base.py | 19 +++++++++-- .../isaaclab/managers/observation_manager.py | 34 +++++++++++++++++++ 3 files changed, 61 insertions(+), 10 deletions(-) diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index b2d0a9a3458..1edb97f023c 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -683,24 +683,27 @@ def _inference(model, images: torch.Tensor) -> torch.Tensor: @configclass -class ActionIODescriptor(IODescriptor): - observation_type: str = "Action" +class GenericIODescriptor(IODescriptor): + observation_type: str = None + -def root_state_io_descriptor(descriptor: RootStateIODescriptor): +def generic_io_descriptor(descriptor: GenericIODescriptor): def decorator(func): - def wrapper(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), inspect: bool = False): + def wrapper(env: ManagerBasedEnv, *args, inspect: bool = False, **kwargs): if inspect: - out = func(env, asset_cfg) + out = func(env, *args, **kwargs) descriptor.shape = (out.shape[-1],) return out else: - return func(env, asset_cfg) + return func(env, *args, **kwargs) descriptor.name = func.__name__ wrapper._has_descriptor = True wrapper._descriptor = descriptor return wrapper return decorator + +@generic_io_descriptor(GenericIODescriptor(description="The last input action to the environment.", dtype=torch.float32, observation_type="Action")) def last_action(env: ManagerBasedEnv, action_name: str | None = None) -> torch.Tensor: """The last input action to the environment. @@ -718,7 +721,8 @@ def last_action(env: ManagerBasedEnv, action_name: str | None = None) -> torch.T """ -def generated_commands(env: ManagerBasedRLEnv, command_name: str) -> torch.Tensor: +@generic_io_descriptor(GenericIODescriptor(description="The generated command term in the command manager.", dtype=torch.float32, observation_type="Command")) +def generated_commands(env: ManagerBasedRLEnv, command_name: str | None = None) -> torch.Tensor: """The generated command from command term in the command manager with the given name.""" return env.command_manager.get_command(command_name) diff --git a/source/isaaclab/isaaclab/managers/manager_base.py b/source/isaaclab/isaaclab/managers/manager_base.py index 081c3e271ec..14d198eb659 100644 --- a/source/isaaclab/isaaclab/managers/manager_base.py +++ b/source/isaaclab/isaaclab/managers/manager_base.py @@ -359,11 +359,24 @@ def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg, args = inspect.signature(func_static).parameters args_with_defaults = [arg for arg in args if args[arg].default is not inspect.Parameter.empty] args_without_defaults = [arg for arg in args if args[arg].default is inspect.Parameter.empty] - args = args_without_defaults + args_with_defaults # ignore first two arguments for env and env_ids - # Think: Check for cases when kwargs are set inside the function? + # Think: Check for cases when kwargs are set inside the function? + if "kwargs" in args_without_defaults: + args_without_defaults.remove("kwargs") + args_with_defaults.append("kwargs") + args = args_without_defaults + args_with_defaults + + print("args", args) + print("args_without_defaults", args_without_defaults) + print("args_with_defaults", args_with_defaults) + print("min_argc", min_argc) + print("term_params", term_params) + print("term_cfg", term_cfg) + if len(args) > min_argc: - if set(args[min_argc:]) != set(term_params + args_with_defaults): + if "kwargs" in args: + pass + elif set(args[min_argc:]) != set(term_params + args_with_defaults): raise ValueError( f"The term '{term_name}' expects mandatory parameters: {args_without_defaults[min_argc:]}" f" and optional parameters: {args_with_defaults}, but received: {term_params}." diff --git a/source/isaaclab/isaaclab/managers/observation_manager.py b/source/isaaclab/isaaclab/managers/observation_manager.py index 8c0e8104a70..b763ab1e471 100644 --- a/source/isaaclab/isaaclab/managers/observation_manager.py +++ b/source/isaaclab/isaaclab/managers/observation_manager.py @@ -225,6 +225,40 @@ def group_obs_concatenate(self) -> dict[str, bool]: """ return self._group_obs_concatenate + @property + def get_IO_descriptors(self): + """Get the IO descriptors for the observation manager. + + Returns: + A dictionary with keys as the group names and values as the IO descriptors. + """ + + data = [] + + for group_name in self._group_obs_term_names: + # check ig group name is valid + if group_name not in self._group_obs_term_names: + raise ValueError( + f"Unable to find the group '{group_name}' in the observation manager." + f" Available groups are: {list(self._group_obs_term_names.keys())}" + ) + # iterate over all the terms in each group + group_term_names = self._group_obs_term_names[group_name] + # buffer to store obs per group + group_obs = dict.fromkeys(group_term_names, None) + # read attributes for each term + obs_terms = zip(group_term_names, self._group_obs_term_cfgs[group_name]) + + for term_name, term_cfg in obs_terms: + # dummy call to cache some values + try: + term_cfg.func(self._env, **term_cfg.params, inspect=True) + data.append(term_cfg.func._descriptor) + except Exception as e: + print(f"Error getting IO descriptor for term '{term_name}' in group '{group_name}': {e}") + + return data + """ Operations. """ From 961defe23c15ac78349176a7c281c97231858f74 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Tue, 10 Jun 2025 14:55:56 +0200 Subject: [PATCH 03/23] Added a new auto-magik decorator that should help streamline the process. There are some caveats so we'll need good doc for it --- .../isaaclab/envs/mdp/observations.py | 278 +++++++++++------- 1 file changed, 172 insertions(+), 106 deletions(-) diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index 1edb97f023c..1133d02a84a 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -11,6 +11,8 @@ from __future__ import annotations +import functools +import inspect import torch from typing import TYPE_CHECKING @@ -27,38 +29,176 @@ from isaaclab.utils import configclass from dataclasses import MISSING + +import dataclasses, functools, inspect +from typing import Any, Callable, ParamSpec, TypeVar, Concatenate + @configclass class IODescriptor: mdp_type: str = "Observation" name: str = None + full_path: str = None description: str = None shape: tuple[int, ...] = None dtype: torch.dtype = None observation_type: str = None - -""" -Root state. -""" @configclass -class RootStateIODescriptor(IODescriptor): - observation_type: str = "RootState" - axes: list[str] = None - units: str = None +class GenericIODescriptor: + mdp_type: str = "Observation" + name: str = None + full_path: str = None + description: str = None + shape: tuple[int, ...] = None + dtype: torch.dtype = None + observation_type: str = None + extras: dict[str, Any] = None + +# These are defined to help with type hinting +P = ParamSpec("P") +R = TypeVar("R") + +# Automatically builds a descriptor from the kwargs +def _make_descriptor(**kwargs: Any) -> "GenericIODescriptor": + """Split *kwargs* into (known dataclass fields) and (extras).""" + field_names = {f.name for f in dataclasses.fields(GenericIODescriptor)} + known = {k: v for k, v in kwargs.items() if k in field_names} + extras = {k: v for k, v in kwargs.items() if k not in field_names} + + desc = GenericIODescriptor(**known) + # User defined extras are stored in the descriptor under the `extras` field + desc.extras = extras + return desc + +# Decorator factory for generic IO descriptors. +def generic_io_descriptor( + _func: Callable[Concatenate[ManagerBasedEnv, P], R] | None = None, + *, + on_inspect: Callable[..., Any] | list[Callable[..., Any]] | None = None, + **descriptor_kwargs: Any, +) -> Callable[[Callable[Concatenate[ManagerBasedEnv, P], R]], + Callable[Concatenate[ManagerBasedEnv, P], R]]: + """ + Decorator factory for generic IO descriptors. + + This decorator can be used in different ways: + 1. The default decorator has all the information I need for my use case: + ..code-block:: python + @generic_io_descriptor(GenericIODescriptor(description="..", dtype="..")) + def my_func(env: ManagerBasedEnv, *args, **kwargs): + ... + ..note:: If description is not set, the function's docstring is used to populate it. + + 2. I need to add more information to the descriptor: + ..code-block:: python + @generic_io_descriptor(description="..", new_var_1="a", new_var_2="b") + def my_func(env: ManagerBasedEnv, *args, **kwargs): + ... + 3. I need to add a hook to the descriptor: + ..code-block:: python + def record_shape(tensor: torch.Tensor, desc: GenericIODescriptor): + desc.shape = (tensor.shape[-1],) + + @generic_io_descriptor(description="..", new_var_1="a", new_var_2="b", on_inspect=[record_shape]) + def my_func(env: ManagerBasedEnv, *args, **kwargs): + ..note:: The hook is called after the function is called, if and only if the `inspect` flag is set when calling the function. + For example: + ..code-block:: python + my_func(env, inspect=True) + + Args: + _func: The function to decorate. + **descriptor_kwargs: Keyword arguments to pass to the descriptor. + + Returns: + A decorator that can be used to decorate a function. + """ + if _func is not None and isinstance(_func, GenericIODescriptor): + descriptor = _func + _func = None + else: + descriptor = _make_descriptor(**descriptor_kwargs) -def root_state_io_descriptor(descriptor: RootStateIODescriptor): - def decorator(func): - def wrapper(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), inspect: bool = False): - return func(env, asset_cfg) + # Ensures the hook is a list + if callable(on_inspect): + inspect_hooks: list[Callable[..., Any]] = [on_inspect] + else: + inspect_hooks: list[Callable[..., Any]] = list(on_inspect or []) # handles None + + def _apply( + func: Callable[Concatenate[ManagerBasedEnv, P], R] + ) -> Callable[Concatenate[ManagerBasedEnv, P], R]: + + # Capture the signature of the function + sig = inspect.signature(func) + + @functools.wraps(func) + def wrapper( + env: ManagerBasedEnv, *args: P.args, **kwargs: P.kwargs + ) -> R: + inspect_flag: bool = kwargs.pop("inspect", False) + out = func(env, *args, **kwargs) + if inspect_flag: + # Injects the function's arguments into the hooks and applies the defaults + bound = sig.bind(env, *args, **kwargs) + bound.apply_defaults() + call_kwargs = { + "output": out, + "descriptor": descriptor, + **bound.arguments, + } + for hook in inspect_hooks: + hook(**call_kwargs) + return out + + # --- Descriptor bookkeeping --- descriptor.name = func.__name__ - wrapper._has_descriptor = True + descriptor.full_path = f"{func.__module__}.{func.__name__}" + descriptor.dtype = str(descriptor.dtype) + # Check if description is set in the descriptor + if descriptor.description is None: + descriptor.description = func.__doc__ + + # Adds the descriptor to the wrapped function as an attribute wrapper._descriptor = descriptor + wrapper._has_descriptor = True + # Alters the signature of the wrapped function to make it match the original function. + # This allows the wrapped functions to pass the checks in the managers. + wrapper.__signature__ = sig return wrapper - return decorator + # If the decorator is used without parentheses, _func will be the function itself. + if callable(_func): + return _apply(_func) + return _apply + + +def record_shape(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): + descriptor.shape = (output.shape[-1],) + +def record_dtype(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): + descriptor.dtype = str(output.dtype) + +def record_joint_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): + asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] + joint_ids = kwargs["asset_cfg"].joint_ids + if joint_ids == slice(None, None, None): + joint_ids = list(range(len(asset.joint_names))) + descriptor.joint_names = [asset.joint_names[i] for i in joint_ids] + +def record_body_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): + asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] + body_ids = kwargs["asset_cfg"].body_ids + if body_ids == slice(None, None, None): + body_ids = list(range(len(asset.body_names))) + descriptor.body_names = [asset.body_names[i] for i in body_ids] + +""" +Root state. +""" -@root_state_io_descriptor(RootStateIODescriptor(description="Root height in the world frame.", units="m", axes=["Z"], shape=(1,), dtype=torch.float32)) +@generic_io_descriptor(units="m", axes=["Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) def base_pos_z(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Root height in the simulation world frame.""" # extract the used quantities (to enable type-hinting) @@ -66,7 +206,7 @@ def base_pos_z(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg( return asset.data.root_pos_w[:, 2].unsqueeze(-1) -@root_state_io_descriptor(RootStateIODescriptor(description="Root linear velocity in the robot's frame.", units="m/s", axes=["X", "Y", "Z"], shape=(3,), dtype=torch.float32)) +@generic_io_descriptor(units="m/s", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) def base_lin_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Root linear velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) @@ -74,7 +214,7 @@ def base_lin_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCf return asset.data.root_lin_vel_b -@root_state_io_descriptor(RootStateIODescriptor(description="Root angular velocity in the robot's frame.", units="rad/s", axes=["X", "Y", "Z"], shape=(3,), dtype=torch.float32)) +@generic_io_descriptor(units="rad/s", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) def base_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Root angular velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) @@ -82,7 +222,7 @@ def base_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCf return asset.data.root_ang_vel_b -@root_state_io_descriptor(RootStateIODescriptor(description="Projection of gravity in the robot's root frame.", units="m/s^2", axes=["X", "Y", "Z"] , shape=(3,), dtype=torch.float32)) +@generic_io_descriptor(units="m/s^2", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) def projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Gravity projection on the asset's root frame.""" # extract the used quantities (to enable type-hinting) @@ -90,7 +230,7 @@ def projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEnt return asset.data.projected_gravity_b -@root_state_io_descriptor(RootStateIODescriptor(description="Root body position in the world frame.", units="m", axes=["X", "Y", "Z"], shape=(3,), dtype=torch.float32)) +@generic_io_descriptor(units="m", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) 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) @@ -98,7 +238,7 @@ def root_pos_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg( return asset.data.root_pos_w - env.scene.env_origins -@root_state_io_descriptor(RootStateIODescriptor(description="Root body orientation in the world frame.", units="unit", axes=["W", "X", "Y", "Z"], shape=(4,), dtype=torch.float32)) +@generic_io_descriptor(units="unit", axes=["W", "X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) def root_quat_w( env: ManagerBasedEnv, make_quat_unique: bool = False, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> torch.Tensor: @@ -116,7 +256,7 @@ def root_quat_w( return math_utils.quat_unique(quat) if make_quat_unique else quat -@root_state_io_descriptor(RootStateIODescriptor(description="Root body linear velocity in the world frame.", units="m/s", axes=["X", "Y", "Z"], shape=(3,), dtype=torch.float32)) +@generic_io_descriptor(units="m/s", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) def root_lin_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Asset root linear velocity in the environment frame.""" # extract the used quantities (to enable type-hinting) @@ -124,7 +264,7 @@ def root_lin_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntity return asset.data.root_lin_vel_w -@root_state_io_descriptor(RootStateIODescriptor(description="Root body angular velocity in the world frame.", units="rad/s", axes=["X", "Y", "Z"], shape=(3,), dtype=torch.float32)) +@generic_io_descriptor(units="rad/s", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) def root_ang_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Asset root angular velocity in the environment frame.""" # extract the used quantities (to enable type-hinting) @@ -136,30 +276,7 @@ def root_ang_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntity Body state """ -@configclass -class BodyStateIODescriptor(IODescriptor): - observation_type: str = "BodyState" - body_ids: list[int] | int = [] - body_names: list[str] | str = [] - -def body_state_io_descriptor(descriptor: BodyStateIODescriptor, inspect: bool = False): - def decorator(func): - descriptor.name = func.__name__ - def wrapper(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): - if inspect: - out = func(env, asset_cfg) - descriptor.shape = (out.shape[-1],) - descriptor.body_ids = asset_cfg.body_ids - descriptor.body_names = asset_cfg.body_names - return out - else: - return func(env, asset_cfg) - wrapper._has_descriptor = True - wrapper._descriptor = descriptor - return wrapper - return decorator - -@body_state_io_descriptor(BodyStateIODescriptor(description="The flattened body poses of the robot in the world frame. The output shape is 7 * num_bodies", dtype=torch.float32)) +@generic_io_descriptor(observation_type="BodyState", body_names=None, on_inspect=[record_shape, record_dtype, record_body_names]) def body_pose_w( env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), @@ -183,7 +300,7 @@ def body_pose_w( return pose.reshape(env.num_envs, -1) -@body_state_io_descriptor(BodyStateIODescriptor(description="The direction of gravity projected on to bodies own frames. The output shape is 3 * num_bodies", dtype=torch.float32)) +@generic_io_descriptor(observation_type="BodyState", body_names=None, on_inspect=[record_shape, record_dtype, record_body_names]) def body_projected_gravity_b( env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), @@ -211,35 +328,7 @@ def body_projected_gravity_b( """ Joint state. """ -@configclass -class JointStateIODescriptor(IODescriptor): - observation_type: str = "JointState" - joint_ids: list[int] | int = [] - joint_names: list[str] | str = [] - -def joint_state_io_descriptor(descriptor: JointStateIODescriptor): - def decorator(func): - descriptor.name = func.__name__ - def wrapper(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), inspect: bool = False): - if inspect: - out = func(env, asset_cfg) - descriptor.shape = (out.shape[-1],) - descriptor.joint_ids = asset_cfg.joint_ids - descriptor.joint_names = asset_cfg.joint_names - - if descriptor.joint_names is None: - asset: Articulation = env.scene[asset_cfg.name] - descriptor.joint_names = asset.joint_names - descriptor.joint_ids = list(range(len(asset.joint_names))) - return out - else: - return func(env, asset_cfg) - wrapper._has_descriptor = True - wrapper._descriptor = descriptor - return wrapper - return decorator - -@joint_state_io_descriptor(JointStateIODescriptor(description="The joint positions of the asset. The output shape is num_joints.", dtype=torch.float32)) +@generic_io_descriptor(observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape]) def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset. @@ -250,7 +339,7 @@ def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" return asset.data.joint_pos[:, asset_cfg.joint_ids] -@joint_state_io_descriptor(JointStateIODescriptor(description="The joint positions of the asset w.r.t. the default joint positions. The output shape is num_joints", dtype=torch.float32)) +@generic_io_descriptor(observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape]) def joint_pos_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset w.r.t. the default joint positions. @@ -261,7 +350,7 @@ def joint_pos_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC return asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids] -@joint_state_io_descriptor(JointStateIODescriptor(description="The joint positions of the asset normalized with the asset's joint limits. The output shape is num_joints", dtype=torch.float32)) +@generic_io_descriptor(observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape]) def joint_pos_limit_normalized( env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> torch.Tensor: @@ -277,7 +366,7 @@ def joint_pos_limit_normalized( asset.data.soft_joint_pos_limits[:, asset_cfg.joint_ids, 1], ) -@joint_state_io_descriptor(JointStateIODescriptor(description="The joint velocities of the asset. The output shape is num_joints", dtype=torch.float32)) +@generic_io_descriptor(observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape]) def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset. @@ -287,7 +376,7 @@ def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" asset: Articulation = env.scene[asset_cfg.name] return asset.data.joint_vel[:, asset_cfg.joint_ids] -@joint_state_io_descriptor(JointStateIODescriptor(description="The joint velocities of the asset w.r.t. the default joint velocities. The output shape is num_joints", dtype=torch.float32)) +@generic_io_descriptor(observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape]) def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset w.r.t. the default joint velocities. @@ -297,7 +386,7 @@ def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC asset: Articulation = env.scene[asset_cfg.name] return asset.data.joint_vel[:, asset_cfg.joint_ids] - asset.data.default_joint_vel[:, asset_cfg.joint_ids] -@joint_state_io_descriptor(JointStateIODescriptor(description="The joint applied effort of the robot. The output shape is num_joints", dtype=torch.float32)) +@generic_io_descriptor(observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape]) def joint_effort(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint applied effort of the robot. @@ -681,29 +770,7 @@ def _inference(model, images: torch.Tensor) -> torch.Tensor: Actions. """ - -@configclass -class GenericIODescriptor(IODescriptor): - observation_type: str = None - - -def generic_io_descriptor(descriptor: GenericIODescriptor): - def decorator(func): - def wrapper(env: ManagerBasedEnv, *args, inspect: bool = False, **kwargs): - if inspect: - out = func(env, *args, **kwargs) - descriptor.shape = (out.shape[-1],) - return out - else: - return func(env, *args, **kwargs) - descriptor.name = func.__name__ - wrapper._has_descriptor = True - wrapper._descriptor = descriptor - return wrapper - return decorator - - -@generic_io_descriptor(GenericIODescriptor(description="The last input action to the environment.", dtype=torch.float32, observation_type="Action")) +@generic_io_descriptor(dtype=torch.float32, observation_type="Action", on_inspect=[record_shape]) def last_action(env: ManagerBasedEnv, action_name: str | None = None) -> torch.Tensor: """The last input action to the environment. @@ -720,8 +787,7 @@ def last_action(env: ManagerBasedEnv, action_name: str | None = None) -> torch.T Commands. """ - -@generic_io_descriptor(GenericIODescriptor(description="The generated command term in the command manager.", dtype=torch.float32, observation_type="Command")) +@generic_io_descriptor(dtype=torch.float32, observation_type="Command", on_inspect=[record_shape]) def generated_commands(env: ManagerBasedRLEnv, command_name: str | None = None) -> torch.Tensor: """The generated command from command term in the command manager with the given name.""" return env.command_manager.get_command(command_name) From 381ff45fb04d76bd44dd4329c5fe8cdd5834128b Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Tue, 10 Jun 2025 14:56:51 +0200 Subject: [PATCH 04/23] removed the custom logic used to bypass the checks inside the managers. The new decorator can mock its signature, so we don't need that anymore --- .../isaaclab/managers/manager_base.py | 21 ++++--------------- 1 file changed, 4 insertions(+), 17 deletions(-) diff --git a/source/isaaclab/isaaclab/managers/manager_base.py b/source/isaaclab/isaaclab/managers/manager_base.py index 14d198eb659..37cb0454f1b 100644 --- a/source/isaaclab/isaaclab/managers/manager_base.py +++ b/source/isaaclab/isaaclab/managers/manager_base.py @@ -353,30 +353,17 @@ def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg, # check if function is callable if not callable(func_static): raise AttributeError(f"The term '{term_name}' is not callable. Received: {term_cfg.func}") - + # check statically if the term's arguments are matched by params term_params = list(term_cfg.params.keys()) args = inspect.signature(func_static).parameters args_with_defaults = [arg for arg in args if args[arg].default is not inspect.Parameter.empty] args_without_defaults = [arg for arg in args if args[arg].default is inspect.Parameter.empty] - # ignore first two arguments for env and env_ids - # Think: Check for cases when kwargs are set inside the function? - if "kwargs" in args_without_defaults: - args_without_defaults.remove("kwargs") - args_with_defaults.append("kwargs") args = args_without_defaults + args_with_defaults - - print("args", args) - print("args_without_defaults", args_without_defaults) - print("args_with_defaults", args_with_defaults) - print("min_argc", min_argc) - print("term_params", term_params) - print("term_cfg", term_cfg) - + # ignore first two arguments for env and env_ids + # Think: Check for cases when kwargs are set inside the function? if len(args) > min_argc: - if "kwargs" in args: - pass - elif set(args[min_argc:]) != set(term_params + args_with_defaults): + if set(args[min_argc:]) != set(term_params + args_with_defaults): raise ValueError( f"The term '{term_name}' expects mandatory parameters: {args_without_defaults[min_argc:]}" f" and optional parameters: {args_with_defaults}, but received: {term_params}." From d6fcf720c051630b539ee2a121146c6ad61b031b Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Tue, 10 Jun 2025 14:57:27 +0200 Subject: [PATCH 05/23] added ugly temporary logic to handle the YAML convertion. Need to clean that up a bit --- .../isaaclab/managers/observation_manager.py | 32 +++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/source/isaaclab/isaaclab/managers/observation_manager.py b/source/isaaclab/isaaclab/managers/observation_manager.py index b763ab1e471..c84bec2894a 100644 --- a/source/isaaclab/isaaclab/managers/observation_manager.py +++ b/source/isaaclab/isaaclab/managers/observation_manager.py @@ -253,12 +253,40 @@ def get_IO_descriptors(self): # dummy call to cache some values try: term_cfg.func(self._env, **term_cfg.params, inspect=True) - data.append(term_cfg.func._descriptor) + desc = term_cfg.func._descriptor.__dict__.copy() + print(f"desc: {desc}") + overloads = {} + for k,v in term_cfg.__dict__.items(): + if k in ["modifiers", "clip", "scale", "history_length", "flatten_history_dim"]: + overloads[k] = v + desc.update(overloads) + data.append(desc) except Exception as e: print(f"Error getting IO descriptor for term '{term_name}' in group '{group_name}': {e}") - return data + # Format the data for YAML export + formatted_data = {} + for item in data: + name = item.pop("name") + formatted_item = {} + formatted_item["overloads"] = {} + formatted_item["extras"] = {} + for k,v in item.items(): + # Check if v is a tuple and convert to list + if isinstance(v, tuple): + v = list(v) + if k in ["scale", "clip", "history_length", "flatten_history_dim"]: + formatted_item["overloads"][k] = v + elif k in ["modifiers", "noise", "description", "units"]: + formatted_item["extras"][k] = v + else: + formatted_item[k] = v + + + formatted_data[name] = formatted_item + return formatted_data + """ Operations. """ From 634bc26db2c9e0454174d667f25f680fc0bb5c38 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Tue, 10 Jun 2025 14:57:55 +0200 Subject: [PATCH 06/23] added ugly temporary logic to handle the YAML convertion. Need to clean that up a bit --- scripts/environments/random_agent.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/scripts/environments/random_agent.py b/scripts/environments/random_agent.py index 9422680e992..b8e7ffee7d4 100644 --- a/scripts/environments/random_agent.py +++ b/scripts/environments/random_agent.py @@ -56,10 +56,16 @@ def main(): env.reset() out = env.unwrapped.get_IO_descriptors - for o in out: - print(f"--- Obs term: {o.name} ---") - for k, v in o.__dict__.items(): - print(f"{k}: {v}") + # Make a yaml file with the output + import yaml + with open("obs_descriptors.yaml", "w") as f: + yaml.safe_dump(out, f) + + for k, v in out.items(): + print(f"--- Obs term: {k} ---") + for k1, v1 in v.items(): + print(f"{k1}: {v1}") + exit(0) # simulate environment while simulation_app.is_running(): # run everything in inference mode From e91bda3bb58a8e71c610f63b3a07eed4a8f6d4a8 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Tue, 10 Jun 2025 15:32:00 +0200 Subject: [PATCH 07/23] passing pre-commits --- scripts/environments/random_agent.py | 3 +- .../isaaclab/envs/manager_based_env.py | 2 +- .../isaaclab/envs/mdp/observations.py | 138 ++++++++++++------ .../isaaclab/managers/manager_base.py | 2 +- .../isaaclab/managers/observation_manager.py | 14 +- 5 files changed, 99 insertions(+), 60 deletions(-) diff --git a/scripts/environments/random_agent.py b/scripts/environments/random_agent.py index b8e7ffee7d4..ab877628620 100644 --- a/scripts/environments/random_agent.py +++ b/scripts/environments/random_agent.py @@ -47,8 +47,6 @@ def main(): # create environment env = gym.make(args_cli.task, cfg=env_cfg) - - # print info (this is vectorized environment) print(f"[INFO]: Gym observation space: {env.observation_space}") print(f"[INFO]: Gym action space: {env.action_space}") @@ -58,6 +56,7 @@ def main(): out = env.unwrapped.get_IO_descriptors # Make a yaml file with the output import yaml + with open("obs_descriptors.yaml", "w") as f: yaml.safe_dump(out, f) diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 20245446a33..5ebaf14403f 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -209,7 +209,7 @@ def step_dt(self) -> float: def device(self): """The device on which the environment is running.""" return self.sim.device - + @property def get_IO_descriptors(self): """Get the IO descriptors for the environment. diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index 1133d02a84a..4ef12726d84 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -14,7 +14,7 @@ import functools import inspect import torch -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Concatenate, ParamSpec, TypeVar import isaaclab.utils.math as math_utils from isaaclab.assets import Articulation, RigidObject @@ -26,22 +26,11 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv -from isaaclab.utils import configclass -from dataclasses import MISSING - +import dataclasses +from collections.abc import Callable -import dataclasses, functools, inspect -from typing import Any, Callable, ParamSpec, TypeVar, Concatenate +from isaaclab.utils import configclass -@configclass -class IODescriptor: - mdp_type: str = "Observation" - name: str = None - full_path: str = None - description: str = None - shape: tuple[int, ...] = None - dtype: torch.dtype = None - observation_type: str = None @configclass class GenericIODescriptor: @@ -54,30 +43,32 @@ class GenericIODescriptor: observation_type: str = None extras: dict[str, Any] = None + # These are defined to help with type hinting P = ParamSpec("P") R = TypeVar("R") + # Automatically builds a descriptor from the kwargs -def _make_descriptor(**kwargs: Any) -> "GenericIODescriptor": +def _make_descriptor(**kwargs: Any) -> GenericIODescriptor: """Split *kwargs* into (known dataclass fields) and (extras).""" field_names = {f.name for f in dataclasses.fields(GenericIODescriptor)} - known = {k: v for k, v in kwargs.items() if k in field_names} - extras = {k: v for k, v in kwargs.items() if k not in field_names} + known = {k: v for k, v in kwargs.items() if k in field_names} + extras = {k: v for k, v in kwargs.items() if k not in field_names} desc = GenericIODescriptor(**known) # User defined extras are stored in the descriptor under the `extras` field desc.extras = extras return desc + # Decorator factory for generic IO descriptors. def generic_io_descriptor( _func: Callable[Concatenate[ManagerBasedEnv, P], R] | None = None, *, on_inspect: Callable[..., Any] | list[Callable[..., Any]] | None = None, **descriptor_kwargs: Any, -) -> Callable[[Callable[Concatenate[ManagerBasedEnv, P], R]], - Callable[Concatenate[ManagerBasedEnv, P], R]]: +) -> Callable[[Callable[Concatenate[ManagerBasedEnv, P], R]], Callable[Concatenate[ManagerBasedEnv, P], R]]: """ Decorator factory for generic IO descriptors. @@ -96,16 +87,32 @@ def my_func(env: ManagerBasedEnv, *args, **kwargs): ... 3. I need to add a hook to the descriptor: ..code-block:: python - def record_shape(tensor: torch.Tensor, desc: GenericIODescriptor): + def record_shape(tensor: torch.Tensor, desc: GenericIODescriptor, **kwargs): desc.shape = (tensor.shape[-1],) - - @generic_io_descriptor(description="..", new_var_1="a", new_var_2="b", on_inspect=[record_shape]) + + @generic_io_descriptor(description="..", new_var_1="a", new_var_2="b", on_inspect=[record_shape, record_dtype]) def my_func(env: ManagerBasedEnv, *args, **kwargs): ..note:: The hook is called after the function is called, if and only if the `inspect` flag is set when calling the function. + For example: ..code-block:: python my_func(env, inspect=True) - + + 4. I need to add a hook to the descriptor and this hook will write to a variable that is not part of the base descriptor. + ..code-block:: python + def record_joint_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): + asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] + joint_ids = kwargs["asset_cfg"].joint_ids + if joint_ids == slice(None, None, None): + joint_ids = list(range(len(asset.joint_names))) + descriptor.joint_names = [asset.joint_names[i] for i in joint_ids] + + @generic_io_descriptor(joint_names=None, new_var_1="a", new_var_2="b", on_inspect=[record_shape, record_dtype, record_joint_names]) + def my_func(env: ManagerBasedEnv, *args, **kwargs): + + ..note:: The hook can access all the variables in the wrapped function's signature. While it is useful, the user should be careful to + access only existing variables. + Args: _func: The function to decorate. **descriptor_kwargs: Keyword arguments to pass to the descriptor. @@ -126,17 +133,13 @@ def my_func(env: ManagerBasedEnv, *args, **kwargs): else: inspect_hooks: list[Callable[..., Any]] = list(on_inspect or []) # handles None - def _apply( - func: Callable[Concatenate[ManagerBasedEnv, P], R] - ) -> Callable[Concatenate[ManagerBasedEnv, P], R]: - + def _apply(func: Callable[Concatenate[ManagerBasedEnv, P], R]) -> Callable[Concatenate[ManagerBasedEnv, P], R]: + # Capture the signature of the function sig = inspect.signature(func) @functools.wraps(func) - def wrapper( - env: ManagerBasedEnv, *args: P.args, **kwargs: P.kwargs - ) -> R: + def wrapper(env: ManagerBasedEnv, *args: P.args, **kwargs: P.kwargs) -> R: inspect_flag: bool = kwargs.pop("inspect", False) out = func(env, *args, **kwargs) if inspect_flag: @@ -177,9 +180,11 @@ def wrapper( def record_shape(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): descriptor.shape = (output.shape[-1],) + def record_dtype(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): descriptor.dtype = str(output.dtype) + def record_joint_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] joint_ids = kwargs["asset_cfg"].joint_ids @@ -187,6 +192,7 @@ def record_joint_names(output: torch.Tensor, descriptor: GenericIODescriptor, ** joint_ids = list(range(len(asset.joint_names))) descriptor.joint_names = [asset.joint_names[i] for i in joint_ids] + def record_body_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] body_ids = kwargs["asset_cfg"].body_ids @@ -194,10 +200,12 @@ def record_body_names(output: torch.Tensor, descriptor: GenericIODescriptor, **k body_ids = list(range(len(asset.body_names))) descriptor.body_names = [asset.body_names[i] for i in body_ids] + """ Root state. """ + @generic_io_descriptor(units="m", axes=["Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) def base_pos_z(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Root height in the simulation world frame.""" @@ -206,7 +214,9 @@ def base_pos_z(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg( return asset.data.root_pos_w[:, 2].unsqueeze(-1) -@generic_io_descriptor(units="m/s", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) +@generic_io_descriptor( + units="m/s", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype] +) def base_lin_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Root linear velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) @@ -214,7 +224,9 @@ def base_lin_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCf return asset.data.root_lin_vel_b -@generic_io_descriptor(units="rad/s", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) +@generic_io_descriptor( + units="rad/s", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype] +) def base_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Root angular velocity in the asset's root frame.""" # extract the used quantities (to enable type-hinting) @@ -222,7 +234,9 @@ def base_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCf return asset.data.root_ang_vel_b -@generic_io_descriptor(units="m/s^2", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) +@generic_io_descriptor( + units="m/s^2", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype] +) def projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Gravity projection on the asset's root frame.""" # extract the used quantities (to enable type-hinting) @@ -230,7 +244,9 @@ def projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEnt return asset.data.projected_gravity_b -@generic_io_descriptor(units="m", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) +@generic_io_descriptor( + units="m", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype] +) 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) @@ -238,7 +254,9 @@ def root_pos_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg( return asset.data.root_pos_w - env.scene.env_origins -@generic_io_descriptor(units="unit", axes=["W", "X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) +@generic_io_descriptor( + units="unit", axes=["W", "X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype] +) def root_quat_w( env: ManagerBasedEnv, make_quat_unique: bool = False, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> torch.Tensor: @@ -256,7 +274,9 @@ def root_quat_w( return math_utils.quat_unique(quat) if make_quat_unique else quat -@generic_io_descriptor(units="m/s", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) +@generic_io_descriptor( + units="m/s", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype] +) def root_lin_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Asset root linear velocity in the environment frame.""" # extract the used quantities (to enable type-hinting) @@ -264,7 +284,9 @@ def root_lin_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntity return asset.data.root_lin_vel_w -@generic_io_descriptor(units="rad/s", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype]) +@generic_io_descriptor( + units="rad/s", axes=["X", "Y", "Z"], observation_type="RootState", on_inspect=[record_shape, record_dtype] +) def root_ang_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """Asset root angular velocity in the environment frame.""" # extract the used quantities (to enable type-hinting) @@ -276,7 +298,10 @@ def root_ang_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntity Body state """ -@generic_io_descriptor(observation_type="BodyState", body_names=None, on_inspect=[record_shape, record_dtype, record_body_names]) + +@generic_io_descriptor( + observation_type="BodyState", body_names=None, on_inspect=[record_shape, record_dtype, record_body_names] +) def body_pose_w( env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), @@ -300,7 +325,9 @@ def body_pose_w( return pose.reshape(env.num_envs, -1) -@generic_io_descriptor(observation_type="BodyState", body_names=None, on_inspect=[record_shape, record_dtype, record_body_names]) +@generic_io_descriptor( + observation_type="BodyState", body_names=None, on_inspect=[record_shape, record_dtype, record_body_names] +) def body_projected_gravity_b( env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), @@ -328,7 +355,11 @@ def body_projected_gravity_b( """ Joint state. """ -@generic_io_descriptor(observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape]) + + +@generic_io_descriptor( + observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape] +) def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset. @@ -339,7 +370,9 @@ def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" return asset.data.joint_pos[:, asset_cfg.joint_ids] -@generic_io_descriptor(observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape]) +@generic_io_descriptor( + observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape] +) def joint_pos_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset w.r.t. the default joint positions. @@ -350,7 +383,9 @@ def joint_pos_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC return asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids] -@generic_io_descriptor(observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape]) +@generic_io_descriptor( + observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape] +) def joint_pos_limit_normalized( env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> torch.Tensor: @@ -366,7 +401,10 @@ def joint_pos_limit_normalized( asset.data.soft_joint_pos_limits[:, asset_cfg.joint_ids, 1], ) -@generic_io_descriptor(observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape]) + +@generic_io_descriptor( + observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape] +) def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset. @@ -376,7 +414,10 @@ def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" asset: Articulation = env.scene[asset_cfg.name] return asset.data.joint_vel[:, asset_cfg.joint_ids] -@generic_io_descriptor(observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape]) + +@generic_io_descriptor( + observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape] +) def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset w.r.t. the default joint velocities. @@ -386,7 +427,10 @@ def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC asset: Articulation = env.scene[asset_cfg.name] return asset.data.joint_vel[:, asset_cfg.joint_ids] - asset.data.default_joint_vel[:, asset_cfg.joint_ids] -@generic_io_descriptor(observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape]) + +@generic_io_descriptor( + observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape] +) def joint_effort(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint applied effort of the robot. @@ -770,6 +814,7 @@ def _inference(model, images: torch.Tensor) -> torch.Tensor: Actions. """ + @generic_io_descriptor(dtype=torch.float32, observation_type="Action", on_inspect=[record_shape]) def last_action(env: ManagerBasedEnv, action_name: str | None = None) -> torch.Tensor: """The last input action to the environment. @@ -787,6 +832,7 @@ def last_action(env: ManagerBasedEnv, action_name: str | None = None) -> torch.T Commands. """ + @generic_io_descriptor(dtype=torch.float32, observation_type="Command", on_inspect=[record_shape]) def generated_commands(env: ManagerBasedRLEnv, command_name: str | None = None) -> torch.Tensor: """The generated command from command term in the command manager with the given name.""" diff --git a/source/isaaclab/isaaclab/managers/manager_base.py b/source/isaaclab/isaaclab/managers/manager_base.py index 37cb0454f1b..081c3e271ec 100644 --- a/source/isaaclab/isaaclab/managers/manager_base.py +++ b/source/isaaclab/isaaclab/managers/manager_base.py @@ -353,7 +353,7 @@ def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg, # check if function is callable if not callable(func_static): raise AttributeError(f"The term '{term_name}' is not callable. Received: {term_cfg.func}") - + # check statically if the term's arguments are matched by params term_params = list(term_cfg.params.keys()) args = inspect.signature(func_static).parameters diff --git a/source/isaaclab/isaaclab/managers/observation_manager.py b/source/isaaclab/isaaclab/managers/observation_manager.py index c84bec2894a..b9e806b367b 100644 --- a/source/isaaclab/isaaclab/managers/observation_manager.py +++ b/source/isaaclab/isaaclab/managers/observation_manager.py @@ -244,8 +244,6 @@ def get_IO_descriptors(self): ) # iterate over all the terms in each group group_term_names = self._group_obs_term_names[group_name] - # buffer to store obs per group - group_obs = dict.fromkeys(group_term_names, None) # read attributes for each term obs_terms = zip(group_term_names, self._group_obs_term_cfgs[group_name]) @@ -254,9 +252,8 @@ def get_IO_descriptors(self): try: term_cfg.func(self._env, **term_cfg.params, inspect=True) desc = term_cfg.func._descriptor.__dict__.copy() - print(f"desc: {desc}") overloads = {} - for k,v in term_cfg.__dict__.items(): + for k, v in term_cfg.__dict__.items(): if k in ["modifiers", "clip", "scale", "history_length", "flatten_history_dim"]: overloads[k] = v desc.update(overloads) @@ -268,10 +265,8 @@ def get_IO_descriptors(self): formatted_data = {} for item in data: name = item.pop("name") - formatted_item = {} - formatted_item["overloads"] = {} - formatted_item["extras"] = {} - for k,v in item.items(): + formatted_item = {"overloads": {}, "extras": {}} + for k, v in item.items(): # Check if v is a tuple and convert to list if isinstance(v, tuple): v = list(v) @@ -281,12 +276,11 @@ def get_IO_descriptors(self): formatted_item["extras"][k] = v else: formatted_item[k] = v - formatted_data[name] = formatted_item return formatted_data - + """ Operations. """ From 5d42ed27231837c893315c9d19e91b94410eae64 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Tue, 10 Jun 2025 16:43:18 +0200 Subject: [PATCH 08/23] WIP added action thingys --- scripts/environments/random_agent.py | 11 +++++-- .../isaaclab/envs/manager_based_env.py | 2 +- .../isaaclab/envs/mdp/actions/actions_cfg.py | 1 - .../envs/mdp/actions/joint_actions.py | 19 +++++++++++ .../isaaclab/managers/action_manager.py | 32 +++++++++++++++++++ .../isaaclab/managers/observation_manager.py | 6 +++- 6 files changed, 66 insertions(+), 5 deletions(-) diff --git a/scripts/environments/random_agent.py b/scripts/environments/random_agent.py index ab877628620..55f5fe92d0a 100644 --- a/scripts/environments/random_agent.py +++ b/scripts/environments/random_agent.py @@ -53,12 +53,19 @@ def main(): # reset environment env.reset() - out = env.unwrapped.get_IO_descriptors + outs = env.unwrapped.get_IO_descriptors + out = outs["observations"] + out_actions = outs["actions"] # Make a yaml file with the output import yaml with open("obs_descriptors.yaml", "w") as f: - yaml.safe_dump(out, f) + yaml.safe_dump(outs, f) + + for k, v in out_actions.items(): + print(f"--- Action term: {k} ---") + for k1, v1 in v.items(): + print(f"{k1}: {v1}") for k, v in out.items(): print(f"--- Obs term: {k} ---") diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 5ebaf14403f..8aef54f14a0 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -217,7 +217,7 @@ def get_IO_descriptors(self): Returns: A dictionary with keys as the group names and values as the IO descriptors. """ - return self.observation_manager.get_IO_descriptors + return {"observations": self.observation_manager.get_IO_descriptors, "actions": self.action_manager.get_IO_descriptors} """ Operations - Setup. diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py b/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py index 9c932d227b0..8a769e4fb82 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py @@ -32,7 +32,6 @@ class JointActionCfg(ActionTermCfg): preserve_order: bool = False """Whether to preserve the order of the joint names in the action output. Defaults to False.""" - @configclass class JointPositionActionCfg(JointActionCfg): """Configuration for the joint position action term. diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py index 0e2689fd254..9e0ffcc513f 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py @@ -122,6 +122,25 @@ def raw_actions(self) -> torch.Tensor: @property def processed_actions(self) -> torch.Tensor: return self._processed_actions + + @property + def IO_descriptor(self): + data = { + "name": self.__class__.__name__, + "mdp_type": "Action", + "full_path": self.__class__.__module__ + "." + self.__class__.__name__, + "description": self.__doc__, + "shape": (self.action_dim,), + "dtype": str(self.raw_actions.dtype), + "action_type": "JointState", + "joint_names": self._joint_names, + "scale": self._scale, + "offset": self._offset[0].detach().cpu().numpy().tolist(), # This seems to be always [4xNum_joints] IDK why. Need to check. + } + + if self.cfg.clip is not None: + data["clip"] = self._clip + return data """ Operations. diff --git a/source/isaaclab/isaaclab/managers/action_manager.py b/source/isaaclab/isaaclab/managers/action_manager.py index 72e78e68cb9..9dce7acd733 100644 --- a/source/isaaclab/isaaclab/managers/action_manager.py +++ b/source/isaaclab/isaaclab/managers/action_manager.py @@ -259,6 +259,38 @@ def has_debug_vis_implementation(self) -> bool: has_debug_vis |= term.has_debug_vis_implementation return has_debug_vis + @property + def get_IO_descriptors(self): + """Get the IO descriptors for the action manager. + + Returns: + A dictionary with keys as the term names and values as the IO descriptors. + """ + + data = [] + + for term in self._terms.values(): + try: + data.append(term.IO_descriptor) + except Exception as e: + print(f"Error getting IO descriptor for term: {e}") + + formatted_data = {} + for item in data: + name = item.pop("name") + formatted_item = {"extras": {}} + for k, v in item.items(): + # Check if v is a tuple and convert to list + if isinstance(v, tuple): + v = list(v) + if k in ["description", "units"]: + formatted_item["extras"][k] = v + else: + formatted_item[k] = v + formatted_data[name] = formatted_item + + return formatted_data + """ Operations. """ diff --git a/source/isaaclab/isaaclab/managers/observation_manager.py b/source/isaaclab/isaaclab/managers/observation_manager.py index b9e806b367b..bc87798e203 100644 --- a/source/isaaclab/isaaclab/managers/observation_manager.py +++ b/source/isaaclab/isaaclab/managers/observation_manager.py @@ -248,12 +248,16 @@ def get_IO_descriptors(self): obs_terms = zip(group_term_names, self._group_obs_term_cfgs[group_name]) for term_name, term_cfg in obs_terms: - # dummy call to cache some values + # Call to the observation function to get the IO descriptor with the inspect flag set to True try: term_cfg.func(self._env, **term_cfg.params, inspect=True) + # Copy the descriptor and update with the term's own extra parameters desc = term_cfg.func._descriptor.__dict__.copy() + # Create a dictionary to store the overloads overloads = {} + # Iterate over the term's own parameters and add them to the overloads dictionary for k, v in term_cfg.__dict__.items(): + # For now we do not add the noise modifier if k in ["modifiers", "clip", "scale", "history_length", "flatten_history_dim"]: overloads[k] = v desc.update(overloads) From 8d1a6f8551fc4678acf854df9ae567dad37f5f3f Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Wed, 11 Jun 2025 15:27:10 +0200 Subject: [PATCH 09/23] Readying code for clean-up --- .../isaaclab/envs/manager_based_env.py | 5 +- .../isaaclab/envs/mdp/actions/actions_cfg.py | 1 + .../envs/mdp/actions/binary_joint_actions.py | 10 + .../envs/mdp/actions/joint_actions.py | 35 +-- .../mdp/actions/joint_actions_to_limits.py | 34 +++ .../envs/mdp/actions/non_holonomic_actions.py | 16 ++ .../mdp/actions/pink_task_space_actions.py | 12 + .../envs/mdp/actions/task_space_actions.py | 35 +++ .../isaaclab/envs/mdp/observations.py | 217 ++---------------- .../isaaclab/managers/action_manager.py | 23 +- .../isaaclab/managers/observation_manager.py | 6 +- 11 files changed, 165 insertions(+), 229 deletions(-) diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 8aef54f14a0..ee642cbf20f 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -217,7 +217,10 @@ def get_IO_descriptors(self): Returns: A dictionary with keys as the group names and values as the IO descriptors. """ - return {"observations": self.observation_manager.get_IO_descriptors, "actions": self.action_manager.get_IO_descriptors} + return { + "observations": self.observation_manager.get_IO_descriptors, + "actions": self.action_manager.get_IO_descriptors, + } """ Operations - Setup. diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py b/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py index 8a769e4fb82..9c932d227b0 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py @@ -32,6 +32,7 @@ class JointActionCfg(ActionTermCfg): preserve_order: bool = False """Whether to preserve the order of the joint names in the action output. Defaults to False.""" + @configclass class JointPositionActionCfg(JointActionCfg): """Configuration for the joint position action term. diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/binary_joint_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/binary_joint_actions.py index f3a4fa37af1..9b8666e4464 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/binary_joint_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/binary_joint_actions.py @@ -17,6 +17,7 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv + from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor from . import actions_cfg @@ -111,6 +112,15 @@ def raw_actions(self) -> torch.Tensor: def processed_actions(self) -> torch.Tensor: return self._processed_actions + @property + def IO_descriptor(self) -> GenericActionIODescriptor: + super().IO_descriptor + self._IO_descriptor.shape = (self.action_dim,) + self._IO_descriptor.dtype = str(self.raw_actions.dtype) + self._IO_descriptor.action_type = "JointAction" + self._IO_descriptor.joint_names = self._joint_names + return self._IO_descriptor + """ Operations. """ diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py index 9e0ffcc513f..8341c59184f 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py @@ -17,6 +17,7 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv + from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor from . import actions_cfg @@ -122,25 +123,25 @@ def raw_actions(self) -> torch.Tensor: @property def processed_actions(self) -> torch.Tensor: return self._processed_actions - + @property - def IO_descriptor(self): - data = { - "name": self.__class__.__name__, - "mdp_type": "Action", - "full_path": self.__class__.__module__ + "." + self.__class__.__name__, - "description": self.__doc__, - "shape": (self.action_dim,), - "dtype": str(self.raw_actions.dtype), - "action_type": "JointState", - "joint_names": self._joint_names, - "scale": self._scale, - "offset": self._offset[0].detach().cpu().numpy().tolist(), # This seems to be always [4xNum_joints] IDK why. Need to check. - } - + def IO_descriptor(self) -> GenericActionIODescriptor: + super().IO_descriptor + self._IO_descriptor.shape = (self.action_dim,) + self._IO_descriptor.dtype = str(self.raw_actions.dtype) + self._IO_descriptor.action_type = "JointAction" + self._IO_descriptor.joint_names = self._joint_names + self._IO_descriptor.scale = self._scale + # This seems to be always [4xNum_joints] IDK why. Need to check. + if isinstance(self._offset, torch.Tensor): + self._IO_descriptor.offset = self._offset[0].detach().cpu().numpy().tolist() + else: + self._IO_descriptor.offset = self._offset if self.cfg.clip is not None: - data["clip"] = self._clip - return data + self._IO_descriptor.clip = self._clip + else: + self._IO_descriptor.clip = None + return self._IO_descriptor """ Operations. diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py index 5398241e15e..d056d2e6ac6 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py @@ -18,6 +18,7 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv + from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor from . import actions_cfg @@ -105,6 +106,25 @@ def raw_actions(self) -> torch.Tensor: def processed_actions(self) -> torch.Tensor: return self._processed_actions + @property + def IO_descriptor(self) -> GenericActionIODescriptor: + super().IO_descriptor + self._IO_descriptor.shape = (self.action_dim,) + self._IO_descriptor.dtype = str(self.raw_actions.dtype) + self._IO_descriptor.action_type = "JointAction" + self._IO_descriptor.joint_names = self._joint_names + self._IO_descriptor.scale = self._scale + # This seems to be always [4xNum_joints] IDK why. Need to check. + if isinstance(self._offset, torch.Tensor): + self._IO_descriptor.offset = self._offset[0].detach().cpu().numpy().tolist() + else: + self._IO_descriptor.offset = self._offset + if self.cfg.clip is not None: + self._IO_descriptor.clip = self._clip + else: + self._IO_descriptor.clip = None + return self._IO_descriptor + """ Operations. """ @@ -195,6 +215,20 @@ def __init__(self, cfg: actions_cfg.EMAJointPositionToLimitsActionCfg, env: Mana # initialize the previous targets self._prev_applied_actions = torch.zeros_like(self.processed_actions) + @property + def IO_descriptor(self) -> GenericActionIODescriptor: + super().IO_descriptor + if isinstance(self._alpha, float): + self._IO_descriptor.alpha = self._alpha + elif isinstance(self._alpha, torch.Tensor): + self._IO_descriptor.alpha = self._alpha[0].detach().cpu().numpy().tolist() + else: + raise ValueError( + f"Unsupported moving average weight type: {type(self._alpha)}. Supported types are float and" + " torch.Tensor." + ) + return self._IO_descriptor + def reset(self, env_ids: Sequence[int] | None = None) -> None: # check if specific environment ids are provided if env_ids is None: diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py index a3e5cebdbf5..5d168da4151 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py @@ -18,6 +18,7 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv + from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor from . import actions_cfg @@ -134,6 +135,21 @@ def raw_actions(self) -> torch.Tensor: def processed_actions(self) -> torch.Tensor: return self._processed_actions + @property + def IO_descriptor(self) -> GenericActionIODescriptor: + super().IO_descriptor + self._IO_descriptor.shape = (self.action_dim,) + self._IO_descriptor.dtype = str(self.raw_actions.dtype) + self._IO_descriptor.action_type = "non holonomic actions" + self._IO_descriptor.scale = self._scale + self._IO_descriptor.offset = self._offset + self._IO_descriptor.clip = self._clip + self._IO_descriptor.body_name = self._body_name + self._IO_descriptor.x_joint_name = self._joint_names[0] + self._IO_descriptor.y_joint_name = self._joint_names[1] + self._IO_descriptor.yaw_joint_name = self._joint_names[2] + return self._IO_descriptor + """ Operations. """ diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py index 11c3ff6cedf..30980fa1fb4 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py @@ -22,6 +22,7 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv + from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor from . import pink_actions_cfg @@ -130,6 +131,17 @@ def processed_actions(self) -> torch.Tensor: """Get the processed actions tensor.""" return self._processed_actions + @property + def IO_descriptor(self) -> GenericActionIODescriptor: + super().IO_descriptor + self._IO_descriptor.shape = (self.action_dim,) + self._IO_descriptor.dtype = str(self.raw_actions.dtype) + self._IO_descriptor.action_type = "PinkInverseKinematicsAction" + self._IO_descriptor.pink_controller_joint_names = self._pink_controlled_joint_names + self._IO_descriptor.hand_joint_names = self._hand_joint_names + self._IO_descriptor.extras["controller_cfg"] = self.cfg.controller.__dict__ + return self._IO_descriptor + # """ # Operations. # """ diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py index 89f51817179..6f9bf0553ec 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py @@ -23,6 +23,7 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv + from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor from . import actions_cfg @@ -148,6 +149,23 @@ def jacobian_b(self) -> torch.Tensor: jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :]) return jacobian + @property + def IO_descriptor(self) -> GenericActionIODescriptor: + super().IO_descriptor + self._IO_descriptor.shape = (self.action_dim,) + self._IO_descriptor.dtype = str(self.raw_actions.dtype) + self._IO_descriptor.action_type = "TaskSpaceAction" + self._IO_descriptor.body_name = self._body_name + self._IO_descriptor.joint_names = self._joint_names + self._IO_descriptor.scale = self._scale + if self.cfg.clip is not None: + self._IO_descriptor.clip = self.cfg.clip + else: + self._IO_descriptor.clip = None + self._IO_descriptor.extras["controller_cfg"] = self.cfg.controller.__dict__ + self._IO_descriptor.extras["body_offset"] = self.cfg.body_offset.__dict__ + return self._IO_descriptor + """ Operations. """ @@ -409,6 +427,23 @@ def jacobian_b(self) -> torch.Tensor: jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :]) return jacobian + @property + def IO_descriptor(self) -> GenericActionIODescriptor: + super().IO_descriptor + self._IO_descriptor.shape = (self.action_dim,) + self._IO_descriptor.dtype = str(self.raw_actions.dtype) + self._IO_descriptor.action_type = "TaskSpaceAction" + self._IO_descriptor.body_name = self._ee_body_name + self._IO_descriptor.joint_names = self._joint_names + self._IO_descriptor.scale = self._scale + if self.cfg.clip is not None: + self._IO_descriptor.clip = self.cfg.clip + else: + self._IO_descriptor.clip = None + self._IO_descriptor.extras["controller_cfg"] = self.cfg.controller.__dict__ + self._IO_descriptor.extras["body_offset"] = self.cfg.body_offset.__dict__ + return self._IO_descriptor + """ Operations. """ diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index 4ef12726d84..f0cc1e13cff 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -11,10 +11,8 @@ from __future__ import annotations -import functools -import inspect import torch -from typing import TYPE_CHECKING, Any, Concatenate, ParamSpec, TypeVar +from typing import TYPE_CHECKING import isaaclab.utils.math as math_utils from isaaclab.assets import Articulation, RigidObject @@ -26,180 +24,13 @@ if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv -import dataclasses -from collections.abc import Callable - -from isaaclab.utils import configclass - - -@configclass -class GenericIODescriptor: - mdp_type: str = "Observation" - name: str = None - full_path: str = None - description: str = None - shape: tuple[int, ...] = None - dtype: torch.dtype = None - observation_type: str = None - extras: dict[str, Any] = None - - -# These are defined to help with type hinting -P = ParamSpec("P") -R = TypeVar("R") - - -# Automatically builds a descriptor from the kwargs -def _make_descriptor(**kwargs: Any) -> GenericIODescriptor: - """Split *kwargs* into (known dataclass fields) and (extras).""" - field_names = {f.name for f in dataclasses.fields(GenericIODescriptor)} - known = {k: v for k, v in kwargs.items() if k in field_names} - extras = {k: v for k, v in kwargs.items() if k not in field_names} - - desc = GenericIODescriptor(**known) - # User defined extras are stored in the descriptor under the `extras` field - desc.extras = extras - return desc - - -# Decorator factory for generic IO descriptors. -def generic_io_descriptor( - _func: Callable[Concatenate[ManagerBasedEnv, P], R] | None = None, - *, - on_inspect: Callable[..., Any] | list[Callable[..., Any]] | None = None, - **descriptor_kwargs: Any, -) -> Callable[[Callable[Concatenate[ManagerBasedEnv, P], R]], Callable[Concatenate[ManagerBasedEnv, P], R]]: - """ - Decorator factory for generic IO descriptors. - - This decorator can be used in different ways: - 1. The default decorator has all the information I need for my use case: - ..code-block:: python - @generic_io_descriptor(GenericIODescriptor(description="..", dtype="..")) - def my_func(env: ManagerBasedEnv, *args, **kwargs): - ... - ..note:: If description is not set, the function's docstring is used to populate it. - - 2. I need to add more information to the descriptor: - ..code-block:: python - @generic_io_descriptor(description="..", new_var_1="a", new_var_2="b") - def my_func(env: ManagerBasedEnv, *args, **kwargs): - ... - 3. I need to add a hook to the descriptor: - ..code-block:: python - def record_shape(tensor: torch.Tensor, desc: GenericIODescriptor, **kwargs): - desc.shape = (tensor.shape[-1],) - - @generic_io_descriptor(description="..", new_var_1="a", new_var_2="b", on_inspect=[record_shape, record_dtype]) - def my_func(env: ManagerBasedEnv, *args, **kwargs): - ..note:: The hook is called after the function is called, if and only if the `inspect` flag is set when calling the function. - - For example: - ..code-block:: python - my_func(env, inspect=True) - - 4. I need to add a hook to the descriptor and this hook will write to a variable that is not part of the base descriptor. - ..code-block:: python - def record_joint_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): - asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] - joint_ids = kwargs["asset_cfg"].joint_ids - if joint_ids == slice(None, None, None): - joint_ids = list(range(len(asset.joint_names))) - descriptor.joint_names = [asset.joint_names[i] for i in joint_ids] - - @generic_io_descriptor(joint_names=None, new_var_1="a", new_var_2="b", on_inspect=[record_shape, record_dtype, record_joint_names]) - def my_func(env: ManagerBasedEnv, *args, **kwargs): - - ..note:: The hook can access all the variables in the wrapped function's signature. While it is useful, the user should be careful to - access only existing variables. - - Args: - _func: The function to decorate. - **descriptor_kwargs: Keyword arguments to pass to the descriptor. - - Returns: - A decorator that can be used to decorate a function. - """ - - if _func is not None and isinstance(_func, GenericIODescriptor): - descriptor = _func - _func = None - else: - descriptor = _make_descriptor(**descriptor_kwargs) - - # Ensures the hook is a list - if callable(on_inspect): - inspect_hooks: list[Callable[..., Any]] = [on_inspect] - else: - inspect_hooks: list[Callable[..., Any]] = list(on_inspect or []) # handles None - - def _apply(func: Callable[Concatenate[ManagerBasedEnv, P], R]) -> Callable[Concatenate[ManagerBasedEnv, P], R]: - - # Capture the signature of the function - sig = inspect.signature(func) - - @functools.wraps(func) - def wrapper(env: ManagerBasedEnv, *args: P.args, **kwargs: P.kwargs) -> R: - inspect_flag: bool = kwargs.pop("inspect", False) - out = func(env, *args, **kwargs) - if inspect_flag: - # Injects the function's arguments into the hooks and applies the defaults - bound = sig.bind(env, *args, **kwargs) - bound.apply_defaults() - call_kwargs = { - "output": out, - "descriptor": descriptor, - **bound.arguments, - } - for hook in inspect_hooks: - hook(**call_kwargs) - return out - - # --- Descriptor bookkeeping --- - descriptor.name = func.__name__ - descriptor.full_path = f"{func.__module__}.{func.__name__}" - descriptor.dtype = str(descriptor.dtype) - # Check if description is set in the descriptor - if descriptor.description is None: - descriptor.description = func.__doc__ - - # Adds the descriptor to the wrapped function as an attribute - wrapper._descriptor = descriptor - wrapper._has_descriptor = True - # Alters the signature of the wrapped function to make it match the original function. - # This allows the wrapped functions to pass the checks in the managers. - wrapper.__signature__ = sig - return wrapper - - # If the decorator is used without parentheses, _func will be the function itself. - if callable(_func): - return _apply(_func) - return _apply - - -def record_shape(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): - descriptor.shape = (output.shape[-1],) - - -def record_dtype(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): - descriptor.dtype = str(output.dtype) - - -def record_joint_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): - asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] - joint_ids = kwargs["asset_cfg"].joint_ids - if joint_ids == slice(None, None, None): - joint_ids = list(range(len(asset.joint_names))) - descriptor.joint_names = [asset.joint_names[i] for i in joint_ids] - - -def record_body_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): - asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] - body_ids = kwargs["asset_cfg"].body_ids - if body_ids == slice(None, None, None): - body_ids = list(range(len(asset.body_names))) - descriptor.body_names = [asset.body_names[i] for i in body_ids] - +from isaaclab.envs.utils.io_descriptors import ( + generic_io_descriptor, + record_body_names, + record_dtype, + record_joint_names, + record_shape, +) """ Root state. @@ -299,9 +130,7 @@ def root_ang_vel_w(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntity """ -@generic_io_descriptor( - observation_type="BodyState", body_names=None, on_inspect=[record_shape, record_dtype, record_body_names] -) +@generic_io_descriptor(observation_type="BodyState", on_inspect=[record_shape, record_dtype, record_body_names]) def body_pose_w( env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), @@ -325,9 +154,7 @@ def body_pose_w( return pose.reshape(env.num_envs, -1) -@generic_io_descriptor( - observation_type="BodyState", body_names=None, on_inspect=[record_shape, record_dtype, record_body_names] -) +@generic_io_descriptor(observation_type="BodyState", on_inspect=[record_shape, record_dtype, record_body_names]) def body_projected_gravity_b( env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), @@ -357,9 +184,7 @@ def body_projected_gravity_b( """ -@generic_io_descriptor( - observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape] -) +@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape]) def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset. @@ -370,9 +195,7 @@ def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" return asset.data.joint_pos[:, asset_cfg.joint_ids] -@generic_io_descriptor( - observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape] -) +@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape]) def joint_pos_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset w.r.t. the default joint positions. @@ -383,9 +206,7 @@ def joint_pos_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC return asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids] -@generic_io_descriptor( - observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape] -) +@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape]) def joint_pos_limit_normalized( env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> torch.Tensor: @@ -402,9 +223,7 @@ def joint_pos_limit_normalized( ) -@generic_io_descriptor( - observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape] -) +@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape]) def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset. @@ -415,9 +234,7 @@ def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" return asset.data.joint_vel[:, asset_cfg.joint_ids] -@generic_io_descriptor( - observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape] -) +@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape]) def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset w.r.t. the default joint velocities. @@ -428,9 +245,7 @@ def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC return asset.data.joint_vel[:, asset_cfg.joint_ids] - asset.data.default_joint_vel[:, asset_cfg.joint_ids] -@generic_io_descriptor( - observation_type="JointState", joint_names=None, on_inspect=[record_joint_names, record_dtype, record_shape] -) +@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape]) def joint_effort(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint applied effort of the robot. diff --git a/source/isaaclab/isaaclab/managers/action_manager.py b/source/isaaclab/isaaclab/managers/action_manager.py index 9dce7acd733..2f994686e5b 100644 --- a/source/isaaclab/isaaclab/managers/action_manager.py +++ b/source/isaaclab/isaaclab/managers/action_manager.py @@ -8,16 +8,18 @@ from __future__ import annotations import inspect +import re import torch import weakref from abc import abstractmethod from collections.abc import Sequence from prettytable import PrettyTable -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import omni.kit.app from isaaclab.assets import AssetBase +from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor from .manager_base import ManagerBase, ManagerTermBase from .manager_term_cfg import ActionTermCfg @@ -50,6 +52,7 @@ def __init__(self, cfg: ActionTermCfg, env: ManagerBasedEnv): super().__init__(cfg, env) # parse config to obtain asset to which the term is applied self._asset: AssetBase = self._env.scene[self.cfg.asset_name] + self._IO_descriptor = GenericActionIODescriptor() # add handle for debug visualization (this is set to a valid handle inside set_debug_vis) self._debug_vis_handle = None @@ -91,6 +94,14 @@ def has_debug_vis_implementation(self) -> bool: source_code = inspect.getsource(self._set_debug_vis_impl) return "NotImplementedError" not in source_code + @property + def IO_descriptor(self) -> GenericActionIODescriptor: + """The IO descriptor for the action term.""" + self._IO_descriptor.name = re.sub(r"([a-z])([A-Z])", r"\1_\2", self.__class__.__name__).lower() + self._IO_descriptor.full_path = f"{self.__class__.__module__}.{self.__class__.__name__}" + self._IO_descriptor.description = self.__class__.__doc__ + return self._IO_descriptor + """ Operations. """ @@ -260,7 +271,7 @@ def has_debug_vis_implementation(self) -> bool: return has_debug_vis @property - def get_IO_descriptors(self): + def get_IO_descriptors(self) -> dict[str, dict[str, Any]]: """Get the IO descriptors for the action manager. Returns: @@ -269,16 +280,16 @@ def get_IO_descriptors(self): data = [] - for term in self._terms.values(): + for term_name, term in self._terms.items(): try: - data.append(term.IO_descriptor) + data.append(term.IO_descriptor.__dict__.copy()) except Exception as e: - print(f"Error getting IO descriptor for term: {e}") + print(f"Error getting IO descriptor for term '{term_name}': {e}") formatted_data = {} for item in data: name = item.pop("name") - formatted_item = {"extras": {}} + formatted_item = {"extras": item.pop("extras")} for k, v in item.items(): # Check if v is a tuple and convert to list if isinstance(v, tuple): diff --git a/source/isaaclab/isaaclab/managers/observation_manager.py b/source/isaaclab/isaaclab/managers/observation_manager.py index bc87798e203..71bf7453e00 100644 --- a/source/isaaclab/isaaclab/managers/observation_manager.py +++ b/source/isaaclab/isaaclab/managers/observation_manager.py @@ -264,23 +264,21 @@ def get_IO_descriptors(self): data.append(desc) except Exception as e: print(f"Error getting IO descriptor for term '{term_name}' in group '{group_name}': {e}") - # Format the data for YAML export formatted_data = {} for item in data: name = item.pop("name") - formatted_item = {"overloads": {}, "extras": {}} + formatted_item = {"overloads": {}, "extras": item.pop("extras")} for k, v in item.items(): # Check if v is a tuple and convert to list if isinstance(v, tuple): v = list(v) if k in ["scale", "clip", "history_length", "flatten_history_dim"]: formatted_item["overloads"][k] = v - elif k in ["modifiers", "noise", "description", "units"]: + elif k in ["modifiers", "description", "units"]: formatted_item["extras"][k] = v else: formatted_item[k] = v - formatted_data[name] = formatted_item return formatted_data From ce14a9820781ea9faf78a8cdd77d6785a6c4b57b Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Wed, 11 Jun 2025 15:50:49 +0200 Subject: [PATCH 10/23] Forgot to add some code... --- scripts/environments/export_IODescriptors.py | 79 +++++++ .../isaaclab/envs/utils/io_descriptors.py | 192 ++++++++++++++++++ 2 files changed, 271 insertions(+) create mode 100644 scripts/environments/export_IODescriptors.py create mode 100644 source/isaaclab/isaaclab/envs/utils/io_descriptors.py diff --git a/scripts/environments/export_IODescriptors.py b/scripts/environments/export_IODescriptors.py new file mode 100644 index 00000000000..9f0f1d83040 --- /dev/null +++ b/scripts/environments/export_IODescriptors.py @@ -0,0 +1,79 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to an environment with random action agent.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Random agent for Isaac Lab environments.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() +args_cli.headless = True + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import gymnasium as gym +import torch + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import parse_env_cfg + +# PLACEHOLDER: Extension template (do not remove this comment) + + +def main(): + """Random actions agent with Isaac Lab environment.""" + # create environment configuration + env_cfg = parse_env_cfg( + args_cli.task, device=args_cli.device, num_envs=1, use_fabric=True + ) + # create environment + env = gym.make(args_cli.task, cfg=env_cfg) + + # print info (this is vectorized environment) + print(f"[INFO]: Gym observation space: {env.observation_space}") + print(f"[INFO]: Gym action space: {env.action_space}") + # reset environment + env.reset() + + outs = env.unwrapped.get_IO_descriptors + out = outs["observations"] + out_actions = outs["actions"] + # Make a yaml file with the output + import yaml + + with open("obs_descriptors.yaml", "w") as f: + yaml.safe_dump(outs, f) + + for k, v in out_actions.items(): + print(f"--- Action term: {k} ---") + for k1, v1 in v.items(): + print(f"{k1}: {v1}") + + for k, v in out.items(): + print(f"--- Obs term: {k} ---") + for k1, v1 in v.items(): + print(f"{k1}: {v1}") + env.step(torch.zeros(env.action_space.shape, device=env.unwrapped.device)) + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py new file mode 100644 index 00000000000..6b17ba644aa --- /dev/null +++ b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +from isaaclab.utils import configclass +from collections.abc import Callable +from typing import Any, Concatenate, ParamSpec, TypeVar, TYPE_CHECKING + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + import torch + +import functools +import inspect +import dataclasses + +@configclass +class GenericActionIODescriptor: + mdp_type: str = "Action" + name: str = None + full_path: str = None + description: str = None + shape: tuple[int, ...] = None + dtype: str = None + action_type: str = None + extras: dict[str, Any] = {} + +@configclass +class GenericIODescriptor: + mdp_type: str = "Observation" + name: str = None + full_path: str = None + description: str = None + shape: tuple[int, ...] = None + dtype: str = None + observation_type: str = None + extras: dict[str, Any] = {} + + +# These are defined to help with type hinting +P = ParamSpec("P") +R = TypeVar("R") + + +# Automatically builds a descriptor from the kwargs +def _make_descriptor(**kwargs: Any) -> GenericIODescriptor: + """Split *kwargs* into (known dataclass fields) and (extras).""" + field_names = {f.name for f in dataclasses.fields(GenericIODescriptor)} + known = {k: v for k, v in kwargs.items() if k in field_names} + extras = {k: v for k, v in kwargs.items() if k not in field_names} + + desc = GenericIODescriptor(**known) + # User defined extras are stored in the descriptor under the `extras` field + desc.extras = extras + return desc + + +# Decorator factory for generic IO descriptors. +def generic_io_descriptor( + _func: Callable[Concatenate[ManagerBasedEnv, P], R] | None = None, + *, + on_inspect: Callable[..., Any] | list[Callable[..., Any]] | None = None, + **descriptor_kwargs: Any, +) -> Callable[[Callable[Concatenate[ManagerBasedEnv, P], R]], Callable[Concatenate[ManagerBasedEnv, P], R]]: + """ + Decorator factory for generic IO descriptors. + + This decorator can be used in different ways: + 1. The default decorator has all the information I need for my use case: + ..code-block:: python + @generic_io_descriptor(GenericIODescriptor(description="..", dtype="..")) + def my_func(env: ManagerBasedEnv, *args, **kwargs): + ... + ..note:: If description is not set, the function's docstring is used to populate it. + + 2. I need to add more information to the descriptor: + ..code-block:: python + @generic_io_descriptor(description="..", new_var_1="a", new_var_2="b") + def my_func(env: ManagerBasedEnv, *args, **kwargs): + ... + 3. I need to add a hook to the descriptor: + ..code-block:: python + def record_shape(tensor: torch.Tensor, desc: GenericIODescriptor, **kwargs): + desc.shape = (tensor.shape[-1],) + + @generic_io_descriptor(description="..", new_var_1="a", new_var_2="b", on_inspect=[record_shape, record_dtype]) + def my_func(env: ManagerBasedEnv, *args, **kwargs): + ..note:: The hook is called after the function is called, if and only if the `inspect` flag is set when calling the function. + + For example: + ..code-block:: python + my_func(env, inspect=True) + + 4. I need to add a hook to the descriptor and this hook will write to a variable that is not part of the base descriptor. + ..code-block:: python + def record_joint_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): + asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] + joint_ids = kwargs["asset_cfg"].joint_ids + if joint_ids == slice(None, None, None): + joint_ids = list(range(len(asset.joint_names))) + descriptor.joint_names = [asset.joint_names[i] for i in joint_ids] + + @generic_io_descriptor(new_var_1="a", new_var_2="b", on_inspect=[record_shape, record_dtype, record_joint_names]) + def my_func(env: ManagerBasedEnv, *args, **kwargs): + + ..note:: The hook can access all the variables in the wrapped function's signature. While it is useful, the user should be careful to + access only existing variables. + + Args: + _func: The function to decorate. + **descriptor_kwargs: Keyword arguments to pass to the descriptor. + + Returns: + A decorator that can be used to decorate a function. + """ + + if _func is not None and isinstance(_func, GenericIODescriptor): + descriptor = _func + _func = None + else: + descriptor = _make_descriptor(**descriptor_kwargs) + + # Ensures the hook is a list + if callable(on_inspect): + inspect_hooks: list[Callable[..., Any]] = [on_inspect] + else: + inspect_hooks: list[Callable[..., Any]] = list(on_inspect or []) # handles None + + def _apply(func: Callable[Concatenate[ManagerBasedEnv, P], R]) -> Callable[Concatenate[ManagerBasedEnv, P], R]: + + # Capture the signature of the function + sig = inspect.signature(func) + + @functools.wraps(func) + def wrapper(env: ManagerBasedEnv, *args: P.args, **kwargs: P.kwargs) -> R: + inspect_flag: bool = kwargs.pop("inspect", False) + out = func(env, *args, **kwargs) + if inspect_flag: + # Injects the function's arguments into the hooks and applies the defaults + bound = sig.bind(env, *args, **kwargs) + bound.apply_defaults() + call_kwargs = { + "output": out, + "descriptor": descriptor, + **bound.arguments, + } + for hook in inspect_hooks: + hook(**call_kwargs) + return out + + # --- Descriptor bookkeeping --- + descriptor.name = func.__name__ + descriptor.full_path = f"{func.__module__}.{func.__name__}" + descriptor.dtype = str(descriptor.dtype) + # Check if description is set in the descriptor + if descriptor.description is None: + descriptor.description = func.__doc__ + + # Adds the descriptor to the wrapped function as an attribute + wrapper._descriptor = descriptor + wrapper._has_descriptor = True + # Alters the signature of the wrapped function to make it match the original function. + # This allows the wrapped functions to pass the checks in the managers. + wrapper.__signature__ = sig + return wrapper + + # If the decorator is used without parentheses, _func will be the function itself. + if callable(_func): + return _apply(_func) + return _apply + + +def record_shape(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): + descriptor.shape = (output.shape[-1],) + + +def record_dtype(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): + descriptor.dtype = str(output.dtype) + + +def record_joint_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): + asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] + joint_ids = kwargs["asset_cfg"].joint_ids + if joint_ids == slice(None, None, None): + joint_ids = list(range(len(asset.joint_names))) + descriptor.joint_names = [asset.joint_names[i] for i in joint_ids] + + +def record_body_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): + asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] + body_ids = kwargs["asset_cfg"].body_ids + if body_ids == slice(None, None, None): + body_ids = list(range(len(asset.body_names))) + descriptor.body_names = [asset.body_names[i] for i in body_ids] \ No newline at end of file From 6ceb35f75836d60998b0c7ad3d890b8807da4f22 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Wed, 11 Jun 2025 15:51:11 +0200 Subject: [PATCH 11/23] Forgot to add some code... --- .../isaaclab/envs/mdp/actions/task_space_actions.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py index 6f9bf0553ec..bde914ccae9 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py @@ -435,12 +435,17 @@ def IO_descriptor(self) -> GenericActionIODescriptor: self._IO_descriptor.action_type = "TaskSpaceAction" self._IO_descriptor.body_name = self._ee_body_name self._IO_descriptor.joint_names = self._joint_names - self._IO_descriptor.scale = self._scale + self._IO_descriptor.position_scale = self.cfg.position_scale + self._IO_descriptor.orientation_scale = self.cfg.orientation_scale + self._IO_descriptor.wrench_scale = self.cfg.wrench_scale + self._IO_descriptor.stiffness_scale = self.cfg.stiffness_scale + self._IO_descriptor.damping_ratio_scale = self.cfg.damping_ratio_scale + self._IO_descriptor.nullspace_joint_pos_target = self.cfg.nullspace_joint_pos_target if self.cfg.clip is not None: self._IO_descriptor.clip = self.cfg.clip else: self._IO_descriptor.clip = None - self._IO_descriptor.extras["controller_cfg"] = self.cfg.controller.__dict__ + self._IO_descriptor.extras["controller_cfg"] = self.cfg.controller_cfg.__dict__ self._IO_descriptor.extras["body_offset"] = self.cfg.body_offset.__dict__ return self._IO_descriptor From b2d2bf1b33398962949aaafd73523a8532fce1e0 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Wed, 11 Jun 2025 15:56:37 +0200 Subject: [PATCH 12/23] undid changes to random action --- scripts/environments/random_agent.py | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/scripts/environments/random_agent.py b/scripts/environments/random_agent.py index 55f5fe92d0a..b3187c3b372 100644 --- a/scripts/environments/random_agent.py +++ b/scripts/environments/random_agent.py @@ -52,26 +52,6 @@ def main(): print(f"[INFO]: Gym action space: {env.action_space}") # reset environment env.reset() - - outs = env.unwrapped.get_IO_descriptors - out = outs["observations"] - out_actions = outs["actions"] - # Make a yaml file with the output - import yaml - - with open("obs_descriptors.yaml", "w") as f: - yaml.safe_dump(outs, f) - - for k, v in out_actions.items(): - print(f"--- Action term: {k} ---") - for k1, v1 in v.items(): - print(f"{k1}: {v1}") - - for k, v in out.items(): - print(f"--- Obs term: {k} ---") - for k1, v1 in v.items(): - print(f"{k1}: {v1}") - exit(0) # simulate environment while simulation_app.is_running(): # run everything in inference mode From 699595e512d9549e474fdda4a211ba5c09f843e4 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Wed, 11 Jun 2025 16:07:02 +0200 Subject: [PATCH 13/23] chanegd naming conventions --- scripts/environments/export_IODescriptors.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/scripts/environments/export_IODescriptors.py b/scripts/environments/export_IODescriptors.py index 9f0f1d83040..b163d4faee8 100644 --- a/scripts/environments/export_IODescriptors.py +++ b/scripts/environments/export_IODescriptors.py @@ -56,7 +56,10 @@ def main(): # Make a yaml file with the output import yaml - with open("obs_descriptors.yaml", "w") as f: + name = args_cli.task.lower().replace("-", "_") + name = name.replace(" ", "_") + + with open(f"{name}_IO_descriptors.yaml", "w") as f: yaml.safe_dump(outs, f) for k, v in out_actions.items(): From d1158278e57e401d037c23dafc1eb9534db17095 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Thu, 12 Jun 2025 11:31:20 +0200 Subject: [PATCH 14/23] added stripping of doc strings + units here and here --- source/isaaclab/isaaclab/envs/mdp/observations.py | 10 +++++----- source/isaaclab/isaaclab/envs/utils/io_descriptors.py | 2 +- source/isaaclab/isaaclab/managers/action_manager.py | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index f0cc1e13cff..fa37a364ccd 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -184,7 +184,7 @@ def body_projected_gravity_b( """ -@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape]) +@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad") def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset. @@ -195,7 +195,7 @@ def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" return asset.data.joint_pos[:, asset_cfg.joint_ids] -@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape]) +@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad") def joint_pos_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset w.r.t. the default joint positions. @@ -223,7 +223,7 @@ def joint_pos_limit_normalized( ) -@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape]) +@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad/s") def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset. @@ -234,7 +234,7 @@ def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" return asset.data.joint_vel[:, asset_cfg.joint_ids] -@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape]) +@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad/s") def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset w.r.t. the default joint velocities. @@ -245,7 +245,7 @@ def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC return asset.data.joint_vel[:, asset_cfg.joint_ids] - asset.data.default_joint_vel[:, asset_cfg.joint_ids] -@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape]) +@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="N.m") def joint_effort(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint applied effort of the robot. diff --git a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py index 6b17ba644aa..2a1573cfa14 100644 --- a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py +++ b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py @@ -152,7 +152,7 @@ def wrapper(env: ManagerBasedEnv, *args: P.args, **kwargs: P.kwargs) -> R: descriptor.dtype = str(descriptor.dtype) # Check if description is set in the descriptor if descriptor.description is None: - descriptor.description = func.__doc__ + descriptor.description = " ".join(func.__doc__.split()) # Adds the descriptor to the wrapped function as an attribute wrapper._descriptor = descriptor diff --git a/source/isaaclab/isaaclab/managers/action_manager.py b/source/isaaclab/isaaclab/managers/action_manager.py index 2f994686e5b..7d2f6f1bd03 100644 --- a/source/isaaclab/isaaclab/managers/action_manager.py +++ b/source/isaaclab/isaaclab/managers/action_manager.py @@ -99,7 +99,7 @@ def IO_descriptor(self) -> GenericActionIODescriptor: """The IO descriptor for the action term.""" self._IO_descriptor.name = re.sub(r"([a-z])([A-Z])", r"\1_\2", self.__class__.__name__).lower() self._IO_descriptor.full_path = f"{self.__class__.__module__}.{self.__class__.__name__}" - self._IO_descriptor.description = self.__class__.__doc__ + self._IO_descriptor.description = " ".join(self.__class__.__doc__.split()) return self._IO_descriptor """ From 35ba3a70f52d8ede89bf174038806fa5b50ab6c3 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Wed, 18 Jun 2025 12:22:41 +0200 Subject: [PATCH 15/23] happy pre-commits --- scripts/environments/export_IODescriptors.py | 4 +--- .../isaaclab/envs/mdp/observations.py | 20 ++++++++++++++----- .../isaaclab/envs/utils/io_descriptors.py | 17 ++++++++++++---- 3 files changed, 29 insertions(+), 12 deletions(-) diff --git a/scripts/environments/export_IODescriptors.py b/scripts/environments/export_IODescriptors.py index b163d4faee8..0dea97b33b8 100644 --- a/scripts/environments/export_IODescriptors.py +++ b/scripts/environments/export_IODescriptors.py @@ -38,9 +38,7 @@ def main(): """Random actions agent with Isaac Lab environment.""" # create environment configuration - env_cfg = parse_env_cfg( - args_cli.task, device=args_cli.device, num_envs=1, use_fabric=True - ) + env_cfg = parse_env_cfg(args_cli.task, device=args_cli.device, num_envs=1, use_fabric=True) # create environment env = gym.make(args_cli.task, cfg=env_cfg) diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index fa37a364ccd..c9d14993c32 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -184,7 +184,9 @@ def body_projected_gravity_b( """ -@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad") +@generic_io_descriptor( + observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad" +) def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset. @@ -195,7 +197,9 @@ def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" return asset.data.joint_pos[:, asset_cfg.joint_ids] -@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad") +@generic_io_descriptor( + observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad" +) def joint_pos_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset w.r.t. the default joint positions. @@ -223,7 +227,9 @@ def joint_pos_limit_normalized( ) -@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad/s") +@generic_io_descriptor( + observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad/s" +) def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset. @@ -234,7 +240,9 @@ def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" return asset.data.joint_vel[:, asset_cfg.joint_ids] -@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad/s") +@generic_io_descriptor( + observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad/s" +) def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset w.r.t. the default joint velocities. @@ -245,7 +253,9 @@ def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityC return asset.data.joint_vel[:, asset_cfg.joint_ids] - asset.data.default_joint_vel[:, asset_cfg.joint_ids] -@generic_io_descriptor(observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="N.m") +@generic_io_descriptor( + observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="N.m" +) def joint_effort(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint applied effort of the robot. diff --git a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py index 2a1573cfa14..4746acc4da8 100644 --- a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py +++ b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py @@ -1,16 +1,24 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + from __future__ import annotations -from isaaclab.utils import configclass from collections.abc import Callable -from typing import Any, Concatenate, ParamSpec, TypeVar, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Concatenate, ParamSpec, TypeVar + +from isaaclab.utils import configclass if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv + from isaaclab.assets.articulation import Articulation import torch +import dataclasses import functools import inspect -import dataclasses + @configclass class GenericActionIODescriptor: @@ -23,6 +31,7 @@ class GenericActionIODescriptor: action_type: str = None extras: dict[str, Any] = {} + @configclass class GenericIODescriptor: mdp_type: str = "Observation" @@ -189,4 +198,4 @@ def record_body_names(output: torch.Tensor, descriptor: GenericIODescriptor, **k body_ids = kwargs["asset_cfg"].body_ids if body_ids == slice(None, None, None): body_ids = list(range(len(asset.body_names))) - descriptor.body_names = [asset.body_names[i] for i in body_ids] \ No newline at end of file + descriptor.body_names = [asset.body_names[i] for i in body_ids] From 8c2539cf0aa9bd86e7a77c9a55ff420dcd80d631 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Tue, 8 Jul 2025 10:37:53 +0200 Subject: [PATCH 16/23] added FIXME --- source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py index 8341c59184f..cb5966c7a3b 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py @@ -137,7 +137,8 @@ def IO_descriptor(self) -> GenericActionIODescriptor: self._IO_descriptor.offset = self._offset[0].detach().cpu().numpy().tolist() else: self._IO_descriptor.offset = self._offset - if self.cfg.clip is not None: + #FIXME: This is not correct. Add list support. + if self.cfg.clip is not None: self._IO_descriptor.clip = self._clip else: self._IO_descriptor.clip = None From f913b7895e9191ab731b2b8d3f44a2596e143c45 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Thu, 10 Jul 2025 10:01:57 +0200 Subject: [PATCH 17/23] Updated based on feedback --- .../envs/mdp/actions/joint_actions.py | 4 +-- .../isaaclab/envs/mdp/observations.py | 6 ++-- .../isaaclab/envs/utils/io_descriptors.py | 31 +++++++++++++++- .../isaaclab/managers/observation_manager.py | 36 ++++++++++--------- 4 files changed, 55 insertions(+), 22 deletions(-) diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py index cb5966c7a3b..d830f411919 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py @@ -138,8 +138,8 @@ def IO_descriptor(self) -> GenericActionIODescriptor: else: self._IO_descriptor.offset = self._offset #FIXME: This is not correct. Add list support. - if self.cfg.clip is not None: - self._IO_descriptor.clip = self._clip + if isinstance(self._clip, torch.Tensor): + self._IO_descriptor.clip = self._clip[0].detach().cpu().numpy().tolist() else: self._IO_descriptor.clip = None return self._IO_descriptor diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index c9d14993c32..f130066dae9 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -30,6 +30,8 @@ record_dtype, record_joint_names, record_shape, + record_joint_pos_offsets, + record_joint_vel_offsets, ) """ @@ -198,7 +200,7 @@ def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" @generic_io_descriptor( - observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad" + observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape, record_joint_pos_offsets], units="rad" ) def joint_pos_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset w.r.t. the default joint positions. @@ -241,7 +243,7 @@ def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" @generic_io_descriptor( - observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape], units="rad/s" + observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape, record_joint_vel_offsets], units="rad/s" ) def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset w.r.t. the default joint velocities. diff --git a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py index 4746acc4da8..89eee03933c 100644 --- a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py +++ b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py @@ -192,10 +192,39 @@ def record_joint_names(output: torch.Tensor, descriptor: GenericIODescriptor, ** joint_ids = list(range(len(asset.joint_names))) descriptor.joint_names = [asset.joint_names[i] for i in joint_ids] - def record_body_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] body_ids = kwargs["asset_cfg"].body_ids if body_ids == slice(None, None, None): body_ids = list(range(len(asset.body_names))) descriptor.body_names = [asset.body_names[i] for i in body_ids] + +def record_joint_pos_offsets(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): + asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] + ids = kwargs["asset_cfg"].joint_ids + # Get the offsets of the joints for the first robot in the scene. + # This assumes that all robots have the same joint offsets. + descriptor.joint_pos_offsets = asset.data.default_joint_pos[:, ids][0] + +def record_joint_vel_offsets(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): + asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] + ids = kwargs["asset_cfg"].joint_ids + # Get the offsets of the joints for the first robot in the scene. + # This assumes that all robots have the same joint offsets. + descriptor.joint_vel_offsets = asset.data.default_joint_vel[:, ids][0] + + +def export_articulations_data(env: ManagerBasedEnv): + articulation_joint_data = {} + for articulation_name, articulation in env.scene.articulations.items(): + articulation_joint_data[articulation_name] = {} + articulation_joint_data[articulation_name]["joint_names"] = articulation.joint_names + articulation_joint_data[articulation_name]["default_joint_pos"] = articulation.data.default_joint_pos[0].detach().cpu().numpy().tolist() + articulation_joint_data[articulation_name]["default_joint_vel"] = articulation.data.default_joint_vel[0].detach().cpu().numpy().tolist() + articulation_joint_data[articulation_name]["default_joint_pos_limits"] = articulation.data.default_joint_pos_limits[0].detach().cpu().numpy().tolist() + articulation_joint_data[articulation_name]["default_joint_damping"] = articulation.data.default_joint_damping[0].detach().cpu().numpy().tolist() + articulation_joint_data[articulation_name]["default_joint_stiffness"] = articulation.data.default_joint_stiffness[0].detach().cpu().numpy().tolist() + articulation_joint_data[articulation_name]["default_joint_friction"] = articulation.data.default_joint_friction[0].detach().cpu().numpy().tolist() + articulation_joint_data[articulation_name]["default_joint_armature"] = articulation.data.default_joint_armature[0].detach().cpu().numpy().tolist() + return articulation_joint_data + \ No newline at end of file diff --git a/source/isaaclab/isaaclab/managers/observation_manager.py b/source/isaaclab/isaaclab/managers/observation_manager.py index 71bf7453e00..942b27770ef 100644 --- a/source/isaaclab/isaaclab/managers/observation_manager.py +++ b/source/isaaclab/isaaclab/managers/observation_manager.py @@ -233,9 +233,10 @@ def get_IO_descriptors(self): A dictionary with keys as the group names and values as the IO descriptors. """ - data = [] + group_data = {} for group_name in self._group_obs_term_names: + group_data[group_name] = [] # check ig group name is valid if group_name not in self._group_obs_term_names: raise ValueError( @@ -261,26 +262,27 @@ def get_IO_descriptors(self): if k in ["modifiers", "clip", "scale", "history_length", "flatten_history_dim"]: overloads[k] = v desc.update(overloads) - data.append(desc) + group_data[group_name].append(desc) except Exception as e: print(f"Error getting IO descriptor for term '{term_name}' in group '{group_name}': {e}") # Format the data for YAML export formatted_data = {} - for item in data: - name = item.pop("name") - formatted_item = {"overloads": {}, "extras": item.pop("extras")} - for k, v in item.items(): - # Check if v is a tuple and convert to list - if isinstance(v, tuple): - v = list(v) - if k in ["scale", "clip", "history_length", "flatten_history_dim"]: - formatted_item["overloads"][k] = v - elif k in ["modifiers", "description", "units"]: - formatted_item["extras"][k] = v - else: - formatted_item[k] = v - formatted_data[name] = formatted_item - + for group_name, data in group_data.items(): + formatted_data[group_name] = [] + for item in data: + name = item.pop("name") + formatted_item = {"name": name, "overloads": {}, "extras": item.pop("extras")} + for k, v in item.items(): + # Check if v is a tuple and convert to list + if isinstance(v, tuple): + v = list(v) + if k in ["scale", "clip", "history_length", "flatten_history_dim"]: + formatted_item["overloads"][k] = v + elif k in ["modifiers", "description", "units"]: + formatted_item["extras"][k] = v + else: + formatted_item[k] = v + formatted_data[group_name].append(formatted_item) return formatted_data """ From 17cd43c58a7d0d295d1ae0fce38861aec9f74c4c Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Thu, 10 Jul 2025 10:48:00 +0200 Subject: [PATCH 18/23] Should be good --- scripts/environments/export_IODescriptors.py | 25 +++++++++++++------ .../isaaclab/envs/manager_based_env.py | 3 ++- .../envs/mdp/actions/joint_actions.py | 7 ++++-- .../isaaclab/managers/action_manager.py | 8 +++--- .../isaaclab/managers/observation_manager.py | 6 ++++- 5 files changed, 34 insertions(+), 15 deletions(-) diff --git a/scripts/environments/export_IODescriptors.py b/scripts/environments/export_IODescriptors.py index 0dea97b33b8..b68574b62fd 100644 --- a/scripts/environments/export_IODescriptors.py +++ b/scripts/environments/export_IODescriptors.py @@ -49,8 +49,9 @@ def main(): env.reset() outs = env.unwrapped.get_IO_descriptors - out = outs["observations"] + out_observations = outs["observations"] out_actions = outs["actions"] + out_articulations = outs["articulations"] # Make a yaml file with the output import yaml @@ -60,15 +61,25 @@ def main(): with open(f"{name}_IO_descriptors.yaml", "w") as f: yaml.safe_dump(outs, f) - for k, v in out_actions.items(): - print(f"--- Action term: {k} ---") - for k1, v1 in v.items(): + for k in out_actions: + print(f"--- Action term: {k['name']} ---") + k.pop("name") + for k1, v1 in k.items(): print(f"{k1}: {v1}") - for k, v in out.items(): - print(f"--- Obs term: {k} ---") - for k1, v1 in v.items(): + for obs_group_name, obs_group in out_observations.items(): + print(f"--- Obs group: {obs_group_name} ---") + for k in obs_group: + print(f"--- Obs term: {k['name']} ---") + k.pop("name") + for k1, v1 in k.items(): + print(f"{k1}: {v1}") + + for articulation_name, articulation_data in out_articulations.items(): + print(f"--- Articulation: {articulation_name} ---") + for k1, v1 in articulation_data.items(): print(f"{k1}: {v1}") + env.step(torch.zeros(env.action_space.shape, device=env.unwrapped.device)) env.close() diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index ee642cbf20f..12668d35a7b 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -21,7 +21,7 @@ from .common import VecEnvObs from .manager_based_env_cfg import ManagerBasedEnvCfg from .ui import ViewportCameraController - +from .utils.io_descriptors import export_articulations_data class ManagerBasedEnv: """The base environment encapsulates the simulation scene and the environment managers for the manager-based workflow. @@ -220,6 +220,7 @@ def get_IO_descriptors(self): return { "observations": self.observation_manager.get_IO_descriptors, "actions": self.action_manager.get_IO_descriptors, + "articulations": export_articulations_data(self), } """ diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py index d830f411919..791a729a459 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py @@ -138,8 +138,11 @@ def IO_descriptor(self) -> GenericActionIODescriptor: else: self._IO_descriptor.offset = self._offset #FIXME: This is not correct. Add list support. - if isinstance(self._clip, torch.Tensor): - self._IO_descriptor.clip = self._clip[0].detach().cpu().numpy().tolist() + if self.cfg.clip is not None: + if isinstance(self._clip, torch.Tensor): + self._IO_descriptor.clip = self._clip[0].detach().cpu().numpy().tolist() + else: + self._IO_descriptor.clip = self._clip else: self._IO_descriptor.clip = None return self._IO_descriptor diff --git a/source/isaaclab/isaaclab/managers/action_manager.py b/source/isaaclab/isaaclab/managers/action_manager.py index 7d2f6f1bd03..aef31f89e62 100644 --- a/source/isaaclab/isaaclab/managers/action_manager.py +++ b/source/isaaclab/isaaclab/managers/action_manager.py @@ -271,7 +271,7 @@ def has_debug_vis_implementation(self) -> bool: return has_debug_vis @property - def get_IO_descriptors(self) -> dict[str, dict[str, Any]]: + def get_IO_descriptors(self) -> list[dict[str, Any]]: """Get the IO descriptors for the action manager. Returns: @@ -286,10 +286,10 @@ def get_IO_descriptors(self) -> dict[str, dict[str, Any]]: except Exception as e: print(f"Error getting IO descriptor for term '{term_name}': {e}") - formatted_data = {} + formatted_data = [] for item in data: name = item.pop("name") - formatted_item = {"extras": item.pop("extras")} + formatted_item = {"name": name, "extras": item.pop("extras")} for k, v in item.items(): # Check if v is a tuple and convert to list if isinstance(v, tuple): @@ -298,7 +298,7 @@ def get_IO_descriptors(self) -> dict[str, dict[str, Any]]: formatted_item["extras"][k] = v else: formatted_item[k] = v - formatted_data[name] = formatted_item + formatted_data.append(formatted_item) return formatted_data diff --git a/source/isaaclab/isaaclab/managers/observation_manager.py b/source/isaaclab/isaaclab/managers/observation_manager.py index 942b27770ef..c649e302131 100644 --- a/source/isaaclab/isaaclab/managers/observation_manager.py +++ b/source/isaaclab/isaaclab/managers/observation_manager.py @@ -226,7 +226,7 @@ def group_obs_concatenate(self) -> dict[str, bool]: return self._group_obs_concatenate @property - def get_IO_descriptors(self): + def get_IO_descriptors(self, group_names_to_export: list[str] = ["policy"]): """Get the IO descriptors for the observation manager. Returns: @@ -276,6 +276,9 @@ def get_IO_descriptors(self): # Check if v is a tuple and convert to list if isinstance(v, tuple): v = list(v) + # Check if v is a tensor and convert to list + if isinstance(v, torch.Tensor): + v = v.detach().cpu().numpy().tolist() if k in ["scale", "clip", "history_length", "flatten_history_dim"]: formatted_item["overloads"][k] = v elif k in ["modifiers", "description", "units"]: @@ -283,6 +286,7 @@ def get_IO_descriptors(self): else: formatted_item[k] = v formatted_data[group_name].append(formatted_item) + formatted_data = {k: v for k, v in formatted_data.items() if k in group_names_to_export} return formatted_data """ From b0a2b5e09bdeccf74a0841a0b8c6be89a772db55 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Thu, 10 Jul 2025 11:29:47 +0200 Subject: [PATCH 19/23] ran pre-commits --- .../isaaclab/envs/manager_based_env.py | 1 + .../envs/mdp/actions/joint_actions.py | 2 +- .../isaaclab/envs/mdp/observations.py | 10 ++++-- .../isaaclab/envs/utils/io_descriptors.py | 32 ++++++++++++++----- 4 files changed, 33 insertions(+), 12 deletions(-) diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 12668d35a7b..f1fa7eb9f2a 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -23,6 +23,7 @@ from .ui import ViewportCameraController from .utils.io_descriptors import export_articulations_data + class ManagerBasedEnv: """The base environment encapsulates the simulation scene and the environment managers for the manager-based workflow. diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py index 791a729a459..dc056f92329 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py @@ -137,7 +137,7 @@ def IO_descriptor(self) -> GenericActionIODescriptor: self._IO_descriptor.offset = self._offset[0].detach().cpu().numpy().tolist() else: self._IO_descriptor.offset = self._offset - #FIXME: This is not correct. Add list support. + # FIXME: This is not correct. Add list support. if self.cfg.clip is not None: if isinstance(self._clip, torch.Tensor): self._IO_descriptor.clip = self._clip[0].detach().cpu().numpy().tolist() diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index f130066dae9..a748a01f629 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -29,9 +29,9 @@ record_body_names, record_dtype, record_joint_names, - record_shape, record_joint_pos_offsets, record_joint_vel_offsets, + record_shape, ) """ @@ -200,7 +200,9 @@ def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" @generic_io_descriptor( - observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape, record_joint_pos_offsets], units="rad" + observation_type="JointState", + on_inspect=[record_joint_names, record_dtype, record_shape, record_joint_pos_offsets], + units="rad", ) def joint_pos_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: """The joint positions of the asset w.r.t. the default joint positions. @@ -243,7 +245,9 @@ def joint_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" @generic_io_descriptor( - observation_type="JointState", on_inspect=[record_joint_names, record_dtype, record_shape, record_joint_vel_offsets], units="rad/s" + observation_type="JointState", + on_inspect=[record_joint_names, record_dtype, record_shape, record_joint_vel_offsets], + units="rad/s", ) def joint_vel_rel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")): """The joint velocities of the asset w.r.t. the default joint velocities. diff --git a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py index 89eee03933c..efe72be9fd8 100644 --- a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py +++ b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py @@ -192,6 +192,7 @@ def record_joint_names(output: torch.Tensor, descriptor: GenericIODescriptor, ** joint_ids = list(range(len(asset.joint_names))) descriptor.joint_names = [asset.joint_names[i] for i in joint_ids] + def record_body_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] body_ids = kwargs["asset_cfg"].body_ids @@ -199,6 +200,7 @@ def record_body_names(output: torch.Tensor, descriptor: GenericIODescriptor, **k body_ids = list(range(len(asset.body_names))) descriptor.body_names = [asset.body_names[i] for i in body_ids] + def record_joint_pos_offsets(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] ids = kwargs["asset_cfg"].joint_ids @@ -206,6 +208,7 @@ def record_joint_pos_offsets(output: torch.Tensor, descriptor: GenericIODescript # This assumes that all robots have the same joint offsets. descriptor.joint_pos_offsets = asset.data.default_joint_pos[:, ids][0] + def record_joint_vel_offsets(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] ids = kwargs["asset_cfg"].joint_ids @@ -219,12 +222,25 @@ def export_articulations_data(env: ManagerBasedEnv): for articulation_name, articulation in env.scene.articulations.items(): articulation_joint_data[articulation_name] = {} articulation_joint_data[articulation_name]["joint_names"] = articulation.joint_names - articulation_joint_data[articulation_name]["default_joint_pos"] = articulation.data.default_joint_pos[0].detach().cpu().numpy().tolist() - articulation_joint_data[articulation_name]["default_joint_vel"] = articulation.data.default_joint_vel[0].detach().cpu().numpy().tolist() - articulation_joint_data[articulation_name]["default_joint_pos_limits"] = articulation.data.default_joint_pos_limits[0].detach().cpu().numpy().tolist() - articulation_joint_data[articulation_name]["default_joint_damping"] = articulation.data.default_joint_damping[0].detach().cpu().numpy().tolist() - articulation_joint_data[articulation_name]["default_joint_stiffness"] = articulation.data.default_joint_stiffness[0].detach().cpu().numpy().tolist() - articulation_joint_data[articulation_name]["default_joint_friction"] = articulation.data.default_joint_friction[0].detach().cpu().numpy().tolist() - articulation_joint_data[articulation_name]["default_joint_armature"] = articulation.data.default_joint_armature[0].detach().cpu().numpy().tolist() + articulation_joint_data[articulation_name]["default_joint_pos"] = ( + articulation.data.default_joint_pos[0].detach().cpu().numpy().tolist() + ) + articulation_joint_data[articulation_name]["default_joint_vel"] = ( + articulation.data.default_joint_vel[0].detach().cpu().numpy().tolist() + ) + articulation_joint_data[articulation_name]["default_joint_pos_limits"] = ( + articulation.data.default_joint_pos_limits[0].detach().cpu().numpy().tolist() + ) + articulation_joint_data[articulation_name]["default_joint_damping"] = ( + articulation.data.default_joint_damping[0].detach().cpu().numpy().tolist() + ) + articulation_joint_data[articulation_name]["default_joint_stiffness"] = ( + articulation.data.default_joint_stiffness[0].detach().cpu().numpy().tolist() + ) + articulation_joint_data[articulation_name]["default_joint_friction"] = ( + articulation.data.default_joint_friction[0].detach().cpu().numpy().tolist() + ) + articulation_joint_data[articulation_name]["default_joint_armature"] = ( + articulation.data.default_joint_armature[0].detach().cpu().numpy().tolist() + ) return articulation_joint_data - \ No newline at end of file From eb6a339b1234c46096c4d8669805f3356136e89c Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Thu, 10 Jul 2025 12:18:47 +0200 Subject: [PATCH 20/23] Added sim info --- scripts/environments/export_IODescriptors.py | 4 ++++ source/isaaclab/isaaclab/envs/manager_based_env.py | 3 ++- source/isaaclab/isaaclab/envs/utils/io_descriptors.py | 8 ++++++++ 3 files changed, 14 insertions(+), 1 deletion(-) diff --git a/scripts/environments/export_IODescriptors.py b/scripts/environments/export_IODescriptors.py index b68574b62fd..5e01692ede1 100644 --- a/scripts/environments/export_IODescriptors.py +++ b/scripts/environments/export_IODescriptors.py @@ -52,6 +52,7 @@ def main(): out_observations = outs["observations"] out_actions = outs["actions"] out_articulations = outs["articulations"] + out_scene = outs["scene"] # Make a yaml file with the output import yaml @@ -80,6 +81,9 @@ def main(): for k1, v1 in articulation_data.items(): print(f"{k1}: {v1}") + for k1, v1 in out_scene.items(): + print(f"{k1}: {v1}") + env.step(torch.zeros(env.action_space.shape, device=env.unwrapped.device)) env.close() diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index f1fa7eb9f2a..c1691b60988 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -21,7 +21,7 @@ from .common import VecEnvObs from .manager_based_env_cfg import ManagerBasedEnvCfg from .ui import ViewportCameraController -from .utils.io_descriptors import export_articulations_data +from .utils.io_descriptors import export_articulations_data, export_scene_data class ManagerBasedEnv: @@ -222,6 +222,7 @@ def get_IO_descriptors(self): "observations": self.observation_manager.get_IO_descriptors, "actions": self.action_manager.get_IO_descriptors, "articulations": export_articulations_data(self), + "scene": export_scene_data(self), } """ diff --git a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py index efe72be9fd8..aca15686313 100644 --- a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py +++ b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py @@ -244,3 +244,11 @@ def export_articulations_data(env: ManagerBasedEnv): articulation.data.default_joint_armature[0].detach().cpu().numpy().tolist() ) return articulation_joint_data + + +def export_scene_data(env: ManagerBasedEnv): + scene_data = {} + scene_data["physics_dt"] = env.physics_dt + scene_data["dt"] = env.step_dt + scene_data["decimation"] = env.cfg.decimation + return scene_data \ No newline at end of file From 60d94802fd34c70228a1205ba7079958fd4fc60d Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Thu, 10 Jul 2025 15:13:04 +0200 Subject: [PATCH 21/23] added param to limit action exports on some terms --- source/isaaclab/isaaclab/envs/utils/io_descriptors.py | 1 + source/isaaclab/isaaclab/managers/action_manager.py | 10 ++++++++++ 2 files changed, 11 insertions(+) diff --git a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py index aca15686313..4a852b4ced9 100644 --- a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py +++ b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py @@ -30,6 +30,7 @@ class GenericActionIODescriptor: dtype: str = None action_type: str = None extras: dict[str, Any] = {} + export: bool = True @configclass diff --git a/source/isaaclab/isaaclab/managers/action_manager.py b/source/isaaclab/isaaclab/managers/action_manager.py index aef31f89e62..9b561ceb6a7 100644 --- a/source/isaaclab/isaaclab/managers/action_manager.py +++ b/source/isaaclab/isaaclab/managers/action_manager.py @@ -53,6 +53,7 @@ def __init__(self, cfg: ActionTermCfg, env: ManagerBasedEnv): # parse config to obtain asset to which the term is applied self._asset: AssetBase = self._env.scene[self.cfg.asset_name] self._IO_descriptor = GenericActionIODescriptor() + self._export_IO_descriptor = True # add handle for debug visualization (this is set to a valid handle inside set_debug_vis) self._debug_vis_handle = None @@ -100,8 +101,14 @@ def IO_descriptor(self) -> GenericActionIODescriptor: self._IO_descriptor.name = re.sub(r"([a-z])([A-Z])", r"\1_\2", self.__class__.__name__).lower() self._IO_descriptor.full_path = f"{self.__class__.__module__}.{self.__class__.__name__}" self._IO_descriptor.description = " ".join(self.__class__.__doc__.split()) + self._IO_descriptor.export = self.export_IO_descriptor return self._IO_descriptor + @property + def export_IO_descriptor(self) -> bool: + """Whether to export the IO descriptor for the action term.""" + return self._export_IO_descriptor + """ Operations. """ @@ -290,6 +297,9 @@ def get_IO_descriptors(self) -> list[dict[str, Any]]: for item in data: name = item.pop("name") formatted_item = {"name": name, "extras": item.pop("extras")} + print(item["export"]) + if not item.pop("export"): + continue for k, v in item.items(): # Check if v is a tuple and convert to list if isinstance(v, tuple): From 92ccb45fdadb97491d95ce95ec58f76fc8936f53 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Fri, 18 Jul 2025 15:22:13 +0200 Subject: [PATCH 22/23] Started working on the doc. --- .../envs/mdp/actions/joint_actions.py | 12 ++ .../mdp/actions/joint_actions_to_limits.py | 25 +++ .../envs/mdp/actions/non_holonomic_actions.py | 15 ++ .../mdp/actions/pink_task_space_actions.py | 14 ++ .../envs/mdp/actions/task_space_actions.py | 33 ++++ .../isaaclab/envs/utils/io_descriptors.py | 148 ++++++++++++++++-- 6 files changed, 233 insertions(+), 14 deletions(-) diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py index dc056f92329..dfc1088e685 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py @@ -126,6 +126,18 @@ def processed_actions(self) -> torch.Tensor: @property def IO_descriptor(self) -> GenericActionIODescriptor: + """The IO descriptor of the action term. + + This descriptor is used to describe the action term of the joint action. + It adds the following information to the base descriptor: + - joint_names: The names of the joints. + - scale: The scale of the action term. + - offset: The offset of the action term. + - clip: The clip of the action term. + + Returns: + The IO descriptor of the action term. + """ super().IO_descriptor self._IO_descriptor.shape = (self.action_dim,) self._IO_descriptor.dtype = str(self.raw_actions.dtype) diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py index d056d2e6ac6..fa220fed20d 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py @@ -108,6 +108,18 @@ def processed_actions(self) -> torch.Tensor: @property def IO_descriptor(self) -> GenericActionIODescriptor: + """The IO descriptor of the action term. + + This descriptor is used to describe the action term of the joint position to limits action. + It adds the following information to the base descriptor: + - joint_names: The names of the joints. + - scale: The scale of the action term. + - offset: The offset of the action term. + - clip: The clip of the action term. + + Returns: + The IO descriptor of the action term. + """ super().IO_descriptor self._IO_descriptor.shape = (self.action_dim,) self._IO_descriptor.dtype = str(self.raw_actions.dtype) @@ -217,6 +229,19 @@ def __init__(self, cfg: actions_cfg.EMAJointPositionToLimitsActionCfg, env: Mana @property def IO_descriptor(self) -> GenericActionIODescriptor: + """The IO descriptor of the action term. + + This descriptor is used to describe the action term of the EMA joint position to limits action. + It adds the following information to the base descriptor: + - joint_names: The names of the joints. + - scale: The scale of the action term. + - offset: The offset of the action term. + - clip: The clip of the action term. + - alpha: The moving average weight. + + Returns: + The IO descriptor of the action term. + """ super().IO_descriptor if isinstance(self._alpha, float): self._IO_descriptor.alpha = self._alpha diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py index 5d168da4151..c8eefc8cf2a 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py @@ -137,6 +137,21 @@ def processed_actions(self) -> torch.Tensor: @property def IO_descriptor(self) -> GenericActionIODescriptor: + """The IO descriptor of the action term. + + This descriptor is used to describe the action term of the non-holonomic action. + It adds the following information to the base descriptor: + - scale: The scale of the action term. + - offset: The offset of the action term. + - clip: The clip of the action term. + - body_name: The name of the body. + - x_joint_name: The name of the x joint. + - y_joint_name: The name of the y joint. + - yaw_joint_name: The name of the yaw joint. + + Returns: + The IO descriptor of the action term. + """ super().IO_descriptor self._IO_descriptor.shape = (self.action_dim,) self._IO_descriptor.dtype = str(self.raw_actions.dtype) diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py index 30980fa1fb4..a1d1af6f18e 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py @@ -133,6 +133,20 @@ def processed_actions(self) -> torch.Tensor: @property def IO_descriptor(self) -> GenericActionIODescriptor: + """The IO descriptor of the action term. + + This descriptor is used to describe the action term of the pink inverse kinematics action. + It adds the following information to the base descriptor: + - scale: The scale of the action term. + - offset: The offset of the action term. + - clip: The clip of the action term. + - pink_controller_joint_names: The names of the pink controller joints. + - hand_joint_names: The names of the hand joints. + - controller_cfg: The configuration of the pink controller. + + Returns: + The IO descriptor of the action term. + """ super().IO_descriptor self._IO_descriptor.shape = (self.action_dim,) self._IO_descriptor.dtype = str(self.raw_actions.dtype) diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py index bde914ccae9..b3953319b34 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py @@ -151,6 +151,20 @@ def jacobian_b(self) -> torch.Tensor: @property def IO_descriptor(self) -> GenericActionIODescriptor: + """The IO descriptor of the action term. + + This descriptor is used to describe the action term of the pink inverse kinematics action. + It adds the following information to the base descriptor: + - body_name: The name of the body. + - joint_names: The names of the joints. + - scale: The scale of the action term. + - clip: The clip of the action term. + - controller_cfg: The configuration of the controller. + - body_offset: The offset of the body. + + Returns: + The IO descriptor of the action term. + """ super().IO_descriptor self._IO_descriptor.shape = (self.action_dim,) self._IO_descriptor.dtype = str(self.raw_actions.dtype) @@ -429,6 +443,25 @@ def jacobian_b(self) -> torch.Tensor: @property def IO_descriptor(self) -> GenericActionIODescriptor: + """The IO descriptor of the action term. + + This descriptor is used to describe the action term of the pink inverse kinematics action. + It adds the following information to the base descriptor: + - body_name: The name of the body. + - joint_names: The names of the joints. + - position_scale: The scale of the position. + - orientation_scale: The scale of the orientation. + - wrench_scale: The scale of the wrench. + - stiffness_scale: The scale of the stiffness. + - damping_ratio_scale: The scale of the damping ratio. + - nullspace_joint_pos_target: The nullspace joint pos target. + - clip: The clip of the action term. + - controller_cfg: The configuration of the controller. + - body_offset: The offset of the body. + + Returns: + The IO descriptor of the action term. + """ super().IO_descriptor self._IO_descriptor.shape = (self.action_dim,) self._IO_descriptor.dtype = str(self.raw_actions.dtype) diff --git a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py index 4a852b4ced9..f5829a525d6 100644 --- a/source/isaaclab/isaaclab/envs/utils/io_descriptors.py +++ b/source/isaaclab/isaaclab/envs/utils/io_descriptors.py @@ -22,19 +22,69 @@ @configclass class GenericActionIODescriptor: + """Generic action IO descriptor. + + This descriptor is used to describe the action space of a policy. + It can be extended as needed to add more information about the action term that is being described. + """ + mdp_type: str = "Action" + """The type of MDP that the action term belongs to.""" + name: str = None + """The name of the action term. + + By default, the name of the action term class is used. + """ + full_path: str = None + """The full path of the action term class. + + By default, python's will retrieve the path from the file that the action term class is defined in + and the name of the action term class. + """ + description: str = None + """The description of the action term. + + By default, the docstring of the action term class is used. + """ + shape: tuple[int, ...] = None + """The shape of the action term. + + This should be populated by the user.""" + dtype: str = None + """The dtype of the action term. + + This should be populated by the user.""" + action_type: str = None + """The type of the action term. + + This attribute is purely informative and should be populated by the user.""" + extras: dict[str, Any] = {} + """Extra information about the action term. + + This attribute is purely informative and should be populated by the user.""" + export: bool = True + """Whether to export the action term. + + Should be set to False if the class is not meant to be exported. + """ @configclass -class GenericIODescriptor: +class GenericObservationIODescriptor: + """Generic observation IO descriptor. + + This descriptor is used to describe the observation space of a policy. + It can be extended as needed to add more information about the observation term that is being described. + """ + mdp_type: str = "Observation" name: str = None full_path: str = None @@ -51,13 +101,13 @@ class GenericIODescriptor: # Automatically builds a descriptor from the kwargs -def _make_descriptor(**kwargs: Any) -> GenericIODescriptor: +def _make_descriptor(**kwargs: Any) -> GenericObservationIODescriptor: """Split *kwargs* into (known dataclass fields) and (extras).""" - field_names = {f.name for f in dataclasses.fields(GenericIODescriptor)} + field_names = {f.name for f in dataclasses.fields(GenericObservationIODescriptor)} known = {k: v for k, v in kwargs.items() if k in field_names} extras = {k: v for k, v in kwargs.items() if k not in field_names} - desc = GenericIODescriptor(**known) + desc = GenericObservationIODescriptor(**known) # User defined extras are stored in the descriptor under the `extras` field desc.extras = extras return desc @@ -121,8 +171,8 @@ def my_func(env: ManagerBasedEnv, *args, **kwargs): Returns: A decorator that can be used to decorate a function. """ - - if _func is not None and isinstance(_func, GenericIODescriptor): + # If the decorator is used with a descriptor, use it as the descriptor. + if _func is not None and isinstance(_func, GenericObservationIODescriptor): descriptor = _func _func = None else: @@ -178,15 +228,38 @@ def wrapper(env: ManagerBasedEnv, *args: P.args, **kwargs: P.kwargs) -> R: return _apply -def record_shape(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): +def record_shape(output: torch.Tensor, descriptor: GenericObservationIODescriptor, **kwargs) -> None: + """Record the shape of the output tensor. + + Args: + output: The output tensor. + descriptor: The descriptor to record the shape to. + **kwargs: Additional keyword arguments. + """ descriptor.shape = (output.shape[-1],) -def record_dtype(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): +def record_dtype(output: torch.Tensor, descriptor: GenericObservationIODescriptor, **kwargs) -> None: + """Record the dtype of the output tensor. + + Args: + output: The output tensor. + descriptor: The descriptor to record the dtype to. + **kwargs: Additional keyword arguments. + """ descriptor.dtype = str(output.dtype) -def record_joint_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): +def record_joint_names(output: torch.Tensor, descriptor: GenericObservationIODescriptor, **kwargs) -> None: + """Record the joint names of the output tensor. + + Expects the `asset_cfg` keyword argument to be set. + + Args: + output: The output tensor. + descriptor: The descriptor to record the joint names to. + **kwargs: Additional keyword arguments. + """ asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] joint_ids = kwargs["asset_cfg"].joint_ids if joint_ids == slice(None, None, None): @@ -194,7 +267,16 @@ def record_joint_names(output: torch.Tensor, descriptor: GenericIODescriptor, ** descriptor.joint_names = [asset.joint_names[i] for i in joint_ids] -def record_body_names(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): +def record_body_names(output: torch.Tensor, descriptor: GenericObservationIODescriptor, **kwargs) -> None: + """Record the body names of the output tensor. + + Expects the `asset_cfg` keyword argument to be set. + + Args: + output: The output tensor. + descriptor: The descriptor to record the body names to. + **kwargs: Additional keyword arguments. + """ asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] body_ids = kwargs["asset_cfg"].body_ids if body_ids == slice(None, None, None): @@ -202,7 +284,16 @@ def record_body_names(output: torch.Tensor, descriptor: GenericIODescriptor, **k descriptor.body_names = [asset.body_names[i] for i in body_ids] -def record_joint_pos_offsets(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): +def record_joint_pos_offsets(output: torch.Tensor, descriptor: GenericObservationIODescriptor, **kwargs): + """Record the joint position offsets of the output tensor. + + Expects the `asset_cfg` keyword argument to be set. + + Args: + output: The output tensor. + descriptor: The descriptor to record the joint position offsets to. + **kwargs: Additional keyword arguments. + """ asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] ids = kwargs["asset_cfg"].joint_ids # Get the offsets of the joints for the first robot in the scene. @@ -210,7 +301,16 @@ def record_joint_pos_offsets(output: torch.Tensor, descriptor: GenericIODescript descriptor.joint_pos_offsets = asset.data.default_joint_pos[:, ids][0] -def record_joint_vel_offsets(output: torch.Tensor, descriptor: GenericIODescriptor, **kwargs): +def record_joint_vel_offsets(output: torch.Tensor, descriptor: GenericObservationIODescriptor, **kwargs): + """Record the joint velocity offsets of the output tensor. + + Expects the `asset_cfg` keyword argument to be set. + + Args: + output: The output tensor. + descriptor: The descriptor to record the joint velocity offsets to. + **kwargs: Additional keyword arguments. + """ asset: Articulation = kwargs["env"].scene[kwargs["asset_cfg"].name] ids = kwargs["asset_cfg"].joint_ids # Get the offsets of the joints for the first robot in the scene. @@ -218,9 +318,20 @@ def record_joint_vel_offsets(output: torch.Tensor, descriptor: GenericIODescript descriptor.joint_vel_offsets = asset.data.default_joint_vel[:, ids][0] -def export_articulations_data(env: ManagerBasedEnv): +def export_articulations_data(env: ManagerBasedEnv) -> dict[str, dict[str, list[float]]]: + """Export the articulations data. + + Args: + env: The environment. + + Returns: + A dictionary containing the articulations data. + """ + # Create a dictionary for all the articulations in the scene. articulation_joint_data = {} for articulation_name, articulation in env.scene.articulations.items(): + # For each articulation, create a dictionary with the articulation's data. + # Some of the data may be redundant with other information provided by the observation descriptors. articulation_joint_data[articulation_name] = {} articulation_joint_data[articulation_name]["joint_names"] = articulation.joint_names articulation_joint_data[articulation_name]["default_joint_pos"] = ( @@ -247,7 +358,16 @@ def export_articulations_data(env: ManagerBasedEnv): return articulation_joint_data -def export_scene_data(env: ManagerBasedEnv): +def export_scene_data(env: ManagerBasedEnv) -> dict[str, Any]: + """Export the scene data. + + Args: + env: The environment. + + Returns: + A dictionary containing the scene data. + """ + # Create a dictionary for the scene data. scene_data = {} scene_data["physics_dt"] = env.physics_dt scene_data["dt"] = env.step_dt From 7e3b4c30d81275317b8eb86c545b4eb81d0466d3 Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Fri, 18 Jul 2025 16:56:06 +0200 Subject: [PATCH 23/23] started working on tutorial --- ...ocity_flat_anymal_d_v0_IO_descriptors.yaml | 344 +++++++++ ...ac_velocity_flat_g1_v0_IO_descriptors.yaml | 719 ++++++++++++++++++ .../01_io_descriptors/io_descriptors_101.rst | 62 ++ 3 files changed, 1125 insertions(+) create mode 100644 docs/source/_static/policy_deployment/01_io_descriptors/isaac_velocity_flat_anymal_d_v0_IO_descriptors.yaml create mode 100644 docs/source/_static/policy_deployment/01_io_descriptors/isaac_velocity_flat_g1_v0_IO_descriptors.yaml create mode 100644 docs/source/policy_deployment/01_io_descriptors/io_descriptors_101.rst diff --git a/docs/source/_static/policy_deployment/01_io_descriptors/isaac_velocity_flat_anymal_d_v0_IO_descriptors.yaml b/docs/source/_static/policy_deployment/01_io_descriptors/isaac_velocity_flat_anymal_d_v0_IO_descriptors.yaml new file mode 100644 index 00000000000..c9573ef6e45 --- /dev/null +++ b/docs/source/_static/policy_deployment/01_io_descriptors/isaac_velocity_flat_anymal_d_v0_IO_descriptors.yaml @@ -0,0 +1,344 @@ +actions: +- action_type: JointAction + clip: null + dtype: torch.float32 + extras: + description: Joint action term that applies the processed actions to the articulation's + joints as position commands. + full_path: isaaclab.envs.mdp.actions.joint_actions.JointPositionAction + joint_names: + - LF_HAA + - LH_HAA + - RF_HAA + - RH_HAA + - LF_HFE + - LH_HFE + - RF_HFE + - RH_HFE + - LF_KFE + - LH_KFE + - RF_KFE + - RH_KFE + mdp_type: Action + name: joint_position_action + offset: + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.4000000059604645 + - -0.4000000059604645 + - 0.4000000059604645 + - -0.4000000059604645 + - -0.800000011920929 + - 0.800000011920929 + - -0.800000011920929 + - 0.800000011920929 + scale: 0.5 + shape: + - 12 +articulations: + robot: + default_joint_armature: + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + default_joint_damping: + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + default_joint_friction: + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + default_joint_pos: + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.4000000059604645 + - -0.4000000059604645 + - 0.4000000059604645 + - -0.4000000059604645 + - -0.800000011920929 + - 0.800000011920929 + - -0.800000011920929 + - 0.800000011920929 + default_joint_pos_limits: + - - -0.7853984236717224 + - 0.6108654141426086 + - - -0.7853984236717224 + - 0.6108654141426086 + - - -0.6108654141426086 + - 0.7853984236717224 + - - -0.6108654141426086 + - 0.7853984236717224 + - - -9.42477798461914 + - 9.42477798461914 + - - -9.42477798461914 + - 9.42477798461914 + - - -9.42477798461914 + - 9.42477798461914 + - - -9.42477798461914 + - 9.42477798461914 + - - -9.42477798461914 + - 9.42477798461914 + - - -9.42477798461914 + - 9.42477798461914 + - - -9.42477798461914 + - 9.42477798461914 + - - -9.42477798461914 + - 9.42477798461914 + default_joint_stiffness: + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + default_joint_vel: + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + joint_names: + - LF_HAA + - LH_HAA + - RF_HAA + - RH_HAA + - LF_HFE + - LH_HFE + - RF_HFE + - RH_HFE + - LF_KFE + - LH_KFE + - RF_KFE + - RH_KFE +observations: + policy: + - dtype: torch.float32 + extras: + axes: + - X + - Y + - Z + description: Root linear velocity in the asset's root frame. + modifiers: null + units: m/s + full_path: isaaclab.envs.mdp.observations.base_lin_vel + mdp_type: Observation + name: base_lin_vel + observation_type: RootState + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 3 + - dtype: torch.float32 + extras: + axes: + - X + - Y + - Z + description: Root angular velocity in the asset's root frame. + modifiers: null + units: rad/s + full_path: isaaclab.envs.mdp.observations.base_ang_vel + mdp_type: Observation + name: base_ang_vel + observation_type: RootState + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 3 + - dtype: torch.float32 + extras: + axes: + - X + - Y + - Z + description: Gravity projection on the asset's root frame. + modifiers: null + units: m/s^2 + full_path: isaaclab.envs.mdp.observations.projected_gravity + mdp_type: Observation + name: projected_gravity + observation_type: RootState + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 3 + - dtype: torch.float32 + extras: + description: The generated command from command term in the command manager + with the given name. + modifiers: null + full_path: isaaclab.envs.mdp.observations.generated_commands + mdp_type: Observation + name: generated_commands + observation_type: Command + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 3 + - dtype: torch.float32 + extras: + description: '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.' + modifiers: null + units: rad + full_path: isaaclab.envs.mdp.observations.joint_pos_rel + joint_names: + - LF_HAA + - LH_HAA + - RF_HAA + - RH_HAA + - LF_HFE + - LH_HFE + - RF_HFE + - RH_HFE + - LF_KFE + - LH_KFE + - RF_KFE + - RH_KFE + joint_pos_offsets: + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.4000000059604645 + - -0.4000000059604645 + - 0.4000000059604645 + - -0.4000000059604645 + - -0.800000011920929 + - 0.800000011920929 + - -0.800000011920929 + - 0.800000011920929 + mdp_type: Observation + name: joint_pos_rel + observation_type: JointState + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 12 + - dtype: torch.float32 + extras: + description: '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.' + modifiers: null + units: rad/s + full_path: isaaclab.envs.mdp.observations.joint_vel_rel + joint_names: + - LF_HAA + - LH_HAA + - RF_HAA + - RH_HAA + - LF_HFE + - LH_HFE + - RF_HFE + - RH_HFE + - LF_KFE + - LH_KFE + - RF_KFE + - RH_KFE + joint_vel_offsets: + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + mdp_type: Observation + name: joint_vel_rel + observation_type: JointState + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 12 + - dtype: torch.float32 + extras: + description: The last input action to the environment. The name of the action + term for which the action is required. If None, the entire action tensor is + returned. + modifiers: null + full_path: isaaclab.envs.mdp.observations.last_action + mdp_type: Observation + name: last_action + observation_type: Action + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 12 +scene: + decimation: 4 + dt: 0.02 + physics_dt: 0.005 diff --git a/docs/source/_static/policy_deployment/01_io_descriptors/isaac_velocity_flat_g1_v0_IO_descriptors.yaml b/docs/source/_static/policy_deployment/01_io_descriptors/isaac_velocity_flat_g1_v0_IO_descriptors.yaml new file mode 100644 index 00000000000..98baf0fdc31 --- /dev/null +++ b/docs/source/_static/policy_deployment/01_io_descriptors/isaac_velocity_flat_g1_v0_IO_descriptors.yaml @@ -0,0 +1,719 @@ +actions: +- action_type: JointAction + clip: null + dtype: torch.float32 + extras: + description: Joint action term that applies the processed actions to the articulation's + joints as position commands. + full_path: isaaclab.envs.mdp.actions.joint_actions.JointPositionAction + joint_names: + - left_hip_pitch_joint + - right_hip_pitch_joint + - torso_joint + - left_hip_roll_joint + - right_hip_roll_joint + - left_shoulder_pitch_joint + - right_shoulder_pitch_joint + - left_hip_yaw_joint + - right_hip_yaw_joint + - left_shoulder_roll_joint + - right_shoulder_roll_joint + - left_knee_joint + - right_knee_joint + - left_shoulder_yaw_joint + - right_shoulder_yaw_joint + - left_ankle_pitch_joint + - right_ankle_pitch_joint + - left_elbow_pitch_joint + - right_elbow_pitch_joint + - left_ankle_roll_joint + - right_ankle_roll_joint + - left_elbow_roll_joint + - right_elbow_roll_joint + - left_five_joint + - left_three_joint + - left_zero_joint + - right_five_joint + - right_three_joint + - right_zero_joint + - left_six_joint + - left_four_joint + - left_one_joint + - right_six_joint + - right_four_joint + - right_one_joint + - left_two_joint + - right_two_joint + mdp_type: Action + name: joint_position_action + offset: + - -0.20000000298023224 + - -0.20000000298023224 + - 0.0 + - 0.0 + - 0.0 + - 0.3499999940395355 + - 0.3499999940395355 + - 0.0 + - 0.0 + - 0.1599999964237213 + - -0.1599999964237213 + - 0.41999998688697815 + - 0.41999998688697815 + - 0.0 + - 0.0 + - -0.23000000417232513 + - -0.23000000417232513 + - 0.8700000047683716 + - 0.8700000047683716 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 1.0 + - 0.0 + - 0.0 + - -1.0 + - 0.5199999809265137 + - -0.5199999809265137 + scale: 0.5 + shape: + - 37 +articulations: + robot: + default_joint_armature: + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.009999999776482582 + - 0.0010000000474974513 + - 0.0010000000474974513 + - 0.0010000000474974513 + - 0.0010000000474974513 + - 0.0010000000474974513 + - 0.0010000000474974513 + - 0.0010000000474974513 + - 0.0010000000474974513 + - 0.0010000000474974513 + - 0.0010000000474974513 + - 0.0010000000474974513 + - 0.0010000000474974513 + - 0.0010000000474974513 + - 0.0010000000474974513 + default_joint_damping: + - 5.0 + - 5.0 + - 5.0 + - 5.0 + - 5.0 + - 10.0 + - 10.0 + - 5.0 + - 5.0 + - 10.0 + - 10.0 + - 5.0 + - 5.0 + - 10.0 + - 10.0 + - 2.0 + - 2.0 + - 10.0 + - 10.0 + - 2.0 + - 2.0 + - 10.0 + - 10.0 + - 10.0 + - 10.0 + - 10.0 + - 10.0 + - 10.0 + - 10.0 + - 10.0 + - 10.0 + - 10.0 + - 10.0 + - 10.0 + - 10.0 + - 10.0 + - 10.0 + default_joint_friction: + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + default_joint_pos: + - -0.20000000298023224 + - -0.20000000298023224 + - 0.0 + - 0.0 + - 0.0 + - 0.3499999940395355 + - 0.3499999940395355 + - 0.0 + - 0.0 + - 0.1599999964237213 + - -0.1599999964237213 + - 0.41999998688697815 + - 0.41999998688697815 + - 0.0 + - 0.0 + - -0.23000000417232513 + - -0.23000000417232513 + - 0.8700000047683716 + - 0.8700000047683716 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 1.0 + - 0.0 + - 0.0 + - -1.0 + - 0.5199999809265137 + - -0.5199999809265137 + default_joint_pos_limits: + - - -2.3499996662139893 + - 3.049999952316284 + - - -2.3499996662139893 + - 3.049999952316284 + - - -2.618000030517578 + - 2.618000030517578 + - - -0.25999996066093445 + - 2.5299997329711914 + - - -2.5299997329711914 + - 0.25999996066093445 + - - -2.967099666595459 + - 2.7924997806549072 + - - -2.967099666595459 + - 2.7924997806549072 + - - -2.749999761581421 + - 2.749999761581421 + - - -2.749999761581421 + - 2.749999761581421 + - - -1.5881999731063843 + - 2.251499652862549 + - - -2.251499652862549 + - 1.5881999731063843 + - - -0.3348899781703949 + - 2.5448997020721436 + - - -0.3348899781703949 + - 2.5448997020721436 + - - -2.618000030517578 + - 2.618000030517578 + - - -2.618000030517578 + - 2.618000030517578 + - - -0.6799999475479126 + - 0.7299999594688416 + - - -0.6799999475479126 + - 0.7299999594688416 + - - -0.22679997980594635 + - 3.420799732208252 + - - -0.22679997980594635 + - 3.420799732208252 + - - -0.26179996132850647 + - 0.26179996132850647 + - - -0.26179996132850647 + - 0.26179996132850647 + - - -2.094299793243408 + - 2.094299793243408 + - - -2.094299793243408 + - 2.094299793243408 + - - -1.8399999141693115 + - 0.30000001192092896 + - - -1.8399999141693115 + - 0.30000001192092896 + - - -0.5235979557037354 + - 0.5235979557037354 + - - -0.30000001192092896 + - 1.8399999141693115 + - - -0.30000001192092896 + - 1.8399999141693115 + - - -0.5235979557037354 + - 0.5235979557037354 + - - -1.8399999141693115 + - 0.0 + - - -1.8399999141693115 + - 0.0 + - - -0.9999999403953552 + - 1.2000000476837158 + - - 0.0 + - 1.8399999141693115 + - - 0.0 + - 1.8399999141693115 + - - -1.2000000476837158 + - 0.9999999403953552 + - - 0.0 + - 1.8399999141693115 + - - -1.8399999141693115 + - 0.0 + default_joint_stiffness: + - 200.0 + - 200.0 + - 200.0 + - 150.0 + - 150.0 + - 40.0 + - 40.0 + - 150.0 + - 150.0 + - 40.0 + - 40.0 + - 200.0 + - 200.0 + - 40.0 + - 40.0 + - 20.0 + - 20.0 + - 40.0 + - 40.0 + - 20.0 + - 20.0 + - 40.0 + - 40.0 + - 40.0 + - 40.0 + - 40.0 + - 40.0 + - 40.0 + - 40.0 + - 40.0 + - 40.0 + - 40.0 + - 40.0 + - 40.0 + - 40.0 + - 40.0 + - 40.0 + default_joint_vel: + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + joint_names: + - left_hip_pitch_joint + - right_hip_pitch_joint + - torso_joint + - left_hip_roll_joint + - right_hip_roll_joint + - left_shoulder_pitch_joint + - right_shoulder_pitch_joint + - left_hip_yaw_joint + - right_hip_yaw_joint + - left_shoulder_roll_joint + - right_shoulder_roll_joint + - left_knee_joint + - right_knee_joint + - left_shoulder_yaw_joint + - right_shoulder_yaw_joint + - left_ankle_pitch_joint + - right_ankle_pitch_joint + - left_elbow_pitch_joint + - right_elbow_pitch_joint + - left_ankle_roll_joint + - right_ankle_roll_joint + - left_elbow_roll_joint + - right_elbow_roll_joint + - left_five_joint + - left_three_joint + - left_zero_joint + - right_five_joint + - right_three_joint + - right_zero_joint + - left_six_joint + - left_four_joint + - left_one_joint + - right_six_joint + - right_four_joint + - right_one_joint + - left_two_joint + - right_two_joint +observations: + policy: + - dtype: torch.float32 + extras: + axes: + - X + - Y + - Z + description: Root linear velocity in the asset's root frame. + modifiers: null + units: m/s + full_path: isaaclab.envs.mdp.observations.base_lin_vel + mdp_type: Observation + name: base_lin_vel + observation_type: RootState + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 3 + - dtype: torch.float32 + extras: + axes: + - X + - Y + - Z + description: Root angular velocity in the asset's root frame. + modifiers: null + units: rad/s + full_path: isaaclab.envs.mdp.observations.base_ang_vel + mdp_type: Observation + name: base_ang_vel + observation_type: RootState + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 3 + - dtype: torch.float32 + extras: + axes: + - X + - Y + - Z + description: Gravity projection on the asset's root frame. + modifiers: null + units: m/s^2 + full_path: isaaclab.envs.mdp.observations.projected_gravity + mdp_type: Observation + name: projected_gravity + observation_type: RootState + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 3 + - dtype: torch.float32 + extras: + description: The generated command from command term in the command manager + with the given name. + modifiers: null + full_path: isaaclab.envs.mdp.observations.generated_commands + mdp_type: Observation + name: generated_commands + observation_type: Command + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 3 + - dtype: torch.float32 + extras: + description: '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.' + modifiers: null + units: rad + full_path: isaaclab.envs.mdp.observations.joint_pos_rel + joint_names: + - left_hip_pitch_joint + - right_hip_pitch_joint + - torso_joint + - left_hip_roll_joint + - right_hip_roll_joint + - left_shoulder_pitch_joint + - right_shoulder_pitch_joint + - left_hip_yaw_joint + - right_hip_yaw_joint + - left_shoulder_roll_joint + - right_shoulder_roll_joint + - left_knee_joint + - right_knee_joint + - left_shoulder_yaw_joint + - right_shoulder_yaw_joint + - left_ankle_pitch_joint + - right_ankle_pitch_joint + - left_elbow_pitch_joint + - right_elbow_pitch_joint + - left_ankle_roll_joint + - right_ankle_roll_joint + - left_elbow_roll_joint + - right_elbow_roll_joint + - left_five_joint + - left_three_joint + - left_zero_joint + - right_five_joint + - right_three_joint + - right_zero_joint + - left_six_joint + - left_four_joint + - left_one_joint + - right_six_joint + - right_four_joint + - right_one_joint + - left_two_joint + - right_two_joint + joint_pos_offsets: + - -0.20000000298023224 + - -0.20000000298023224 + - 0.0 + - 0.0 + - 0.0 + - 0.3499999940395355 + - 0.3499999940395355 + - 0.0 + - 0.0 + - 0.1599999964237213 + - -0.1599999964237213 + - 0.41999998688697815 + - 0.41999998688697815 + - 0.0 + - 0.0 + - -0.23000000417232513 + - -0.23000000417232513 + - 0.8700000047683716 + - 0.8700000047683716 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 1.0 + - 0.0 + - 0.0 + - -1.0 + - 0.5199999809265137 + - -0.5199999809265137 + mdp_type: Observation + name: joint_pos_rel + observation_type: JointState + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 37 + - dtype: torch.float32 + extras: + description: '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.' + modifiers: null + units: rad/s + full_path: isaaclab.envs.mdp.observations.joint_vel_rel + joint_names: + - left_hip_pitch_joint + - right_hip_pitch_joint + - torso_joint + - left_hip_roll_joint + - right_hip_roll_joint + - left_shoulder_pitch_joint + - right_shoulder_pitch_joint + - left_hip_yaw_joint + - right_hip_yaw_joint + - left_shoulder_roll_joint + - right_shoulder_roll_joint + - left_knee_joint + - right_knee_joint + - left_shoulder_yaw_joint + - right_shoulder_yaw_joint + - left_ankle_pitch_joint + - right_ankle_pitch_joint + - left_elbow_pitch_joint + - right_elbow_pitch_joint + - left_ankle_roll_joint + - right_ankle_roll_joint + - left_elbow_roll_joint + - right_elbow_roll_joint + - left_five_joint + - left_three_joint + - left_zero_joint + - right_five_joint + - right_three_joint + - right_zero_joint + - left_six_joint + - left_four_joint + - left_one_joint + - right_six_joint + - right_four_joint + - right_one_joint + - left_two_joint + - right_two_joint + joint_vel_offsets: + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + - 0.0 + mdp_type: Observation + name: joint_vel_rel + observation_type: JointState + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 37 + - dtype: torch.float32 + extras: + description: The last input action to the environment. The name of the action + term for which the action is required. If None, the entire action tensor is + returned. + modifiers: null + full_path: isaaclab.envs.mdp.observations.last_action + mdp_type: Observation + name: last_action + observation_type: Action + overloads: + clip: null + flatten_history_dim: true + history_length: 0 + scale: null + shape: + - 37 +scene: + decimation: 4 + dt: 0.02 + physics_dt: 0.005 diff --git a/docs/source/policy_deployment/01_io_descriptors/io_descriptors_101.rst b/docs/source/policy_deployment/01_io_descriptors/io_descriptors_101.rst new file mode 100644 index 00000000000..0a367c51c89 --- /dev/null +++ b/docs/source/policy_deployment/01_io_descriptors/io_descriptors_101.rst @@ -0,0 +1,62 @@ +IO Descriptors 101 +================== + +In this tutorial, we will learn about IO descriptors, what they are, how to export them, and how to add them to your environment. +We will use the Anymal-D robot as an example to demonstrate how to export IO descriptors an environment, and the Spot robot to +explain how to attach IO descriptors to custom action and observation terms. + + +What are IO Descriptors? +------------------------ + +Before we dive into IO descriptors, let's first understand what they are and why they are useful. + +IO descriptors are a way to describe the input and output of a policy trained using the ManagerBasedRLEnv in Isaac Lab. In other words, +they describe the action and observation terms of a policy. This description is used to generate a YAML file that can be loaded in an +external tool to run the policies without having to manually input the configuration of the action and observation terms. + +In addition to this the IO Descriptors provide the following information: +- The parameters of all the joints in the articulation. +- Some simulation parameters including the simulation time step, and the policy time step. +- For some action and observation terms, it provides the joint names or body names in the same order as they appear in the action terms. +- For both the observation and action terms, it provides the terms in the exact same order as they appear in the managers. Making it easy to + reconstruct them from the YAML file. + +Here is an example of what the action part of the YAML generated from the IO descriptors looks like for the Anymal-D robot: +.. literalinclude:: ../../_static/policy_deployment/01_io_descriptors/isaac_velocity_flat_anymal_d_v0_IO_descriptors.yaml + :language: yaml + :lines: 1-39 + +Here is an example of what a portion of the observation part of the YAML generated from the IO descriptors looks like for the Anymal-D robot: +.. literalinclude:: ../../_static/policy_deployment/01_io_descriptors/isaac_velocity_flat_anymal_d_v0_IO_descriptors.yaml + :language: yaml + :lines: 158-199 + +.. literalinclude:: ../../_static/policy_deployment/01_io_descriptors/isaac_velocity_flat_anymal_d_v0_IO_descriptors.yaml + :language: yaml + :lines: 236-279 + +Something to note here is that both the action and observation terms are returned as list of dictionaries, and not a dictionary of dictionaries. +This is done to ensure the order of the terms is preserved. Hence, to retrive the action or observation term, the users need to look for the +`name` key in the dictionaries. + +For example, in the following snippet, we are looking at the `projected_gravity` observation term. The `name` key is used to identify the term. +The `full_path` key is used to provide an explicit path to the function in Isaac Lab's source code. Some flags like `mdp_type` and `observation_type` +are also provided, these don't have any functional impact. They are here to inform the user that this is the category this term belongs to. + +.. literalinclude:: ../../_static/policy_deployment/01_io_descriptors/isaac_velocity_flat_anymal_d_v0_IO_descriptors.yaml + :language: yaml + :lines: 200-219 + :emphasize-lines: 211,209 + + + + +Exporting IO Descriptors from an Environment +-------------------------------------------- + + + + +IO descriptors are a way to describe the input and output of a policy. +They are used to describe the action and observation spaces of a policy.icy. \ No newline at end of file