Add a close() method to mujoco.Renderer which explicitly frees the gl context.

PiperOrigin-RevId: 590221484
Change-Id: I07392f9c53f662a37a57ab23acfef9aa5a98a49b

python/mujoco/cgl/__init__.py

@@ -57,15 +57,18 @@ class GLContext:
 
   def free(self):
     """Frees resources associated with this context."""
-    if self._context:
-      cgl.CGLUnlockContext(self._context)
-      cgl.CGLSetCurrentContext(None)
-      cgl.CGLReleaseContext(self._context)
-    self._context = None
+    try:
+      if self._context:
+        cgl.CGLUnlockContext(self._context)
+        cgl.CGLSetCurrentContext(None)
+        cgl.CGLReleaseContext(self._context)
+      self._context = None
 
-    if self._pix:
-      cgl.CGLReleasePixelFormat(self._pix)
-    self._context = None
+      if self._pix:
+        cgl.CGLReleasePixelFormat(self._pix)
+      self._pix = None
+    except Exception:  # pylint: disable=broad-exception-caught
+      pass
 
   def __del__(self):
     self.free()

python/mujoco/glfw/__init__.py

@@ -35,7 +35,7 @@ class GLContext:
       if glfw.get_current_context() == self._context:
         glfw.make_context_current(None)
       glfw.destroy_window(self._context)
-      self._context = None
+    self._context = None
 
   def __del__(self):
     self.free()

python/mujoco/renderer.py

@@ -136,6 +136,9 @@ the clause:
       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()
 
@@ -145,6 +148,8 @@ the clause:
       self._scene.flags[_enums.mjtRndFlag.mjRND_SEGMENT] = True
       self._scene.flags[_enums.mjtRndFlag.mjRND_IDCOLOR] = True
 
+    if self._gl_context is None:
+      raise RuntimeError('render cannot be called after close.')
     self._gl_context.make_current()
 
     if self._depth_rendering:
@@ -288,3 +293,35 @@ the clause:
         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
+
+  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/renderer_test.py

@@ -33,10 +33,10 @@ class MuJoCoRendererTest(parameterized.TestCase):
 """
     model = mujoco.MjModel.from_xml_string(xml)
     data = mujoco.MjData(model)
-    renderer = mujoco.Renderer(model, 50, 50)
-    mujoco.mj_forward(model, data)
-    with self.assertRaisesRegex(ValueError, r'camera "b" does not exist'):
-      renderer.update_scene(data, 'b')
+    with mujoco.Renderer(model, 50, 50) as renderer:
+      mujoco.mj_forward(model, data)
+      with self.assertRaisesRegex(ValueError, r'camera "b" does not exist'):
+        renderer.update_scene(data, 'b')
 
   def test_renderer_camera_under_range(self):
     xml = """
@@ -48,10 +48,10 @@ class MuJoCoRendererTest(parameterized.TestCase):
 """
     model = mujoco.MjModel.from_xml_string(xml)
     data = mujoco.MjData(model)
-    renderer = mujoco.Renderer(model, 50, 50)
-    mujoco.mj_forward(model, data)
-    with self.assertRaisesRegex(ValueError, '-2 is out of range'):
-      renderer.update_scene(data, -2)
+    with mujoco.Renderer(model, 50, 50) as renderer:
+      mujoco.mj_forward(model, data)
+      with self.assertRaisesRegex(ValueError, '-2 is out of range'):
+        renderer.update_scene(data, -2)
 
   def test_renderer_camera_over_range(self):
     xml = """
@@ -63,10 +63,10 @@ class MuJoCoRendererTest(parameterized.TestCase):
 """
     model = mujoco.MjModel.from_xml_string(xml)
     data = mujoco.MjData(model)
-    renderer = mujoco.Renderer(model, 50, 50)
-    mujoco.mj_forward(model, data)
-    with self.assertRaisesRegex(ValueError, '1 is out of range'):
-      renderer.update_scene(data, 1)
+    with mujoco.Renderer(model, 50, 50) as renderer:
+      mujoco.mj_forward(model, data)
+      with self.assertRaisesRegex(ValueError, '1 is out of range'):
+        renderer.update_scene(data, 1)
 
   def test_renderer_renders_scene(self):
     xml = """
