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

Fix #3435. Add token to ensure sequential calls for mjx-warp refit and render.

Browse Source 
PiperOrigin-RevId: 960553589
Change-Id: I76caca5a82dd7f96b39e51c5b67b7382ebd1726d
This commit is contained in:
Baruch Tabanpour
2026-08-06 16:05:05 -07:00
committed by Copybara-Service
parent a1d772c9ad
commit 5e3464f475
13 changed files with 245 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+13
View File
@@ -90,6 +90,19 @@ Rendering
**Migration:** Set :at:`softness` to 1 to reproduce the previous appearance of existing models.
MJX
^^^
.. admonition:: Breaking API changes
:class: attention
- :func:`mjx.render` and :func:`mjx.render_with_segmentation` now return the updated :class:`mjx.Data` as the last
element in their return tuple (i.e. ``(rgb, depth, d)`` and ``(rgb, depth, seg, d)``). This ensures JAX/XLA
strictly enforces causal scheduling between sequential ``refit_bvh`` and ``render`` calls.
**Migration:** Update unpacking calls from ``pixels, depth = mjx.render(mx, d, rc)`` to
``pixels, depth, d = mjx.render(mx, d, rc)``.
Bug fixes
^^^^^^^^^
doc/mjx.rst
+7 -1
View File
@@ -267,7 +267,7 @@ volume hierarchy (BVH) and executing the raycaster:
d = mjx.refit_bvh(mx, d, rc_pytree)
# 2. Render all configured cameras
pixels, _ = mjx.render(mx, d, rc_pytree)
pixels, _, d = mjx.render(mx, d, rc_pytree)
# 3. Extract the RGB tensor for the first camera (index 0)
rgb = get_rgb(rc_pytree, 0, pixels)
@@ -276,6 +276,12 @@ volume hierarchy (BVH) and executing the raycaster:
rgb, d = render_fn(mx, d, rc.pytree())
.. NOTE::
:func:`~mujoco.mjx.refit_bvh` and :func:`~mujoco.mjx.render` update an internal execution token
(``d._impl._jax_token``) within :class:`~mujoco.mjx.Data`. Passing ``d`` sequentially through
``refit_bvh`` and ``render`` creates an explicit data dependency, preventing XLA from reordering BVH
updates and raycasting passes across iterations or unrolled loops.
.. WARNING::
The batch dimension ``nworld`` is fixed when the render context is created via
:func:`~mujoco.mjx.create_render_context` since the underlying Warp render context allocates
mjx/mujoco/mjx/_src/io.py
+4
View File
@@ -786,6 +786,8 @@ def _make_data_warp(
impl_fields = {}
for k in mjxw.types.DataWarp.__annotations__.keys():
if k == '_jax_token': # custom token to force sequential calls in JAX
continue
field = _get_nested_attr(dw, k, split='__')
field = _wp_to_np_type(field)
if mjxw.types._BATCH_DIM['Data'][k]: # pylint: disable=protected-access
@@ -1195,6 +1197,8 @@ def _put_data_warp(
impl_fields = {}
for k in mjxw.types.DataWarp.__annotations__.keys():
if k == '_jax_token': # custom token to force sequential calls in JAX
continue
field = _get_nested_attr(dw, k, split='__')
field = _wp_to_np_type(field)
if mjxw.types._BATCH_DIM['Data'][k]: # pylint: disable=protected-access
mjx/mujoco/mjx/_src/render.py
+33 -24
View File
@@ -36,39 +36,48 @@ def _require_segmentation_enabled(warp_rc) -> None:
)
def render(m: Model, d: Data, ctx: Any) -> tuple[jax.Array, jax.Array]:
"""Render packed RGB and depth buffers."""
def _call_render(
m: Model,
d: Data,
ctx: Any,
require_seg: bool = False,
) -> tuple[jax.Array, jax.Array, jax.Array, Data]:
if m.impl == Impl.WARP and d.impl == Impl.WARP and mjxw.WARP_INSTALLED:
from mujoco.mjx.warp import render as mjxw_render # pytype: disable=import-error
from mujoco.mjx.warp import render_context # pytype: disable=import-error
from mujoco.mjx.warp import render as mjxw_render # pylint: disable=g-import-not-at-top # pytype: disable=import-error
from mujoco.mjx.warp import render_context # pylint: disable=g-import-not-at-top # pytype: disable=import-error
render_context.get(ctx)
out = mjxw_render.render(m, d, ctx)
return out[0], out[1]
warp_rc = render_context.get(ctx)
if require_seg:
_require_segmentation_enabled(warp_rc)
rgb, depth, seg, token_array = mjxw_render.render(m, d, ctx)
token = token_array.reshape(
d._impl._jax_token.shape # pytype: disable=attribute-error
)
d = d.tree_replace({'_impl._jax_token': token})
return rgb, depth, seg, d
raise NotImplementedError('render only implemented for MuJoCo Warp.')
def render(m: Model, d: Data, ctx: Any) -> tuple[jax.Array, jax.Array, Data]:
"""Render packed RGB and depth buffers.
Returns:
A tuple ``(rgb, depth, d)`` where ``rgb`` and ``depth`` are packed buffers
and ``d`` is the updated ``Data`` carrying the post-render execution token.
"""
rgb, depth, unused_seg, d = _call_render(m, d, ctx)
return rgb, depth, d
def render_with_segmentation(
m: Model, d: Data, ctx: Any
) -> tuple[jax.Array, jax.Array, jax.Array]:
) -> tuple[jax.Array, jax.Array, jax.Array, Data]:
"""Render and return RGB, depth, and packed segmentation outputs.
Returns:
A tuple ``(rgb, depth, seg)`` of packed buffers. The segmentation buffer
stores per-pixel ``(object_id, object_type)`` pairs matching the
``mujoco_warp`` convention.
A tuple ``(rgb, depth, seg, d)`` where the first three are packed buffers
and ``d`` is the updated ``Data`` carrying the post-render execution token.
"""
if m.impl == Impl.WARP and d.impl == Impl.WARP and mjxw.WARP_INSTALLED:
from mujoco.mjx.warp import render as mjxw_render # pytype: disable=import-error
from mujoco.mjx.warp import render_context # pytype: disable=import-error
warp_rc = render_context.get(ctx)
_require_segmentation_enabled(warp_rc)
out = mjxw_render.render(m, d, ctx)
return out[0], out[1], out[2]
raise NotImplementedError(
'render_with_segmentation only implemented for MuJoCo Warp.'
)
rgb, depth, seg, d = _call_render(m, d, ctx, require_seg=True)
return rgb, depth, seg, d
mjx/mujoco/mjx/_src/render_test.py
+2 -2
View File
@@ -83,7 +83,7 @@ class RenderIntegrationTest(parameterized.TestCase):
self._maybe_skip()
mx, dx_batch, rc = _setup(batch_size)
rgb_packed, depth_packed, seg_packed = jax.jit(
rgb_packed, depth_packed, seg_packed, dx_batch = jax.jit(
mjx.render_with_segmentation
)(mx, dx_batch, rc.pytree())
@@ -105,7 +105,7 @@ class RenderIntegrationTest(parameterized.TestCase):
self._maybe_skip()
mx, dx_batch, rc = _setup(batch_size)
rgb_packed, depth_packed, seg_packed = jax.jit(
rgb_packed, depth_packed, seg_packed, dx_batch = jax.jit(
mjx.render_with_segmentation
)(mx, dx_batch, rc.pytree())
mjx/mujoco/mjx/codegen/generate_warp_shim.py
+25 -15
View File
@@ -249,20 +249,26 @@ def _warp_function(
render_context_args.append('rgb: wp.array2d[wp.uint32],')
render_context_args.append('depth: wp.array2d[wp.float32],')
render_context_args.append('seg: wp.array2d[wp.vec2i],')
render_context_args.append('output_token: wp.array[int],')
fn_assignments.append(' render_context.rgb_data = rgb')
fn_assignments.append(' render_context.depth_data = depth')
fn_assignments.append(' render_context.seg_data = seg')
fn_assignments.append(' output_token.zero_()')
else:
fn_assignments.append(' dummy.zero_()')
fn_assignments.append(' output_token.zero_()')
token_args = []
if field_usage.render_context_in_caller:
token_args.append('_jax_token: wp.array[int],') # pyrefly: ignore[bad-argument-type]
fn_call = f'mjwarp.{fn_name}(_m, _d{render_context_call_arg})'
fn_args_raw = fn_args_model + fn_args_data + render_context_args
fn_args_raw = fn_args_model + fn_args_data + token_args + render_context_args
# create a dummy output if there are no output fields
needs_dummy_output = not field_usage.data_out_fields
if needs_dummy_output and fn_name != 'render':
fn_args_raw.append('# Dummy output') # pyrefly: ignore[bad-argument-type]
fn_args_raw.append('dummy: wp.array[int],') # pyrefly: ignore[bad-argument-type]
# create an output token if there are no output fields
needs_output_token = not field_usage.data_out_fields
if needs_output_token and fn_name != 'render':
fn_args_raw.append('# Output token') # pyrefly: ignore[bad-argument-type]
fn_args_raw.append('output_token: wp.array[int],') # pyrefly: ignore[bad-argument-type]
return fn_args_raw, fn_assignments, fn_call
@@ -292,7 +298,7 @@ def _jax_shim_fn(
jax_args.append('d.qpos.shape[0]')
continue
if arg in ('rc_id', 'dummy'):
if arg in ('rc_id', 'output_token', '_jax_token'):
continue
if arg in ('rgb', 'depth', 'seg') and fn_name == 'render':
@@ -344,28 +350,31 @@ def _jax_shim_fn(
jax_args.append(arg_jax)
if field_usage.render_context_in_caller:
jax_args.append('d._impl._jax_token')
jax_args.append('ctx.key')
# If there are no Warp array outputs, we need a dummy output for JAX FFI.
needs_dummy_output = not any(
# If there are no Warp array outputs, we need an output token for JAX FFI.
needs_output_token = not any(
'array' in mjwarp_field_info[f].expected_type
for f in field_usage.data_out_fields
)
if needs_dummy_output and fn_name != 'render':
if needs_output_token and fn_name != 'render':
num_outputs = 1
output_dims = ["'output_token': (d.qpos.shape[0],)"]
if field_usage.render_context_in_caller:
output_dims = ["'dummy': (render_ctx.nworld,)"]
else:
output_dims = ["'dummy': (d.qpos.shape[0],)"]
tree_replace = ['"_impl._jax_token": out[0]']
has_side_effect = True
if fn_name == 'render':
num_outputs = 4
output_dims = [
"'rgb': render_ctx.rgb_data_shape",
"'depth': render_ctx.depth_data_shape",
"'seg': render_ctx.seg_data_shape",
"'output_token': (d.qpos.shape[0],)",
]
tree_replace = []
has_side_effect = True
render_ctx_param = (
'ctx: RenderContextPytree' if field_usage.render_context_in_caller else ''
@@ -495,7 +504,8 @@ def _{fn_name}_jax_impl({','.join(fn_args)}):
'@functools.partial(ffi.marshal_custom_vmap, tree_map_output=True)'
)
vmap_return_stmt = (
f'out = {fn_name}({fn_call_str})\n return out, [True, True, True]'
f'out = {fn_name}({fn_call_str})\n return out, [True, True, True,'
' is_batched[1]._impl._jax_token]'
)
src += f"""
mjx/mujoco/mjx/codegen/generate_warp_types.py
+33 -4
View File
@@ -204,7 +204,11 @@ def _build_new_class_body_ast(
new_body_nodes.append(ast.Expr(value=ast.Constant(value=docstring)))
if sort_keys:
maybe_sorted_keys = sorted(list(keys))
non_default_keys = sorted(
[k for k in keys if not defaults or k not in defaults]
)
default_keys = sorted([k for k in keys if defaults and k in defaults])
maybe_sorted_keys = non_default_keys + default_keys
else:
maybe_sorted_keys = list(keys)
@@ -213,9 +217,14 @@ def _build_new_class_body_ast(
value_node = None
if defaults and key in defaults:
value_node = ast.Constant(value=defaults[key])
val = defaults[key]
value_node = val if isinstance(val, ast.AST) else ast.Constant(value=val)
if value_node and value_node.value is None:
if (
value_node
and isinstance(value_node, ast.Constant)
and value_node.value is None
):
annotation_node = _ast_parse_type(
f'typing.Optional[{ast.unparse(annotation_node)}]'
)
@@ -402,6 +411,7 @@ def write_core_cls(
flatten_fields: bool = False,
set_diff: bool = True,
extra_annotations: dict[str, type] | None = None,
defaults: Optional[Dict[str, Any]] = None,
):
"""Writes a core API class (e.g. Model/Data/Option/Statistic)."""
cls = _CLS_MAP[cls_name]
@@ -442,6 +452,7 @@ def write_core_cls(
keys, # pyrefly: ignore[bad-argument-type]
cls_name,
annotations,
defaults=defaults,
shape_property=shape_property,
)
@@ -513,6 +524,9 @@ def write_ndim_annotations(target_fpath: epath.Path):
ndim = _to_jax_ndim(name, type_)
ndim_annotations[cls_name][name] = ndim
# custom token to force sequential calls in JAX
ndim_annotations['Data']['_jax_token'] = 1
with target_fpath.open('a') as f:
f.write('\n_NDIM = ' + _to_py_string(ndim_annotations))
@@ -529,6 +543,9 @@ def write_nworld_leading_dim(target_fpath: epath.Path):
# Same batch check as mujoco_warp._src.io._mark_batched.
batched[cls_name][name] = _is_batched_field(type_) is True
# custom token to force sequential calls in JAX
batched['Data']['_jax_token'] = True
with target_fpath.open('a') as f:
f.write('\n_BATCH_DIM = ' + _to_py_string(batched))
@@ -551,7 +568,19 @@ def main(argv):
extra_annotations={'graph_mode': GraphMode},
)
write_core_cls('Model', target_fpath, mjx_types_fpath)
write_core_cls('Data', target_fpath, mjx_types_fpath, flatten_fields=True)
token_default = ast.parse(
'dataclasses.field(default_factory=lambda: jax.numpy.zeros((),'
' dtype=jax.numpy.int32))',
mode='eval',
).body
write_core_cls(
'Data',
target_fpath,
mjx_types_fpath,
flatten_fields=True,
extra_annotations={'_jax_token': wp.array(dtype=int)},
defaults={'_jax_token': token_default},
)
write_register_vmappable(target_fpath)
write_ndim_annotations(target_fpath)
write_nworld_leading_dim(target_fpath)
mjx/mujoco/mjx/warp/bvh.py
+8 -5
View File
@@ -48,6 +48,7 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
)
@ffi.format_args_for_warp
def _refit_bvh_shim(
# Model
@@ -72,10 +73,11 @@ def _refit_bvh_shim(
flexvert_xpos: wp.array2d[wp.vec3],
geom_xmat: wp.array2d[wp.mat33],
geom_xpos: wp.array2d[wp.vec3],
_jax_token: wp.array[int],
# Registry
rc_id: int,
# Dummy output
dummy: wp.array[int],
# Output token
output_token: wp.array[int],
):
_m.stat = _s
_m.opt = _o
@@ -103,7 +105,7 @@ def _refit_bvh_shim(
_d.geom_xpos = geom_xpos
_d.nworld = nworld
render_context = _MJX_RENDER_CONTEXT_BUFFERS[(rc_id, wp.get_device().ordinal)]
dummy.zero_()
output_token.zero_()
mjwarp.refit_bvh(_m, _d, render_context)
@@ -111,7 +113,7 @@ def _refit_bvh_jax_impl(
m: types.Model, d: types.Data, ctx: RenderContextPytree
):
render_ctx = _MJX_RENDER_CONTEXT_BUFFERS[(ctx.key, None)]
output_dims = {'dummy': (render_ctx.nworld,)}
output_dims = {'output_token': (d.qpos.shape[0],)}
jf = ffi.jax_callable_variadic_tuple(
_refit_bvh_shim,
num_outputs=1,
@@ -146,9 +148,10 @@ def _refit_bvh_jax_impl(
d._impl.flexvert_xpos,
d.geom_xmat,
d.geom_xpos,
d._impl._jax_token,
ctx.key,
)
d = d.tree_replace({})
d = d.tree_replace({'_impl._jax_token': out[0]})
return d
mjx/mujoco/mjx/warp/forward.py
-1
View File
@@ -46,7 +46,6 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
)
@ffi.format_args_for_warp
def _forward_shim(
# Model
mjx/mujoco/mjx/warp/render.py
+9 -3
View File
@@ -47,6 +47,7 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
)
@ffi.format_args_for_warp
def _render_shim(
# Model
@@ -89,11 +90,13 @@ def _render_shim(
geom_xpos: wp.array2d[wp.vec3],
light_xdir: wp.array2d[wp.vec3],
light_xpos: wp.array2d[wp.vec3],
_jax_token: wp.array[int],
# Registry
rc_id: int,
rgb: wp.array2d[wp.uint32],
depth: wp.array2d[wp.float32],
seg: wp.array2d[wp.vec2i],
output_token: wp.array[int],
):
_m.stat = _s
_m.opt = _o
@@ -142,6 +145,7 @@ def _render_shim(
render_context.rgb_data = rgb
render_context.depth_data = depth
render_context.seg_data = seg
output_token.zero_()
mjwarp.render(_m, _d, render_context)
@@ -151,10 +155,11 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
'rgb': render_ctx.rgb_data_shape,
'depth': render_ctx.depth_data_shape,
'seg': render_ctx.seg_data_shape,
'output_token': (d.qpos.shape[0],),
}
jf = ffi.jax_callable_variadic_tuple(
_render_shim,
num_outputs=3,
num_outputs=4,
output_dims=output_dims,
vmap_method=None,
in_out_argnames=set([]),
@@ -189,7 +194,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
]),
stage_out_argnames=set([]),
graph_mode=m.opt._impl.graph_mode,
has_side_effect=False,
has_side_effect=True,
)
out = jf(
render_ctx.nworld,
@@ -230,6 +235,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
d.geom_xpos,
d._impl.light_xdir,
d._impl.light_xpos,
d._impl._jax_token,
ctx.key,
)
d = d.tree_replace({})
@@ -252,4 +258,4 @@ def render_vmap(
ctx: RenderContextPytree,
):
out = render(m, d, ctx)
return out, [True, True, True]
return out, [True, True, True, is_batched[1]._impl._jax_token]
mjx/mujoco/mjx/warp/render_test.py
+105 -10
View File
@@ -105,10 +105,12 @@ class RenderTest(parameterized.TestCase):
mx, dx_batch, rc = _get_model_data_rc(xml, batch_size)
dx_batch = jax.jit(mjx.refit_bvh)(mx, dx_batch, rc.pytree())
out_batch = jax.jit(mjx.render)(mx, dx_batch, rc.pytree())
rgb_arr, depth_arr, dx_batch = jax.jit(mjx.render)(
mx, dx_batch, rc.pytree()
)
rgb = np.asarray(out_batch[0])
depth = np.asarray(out_batch[1])
rgb = np.asarray(rgb_arr)
depth = np.asarray(depth_arr)
self.assertGreater(np.count_nonzero(rgb), 0)
self.assertGreater(np.count_nonzero(depth), 0)
@@ -148,7 +150,9 @@ class RenderTest(parameterized.TestCase):
dx_2d = jax.tree.map(_reshape_batched, dx_batch)
out_batch = jax.vmap(inner, in_axes=(None, 0, None))(mx, dx_2d, rc.pytree())
out_batch = jax.tree.map(lambda x: x.reshape(-1, *x.shape[2:]), out_batch)
out_batch = jax.tree.map(
lambda x: x.reshape(-1, *x.shape[2:]) if x.size > 0 else x, out_batch
)
rgb = np.asarray(out_batch[0])
depth = np.asarray(out_batch[1])
@@ -169,11 +173,13 @@ class RenderTest(parameterized.TestCase):
mx, dx_batch, rc = _get_model_data_rc(xml, batch_size, render_seg=True)
dx_batch = jax.jit(mjx.refit_bvh)(mx, dx_batch, rc.pytree())
out_batch = jax.jit(mjx.render_with_segmentation)(mx, dx_batch, rc.pytree())
rgb_arr, depth_arr, seg_arr, dx_batch = jax.jit(
mjx.render_with_segmentation
)(mx, dx_batch, rc.pytree())
rgb = np.asarray(out_batch[0])
depth = np.asarray(out_batch[1])
seg = np.asarray(out_batch[2])
rgb = np.asarray(rgb_arr)
depth = np.asarray(depth_arr)
seg = np.asarray(seg_arr)
self.assertGreater(np.count_nonzero(rgb), 0)
self.assertGreater(np.count_nonzero(depth), 0)
@@ -181,7 +187,7 @@ class RenderTest(parameterized.TestCase):
self.assertGreater(np.unique(seg[..., 0]).shape[0], 1)
unpacked_seg = jax.vmap(mjx.get_segmentation, in_axes=(None, None, 0))(
rc.pytree(), 0, out_batch[2]
rc.pytree(), 0, seg_arr
)
unpacked_seg = np.asarray(unpacked_seg)
width, height = rc._default.cam_res.numpy()[
@@ -248,7 +254,9 @@ class RenderTest(parameterized.TestCase):
dx_2d = jax.tree.map(_reshape_batched, dx_batch)
out_batch = jax.vmap(inner, in_axes=(None, 0, None))(mx, dx_2d, rc.pytree())
out_batch = jax.tree.map(lambda x: x.reshape(-1, *x.shape[2:]), out_batch)
out_batch = jax.tree.map(
lambda x: x.reshape(-1, *x.shape[2:]) if x.size > 0 else x, out_batch
)
rgb = np.asarray(out_batch[0])
depth = np.asarray(out_batch[1])
seg = np.asarray(out_batch[2])
@@ -258,6 +266,93 @@ class RenderTest(parameterized.TestCase):
np.testing.assert_array_equal(seg, ref_seg)
self.assertTrue(np.any(seg[..., 0] != -1))
def test_sliding_box_poses(self):
"""Tests that refit_bvh executes before render for multiple poses in a single JIT."""
self._maybe_skip()
xml = """\
<mujoco>
<option gravity="0 0 0"/>
<asset>
<material name="mat" rgba="1 1 1 1"/>
</asset>
<worldbody>
<camera pos="0 0 1" resolution="64 64"
sensorsize="0.036 0.036" focal="0.012 0.012"/>
<geom type="plane" size="2 2 0.1" material="mat"/>
<body pos="0 0 0.1" mocap="true">
<geom type="box" size="0.1 0.1 0.1" material="mat"/>
</body>
<body pos="0 0 5">
<joint type="free"/>
<geom type="sphere" size="0.01" group="3"
contype="0" conaffinity="0"/>
</body>
</worldbody>
</mujoco>
"""
m = mujoco.MjModel.from_xml_string(xml)
mx = mjx.put_model(m, impl='warp')
dx = mjx.put_data(m, mujoco.MjData(m), impl='warp')
rc = mjx.create_render_context(
mjm=m, nworld=1, cam_res=(64, 64), render_rgb=True, render_depth=True
)
x_positions = (-0.8, -0.4, 0.0, 0.4, 0.8)
def step_and_render(d, x_pos):
mocap_pos = d.mocap_pos.at[0, 0].set(x_pos)
d = forward.forward(mx, d.replace(mocap_pos=mocap_pos))
d = mjx.refit_bvh(mx, d, rc.pytree())
_, depth, d = mjx.render(mx, d, rc.pytree())
return d, depth
# 1. Step-by-step execution.
ref_depths = []
d_curr = dx
for x in x_positions:
d_curr, depth = jax.jit(step_and_render)(d_curr, x)
ref_depths.append(np.asarray(depth))
# 2. Multi-step rollout in a single compiled jax.jit via jax.lax.scan.
def rollout_scan(d0):
def scan_fn(d, x):
d, depth = step_and_render(d, x)
return d, depth
_, scan_depths = jax.lax.scan(
scan_fn, d0, np.array(x_positions, dtype=np.float32)
)
return scan_depths
scanned_depths = [np.asarray(d) for d in jax.jit(rollout_scan)(dx)]
# 3. Explicitly unrolled sequence in a single compiled jax.jit.
def rollout_unroll(d0):
depths = []
d = d0
for x in x_positions:
mocap_pos = d.mocap_pos.at[0, 0].set(x)
d = forward.forward(mx, d.replace(mocap_pos=mocap_pos))
d = mjx.refit_bvh(mx, d, rc.pytree())
_, depth, d = mjx.render(mx, d, rc.pytree())
depths.append(depth)
return depths
unrolled_depths = [np.asarray(d) for d in jax.jit(rollout_unroll)(dx)]
# Assert exact equality across all frames for both scan and unroll.
for i, (want, scan_got, unroll_got) in enumerate(
zip(ref_depths, scanned_depths, unrolled_depths)
):
np.testing.assert_array_equal(
scan_got, want, err_msg=f'Scan mismatch at frame {i}'
)
np.testing.assert_array_equal(
unroll_got, want, err_msg=f'Unroll mismatch at frame {i}'
)
del rc
class RenderContextGarbageCollectionTest(absltest.TestCase):
"""Tests that RenderContext cleans up buffers on deletion."""
mjx/mujoco/mjx/warp/testspeed.py
+1 -1
View File
@@ -154,7 +154,7 @@ def benchmark(
def render_fn(mx, d, rc):
d = mjx.refit_bvh(mx, d, rc)
pixels, _ = mjx.render(mx, d, rc)
pixels, _, d = mjx.render(mx, d, rc)
return render_util.get_rgb(rc, 0, pixels), d
@jax_jit
mjx/mujoco/mjx/warp/types.py
+5
View File
@@ -526,6 +526,9 @@ class DataWarp(PyTreeNode):
tree_island: jax.Array
wrap_obj: jax.Array
wrap_xpos: jax.Array
_jax_token: jax.Array = dataclasses.field(
default_factory=lambda: jax.numpy.zeros((), dtype=jax.numpy.int32)
)
shape = property(lambda self: self.cacc.shape)
DATA_NON_VMAP = {
'cdof_tri_col',
@@ -588,6 +591,7 @@ batching.register_vmappable(DataWarp, int, int, _to_elt, _from_elt, None)
_NDIM = {
'Data': {
'M': 2,
'_jax_token': 1,
'act': 2,
'act_dot': 2,
'actuator_force': 2,
@@ -1312,6 +1316,7 @@ _NDIM = {
_BATCH_DIM = {
'Data': {
'M': True,
'_jax_token': True,
'act': True,
'act_dot': True,
'actuator_force': True,