diff --git a/dm_control/mujoco/wrapper/core.py b/dm_control/mujoco/wrapper/core.py index 95fa150a..23093756 100644 --- a/dm_control/mujoco/wrapper/core.py +++ b/dm_control/mujoco/wrapper/core.py @@ -619,16 +619,19 @@ def __init__(self, if not isinstance(font_scale, mujoco.mjtFontScale): font_scale = mujoco.mjtFontScale(font_scale) self._gl_context = gl_context + self._ptr = None + self._gl_context_refcounted = False with gl_context.make_current() as ctx: ptr = ctx.call(mujoco.MjrContext, model.ptr, font_scale) ctx.call(mujoco.mjr_setBuffer, mujoco.mjtFramebuffer.mjFB_OFFSCREEN, ptr) gl_context.keep_alive(ptr) gl_context.increment_refcount() + self._gl_context_refcounted = True self._ptr = weakref.ref(ptr) @property def ptr(self): - return self._ptr() + return self._ptr() if self._ptr is not None else None def free(self): """Frees the native resources held by this MjrContext. @@ -637,17 +640,24 @@ def free(self): necessary. This MjrContext object MUST NOT be used after this function has been called. """ - if self._gl_context and not self._gl_context.terminated: - ptr = self.ptr - if ptr: - self._gl_context.dont_keep_alive(ptr) - with self._gl_context.make_current() as ctx: - ctx.call(ptr.free) - - if self._gl_context: - self._gl_context.decrement_refcount() - self._gl_context.free() - self._gl_context = None + gl_context = self._gl_context + if gl_context is None: + return + + if self._gl_context_refcounted: + if not gl_context.terminated: + ptr = self.ptr + if ptr: + gl_context.dont_keep_alive(ptr) + with gl_context.make_current() as ctx: + ctx.call(ptr.free) + + gl_context.decrement_refcount() + gl_context.free() + self._gl_context_refcounted = False + + self._ptr = None + self._gl_context = None def __del__(self): self.free() diff --git a/dm_control/mujoco/wrapper/mjr_context_failure_test.py b/dm_control/mujoco/wrapper/mjr_context_failure_test.py new file mode 100644 index 00000000..cabe45a6 --- /dev/null +++ b/dm_control/mujoco/wrapper/mjr_context_failure_test.py @@ -0,0 +1,42 @@ +# Copyright 2026 The dm_control Authors. +# +# 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 cleanup of partially initialized MuJoCo rendering contexts.""" + +from unittest import mock + +from absl.testing import absltest +from dm_control.mujoco.wrapper import core + + +class MjrContextFailureTest(absltest.TestCase): + + def testFreePartiallyInitializedContextDoesNotReleaseGlContext(self): + gl_context = mock.MagicMock() + gl_context.terminated = False + mjr_context = core.MjrContext.__new__(core.MjrContext) + mjr_context._gl_context = gl_context + mjr_context._ptr = None + mjr_context._gl_context_refcounted = False + + mjr_context.free() + + gl_context.decrement_refcount.assert_not_called() + gl_context.free.assert_not_called() + self.assertIsNone(mjr_context.ptr) + + +if __name__ == "__main__": + absltest.main()