From 07cda5a3c490774dfe586a911d73e65c0bc22261 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Thu, 26 Feb 2026 14:48:56 -0800 Subject: [PATCH] Refactor MJX warp RenderContext, get rid of brittle pytree owner pattern. PiperOrigin-RevId: 875901011 Change-Id: I2ba0734d9ca001ff91361aab6412fe370919344b --- doc/mjx.rst | 23 +++-- mjx/mujoco/mjx/_src/bvh.py | 8 ++ mjx/mujoco/mjx/_src/render.py | 8 ++ mjx/mujoco/mjx/_src/render_util.py | 42 ++++++--- mjx/mujoco/mjx/_src/render_util_test.py | 17 ++-- mjx/mujoco/mjx/warp/__init__.py | 1 + mjx/mujoco/mjx/warp/bvh.py | 14 +-- mjx/mujoco/mjx/warp/collision_driver.py | 2 - mjx/mujoco/mjx/warp/forward.py | 2 - mjx/mujoco/mjx/warp/io.py | 34 +++++--- mjx/mujoco/mjx/warp/render.py | 12 ++- mjx/mujoco/mjx/warp/render_context.py | 70 +++++++++++++++ mjx/mujoco/mjx/warp/render_test.py | 110 ++++++++++++++++++++++-- mjx/mujoco/mjx/warp/smooth.py | 2 - mjx/mujoco/mjx/warp/testspeed.py | 4 +- mjx/mujoco/mjx/warp/types.py | 52 ----------- mjx/mujoco/mjx/warp/visualize_render.py | 14 +-- 17 files changed, 284 insertions(+), 131 deletions(-) create mode 100644 mjx/mujoco/mjx/warp/render_context.py diff --git a/doc/mjx.rst b/doc/mjx.rst index a68d57b1..b6a79671 100644 --- a/doc/mjx.rst +++ b/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`` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/mjx/mujoco/mjx/_src/bvh.py b/mjx/mujoco/mjx/_src/bvh.py index a35f4f6e..1a75a8a0 100644 --- a/mjx/mujoco/mjx/_src/bvh.py +++ b/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.') diff --git a/mjx/mujoco/mjx/_src/render.py b/mjx/mujoco/mjx/_src/render.py index 5577ffa6..45f72aed 100644 --- a/mjx/mujoco/mjx/_src/render.py +++ b/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.') diff --git a/mjx/mujoco/mjx/_src/render_util.py b/mjx/mujoco/mjx/_src/render_util.py index cfe6bf0b..54ab6af7 100644 --- a/mjx/mujoco/mjx/_src/render_util.py +++ b/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]) diff --git a/mjx/mujoco/mjx/_src/render_util_test.py b/mjx/mujoco/mjx/_src/render_util_test.py index ed82f5c0..46c74420 100644 --- a/mjx/mujoco/mjx/_src/render_util_test.py +++ b/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( diff --git a/mjx/mujoco/mjx/warp/__init__.py b/mjx/mujoco/mjx/warp/__init__.py index 92d83d4a..cc582e4a 100644 --- a/mjx/mujoco/mjx/warp/__init__.py +++ b/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: diff --git a/mjx/mujoco/mjx/warp/bvh.py b/mjx/mujoco/mjx/warp/bvh.py index 529145ee..0e9287d4 100644 --- a/mjx/mujoco/mjx/warp/bvh.py +++ b/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] diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index 2ae01123..b5822eb2 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/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 diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index e48219a8..11f7dd9f 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/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 diff --git a/mjx/mujoco/mjx/warp/io.py b/mjx/mujoco/mjx/warp/io.py index db0f2df9..b32d7122 100644 --- a/mjx/mujoco/mjx/warp/io.py +++ b/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 + ) diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py index 3b2d3aee..446431e7 100644 --- a/mjx/mujoco/mjx/warp/render.py +++ b/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] diff --git a/mjx/mujoco/mjx/warp/render_context.py b/mjx/mujoco/mjx/warp/render_context.py new file mode 100644 index 00000000..774e6030 --- /dev/null +++ b/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 diff --git a/mjx/mujoco/mjx/warp/render_test.py b/mjx/mujoco/mjx/warp/render_test.py index 7157f0cc..d4864d9c 100644 --- a/mjx/mujoco/mjx/warp/render_test.py +++ b/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() diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py index 0e1ec813..310b3562 100644 --- a/mjx/mujoco/mjx/warp/smooth.py +++ b/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 diff --git a/mjx/mujoco/mjx/warp/testspeed.py b/mjx/mujoco/mjx/warp/testspeed.py index e2525a8d..d8869263 100644 --- a/mjx/mujoco/mjx/warp/testspeed.py +++ b/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 diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index 03d403a4..d86108f6 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/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: ( diff --git a/mjx/mujoco/mjx/warp/visualize_render.py b/mjx/mujoco/mjx/warp/visualize_render.py index 6a3b5afa..7e470856 100644 --- a/mjx/mujoco/mjx/warp/visualize_render.py +++ b/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)