Adds different render modes to `IsaacEnv` (#63)

* updates envs to pass all keyword args to isaac-env
* adds rendering modes to IsaacEnv
* fixes loading of extensions for viewport rendering
* adds unit test for recording videos of all envs
* adds physx spellings to vscode settings.json
* adds video logging to all workflow scripts
* updates changelog and version
* adds function to print callables in dict
* adds documentation on env wrappers

---------
Signed-off-by: Mayank Mittal <12863862+Mayankm96@users.noreply.github.com>

.vscode/tools/settings.template.json

@@ -40,6 +40,8 @@
         "teleoperation",
         "xform",
         "numpy",
+        "flatcache",
+        "physx",
         "dpad",
         "gamepad"
     ],

docs/index.rst

@@ -58,6 +58,7 @@ using the following BibTeX entry:
 
    source/tutorials_envs/00_gym_env
    source/tutorials_envs/01_create_env
+   source/tutorials_envs/02_wrappers
 
 .. toctree::
    :maxdepth: 2

docs/source/api/orbit_envs.isaac_env.rst

@@ -52,6 +52,7 @@ Base Environment
    :undoc-members:
    :show-inheritance:
    :private-members:
+   :exclude-members: _configure_simulation_flags, _create_viewport_render_product, _last_obs_buf
 
 Base Configuration
 ---------------------

docs/source/setup/installation.rst

@@ -190,6 +190,9 @@ use the following command:
       .. code:: bash
 
          conda activate orbit
+      
+      Once you are in the virtual environment, you do not need to use ``./orbit.sh -p``
+      to run python scripts. You can use the default python executable in your environment.
 
 
 Building extensions

docs/source/tutorials_envs/01_create_env.rst

@@ -223,17 +223,17 @@ As examples of this in the ``omni.isaac.orbit_envs`` package, we have the follow
 
    .. literalinclude:: ../../../source/extensions/omni.isaac.orbit_envs/omni/isaac/orbit_envs/__init__.py
       :language: python
-      :lines: 45-49
+      :lines: 52-56
       :linenos:
-      :lineno-start: 45
+      :lineno-start: 52
 
 .. dropdown:: :fa:`eye,mr-1` Registering an environment with a Python dataclass configuration file
 
    .. literalinclude:: ../../../source/extensions/omni.isaac.orbit_envs/omni/isaac/orbit_envs/__init__.py
       :language: python
-      :lines: 83-87
+      :lines: 84-88
       :linenos:
-      :lineno-start: 83
+      :lineno-start: 84
 
 
 The Code Execution

docs/source/tutorials_envs/02_wrappers.rst

