From 53883ce86a7da0f7df7016c009d4ce79eefc283e Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Tue, 12 Dec 2023 08:58:49 -0800 Subject: [PATCH] Add a `close()` method to `mujoco.Renderer` which explicitly frees the gl context. PiperOrigin-RevId: 590221484 Change-Id: I07392f9c53f662a37a57ab23acfef9aa5a98a49b --- python/mujoco/cgl/__init__.py | 19 +++--- python/mujoco/glfw/__init__.py | 2 +- python/mujoco/renderer.py | 37 ++++++++++++ python/mujoco/renderer_test.py | 104 ++++++++++++++++----------------- 4 files changed, 100 insertions(+), 62 deletions(-) diff --git a/python/mujoco/cgl/__init__.py b/python/mujoco/cgl/__init__.py index ba1e36a6..36037116 100644 --- a/python/mujoco/cgl/__init__.py +++ b/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() diff --git a/python/mujoco/glfw/__init__.py b/python/mujoco/glfw/__init__.py index a5693bec..0fd8bd9a 100644 --- a/python/mujoco/glfw/__init__.py +++ b/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() diff --git a/python/mujoco/renderer.py b/python/mujoco/renderer.py index f34c8243..dfb387c6 100644 --- a/python/mujoco/renderer.py +++ b/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() diff --git a/python/mujoco/renderer_test.py b/python/mujoco/renderer_test.py index 8aa44a97..e8022006 100644 --- a/python/mujoco/renderer_test.py +++ b/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__':