From be51d6547154ec149b92d71722db9b919f9a36c9 Mon Sep 17 00:00:00 2001 From: Horde Date: Mon, 15 Jun 2026 22:22:14 +0000 Subject: [PATCH 01/11] Add deploy LEAPP inputs for gear assembly --- .../config/rizon_4s/joint_pos_env_cfg.py | 2 +- .../manipulation/deploy/mdp/__init__.pyi | 2 + .../manipulation/deploy/mdp/actions.py | 113 +++++++++++++ .../manipulation/deploy/mdp/actions_cfg.py | 18 ++ .../manipulation/deploy/mdp/observations.py | 156 +++++++++++++++++- 5 files changed, 282 insertions(+), 9 deletions(-) create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions_cfg.py diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py index c93a78ee6da9..068d568a1b55 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py @@ -276,7 +276,7 @@ def __post_init__(self): # Action configuration for Rizon 4s arm # Using smaller action scale for stability self.joint_action_scale = 0.025 - self.actions.arm_action = mdp.RelativeJointPositionActionCfg( + self.actions.arm_action = mdp.DeployRelativeJointPositionActionCfg( asset_name="robot", joint_names=[ "joint1", diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi index ccfb35c62f19..396e12a1f04c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi @@ -4,6 +4,7 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "DeployRelativeJointPositionActionCfg", "randomize_gear_type", "randomize_gears_and_base_pose", "set_robot_to_grasp_pose", @@ -50,6 +51,7 @@ from .delayed_joint_actions_cfg import ( ShapedDelayedRelativeJointPositionActionCfg, FlexivDynamicsAwareRelativeJointPositionActionCfg, ) +from .actions_cfg import DeployRelativeJointPositionActionCfg from .events import ( randomize_gear_type, randomize_gears_and_base_pose, diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions.py new file mode 100644 index 000000000000..82c870f10da3 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions.py @@ -0,0 +1,113 @@ +# 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 + +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): + real_env = object.__getattribute__(env, "_real_env") if type(env).__name__ == "_EnvProxy" else env + return real_env + + +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: + real_env = _leapp_real_env(env) + return name in getattr(real_env, _LEAPP_CONSUMED_OBSERVATION_INPUTS, set()) + + +class DeployRelativeJointPositionAction(RelativeJointPositionAction): + """Relative joint action that reuses traced current joint observations during LEAPP export.""" + + 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 = real_asset.data.joint_pos.torch[:, self._joint_ids] + else: + current_joint_pos = asset.data.joint_pos.torch[:, 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/manager_based/manipulation/deploy/mdp/actions_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions_cfg.py new file mode 100644 index 000000000000..970dbe45b65b --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/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/manager_based/manipulation/deploy/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py index fbe42cf647f4..1e5f648ed6c8 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py @@ -13,6 +13,11 @@ 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, matrix_from_quat if TYPE_CHECKING: @@ -22,6 +27,106 @@ 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) + + +# These wrappers intentionally shadow the generic Isaac Lab joint observations +# for the deploy gear-assembly MDP. The generic terms read +# ``asset.data.joint_pos/vel.torch`` before slicing by ``asset_cfg.joint_ids``, +# so LEAPP sees and exports the full articulation tensor. Here we create the +# sliced arm-joint tensor first and annotate that tensor as the LEAPP input. +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 = asset.data.joint_pos.torch[:, 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 = asset.data.joint_vel.torch[:, 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. @@ -117,19 +222,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 = wp.to_torch(self.asset.data.root_pos_w) - base_quat = wp.to_torch(self.asset.data.root_quat_w) + 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] @@ -137,7 +245,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): @@ -181,8 +304,25 @@ def __call__( Returns: Gear shaft orientation tensor of shape (num_envs, 4) """ - # Get base quaternion - base_quat = wp.to_torch(self.asset.data.root_quat_w) + # 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 From f85af21c78b3ffa8f10ca44a1cce096026e01de7 Mon Sep 17 00:00:00 2001 From: Horde Date: Mon, 15 Jun 2026 22:35:14 +0000 Subject: [PATCH 02/11] Expose deploy joint observations in MDP namespace --- .../manager_based/manipulation/deploy/mdp/__init__.pyi | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi index 396e12a1f04c..16ba65036cad 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.pyi @@ -16,6 +16,8 @@ __all__ = [ "gear_quat_w", "gear_shaft_pos_w", "gear_shaft_quat_w", + "joint_pos", + "joint_vel", "rigid_object_pos_w", "rigid_object_quat_w", "rigid_object_rot_6d_w", @@ -65,6 +67,8 @@ from .observations import ( gear_quat_w, gear_shaft_pos_w, gear_shaft_quat_w, + joint_pos, + joint_vel, rigid_object_pos_w, rigid_object_quat_w, rigid_object_rot_6d_w, From 08650dbb2e7aa400bbf39cbb48d2ee90942051e9 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Sun, 12 Jul 2026 07:26:12 -0700 Subject: [PATCH 03/11] Add LEAPP export support for DisplayPort --- .../reinforcement_learning/leapp/deploy.py | 82 ++ .../leapp/rsl_rl/export.py | 437 +++++++++++ .../isaaclab/isaaclab/utils/leapp/__init__.py | 10 + .../isaaclab/utils/leapp/__init__.pyi | 61 ++ .../isaaclab/utils/leapp/export_annotator.py | 740 ++++++++++++++++++ .../isaaclab/utils/leapp/leapp_semantics.py | 145 ++++ source/isaaclab/isaaclab/utils/leapp/proxy.py | 521 ++++++++++++ source/isaaclab/isaaclab/utils/leapp/utils.py | 87 ++ 8 files changed, 2083 insertions(+) create mode 100644 scripts/reinforcement_learning/leapp/deploy.py create mode 100644 scripts/reinforcement_learning/leapp/rsl_rl/export.py create mode 100644 source/isaaclab/isaaclab/utils/leapp/__init__.py create mode 100644 source/isaaclab/isaaclab/utils/leapp/__init__.pyi create mode 100644 source/isaaclab/isaaclab/utils/leapp/export_annotator.py create mode 100644 source/isaaclab/isaaclab/utils/leapp/leapp_semantics.py create mode 100644 source/isaaclab/isaaclab/utils/leapp/proxy.py create mode 100644 source/isaaclab/isaaclab/utils/leapp/utils.py diff --git a/scripts/reinforcement_learning/leapp/deploy.py b/scripts/reinforcement_learning/leapp/deploy.py new file mode 100644 index 000000000000..5b032daa9dda --- /dev/null +++ b/scripts/reinforcement_learning/leapp/deploy.py @@ -0,0 +1,82 @@ +# 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 a LEAPP-exported policy in an Isaac Lab simulation.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import sys + +from isaaclab.app import AppLauncher + +parser = argparse.ArgumentParser(description="Deploy a LEAPP-exported policy in simulation.") +parser.add_argument("--task", type=str, default=None, help="Name of the registered Isaac Lab task.") +parser.add_argument("--leapp_model", type=str, default=None, help="Path to the LEAPP .yaml pipeline description.") +parser.add_argument("--seed", type=int, default=None, help="Seed for the environment.") +AppLauncher.add_app_launcher_args(parser) +args_cli, hydra_args = parser.parse_known_args() + +if args_cli.task is None or args_cli.leapp_model is None: + missing_args = [] + if args_cli.task is None: + missing_args.append("--task") + if args_cli.leapp_model is None: + missing_args.append("--leapp_model") + parser.error(f"the following arguments are required: {', '.join(missing_args)}") + +sys.argv = [sys.argv[0]] + hydra_args + +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import torch + +from isaaclab.envs.leapp_deployment_env import LeappDeploymentEnv + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry + + +def main(): + # ── Load env config from gym registry ───────────────────────── + task_name = args_cli.task.split(":")[-1] + env_cfg = load_cfg_from_registry(task_name, "env_cfg_entry_point") + + if args_cli.seed is not None: + env_cfg.seed = args_cli.seed + if args_cli.device is not None: + env_cfg.sim.device = args_cli.device + + # ── Create deploy env ───────────────────────────────────────── + env = LeappDeploymentEnv(env_cfg, args_cli.leapp_model) + + if getattr(args_cli, "headless", False): + print( + "[WARN]: Running deploy without a viewport. This happens when headless mode is active, " + "including the default case where no visualizer was selected. The policy may be " + "stepping normally, but no viewport will appear unless you specify the " + "`--visualizer` field." + ) + + print(f"[INFO]: Deploying task '{task_name}' with LEAPP model: {args_cli.leapp_model}") + print(f"[INFO]: Num envs: {env.num_envs}, decimation: {env.cfg.decimation}, step_dt: {env.step_dt:.4f}s") + + # ── Run loop ────────────────────────────────────────────────── + env.reset() + try: + with torch.inference_mode(): + while simulation_app.is_running(): + env.step() + env.close() + except KeyboardInterrupt: + pass + + +if __name__ == "__main__": + main() + simulation_app.close() diff --git a/scripts/reinforcement_learning/leapp/rsl_rl/export.py b/scripts/reinforcement_learning/leapp/rsl_rl/export.py new file mode 100644 index 000000000000..76585046602f --- /dev/null +++ b/scripts/reinforcement_learning/leapp/rsl_rl/export.py @@ -0,0 +1,437 @@ +# 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 + +"""Script to export a checkpoint if an RL agent from RSL-RL.""" + +from __future__ import annotations + +import argparse +import contextlib +import importlib.metadata as metadata +import os +import sys +import time +from collections.abc import Mapping +from pathlib import Path + +from isaaclab.app import AppLauncher + +from isaaclab_tasks.utils import setup_preset_cli + +_RSL_RL_SCRIPTS_DIR = Path(__file__).resolve().parents[2] / "rsl_rl" +if str(_RSL_RL_SCRIPTS_DIR) not in sys.path: + sys.path.insert(0, str(_RSL_RL_SCRIPTS_DIR)) +import cli_args # isort: skip + + +RSL_RL_MIN_VERSION = "5.0.1" +_RUNTIME_IMPORTS_LOADED = False + +# Keep heavy/runtime-sensitive imports out of module import time. The CLI needs +# to parse launcher arguments and start Isaac Sim/Kit before importing torch, +# LEAPP, RSL-RL, and task modules; importing them earlier has caused launcher +# import-order failures. ``_load_runtime_dependencies()`` populates these +# globals immediately before export execution. +torch = None +leapp = None +annotate = None +gym = None +DistillationRunner = None +OnPolicyRunner = None +ManagerBasedRLEnv = None +RslRlVecEnvWrapper = None +handle_deprecated_rsl_rl_cfg = None +retrieve_file_path = None +patch_env_for_export = None +ensure_env_spec_id = None +get_published_pretrained_checkpoint = None +get_checkpoint_path = None +hydra_task_config = None +installed_version = None + + +def create_arg_parser() -> argparse.ArgumentParser: + """Create the command-line parser for RSL-RL policy export.""" + parser = argparse.ArgumentParser(description="Export an RL agent with RSL-RL.") + parser.add_argument( + "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations." + ) + parser.add_argument("--task", type=str, default=None, help="Name of the task.") + parser.add_argument( + "--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point." + ) + parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") + parser.add_argument( + "--use_pretrained_checkpoint", + action="store_true", + help="Use the pre-trained checkpoint from Nucleus.", + ) + + # LEAPP arguments + parser.add_argument( + "--export_task_name", + type=str, + default=None, + help="Name of the exported graph. Defaults to the task name.", + ) + parser.add_argument( + "--export_method", + type=str, + default="onnx-dynamo", + choices=["onnx-dynamo", "onnx-torchscript", "jit-script", "jit-trace"], + help="Method to export the policy", + ) + parser.add_argument( + "--export_save_path", + type=str, + default=None, + help="Path to save the exported model", + ) + parser.add_argument( + "--validation_steps", + type=int, + default=5, + help="Number of steps to validate the exported model", + ) + parser.add_argument( + "--disable_graph_visualization", + action="store_true", + default=False, + help="Disable LEAPP graph visualization during compile_graph().", + ) + + cli_args.add_rsl_rl_args(parser) + AppLauncher.add_app_launcher_args(parser) + return parser + + +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) + args_cli.headless = True + return args_cli, hydra_args + + +def _load_runtime_dependencies() -> None: + """Import runtime dependencies after Isaac Sim has been launched.""" + global _RUNTIME_IMPORTS_LOADED + global annotate, leapp, torch + global DistillationRunner, ManagerBasedRLEnv, OnPolicyRunner, RslRlVecEnvWrapper, get_checkpoint_path, gym + global ensure_env_spec_id, get_published_pretrained_checkpoint, handle_deprecated_rsl_rl_cfg, hydra_task_config + global installed_version + global patch_env_for_export, retrieve_file_path + + if _RUNTIME_IMPORTS_LOADED: + return + + try: + import leapp as leapp_module + except ImportError as e: + raise ImportError("LEAPP package is required for policy export. Install with: pip install leapp") from e + annotate_module = getattr(leapp_module, "annotate") + + import gymnasium as gym_module + import torch as torch_module + from packaging import version as packaging_version_module + from rsl_rl.runners import DistillationRunner as DistillationRunnerCls + from rsl_rl.runners import OnPolicyRunner as OnPolicyRunnerCls + + # Disable TorchScript before importing task/environment modules so any + # @torch.jit.script helpers resolve to plain Python functions during export. + torch_module.jit._state.disable() + + from isaaclab.envs import ManagerBasedRLEnv as ManagerBasedRLEnvCls + from isaaclab.utils.assets import retrieve_file_path as retrieve_file_path_fn + from isaaclab.utils.leapp import patch_env_for_export as patch_env_for_export_fn + from isaaclab.utils.leapp.utils import ensure_env_spec_id as ensure_env_spec_id_fn + + from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper as RslRlVecEnvWrapperCls + from isaaclab_rl.rsl_rl import handle_deprecated_rsl_rl_cfg as handle_deprecated_rsl_rl_cfg_fn + from isaaclab_rl.utils.pretrained_checkpoint import ( + get_published_pretrained_checkpoint as get_published_pretrained_checkpoint_fn, + ) + + __import__("isaaclab_tasks") + from isaaclab_tasks.utils import get_checkpoint_path as get_checkpoint_path_fn + from isaaclab_tasks.utils.hydra import hydra_task_config as hydra_task_config_fn + + installed_version = metadata.version("rsl-rl-lib") + if packaging_version_module.parse(installed_version) < packaging_version_module.parse(RSL_RL_MIN_VERSION): + print( + f"[WARNING] LEAPP RSL-RL export is validated with rsl-rl-lib {RSL_RL_MIN_VERSION} or newer. " + f"Installed version is '{installed_version}'." + ) + + torch = torch_module + leapp = leapp_module + annotate = annotate_module + gym = gym_module + DistillationRunner = DistillationRunnerCls + OnPolicyRunner = OnPolicyRunnerCls + ManagerBasedRLEnv = ManagerBasedRLEnvCls + RslRlVecEnvWrapper = RslRlVecEnvWrapperCls + handle_deprecated_rsl_rl_cfg = handle_deprecated_rsl_rl_cfg_fn + retrieve_file_path = retrieve_file_path_fn + patch_env_for_export = patch_env_for_export_fn + ensure_env_spec_id = ensure_env_spec_id_fn + get_published_pretrained_checkpoint = get_published_pretrained_checkpoint_fn + get_checkpoint_path = get_checkpoint_path_fn + hydra_task_config = hydra_task_config_fn + _RUNTIME_IMPORTS_LOADED = True + + +def get_actor_memory_module(policy): + """Return the actor-side RNN module for supported RSL-RL recurrent policies.""" + if hasattr(policy, "rnn"): + return policy.rnn + return None + + +def is_actor_recurrent_policy(policy) -> bool: + """Return whether the actor policy has a supported recurrent state container.""" + return bool(getattr(policy, "is_recurrent", False) and get_actor_memory_module(policy) is not None) + + +def get_actor_hidden_state(policy): + """Return the actor-side recurrent hidden state for supported RSL-RL policy APIs.""" + if hasattr(policy, "get_hidden_state"): + return policy.get_hidden_state() + memory = get_actor_memory_module(policy) + return None if memory is None else getattr(memory, "hidden_state", None) + + +def set_actor_hidden_state(policy, actor_hidden) -> None: + """Assign the actor-side recurrent hidden state for supported RSL-RL policy APIs.""" + memory = get_actor_memory_module(policy) + if memory is not None: + memory.hidden_state = actor_hidden + + +def ensure_actor_hidden_state_initialized(policy, batch_size: int, device, dtype): + """Initialize and return the actor hidden state when a recurrent policy has not created it yet.""" + # ``torch`` is a lazy runtime global populated by ``_load_runtime_dependencies()`` + # after Isaac Sim launches and before export calls this helper. + assert torch is not None + actor_state = get_actor_hidden_state(policy) + if actor_state is not None: + return actor_state + + memory = get_actor_memory_module(policy) + if memory is None or not hasattr(memory, "rnn"): + return None + + num_layers = memory.rnn.num_layers + hidden_size = memory.rnn.hidden_size + zeros = torch.zeros(num_layers, batch_size, hidden_size, device=device, dtype=dtype) + if isinstance(memory.rnn, torch.nn.LSTM): + actor_state = (zeros.clone(), zeros.clone()) + else: + actor_state = zeros + set_actor_hidden_state(policy, actor_state) + return actor_state + + +def state_dict_from_actor_hidden(actor_hidden): + """Convert the actor hidden state into the named tensor mapping expected by LEAPP state APIs.""" + if actor_hidden is None: + return {} + if isinstance(actor_hidden, tuple): + return {f"actor_state_{idx}": tensor for idx, tensor in enumerate(actor_hidden)} + return {"actor_state": actor_hidden} + + +def actor_hidden_from_registered(registered_state, original_hidden): + """Restore the registered LEAPP state to the hidden-state structure expected by the actor memory module.""" + if isinstance(original_hidden, tuple): + if isinstance(registered_state, tuple): + return registered_state + return (registered_state,) + 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, + agent_cfg, + simulation_app=None, +) -> bool: + """Export a RSL-RL agent.""" + _load_runtime_dependencies() + + task_name = args_cli.task.split(":")[-1] + checkpoint_task_name = task_name.replace("-Play", "") + + agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli) + env_cfg.scene.num_envs = 1 + + agent_cfg = handle_deprecated_rsl_rl_cfg(agent_cfg, installed_version) + + # note: certain randomizations occur in the environment initialization so we set the seed here + env_cfg.seed = agent_cfg.seed + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + + log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name) + log_root_path = os.path.abspath(log_root_path) + print(f"[INFO] Loading checkpoint search path from directory: {log_root_path}") + if args_cli.use_pretrained_checkpoint: + resume_path = get_published_pretrained_checkpoint("rsl_rl", checkpoint_task_name) + if not resume_path: + print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.") + return False + elif args_cli.checkpoint: + resume_path = retrieve_file_path(args_cli.checkpoint) + else: + resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) + + if not resume_path: + print(f"[INFO] No checkpoint found for task: {checkpoint_task_name} in directory: {log_root_path}") + return False + + log_dir = os.path.dirname(resume_path) + + env_cfg.log_dir = log_dir + + env = None + 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) + + graph_name = args_cli.export_task_name if args_cli.export_task_name is not None else task_name + + if isinstance(env.unwrapped, ManagerBasedRLEnv): + # Patch only the observation groups consumed by the actor policy. + # This filters out the critic and teacher observation groups. + obs_groups_cfg = getattr(agent_cfg, "obs_groups", None) + if isinstance(obs_groups_cfg, Mapping): + required_obs_groups = set(obs_groups_cfg.get("actor", ["policy"])) + else: + required_obs_groups = {"policy"} + patch_env_for_export( + env, + export_method=args_cli.export_method, + required_obs_groups=required_obs_groups, + ) + + env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions) + + print(f"[INFO]: Loading model checkpoint from: {resume_path}") + if agent_cfg.class_name == "OnPolicyRunner": + runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device) + elif agent_cfg.class_name == "DistillationRunner": + runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device) + else: + raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}") + runner.load(resume_path) + + policy = runner.get_inference_policy(device=env.unwrapped.device) + + if args_cli.export_save_path is not None: + save_path = args_cli.export_save_path + elif args_cli.use_pretrained_checkpoint: + # Use a predictable path independent of the Nucleus mirror directory structure. + save_path = os.path.join(".pretrained_checkpoints", "rsl_rl", checkpoint_task_name) + else: + save_path = log_dir + leapp.start(graph_name, save_path=save_path, max_cached_io=max(args_cli.validation_steps, 2)) + leapp_started = True + obs = env.reset()[0] + if simulation_app is not None: + while not simulation_app.is_running(): + time.sleep(0.5) + + for _ in range(max(args_cli.validation_steps, 2)): + with torch.inference_mode(): + if is_actor_recurrent_policy(policy): + actor_hidden = ensure_actor_hidden_state_initialized( + policy, + batch_size=env.num_envs, + device=env.unwrapped.device, + dtype=next(policy.parameters()).dtype, + ) + registered_state = annotate.state_tensors( + policy_node_name, + state_dict_from_actor_hidden(actor_hidden), + ) + set_actor_hidden_state(policy, actor_hidden_from_registered(registered_state, actor_hidden)) + + actions = policy(obs) + + if is_actor_recurrent_policy(policy): + actor_hidden_after = get_actor_hidden_state(policy) + annotate.update_state( + policy_node_name, + state_dict_from_actor_hidden(actor_hidden_after), + ) + + obs, _, _, _ = env.step(actions) + + leapp.stop() + leapp_started = False + validate = args_cli.validation_steps > 0 + leapp.compile_graph(visualize=not args_cli.disable_graph_visualization, validate=validate) + finally: + if leapp_started: + with contextlib.suppress(Exception): + leapp.stop() + if env is not None: + env.close() + + return True + + +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 + + from isaaclab_tasks.utils.hydra import hydra_task_config + + original_argv = sys.argv + # Hydra reads the preset tokens (physics=/renderer=/presets=) from sys.argv directly. + sys.argv = [sys.argv[0]] + hydra_args + exported = False + + try: + + @hydra_task_config(args_cli.task, args_cli.agent) + def _main(env_cfg, agent_cfg) -> None: + nonlocal exported + with launch_simulation(env_cfg, args_cli): + exported = export_rsl_rl_agent(args_cli, env_cfg, agent_cfg) + + _main() + finally: + sys.argv = original_argv + + return exported + + +def main_cli(argv: list[str] | None = None) -> bool: + """Run the command-line export flow.""" + args_cli, hydra_args = parse_export_args(argv) + return run_export_with_hydra(args_cli, hydra_args) + + +if __name__ == "__main__": + main_cli() diff --git a/source/isaaclab/isaaclab/utils/leapp/__init__.py b/source/isaaclab/isaaclab/utils/leapp/__init__.py new file mode 100644 index 000000000000..f39b0e4d7eea --- /dev/null +++ b/source/isaaclab/isaaclab/utils/leapp/__init__.py @@ -0,0 +1,10 @@ +# 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 + +"""Sub-module for LEAPP export annotation and proxy-based policy tracing.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab/isaaclab/utils/leapp/__init__.pyi b/source/isaaclab/isaaclab/utils/leapp/__init__.pyi new file mode 100644 index 000000000000..e2b8f497b5f3 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/leapp/__init__.pyi @@ -0,0 +1,61 @@ +# 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 + +__all__ = [ + "ExportPatcher", + "InputKindEnum", + "LeappTensorSemantics", + "OutputKindEnum", + "POSE6_ELEMENT_NAMES", + "POSE7_ELEMENT_NAMES", + "QUAT_XYZW_ELEMENT_NAMES", + "WRENCH6_ELEMENT_NAMES", + "XYZ_ELEMENT_NAMES", + "body_names_resolver", + "body_pose6_resolver", + "body_pose_resolver", + "body_quat_resolver", + "body_wrench_resolver", + "body_xyz_resolver", + "build_command_connection", + "build_state_connection", + "build_write_connection", + "joint_names_resolver", + "leapp_tensor_semantics", + "patch_env_for_export", + "resolve_leapp_element_names", + "target_frame_pose_resolver", + "target_frame_quat_resolver", + "target_frame_xyz_resolver", +] + +from .export_annotator import ExportPatcher, patch_env_for_export +from .leapp_semantics import ( + InputKindEnum, + OutputKindEnum, + POSE6_ELEMENT_NAMES, + POSE7_ELEMENT_NAMES, + QUAT_XYZW_ELEMENT_NAMES, + WRENCH6_ELEMENT_NAMES, + XYZ_ELEMENT_NAMES, + LeappTensorSemantics, + body_names_resolver, + body_pose6_resolver, + body_pose_resolver, + body_quat_resolver, + body_wrench_resolver, + body_xyz_resolver, + joint_names_resolver, + leapp_tensor_semantics, + resolve_leapp_element_names, + target_frame_pose_resolver, + target_frame_quat_resolver, + target_frame_xyz_resolver, +) +from .utils import ( + build_command_connection, + build_state_connection, + build_write_connection, +) diff --git a/source/isaaclab/isaaclab/utils/leapp/export_annotator.py b/source/isaaclab/isaaclab/utils/leapp/export_annotator.py new file mode 100644 index 000000000000..df695c448f02 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/leapp/export_annotator.py @@ -0,0 +1,740 @@ +# 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 + +"""Export annotations for Isaac Lab policies using proxy-based patching. + +Observation and action annotation share a unified dedup cache so that a +state property (e.g. ``joint_pos``) read by both an observation term and +an action term resolves to one LEAPP input edge. + +- Observation term functions see an ``_EnvProxy`` whose scene returns + ``_EntityProxy`` objects with annotating data proxies. + +- Action terms have their ``_asset`` attribute replaced with an + _ArticulationWriteProxy that intercepts ``_leapp_semantics``-decorated + write methods **and** routes ``.data`` reads through the same annotating + data proxy used by observations. + +Cache lifecycle (assuming single-env play-mode export): + + compute() clear cache → obs terms populate cache + policy inference TracedTensors propagate through NN + process_action() register_buffer for raw_actions + apply_action() [tracing] reuse cached TracedTensors for state reads, + capture write outputs, call output_tensors(), + then clear cache + apply_action() [decim.] clear cache → fresh reads for simulation + ... + compute() clear cache → fresh reads for next obs +""" + +from __future__ import annotations + +import inspect +import logging +from collections.abc import Callable +from contextlib import suppress +from typing import TYPE_CHECKING, Any + +import torch +from leapp import annotate +from leapp.utils.tensor_description import TensorSemantics + +from isaaclab.assets.articulation.base_articulation import BaseArticulation +from isaaclab.managers import ManagerTermBase + +from .leapp_semantics import select_element_names +from .proxy import _ArticulationWriteProxy, _DataProxy, _EnvProxy, _ManagerTermProxy +from .utils import ( + TracedProxyArray, + build_command_connection, + build_write_connection, +) + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + +VARIABLE_IMPEDANCE_MODES = frozenset({"variable", "variable_kp"}) + + +# ══════════════════════════════════════════════════════════════════ +# ExportPatcher +# ══════════════════════════════════════════════════════════════════ + + +class ExportPatcher: + """Unified patcher that annotates observation inputs and action outputs for LEAPP export. + + Observation-side property semantics are resolved lazily inside + ``_DataProxy`` by combining: + + - the concrete runtime getter from the backend data class + - the nearest ``_leapp_semantics`` metadata found while walking the MRO + + This lets backends override property implementations without duplicating + decorators from the abstract API. + + - The observation proxy chain (``_EnvProxy`` → ``_SceneProxy`` → + ``_EntityProxy`` → ``_DataProxy``) for state reads + by observation term functions. + - The ``_ArticulationWriteProxy`` on each action term, which intercepts + target writes **and** routes ``.data`` reads through the same + ``_DataProxy`` / cache. + + """ + + def __init__(self, export_method: str, required_obs_groups: set[str] | None = None): + """Initialize the export patcher. + + Args: + export_method: LEAPP export backend passed to + :func:`annotate.output_tensors`. + required_obs_groups: Observation groups that should be patched, or + ``None`` to patch all groups. + """ + self.task_name: str | None = None + self.export_method = export_method + self.required_obs_groups = required_obs_groups + self._annotated_tensor_cache: dict[tuple[int, str], TracedProxyArray] = {} + self._data_property_resolution_cache: dict[tuple[type, str], tuple[Callable, object] | None] = {} + self._write_method_resolution_cache: dict[ + tuple[type, str], tuple[Callable, object, inspect.Signature] | None + ] = {} + self._action_output_cache: list[TensorSemantics] = [] + self._captured_write_term_names: set[str] = set() + self._fallback_term_names: set[str] = set() + self._pending_action_output_export: bool = False + self._uses_last_action_state: bool = False + self._action_term_scene_keys: dict[str, str] = {} + + def setup(self, env): + """Patch the environment in place for LEAPP-aware export. + + Args: + env: Wrapped manager-based environment whose unwrapped instance + should be patched. + """ + unwrapped = env.env.unwrapped + task_name = str(unwrapped.spec.id) + self.task_name = task_name + + proxy_env = _EnvProxy( + unwrapped, + task_name, + self._data_property_resolution_cache, + self._annotated_tensor_cache, + ) + + self._disable_training_managers(unwrapped) + self._patch_observation_manager(unwrapped.observation_manager, proxy_env) + self._patch_history_buffers(unwrapped.observation_manager) + self._patch_action_manager( + unwrapped.action_manager, + self._annotated_tensor_cache, + ) + + # ── Disable training-only managers ───────────────────────────── + + @staticmethod + def _disable_training_managers(unwrapped): + """Replace training-only manager methods with no-ops. + + During export the curriculum, reward, termination, and recorder + managers serve no purpose. Disabling them avoids side-effect + crashes (e.g. ADR curriculum terms accessing nullified noise + configs) and removes unnecessary computation. + + Args: + unwrapped: Unwrapped environment whose training-only managers + should be disabled. + """ + num_envs = unwrapped.num_envs + device = unwrapped.device + _zero_reward = torch.zeros(num_envs, device=device) + _no_termination = torch.zeros(num_envs, dtype=torch.bool, device=device) + + def _noop_curriculum(env_ids=None): + return None + + def _zero_reward_compute(dt): + return _zero_reward + + def _no_termination_compute(): + return _no_termination + + def _noop(*args, **kwargs): + return None + + if hasattr(unwrapped, "curriculum_manager"): + unwrapped.curriculum_manager.compute = _noop_curriculum + + if hasattr(unwrapped, "reward_manager"): + unwrapped.reward_manager.compute = _zero_reward_compute + + if hasattr(unwrapped, "termination_manager"): + unwrapped.termination_manager.compute = _no_termination_compute + + if hasattr(unwrapped, "recorder_manager"): + rm = unwrapped.recorder_manager + + rm.record_pre_step = _noop + rm.record_post_step = _noop + rm.record_pre_reset = _noop + rm.record_post_reset = _noop + rm.record_post_physics_decimation_step = _noop + + @staticmethod + def _resolve_scene_entity_key(scene, entity: Any) -> str | None: + """Return the scene dictionary key for an entity. + + Args: + scene: Scene object that stores entity dictionaries. + entity: Entity instance to locate. + + Returns: + The scene key for ``entity`` if found, otherwise ``None``. + """ + for attr_value in vars(scene).values(): + if not isinstance(attr_value, dict): + continue + for key, candidate in attr_value.items(): + if candidate is entity: + return key + return None + + # ── Observation manager patches ─────────────────────────────── + + def _patch_history_buffers(self, obs_manager): + """Patch history-enabled observation buffers to export as LEAPP state. + + Args: + obs_manager: Observation manager whose history buffers should be + wrapped. + """ + history_buffers = getattr(obs_manager, "_group_obs_term_history_buffer", {}) + term_names_by_group = getattr(obs_manager, "_group_obs_term_names", {}) + + for group_name, term_cfgs in obs_manager._group_obs_term_cfgs.items(): + if self.required_obs_groups is not None and group_name not in self.required_obs_groups: + continue + group_buffers = history_buffers.get(group_name, {}) + group_term_names = term_names_by_group.get(group_name, []) + + for index, term_cfg in enumerate(term_cfgs): + history_length = getattr(term_cfg, "history_length", 0) or 0 + if history_length <= 0: + continue + + if index >= len(group_term_names): + continue + + term_name = group_term_names[index] + circular_buffer = group_buffers.get(term_name) + if circular_buffer is None: + continue + + state_name = f"h_{group_name}_{term_name}" + self._patch_history_buffer_append(circular_buffer, state_name) + + def _patch_history_buffer_append(self, circular_buffer, state_name: str): + """Wrap ``_append`` so history buffers become explicit LEAPP state. + + Args: + circular_buffer: Circular buffer instance to patch. + state_name: LEAPP state tensor name for the buffer contents. + """ + if hasattr(circular_buffer, "_leapp_original_append"): + return + + task_name = self.task_name + original_append = circular_buffer._append + + def patched_append(data: torch.Tensor): + """Annotate history buffer updates as LEAPP state transitions. + + Args: + data: New observation slice appended to the buffer. + + Returns: + ``None``. + """ + if circular_buffer._buffer is not None: + circular_buffer._buffer = annotate.state_tensors(task_name, {state_name: circular_buffer._buffer}) + + original_append(data) + + if circular_buffer._buffer is not None: + circular_buffer._buffer = annotate.update_state(task_name, {state_name: circular_buffer._buffer}) + + circular_buffer._leapp_original_append = original_append + circular_buffer._append = patched_append + + def _patch_observation_manager(self, obs_manager, proxy_env): + """Patch observation terms to use annotating proxies and disable noise. + + Args: + obs_manager: Observation manager instance to patch. + proxy_env: Proxy environment routed into observation terms. + """ + for group_name, term_cfgs in obs_manager._group_obs_term_cfgs.items(): + if self.required_obs_groups is not None and group_name not in self.required_obs_groups: + continue + for term_cfg in term_cfgs: + original_func = term_cfg.func + func_name = getattr(original_func, "__name__", None) + + if func_name == "last_action": + self._uses_last_action_state = True + term_cfg.func = self._wrap_last_action(original_func) + elif func_name == "generated_commands": + term_cfg.func = self._wrap_generated_commands(original_func, term_cfg) + elif func_name == "projected_gravity": + term_cfg.func = self._wrap_projected_gravity(original_func, proxy_env) + else: + term_cfg.func = self._wrap_with_proxy(original_func, proxy_env) + + term_cfg.noise = None + + original_compute = obs_manager.compute + cache = self._annotated_tensor_cache + + 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 + + # ── Action manager patches ──────────────────────────────────── + + def _patch_action_manager(self, action_manager, cache): + """Patch action terms with write/read proxies and manager hooks. + + Args: + action_manager: Action manager instance to patch. + cache: Shared tensor dedup cache for annotated state reads. + """ + assert self.task_name is not None + scene = action_manager._env.scene + for term_name, term in action_manager._terms.items(): + asset = getattr(term, "_asset", None) + if isinstance(asset, BaseArticulation): + real_asset: BaseArticulation = asset + scene_key = self._resolve_scene_entity_key(scene, real_asset) or "ego" + data_proxy = _DataProxy( + real_asset.data, + scene_key, + self.task_name, + self._data_property_resolution_cache, + cache, + input_name_resolver=lambda prop_name, k=scene_key: f"{k}_{prop_name}", + ) + term._asset = _ArticulationWriteProxy( + real_asset=real_asset, + entity_name=scene_key, + term_name=term_name, + output_cache=self._action_output_cache, + method_resolution_cache=self._write_method_resolution_cache, + captured_write_term_names=self._captured_write_term_names, + data_proxy=data_proxy, + ) + self._action_term_scene_keys[term_name] = scene_key + + self._patch_action_manager_methods(action_manager) + + def _patch_action_manager_methods(self, action_manager): + """Patch ``process_action`` and ``apply_action`` on the action manager instance. + + ``process_action`` registers raw_action buffers for LEAPP tracing and + preserves the action tensor clone. + + ``apply_action`` coordinates the cache and output lifecycle: + + - **Tracing pass** (first ``apply_action`` after ``process_action``): + The cache still holds TracedTensors populated by ``compute_group``. + Action terms that read state (e.g. ``RelativeJointPositionAction`` + reading ``joint_pos``) get those TracedTensors from the cache, + keeping the LEAPP graph connected. After ``output_tensors()`` the + cache is cleared so subsequent decimation sub-steps read fresh values. + + - **Non-tracing passes** (remaining decimation sub-steps and all + subsequent iterations): The cache is cleared **before** running + action terms so every ``.data`` read returns the current simulator + value, preserving simulation correctness. + + Args: + action_manager: Action manager whose instance methods should be + wrapped. + """ + original_process = action_manager.process_action + original_apply = action_manager.apply_action + task_name = self.task_name + cache = self._annotated_tensor_cache + + def patched_process_action(action: torch.Tensor): + """Register raw_action buffers, call real process_action, preserve action clone.""" + original_process(action) + action_manager._action = action.clone() + self._pending_action_output_export = True + + def patched_apply_action(): + """Coordinate cache lifecycle and LEAPP output annotation.""" + if not self._pending_action_output_export: + cache.clear() + return original_apply() + + # Tracing pass: cache still holds TracedTensors from compute_group. + self._action_output_cache.clear() + self._captured_write_term_names.clear() + original_apply() + + self._action_output_cache.extend(self._collect_action_outputs(action_manager)) + self._action_output_cache.extend(self._collect_processed_action_fallbacks(action_manager)) + if self._uses_last_action_state: + annotate.update_state(task_name, {"last_action": action_manager._action}) + fallback_terms = self._fallback_term_names + static_values = self._collect_action_static_outputs(action_manager, fallback_terms) + annotate.output_tensors( + task_name, + self._action_output_cache, + static_outputs=static_values or None, + export_with=self.export_method, + ) + self._pending_action_output_export = False + self._action_output_cache.clear() + cache.clear() + return None + + action_manager.process_action = patched_process_action + action_manager.apply_action = patched_apply_action + + # ── Observation term wrappers ───────────────────────────────── + + @staticmethod + def _wrap_with_proxy(original_func, proxy_env): + """Wrap a term function so it receives the proxy env. + + Args: + original_func: Original observation term function or manager term. + proxy_env: Proxy environment routed into the wrapped callable. + + Returns: + Wrapped callable that substitutes ``proxy_env`` for the real env. + """ + + if isinstance(original_func, ManagerTermBase): + return _ManagerTermProxy(original_func, proxy_env) + + def wrapped(*args, **kwargs): + """Invoke the original function with the proxy environment. + + Args: + *args: Original positional arguments. + **kwargs: Original keyword arguments. + + Returns: + Result of the wrapped observation term. + """ + if args: + args = (proxy_env, *args[1:]) + else: + args = (proxy_env,) + return original_func(*args, **kwargs) + + wrapped.__name__ = getattr(original_func, "__name__", "unknown") + return wrapped + + @staticmethod + def _wrap_projected_gravity(original_func, proxy_env): + """Wrap projected gravity as root-quaternion input plus fixed gravity projection. + + Deployment backends generally provide body orientation, not an already + projected gravity vector. During export, keep the policy observation as + projected gravity while exposing ``root_quat_w`` at the LEAPP graph + boundary. + """ + + def wrapped(*args, **kwargs): + """Compute projected gravity from an annotated root quaternion input.""" + kwargs.pop("inspect", None) + asset_cfg = kwargs.get("asset_cfg") + if asset_cfg is None and len(args) > 1: + asset_cfg = args[1] + asset_name = getattr(asset_cfg, "name", "robot") + root_quat_w = proxy_env.scene[asset_name].data.root_quat_w.torch + gravity_w = torch.zeros((*root_quat_w.shape[:-1], 3), dtype=root_quat_w.dtype, device=root_quat_w.device) + gravity_w[..., 2] = -1.0 + quat_xyz = root_quat_w[..., :3] + quat_w = root_quat_w[..., 3:4] + t = torch.cross(quat_xyz, gravity_w, dim=-1) * 2.0 + return gravity_w - quat_w * t + torch.cross(quat_xyz, t, dim=-1) + + wrapped.__name__ = getattr(original_func, "__name__", "unknown") + return wrapped + + def _wrap_last_action(self, original_func): + """Wrap ``last_action`` as a LEAPP state tensor. + + ``last_action`` is feedback state, not a regular dangling input. We + therefore register it through ``annotate.state_tensors(...)`` on the + observation side and update it through ``annotate.update_state(...)`` + after the traced action pass. + + Args: + original_func: Original ``last_action`` observation term. + + Returns: + Wrapped callable that exports ``last_action`` as LEAPP state. + """ + task_name = self.task_name + + def wrapped(env, action_name=None, **kwargs): + """Run the wrapped ``last_action`` term and annotate its output. + + Args: + env: Environment passed by the observation manager. + action_name: Optional action term name. + **kwargs: Additional keyword arguments for the term. + + Returns: + Annotated last-action tensor. + """ + result = original_func(env, action_name, **kwargs) + return annotate.state_tensors(task_name, {"last_action": result}) + + wrapped.__name__ = original_func.__name__ + return wrapped + + def _wrap_generated_commands(self, original_func, term_cfg): + """Wrap the ``generated_commands`` observation term to annotate its output as a LEAPP input. + + Resolves command semantics (kind, element_names) from the command manager + configuration when available. + + Args: + original_func: Original ``generated_commands`` observation term. + term_cfg: Observation term config used to resolve the command name. + + Returns: + Wrapped callable that exports generated commands as LEAPP inputs. + """ + task_name = self.task_name + command_name_from_cfg = term_cfg.params.get("command_name") + + def wrapped(env, command_name=None, **kwargs): + """Run the wrapped command term and annotate its output. + + Args: + env: Environment passed by the observation manager. + command_name: Optional command term name override. + **kwargs: Additional keyword arguments for the term. + + Returns: + Annotated command tensor. + """ + result = original_func(env, command_name, **kwargs) + leapp_input_name = command_name or command_name_from_cfg or "commands" + command_cfg = None + with suppress(AttributeError, KeyError): + command_cfg = env.command_manager.get_term(leapp_input_name).cfg + sem = TensorSemantics( + name=leapp_input_name, + ref=result, + kind=getattr(command_cfg, "cmd_kind", None), + element_names=getattr(command_cfg, "element_names", None), + extra=build_command_connection(leapp_input_name), + ) + return annotate.input_tensors(task_name, sem) + + wrapped.__name__ = original_func.__name__ + return wrapped + + # ── Output collection ───────────────────────────────────────── + + def _collect_action_outputs(self, action_manager) -> list[TensorSemantics]: + """Collect non-writer action tensors that should be exported. + + Args: + action_manager: Action manager whose terms should be inspected. + + Returns: + Exportable tensor semantics for dynamic action outputs such as OSC + gains. + """ + tensors: list[TensorSemantics] = [] + for term_name, term in action_manager._terms.items(): + osc = getattr(term, "_osc", None) + if osc and hasattr(osc, "cfg") and osc.cfg.impedance_mode in VARIABLE_IMPEDANCE_MODES: + asset = getattr(term, "_asset", None) + real_asset = getattr(asset, "_real_asset", asset) + joint_ids = getattr(term, "_joint_ids", None) + joint_names = getattr(real_asset, "joint_names", None) if real_asset else None + scene_key = self._action_term_scene_keys.get(term_name, "ego") + tensors.append( + TensorSemantics( + name=f"{term_name}_kp_gains", + ref=torch.diagonal(osc._motion_p_gains_task, dim1=-2, dim2=-1), + kind="kp", + element_names=select_element_names(joint_names, joint_ids), + extra=build_write_connection(scene_key, "write_joint_stiffness_to_sim_index"), + ) + ) + tensors.append( + TensorSemantics( + name=f"{term_name}_kd_gains", + ref=torch.diagonal(osc._motion_d_gains_task, dim1=-2, dim2=-1), + kind="kd", + element_names=select_element_names(joint_names, joint_ids), + extra=build_write_connection(scene_key, "write_joint_damping_to_sim_index"), + ) + ) + return tensors + + def _collect_processed_action_fallbacks(self, action_manager) -> list[TensorSemantics]: + """Fallback: use ``term.processed_actions`` for terms that produced no write outputs. + + When an action term does not call any ``_leapp_semantics``-decorated write method + (e.g. ``PreTrainedPolicyAction`` which delegates writes to a nested sub-policy), + we fall back to capturing ``term.processed_actions`` as the output tensor. + + Args: + action_manager: Action manager whose terms should be inspected. + + Returns: + Fallback tensor semantics built from ``processed_actions``. + """ + logger = logging.getLogger(__name__) + fallback_terms: set[str] = set() + tensors: list[TensorSemantics] = [] + for term_name, term in action_manager._terms.items(): + if term_name in self._captured_write_term_names: + continue + processed = getattr(term, "processed_actions", None) + if processed is None: + continue + if isinstance(processed, torch.Tensor): + logger.warning( + "Action term '%s' did not write to any asset directly. Falling back to processed_actions as the" + " export output.\nIf you wish to add semantic data to this policy, you need to manually annotate it" + " with output_tensors.", + term_name, + ) + tensors.append( + TensorSemantics( + name=term_name, + ref=processed.clone(), + kind=None, + element_names=None, + ) + ) + fallback_terms.add(term_name) + self._fallback_term_names = fallback_terms + return tensors + + def _collect_action_static_outputs( + self, action_manager, skip_terms: set[str] | None = None + ) -> list[TensorSemantics]: + """Collect static kp/kd gain values from action terms for export metadata. + + Terms in ``skip_terms`` are excluded — these are terms that fell back + to ``processed_actions`` and whose static gains (kp/kd) belong to a + lower abstraction level that is not part of the exported policy. + + Args: + action_manager: Action manager whose terms should be inspected. + skip_terms: Action term names whose static outputs should be + skipped. + + Returns: + Static tensor semantics for action gains exported as metadata. + """ + static_values: list[TensorSemantics] = [] + for term_name, term in action_manager._terms.items(): + if skip_terms and term_name in skip_terms: + continue + osc = getattr(term, "_osc", None) + if osc and hasattr(osc, "cfg") and osc.cfg.impedance_mode in VARIABLE_IMPEDANCE_MODES: + continue + asset = getattr(term, "_asset", None) + real_asset = getattr(asset, "_real_asset", asset) + if real_asset and hasattr(real_asset, "data"): + data = real_asset.data + joint_ids = getattr(term, "_joint_ids", None) + 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 + static_values.append( + TensorSemantics( + name=f"{term_name}_kp_gains", + ref=gains[:, joint_ids] if joint_ids else gains, + kind="kp", + element_names=select_element_names(joint_names, joint_ids), + extra=build_write_connection(scene_key, "write_joint_stiffness_to_sim_index"), + ) + ) + if hasattr(data, "default_joint_damping") and data.default_joint_damping is not None: + gains = data.default_joint_damping.torch + static_values.append( + TensorSemantics( + name=f"{term_name}_kd_gains", + ref=gains[:, joint_ids] if joint_ids else gains, + kind="kd", + element_names=select_element_names(joint_names, joint_ids), + extra=build_write_connection(scene_key, "write_joint_damping_to_sim_index"), + ) + ) + return static_values + + +# ══════════════════════════════════════════════════════════════════ +# Public entry point +# ══════════════════════════════════════════════════════════════════ + + +def patch_env_for_export( + env: ManagerBasedEnv, + export_method: str, + required_obs_groups: set[str] | None = None, +) -> None: + """Patch the env's observation and action managers for LEAPP export. + + This is a thin public entry point around ``ExportPatcher``. It mutates + the provided env instance in-place so that: + + - Observation terms route through proxy objects that annotate tensor + reads from **any** scene entity data class (articulations, rigid + objects, sensors, etc.). + - Action terms route through proxy objects that annotate both data + reads **and** ``Articulation`` write methods. + + Data properties are resolved lazily through proxies — no hardcoded + class list is required. To produce LEAPP input annotations, the + accessed data property getter must carry ``_leapp_semantics``. + Likewise, action-side write methods must be annotated to produce + semantic LEAPP outputs. Undecorated reads and writes are forwarded + as normal runtime access, but they do not gain semantic annotation + metadata through this patching path. + + State reads are deduplicated across observation and action paths via a + shared cache, so a property like ``joint_pos`` that is read by both an + observation term and a relative-position action term appears as a single + LEAPP input edge. + + The underlying env, scene, assets, and tensors remain shared with the rest + of the pipeline; only the manager call paths are redirected. + + Args: + env: Manager-based environment to patch in place. + export_method: LEAPP export backend passed to + :func:`annotate.output_tensors`. + required_obs_groups: Observation groups that should be patched, or + ``None`` to patch all groups. + """ + patcher = ExportPatcher(export_method, required_obs_groups=required_obs_groups) + patcher.setup(env) diff --git a/source/isaaclab/isaaclab/utils/leapp/leapp_semantics.py b/source/isaaclab/isaaclab/utils/leapp/leapp_semantics.py new file mode 100644 index 000000000000..340291de16ad --- /dev/null +++ b/source/isaaclab/isaaclab/utils/leapp/leapp_semantics.py @@ -0,0 +1,145 @@ +# 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 + +"""LEAPP semantic metadata helpers for raw tensor-producing functions.""" + +from __future__ import annotations + +from collections.abc import Callable +from contextlib import suppress +from dataclasses import dataclass +from typing import Any + +try: + from leapp import InputKindEnum, OutputKindEnum +except ImportError: + + class _LeappEnumSentinel: + """Stand-in when leapp is not installed. + + Any attribute access returns ``None`` so that + ``@leapp_tensor_semantics(kind=InputKindEnum.BODY_POSE)`` + silently stores ``kind=None`` instead of crashing at import time. + The real enum values are only needed at export time, when leapp + *is* guaranteed to be available. + """ + + def __getattr__(self, name: str): + return None + + InputKindEnum = _LeappEnumSentinel() # type: ignore[assignment,misc] + OutputKindEnum = _LeappEnumSentinel() # type: ignore[assignment,misc] + + +@dataclass(frozen=True) +class LeappTensorSemantics: + """Semantic metadata attached directly to a raw tensor-producing function.""" + + kind: Any = None + element_names: list[str] | list[list[str]] | None = None + element_names_resolver: Callable | None = None + const: bool = False + + +XYZ_ELEMENT_NAMES: list[str] = ["x", "y", "z"] +QUAT_XYZW_ELEMENT_NAMES: list[str] = ["qx", "qy", "qz", "qw"] +POSE7_ELEMENT_NAMES: list[str] = ["x", "y", "z", "qx", "qy", "qz", "qw"] +POSE6_ELEMENT_NAMES: list[str] = ["x", "y", "z", "angular_x", "angular_y", "angular_z"] +WRENCH6_ELEMENT_NAMES: list[str] = ["fx", "fy", "fz", "tx", "ty", "tz"] + + +def select_element_names(names: list[str] | None, indices: Any = None) -> list[str] | None: + """Select element names using optional runtime indices.""" + if names is None: + return None + if indices is None or indices == slice(None): + return list(names) + if isinstance(indices, slice): + return list(names[indices]) + with suppress(AttributeError): + indices = indices.tolist() + if isinstance(indices, (list, tuple)): + return [names[int(index)] for index in indices] + if isinstance(indices, int): + return [names[indices]] + return None + + +def leapp_tensor_semantics( + *, + kind: Any = None, + element_names: list[str] | list[list[str]] | None = None, + element_names_resolver: Callable | None = None, + const: bool = False, +) -> Callable: + """Attach LEAPP semantic metadata to a raw tensor-producing function.""" + + semantics = LeappTensorSemantics( + kind=kind, + element_names=element_names, + element_names_resolver=element_names_resolver, + const=const, + ) + + def _apply(func: Callable) -> Callable: + func._leapp_semantics = semantics + return func + + return _apply + + +def resolve_leapp_element_names(semantics: LeappTensorSemantics | None, data_self) -> list | None: + """Resolve element names from attached semantics and a tensor-producing object.""" + if semantics is None: + return None + if semantics.element_names is not None: + return semantics.element_names + if semantics.element_names_resolver is not None: + return semantics.element_names_resolver(data_self) + return None + + +# ── Predefined element-name resolvers ───────────────────────────── + + +def joint_names_resolver(data_self) -> list[str] | None: + """Resolve joint element names from the data object at trace time.""" + return select_element_names( + getattr(data_self, "joint_names", getattr(data_self, "_joint_names", None)), + getattr(data_self, "_joint_ids", None), + ) + + +def body_names_resolver(data_self) -> list[str] | None: + """Resolve body element names from the data object at trace time.""" + return select_element_names( + getattr(data_self, "body_names", getattr(data_self, "_body_names", None)), + getattr(data_self, "_body_ids", None), + ) + + +def _compound_resolver(outer_fn: Callable, inner_names: list[str]) -> Callable: + """Build a 2D resolver: ``[outer_names, inner_constant_names]``.""" + + def resolver(data_self) -> list | None: + outer = outer_fn(data_self) + return [outer, inner_names] if outer else None + + return resolver + + +def _target_frame_names(data_self) -> list[str] | None: + names = getattr(data_self, "target_frame_names", None) + return list(names) if names is not None else None + + +body_xyz_resolver = _compound_resolver(body_names_resolver, XYZ_ELEMENT_NAMES) +body_pose_resolver = _compound_resolver(body_names_resolver, POSE7_ELEMENT_NAMES) +body_pose6_resolver = _compound_resolver(body_names_resolver, POSE6_ELEMENT_NAMES) +body_quat_resolver = _compound_resolver(body_names_resolver, QUAT_XYZW_ELEMENT_NAMES) +body_wrench_resolver = _compound_resolver(body_names_resolver, WRENCH6_ELEMENT_NAMES) +target_frame_xyz_resolver = _compound_resolver(_target_frame_names, XYZ_ELEMENT_NAMES) +target_frame_quat_resolver = _compound_resolver(_target_frame_names, QUAT_XYZW_ELEMENT_NAMES) +target_frame_pose_resolver = _compound_resolver(_target_frame_names, POSE7_ELEMENT_NAMES) diff --git a/source/isaaclab/isaaclab/utils/leapp/proxy.py b/source/isaaclab/isaaclab/utils/leapp/proxy.py new file mode 100644 index 000000000000..49ef86e7265f --- /dev/null +++ b/source/isaaclab/isaaclab/utils/leapp/proxy.py @@ -0,0 +1,521 @@ +# 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 + +from __future__ import annotations + +import inspect +from collections.abc import Callable +from typing import Any, cast + +import torch +from leapp.utils.tensor_description import TensorSemantics + +from isaaclab.managers import ManagerTermBase +from isaaclab.utils.warp.proxy_array import ProxyArray + +from .leapp_semantics import resolve_leapp_element_names +from .utils import TracedProxyArray, build_write_connection + + +def _resolve_annotated_property( + property_resolution_cache: dict[tuple[type, str], tuple[Callable, Any] | None], + real_data: Any, + name: str, +) -> tuple[Callable, Any] | None: + """Resolve a concrete property getter and inherited semantics metadata. + + The execution getter always comes from the concrete runtime class. Semantic + metadata is resolved independently by walking the MRO until a property + definition with ``_leapp_semantics`` is found. This mirrors the output-side + export path, where semantics are authored on the base API while concrete + backends provide the runtime implementation. + """ + cache_key = (type(real_data), name) + if cache_key in property_resolution_cache: + return property_resolution_cache[cache_key] + + execution_prop = getattr(type(real_data), name, None) + if not isinstance(execution_prop, property) or execution_prop.fget is None: + property_resolution_cache[cache_key] = None + return None + + semantics_meta = None + for data_cls in type(real_data).__mro__: + prop = data_cls.__dict__.get(name) + if not isinstance(prop, property) or prop.fget is None: + continue + candidate = getattr(prop.fget, "_leapp_semantics", None) + if candidate is None: + continue + if getattr(candidate, "const", False): + property_resolution_cache[cache_key] = None + return None + semantics_meta = candidate + break + + if semantics_meta is None: + property_resolution_cache[cache_key] = None + return None + + resolution = (execution_prop.fget, semantics_meta) + property_resolution_cache[cache_key] = resolution + return resolution + + +def _resolve_annotated_method( + method_resolution_cache: dict[tuple[type, str], tuple[Callable, Any, inspect.Signature] | None], + real_asset: Any, + name: str, +) -> tuple[Callable, Any, inspect.Signature] | None: + """Resolve a concrete bound method and inherited semantics metadata.""" + cache_key = (type(real_asset), name) + if cache_key in method_resolution_cache: + return method_resolution_cache[cache_key] + + original_method = getattr(real_asset, name, None) + if not callable(original_method): + method_resolution_cache[cache_key] = None + return None + + for asset_cls in type(real_asset).__mro__: + candidate = asset_cls.__dict__.get(name) + if not callable(candidate): + continue + semantics_meta = getattr(candidate, "_leapp_semantics", None) + if semantics_meta is None: + continue + resolution = (original_method, semantics_meta, inspect.signature(candidate)) + method_resolution_cache[cache_key] = resolution + return resolution + + method_resolution_cache[cache_key] = None + return None + + +class _WriteJointNameContext: + """Resolve runtime joint-name subsets for lazy write interception.""" + + __slots__ = ("joint_names", "_joint_ids") + + def __init__(self, joint_names: list[str], joint_ids): + self.joint_names = joint_names + self._joint_ids = joint_ids + + +def _unique_output_name(term_name: str, method_name: str, output_cache: list[TensorSemantics]) -> str: + """Return a stable, unique output name for an action write entry.""" + existing = {t.name for t in output_cache} + candidate = term_name + if candidate in existing: + candidate = f"{term_name}_{method_name}" + suffix = 2 + while candidate in existing: + candidate = f"{term_name}_{method_name}_{suffix}" + suffix += 1 + return candidate + + +class _DataProxy: + """Proxy around a real data object that intercepts tensor-returning property reads. + + The real data object may be any scene entity data class (``ArticulationData``, + ``RigidObjectData``, sensor data classes, etc.). The proxy resolves property + semantics lazily on first access by walking the runtime class MRO. This lets + concrete backend overrides reuse semantic metadata authored on abstract base + properties without copying decorators onto every implementation. + + When a semantic property returns a :class:`~isaaclab.utils.warp.ProxyArray`, + the result is wrapped in a ``TracedProxyArray`` and cached for + deduplication. Non-proxy results and ordinary attributes are forwarded + transparently. + + All other attribute access is forwarded transparently to the real object. + """ + + def __init__( + self, + real_data: Any, + entity_name: str, + task_name: str, + property_resolution_cache: dict[tuple[type, str], tuple[Callable, Any] | None], + cache: dict, + input_name_resolver: Callable, + ): + object.__setattr__(self, "_real_data", real_data) + object.__setattr__(self, "_entity_name", entity_name) + object.__setattr__(self, "_task_name", task_name) + object.__setattr__(self, "_property_resolution_cache", property_resolution_cache) + object.__setattr__(self, "_cache", cache) + object.__setattr__(self, "_input_name_resolver", input_name_resolver) + + def __getattr__(self, name): + """Intercept semantic property reads; forward everything else.""" + real_data = object.__getattribute__(self, "_real_data") + resolution = _resolve_annotated_property( + object.__getattribute__(self, "_property_resolution_cache"), real_data, name + ) + if resolution is None: + return getattr(real_data, name) + + cache = object.__getattribute__(self, "_cache") + cache_key = (id(real_data), name) + if cache_key in cache: + return cache[cache_key] + + execution_fget, semantics_meta = resolution + result = execution_fget(real_data) + if not isinstance(result, ProxyArray): + return result + + input_name = object.__getattribute__(self, "_input_name_resolver")(name) + traced = TracedProxyArray( + result, + input_name=input_name, + semantics_meta=semantics_meta, + real_data=real_data, + entity_name=object.__getattribute__(self, "_entity_name"), + property_name=name, + task_name=object.__getattribute__(self, "_task_name"), + ) + cache[cache_key] = traced + return traced + + +class _EntityProxy: + """Proxy around a real scene entity that returns a ``_DataProxy`` for ``.data``. + + All other attribute access is forwarded transparently to the real asset. + """ + + def __init__(self, real_entity: Any, data_proxy: _DataProxy): + object.__setattr__(self, "_real_entity", real_entity) + object.__setattr__(self, "_data_proxy", data_proxy) + + @property + def data(self): + """Return the annotating data proxy instead of the real data object.""" + return object.__getattribute__(self, "_data_proxy") + + def __getattr__(self, name): + """Forward all non-data attribute access to the real scene entity.""" + return getattr(object.__getattribute__(self, "_real_entity"), name) + + +class _EntityMappingProxy: + """Proxy around a mapping of scene entities that lazily wraps data-producing entries.""" + + def __init__( + self, + real_mapping, + task_name: str, + property_resolution_cache: dict[tuple[type, str], tuple[Callable, Any] | None], + cache: dict, + ): + object.__setattr__(self, "_real_mapping", real_mapping) + object.__setattr__(self, "_task_name", task_name) + object.__setattr__(self, "_property_resolution_cache", property_resolution_cache) + object.__setattr__(self, "_cache", cache) + object.__setattr__(self, "_proxied", {}) + + def __getitem__(self, key): + """Return a proxied entity when it has a ``.data`` attribute.""" + proxied = object.__getattribute__(self, "_proxied") + if key in proxied: + return proxied[key] + real_mapping = object.__getattribute__(self, "_real_mapping") + entity = real_mapping[key] + data = getattr(entity, "data", None) + if data is None: + return entity + data_proxy = _DataProxy( + data, + key, + object.__getattribute__(self, "_task_name"), + object.__getattribute__(self, "_property_resolution_cache"), + object.__getattribute__(self, "_cache"), + input_name_resolver=lambda prop_name: f"{key}_{prop_name}", + ) + proxy = _EntityProxy(entity, data_proxy) + proxied[key] = proxy + return proxy + + def get(self, key, default=None): + """Return a proxied entity when present, default otherwise.""" + real_mapping = object.__getattribute__(self, "_real_mapping") + if key not in real_mapping: + return default + return self[key] + + def __iter__(self): + return iter(object.__getattribute__(self, "_real_mapping")) + + def __len__(self): + return len(object.__getattribute__(self, "_real_mapping")) + + def __getattr__(self, name): + """Forward all other mapping access to the real mapping.""" + return getattr(object.__getattribute__(self, "_real_mapping"), name) + + +class _SceneProxy: + """Proxy around the real InteractiveScene. + + When an observation term looks up a scene entity by name, this proxy lazily + wraps any entity that has a ``.data`` attribute. All tensor-returning + properties on the data object are intercepted for LEAPP annotation. This + covers articulations, rigid objects, and sensors through both + ``scene["name"]`` and ``scene.sensors["name"]`` access paths. + """ + + def __init__( + self, + real_scene, + task_name: str, + property_resolution_cache: dict[tuple[type, str], tuple[Callable, Any] | None], + cache: dict, + ): + object.__setattr__(self, "_real_scene", real_scene) + object.__setattr__(self, "_task_name", task_name) + object.__setattr__(self, "_property_resolution_cache", property_resolution_cache) + object.__setattr__(self, "_cache", cache) + object.__setattr__(self, "_proxied", {}) + object.__setattr__(self, "_sensor_mapping_proxy", None) + + def _maybe_proxy_entity(self, key: str, entity: Any): + """Return a proxy for any entity that has a ``.data`` attribute.""" + proxied = object.__getattribute__(self, "_proxied") + if key in proxied: + return proxied[key] + + data = getattr(entity, "data", None) + if data is None: + return entity + + cache = object.__getattribute__(self, "_cache") + data_proxy = _DataProxy( + data, + key, + object.__getattribute__(self, "_task_name"), + object.__getattribute__(self, "_property_resolution_cache"), + cache, + input_name_resolver=lambda prop_name, k=key: f"{k}_{prop_name}", + ) + proxy = _EntityProxy(entity, data_proxy) + proxied[key] = proxy + return proxy + + def __getitem__(self, key): + """Return a proxied entity when it exposes annotated data getters.""" + real_scene = object.__getattribute__(self, "_real_scene") + entity = real_scene[key] + return self._maybe_proxy_entity(key, entity) + + @property + def sensors(self): + """Return a mapping proxy for scene sensors.""" + sensor_mapping_proxy = object.__getattribute__(self, "_sensor_mapping_proxy") + if sensor_mapping_proxy is None: + real_scene = object.__getattribute__(self, "_real_scene") + sensor_mapping_proxy = _EntityMappingProxy( + real_scene.sensors, + object.__getattribute__(self, "_task_name"), + object.__getattribute__(self, "_property_resolution_cache"), + object.__getattribute__(self, "_cache"), + ) + object.__setattr__(self, "_sensor_mapping_proxy", sensor_mapping_proxy) + return sensor_mapping_proxy + + def __getattr__(self, name): + """Forward all other scene access to the real scene.""" + return getattr(object.__getattribute__(self, "_real_scene"), name) + + +class _EnvProxy: + """Proxy around the real env that returns a _SceneProxy for ``.scene``. + + All other attribute access (``num_envs``, ``command_manager``, etc.) + is forwarded transparently to the real env. + """ + + def __init__( + self, + real_env, + task_name: str, + property_resolution_cache: dict[tuple[type, str], tuple[Callable, Any] | None], + cache: dict, + ): + object.__setattr__(self, "_real_env", real_env) + object.__setattr__( + self, + "_scene_proxy", + _SceneProxy(real_env.scene, task_name, property_resolution_cache, cache), + ) + + @property + def scene(self): + """Return the scene proxy instead of the real scene.""" + return object.__getattribute__(self, "_scene_proxy") + + def __getattr__(self, name): + """Forward all non-scene attribute access to the real env.""" + return getattr(object.__getattribute__(self, "_real_env"), name) + + +def _build_scene_entity_lookup(real_scene) -> dict[int, tuple[str, str]]: + """Map real scene entity object ids to their lookup path.""" + lookup: dict[int, tuple[str, str]] = {} + for attr_name, attr_value in vars(real_scene).items(): + if not isinstance(attr_value, dict): + continue + container_kind = "sensors" if attr_name == "sensors" else "scene" + for key, entity in attr_value.items(): + lookup[id(entity)] = (container_kind, key) + return lookup + + +class _ManagerTermProxy(ManagerTermBase): + """Proxy a class-based manager term while preserving its lifecycle methods. + + Observation manager terms can be stateful ``ManagerTermBase`` instances that + expose ``reset()`` and ``serialize()`` in addition to being callable. This + proxy preserves that interface while swapping the env argument passed into + ``__call__`` for the observation-side proxy env. + """ + + def __init__(self, target: ManagerTermBase, proxy_env: _EnvProxy): + super().__init__(target.cfg, target._env) + self._target = target + self._proxy_env = proxy_env + self._entity_lookup = _build_scene_entity_lookup(target._env.scene) + + @property + def __name__(self) -> str: + """Expose the wrapped term name for compatibility and debugging.""" + return getattr(self._target, "__name__", self._target.__class__.__name__) + + def reset(self, env_ids=None) -> None: + """Forward resets to the wrapped term instance.""" + self._target.reset(env_ids=env_ids) + + def serialize(self) -> dict: + """Forward serialization to the wrapped term instance.""" + return self._target.serialize() + + def __call__(self, *args, **kwargs): + """Call the wrapped term with the proxy env in place of the real env.""" + if args: + args = (self._proxy_env, *args[1:]) + else: + args = (self._proxy_env,) + swapped_attrs: list[tuple[str, Any]] = [] + for attr_name, attr_value in vars(self._target).items(): + lookup = self._entity_lookup.get(id(attr_value)) + if lookup is None: + continue + + container_kind, key = lookup + proxy_entity = ( + self._proxy_env.scene.sensors[key] if container_kind == "sensors" else self._proxy_env.scene[key] + ) + swapped_attrs.append((attr_name, attr_value)) + setattr(self._target, attr_name, proxy_entity) + + try: + return self._target(*args, **kwargs) + finally: + for attr_name, attr_value in swapped_attrs: + setattr(self._target, attr_name, attr_value) + + def __getattr__(self, name): + """Forward all other attribute access to the wrapped term instance.""" + return getattr(self._target, name) + + +# ══════════════════════════════════════════════════════════════════ +# Action-side proxy +# ══════════════════════════════════════════════════════════════════ + + +class _ArticulationWriteProxy: + """Proxy around a real articulation implementation for action terms. + + Intercepts ``_leapp_semantics``-decorated write methods **and** routes + ``.data`` reads through a shared ``_DataProxy`` so that + action-side state reads (e.g. ``self._asset.data.joint_pos`` inside + ``RelativeJointPositionAction``) participate in LEAPP annotation and + share the dedup cache with observation-side reads. + + All other attribute access is forwarded transparently to the real asset. + """ + + def __init__( + self, + real_asset: Any, + entity_name: str, + term_name: str, + output_cache: list[TensorSemantics], + method_resolution_cache: dict[tuple[type, str], tuple[Callable, Any, inspect.Signature] | None], + captured_write_term_names: set[str], + data_proxy: _DataProxy, + ): + object.__setattr__(self, "_real_asset", real_asset) + object.__setattr__(self, "_entity_name", entity_name) + object.__setattr__(self, "_term_name", term_name) + object.__setattr__(self, "_output_cache", output_cache) + object.__setattr__(self, "_method_resolution_cache", method_resolution_cache) + object.__setattr__(self, "_captured_write_term_names", captured_write_term_names) + object.__setattr__(self, "_data_proxy", data_proxy) + + @property + def data(self): + """Return the shared annotating data proxy.""" + return object.__getattribute__(self, "_data_proxy") + + def __getattr__(self, name): + """Return an annotating wrapper for semantic write methods; forward everything else.""" + real_asset = object.__getattribute__(self, "_real_asset") + resolution = _resolve_annotated_method( + object.__getattribute__(self, "_method_resolution_cache"), + real_asset, + name, + ) + if resolution is None: + return getattr(real_asset, name) + + original_method, semantics_meta, signature = resolution + term_name = object.__getattribute__(self, "_term_name") + output_cache = object.__getattribute__(self, "_output_cache") + captured_write_term_names = object.__getattribute__(self, "_captured_write_term_names") + + def interceptor(*args, **kwargs): + result = original_method(*args, **kwargs) + bound_args = signature.bind_partial(real_asset, *args, **kwargs) + target = bound_args.arguments.get("target") + + if not isinstance(target, torch.Tensor): + return result + + target_tensor = cast(torch.Tensor, target) + joint_ids = bound_args.arguments.get("joint_ids") + output_cache.append( + TensorSemantics( + name=_unique_output_name(term_name, name, output_cache), + ref=target_tensor.clone(), + kind=semantics_meta.kind, + element_names=resolve_leapp_element_names( + semantics_meta, + _WriteJointNameContext(real_asset.joint_names, joint_ids), + ), + extra=build_write_connection( + object.__getattribute__(self, "_entity_name"), + name, + ), + ) + ) + captured_write_term_names.add(term_name) + + return result + + return interceptor diff --git a/source/isaaclab/isaaclab/utils/leapp/utils.py b/source/isaaclab/isaaclab/utils/leapp/utils.py new file mode 100644 index 000000000000..2308f662ab80 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/leapp/utils.py @@ -0,0 +1,87 @@ +# 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 + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import torch +from leapp import annotate +from leapp.utils.tensor_description import TensorSemantics + +from isaaclab.utils.warp.proxy_array import ProxyArray + +from .leapp_semantics import LeappTensorSemantics, resolve_leapp_element_names + + +class TracedProxyArray(ProxyArray): + _traced_array: torch.Tensor + + def __init__( + self, + proxy_array: ProxyArray, + *, + input_name: str, + semantics_meta: LeappTensorSemantics, + real_data: Any, + entity_name: str, + property_name: str, + task_name: str, + ) -> None: + super().__init__(proxy_array.warp) + astorch = super().torch + sem = TensorSemantics( + name=input_name, + ref=astorch, + kind=semantics_meta.kind, + element_names=resolve_leapp_element_names(semantics_meta, real_data), + extra=build_state_connection(entity_name, property_name), + ) + annotated = annotate.input_tensors(task_name, sem) + object.__setattr__(self, "_traced_array", annotated) + + @property + def torch(self) -> torch.Tensor: + return self._traced_array + + @property + def warp(self) -> Any: + raise AttributeError("warp arrays are not supported for leapp export") + + +def ensure_env_spec_id(env, fallback_task_name: str = "policy") -> str: + """Return ``env.unwrapped.spec.id``, creating a fallback spec when needed.""" + spec = getattr(env.unwrapped, "spec", None) + if spec is None: + env.unwrapped.spec = SimpleNamespace(id=fallback_task_name) + return fallback_task_name + + task_name = getattr(spec, "id", None) + if task_name is None: + spec.id = fallback_task_name + return fallback_task_name + + return task_name + + +# ══════════════════════════════════════════════════════════════════ +# Connection Builders +# ══════════════════════════════════════════════════════════════════ + + +def build_state_connection(entity_name: str, property_name: str) -> dict[str, str]: + """Return a compact deployment connection string for a state property.""" + return {"isaaclab_connection": f"state:{entity_name}:{property_name}"} + + +def build_command_connection(command_name: str) -> dict[str, str]: + """Return a compact deployment connection string for a command term.""" + return {"isaaclab_connection": f"command:{command_name}"} + + +def build_write_connection(entity_name: str, method_name: str) -> dict[str, str]: + """Return a compact deployment connection string for an articulation write target.""" + return {"isaaclab_connection": f"write:{entity_name}:{method_name}"} From f9b156943ff5eeda1be00c7d9cc10d959fbf0ffc Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Sun, 12 Jul 2026 07:28:43 -0700 Subject: [PATCH 04/11] Support older preset CLI in LEAPP export --- .../leapp/rsl_rl/export.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/scripts/reinforcement_learning/leapp/rsl_rl/export.py b/scripts/reinforcement_learning/leapp/rsl_rl/export.py index 76585046602f..e6e9b3343fb7 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 From 62ebacf0145f4cf60d1a5e941ee134c65053fbca Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Sun, 12 Jul 2026 07:29:23 -0700 Subject: [PATCH 05/11] Use legacy launch helper for LEAPP export --- scripts/reinforcement_learning/leapp/rsl_rl/export.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/scripts/reinforcement_learning/leapp/rsl_rl/export.py b/scripts/reinforcement_learning/leapp/rsl_rl/export.py index e6e9b3343fb7..b5f2c4387d3c 100644 --- a/scripts/reinforcement_learning/leapp/rsl_rl/export.py +++ b/scripts/reinforcement_learning/leapp/rsl_rl/export.py @@ -409,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 From 7b969733599e331891a0d4b989e3cd8a04a956c8 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Sun, 12 Jul 2026 07:29:58 -0700 Subject: [PATCH 06/11] Add LEAPP warp proxy array helper --- .../isaaclab/utils/warp/proxy_array.py | 356 ++++++++++++++++++ 1 file changed, 356 insertions(+) create mode 100644 source/isaaclab/isaaclab/utils/warp/proxy_array.py diff --git a/source/isaaclab/isaaclab/utils/warp/proxy_array.py b/source/isaaclab/isaaclab/utils/warp/proxy_array.py new file mode 100644 index 000000000000..2a56c9bebc4b --- /dev/null +++ b/source/isaaclab/isaaclab/utils/warp/proxy_array.py @@ -0,0 +1,356 @@ +# 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 + +"""Warp-first dual-access array wrapper with explicit ``.torch`` and ``.warp`` accessors. + +Inspired by ProxyArray from mujocolab/mjlab (BSD-3-Clause). +""" + +from __future__ import annotations + +import os +import warnings +from typing import ClassVar + +import torch +import warp as wp + +_QUATF_ACCESS_WARN_ENV = "WARN_ON_TORCH_QUATF_ACCESS" +"""Environment variable that, when set to ``"1"``, makes :attr:`ProxyArray.torch` +emit a :class:`UserWarning` on every read of a ``wp.quatf``-typed array. Used as a +runtime aid for tracking down call sites that may still assume Isaac Lab 2.x's +``(w, x, y, z)`` quaternion convention after the migration to Isaac Lab 3.x's +``(x, y, z, w)`` convention. See the Isaac Lab 3.0 migration guide for details.""" + + +class ProxyArray: + """Warp-first array wrapper providing cached zero-copy ``.torch`` and ``.warp`` accessors. + + This class wraps a :class:`warp.array` and provides: + + * A ``.warp`` property that returns the original warp array (for kernel interop). + * A ``.torch`` property that returns a cached, zero-copy :class:`torch.Tensor` view + (via :func:`warp.to_torch`). + * Convenience properties (``shape``, ``dtype``, ``device``) delegated to the warp array. + * A deprecation bridge for common torch functions, indexing, and arithmetic/comparison + operators while emitting a one-time :class:`DeprecationWarning`. Tensor instance methods + such as ``clone()`` are not forwarded; use explicit ``.torch`` access for those. + + Example: + + .. code-block:: python + + import warp as wp + from isaaclab.utils.warp.proxy_array import ProxyArray + + arr = wp.zeros(100, dtype=wp.vec3f, device="cuda:0") + ta = ProxyArray(arr) + + # Explicit access (preferred) + ta.warp # -> wp.array, shape (100,), dtype vec3f + ta.torch # -> torch.Tensor, shape (100, 3) + + # Deprecation bridge (warns once, then silent) + result = ta + 1.0 # works, emits DeprecationWarning + """ + + _deprecation_warned: ClassVar[bool] = False + """Class-level flag ensuring the deprecation warning is emitted at most once.""" + + def __init__(self, wp_array: wp.array) -> None: + """Initialize the ProxyArray wrapper. + + The instance is immutable after construction: the wrapped ``wp.array`` cannot + be reassigned. If the underlying simulation memory is re-allocated, construct + a new :class:`ProxyArray` instead of mutating an existing one. + + Args: + wp_array: The warp array to wrap. + + Raises: + TypeError: If ``wp_array`` is not a :class:`warp.array`. + """ + if not isinstance(wp_array, wp.array): + raise TypeError( + f"ProxyArray expects a warp.array, got {type(wp_array).__name__}." + " If you have a ProxyArray, use it directly instead of wrapping it again." + ) + # Bypass __setattr__ for the two internal fields — everything else raises. + object.__setattr__(self, "_warp", wp_array) + object.__setattr__(self, "_torch_cache", None) + # Cached once at construction so the .torch read path stays a constant-time + # check; only used when the WARN_ON_TORCH_QUATF_ACCESS env var is set. + object.__setattr__(self, "_is_quatf", wp_array.dtype is wp.quatf) + + def __setattr__(self, name: str, value) -> None: + """Forbid mutation of ProxyArray instances except for the internal torch cache. + + The torch view is populated lazily on first ``.torch`` access; that is the + only allowed post-init state change. Every other write raises + :class:`AttributeError` so callers don't accidentally re-point the wrapper. + """ + if name == "_torch_cache": + object.__setattr__(self, name, value) + return + raise AttributeError( + f"ProxyArray is immutable; cannot set attribute {name!r}." + " Construct a new ProxyArray instead of mutating an existing one." + ) + + @staticmethod + def _quatf_access_warning_enabled() -> bool: + """Return ``True`` when the ``WARN_ON_TORCH_QUATF_ACCESS`` env var is set to ``"1"``. + + Read on every :attr:`torch` access to keep the flag dynamic — a single + ``os.environ`` lookup is cheap relative to the warp/torch interop work + that follows. + """ + return os.environ.get(_QUATF_ACCESS_WARN_ENV, "0") == "1" + + # ------------------------------------------------------------------ + # Core accessors + # ------------------------------------------------------------------ + + @property + def warp(self) -> wp.array: + """The underlying warp array.""" + return self._warp + + @property + def torch(self) -> torch.Tensor: + """A cached, zero-copy :class:`torch.Tensor` view of the warp array. + + The tensor is created on first access via :func:`warp.to_torch` and cached + for subsequent calls. Since this is a zero-copy view, modifications to the + tensor are visible through the warp array and vice versa. + + When the underlying warp array has dtype ``wp.quatf`` and the + ``WARN_ON_TORCH_QUATF_ACCESS`` environment variable is set to ``"1"``, + each read emits a :class:`UserWarning` pointing at the call site. This + is a runtime aid for migrating Isaac Lab 2.x code (which used the + ``(w, x, y, z)`` quaternion convention) to Isaac Lab 3.x's + ``(x, y, z, w)`` convention. + """ + if self._is_quatf and self._quatf_access_warning_enabled(): + warnings.warn( + "Reading .torch on a wp.quatf-typed ProxyArray. The Isaac Lab" + " quaternion convention changed from (w, x, y, z) in 2.x to" + " (x, y, z, w) in 3.x. If your code assumes the old order," + " this is likely the source of incorrect rotations." + f" Unset {_QUATF_ACCESS_WARN_ENV} to silence this warning.", + UserWarning, + stacklevel=2, + ) + if self._torch_cache is None: + self._torch_cache = wp.to_torch(self._warp) + return self._torch_cache + + # ------------------------------------------------------------------ + # Convenience properties + # ------------------------------------------------------------------ + + @property + def shape(self) -> tuple[int, ...]: + """Shape of the underlying warp array.""" + return self._warp.shape + + @property + def dtype(self): + """Warp dtype of the underlying array.""" + return self._warp.dtype + + @property + def device(self) -> str: + """Device string of the underlying warp array.""" + return self._warp.device + + def __len__(self) -> int: + """Return the size of the first dimension.""" + return self._warp.shape[0] + + def __repr__(self) -> str: + """Return a string representation of the ProxyArray.""" + return f"ProxyArray(shape={self.shape}, dtype={self.dtype}, device={self.device})" + + # ------------------------------------------------------------------ + # Warp kernel interop + # ------------------------------------------------------------------ + + @property + def __cuda_array_interface__(self): + """Delegate the CUDA array interface to the underlying warp array. + + This allows a ``ProxyArray`` to be passed directly as an argument to + :func:`warp.launch` without explicitly accessing ``.warp``. + + Raises: + AttributeError: If the underlying warp array is not on a CUDA device. + """ + return self._warp.__cuda_array_interface__ + + @property + def __array_interface__(self): + """Delegate the NumPy array interface to the underlying warp array. + + This allows a ``ProxyArray`` to be passed directly as an argument to + :func:`warp.launch` on CPU without explicitly accessing ``.warp``. + + Raises: + AttributeError: If the underlying warp array is not on a CPU device. + """ + return self._warp.__array_interface__ + + # ------------------------------------------------------------------ + # Attribute forwarding (deprecation bridge — delegates to .torch) + # ------------------------------------------------------------------ + + def __getattr__(self, name: str): + """Forward unknown attribute access to the torch view (deprecation bridge). + + Called only when normal attribute lookup fails (i.e. the attribute is not + defined on :class:`ProxyArray` itself), so explicit properties such as + ``shape``, ``dtype``, ``device``, ``warp``, and ``torch`` are unaffected. + + This allows tensor instance methods (``float()``, ``clone()``, ``cpu()``, + ``permute()``, etc.) to be called on a :class:`ProxyArray` without an + explicit ``.torch`` accessor, emitting a one-time :class:`DeprecationWarning`. + """ + self._warn_implicit() + return getattr(self.torch, name) + + # ------------------------------------------------------------------ + # Indexing (deprecation bridge — delegates to .torch) + # ------------------------------------------------------------------ + + def __getitem__(self, key): + """Index into the torch view of this array. + + Supports all torch indexing: ``int``, ``slice``, ``tuple``, + boolean masks, and fancy indexing (multi-dimensional). + """ + self._warn_implicit() + return self.torch[key] + + def __setitem__(self, key, value): + """Write through the torch view into the shared warp memory. + + Supports all torch indexing: ``int``, ``slice``, ``tuple``, + boolean masks, and fancy indexing (multi-dimensional). + """ + self._warn_implicit() + self.torch[key] = value + + # ------------------------------------------------------------------ + # Deprecation bridge + # ------------------------------------------------------------------ + + @classmethod + def _warn_implicit(cls) -> None: + """Emit a one-time deprecation warning for implicit torch usage.""" + if not cls._deprecation_warned: + cls._deprecation_warned = True + warnings.warn( + "Implicit use of ProxyArray as a torch.Tensor is deprecated. " + "Use the explicit .torch property instead (e.g., array.torch).", + DeprecationWarning, + stacklevel=3, + ) + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + """Enable torch operations on ProxyArray by unwrapping to ``.torch``. + + This method is called by PyTorch when a torch function receives a + ``ProxyArray`` as an argument. It unwraps all ``ProxyArray`` instances + to their ``.torch`` tensors and delegates to the original function. + """ + if kwargs is None: + kwargs = {} + cls._warn_implicit() + + def unwrap(x): + if isinstance(x, ProxyArray): + return x.torch + if isinstance(x, (list, tuple)): + return type(x)(unwrap(i) for i in x) + return x + + args = unwrap(args) + kwargs = {k: unwrap(v) for k, v in kwargs.items()} + return func(*args, **kwargs) + + # ------------------------------------------------------------------ + # Arithmetic operators + # ------------------------------------------------------------------ + + def _binop(self, other, op: str) -> torch.Tensor: + """Helper for binary and reflected binary operations.""" + self._warn_implicit() + other_val = other.torch if isinstance(other, ProxyArray) else other + return getattr(self.torch, op)(other_val) + + def __add__(self, other) -> torch.Tensor: + return self._binop(other, "__add__") + + def __radd__(self, other) -> torch.Tensor: + return self._binop(other, "__radd__") + + def __sub__(self, other) -> torch.Tensor: + return self._binop(other, "__sub__") + + def __rsub__(self, other) -> torch.Tensor: + return self._binop(other, "__rsub__") + + def __mul__(self, other) -> torch.Tensor: + return self._binop(other, "__mul__") + + def __rmul__(self, other) -> torch.Tensor: + return self._binop(other, "__rmul__") + + def __truediv__(self, other) -> torch.Tensor: + return self._binop(other, "__truediv__") + + def __rtruediv__(self, other) -> torch.Tensor: + return self._binop(other, "__rtruediv__") + + def __pow__(self, other) -> torch.Tensor: + return self._binop(other, "__pow__") + + def __rpow__(self, other) -> torch.Tensor: + return self._binop(other, "__rpow__") + + def __neg__(self) -> torch.Tensor: + self._warn_implicit() + return -self.torch + + def __pos__(self) -> torch.Tensor: + self._warn_implicit() + return +self.torch + + def __abs__(self) -> torch.Tensor: + self._warn_implicit() + return abs(self.torch) + + # ------------------------------------------------------------------ + # Comparison operators + # ------------------------------------------------------------------ + + def __eq__(self, other) -> torch.Tensor: + return self._binop(other, "__eq__") + + def __ne__(self, other) -> torch.Tensor: + return self._binop(other, "__ne__") + + def __lt__(self, other) -> torch.Tensor: + return self._binop(other, "__lt__") + + def __le__(self, other) -> torch.Tensor: + return self._binop(other, "__le__") + + def __gt__(self, other) -> torch.Tensor: + return self._binop(other, "__gt__") + + def __ge__(self, other) -> torch.Tensor: + return self._binop(other, "__ge__") From f1be5e9a467f59912341c19b699f5cb0d7762549 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Sun, 12 Jul 2026 07:31:28 -0700 Subject: [PATCH 07/11] Support Warp-backed LEAPP deploy tensors --- .../manager_based/manipulation/deploy/mdp/actions.py | 11 +++++++++-- .../manipulation/deploy/mdp/observations.py | 4 ++-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions.py index 82c870f10da3..87c5ec129001 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/actions.py @@ -9,6 +9,8 @@ from typing import TYPE_CHECKING +import warp as wp + from isaaclab.envs.mdp.actions.joint_actions import RelativeJointPositionAction if TYPE_CHECKING: @@ -25,6 +27,11 @@ def _leapp_real_env(env): return real_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) @@ -105,9 +112,9 @@ def apply_actions(self): f"'{self.cfg.asset_name}_joint_pos' observation during LEAPP export." ) real_asset = object.__getattribute__(asset, "_real_asset") - current_joint_pos = real_asset.data.joint_pos.torch[:, self._joint_ids] + current_joint_pos = _tensor_data_to_torch(real_asset.data.joint_pos)[:, self._joint_ids] else: - current_joint_pos = asset.data.joint_pos.torch[:, self._joint_ids] + 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/manager_based/manipulation/deploy/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py index 1e5f648ed6c8..4bec82cd16a3 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py @@ -84,7 +84,7 @@ def joint_pos(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg """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 = asset.data.joint_pos.torch[:, asset_cfg.joint_ids] + 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 @@ -108,7 +108,7 @@ def joint_vel(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg """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 = asset.data.joint_vel.torch[:, asset_cfg.joint_ids] + 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 b0343ec7b09449083c850443ebfc6fc8d3e10ed4 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Sun, 12 Jul 2026 07:32:48 -0700 Subject: [PATCH 08/11] Annotate deploy object pose inputs for LEAPP --- .../manipulation/deploy/mdp/observations.py | 52 +++++++++++++++++-- 1 file changed, 47 insertions(+), 5 deletions(-) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py index 4bec82cd16a3..ef0c30edf7b7 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py @@ -75,6 +75,14 @@ def _set_leapp_traced_observation_input(env, name: str, tensor: torch.Tensor) -> getattr(real_env, _LEAPP_CONSUMED_OBSERVATION_INPUTS, set()).discard(name) +def _deploy_object_input_base_name(asset_name: str) -> str: + """Return a stable deploy input base name for common socket/plug assets.""" + for prefix in ("dp_", "gb300_", "factory_"): + if asset_name.startswith(prefix): + return asset_name[len(prefix) :] + return asset_name + + # These wrappers intentionally shadow the generic Isaac Lab joint observations # for the deploy gear-assembly MDP. The generic terms read # ``asset.data.joint_pos/vel.torch`` before slicing by ``asset_cfg.joint_ids``, @@ -522,14 +530,32 @@ def __call__( asset_cfg: SceneEntityCfg | None = None, offset: list | None = None, ) -> torch.Tensor: - obj_pos = wp.to_torch(self.asset.data.root_pos_w) - obj_quat = wp.to_torch(self.asset.data.root_quat_w) + real_env = _leapp_real_env(env) + asset = real_env.scene[self.asset_cfg.name] + obj_pos = _tensor_data_to_torch(asset.data.root_pos_w) + obj_quat = _tensor_data_to_torch(asset.data.root_quat_w) if torch.any(self.offset_tensor != 0): - offset_repeated = self.offset_tensor.unsqueeze(0).repeat(env.num_envs, 1) + offset_repeated = self.offset_tensor.unsqueeze(0).repeat(real_env.num_envs, 1) obj_pos, _ = combine_frame_transforms(obj_pos, obj_quat, offset_repeated, self.identity_quat) - return obj_pos - env.scene.env_origins + obj_pos = obj_pos - real_env.scene.env_origins + if _is_leapp_export_env(env): + from leapp import annotate + from leapp.utils.tensor_description import TensorSemantics + + input_name = f"{_deploy_object_input_base_name(self.asset_cfg.name)}_pos" + obj_pos = annotate.input_tensors( + env.unwrapped.spec.id, + TensorSemantics( + name=input_name, + ref=obj_pos, + kind=InputKindEnum.BODY_POSITION, + element_names=XYZ_ELEMENT_NAMES, + extra={"isaaclab_connection": f"observation:policy:{input_name}"}, + ), + ) + return obj_pos class rigid_object_quat_w(ManagerTermBase): @@ -560,7 +586,23 @@ def __call__( env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg | None = None, ) -> torch.Tensor: - obj_quat = wp.to_torch(self.asset.data.root_quat_w) + real_env = _leapp_real_env(env) + obj_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 + + input_name = f"{_deploy_object_input_base_name(self.asset_cfg.name)}_quat" + obj_quat = annotate.input_tensors( + env.unwrapped.spec.id, + TensorSemantics( + name=input_name, + ref=obj_quat, + kind=InputKindEnum.BODY_ROTATION, + element_names=QUAT_XYZW_ELEMENT_NAMES, + extra={"isaaclab_connection": f"observation:policy:{input_name}"}, + ), + ) w_negative = obj_quat[:, 3] < 0 positive_quat = obj_quat.clone() From 4566e67d88e1b1452071dcc1bd2a6b0c8fb1ff8c Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Sun, 12 Jul 2026 07:33:28 -0700 Subject: [PATCH 09/11] Use trace-safe deploy quaternion canonicalization --- .../manipulation/deploy/mdp/observations.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py index ef0c30edf7b7..9b0b67e0a486 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py @@ -334,11 +334,8 @@ def __call__( # Ensure w component is positive (q and -q represent the same rotation) # Pick one canonical form to reduce observation variation seen by the policy - w_negative = base_quat[:, 3] < 0 - positive_quat = base_quat.clone() - positive_quat[w_negative] = -base_quat[w_negative] - - return positive_quat + w_negative = base_quat[:, 3:4] < 0 + return torch.where(w_negative, -base_quat, base_quat) class gear_pos_w(ManagerTermBase): @@ -604,11 +601,8 @@ def __call__( ), ) - w_negative = obj_quat[:, 3] < 0 - positive_quat = obj_quat.clone() - positive_quat[w_negative] = -obj_quat[w_negative] - - return positive_quat + w_negative = obj_quat[:, 3:4] < 0 + return torch.where(w_negative, -obj_quat, obj_quat) def _quat_to_rot_6d(quat: torch.Tensor) -> torch.Tensor: From 83e00bbee87b8a72f6dc36d58a3f59141ee3e877 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Mon, 13 Jul 2026 20:51:38 -0700 Subject: [PATCH 10/11] Support DisplayPort NoJointVel LEAPP deploy export --- .../isaaclab/assets/articulation/base_articulation.py | 3 +++ .../isaaclab/isaaclab/utils/leapp/export_annotator.py | 10 ++++++++-- .../config/displayport_rizon_4s/joint_pos_env_cfg.py | 6 ++++-- 3 files changed, 15 insertions(+), 4 deletions(-) diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index 25ca2c4ceaf0..9bffdd0cbc9a 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -16,6 +16,7 @@ import torch import warp as wp +from ...utils.leapp.leapp_semantics import OutputKindEnum, joint_names_resolver, leapp_tensor_semantics from ..asset_base import AssetBase if TYPE_CHECKING: @@ -1266,6 +1267,7 @@ def set_inertias_mask( raise NotImplementedError() @abstractmethod + @leapp_tensor_semantics(kind=OutputKindEnum.JOINT_POSITION, element_names_resolver=joint_names_resolver) def set_joint_position_target_index( self, *, @@ -1293,6 +1295,7 @@ def set_joint_position_target_index( raise NotImplementedError() @abstractmethod + @leapp_tensor_semantics(kind=OutputKindEnum.JOINT_POSITION, element_names_resolver=joint_names_resolver) def set_joint_position_target_mask( self, *, diff --git a/source/isaaclab/isaaclab/utils/leapp/export_annotator.py b/source/isaaclab/isaaclab/utils/leapp/export_annotator.py index df695c448f02..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 # ══════════════════════════════════════════════════════════════════ @@ -668,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", @@ -679,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/manager_based/manipulation/deploy/cable_insertion/config/displayport_rizon_4s/joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/cable_insertion/config/displayport_rizon_4s/joint_pos_env_cfg.py index f2d509660719..da1a368557b2 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/cable_insertion/config/displayport_rizon_4s/joint_pos_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/cable_insertion/config/displayport_rizon_4s/joint_pos_env_cfg.py @@ -11,6 +11,7 @@ """ import math +import os import torch @@ -63,7 +64,8 @@ # switches to a specific combination; see the experiment matrix / commit log. # The block between the START/END markers is what the experiment commits edit. # --- EXP TOGGLES START --- -EXP_SYSID = True # enable sim2real (sysid) action model + PhysX SysID gains +EXP_SYSID = os.getenv("DP_CABLE_EXP_SYSID", "1").lower() not in {"0", "false", "no", "off"} +# enable sim2real (sysid) action model + PhysX SysID gains EXP_SOCKET_POS_RANGE = [0.01, 0.01, 0.02] # socket position randomization, +/- m per axis [x, y, z] EXP_SOCKET_ORN_DEG = 2.0 # socket orientation randomization, +/- deg on roll/pitch/yaw EXP_CURRICULUM = "anneal_80_0_500" # disabled|fixed80|anneal_80_0_1000|anneal_80_20_1000|anneal_80_20_500|anneal_80_0_500 @@ -483,7 +485,7 @@ def __post_init__(self): command_acceleration_limit=USE_SIM2REAL_COMMAND_ACCELERATION_LIMIT, ) else: - self.actions.arm_action = mdp.RelativeJointPositionActionCfg( + self.actions.arm_action = mdp.DeployRelativeJointPositionActionCfg( asset_name="robot", joint_names=_arm_joint_names, scale=self.joint_action_scale, From ee09a800ea1cca3a60d7a4db99ca2d01cfb9ee13 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Mon, 13 Jul 2026 21:38:24 -0700 Subject: [PATCH 11/11] Remove unrelated gear assembly LEAPP action change --- .../deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py index 068d568a1b55..c93a78ee6da9 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/gear_assembly/config/rizon_4s/joint_pos_env_cfg.py @@ -276,7 +276,7 @@ def __post_init__(self): # Action configuration for Rizon 4s arm # Using smaller action scale for stability self.joint_action_scale = 0.025 - self.actions.arm_action = mdp.DeployRelativeJointPositionActionCfg( + self.actions.arm_action = mdp.RelativeJointPositionActionCfg( asset_name="robot", joint_names=[ "joint1",