Merge pull request #3216 from tkelestemur:mjx-render-batched-unpackers

PiperOrigin-RevId: 899017204
Change-Id: I7409a00cff81a1d24afe2ca50ec14c2f25c103ea

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

@@ -18,11 +18,12 @@ from unittest import mock
 from absl.testing import absltest
 import jax
 import jax.numpy as jnp
+import numpy as np
+
 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'
 
@@ -63,6 +64,26 @@ class RenderUtilTest(absltest.TestCase):
 
     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.patch.dict(
+        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
+        {(0, None): 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)
@@ -96,6 +117,26 @@ class RenderUtilTest(absltest.TestCase):
 
     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.patch.dict(
+        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
+        {(0, None): 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)