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/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)