Skip to content

Commit

Permalink
Headless (#132)
Browse files Browse the repository at this point in the history
* add offscreen render capability

* add viewer changes for cluster

* refactor headless feature

* update glfw init order
  • Loading branch information
mdrolet01 authored Oct 25, 2023
1 parent 8ee07d3 commit 544b6a3
Showing 1 changed file with 110 additions and 31 deletions.
141 changes: 110 additions & 31 deletions mushroom_rl/utils/mujoco/viewer.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@
import mujoco
import time
from itertools import cycle

import numpy as np
import os


class MujocoGlfwViewer:
Expand All @@ -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.
Expand Down Expand Up @@ -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:
Expand All @@ -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 = {}
Expand Down Expand Up @@ -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.
Expand All @@ -279,15 +304,17 @@ def render(self, data, record):

def render_inner_loop(self):

self._create_overlay()
if not self._headless:
self._create_overlay()

render_start = time.time()

mujoco.mjv_updateScene(self._model, data, self._scene_option, None, self._camera,
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)

Expand All @@ -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)

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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

0 comments on commit 544b6a3

Please sign in to comment.