Add camera pose editors plugins

mani_skill/envs/sapien_env.py

@@ -39,6 +39,9 @@
 from mani_skill.utils.structs import Actor, Articulation
 from mani_skill.utils.structs.pose import Pose
 from mani_skill.utils.structs.types import Array, SimConfig
+from mani_skill.utils.visualization.viewer_camera_control import (
+    ViewerCameraControlPlugin,
+)
 from mani_skill.utils.visualization.misc import tile_images
 
 
@@ -115,6 +118,9 @@ class BaseEnv(gym.Env):
         enhanced_determinism (bool): By default this is False and env resets will reset the episode RNG only when a seed / seed list is given.
             If True, the environment will reset the episode RNG upon each reset regardless of whether a seed is provided.
             Generally enhanced_determinisim is not needed and users are recommended to pass seeds into the env reset function instead.
+
+        enable_camera_pose_editing (bool): If True, the human viewer exposes an in-window camera editor that lets you click camera frustums
+            and drag a gizmo to update their poses live.
     """
 
     # fmt: off
@@ -208,8 +214,10 @@ def __init__(
         render_backend: str = "gpu",
         parallel_in_single_scene: bool = False,
         enhanced_determinism: bool = False,
+        enable_camera_pose_editing: bool = False,
     ):
         self._enhanced_determinism = enhanced_determinism
+        self._enable_camera_pose_editing = enable_camera_pose_editing
 
         self.num_envs = num_envs
         self.reconfiguration_freq = reconfiguration_freq if reconfiguration_freq is not None else 0
@@ -312,6 +320,7 @@ def __init__(
         # Render mode
         self.render_mode = render_mode
         self._viewer = None
+        self._viewer_camera_control_plugin = None
 
         # Lighting
         self.enable_shadow = enable_shadow
@@ -1250,6 +1259,7 @@ def _close_viewer(self):
             return
         self._viewer.close()
         self._viewer = None
+        self._viewer_camera_control_plugin = None
 
     @cached_property
     def segmentation_id_map(self):
@@ -1331,6 +1341,22 @@ def set_state(self, state: Array, env_idx: torch.Tensor = None):
     def viewer(self):
         return self._viewer
 
+    def _ensure_viewer_camera_control_plugin(self):
+        if self._viewer is None:
+            return
+        existing_plugin = sapien_utils.get_obj_by_type(
+            self._viewer.plugins, ViewerCameraControlPlugin, is_unique=False
+        )
+        if isinstance(existing_plugin, list):
+            self._viewer_camera_control_plugin = existing_plugin[0]
+        elif existing_plugin is not None:
+            self._viewer_camera_control_plugin = existing_plugin
+        if self._viewer_camera_control_plugin is None:
+            self._viewer_camera_control_plugin = ViewerCameraControlPlugin(self)
+            self._viewer.plugins.append(self._viewer_camera_control_plugin)
+            self._viewer_camera_control_plugin.init(self._viewer)
+        self._viewer_camera_control_plugin.notify_scene_change()
+
     def _setup_viewer(self):
         """Setup the interactive viewer.
 
@@ -1347,6 +1373,8 @@ def _setup_viewer(self):
         )
         control_window.show_joint_axes = False
         control_window.show_camera_linesets = False
+        if self._enable_camera_pose_editing:
+            self._ensure_viewer_camera_control_plugin()
         if "render_camera" in self._human_render_cameras:
             self._viewer.set_camera_pose(
                 self._human_render_cameras["render_camera"].camera.global_pose[0].sp

mani_skill/examples/demo_random_action.py

@@ -49,6 +49,9 @@ class Args:
     pause: Annotated[bool, tyro.conf.arg(aliases=["-p"])] = False
     """If using human render mode, auto pauses the simulation upon loading"""
 
+    edit_camera_poses: bool = False
+    """Enable in-viewer camera pose editing when using the human render mode"""
+
     quiet: bool = False
     """Disable verbose output."""
 
@@ -83,6 +86,7 @@ def main(args: Args):
         render_backend=args.render_backend,
         enable_shadow=True,
         parallel_in_single_scene=parallel_in_single_scene,
+        enable_camera_pose_editing=args.edit_camera_poses,
     )
     if args.robot_uids is not None:
         env_kwargs["robot_uids"] = tuple(args.robot_uids.split(","))
@@ -107,21 +111,26 @@ def main(args: Args):
     obs, _ = env.reset(seed=args.seed, options=dict(reconfigure=True))
     if args.seed is not None and env.action_space is not None:
             env.action_space.seed(args.seed[0])
+    pause_after_viewer_boot = False
     if args.render_mode == "human":
         viewer = env.render()
-        if isinstance(viewer, sapien.utils.Viewer):
-            viewer.paused = args.pause
-        env.render()
+        pause_after_viewer_boot = args.pause and isinstance(viewer, sapien.utils.Viewer)
+        if args.edit_camera_poses:
+            print("Camera pose editing enabled: click a camera frustum in the viewer and drag the gizmo in the Camera Editor window.")
     while True:
+        if args.render_mode == "human":
+            viewer = env.render()
+            if isinstance(viewer, sapien.utils.Viewer) and pause_after_viewer_boot:
+                viewer.paused = True
+                pause_after_viewer_boot = False
+                continue
         action = env.action_space.sample() if env.action_space is not None else None
         obs, reward, terminated, truncated, info = env.step(action)
         if verbose:
             print("reward", reward)
             print("terminated", terminated)
             print("truncated", truncated)
             print("info", info)
-        if args.render_mode == "human":
-            env.render()
         if args.render_mode is None or args.render_mode != "human":
             if (terminated | truncated).any():
                 break