Separates out the UI window from RLEnv class (#214)

# Description

This MR implements a long-awaited feature of separating out the UI and
UI-related operations outside the environment class.

The user can also set the UI window class themselves through the
configuration. This makes the UI window modular and provides more
freedom for future additions.

## Type of change

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

## Screenshots

Also updated the GUI with various functionalities. This was kind of the
motivation to separate out the UI into a different window.


![new-gui](https://github.com/isaac-orbit/orbit/assets/12863862/813bc39f-a968-4222-a1d9-d5c4df72653d)

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./orbit.sh --format`
- [ ] I have made corresponding changes to the documentation
- [x] 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

---------
Signed-off-by: Mayank Mittal <12863862+Mayankm96@users.noreply.github.com>
Co-authored-by: David Hoeller <dhoeller@ethz.ch>

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

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

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

 Changelog
 ---------
 
+0.9.25 (2023-10-27)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added the :mod:`omni.isaac.orbit.envs.ui` module to put all the UI-related classes in one place. This currently
+  implements the :class:`omni.isaac.orbit.envs.ui.BaseEnvWindow` and :class:`omni.isaac.orbit.envs.ui.RLEnvWindow`
+  classes. Users can inherit from these classes to create their own UI windows.
+* Added the attribute :attr:`omni.isaac.orbit.envs.BaseEnvCfg.ui_window_class_type` to specify the UI window class
+  to be used for the environment. This allows the user to specify their own UI window class to be used for the
+  environment.
+
+
 0.9.24 (2023-10-27)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/command_generators/velocity_command_generator.py

@@ -150,12 +150,12 @@ class UniformVelocityCommandGenerator(CommandGeneratorBase):
                 # -- goal
                 marker_cfg = GREEN_ARROW_X_MARKER_CFG.copy()
                 marker_cfg.prim_path = "/Visuals/Command/velocity_goal"
-                marker_cfg.markers["arrow"].scale = (2.5, 0.1, 0.1)
+                marker_cfg.markers["arrow"].scale = (0.5, 0.5, 0.5)
                 self.base_vel_goal_visualizer = VisualizationMarkers(marker_cfg)
                 # -- current
                 marker_cfg = BLUE_ARROW_X_MARKER_CFG.copy()
                 marker_cfg.prim_path = "/Visuals/Command/velocity_current"
-                marker_cfg.markers["arrow"].scale = (2.5, 0.1, 0.1)
+                marker_cfg.markers["arrow"].scale = (0.5, 0.5, 0.5)
                 self.base_vel_visualizer = VisualizationMarkers(marker_cfg)
             # set their visibility to true
             self.base_vel_goal_visualizer.set_visibility(True)
@@ -187,7 +187,7 @@ class UniformVelocityCommandGenerator(CommandGeneratorBase):
         default_scale = self.base_vel_goal_visualizer.cfg.markers["arrow"].scale
         # arrow-scale
         arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1)
-        arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 2.5
+        arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0
         # arrow-direction
         heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0])
         zeros = torch.zeros_like(heading_angle)

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/base_env.py

@@ -101,6 +101,15 @@ class BaseEnv:
             # add timeline event to load managers
             self.load_managers()
 
+        # extend UI elements
+        # we need to do this here after all the managers are initialized
+        # this is because they dictate the sensors and commands right now
+        if self.sim.has_gui() and self.cfg.ui_window_class_type is not None:
+            self._window = self.cfg.ui_window_class_type(self, window_name="Orbit")
+        else:
+            # if no window, then we don't need to store the window
+            self._window = None
+
     def __del__(self):
         """Cleanup for the environment."""
         self.close()
@@ -181,5 +190,8 @@ class BaseEnv:
             # clear callbacks and instance
             self.sim.clear_all_callbacks()
             self.sim.clear_instance()
+            # destroy the window
+            if self._window is not None:
+                self._window = None
             # update closing status
             self._is_closed = True

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/base_env_cfg.py

@@ -9,8 +9,9 @@ This module defines the general configuration of the environment. It includes pa
 configuring the environment instances, viewer settings, and simulation parameters.
 """
 
+from __future__ import annotations
+
 from dataclasses import MISSING
-from typing import Tuple
 
 import omni.isaac.orbit.envs.mdp as mdp
 from omni.isaac.orbit.managers import RandomizationTermCfg as RandTerm
@@ -18,6 +19,8 @@ from omni.isaac.orbit.scene import InteractiveSceneCfg
 from omni.isaac.orbit.sim import SimulationCfg
 from omni.isaac.orbit.utils import configclass
 
+from .ui import BaseEnvWindow
+
 __all__ = ["BaseEnvCfg", "ViewerCfg"]
 
 
@@ -25,15 +28,15 @@ __all__ = ["BaseEnvCfg", "ViewerCfg"]
 class ViewerCfg:
     """Configuration of the scene viewport camera."""
 
-    eye: Tuple[float, float, float] = (7.5, 7.5, 7.5)
+    eye: tuple[float, float, float] = (7.5, 7.5, 7.5)
     """Initial camera position (in m). Default is (7.5, 7.5, 7.5)."""
-    lookat: Tuple[float, float, float] = (0.0, 0.0, 0.0)
+    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)."""
     cam_prim_path: str = "/OmniverseKit_Persp"
     """The camera prim path to record images from. Default is "/OmniverseKit_Persp", which is the
     default camera in the default viewport.
     """
-    resolution: Tuple[int, int] = (1280, 720)
+    resolution: tuple[int, int] = (1280, 720)
     """The resolution (width, height) of the camera specified using :attr:`cam_prim_path`.
     Default is (1280, 720).
     """
@@ -59,6 +62,17 @@ class BaseEnvCfg:
     """Viewer configuration. Default is ViewerCfg()."""
     sim: SimulationCfg = SimulationCfg()
     """Physics simulation configuration. Default is SimulationCfg()."""
+    # ui settings
+    ui_window_class_type: type | None = BaseEnvWindow
+    """The class type of the UI window. Defaults to None.
+
+    If None, then no UI window is created.
+
+    Note:
+        If you want to make your own UI window, you can create a class that inherits from
+        from :class:`omni.isaac.orbit.envs.ui.base_env_window.BaseEnvWindow`. Then, you can set
+        this attribute to your class type.
+    """
 
     # general settings
     decimation: int = MISSING

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/rl_env.py

@@ -5,16 +5,12 @@
 
 from __future__ import annotations
 
-import asyncio
 import gym
 import math
 import numpy as np
 import torch
-import weakref
 from typing import Any, ClassVar, Dict, Sequence, Tuple, Union
 
-import omni.usd
-
 from omni.isaac.orbit.command_generators import CommandGeneratorBase
 from omni.isaac.orbit.managers import CurriculumManager, RewardManager, TerminationManager
 
@@ -110,15 +106,6 @@ class RLEnv(BaseEnv, gym.Env):
         # perform randomization at the start of the simulation
         if "startup" in self.randomization_manager.available_modes:
             self.randomization_manager.randomize(mode="startup")
-        # extend UI elements
-        # we need to do this here after all the managers are initialized
-        # this is because they dictate the sensors and commands right now
-        if self.sim.has_gui():
-            self._build_ui()
-        else:
-            # if no window, then we don't need to store the window
-            self._orbit_window = None
-            self._orbit_window_elements = dict()
 
     """
     Properties.
@@ -298,16 +285,6 @@ class RLEnv(BaseEnv, gym.Env):
                 f"Render mode '{mode}' is not supported. Please use: {self.metadata['render.modes']}."
             )
 
-    def close(self):
-        if not self._is_closed:
-            # destroy the window
-            if self._orbit_window is not None:
-                self._orbit_window.visible = False
-                self._orbit_window.destroy()
-                self._orbit_window = None
-            # update closing status
-            super().close()
-
     """
     Implementation specifics.
     """
@@ -357,147 +334,3 @@ class RLEnv(BaseEnv, gym.Env):
         #  -- add information to extra if timeout occurred due to episode length
         # Note: this is used by algorithms like PPO where time-outs are handled differently
         self.extras["time_outs"] = self.termination_manager.time_outs
-
-    """
-    Helper functions - GUI.
-    """
-
-    def _build_ui(self):
-        """Constructs the GUI for the environment."""
-        # need to import here to wait for the GUI extension to be loaded
-        import omni.isaac.ui.ui_utils as ui_utils
-        import omni.ui as ui
-        from omni.kit.window.extensions import SimpleCheckBox
-
-        # create window for UI
-        self._orbit_window = omni.ui.Window(
-            "Orbit", width=400, height=500, visible=True, dock_preference=ui.DockPreference.RIGHT_TOP
-        )
-        # dock next to properties window
-        asyncio.ensure_future(self._dock_window(window_title=self._orbit_window.title))
-
-        # keep a dictionary of stacks so that child environments can add their own UI elements
-        # this can be done by using the `with` context manager
-        self._orbit_window_elements = dict()
-        # create main frame
-        self._orbit_window_elements["main_frame"] = self._orbit_window.frame
-        with self._orbit_window_elements["main_frame"]:
-            # create main stack
-            self._orbit_window_elements["main_vstack"] = ui.VStack(spacing=5, height=0)
-            with self._orbit_window_elements["main_vstack"]:
-                # create collapsable frame for controls
-                self._orbit_window_elements["control_frame"] = ui.CollapsableFrame(
-                    title="Controls",
-                    width=ui.Fraction(1),
-                    height=0,
-                    collapsed=False,
-                    style=ui_utils.get_style(),
-                    horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED,
-                    vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_ON,
-                )
-                with self._orbit_window_elements["control_frame"]:
-                    # create stack for controls
-                    self._orbit_window_elements["controls_vstack"] = ui.VStack(spacing=5, height=0)
-                    with self._orbit_window_elements["controls_vstack"]:
-                        # create rendering mode dropdown
-                        render_mode_cfg = {
-                            "label": "Rendering Mode",
-                            "type": "dropdown",
-                            "default_val": self.sim.render_mode.value,
-                            "items": [member.name for member in self.sim.RenderMode if member.value >= 0],
-                            "tooltip": "Select a rendering mode\n" + self.sim.RenderMode.__doc__,
-                            "on_clicked_fn": lambda value: self.sim.set_render_mode(self.sim.RenderMode[value]),
-                        }
-                        self._orbit_window_elements["render_dropdown"] = ui_utils.dropdown_builder(**render_mode_cfg)
-
-                        # create a number slider to move to environment origin
-                        def viewport_camera_origin_fn(model: ui.SimpleIntModel):
-                            """Moves the viewport to the origin of the environment."""
-                            # obtain the origin of the environment
-                            origin = self.scene.env_origins[model.as_int - 1].detach().cpu().numpy()
-                            cam_eye = origin + np.asarray(self.cfg.viewer.eye)
-                            cam_target = origin + np.asarray(self.cfg.viewer.lookat)
-                            # set the camera view
-                            self.sim.set_camera_view(eye=cam_eye, target=cam_target)
-
-                        viewport_origin_cfg = {
-                            "label": "View Environment",
-                            "type": "button",
-                            "default_val": 1,
-                            "min": 1,
-                            "max": self.num_envs,
-                            "tooltip": "Move the viewport to the origin of the environment",
-                        }
-                        self._orbit_window_elements["viewport_btn"] = ui_utils.int_builder(**viewport_origin_cfg)
-                        # create a number slider to move to environment origin
-                        self._orbit_window_elements["viewport_btn"].add_value_changed_fn(viewport_camera_origin_fn)
-
-                # create collapsable frame for debug visualization
-                self._orbit_window_elements["debug_frame"] = ui.CollapsableFrame(
-                    title="Debug Visualization",
-                    width=ui.Fraction(1),
-                    height=0,
-                    collapsed=False,
-                    style=ui_utils.get_style(),
-                    horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED,
-                    vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_ON,
-                )
-                with self._orbit_window_elements["debug_frame"]:
-                    # create stack for debug visualization
-                    self._orbit_window_elements["debug_vstack"] = ui.VStack(spacing=5, height=0)
-                    with self._orbit_window_elements["debug_vstack"]:
-                        elements = [
-                            self.scene.terrain,
-                            self.command_manager,
-                            *self.scene.rigid_objects.values(),
-                            *self.scene.articulations.values(),
-                            *self.scene.sensors.values(),
-                        ]
-                        names = [
-                            "terrain",
-                            "commands",
-                            *self.scene.rigid_objects.keys(),
-                            *self.scene.articulations.keys(),
-                            *self.scene.sensors.keys(),
-                        ]
-                        # create one for the terrain
-                        for elem, name in zip(elements, names):
-                            if elem is not None:
-                                with ui.HStack():
-                                    # create the UI element
-                                    text = (
-                                        "Toggle debug visualization."
-                                        if elem.has_debug_vis_implementation
-                                        else "Debug visualization not implemented."
-                                    )
-                                    ui.Label(
-                                        name.replace("_", " ").title(),
-                                        width=ui_utils.LABEL_WIDTH - 12,
-                                        alignment=ui.Alignment.LEFT_CENTER,
-                                        tooltip=text,
-                                    )
-                                    self._orbit_window_elements[f"{name}_cb"] = SimpleCheckBox(
-                                        model=ui.SimpleBoolModel(),
-                                        enabled=elem.has_debug_vis_implementation,
-                                        checked=elem.cfg.debug_vis,
-                                        on_checked_fn=lambda value, e=weakref.proxy(elem): e.set_debug_vis(value),
-                                    )
-                                    ui_utils.add_line_rect_flourish()
-
-    async def _dock_window(self, window_title: str):
-        """Docks the orbit window to the property window."""
-        # need to import here to wait for the GUI extension to be loaded
-        import omni.ui as ui
-
-        for _ in range(5):
-            if ui.Workspace.get_window(window_title):
-                break
-            await self.sim.app.next_update_async()
-
-        # dock next to properties window
-        orbit_window = ui.Workspace.get_window(window_title)
-        property_window = ui.Workspace.get_window("Property")
-
-        if orbit_window and property_window:
-            orbit_window.dock_in(property_window, ui.DockPosition.SAME, 1.0)
-            orbit_window.focus()

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/rl_env_cfg.py

@@ -11,12 +11,16 @@ from omni.isaac.orbit.command_generators import CommandGeneratorBaseCfg
 from omni.isaac.orbit.utils import configclass
 
 from .base_env_cfg import BaseEnvCfg
+from .ui import RLEnvWindow
 
 
 @configclass
 class RLEnvCfg(BaseEnvCfg):
     """Configuration for a reinforcement learning environment."""
 
+    # ui settings
+    ui_window_class_type: type | None = RLEnvWindow
+
     # general settings
     episode_length_s: float = MISSING
     """Duration of an episode (in seconds)."""

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/ui/__init__.py

+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+This sub-package contains implementations of UI elements for the environments.
+
+The UI elements are used to control the environment and visualize the state of the environment.
+"""
+
+from __future__ import annotations
+
+# enable the extension for UI elements
+# this only needs to be done once
+from omni.isaac.core.utils.extensions import enable_extension
+
+enable_extension("omni.isaac.ui")
+
+# import all UI elements here
+from .base_env_window import BaseEnvWindow
+from .rl_env_window import RLEnvWindow
+
+__all__ = ["BaseEnvWindow", "RLEnvWindow"]

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/ui/base_env_window.py

+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import asyncio
+import numpy as np
+import weakref
+from typing import TYPE_CHECKING
+
+import omni.isaac.ui.ui_utils as ui_utils
+import omni.ui
+from omni.kit.window.extensions import SimpleCheckBox
+
+if TYPE_CHECKING:
+    from ..base_env import BaseEnv
+
+
+class BaseEnvWindow:
+    """Window manager for the basic environment.
+
+    This class creates a window that is used to control the environment. The window
+    contains controls for rendering, debug visualization, and other environment-specific
+    UI elements.
+
+    Users can add their own UI elements to the window by using the `with` context manager.
+    This can be done either be inheriting the class or by using the `env.window` object
+    directly from the standalone execution script.
+
+    Example for adding a UI element from the standalone execution script:
+        >>> with env.window.ui_window_elements["main_vstack"]:
+        >>>     ui.Label("My UI element")
+
+    """
+
+    def __init__(self, env: BaseEnv, window_name: str = "Orbit"):
+        """Initialize the window.
+
+        Args:
+            env: The environment object.
+            window_name: The name of the window. Defaults to "Orbit".
+        """
+        # store inputs
+        self.env = env
+        # store the viewer related variables
+        # -- locations of the camera
+        self._viewer_eye = np.array(self.env.cfg.viewer.eye)
+        self._viewer_lookat = np.array(self.env.cfg.viewer.lookat)
+        # -- which environment to lookat
+        self._viewer_env_index = 0
+        self._viewer_origin = self.env.scene.env_origins[self._viewer_env_index].detach().cpu().numpy()
+        # -- whether to follow an asset and which one
+        self._viewer_follow_enabled = False
+        self._viewer_assets_options = [
+            *self.env.scene.rigid_objects.keys(),
+            *self.env.scene.articulations.keys(),
+        ]
+        self._viewer_asset_name = self._viewer_assets_options[0]
+        # create a handle to the camera callback
+        self._viewer_follow_cam_handle = None
+
+        print("Creating window for environment.")
+        # create window for UI
+        self.ui_window = omni.ui.Window(
+            window_name, width=400, height=500, visible=True, dock_preference=omni.ui.DockPreference.RIGHT_TOP
+        )
+        # dock next to properties window
+        asyncio.ensure_future(self._dock_window(window_title=self.ui_window.title))
+
+        # keep a dictionary of stacks so that child environments can add their own UI elements
+        # this can be done by using the `with` context manager
+        self.ui_window_elements = dict()
+        # create main frame
+        self.ui_window_elements["main_frame"] = self.ui_window.frame
+        with self.ui_window_elements["main_frame"]:
+            # create main stack
+            self.ui_window_elements["main_vstack"] = omni.ui.VStack(spacing=5, height=0)
+            with self.ui_window_elements["main_vstack"]:
+                # create collapsable frame for simulation
+                self._build_sim_frame()
+                # create collapsable frame for viewer
+                self._build_viewer_frame()
+                # create collapsable frame for debug visualization
+                self._build_debug_vis_frame()
+
+    def __del__(self):
+        """Destructor for the window."""
+        # destroy the window
+        if self.ui_window is not None:
+            self.ui_window.visible = False
+            self.ui_window.destroy()
+            self.ui_window = None
+
+    """
+    Build sub-sections of the UI.
+    """
+
+    def _build_sim_frame(self):
+        """Builds the sim-related controls frame for the UI."""
+        # create collapsable frame for controls
+        self.ui_window_elements["sim_frame"] = omni.ui.CollapsableFrame(
+            title="Simulation Settings",
+            width=omni.ui.Fraction(1),
+            height=0,
+            collapsed=False,
+            style=ui_utils.get_style(),
+            horizontal_scrollbar_policy=omni.ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED,
+            vertical_scrollbar_policy=omni.ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_ON,
+        )
+        with self.ui_window_elements["sim_frame"]:
+            # create stack for controls
+            self.ui_window_elements["sim_vstack"] = omni.ui.VStack(spacing=5, height=0)
+            with self.ui_window_elements["sim_vstack"]:
+                # create rendering mode dropdown
+                render_mode_cfg = {
+                    "label": "Rendering Mode",
+                    "type": "dropdown",
+                    "default_val": self.env.sim.render_mode.value,
+                    "items": [member.name for member in self.env.sim.RenderMode if member.value >= 0],
+                    "tooltip": "Select a rendering mode\n" + self.env.sim.RenderMode.__doc__,
+                    "on_clicked_fn": lambda value: self.env.sim.set_render_mode(self.env.sim.RenderMode[value]),
+                }
+                self.ui_window_elements["render_dropdown"] = ui_utils.dropdown_builder(**render_mode_cfg)
+
+    def _build_viewer_frame(self):
+        """Build the viewer-related control frame for the UI."""
+        # create collapsable frame for viewer
+        self.ui_window_elements["viewer_frame"] = omni.ui.CollapsableFrame(
+            title="Viewer Settings",
+            width=omni.ui.Fraction(1),
+            height=0,
+            collapsed=False,
+            style=ui_utils.get_style(),
+            horizontal_scrollbar_policy=omni.ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED,
+            vertical_scrollbar_policy=omni.ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_ON,
+        )
+        with self.ui_window_elements["viewer_frame"]:
+            # create stack for controls
+            self.ui_window_elements["viewer_vstack"] = omni.ui.VStack(spacing=5, height=0)
+            with self.ui_window_elements["viewer_vstack"]:
+                # create a number slider to move to environment origin
+                viewport_origin_cfg = {
+                    "label": "Environment Index",
+                    "type": "button",
+                    "default_val": self._viewer_env_index + 1,
+                    "min": 1,
+                    "max": self.env.num_envs,
+                    "tooltip": "Move the viewport to the origin of the environment",
+                }
+                self.ui_window_elements["viewer_env_index"] = ui_utils.int_builder(**viewport_origin_cfg)
+                # create a number slider to move to environment origin
+                self.ui_window_elements["viewer_env_index"].add_value_changed_fn(self._set_viewer_env_index_fn)
+
+                # create a tracker for the camera location
+                viewer_follow_cfg = {
+                    "label": "Follow Mode",
+                    "type": "checkbox_dropdown",
+                    "default_val": [False, 0],
+                    "items": [name.replace("_", " ").title() for name in self._viewer_assets_options],
+                    "tooltip": "Follow an asset in the scene.",
+                    "on_clicked_fn": [self._toggle_viewer_follow_fn, self._set_viewer_follow_asset_fn],
+                }
+                self.ui_window_elements["viewer_follow"] = ui_utils.combo_cb_dropdown_builder(**viewer_follow_cfg)
+
+                # add viewer default eye and lookat locations
+                self.ui_window_elements["viewer_eye"] = ui_utils.xyz_builder(
+                    label="Camera Eye",
+                    tooltip="Modify the XYZ location of the viewer eye",
+                    default_val=self.env.cfg.viewer.eye,
+                    step=0.1,
+                    on_value_changed_fn=[self._set_viewer_location_fn] * 3,
+                )
+                self.ui_window_elements["viewer_lookat"] = ui_utils.xyz_builder(
+                    label="Camera Target",
+                    tooltip="Modify the XYZ location of the viewer target",
+                    default_val=self.env.cfg.viewer.lookat,
+                    step=0.1,
+                    on_value_changed_fn=[self._set_viewer_location_fn] * 3,
+                )
+
+    def _build_debug_vis_frame(self):
+        """Builds the debug visualization frame for various scene elements.
+
+        This function inquires the scene for all elements that have a debug visualization
+        implemented and creates a checkbox to toggle the debug visualization for each element
+        that has it implemented. If the element does not have a debug visualization implemented,
+        a label is created instead.
+        """
+        # create collapsable frame for debug visualization
+        self.ui_window_elements["debug_frame"] = omni.ui.CollapsableFrame(
+            title="Scene Debug Visualization",
+            width=omni.ui.Fraction(1),
+            height=0,
+            collapsed=False,
+            style=ui_utils.get_style(),
+            horizontal_scrollbar_policy=omni.ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED,
+            vertical_scrollbar_policy=omni.ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_ON,
+        )
+        with self.ui_window_elements["debug_frame"]:
+            # create stack for debug visualization
+            self.ui_window_elements["debug_vstack"] = omni.ui.VStack(spacing=5, height=0)
+            with self.ui_window_elements["debug_vstack"]:
+                elements = [
+                    self.env.scene.terrain,
+                    *self.env.scene.rigid_objects.values(),
+                    *self.env.scene.articulations.values(),
+                    *self.env.scene.sensors.values(),
+                ]
+                names = [
+                    "terrain",
+                    *self.env.scene.rigid_objects.keys(),
+                    *self.env.scene.articulations.keys(),
+                    *self.env.scene.sensors.keys(),
+                ]
+                # create one for the terrain
+                for elem, name in zip(elements, names):
+                    if elem is not None:
+                        self._create_debug_vis_ui_element(name, elem)
+
+    """
+    Custom callbacks for UI elements.
+    """
+
+    def _toggle_viewer_follow_fn(self, value: bool):
+        """Toggles the viewer follow mode."""
+        # store the desired env index
+        self._viewer_follow_enabled = value
+        # register the camera callback for the follow mode
+        if self._viewer_follow_enabled:
+            # create a subscriber for the post update event if it doesn't exist
+            if self._viewer_follow_cam_handle is None:
+                app_interface = omni.kit.app.get_app_interface()
+                self._viewer_follow_cam_handle = (
+                    app_interface.get_post_update_event_stream().create_subscription_to_pop(
+                        lambda event, obj=weakref.proxy(self): obj._viewer_follow_asset_callback(event)
+                    )
+                )
+        else:
+            # remove the subscriber if it exists
+            if self._viewer_follow_cam_handle is not None:
+                self._viewer_follow_cam_handle.unsubscribe()
+                self._viewer_follow_cam_handle = None
+        # update the camera view
+        self._update_camera_view()
+
+    def _set_viewer_follow_asset_fn(self, value: str):
+        """Sets the asset to follow."""
+        # find which index the asset is
+        fancy_names = [name.replace("_", " ").title() for name in self._viewer_assets_options]
+        # store the desired env index
+        self._viewer_asset_name = self._viewer_assets_options[fancy_names.index(value)]
+        # update the camera view
+        self._update_camera_view()
+
+    def _set_viewer_location_fn(self, model: omni.ui.SimpleFloatModel):
+        """Sets the viewer location based on the UI."""
+        # obtain the camera locations
+        for i in range(3):
+            self._viewer_eye[i] = self.ui_window_elements["viewer_eye"][i].get_value_as_float()
+            self._viewer_lookat[i] = self.ui_window_elements["viewer_lookat"][i].get_value_as_float()
+        # update the camera view
+        self._update_camera_view()
+
+    def _set_viewer_env_index_fn(self, model: omni.ui.SimpleIntModel):
+        """Moves the viewport to the origin of the environment."""
+        # store the desired env index
+        self._viewer_env_index = model.as_int - 1
+        # obtain the origin of the environment
+        if self._viewer_follow_enabled:
+            self._viewer_origin = self.env.scene[self._viewer_asset_name].data.root_pos_w[self._viewer_env_index]
+        else:
+            self._viewer_origin = self.env.scene.env_origins[self._viewer_env_index]
+        # origin
+        self._viewer_origin = self._viewer_origin.detach().cpu().numpy()
+        # update the camera view
+        self._update_camera_view()
+
+    def _viewer_follow_asset_callback(self, event):
+        # update the camera origins
+        self._viewer_origin = self.env.scene[self._viewer_asset_name].data.root_pos_w[self._viewer_env_index]
+        # origin
+        self._viewer_origin = self._viewer_origin.detach().cpu().numpy()
+        # update the camera view
+        self._update_camera_view()
+
+    """
+    Helper functions - UI updates.
+    """
+
+    def _update_camera_view(self, event=None):
+        """Updates the camera view based on the current settings."""
+        # set the camera locations
+        cam_eye = self._viewer_origin + self._viewer_eye
+        cam_target = self._viewer_origin + self._viewer_lookat
+        # set the camera view
+        self.env.sim.set_camera_view(eye=cam_eye, target=cam_target)
+
+    """
+    Helper functions - UI building.
+    """
+
+    def _create_debug_vis_ui_element(self, name: str, elem: object):
+        """Create a checkbox for toggling debug visualization for the given element."""
+        with omni.ui.HStack():
+            # create the UI element
+            text = (
+                "Toggle debug visualization."
+                if elem.has_debug_vis_implementation
+                else "Debug visualization not implemented."
+            )
+            omni.ui.Label(
+                name.replace("_", " ").title(),
+                width=ui_utils.LABEL_WIDTH - 12,
+                alignment=omni.ui.Alignment.LEFT_CENTER,
+                tooltip=text,
+            )
+            self.ui_window_elements[f"{name}_cb"] = SimpleCheckBox(
+                model=omni.ui.SimpleBoolModel(),
+                enabled=elem.has_debug_vis_implementation,
+                checked=elem.cfg.debug_vis,
+                on_checked_fn=lambda value, e=weakref.proxy(elem): e.set_debug_vis(value),
+            )
+            ui_utils.add_line_rect_flourish()
+
+    async def _dock_window(self, window_title: str):
+        """Docks the custom UI window to the property window."""
+        # wait for the window to be created
+        for _ in range(5):
+            if omni.ui.Workspace.get_window(window_title):
+                break
+            await self.env.sim.app.next_update_async()
+
+        # dock next to properties window
+        custom_window = omni.ui.Workspace.get_window(window_title)
+        property_window = omni.ui.Workspace.get_window("Property")
+        if custom_window and property_window:
+            custom_window.dock_in(property_window, omni.ui.DockPosition.SAME, 1.0)
+            custom_window.focus()

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/ui/rl_env_window.py

+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from .base_env_window import BaseEnvWindow
+
+if TYPE_CHECKING:
+    from ..rl_env import RLEnv
+
+
+class RLEnvWindow(BaseEnvWindow):
+    """Window manager for the RL environment.
+
+    On top of the basic environment window, this class adds controls for the RL environment.
+    This includes visualization of the command manager.
+    """
+
+    def __init__(self, env: RLEnv, window_name: str = "Orbit"):
+        """Initialize the window.
+
+        Args:
+            env: The environment object.
+            window_name: The name of the window. Defaults to "Orbit".
+        """
+        # initialize base window
+        super().__init__(env, window_name)
+
+        # add custom UI elements
+        with self.ui_window_elements["main_vstack"]:
+            with self.ui_window_elements["debug_frame"]:
+                with self.ui_window_elements["debug_vstack"]:
+                    # add command manager visualization
+                    self._create_debug_vis_ui_element("commands", self.env.command_manager)