Update environment loader and wrapper for Isaac Lab 2.0

Toni-SM · Jan 7, 2025 · 8998421 · 8998421
1 parent deb28ff
commit 8998421
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 5 deletions.
diff --git a/skrl/envs/loaders/torch/isaaclab_envs.py b/skrl/envs/loaders/torch/isaaclab_envs.py
@@ -139,7 +139,10 @@ def load_isaaclab_env(
     )
 
     # launch the simulation app
-    from omni.isaac.lab.app import AppLauncher
+    try:
+        from omni.isaac.lab.app import AppLauncher
+    except ModuleNotFoundError:
+        from isaaclab.app import AppLauncher
 
     AppLauncher.add_app_launcher_args(parser)
     args = parser.parse_args()
@@ -149,8 +152,12 @@ def load_isaaclab_env(
     def close_the_simulator():
         app_launcher.app.close()
 
-    import omni.isaac.lab_tasks  # type: ignore
-    from omni.isaac.lab_tasks.utils import parse_env_cfg  # type: ignore
+    try:
+        import omni.isaac.lab_tasks  # type: ignore
+        from omni.isaac.lab_tasks.utils import parse_env_cfg  # type: ignore
+    except ModuleNotFoundError:
+        import isaaclab_tasks  # type: ignore
+        from isaaclab_tasks.utils import parse_env_cfg  # type: ignore
 
     cfg = parse_env_cfg(args.task, device=args.device, num_envs=args.num_envs, use_fabric=not args.disable_fabric)
 

diff --git a/skrl/envs/wrappers/jax/__init__.py b/skrl/envs/wrappers/jax/__init__.py
@@ -97,7 +97,7 @@ def _in(values, container):
         if verbose:
             logger.info(f"Environment wrapper: 'auto' (class: {', '.join(base_classes)})")
 
-        if _in("omni.isaac.lab.envs..*", base_classes):
+        if _in(["omni.isaac.lab.*", "isaaclab.*"], base_classes):
             return "isaaclab-*"
         elif _in("omni.isaac.gym..*", base_classes):
             return "omniverse-isaacgym"

diff --git a/skrl/envs/wrappers/torch/__init__.py b/skrl/envs/wrappers/torch/__init__.py
@@ -103,7 +103,7 @@ def _in(values, container):
         if verbose:
             logger.info(f"Environment wrapper: 'auto' (class: {', '.join(base_classes)})")
 
-        if _in("omni.isaac.lab.envs..*", base_classes):
+        if _in(["omni.isaac.lab.*", "isaaclab.*"], base_classes):
             return "isaaclab-*"
         elif _in("omni.isaac.gym..*", base_classes):
             return "omniverse-isaacgym"