RENDERER: FEATURE: Adding the new mujoco renderer. Making MUJOCO as default backend

Author: vikashplus

robohive/envs/hands/hammer_v1.py

@@ -126,7 +126,6 @@ def reset(self, reset_qpos=None, reset_qvel=None, **kwargs):
         qv = self.init_qvel.copy() if reset_qvel==None else reset_qvel
         self.robot.reset(reset_pos=qp, reset_vel=qv, **kwargs)
 
-
         self.sim.model.body_pos[self.target_bid,2] = self.np_random.uniform(low=0.1, high=0.25)
         self.sim.forward()
         return self.get_obs()

robohive/physics/mj_sim_scene.py

@@ -16,7 +16,7 @@
 import_utils.mujoco_isavailable()
 import dm_control.mujoco as dm_mujoco
 
-from robohive.renderer.mj_renderer import DMRenderer
+from robohive.renderer.mj_renderer import MJRenderer
 from robohive.physics.sim_scene import SimScene
 
 
@@ -52,10 +52,9 @@ def advance(self, substeps: int = 1, render:bool = True):
             # self.renderer.refresh_window()
             self.renderer.render_to_window()
 
-
-    def _create_renderer(self, sim: Any) -> DMRenderer:
+    def _create_renderer(self, sim: Any) -> MJRenderer:
         """Creates a renderer for the given simulation."""
-        return DMRenderer(sim)
+        return MJRenderer(sim)
 
     def copy_model(self) -> Any:
         """Returns a copy of the MjModel object."""

robohive/physics/sim_scene.py

@@ -52,9 +52,9 @@ def create(*args, backend: Union[SimBackend, int], **kwargs) -> 'SimScene':
     @staticmethod
     def get_sim(model_handle: Any) -> 'SimScene':
         sim_backend = os.getenv('sim_backend')
-        if sim_backend == 'MUJOCO_PY' or sim_backend == None:
+        if sim_backend == 'MUJOCO_PY':
             return SimScene.create(model_handle=model_handle, backend=SimBackend.MUJOCO_PY)
-        elif sim_backend == 'MUJOCO':
+        elif sim_backend == 'MUJOCO' or sim_backend == None:
             return SimScene.create(model_handle=model_handle, backend=SimBackend.DM_CONTROL)
         else:
             raise ValueError("Unknown sim_backend: {}. Available choices: MUJOCO_PY, MUJOCO")

robohive/renderer/mj_renderer.py

@@ -7,34 +7,33 @@
 
 """Rendering simulation using dm_control."""
 
-import sys
-
 import numpy as np
-import dm_control.mujoco as dm_mujoco
-import dm_control.viewer as dm_viewer
-import dm_control._render as dm_render
+import mujoco
+from mujoco import viewer
 
 from typing import Union
-
 from robohive.renderer.renderer import Renderer, RenderMode
 
 # Default window dimensions.
 DEFAULT_WINDOW_WIDTH = 640
 DEFAULT_WINDOW_HEIGHT = 480
 
 # Default window title.
-DEFAULT_WINDOW_TITLE = 'MuJoCo Viewer'
+DEFAULT_WINDOW_TITLE = 'RoboHive Viewer'
 
-# Internal renderbuffer size, in pixels.
-_MAX_RENDERBUFFER_SIZE = 2048
 
-
-class DMRenderer(Renderer):
+class MJRenderer(Renderer):
     """Renders DM Control Physics objects."""
 
-    def __init__(self, sim: dm_mujoco.Physics):
+    def __init__(self, sim):
         super().__init__(sim)
         self._window = None
+        self._renderer = None
+
+
+    def setup_renderer(self, model, height, width):
+        self._renderer = mujoco.Renderer(model, height=height, width=width)
+
 
     def render_to_window(self):
         """Renders the Physics object to a window.
@@ -44,22 +43,23 @@ def render_to_window(self):
         This function is a no-op if the window was already created.
         """
         if not self._window:
-            self._window = DMRenderWindow()
-            self._window.load_model(self._sim)
-            self._update_camera_properties(self._window.camera)
+            self._window = viewer.launch_passive(self._sim.model.ptr, self._sim.data.ptr)
+            self._update_camera_properties(self._window.cam)
+
+        self._window.cam.azimuth+=.1 # trick to rotate camera for 360 videos
+        self.refresh_window()
 
-        self._window.run_frame()
 
     def refresh_window(self):
         """Refreshes the rendered window if one is present."""
         if self._window is None:
             return
