fix render batching

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))

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()
+class RenderUtilTest(absltest.TestCase):
-    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.')
 
   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(
+    with _mock_render_runtime(warp_rc):
-        '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)
 
     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(
+    with _mock_render_runtime(warp_rc):
-        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
-        {(0, None): 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(
+    with _mock_render_runtime(warp_rc):
-        '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))(
           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(
+    with _mock_render_runtime(warp_rc):
-        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
-        {(0, None): warp_rc},
-    ):
       depth = jax.jit(
           jax.vmap(render_util.get_depth, in_axes=(None, None, 0, None)),
           static_argnums=(0, 1, 3),