Adds live plots to managers

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

@@ -11,7 +11,7 @@
 import carb
 import omni.isaac.core.utils.torch as torch_utils
 
-from omni.isaac.lab.managers import ActionManager, EventManager, ObservationManager
+from omni.isaac.lab.managers import ActionManager, EnvLiveVisualizer, EventManager, ObservationManager
 from omni.isaac.lab.scene import InteractiveScene
 from omni.isaac.lab.sim import SimulationContext
 from omni.isaac.lab.utils.timer import Timer
@@ -143,7 +143,10 @@ 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:
+            self.setup_manager_visualizers()
             self._window = self.cfg.ui_window_class_type(self, window_name="IsaacLab")
+            # setup live visualizers
+
         else:
             # if no window, then we don't need to store the window
             self._window = None
@@ -221,6 +224,17 @@ def load_managers(self):
         # when all the other managers are created
         if self.__class__ == ManagerBasedEnv and "startup" in self.event_manager.available_modes:
             self.event_manager.apply(mode="startup")
+            # setup live visualizers
+            self.env_vis_manager = EnvLiveVisualizer(
+                cfg=self.cfg.live_visualizer,
+                managers={"action_manager": self.action_manager, "observation_manager": self.observation_manager},
+            )
+
+    def setup_manager_visualizers(self):
+        self.env_vis_manager = EnvLiveVisualizer(
+            cfg=self.cfg.live_visualizer,
+            managers={"action_manager": self.action_manager, "observation_manager": self.observation_manager},
+        )
 
     """
     Operations - MDP.
@@ -368,3 +382,7 @@ def _reset_idx(self, env_ids: Sequence[int]):
         # -- event manager
         info = self.event_manager.reset(env_ids)
         self.extras["log"].update(info)
+
+    @property
+    def manager_visualizers(self):
+        return self.env_vis_manager.manager_visualizers

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

@@ -12,6 +12,7 @@
 from dataclasses import MISSING
 
 import omni.isaac.lab.envs.mdp as mdp
+from omni.isaac.lab.managers import DefaultManagerBasedEnvLiveVisCfg
 from omni.isaac.lab.managers import EventTermCfg as EventTerm
 from omni.isaac.lab.scene import InteractiveSceneCfg
 from omni.isaac.lab.sim import SimulationCfg
@@ -95,3 +96,8 @@ class ManagerBasedEnvCfg:
 
     Please refer to the :class:`omni.isaac.lab.managers.EventManager` class for more details.
     """
+
+    live_visualizer: object = DefaultManagerBasedEnvLiveVisCfg()
+    """Live visualization settings. Defaults to showing all available actiona and observation terms.
+
+    Please refer to the :class:`omni.isaac.lab.managers.EnvLiveVisualizer`"""

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

@@ -15,7 +15,13 @@
 
 from omni.isaac.version import get_version
 
-from omni.isaac.lab.managers import CommandManager, CurriculumManager, RewardManager, TerminationManager
+from omni.isaac.lab.managers import (
+    CommandManager,
+    CurriculumManager,
+    EnvLiveVisualizer,
+    RewardManager,
+    TerminationManager,
+)
 
 from .common import VecEnvStepReturn
 from .manager_based_env import ManagerBasedEnv
@@ -132,6 +138,19 @@ def load_managers(self):
         if "startup" in self.event_manager.available_modes:
             self.event_manager.apply(mode="startup")
 
+    def setup_manager_visualizers(self):
+        self.env_vis_manager = EnvLiveVisualizer(
+            cfg=self.cfg.live_visualizer,
+            managers={
+                "action_manager": self.action_manager,
+                "observation_manager": self.observation_manager,
+                "command_manager": self.command_manager,
+                "termination_manager": self.termination_manager,
+                "reward_manager": self.reward_manager,
+                "curriculum_manager": self.curriculum_manager,
+            },
+        )
+
     """
     Operations - MDP
     """

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

@@ -5,6 +5,7 @@
 
 from dataclasses import MISSING
 
+from omni.isaac.lab.managers import DefaultManagerBasedRLEnvLiveVisCfg
 from omni.isaac.lab.utils import configclass
 
 from .manager_based_env_cfg import ManagerBasedEnvCfg
