Fixes callbacks with stage in memory and organize environment tests (#519)

# Description

<!--
Thank you for your interest in sending a pull request. Please make sure
to check the contribution guidelines.

Link:
https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html
-->

Fixes callbacks being cleared when stage in memory is attached.
- uses this change Isaac Sim -
[link](https://gitlab-master.nvidia.com/omniverse/isaac/omni_isaac_sim/-/merge_requests/6598)
- removes some of the old callbacks fix, which is no longer needed with
the Isaac Sim change

Reorganizes environment unit tests
- split test_environments -> test_environments and
test_environments_with_stage_in_memory for separate test reporting
- add a utils file for shared functions in isaaclab_tasks/tests
- removes repose cube allegro direct unit test which can cause an OOM
during the CI test
- improves flakey env tests

<!-- As a practice, it is recommended to open an issue to have
discussions on the proposed pull request.
This makes it easier for the community to keep track of what is being
developed or added, and if a given feature
is demanded by more than one party. -->

## Type of change

<!-- As you go through the list, delete the ones that are not
applicable. -->

- Bug fix (non-breaking change which fixes an issue)
- New feature (non-breaking change which adds functionality)
- Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- This change requires a documentation update

## Screenshots

Please attach before and after screenshots of the change if applicable.

<!--
Example:

| Before | After |
| ------ | ----- |
| _gif/png before_ | _gif/png after_ |

To upload images to a PR -- simply drag and drop an image while in edit
mode and it should upload the image directly. You can then paste that
source into the above before/after sections.
-->

## Checklist

- [X] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [X] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [ ] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

---------
Signed-off-by: rwiltz <165190220+rwiltz@users.noreply.github.com>
Signed-off-by: Kelly Guo <kellyguo123@hotmail.com>
Signed-off-by: Ashwin Varghese Kuruttukulam <123109010+ashwinvkNV@users.noreply.github.com>
Signed-off-by: Kelly Guo <kellyg@nvidia.com>
Signed-off-by: Michael Gussert <michael@gussert.com>
Signed-off-by: samibouziri <79418773+samibouziri@users.noreply.github.com>
Signed-off-by: Mayank Mittal <12863862+Mayankm96@users.noreply.github.com>
Signed-off-by: Kyle Morgenstein <34984693+KyleM73@users.noreply.github.com>
Signed-off-by: Hongyu Li <lihongyu0807@icloud.com>
Signed-off-by: Toni-SM <toni.semu@gmail.com>
Signed-off-by: James Tigue <166445701+jtigue-bdai@users.noreply.github.com>
Signed-off-by: Pascal Roth <57946385+pascal-roth@users.noreply.github.com>
Signed-off-by: Victor Khaustov <3192677+vi3itor@users.noreply.github.com>
Signed-off-by: AlvinC <alvincny529@gmail.com>
Signed-off-by: Tyler Lum <tylergwlum@gmail.com>
Signed-off-by: Miguel Alonso Jr. <76960110+miguelalonsojr@users.noreply.github.com>
Signed-off-by: renaudponcelet <renaud.poncelet@gmail.com>
Signed-off-by: jiehanw <jiehanw@nvidia.com>
Signed-off-by: matthewtrepte <mtrepte@nvidia.com>
Co-authored-by: lotusl-code <lotusl@nvidia.com>
Co-authored-by: Kelly Guo <kellyg@nvidia.com>
Co-authored-by: jaczhangnv <jaczhang@nvidia.com>
Co-authored-by: rwiltz <165190220+rwiltz@users.noreply.github.com>
Co-authored-by: Yanzi Zhu <yanziz@nvidia.com>
Co-authored-by: nv-mhaselton <mhaselton@nvidia.com>
Co-authored-by: cosmith-nvidia <141183495+cosmith-nvidia@users.noreply.github.com>
Co-authored-by: Michael Gussert <michael@gussert.com>
Co-authored-by: CY Chen <cyc@nvidia.com>
Co-authored-by: oahmednv <oahmed@Nvidia.com>
Co-authored-by: Ashwin Varghese Kuruttukulam <123109010+ashwinvkNV@users.noreply.github.com>
Co-authored-by: Rafael Wiltz <rwiltz@nvidia.com>
Co-authored-by: Peter Du <peterd@nvidia.com>
Co-authored-by: chengronglai <chengrongl@nvidia.com>
Co-authored-by: pulkitg01 <pulkitg@nvidia.com>
Co-authored-by: Connor Smith <cosmith@nvidia.com>
Co-authored-by: Ashwin Varghese Kuruttukulam <ashwinvk@nvidia.com>
Co-authored-by: Kelly Guo <kellyguo123@hotmail.com>
Co-authored-by: shauryadNv <shauryad@nvidia.com>
Co-authored-by: Mayank Mittal <12863862+Mayankm96@users.noreply.github.com>
Co-authored-by: samibouziri <79418773+samibouziri@users.noreply.github.com>
Co-authored-by: James Smith <142246516+jsmith-bdai@users.noreply.github.com>
Co-authored-by: Shundo Kishi <syundo0730@gmail.com>
Co-authored-by: Sheikh Dawood <sabdulajees@nvidia.com>
Co-authored-by: Toni-SM <aserranomuno@nvidia.com>
Co-authored-by: Gonglitian <70052908+Gonglitian@users.noreply.github.com>
Co-authored-by: James Tigue <166445701+jtigue-bdai@users.noreply.github.com>
Co-authored-by: Mayank Mittal <mittalma@leggedrobotics.com>
Co-authored-by: Kyle Morgenstein <34984693+KyleM73@users.noreply.github.com>
Co-authored-by: Johnson Sun <20457146+j3soon@users.noreply.github.com>
Co-authored-by: Pascal Roth <57946385+pascal-roth@users.noreply.github.com>
Co-authored-by: Hongyu Li <lihongyu0807@icloud.com>
Co-authored-by: Jean-Francois-Lafleche <57650687+Jean-Francois-Lafleche@users.noreply.github.com>
Co-authored-by: Wei Jinqi <changshanshi@outlook.com>
Co-authored-by: Louis LE LAY <le.lay.louis@gmail.com>
Co-authored-by: Harsh Patel <hapatel@theaiinstitute.com>
Co-authored-by: Kousheek Chakraborty <kousheekc@gmail.com>
Co-authored-by: Victor Khaustov <3192677+vi3itor@users.noreply.github.com>
Co-authored-by: AlvinC <alvincny529@gmail.com>
Co-authored-by: Felipe Mohr <50018670+felipemohr@users.noreply.github.com>
Co-authored-by: AdAstra7 <87345760+likecanyon@users.noreply.github.com>
Co-authored-by: gao <ziqi.gao@iff-extern.fraunhofer.de>
Co-authored-by: Tyler Lum <tylergwlum@gmail.com>
Co-authored-by: -T.K.- <t_k_233@outlook.com>
Co-authored-by: Clemens Schwarke <96480707+ClemensSchwarke@users.noreply.github.com>
Co-authored-by: Miguel Alonso Jr. <76960110+miguelalonsojr@users.noreply.github.com>
Co-authored-by: Miguel Alonso Jr. <miguel.alonso@nfinite.app>
Co-authored-by: renaudponcelet <renaud.poncelet@gmail.com>
Co-authored-by: Ales Borovicka <aborovicka@nvidia.com>
Co-authored-by: nv-mm <mmagdics@nvidia.com>
Co-authored-by: Antoine RICHARD <antoiner@nvidia.com>
Co-authored-by: dvangelder-nvidia <dvangelder@nvidia.com>
Co-authored-by: Jiehan Wang <33852873+JerryJiehanWang@users.noreply.github.com>
Co-authored-by: Octi Zhang <zhengyuz@nvidia.com>

source/isaaclab/config/extension.toml

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.42.27"
+version = "0.42.28"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

source/isaaclab/docs/CHANGELOG.rst

 Changelog
 ---------
 
+0.42.28 (2025-07-18)
+~~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added safe callbacks for stage in memory attaching.
+* Remove on prim deletion callback workaround
+
+
 0.42.27 (2025-07-21)
 ~~~~~~~~~~~~~~~~~~~~
 

source/isaaclab/isaaclab/assets/asset_base.py

@@ -21,7 +21,6 @@ from isaacsim.core.simulation_manager import IsaacEvents, SimulationManager
 from isaacsim.core.utils.stage import get_current_stage
 
 import isaaclab.sim as sim_utils
-from isaaclab.sim import SimulationContext
 
 if TYPE_CHECKING:
     from .asset_base_cfg import AssetBaseCfg
@@ -93,39 +92,8 @@ class AssetBase(ABC):
         if len(matching_prims) == 0:
             raise RuntimeError(f"Could not find prim with path {self.cfg.prim_path}.")
 
-        # register simulator callbacks (with weakref safety to avoid crashes on deletion)
-        def safe_callback(callback_name, event, obj_ref):
-            """Safely invoke a callback on a weakly-referenced object, ignoring ReferenceError if deleted."""
-            try:
-                obj = obj_ref
-                getattr(obj, callback_name)(event)
-            except ReferenceError:
-                # Object has been deleted; ignore.
-                pass
-
-        # note: use weakref on callbacks to ensure that this object can be deleted when its destructor is called.
-        # add callbacks for stage play/stop
-        obj_ref = weakref.proxy(self)
-        timeline_event_stream = omni.timeline.get_timeline_interface().get_timeline_event_stream()
-
-        # the order is set to 10 which is arbitrary but should be lower priority than the default order of 0
-        # register timeline PLAY event callback (lower priority with order=10)
-        self._initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
-            int(omni.timeline.TimelineEventType.PLAY),
-            lambda event, obj_ref=obj_ref: safe_callback("_initialize_callback", event, obj_ref),
-            order=10,
-        )
-        # register timeline STOP event callback (lower priority with order=10)
-        self._invalidate_initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
-            int(omni.timeline.TimelineEventType.STOP),
-            lambda event, obj_ref=obj_ref: safe_callback("_invalidate_initialize_callback", event, obj_ref),
-            order=10,
-        )
-        # register prim deletion callback
-        self._prim_deletion_callback_id = SimulationManager.register_callback(
-            lambda event, obj_ref=obj_ref: safe_callback("_on_prim_deletion", event, obj_ref),
-            event=IsaacEvents.PRIM_DELETION,
-        )
+        # register various callback functions
+        self._register_callbacks()
 
         # add handle for debug visualization (this is set to a valid handle inside set_debug_vis)
         self._debug_vis_handle = None
