Adds live plots to managers (#893)

# Description

This adds a UI interface to the Managers in the `ManagerBasedEnv` and
The `MangerBasedRLEnv`. Additions include:
- UI widgets for `LiveLinePlot` and `ImagePlot`
- `ManagerLiveVisualizer/Cfg`: Given a `ManagerBase` (i.e.
action_manager, observation_manager, etc) and a config file this class
creates the the interface between managers and the UI.
- `EnvLiveVisualizer`: A 'manager' of `ManagerLiveVisualizer`. This is
added to the `ManagerBasedEnv` but is only called during the
initialization of the managers in `load_managers`
- Adds `get_active_iterable_terms` implementation methods to
ActionManager, ObservationManager, CommandsManager, CurriculumManager,
RewardManager, and TerminationManager. This method exports the active
term data and labels for each manager and is called by
ManagerLiveVisualizer.
- Additions to `BaseEnvWindow` and `RLEnvWindow` to register
`ManagerLiveVisualizer` UI interfaces for the chosen managers.

## Screenshots
[Screencast from 09-06-2024 01:20:18
PM.webm](https://github.com/user-attachments/assets/3ef0191d-5446-41bd-b274-43d886fb2d70)

## Implementation

![image](https://github.com/user-attachments/assets/49bd5493-3311-4c5c-a87c-6bbcd76a60fe)

## Type of change

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

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] 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
- [ ] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

---------
Signed-off-by: jtigue-bdai <166445701+jtigue-bdai@users.noreply.github.com>
Signed-off-by: Pascal Roth <57946385+pascal-roth@users.noreply.github.com>
Signed-off-by: David Hoeller <dhoeller@nvidia.com>
Co-authored-by: zrene <zrene@ethz.ch>
Co-authored-by: James Tigue <jtigue@theaiinstitute.com>
Co-authored-by: jtigue-bdai <166445701+jtigue-bdai@users.noreply.github.com>
Co-authored-by: David Hoeller <dhoeller@nvidia.com>
Co-authored-by: Aravind EV <aravindev@live.in>
Co-authored-by: Kelly Guo <kellyg@nvidia.com>
Co-authored-by: Kelly Guo <kellyguo123@hotmail.com>

.pre-commit-config.yaml

@@ -38,7 +38,7 @@ repos:
       - id: pyupgrade
         args: ["--py310-plus"]
         # FIXME: This is a hack because Pytorch does not like: torch.Tensor | dict aliasing
-        exclude: "source/extensions/omni.isaac.lab/omni/isaac/lab/envs/common.py"
+        exclude: "source/extensions/omni.isaac.lab/omni/isaac/lab/envs/common.py|source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/image_plot.py"
   - repo: https://github.com/codespell-project/codespell
     rev: v2.2.6
     hooks:

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

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.28.0"
+version = "0.29.0"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

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

 Changelog
 ---------
 
+0.29.0 (2024-12-15)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added UI interface to the Managers in the ManagerBasedEnv and MangerBasedRLEnv classes.
+* Added UI widgets for :class:`LiveLinePlot` and :class:`ImagePlot`.
+* Added ``ManagerLiveVisualizer/Cfg``: Given a ManagerBase (i.e. action_manager, observation_manager, etc) and a config file this class creates the the interface between managers and the UI.
+* Added :class:`EnvLiveVisualizer`: A 'manager' of ManagerLiveVisualizer. This is added to the ManagerBasedEnv but is only called during the initialization of the managers in load_managers
+* Added ``get_active_iterable_terms`` implementation methods to ActionManager, ObservationManager, CommandsManager, CurriculumManager, RewardManager, and TerminationManager. This method exports the active term data and labels for each manager and is called by ManagerLiveVisualizer.
+* Additions to :class:`BaseEnvWindow` and :class:`RLEnvWindow` to register ManagerLiveVisualizer UI interfaces for the chosen managers.
+
+
 0.28.0 (2024-12-15)
 ~~~~~~~~~~~~~~~~~~~
 
@@ -93,7 +107,7 @@ Changed
 
 
 0.27.21 (2024-12-06)
-~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~
 
 Fixed
 ^^^^^

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_env.py

@@ -14,6 +14,7 @@ import omni.log
 from omni.isaac.lab.managers import ActionManager, EventManager, ObservationManager, RecorderManager
 from omni.isaac.lab.scene import InteractiveScene
 from omni.isaac.lab.sim import SimulationContext
+from omni.isaac.lab.ui.widgets import ManagerLiveVisualizer
 from omni.isaac.lab.utils.timer import Timer
 
 from .common import VecEnvObs
@@ -148,6 +149,8 @@ class ManagerBasedEnv:
         # 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:
+            # setup live visualizers
+            self.setup_manager_visualizers()
             self._window = self.cfg.ui_window_class_type(self, window_name="IsaacLab")
         else:
             # if no window, then we don't need to store the window
@@ -233,6 +236,14 @@ class ManagerBasedEnv:
         if self.__class__ == ManagerBasedEnv and "startup" in self.event_manager.available_modes:
             self.event_manager.apply(mode="startup")
 
+    def setup_manager_visualizers(self):
+        """Creates live visualizers for manager terms."""
+
+        self.manager_visualizers = {
+            "action_manager": ManagerLiveVisualizer(manager=self.action_manager),
+            "observation_manager": ManagerLiveVisualizer(manager=self.observation_manager),
+        }
+
     """
     Operations - MDP.
     """

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env.py

@@ -16,6 +16,7 @@ from typing import Any, ClassVar
 from omni.isaac.version import get_version
 
 from omni.isaac.lab.managers import CommandManager, CurriculumManager, RewardManager, TerminationManager
+from omni.isaac.lab.ui.widgets import ManagerLiveVisualizer
 
 from .common import VecEnvStepReturn
 from .manager_based_env import ManagerBasedEnv
@@ -132,6 +133,18 @@ class ManagerBasedRLEnv(ManagerBasedEnv, gym.Env):
         if "startup" in self.event_manager.available_modes:
             self.event_manager.apply(mode="startup")
 
+    def setup_manager_visualizers(self):
+        """Creates live visualizers for manager terms."""
+
+        self.manager_visualizers = {
+            "action_manager": ManagerLiveVisualizer(manager=self.action_manager),
+            "observation_manager": ManagerLiveVisualizer(manager=self.observation_manager),
+            "command_manager": ManagerLiveVisualizer(manager=self.command_manager),
+            "termination_manager": ManagerLiveVisualizer(manager=self.termination_manager),
+            "reward_manager": ManagerLiveVisualizer(manager=self.reward_manager),
+            "curriculum_manager": ManagerLiveVisualizer(manager=self.curriculum_manager),
+        }
+
     """
     Operations - MDP
     """

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

@@ -16,6 +16,8 @@ import omni.kit.commands
 import omni.usd
 from pxr import PhysxSchema, Sdf, Usd, UsdGeom, UsdPhysics
 
+from omni.isaac.lab.ui.widgets import ManagerLiveVisualizer
+
 if TYPE_CHECKING:
     import omni.ui
 
@@ -57,6 +59,9 @@ class BaseEnvWindow:
             *self.env.scene.articulations.keys(),
         ]
 
+        # Listeners for environment selection changes
+        self._ui_listeners: list[ManagerLiveVisualizer] = []
+
         print("Creating window for environment.")
         # create window for UI
         self.ui_window = omni.ui.Window(
@@ -80,6 +85,10 @@ class BaseEnvWindow:
                 self._build_viewer_frame()
                 # create collapsable frame for debug visualization
                 self._build_debug_vis_frame()
+                with self.ui_window_elements["debug_frame"]:
+                    with self.ui_window_elements["debug_vstack"]:
+                        self._visualize_manager(title="Actions", class_name="action_manager")
+                        self._visualize_manager(title="Observations", class_name="observation_manager")
 
     def __del__(self):
         """Destructor for the window."""
@@ -200,9 +209,6 @@ class BaseEnvWindow:
         that has it implemented. If the element does not have a debug visualization implemented,
         a label is created instead.
         """
-        # import omni.isaac.ui.ui_utils as ui_utils
-        # import omni.ui
-
         # create collapsable frame for debug visualization
         self.ui_window_elements["debug_frame"] = omni.ui.CollapsableFrame(
             title="Scene Debug Visualization",
@@ -234,6 +240,26 @@ class BaseEnvWindow:
                     if elem is not None:
                         self._create_debug_vis_ui_element(name, elem)
 
+    def _visualize_manager(self, title: str, class_name: str) -> None:
+        """Checks if the attribute with the name 'class_name' can be visualized. If yes, create vis interface.
+
+        Args:
+            title: The title of the manager visualization frame.
+            class_name: The name of the manager to visualize.
+        """
+
+        if hasattr(self.env, class_name) and class_name in self.env.manager_visualizers:
+            manager = self.env.manager_visualizers[class_name]
+            if hasattr(manager, "has_debug_vis_implementation"):
+                self._create_debug_vis_ui_element(title, manager)
+            else:
+                print(
+                    f"ManagerLiveVisualizer cannot be created for manager: {class_name}, has_debug_vis_implementation"
+                    " does not exist"
+                )
+        else:
+            print(f"ManagerLiveVisualizer cannot be created for manager: {class_name}, Manager does not exist")
+
     """
     Custom callbacks for UI elements.
     """
@@ -357,6 +383,9 @@ class BaseEnvWindow:
             raise ValueError("Viewport camera controller is not initialized! Please check the rendering mode.")
         # store the desired env index, UI is 1-indexed
         vcc.set_view_env_index(model.as_int - 1)
+        # notify additional listeners
+        for listener in self._ui_listeners:
+            listener.set_env_selection(model.as_int - 1)
 
     """
     Helper functions - UI building.
@@ -379,14 +408,30 @@ class BaseEnvWindow:
                 alignment=omni.ui.Alignment.LEFT_CENTER,
                 tooltip=text,
             )
+            has_cfg = hasattr(elem, "cfg") and elem.cfg is not None
+            is_checked = False
+            if has_cfg:
+                is_checked = (hasattr(elem.cfg, "debug_vis") and elem.cfg.debug_vis) or (
+                    hasattr(elem, "debug_vis") and elem.debug_vis
+                )
             self.ui_window_elements[f"{name}_cb"] = SimpleCheckBox(
                 model=omni.ui.SimpleBoolModel(),
                 enabled=elem.has_debug_vis_implementation,
-                checked=elem.cfg.debug_vis if elem.cfg else False,
+                checked=is_checked,
                 on_checked_fn=lambda value, e=weakref.proxy(elem): e.set_debug_vis(value),
             )
             omni.isaac.ui.ui_utils.add_line_rect_flourish()
 
+        # Create a panel for the debug visualization
+        if isinstance(elem, ManagerLiveVisualizer):
+            self.ui_window_elements[f"{name}_panel"] = omni.ui.Frame(width=omni.ui.Fraction(1))
+            if not elem.set_vis_frame(self.ui_window_elements[f"{name}_panel"]):
+                print(f"Frame failed to set for ManagerLiveVisualizer: {name}")
+
+        # Add listener for environment selection changes
+        if isinstance(elem, ManagerLiveVisualizer):
+            self._ui_listeners.append(elem)
+
     async def _dock_window(self, window_title: str):
         """Docks the custom UI window to the property window."""
         # wait for the window to be created

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/ui/manager_based_rl_env_window.py

@@ -34,5 +34,7 @@ class ManagerBasedRLEnvWindow(BaseEnvWindow):
         with self.ui_window_elements["main_vstack"]:
             with self.ui_window_elements["debug_frame"]:
                 with self.ui_window_elements["debug_vstack"]:
-                    self._create_debug_vis_ui_element("commands", self.env.command_manager)
-                    self._create_debug_vis_ui_element("actions", self.env.action_manager)
+                    self._visualize_manager(title="Commands", class_name="command_manager")
+                    self._visualize_manager(title="Rewards", class_name="reward_manager")
+                    self._visualize_manager(title="Curriculum", class_name="curriculum_manager")
+                    self._visualize_manager(title="Termination", class_name="termination_manager")

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/action_manager.py

@@ -106,6 +106,7 @@ class ActionTerm(ManagerTermBase):
         # check if debug visualization is supported
         if not self.has_debug_vis_implementation:
             return False
+
         # toggle debug visualization objects
         self._set_debug_vis_impl(debug_vis)
         # toggle debug visualization handles
@@ -262,7 +263,26 @@ class ActionManager(ManagerBase):
     Operations.
     """
 
-    def set_debug_vis(self, debug_vis: bool) -> bool:
+    def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]:
+        """Returns the active terms as iterable sequence of tuples.
+
+        The first element of the tuple is the name of the term and the second element is the raw value(s) of the term.
+
+        Args:
+            env_idx: The specific environment to pull the active terms from.
+
+        Returns:
+            The active terms.
+        """
+        terms = []
+        idx = 0
+        for name, term in self._terms.items():
+            term_actions = self._action[env_idx, idx : idx + term.action_dim].cpu()
+            terms.append((name, term_actions.tolist()))
+            idx += term.action_dim
+        return terms
+
+    def set_debug_vis(self, debug_vis: bool):
         """Sets whether to visualize the action data.
         Args:
             debug_vis: Whether to visualize the action data.

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/command_manager.py

@@ -296,7 +296,26 @@ class CommandManager(ManagerBase):
     Operations.
     """
 
-    def set_debug_vis(self, debug_vis: bool) -> bool:
+    def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]:
+        """Returns the active terms as iterable sequence of tuples.
+
+        The first element of the tuple is the name of the term and the second element is the raw value(s) of the term.
+
+        Args:
+            env_idx: The specific environment to pull the active terms from.
+
+        Returns:
+            The active terms.
+        """
+
+        terms = []
+        idx = 0
+        for name, term in self._terms.items():
+            terms.append((name, term.command[env_idx].cpu().tolist()))
+            idx += term.command.shape[1]
+        return terms
+
+    def set_debug_vis(self, debug_vis: bool):
         """Sets whether to visualize the command data.
 
         Args:

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/curriculum_manager.py

@@ -138,6 +138,40 @@ class CurriculumManager(ManagerBase):
             state = term_cfg.func(self._env, env_ids, **term_cfg.params)
             self._curriculum_state[name] = state
 
+    def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]:
+        """Returns the active terms as iterable sequence of tuples.
+
+        The first element of the tuple is the name of the term and the second element is the raw value(s) of the term.
+
+        Args:
+            env_idx: The specific environment to pull the active terms from.
+
+        Returns:
+            The active terms.
+        """
+
+        terms = []
+
+        for term_name, term_state in self._curriculum_state.items():
+            if term_state is not None:
+                # deal with dict
+                data = []
+
+                if isinstance(term_state, dict):
+                    # each key is a separate state to log
+                    for key, value in term_state.items():
+                        if isinstance(value, torch.Tensor):
+                            value = value.item()
+                        terms[term_name].append(value)
+                else:
+                    # log directly if not a dict
+                    if isinstance(term_state, torch.Tensor):
+                        term_state = term_state.item()
+                    data.append(term_state)
+                terms.append((term_name, data))
+
+        return terms
+
     """
     Helper functions.
     """

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/manager_base.py

