Fix bug in mjx-warp render vmap.

PiperOrigin-RevId: 871901512
Change-Id: Ifc25e98ffd89284cc6db9248bcca352bb8811e7e

mjx/mujoco/mjx/warp/ffi.py

@@ -252,7 +252,7 @@ def _squeeze_dim(leaf_expanded: Any, leaf: Any) -> Any:
   return leaf_expanded
 
 
-def marshal_jax_warp_callable(func, skip_output_dim_reshape: bool = False):
+def marshal_jax_warp_callable(func, tree_map_output: bool = False):
   """Marshal fields into a MuJoCo Warp function."""
 
   @functools.wraps(func)
@@ -273,7 +273,7 @@ def marshal_jax_warp_callable(func, skip_output_dim_reshape: bool = False):
     )
     d_expanded_result = func(m_expanded, d_expanded, *extra_args)
 
-    if skip_output_dim_reshape:
+    if tree_map_output:
       return d_expanded_result
     d_result = jax.tree.map(_squeeze_dim, d_expanded_result, d)
     return d_result
@@ -360,7 +360,10 @@ def _check_leading_dim(
     )
 
 
-def marshal_custom_vmap(vmap_func, skip_output_dim_reshape: bool = False):
+def marshal_custom_vmap(
+    vmap_func,
+    tree_map_output: bool = False,
+):
   """Marshal fields for a custom vmap into an MuJoCo Warp function."""
 
   @functools.wraps(vmap_func)
@@ -397,8 +400,11 @@ def marshal_custom_vmap(vmap_func, skip_output_dim_reshape: bool = False):
     d_broadcast_flat_result, out_batched = vmap_func(
         axis_size, is_batched, m_flat, d_broadcast_flat, *extra_args
     )
-    if skip_output_dim_reshape:
+    if tree_map_output:
-      return d_broadcast_flat_result, out_batched
+      out = jax.tree.map(
+          lambda x: x.reshape(axis_size, -1), d_broadcast_flat_result
+      )
+      return out, out_batched
 
     # Explicitly mark MuJoCo Warp data fields as batched after vmapping is done.
     out_batched = jax.tree.map_with_path(

mjx/mujoco/mjx/warp/render.py

@@ -179,13 +179,13 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext):
 
 
 @jax.custom_batching.custom_vmap
-@functools.partial(ffi.marshal_jax_warp_callable, skip_output_dim_reshape=True)
+@functools.partial(ffi.marshal_jax_warp_callable, tree_map_output=True)
 def render(m: types.Model, d: types.Data, ctx: RenderContext):
   return _render_jax_impl(m, d, ctx)
 
 
 @render.def_vmap
-@functools.partial(ffi.marshal_custom_vmap, skip_output_dim_reshape=True)
+@functools.partial(ffi.marshal_custom_vmap, tree_map_output=True)
 def render_vmap(
     unused_axis_size,
     is_batched,

mjx/mujoco/mjx/warp/render_test.py

@@ -21,8 +21,10 @@ import jax
 from jax import numpy as jp
 import mujoco
 from mujoco import mjx
+from mujoco.mjx._src import bvh
 from mujoco.mjx._src import forward
 from mujoco.mjx._src import io
+from mujoco.mjx._src import render
 import mujoco.mjx.warp as mjxw
 from mujoco.mjx.warp import test_util as tu
 from mujoco.mjx.warp import warp as wp  # pylint: disable=g-importing-member
@@ -32,6 +34,40 @@ import numpy as np
 _FORCE_TEST = os.environ.get('MJX_WARP_FORCE_TEST', '0') == '1'
 
 
+def _get_model_data_rc(xml, batch_size):
+  m = tu.load_test_file(xml)
+  d = mujoco.MjData(m)
+  mujoco.mj_forward(m, d)
+
+  mx = mjx.put_model(m, impl='warp')
+
+  worldids = jp.arange(batch_size)
+  dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids)
+
+  key = jax.random.PRNGKey(0)
+  keys = jax.random.split(key, batch_size)
+  qpos0 = jp.array(m.qpos0)
+  rand_qpos = jax.vmap(
+      lambda k: qpos0 + jax.random.uniform(k, (m.nq,), minval=-0.2, maxval=0.05)
+  )(keys)
+  dx_batch = jax.vmap(lambda dx, q: dx.replace(qpos=q))(dx_batch, rand_qpos)
+
+  dx_batch = jax.jit(jax.vmap(forward.forward, in_axes=(None, 0)))(mx, dx_batch)
+
+  width, height = 32, 32
+  rc = mjx.create_render_context(
+      mjm=m,
+      nworld=batch_size,
+      cam_res=(width, height),
+      use_textures=True,
+      use_shadows=True,
+      render_rgb=True,
+      render_depth=True,
+      enabled_geom_groups=[0, 1, 2],
+  )
+  return mx, dx_batch, rc
+
+
 class RenderTest(parameterized.TestCase):
 
   def setUp(self):
