Refactor MJX warp RenderContext, get rid of brittle pytree owner pattern.

PiperOrigin-RevId: 875901011
Change-Id: I2ba0734d9ca001ff91361aab6412fe370919344b

doc/mjx.rst

@@ -215,8 +215,11 @@ MJX-Warp Batch Rendering
 MJX-Warp includes a hardware-accelerated batch renderer for generating pixel observations (such as RGB and depth)
 across multiple parallel environments.
 
-To use the batch renderer, you must first create a specialized render context that allocates the necessary buffers.
-Note that the number of parallel worlds (``nworld``) is fixed when creating the context:
+To use the batch renderer, you must first create a render context that allocates the necessary buffers.
+Note that the number of parallel worlds (``nworld``) is fixed when creating the context.
+``create_render_context`` returns a render context object that provides direct access to buffer
+metadata (camera resolution, addresses, etc.). Call ``.pytree()`` to obtain the lightweight JAX
+pytree that should be passed into ``jit``/``vmap``-compiled functions:
 
 .. code-block:: python
 
@@ -233,6 +236,10 @@ Note that the number of parallel worlds (``nworld``) is fixed when creating the
         enabled_geom_groups=[0, 1, 2],
     )
 
+Hold a reference to ``rc`` for the lifetime of your program and pass ``rc.pytree()`` to
+downstream JAX functions.  The pytree is a lightweight handle that refers back to the
+context via an internal registry.
+
 Once the context is created, you can render images within a compiled JAX function. This involves updating the bounding
 volume hierarchy (BVH) and executing the raycaster:
 
@@ -241,20 +248,20 @@ volume hierarchy (BVH) and executing the raycaster:
     from mujoco.mjx import get_rgb
 
     @jax.jit
-    def render_fn(mx, d, rc):
+    def render_fn(mx, d, rc_pytree):
         # 1. Update the BVH for the current scene state
-        d = mjx.refit_bvh(mx, d, rc)
+        d = mjx.refit_bvh(mx, d, rc_pytree)
 
         # 2. Render all configured cameras
-        pixels, _ = mjx.render(mx, d, rc)
+        pixels, _ = mjx.render(mx, d, rc_pytree)
 
         # 3. Extract the RGB tensor for the first camera (index 0)
-        rgb = get_rgb(rc, 0, pixels)
+        rgb = get_rgb(rc_pytree, 0, pixels)
 
-        # CAVEAT: Always return or use the updated `d` in your computation graph.
-        # Otherwise, JAX's dead-code elimination will optimize away the refit_bvh call!
         return rgb, d
 
+    rgb, d = render_fn(mx, d, rc.pytree())
+
 Multi-GPU rendering with ``pmap``
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

mjx/mujoco/mjx/_src/bvh.py

@@ -26,7 +26,15 @@ import mujoco.mjx.warp as mjxw
 def refit_bvh(m: Model, d: Data, ctx: Any):
   """Refit the scene BVH for the current pose."""
   if m.impl == Impl.WARP and d.impl == Impl.WARP and mjxw.WARP_INSTALLED:
+    import mujoco.mjx.warp.render_context as mjxw_rc  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
     from mujoco.mjx.warp import bvh as mjxw_bvh  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+
+    if not isinstance(ctx, mjxw_rc.RenderContextPytree):
+      raise TypeError(
+          f'Expected RenderContextPytree, got {type(ctx).__name__}.'
+          ' Use rc.pytree() to get the JAX-compatible handle.'
+      )
+
     return mjxw_bvh.refit_bvh(m, d, ctx)
 
   raise NotImplementedError('refit_bvh only implemented for MuJoCo Warp.')

mjx/mujoco/mjx/_src/render.py

@@ -26,7 +26,15 @@ import mujoco.mjx.warp as mjxw
 def render(m: Model, d: Data, ctx: Any) -> Data:
   """Render."""
   if m.impl == Impl.WARP and d.impl == Impl.WARP and mjxw.WARP_INSTALLED:
+    import mujoco.mjx.warp.render_context as mjxw_rc  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
     from mujoco.mjx.warp import render as mjxw_render  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+