@@ -193,6 +193,16 @@ class ManagerBase(ABC):
         # return the matching names
         return string_utils.resolve_matching_names(name_keys, list_of_strings)[1]
 
+    def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]:
+        """Returns the active terms as iterable sequence of tuples.
+
+        The first element of the tuple is the name of the term and the second element is the raw value(s) of the term.
+
+        Returns:
+            The active terms.
+        """
+        raise NotImplementedError
+
     """
     Implementation specific.
     """

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/observation_manager.py

@@ -100,6 +100,9 @@ class ObservationManager(ManagerBase):
             else:
                 self._group_obs_dim[group_name] = group_term_dims
 
+        # Stores the latest observations.
+        self._obs_buffer: dict[str, torch.Tensor | dict[str, torch.Tensor]] | None = None
+
     def __str__(self) -> str:
         """Returns: A string representation for the observation manager."""
         msg = f"<ObservationManager> contains {len(self._group_obs_term_names)} groups.\n"
@@ -130,6 +133,43 @@ class ObservationManager(ManagerBase):
 
         return msg
 
+    def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]:
+        """Returns the active terms as iterable sequence of tuples.
+
+        The first element of the tuple is the name of the term and the second element is the raw value(s) of the term.
+
+        Args:
+            env_idx: The specific environment to pull the active terms from.
+
+        Returns:
+            The active terms.
+        """
+        terms = []
+
+        if self._obs_buffer is None:
+            self.compute()
+        obs_buffer: dict[str, torch.Tensor | dict[str, torch.Tensor]] = self._obs_buffer
+
+        for group_name, _ in self._group_obs_dim.items():
+            if not self.group_obs_concatenate[group_name]:
+                for name, term in obs_buffer[group_name].items():
+                    terms.append((group_name + "-" + name, term[env_idx].cpu().tolist()))
+                continue
+
+            idx = 0
+            # add info for each term
+            data = obs_buffer[group_name]
+            for name, shape in zip(
+                self._group_obs_term_names[group_name],
+                self._group_obs_term_dim[group_name],
+            ):
+                data_length = np.prod(shape)
+                term = data[env_idx, idx : idx + data_length]
+                terms.append((group_name + "-" + name, term.cpu().tolist()))
+                idx += data_length
+
+        return terms
+
     """
     Properties.
     """
