chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

Fix bug in mjx-warp render vmap.

Browse Source 
PiperOrigin-RevId: 871901512
Change-Id: Ifc25e98ffd89284cc6db9248bcca352bb8811e7e
This commit is contained in:
Baruch Tabanpour
2026-02-18 09:09:32 -08:00
committed by Copybara-Service
parent 7d3e1abb21
commit 485e2eb028
3 changed files with 100 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/warp/ffi.py
+11 -5
View File
@@ -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(
mjx/mujoco/mjx/warp/render.py
+2 -2
View File
@@ -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
+87 -44
View File
@@ -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()