+    if not isinstance(ctx, mjxw_rc.RenderContextPytree):
+      raise TypeError(
+          f'Expected RenderContextPytree, got {type(ctx).__name__}.'
+          ' Use rc.pytree() to get the JAX-compatible handle.'
+      )
+
     return mjxw_render.render(m, d, ctx)
 
   raise NotImplementedError('render only implemented for MuJoCo Warp.')

mjx/mujoco/mjx/_src/render_util.py

@@ -14,22 +14,25 @@
 # ==============================================================================
 """JAX render utilities for unpacking render output from MuJoCo Warp."""
 
-from typing import Any
+from typing import TYPE_CHECKING
 
 import jax
 import jax.numpy as jnp
 import mujoco.mjx.warp as mjxw
 
+if TYPE_CHECKING:
+  from mujoco.mjx.warp.render_context import RenderContextPytree
+
 
 def get_rgb(
-    rc: Any,
+    rc: 'RenderContextPytree',
     cam_id: int,
     rgb_data: jax.Array,
 ) -> jax.Array:
   """Unpack uint32 ABGR pixel data into float32 RGB.
 
   Args:
-    rc: The RenderContext handle.
+    rc: RenderContext or RenderContextPytree (both have .key).
     cam_id: Camera index to extract.
     rgb_data: Packed render output, shape (total_pixels,) as uint32.
 
@@ -39,12 +42,18 @@ def get_rgb(
   Raises:
     RuntimeError: If Warp is not installed.
   """
-  if mjxw.WARP_INSTALLED:
-    import mujoco.mjx.warp.render as mjxw_render  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
-  else:
+  if not mjxw.WARP_INSTALLED:
     raise RuntimeError('Warp not installed.')
 
-  warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[(rc.key, None)]
+  import mujoco.mjx.warp.render_context as mjxw_rc  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+
+  if not isinstance(rc, mjxw_rc.RenderContextPytree):
+    raise TypeError(
+        f'Expected RenderContextPytree, got {type(rc).__name__}.'
+        ' Use rc.pytree() to get the JAX-compatible handle.'
+    )
+
+  warp_rc = mjxw_rc._MJX_RENDER_CONTEXT_BUFFERS[(rc.key, None)]  # pylint: disable=protected-access
   rgb_adr = int(warp_rc.rgb_adr.numpy()[cam_id])
   width = int(warp_rc.cam_res.numpy()[cam_id][0])
   height = int(warp_rc.cam_res.numpy()[cam_id][1])