@@ -212,6 +252,9 @@ class ObservationManager(ManagerBase):
         for group_name in self._group_obs_term_names:
             obs_buffer[group_name] = self.compute_group(group_name)
         # otherwise return a dict with observations of all groups
+
+        # Cache the observations.
+        self._obs_buffer = obs_buffer
         return obs_buffer
 
     def compute_group(self, group_name: str) -> torch.Tensor | dict[str, torch.Tensor]:

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/reward_manager.py

@@ -61,6 +61,9 @@ class RewardManager(ManagerBase):
         # create buffer for managing reward per environment
         self._reward_buf = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
 
+        # Buffer which stores the current step reward for each term for each environment
+        self._step_reward = torch.zeros((self.num_envs, len(self._term_names)), dtype=torch.float, device=self.device)
+
     def __str__(self) -> str:
         """Returns: A string representation for reward manager."""
         msg = f"<RewardManager> contains {len(self._term_names)} active terms.\n"
@@ -148,6 +151,9 @@ class RewardManager(ManagerBase):
             # update episodic sum
             self._episode_sums[name] += value
 
+            # Update current reward for this step.
+            self._step_reward[:, self._term_names.index(name)] = value / dt
+
         return self._reward_buf
 
     """
@@ -186,6 +192,22 @@ class RewardManager(ManagerBase):
         # return the configuration
         return self._term_cfgs[self._term_names.index(term_name)]
 
+    def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]:
+        """Returns the active terms as iterable sequence of tuples.
+
+        The first element of the tuple is the name of the term and the second element is the raw value(s) of the term.
+
+        Args:
+            env_idx: The specific environment to pull the active terms from.
+
+        Returns:
+            The active terms.
+        """
+        terms = []
+        for idx, name in enumerate(self._term_names):
+            terms.append((name, [self._step_reward[env_idx, idx].cpu().item()]))
+        return terms
+
     """
     Helper functions.
     """

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/termination_manager.py

@@ -184,6 +184,22 @@ class TerminationManager(ManagerBase):
         """
         return self._term_dones[name]
 
