diff --git a/mjx/mujoco/mjx/warp/ffi.py b/mjx/mujoco/mjx/warp/ffi.py index 18632162..6d0b6d7e 100644 --- a/mjx/mujoco/mjx/warp/ffi.py +++ b/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: - return d_broadcast_flat_result, out_batched + if tree_map_output: + 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( diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py index dbc783a2..091f72d0 100644 --- a/mjx/mujoco/mjx/warp/render.py +++ b/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, diff --git a/mjx/mujoco/mjx/warp/render_test.py b/mjx/mujoco/mjx/warp/render_test.py index db49fa32..7157f0cc 100644 --- a/mjx/mujoco/mjx/warp/render_test.py +++ b/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( - xml=( - 'humanoid/humanoid.xml', - ), - batch_size=(1, 16), - ) - def test_render(self, xml: str, batch_size: int): - """Tests MJX render pipeline.""" + def _skip_if_no_warp(self): 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) - 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], - ) + @parameterized.product( + xml=('humanoid/humanoid.xml',), + batch_size=(1, 16), + ) + def test_render(self, xml: str, batch_size: int): + """Tests MJX render pipeline.""" + self._skip_if_no_warp() + 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) @@ -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()