+Using environment wrappers
+==========================
+
+Environment wrappers are a way to modify the behavior of an environment without modifying the environment itself.
+This can be used to apply functions to modify observations or rewards, record videos, enforce time limits, etc.
+A detailed description of the API is available in the `gym.Wrapper <https://gymnasium.farama.org/api/wrappers/>`_ class.
+
+
+At present, all environments inheriting from the :class:`omni.isaac.orbit_envs.isaac_env.IsaacEnv` class
+are compatible with ``gym.Wrapper``, since the base class implements the ``gym.Env`` interface.
+In order to wrap an environment, you need to first initialize the base environment. After that, you can
+wrap it with as many wrappers as you want by calling `env = wrapper(env, *args, **kwargs)` repeatedly.
+
+For example, here is how you would wrap an environment to enforce that reset is called before step or render:
+
+.. code-block:: python
+
+    import gym
+
+    import omni.isaac.orbit_envs  # noqa: F401
+    from omni.isaac.orbit_envs.utils import load_default_env_cfg
+
+    # create base environment
+    cfg = load_default_env_cfg("Isaac-Reach-Franka-v0")
+    env = gym.make("Isaac-Reach-Franka-v0", cfg=cfg, headless=True)
+    # wrap environment to enforce that reset is called before step
+    env = gym.wrappers.OrderEnforcing(env)
+
+
+Wrapper for recording videos
+----------------------------
+
+The :class:`gym.wrappers.RecordVideo <gym:RecordVideo>` wrapper can be used to record videos of the environment.
+The wrapper takes a ``video_dir`` argument, which specifies where to save the videos. The videos are saved in
+`mp4 <https://en.wikipedia.org/wiki/MP4_file_format>`__ format at specified intervals for specified
+number of environment steps or episodes.
+
+To use the wrapper, you need to first install ``ffmpeg``. On Ubuntu, you can install it by running:
+
+.. code-block:: bash
+
+    sudo apt-get install ffmpeg
+
+The :class:`IsaacEnv` supports different rendering modes: "human" and "rgb_array". The "human" mode
+is used when you want to render the environment to the screen. The "rgb_array" mode is used when you want
+to get the rendered image as a numpy array. However, the "rgb_array" mode is allowed only when viewport
+is enabled. This can be done by setting the ``viewport`` argument to ``True`` when creating the environment.
+
+.. code-block:: python
+
+    import gym
+
+    import omni.isaac.orbit_envs  # noqa: F401
+    from omni.isaac.orbit_envs.utils import load_default_env_cfg
+
+    # create base environment
+    cfg = load_default_env_cfg("Isaac-Reach-Franka-v0")
+    env = gym.make("Isaac-Reach-Franka-v0", cfg=cfg, headless=True, viewport=True)
+    # wrap environment to enforce that reset is called before step
+    env = gym.wrappers.OrderEnforcing(env)
+
+.. attention::
+
+  By default, when running an environment in headless mode, the Omniverse viewport is disabled. This is done to
+  improve performance by avoiding unnecessary rendering.
+
+  We notice the following performance in different rendering modes with the  ``Isaac-Reach-Franka-v0`` environment
+  using an RTX 3090 GPU:
+
+  * Headless execution without viewport enabled: ~65,000 FPS
+  * Headless execution with viewport enabled: ~57,000 FPS
+  * Non-headless execution (i.e. with GUI): ~13,000 FPS
+
+
+The viewport camera used for rendering is the default camera in the scene called ``"/OmniverseKit_Persp"``.
+The camera's pose and image resolution can be configured through the
+:class:`omni.isaac.orbit_env.isaac_env_cfg.ViewerCfg` class.
+
+
+.. dropdown:: :fa:`eye,mr-1` Default parameters of the :class:`ViewerCfg` in the ``isaac_env_cfg.py`` file:
+
+   .. literalinclude:: ../../../source/extensions/omni.isaac.orbit_envs/omni/isaac/orbit_envs/isaac_env_cfg.py
+      :language: python
+      :lines: 37-52
+      :linenos:
+      :lineno-start: 37
+
+
+After adjusting the parameters, you can record videos by wrapping the environment with the
+:class:`gym.wrappers.RecordVideo <gym:RecordVideo>` wrapper. As an example, the following code
+records a video of the ``Isaac-Reach-Franka-v0`` environment for 200 steps, and saves it in the
+``videos`` folder at a step interval of 1500 steps.
+
+.. code:: python
+
+    import gym
+
+    # adjust camera resolution and pose
+    env_cfg.viewer.resolution = (640, 480)
+    env_cfg.viewer.eye = (1.0, 1.0, 1.0)
+    env_cfg.viewer.lookat = (0.0, 0.0, 0.0)
+    # create isaac-env instance
+    env = gym.make(task_name, cfg=env_cfg, headless=headless, viewport=True)
+    # wrap for video recording
+    video_kwargs = {
+        "video_folder": "videos",
+        "step_trigger": lambda step: step % 1500 == 0,
+        "video_length": 200,
+    }
+    env = gym.wrappers.RecordVideo(env, **video_kwargs)
+
+
+Wrapper for learning frameworks
+-------------------------------
+
+Every learning framework has its own API for interacting with environments. For example, the
+`Stable Baselines3 <https://stable-baselines3.readthedocs.io/en/master/>`__ library uses the
+`gym.Env <https://gymnasium.farama.org/api/env/>`__ interface to interact with environments.
+However, libraries like `RL-Games <https://github.com/Denys88/rl_games>`__ or 
+`RSL-RL <https://github.com/leggedrobotics/rsl_rl>`__ use their own API for interfacing with a
+learning environments. Since there is no one-size-fits-all solution, we do not base the :class:`IsaacEnv`
+class on any particular learning framework's environment definition. Instead, we implement
+wrappers to make it compatible with the learning framework's environment definition.
+
+As an example of how to use the :class:`IsaacEnv` with Stable-Baselines3:
+
+.. code:: python
+
+    from omni.isaac.orbit_envs.utils.wrappers.sb3 import Sb3VecEnvWrapper
+
+    # create isaac-env instance
+    env = gym.make(task_name, cfg=env_cfg, headless=headless)
+    # wrap around environment for stable baselines
+    env = Sb3VecEnvWrapper(env)
+
+
+.. caution::
+
+  Wrapping the environment with the respective learning framework's wrapper should happen in the end,
+  i.e. after all other wrappers have been applied. This is because the learning framework's wrapper
+  modifies the interpretation of environment's APIs which may no longer be compatible with ``gym.Env``.
+
+
+To add support for a new learning framework, you need to implement a wrapper class that
+converts the :class:`IsaacEnv` to the learning framework's environment definition. This
+wrapper class should typically inherit from the ``gym.Wrapper`` class. We include a
+set of these wrappers in the :mod:`omni.isaac.orbit_envs.utils.wrappers` module. You can
+use these wrappers as a reference to implement your own wrapper for a new learning framework.

source/extensions/omni.isaac.orbit/config/extension.toml

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.3.1"
+version = "0.3.2"
 
 # Description
 title = "ORBIT framework for Robot Learning"

source/extensions/omni.isaac.orbit/docs/CHANGELOG.rst

 Changelog
 ---------
 
+0.3.2 (2023-04-27)
+~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added safe-printing of functions while using the :meth:`omni.isaac.orbit.utils.dict.print_dict` function.
+
+
 0.3.1 (2023-04-23)
 ~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/utils/dict.py

@@ -8,6 +8,7 @@
 
 import collections.abc
 import importlib
+import inspect
 from typing import Any, Callable, Dict, Iterable, Mapping
 
 from .array import TENSOR_TYPE_CONVERSIONS, TENSOR_TYPES
@@ -223,7 +224,13 @@ def print_dict(val, nesting: int = -4, start: bool = True):
             print(k, end=": ")
             print_dict(val[k], nesting, start=False)
     else:
