initial implementation of mjx-warp render.

PiperOrigin-RevId: 869822624
Change-Id: Ic20db1dcc2e7287752b273c4ad986066cceccad9

mjx/mujoco/mjx/__init__.py

@@ -22,6 +22,7 @@ from mujoco.mjx._src.types import Data
 
 # pylint:disable=g-importing-member
 from mujoco.mjx._src.collision_driver import collision
+from mujoco.mjx._src.bvh import refit_bvh
 from mujoco.mjx._src.constraint import make_constraint
 from mujoco.mjx._src.derivative import deriv_smooth_vel
 from mujoco.mjx._src.forward import euler
@@ -33,6 +34,7 @@ from mujoco.mjx._src.forward import fwd_velocity
 from mujoco.mjx._src.forward import implicit
 from mujoco.mjx._src.forward import rungekutta4
 from mujoco.mjx._src.inverse import inverse
+from mujoco.mjx._src.io import create_render_context
 from mujoco.mjx._src.io import get_data
 from mujoco.mjx._src.io import get_data_into
 from mujoco.mjx._src.io import get_state
@@ -43,6 +45,7 @@ from mujoco.mjx._src.io import set_state
 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.sensor import sensor_acc
 from mujoco.mjx._src.sensor import sensor_pos
 from mujoco.mjx._src.sensor import sensor_vel

mjx/mujoco/mjx/_src/bvh.py

@@ -0,0 +1,32 @@
+# 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.
+# ==============================================================================
+"""BVH helpers for MJX."""
+
+from typing import Any
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import Impl
+from mujoco.mjx._src.types import Model
+# pylint: enable=g-importing-member
+import mujoco.mjx.warp as mjxw
+
+
+def refit_bvh(m: Model, d: Data, ctx: Any):
+  """Refit the scene BVH for the current pose."""
+  if m.impl == Impl.WARP and d.impl == Impl.WARP and mjxw.WARP_INSTALLED:
+    from mujoco.mjx.warp import bvh as mjxw_bvh  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+    return mjxw_bvh.refit_bvh(m, d, ctx)
+
+  raise NotImplementedError('refit_bvh only implemented for MuJoCo Warp.')

mjx/mujoco/mjx/_src/io.py

@@ -17,7 +17,7 @@
 import copy
 import logging
 import os
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 import warnings
 
 import jax
@@ -52,7 +52,9 @@ def _is_cuda_gpu_device(device: jax.Device) -> bool:
 
 def _check_warp_installed():
   if not mjxw.WARP_INSTALLED:
-    raise RuntimeError('warp-lang is not installed. Cannot use WARP implementation of MJX.')
+    raise RuntimeError(
+        'warp-lang is not installed. Cannot use WARP implementation of MJX.'
+    )
 
 
 def _resolve_impl(
@@ -121,8 +123,7 @@ def _check_impl_device_compatibility(
   if impl == types.Impl.WARP:
     if not _is_cuda_gpu_device(device):
       raise AssertionError(
-          'Warp implementation requires a CUDA GPU device, got '
-          f'{device}.'
+          f'Warp implementation requires a CUDA GPU device, got {device}.'
       )
     _check_warp_installed()
 
@@ -999,8 +1000,7 @@ def make_data(
 
   if isinstance(m, types.Model) and m.impl != impl:
     raise ValueError(
-        f'Model impl {m.impl} does not match make_data '
-        f'implementation {impl}.'
+        f'Model impl {m.impl} does not match make_data implementation {impl}.'
     )
 
   if impl == types.Impl.JAX:
@@ -1546,8 +1546,7 @@ def _get_data_into(
       all_fields = types.Data.fields() + types.DataC.fields()
     else:
       raise NotImplementedError(
-          f'get_data_into for implementation "{d.impl}" not implemented'
-          ' yet.'
+          f'get_data_into for implementation "{d.impl}" not implemented yet.'
       )
 
     for field in all_fields:
@@ -1699,8 +1698,9 @@ def _get_data_into_cpp(
   # mjx.Model which we don't have access to in this function.
   fields_to_check = ['qpos', 'qvel', 'act', 'mocap_pos', 'mocap_quat']
   for i in range(batch_size):
-    d_i: types.Data = jax.tree_util.tree_map(
-        lambda x, i=i: x[i], d) if batched else d
+    d_i: types.Data = (
+        jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d
+    )
     src_data = mj_data_list[i]
 
     needs_syncing = False
@@ -1939,3 +1939,26 @@ def set_state(
       offset += size
 
   return d.replace(**updates)
+
+
+def create_render_context(
+    mjm: mujoco.MjModel,
+    nworld: int,
+    **kwargs,
+):
+  """Creates a render context.
+
+  Args:
+    mjm: the MuJoCo model
+    nworld: number of worlds to render. We must hardcode the nworld
+      because Warp creates arrays of size nworld that are not exposed
+      to JAX. Thus we cannot use JAX transforms like vmap with the
+      render context.
+    **kwargs: forwarded to the render context constructor.
+
+  Returns:
+    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
+  return mjxw_render.create_render_context(mjm, nworld=nworld, **kwargs)

mjx/mujoco/mjx/_src/render.py

@@ -0,0 +1,32 @@
+# 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.
+# ==============================================================================
+"""Render helpers for MJX."""
+
+from typing import Any
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import Impl
+from mujoco.mjx._src.types import Model
+# pylint: enable=g-importing-member
+import mujoco.mjx.warp as mjxw
+
+
+def render(m: Model, d: Data, ctx: Any) -> Data:
+  """Render."""
+  if m.impl == Impl.WARP and d.impl == Impl.WARP and mjxw.WARP_INSTALLED:
+    from mujoco.mjx.warp import render as mjxw_render  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+    return mjxw_render.render(m, d, ctx)
+
+  raise NotImplementedError('render only implemented for MuJoCo Warp.')

mjx/mujoco/mjx/_src/render_util.py

@@ -0,0 +1,73 @@
+# 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.
+# ==============================================================================
+"""JAX render utilities for unpacking render output from MuJoCo Warp."""
+
+import jax
+import jax.numpy as jnp
+
+
+def get_rgb(
+    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)
+      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]
+  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
+  rgb = jnp.stack([r, g, b], axis=-1)
+  nworld = rgb_data.shape[0]
+  return rgb.reshape(nworld, height, width, 3)
+
+
+def get_depth(
+    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)
+      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]
+  nworld = depth_data.shape[0]
+  depth = jnp.clip(raw / depth_scale, 0.0, 1.0)
+  return depth.reshape(nworld, height, width)

mjx/mujoco/mjx/_src/types.py

@@ -917,7 +917,7 @@ class Model(PyTreeNode):
   tex_adr: np.ndarray
   tex_data: np.ndarray
   mat_rgba: jax.Array
-  mat_texid: np.ndarray
+  mat_texid: jax.Array
   pair_dim: np.ndarray
   pair_geom1: np.ndarray
   pair_geom2: np.ndarray

mjx/mujoco/mjx/warp/bvh.py

@@ -0,0 +1,117 @@
+# 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.
+# ==============================================================================
+"""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.render import _MJX_RENDER_CONTEXT_BUFFERS
+from mujoco.mjx.warp.types import RenderContext
+import mujoco.mjx.third_party.mujoco_warp as mjwarp
+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}
+)
+
+
+@ffi.format_args_for_warp
+def _refit_bvh_shim(
+    # Model
+    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
+    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),
+):
+  _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
+  _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 = geom_xpos.shape[0]
+
+  render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id]
+  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]
+  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)},
+      vmap_method=None,
+  )
+  out = jf(
+      m.geom_dataid,
+      m.geom_size,
+      m.geom_type,
+      m.nflex,
+      m.nflexelemdata,
+      m.nflexvert,
+      m.flex_dim,
+      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])
+
+
+@jax.custom_batching.custom_vmap
+@functools.partial(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)
+def refit_bvh_vmap(unused_axis_size, is_batched, m, d, ctx):
+  out = refit_bvh(m, d, ctx)
+  return out, is_batched[1]

