Add Newton gear LEAPP export support

scripts/reinforcement_learning/leapp/rsl_rl/export.py

@@ -18,7 +18,10 @@
 
 from isaaclab.app import AppLauncher
 
-from isaaclab_tasks.utils import setup_preset_cli
+try:
+    from isaaclab_tasks.utils import setup_preset_cli
+except ImportError:
+    setup_preset_cli = None
 
 _RSL_RL_SCRIPTS_DIR = Path(__file__).resolve().parents[2] / "rsl_rl"
 if str(_RSL_RL_SCRIPTS_DIR) not in sys.path:
@@ -110,10 +113,13 @@ def create_arg_parser() -> argparse.ArgumentParser:
 def parse_export_args(argv: list[str] | None = None) -> tuple[argparse.Namespace, list[str]]:
     """Parse export arguments and return remaining Hydra overrides."""
     parser = create_arg_parser()
-    # setup_preset_cli attaches the preset-selection help group then parses;
-    # remainder still carries typed selectors (physics=/renderer=/presets=)
-    # verbatim for run_export_with_hydra to fold before invoking Hydra.
-    args_cli, hydra_args = setup_preset_cli(parser, argv)
+    if setup_preset_cli is not None:
+        # setup_preset_cli attaches the preset-selection help group then parses;
+        # remainder still carries typed selectors (physics=/renderer=/presets=)
+        # verbatim for run_export_with_hydra to fold before invoking Hydra.
+        args_cli, hydra_args = setup_preset_cli(parser, argv)
+    else:
+        args_cli, hydra_args = parser.parse_known_args(argv)
     args_cli.headless = True
     return args_cli, hydra_args
 
@@ -255,6 +261,19 @@ def actor_hidden_from_registered(registered_state, original_hidden):
     return registered_state
 
 
