Move renderer.py and gl_context.py into a submodule.

PiperOrigin-RevId: 870789311
Change-Id: Iaa0fcc764a770eeb4570a6a189a5ad032d41d64c

python/mujoco/__init__.py

@@ -59,8 +59,16 @@ from mujoco._functions import *
 from mujoco._render import *
 from mujoco._specs import *
 from mujoco._structs import *
-from mujoco.gl_context import *
+
-from mujoco.renderer import Renderer
+try:
+  from mujoco.rendering.classic.renderer import Renderer  # pytype: disable=import-error
+except ImportError:
+  pass
+
+try:
+  from mujoco.rendering.classic.gl_context import *  # pytype: disable=import-error
+except ImportError:
+  pass
 
 MjStruct: TypeAlias = Union[
     _specs.MjsBody,

python/mujoco/gl_context.py

@@ -12,37 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Exports GLContext for MuJoCo Python bindings."""
+"""Wrapper around mujoco.rendering.classic.gl_context for backwards compatibility."""
 
-import ctypes
-import ctypes.util
-import os
-import platform
 
-# pylint: disable=g-import-not-at-top
+try:
-_SYSTEM = platform.system()
+  from mujoco.rendering.classic.gl_context import *  # pytype: disable=import-error
-_MUJOCO_GL = os.environ.get('MUJOCO_GL', '').lower().strip()
+except ImportError:
-if _MUJOCO_GL not in ('disable', 'disabled', 'off', 'false', '0'):
+  pass
-  _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
-  elif _SYSTEM == 'Darwin':
-    from mujoco.cgl import GLContext as _GLContext
-    GLContext = _GLContext
-  else:
-    from mujoco.glfw import GLContext as _GLContext
-    GLContext = _GLContext

python/mujoco/renderer.py

@@ -12,329 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Defines a renderer class for the MuJoCo Python native bindings."""
+"""Wrapper around mujoco.rendering.classic.renderer for backwards compatibility."""
-
-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:
+try:
-  """Renders MuJoCo scenes."""
+  from mujoco.rendering.classic.renderer import Renderer  # pytype: disable=import-error
-
+except ImportError:
-  def __init__(
+  pass
-      self,
-      model: _structs.MjModel,
-      height: int = 240,
-      width: int = 320,
-      max_geom: int = 10000,
-      font_scale: _enums.mjtFontScale = _enums.mjtFontScale.mjFONTSCALE_150,
-  ) -> 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.
-      font_scale: Optional enum specifying the font scale for text.
-
-    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)
-
-    # Create render contexts.
-    # TODO(nimrod): Figure out why pytype doesn't like gl_context.GLContext
-    self._gl_context = None  # type: ignore
-    if gl_context.GLContext is not None:
-      self._gl_context = gl_context.GLContext(width, height)
-    if self._gl_context:
-      self._gl_context.make_current()
-    self._mjr_context = _render.MjrContext(model, font_scale.value)
-    _render.mjr_setBuffer(
-        _enums.mjtFramebuffer.mjFB_OFFSCREEN.value, self._mjr_context
-    )
-    self._mjr_context.readDepthMap = _enums.mjtDepthMap.mjDEPTH_ZEROFAR
-
-    # 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, *, out: Optional[np.ndarray] = None) -> np.ndarray:
-    """Renders the scene as a numpy array of pixel values.
-
-    Args:
-      out: Alternative output array in which to place the resulting pixels. It
-        must have the same shape as the expected output but the type will be
-        cast if necessary. The expted shape depends on the value of
-        `self._depth_rendering`: when `True`, we expect `out.shape == (width,
-        height)`, and `out.shape == (width, height, 3)` when `False`.
-
-    Returns:
-      A new numpy array holding the pixels with shape `(H, W)` or `(H, W, 3)`,
-      depending on the value of `self._depth_rendering` unless
-      `out is None`, in which case a reference to `out` is returned.
-
-    Raises:
-      RuntimeError: if this method is called after the close method.
-    """
-    original_flags = self._scene.flags.copy()
-
-    # Using segmented rendering for depth makes the calculated depth more
-    # accurate at far distances.
-    if self._depth_rendering or self._segmentation_rendering:
-      self._scene.flags[_enums.mjtRndFlag.mjRND_SEGMENT] = True
-      self._scene.flags[_enums.mjtRndFlag.mjRND_IDCOLOR] = True
-
-    if self._mjr_context is None:
-      raise RuntimeError('render cannot be called after close.')
-
-    if self._gl_context:
-      self._gl_context.make_current()
-
-    if self._depth_rendering:
-      out_shape = (self._height, self._width)
-      out_dtype = np.float32
-    else:
-      out_shape = (self._height, self._width, 3)
-      out_dtype = np.uint8
-
-    if out is None:
-      out = np.empty(out_shape, dtype=out_dtype)
-    else:
-      if out.shape != out_shape:
-        raise ValueError(
-            f'Expected `out.shape == {out_shape}`. Got `out.shape={out.shape}`'
-            ' instead. When using depth rendering, the out array should be of'
-            ' shape `(width, height)` and otherwise (width, height, 3).'
-            f' Got `(self.height, self.width)={(self.height, self.width)}` and'
-            f' `self._depth_rendering={self._depth_rendering}`.'
-        )
-
-    # Render scene and read contents of RGB and depth buffers.
-    _render.mjr_render(self._rect, self._scene, self._mjr_context)
-    if self._depth_rendering:
-      _render.mjr_readPixels(None, out, self._rect, self._mjr_context)
-
-      # 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
-
-      # Calculate OpenGL perspective matrix values in float32 precision
-      # so they are close to what glFrustum returns
-      # https://registry.khronos.org/OpenGL-Refpages/gl2.1/xhtml/glFrustum.xml
-      zfar = np.float32(far)
-      znear = np.float32(near)
-      c_coef = -(zfar + znear) / (zfar - znear)
-      d_coef = -(np.float32(2) * zfar * znear) / (zfar - znear)
-
-      # In reverse Z mode the perspective matrix is transformed by the following
-      c_coef = np.float32(-0.5) * c_coef - np.float32(0.5)
-      d_coef = np.float32(-0.5) * d_coef
-
-      # We need 64 bits to convert Z from ndc to metric depth without noticeable
-      # losses in precision
-      out_64 = out.astype(np.float64)
-
-      # Undo OpenGL projection
-      # Note: We do not need to take action to convert from window coordinates
-      # to normalized device coordinates because in reversed Z mode the mapping
-      # is identity
-      out_64 = d_coef / (out_64 + c_coef)
-
-      # Cast result back to float32 for backwards compatibility
-      # This has a small accuracy cost
-      out[:] = out_64.astype(np.float32)
-
-      # Reset scene flags.
-      np.copyto(self._scene.flags, original_flags)
-    elif self._segmentation_rendering:
-      _render.mjr_readPixels(out, None, self._rect, self._mjr_context)
-
-      # Convert 3-channel uint8 to 1-channel uint32.
-      image3 = out.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
-      out = segid2output[segimage]
-
-      # Reset scene flags.
-      np.copyto(self._scene.flags, original_flags)
-    else:
-      _render.mjr_readPixels(out, None, self._rect, self._mjr_context)
-
-    if self._gl_context:
-      # If using EGL, OSMesa, or GLFW, the output image is flipped vertically.
-      # No flip is needed for Filament.
-      out[:] = np.flipud(out)
-
-    return out
-
-  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.
-
-    Raises:
-      ValueError: If `camera_id` is outside the valid range, or if camera does
-        not exist.
-    """
-    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.value, camera_id
-        )
-        if camera_id == -1:
-          raise ValueError(f'The camera "{camera}" does not exist.')
-      if camera_id < -1 or camera_id >= self._model.ncam:
-        raise ValueError(
-            f'The camera id {camera_id} is out of'
-            f' range [-1, {self._model.ncam}).'
-        )
-
-      # 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.value,
-        self._scene,
-    )
-
-  def close(self) -> None:
-    """Frees the resources used by the renderer.
-
-    This method can be used directly:
-
-    ```python
-    renderer = Renderer(...)
-    # Use renderer.
-    renderer.close()
-    ```
-
-    or via a context manager:
-
-    ```python
-    with Renderer(...) as renderer:
-      # Use renderer.
-    ```
-    """
-    if self._gl_context:
-      self._gl_context.free()
-    self._gl_context = None
-    if self._mjr_context:
-      self._mjr_context.free()
-    self._mjr_context = None
-
-  def __enter__(self):
-    return self
-
-  def __exit__(self, exc_type, exc_value, traceback):
-    del exc_type, exc_value, traceback  # Unused.
-    self.close()
-
-  def __del__(self) -> None:
-    self.close()

python/mujoco/rendering/classic/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2026 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.
+# ==============================================================================
+
+# Empty.

python/mujoco/rendering/classic/gl_context.py

@@ -0,0 +1,48 @@
+# 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
+  elif _SYSTEM == 'Darwin':
+    from mujoco.cgl import GLContext as _GLContext
+    GLContext = _GLContext
+  else:
+    from mujoco.glfw import GLContext as _GLContext
+    GLContext = _GLContext

python/mujoco/rendering/classic/renderer.py

@@ -0,0 +1,339 @@
+# 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
+
+import mujoco
+from mujoco.rendering.classic import gl_context
+import numpy as np
+
+GLContext = gl_context.GLContext if hasattr(gl_context, 'GLContext') else None
+
+
+class Renderer:
+  """Renders MuJoCo scenes."""
+
+  def __init__(
+      self,
+      model: mujoco.MjModel,
+      height: int = 240,
+      width: int = 320,
+      max_geom: int = 10000,
+      font_scale: mujoco.mjtFontScale = mujoco.mjtFontScale.mjFONTSCALE_150,
+  ) -> 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.
+      font_scale: Optional enum specifying the font scale for text.
+
+    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 = mujoco.MjvScene(model=model, maxgeom=max_geom)
+    self._scene_option = mujoco.MjvOption()
+
+    self._rect = mujoco.MjrRect(0, 0, self._width, self._height)
+
+    # Create render contexts.
+    # TODO(nimrod): Figure out why pytype doesn't like gl_context.GLContext
+    self._gl_context = None  # type: ignore
+    if gl_context.GLContext is not None:
+      self._gl_context = gl_context.GLContext(width, height)
+    if self._gl_context:
+      self._gl_context.make_current()
+    self._mjr_context = mujoco.MjrContext(model, font_scale.value)
+    mujoco.mjr_setBuffer(
+        mujoco.mjtFramebuffer.mjFB_OFFSCREEN.value, self._mjr_context
+    )
+    self._mjr_context.readDepthMap = mujoco.mjtDepthMap.mjDEPTH_ZEROFAR
+
+    # Default render flags.
+    self._depth_rendering = False
+    self._segmentation_rendering = False
+
+  @property
+  def model(self):
+    return self._model
+
+  @property
+  def scene(self) -> mujoco.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, *, out: Optional[np.ndarray] = None) -> np.ndarray:
+    """Renders the scene as a numpy array of pixel values.
+
+    Args:
+      out: Alternative output array in which to place the resulting pixels. It
+        must have the same shape as the expected output but the type will be
+        cast if necessary. The expted shape depends on the value of
+        `self._depth_rendering`: when `True`, we expect `out.shape == (width,
+        height)`, and `out.shape == (width, height, 3)` when `False`.
+
+    Returns:
+      A new numpy array holding the pixels with shape `(H, W)` or `(H, W, 3)`,
+      depending on the value of `self._depth_rendering` unless
+      `out is None`, in which case a reference to `out` is returned.
+
+    Raises:
+      RuntimeError: if this method is called after the close method.
+    """
+    original_flags = self._scene.flags.copy()
+
+    # Using segmented rendering for depth makes the calculated depth more
+    # accurate at far distances.
+    if self._depth_rendering or self._segmentation_rendering:
+      self._scene.flags[mujoco.mjtRndFlag.mjRND_SEGMENT] = True
+      self._scene.flags[mujoco.mjtRndFlag.mjRND_IDCOLOR] = True
+
+    if self._mjr_context is None:
+      raise RuntimeError('render cannot be called after close.')
+
+    if self._gl_context:
+      self._gl_context.make_current()
+
+    if self._depth_rendering:
+      out_shape = (self._height, self._width)
+      out_dtype = np.float32
+    else:
+      out_shape = (self._height, self._width, 3)
+      out_dtype = np.uint8
+
+    if out is None:
+      out = np.empty(out_shape, dtype=out_dtype)
+    else:
+      if out.shape != out_shape:
+        raise ValueError(
+            f'Expected `out.shape == {out_shape}`. Got `out.shape={out.shape}`'
+            ' instead. When using depth rendering, the out array should be of'
+            ' shape `(width, height)` and otherwise (width, height, 3).'
+            f' Got `(self.height, self.width)={(self.height, self.width)}` and'
+            f' `self._depth_rendering={self._depth_rendering}`.'
+        )
+
+    # Render scene and read contents of RGB and depth buffers.
+    mujoco.mjr_render(self._rect, self._scene, self._mjr_context)
+    if self._depth_rendering:
+      mujoco.mjr_readPixels(None, out, self._rect, self._mjr_context)
+
+      # 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
+
+      # Calculate OpenGL perspective matrix values in float32 precision
+      # so they are close to what glFrustum returns
+      # https://registry.khronos.org/OpenGL-Refpages/gl2.1/xhtml/glFrustum.xml
+      zfar = np.float32(far)
+      znear = np.float32(near)
+      c_coef = -(zfar + znear) / (zfar - znear)
+      d_coef = -(np.float32(2) * zfar * znear) / (zfar - znear)
+
+      # In reverse Z mode the perspective matrix is transformed by the following
+      c_coef = np.float32(-0.5) * c_coef - np.float32(0.5)
+      d_coef = np.float32(-0.5) * d_coef
+
+      # We need 64 bits to convert Z from ndc to metric depth without noticeable
+      # losses in precision
+      out_64 = out.astype(np.float64)
+
+      # Undo OpenGL projection
+      # Note: We do not need to take action to convert from window coordinates
+      # to normalized device coordinates because in reversed Z mode the mapping
+      # is identity
+      out_64 = d_coef / (out_64 + c_coef)
+
+      # Cast result back to float32 for backwards compatibility
+      # This has a small accuracy cost
+      out[:] = out_64.astype(np.float32)
+
+      # Reset scene flags.
+      np.copyto(self._scene.flags, original_flags)
+    elif self._segmentation_rendering:
+      mujoco.mjr_readPixels(out, None, self._rect, self._mjr_context)
+
+      # Convert 3-channel uint8 to 1-channel uint32.
+      image3 = out.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
+      out = segid2output[segimage]
+
+      # Reset scene flags.
+      np.copyto(self._scene.flags, original_flags)
+    else:
+      mujoco.mjr_readPixels(out, None, self._rect, self._mjr_context)
+
+    if self._gl_context:
+      # If using EGL, OSMesa, or GLFW, the output image is flipped vertically.
+      # No flip is needed for Filament.
+      out[:] = np.flipud(out)
+
+    return out
+
+  def update_scene(
+      self,
+      data: mujoco.MjData,
+      camera: Union[int, str, mujoco.MjvCamera] = -1,
+      scene_option: Optional[mujoco.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.
+
+    Raises:
+      ValueError: If `camera_id` is outside the valid range, or if camera does
+        not exist.
+    """
+    if not isinstance(camera, mujoco.MjvCamera):
+      camera_id = camera
+      if isinstance(camera_id, str):
+        camera_id = mujoco.mj_name2id(
+            self._model, mujoco.mjtObj.mjOBJ_CAMERA.value, camera_id
+        )
+        if camera_id == -1:
+          raise ValueError(f'The camera "{camera}" does not exist.')
+      if camera_id < -1 or camera_id >= self._model.ncam:
+        raise ValueError(
+            f'The camera id {camera_id} is out of'
+            f' range [-1, {self._model.ncam}).'
+        )
+
+      # Render camera.
+      camera = mujoco.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 = mujoco.mjtCamera.mjCAMERA_FREE
+        mujoco.mjv_defaultFreeCamera(self._model, camera)
+      else:
+        camera.type = mujoco.mjtCamera.mjCAMERA_FIXED
+
+    scene_option = scene_option or self._scene_option
+    mujoco.mjv_updateScene(
+        self._model,
+        data,
+        scene_option,
+        None,
+        camera,
+        mujoco.mjtCatBit.mjCAT_ALL.value,
+        self._scene,
+    )
+
+  def close(self) -> None:
+    """Frees the resources used by the renderer.
+
+    This method can be used directly:
+
+    ```python
+    renderer = Renderer(...)
+    # Use renderer.
+    renderer.close()
+    ```
+
+    or via a context manager:
+
+    ```python
+    with Renderer(...) as renderer:
+      # Use renderer.
+    ```
+    """
+    if self._gl_context:
+      self._gl_context.free()
+    self._gl_context = None
+    if self._mjr_context:
+      self._mjr_context.free()
+    self._mjr_context = None
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    del exc_type, exc_value, traceback  # Unused.
+    self.close()
+
+  def __del__(self) -> None:
+    self.close()