Adds context manager to reduce code duplication in scripts that need a sim context (#436)

# Description

This PR introduces a context manager wrapper to `SimulationContext` that
can be reused by any future simulation scripts that require a similar
set up and tear down functionality. This enables using the `with`
keyword like:
```
for i in range(10):
    with build_sim_context(sim_cfg) as sim:
        # Do something with sim context
        # State will be cleaned up at the end of this with block
```

Fixes #130 

## Type of change

- New feature (non-breaking change which adds functionality)

## Checklist

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

---------
Signed-off-by: jsmith-bdai <142246516+jsmith-bdai@users.noreply.github.com>
Signed-off-by: James Smith <142246516+jsmith-bdai@users.noreply.github.com>
Co-authored-by: Mayank Mittal <12863862+Mayankm96@users.noreply.github.com>

docs/source/api/orbit/omni.isaac.orbit.sim.rst

@@ -11,6 +11,7 @@
     schemas
     spawners
     utils
+    runners
 
   .. rubric:: Classes
 
@@ -20,6 +21,12 @@
     SimulationCfg
     PhysxCfg
 
+  .. rubric:: Functions
+
+  .. autosummary::
+
+    simulation_context.build_simulation_context
+
 Simulation Context
 ------------------
 
@@ -40,6 +47,10 @@ Simulation Configuration
   :show-inheritance:
   :exclude-members: __init__
 
+Simulation Context Builder
+--------------------------
+
+.. automethod:: simulation_context.build_simulation_context
 
 Utilities
 ---------
@@ -47,3 +58,10 @@ Utilities
 .. automodule:: omni.isaac.orbit.sim.utils
   :members:
   :show-inheritance:
+
+Runners
+-------
+
+.. automodule:: omni.isaac.orbit.sim.runners
+  :members:
+  :show-inheritance:

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/__init__.py

@@ -27,8 +27,9 @@ To make it convenient to use the module, we recommend importing the module as fo
 """
 
 from .converters import *  # noqa: F401, F403
+from .runners import *  # noqa: F401, F403
 from .schemas import *  # noqa: F401, F403
 from .simulation_cfg import PhysxCfg, SimulationCfg  # noqa: F401, F403
-from .simulation_context import SimulationContext  # noqa: F401, F403
+from .simulation_context import SimulationContext, build_simulation_context  # noqa: F401, F403
 from .spawners import *  # noqa: F401, F403
 from .utils import *  # noqa: F401, F403

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/runners.py

+# Copyright (c) 2022-2024, The ORBIT Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Sub-module with runners to simplify running main and unittests."""
+
+from __future__ import annotations
+
+import traceback
+import unittest
+
+import carb
+from omni.isaac.kit import SimulationApp
+
+__all__ = ["run_tests", "run_main"]
+
+
+def run_tests(simulation_app: SimulationApp, verbosity: int = 2, try_except: bool = True):
+    """Wrapper for running tests via ``unittest`` with the provided simulation app.
+
+    Args:
+        simulation_app: An instance of the :class:`omni.isaac.kit.SimulationApp`.
+        verbosity: Verbosity level for the test runner. Defaults to 2.
+        try_except: Whether to wrap ``unittest.main()`` in a try-except block. Defaults to True.
+            This is useful to remove the try-except block, as it causes issues with VSCode's debugger.
+            When False, there is a bit more console spam, but the debugger works as expected.
+    """
+    _run_function(
+        fn=unittest.main, kwargs={"verbosity": verbosity}, simulation_app=simulation_app, try_except=try_except
+    )
+
+
+def run_main(main_fn: callable[[], None], simulation_app: SimulationApp, try_except: bool = False):
+    """Wrapper for running ``main`` with the provided simulation app.
+
+    Args:
+        main_fn: Main function to run.
+        simulation_app: An instance of the :class:`omni.isaac.kit.SimulationApp`.
+        try_except: Whether to wrap ``main()`` in a try-except block. Defaults to False.
+            This is useful to remove the try-except block, as it causes issues with VSCode's debugger.
+            When False, there is a bit more console spam, but the debugger works as expected.
+
+    """
+    _run_function(fn=main_fn, kwargs={}, simulation_app=simulation_app, try_except=try_except)
+
+
+"""
+Private methods.
+"""
+
+
+def _run_function(fn: callable, kwargs: dict, simulation_app: SimulationApp, try_except: bool = False):
+    """Wrapper for running a function with the provided simulation app.
+
+    Args:
+        fn: Function to run.
+        kwargs: Keyword arguments to pass to the function.
+        simulation_app: An instance of the :class:`omni.isaac.kit.SimulationApp`.
+        try_except: Whether to wrap ``fn()`` in a try-except block. Defaults to False.
+            This is useful to remove the try-except block, as it causes issues with VSCode's debugger.
+            When False, there is a bit more console spam, but the debugger works as expected.
+    """
+    if try_except:
+        try:
+            fn(**kwargs)
+        except Exception as err:
+            # Log out exceptions before re-raising them
+            carb.log_error(err)
+            carb.log_error(traceback.format_exc())
+            raise
+        finally:
+            # Close sim app
+            simulation_app.close()
+    else:
+        fn(**kwargs)
+        simulation_app.close()

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/simulation_context.py

@@ -9,7 +9,10 @@ import builtins
 import enum
 import numpy as np
 import sys
+import traceback
 import weakref
+from collections.abc import Iterator
+from contextlib import contextmanager
 from typing import Any
 
 import carb
@@ -21,6 +24,7 @@ from omni.isaac.version import get_version
 from pxr import Gf, Usd
 
 from .simulation_cfg import SimulationCfg
+from .spawners import DomeLightCfg, GroundPlaneCfg
 from .utils import bind_physics_material
 
 
@@ -77,7 +81,7 @@ class SimulationContext(_SimulationContext):
            extensions that are running in the background that need to be updated when the simulation is running.
         2. **Cameras**: These are typically based on Hydra textures and are used to render the scene from different
            viewpoints. They can be attached to a viewport or be used independently to render the scene.
-        3. **`Viewports`**: These are windows where you can see the rendered scene.
+        3. **Viewports**: These are windows where you can see the rendered scene.
 
         Updating each of the above components has a different overhead. For example, updating the viewports is
         computationally expensive compared to updating the UI elements. Therefore, it is useful to be able to
@@ -593,3 +597,101 @@ class SimulationContext(_SimulationContext):
             self.app.shutdown()
             # disabled on linux to avoid a crash
             carb.get_framework().unload_all_plugins()
+
+
+@contextmanager
+def build_simulation_context(
+    create_new_stage: bool = True,
+    gravity_enabled: bool = True,
+    device: str = "cuda:0",
+    dt: float = 0.01,
+    sim_cfg: SimulationCfg | None = None,
+    add_ground_plane: bool = False,
+    add_lighting: bool = False,
+    auto_add_lighting: bool = False,
+) -> Iterator[SimulationContext]:
+    """Context manager to build a simulation context with the provided settings.
+
+    This function facilitates the creation of a simulation context and provides flexibility in configuring various
+    aspects of the simulation, such as time step, gravity, device, and scene elements like ground plane and
+    lighting.
+
+    If :attr:`sim_cfg` is None, then an instance of :class:`SimulationCfg` is created with default settings, with parameters
+    overwritten based on arguments to the function.
+
+    An example usage of the context manager function:
+
+    ..  code-block:: python
+
+        with build_simulation_context() as sim:
+             # Design the scene
+
+             # Play the simulation
+             sim.reset()
+             while sim.is_playing():
+                 sim.step()
+
+    Args:
+        create_new_stage: Whether to create a new stage. Defaults to True.
+        gravity_enabled: Whether to enable gravity in the simulation. Defaults to True.
+        device: Device to run the simulation on. Defaults to "cuda:0".
+        dt: Time step for the simulation: Defaults to 0.01.
+        sim_cfg: :class:`omni.isaac.orbit.sim.SimulationCfg` to use for the simulation. Defaults to None.
+        add_ground_plane: Whether to add a ground plane to the simulation. Defaults to False.
+        add_lighting: Whether to add a dome light to the simulation. Defaults to False.
+        auto_add_lighting: Whether to automatically add a dome light to the simulation if the simulation has a GUI.
+            Defaults to False. This is useful for debugging tests in the GUI.
+
+    Yields:
+        The simulation context to use for the simulation.
+
+    """
+    try:
+        if create_new_stage:
+            stage_utils.create_new_stage()
+
+        if sim_cfg is None:
+            # Construct one and overwrite the dt, gravity, and device
+            sim_cfg = SimulationCfg(dt=dt)
+
+            # Set up gravity
+            if gravity_enabled:
+                sim_cfg.gravity = (0.0, 0.0, -9.81)
+            else:
+                sim_cfg.gravity = (0.0, 0.0, 0.0)
+
+            # Set device
+            sim_cfg.device = device
+
+        # Construct simulation context
+        sim = SimulationContext(sim_cfg)
+
+        if add_ground_plane:
+            # Ground-plane
+            cfg = GroundPlaneCfg()
+            cfg.func("/World/defaultGroundPlane", cfg)
+
+        if add_lighting or (auto_add_lighting and sim.has_gui()):
+            # Lighting
+            cfg = DomeLightCfg(
+                color=(0.1, 0.1, 0.1),
+                enable_color_temperature=True,
+                color_temperature=5500,
+                intensity=10000,
+            )
+            # Dome light named specifically to avoid conflicts
+            cfg.func(prim_path="/World/defaultDomeLight", cfg=cfg, translation=(0.0, 0.0, 10.0))
+
+        yield sim
+
+    except Exception:
+        carb.log_error(traceback.format_exc())
+        raise
+    finally:
+        if not sim.has_gui():
+            # Stop simulation only if we aren't rendering otherwise the app will hang indefinitely
+            sim.stop()
+
+        # Clear the stage
+        sim.clear_all_callbacks()
+        sim.clear_instance()

source/extensions/omni.isaac.orbit/test/sim/test_build_simulation_context_headless.py

+# Copyright (c) 2022-2024, The ORBIT Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+"""
+This test has a lot of duplication with ``test_build_simulation_context_nonheadless.py``. This is intentional to ensure that the
+tests are run in both headless and non-headless modes, and we currently can't re-build the simulation app in a script.
+
+If you need to make a change to this test, please make sure to also make the same change to ``test_build_simulation_context_nonheadless.py``.
+
+"""
+
+from __future__ import annotations
+
+"""Launch Isaac Sim Simulator first."""
+
+from omni.isaac.orbit.app import AppLauncher
+
+# launch omniverse app
+app_launcher = AppLauncher(headless=True)
+simulation_app = app_launcher.app
+
+"""Rest everything follows."""
+import unittest
+
+from omni.isaac.core.utils.prims import is_prim_path_valid
+
+from omni.isaac.orbit.sim.runners import run_tests
+from omni.isaac.orbit.sim.simulation_cfg import SimulationCfg
+from omni.isaac.orbit.sim.simulation_context import build_simulation_context
+
+
+class TestBuildSimulationContextHeadless(unittest.TestCase):
+    """Tests for simulation context builder with headless usecase."""
+
+    """
+    Tests
+    """
+
+    def test_build_simulation_context_no_cfg(self):
+        """Test that the simulation context is built when no simulation cfg is passed in."""
+        for gravity_enabled in (True, False):
+            for device in ("cuda:0", "cpu"):
+                for dt in (0.01, 0.1):
+                    with self.subTest(gravity_enabled=gravity_enabled, device=device, dt=dt):
+                        with build_simulation_context(gravity_enabled=gravity_enabled, device=device, dt=dt) as sim:
+                            if gravity_enabled:
+                                self.assertEqual(sim.cfg.gravity, (0.0, 0.0, -9.81))
+                            else:
+                                self.assertEqual(sim.cfg.gravity, (0.0, 0.0, 0.0))
+
+                            if device == "cuda:0":
+                                self.assertEqual(sim.cfg.device, "cuda:0")
+                            else:
+                                self.assertEqual(sim.cfg.device, "cpu")
+
+                            self.assertEqual(sim.cfg.dt, dt)
+
+                            # Ensure that dome light didn't get added automatically as we are headless
+                            self.assertFalse(is_prim_path_valid("/World/defaultDomeLight"))
+
+    def test_build_simulation_context_ground_plane(self):
+        """Test that the simulation context is built with the correct ground plane."""
+        for add_ground_plane in (True, False):
+            with self.subTest(add_ground_plane=add_ground_plane):
+                with build_simulation_context(add_ground_plane=add_ground_plane) as _:
+                    # Ensure that ground plane got added
+                    self.assertEqual(is_prim_path_valid("/World/defaultGroundPlane"), add_ground_plane)
+
+    def test_build_simulation_context_auto_add_lighting(self):
+        """Test that the simulation context is built with the correct lighting."""
+        for add_lighting in (True, False):
+            for auto_add_lighting in (True, False):
+                with self.subTest(add_lighting=add_lighting, auto_add_lighting=auto_add_lighting):
+                    with build_simulation_context(add_lighting=add_lighting, auto_add_lighting=auto_add_lighting) as _:
+                        if add_lighting:
+                            # Ensure that dome light got added
+                            self.assertTrue(is_prim_path_valid("/World/defaultDomeLight"))
+                        else:
+                            # Ensure that dome light didn't get added as there's no GUI
+                            self.assertFalse(is_prim_path_valid("/World/defaultDomeLight"))
+
+    def test_build_simulation_context_cfg(self):
+        """Test that the simulation context is built with the correct cfg and values don't get overridden."""
+        dt = 0.001
+        # Non-standard gravity
+        gravity = (0.0, 0.0, -1.81)
+        device = "cuda:0"
+
+        cfg = SimulationCfg(
+            gravity=gravity,
+            device=device,
+            dt=dt,
+        )
+
+        with build_simulation_context(sim_cfg=cfg, gravity_enabled=False, dt=0.01, device="cpu") as sim:
+            self.assertEqual(sim.cfg.gravity, gravity)
+            self.assertEqual(sim.cfg.device, device)
+            self.assertEqual(sim.cfg.dt, dt)
+
+
+if __name__ == "__main__":
+    run_tests(simulation_app=simulation_app)

source/extensions/omni.isaac.orbit/test/sim/test_build_simulation_context_nonheadless.py

+# Copyright (c) 2022-2024, The ORBIT Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+"""This test has a lot of duplication with ``test_build_simulation_context_headless.py``. This is intentional to ensure that the
+tests are run in both headless and non-headless modes, and we currently can't re-build the simulation app in a script.
+
+If you need to make a change to this test, please make sure to also make the same change to ``test_build_simulation_context_headless.py``.
+
+"""
+
+from __future__ import annotations
+
+"""Launch Isaac Sim Simulator first."""
+
+from omni.isaac.orbit.app import AppLauncher
+
+# launch omniverse app
+app_launcher = AppLauncher(headless=False)
+simulation_app = app_launcher.app
+
+"""Rest everything follows."""
+import unittest
+
+from omni.isaac.core.utils.prims import is_prim_path_valid
+
+from omni.isaac.orbit.sim.runners import run_tests
+from omni.isaac.orbit.sim.simulation_cfg import SimulationCfg
+from omni.isaac.orbit.sim.simulation_context import build_simulation_context
+
+
+class TestBuildSimulationContextNonheadless(unittest.TestCase):
+    """Tests for simulation context builder with non-headless usecase."""
+
+    """
+    Tests
+    """
+
+    def test_build_simulation_context_no_cfg(self):
+        """Test that the simulation context is built when no simulation cfg is passed in."""
+        for gravity_enabled in (True, False):
+            for device in ("cuda:0", "cpu"):
+                for dt in (0.01, 0.1):
+                    with self.subTest(gravity_enabled=gravity_enabled, device=device, dt=dt):
+                        with build_simulation_context(
+                            gravity_enabled=gravity_enabled,
+                            device=device,
+                            dt=dt,
+                        ) as sim:
+                            if gravity_enabled:
+                                self.assertEqual(sim.cfg.gravity, (0.0, 0.0, -9.81))
+                            else:
+                                self.assertEqual(sim.cfg.gravity, (0.0, 0.0, 0.0))
+
+                            if device == "cuda:0":
+                                self.assertEqual(sim.cfg.device, "cuda:0")
+                            else:
+                                self.assertEqual(sim.cfg.device, "cpu")
+
+                            self.assertEqual(sim.cfg.dt, dt)
+
+    def test_build_simulation_context_ground_plane(self):
+        """Test that the simulation context is built with the correct ground plane."""
+        for add_ground_plane in (True, False):
+            with self.subTest(add_ground_plane=add_ground_plane):
+                with build_simulation_context(add_ground_plane=add_ground_plane) as _:
+                    # Ensure that ground plane got added
+                    self.assertEqual(is_prim_path_valid("/World/defaultGroundPlane"), add_ground_plane)
+
+    def test_build_simulation_context_auto_add_lighting(self):
+        """Test that the simulation context is built with the correct lighting."""
+        for add_lighting in (True, False):
+            for auto_add_lighting in (True, False):
+                with self.subTest(add_lighting=add_lighting, auto_add_lighting=auto_add_lighting):
+                    with build_simulation_context(add_lighting=add_lighting, auto_add_lighting=auto_add_lighting) as _:
+                        if auto_add_lighting or add_lighting:
+                            # Ensure that dome light got added
+                            self.assertTrue(is_prim_path_valid("/World/defaultDomeLight"))
+                        else:
+                            # Ensure that dome light didn't get added
+                            self.assertFalse(is_prim_path_valid("/World/defaultDomeLight"))
+
+    def test_build_simulation_context_cfg(self):
+        """Test that the simulation context is built with the correct cfg and values don't get overridden."""
+        dt = 0.001
+        # Non-standard gravity
+        gravity = (0.0, 0.0, -1.81)
+        device = "cuda:0"
+
+        cfg = SimulationCfg(
+            gravity=gravity,
+            device=device,
+            dt=dt,
+        )
+
+        with build_simulation_context(sim_cfg=cfg, gravity_enabled=False, dt=0.01, device="cpu") as sim:
+            self.assertEqual(sim.cfg.gravity, gravity)
+            self.assertEqual(sim.cfg.device, device)
+            self.assertEqual(sim.cfg.dt, dt)
+
+
+if __name__ == "__main__":
+    run_tests(simulation_app=simulation_app)