Add Renderer class to Python native bindings.

PiperOrigin-RevId: 492994751
Change-Id: Iebbe5980668a15b1cd62c0076cc131255462b2eb

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<PyViewer>` for details. (Contribution by `Levi Burner <https://github.com/aftersomemath>`_.)
+- The ``Renderer`` class from the MuJoCo tutorial Colab is now available directly in the native Python bindings.
 
 General
 ^^^^^^^

python/mujoco/__init__.py

@@ -43,27 +43,8 @@ from mujoco._errors import *
 from mujoco._functions import *
 from mujoco._render import *
 from mujoco._structs import *
-
+from mujoco.gl_context import *
-# pylint: disable=g-import-not-at-top
+from mujoco.renderer import Renderer
-_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
 
 HEADERS_DIR = os.path.join(os.path.dirname(__file__), 'include/mujoco')
 PLUGINS_DIR = os.path.join(os.path.dirname(__file__), 'plugin')

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

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:
+<visual>
+  <global offwidth="my_width"/>
+</visual>""".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:
+<visual>
+  <global offheight="my_height"/>
+</visual>""".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,
+    )

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 = """
+ <mujoco>
+  <worldbody>
+    <camera name="closeup" pos="0 -6 0" xyaxes="1 0 0 0 1 100"/>
+    <geom name="white_box" type="box" size="1 1 1" rgba="1 1 1 1"/>
+  </worldbody>
+</mujoco>
+"""
+    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()