modify opencv renderer to support multiple camera view

robosuite/demos/demo_gripper_interaction.py

@@ -14,7 +14,7 @@
 from robosuite.models.arenas.table_arena import TableArena
 from robosuite.models.grippers import PandaGripper, RethinkGripper
 from robosuite.models.objects import BoxObject
-from robosuite.utils import OpenCVRenderer
+from robosuite.renderers.viewer import OpenCVViewer
 from robosuite.utils.binding_utils import MjRenderContextOffscreen, MjSim
 from robosuite.utils.mjcf_utils import new_actuator, new_joint
 
@@ -66,7 +66,7 @@
     model = world.get_model(mode="mujoco")
 
     sim = MjSim(model)
-    viewer = OpenCVRenderer(sim)
+    viewer = OpenCVViewer(sim)
     render_context = MjRenderContextOffscreen(sim, device_id=-1)
     sim.add_render_context(render_context)
 

robosuite/demos/demo_multi_camera.py

@@ -0,0 +1,70 @@
+import time
+
+from robosuite.robots import MobileRobot
+from robosuite.utils.input_utils import *
+
+MAX_FR = 25  # max frame rate for running simluation
+
+if __name__ == "__main__":
+
+    # Create dict to hold options that will be passed to env creation call
+    options = {}
+
+    # print welcome info
+    print("Welcome to robosuite v{}!".format(suite.__version__))
+    print(suite.__logo__)
+
+    # Choose environment and add it to options
+    options["env_name"] = choose_environment()
+
+    # If a multi-arm environment has been chosen, choose configuration and appropriate robot(s)
+    if "TwoArm" in options["env_name"]:
+        # Choose env config and add it to options
+        options["env_configuration"] = choose_multi_arm_config()
+
+        # If chosen configuration was bimanual, the corresponding robot must be Baxter. Else, have user choose robots
+        if options["env_configuration"] == "single-robot":
+            options["robots"] = choose_robots(exclude_bimanual=False, use_humanoids=True, exclude_single_arm=True)
+        else:
+            options["robots"] = []
+
+            # Have user choose two robots
+            for i in range(2):
+                print("Please choose Robot {}...\n".format(i))
+                options["robots"].append(choose_robots(exclude_bimanual=False, use_humanoids=True))
+    # If a humanoid environment has been chosen, choose humanoid robots
+    elif "Humanoid" in options["env_name"]:
+        options["robots"] = choose_robots(use_humanoids=True)
+    else:
+        options["robots"] = choose_robots(exclude_bimanual=False, use_humanoids=True)
+
+    # initialize the task
+    env = suite.make(
+        **options,
+        has_renderer=True,
+        has_offscreen_renderer=False,
+        ignore_done=True,
+        use_camera_obs=False,
+        control_freq=20,
+        renderer="mujoco",
+    )
+    env.reset()
+
+    camera_name = ["agentview", "birdview"]
+    env.viewer.set_camera(camera_name=camera_name)
+    for robot in env.robots:
+        if isinstance(robot, MobileRobot):
+            robot.enable_parts(legs=False, base=False)
+
+    # do visualization
+    for i in range(10000):
+        start = time.time()
+        action = np.random.randn(*env.action_spec[0].shape)
+        obs, reward, done, _ = env.step(action)
+        env.render()
+
+        # limit frame rate if necessary
+        elapsed = time.time() - start
+        diff = 1 / MAX_FR - elapsed
+        if diff > 0:
+            time.sleep(diff)

robosuite/environments/base.py

@@ -9,7 +9,8 @@
 import robosuite.macros as macros
 import robosuite.utils.sim_utils as SU
 from robosuite.renderers.base import load_renderer_config
-from robosuite.utils import OpenCVRenderer, SimulationError, XMLError
+from robosuite.renderers.viewer import OpenCVViewer
+from robosuite.utils import SimulationError, XMLError
 from robosuite.utils.binding_utils import MjRenderContextOffscreen, MjSim
 
 REGISTERED_ENVS = {}
@@ -175,7 +176,7 @@ def initialize_renderer(self):
         if self.renderer == "mujoco":
             pass
         elif self.renderer == "mjviewer":
-            from robosuite.renderers.mjviewer.mjviewer_renderer import MjviewerRenderer
+            from robosuite.renderers.viewer import MjviewerRenderer
 
             if self.render_camera is not None:
                 camera_id = self.sim.model.camera_name2id(self.render_camera)
@@ -318,7 +319,7 @@ def _reset_internal(self):
         # create visualization screen or renderer
         if self.has_renderer and self.viewer is None:
             if self.renderer == "mujoco":
-                self.viewer = OpenCVRenderer(self.sim)
+                self.viewer = OpenCVViewer(self.sim)
 
                 # Set the camera angle for viewing
                 if self.render_camera is not None:

robosuite/models/grippers/gripper_tester.py

@@ -9,7 +9,7 @@
 from robosuite.models.arenas.table_arena import TableArena
 from robosuite.models.objects import BoxObject
 from robosuite.models.world import MujocoWorldBase