mjx/mujoco/mjx/warp/collision_driver.py

@@ -1,4 +1,4 @@
-# Copyright 2025 DeepMind Technologies Limited
+# 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.
@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
+
 import dataclasses
 import jax
 from mujoco.mjx._src import types
@@ -399,6 +400,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
 @ffi.marshal_jax_warp_callable
 def collision(m: types.Model, d: types.Data):
   return _collision_jax_impl(m, d)
+
+
 @collision.def_vmap
 @ffi.marshal_custom_vmap
 def collision_vmap(unused_axis_size, is_batched, m, d):

mjx/mujoco/mjx/warp/ffi.py

@@ -252,11 +252,11 @@ def _squeeze_dim(leaf_expanded: Any, leaf: Any) -> Any:
   return leaf_expanded
 
 
-def marshal_jax_warp_callable(func, raw_output: bool = False):
+def marshal_jax_warp_callable(func, skip_output_dim_reshape: bool = False):
   """Marshal fields into a MuJoCo Warp function."""
 
   @functools.wraps(func)
-  def wrapper(m, d):
+  def wrapper(m, d, *extra_args):
     # Expand dims for Warp implicit vmap before calling into the FFI wrapped
     # function.
     m_expanded = jax.tree.map_with_path(
@@ -271,9 +271,9 @@ def marshal_jax_warp_callable(func, raw_output: bool = False):
         ),
         d,
     )
-    d_expanded_result = func(m_expanded, d_expanded)
+    d_expanded_result = func(m_expanded, d_expanded, *extra_args)
 
-    if raw_output:
+    if skip_output_dim_reshape:
       return d_expanded_result
     d_result = jax.tree.map(_squeeze_dim, d_expanded_result, d)
     return d_result
@@ -360,11 +360,11 @@ def _check_leading_dim(
     )
 
 
-def marshal_custom_vmap(vmap_func, raw_output: bool = False):
+def marshal_custom_vmap(vmap_func, skip_output_dim_reshape: bool = False):
   """Marshal fields for a custom vmap into an MuJoCo Warp function."""
 
   @functools.wraps(vmap_func)
-  def wrapper(axis_size, is_batched, m, d):
+  def wrapper(axis_size, is_batched, m, d, *extra_args):
     # Vmappable data fields may not have been broadcasted if vmap_func is called
     # within a vmap trace. Since data fields are read/write in warp, we need to
     # explicitly broadcast them here.
@@ -395,9 +395,9 @@ def marshal_custom_vmap(vmap_func, raw_output: bool = False):
         d_broadcast,
     )
     d_broadcast_flat_result, out_batched = vmap_func(
-        axis_size, is_batched, m_flat, d_broadcast_flat
+        axis_size, is_batched, m_flat, d_broadcast_flat, *extra_args
     )
-    if raw_output:
+    if skip_output_dim_reshape:
       return d_broadcast_flat_result, out_batched
 
     # Explicitly mark MuJoCo Warp data fields as batched after vmapping is done.

mjx/mujoco/mjx/warp/forward.py

@@ -1,4 +1,4 @@
-# Copyright 2025 DeepMind Technologies Limited
+# 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.
@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
+
 import dataclasses
 import jax
 from mujoco.mjx._src import types
@@ -1809,6 +1810,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
 @ffi.marshal_jax_warp_callable
 def forward(m: types.Model, d: types.Data):
   return _forward_jax_impl(m, d)
+
+
 @forward.def_vmap
 @ffi.marshal_custom_vmap
 def forward_vmap(unused_axis_size, is_batched, m, d):
@@ -3604,6 +3607,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
 @ffi.marshal_jax_warp_callable
 def step(m: types.Model, d: types.Data):
   return _step_jax_impl(m, d)
+
+
 @step.def_vmap
 @ffi.marshal_custom_vmap
 def step_vmap(unused_axis_size, is_batched, m, d):

