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,6 +57,7 @@ class GLContext:
 
   def free(self):
     """Frees resources associated with this context."""
+    try:
       if self._context:
         cgl.CGLUnlockContext(self._context)
         cgl.CGLSetCurrentContext(None)
@@ -65,7 +66,9 @@ class GLContext:
 
       if self._pix:
         cgl.CGLReleasePixelFormat(self._pix)
-    self._context = None
+      self._pix = None
+    except Exception:  # pylint: disable=broad-exception-caught
+      pass
 
   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,7 +33,7 @@ class MuJoCoRendererTest(parameterized.TestCase):
 """
     model = mujoco.MjModel.from_xml_string(xml)
     data = mujoco.MjData(model)
-    renderer = mujoco.Renderer(model, 50, 50)
+    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')
@@ -48,7 +48,7 @@ class MuJoCoRendererTest(parameterized.TestCase):
 """
     model = mujoco.MjModel.from_xml_string(xml)
     data = mujoco.MjData(model)
-    renderer = mujoco.Renderer(model, 50, 50)
+    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)
@@ -63,7 +63,7 @@ class MuJoCoRendererTest(parameterized.TestCase):
 """
     model = mujoco.MjModel.from_xml_string(xml)
     data = mujoco.MjData(model)
-    renderer = mujoco.Renderer(model, 50, 50)
+    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)
@@ -79,7 +79,7 @@ class MuJoCoRendererTest(parameterized.TestCase):
 """
     model = mujoco.MjModel.from_xml_string(xml)
     data = mujoco.MjData(model)
-    renderer = mujoco.Renderer(model, 50, 50)
+    with mujoco.Renderer(model, 50, 50) as renderer:
       mujoco.mj_forward(model, data)
       renderer.update_scene(data, 'closeup')
 
@@ -105,7 +105,7 @@ 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)
+    with mujoco.Renderer(model, 50, 50) as renderer:
       renderer.update_scene(data, 'closeup')
       pixels = [renderer.render()]
 
@@ -139,7 +139,7 @@ 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)
+    with mujoco.Renderer(model, *render_size) as renderer:
       renderer.update_scene(data, 'closeup')
 
       self.assertTrue(np.all(render_out == 0))
@@ -153,9 +153,7 @@ class MuJoCoRendererTest(parameterized.TestCase):
       failing_render_size = (10, 10)
       self.assertNotEqual(failing_render_size, render_size)
       with self.assertRaises(ValueError):
-      pixels = renderer.render(
+        renderer.render(out=np.zeros((*failing_render_size, 3), np.uint8))
-          out=np.zeros((*failing_render_size, 3), np.uint8)
-      )
 
 
 if __name__ == '__main__':