@@ -78,3 +79,8 @@ class ManagerBasedRLEnvCfg(ManagerBasedEnvCfg):
 
     Please refer to the :class:`omni.isaac.lab.managers.CommandManager` class for more details.
     """
+
+    live_visualizer: object = DefaultManagerBasedRLEnvLiveVisCfg()
+    """Live visualization settings. Defaults to showing all available actiona and observation terms.
+
+    Please refer tot he :class:`omni.isaac.lab.managers.EnvLiveVisualizer`"""

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py

@@ -17,7 +17,7 @@
 import omni.isaac.lab.utils.math as math_utils
 from omni.isaac.lab.assets import Articulation, RigidObject
 from omni.isaac.lab.managers import SceneEntityCfg
-from omni.isaac.lab.sensors import RayCaster
+from omni.isaac.lab.sensors import Camera, RayCaster, TiledCamera
 
 if TYPE_CHECKING:
     from omni.isaac.lab.envs import ManagerBasedEnv, ManagerBasedRLEnv
@@ -182,6 +182,40 @@ def body_incoming_wrench(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> tor
     return link_incoming_forces.view(env.num_envs, -1)
 
 
+def camera_rgb_image(env: ManagerBasedEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
+    """Camera sensor rgb image.
+
+    Args:
+        env: The IsaacLab environment.
+        sensor_cfg: The config referring to the Camera or TiledCamera sensor
+
+    Returns:
+        RGB camera data if available. Shape is determined by the camera.data.image_shape.
+    """
+
+    # extract camera or tiled camera sensor
+    camera: Camera | TiledCamera = env.scene.sensors[sensor_cfg.name]
+    # rgb data
+    return camera.data.output["rgb"]
+
+
+def camera_depth_image(env: ManagerBasedEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
+    """Camera sensor depth image, this refers to the distance_to_image_plane annotator.
+
+    Args:
+        env: The IsaacLab environment.
+        sensor_cfg: The config referring to the Camera or TiledCamera sensor
+
+    Returns:
+        Depth image camera data in meters if available. Shape is determined by the camera.data.image_shape.
+    """
+
+    # extract camera or tiled camera sensor
+    camera: Camera | TiledCamera = env.scene.sensors[sensor_cfg.name]
+    # depth data
+    return camera.data.output["distance_to_image_plane"]
+
+
 """
 Actions.
 """

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

@@ -16,6 +16,9 @@
 import omni.usd
 from pxr import PhysxSchema, Sdf, Usd, UsdGeom, UsdPhysics
 
+# from omni.isaac.lab.ui.widgets.ui_visualizer_mixin import UiVisualizerMixin
+from omni.isaac.lab.managers import ManagerLiveVisualizer
+
 if TYPE_CHECKING:
     import omni.ui
 
@@ -57,6 +60,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 +86,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 +210,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 +241,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 +384,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.
@@ -382,11 +412,22 @@ def _create_debug_vis_ui_element(self, name: str, elem: object):
             self.ui_window_elements[f"{name}_cb"] = SimpleCheckBox(
                 model=omni.ui.SimpleBoolModel(),
                 enabled=elem.has_debug_vis_implementation,
-                checked=elem.cfg.debug_vis,
+                checked=(hasattr(elem.cfg, "debug_vis") and elem.cfg.debug_vis)
+                or (hasattr(elem, "debug_vis") and elem.debug_vis),
                 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/__init__.py

@@ -15,6 +15,12 @@
 from .curriculum_manager import CurriculumManager
 from .event_manager import EventManager
 from .manager_base import ManagerBase, ManagerTermBase
+from .manager_live_visualizer import (
+    DefaultManagerBasedEnvLiveVisCfg,
+    DefaultManagerBasedRLEnvLiveVisCfg,
+    EnvLiveVisualizer,
+    ManagerLiveVisualizer,
+)
 from .manager_term_cfg import (
     ActionTermCfg,
     CommandTermCfg,

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
@@ -253,6 +254,20 @@ def has_debug_vis_implementation(self) -> bool:
     Operations.
     """
 
+    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.
+        """
+        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) -> bool:
         """Sets whether to visualize the action data.
         Args:

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

@@ -291,6 +291,21 @@ def has_debug_vis_implementation(self) -> bool:
     Operations.
     """
 
+    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.
+        """
+
+        terms = []
+        idx = 0
+        for name, term in self._terms.items():
+            term_commands = self._commands[env_idx, idx : idx + term.command.shape[1]].cpu()
+            terms.append((name, term_commands.tolist()))
+            idx += term.command.shape[1]
+        return terms
+
     def set_debug_vis(self, debug_vis: bool) -> bool:
         """Sets whether to visualize the command data.
 

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

@@ -131,6 +131,35 @@ 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.
+        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

@@ -192,6 +192,14 @@ def find_terms(self, name_keys: str | Sequence[str]) -> list[str]:
         # 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.
     """