mjx/mujoco/mjx/warp/render.py

@@ -0,0 +1,263 @@
+# 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.
+# ==============================================================================
+"""DO NOT EDIT. This file is auto-generated."""
+
+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.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}
+)
+_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 _render_shim(
+    # Model
+    geom_dataid: wp.array(dtype=int),
+    geom_matid: wp.array2d(dtype=int),
+    geom_rgba: wp.array2d(dtype=wp.vec4),
+    geom_size: wp.array2d(dtype=wp.vec3),
+    geom_type: wp.array(dtype=int),
+    light_active: wp.array2d(dtype=bool),
+    light_castshadow: wp.array2d(dtype=bool),
+    light_type: wp.array2d(dtype=int),
+    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),
+    # Data
+    cam_xmat: wp.array2d(dtype=wp.mat33),
+    cam_xpos: wp.array2d(dtype=wp.vec3),
+    geom_xmat: wp.array2d(dtype=wp.mat33),
+    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.array3d(dtype=wp.uint32),
+    depth: wp.array3d(dtype=wp.float32),
+):
+  _m.stat = _s
+  _m.opt = _o
+  _d.efc = _e
+  _d.contact = _c
+  _m.geom_dataid = geom_dataid
+  _m.geom_matid = geom_matid
+  _m.geom_rgba = geom_rgba
+  _m.geom_size = geom_size
+  _m.geom_type = geom_type
+  _m.light_active = light_active
+  _m.light_castshadow = light_castshadow
+  _m.light_type = light_type
+  _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
+
+  _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]
+
+  render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id]
+  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),
+      'depth': (nworld, nrender, height * width),
+  }
+
+  jf = ffi.jax_callable_variadic_tuple(
+      _render_shim,
+      num_outputs=2,
+      output_dims=output_dims,
+      vmap_method=None,
+  )
+  out = jf(
+      m.geom_dataid,
+      m.geom_matid,
+      m.geom_rgba,
+      m.geom_size,
+      m.geom_type,
+      m.light_active,
+      m.light_castshadow,
+      m.light_type,
+      m.mat_rgba,
+      m.mat_texid,
+      m.mat_texrepeat,
+      m.mesh_face,
+      m.mesh_faceadr,
+      m.ncam,
+      m.ngeom,
+      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,
+      ctx.key,
+  )
+  return out
+
+
+@jax.custom_batching.custom_vmap
+@functools.partial(ffi.marshal_jax_warp_callable, skip_output_dim_reshape=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)
+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/render_test.py

@@ -0,0 +1,111 @@
+# 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 functools
+import os
+
+from absl.testing import absltest
+from absl.testing import parameterized
+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
+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'
+
+
+class RenderTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    if mjxw.WARP_INSTALLED:
+      tempdir = '/tmp/wp_kernel_cache_dir_RenderTest'
+      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."""
+    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],
+    )
+
+    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)
+
+    rgb = np.asarray(out_batch[0])
+    depth = np.asarray(out_batch[1])
+
+    self.assertGreater(np.count_nonzero(rgb), 0)
+    self.assertGreater(np.count_nonzero(depth), 0)
+    self.assertNotEqual(np.unique(rgb).shape[0], 1)
+    self.assertNotEqual(np.unique(depth).shape[0], 1)
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/warp/smooth.py

@@ -1,4 +1,4 @@
-# Copyright 2025 DeepMind Technologies Limited
+# 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.
@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
+
 import dataclasses
 import jax
 from mujoco.mjx._src import types
@@ -277,6 +278,8 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
 @ffi.marshal_jax_warp_callable
 def kinematics(m: types.Model, d: types.Data):
   return _kinematics_jax_impl(m, d)
+
+
 @kinematics.def_vmap
 @ffi.marshal_custom_vmap
 def kinematics_vmap(unused_axis_size, is_batched, m, d):
@@ -456,6 +459,8 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
 @ffi.marshal_jax_warp_callable
 def tendon(m: types.Model, d: types.Data):
   return _tendon_jax_impl(m, d)
+
+
 @tendon.def_vmap
 @ffi.marshal_custom_vmap
 def tendon_vmap(unused_axis_size, is_batched, m, d):

mjx/mujoco/mjx/warp/testspeed.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Run benchmarks."""
+"""Run benchmarks for MJX-Warp."""
 
 import functools
 import os
