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

PiperOrigin-RevId: 960553589
Change-Id: I76caca5a82dd7f96b39e51c5b67b7382ebd1726d

doc/changelog.rst

@@ -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

@@ -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

@@ -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

@@ -36,39 +36,48 @@ def _require_segmentation_enabled(warp_rc) -> None:
     )
 
 
-def render(m: Model, d: Data, ctx: Any) -> tuple[jax.Array, jax.Array]:
+def _call_render(
-  """Render packed RGB and depth buffers."""
+    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 as mjxw_render  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
-    from mujoco.mjx.warp import render_context  # 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)
+    warp_rc = render_context.get(ctx)
-    out = mjxw_render.render(m, d, ctx)
+    if require_seg:
-    return out[0], out[1]
+      _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
+    A tuple ``(rgb, depth, seg, d)`` where the first three are packed buffers
-    stores per-pixel ``(object_id, object_type)`` pairs matching the
+    and ``d`` is the updated ``Data`` carrying the post-render execution token.
-    ``mujoco_warp`` convention.
   """
-  if m.impl == Impl.WARP and d.impl == Impl.WARP and mjxw.WARP_INSTALLED:
+  rgb, depth, seg, d = _call_render(m, d, ctx, require_seg=True)
-    from mujoco.mjx.warp import render as mjxw_render  # pytype: disable=import-error
+  return rgb, depth, seg, d
-    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.'
-  )

mjx/mujoco/mjx/_src/render_test.py

@@ -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

@@ -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
+  # create an output token if there are no output fields
-  needs_dummy_output = not field_usage.data_out_fields
+  needs_output_token = not field_usage.data_out_fields
-  if needs_dummy_output and fn_name != 'render':
+  if needs_output_token and fn_name != 'render':
-    fn_args_raw.append('# Dummy output')  # pyrefly: ignore[bad-argument-type]
+    fn_args_raw.append('# Output token')  # pyrefly: ignore[bad-argument-type]
-    fn_args_raw.append('dummy: wp.array[int],')  # 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.
+  # If there are no Warp array outputs, we need an output token for JAX FFI.
-  needs_dummy_output = not any(
+  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,)"]
+      tree_replace = ['"_impl._jax_token": out[0]']
-    else:
-      output_dims = ["'dummy': (d.qpos.shape[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

@@ -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

@@ -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
+    # Output token
-    dummy: wp.array[int],
+    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

@@ -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

@@ -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 +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])
+    rgb = np.asarray(rgb_arr)
-    depth = np.asarray(out_batch[1])
+    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])
+    rgb = np.asarray(rgb_arr)
-    depth = np.asarray(out_batch[1])
+    depth = np.asarray(depth_arr)
-    seg = np.asarray(out_batch[2])
+    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

@@ -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

@@ -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,