@@ -41,56 +77,21 @@ class RenderTest(parameterized.TestCase):
       wp.config.kernel_cache_dir = tempdir
     np.random.seed(0)
 
-  @parameterized.product(
+  def _skip_if_no_warp(self):
-      xml=(
-          'humanoid/humanoid.xml',
-      ),
-      batch_size=(1, 16),
-  )
-  def test_render(self, xml: str, batch_size: int):
-    """Tests MJX render pipeline."""
     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.')
 
-    m = tu.load_test_file(xml)
+  @parameterized.product(
-    d = mujoco.MjData(m)
+      xml=('humanoid/humanoid.xml',),
-    mujoco.mj_forward(m, d)
+      batch_size=(1, 16),
-
+  )
-    mx = mjx.put_model(m, impl='warp')
+  def test_render(self, xml: str, batch_size: int):
-
+    """Tests MJX render pipeline."""
-    worldids = jp.arange(batch_size)
+    self._skip_if_no_warp()
-    dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids)
+    mx, dx_batch, rc = _get_model_data_rc(xml, batch_size)
-
-    key = jax.random.PRNGKey(0)
-    keys = jax.random.split(key, batch_size)
-    qpos0 = jp.array(m.qpos0)
-    rand_qpos = jax.vmap(
-        lambda k: qpos0 + jax.random.uniform(
-            k, (m.nq,), minval=-0.2, maxval=0.05
-        )
-    )(keys)
-    dx_batch = jax.vmap(
-        lambda dx, q: dx.replace(qpos=q)
-    )(dx_batch, rand_qpos)
-
-    dx_batch = jax.jit(
-        jax.vmap(forward.forward, in_axes=(None, 0))
-    )(mx, dx_batch)
-
-    width, height = 32, 32
-    rc = mjx.create_render_context(
-        mjm=m,
-        nworld=batch_size,
-        cam_res=(width, height),
-        use_textures=True,
-        use_shadows=True,
-        render_rgb=True,
-        render_depth=True,
-        enabled_geom_groups=[0, 1, 2],
-    )
 
     dx_batch = jax.jit(mjx.refit_bvh)(mx, dx_batch, rc)
     out_batch = jax.jit(mjx.render)(mx, dx_batch, rc)
@@ -103,6 +104,48 @@ class RenderTest(parameterized.TestCase):
     self.assertNotEqual(np.unique(rgb).shape[0], 1)
     self.assertNotEqual(np.unique(depth).shape[0], 1)
 
+  @parameterized.product(
+      xml=('humanoid/humanoid.xml',),
+      batch_size=(4, 16),
+  )
+  def test_render_nested_vmap(self, xml: str, batch_size: int):
+    """Tests MJX render pipeline with nested vmap."""
+    self._skip_if_no_warp()
+    mx, dx_batch, rc = _get_model_data_rc(xml, batch_size)
+
+    def inner(mx, dx, rc):
+      dx = jax.vmap(bvh.refit_bvh, in_axes=(None, 0, None))(mx, dx, rc)
+      out = jax.vmap(render.render, in_axes=(None, 0, None))(mx, dx, rc)
+      return out
+
+    # get reference with single vmap
+    dx_batch = jax.vmap(bvh.refit_bvh, in_axes=(None, 0, None))(
+        mx, dx_batch, rc
+    )
+    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])
+
+    # call with nested vmap
+    def _reshape_batched(x):
+      if x.shape[0] == batch_size:
+        return x.reshape(2, batch_size // 2, *x.shape[1:])
+      return x
+
+    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.tree.map(lambda x: x.reshape(-1, *x.shape[2:]), out_batch)
+    rgb = np.asarray(out_batch[0])
+    depth = np.asarray(out_batch[1])
+
+    np.testing.assert_array_equal(rgb, ref_rgb)
+    np.testing.assert_array_equal(depth, ref_depth)
+    self.assertGreater(np.count_nonzero(rgb), 0)
+    self.assertNotEqual(np.unique(rgb).shape[0], 1)
+    self.assertGreater(np.count_nonzero(depth), 0)
+    self.assertNotEqual(np.unique(depth).shape[0], 1)
+
 
 if __name__ == '__main__':
   absltest.main()