@@ -296,6 +264,43 @@ class AssetBase(ABC):
     Internal simulation callbacks.
     """
 
+    def _register_callbacks(self):
+        """Registers the timeline and prim deletion callbacks."""
+
+        # register simulator callbacks (with weakref safety to avoid crashes on deletion)
+        def safe_callback(callback_name, event, obj_ref):
+            """Safely invoke a callback on a weakly-referenced object, ignoring ReferenceError if deleted."""
+            try:
+                obj = obj_ref
+                getattr(obj, callback_name)(event)
+            except ReferenceError:
+                # Object has been deleted; ignore.
+                pass
+
+        # note: use weakref on callbacks to ensure that this object can be deleted when its destructor is called.
+        # add callbacks for stage play/stop
+        obj_ref = weakref.proxy(self)
+        timeline_event_stream = omni.timeline.get_timeline_interface().get_timeline_event_stream()
+
+        # the order is set to 10 which is arbitrary but should be lower priority than the default order of 0
+        # register timeline PLAY event callback (lower priority with order=10)
+        self._initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
+            int(omni.timeline.TimelineEventType.PLAY),
+            lambda event, obj_ref=obj_ref: safe_callback("_initialize_callback", event, obj_ref),
+            order=10,
+        )
+        # register timeline STOP event callback (lower priority with order=10)
+        self._invalidate_initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
+            int(omni.timeline.TimelineEventType.STOP),
+            lambda event, obj_ref=obj_ref: safe_callback("_invalidate_initialize_callback", event, obj_ref),
+            order=10,
+        )
+        # register prim deletion callback
+        self._prim_deletion_callback_id = SimulationManager.register_callback(
+            lambda event, obj_ref=obj_ref: safe_callback("_on_prim_deletion", event, obj_ref),
+            event=IsaacEvents.PRIM_DELETION,
+        )
+
     def _initialize_callback(self, event):
         """Initializes the scene elements.
 
@@ -332,9 +337,6 @@ class AssetBase(ABC):
         Note:
             This function is called when the prim is deleted.
         """
-        # skip callback if required
-        if getattr(SimulationContext.instance(), "_skip_next_prim_deletion_callback_fn", False):
-            return
         if prim_path == "/":
             self._clear_callbacks()
             return

source/isaaclab/isaaclab/assets/rigid_object_collection/rigid_object_collection.py

@@ -7,7 +7,6 @@ from __future__ import annotations
 
 import re
 import torch
-import weakref
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
@@ -15,13 +14,12 @@ import omni.kit.app
 import omni.log
 import omni.physics.tensors.impl.api as physx
 import omni.timeline
-from isaacsim.core.simulation_manager import IsaacEvents, SimulationManager
+from isaacsim.core.simulation_manager import SimulationManager
 from pxr import UsdPhysics
 
 import isaaclab.sim as sim_utils
 import isaaclab.utils.math as math_utils
 import isaaclab.utils.string as string_utils
-from isaaclab.sim import SimulationContext
 
 from ..asset_base import AssetBase
 from .rigid_object_collection_data import RigidObjectCollectionData
@@ -93,40 +91,8 @@ class RigidObjectCollection(AssetBase):
         # stores object names
         self._object_names_list = []
 
-        # register simulator callbacks (with weakref safety to avoid crashes on deletion)
-        def safe_callback(callback_name, event, obj_ref):
-            """Safely invoke a callback on a weakly-referenced object, ignoring ReferenceError if deleted."""
-            try:
-                obj = obj_ref
-                getattr(obj, callback_name)(event)
-            except ReferenceError:
-                # Object has been deleted; ignore.
-                pass
-
-        # note: use weakref on callbacks to ensure that this object can be deleted when its destructor is called.
-        # add callbacks for stage play/stop
-        obj_ref = weakref.proxy(self)
-        timeline_event_stream = omni.timeline.get_timeline_interface().get_timeline_event_stream()
-
-        # the order is set to 10 which is arbitrary but should be lower priority than the default order of 0
-        # register timeline PLAY event callback (lower priority with order=10)
-        self._initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
-            int(omni.timeline.TimelineEventType.PLAY),
-            lambda event, obj_ref=obj_ref: safe_callback("_initialize_callback", event, obj_ref),
-            order=10,
-        )
-        # register timeline STOP event callback (lower priority with order=10)
-        self._invalidate_initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
-            int(omni.timeline.TimelineEventType.STOP),
-            lambda event, obj_ref=obj_ref: safe_callback("_invalidate_initialize_callback", event, obj_ref),
-            order=10,
-        )
-        # register prim deletion callback
-        self._prim_deletion_callback_id = SimulationManager.register_callback(
-            lambda event, obj_ref=obj_ref: safe_callback("_on_prim_deletion", event, obj_ref),
-            event=IsaacEvents.PRIM_DELETION,
-        )
-
+        # register various callback functions
+        self._register_callbacks()
         self._debug_vis_handle = None
 
     """
