Update mjx-warp render with partial codegen.

PiperOrigin-RevId: 869874195
Change-Id: I71a16a6c192873786a09211dc4d326c3a4115ad9

mjx/mujoco/mjx/_src/io.py

@@ -1960,5 +1960,5 @@ def create_render_context(
     Render context object that is JAX compatible.
   """
   _check_warp_installed()
-  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 io as mjxw_io  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
-  return mjxw_render.create_render_context(mjm, nworld=nworld, **kwargs)
+  return mjxw_io.create_render_context(mjm, nworld=nworld, **kwargs)

mjx/mujoco/mjx/_src/render_util.py

@@ -14,30 +14,40 @@
 # ==============================================================================
 """JAX render utilities for unpacking render output from MuJoCo Warp."""
 
+import typing
+from typing import Any
+
 import jax
 import jax.numpy as jnp
 
 
 def get_rgb(
+    rc: Any,
     rgb_data: jax.Array,
     cam_id: int,
-    width: int,
-    height: int,
 ) -> jax.Array:
   """Unpack uint32 ABGR pixel data into float32 RGB.
 
   Args:
-    rgb_data: Packed render output, shape (nworld, ncam, H*W)
+    rc: The RenderContext handle.
+    rgb_data: Packed render output, shape (nworld, total_pixels)
       as uint32.
     cam_id: Camera index to extract.
-    width: Image width.
-    height: Image height.
 
   Returns:
     Float32 RGB array with shape (nworld, H, W, 3), values
     in [0, 1].
   """
-  packed = rgb_data[:, cam_id]
+  import mujoco.mjx.warp.render as mjxw_render  # pylint: disable=g-import-not-at-top
+  warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[rc.key]
+  rgb_adr = int(warp_rc.rgb_adr.numpy()[cam_id])
+  width = int(warp_rc.cam_res.numpy()[cam_id][0])
+  height = int(warp_rc.cam_res.numpy()[cam_id][1])
+
+  packed = jax.lax.dynamic_slice_in_dim(
+      rgb_data, rgb_adr, width * height, axis=1
+  )
+
   r = (packed & 0xFF).astype(jnp.float32) / 255.0
   g = ((packed >> 8) & 0xFF).astype(jnp.float32) / 255.0
   b = ((packed >> 16) & 0xFF).astype(jnp.float32) / 255.0
@@ -47,27 +57,35 @@ def get_rgb(
 
 
 def get_depth(
+    rc: Any,
     depth_data: jax.Array,
     cam_id: int,
-    width: int,
-    height: int,
     depth_scale: float,
 ) -> jax.Array:
   """Extract and normalize depth data for a camera.
 
   Args:
-    depth_data: Raw depth output, shape (nworld, ncam, H*W)
+    rc: The RenderContext handle.
+    depth_data: Raw depth output, shape (nworld, total_pixels)
       as float32.
     cam_id: Camera index to extract.
-    width: Image width.
-    height: Image height.
     depth_scale: Scale factor for normalizing depth values.
 
   Returns:
     Float32 depth array with shape (nworld, H, W), clamped
     to [0, 1].
   """
-  raw = depth_data[:, cam_id]
+  import mujoco.mjx.warp.render as mjxw_render  # pylint: disable=g-import-not-at-top
+  warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[rc.key]
+  depth_adr = int(warp_rc.depth_adr.numpy()[cam_id])
+  width = int(warp_rc.cam_res.numpy()[cam_id][0])
+  height = int(warp_rc.cam_res.numpy()[cam_id][1])
+
+  raw = jax.lax.dynamic_slice_in_dim(
+      depth_data, depth_adr, width * height, axis=1
+  )
+
   nworld = depth_data.shape[0]
   depth = jnp.clip(raw / depth_scale, 0.0, 1.0)
   return depth.reshape(nworld, height, width)
+

mjx/mujoco/mjx/warp/bvh.py

@@ -12,106 +12,128 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+
 """DO NOT EDIT. This file is auto-generated."""
 
 import dataclasses
-import functools
 import jax
-import mujoco
 from mujoco.mjx._src import types
 from mujoco.mjx.warp import ffi
-# Re-use the render registry
+from mujoco.mjx.warp.io import _MJX_RENDER_CONTEXT_BUFFERS
-from mujoco.mjx.warp.render import _MJX_RENDER_CONTEXT_BUFFERS
 from mujoco.mjx.warp.types import RenderContext
 import mujoco.mjx.third_party.mujoco_warp as mjwarp
+from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 import warp as wp
 
+
 _m = mjwarp.Model(
     **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
 )
 _d = mjwarp.Data(
     **{f.name: None for f in dataclasses.fields(mjwarp.Data) if f.init}
 )
+_o = mjwarp.Option(
+    **{f.name: None for f in dataclasses.fields(mjwarp.Option) if f.init}
+)
+_s = mjwarp.Statistic(
+    **{f.name: None for f in dataclasses.fields(mjwarp.Statistic) if f.init}
+)
+_c = mjwarp.Contact(
+    **{f.name: None for f in dataclasses.fields(mjwarp.Contact) if f.init}
+)
+_e = mjwarp.Constraint(
+    **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
+)
 
 
 @ffi.format_args_for_warp
 def _refit_bvh_shim(
     # Model
+    nworld: int,
+    flex_dim: wp.array(dtype=int),
+    flex_elem: wp.array(dtype=int),
+    flex_elemnum: wp.array(dtype=int),
+    flex_vertadr: wp.array(dtype=int),
     geom_dataid: wp.array(dtype=int),
     geom_size: wp.array2d(dtype=wp.vec3),
     geom_type: wp.array(dtype=int),
     nflex: int,
     nflexelemdata: int,
     nflexvert: int,
-    flex_dim: wp.array(dtype=int),
-    flex_elem: wp.array(dtype=int),
-    flex_elemnum: wp.array(dtype=int),
-    flex_vertadr: wp.array(dtype=int),
     # Data
+    flexvert_xpos: wp.array2d(dtype=wp.vec3),
     geom_xmat: wp.array2d(dtype=wp.mat33),
     geom_xpos: wp.array2d(dtype=wp.vec3),
-    flexvert_xpos: wp.array2d(dtype=wp.vec3),
     # Registry
     rc_id: int,
-    geom_xpos_out: wp.array2d(dtype=wp.vec3),
+    # Dummy output
+    dummy: wp.array(dtype=int),
 ):
+  _m.stat = _s
+  _m.opt = _o
+  _d.efc = _e
+  _d.contact = _c
+  _m.flex_dim = flex_dim
+  _m.flex_elem = flex_elem
+  _m.flex_elemnum = flex_elemnum
+  _m.flex_vertadr = flex_vertadr
   _m.geom_dataid = geom_dataid
   _m.geom_size = geom_size
   _m.geom_type = geom_type
   _m.nflex = nflex
   _m.nflexelemdata = nflexelemdata
   _m.nflexvert = nflexvert
-  _m.flex_dim = flex_dim
+  _d.flexvert_xpos = flexvert_xpos
-  _m.flex_elem = flex_elem
-  _m.flex_elemnum = flex_elemnum
-  _m.flex_vertadr = flex_vertadr
   _d.geom_xmat = geom_xmat
   _d.geom_xpos = geom_xpos
-  _d.flexvert_xpos = flexvert_xpos
+  _d.nworld = nworld
-  _d.nworld = geom_xpos.shape[0]
-
   render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id]
+
+  dummy.zero_()
   mjwarp.refit_bvh(_m, _d, render_context)
-  wp.copy(geom_xpos_out, geom_xpos)
 
 
 def _refit_bvh_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext):
-  nworld = d.qpos.shape[0]
+  output_dims = {'dummy': (d.qpos.shape[0],)}
-  ngeom = d.geom_xpos.shape[1]
-
   jf = ffi.jax_callable_variadic_tuple(
       _refit_bvh_shim,
       num_outputs=1,
-      output_dims={'geom_xpos_out': (nworld, ngeom, 3)},
+      output_dims=output_dims,
       vmap_method=None,
+      in_out_argnames=set([]),
+      stage_in_argnames=set(['geom_size', 'geom_xmat', 'geom_xpos']),
+      stage_out_argnames=set([]),
+      graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(
+      d.qpos.shape[0],
+      m._impl.flex_dim,
+      m._impl.flex_elem,
+      m._impl.flex_elemnum,
+      m._impl.flex_vertadr,
       m.geom_dataid,
       m.geom_size,
       m.geom_type,
-      m.nflex,
+      m._impl.nflex,
-      m.nflexelemdata,
+      m._impl.nflexelemdata,
-      m.nflexvert,
+      m._impl.nflexvert,
-      m.flex_dim,
+      d._impl.flexvert_xpos,
-      m.flex_elem,
-      m.flex_elemnum,
-      m.flex_vertadr,
       d.geom_xmat,
       d.geom_xpos,
-      d.flexvert_xpos,
       ctx.key,
   )
-  return d.replace(geom_xpos=out[0])
+  d = d.tree_replace({'time': d.time + out[0]})
+  return d
 
 
 @jax.custom_batching.custom_vmap
-@functools.partial(ffi.marshal_jax_warp_callable)
+@ffi.marshal_jax_warp_callable
 def refit_bvh(m: types.Model, d: types.Data, ctx: RenderContext):
   return _refit_bvh_jax_impl(m, d, ctx)
 
 
 @refit_bvh.def_vmap
-@functools.partial(ffi.marshal_custom_vmap)
+@ffi.marshal_custom_vmap
-def refit_bvh_vmap(unused_axis_size, is_batched, m, d, ctx):
+def refit_bvh_vmap(unused_axis_size, is_batched, m, d, ctx: RenderContext):
-  out = refit_bvh(m, d, ctx)
+  d = refit_bvh(m, d, ctx)
-  return out, is_batched[1]
+  return d, is_batched[1]

mjx/mujoco/mjx/warp/collision_driver.py

@@ -223,6 +223,7 @@ def _collision_shim(
   _d.naconmax = naconmax
   _d.ncollision = ncollision
   _d.nworld = nworld
+
   mjwarp.collision(_m, _d)
 
 
@@ -249,7 +250,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       num_outputs=15,
       output_dims=output_dims,
       vmap_method=None,
-      in_out_argnames={
+      in_out_argnames=set([
           'nacon',
           'ncollision',
           'contact__dim',
@@ -265,8 +266,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
           'contact__solreffriction',
           'contact__type',
           'contact__worldid',
-      },
+      ]),
-      stage_in_argnames={
+      stage_in_argnames=set([
           'geom_aabb',
           'geom_friction',
           'geom_gap',
@@ -285,8 +286,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
           'pair_solimp',
           'pair_solref',
           'pair_solreffriction',
-      },
+      ]),
-      stage_out_argnames={},
+      stage_out_argnames=set([]),
       graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(

mjx/mujoco/mjx/warp/forward.py

@@ -909,6 +909,7 @@ def _forward_shim(
   _d.xpos = xpos
   _d.xquat = xquat
   _d.nworld = nworld
+
   mjwarp.forward(_m, _d)
 
 
@@ -1012,7 +1013,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       num_outputs=92,
       output_dims=output_dims,
       vmap_method=None,
-      in_out_argnames={
+      in_out_argnames=set([
           'act_dot',
           'actuator_force',
           'actuator_length',
@@ -1105,8 +1106,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'efc__state',
           'efc__type',
           'efc__vel',
-      },
+      ]),
-      stage_in_argnames={
+      stage_in_argnames=set([
           'act',
           'act_dot',
           'actuator_acc0',
@@ -1242,8 +1243,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'xmat',
           'xpos',
           'xquat',
-      },
+      ]),
-      stage_out_argnames={
+      stage_out_argnames=set([
           'act_dot',
           'actuator_force',
           'actuator_length',
@@ -1276,7 +1277,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'xmat',
           'xpos',
           'xquat',
-      },
+      ]),
       graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(
@@ -2688,6 +2689,7 @@ def _step_shim(
   _d.xpos = xpos
   _d.xquat = xquat
   _d.nworld = nworld
+
   mjwarp.step(_m, _d)
 
 
@@ -2795,7 +2797,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       num_outputs=96,
       output_dims=output_dims,
       vmap_method=None,
-      in_out_argnames={
+      in_out_argnames=set([
           'act',
           'act_dot',
           'actuator_force',
@@ -2892,8 +2894,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'efc__state',
           'efc__type',
           'efc__vel',
-      },
+      ]),
-      stage_in_argnames={
+      stage_in_argnames=set([
           'act',
           'act_dot',
           'actuator_acc0',
@@ -3029,8 +3031,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'xmat',
           'xpos',
           'xquat',
-      },
+      ]),
-      stage_out_argnames={
+      stage_out_argnames=set([
           'act',
           'act_dot',
           'actuator_force',
@@ -3067,7 +3069,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'xmat',
           'xpos',
           'xquat',
-      },
+      ]),
       graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(

mjx/mujoco/mjx/warp/io.py

@@ -0,0 +1,60 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""I/O functions for MJX Warp."""
+
+import threading
+
+import mujoco
+from mujoco.mjx.warp.types import RenderContext
+import mujoco.mjx.third_party.mujoco_warp as mjw
+
+_MJX_RENDER_CONTEXT_COUNTER = 0
+_MJX_RENDER_CONTEXT_LOCK = threading.Lock()
+_MJX_RENDER_CONTEXT_BUFFERS = {}
+
+
+def create_render_context(
+    mjm: mujoco.MjModel,
+    nworld: int,
+    **kwargs,
+):
+  # NOTE: MuJoCo Warp render context expects a Warp Model and Data.
+  # We create them here but throw them away right after. Preferably,
+  # the render context should only rely on mujoco.MjModel so we
+  # do not have to pay the cost of creating dummy Warp Model and Data.
+  # Some assumptions may be violated if the downstream render context
+  # builder holds onto the memory of m and d. The API on the MuJoCo
+  # Warp side needs to be cleaned up.
+  m = mjw.put_model(mjm)
+  d = mjw.make_data(mjm, nworld=nworld)
+  mjw.forward(m, d)
+
+  rc = mjw.create_render_context(
+      mjm=mjm,
+      m=m,
+      d=d,
+      **kwargs,
+  )
+  rc.rgb_data_shape = rc.rgb_data.shape
+  rc.depth_data_shape = rc.depth_data.shape
+  rc.rgb_data = None
+  rc.depth_data = None
+
+  global _MJX_RENDER_CONTEXT_COUNTER
+  with _MJX_RENDER_CONTEXT_LOCK:
+    _MJX_RENDER_CONTEXT_COUNTER += 1
+    key = _MJX_RENDER_CONTEXT_COUNTER
+    _MJX_RENDER_CONTEXT_BUFFERS[key] = rc
+  return RenderContext(key, _owner=True)

mjx/mujoco/mjx/warp/render.py

@@ -16,19 +16,15 @@
 
 import dataclasses
 import functools
-import threading
+
 import jax
-import mujoco
 from mujoco.mjx._src import types
 from mujoco.mjx.warp import ffi
 from mujoco.mjx.warp import mujoco_warp as mjwarp
+from mujoco.mjx.warp.io import _MJX_RENDER_CONTEXT_BUFFERS
 from mujoco.mjx.warp.types import RenderContext
 import warp as wp
 
-_MJX_RENDER_CONTEXT_BUFFERS = {}
-_MJX_RENDER_CONTEXT_LOCK = threading.Lock()
-_MJX_RENDER_CONTEXT_COUNTER = 0
-
 
 _m = mjwarp.Model(
     **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -90,8 +86,8 @@ def _render_shim(
     flexvert_xpos: wp.array2d(dtype=wp.vec3),
     # Registry
     rc_id: int,
-    rgb: wp.array3d(dtype=wp.uint32),
+    rgb: wp.array2d(dtype=wp.uint32),
-    depth: wp.array3d(dtype=wp.float32),
+    depth: wp.array2d(dtype=wp.float32),
 ):
   _m.stat = _s
   _m.opt = _o
@@ -136,25 +132,17 @@ def _render_shim(
   _d.nworld = cam_xpos.shape[0]
 
   render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id]
+  render_context.rgb_data = rgb
+  render_context.depth_data = depth
   mjwarp.render(_m, _d, render_context)
 
-  wp.copy(rgb, render_context.rgb_data)
-  wp.copy(depth, render_context.depth_data)
-
 
 def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext):
   render_ctx = _MJX_RENDER_CONTEXT_BUFFERS[ctx.key]
 
-  nrender = render_ctx.nrender
-
-  nworld = d.qpos.shape[0]
-  # get width, height from cam_res
-  width = int(render_ctx.cam_res.numpy()[0][0])
-  height = int(render_ctx.cam_res.numpy()[0][1])
-
   output_dims = {
-      'rgb': (nworld, nrender, height * width),
+      'rgb': render_ctx.rgb_data_shape,
-      'depth': (nworld, nrender, height * width),
+      'depth': render_ctx.depth_data_shape,
   }
 
   jf = ffi.jax_callable_variadic_tuple(
@@ -214,50 +202,3 @@ def render(m: types.Model, d: types.Data, ctx: RenderContext):
 def render_vmap(unused_axis_size, is_batched, m, d, ctx):
   out = render(m, d, ctx)
   return out, [True, True]
-
-
-def create_render_context(
-    mjm: mujoco.MjModel,
-    nworld: int,
-    cam_res: list[tuple[int, int]] | tuple[int, int] | None = None,
-    render_rgb: list[bool] | bool | None = None,
-    render_depth: list[bool] | bool | None = None,
-    use_textures: bool = True,
-    use_shadows: bool = False,
-    enabled_geom_groups: list[int] = [0, 1, 2],
-    cam_active: list[bool] | None = None,
-    flex_render_smooth: bool = True,
-):
-  from mujoco.mjx.warp import mujoco_warp as mjw
-
-  # NOTE: MuJoCo Warp render context expects a Warp Model and Data.
-  # We create them here but throw them away right after. Preferably,
-  # the render context should only rely on mujoco.MjModel so we
-  # do not have to pay the cost of creating dummy Warp Model and Data.
-  # Some assumptions may be violated if the downstream render context
-  # builder holds onto the memory of m and d. The API on the MuJoCo
-  # Warp side needs to be cleaned up.
-  m = mjw.put_model(mjm)
-  d = mjw.make_data(mjm, nworld=nworld)
-  mjw.forward(m, d)
-
-  rc = mjw.create_render_context(
-      mjm=mjm,
-      m=m,
-      d=d,
-      cam_res=cam_res,
-      use_textures=use_textures,
-      use_shadows=use_shadows,
-      render_rgb=render_rgb,
-      render_depth=render_depth,
-      enabled_geom_groups=enabled_geom_groups,
-      cam_active=cam_active,
-      flex_render_smooth=flex_render_smooth,
-  )
-
-  global _MJX_RENDER_CONTEXT_COUNTER
-  with _MJX_RENDER_CONTEXT_LOCK:
-    _MJX_RENDER_CONTEXT_COUNTER += 1
-    key = _MJX_RENDER_CONTEXT_COUNTER
-    _MJX_RENDER_CONTEXT_BUFFERS[key] = rc
-  return RenderContext(key, _owner=True)

mjx/mujoco/mjx/warp/smooth.py

@@ -137,6 +137,7 @@ def _kinematics_shim(
   _d.xpos = xpos
   _d.xquat = xquat
   _d.nworld = nworld
+
   mjwarp.kinematics(_m, _d)
 
 
@@ -159,7 +160,7 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
       num_outputs=11,
       output_dims=output_dims,
       vmap_method=None,
-      in_out_argnames={
+      in_out_argnames=set([
           'geom_xmat',
           'geom_xpos',
           'site_xmat',
@@ -171,8 +172,8 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
           'xmat',
           'xpos',
           'xquat',
-      },
+      ]),
-      stage_in_argnames={
+      stage_in_argnames=set([
           'body_ipos',
           'body_iquat',
           'body_pos',
@@ -198,8 +199,8 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
           'xmat',
           'xpos',
           'xquat',
-      },
+      ]),
-      stage_out_argnames={
+      stage_out_argnames=set([
           'geom_xmat',
           'geom_xpos',
           'site_xmat',
@@ -211,7 +212,7 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
           'xmat',
           'xpos',
           'xquat',
-      },
+      ]),
       graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(
@@ -368,6 +369,7 @@ def _tendon_shim(
   _d.wrap_obj = wrap_obj
   _d.wrap_xpos = wrap_xpos
   _d.nworld = nworld
+
   mjwarp.tendon(_m, _d)
 
 
@@ -385,15 +387,15 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
       num_outputs=6,
       output_dims=output_dims,
       vmap_method=None,
-      in_out_argnames={
+      in_out_argnames=set([
           'ten_J',
           'ten_length',
           'ten_wrapadr',
           'ten_wrapnum',
           'wrap_obj',
           'wrap_xpos',
-      },
+      ]),
-      stage_in_argnames={
+      stage_in_argnames=set([
           'cdof',
           'geom_size',
           'geom_xmat',
@@ -402,8 +404,8 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
           'site_xpos',
           'subtree_com',
           'ten_length',
-      },
+      ]),
-      stage_out_argnames={'ten_length'},
+      stage_out_argnames=set(['ten_length']),
       graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(

mjx/mujoco/mjx/warp/testspeed.py

@@ -26,7 +26,7 @@ import jax.numpy as jnp
 import jax.tree_util
 import mujoco
 from mujoco import mjx
-from mujoco.mjx._src import io
+from mujoco.mjx.warp import io
 from mujoco.mjx._src import test_util
 from mujoco.mjx.warp import collision_driver as wp_collision
 from mujoco.mjx.warp import forward as wp_forward
@@ -161,7 +161,7 @@ def benchmark(
         d = mjx.refit_bvh(mx, d, rc)
         pixels = mjx.render(mx, d, rc)
         leaves = jax.tree_util.tree_leaves(pixels)
-        accum += sum(x[0, 0, 0] for x in leaves) if leaves else 0.0
+        accum += sum(x[0, 0] for x in leaves if x.size > 0) if leaves else 0.0
 
       return (d, accum), None
 

mjx/mujoco/mjx/warp/visualize_render.py

@@ -31,6 +31,7 @@ from mujoco.mjx._src import io
 from mujoco.mjx._src import render
 from mujoco.mjx._src import render_util
 from mujoco.mjx._src import test_util
+from mujoco.mjx.warp import io
 import numpy as np
 import warp as wp
 
@@ -171,7 +172,7 @@ def _main(_: Sequence[str]):
   print(f'  rgb shape: {rgb_packed.shape}\n')
 
   rgb = render_util.get_rgb(
-      rgb_packed, _CAMERA_ID.value, _WIDTH.value, _HEIGHT.value
+      rc, rgb_packed, _CAMERA_ID.value
   )
 
   single_path = os.path.join(