-        self._window.run_frame()
+        self._window.sync()
+
 
     def render_offscreen(self,
-                         width: int,
-                         height: int,
-                        #  mode: RenderMode = RenderMode.RGB,
+                         width: int = DEFAULT_WINDOW_WIDTH,
+                         height: int = DEFAULT_WINDOW_HEIGHT,
                          rgb: bool = True,
                          depth: bool = False,
                          segmentation: bool = False,
@@ -77,36 +77,29 @@ def render_offscreen(self,
         Returns:
             A numpy array of the pixels.
         """
-        assert width > 0 and height > 0
-
-        if isinstance(camera_id, str):
-            camera_id = self._sim.model.name2id(camera_id, 'camera')
-        elif camera_id == None:
-            camera_id = -1
-
-        # TODO(michaelahn): Consider caching the camera.
-        camera = dm_mujoco.Camera(
-            physics=self._sim, height=height, width=width, camera_id=camera_id)
-
-        # Update the camera configuration for the free-camera.
-        if camera_id == -1:
-            self._update_camera_properties(camera._render_camera)  # pylint: disable=protected-access
+        if camera_id == None:
+            camera_id = 0
+        if self._renderer is None:
+            self.setup_renderer(self._sim.model.ptr, width=width, height=height)
 
-        # image = camera.render(
-        #    depth=(mode == RenderMode.DEPTH),
-        #    segmentation=(mode == RenderMode.SEGMENTATION))
         rgb_arr = None; dpt_arr = None; seg_arr = None
         if rgb:
-            rgb_arr = camera.render()
-            # rgb_arr = rgb_arr[::-1, :, :]
+            self._renderer.update_scene(self._sim.data.ptr, camera=camera_id)
+            rgb_arr = self._renderer.render()
         if depth:
-            dpt_arr = camera.render(depth=True)
+            self._renderer.enable_depth_rendering()
+            self._renderer.update_scene(self._sim.data.ptr, camera=camera_id)
+            dpt_arr = self._renderer.render()
             dpt_arr = dpt_arr[::-1, :]
+            self._renderer.disable_depth_rendering()
         if segmentation:
-            seg_arr = camera.render(segmentation=True)
+            self._renderer.enable_segmentation_rendering()
+            self._renderer.update_scene(self._sim.data.ptr, camera=camera_id)
+            dpt_arr = self._renderer.render()
+            dpt_arr = dpt_arr[::-1, :]
+            self._renderer.disable_segmentation_rendering()
             seg_arr = seg_arr[::-1, :]
 
-        camera._scene.free()  # pylint: disable=protected-access
         if depth and segmentation:
             return rgb_arr, dpt_arr, seg_arr
         elif depth:
@@ -121,69 +114,4 @@ def close(self):
         """Cleans up any resources being used by the renderer."""
         if self._window:
             self._window.close()
-            self._window = None
-
-
-class DMRenderWindow:
-    """Class that encapsulates a graphical window."""
-
-    def __init__(self,
-                 width: int = DEFAULT_WINDOW_WIDTH,
-                 height: int = DEFAULT_WINDOW_HEIGHT,
-                 title: str = DEFAULT_WINDOW_TITLE):
-        """Creates a graphical render window.
-
-        Args:
-            width: The width of the window.
-            height: The height of the window.
-            title: The title of the window.
-        """
-        self._viewport = dm_viewer.renderer.Viewport(width, height)
-        self._window = dm_viewer.gui.RenderWindow(width, height, title)
-        self._viewer = dm_viewer.viewer.Viewer(
-            self._viewport, self._window.mouse, self._window.keyboard)
-        self._draw_surface = None
-        self._renderer = dm_viewer.renderer.NullRenderer()
-
-    @property
-    def camera(self):
-        return self._viewer._camera._camera  # pylint: disable=protected-access
-
-    def close(self):
-        self._viewer.deinitialize()
-        self._renderer.release()
-        self._draw_surface.free()
-        self._window.close()
-
-    def load_model(self, physics):
-        """Loads the given Physics object to render."""
-        self._viewer.deinitialize()
-
-        self._draw_surface = dm_render.Renderer(
-            max_width=_MAX_RENDERBUFFER_SIZE, max_height=_MAX_RENDERBUFFER_SIZE)
-        self._renderer = dm_viewer.renderer.OffScreenRenderer(
-            physics.model, self._draw_surface)
-
-        self._viewer.initialize(physics, self._renderer, touchpad=False)
-
-    def run_frame(self):
-        """Renders one frame of the simulation.
-
-        NOTE: This is extremely slow at the moment.
-        """
-        # pylint: disable=protected-access
-        glfw = dm_viewer.gui.glfw_gui.glfw
-        glfw_window = self._window._context.window
-        if glfw.window_should_close(glfw_window):
-            sys.exit(0)
-
-        self._viewport.set_size(*self._window.shape)
-        self._viewer.render()
-        pixels = self._renderer.pixels
-
-        with self._window._context.make_current() as ctx:
-            ctx.call(self._window._update_gui_on_render_thread, glfw_window,
-                     pixels)
-        self._window._mouse.process_events()
-        self._window._keyboard.process_events()
-        # pylint: enable=protected-access
+            self._window = None

robohive/renderer/mjpy_renderer.py

@@ -23,21 +23,24 @@ def __init__(self, sim):
         self._onscreen_renderer = None
         self._offscreen_renderer = None
 
+
     def render_to_window(self):
         """Renders the simulation to a window."""
         if not self._onscreen_renderer:
             self._onscreen_renderer = mujoco_py.MjViewer(self._sim)
             self._update_camera_properties(self._onscreen_renderer.cam)
 
-        self._onscreen_renderer.render()
+        self.refresh_window()
         # self._onscreen_renderer.cam.azimuth+=.1 # trick to rotate camera for 360 videos
 
+
     def refresh_window(self):
         """Refreshes the rendered window if one is present."""
         if self._onscreen_renderer is None:
             return
         self._onscreen_renderer.render()
 
+
     def render_offscreen(self,
                          width: int,
                          height: int,

setup.py

@@ -53,7 +53,7 @@ def package_files(directory):
     extras_require={
       # To use mujoco bindings, run (pip install -e ".[mujoco]") and set sim_backend=MUJOCO
       'mujoco':[
-        'mujoco==2.3.3',
+        'mujoco==2.3.5',
         'dm-control==1.0.11'
         ],
       # To use hardware dependencies, run (pip install -e ".[a0]") and follow install instructions inside robot