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>

Author: 6 people

.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

@@ -1,7 +1,7 @@
 [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

@@ -1,6 +1,20 @@
 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 @@
 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 @@ def __init__(self, cfg: ManagerBasedEnvCfg):
         # 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 @@ def load_managers(self):
         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 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 @@ def load_managers(self):
         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.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 @@ def __init__(self, env: ManagerBasedEnv, window_name: str = "IsaacLab"):
             *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 @@ def __init__(self, env: ManagerBasedEnv, window_name: str = "IsaacLab"):
                 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 @@ def _build_debug_vis_frame(self):
         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 @@ def _build_debug_vis_frame(self):
                     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 @@ def _set_viewer_env_index_fn(self, model: omni.ui.SimpleIntModel):
             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 @@ def _create_debug_vis_ui_element(self, name: str, elem: object):
                 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 @@ def __init__(self, env: ManagerBasedRLEnv, window_name: str = "IsaacLab"):
         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 @@ def set_debug_vis(self, debug_vis: bool) -> bool:
         # 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 @@ def has_debug_vis_implementation(self) -> bool:
     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 @@ def has_debug_vis_implementation(self) -> bool:
     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 @@ def compute(self, env_ids: Sequence[int] | None = None):
             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.
     """