From de3912811408d61b661e82cf5d4bb64549b4f72f Mon Sep 17 00:00:00 2001 From: Tarik Kelestemur Date: Wed, 8 Apr 2026 14:09:45 -0400 Subject: [PATCH] fix render batching --- mjx/mujoco/mjx/_src/render_util.py | 19 +++--- mjx/mujoco/mjx/_src/render_util_test.py | 81 ++++++++++++++++--------- 2 files changed, 62 insertions(+), 38 deletions(-) diff --git a/mjx/mujoco/mjx/_src/render_util.py b/mjx/mujoco/mjx/_src/render_util.py index d0bd1110..098291c3 100644 --- a/mjx/mujoco/mjx/_src/render_util.py +++ b/mjx/mujoco/mjx/_src/render_util.py @@ -18,6 +18,7 @@ from typing import TYPE_CHECKING import jax import jax.numpy as jnp + import mujoco.mjx.warp as mjxw if TYPE_CHECKING: @@ -34,10 +35,11 @@ def get_rgb( Args: rc: RenderContextPytree. cam_id: Camera index to extract. - rgb_data: Packed render output, shape (total_pixels,) as uint32. + rgb_data: Packed render output, shape (..., total_pixels) as uint32. Returns: - Float32 RGB array with shape (H, W, 3), values in [0, 1]. + Float32 RGB array with shape (..., H, W, 3), values in [0, 1]. + Any leading batch axes in `rgb_data` are preserved. Raises: RuntimeError: If Warp is not installed. @@ -59,14 +61,14 @@ def get_rgb( height = int(warp_rc.cam_res.numpy()[cam_id][1]) packed = jax.lax.dynamic_slice_in_dim( - rgb_data, rgb_adr, width * height, axis=0 + rgb_data, rgb_adr, width * height, axis=rgb_data.ndim - 1 ) b = (packed & 0xFF).astype(jnp.float32) / 255.0 g = ((packed >> 8) & 0xFF).astype(jnp.float32) / 255.0 r = ((packed >> 16) & 0xFF).astype(jnp.float32) / 255.0 rgb = jnp.stack([r, g, b], axis=-1) - return rgb.reshape(height, width, 3) + return rgb.reshape(packed.shape[:-1] + (height, width, 3)) def get_depth( @@ -80,11 +82,12 @@ def get_depth( Args: rc: RenderContextPytree. cam_id: Camera index to extract. - depth_data: Raw depth output, shape (total_pixels,) as float32. + depth_data: Raw depth output, shape (..., total_pixels) as float32. depth_scale: Scale factor for normalizing depth values. Returns: - Float32 depth array with shape (H, W), clamped to [0, 1]. + Float32 depth array with shape (..., H, W, 1), clamped to [0, 1]. + Any leading batch axes in `depth_data` are preserved. Raises: RuntimeError: If Warp is not installed. @@ -106,8 +109,8 @@ def get_depth( height = int(warp_rc.cam_res.numpy()[cam_id][1]) raw = jax.lax.dynamic_slice_in_dim( - depth_data, depth_adr, width * height, axis=0 + depth_data, depth_adr, width * height, axis=depth_data.ndim - 1 ) depth = jnp.clip(raw / depth_scale, 0.0, 1.0) - return depth.reshape(height, width, 1) + return depth.reshape(raw.shape[:-1] + (height, width, 1)) diff --git a/mjx/mujoco/mjx/_src/render_util_test.py b/mjx/mujoco/mjx/_src/render_util_test.py index 7549b0af..b09ed7a3 100644 --- a/mjx/mujoco/mjx/_src/render_util_test.py +++ b/mjx/mujoco/mjx/_src/render_util_test.py @@ -12,19 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -import os +import contextlib from unittest import mock from absl.testing import absltest import jax 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' +from mujoco.mjx._src import render_util +from mujoco.mjx.warp.render_context import RenderContextPytree def _fake_render_context(ncam, width, height): @@ -39,15 +36,17 @@ def _fake_render_context(ncam, width, height): return rc -class RenderUtilTest(absltest.TestCase): +@contextlib.contextmanager +def _mock_render_runtime(warp_rc): + with mock.patch.object(render_util.mjxw, 'WARP_INSTALLED', True): + with mock.patch.dict( + 'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS', + {(0, None): warp_rc}, + ): + yield - def setUp(self): - super().setUp() - if not _FORCE_TEST: - if not mjxw.WARP_INSTALLED: - self.skipTest('Warp not installed.') - if not io.has_cuda_gpu_device(): - self.skipTest('No CUDA GPU device available.') + +class RenderUtilTest(absltest.TestCase): def test_get_rgb(self): width, height = 4, 4 @@ -55,24 +54,35 @@ class RenderUtilTest(absltest.TestCase): rc = mock.MagicMock(spec=RenderContextPytree, key=0) rgb_data = jnp.zeros((width * height,), dtype=jnp.uint32) - with mock.patch.dict( - 'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS', - {(0, None): warp_rc}, - ): + with _mock_render_runtime(warp_rc): rgb = jax.jit(render_util.get_rgb, static_argnums=(0, 1))(rc, 0, rgb_data) self.assertEqual(rgb.shape, (height, width, 3)) + def test_get_rgb_preserves_leading_dims(self): + width, height = 4, 4 + warp_rc = _fake_render_context(1, width, height) + rc = mock.MagicMock(spec=RenderContextPytree, key=0) + + with _mock_render_runtime(warp_rc): + for leading_shape in ((1,), (3,), (2, 3)): + with self.subTest(leading_shape=leading_shape): + rgb_data = jnp.zeros( + leading_shape + (width * height,), dtype=jnp.uint32 + ) + rgb = jax.jit(render_util.get_rgb, static_argnums=(0, 1))( + rc, 0, rgb_data + ) + + self.assertEqual(rgb.shape, leading_shape + (height, width, 3)) + def test_get_rgb_vmap(self): nworld, width, height = 3, 4, 4 warp_rc = _fake_render_context(1, width, height) 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_context._MJX_RENDER_CONTEXT_BUFFERS', - {(0, None): warp_rc}, - ): + with _mock_render_runtime(warp_rc): rgb = jax.jit( jax.vmap(render_util.get_rgb, in_axes=(None, None, 0)), static_argnums=(0, 1), @@ -86,26 +96,37 @@ class RenderUtilTest(absltest.TestCase): rc = mock.MagicMock(spec=RenderContextPytree, key=0) depth_data = jnp.zeros((width * height,), dtype=jnp.float32) - with mock.patch.dict( - 'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS', - {(0, None): warp_rc}, - ): + with _mock_render_runtime(warp_rc): depth = jax.jit(render_util.get_depth, static_argnums=(0, 1, 3))( rc, 0, depth_data, 5.0 ) self.assertEqual(depth.shape, (height, width, 1)) + def test_get_depth_preserves_leading_dims(self): + width, height = 4, 4 + warp_rc = _fake_render_context(1, width, height) + rc = mock.MagicMock(spec=RenderContextPytree, key=0) + + with _mock_render_runtime(warp_rc): + for leading_shape in ((1,), (3,), (2, 3)): + with self.subTest(leading_shape=leading_shape): + depth_data = jnp.zeros( + leading_shape + (width * height,), dtype=jnp.float32 + ) + depth = jax.jit(render_util.get_depth, static_argnums=(0, 1, 3))( + rc, 0, depth_data, 5.0 + ) + + self.assertEqual(depth.shape, leading_shape + (height, width, 1)) + def test_get_depth_vmap(self): nworld, width, height = 3, 4, 4 warp_rc = _fake_render_context(1, width, height) 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_context._MJX_RENDER_CONTEXT_BUFFERS', - {(0, None): warp_rc}, - ): + with _mock_render_runtime(warp_rc): depth = jax.jit( jax.vmap(render_util.get_depth, in_axes=(None, None, 0, None)), static_argnums=(0, 1, 3),