@@ -23,8 +23,10 @@ from absl import app
 from absl import flags
 import jax
 import jax.numpy as jnp
+import jax.tree_util
 import mujoco
 from mujoco import mjx
+from mujoco.mjx._src 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
@@ -64,6 +66,20 @@ _BENCHMARK = flags.DEFINE_enum(
     ['jax_warp', 'jax', 'warp'],
     'Which benchmark to run.',
 )
+# Render flags
+_RENDER_WIDTH = flags.DEFINE_integer('render_width', 64, 'render width')
+_RENDER_HEIGHT = flags.DEFINE_integer('render_height', 64, 'render height')
+_RENDER_RGB = flags.DEFINE_boolean('render_rgb', False, 'render RGB image')
+_RENDER_DEPTH = flags.DEFINE_boolean(
+    'render_depth', False, 'render depth image'
+)
+_RENDER_USE_TEXTURES = flags.DEFINE_boolean(
+    'render_textures', True, 'use textures'
+)
+_RENDER_USE_SHADOWS = flags.DEFINE_boolean(
+    'render_shadows', False, 'use shadows'
+)
+
 _COMPILER_OPTIONS = {'xla_gpu_graph_min_graph_size': 1}
 jax_jit = functools.partial(jax.jit, compiler_options=_COMPILER_OPTIONS)
 
@@ -101,6 +117,7 @@ def benchmark(
     nstep: int = 1000,
     nenv: int = 8192,
     unroll_steps: int = 4,
+    render: bool = False,
 ) -> Tuple[float, float, int]:
   """Benchmark a model."""
 
@@ -119,13 +136,39 @@ def benchmark(
 
   jax.block_until_ready(d)
 
+  rc = None
+  if render:
+    ncam = max(1, m.ncam)
+    rc = io.create_render_context(
+        mjm=m,
+        nworld=nenv,
+        cam_res=(_RENDER_WIDTH.value, _RENDER_HEIGHT.value),
+        use_textures=_RENDER_USE_TEXTURES.value,
+        use_shadows=_RENDER_USE_SHADOWS.value,
+        render_rgb=[_RENDER_RGB.value] * ncam,
+        render_depth=[_RENDER_DEPTH.value] * ncam,
+        enabled_geom_groups=[0, 1, 2],
+    )
+
   @jax_jit
   def unroll(d):
-    def fn(d, _):
+    def fn(carry, _):
+      d, accum = carry
       d = d.replace(qpos=d.qpos + 0 * d.qpos)
-      return step_fn(mx, d), None
+      d = step_fn(mx, d)
 
-    return jax.lax.scan(fn, d, None, length=nstep, unroll=unroll_steps)
+      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, 0] for x in leaves) if leaves else 0.0
+
+      return (d, accum), None
+
+    (d_final, accum_final), _ = jax.lax.scan(
+        fn, (d, jnp.array(0.0)), None, length=nstep, unroll=unroll_steps
+    )
+    return d_final, accum_final
 
   jit_time, run_time = _measure(unroll, d)
   steps = nstep * nenv
@@ -147,6 +190,7 @@ def benchmark_raw_warp(
     nenv: int = 8192,
     unroll_steps: int = 4,
     function: str = 'kinematics',
+    render: bool = False,
 ):
   """Benchmarks raw warp."""
   del unroll_steps
