Source code for omni.isaac.orbit.envs.ui.rl_task_env_window

# Copyright (c) 2022-2024, The ORBIT Project Developers.
# 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_task_env import RLTaskEnv


class RLTaskEnvWindow(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: RLTaskEnv, 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)