+    def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]:
+        """Returns the active terms as iterable sequence of tuples.
+
+        The first element of the tuple is the name of the term and the second element is the raw value(s) of the term.
+
+        Args:
+            env_idx: The specific environment to pull the active terms from.
+
+        Returns:
+            The active terms.
+        """
+        terms = []
+        for key in self._term_dones.keys():
+            terms.append((key, [self._term_dones[key][env_idx].float().cpu().item()]))
+        return terms
+
     """
     Operations - Term settings.
     """

source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/__init__.py

+# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from .image_plot import ImagePlot
+from .line_plot import LiveLinePlot
+from .manager_live_visualizer import ManagerLiveVisualizer
+from .ui_visualizer_base import UiVisualizerBase

source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/image_plot.py

+# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+import numpy as np
+from matplotlib import cm
+from typing import TYPE_CHECKING, Optional
+
+import carb
+import omni
+import omni.log
+
+from .ui_widget_wrapper import UIWidgetWrapper
+
+if TYPE_CHECKING:
+    import omni.isaac.ui
+    import omni.ui
+
+
+class ImagePlot(UIWidgetWrapper):
+    """An image plot widget to display live data.
+
+    It has the following Layout where the mode frame is only useful for depth images:
+    +-------------------------------------------------------+
+    |                  containing_frame                     |
+    |+-----------------------------------------------------+|
+    |                   main_plot_frame                     |
+    ||+---------------------------------------------------+||
+    |||                    plot_frames                    |||
+    |||                                                   |||
+    |||                                                   |||
+    |||               (Image Plot Data)                   |||
+    |||                                                   |||
+    |||                                                   |||
+    |||+-------------------------------------------------+|||
+    |||                   mode_frame                      |||
+    |||                                                   |||
+    |||    [x][Absolute] [x][Grayscaled] [ ][Colorized]   |||
+    |+-----------------------------------------------------+|
+    +-------------------------------------------------------+
+
+    """
+
+    def __init__(
+        self,
+        image: Optional[np.ndarray] = None,
+        label: str = "",
+        widget_height: int = 200,
+        show_min_max: bool = True,
+        unit: tuple[float, str] = (1, ""),
+    ):
+        """Create an XY plot UI Widget with axis scaling, legends, and support for multiple plots.
+
+        Overlapping data is most accurately plotted when centered in the frame with reasonable axis scaling.
+        Pressing down the mouse gives the x and y values of each function at an x coordinate.
+
+        Args:
+            image: Image to display
+            label: Short descriptive text to the left of the plot
+            widget_height: Height of the plot in pixels
+            show_min_max: Whether to show the min and max values of the image
+            unit: Tuple of (scale, name) for the unit of the image
+        """
+        self._show_min_max = show_min_max
+        self._unit_scale = unit[0]
+        self._unit_name = unit[1]
+
+        self._curr_mode = "None"
+
+        self._has_built = False
+
+        self._enabled = True
+
+        self._byte_provider = omni.ui.ByteImageProvider()
+        if image is None:
+            carb.log_warn("image is NONE")
+            image = np.ones((480, 640, 3), dtype=np.uint8) * 255
+            image[:, :, 0] = 0
+            image[:, :240, 1] = 0
+
+        # if image is channel first, convert to channel last
+        if image.ndim == 3 and image.shape[0] in [1, 3, 4]:
+            image = np.moveaxis(image, 0, -1)
+
+        self._aspect_ratio = image.shape[1] / image.shape[0]
+        self._widget_height = widget_height
+        self._label = label
+        self.update_image(image)
+
+        plot_frame = self._create_ui_widget()
+
+        super().__init__(plot_frame)
+
+    def setEnabled(self, enabled: bool):
+        self._enabled = enabled
+
+    def update_image(self, image: np.ndarray):
+        if not self._enabled:
+            return
+
+        # if image is channel first, convert to channel last
+        if image.ndim == 3 and image.shape[0] in [1, 3, 4]:
+            image = np.moveaxis(image, 0, -1)
+
+        height, width = image.shape[:2]
+
+        if self._curr_mode == "Normalization":
+            image = (image - image.min()) / (image.max() - image.min())
+            image = (image * 255).astype(np.uint8)
+        elif self._curr_mode == "Colorization":
+            if image.ndim == 3 and image.shape[2] == 3:
+                omni.log.warn("Colorization mode is only available for single channel images")
+            else:
+                image = (image - image.min()) / (image.max() - image.min())
+                colormap = cm.get_cmap("jet")
+                if image.ndim == 3 and image.shape[2] == 1:
+                    image = (colormap(image).squeeze(2) * 255).astype(np.uint8)
+                else:
+                    image = (colormap(image) * 255).astype(np.uint8)
+
+        # convert image to 4-channel RGBA
+        if image.ndim == 2 or (image.ndim == 3 and image.shape[2] == 1):
+            image = np.dstack((image, image, image, np.full((height, width, 1), 255, dtype=np.uint8)))
+
+        elif image.ndim == 3 and image.shape[2] == 3:
+            image = np.dstack((image, np.full((height, width, 1), 255, dtype=np.uint8)))
+
+        self._byte_provider.set_bytes_data(image.flatten().data, [width, height])
+
+    def update_min_max(self, image: np.ndarray):
+        if self._show_min_max and hasattr(self, "_min_max_label"):
+            non_inf = image[np.isfinite(image)].flatten()
+            if len(non_inf) > 0:
+                self._min_max_label.text = self._get_unit_description(
+                    np.min(non_inf), np.max(non_inf), np.median(non_inf)
+                )
+            else:
+                self._min_max_label.text = self._get_unit_description(0, 0)
+
+    def _create_ui_widget(self):
+        containing_frame = omni.ui.Frame(build_fn=self._build_widget)
+        return containing_frame
+
+    def _get_unit_description(self, min_value: float, max_value: float, median_value: float = None):
+        return (
+            f"Min: {min_value * self._unit_scale:.2f} {self._unit_name} Max:"
+            f" {max_value * self._unit_scale:.2f} {self._unit_name}"
+            + (f" Median: {median_value * self._unit_scale:.2f} {self._unit_name}" if median_value is not None else "")
+        )
+
+    def _build_widget(self):
+
+        with omni.ui.VStack(spacing=3):
+            with omni.ui.HStack():
+                # Write the leftmost label for what this plot is
+                omni.ui.Label(
+                    self._label, width=omni.isaac.ui.ui_utils.LABEL_WIDTH, alignment=omni.ui.Alignment.LEFT_TOP
+                )
+                with omni.ui.Frame(width=self._aspect_ratio * self._widget_height, height=self._widget_height):
+                    self._base_plot = omni.ui.ImageWithProvider(self._byte_provider)
+
+            if self._show_min_max:
+                self._min_max_label = omni.ui.Label(self._get_unit_description(0, 0))
+
+            omni.ui.Spacer(height=8)
+            self._mode_frame = omni.ui.Frame(build_fn=self._build_mode_frame)
+
+            omni.ui.Spacer(width=5)
+        self._has_built = True
+
+    def _build_mode_frame(self):
+        """Build the frame containing the mode selection for the plots.
+
+        This is an internal function to build the frame containing the mode selection for the plots. This function
+        should only be called from within the build function of a frame.
+
+        The built widget has the following layout:
+        +-------------------------------------------------------+
+        |                   legends_frame                       |
+        ||+---------------------------------------------------+||
+        |||                                                   |||
+        |||    [x][Series 1] [x][Series 2] [ ][Series 3]      |||
+        |||                                                   |||
+        |||+-------------------------------------------------+|||
+        |+-----------------------------------------------------+|
+        +-------------------------------------------------------+
+        """
+        with omni.ui.HStack():
+            with omni.ui.HStack():
+
+                def _change_mode(value):
+                    self._curr_mode = value
+
+                omni.isaac.ui.ui_utils.dropdown_builder(
+                    label="Mode",
+                    type="dropdown",
+                    items=["Original", "Normalization", "Colorization"],
+                    tooltip="Select a mode",
+                    on_clicked_fn=_change_mode,
+                )