-        print(val)
+        # deal with functions in print statements
+        if callable(val) and val.__name__ == "<lambda>":
+            print("lambda", inspect.getsourcelines(val)[0][0].strip().split("lambda")[1].strip()[:-1])
+        elif callable(val):
+            print(f"{val.__module__}:{val.__name__}")
+        else:
+            print(val)
 
 
 """

source/extensions/omni.isaac.orbit/test/test_dict_utils.py

+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+import unittest
+
+import omni.isaac.orbit.utils.dict as dict_utils
+
+
+class TestDictUtilities(unittest.TestCase):
+    """Test fixture for checking Kit utilities in Orbit."""
+
+    def test_print_dict(self):
+        """Test printing of dictionary."""
+        # create a complex nested dictionary
+        test_dict = {
+            "a": 1,
+            "b": 2,
+            "c": {"d": 3, "e": 4, "f": {"g": 5, "h": 6}},
+            "i": 7,
+            "j": lambda x: x**2,
+            "k": dict_utils.class_to_dict,
+        }
+        # print the dictionary
+        dict_utils.print_dict(test_dict)
+
+
+if __name__ == "__main__":
+    unittest.main()

source/extensions/omni.isaac.orbit_envs/config/extension.toml

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.2.3"
+version = "0.3.0"
 
 # Description
 title = "ORBIT Environments"

source/extensions/omni.isaac.orbit_envs/docs/CHANGELOG.rst

 Changelog
 ---------
 
+0.3.0 (2023-04-14)
+~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added a new flag ``viewport`` to the :class:`IsaacEnv` class to enable/disable rendering of the viewport.
+  If the flag is set to ``True``, the viewport is enabled and the environment is rendered in the background.
+* Updated the training scripts in the ``source/standalone/workflows`` directory to use the new flag ``viewport``.
+  If the CLI argument ``--video`` is passed, videos are recorded in the ``videos`` directory using the
+  :class:`gym.wrappers.RecordVideo` wrapper.
+
+Changed
+^^^^^^^
+
+* The :class:`IsaacEnv` class supports different rendering mode as referenced in OpenAI Gym's ``render`` method.
+  These modes are:
+
+  * ``rgb_array``: Renders the environment in the background and returns the rendered image as a numpy array.
+  * ``human``: Renders the environment in the background and displays the rendered image in a window.
+
+* Changed the constructor in the classes inheriting from :class:`IsaacEnv` to pass all the keyword arguments to the
+  constructor of :class:`IsaacEnv` class.
+
+Fixed
+^^^^^
+
+* Clarified the documentation of ``headless`` flag in the :class:`IsaacEnv` class. It refers to whether or not
+  to render at every sim step, not whether to render the viewport or not.
+* Fixed the unit tests for running random agent on included environments.
+
 0.2.3 (2023-03-06)
 ~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit_envs/omni/isaac/orbit_envs/classic/ant/ant_env.py

@@ -26,12 +26,12 @@ class AntEnv(IsaacEnv):
         https://github.com/openai/gym/blob/master/gym/envs/mujoco/ant_v3.py
     """
 
-    def __init__(self, cfg: dict, headless: bool = False):
+    def __init__(self, cfg: dict, **kwargs):
         """Initializes the environment.
 
         Args:
             cfg (dict): The configuration dictionary.
-            headless (bool, optional): Whether to enable rendering or not. Defaults to False.
+            kwargs (dict): Additional keyword arguments. See IsaacEnv for more details.
         """
         # copy configuration
         self.cfg_dict = cfg.copy()
@@ -41,7 +41,7 @@ class AntEnv(IsaacEnv):
         )
         isaac_cfg.sim.from_dict(self.cfg_dict["sim"])
         # initialize the base class to setup the scene.
-        super().__init__(isaac_cfg, headless=headless)
+        super().__init__(isaac_cfg, **kwargs)
 
         # define views over instances
         self.ants = ArticulationView(prim_paths_expr=self.env_ns + "/.*/Ant/torso", reset_xform_properties=False)

source/extensions/omni.isaac.orbit_envs/omni/isaac/orbit_envs/classic/cartpole/cartpole_env.py

@@ -25,12 +25,12 @@ class CartpoleEnv(IsaacEnv):
         https://github.com/openai/gym/blob/master/gym/envs/classic_control/cartpole.py
     """
 
-    def __init__(self, cfg: dict, headless: bool = False):
+    def __init__(self, cfg: dict, **kwargs):
         """Initializes the environment.
 
         Args:
             cfg (dict): The configuration dictionary.
-            headless (bool, optional): Whether to enable rendering or not. Defaults to False.
+            kwargs (dict): Additional keyword arguments. See IsaacEnv for more details.
         """
         # copy configuration
         self.cfg_dict = cfg.copy()
@@ -40,7 +40,7 @@ class CartpoleEnv(IsaacEnv):
         )
         isaac_cfg.sim.from_dict(self.cfg_dict["sim"])
         # initialize the base class to setup the scene.
-        super().__init__(isaac_cfg, headless=headless)
+        super().__init__(isaac_cfg, **kwargs)
 
         # define views over instances
         self.cartpoles = ArticulationView(prim_paths_expr=self.env_ns + "/.*/Cartpole", reset_xform_properties=False)

source/extensions/omni.isaac.orbit_envs/omni/isaac/orbit_envs/classic/humanoid/humanoid_env.py

@@ -26,12 +26,12 @@ class HumanoidEnv(IsaacEnv):
         https://github.com/openai/gym/blob/master/gym/envs/mujoco/humanoid_v3.py
     """
 