+def resolve_env_action_class_types(env_cfg) -> None:
+    """Resolve Hydra-overridden action ``class_type`` strings before constructing the env."""
+    from isaaclab.utils.string import string_to_callable
+
+    actions_cfg = getattr(env_cfg, "actions", None)
+    if actions_cfg is None:
+        return
+    for action_cfg in vars(actions_cfg).values():
+        class_type = getattr(action_cfg, "class_type", None)
+        if isinstance(class_type, str) and "{DIR}" not in class_type:
+            action_cfg.class_type = string_to_callable(class_type)
+
+
 def export_rsl_rl_agent(
     args_cli: argparse.Namespace,
     env_cfg,
@@ -301,6 +320,7 @@ def export_rsl_rl_agent(
     leapp_started = False
 
     try:
+        resolve_env_action_class_types(env_cfg)
         env = gym.make(args_cli.task, cfg=env_cfg, render_mode=None)
         policy_node_name = ensure_env_spec_id(env)
 
@@ -389,7 +409,10 @@ def export_rsl_rl_agent(
 
 def run_export_with_hydra(args_cli: argparse.Namespace, hydra_args: list[str]) -> bool:
     """Resolve Hydra task configuration and export one RSL-RL policy."""
-    from isaaclab.app import launch_simulation
+    try:
+        from isaaclab.app import launch_simulation
+    except ImportError:
+        from isaaclab_tasks.utils import launch_simulation
 
     from isaaclab_tasks.utils.hydra import hydra_task_config
 

source/isaaclab/isaaclab/utils/leapp/export_annotator.py

@@ -39,6 +39,7 @@
 from typing import TYPE_CHECKING, Any
 
 import torch
+import warp as wp
 from leapp import annotate
 from leapp.utils.tensor_description import TensorSemantics
 
@@ -60,6 +61,11 @@
 VARIABLE_IMPEDANCE_MODES = frozenset({"variable", "variable_kp"})
 
 
+def _tensor_data_to_torch(data):
+    """Return a torch tensor view for Isaac Lab data stored as torch or Warp-backed data."""
+    return data.torch if hasattr(data, "torch") else wp.to_torch(data)
+
+
 # ══════════════════════════════════════════════════════════════════
 # ExportPatcher
 # ══════════════════════════════════════════════════════════════════
@@ -304,6 +310,8 @@ def _patch_observation_manager(self, obs_manager, proxy_env):
         def patched_compute(*args, **kwargs):
             """Clear the tensor dedup cache once per full observation pass."""
             cache.clear()
+            with suppress(Exception):
+                setattr(obs_manager._env, "_leapp_input_annotation_cache", {})
             return original_compute(*args, **kwargs)
 
         obs_manager.compute = patched_compute
@@ -666,7 +674,7 @@ def _collect_action_static_outputs(
                 joint_names = getattr(real_asset, "joint_names", None)
                 scene_key = self._action_term_scene_keys.get(term_name, "ego")
                 if hasattr(data, "default_joint_stiffness") and data.default_joint_stiffness is not None:
-                    gains = data.default_joint_stiffness.torch
+                    gains = _tensor_data_to_torch(data.default_joint_stiffness)
                     static_values.append(
                         TensorSemantics(
                             name=f"{term_name}_kp_gains",
@@ -677,7 +685,7 @@ def _collect_action_static_outputs(
                         )
                     )
                 if hasattr(data, "default_joint_damping") and data.default_joint_damping is not None:
-                    gains = data.default_joint_damping.torch
+                    gains = _tensor_data_to_torch(data.default_joint_damping)
                     static_values.append(
                         TensorSemantics(
                             name=f"{term_name}_kd_gains",

source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/gear_assembly/config/rizon_4s/ros_inference_env_cfg.py

@@ -11,6 +11,8 @@
 from isaaclab.managers import ObservationTermCfg as ObsTerm
 from isaaclab.utils.configclass import configclass
 
+import isaaclab_tasks.contrib.deploy.mdp as mdp
+
 from .joint_pos_env_cfg import Rizon4sGearAssemblyEnvCfg
 
 
@@ -50,6 +52,12 @@ def __post_init__(self):
 
         # Set joint_action_scale from the existing arm_action.scale
         self.joint_action_scale = self.actions.arm_action.scale
+        self.actions.arm_action = mdp.DeployRelativeJointPositionActionCfg(
+            asset_name="robot",
+            joint_names=self.arm_joint_names,
+            scale=self.joint_action_scale,
+            use_zero_offset=True,
+        )
 
         # Dynamically generate action_scale_joint_space based on action_space
         self.action_scale_joint_space = [self.joint_action_scale] * self.action_space

source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/__init__.pyi

@@ -4,6 +4,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 __all__ = [
+    "DeployRelativeJointPositionActionCfg",
     "randomize_gear_type",
     "randomize_gears_and_base_pose",
     "pin_unselected_gears_to_shafts",
@@ -14,6 +15,8 @@ __all__ = [
     "gear_quat_w",
     "gear_shaft_pos_w",
     "gear_shaft_quat_w",
+    "joint_pos",
+    "joint_vel",
     "keypoint_command_error",
     "keypoint_command_error_exp",
     "keypoint_entity_error",
@@ -30,8 +33,9 @@ from .events import (
     randomize_gears_and_base_pose,
     set_robot_to_grasp_pose,
 )
+from .actions_cfg import DeployRelativeJointPositionActionCfg
 from .noise_models import ResetSampledConstantNoiseModel, ResetSampledConstantNoiseModelCfg
-from .observations import gear_pos_w, gear_quat_w, gear_shaft_pos_w, gear_shaft_quat_w
+from .observations import gear_pos_w, gear_quat_w, gear_shaft_pos_w, gear_shaft_quat_w, joint_pos, joint_vel
 from .rewards import (
     keypoint_command_error,
     keypoint_command_error_exp,

source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions.py

@@ -0,0 +1,124 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Deploy-specific action terms for LEAPP export workflows."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import warp as wp
+
+from isaaclab.envs.mdp.actions.joint_actions import RelativeJointPositionAction
+
+if TYPE_CHECKING:
+    from isaaclab.envs import ManagerBasedEnv
+
+    from .actions_cfg import DeployRelativeJointPositionActionCfg
+
+
+_LEAPP_TRACED_OBSERVATION_INPUTS = "_leapp_traced_observation_inputs"
+_LEAPP_CONSUMED_OBSERVATION_INPUTS = "_leapp_consumed_observation_inputs"
+
+
+def _leapp_real_env(env):
+    """Return the wrapped Isaac Lab env when LEAPP passes an export proxy."""
+    return object.__getattribute__(env, "_real_env") if type(env).__name__ == "_EnvProxy" else env
+
+
+def _tensor_data_to_torch(data):
+    """Return a torch tensor view for Isaac Lab data stored as torch or Warp-backed data."""
+    return data.torch if hasattr(data, "torch") else wp.to_torch(data)
+
+
+def _get_observation_term_from_buffer(env, group_name: str, term_name: str):
+    """Return a term slice from the cached observation buffer."""
+    obs_buffer = getattr(env, "obs_buf", None)
+    if obs_buffer is None:
+        obs_buffer = getattr(getattr(env, "observation_manager", None), "_obs_buffer", None)
+    if not obs_buffer or group_name not in obs_buffer:
+        return None
+
+    group_obs = obs_buffer[group_name]
+    if isinstance(group_obs, dict):
+        return group_obs.get(term_name)
+
+    obs_manager = getattr(env, "observation_manager", None)
+    if obs_manager is None:
+        return None
+
+    term_names = obs_manager.active_terms.get(group_name, [])
+    if term_name not in term_names:
+        return None
+
+    term_index = term_names.index(term_name)
+    term_dims = obs_manager.group_obs_term_dim[group_name]
+    concat_dim = obs_manager._group_obs_concatenate_dim[group_name]
+    if concat_dim > 0:
+        concat_dim -= 1
+
+    start = sum(dim[concat_dim] for dim in term_dims[:term_index])
+    length = term_dims[term_index][concat_dim]
+    return group_obs.narrow(dim=concat_dim, start=start, length=length)
+
+
+def _pop_leapp_traced_observation_input(env, name: str, *, group_name: str, term_name: str):
+    """Consume one traced observation tensor for the current LEAPP action trace."""
+    real_env = _leapp_real_env(env)
+    consumed_inputs = getattr(real_env, _LEAPP_CONSUMED_OBSERVATION_INPUTS, None)
+    if consumed_inputs is None:
+        consumed_inputs = set()
+        setattr(real_env, _LEAPP_CONSUMED_OBSERVATION_INPUTS, consumed_inputs)
+
+    if name in consumed_inputs:
+        return None
+
+    traced_inputs = getattr(real_env, _LEAPP_TRACED_OBSERVATION_INPUTS, {})
+    traced_tensor = traced_inputs.pop(name, None)
+    if traced_tensor is None:
+        traced_tensor = _get_observation_term_from_buffer(real_env, group_name, term_name)
+
+    if traced_tensor is not None:
+        consumed_inputs.add(name)
+    return traced_tensor
+
+
+def _is_leapp_observation_input_consumed(env, name: str) -> bool:
+    """Return whether an export input was already consumed by this action trace."""
+    real_env = _leapp_real_env(env)
+    return name in getattr(real_env, _LEAPP_CONSUMED_OBSERVATION_INPUTS, set())
+
+
+class DeployRelativeJointPositionAction(RelativeJointPositionAction):
+    """Relative joint action that exports absolute joint-position targets for deployment."""
+
+    cfg: DeployRelativeJointPositionActionCfg
+
+    def __init__(self, cfg: DeployRelativeJointPositionActionCfg, env: ManagerBasedEnv):
+        super().__init__(cfg, env)
+
+    def apply_actions(self):
+        asset = self._asset
+        if type(asset).__name__ == "_ArticulationWriteProxy":
+            observation_input_name = f"{self.cfg.asset_name}_joint_pos"
+            current_joint_pos = _pop_leapp_traced_observation_input(
+                self._env,
+                observation_input_name,
+                group_name="policy",
+                term_name="joint_pos",
+            )
+            if current_joint_pos is None:
+                if not _is_leapp_observation_input_consumed(self._env, observation_input_name):
+                    raise RuntimeError(
+                        "DeployRelativeJointPositionAction requires the traced "
+                        f"'{self.cfg.asset_name}_joint_pos' observation during LEAPP export."
+                    )
+                real_asset = object.__getattribute__(asset, "_real_asset")
+                current_joint_pos = _tensor_data_to_torch(real_asset.data.joint_pos)[:, self._joint_ids]
+        else:
+            current_joint_pos = _tensor_data_to_torch(asset.data.joint_pos)[:, self._joint_ids]
+
+        current_actions = self.processed_actions + current_joint_pos
+        self._asset.set_joint_position_target_index(target=current_actions, joint_ids=self._joint_ids)

source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/actions_cfg.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Deploy-specific action configuration classes."""
+
+from __future__ import annotations
+
+from isaaclab.envs.mdp.actions.actions_cfg import RelativeJointPositionActionCfg
+from isaaclab.utils.configclass import configclass
+
+
+@configclass
+class DeployRelativeJointPositionActionCfg(RelativeJointPositionActionCfg):
+    """Configuration for deploy relative joint actions with explicit LEAPP current-joint input."""
+
+    class_type: type | str = "{DIR}.actions:DeployRelativeJointPositionAction"