From 544b6a37e22c94ae76602d9ea9801b4cefae7e10 Mon Sep 17 00:00:00 2001 From: Michael Drolet <32950163+mdrolet01@users.noreply.github.com> Date: Wed, 25 Oct 2023 15:55:24 +0200 Subject: [PATCH] Headless (#132) * add offscreen render capability * add viewer changes for cluster * refactor headless feature * update glfw init order --- mushroom_rl/utils/mujoco/viewer.py | 141 ++++++++++++++++++++++------- 1 file changed, 110 insertions(+), 31 deletions(-) diff --git a/mushroom_rl/utils/mujoco/viewer.py b/mushroom_rl/utils/mujoco/viewer.py index 6c7ed339..8bab2c30 100644 --- a/mushroom_rl/utils/mujoco/viewer.py +++ b/mushroom_rl/utils/mujoco/viewer.py @@ -2,8 +2,8 @@ import mujoco import time from itertools import cycle - import numpy as np +import os class MujocoGlfwViewer: @@ -15,7 +15,7 @@ class MujocoGlfwViewer: def __init__(self, model, dt, width=1920, height=1080, start_paused=False, custom_render_callback=None, record=False, camera_params=None, default_camera_mode="static", hide_menu_on_startup=False, - geom_group_visualization_on_startup=None): + geom_group_visualization_on_startup=None, headless=False): """ Constructor. @@ -46,37 +46,42 @@ def __init__(self, model, dt, width=1920, height=1080, start_paused=False, self.frames = 0 self.start_time = time.time() - glfw.init() - glfw.window_hint(glfw.COCOA_RETINA_FRAMEBUFFER, 0) + self._headless = headless + self._model = model + self._font_scale = 100 + + if headless: + self._opengl_context = self.get_opengl_backend(width, height) + self._opengl_context.make_current() + self._width, self._height = self.update_headless_size(width, height) + else: + self._width, self._height = width, height + glfw.init() + glfw.window_hint(glfw.COCOA_RETINA_FRAMEBUFFER, 0) + self._window = glfw.create_window(width=self._width, height=self._height, title="MuJoCo", monitor=None, share=None) + glfw.make_context_current(self._window) + glfw.set_mouse_button_callback(self._window, self.mouse_button) + glfw.set_cursor_pos_callback(self._window, self.mouse_move) + glfw.set_key_callback(self._window, self.keyboard) + glfw.set_scroll_callback(self._window, self.scroll) + self._set_mujoco_buffers() + if record: # dont allow to change the window size to have equal frame size during recording glfw.window_hint(glfw.RESIZABLE, False) + self._viewport = mujoco.MjrRect(0, 0, self._width, self._height) self._loop_count = 0 self._time_per_render = 1 / 60. self._run_speed_factor = 1.0 self._paused = start_paused - self._window = glfw.create_window(width=width, height=height, title="MuJoCo", monitor=None, share=None) - glfw.make_context_current(self._window) - - self._width = width - self._height = height - # Disable v_sync, so swap_buffers does not block # glfw.swap_interval(0) - glfw.set_mouse_button_callback(self._window, self.mouse_button) - glfw.set_cursor_pos_callback(self._window, self.mouse_move) - glfw.set_key_callback(self._window, self.keyboard) - glfw.set_scroll_callback(self._window, self.scroll) - - self._model = model - - self._scene = mujoco.MjvScene(model, 1000) + self._scene = mujoco.MjvScene(self._model, 1000) self._scene_option = mujoco.MjvOption() - self._camera = mujoco.MjvCamera() mujoco.mjv_defaultFreeCamera(model, self._camera) if camera_params is None: @@ -92,10 +97,6 @@ def __init__(self, model, dt, width=1920, height=1080, start_paused=False, self._camera_mode_target = next(self._camera_mode_iter) self._set_camera() - self._viewport = mujoco.MjrRect(0, 0, width, height) - self._font_scale = 100 - self._context = mujoco.MjrContext(model, mujoco.mjtFontScale(self._font_scale)) - self.custom_render_callback = custom_render_callback self._overlay = {} @@ -264,6 +265,30 @@ def scroll(self, window, x_offset, y_offset): mujoco.mjv_moveCamera(self._model, mujoco.mjtMouse.mjMOUSE_ZOOM, 0, 0.05 * y_offset, self._scene, self._camera) + def _set_mujoco_buffers(self): + self._context = mujoco.MjrContext(self._model, mujoco.mjtFontScale(self._font_scale)) + if self._headless: + mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_OFFSCREEN, self._context) + if self._context.currentBuffer != mujoco.mjtFramebuffer.mjFB_OFFSCREEN: + raise RuntimeError("Offscreen rendering not supported") + else: + mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_WINDOW, self._context) + if self._context.currentBuffer != mujoco.mjtFramebuffer.mjFB_WINDOW: + raise RuntimeError("Window rendering not supported") + + def update_headless_size(self, width, height): + _context = mujoco.MjrContext(self._model, mujoco.mjtFontScale(self._font_scale)) + if width > _context.offWidth or height > _context.offHeight: + width = max(width, self._model.vis.global_.offwidth) + height = max(height, self._model.vis.global_.offheight) + + if width != _context.offWidth or height != _context.offHeight: + self._model.vis.global_.offwidth = width + self._model.vis.global_.offheight = height + + self._set_mujoco_buffers() + return width, height + def render(self, data, record): """ Main rendering function. @@ -279,7 +304,8 @@ def render(self, data, record): def render_inner_loop(self): - self._create_overlay() + if not self._headless: + self._create_overlay() render_start = time.time() @@ -287,7 +313,8 @@ def render_inner_loop(self): mujoco.mjtCatBit.mjCAT_ALL, self._scene) - self._viewport.width, self._viewport.height = glfw.get_window_size(self._window) + if not self._headless: + self._viewport.width, self._viewport.height = glfw.get_window_size(self._window) mujoco.mjr_render(self._viewport, self._scene, self._context) @@ -307,16 +334,18 @@ def render_inner_loop(self): if self.custom_render_callback is not None: self.custom_render_callback(self._viewport, self._context) - glfw.swap_buffers(self._window) - glfw.poll_events() + if not self._headless: + glfw.swap_buffers(self._window) + glfw.poll_events() self.frames += 1 self._overlay.clear() - if glfw.window_should_close(self._window): - self.stop() - exit(0) + if not self._headless: + if glfw.window_should_close(self._window): + self.stop() + exit(0) self._time_per_render = 0.9 * self._time_per_render + 0.1 * (time.time() - render_start) @@ -349,7 +378,10 @@ def read_pixels(self, depth=False): """ - shape = glfw.get_framebuffer_size(self._window) + if self._headless: + shape = [self._width, self._height] + else: + shape = glfw.get_framebuffer_size(self._window) if depth: rgb_img = np.zeros((shape[1], shape[0], 3), dtype=np.uint8) @@ -531,3 +563,50 @@ def get_default_camera_params(): return dict(static=dict(distance=15.0, elevation=-45.0, azimuth=90.0, lookat=np.array([0.0, 0.0, 0.0])), follow=dict(distance=3.5, elevation=0.0, azimuth=90.0), top_static=dict(distance=5.0, elevation=-90.0, azimuth=90.0, lookat=np.array([0.0, 0.0, 0.0]))) + + + def get_opengl_backend(self, width, height): + # Reference: https://github.com/openai/gym/blob/master/gym/envs/mujoco/mujoco_rendering.py + def _import_egl(width, height): + from mujoco.egl import GLContext + return GLContext(width, height) + + def _import_glfw(width, height): + from mujoco.glfw import GLContext + return GLContext(width, height) + + def _import_osmesa(width, height): + from mujoco.osmesa import GLContext + return GLContext(width, height) + + _ALL_RENDERERS = { + "glfw" : _import_glfw, + "osmesa" : _import_osmesa, + "egl" : _import_egl + } + + backend = os.environ.get("MUJOCO_GL") + if backend is not None: + try: + opengl_context = _ALL_RENDERERS[backend](width, height) + except KeyError: + raise RuntimeError( + "Environment variable {} must be one of {!r}: got {!r}.".format( + "MUJOCO_GL", _ALL_RENDERERS.keys(), backend + ) + ) + + else: + for name, _ in _ALL_RENDERERS.items(): + try: + opengl_context = _ALL_RENDERERS[name](width, height) + backend = name + break + except: # noqa:E722 + pass + if backend is None: + raise RuntimeError( + "No OpenGL backend could be imported. Attempting to create a " + "rendering context will result in a RuntimeError." + ) + return opengl_context \ No newline at end of file