-from robosuite.utils import OpenCVRenderer
+from robosuite.renderers.viewer import OpenCVViewer
 from robosuite.utils.binding_utils import MjRenderContextOffscreen, MjSim
 from robosuite.utils.mjcf_utils import array_to_string, new_actuator, new_joint
 
@@ -116,7 +116,7 @@ def start_simulation(self):
         self.sim = MjSim.from_xml_string(model_xml)
 
         if self.render:
-            self.viewer = OpenCVRenderer(self.sim)
+            self.viewer = OpenCVViewer(self.sim)
             # We also need to add the offscreen context
             if self.sim._render_context_offscreen is None:
                 render_context = MjRenderContextOffscreen(self.sim, device_id=-1)

robosuite/renderers/viewer/__init__.py

@@ -0,0 +1,2 @@
+from .mjviewer_renderer import MjviewerRenderer
+from .opencv_renderer import OpenCVViewer

robosuite/renderers/viewer/opencv_renderer.py

@@ -0,0 +1,90 @@
+"""
+opencv renderer class.
+"""
+import platform
+
+import cv2
+import numpy as np
+
+
+class OpenCVViewer:
+    def __init__(self, sim):
+        # TODO: update this appropriately - need to get screen dimensions
+        self.width = 1280
+        self.height = 800
+
+        self.sim = sim
+        self.set_camera(camera_id=0)
+        self._window_name = "offscreen render"
+        self._has_window = False
+        self.keypress_callback = None
+
+    def set_camera(self, camera_id=None, camera_name=None, width=None, height=None):
+        """
+        Set the camera view to the specified camera ID.
+
+        Args:
+            camera_id (int or list): id(s) of the camera to set the current viewer to
+            camera_name (str or list or None): name(s) of the camera to set the current viewer to
+        """
+
+        # enforce exactly one arg
+        assert (camera_id is not None) or (camera_name is not None)
+        assert (camera_id is None) or (camera_name is None)
+
+        # set width and height
+        if width is not None:
+            self.width = width
+        if height is not None:
+            self.height = height
+
+        if camera_id is not None:
+            if isinstance(camera_id, int):
+                camera_id = [camera_id]
+            self.camera_names = [self.sim.model.camera_id2name(cam_id) for cam_id in camera_id]
+        else:
+            if isinstance(camera_name, str):
+                camera_name = [camera_name]
+            self.camera_names = list(camera_name)
+
+    def render(self):
+        # get frame with offscreen renderer (assumes that the renderer already exists)
+        im = [
+            self.sim.render(camera_name=cam_name, height=self.height, width=self.width)[..., ::-1]
+            for cam_name in self.camera_names
+        ]
+        im = np.concatenate(im, axis=1)  # concatenate horizontally
+
+        # write frame to window
+        im = np.flip(im, axis=0)
+        cv2.imshow(self._window_name, im)
+        if (platform.system() != "Darwin") and (not self._has_window):
+            # move window to top left of screen, and ensure we only move window on creation
+            cv2.moveWindow(self._window_name, 0, 0)
+        key = cv2.waitKey(1)
+        if self.keypress_callback is not None:
+            self.keypress_callback(key)
+        self._has_window = True
+
+    def add_keypress_callback(self, keypress_callback):
+        self.keypress_callback = keypress_callback
+
+    def close_window(self):
+        """
+        Helper method to close the active window.
+        """
+        if self._has_window:
+            cv2.destroyWindow(self._window_name)
+            cv2.waitKey(1)
+        self._has_window = False
+
+    def close(self):
+        """
+        Any cleanup to close renderer.
+        """
+
+        # NOTE: assume that @sim will get cleaned up outside the renderer - just delete the reference
+        self.sim = None
+
+        # close window
+        self.close_window()

robosuite/utils/__init__.py

@@ -1,3 +1 @@
 from .errors import robosuiteError, XMLError, SimulationError, RandomizationError
-
-from .opencv_renderer import OpenCVRenderer

tests/test_renderers/test_all_renderers.py

@@ -35,6 +35,31 @@ def test_mujoco_renderer():
         env.render()
 
 
+@pytest.mark.skipif(not is_display_available(), reason="No display available for on-screen rendering.")
+def test_multiple_mujoco_renderer():
+    env = suite.make(
+        env_name="Lift",
+        robots="Panda",
+        controller_configs=load_composite_controller_config(controller="BASIC"),
+        has_renderer=True,
+        has_offscreen_renderer=False,
+        ignore_done=True,
+        use_camera_obs=False,
+        control_freq=20,
+        renderer="mujoco",
+    )
+
+    env.reset()
+    camera_name = ["agentview", "birdview"]
+    env.viewer.set_camera(camera_name=camera_name, width=1080, height=720)
+    low, high = env.action_spec
+
+    for i in range(10):
+        action = np.random.uniform(low, high)
+        obs, reward, done, _ = env.step(action)
+        env.render()
+
+
 @pytest.mark.skipif(not is_display_available(), reason="No display available for on-screen rendering.")
 def test_mjviewer_renderer():
     env = suite.make(