added some import gaurds, improved code readability

Author: frlai

scripts/reinforcement_learning/rsl_rl/export.py

@@ -15,9 +15,13 @@
 import time
 from collections.abc import Mapping
 
-import leapp
 import torch
-from leapp import annotate
+
+try:
+    import leapp
+    from leapp import annotate
+except ImportError as e:
+    raise ImportError("LEAPP package is required for policy export. Install with: pip install leapp") from e
 
 # Disable TorchScript before importing task/environment modules so any
 # @torch.jit.script helpers resolve to plain Python functions during export.

source/isaaclab/isaaclab/envs/direct_deployment_env.py

@@ -21,7 +21,11 @@
 
 import torch
 import yaml
-from leapp import InferenceManager
+
+try:
+    from leapp import InferenceManager
+except ImportError as e:
+    raise ImportError("LEAPP package is required for policy deployment testing.Install with: pip install leapp") from e
 
 from isaaclab.managers import CommandManager, EventManager
 from isaaclab.scene import InteractiveScene

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

@@ -58,6 +58,9 @@
     from isaaclab.envs import ManagerBasedEnv
 
 
+VARIABLE_IMPEDANCE_MODES = frozenset({"variable", "variable_kp"})
+
+
 # ══════════════════════════════════════════════════════════════════
 # ExportPatcher
 # ══════════════════════════════════════════════════════════════════
@@ -139,24 +142,33 @@ def _disable_training_managers(unwrapped):
         """
         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 = lambda env_ids=None: None
+            unwrapped.curriculum_manager.compute = _noop_curriculum
 
         if hasattr(unwrapped, "reward_manager"):
-            _zero_reward = torch.zeros(num_envs, device=device)
-            unwrapped.reward_manager.compute = lambda dt: _zero_reward
+            unwrapped.reward_manager.compute = _zero_reward_compute
 
         if hasattr(unwrapped, "termination_manager"):
-            _no_termination = torch.zeros(num_envs, dtype=torch.bool, device=device)
-            unwrapped.termination_manager.compute = lambda: _no_termination
+            unwrapped.termination_manager.compute = _no_termination_compute
 
         if hasattr(unwrapped, "recorder_manager"):
             rm = unwrapped.recorder_manager
 
-            def _noop(*args, **kwargs):
-                return None
-
             rm.record_pre_step = _noop
             rm.record_post_step = _noop
             rm.record_pre_reset = _noop
@@ -415,7 +427,7 @@ def _collect_action_outputs(self, action_manager) -> list[TensorSemantics]:
         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", "variable_kp"]:
+            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)
@@ -490,7 +502,7 @@ def _collect_action_static_outputs(
             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", "variable_kp"]:
+            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)

source/isaaclab_rl/test/export/test_rsl_rl_direct_export_flow.py

@@ -170,7 +170,7 @@ def test_direct_env_export_flow():
         cwd=_REPO_ROOT,
         capture_output=True,
         text=True,
-        timeout=600,
+        timeout=6000,
     )
 
     if result.returncode != 0: