remove attach_stage_to_usd_context, and simply attach in sim context

Author: ooctipus

docs/source/how-to/optimize_stage_creation.rst

@@ -47,25 +47,23 @@ Stage in memory can be toggled by setting the :attr:`isaaclab.sim.SimulationCfg.
     # create env with stage in memory
     env = ManagerBasedRLEnv(cfg=cfg)
 
-Note, if stage in memory is enabled without using an existing RL environment class, a few more steps are need.
-The stage creation steps should be wrapped in a :py:keyword:`with` statement to set the stage context.
-If the stage needs to be attached, the :meth:`~isaaclab.sim.utils.attach_stage_to_usd_context` function should
-be called after the stage is created.
+When using stage in memory without an existing RL environment class, wrap the stage creation steps
+in a :py:keyword:`with` statement to set the stage context. The stage is automatically attached
+to the USD context when ``SimulationContext`` is created with ``create_stage_in_memory=True``.
 
 **Using Stage in Memory with a manual scene setup**
 
 .. code-block:: python
 
     # init simulation context with stage in memory
+    # Note: stage is automatically attached to USD context
     sim = SimulationContext(cfg=SimulationCfg(create_stage_in_memory=True))
 
-    # grab stage in memory and set stage context
+    # grab stage and set stage context
     with stage_utils.use_stage(sim.stage):
         # create cartpole scene
         scene_cfg = CartpoleSceneCfg(num_envs=1024)
         scene = InteractiveScene(scene_cfg)
-        # attach stage to memory after stage is created
-        sim_utils.attach_stage_to_usd_context()
 
     sim.play()
 
@@ -117,13 +115,11 @@ Limitations
 
 - Cannot be currently enabled at the same time as **Fabric Cloning**.
 
-- Attaching stage in memory to the USD context can be slow, offsetting some or all of the performance benefits.
-
-  - Note, attaching is only necessary when rendering is enabled. For example, in headless mode, attachment is not required.
+- The stage is automatically attached to the USD context at ``SimulationContext`` creation, ensuring proper
+  lifecycle events for viewport and physics systems.
 
 - Certain low-level Kit APIs do not yet support stage in memory.
 
-  - In most cases, when these APIs are hit, existing scripts will automatically early attach the stage and print a warning message.
   - In one particular case, for some environments, the API call to color the ground plane is skipped, when stage in memory is enabled.
 
 

source/isaaclab/isaaclab/envs/direct_marl_env.py

@@ -25,7 +25,7 @@
 from isaaclab.managers import EventManager
 from isaaclab.scene import InteractiveScene
 from isaaclab.sim import SimulationContext
-from isaaclab.sim.utils.stage import attach_stage_to_usd_context, use_stage
+from isaaclab.sim.utils.stage import use_stage
 from isaaclab.utils.noise import NoiseModel
 from isaaclab.utils.seed import configure_seed
 from isaaclab.utils.timer import Timer
@@ -125,7 +125,6 @@ def __init__(self, cfg: DirectMARLEnvCfg, render_mode: str | None = None, **kwar
             with use_stage(self.sim.stage):
                 self.scene = InteractiveScene(self.cfg.scene)
                 self._setup_scene()
-                attach_stage_to_usd_context()
         print("[INFO]: Scene manager: ", self.scene)
 
         # set up camera viewport controller

source/isaaclab/isaaclab/envs/direct_rl_env.py

@@ -26,7 +26,7 @@
 from isaaclab.managers import EventManager
 from isaaclab.scene import InteractiveScene
 from isaaclab.sim import SimulationContext
-from isaaclab.sim.utils.stage import attach_stage_to_usd_context, use_stage
+from isaaclab.sim.utils.stage import use_stage
 from isaaclab.utils.noise import NoiseModel
 from isaaclab.utils.seed import configure_seed
 from isaaclab.utils.timer import Timer
@@ -131,7 +131,6 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs
             with use_stage(self.sim.stage):
                 self.scene = InteractiveScene(self.cfg.scene)
                 self._setup_scene()
-                attach_stage_to_usd_context()
         print("[INFO]: Scene manager: ", self.scene)
 
         # set up camera viewport controller

source/isaaclab/isaaclab/envs/manager_based_env.py

@@ -16,7 +16,7 @@
 from isaaclab.managers import ActionManager, EventManager, ObservationManager, RecorderManager
 from isaaclab.scene import InteractiveScene
 from isaaclab.sim import SimulationContext
-from isaaclab.sim.utils.stage import attach_stage_to_usd_context, use_stage
+from isaaclab.sim.utils.stage import use_stage
 from isaaclab.ui.widgets import ManagerLiveVisualizer
 from isaaclab.utils.seed import configure_seed
 from isaaclab.utils.timer import Timer
@@ -138,7 +138,6 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
             # set the stage context for scene creation steps which use the stage
             with use_stage(self.sim.stage):
                 self.scene = InteractiveScene(self.cfg.scene)
-                attach_stage_to_usd_context()
         print("[INFO]: Scene manager: ", self.scene)
 
         # set up camera viewport controller

source/isaaclab/isaaclab/sim/utils/stage.py

@@ -93,9 +93,9 @@ def get_context_stage() -> Usd.Stage | None:
     * Most Isaac Sim/Omniverse systems operate on by default
 
     This is different from an "in-memory stage" created via
-    :func:`create_new_stage_in_memory`, which exists only in RAM and is not
-    attached to the context (invisible to viewport/UI until explicitly attached
-    via :func:`attach_stage_to_usd_context`).
+    :func:`create_new_stage_in_memory`, which exists only in RAM. Note that when
+    using ``SimulationCfg(create_stage_in_memory=True)``, the in-memory stage is
+    automatically attached to the USD context at ``SimulationContext`` creation.
 
     Returns:
         The stage attached to the USD context, or None if no stage is attached.
@@ -513,56 +513,3 @@ def get_current_stage_id() -> int:
         stage_id = stage_cache.Insert(stage).ToLongInt()
     # return stage ID
     return stage_id
-
-
-def attach_stage_to_usd_context(attaching_early: bool = False):
-    """Attaches the current USD stage in memory to the USD context.
-
-    This function should be called during or after scene is created and before stage is simulated or rendered.
-    If the stage is not in memory or rendering is not enabled, this function will return without attaching.
-
-    .. versionadded:: 2.3.0
-        This function is available in Isaac Sim 5.0 and later. For backwards
-        compatibility, it returns without attaching to the USD context.
-
-    Args:
-        attaching_early: Whether to attach the stage to the usd context before stage is created. Defaults to False.
-    """
-
-    import carb
-    import omni.physx
-    import omni.usd
-
-    from isaaclab.sim.simulation_context import SimulationContext
-
-    # if Isaac Sim version is less than 5.0, stage in memory is not supported
-    if get_isaac_sim_version().major < 5:
-        return
-
-    # if stage is not in memory, we can return early
-    if not is_current_stage_in_memory():
-        return
-
-    # attach stage to physx
-    stage_id = get_current_stage_id()
-    physx_sim_interface = omni.physx.get_physx_simulation_interface()
-    physx_sim_interface.attach_stage(stage_id)
-
-    # this carb flag is equivalent to if rendering is enabled
-    carb_setting = carb.settings.get_settings()  # type: ignore
-    is_rendering_enabled = carb_setting.get("/physics/fabricUpdateTransformations")
-
-    # if rendering is not enabled, we don't need to attach it
-    if not is_rendering_enabled:
-        return
-
-    # early attach warning msg
-    if attaching_early:
-        logger.warning(
-            "Attaching stage in memory to USD context early to support an operation which"
-            " does not support stage in memory."
-        )
-
-    # Enable physics fabric and attach stage to usd context for rendering
-    SimulationContext.instance().set_setting("/isaaclab/fabric_enabled", True)
-    omni.usd.get_context().attach_stage_with_callback(stage_id)

source/isaaclab/test/sim/test_simulation_stage_in_memory.py

@@ -55,7 +55,10 @@ def test_stage_in_memory_with_shapes(sim):
     # define parameters
     num_clones = 10
 
-    # grab stage in memory and set as current stage via the with statement
+    # verify stage is attached to USD context (happens automatically now with create_stage_in_memory)
+    assert not sim_utils.is_current_stage_in_memory()
+
+    # grab stage and set as current stage via the with statement
     stage_in_memory = sim.stage
     with sim_utils.use_stage(stage_in_memory):
         # create cloned cone stage
@@ -108,27 +111,11 @@ def test_stage_in_memory_with_shapes(sim):
         prim_path_regex = "/World/env_.*/Cone"
         cfg.func(prim_path_regex, cfg)
 
-        # verify stage is in memory
-        assert sim_utils.is_current_stage_in_memory()
-
-        # verify prims exist in stage in memory
+        # verify prims exist in stage
         prims = sim_utils.find_matching_prim_paths(prim_path_regex)
         assert len(prims) == num_clones
 
-        # verify prims do not exist in context stage (if one exists)
-        context_stage = sim_utils.get_context_stage()
-        if context_stage is not None:
-            with sim_utils.use_stage(context_stage):
-                prims = sim_utils.find_matching_prim_paths(prim_path_regex)
-                assert len(prims) != num_clones
-
-        # attach stage to context
-        sim_utils.attach_stage_to_usd_context()
-
-    # verify stage is no longer in memory
-    assert not sim_utils.is_current_stage_in_memory()
-
-    # verify prims now exist in context stage
+    # verify prims exist in context stage
     prims = sim_utils.find_matching_prim_paths(prim_path_regex)
     assert len(prims) == num_clones
 
@@ -147,7 +134,10 @@ def test_stage_in_memory_with_usds(sim):
         f"{ISAACLAB_NUCLEUS_DIR}/Robots/ANYbotics/ANYmal-D/anymal_d.usd",
     ]
 
-    # grab stage in memory and set as current stage via the with statement
+    # verify stage is attached to USD context (happens automatically now with create_stage_in_memory)
+    assert not sim_utils.is_current_stage_in_memory()
+
+    # grab stage and set as current stage via the with statement
     stage_in_memory = sim.stage
     with sim_utils.use_stage(stage_in_memory):
         # create cloned robot stage
@@ -174,27 +164,11 @@ def test_stage_in_memory_with_usds(sim):
         prim_path_regex = "/World/env_.*/Robot"
         cfg.func(prim_path_regex, cfg)
 
-        # verify stage is in memory
-        assert sim_utils.is_current_stage_in_memory()
-
-        # verify prims exist in stage in memory
+        # verify prims exist in stage
         prims = sim_utils.find_matching_prim_paths(prim_path_regex)
         assert len(prims) == num_clones
 
-        # verify prims do not exist in context stage (if one exists)
-        context_stage = sim_utils.get_context_stage()
-        if context_stage is not None:
-            with sim_utils.use_stage(context_stage):
-                prims = sim_utils.find_matching_prim_paths(prim_path_regex)
-                assert len(prims) != num_clones
-
-        # attach stage to context
-        sim_utils.attach_stage_to_usd_context()
-
-    # verify stage is no longer in memory
-    assert not sim_utils.is_current_stage_in_memory()
-
-    # verify prims now exist in context stage
+    # verify prims exist in context stage
     prims = sim_utils.find_matching_prim_paths(prim_path_regex)
     assert len(prims) == num_clones
 
@@ -210,7 +184,10 @@ def test_stage_in_memory_with_clone_in_fabric(sim):
     usd_path = f"{ISAACLAB_NUCLEUS_DIR}/Robots/ANYbotics/ANYmal-C/anymal_c.usd"
     num_clones = 100
 
-    # grab stage in memory and set as current stage via the with statement
+    # verify stage is attached to USD context (happens automatically now with create_stage_in_memory)
+    assert not sim_utils.is_current_stage_in_memory()
+
+    # grab stage and set as current stage via the with statement
     stage_in_memory = sim.stage
     with sim_utils.use_stage(stage_in_memory):
         # set up paths
@@ -235,22 +212,8 @@ def test_stage_in_memory_with_clone_in_fabric(sim):
             replicate_physics=True,
             clone_in_fabric=True,
         )
-        prim_path_regex = "/World/envs/env_.*"
-
-        # verify prims do not exist in context stage (if one exists)
-        context_stage = sim_utils.get_context_stage()
-        if context_stage is not None:
-            with sim_utils.use_stage(context_stage):
-                prims = sim_utils.find_matching_prim_paths(prim_path_regex)
-                assert len(prims) != num_clones
-
-        # attach stage to context
-        sim_utils.attach_stage_to_usd_context()
-
-    # verify stage is no longer in memory
-    assert not sim_utils.is_current_stage_in_memory()
 
-    # verify prims now exist in fabric stage using usdrt apis
+    # verify prims exist in fabric stage using usdrt apis
     stage_id = sim_utils.get_current_stage_id()
     usdrt_stage = usdrt.Usd.Stage.Attach(stage_id)
     for i in range(num_clones):