@@ -729,9 +695,6 @@ class RigidObjectCollection(AssetBase):
         Note:
             This function is called when the prim is deleted.
         """
-        # skip callback if required
-        if getattr(SimulationContext.instance(), "_skip_next_prim_deletion_callback_fn", False):
-            return
         if prim_path == "/":
             self._clear_callbacks()
             return

source/isaaclab/isaaclab/assets/surface_gripper/surface_gripper.py

@@ -7,11 +7,8 @@ from __future__ import annotations
 
 import torch
 import warnings
-import weakref
 from typing import TYPE_CHECKING
 
-import omni.timeline
-from isaacsim.core.simulation_manager import IsaacEvents, SimulationManager
 from isaacsim.core.utils.extensions import enable_extension
 from isaacsim.core.version import get_version
 
@@ -69,39 +66,8 @@ class SurfaceGripper(AssetBase):
         self._is_initialized = False
         self._debug_vis_handle = None
 
-        # Register simulator callbacks (with weakref safety to avoid crashes on deletion)
-        def safe_callback(callback_name, event, obj_ref):
-            """Safely invoke a callback on a weakly-referenced object, ignoring ReferenceError if deleted."""
-            try:
-                obj = obj_ref
-                getattr(obj, callback_name)(event)
-            except ReferenceError:
-                # Object has been deleted; ignore.
-                pass
-
-        # note: Use weakref on callbacks to ensure that this object can be deleted when its destructor is called.
-        # add callbacks for stage play/stop
-        obj_ref = weakref.proxy(self)
-        timeline_event_stream = omni.timeline.get_timeline_interface().get_timeline_event_stream()
-
-        # The order is set to 10 which is arbitrary but should be lower priority than the default order of 0
-        # Register timeline PLAY event callback (lower priority with order=10)
-        self._initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
-            int(omni.timeline.TimelineEventType.PLAY),
-            lambda event, obj_ref=obj_ref: safe_callback("_initialize_callback", event, obj_ref),
-            order=10,
-        )
-        # Register timeline STOP event callback (lower priority with order=10)
-        self._invalidate_initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
-            int(omni.timeline.TimelineEventType.STOP),
-            lambda event, obj_ref=obj_ref: safe_callback("_invalidate_initialize_callback", event, obj_ref),
-            order=10,
-        )
-        # Register prim deletion callback
-        self._prim_deletion_callback_id = SimulationManager.register_callback(
-            lambda event, obj_ref=obj_ref: safe_callback("_on_prim_deletion", event, obj_ref),
-            event=IsaacEvents.PRIM_DELETION,
-        )
+        # register various callback functions
+        self._register_callbacks()
 
     """
     Properties

source/isaaclab/isaaclab/sensors/sensor_base.py

@@ -26,7 +26,6 @@ from isaacsim.core.simulation_manager import IsaacEvents, SimulationManager
 from isaacsim.core.utils.stage import get_current_stage
 
 import isaaclab.sim as sim_utils
-from isaaclab.sim import SimulationContext
 
 if TYPE_CHECKING:
     from .sensor_base_cfg import SensorBaseCfg
@@ -64,39 +63,8 @@ class SensorBase(ABC):
         # get stage handle
         self.stage = get_current_stage()
 
-        # register simulator callbacks (with weakref safety to avoid crashes on deletion)
-        def safe_callback(callback_name, event, obj_ref):
-            """Safely invoke a callback on a weakly-referenced object, ignoring ReferenceError if deleted."""
-            try:
-                obj = obj_ref
-                getattr(obj, callback_name)(event)
-            except ReferenceError:
-                # Object has been deleted; ignore.
-                pass
-
-        # note: use weakref on callbacks to ensure that this object can be deleted when its destructor is called.
-        # add callbacks for stage play/stop
-        obj_ref = weakref.proxy(self)
-        timeline_event_stream = omni.timeline.get_timeline_interface().get_timeline_event_stream()
-
-        # the order is set to 10 which is arbitrary but should be lower priority than the default order of 0
-        # register timeline PLAY event callback (lower priority with order=10)
-        self._initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
-            int(omni.timeline.TimelineEventType.PLAY),
-            lambda event, obj_ref=obj_ref: safe_callback("_initialize_callback", event, obj_ref),
-            order=10,
-        )
-        # register timeline STOP event callback (lower priority with order=10)
-        self._invalidate_initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
-            int(omni.timeline.TimelineEventType.STOP),
-            lambda event, obj_ref=obj_ref: safe_callback("_invalidate_initialize_callback", event, obj_ref),
-            order=10,
-        )
-        # register prim deletion callback
-        self._prim_deletion_callback_id = SimulationManager.register_callback(
-            lambda event, obj_ref=obj_ref: safe_callback("_on_prim_deletion", event, obj_ref),
-            event=IsaacEvents.PRIM_DELETION,
-        )
+        # register various callback functions
+        self._register_callbacks()
 
         # add handle for debug visualization (this is set to a valid handle inside set_debug_vis)
         self._debug_vis_handle = None
@@ -280,6 +248,43 @@ class SensorBase(ABC):
     Internal simulation callbacks.
     """
 
+    def _register_callbacks(self):
+        """Registers the timeline and prim deletion callbacks."""
+
+        # register simulator callbacks (with weakref safety to avoid crashes on deletion)
+        def safe_callback(callback_name, event, obj_ref):
+            """Safely invoke a callback on a weakly-referenced object, ignoring ReferenceError if deleted."""
+            try:
+                obj = obj_ref
+                getattr(obj, callback_name)(event)
+            except ReferenceError:
+                # Object has been deleted; ignore.
+                pass
+
+        # note: use weakref on callbacks to ensure that this object can be deleted when its destructor is called.
+        # add callbacks for stage play/stop
+        obj_ref = weakref.proxy(self)
+        timeline_event_stream = omni.timeline.get_timeline_interface().get_timeline_event_stream()
+
+        # the order is set to 10 which is arbitrary but should be lower priority than the default order of 0
+        # register timeline PLAY event callback (lower priority with order=10)
+        self._initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
+            int(omni.timeline.TimelineEventType.PLAY),
+            lambda event, obj_ref=obj_ref: safe_callback("_initialize_callback", event, obj_ref),
+            order=10,
+        )
+        # register timeline STOP event callback (lower priority with order=10)
+        self._invalidate_initialize_handle = timeline_event_stream.create_subscription_to_pop_by_type(
+            int(omni.timeline.TimelineEventType.STOP),
+            lambda event, obj_ref=obj_ref: safe_callback("_invalidate_initialize_callback", event, obj_ref),
+            order=10,
+        )
+        # register prim deletion callback
+        self._prim_deletion_callback_id = SimulationManager.register_callback(
+            lambda event, obj_ref=obj_ref: safe_callback("_on_prim_deletion", event, obj_ref),
+            event=IsaacEvents.PRIM_DELETION,
+        )
+
     def _initialize_callback(self, event):
         """Initializes the scene elements.
 
@@ -311,9 +316,6 @@ class SensorBase(ABC):
         Note:
             This function is called when the prim is deleted.
         """
-        # skip callback if required
-        if getattr(SimulationContext.instance(), "_skip_next_prim_deletion_callback_fn", False):
-            return
         if prim_path == "/":
             self._clear_callbacks()
             return

source/isaaclab/isaaclab/sim/simulation_context.py

@@ -243,10 +243,6 @@ class SimulationContext(_SimulationContext):
             self._app_control_on_stop_handle = None
         self._disable_app_control_on_stop_handle = False
 
-        # flag for skipping prim deletion callback
-        # when stage in memory is attached
-        self._skip_next_prim_deletion_callback_fn = False
-
         # flatten out the simulation dictionary
         sim_params = self.cfg.to_dict()
         if sim_params is not None:

source/isaaclab/isaaclab/sim/utils.py

@@ -844,6 +844,8 @@ def attach_stage_to_usd_context(attaching_early: bool = False):
         attaching_early: Whether to attach the stage to the usd context before stage is created. Defaults to False.
     """
 
+    from isaacsim.core.simulation_manager import SimulationManager
+
     from isaaclab.sim.simulation_context import SimulationContext
 
     # if Isaac Sim version is less than 5.0, stage in memory is not supported
@@ -878,8 +880,8 @@ def attach_stage_to_usd_context(attaching_early: bool = False):
     # skip this callback to avoid wiping the stage after attachment
     SimulationContext.instance().skip_next_stage_open_callback()
 
-    # skip this callback to avoid clearing the prims
-    SimulationContext.instance()._skip_next_prim_deletion_callback_fn = True
+    # disable stage open callback to avoid clearing callbacks
+    SimulationManager.enable_stage_open_callback(False)
 
     # enable physics fabric
     SimulationContext.instance()._physics_context.enable_fabric(True)
@@ -891,7 +893,8 @@ def attach_stage_to_usd_context(attaching_early: bool = False):
     physx_sim_interface = omni.physx.get_physx_simulation_interface()
     physx_sim_interface.attach_stage(stage_id)
 
-    SimulationContext.instance()._skip_next_prim_deletion_callback_fn = False
+    # re-enable stage open callback
+    SimulationManager.enable_stage_open_callback(True)
 
 
 def is_current_stage_in_memory() -> bool:

source/isaaclab_tasks/config/extension.toml

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.10.41"
+version = "0.10.42"
 
 # Description
 title = "Isaac Lab Environments"

source/isaaclab_tasks/docs/CHANGELOG.rst

 Changelog
 ---------
 
+0.10.42 (2025-07-11)
+~~~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+* Organized environment unit tests
+
+
 0.10.41 (2025-07-01)
 ~~~~~~~~~~~~~~~~~~~~
 

source/isaaclab_tasks/test/benchmarking/configs.yaml

@@ -67,14 +67,6 @@ fast:
       episode_length: 900
     upper_thresholds:
       duration: 1800
-  rl_games:Isaac-Repose-Cube-Allegro-Direct-v0:
-    max_iterations: 500
-    lower_thresholds:
-      reward: 200
-      episode_length: 150
-    upper_thresholds:
-      duration: 1500
-
 
 # mode for weekly CI
 full:

source/isaaclab_tasks/test/benchmarking/conftest.py

@@ -5,8 +5,8 @@
 
 import json
 
+import env_benchmark_test_utils as utils
 import pytest
-import test_utils as utils
 
 # Global variable for storing KPI data
 GLOBAL_KPI_STORE = {}

source/isaaclab_tasks/test/benchmarking/test_environments_training.py

@@ -18,8 +18,8 @@ import sys
 import time
 
 import carb
+import env_benchmark_test_utils as utils
 import pytest
-import test_utils as utils
 
 from isaaclab.utils.pretrained_checkpoint import WORKFLOW_EXPERIMENT_NAME_VARIABLE, WORKFLOW_TRAINER
 

source/isaaclab_tasks/test/env_test_utils.py

+# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Shared test utilities for Isaac Lab environments."""
+
+import gymnasium as gym
+import os
+import torch
+
+import carb
+import omni.usd
+import pytest
+from isaacsim.core.version import get_version
+
+from isaaclab.envs.utils.spaces import sample_space
+
+from isaaclab_tasks.utils.parse_cfg import parse_env_cfg
+
+
+def setup_environment(
+    include_play: bool = False,
+    factory_envs: bool | None = None,
+    multi_agent: bool | None = None,
+) -> list[str]:
+    """
+    Acquire all registered Isaac environment task IDs with optional filters.
+
+    Args:
+        include_play: If True, include environments ending in 'Play-v0'.
+        factory_envs:
+            - True: include only Factory environments
+            - False: exclude Factory environments
+            - None: include both Factory and non-Factory environments
+        multi_agent:
+            - True: include only multi-agent environments
+            - False: include only single-agent environments
+            - None: include all environments regardless of agent type
+
+    Returns:
+        A sorted list of task IDs matching the selected filters.
+    """
+    # disable interactive mode for wandb for automate environments
+    os.environ["WANDB_DISABLED"] = "true"
+
+    # acquire all Isaac environment names
+    registered_tasks = []
+    for task_spec in gym.registry.values():
+        # only consider Isaac environments
+        if "Isaac" not in task_spec.id:
+            continue
+
+        # filter Play environments, if needed
+        if not include_play and task_spec.id.endswith("Play-v0"):
+            continue
+
+        # TODO: factory environments cause tests to fail if run together with other envs,
+        # so we collect these environments separately to run in a separate unit test.
+        # apply factory filter
+        if (factory_envs is True and "Factory" not in task_spec.id) or (
+            factory_envs is False and "Factory" in task_spec.id
+        ):
+            continue
+        # if None: no filter
+
+        # apply multi agent filter
+        if multi_agent is not None:
+            # parse config
+            env_cfg = parse_env_cfg(task_spec.id)
+            if (multi_agent is True and not hasattr(env_cfg, "possible_agents")) or (
+                multi_agent is False and hasattr(env_cfg, "possible_agents")
+            ):
+                continue
+        # if None: no filter
+
+        registered_tasks.append(task_spec.id)
+
+    # sort environments alphabetically
+    registered_tasks.sort()
+
+    # this flag is necessary to prevent a bug where the simulation gets stuck randomy when running many environments
+    carb.settings.get_settings().set_bool("/physics/cooking/ujitsoCollisionCooking", False)
+
+    print(">>> All registered environments:", registered_tasks)
+
+    return registered_tasks
+
+
+def _run_environments(
+    task_name,
+    device,
+    num_envs,
+    num_steps=100,
+    multi_agent=False,
+    create_stage_in_memory=False,
+    disable_clone_in_fabric=False,
+):
+    """Run all environments and check environments return valid signals.
+
+    Args:
+        task_name: Name of the environment.
+        device: Device to use (e.g., 'cuda').
+        num_envs: Number of environments.
+        num_steps: Number of simulation steps.
+        multi_agent: Whether the environment is multi-agent.
+        create_stage_in_memory: Whether to create stage in memory.
+        disable_clone_in_fabric: Whether to disable fabric cloning.
+    """
+
+    # skip test if stage in memory is not supported
+    isaac_sim_version = float(".".join(get_version()[2]))
+    if isaac_sim_version < 5 and create_stage_in_memory:
+        pytest.skip("Stage in memory is not supported in this version of Isaac Sim")
+
+    # skip these environments as they cannot be run with 32 environments within reasonable VRAM
+    if num_envs == 32 and task_name in [
+        "Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-v0",
+        "Isaac-Stack-Cube-Instance-Randomize-Franka-IK-Rel-v0",
+        "Isaac-Stack-Cube-Instance-Randomize-Franka-v0",
+        "Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-v0",
+        "Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos-v0",
+    ]:
+        return
+
+    # skip automate environments as they require cuda installation
+    if task_name in ["Isaac-AutoMate-Assembly-Direct-v0", "Isaac-AutoMate-Disassembly-Direct-v0"]:
+        return
+
+    # skipping this test for now as it requires torch 2.6 or newer
+    if task_name == "Isaac-Cartpole-RGB-TheiaTiny-v0":
+        return
+
+    # TODO: why is this failing in Isaac Sim 5.0??? but the environment itself can run.
+    if task_name == "Isaac-Lift-Teddy-Bear-Franka-IK-Abs-v0":
+        return
+
+    print(f""">>> Running test for environment: {task_name}""")
+    _check_random_actions(
+        task_name,
+        device,
+        num_envs,
+        num_steps=num_steps,
+        multi_agent=multi_agent,
+        create_stage_in_memory=create_stage_in_memory,
+        disable_clone_in_fabric=disable_clone_in_fabric,
+    )
+    print(f""">>> Closing environment: {task_name}""")
+    print("-" * 80)
+
+
+def _check_random_actions(
+    task_name: str,
+    device: str,
+    num_envs: int,
+    num_steps: int = 100,
+    multi_agent: bool = False,
+    create_stage_in_memory: bool = False,
+    disable_clone_in_fabric: bool = False,
+):
+    """Run random actions and check environments return valid signals.
+
+    Args:
+        task_name: Name of the environment.
+        device: Device to use (e.g., 'cuda').
+        num_envs: Number of environments.
+        num_steps: Number of simulation steps.
+        multi_agent: Whether the environment is multi-agent.
+        create_stage_in_memory: Whether to create stage in memory.
+        disable_clone_in_fabric: Whether to disable fabric cloning.
+    """
+    # create a new context stage, if stage in memory is not enabled
+    if not create_stage_in_memory:
+        omni.usd.get_context().new_stage()
+
+    # reset the rtx sensors carb setting to False
+    carb.settings.get_settings().set_bool("/isaaclab/render/rtx_sensors", False)
+    try:
+        # parse config
+        env_cfg = parse_env_cfg(task_name, device=device, num_envs=num_envs)
+        # set config args
+        env_cfg.sim.create_stage_in_memory = create_stage_in_memory
+        if disable_clone_in_fabric:
+            env_cfg.scene.clone_in_fabric = False
+
+        # filter based off multi agents mode and create env
+        if multi_agent:
+            if not hasattr(env_cfg, "possible_agents"):
+                print(f"[INFO]: Skipping {task_name} as it is not a multi-agent task")
+                return
+            env = gym.make(task_name, cfg=env_cfg)
+        else:
+            if hasattr(env_cfg, "possible_agents"):
+                print(f"[INFO]: Skipping {task_name} as it is a multi-agent task")
+                return
+            env = gym.make(task_name, cfg=env_cfg)
+
+    except Exception as e:
+        # try to close environment on exception
+        if "env" in locals() and hasattr(env, "_is_closed"):
+            env.close()
+        else:
+            if hasattr(e, "obj") and hasattr(e.obj, "_is_closed"):
+                e.obj.close()
+        pytest.fail(f"Failed to set-up the environment for task {task_name}. Error: {e}")
+
+    # disable control on stop
+    env.unwrapped.sim._app_control_on_stop_handle = None  # type: ignore
+
+    # override action space if set to inf for `Isaac-Lift-Teddy-Bear-Franka-IK-Abs-v0`
+    if task_name == "Isaac-Lift-Teddy-Bear-Franka-IK-Abs-v0":
+        for i in range(env.unwrapped.single_action_space.shape[0]):
+            if env.unwrapped.single_action_space.low[i] == float("-inf"):
+                env.unwrapped.single_action_space.low[i] = -1.0
+            if env.unwrapped.single_action_space.high[i] == float("inf"):
+                env.unwrapped.single_action_space.low[i] = 1.0
+
+    # reset environment
+    obs, _ = env.reset()
+
+    # check signal
+    assert _check_valid_tensor(obs)
+
+    # simulate environment for num_steps
+    with torch.inference_mode():
+        for _ in range(num_steps):
+            # sample actions according to the defined space
+            if multi_agent:
+                actions = {
+                    agent: sample_space(
+                        env.unwrapped.action_spaces[agent], device=env.unwrapped.device, batch_size=num_envs
+                    )
+                    for agent in env.unwrapped.possible_agents
+                }
+            else:
+                actions = sample_space(
+                    env.unwrapped.single_action_space, device=env.unwrapped.device, batch_size=num_envs
+                )
+            # apply actions
+            transition = env.step(actions)
+            # check signals
+            for data in transition[:-1]:  # exclude info
+                if multi_agent:
+                    for agent, agent_data in data.items():
+                        assert _check_valid_tensor(agent_data), f"Invalid data ('{agent}'): {agent_data}"
+                else:
+                    assert _check_valid_tensor(data), f"Invalid data: {data}"
+
+    # close environment
+    env.close()
+
+
+def _check_valid_tensor(data: torch.Tensor | dict) -> bool:
+    """Checks if given data does not have corrupted values.
+
+    Args:
+        data: Data buffer.
+
+    Returns:
+        True if the data is valid.
+    """
+    if isinstance(data, torch.Tensor):
+        return not torch.any(torch.isnan(data))
+    elif isinstance(data, (tuple, list)):
+        return all(_check_valid_tensor(value) for value in data)
+    elif isinstance(data, dict):
+        return all(_check_valid_tensor(value) for value in data.values())
+    else:
+        raise ValueError(f"Input data of invalid type: {type(data)}.")

source/isaaclab_tasks/test/test_environments.py

@@ -7,9 +7,6 @@
 
 import sys
 
-# Omniverse logger
-import omni.log
-
 # Import pinocchio in the main script to force the use of the dependencies installed by IsaacLab and not the one installed by Isaac Sim
 # pinocchio is required by the Pink IK controller
 if sys.platform != "win32":
@@ -24,170 +21,14 @@ simulation_app = app_launcher.app
 
 """Rest everything follows."""
 
-import gymnasium as gym
-import os
-import torch
-
-import carb
-import omni.usd
 import pytest
-from isaacsim.core.version import get_version
-
-from isaaclab.envs import ManagerBasedRLEnvCfg
-from isaaclab.envs.utils.spaces import sample_space
+from env_test_utils import _run_environments, setup_environment
 
 import isaaclab_tasks  # noqa: F401
-from isaaclab_tasks.utils.parse_cfg import parse_env_cfg
-
-
-# @pytest.fixture(scope="module", autouse=True)
-def setup_environment():
-    # disable interactive mode for wandb for automate environments
-    os.environ["WANDB_DISABLED"] = "true"
-    # acquire all Isaac environments names
-    registered_tasks = list()
-    for task_spec in gym.registry.values():
-        # TODO: Factory environments causes test to fail if run together with other envs
-        if "Isaac" in task_spec.id and not task_spec.id.endswith("Play-v0") and "Factory" not in task_spec.id:
-            registered_tasks.append(task_spec.id)
-    # sort environments by name
-    registered_tasks.sort()
-    # this flag is necessary to prevent a bug where the simulation gets stuck randomly when running the
-    # test on many environments.
-    carb_settings_iface = carb.settings.get_settings()
-    carb_settings_iface.set_bool("/physics/cooking/ujitsoCollisionCooking", False)
-
-    return registered_tasks
-
-
-# note, running an env test without stage in memory then
-# running an env test with stage in memory causes IsaacLab to hang.
-# so, here we run all envs with stage in memory first, then run
-# all envs without stage in memory.
-@pytest.mark.order(1)
-@pytest.mark.parametrize("num_envs, device", [(32, "cuda"), (1, "cuda")])
-@pytest.mark.parametrize("task_name", setup_environment())
-def test_environments_with_stage_in_memory(task_name, num_envs, device):
-    # run environments with stage in memory
-    _run_environments(task_name, device, num_envs, num_steps=100, create_stage_in_memory=True)
 
 
-@pytest.mark.order(2)
 @pytest.mark.parametrize("num_envs, device", [(32, "cuda"), (1, "cuda")])
-@pytest.mark.parametrize("task_name", setup_environment())
+@pytest.mark.parametrize("task_name", setup_environment(include_play=False, factory_envs=False, multi_agent=False))
 def test_environments(task_name, num_envs, device):
     # run environments without stage in memory
-    _run_environments(task_name, device, num_envs, num_steps=100, create_stage_in_memory=False)
-
-
-def _run_environments(task_name, device, num_envs, num_steps, create_stage_in_memory):
-    """Run all environments and check environments return valid signals."""
-
-    # skip test if stage in memory is not supported
-    isaac_sim_version = float(".".join(get_version()[2]))
-    if isaac_sim_version < 5 and create_stage_in_memory:
-        pytest.skip("Stage in memory is not supported in this version of Isaac Sim")
-
-    # skip these environments as they cannot be run with 32 environments within reasonable VRAM
-    if num_envs == 32 and task_name in [
-        "Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-v0",
-        "Isaac-Stack-Cube-Instance-Randomize-Franka-IK-Rel-v0",
-        "Isaac-Stack-Cube-Instance-Randomize-Franka-v0",
-        "Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-v0",
-        "Isaac-Stack-Cube-Franka-IK-Rel-Visuomotor-Cosmos-v0",
-    ]:
-        return
-    # skip automate environments as they require cuda installation
-    if task_name in ["Isaac-AutoMate-Assembly-Direct-v0", "Isaac-AutoMate-Disassembly-Direct-v0"]:
-        return
-    # skipping this test for now as it requires torch 2.6 or newer
-
-    if task_name == "Isaac-Cartpole-RGB-TheiaTiny-v0":
-        return
-    # TODO: why is this failing in Isaac Sim 5.0??? but the environment itself can run.
-    if task_name == "Isaac-Lift-Teddy-Bear-Franka-IK-Abs-v0":
-        return
-    print(f">>> Running test for environment: {task_name}")
-    _check_random_actions(task_name, device, num_envs, num_steps=100, create_stage_in_memory=create_stage_in_memory)
-    print(f">>> Closing environment: {task_name}")
-    print("-" * 80)
-
-
-def _check_random_actions(
-    task_name: str, device: str, num_envs: int, num_steps: int = 1000, create_stage_in_memory: bool = False
-):
-    """Run random actions and check environments returned signals are valid."""
-
-    if not create_stage_in_memory:
-        # create a new context stage
-        omni.usd.get_context().new_stage()
-
-    # reset the rtx sensors carb setting to False
-    carb.settings.get_settings().set_bool("/isaaclab/render/rtx_sensors", False)
-    try:
-        # parse configuration
-        env_cfg: ManagerBasedRLEnvCfg = parse_env_cfg(task_name, device=device, num_envs=num_envs)
-        env_cfg.sim.create_stage_in_memory = create_stage_in_memory
-
-        # skip test if the environment is a multi-agent task
-        if hasattr(env_cfg, "possible_agents"):
-            print(f"[INFO]: Skipping {task_name} as it is a multi-agent task")
-            return
-
-        # create environment
-        env = gym.make(task_name, cfg=env_cfg)
-    except Exception as e:
-        if "env" in locals() and hasattr(env, "_is_closed"):
-            env.close()
-        else:
-            if hasattr(e, "obj") and hasattr(e.obj, "_is_closed"):
-                e.obj.close()
-        pytest.fail(f"Failed to set-up the environment for task {task_name}. Error: {e}")
-
-    # disable control on stop
-    env.unwrapped.sim._app_control_on_stop_handle = None  # type: ignore
-
-    # override action space if set to inf for `Isaac-Lift-Teddy-Bear-Franka-IK-Abs-v0`
-    if task_name == "Isaac-Lift-Teddy-Bear-Franka-IK-Abs-v0":
-        for i in range(env.unwrapped.single_action_space.shape[0]):
-            if env.unwrapped.single_action_space.low[i] == float("-inf"):
-                env.unwrapped.single_action_space.low[i] = -1.0
-            if env.unwrapped.single_action_space.high[i] == float("inf"):
-                env.unwrapped.single_action_space.low[i] = 1.0
-
-    # reset environment
-    obs, _ = env.reset()
-    # check signal
-    assert _check_valid_tensor(obs)
-    # simulate environment for num_steps steps
-    with torch.inference_mode():
-        for _ in range(num_steps):
-            # sample actions according to the defined space
-            actions = sample_space(env.unwrapped.single_action_space, device=env.unwrapped.device, batch_size=num_envs)
-            # apply actions
-            transition = env.step(actions)
-            # check signals
-            for data in transition[:-1]:  # exclude info
-                assert _check_valid_tensor(data), f"Invalid data: {data}"
-
-    # close the environment
-    env.close()
-
-
-def _check_valid_tensor(data: torch.Tensor | dict) -> bool:
-    """Checks if given data does not have corrupted values.
-
-    Args:
-        data: Data buffer.
-
-    Returns:
-        True if the data is valid.
-    """
-    if isinstance(data, torch.Tensor):
-        return not torch.any(torch.isnan(data))
-    elif isinstance(data, (tuple, list)):
-        return all(_check_valid_tensor(value) for value in data)
-    elif isinstance(data, dict):
-        return all(_check_valid_tensor(value) for value in data.values())
-    else:
-        raise ValueError(f"Input data of invalid type: {type(data)}.")
+    _run_environments(task_name, device, num_envs, create_stage_in_memory=False)