-    def __init__(self, cfg: dict, headless: bool = False):
+    def __init__(self, cfg: dict, **kwargs):
         """Initializes the environment.
 
         Args:
             cfg (dict): The configuration dictionary.
-            headless (bool, optional): Whether to enable rendering or not. Defaults to False.
+            kwargs (dict): Additional keyword arguments. See IsaacEnv for more details.
         """
         # copy configuration
         self.cfg_dict = cfg.copy()
@@ -41,7 +41,7 @@ class HumanoidEnv(IsaacEnv):
         )
         isaac_cfg.sim.from_dict(self.cfg_dict["sim"])
         # initialize the base class to setup the scene.
-        super().__init__(isaac_cfg, headless=headless)
+        super().__init__(isaac_cfg, **kwargs)
 
         # define views over instances
         self.humanoids = ArticulationView(

source/extensions/omni.isaac.orbit_envs/omni/isaac/orbit_envs/isaac_env.py

@@ -8,8 +8,9 @@
 
 import abc
 import gym
+import numpy as np
 import torch
-from typing import Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, ClassVar, Dict, Iterable, List, Optional, Tuple, Union
 
 import carb
 import omni.isaac.core.utils.prims as prim_utils
@@ -65,15 +66,27 @@ class IsaacEnv(gym.Env):
     environment for their agents.
     """
 
-    def __init__(self, cfg: IsaacEnvCfg, headless: bool = False):
+    is_vector_env: ClassVar[bool] = True
+    """Whether the environment is a vectorized environment."""
+    metadata: ClassVar[Dict[str, Any]] = {"render.modes": ["human", "rgb_array"]}
+    """Metadata for the environment."""
+
+    def __init__(self, cfg: IsaacEnvCfg, headless: bool = False, viewport: bool = False, **kwargs):
         """Initialize the environment.
 
         We currently support only PyTorch backend for the environment. In the future, we plan to extend this to use
         other backends such as Warp.
 
+        If the environment is not headless and viewport is enabled, then the viewport will be rendered in the GUI.
+        This allows us to render the environment even in the headless mode. However, it will significantly slow
+        down the simulation. Thus, it is recommended to set both ``headless`` and ``viewport``
+        to ``False`` when training an environment (unless it uses perception sensors).
+
         Args:
             cfg (IsaacEnvCfg): Instance of the environment configuration.
-            headless (bool, optional): Whether to run with GUI or without GUI. Defaults to False.
+            headless (bool, optional): Whether to render at every simulation step. Defaults to False.
+            viewport (bool, optional): Whether to enable the GUI viewport. If True, then the viewport
+                will be rendered in the GUI (even in the headless mode). Defaults to False.
 
         Raises:
             RuntimeError: No stage is found in the simulation.
@@ -81,6 +94,7 @@ class IsaacEnv(gym.Env):
         # store inputs to class
         self.cfg = cfg
         self.enable_render = not headless
+        self.enable_viewport = viewport or self.enable_render
         # extract commonly used parameters
         self.num_envs = self.cfg.env.num_envs
         self.device = self.cfg.sim.device
@@ -115,6 +129,8 @@ class IsaacEnv(gym.Env):
             physics_prim_path="/physicsScene",
             device=self.device,
         )
+        # create renderer and set camera view
+        self._create_viewport_render_product()
         # add flag for checking closing status
         self._is_closed = False
 
@@ -129,8 +145,6 @@ class IsaacEnv(gym.Env):
         #   physics handles become invalid. So it is not possible to call :meth:`_get_observations()`
         self._last_obs_buf: VecEnvObs = None
 
-        # set camera view
-        set_camera_view(eye=self.cfg.viewer.eye, target=self.cfg.viewer.lookat)
         # create cloner for duplicating the scenes
         cloner = GridCloner(spacing=self.cfg.env.env_spacing)
         cloner.define_base_env(self.env_ns)
@@ -274,9 +288,46 @@ class IsaacEnv(gym.Env):
         # return observations, rewards, resets and extras
         return self._last_obs_buf, self.reward_buf, self.reset_buf, self.extras
 
-    def render(self):
-        """Run rendering without stepping through the physics."""
-        self.sim.render()
+    def render(self, mode: str = "human") -> Optional[np.ndarray]:
+        """Run rendering without stepping through the physics.
+
+        By convention, if mode is:
+
+        - **human**: render to the current display and return nothing. Usually for human consumption.
+        - **rgb_array**: Return an numpy.ndarray with shape (x, y, 3), representing RGB values for an
+          x-by-y pixel image, suitable for turning into a video.
+
+        Args:
+            mode (str, optional): The mode to render with. Defaults to "human".
+        """
+        # render the scene only if rendering at every step is disabled
+        # this is because we do not want to render the scene twice
+        if not self.enable_render:
+            # manually flush the flatcache data to update Hydra textures
+            if self.sim.get_physics_context().use_flatcache:
+                self._flatcache_iface.update(0.0, 0.0)
+            # render the scene
+            self.sim.render()
+        # decide the rendering mode
+        if mode == "human":
+            return None
+        elif mode == "rgb_array":
+            # check if viewport is enabled -- if not, then complain because we won't get any data
+            if not self.enable_viewport:
+                raise RuntimeError(
+                    f"Cannot render '{mode}' when enable viewport is False. Please check the provided"
+                    "arguments to the environment class at initialization."
+                )
+            # obtain the rgb data
+            rgb_data = self._rgb_annotator.get_data()
+            # convert to numpy array
+            rgb_data = np.frombuffer(rgb_data, dtype=np.uint8).reshape(*rgb_data.shape)
+            # return the rgb data
+            return rgb_data[:, :, :3]
+        else:
+            raise NotImplementedError(
+                f"Render mode '{mode}' is not supported. Please use: {self.metadata['render.modes']}."
+            )
 
     def close(self):
         """Cleanup for the environment."""