@@ -178,8 +222,27 @@ def benchmark_raw_warp(
   else:
     raise NotImplementedError(f'{function} not implemented in speed test.')
 
+  ncam = max(1, m.ncam)
+  rc = None
+  if render:
+    mjwarp.forward(mw, dw)
+    rc = mjwarp.create_render_context(
+        m, mw, dw,
+        (_RENDER_WIDTH.value, _RENDER_HEIGHT.value),
+        [_RENDER_RGB.value] * ncam,
+        [_RENDER_DEPTH.value] * ncam,
+        _RENDER_USE_TEXTURES.value,
+        _RENDER_USE_SHADOWS.value,
+    )
+
+  def run_fn(m, d):
+    fn(m, d)
+    if render:
+      mjwarp.refit_bvh(m, d, rc)
+      mjwarp.render(m, d, rc)
+
   start = time.time()
-  graph = _compile_fn(fn, mw, dw)
+  graph = _compile_fn(run_fn, mw, dw)
   jit_time = time.time() - start
 
   start = time.time()
@@ -241,17 +304,27 @@ def _main(_: Sequence[str]):
   print(f' timestep             : {m.opt.timestep}')
   print(f' unroll               : {unroll}')
   print(f' benchmark            : {benchmark_type}')
-  print(f' graph_mode           : {_GRAPH_MODE.value}\n')
+  print(f' graph_mode           : {_GRAPH_MODE.value}')
+  is_render = _RENDER_DEPTH.value or _RENDER_RGB.value
+  if is_render:
+    print(
+        f' render resolution    : {_RENDER_WIDTH.value}x{_RENDER_HEIGHT.value}'
+    )
+    print(f' textures             : {_RENDER_USE_TEXTURES.value}')
+    print(f' shadows              : {_RENDER_USE_SHADOWS.value}')
+  print()
 
   if benchmark_type == 'jax_warp':
-    jit_time, run_time, steps = benchmark(m, mw, func_warp, nstep, nenv, unroll)
+    jit_time, run_time, steps = benchmark(
+        m, mw, func_warp, nstep, nenv, unroll, is_render
+    )
     print(f' JAX WARP FFI (GraphMode: {_GRAPH_MODE.value}):')
   elif benchmark_type == 'jax':
     jit_time, run_time, steps = benchmark(m, mx, func_jax, nstep, nenv, unroll)
     print(' Pure JAX:')
   elif benchmark_type == 'warp':
     jit_time, run_time, steps = benchmark_raw_warp(
-        m, nstep, nenv, unroll, function=function_
+        m, nstep, nenv, unroll, function=function_, render=is_render
     )
     print(' Pure WARP:')
   else:

mjx/mujoco/mjx/warp/types.py

@@ -1,4 +1,4 @@
-# Copyright 2025 DeepMind Technologies Limited
+# 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.
@@ -13,8 +13,10 @@
 # limitations under the License.
 # ==============================================================================
 """MJX Warp types.
+
 DO NOT EDIT. This file is auto-generated.
 """
+
 import dataclasses
 import typing
 from typing import Tuple
@@ -23,6 +25,7 @@ from jax import tree_util
 from jax.interpreters import batching
 from mujoco.mjx._src import dataclasses as mjx_dataclasses
 import numpy as np
+
 if typing.TYPE_CHECKING:
   GraphMode = int
 else:
@@ -32,6 +35,7 @@ else:
     GraphMode = int
 PyTreeNode = mjx_dataclasses.PyTreeNode
 
+
 @dataclasses.dataclass(frozen=True)
 @tree_util.register_pytree_node_class
 class TileSet:
@@ -64,6 +68,7 @@ class BlockDim:
 
   TODO(team): experimental and may be removed
   """
+
   actuator_velocity: int
   cholesky_factorize: int
   cholesky_factorize_solve: int
@@ -91,12 +96,48 @@ class BlockDim:
     return cls(*children)
 
 
+@dataclasses.dataclass(frozen=True)
+@tree_util.register_pytree_node_class
+class RenderContext:
+  """Render context handle with automatic cleanup.
+
+  The key is static (aux_data) so JAX doesn't trace it, allowing the Warp FFI
+  to receive a concrete int value. Only the original instance (_owner=True)
+  cleans up on deletion; JAX copies have _owner=False.
+  """
+
+  key: int
+  _owner: bool = True
+
+  def tree_flatten(self):
+    return ((), (self.key, False))
+
+  @classmethod
+  def tree_unflatten(cls, aux_data, children):
+    del children
+    key, owner = aux_data
+    return cls(key, owner)
+
+  def __del__(self):
+    if not self._owner:
+      return
+    lock = globals().get('_MJX_RENDER_CONTEXT_LOCK')
+    buffers = globals().get('_MJX_RENDER_CONTEXT_BUFFERS')
+    if lock is None or buffers is None:
+      return
+    with lock:
+      buffers.pop(self.key, None)
+
+
 class StatisticWarp(PyTreeNode):
   """Derived fields from Statistic."""
+
   meaninertia: jax.Array
 
+
 class OptionWarp(PyTreeNode):
   """Derived fields from Option."""
+
   broadphase: int
   broadphase_filter: int
   ccd_iterations: int
@@ -111,8 +152,10 @@ class OptionWarp(PyTreeNode):
   sdf_initpoints: int
   sdf_iterations: int
 
+
 class ModelWarp(PyTreeNode):
   """Derived fields from Model."""
+
   M_colind: np.ndarray
   M_rowadr: np.ndarray
   M_rownnz: np.ndarray
@@ -245,8 +288,10 @@ class ModelWarp(PyTreeNode):
   wrap_site_adr: np.ndarray
   wrap_site_pair_adr: np.ndarray
 
+
 class DataWarp(PyTreeNode):
   """Derived fields from Data."""
+
   actuator_moment: jax.Array
   actuator_velocity: jax.Array
   cacc: jax.Array
@@ -312,6 +357,8 @@ class DataWarp(PyTreeNode):
   wrap_obj: jax.Array
   wrap_xpos: jax.Array
   shape = property(lambda self: self.cacc.shape)
+
+
 DATA_NON_VMAP = {
     'contact__dim',
     'contact__dist',
@@ -335,6 +382,7 @@ DATA_NON_VMAP = {
     'nworld',
 }
 
+
 def _to_elt(cont, _, d, axis):
   return DataWarp(**{
       f.name: (

mjx/mujoco/mjx/warp/visualize_render.py

@@ -0,0 +1,196 @@
+# 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.
+# ==============================================================================
+"""Visualize batch rendering output from MuJoCo Warp for debugging."""
+
+import functools
+import os
+from typing import Sequence
+
+from absl import app
+from absl import flags
+import jax
+import jax.numpy as jp
+import mediapy as media
+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
+import numpy as np
+import warp as wp
+
+
+_MODELFILE = flags.DEFINE_string(
+    'modelfile',
+    'humanoid/humanoid.xml',
+    'path to model',
+)
+_NWORLD = flags.DEFINE_integer(
+    'nworld', 4, 'number of worlds to render'
+)
+_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'
+)
+_OUTPUT_DIR = flags.DEFINE_string(
+    'output_dir', '/tmp/visualize_render', 'output directory'
+)
+_USE_TEXTURES = flags.DEFINE_boolean(
+    'use_textures', True, 'enable textures'
+)
+_USE_SHADOWS = flags.DEFINE_boolean(
+    'use_shadows', True, 'enable shadows'
+)
+_WP_KERNEL_CACHE_DIR = flags.DEFINE_string(
+    'wp_kernel_cache_dir',
+    '/tmp/wp_kernel_cache_dir_visualize_render',
+    'warp kernel cache directory',
+)
+
+_COMPILER_OPTIONS = {'xla_gpu_graph_min_graph_size': 1}
+jax_jit = functools.partial(
+    jax.jit, compiler_options=_COMPILER_OPTIONS
+)
+
+
+def _save_single(rgb, out_path):
+  """Save first world as a single image."""
+  img = np.asarray(rgb[0])
+  img_uint8 = (img * 255).astype(np.uint8)
+  media.write_image(out_path, img_uint8)
+  print(f'  single image: {out_path}')
+
+
+def _save_tiled(rgb, out_path):
+  """Save all worlds as a tiled grid."""
+  nworld, height, width, _ = rgb.shape
+  cols = int(np.ceil(np.sqrt(nworld)))
+  rows = int(np.ceil(nworld / cols))
+  canvas = np.zeros(
+      (rows * height, cols * width, 3), dtype=np.uint8
+  )
+
+  for w in range(nworld):
+    img_uint8 = (np.asarray(rgb[w]) * 255).astype(
+        np.uint8
+    )
+    r, c = w // cols, w % cols
+    y0, y1 = r * height, (r + 1) * height
+    x0, x1 = c * width, (c + 1) * width
+    canvas[y0:y1, x0:x1, :] = img_uint8
+
+  media.write_image(out_path, canvas)
+  print(f'  tiled image:  {out_path}')
+
+
+def _main(_: Sequence[str]):
+  os.environ['MJX_WARP_ENABLED'] = 'true'
+
+  wp.config.kernel_cache_dir = _WP_KERNEL_CACHE_DIR.value
+
+  os.makedirs(_OUTPUT_DIR.value, exist_ok=True)
+
+  try:
+    m = test_util.load_test_file(_MODELFILE.value)
+  except Exception:
+    m = mujoco.MjModel.from_xml_path(_MODELFILE.value)
+
+  print('visualize_render.py:\n')
+  print(f'  modelfile   : {_MODELFILE.value}')
+  print(f'  nworld      : {_NWORLD.value}')
+  print(f'  resolution  : {_WIDTH.value}x{_HEIGHT.value}')
+  print(f'  camera_id   : {_CAMERA_ID.value}')
+  print(f'  use_textures: {_USE_TEXTURES.value}')
+  print(f'  use_shadows : {_USE_SHADOWS.value}')
+  print(f'  output_dir  : {_OUTPUT_DIR.value}\n')
+
+  mx = mjx.put_model(m, impl='warp')
+
+  worldids = jp.arange(_NWORLD.value)
+
+  @jax.vmap
+  def init(worldid):
+    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
+    )
+    return dx.replace(qpos=qpos)
+
+  print('initializing data...')
+  dx_batch = jax_jit(init)(worldids)
+
+  print('running forward...')
+  dx_batch = jax_jit(
+      jax.vmap(forward.forward, in_axes=(None, 0))
+  )(mx, dx_batch)
+
+  print('creating render context...')
+  rc = io.create_render_context(
+      mjm=m,
+      nworld=_NWORLD.value,
+      cam_res=(_WIDTH.value, _HEIGHT.value),
+      use_textures=_USE_TEXTURES.value,
+      use_shadows=_USE_SHADOWS.value,
+      render_rgb=True,
+      render_depth=True,
+      enabled_geom_groups=[0, 1, 2],
+  )
+
+  print('rendering...')
+  dx_batch = jax_jit(
+      jax.vmap(
+          bvh.refit_bvh, in_axes=(None, 0, None)
+      )
+  )(mx, dx_batch, rc)
+
+  out_batch = jax_jit(
+      jax.vmap(
+          render.render, in_axes=(None, 0, None)
+      )
+  )(mx, dx_batch, rc)
+
+  rgb_packed = out_batch[0]
+  print(f'  rgb shape: {rgb_packed.shape}\n')
+
+  rgb = render_util.get_rgb(
+      rgb_packed, _CAMERA_ID.value, _WIDTH.value, _HEIGHT.value
+  )
+
+  single_path = os.path.join(
+      _OUTPUT_DIR.value, f'camera_{_CAMERA_ID.value}.png'
+  )
+  _save_single(rgb, single_path)
+
+  if _NWORLD.value > 1:
+    tiled_path = os.path.join(
+        _OUTPUT_DIR.value, f'tiled_{_CAMERA_ID.value}.png'
+    )
+    _save_tiled(rgb, tiled_path)
+
+  print('\ndone.')
+
+
+def main():
+  app.run(_main)
+
+
+if __name__ == '__main__':
+  main()