Last bits of mjx-warp render codegen.

PiperOrigin-RevId: 869900622
Change-Id: If86942fddc45b9d71505d7fc4e2345dc120ccab2

mjx/mujoco/mjx/__init__.py

@@ -46,6 +46,8 @@ from mujoco.mjx._src.io import state_size
 from mujoco.mjx._src.passive import passive
 from mujoco.mjx._src.ray import ray
 from mujoco.mjx._src.render import render
+from mujoco.mjx._src.render_util import get_depth
+from mujoco.mjx._src.render_util import get_rgb
 from mujoco.mjx._src.sensor import sensor_acc
 from mujoco.mjx._src.sensor import sensor_pos
 from mujoco.mjx._src.sensor import sensor_vel

mjx/mujoco/mjx/_src/render_util.py

@@ -14,11 +14,11 @@
 # ==============================================================================
 """JAX render utilities for unpacking render output from MuJoCo Warp."""
 
-import typing
 from typing import Any
 
 import jax
 import jax.numpy as jnp
+import mujoco.mjx.warp as mjxw
 
 
 def get_rgb(
@@ -30,30 +30,34 @@ def get_rgb(
 
   Args:
     rc: The RenderContext handle.
-    rgb_data: Packed render output, shape (nworld, total_pixels)
-      as uint32.
+    rgb_data: Packed render output, shape (total_pixels,) as uint32.
     cam_id: Camera index to extract.
 
   Returns:
-    Float32 RGB array with shape (nworld, H, W, 3), values
-    in [0, 1].
+    Float32 RGB array with shape (H, W, 3), values in [0, 1].
+
+  Raises:
+    RuntimeError: If Warp is not installed.
   """
-  import mujoco.mjx.warp.render as mjxw_render  # pylint: disable=g-import-not-at-top
+  if mjxw.WARP_INSTALLED:
+    import mujoco.mjx.warp.render as mjxw_render  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+  else:
+    raise RuntimeError('Warp not installed.')
+
   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
+      rgb_data, rgb_adr, width * height, axis=0
   )
 
-  r = (packed & 0xFF).astype(jnp.float32) / 255.0
+  b = (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
+  r = ((packed >> 16) & 0xFF).astype(jnp.float32) / 255.0
   rgb = jnp.stack([r, g, b], axis=-1)
-  nworld = rgb_data.shape[0]
-  return rgb.reshape(nworld, height, width, 3)
+  return rgb.reshape(height, width, 3)
 
 
 def get_depth(
@@ -66,26 +70,28 @@ def get_depth(
 
   Args:
     rc: The RenderContext handle.
-    depth_data: Raw depth output, shape (nworld, total_pixels)
-      as float32.
+    depth_data: Raw depth output, shape (total_pixels,) as float32.
     cam_id: Camera index to extract.
     depth_scale: Scale factor for normalizing depth values.
 
   Returns:
-    Float32 depth array with shape (nworld, H, W), clamped
-    to [0, 1].
+    Float32 depth array with shape (H, W), clamped to [0, 1].
+
+  Raises:
+    RuntimeError: If Warp is not installed.
   """
-  import mujoco.mjx.warp.render as mjxw_render  # pylint: disable=g-import-not-at-top
+  if mjxw.WARP_INSTALLED:
+    import mujoco.mjx.warp.render as mjxw_render  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+  else:
+    raise RuntimeError('Warp not installed.')
   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
+      depth_data, depth_adr, width * height, axis=0
   )
 
-  nworld = depth_data.shape[0]
   depth = jnp.clip(raw / depth_scale, 0.0, 1.0)
-  return depth.reshape(nworld, height, width)
-
+  return depth.reshape(height, width)

mjx/mujoco/mjx/_src/render_util_test.py

@@ -0,0 +1,123 @@
+# 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.
+# ==============================================================================
+import os
+from unittest import mock
+
+from absl.testing import absltest
+import jax
+import jax.numpy as jnp
+from mujoco.mjx._src import io
+from mujoco.mjx._src import render_util
+import mujoco.mjx.warp as mjxw
+import numpy as np
+
+_FORCE_TEST = os.environ.get('MJX_WARP_FORCE_TEST', '0') == '1'
+
+
+def _fake_render_context(ncam, width, height):
+  """Fake RenderContext for testing."""
+  rc = mock.MagicMock()
+  rgb_adr = np.arange(ncam, dtype=np.int32) * width * height
+  depth_adr = np.arange(ncam, dtype=np.int32) * width * height
+  cam_res = np.tile([width, height], (ncam, 1)).astype(np.int32)
+  rc.rgb_adr.numpy.return_value = rgb_adr
+  rc.depth_adr.numpy.return_value = depth_adr
+  rc.cam_res.numpy.return_value = cam_res
+  return rc
+
+
+class RenderUtilTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    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.')
+
+  def test_get_rgb(self):
+    width, height = 4, 4
+    warp_rc = _fake_render_context(1, width, height)
+    rc = mock.MagicMock(key=0)
+    rgb_data = jnp.zeros((width * height,), dtype=jnp.uint32)
+
+    with mock.patch.dict(
+        'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS',
+        {0: warp_rc},
+    ):
+      rgb = jax.jit(render_util.get_rgb, static_argnums=(0, 2))(
+          rc, rgb_data, 0
+      )
+
+    self.assertEqual(rgb.shape, (height, width, 3))
+
+  def test_get_rgb_vmap(self):
+    nworld, width, height = 3, 4, 4
+    warp_rc = _fake_render_context(1, width, height)
+    rc = mock.MagicMock(key=0)
+    rgb_data = jnp.zeros(
+        (nworld, width * height), dtype=jnp.uint32
+    )
+
+    with mock.patch.dict(
+        'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS',
+        {0: warp_rc},
+    ):
+      rgb = jax.jit(
+          jax.vmap(render_util.get_rgb, in_axes=(None, 0, None)),
+          static_argnums=(0, 2),
+      )(rc, rgb_data, 0)
+
+    self.assertEqual(rgb.shape, (nworld, height, width, 3))
+
+  def test_get_depth(self):
+    width, height = 4, 4
+    warp_rc = _fake_render_context(1, width, height)
+    rc = mock.MagicMock(key=0)
+    depth_data = jnp.zeros((width * height,), dtype=jnp.float32)
+
+    with mock.patch.dict(
+        'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS',
+        {0: warp_rc},
+    ):
+      depth = jax.jit(
+          render_util.get_depth, static_argnums=(0, 2, 3)
+      )(rc, depth_data, 0, 5.0)
+
+    self.assertEqual(depth.shape, (height, width))
+
+  def test_get_depth_vmap(self):
+    nworld, width, height = 3, 4, 4
+    warp_rc = _fake_render_context(1, width, height)
+    rc = mock.MagicMock(key=0)
+    depth_data = jnp.zeros(
+        (nworld, width * height), dtype=jnp.float32
+    )
+
+    with mock.patch.dict(
+        'mujoco.mjx.warp.render._MJX_RENDER_CONTEXT_BUFFERS',
+        {0: warp_rc},
+    ):
+      depth = jax.jit(
+          jax.vmap(render_util.get_depth, in_axes=(None, 0, None, None)),
+          static_argnums=(0, 2, 3),
+      )(rc, depth_data, 0, 5.0)
+
+    self.assertEqual(depth.shape, (nworld, height, width))
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/warp/bvh.py

@@ -16,6 +16,7 @@
 """DO NOT EDIT. This file is auto-generated."""
 
 import dataclasses
+import functools
 import jax
 from mujoco.mjx._src import types
 from mujoco.mjx.warp import ffi
@@ -88,7 +89,6 @@ def _refit_bvh_shim(
   _d.geom_xpos = geom_xpos
   _d.nworld = nworld
   render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id]
-
   dummy.zero_()
   mjwarp.refit_bvh(_m, _d, render_context)
 
@@ -134,6 +134,12 @@ def refit_bvh(m: types.Model, d: types.Data, ctx: RenderContext):
 
 @refit_bvh.def_vmap
 @ffi.marshal_custom_vmap
-def refit_bvh_vmap(unused_axis_size, is_batched, m, d, ctx: RenderContext):
+def refit_bvh_vmap(
+    unused_axis_size,
+    is_batched,
+    m: types.Model,
+    d: types.Data,
+    ctx: RenderContext,
+):
   d = refit_bvh(m, d, ctx)
   return d, is_batched[1]

mjx/mujoco/mjx/warp/collision_driver.py

@@ -16,6 +16,7 @@
 """DO NOT EDIT. This file is auto-generated."""
 
 import dataclasses
+import functools
 import jax
 from mujoco.mjx._src import types
 from mujoco.mjx.warp import ffi
@@ -223,7 +224,6 @@ def _collision_shim(
   _d.naconmax = naconmax
   _d.ncollision = ncollision
   _d.nworld = nworld
-
   mjwarp.collision(_m, _d)
 
 
@@ -405,6 +405,6 @@ def collision(m: types.Model, d: types.Data):
 
 @collision.def_vmap
 @ffi.marshal_custom_vmap
-def collision_vmap(unused_axis_size, is_batched, m, d):
+def collision_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data):
   d = collision(m, d)
   return d, is_batched[1]

mjx/mujoco/mjx/warp/forward.py

@@ -16,6 +16,7 @@
 """DO NOT EDIT. This file is auto-generated."""
 
 import dataclasses
+import functools
 import jax
 from mujoco.mjx._src import types
 from mujoco.mjx.warp import ffi
@@ -909,7 +910,6 @@ def _forward_shim(
   _d.xpos = xpos
   _d.xquat = xquat
   _d.nworld = nworld
-
   mjwarp.forward(_m, _d)
 
 
@@ -1815,7 +1815,7 @@ def forward(m: types.Model, d: types.Data):
 
 @forward.def_vmap
 @ffi.marshal_custom_vmap
-def forward_vmap(unused_axis_size, is_batched, m, d):
+def forward_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data):
   d = forward(m, d)
   return d, is_batched[1]
 
@@ -2689,7 +2689,6 @@ def _step_shim(
   _d.xpos = xpos
   _d.xquat = xquat
   _d.nworld = nworld
-
   mjwarp.step(_m, _d)
 
 
@@ -3613,6 +3612,6 @@ def step(m: types.Model, d: types.Data):
 
 @step.def_vmap
 @ffi.marshal_custom_vmap
-def step_vmap(unused_axis_size, is_batched, m, d):
+def step_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data):
   d = step(m, d)
   return d, is_batched[1]

mjx/mujoco/mjx/warp/render.py

@@ -12,17 +12,18 @@
 # 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
 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 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
 
 
@@ -49,6 +50,11 @@ _e = mjwarp.Constraint(
 @ffi.format_args_for_warp
 def _render_shim(
     # Model
+    nworld: int,
+    cam_fovy: wp.array2d(dtype=float),
+    cam_intrinsic: wp.array2d(dtype=wp.vec4),
+    cam_projection: wp.array(dtype=int),
+    cam_sensorsize: wp.array(dtype=wp.vec2),
     geom_dataid: wp.array(dtype=int),
     geom_matid: wp.array2d(dtype=int),
     geom_rgba: wp.array2d(dtype=wp.vec4),
@@ -60,22 +66,9 @@ def _render_shim(
     mat_rgba: wp.array2d(dtype=wp.vec4),
     mat_texid: wp.array3d(dtype=int),
     mat_texrepeat: wp.array2d(dtype=wp.vec2),
-    mesh_face: wp.array(dtype=wp.vec3i),
     mesh_faceadr: wp.array(dtype=int),
-    ncam: int,
-    ngeom: int,
-    nlight: 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),
-    cam_projection: wp.array(dtype=int),
-    cam_fovy: wp.array2d(dtype=wp.float32),
-    cam_sensorsize: wp.array(dtype=wp.vec2),
-    cam_intrinsic: wp.array2d(dtype=wp.vec4),
+    nlight: int,
     # Data
     cam_xmat: wp.array2d(dtype=wp.mat33),
     cam_xpos: wp.array2d(dtype=wp.vec3),
@@ -83,7 +76,6 @@ def _render_shim(
     geom_xpos: wp.array2d(dtype=wp.vec3),
     light_xdir: wp.array2d(dtype=wp.vec3),
     light_xpos: wp.array2d(dtype=wp.vec3),
-    flexvert_xpos: wp.array2d(dtype=wp.vec3),
     # Registry
     rc_id: int,
     rgb: wp.array2d(dtype=wp.uint32),
@@ -93,6 +85,10 @@ def _render_shim(
   _m.opt = _o
   _d.efc = _e
   _d.contact = _c
+  _m.cam_fovy = cam_fovy
+  _m.cam_intrinsic = cam_intrinsic
+  _m.cam_projection = cam_projection
+  _m.cam_sensorsize = cam_sensorsize
   _m.geom_dataid = geom_dataid
   _m.geom_matid = geom_matid
   _m.geom_rgba = geom_rgba
@@ -104,33 +100,16 @@ def _render_shim(
   _m.mat_rgba = mat_rgba
   _m.mat_texid = mat_texid
   _m.mat_texrepeat = mat_texrepeat
-  _m.mesh_face = mesh_face
   _m.mesh_faceadr = mesh_faceadr
-  _m.ncam = ncam
-  _m.ngeom = ngeom
-  _m.nlight = nlight
-  _m.cam_projection = cam_projection
-  _m.cam_fovy = cam_fovy
-  _m.cam_sensorsize = cam_sensorsize
-  _m.cam_intrinsic = cam_intrinsic
   _m.nflex = nflex
-  _m.nflexelemdata = nflexelemdata
-  _m.nflexvert = nflexvert
-  _m.flex_dim = flex_dim
-  _m.flex_elem = flex_elem
-  _m.flex_elemnum = flex_elemnum
-  _m.flex_vertadr = flex_vertadr
-
+  _m.nlight = nlight
   _d.cam_xmat = cam_xmat
   _d.cam_xpos = cam_xpos
   _d.geom_xmat = geom_xmat
   _d.geom_xpos = geom_xpos
   _d.light_xdir = light_xdir
   _d.light_xpos = light_xpos
-  _d.flexvert_xpos = flexvert_xpos
-
-  _d.nworld = cam_xpos.shape[0]
-
+  _d.nworld = nworld
   render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id]
   render_context.rgb_data = rgb
   render_context.depth_data = depth
@@ -139,55 +118,63 @@ def _render_shim(
 
 def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext):
   render_ctx = _MJX_RENDER_CONTEXT_BUFFERS[ctx.key]
-
   output_dims = {
       'rgb': render_ctx.rgb_data_shape,
       'depth': render_ctx.depth_data_shape,
   }
-
   jf = ffi.jax_callable_variadic_tuple(
       _render_shim,
       num_outputs=2,
       output_dims=output_dims,
       vmap_method=None,
+      in_out_argnames=set([]),
+      stage_in_argnames=set([
+          'cam_fovy',
+          'cam_intrinsic',
+          'cam_xmat',
+          'cam_xpos',
+          'geom_matid',
+          'geom_rgba',
+          'geom_size',
+          'geom_xmat',
+          'geom_xpos',
+          'light_castshadow',
+          'light_type',
+          'mat_rgba',
+          'mat_texid',
+      ]),
+      stage_out_argnames=set([]),
+      graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(
+      d.qpos.shape[0],
+      m.cam_fovy,
+      m.cam_intrinsic,
+      m._impl.cam_projection,
+      m.cam_sensorsize,
       m.geom_dataid,
       m.geom_matid,
       m.geom_rgba,
       m.geom_size,
       m.geom_type,
-      m.light_active,
+      m._impl.light_active,
       m.light_castshadow,
       m.light_type,
       m.mat_rgba,
       m.mat_texid,
-      m.mat_texrepeat,
-      m.mesh_face,
+      m._impl.mat_texrepeat,
       m.mesh_faceadr,
-      m.ncam,
-      m.ngeom,
+      m._impl.nflex,
       m.nlight,
-      m.nflex,
-      m.nflexelemdata,
-      m.nflexvert,
-      m.flex_dim,
-      m.flex_elem,
-      m.flex_elemnum,
-      m.flex_vertadr,
-      m.cam_projection,
-      m.cam_fovy,
-      m.cam_sensorsize,
-      m.cam_intrinsic,
       d.cam_xmat,
       d.cam_xpos,
       d.geom_xmat,
       d.geom_xpos,
-      d.light_xdir,
-      d.light_xpos,
-      d.flexvert_xpos,
+      d._impl.light_xdir,
+      d._impl.light_xpos,
       ctx.key,
   )
+  d = d.tree_replace({})
   return out
 
 
@@ -199,6 +186,12 @@ def render(m: types.Model, d: types.Data, ctx: RenderContext):
 
 @render.def_vmap
 @functools.partial(ffi.marshal_custom_vmap, skip_output_dim_reshape=True)
-def render_vmap(unused_axis_size, is_batched, m, d, ctx):
+def render_vmap(
+    unused_axis_size,
+    is_batched,
+    m: types.Model,
+    d: types.Data,
+    ctx: RenderContext,
+):
   out = render(m, d, ctx)
   return out, [True, True]

mjx/mujoco/mjx/warp/render_test.py

@@ -21,13 +21,14 @@ import jax
 from jax import numpy as jp
 import mujoco
 from mujoco import mjx
-from mujoco.mjx._src import io
 from mujoco.mjx._src import forward
+from mujoco.mjx._src import io
 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
 import numpy as np
 
+
 _FORCE_TEST = os.environ.get('MJX_WARP_FORCE_TEST', '0') == '1'
 
 
@@ -91,12 +92,8 @@ class RenderTest(parameterized.TestCase):
         enabled_geom_groups=[0, 1, 2],
     )
 
-    dx_batch = jax.jit(
-        jax.vmap(mjx.refit_bvh, in_axes=(None, 0, None))
-    )(mx, dx_batch, rc)
-    out_batch = jax.jit(
-        jax.vmap(mjx.render, in_axes=(None, 0, None))
-    )(mx, dx_batch, rc)
+    dx_batch = jax.jit(mjx.refit_bvh)(mx, dx_batch, rc)
+    out_batch = jax.jit(mjx.render)(mx, dx_batch, rc)
 
     rgb = np.asarray(out_batch[0])
     depth = np.asarray(out_batch[1])

mjx/mujoco/mjx/warp/smooth.py

@@ -16,6 +16,7 @@
 """DO NOT EDIT. This file is auto-generated."""
 
 import dataclasses
+import functools
 import jax
 from mujoco.mjx._src import types
 from mujoco.mjx.warp import ffi
@@ -137,7 +138,6 @@ def _kinematics_shim(
   _d.xpos = xpos
   _d.xquat = xquat
   _d.nworld = nworld
-
   mjwarp.kinematics(_m, _d)
 
 
@@ -283,7 +283,9 @@ def kinematics(m: types.Model, d: types.Data):
 
 @kinematics.def_vmap
 @ffi.marshal_custom_vmap
-def kinematics_vmap(unused_axis_size, is_batched, m, d):
+def kinematics_vmap(
+    unused_axis_size, is_batched, m: types.Model, d: types.Data
+):
   d = kinematics(m, d)
   return d, is_batched[1]
 
@@ -369,7 +371,6 @@ def _tendon_shim(
   _d.wrap_obj = wrap_obj
   _d.wrap_xpos = wrap_xpos
   _d.nworld = nworld
-
   mjwarp.tendon(_m, _d)
 
 
@@ -465,6 +466,6 @@ def tendon(m: types.Model, d: types.Data):
 
 @tendon.def_vmap
 @ffi.marshal_custom_vmap
-def tendon_vmap(unused_axis_size, is_batched, m, d):
+def tendon_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data):
   d = tendon(m, d)
   return d, is_batched[1]

mjx/mujoco/mjx/warp/testspeed.py

@@ -26,16 +26,18 @@ import jax.numpy as jnp
 import jax.tree_util
 import mujoco
 from mujoco import mjx
-from mujoco.mjx.warp import io
+from mujoco.mjx._src import render_util
 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
+from mujoco.mjx.warp import io
 from mujoco.mjx.warp import smooth as wp_smooth
 import mujoco.mjx.third_party.mujoco_warp as mjwarp
 import numpy as np
 import warp as wp
 from mujoco.mjx.third_party.warp._src.jax_experimental import ffi as warp_ffi
 
+
 _MODELFILE = flags.DEFINE_string(
     'modelfile',
     'humanoid/humanoid.xml',
@@ -150,6 +152,11 @@ def benchmark(
         enabled_geom_groups=[0, 1, 2],
     )
 
+  def render_fn(mx, d, rc):
+    d = mjx.refit_bvh(mx, d, rc)
+    pixels, _ = mjx.render(mx, d, rc)
+    return render_util.get_rgb(rc, pixels, 0), d
+
   @jax_jit
   def unroll(d):
     def fn(carry, _):
@@ -158,10 +165,8 @@ def benchmark(
       d = step_fn(mx, d)
 
       if render:
-        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] for x in leaves if x.size > 0) if leaves else 0.0
+        rgb, d = jax.vmap(render_fn, in_axes=(None, 0, None))(mx, d, rc)
+        accum += rgb[0, 0, 0, 0]
 
       return (d, accum), None
 

mjx/mujoco/mjx/warp/visualize_render.py

@@ -27,7 +27,6 @@ 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
 from mujoco.mjx._src import render_util
 from mujoco.mjx._src import test_util
@@ -47,11 +46,14 @@ _NWORLD = flags.DEFINE_integer(
 _WIDTH = flags.DEFINE_integer('width', 512, 'image width')
 _HEIGHT = flags.DEFINE_integer('height', 512, 'image height')
 _CAMERA_ID = flags.DEFINE_integer(
-    'camera_id', 1, 'camera id to visualize'
+    'camera_id', 0, 'camera id to visualize'
 )
 _OUTPUT_DIR = flags.DEFINE_string(
     'output_dir', '/tmp/visualize_render', 'output directory'
 )
+_RANDOMIZE_QPOS = flags.DEFINE_boolean(
+    'randomize_qpos', False, 'randomize initial qpos'
+)
 _USE_TEXTURES = flags.DEFINE_boolean(
     'use_textures', True, 'enable textures'
 )
@@ -130,9 +132,13 @@ def _main(_: Sequence[str]):
     dx = mjx.make_data(m, impl='warp')
     rng = jax.random.PRNGKey(worldid)
     qpos0 = jp.array(m.qpos0)
-    qpos = qpos0 + jax.random.uniform(
-        rng, (m.nq,), minval=-0.2, maxval=0.05
-    )
+    qpos = qpos0
+    if _RANDOMIZE_QPOS.value:
+      # TODO(robotics-team): consider integrating velocity if there are free
+      # joints.
+      qpos = qpos0 + jax.random.uniform(
+          rng, (m.nq,), minval=-0.2, maxval=0.05
+      )
     return dx.replace(qpos=qpos)
 
   print('initializing data...')
@@ -171,9 +177,9 @@ def _main(_: Sequence[str]):
   rgb_packed = out_batch[0]
   print(f'  rgb shape: {rgb_packed.shape}\n')
 
-  rgb = render_util.get_rgb(
-      rc, rgb_packed, _CAMERA_ID.value
-  )
+  rgb = jax.vmap(
+      render_util.get_rgb, in_axes=(None, 0, None)
+  )(rc, rgb_packed, _CAMERA_ID.value)
 
   single_path = os.path.join(
       _OUTPUT_DIR.value, f'camera_{_CAMERA_ID.value}.png'