Adds a viewer camera controller to base environment (#411)

# Description

This PR adds the functionality to set up the view camera with respect to
different frames: world center, Environment Center, and an assets root
frame. In the last case, if the asset moves, the camera will track it.
This is a useful feature for recording videos during the training
logging.

## Type of change

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

## Screenshots

Base Env:

https://github.com/isaac-orbit/orbit/assets/142246516/9b4f1580-2bd8-4d6c-b405-843b66a2d0b0

RL Env:

https://github.com/isaac-orbit/orbit/assets/3355155/4f51d6fa-3dfb-4d58-a83b-2845a5a80fe2


## 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
- [ ] I have run all the tests with `./orbit.sh --test` and they pass
- [ ] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [ ] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

---------
Signed-off-by: Farbod Farshidian <ffarshidian@theaiinstitute.com>
Co-authored-by: James Smith <jsmith@theaiinstitute.com>
Co-authored-by: Mayank Mittal <mittalma@leggedrobotics.com>

docs/source/api/orbit/omni.isaac.orbit.envs.ui.rst

@@ -9,6 +9,7 @@
 
         BaseEnvWindow
         RLTaskEnvWindow
+        ViewportCameraController
 
 Base Environment UI
 -------------------
@@ -22,3 +23,9 @@ RL Task Environment UI
 .. autoclass:: RLTaskEnvWindow
     :members:
     :show-inheritance:
+
+Viewport Camera Controller
+--------------------------
+
+.. autoclass:: ViewportCameraController
+    :members:

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

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.10.25"
+version = "0.10.26"
 
 # Description
 title = "ORBIT framework for Robot Learning"

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

 Changelog
 ---------
 
+0.10.26 (2024-02-26)
+~~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added a viewport camera controller class to the :class:`omni.isaac.orbit.envs.BaseEnv`. This is useful
+  for applications where the user wants to render the viewport from different perspectives even when the
+  simulation is running in headless mode.
+
+
 0.10.25 (2024-02-26)
 ~~~~~~~~~~~~~~~~~~~~
 
 Fixed
 ^^^^^
 
-* Ensures that all path arguments in :mod:`omni.isaac.orbit.sim.utils` are cast to ``str``. Previously we had handled path types as strings without casting.
+* Ensures that all path arguments in :mod:`omni.isaac.orbit.sim.utils` are cast to ``str``. Previously,
+  we had handled path types as strings without casting.
 
 
 0.10.24 (2024-02-26)

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/base_env.py

@@ -18,6 +18,7 @@ from omni.isaac.orbit.sim import SimulationContext
 from omni.isaac.orbit.utils.timer import Timer
 
 from .base_env_cfg import BaseEnvCfg
+from .ui import ViewportCameraController
 
 VecEnvObs = dict[str, torch.Tensor | dict[str, torch.Tensor]]
 """Observation returned by the environment.
@@ -101,10 +102,6 @@ class BaseEnv:
             self.sim = SimulationContext(self.cfg.sim)
         else:
             raise RuntimeError("Simulation context already exists. Cannot create a new one.")
-        # set camera view for "/OmniverseKit_Persp" camera
-        # viewport is not available in other rendering modes so the function will throw a warning
-        if self.sim.render_mode >= self.sim.RenderMode.PARTIAL_RENDERING:
-            self.sim.set_camera_view(eye=self.cfg.viewer.eye, target=self.cfg.viewer.lookat)
 
         # print useful information
         print("[INFO]: Base environment:")
@@ -120,6 +117,15 @@ class BaseEnv:
             self.scene = InteractiveScene(self.cfg.scene)
         print("[INFO]: Scene manager: ", self.scene)
 
+        # set up camera viewport controller
+        # viewport is not available in other rendering modes so the function will throw a warning
+        # FIXME: This needs to be fixed in the future when we unify the UI functionalities even for
+        # non-rendering modes.
+        if self.sim.render_mode >= self.sim.RenderMode.PARTIAL_RENDERING:
+            self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer)
+        else:
+            self.viewport_camera_controller = None
+
         # play the simulator to activate physics handles
         # note: this activates the physics simulation view that exposes TensorAPIs
         # note: when started in extension mode, first call sim.reset_async() and then initialize the managers
@@ -299,6 +305,7 @@ class BaseEnv:
             del self.observation_manager
             del self.randomization_manager
             del self.scene
+            del self.viewport_camera_controller
             # clear callbacks and instance
             self.sim.clear_all_callbacks()
             self.sim.clear_instance()

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/base_env_cfg.py

@@ -12,6 +12,7 @@ configuring the environment instances, viewer settings, and simulation parameter
 from __future__ import annotations
 
 from dataclasses import MISSING
+from typing import Literal
 
 import omni.isaac.orbit.envs.mdp as mdp
 from omni.isaac.orbit.managers import RandomizationTermCfg as RandTerm
@@ -28,17 +29,42 @@ class ViewerCfg:
 
     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)
     """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.
+    """The camera prim path to record images from. Default is "/OmniverseKit_Persp",
+    which is the default camera in the viewport.
     """
+
     resolution: tuple[int, int] = (1280, 720)
     """The resolution (width, height) of the camera specified using :attr:`cam_prim_path`.
     Default is (1280, 720).
     """
 
+    origin_type: Literal["world", "env", "asset_root"] = "world"
+    """The frame in which the camera position (eye) and target (lookat) are defined in. Default is "world".
+
+    Available options are:
+
+    * ``"world"``: The origin of the world.
+    * ``"env"``: The origin of the environment defined by :attr:`env_index`.
+    * ``"asset_root"``: The center of the asset defined by :attr:`asset_name` in environment :attr:`env_index`.
+    """
+
+    env_index: int = 0
+    """The environment index for frame origin. Default is 0.
+
+    This quantity is only effective if :attr:`origin` is set to "env" or "asset_root".
+    """
+
+    asset_name: str | None = None
+    """The asset name in the interactive scene for the frame origin. Default is None.
+
+    This quantity is only effective if :attr:`origin` is set to "asset_root".
+    """
+
 
 @configclass
 class DefaultRandomizationManagerCfg:
@@ -62,7 +88,7 @@ class BaseEnvCfg:
     """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.
+    """The class type of the UI window. Default is None.
 
     If None, then no UI window is created.
 
@@ -84,5 +110,5 @@ class BaseEnvCfg:
     actions: object = MISSING
     """Action space settings."""
     randomization: object = DefaultRandomizationManagerCfg()
-    """Randomization settings. Defaults to the default randomization manager, which resets
+    """Randomization settings. Default is the default randomization manager, which resets
     the scene to its default state."""

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/ui/__init__.py

@@ -12,3 +12,4 @@ toggling different debug visualization tools, and other user-defined functionali
 
 from .base_env_window import BaseEnvWindow
 from .rl_task_env_window import RLTaskEnvWindow
+from .viewport_camera_controller import ViewportCameraController

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

@@ -6,7 +6,6 @@
 from __future__ import annotations
 
 import asyncio
-import numpy as np
 import weakref
 from typing import TYPE_CHECKING
 
@@ -44,22 +43,14 @@ class BaseEnvWindow:
         """
         # store inputs
         self.env = env
-        # store the viewer related variables
-        # -- locations of the camera
-        self._viewer_eye = np.array(self.env.cfg.viewer.eye)
-        self._viewer_lookat = np.array(self.env.cfg.viewer.lookat)
-        # -- which environment to lookat
-        self._viewer_env_index = 0
-        self._viewer_origin = self.env.scene.env_origins[self._viewer_env_index].detach().cpu().numpy()
-        # -- whether to follow an asset and which one
-        self._viewer_follow_enabled = False
+        # prepare the list of assets that can be followed by the viewport camera
+        # note that the first two options are "World" and "Env" which are special cases
         self._viewer_assets_options = [
+            "World",
+            "Env",
             *self.env.scene.rigid_objects.keys(),
             *self.env.scene.articulations.keys(),
         ]
-        self._viewer_asset_name = self._viewer_assets_options[0]
-        # create a handle to the camera callback
-        self._viewer_follow_cam_handle = None
 
         print("Creating window for environment.")
         # create window for UI
@@ -141,13 +132,14 @@ class BaseEnvWindow:
             self.ui_window_elements["viewer_vstack"] = omni.ui.VStack(spacing=5, height=0)
             with self.ui_window_elements["viewer_vstack"]:
                 # create a number slider to move to environment origin
+                # NOTE: slider is 1-indexed, whereas the env index is 0-indexed
                 viewport_origin_cfg = {
                     "label": "Environment Index",
                     "type": "button",
-                    "default_val": self._viewer_env_index + 1,
+                    "default_val": self.env.cfg.viewer.env_index + 1,
                     "min": 1,
                     "max": self.env.num_envs,
-                    "tooltip": "Move the viewport to the origin of the environment",
+                    "tooltip": "The environment index to follow. Only effective if follow mode is not 'World'.",
                 }
                 self.ui_window_elements["viewer_env_index"] = ui_utils.int_builder(**viewport_origin_cfg)
                 # create a number slider to move to environment origin
@@ -156,25 +148,25 @@ class BaseEnvWindow:
                 # create a tracker for the camera location
                 viewer_follow_cfg = {
                     "label": "Follow Mode",
-                    "type": "checkbox_dropdown",
-                    "default_val": [False, 0],
+                    "type": "dropdown",
+                    "default_val": 0,
                     "items": [name.replace("_", " ").title() for name in self._viewer_assets_options],
-                    "tooltip": "Follow an asset in the scene.",
-                    "on_clicked_fn": [self._toggle_viewer_follow_fn, self._set_viewer_follow_asset_fn],
+                    "tooltip": "Select the viewport camera following mode.",
+                    "on_clicked_fn": self._set_viewer_origin_type_fn,
                 }
-                self.ui_window_elements["viewer_follow"] = ui_utils.combo_cb_dropdown_builder(**viewer_follow_cfg)
+                self.ui_window_elements["viewer_follow"] = ui_utils.dropdown_builder(**viewer_follow_cfg)
 
                 # add viewer default eye and lookat locations
                 self.ui_window_elements["viewer_eye"] = ui_utils.xyz_builder(
                     label="Camera Eye",
-                    tooltip="Modify the XYZ location of the viewer eye",
+                    tooltip="Modify the XYZ location of the viewer eye.",
                     default_val=self.env.cfg.viewer.eye,
                     step=0.1,
                     on_value_changed_fn=[self._set_viewer_location_fn] * 3,
                 )
                 self.ui_window_elements["viewer_lookat"] = ui_utils.xyz_builder(
                     label="Camera Target",
-                    tooltip="Modify the XYZ location of the viewer target",
+                    tooltip="Modify the XYZ location of the viewer target.",
                     default_val=self.env.cfg.viewer.lookat,
                     step=0.1,
                     on_value_changed_fn=[self._set_viewer_location_fn] * 3,
@@ -223,79 +215,46 @@ class BaseEnvWindow:
     Custom callbacks for UI elements.
     """
 
-    def _toggle_viewer_follow_fn(self, value: bool):
-        """Toggles the viewer follow mode."""
-        # store the desired env index
-        self._viewer_follow_enabled = value
-        # register the camera callback for the follow mode
-        if self._viewer_follow_enabled:
-            # create a subscriber for the post update event if it doesn't exist
-            if self._viewer_follow_cam_handle is None:
-                app_interface = omni.kit.app.get_app_interface()
-                self._viewer_follow_cam_handle = (
-                    app_interface.get_post_update_event_stream().create_subscription_to_pop(
-                        lambda event, obj=weakref.proxy(self): obj._viewer_follow_asset_callback(event)
-                    )
-                )
+    def _set_viewer_origin_type_fn(self, value: str):
+        """Sets the origin of the viewport's camera. This is based on the drop-down menu in the UI."""
+        # Extract the viewport camera controller from environment
+        vcc = self.env.viewport_camera_controller
+        if vcc is None:
+            raise ValueError("Viewport camera controller is not initialized! Please check the rendering mode.")
+
+        # Based on origin type, update the camera view
+        if value == "World":
+            vcc.update_view_to_world()
+        elif value == "Env":
+            vcc.update_view_to_env()
         else:
-            # remove the subscriber if it exists
-            if self._viewer_follow_cam_handle is not None:
-                self._viewer_follow_cam_handle.unsubscribe()
-                self._viewer_follow_cam_handle = None
-        # update the camera view
-        self._update_camera_view()
-
-    def _set_viewer_follow_asset_fn(self, value: str):
-        """Sets the asset to follow."""
             # find which index the asset is
             fancy_names = [name.replace("_", " ").title() for name in self._viewer_assets_options]
             # store the desired env index
-        self._viewer_asset_name = self._viewer_assets_options[fancy_names.index(value)]
+            viewer_asset_name = self._viewer_assets_options[fancy_names.index(value)]
             # update the camera view
-        self._update_camera_view()
+            vcc.update_view_to_asset_root(viewer_asset_name)
 
     def _set_viewer_location_fn(self, model: omni.ui.SimpleFloatModel):
-        """Sets the viewer location based on the UI."""
-        # obtain the camera locations
-        for i in range(3):
-            self._viewer_eye[i] = self.ui_window_elements["viewer_eye"][i].get_value_as_float()
-            self._viewer_lookat[i] = self.ui_window_elements["viewer_lookat"][i].get_value_as_float()
+        """Sets the viewport camera location based on the UI."""
+        # access the viewport camera controller (for brevity)
+        vcc = self.env.viewport_camera_controller
+        if vcc is None:
+            raise ValueError("Viewport camera controller is not initialized! Please check the rendering mode.")
+        # obtain the camera locations and set them in the viewpoint camera controller
+        eye = [self.ui_window_elements["viewer_eye"][i].get_value_as_float() for i in range(3)]
+        lookat = [self.ui_window_elements["viewer_lookat"][i].get_value_as_float() for i in range(3)]
         # update the camera view
-        self._update_camera_view()
+        vcc.update_view_location(eye, lookat)
 
     def _set_viewer_env_index_fn(self, model: omni.ui.SimpleIntModel):
-        """Moves the viewport to the origin of the environment."""
-        # store the desired env index
-        self._viewer_env_index = model.as_int - 1
-        # obtain the origin of the environment
-        if self._viewer_follow_enabled:
-            self._viewer_origin = self.env.scene[self._viewer_asset_name].data.root_pos_w[self._viewer_env_index]
-        else:
-            self._viewer_origin = self.env.scene.env_origins[self._viewer_env_index]
-        # origin
-        self._viewer_origin = self._viewer_origin.detach().cpu().numpy()
-        # update the camera view
-        self._update_camera_view()
-
-    def _viewer_follow_asset_callback(self, event):
-        # update the camera origins
-        self._viewer_origin = self.env.scene[self._viewer_asset_name].data.root_pos_w[self._viewer_env_index]
-        # origin
-        self._viewer_origin = self._viewer_origin.detach().cpu().numpy()
-        # update the camera view
-        self._update_camera_view()
-
-    """
-    Helper functions - UI updates.
-    """
-
-    def _update_camera_view(self, event=None):
-        """Updates the camera view based on the current settings."""
-        # set the camera locations
-        cam_eye = self._viewer_origin + self._viewer_eye
-        cam_target = self._viewer_origin + self._viewer_lookat
-        # set the camera view
-        self.env.sim.set_camera_view(eye=cam_eye, target=cam_target)
+        """Sets the environment index and updates the camera if in 'env' origin mode."""
+        # access the viewport camera controller (for brevity)
+        vcc = self.env.viewport_camera_controller
+        if vcc is None:
+            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)
 
     """
     Helper functions - UI building.

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/ui/viewport_camera_controller.py

+# Copyright (c) 2022-2024, The ORBIT Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import copy
+import numpy as np
+import torch
+import weakref
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+import omni.kit.app
+import omni.timeline
+
+if TYPE_CHECKING:
+    from omni.isaac.orbit.envs import BaseEnv, ViewerCfg
+
+
+class ViewportCameraController:
+    """This class handles controlling the camera associated with a viewport in the simulator.
+
+    It can be used to set the viewpoint camera to track different origin types:
+
+    - **world**: the center of the world (static)
+    - **env**: the center of an environment (static)
+    - **asset_root**: the root of an asset in the scene (e.g. tracking a robot moving in the scene)
+
+    On creation, the camera is set to track the origin type specified in the configuration.
+
+    For the :attr:`asset_root` origin type, the camera is updated at each rendering step to track the asset's
+    root position. For this, it registers a callback to the post update event stream from the simulation app.
+    """
+
+    def __init__(self, env: BaseEnv, cfg: ViewerCfg):
+        """Initialize the ViewportCameraController.
+
+        Args:
+            env: The environment.
+            cfg: The configuration for the viewport camera controller.
+
+        Raises:
+            ValueError: If origin type is configured to be "env" but :attr:`cfg.env_index` is out of bounds.
+            ValueError: If origin type is configured to be "asset_root" but :attr:`cfg.asset_name` is unset.
+
+        """
+        # store inputs
+        self._env = env
+        self._cfg = copy.deepcopy(cfg)
+        # cast viewer eye and look-at to numpy arrays
+        self.default_cam_eye = np.array(self._cfg.eye)
+        self.default_cam_lookat = np.array(self._cfg.lookat)
+
+        # set the camera origins
+        if self.cfg.origin_type == "env":
+            # check that the env_index is within bounds
+            self.set_view_env_index(self.cfg.env_index)
+            # set the camera origin to the center of the environment
+            self.update_view_to_env()
+        elif self.cfg.origin_type == "asset_root":
+            # note: we do not yet update camera for tracking an asset origin, as the asset may not yet be
+            # in the scene when this is called. Instead, we subscribe to the post update event to update the camera
+            # at each rendering step.
+            if self.cfg.asset_name is None:
+                raise ValueError(f"No asset name provided for viewer with origin type: '{self.cfg.origin_type}'.")
+
+        # subscribe to post update event so that camera view can be updated at each rendering step
+        app_interface = omni.kit.app.get_app_interface()
+        app_event_stream = app_interface.get_post_update_event_stream()
+        self._viewport_camera_update_handle = app_event_stream.create_subscription_to_pop(
+            lambda event, obj=weakref.proxy(self): obj._update_tracking_callback(event)
+        )
+
+    def __del__(self):
+        """Unsubscribe from the callback."""
+        # use hasattr to handle case where __init__ has not completed before __del__ is called
+        if hasattr(self, "_viewport_camera_update_handle") and self._viewport_camera_update_handle is not None:
+            self._viewport_camera_update_handle.unsubscribe()
+            self._viewport_camera_update_handle = None
+
+    """
+    Properties
+    """
+
+    @property
+    def cfg(self) -> ViewerCfg:
+        """The configuration for the viewer."""
+        return self._cfg
+
+    """
+    Public Functions
+    """
+
+    def set_view_env_index(self, env_index: int):
+        """Sets the environment index for the camera view.
+
+        Args:
+            env_index: The index of the environment to set the camera view to.
+
+        Raises:
+            ValueError: If the environment index is out of bounds. It should be between 0 and num_envs - 1.
+        """
+        # check that the env_index is within bounds
+        if env_index < 0 or env_index >= self._env.num_envs:
+            raise ValueError(
+                f"Out of range value for attribute 'env_index': {env_index}."
+                f" Expected a value between 0 and {self._env.num_envs - 1} for the current environment."
+            )
+        # update the environment index
+        self.cfg.env_index = env_index
+        # update the camera view if the origin is set to env type (since, the camera view is static)
+        # note: for assets, the camera view is updated at each rendering step
+        if self.cfg.origin_type == "env":
+            self.update_view_to_env()
+
+    def update_view_to_world(self):
+        """Updates the viewer's origin to the origin of the world which is (0, 0, 0)."""
+        # set origin type to world
+        self.cfg.origin_type = "world"
+        # update the camera origins
+        self.viewer_origin = torch.zeros(3)
+        # update the camera view
+        self.update_view_location()
+
+    def update_view_to_env(self):
+        """Updates the viewer's origin to the origin of the selected environment."""
+        # set origin type to world
+        self.cfg.origin_type = "env"
+        # update the camera origins
+        self.viewer_origin = self._env.scene.env_origins[self.cfg.env_index]
+        # update the camera view
+        self.update_view_location()
+
+    def update_view_to_asset_root(self, asset_name: str):
+        """Updates the viewer's origin based upon the root of an asset in the scene.
+
+        Args:
+            asset_name: The name of the asset in the scene. The name should match the name of the
+                asset in the scene.
+
+        Raises:
+            ValueError: If the asset is not in the scene.
+        """
+        # check if the asset is in the scene
+        if self.cfg.asset_name != asset_name:
+            asset_entities = [*self._env.scene.rigid_objects.keys(), *self._env.scene.articulations.keys()]
+            if asset_name not in asset_entities:
+                raise ValueError(f"Asset '{asset_name}' is not in the scene. Available entities: {asset_entities}.")
+        # update the asset name
+        self.cfg.asset_name = asset_name
+        # set origin type to asset_root
+        self.cfg.origin_type = "asset_root"
+        # update the camera origins
+        self.viewer_origin = self._env.scene[self.cfg.asset_name].data.root_pos_w[self.cfg.env_index]
+        # update the camera view
+        self.update_view_location()
+
+    def update_view_location(self, eye: Sequence[float] | None = None, lookat: Sequence[float] | None = None):
+        """Updates the camera view pose based on the current viewer origin and the eye and lookat positions.
+
+        Args:
+            eye: The eye position of the camera. If None, the current eye position is used.
+            lookat: The lookat position of the camera. If None, the current lookat position is used.
+        """
+        # store the camera view pose for later use
+        if eye is not None:
+            self.default_cam_eye = np.asarray(eye)
+        if lookat is not None:
+            self.default_cam_lookat = np.asarray(lookat)
+        # set the camera locations
+        viewer_origin = self.viewer_origin.detach().cpu().numpy()
+        cam_eye = viewer_origin + self.default_cam_eye
+        cam_target = viewer_origin + self.default_cam_lookat
+
+        # set the camera view
+        self._env.sim.set_camera_view(eye=cam_eye, target=cam_target)
+
+    """
+    Private Functions
+    """
+
+    def _update_tracking_callback(self, event):
+        """Updates the camera view at each rendering step."""
+        # update the camera view if the origin is set to asset_root
+        # in other cases, the camera view is static and does not need to be updated continuously
+        if self.cfg.origin_type == "asset_root" and self.cfg.asset_name is not None:
+            self.update_view_to_asset_root(self.cfg.asset_name)