diff --git a/doc/changelog.rst b/doc/changelog.rst index 06658234..d572a524 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -11,6 +11,7 @@ Python bindings - The ``simulate`` GUI is now available through the ``mujoco`` Python package as ``mujoco.viewer``. See :ref:`documentation` for details. (Contribution by `Levi Burner `_.) +- The ``Renderer`` class from the MuJoCo tutorial Colab is now available directly in the native Python bindings. General ^^^^^^^ diff --git a/python/mujoco/__init__.py b/python/mujoco/__init__.py index 191a4ff4..3b2904b9 100644 --- a/python/mujoco/__init__.py +++ b/python/mujoco/__init__.py @@ -43,27 +43,8 @@ from mujoco._errors import * from mujoco._functions import * from mujoco._render import * from mujoco._structs import * - -# pylint: disable=g-import-not-at-top -_MUJOCO_GL = os.environ.get('MUJOCO_GL', '').lower().strip() -if _MUJOCO_GL not in ('disable', 'disabled', 'off', 'false', '0'): - _VALID_MUJOCO_GL = ('enable', 'enabled', 'on', 'true', '1' , 'glfw', '') - if _SYSTEM == 'Linux': - _VALID_MUJOCO_GL += ('glx', 'egl', 'osmesa') - elif _SYSTEM == 'Windows': - _VALID_MUJOCO_GL += ('wgl',) - elif _SYSTEM == 'Darwin': - _VALID_MUJOCO_GL += ('cgl',) - if _MUJOCO_GL not in _VALID_MUJOCO_GL: - raise RuntimeError( - f'invalid value for environment variable MUJOCO_GL: {_MUJOCO_GL}') - - if _SYSTEM == 'Linux' and _MUJOCO_GL == 'osmesa': - from mujoco.osmesa import GLContext - elif _SYSTEM == 'Linux' and _MUJOCO_GL == 'egl': - from mujoco.egl import GLContext - else: - from mujoco.glfw import GLContext +from mujoco.gl_context import * +from mujoco.renderer import Renderer HEADERS_DIR = os.path.join(os.path.dirname(__file__), 'include/mujoco') PLUGINS_DIR = os.path.join(os.path.dirname(__file__), 'plugin') diff --git a/python/mujoco/gl_context.py b/python/mujoco/gl_context.py new file mode 100644 index 00000000..3a5335ea --- /dev/null +++ b/python/mujoco/gl_context.py @@ -0,0 +1,45 @@ +# Copyright 2022 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Exports GLContext for MuJoCo Python bindings.""" + +import ctypes +import ctypes.util +import os +import platform + +# pylint: disable=g-import-not-at-top +_SYSTEM = platform.system() +_MUJOCO_GL = os.environ.get('MUJOCO_GL', '').lower().strip() +if _MUJOCO_GL not in ('disable', 'disabled', 'off', 'false', '0'): + _VALID_MUJOCO_GL = ('enable', 'enabled', 'on', 'true', '1' , 'glfw', '') + if _SYSTEM == 'Linux': + _VALID_MUJOCO_GL += ('glx', 'egl', 'osmesa') + elif _SYSTEM == 'Windows': + _VALID_MUJOCO_GL += ('wgl',) + elif _SYSTEM == 'Darwin': + _VALID_MUJOCO_GL += ('cgl',) + if _MUJOCO_GL not in _VALID_MUJOCO_GL: + raise RuntimeError( + f'invalid value for environment variable MUJOCO_GL: {_MUJOCO_GL}') + + if _SYSTEM == 'Linux' and _MUJOCO_GL == 'osmesa': + from mujoco.osmesa import GLContext as _GLContext + GLContext = _GLContext + elif _SYSTEM == 'Linux' and _MUJOCO_GL == 'egl': + from mujoco.egl import GLContext as _GLContext + GLContext = _GLContext + else: + from mujoco.glfw import GLContext as _GLContext + GLContext = _GLContext diff --git a/python/mujoco/renderer.py b/python/mujoco/renderer.py new file mode 100644 index 00000000..fa6574dd --- /dev/null +++ b/python/mujoco/renderer.py @@ -0,0 +1,230 @@ +# Copyright 2022 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Defines a renderer class for the MuJoCo Python native bindings.""" + +from typing import Optional, Union + +from mujoco import _enums +from mujoco import _functions +from mujoco import _render +from mujoco import _structs +from mujoco import gl_context +import numpy as np + + +class Renderer: + """Renders MuJoCo scenes.""" + + def __init__( + self, + model: _structs.MjModel, + height: int = 240, + width: int = 320, + max_geom: int = 10000 + ) -> None: + """Initializes a new `Renderer`. + + Args: + model: an MjModel instance. + height: image height in pixels. + width: image width in pixels. + max_geom: Optional integer specifying the maximum number of geoms that can + be rendered in the same scene. If None this will be chosen automatically + based on the estimated maximum number of renderable geoms in the model. + Raises: + ValueError: If `camera_id` is outside the valid range, or if `width` or + `height` exceed the dimensions of MuJoCo's offscreen framebuffer. + """ + buffer_width = model.vis.global_.offwidth + buffer_height = model.vis.global_.offheight + if width > buffer_width: + raise ValueError(f""" +Image width {width} > framebuffer width {buffer_width}. Either reduce the image +width or specify a larger offscreen framebuffer in the model XML using the +clause: + + +""".lstrip()) + + if height > buffer_height: + raise ValueError(f""" +Image height {height} > framebuffer height {buffer_height}. Either reduce the +image height or specify a larger offscreen framebuffer in the model XML using +the clause: + + +""".lstrip()) + + self._width = width + self._height = height + self._model = model + + self._scene = _structs.MjvScene(model=model, maxgeom=max_geom) + self._scene_option = _structs.MjvOption() + + self._rect = _render.MjrRect(0, 0, self._width, self._height) + + # Internal buffers. + self._rgb_buffer = np.empty((self._height, self._width, 3), dtype=np.uint8) + self._depth_buffer = np.empty((self._height, self._width), dtype=np.float32) + + # Create render contexts. + self._gl_context = gl_context.GLContext(width, height) + self._gl_context.make_current() + self._mjr_context = _render.MjrContext( + model, _enums.mjtFontScale.mjFONTSCALE_150 + ) + _render.mjr_setBuffer( + _enums.mjtFramebuffer.mjFB_OFFSCREEN, self._mjr_context + ) + + # Default render flags. + self._depth_rendering = False + self._segmentation_rendering = False + + @property + def model(self): + return self._model + + @property + def scene(self) -> _structs.MjvScene: + return self._scene + + @property + def height(self): + return self._height + + @property + def width(self): + return self._width + + def enable_depth_rendering(self): + self._segmentation_rendering = False + self._depth_rendering = True + + def disable_depth_rendering(self): + self._depth_rendering = False + + def enable_segmentation_rendering(self): + self._segmentation_rendering = True + self._depth_rendering = False + + def disable_segmentation_rendering(self): + self._segmentation_rendering = False + + def render(self) -> np.ndarray: + """Renders the scene as a numpy array of pixel values. + + Returns: + A numpy array of pixels with dimensions (H, W, 3). The array will be + mutated by future calls to `render`. + """ + original_flags = self._scene.flags.copy() + + if self._segmentation_rendering: + self._scene.flags[_enums.mjtRndFlag.mjRND_SEGMENT] = True + self._scene.flags[_enums.mjtRndFlag.mjRND_IDCOLOR] = True + + self._gl_context.make_current() + + # Render scene and read contents of RGB and depth buffers. + _render.mjr_render(self._rect, self._scene, self._mjr_context) + _render.mjr_readPixels(self._rgb_buffer, self._depth_buffer, self._rect, + self._mjr_context) + + if self._depth_rendering: + # Get the distances to the near and far clipping planes. + extent = self._model.stat.extent + near = self._model.vis.map.znear * extent + far = self._model.vis.map.zfar * extent + + # Convert from [0 1] to depth in meters, see links below: + # http://stackoverflow.com/a/6657284/1461210 + # https://www.khronos.org/opengl/wiki/Depth_Buffer_Precision + pixels = near / (1 - self._depth_buffer * (1 - near / far)) + + elif self._segmentation_rendering: + # Convert 3-channel uint8 to 1-channel uint32. + image3 = self._rgb_buffer.astype(np.uint32) + segimage = (image3[:, :, 0] + + image3[:, :, 1] * (2**8) + + image3[:, :, 2] * (2**16)) + # Remap segid to 2-channel (object ID, object type) pair. + # Seg ID 0 is background -- will be remapped to (-1, -1). + ngeoms = self._scene.ngeom + segid2output = np.full((ngeoms + 1, 2), fill_value=-1, + dtype=np.int32) # Seg id cannot be > ngeom + 1. + visible_geoms = [g for g in self._scene.geoms[:ngeoms] if g.segid != -1] + visible_segids = np.array([g.segid + 1 for g in visible_geoms], np.int32) + visible_objid = np.array([g.objid for g in visible_geoms], np.int32) + visible_objtype = np.array([g.objtype for g in visible_geoms], np.int32) + segid2output[visible_segids, 0] = visible_objid + segid2output[visible_segids, 1] = visible_objtype + pixels = segid2output[segimage] + + # Reset scene flags. + np.copyto(self._scene.flags, original_flags) + else: + pixels = self._rgb_buffer + return np.flipud(pixels) + + def update_scene( + self, + data: _structs.MjData, + camera: Union[int, str, _structs.MjvCamera] = -1, + scene_option: Optional[_structs.MjvOption] = None + ): + """Updates geometry used for rendering. + + Args: + data: An instance of `MjData`. + camera: An instance of `MjvCamera`, a string or an integer + scene_option: A custom `MjvOption` instance to use to render + the scene instead of the default. + """ + if not isinstance(camera, _structs.MjvCamera): + camera_id = camera + if isinstance(camera_id, str): + camera_id = _functions.mj_name2id(self._model, + _enums.mjtObj.mjOBJ_CAMERA, camera_id) + if camera_id < -1: + raise ValueError('camera_id cannot be smaller than -1.') + if camera_id >= self._model.ncam: + raise ValueError( + f'model has {self._model.ncam} fixed cameras. ' + f'camera_id={camera_id} is invalid.' + ) + + # Render camera. + camera = _structs.MjvCamera() + camera.fixedcamid = camera_id + + # Defaults to mjCAMERA_FREE, otherwise mjCAMERA_FIXED refers to a + # camera explicitly defined in the model. + if camera_id == -1: + camera.type = _enums.mjtCamera.mjCAMERA_FREE + _functions.mjv_defaultFreeCamera(self._model, camera) + else: + camera.type = _enums.mjtCamera.mjCAMERA_FIXED + + scene_option = scene_option or self._scene_option + _functions.mjv_updateScene( + self._model, + data, + scene_option, + None, + camera, _enums.mjtCatBit.mjCAT_ALL, + self._scene, + ) diff --git a/python/mujoco/renderer_test.py b/python/mujoco/renderer_test.py new file mode 100644 index 00000000..295de178 --- /dev/null +++ b/python/mujoco/renderer_test.py @@ -0,0 +1,50 @@ +# Copyright 2022 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for the MuJoCo renderer.""" + +from absl.testing import absltest +from absl.testing import parameterized +import mujoco + + +@absltest.skipUnless(hasattr(mujoco, 'GLContext'), + 'MuJoCo rendering is disabled') +class MuJoCoRendererTest(parameterized.TestCase): + def test_renderer_renders_scene(self): + xml = """ + + + + + + +""" + model = mujoco.MjModel.from_xml_string(xml) + data = mujoco.MjData(model) + renderer = mujoco.Renderer(model, 50, 50) + mujoco.mj_forward(model, data) + renderer.update_scene(data, 'closeup') + + pixels = renderer.render().flatten() + not_all_black = False + + # Pixels should all be a neutral color. + for pixel in pixels: + if pixel > 0: + not_all_black = True + break + self.assertTrue(not_all_black) +if __name__ == '__main__': + absltest.main()