@@ -390,7 +441,7 @@ class IsaacEnv(gym.Env):
     """
 
     def _configure_simulation_flags(self, sim_params: dict = None):
-        """Configure various simulation flags for performance improvements at load and run time."""
+        """Configure simulation flags and extensions at load and run time."""
         # acquire settings interface
         carb_settings_iface = carb.settings.get_settings()
         # enable hydra scene-graph instancing
@@ -405,12 +456,51 @@ class IsaacEnv(gym.Env):
             carb_settings_iface.set_bool("/physics/disableContactProcessing", True)
 
         # set flags based on whether rendering is enabled or not
-        if self.enable_render:
+        # note: enabling extensions is order-sensitive. please do not change the order.
+        if self.enable_render or self.enable_viewport:
             # enable scene querying if rendering is enabled
             # this is needed for some GUI features
             sim_params["enable_scene_query_support"] = True
+            # load extra viewport extensions if requested
+            if self.enable_viewport:
+                # extension to enable UI buttons (otherwise we get attribute errors)
+                enable_extension("omni.kit.window.toolbar")
+                # extension to make RTX realtime and path-traced renderers
+                enable_extension("omni.kit.viewport.rtx")
+                # extension to make HydraDelegate renderers
+                enable_extension("omni.kit.viewport.pxr")
             # enable viewport extension if not running in headless mode
             enable_extension("omni.kit.viewport.bundle")
+            # load extra render extensions if requested
+            if self.enable_viewport:
+                # extension for window status bar
+                enable_extension("omni.kit.window.status_bar")
         # enable isaac replicator extension
         # note: moved here since it requires to have the viewport extension to be enabled first.
         enable_extension("omni.replicator.isaac")
+
+    def _create_viewport_render_product(self):
+        """Create a render product of the viewport for rendering."""
+        # set camera view for "/OmniverseKit_Persp" camera
+        set_camera_view(eye=self.cfg.viewer.eye, target=self.cfg.viewer.lookat)
+
+        # check if flatcache is enabled
+        # this is needed to flush the flatcache data into Hydra manually when calling `env.render()`
+        # ref: https://docs.omniverse.nvidia.com/prod_extensions/prod_extensions/ext_physics.html
+        if not self.enable_render and self.sim.get_physics_context().use_flatcache:
+            from omni.physxflatcache import get_physx_flatcache_interface
+
+            # acquire flatcache interface
+            self._flatcache_iface = get_physx_flatcache_interface()
+
+        # check if viewport is enabled before creating render product
+        if self.enable_viewport:
+            import omni.replicator.core as rep
+
+            # create render product
+            self._render_product = rep.create.render_product("/OmniverseKit_Persp", self.cfg.viewer.resolution)
+            # create rgb annotator -- used to read data from the render product
+            self._rgb_annotator = rep.AnnotatorRegistry.get_annotator("rgb", device="cpu")
+            self._rgb_annotator.attach([self._render_product])
+        else:
+            carb.log_info("Viewport is disabled. Skipping creation of render product.")

source/extensions/omni.isaac.orbit_envs/omni/isaac/orbit_envs/isaac_env_cfg.py

@@ -44,6 +44,12 @@ class ViewerCfg:
     """Initial camera position (in m). Default is (7.5, 7.5, 7.5)."""
     lookat: Tuple[float, float, float] = (0.0, 0.0, 0.0)
     """Initial camera target position (in m). Default is (0.0, 0.0, 0.0)."""
+    resolution: Tuple[int, int] = (1280, 720)
+    """The resolution (width, height) of the default viewport (in pixels). Default is (1280, 720).
+
+    This is the resolution of the camera "/OmniverseKit_Persp", that is used in default viewport.
+    The camera is also used for rendering RGB images of the simulation.
+    """
 
 
 ##

source/extensions/omni.isaac.orbit_envs/omni/isaac/orbit_envs/locomotion/velocity/velocity_env.py

@@ -29,7 +29,7 @@ from .velocity_cfg import VelocityEnvCfg
 class VelocityEnv(IsaacEnv):
     """Environment for tracking a base SE(2) velocity command for a legged robot."""
 
-    def __init__(self, cfg: VelocityEnvCfg = None, headless: bool = False):
+    def __init__(self, cfg: VelocityEnvCfg = None, **kwargs):
         # copy configuration
         self.cfg = cfg
 
@@ -37,7 +37,7 @@ class VelocityEnv(IsaacEnv):
         self.robot = LeggedRobot(cfg=self.cfg.robot)
 
         # initialize the base class to setup the scene.
-        super().__init__(self.cfg, headless=headless)
+        super().__init__(self.cfg, **kwargs)
         # parse the configuration for information
         self._process_cfg()
         # initialize views for the cloned scenes

source/extensions/omni.isaac.orbit_envs/omni/isaac/orbit_envs/manipulation/lift/lift_env.py

@@ -27,7 +27,7 @@ from .lift_cfg import LiftEnvCfg, RandomizationCfg
 class LiftEnv(IsaacEnv):
     """Environment for lifting an object off a table with a single-arm manipulator."""
 
-    def __init__(self, cfg: LiftEnvCfg = None, headless: bool = False):
+    def __init__(self, cfg: LiftEnvCfg = None, **kwargs):
         # copy configuration
         self.cfg = cfg
         # parse the configuration for controller configuration
@@ -38,7 +38,7 @@ class LiftEnv(IsaacEnv):
         self.object = RigidObject(cfg=self.cfg.object)
 
         # initialize the base class to setup the scene.
-        super().__init__(self.cfg, headless=headless)
+        super().__init__(self.cfg, **kwargs)
         # parse the configuration for information
         self._process_cfg()
         # initialize views for the cloned scenes

source/extensions/omni.isaac.orbit_envs/omni/isaac/orbit_envs/manipulation/reach/reach_env.py

@@ -25,7 +25,7 @@ from .reach_cfg import RandomizationCfg, ReachEnvCfg
 class ReachEnv(IsaacEnv):
     """Environment for reaching to desired pose for a single-arm manipulator."""
 
-    def __init__(self, cfg: ReachEnvCfg = None, headless: bool = False):
+    def __init__(self, cfg: ReachEnvCfg = None, **kwargs):
         # copy configuration
         self.cfg = cfg
         # parse the configuration for controller configuration
@@ -35,7 +35,7 @@ class ReachEnv(IsaacEnv):
         self.robot = SingleArmManipulator(cfg=self.cfg.robot)
 
         # initialize the base class to setup the scene.
-        super().__init__(self.cfg, headless=headless)
+        super().__init__(self.cfg, **kwargs)
         # parse the configuration for information
         self._process_cfg()
         # initialize views for the cloned scenes

source/extensions/omni.isaac.orbit_envs/omni/isaac/orbit_envs/utils/parse_cfg.py

@@ -27,7 +27,13 @@ def load_default_env_cfg(task_name: str) -> Union[dict, Any]:
 
     Returns:
         Union[dict, Any]: The parsed configuration object.
+
+    Raises:
+        ValueError: If the task name is not provided, i.e. None.
     """