source/isaaclab_tasks/test/test_environments_with_stage_in_memory.py

+# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Launch Isaac Sim Simulator first."""
+
+import sys
+
+# Import pinocchio in the main script to force the use of the dependencies installed by IsaacLab and not the one installed by Isaac Sim
+# pinocchio is required by the Pink IK controller
+if sys.platform != "win32":
+    import pinocchio  # noqa: F401
+
+from isaaclab.app import AppLauncher
+
+# launch the simulator
+app_launcher = AppLauncher(headless=True, enable_cameras=True)
+simulation_app = app_launcher.app
+
+from isaacsim.core.version import get_version
+
+"""Rest everything follows."""
+
+import pytest
+from env_test_utils import _run_environments, setup_environment
+
+import isaaclab_tasks  # noqa: F401
+
+# note, running an env test without stage in memory then
+# running an env test with stage in memory causes IsaacLab to hang.
+# so, here we run all envs with stage in memory separately
+
+# TODO(mtrepte): re-enable with fabric cloning fix
+# @pytest.mark.parametrize("num_envs, device", [(2, "cuda")])
+# @pytest.mark.parametrize("task_name", setup_environment(include_play=False,factory_envs=False, multi_agent=False))
+# def test_environments_with_stage_in_memory_and_clone_in_fabric_disabled(task_name, num_envs, device):
+# # skip test if stage in memory is not supported
+# isaac_sim_version = float(".".join(get_version()[2]))
+# if isaac_sim_version < 5:
+#     pytest.skip("Stage in memory is not supported in this version of Isaac Sim")
+
+# # run environments with stage in memory
+# _run_environments(task_name, device, num_envs, create_stage_in_memory=True)
+
+
+@pytest.mark.parametrize("num_envs, device", [(2, "cuda")])
+@pytest.mark.parametrize("task_name", setup_environment(include_play=False, factory_envs=False, multi_agent=False))
+def test_environments_with_stage_in_memory_and_clone_in_fabric_disabled(task_name, num_envs, device):
+    # skip test if stage in memory is not supported
+    isaac_sim_version = float(".".join(get_version()[2]))
+    if isaac_sim_version < 5:
+        pytest.skip("Stage in memory is not supported in this version of Isaac Sim")
+
+    # run environments with stage in memory
+    _run_environments(task_name, device, num_envs, create_stage_in_memory=True, disable_clone_in_fabric=True)

source/isaaclab_tasks/test/test_factory_environments.py

@@ -11,119 +11,19 @@ from isaaclab.app import AppLauncher
 app_launcher = AppLauncher(headless=True, enable_cameras=True)
 simulation_app = app_launcher.app
 
-
 """Rest everything follows."""
 
-import gymnasium as gym
-import torch
-
-import carb
-import omni.usd
 import pytest
-
-from isaaclab.envs import ManagerBasedRLEnvCfg
-from isaaclab.envs.utils.spaces import sample_space
+from env_test_utils import _check_random_actions, setup_environment
 
 import isaaclab_tasks  # noqa: F401
-from isaaclab_tasks.utils.parse_cfg import parse_env_cfg
-
-
-def setup_environment():
-    # acquire all Isaac environments names
-    registered_tasks = list()
-    for task_spec in gym.registry.values():
-        if "Isaac" in task_spec.id and not task_spec.id.endswith("Play-v0") and "Factory" in task_spec.id:
-            registered_tasks.append(task_spec.id)
-    # sort environments by name
-    registered_tasks.sort()
-
-    # this flag is necessary to prevent a bug where the simulation gets stuck randomly when running the
-    # test on many environments.
-    carb_settings_iface = carb.settings.get_settings()
-    carb_settings_iface.set_bool("/physics/cooking/ujitsoCollisionCooking", False)
-
-    return registered_tasks
 
 
 @pytest.mark.parametrize("num_envs, device", [(32, "cuda"), (1, "cuda")])
-@pytest.mark.parametrize("task_name", setup_environment())
+@pytest.mark.parametrize("task_name", setup_environment(factory_envs=True, multi_agent=False))
 def test_factory_environments(task_name, num_envs, device):
     """Run all factory environments and check environments return valid signals."""
     print(f">>> Running test for environment: {task_name}")
-    _check_random_actions(task_name, device, num_envs, num_steps=100)
+    _check_random_actions(task_name, device, num_envs)
     print(f">>> Closing environment: {task_name}")
     print("-" * 80)
-
-
-def _check_random_actions(task_name: str, device: str, num_envs: int, num_steps: int = 1000):
-    """Run random actions and check environments returned signals are valid."""
-    # create a new stage
-    omni.usd.get_context().new_stage()
-    # reset the rtx sensors carb setting to False
-    carb.settings.get_settings().set_bool("/isaaclab/render/rtx_sensors", False)
-    try:
-        # parse configuration
-        env_cfg: ManagerBasedRLEnvCfg = parse_env_cfg(task_name, device=device, num_envs=num_envs)
-
-        # skip test if the environment is a multi-agent task
-        if hasattr(env_cfg, "possible_agents"):
-            print(f"[INFO]: Skipping {task_name} as it is a multi-agent task")
-            return
-
-        # create environment
-        env = gym.make(task_name, cfg=env_cfg)
-    except Exception as e:
-        if "env" in locals() and hasattr(env, "_is_closed"):
-            env.close()
-        else:
-            if hasattr(e, "obj") and hasattr(e.obj, "_is_closed"):
-                e.obj.close()
-        pytest.fail(f"Failed to set-up the environment for task {task_name}. Error: {e}")
-
-    # disable control on stop
-    env.unwrapped.sim._app_control_on_stop_handle = None  # type: ignore
-
-    # override action space if set to inf for `Isaac-Lift-Teddy-Bear-Franka-IK-Abs-v0`
-    if task_name == "Isaac-Lift-Teddy-Bear-Franka-IK-Abs-v0":
-        for i in range(env.unwrapped.single_action_space.shape[0]):
-            if env.unwrapped.single_action_space.low[i] == float("-inf"):
-                env.unwrapped.single_action_space.low[i] = -1.0
-            if env.unwrapped.single_action_space.high[i] == float("inf"):
-                env.unwrapped.single_action_space.low[i] = 1.0
-
-    # reset environment
-    obs, _ = env.reset()
-    # check signal
-    assert _check_valid_tensor(obs)
-    # simulate environment for num_steps steps
-    with torch.inference_mode():
-        for _ in range(num_steps):
-            # sample actions according to the defined space
-            actions = sample_space(env.unwrapped.single_action_space, device=env.unwrapped.device, batch_size=num_envs)
-            # apply actions
-            transition = env.step(actions)
-            # check signals
-            for data in transition[:-1]:  # exclude info
-                assert _check_valid_tensor(data), f"Invalid data: {data}"
-
-    # close the environment
-    env.close()
-
-
-def _check_valid_tensor(data: torch.Tensor | dict) -> bool:
-    """Checks if given data does not have corrupted values.
-
-    Args:
-        data: Data buffer.
-
-    Returns:
-        True if the data is valid.
-    """
-    if isinstance(data, torch.Tensor):
-        return not torch.any(torch.isnan(data))
-    elif isinstance(data, (tuple, list)):
-        return all(_check_valid_tensor(value) for value in data)
-    elif isinstance(data, dict):
-        return all(_check_valid_tensor(value) for value in data.values())
-    else:
-        raise ValueError(f"Input data of invalid type: {type(data)}.")