@@ -79,19 +79,19 @@ class MuJoCoRendererTest(parameterized.TestCase):
 """
     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')
+    with mujoco.Renderer(model, 50, 50) as renderer:
+      mujoco.mj_forward(model, data)
+      renderer.update_scene(data, 'closeup')
 
-    pixels = renderer.render().flatten()
-    not_all_black = False
+      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)
+      # Pixels should all be a neutral color.
+      for pixel in pixels:
+        if pixel > 0:
+          not_all_black = True
+          break
+      self.assertTrue(not_all_black)
 
   def test_renderer_output_without_out(self):
     xml = """
@@ -105,25 +105,25 @@ class MuJoCoRendererTest(parameterized.TestCase):
     model = mujoco.MjModel.from_xml_string(xml)
     data = mujoco.MjData(model)
     mujoco.mj_forward(model, data)
-    renderer = mujoco.Renderer(model, 50, 50)
-    renderer.update_scene(data, 'closeup')
-    pixels = [renderer.render()]
-
-    colors = (
-        (1.0, 0.0, 0.0, 1.0),
-        (0.0, 1.0, 0.0, 1.0),
-        (0.0, 0.0, 1.0, 1.0),
-    )
-
-    for i, color in enumerate(colors):
-      model.geom_rgba[0, :] = color
-      mujoco.mj_forward(model, data)
+    with mujoco.Renderer(model, 50, 50) as renderer:
       renderer.update_scene(data, 'closeup')
-      pixels.append(renderer.render())
-      self.assertIsNot(pixels[-2], pixels[-1])
+      pixels = [renderer.render()]
 
-      # Pixels should change over steps.
-      self.assertFalse((pixels[i + 1] == pixels[i]).all())
+      colors = (
+          (1.0, 0.0, 0.0, 1.0),
+          (0.0, 1.0, 0.0, 1.0),
+          (0.0, 0.0, 1.0, 1.0),
+      )
+
+      for i, color in enumerate(colors):
+        model.geom_rgba[0, :] = color
+        mujoco.mj_forward(model, data)
+        renderer.update_scene(data, 'closeup')
+        pixels.append(renderer.render())
+        self.assertIsNot(pixels[-2], pixels[-1])
+
+        # Pixels should change over steps.
+        self.assertFalse((pixels[i + 1] == pixels[i]).all())
 
   def test_renderer_output_with_out(self):
     xml = """
@@ -139,23 +139,21 @@ class MuJoCoRendererTest(parameterized.TestCase):
     model = mujoco.MjModel.from_xml_string(xml)
     data = mujoco.MjData(model)
     mujoco.mj_forward(model, data)
-    renderer = mujoco.Renderer(model, *render_size)
-    renderer.update_scene(data, 'closeup')
+    with mujoco.Renderer(model, *render_size) as renderer:
+      renderer.update_scene(data, 'closeup')
 
-    self.assertTrue(np.all(render_out == 0))
+      self.assertTrue(np.all(render_out == 0))
 
-    pixels = renderer.render(out=render_out)
+      pixels = renderer.render(out=render_out)
 
-    # Pixels should always refer to the same `render_out` array.
-    self.assertIs(pixels, render_out)
-    self.assertFalse(np.all(render_out == 0))
+      # Pixels should always refer to the same `render_out` array.
+      self.assertIs(pixels, render_out)
+      self.assertFalse(np.all(render_out == 0))
 
-    failing_render_size = (10, 10)
-    self.assertNotEqual(failing_render_size, render_size)
-    with self.assertRaises(ValueError):
-      pixels = renderer.render(
-          out=np.zeros((*failing_render_size, 3), np.uint8)
-      )
+      failing_render_size = (10, 10)
+      self.assertNotEqual(failing_render_size, render_size)
+      with self.assertRaises(ValueError):
+        renderer.render(out=np.zeros((*failing_render_size, 3), np.uint8))
 
 
 if __name__ == '__main__':