+    # check if a task name is provided
+    if task_name is None:
+        raise ValueError("Please provide a valid task name. Hint: Use --task <task_name>.")
     # retrieve the configuration file to load
     cfg_entry_point: str = gym.spec(task_name)._kwargs.pop("cfg_entry_point")
 
@@ -74,7 +80,13 @@ def parse_env_cfg(task_name: str, use_gpu: bool = True, num_envs: int = None, **
 
     Returns:
         Union[dict, Any]: The parsed configuration object.
+
+    Raises:
+        ValueError: If the task name is not provided, i.e. None.
     """
+    # check if a task name is provided
+    if task_name is None:
+        raise ValueError("Please provide a valid task name. Hint: Use --task <task_name>.")
     # create a dictionary to update from
     args_cfg = {"sim": {"physx": dict()}, "env": dict()}
     # resolve pipeline to use (based on input)

source/extensions/omni.isaac.orbit_envs/test/test_environments.py

@@ -5,12 +5,15 @@
 
 """Launch Isaac Sim Simulator first."""
 
+import os
 
 from omni.isaac.kit import SimulationApp
 
 # launch the simulator
+app_experience = f"{os.environ['EXP_PATH']}/omni.isaac.sim.python.gym.headless.kit"
 config = {"headless": True}
-simulation_app = SimulationApp(config)
+simulation_app = SimulationApp(config, experience=app_experience)
+
 
 """Rest everything follows."""
 
@@ -21,6 +24,8 @@ import torch
 import unittest
 from typing import Dict, Union
 
+import omni.usd
+
 import omni.isaac.orbit_envs  # noqa: F401
 from omni.isaac.orbit_envs.utils.parse_cfg import parse_env_cfg
 
@@ -34,14 +39,15 @@ class TestEnvironments(unittest.TestCase):
         simulation_app.close()
 
     def setUp(self) -> None:
-        self.use_gpu = False
-        self.num_envs = 20
+        self.num_envs = 512
         self.headless = simulation_app.config["headless"]
         # acquire all Isaac environments names
         self.registered_tasks = list()
         for task_spec in gym.envs.registry.all():
             if "Isaac" in task_spec.id:
                 self.registered_tasks.append(task_spec.id)
+        # sort environments by name
+        self.registered_tasks.sort()
         # print all existing task names
         print(">>> All registered environments:", self.registered_tasks)
 
@@ -50,8 +56,10 @@ class TestEnvironments(unittest.TestCase):
 
         for task_name in self.registered_tasks:
             print(f">>> Running test for environment: {task_name}")
+            # create a new stage
+            omni.usd.get_context().new_stage()
             # parse configuration
-            env_cfg = parse_env_cfg(task_name, use_gpu=self.use_gpu, num_envs=self.num_envs)
+            env_cfg = parse_env_cfg(task_name, use_gpu=True, num_envs=self.num_envs)
             # create environment
             env = gym.make(task_name, cfg=env_cfg, headless=self.headless)
 

source/extensions/omni.isaac.orbit_envs/test/test_record_video.py

+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Launch Isaac Sim Simulator first."""
+
+import os
+
+from omni.isaac.kit import SimulationApp
+
+# launch the simulator
+app_experience = f"{os.environ['EXP_PATH']}/omni.isaac.sim.python.gym.headless.render.kit"
+config = {"headless": True}
+simulation_app = SimulationApp(config, experience=app_experience)
+
+"""Rest everything follows."""
+
+
+import gym
+import os
+import torch
+import unittest
+
+import omni.isaac.contrib_envs  # noqa: F401
+import omni.isaac.orbit_envs  # noqa: F401
+from omni.isaac.orbit_envs.utils import parse_env_cfg
+
+
+class TestRecordVideoWrapper(unittest.TestCase):
+    """Test recording videos using the RecordVideo wrapper."""
+
+    @classmethod
+    def tearDownClass(cls):
+        """Closes simulator after running all test fixtures."""
+        simulation_app.close()
+
+    def setUp(self) -> None:
+        # common parameters
+        self.num_envs = 64
+        self.use_gpu = True
+        self.headless = simulation_app.config["headless"]
+        # directory to save videos
+        self.videos_dir = os.path.join(os.path.dirname(__file__), "videos")
+        self.step_trigger = lambda step: step % 225 == 0
+        self.video_length = 200
+        # acquire all Isaac environments names
+        self.registered_tasks = list()
+        for task_spec in gym.envs.registry.all():
+            if "Isaac" in task_spec.id:
+                self.registered_tasks.append(task_spec.id)
+        # sort environments by name
+        self.registered_tasks.sort()
+        # print all existing task names
+        print(">>> All registered environments:", self.registered_tasks)
+
+    def test_record_video(self):
+        """Run random actions agent with recording of videos."""
+        import omni.usd
+
+        for task_name in self.registered_tasks:
+            print(f">>> Running test for environment: {task_name}")
+            # create a new stage
+            omni.usd.get_context().new_stage()
+            # parse configuration
+            env_cfg = parse_env_cfg(task_name, use_gpu=self.use_gpu, num_envs=self.num_envs)
+            # create environment
+            env = gym.make(task_name, cfg=env_cfg, headless=self.headless, viewport=True)
+
+            # directory to save videos
+            videos_dir = os.path.join(self.videos_dir, task_name)
+            # wrap environment to record videos
+            env = gym.wrappers.RecordVideo(
+                env, videos_dir, step_trigger=self.step_trigger, video_length=self.video_length
+            )
+
+            # reset environment
+            env.reset()
+            # simulate environment
+            for _ in range(500):
+                # compute zero actions
+                actions = 2 * torch.rand((env.num_envs, env.action_space.shape[0]), device=env.device) - 1
+                # apply actions
+                _, _, _, _ = env.step(actions)
+                # render environment
+                env.render(mode="human")
+                # check if simulator is stopped
+                if env.unwrapped.sim.is_stopped():
+                    break
+
+            # close the simulator
+            env.close()
+
+
+if __name__ == "__main__":
+    unittest.main()

source/standalone/workflows/rl_games/train.py

@@ -16,6 +16,9 @@ from omni.isaac.kit import SimulationApp
 # add argparse arguments
 parser = argparse.ArgumentParser("Welcome to Orbit: Omniverse Robotics Environments!")
 parser.add_argument("--headless", action="store_true", default=False, help="Force display off at all times.")
+parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
+parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
+parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).")
 parser.add_argument("--cpu", action="store_true", default=False, help="Use CPU pipeline.")
 parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
 parser.add_argument("--task", type=str, default=None, help="Name of the task.")