source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/ui_visualizer_base.py

+# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import inspect
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import omni.ui
+
+
+class UiVisualizerBase:
+    """Base Class for components that support debug visualizations that requires access to some UI elements.
+
+    This class provides a set of functions that can be used to assign ui interfaces.
+
+    The following functions are provided:
+
+    * :func:`set_debug_vis`: Assigns a debug visualization interface. This function is called by the main UI
+        when the checkbox for debug visualization is toggled.
+    * :func:`set_vis_frame`: Assigns a small frame within the isaac lab tab that can be used to visualize debug
+        information. Such as e.g. plots or images. It is called by the main UI on startup to create the frame.
+    * :func:`set_window`: Assigngs the main window that is used by the main UI. This allows the user
+        to have full controller over all UI elements. But be warned, with great power comes great responsibility.
+    """
+
+    """
+    Exposed Properties
+    """
+
+    @property
+    def has_debug_vis_implementation(self) -> bool:
+        """Whether the component has a debug visualization implemented."""
+        # check if function raises NotImplementedError
+        source_code = inspect.getsource(self._set_debug_vis_impl)
+        return "NotImplementedError" not in source_code
+
+    @property
+    def has_vis_frame_implementation(self) -> bool:
+        """Whether the component has a debug visualization implemented."""
+        # check if function raises NotImplementedError
+        source_code = inspect.getsource(self._set_vis_frame_impl)
+        return "NotImplementedError" not in source_code
+
+    @property
+    def has_window_implementation(self) -> bool:
+        """Whether the component has a debug visualization implemented."""
+        # check if function raises NotImplementedError
+        source_code = inspect.getsource(self._set_window_impl)
+        return "NotImplementedError" not in source_code
+
+    @property
+    def has_env_selection_implementation(self) -> bool:
+        """Whether the component has a debug visualization implemented."""
+        # check if function raises NotImplementedError
+        source_code = inspect.getsource(self._set_env_selection_impl)
+        return "NotImplementedError" not in source_code
+
+    """
+    Exposed Setters
+    """
+
+    def set_env_selection(self, env_selection: int) -> bool:
+        """Sets the selected environment id.
+
+        This function is called by the main UI when the user selects a different environment.
+
+        Args:
+            env_selection: The currently selected environment id.
+
+        Returns:
+            Whether the environment selection was successfully set. False if the component
+            does not support environment selection.
+        """
+        # check if environment selection is supported
+        if not self.has_env_selection_implementation:
+            return False
+        # set environment selection
+        self._set_env_selection_impl(env_selection)
+        return True
+
+    def set_window(self, window: omni.ui.Window) -> bool:
+        """Sets the current main ui window.
+
+        This function is called by the main UI when the window is created. It allows the component
+        to add custom UI elements to the window or to control the window and its elements.
+
+        Args:
+            window: The ui window.
+
+        Returns:
+            Whether the window was successfully set. False if the component
+            does not support this functionality.
+        """
+        # check if window is supported
+        if not self.has_window_implementation:
+            return False
+        # set window
+        self._set_window_impl(window)
+        return True
+
+    def set_vis_frame(self, vis_frame: omni.ui.Frame) -> bool:
+        """Sets the debug visualization frame.
+
+        This function is called by the main UI when the window is created. It allows the component
+        to modify a small frame within the orbit tab that can be used to visualize debug information.
+
+        Args:
+            vis_frame: The debug visualization frame.
+
+        Returns:
+            Whether the debug visualization frame was successfully set. False if the component
+            does not support debug visualization.
+        """
+        # check if debug visualization is supported
+        if not self.has_vis_frame_implementation:
+            return False
+        # set debug visualization frame
+        self._set_vis_frame_impl(vis_frame)
+        return True
+
+    """
+    Internal Implementation
+    """
+
+    def _set_env_selection_impl(self, env_idx: int):
+        """Set the environment selection."""
+        raise NotImplementedError(f"Environment selection is not implemented for {self.__class__.__name__}.")
+
+    def _set_window_impl(self, window: omni.ui.Window):
+        """Set the window."""
+        raise NotImplementedError(f"Window is not implemented for {self.__class__.__name__}.")
+
+    def _set_debug_vis_impl(self, debug_vis: bool):
+        """Set debug visualization state."""
+        raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.")
+
+    def _set_vis_frame_impl(self, vis_frame: omni.ui.Frame):
+        """Set debug visualization into visualization objects.
+
+        This function is responsible for creating the visualization objects if they don't exist
+        and input ``debug_vis`` is True. If the visualization objects exist, the function should
+        set their visibility into the stage.
+        """
+        raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.")

source/extensions/omni.isaac.lab/omni/isaac/lab/ui/widgets/ui_widget_wrapper.py

+# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+# This file has been adapted from _isaac_sim/exts/omni.isaac.ui/omni/isaac/ui/element_wrappers/base_ui_element_wrappers.py
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import omni
+
+if TYPE_CHECKING:
+    import omni.ui
+
+
+class UIWidgetWrapper:
+    """
+    Base class for creating wrappers around any subclass of omni.ui.Widget in order to provide an easy interface
+    for creating and managing specific types of widgets such as state buttons or file pickers.
+    """
+
+    def __init__(self, container_frame: omni.ui.Frame):
+        self._container_frame = container_frame
+
+    @property
+    def container_frame(self) -> omni.ui.Frame:
+        return self._container_frame
+
+    @property
+    def enabled(self) -> bool:
+        return self.container_frame.enabled
+
+    @enabled.setter
+    def enabled(self, value: bool):
+        self.container_frame.enabled = value
+
+    @property
+    def visible(self) -> bool:
+        return self.container_frame.visible
+
+    @visible.setter
+    def visible(self, value: bool):
+        self.container_frame.visible = value
+
+    def cleanup(self):
+        """
+        Perform any necessary cleanup
+        """
+        pass