@@ -61,7 +70,7 @@ def get_rgb(
 
 
 def get_depth(
-    rc: Any,
+    rc: 'RenderContextPytree',
     cam_id: int,
     depth_data: jax.Array,
     depth_scale: float,
@@ -69,7 +78,7 @@ def get_depth(
   """Extract and normalize depth data for a camera.
 
   Args:
-    rc: The RenderContext handle.
+    rc: RenderContext or RenderContextPytree (both have .key).
     cam_id: Camera index to extract.
     depth_data: Raw depth output, shape (total_pixels,) as float32.
     depth_scale: Scale factor for normalizing depth values.
@@ -80,11 +89,18 @@ def get_depth(
   Raises:
     RuntimeError: If Warp is not installed.
   """
-  if mjxw.WARP_INSTALLED:
-    import mujoco.mjx.warp.render as mjxw_render  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
-  else:
+  if not mjxw.WARP_INSTALLED:
     raise RuntimeError('Warp not installed.')
-  warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[(rc.key, None)]
+
+  import mujoco.mjx.warp.render_context as mjxw_rc  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+
+  if not isinstance(rc, mjxw_rc.RenderContextPytree):
+    raise TypeError(
+        f'Expected RenderContextPytree, got {type(rc).__name__}.'
+        ' Use rc.pytree() to get the JAX-compatible handle.'
+    )
+
+  warp_rc = mjxw_rc._MJX_RENDER_CONTEXT_BUFFERS[(rc.key, None)]  # pylint: disable=protected-access
   depth_adr = int(warp_rc.depth_adr.numpy()[cam_id])
   width = int(warp_rc.cam_res.numpy()[cam_id][0])
   height = int(warp_rc.cam_res.numpy()[cam_id][1])

mjx/mujoco/mjx/_src/render_util_test.py

@@ -21,6 +21,7 @@ import jax.numpy as jnp
 from mujoco.mjx._src import io
 from mujoco.mjx._src import render_util
 import mujoco.mjx.warp as mjxw
+from mujoco.mjx.warp.render_context import RenderContextPytree
 import numpy as np
 
 _FORCE_TEST = os.environ.get('MJX_WARP_FORCE_TEST', '0') == '1'
@@ -51,11 +52,11 @@ class RenderUtilTest(absltest.TestCase):
   def test_get_rgb(self):
     width, height = 4, 4
     warp_rc = _fake_render_context(1, width, height)
-    rc = mock.MagicMock(key=0)
+    rc = mock.MagicMock(spec=RenderContextPytree, key=0)
     rgb_data = jnp.zeros((width * height,), dtype=jnp.uint32)
 
     with mock.patch.dict(
-        'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS',
+        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
         {(0, None): warp_rc},
     ):
       rgb = jax.jit(render_util.get_rgb, static_argnums=(0, 1))(rc, 0, rgb_data)
@@ -65,11 +66,11 @@ class RenderUtilTest(absltest.TestCase):
   def test_get_rgb_vmap(self):
     nworld, width, height = 3, 4, 4
     warp_rc = _fake_render_context(1, width, height)
-    rc = mock.MagicMock(key=0)
+    rc = mock.MagicMock(spec=RenderContextPytree, key=0)
     rgb_data = jnp.zeros((nworld, width * height), dtype=jnp.uint32)
 
     with mock.patch.dict(
-        'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS',
+        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
         {(0, None): warp_rc},
     ):
       rgb = jax.jit(
@@ -82,11 +83,11 @@ class RenderUtilTest(absltest.TestCase):
   def test_get_depth(self):
     width, height = 4, 4
     warp_rc = _fake_render_context(1, width, height)
-    rc = mock.MagicMock(key=0)
+    rc = mock.MagicMock(spec=RenderContextPytree, key=0)
     depth_data = jnp.zeros((width * height,), dtype=jnp.float32)
 
     with mock.patch.dict(
-        'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS',
+        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
         {(0, None): warp_rc},
     ):
       depth = jax.jit(render_util.get_depth, static_argnums=(0, 1, 3))(
@@ -98,11 +99,11 @@ class RenderUtilTest(absltest.TestCase):
   def test_get_depth_vmap(self):
     nworld, width, height = 3, 4, 4
     warp_rc = _fake_render_context(1, width, height)
-    rc = mock.MagicMock(key=0)
+    rc = mock.MagicMock(spec=RenderContextPytree, key=0)
     depth_data = jnp.zeros((nworld, width * height), dtype=jnp.float32)
 
     with mock.patch.dict(
-        'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS',
+        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
         {(0, None): warp_rc},
     ):
       depth = jax.jit(

mjx/mujoco/mjx/warp/__init__.py

@@ -15,6 +15,7 @@
 
 import typing
 from typing import Any
+from mujoco.mjx.warp import render_context
 from mujoco.mjx.warp import types
 
 if not typing.TYPE_CHECKING:

mjx/mujoco/mjx/warp/bvh.py

@@ -14,14 +14,13 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
-
 import dataclasses
 import functools
 import jax
 from mujoco.mjx._src import types
 from mujoco.mjx.warp import ffi
-from mujoco.mjx.warp.io import _MJX_RENDER_CONTEXT_BUFFERS
-from mujoco.mjx.warp.types import RenderContext
+from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS
+from mujoco.mjx.warp.render_context import RenderContextPytree
 import mujoco.mjx.third_party.mujoco_warp as mjwarp
 from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 import warp as wp
@@ -46,7 +45,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _refit_bvh_shim(
     # Model
@@ -93,7 +91,9 @@ def _refit_bvh_shim(
   mjwarp.refit_bvh(_m, _d, render_context)
 
 
-def _refit_bvh_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext):
+def _refit_bvh_jax_impl(
+    m: types.Model, d: types.Data, ctx: RenderContextPytree
+):
   output_dims = {'dummy': (d.qpos.shape[0],)}
   jf = ffi.jax_callable_variadic_tuple(
       _refit_bvh_shim,
@@ -129,7 +129,7 @@ def _refit_bvh_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext):
 
 @jax.custom_batching.custom_vmap
 @ffi.marshal_jax_warp_callable
-def refit_bvh(m: types.Model, d: types.Data, ctx: RenderContext):
+def refit_bvh(m: types.Model, d: types.Data, ctx: RenderContextPytree):
   return _refit_bvh_jax_impl(m, d, ctx)
 
 
@@ -140,7 +140,7 @@ def refit_bvh_vmap(
     is_batched,
     m: types.Model,
     d: types.Data,
-    ctx: RenderContext,
+    ctx: RenderContextPytree,
 ):
   d = refit_bvh(m, d, ctx)
   return d, is_batched[1]

mjx/mujoco/mjx/warp/collision_driver.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
-
 import dataclasses
 import functools
 import jax
@@ -44,7 +43,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _collision_shim(
     # Model

mjx/mujoco/mjx/warp/forward.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
-
 import dataclasses
 import functools
 import jax
@@ -44,7 +43,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _forward_shim(
     # Model

mjx/mujoco/mjx/warp/io.py

@@ -14,16 +14,12 @@
 # ==============================================================================
 """I/O functions for MJX Warp."""
 
-import threading
-
 import mujoco
-from mujoco.mjx.warp.types import RenderContext
+from mujoco.mjx.warp import render_context
 import mujoco.mjx.third_party.mujoco_warp as mjw
 import warp as wp
 
 _MJX_RENDER_CONTEXT_COUNTER = 0
-_MJX_RENDER_CONTEXT_LOCK = threading.Lock()
-_MJX_RENDER_CONTEXT_BUFFERS = {}
 
 
 def _create_context(mjm, nworld, device, **kwargs):
@@ -42,20 +38,30 @@ def create_render_context(
     devices: list[str | None] | None = None,
     **kwargs,
 ):
+  """Creates a render context using mujoco_warp.create_render_context."""
   global _MJX_RENDER_CONTEXT_COUNTER
 
   if not devices:
     devices = [None]
 
-  contexts = [_create_context(mjm, nworld, d, **kwargs) for d in devices]
+  contexts = {}
+  default = None
+  for d in devices:
+    ctx = _create_context(mjm, nworld, d, **kwargs)
+    ordinal = wp.get_device(d).ordinal
+    contexts[ordinal] = ctx
+    if default is None:
+      default = ctx
 
-  with _MJX_RENDER_CONTEXT_LOCK:
+  # pylint: disable=protected-access
+  with render_context._MJX_RENDER_CONTEXT_LOCK:
     _MJX_RENDER_CONTEXT_COUNTER += 1
     key = _MJX_RENDER_CONTEXT_COUNTER
-    for d, ctx in zip(devices, contexts):
-      ordinal = wp.get_device(d).ordinal
-      _MJX_RENDER_CONTEXT_BUFFERS[(key, ordinal)] = ctx
-    if (key, None) not in _MJX_RENDER_CONTEXT_BUFFERS:
-      # save the first context as the default context
-      _MJX_RENDER_CONTEXT_BUFFERS[(key, None)] = contexts[0]
-  return RenderContext(key, _owner=True)
+    for ordinal, ctx in contexts.items():
+      render_context._MJX_RENDER_CONTEXT_BUFFERS[(key, ordinal)] = ctx
+    render_context._MJX_RENDER_CONTEXT_BUFFERS[(key, None)] = default
+  # pylint: enable=protected-access
+
+  return render_context.RenderContext(
+      key=key, contexts=contexts, default=default
+  )

mjx/mujoco/mjx/warp/render.py

@@ -14,14 +14,13 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
-
 import dataclasses
 import functools
 import jax
 from mujoco.mjx._src import types
 from mujoco.mjx.warp import ffi
-from mujoco.mjx.warp.io import _MJX_RENDER_CONTEXT_BUFFERS
-from mujoco.mjx.warp.types import RenderContext
+from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS
+from mujoco.mjx.warp.render_context import RenderContextPytree
 import mujoco.mjx.third_party.mujoco_warp as mjwarp
 from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 import warp as wp
@@ -46,7 +45,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _render_shim(
     # Model
@@ -116,7 +114,7 @@ def _render_shim(
   mjwarp.render(_m, _d, render_context)
 
 
-def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext):
+def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
   render_ctx = _MJX_RENDER_CONTEXT_BUFFERS[(ctx.key, None)]
   output_dims = {
       'rgb': render_ctx.rgb_data_shape,
@@ -181,7 +179,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext):
 
 @jax.custom_batching.custom_vmap
 @functools.partial(ffi.marshal_jax_warp_callable, tree_map_output=True)
-def render(m: types.Model, d: types.Data, ctx: RenderContext):
+def render(m: types.Model, d: types.Data, ctx: RenderContextPytree):
   return _render_jax_impl(m, d, ctx)
 
 
@@ -192,7 +190,7 @@ def render_vmap(
     is_batched,
     m: types.Model,
     d: types.Data,
-    ctx: RenderContext,
+    ctx: RenderContextPytree,
 ):
   out = render(m, d, ctx)
   return out, [True, True]

mjx/mujoco/mjx/warp/render_context.py

@@ -0,0 +1,70 @@
+# 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.
+# ==============================================================================
+"""MJX Warp render context types and buffer registry."""
+
+import threading
+
+from mujoco.mjx._src import dataclasses as mjx_dataclasses
+
+
+_MJX_RENDER_CONTEXT_LOCK = threading.Lock()
+_MJX_RENDER_CONTEXT_BUFFERS = {}
+
+
+class RenderContext:
+  """MJX render context wrapping one or more warp render contexts.
+
+  Returned by ``io.create_render_context``.  Holds the raw warp contexts
+  directly so callers can read camera resolution, buffer addresses, etc.
+
+  Use :meth:`pytree` to obtain the lightweight JAX-compatible handle
+  that should be passed into ``jit``/``vmap``-compiled functions such as
+  ``render`` and ``refit_bvh``.
+  """
+
+  def __init__(self, key, contexts, default):
+    self.key = key
+    self._contexts = contexts  # {device_ordinal: warp RenderContext}
+    self._default = default  # the first warp RenderContext
+
+  def pytree(self):
+    """Returns a lightweight JAX pytree for use in jit/vmap."""
+    return RenderContextPytree(self.key)
+
+  def __getattr__(self, name):
+    """Delegate attribute access to the default warp context."""
+    return getattr(self._default, name)
+
+  def __del__(self):
+    lock = _MJX_RENDER_CONTEXT_LOCK
+    buffers = _MJX_RENDER_CONTEXT_BUFFERS
+    if lock is None or buffers is None:
+      return
+    with lock:
+      keys_to_remove = [
+          k for k in buffers.keys() if isinstance(k, tuple) and k[0] == self.key
+      ]
+      for k in keys_to_remove:
+        buffers.pop(k, None)
+
+
+class RenderContextPytree(mjx_dataclasses.PyTreeNode):
+  """Minimal JAX pytree holding just the render context key.
+
+  The key is static (aux_data) so JAX doesn't trace it, allowing the
+  Warp FFI to receive a concrete int value.
+  """
+
+  key: int

mjx/mujoco/mjx/warp/render_test.py

@@ -77,7 +77,7 @@ class RenderTest(parameterized.TestCase):
       wp.config.kernel_cache_dir = tempdir
     np.random.seed(0)
 
-  def _skip_if_no_warp(self):
+  def _maybe_skip(self):
     if not _FORCE_TEST:
       if not mjxw.WARP_INSTALLED:
         self.skipTest('Warp not installed.')
@@ -90,11 +90,11 @@ class RenderTest(parameterized.TestCase):
   )
   def test_render(self, xml: str, batch_size: int):
     """Tests MJX render pipeline."""
-    self._skip_if_no_warp()
+    self._maybe_skip()
     mx, dx_batch, rc = _get_model_data_rc(xml, batch_size)
 
-    dx_batch = jax.jit(mjx.refit_bvh)(mx, dx_batch, rc)
-    out_batch = jax.jit(mjx.render)(mx, dx_batch, rc)
+    dx_batch = jax.jit(mjx.refit_bvh)(mx, dx_batch, rc.pytree())
+    out_batch = jax.jit(mjx.render)(mx, dx_batch, rc.pytree())
 
     rgb = np.asarray(out_batch[0])
     depth = np.asarray(out_batch[1])
@@ -110,7 +110,7 @@ class RenderTest(parameterized.TestCase):
   )
   def test_render_nested_vmap(self, xml: str, batch_size: int):
     """Tests MJX render pipeline with nested vmap."""
-    self._skip_if_no_warp()
+    self._maybe_skip()
     mx, dx_batch, rc = _get_model_data_rc(xml, batch_size)
 
     def inner(mx, dx, rc):
@@ -120,9 +120,11 @@ class RenderTest(parameterized.TestCase):
 
     # get reference with single vmap
     dx_batch = jax.vmap(bvh.refit_bvh, in_axes=(None, 0, None))(
-        mx, dx_batch, rc
+        mx, dx_batch, rc.pytree()
+    )
+    ref = jax.vmap(render.render, in_axes=(None, 0, None))(
+        mx, dx_batch, rc.pytree()
     )
-    ref = jax.vmap(render.render, in_axes=(None, 0, None))(mx, dx_batch, rc)
     ref_rgb = np.asarray(ref[0])
     ref_depth = np.asarray(ref[1])
 
@@ -134,7 +136,7 @@ class RenderTest(parameterized.TestCase):
 
     dx_2d = jax.tree.map(_reshape_batched, dx_batch)
 
-    out_batch = jax.vmap(inner, in_axes=(None, 0, None))(mx, dx_2d, rc)
+    out_batch = jax.vmap(inner, in_axes=(None, 0, None))(mx, dx_2d, rc.pytree())
     out_batch = jax.tree.map(lambda x: x.reshape(-1, *x.shape[2:]), out_batch)
     rgb = np.asarray(out_batch[0])
     depth = np.asarray(out_batch[1])
@@ -147,5 +149,97 @@ class RenderTest(parameterized.TestCase):
     self.assertNotEqual(np.unique(depth).shape[0], 1)
 
 
+class RenderContextGarbageCollectionTest(absltest.TestCase):
+  """Tests that RenderContext cleans up buffers on deletion."""
+
+  def setUp(self):
+    super().setUp()
+    if mjxw.WARP_INSTALLED:
+      tempdir = '/tmp/wp_kernel_cache_dir_RenderContextGCTest'
+      wp.config.kernel_cache_dir = tempdir
+
+  def _maybe_skip(self):
+    if not mjxw.WARP_INSTALLED:
+      self.skipTest('Warp not installed.')
+    if not io.has_cuda_gpu_device():
+      self.skipTest('No CUDA GPU device available.')
+
+  def test_render_context_gc(self):
+    """Verifies __del__ removes entries from _MJX_RENDER_CONTEXT_BUFFERS."""
+    self._maybe_skip()
+    from mujoco.mjx.warp import render_context as rc_module  # pylint: disable=g-import-not-at-top
+
+    self.assertEmpty(rc_module._MJX_RENDER_CONTEXT_BUFFERS)
+
+    _, _, rc = _get_model_data_rc('humanoid/humanoid.xml', 1)
+    key = rc.key
+
+    # Sanity check: buffers exist for this key.
+    matching = [
+        k
+        for k in rc_module._MJX_RENDER_CONTEXT_BUFFERS
+        if isinstance(k, tuple) and k[0] == key
+    ]
+    self.assertNotEmpty(matching)
+
+    # Delete the RenderContext and verify cleanup.
+    del rc
+
+    matching = [
+        k
+        for k in rc_module._MJX_RENDER_CONTEXT_BUFFERS
+        if isinstance(k, tuple) and k[0] == key
+    ]
+    self.assertEmpty(matching)
+
+  def test_render_context_gc_multi_keys(self):
+    """Verifies deleting one context doesn't remove another's buffers."""
+    self._maybe_skip()
+    from mujoco.mjx.warp import render_context as rc_module  # pylint: disable=g-import-not-at-top
+
+    self.assertEmpty(rc_module._MJX_RENDER_CONTEXT_BUFFERS)
+
+    _, _, rc_a = _get_model_data_rc('humanoid/humanoid.xml', 1)
+    _, _, rc_b = _get_model_data_rc('humanoid/humanoid.xml', 1)
+    key_a = rc_a.key
+    key_b = rc_b.key
+
+    # rc_a's buffers should be present.
+    matching_a = [
+        k
+        for k in rc_module._MJX_RENDER_CONTEXT_BUFFERS
+        if isinstance(k, tuple) and k[0] == key_a
+    ]
+    self.assertNotEmpty(matching_a)
+
+    # Delete only rc_a.
+    del rc_a
+
+    # rc_a's buffers should be gone.
+    matching_a = [
+        k
+        for k in rc_module._MJX_RENDER_CONTEXT_BUFFERS
+        if isinstance(k, tuple) and k[0] == key_a
+    ]
+    self.assertEmpty(matching_a)
+
+    # rc_b's buffers should still be present.
+    matching_b = [
+        k
+        for k in rc_module._MJX_RENDER_CONTEXT_BUFFERS
+        if isinstance(k, tuple) and k[0] == key_b
+    ]
+    self.assertNotEmpty(matching_b)
+
+    # Clean up rc_b.
+    del rc_b
+    matching_b = [
+        k
+        for k in rc_module._MJX_RENDER_CONTEXT_BUFFERS
+        if isinstance(k, tuple) and k[0] == key_b
+    ]
+    self.assertEmpty(matching_b)
+
+
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/warp/smooth.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
-
 import dataclasses
 import functools
 import jax
@@ -44,7 +43,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _kinematics_shim(
     # Model

mjx/mujoco/mjx/warp/testspeed.py

@@ -165,7 +165,9 @@ def benchmark(
       d = step_fn(mx, d)
 
       if render:
-        rgb, d = jax.vmap(render_fn, in_axes=(None, 0, None))(mx, d, rc)
+        rgb, d = jax.vmap(render_fn, in_axes=(None, 0, None))(
+            mx, d, rc.pytree()
+        )
         accum += rgb[0, 0, 0, 0]
 
       return (d, accum), None

mjx/mujoco/mjx/warp/types.py

@@ -13,10 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 """MJX Warp types.
-
 DO NOT EDIT. This file is auto-generated.
 """
-
 import dataclasses
 import typing
 from typing import Tuple
@@ -25,7 +23,6 @@ from jax import tree_util
 from jax.interpreters import batching
 from mujoco.mjx._src import dataclasses as mjx_dataclasses
 import numpy as np
-
 if typing.TYPE_CHECKING:
   GraphMode = int
 else:
@@ -35,7 +32,6 @@ else:
     GraphMode = int
 PyTreeNode = mjx_dataclasses.PyTreeNode
 
-
 @dataclasses.dataclass(frozen=True)
 @tree_util.register_pytree_node_class
 class TileSet:
@@ -68,7 +64,6 @@ class BlockDim:
 
   TODO(team): experimental and may be removed
   """
-
   actuator_velocity: int
   cholesky_factorize: int
   cholesky_factorize_solve: int
@@ -96,52 +91,12 @@ class BlockDim:
     return cls(*children)
 
 
-@dataclasses.dataclass(frozen=True)
-@tree_util.register_pytree_node_class
-class RenderContext:
-  """Render context handle with automatic cleanup.
-
-  The key is static (aux_data) so JAX doesn't trace it, allowing the Warp FFI
-  to receive a concrete int value. Only the original instance (_owner=True)
-  cleans up on deletion; JAX copies have _owner=False.
-  """
-
-  key: int
-  _owner: bool = True
-
-  def tree_flatten(self):
-    return ((), (self.key, False))
-
-  @classmethod
-  def tree_unflatten(cls, aux_data, children):
-    del children
-    key, owner = aux_data
-    return cls(key, owner)
-
-  def __del__(self):
-    if not self._owner:
-      return
-    lock = globals().get('_MJX_RENDER_CONTEXT_LOCK')
-    buffers = globals().get('_MJX_RENDER_CONTEXT_BUFFERS')
-    if lock is None or buffers is None:
-      return
-    with lock:
-      keys_to_remove = [
-          k for k in buffers.keys() if isinstance(k, tuple) and k[0] == self.key
-      ]
-      for k in keys_to_remove:
-        buffers.pop(k, None)
-
-
 class StatisticWarp(PyTreeNode):
   """Derived fields from Statistic."""
-
   meaninertia: jax.Array
 
-
 class OptionWarp(PyTreeNode):
   """Derived fields from Option."""
-
   broadphase: int
   broadphase_filter: int
   ccd_iterations: int
@@ -156,10 +111,8 @@ class OptionWarp(PyTreeNode):
   sdf_initpoints: int
   sdf_iterations: int
 
-
 class ModelWarp(PyTreeNode):
   """Derived fields from Model."""
-
   M_colind: np.ndarray
   M_rowadr: np.ndarray
   M_rownnz: np.ndarray
@@ -292,10 +245,8 @@ class ModelWarp(PyTreeNode):
   wrap_site_adr: np.ndarray
   wrap_site_pair_adr: np.ndarray
 
-
 class DataWarp(PyTreeNode):
   """Derived fields from Data."""
-
   actuator_moment: jax.Array
   actuator_velocity: jax.Array
   cacc: jax.Array
@@ -361,8 +312,6 @@ class DataWarp(PyTreeNode):
   wrap_obj: jax.Array
   wrap_xpos: jax.Array
   shape = property(lambda self: self.cacc.shape)
-
-
 DATA_NON_VMAP = {
     'contact__dim',
     'contact__dist',
@@ -386,7 +335,6 @@ DATA_NON_VMAP = {
     'nworld',
 }
 
-
 def _to_elt(cont, _, d, axis):
   return DataWarp(**{
       f.name: (

mjx/mujoco/mjx/warp/visualize_render.py

@@ -148,11 +148,11 @@ def _main(_: Sequence[str]):
   )
 
   dx_batch = jax_jit(jax.vmap(bvh.refit_bvh, in_axes=(None, 0, None)))(
-      mx, dx_batch, rc
+      mx, dx_batch, rc.pytree()
   )
 
   out_batch = jax_jit(jax.vmap(render.render, in_axes=(None, 0, None)))(
-      mx, dx_batch, rc
+      mx, dx_batch, rc.pytree()
   )
 
   rgb_packed = out_batch[0]
@@ -161,11 +161,11 @@ def _main(_: Sequence[str]):
   print(f'  depth shape: {depth_packed.shape}\n')
 
   rgb = jax.vmap(render_util.get_rgb, in_axes=(None, None, 0))(
-      rc, _CAMERA_ID.value, rgb_packed
+      rc.pytree(), _CAMERA_ID.value, rgb_packed
   )
 
   depth = jax.vmap(render_util.get_depth, in_axes=(None, None, 0, None))(
-      rc, _CAMERA_ID.value, depth_packed, 10.0
+      rc.pytree(), _CAMERA_ID.value, depth_packed, 10.0
   )
 
   single_path = os.path.join(
@@ -232,9 +232,9 @@ def _main(_: Sequence[str]):
     mx_pmap = jax.tree.map(lambda x: safe_shard(x, sharded), mx)
 
     def inner(mx, dx):
-      dx = bvh.refit_bvh(mx, dx, pmap_rc)
-      out = render.render(mx, dx, pmap_rc)
-      return render_util.get_rgb(pmap_rc, _CAMERA_ID.value, out[0])
+      dx = bvh.refit_bvh(mx, dx, pmap_rc.pytree())
+      out = render.render(mx, dx, pmap_rc.pytree())
+      return render_util.get_rgb(pmap_rc.pytree(), _CAMERA_ID.value, out[0])
 
     inner = jax.vmap(inner, in_axes=(None, 0))
     out = jax.pmap(inner)(mx_pmap, dx_pmap)