@@ -44,6 +47,7 @@ from rl_games.common import env_configurations, vecenv
 from rl_games.common.algo_observer import IsaacAlgoObserver
 from rl_games.torch_runner import Runner
 
+from omni.isaac.orbit.utils.dict import print_dict
 from omni.isaac.orbit.utils.io import dump_pickle, dump_yaml
 
 import omni.isaac.contrib_envs  # noqa: F401
@@ -89,7 +93,17 @@ def main():
     clip_actions = agent_cfg["params"]["env"].get("clip_actions", math.inf)
 
     # create isaac environment
-    env = gym.make(args_cli.task, cfg=env_cfg, headless=args_cli.headless)
+    env = gym.make(args_cli.task, cfg=env_cfg, headless=args_cli.headless, viewport=args_cli.video)
+    # wrap for video recording
+    if args_cli.video:
+        video_kwargs = {
+            "video_folder": os.path.join(log_dir, "videos"),
+            "step_trigger": lambda step: step % args_cli.video_interval == 0,
+            "video_length": args_cli.video_length,
+        }
+        print("[INFO] Recording videos during training.")
+        print_dict(video_kwargs, nesting=4)
+        env = gym.wrappers.RecordVideo(env, **video_kwargs)
     # wrap around environment for rl-games
     env = RlGamesVecEnvWrapper(env, rl_device, clip_obs, clip_actions)
 

source/standalone/workflows/rsl_rl/train.py

@@ -16,6 +16,9 @@ from omni.isaac.kit import SimulationApp
 # add argparse arguments
 parser = argparse.ArgumentParser("Welcome to Orbit: Omniverse Robotics Environments!")
 parser.add_argument("--headless", action="store_true", default=False, help="Force display off at all times.")
+parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
+parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
+parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).")
 parser.add_argument("--cpu", action="store_true", default=False, help="Use CPU pipeline.")
 parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
 parser.add_argument("--task", type=str, default=None, help="Name of the task.")
@@ -40,6 +43,7 @@ from datetime import datetime
 
 from rsl_rl.runners import OnPolicyRunner
 
+from omni.isaac.orbit.utils.dict import print_dict
 from omni.isaac.orbit.utils.io import dump_pickle, dump_yaml
 
 import omni.isaac.contrib_envs  # noqa: F401
@@ -76,7 +80,17 @@ def main():
     dump_pickle(os.path.join(log_dir, "params", "agent.pkl"), agent_cfg)
 
     # create isaac environment
-    env = gym.make(args_cli.task, cfg=env_cfg, headless=args_cli.headless)
+    env = gym.make(args_cli.task, cfg=env_cfg, headless=args_cli.headless, viewport=args_cli.video)
+    # wrap for video recording
+    if args_cli.video:
+        video_kwargs = {
+            "video_folder": os.path.join(log_dir, "videos"),
+            "step_trigger": lambda step: step % args_cli.video_interval == 0,
+            "video_length": args_cli.video_length,
+        }
+        print("[INFO] Recording videos during training.")
+        print_dict(video_kwargs, nesting=4)
+        env = gym.wrappers.RecordVideo(env, **video_kwargs)
     # wrap around environment for rsl-rl
     env = RslRlVecEnvWrapper(env)
 

source/standalone/workflows/sb3/train.py

@@ -17,6 +17,9 @@ from omni.isaac.kit import SimulationApp
 # add argparse arguments
 parser = argparse.ArgumentParser("Welcome to Orbit: Omniverse Robotics Environments!")
 parser.add_argument("--headless", action="store_true", default=False, help="Force display off at all times.")
+parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
+parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
+parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).")
 parser.add_argument("--cpu", action="store_true", default=False, help="Use CPU pipeline.")
 parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
 parser.add_argument("--task", type=str, default=None, help="Name of the task.")
@@ -45,6 +48,7 @@ from stable_baselines3.common.callbacks import CheckpointCallback
 from stable_baselines3.common.logger import configure
 from stable_baselines3.common.vec_env import VecNormalize
 
+from omni.isaac.orbit.utils.dict import print_dict
 from omni.isaac.orbit.utils.io import dump_pickle, dump_yaml
 
 import omni.isaac.contrib_envs  # noqa: F401
@@ -77,7 +81,17 @@ def main():
     n_timesteps = agent_cfg.pop("n_timesteps")
 
     # create isaac environment
-    env = gym.make(args_cli.task, cfg=env_cfg, headless=args_cli.headless)
+    env = gym.make(args_cli.task, cfg=env_cfg, headless=args_cli.headless, viewport=args_cli.video)
+    # wrap for video recording
+    if args_cli.video:
+        video_kwargs = {
+            "video_folder": os.path.join(log_dir, "videos"),
+            "step_trigger": lambda step: step % args_cli.video_interval == 0,
+            "video_length": args_cli.video_length,
+        }
+        print("[INFO] Recording videos during training.")
+        print_dict(video_kwargs, nesting=4)
+        env = gym.wrappers.RecordVideo(env, **video_kwargs)
     # wrap around environment for stable baselines
     env = Sb3VecEnvWrapper(env)
     # set the seed

source/standalone/workflows/skrl/train.py

@@ -21,6 +21,9 @@ from omni.isaac.kit import SimulationApp
 # add argparse arguments
 parser = argparse.ArgumentParser("Welcome to Orbit: Omniverse Robotics Environments!")
 parser.add_argument("--headless", action="store_true", default=False, help="Force display off at all times.")
+parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
+parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
+parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).")
 parser.add_argument("--cpu", action="store_true", default=False, help="Use CPU pipeline.")
 parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
 parser.add_argument("--task", type=str, default=None, help="Name of the task.")
@@ -48,6 +51,7 @@ from skrl.memories.torch import RandomMemory
 from skrl.utils import set_seed
 from skrl.utils.model_instantiators import deterministic_model, gaussian_model, shared_model
 
+from omni.isaac.orbit.utils.dict import print_dict
 from omni.isaac.orbit.utils.io import dump_pickle, dump_yaml
 
 import omni.isaac.contrib_envs  # noqa: F401
@@ -88,7 +92,17 @@ def main():
     dump_pickle(os.path.join(log_dir, "params", "agent.pkl"), experiment_cfg)
 
     # create isaac environment
-    env = gym.make(args_cli.task, cfg=env_cfg, headless=args_cli.headless)
+    env = gym.make(args_cli.task, cfg=env_cfg, headless=args_cli.headless, viewport=args_cli.video)
+    # wrap for video recording
+    if args_cli.video:
+        video_kwargs = {
+            "video_folder": os.path.join(log_dir, "videos"),
+            "step_trigger": lambda step: step % args_cli.video_interval == 0,
+            "video_length": args_cli.video_length,
+        }
+        print("[INFO] Recording videos during training.")
+        print_dict(video_kwargs, nesting=4)
+        env = gym.wrappers.RecordVideo(env, **video_kwargs)
     # wrap around environment for skrl
     env = SkrlVecEnvWrapper(env)  # same as: `wrap_env(env, wrapper="isaac-orbit")`