diff --git a/scripts/reinforcement_learning/leapp/rsl_rl/export.py b/scripts/reinforcement_learning/leapp/rsl_rl/export.py index c32f08b9f835..b5f2c4387d3c 100644 --- a/scripts/reinforcement_learning/leapp/rsl_rl/export.py +++ b/scripts/reinforcement_learning/leapp/rsl_rl/export.py @@ -18,7 +18,10 @@ from isaaclab.app import AppLauncher -from isaaclab_tasks.utils import setup_preset_cli +try: + from isaaclab_tasks.utils import setup_preset_cli +except ImportError: + setup_preset_cli = None _RSL_RL_SCRIPTS_DIR = Path(__file__).resolve().parents[2] / "rsl_rl" if str(_RSL_RL_SCRIPTS_DIR) not in sys.path: @@ -110,10 +113,13 @@ def create_arg_parser() -> argparse.ArgumentParser: def parse_export_args(argv: list[str] | None = None) -> tuple[argparse.Namespace, list[str]]: """Parse export arguments and return remaining Hydra overrides.""" parser = create_arg_parser() - # setup_preset_cli attaches the preset-selection help group then parses; - # remainder still carries typed selectors (physics=/renderer=/presets=) - # verbatim for run_export_with_hydra to fold before invoking Hydra. - args_cli, hydra_args = setup_preset_cli(parser, argv) + if setup_preset_cli is not None: + # setup_preset_cli attaches the preset-selection help group then parses; + # remainder still carries typed selectors (physics=/renderer=/presets=) + # verbatim for run_export_with_hydra to fold before invoking Hydra. + args_cli, hydra_args = setup_preset_cli(parser, argv) + else: + args_cli, hydra_args = parser.parse_known_args(argv) args_cli.headless = True return args_cli, hydra_args @@ -255,6 +261,19 @@ def actor_hidden_from_registered(registered_state, original_hidden): return registered_state +def resolve_env_action_class_types(env_cfg) -> None: + """Resolve Hydra-overridden action ``class_type`` strings before constructing the env.""" + from isaaclab.utils.string import string_to_callable + + actions_cfg = getattr(env_cfg, "actions", None) + if actions_cfg is None: + return + for action_cfg in vars(actions_cfg).values(): + class_type = getattr(action_cfg, "class_type", None) + if isinstance(class_type, str) and "{DIR}" not in class_type: + action_cfg.class_type = string_to_callable(class_type) + + def export_rsl_rl_agent( args_cli: argparse.Namespace, env_cfg, @@ -301,6 +320,7 @@ def export_rsl_rl_agent( leapp_started = False try: + resolve_env_action_class_types(env_cfg) env = gym.make(args_cli.task, cfg=env_cfg, render_mode=None) policy_node_name = ensure_env_spec_id(env) @@ -389,7 +409,10 @@ def export_rsl_rl_agent( def run_export_with_hydra(args_cli: argparse.Namespace, hydra_args: list[str]) -> bool: """Resolve Hydra task configuration and export one RSL-RL policy.""" - from isaaclab.app import launch_simulation + try: + from isaaclab.app import launch_simulation + except ImportError: + from isaaclab_tasks.utils import launch_simulation from isaaclab_tasks.utils.hydra import hydra_task_config diff --git a/source/isaaclab/isaaclab/utils/leapp/export_annotator.py b/source/isaaclab/isaaclab/utils/leapp/export_annotator.py index 3ca82a66c082..b2779742df04 100644 --- a/source/isaaclab/isaaclab/utils/leapp/export_annotator.py +++ b/source/isaaclab/isaaclab/utils/leapp/export_annotator.py @@ -39,6 +39,7 @@ from typing import TYPE_CHECKING, Any import torch +import warp as wp from leapp import annotate from leapp.utils.tensor_description import TensorSemantics @@ -60,6 +61,11 @@ VARIABLE_IMPEDANCE_MODES = frozenset({"variable", "variable_kp"}) +def _tensor_data_to_torch(data): + """Return a torch tensor view for Isaac Lab data stored as torch or Warp-backed data.""" + return data.torch if hasattr(data, "torch") else wp.to_torch(data) + + # ══════════════════════════════════════════════════════════════════ # ExportPatcher # ══════════════════════════════════════════════════════════════════ @@ -304,6 +310,8 @@ def _patch_observation_manager(self, obs_manager, proxy_env): def patched_compute(*args, **kwargs): """Clear the tensor dedup cache once per full observation pass.""" cache.clear() + with suppress(Exception): + setattr(obs_manager._env, "_leapp_input_annotation_cache", {}) return original_compute(*args, **kwargs) obs_manager.compute = patched_compute @@ -666,7 +674,7 @@ def _collect_action_static_outputs( joint_names = getattr(real_asset, "joint_names", None) scene_key = self._action_term_scene_keys.get(term_name, "ego") if hasattr(data, "default_joint_stiffness") and data.default_joint_stiffness is not None: - gains = data.default_joint_stiffness.torch + gains = _tensor_data_to_torch(data.default_joint_stiffness) static_values.append( TensorSemantics( name=f"{term_name}_kp_gains", @@ -677,7 +685,7 @@ def _collect_action_static_outputs( ) ) if hasattr(data, "default_joint_damping") and data.default_joint_damping is not None: - gains = data.default_joint_damping.torch + gains = _tensor_data_to_torch(data.default_joint_damping) static_values.append( TensorSemantics( name=f"{term_name}_kd_gains", diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/ros_inference_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/ros_inference_env_cfg.py index ccbf550d9199..65973168d228 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/ros_inference_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/ros_inference_env_cfg.py @@ -11,6 +11,8 @@ from isaaclab.managers import ObservationTermCfg as ObsTerm from isaaclab.utils.configclass import configclass +import isaaclab_tasks.contrib.deploy.mdp as mdp + from .joint_pos_env_cfg import Rizon4sGearAssemblyEnvCfg @@ -50,6 +52,12 @@ def __post_init__(self): # Set joint_action_scale from the existing arm_action.scale self.joint_action_scale = self.actions.arm_action.scale + self.actions.arm_action = mdp.DeployRelativeJointPositionActionCfg( + asset_name="robot", + joint_names=self.arm_joint_names, + scale=self.joint_action_scale, + use_zero_offset=True, + ) # Dynamically generate action_scale_joint_space based on action_space self.action_scale_joint_space = [self.joint_action_scale] * self.action_space diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/__init__.pyi index 6178634ef3cd..cb912a3ba53b 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/__init__.pyi @@ -4,6 +4,7 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "DeployRelativeJointPositionActionCfg", "randomize_gear_type", "randomize_gears_and_base_pose", "pin_unselected_gears_to_shafts", @@ -14,6 +15,8 @@ __all__ = [ "gear_quat_w", "gear_shaft_pos_w", "gear_shaft_quat_w", + "joint_pos", + "joint_vel", "keypoint_command_error", "keypoint_command_error_exp", "keypoint_entity_error", @@ -30,8 +33,9 @@ from .events import ( randomize_gears_and_base_pose, set_robot_to_grasp_pose, ) +from .actions_cfg import DeployRelativeJointPositionActionCfg from .noise_models import ResetSampledConstantNoiseModel, ResetSampledConstantNoiseModelCfg -from .observations import gear_pos_w, gear_quat_w, gear_shaft_pos_w, gear_shaft_quat_w +from .observations import gear_pos_w, gear_quat_w, gear_shaft_pos_w, gear_shaft_quat_w, joint_pos, joint_vel from .rewards import ( keypoint_command_error, keypoint_command_error_exp, diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions.py new file mode 100644 index 000000000000..59b3068bb5cd --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions.py @@ -0,0 +1,124 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Deploy-specific action terms for LEAPP export workflows.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import warp as wp + +from isaaclab.envs.mdp.actions.joint_actions import RelativeJointPositionAction + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from .actions_cfg import DeployRelativeJointPositionActionCfg + + +_LEAPP_TRACED_OBSERVATION_INPUTS = "_leapp_traced_observation_inputs" +_LEAPP_CONSUMED_OBSERVATION_INPUTS = "_leapp_consumed_observation_inputs" + + +def _leapp_real_env(env): + """Return the wrapped Isaac Lab env when LEAPP passes an export proxy.""" + return object.__getattribute__(env, "_real_env") if type(env).__name__ == "_EnvProxy" else env + + +def _tensor_data_to_torch(data): + """Return a torch tensor view for Isaac Lab data stored as torch or Warp-backed data.""" + return data.torch if hasattr(data, "torch") else wp.to_torch(data) + + +def _get_observation_term_from_buffer(env, group_name: str, term_name: str): + """Return a term slice from the cached observation buffer.""" + obs_buffer = getattr(env, "obs_buf", None) + if obs_buffer is None: + obs_buffer = getattr(getattr(env, "observation_manager", None), "_obs_buffer", None) + if not obs_buffer or group_name not in obs_buffer: + return None + + group_obs = obs_buffer[group_name] + if isinstance(group_obs, dict): + return group_obs.get(term_name) + + obs_manager = getattr(env, "observation_manager", None) + if obs_manager is None: + return None + + term_names = obs_manager.active_terms.get(group_name, []) + if term_name not in term_names: + return None + + term_index = term_names.index(term_name) + term_dims = obs_manager.group_obs_term_dim[group_name] + concat_dim = obs_manager._group_obs_concatenate_dim[group_name] + if concat_dim > 0: + concat_dim -= 1 + + start = sum(dim[concat_dim] for dim in term_dims[:term_index]) + length = term_dims[term_index][concat_dim] + return group_obs.narrow(dim=concat_dim, start=start, length=length) + + +def _pop_leapp_traced_observation_input(env, name: str, *, group_name: str, term_name: str): + """Consume one traced observation tensor for the current LEAPP action trace.""" + real_env = _leapp_real_env(env) + consumed_inputs = getattr(real_env, _LEAPP_CONSUMED_OBSERVATION_INPUTS, None) + if consumed_inputs is None: + consumed_inputs = set() + setattr(real_env, _LEAPP_CONSUMED_OBSERVATION_INPUTS, consumed_inputs) + + if name in consumed_inputs: + return None + + traced_inputs = getattr(real_env, _LEAPP_TRACED_OBSERVATION_INPUTS, {}) + traced_tensor = traced_inputs.pop(name, None) + if traced_tensor is None: + traced_tensor = _get_observation_term_from_buffer(real_env, group_name, term_name) + + if traced_tensor is not None: + consumed_inputs.add(name) + return traced_tensor + + +def _is_leapp_observation_input_consumed(env, name: str) -> bool: + """Return whether an export input was already consumed by this action trace.""" + real_env = _leapp_real_env(env) + return name in getattr(real_env, _LEAPP_CONSUMED_OBSERVATION_INPUTS, set()) + + +class DeployRelativeJointPositionAction(RelativeJointPositionAction): + """Relative joint action that exports absolute joint-position targets for deployment.""" + + cfg: DeployRelativeJointPositionActionCfg + + def __init__(self, cfg: DeployRelativeJointPositionActionCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + def apply_actions(self): + asset = self._asset + if type(asset).__name__ == "_ArticulationWriteProxy": + observation_input_name = f"{self.cfg.asset_name}_joint_pos" + current_joint_pos = _pop_leapp_traced_observation_input( + self._env, + observation_input_name, + group_name="policy", + term_name="joint_pos", + ) + if current_joint_pos is None: + if not _is_leapp_observation_input_consumed(self._env, observation_input_name): + raise RuntimeError( + "DeployRelativeJointPositionAction requires the traced " + f"'{self.cfg.asset_name}_joint_pos' observation during LEAPP export." + ) + real_asset = object.__getattribute__(asset, "_real_asset") + current_joint_pos = _tensor_data_to_torch(real_asset.data.joint_pos)[:, self._joint_ids] + else: + current_joint_pos = _tensor_data_to_torch(asset.data.joint_pos)[:, self._joint_ids] + + current_actions = self.processed_actions + current_joint_pos + self._asset.set_joint_position_target_index(target=current_actions, joint_ids=self._joint_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions_cfg.py new file mode 100644 index 000000000000..970dbe45b65b --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions_cfg.py @@ -0,0 +1,18 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Deploy-specific action configuration classes.""" + +from __future__ import annotations + +from isaaclab.envs.mdp.actions.actions_cfg import RelativeJointPositionActionCfg +from isaaclab.utils.configclass import configclass + + +@configclass +class DeployRelativeJointPositionActionCfg(RelativeJointPositionActionCfg): + """Configuration for deploy relative joint actions with explicit LEAPP current-joint input.""" + + class_type: type | str = "{DIR}.actions:DeployRelativeJointPositionAction" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/observations.py index ac12d8b22f7f..212e699323bc 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/observations.py @@ -10,8 +10,10 @@ from typing import TYPE_CHECKING import torch +import warp as wp from isaaclab.managers import ManagerTermBase, ObservationTermCfg, SceneEntityCfg +from isaaclab.utils.leapp import QUAT_XYZW_ELEMENT_NAMES, XYZ_ELEMENT_NAMES, InputKindEnum from isaaclab.utils.math import combine_frame_transforms if TYPE_CHECKING: @@ -21,6 +23,101 @@ from .events import randomize_gear_type +_LEAPP_TRACED_OBSERVATION_INPUTS = "_leapp_traced_observation_inputs" +_LEAPP_CONSUMED_OBSERVATION_INPUTS = "_leapp_consumed_observation_inputs" + + +def _tensor_data_to_torch(data) -> torch.Tensor: + """Return a torch tensor view for Isaac Lab data stored as torch or Warp-backed data.""" + return data.torch if hasattr(data, "torch") else wp.to_torch(data) + + +def _selected_joint_names(asset, joint_ids) -> list[str] | None: + """Return joint names selected by the observation config.""" + joint_names = getattr(asset, "joint_names", None) + if joint_names is None: + return None + if joint_ids is None or joint_ids == slice(None): + return list(joint_names) + if isinstance(joint_ids, slice): + return list(joint_names[joint_ids]) + if hasattr(joint_ids, "tolist"): + joint_ids = joint_ids.tolist() + return [joint_names[int(joint_id)] for joint_id in joint_ids] + + +def _is_leapp_export_env(env) -> bool: + """Return whether the observation is running under the LEAPP export proxy.""" + return type(env).__name__ == "_EnvProxy" + + +def _leapp_real_env(env): + """Return the wrapped Isaac Lab env when LEAPP passes an export proxy.""" + if _is_leapp_export_env(env): + return object.__getattribute__(env, "_real_env") + return env + + +def _set_leapp_traced_observation_input(env, name: str, tensor: torch.Tensor) -> None: + """Store a traced observation tensor for later export-only reuse.""" + if not _is_leapp_export_env(env): + return + real_env = _leapp_real_env(env) + traced_inputs = getattr(real_env, _LEAPP_TRACED_OBSERVATION_INPUTS, None) + if traced_inputs is None: + traced_inputs = {} + setattr(real_env, _LEAPP_TRACED_OBSERVATION_INPUTS, traced_inputs) + traced_inputs[name] = tensor + getattr(real_env, _LEAPP_CONSUMED_OBSERVATION_INPUTS, set()).discard(name) + + +def joint_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Joint positions for the configured joints, exposed as the LEAPP input boundary.""" + real_env = _leapp_real_env(env) + asset = real_env.scene[asset_cfg.name] + selected_joint_pos = _tensor_data_to_torch(asset.data.joint_pos)[:, asset_cfg.joint_ids] + joint_names = _selected_joint_names(asset, asset_cfg.joint_ids) + if _is_leapp_export_env(env): + from leapp import annotate + from leapp.utils.tensor_description import TensorSemantics + + selected_joint_pos = annotate.input_tensors( + env.unwrapped.spec.id, + TensorSemantics( + name=f"{asset_cfg.name}_joint_pos", + ref=selected_joint_pos, + kind=InputKindEnum.JOINT_POSITION, + element_names=joint_names, + extra={"isaaclab_connection": f"state:{asset_cfg.name}:joint_pos"}, + ), + ) + _set_leapp_traced_observation_input(env, f"{asset_cfg.name}_joint_pos", selected_joint_pos) + return selected_joint_pos + + +def joint_vel(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Joint velocities for the configured joints, exposed as the LEAPP input boundary.""" + real_env = _leapp_real_env(env) + asset = real_env.scene[asset_cfg.name] + selected_joint_vel = _tensor_data_to_torch(asset.data.joint_vel)[:, asset_cfg.joint_ids] + joint_names = _selected_joint_names(asset, asset_cfg.joint_ids) + if _is_leapp_export_env(env): + from leapp import annotate + from leapp.utils.tensor_description import TensorSemantics + + selected_joint_vel = annotate.input_tensors( + env.unwrapped.spec.id, + TensorSemantics( + name=f"{asset_cfg.name}_joint_vel", + ref=selected_joint_vel, + kind=InputKindEnum.JOINT_VELOCITY, + element_names=joint_names, + extra={"isaaclab_connection": f"state:{asset_cfg.name}:joint_vel"}, + ), + ) + return selected_joint_vel + + class gear_shaft_pos_w(ManagerTermBase): """Gear shaft position in world frame with offset applied. @@ -116,19 +213,22 @@ def __call__( Returns: Gear shaft position tensor of shape (num_envs, 3) """ + real_env = _leapp_real_env(env) + # Check if gear type manager exists # During initialization (shape checking), the manager may not exist yet - if not hasattr(env, "_gear_type_manager"): + if not hasattr(real_env, "_gear_type_manager"): # Return default shape during initialization - return torch.zeros(env.num_envs, 3, device=env.device) + return torch.zeros(real_env.num_envs, 3, device=real_env.device) - gear_type_manager: randomize_gear_type = env._gear_type_manager + gear_type_manager: randomize_gear_type = real_env._gear_type_manager # Get gear type indices directly as tensor (no Python loops!) gear_type_indices = gear_type_manager.get_all_gear_type_indices() # Get base gear position and orientation - base_pos = self.asset.data.root_pos_w.torch - base_quat = self.asset.data.root_quat_w.torch + asset = real_env.scene[self.asset_cfg.name] + base_pos = _tensor_data_to_torch(asset.data.root_pos_w) + base_quat = _tensor_data_to_torch(asset.data.root_quat_w) # Update offsets using vectorized indexing self.offsets_buffer = self.gear_offsets_stacked[gear_type_indices] @@ -136,7 +236,22 @@ def __call__( # Transform offsets shaft_pos, _ = combine_frame_transforms(base_pos, base_quat, self.offsets_buffer, self.identity_quat) - return shaft_pos - env.scene.env_origins + shaft_pos = shaft_pos - real_env.scene.env_origins + if _is_leapp_export_env(env): + from leapp import annotate + from leapp.utils.tensor_description import TensorSemantics + + shaft_pos = annotate.input_tensors( + env.unwrapped.spec.id, + TensorSemantics( + name="gear_shaft_pos", + ref=shaft_pos, + kind=InputKindEnum.BODY_POSITION, + element_names=XYZ_ELEMENT_NAMES, + extra={"isaaclab_connection": "observation:policy:gear_shaft_pos"}, + ), + ) + return shaft_pos class gear_shaft_quat_w(ManagerTermBase): @@ -180,8 +295,25 @@ def __call__( Returns: Gear shaft orientation tensor of shape (num_envs, 4) """ - # Get base quaternion - base_quat = self.asset.data.root_quat_w.torch + # Get the raw shaft/base quaternion. During LEAPP export, read it from + # the real env and annotate it before canonicalization so the ONNX graph + # owns the ``qw >= 0`` sign convention. + real_env = _leapp_real_env(env) + base_quat = _tensor_data_to_torch(real_env.scene[self.asset_cfg.name].data.root_quat_w) + if _is_leapp_export_env(env): + from leapp import annotate + from leapp.utils.tensor_description import TensorSemantics + + base_quat = annotate.input_tensors( + env.unwrapped.spec.id, + TensorSemantics( + name="gear_shaft_quat", + ref=base_quat, + kind=InputKindEnum.BODY_ROTATION, + element_names=QUAT_XYZW_ELEMENT_NAMES, + extra={"isaaclab_connection": "observation:policy:gear_shaft_quat"}, + ), + ) # Ensure w component is positive (q and -q represent the same rotation) # Pick one canonical form to reduce observation variation seen by the policy