source/isaaclab_tasks/test/test_multi_agent_environments.py

@@ -14,114 +14,19 @@ simulation_app = app_launcher.app
 
 """Rest everything follows."""
 
-import gymnasium as gym
-import torch
-
-import omni.usd
 import pytest
-
-from isaaclab.envs import DirectMARLEnv, DirectMARLEnvCfg
-from isaaclab.envs.utils.spaces import sample_space
+from env_test_utils import _check_random_actions, setup_environment
 
 import isaaclab_tasks  # noqa: F401
-from isaaclab_tasks.utils.parse_cfg import parse_env_cfg
-
-
-# @pytest.fixture(scope="module", autouse=True)
-def setup_environment():
-    # acquire all Isaac environments names
-    registered_tasks = list()
-    for task_spec in gym.registry.values():
-        if "Isaac" in task_spec.id and not task_spec.id.endswith("Play-v0"):
-            registered_tasks.append(task_spec.id)
-    # sort environments by name
-    registered_tasks.sort()
-    # print all existing task names
-    print(">>> All registered environments:", registered_tasks)
-    return registered_tasks
 
 
 @pytest.mark.parametrize("num_envs, device", [(32, "cuda"), (1, "cuda")])
-@pytest.mark.parametrize("task_name", setup_environment())
+@pytest.mark.parametrize("task_name", setup_environment(multi_agent=True))
 def test_environments(task_name, num_envs, device):
     """Run all environments with given parameters and check environments return valid signals."""
     print(f">>> Running test for environment: {task_name} with num_envs={num_envs} and device={device}")
     # check environment
-    _check_random_actions(task_name, device, num_envs, num_steps=100)
+    _check_random_actions(task_name, device, num_envs, multi_agent=True)
     # close the environment
     print(f">>> Closing environment: {task_name}")
     print("-" * 80)
-
-
-def _check_random_actions(task_name: str, device: str, num_envs: int, num_steps: int = 1000):
-    """Run random actions and check environments returned signals are valid."""
-    # create a new stage
-    omni.usd.get_context().new_stage()
-    try:
-        # parse configuration
-        env_cfg: DirectMARLEnvCfg = parse_env_cfg(task_name, device=device, num_envs=num_envs)
-
-        # skip test if the environment is not a multi-agent task
-        if not hasattr(env_cfg, "possible_agents"):
-            print(f"[INFO]: Skipping {task_name} as it is not a multi-agent task")
-            return
-
-        # create environment
-        env: DirectMARLEnv = gym.make(task_name, cfg=env_cfg)
-    except Exception as e:
-        if "env" in locals() and hasattr(env, "_is_closed"):
-            env.close()
-        else:
-            if hasattr(e, "obj") and hasattr(e.obj, "_is_closed"):
-                e.obj.close()
-        pytest.fail(f"Failed to set-up the environment for task {task_name}. Error: {e}")
-
-    # this flag is necessary to prevent a bug where the simulation gets stuck randomly when running the
-    # test on many environments.
-    env.unwrapped.sim.set_setting("/physics/cooking/ujitsoCollisionCooking", False)
-
-    # avoid shutdown of process on simulation stop
-    env.unwrapped.sim._app_control_on_stop_handle = None
-
-    # reset environment
-    obs, _ = env.reset()
-    # check signal
-    assert _check_valid_tensor(obs)
-    # simulate environment for num_steps steps
-    with torch.inference_mode():
-        for _ in range(num_steps):
-            # sample actions according to the defined space
-            actions = {
-                agent: sample_space(
-                    env.unwrapped.action_spaces[agent], device=env.unwrapped.device, batch_size=num_envs
-                )
-                for agent in env.unwrapped.possible_agents
-            }
-            # apply actions
-            transition = env.step(actions)
-            # check signals
-            for item in transition[:-1]:  # exclude info
-                for agent, data in item.items():
-                    assert _check_valid_tensor(data), f"Invalid data ('{agent}'): {data}"
-
-    # close the environment
-    env.close()
-
-
-def _check_valid_tensor(data: torch.Tensor | dict) -> bool:
-    """Checks if given data does not have corrupted values.
-
-    Args:
-        data: Data buffer.
-
-    Returns:
-        True if the data is valid.
-    """
-    if isinstance(data, torch.Tensor):
-        return not torch.any(torch.isnan(data))
-    elif isinstance(data, (tuple, list)):
-        return all(_check_valid_tensor(value) for value in data)
-    elif isinstance(data, dict):
-        return all(_check_valid_tensor(value) for value in data.values())
-    else:
-        raise ValueError(f"Input data of invalid type: {type(data)}.")

source/isaaclab_tasks/test/test_record_video.py

@@ -19,25 +19,12 @@ import torch
 
 import omni.usd
 import pytest
+from env_test_utils import setup_environment
 
 import isaaclab_tasks  # noqa: F401
 from isaaclab_tasks.utils import parse_env_cfg
 
 
-# @pytest.fixture(scope="module", autouse=True)
-def setup_environment():
-    # acquire all Isaac environments names
-    registered_tasks = list()
-    for task_spec in gym.registry.values():
-        if "Isaac" in task_spec.id:
-            registered_tasks.append(task_spec.id)
-    # sort environments by name
-    registered_tasks.sort()
-    # print all existing task names
-    print(">>> All registered environments:", registered_tasks)
-    return registered_tasks
-
-
 @pytest.fixture(scope="function")
 def setup_video_params():
     # common parameters
@@ -49,7 +36,7 @@ def setup_video_params():
     return num_envs, device, step_trigger, video_length
 
 
-@pytest.mark.parametrize("task_name", setup_environment())
+@pytest.mark.parametrize("task_name", setup_environment(include_play=True))
 def test_record_video(task_name, setup_video_params):
     """Run random actions agent with recording of videos."""
     num_envs, device, step_trigger, video_length = setup_video_params

tools/test_settings.py

@@ -18,7 +18,11 @@ DEFAULT_TIMEOUT = 200
 PER_TEST_TIMEOUTS = {
     "test_articulation.py": 500,
     "test_rigid_object.py": 300,
-    "test_environments.py": 1500,  # This test runs through all the environments for 100 steps each
+    "test_stage_in_memory.py": 500,
+    "test_environments.py": 2000,  # This test runs through all the environments for 100 steps each
+    "test_environments_with_stage_in_memory.py": (
+        1000
+    ),  # Like the above, with stage in memory and with and without fabric cloning
     "test_environment_determinism.py": 500,  # This test runs through many the environments for 100 steps each
     "test_factory_environments.py": 300,  # This test runs through Factory environments for 100 steps each
     "test_env_rendering_logic.py": 300,