Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 871338519
Change-Id: I6ef9b4a4c3b623a87f4c43efa798b6757ac534ae

mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py

@@ -60,6 +60,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.passive import passive as passive
 from mujoco.mjx.third_party.mujoco_warp._src.ray import ray as ray
 from mujoco.mjx.third_party.mujoco_warp._src.ray import rays as rays
 from mujoco.mjx.third_party.mujoco_warp._src.render import render as render
+from mujoco.mjx.third_party.mujoco_warp._src.render_util import get_depth as get_depth
+from mujoco.mjx.third_party.mujoco_warp._src.render_util import get_rgb as get_rgb
 from mujoco.mjx.third_party.mujoco_warp._src.sensor import energy_pos as energy_pos
 from mujoco.mjx.third_party.mujoco_warp._src.sensor import energy_vel as energy_vel
 from mujoco.mjx.third_party.mujoco_warp._src.sensor import sensor_acc as sensor_acc

mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py

@@ -179,7 +179,6 @@ def _compute_bvh_bounds(
   # Data in:
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
   geom_xmat_in: wp.array2d(dtype=wp.mat33),
-  nworld_in: int,
   # In:
   bvh_ngeom: int,
   enabled_geom_ids: wp.array(dtype=int),
@@ -235,18 +234,23 @@ def compute_bvh_group_roots(
   group_root_out[tid] = root
 
 
-def build_scene_bvh(m: Model, d: Data, rc: RenderContext):
+def build_scene_bvh(mjm: mujoco.MjModel, mjd: mujoco.MjData, rc: RenderContext, nworld: int):
   """Build a global BVH for all geometries in all worlds."""
+  geom_type = wp.array(mjm.geom_type, dtype=int)
+  geom_dataid = wp.array(mjm.geom_dataid, dtype=int)
+  geom_size = wp.array(np.tile(mjm.geom_size[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3)
+  geom_xpos = wp.array(np.tile(mjd.geom_xpos[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3)
+  geom_xmat = wp.array(np.tile(mjd.geom_xmat.reshape(mjm.ngeom, 3, 3)[np.newaxis, :, :, :], (nworld, 1, 1, 1)), dtype=wp.mat33)
+
   wp.launch(
     kernel=_compute_bvh_bounds,
-    dim=(d.nworld, rc.bvh_ngeom),
+    dim=(nworld, rc.bvh_ngeom),
     inputs=[
-      m.geom_type,
-      m.geom_dataid,
-      m.geom_size,
-      d.geom_xpos,
-      d.geom_xmat,
-      d.nworld,
+      geom_type,
+      geom_dataid,
+      geom_size,
+      geom_xpos,
+      geom_xmat,
       rc.bvh_ngeom,
       rc.enabled_geom_ids,
       rc.mesh_bounds_size,
@@ -265,7 +269,7 @@ def build_scene_bvh(m: Model, d: Data, rc: RenderContext):
 
   wp.launch(
     kernel=compute_bvh_group_roots,
-    dim=d.nworld,
+    dim=nworld,
     inputs=[bvh.id],
     outputs=[rc.group_root],
   )
@@ -281,7 +285,6 @@ def refit_scene_bvh(m: Model, d: Data, rc: RenderContext):
       m.geom_size,
       d.geom_xpos,
       d.geom_xmat,
-      d.nworld,
       rc.bvh_ngeom,
       rc.enabled_geom_ids,
       rc.mesh_bounds_size,
@@ -847,13 +850,18 @@ def _update_flex_face_points(
 
 
 def build_flex_bvh(
-  mjm: mujoco.MjModel, m: Model, d: Data, constructor: str = "sah", leaf_size: int = 2
+  mjm: mujoco.MjModel, mjd: mujoco.MjData, nworld: int, constructor: str = "sah", leaf_size: int = 2
 ) -> tuple[wp.Mesh, wp.array, wp.array, wp.array, wp.array, wp.array, int]:
   """Create a Warp mesh BVH from flex data."""
   if (mjm.flex_dim == 1).any():
     raise ValueError("1D Flex objects are not currently supported.")
 
   nflex = mjm.nflex
+  nflexvert = mjm.nflexvert
+  nflexelemdata = len(mjm.flex_elem)
+
+  flex_elem = wp.array(mjm.flex_elem, dtype=int)
+  flexvert_xpos = wp.array(np.tile(mjd.flexvert_xpos[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3)
 
   flex_faceadr = [0]
   for f in range(nflex):
@@ -865,23 +873,23 @@ def build_flex_bvh(
   nface = int(flex_faceadr[-1])
   flex_faceadr = flex_faceadr[:-1]
 
-  face_point = wp.zeros(nface * 3 * d.nworld, dtype=wp.vec3)
-  face_index = wp.zeros(nface * 3 * d.nworld, dtype=wp.int32)
-  group = wp.zeros(nface * d.nworld, dtype=int)
+  face_point = wp.empty(nface * 3 * nworld, dtype=wp.vec3)
+  face_index = wp.empty(nface * 3 * nworld, dtype=wp.int32)
+  group = wp.empty(nface * nworld, dtype=int)
 
-  flexvert_norm = wp.zeros(d.flexvert_xpos.shape, dtype=wp.vec3)
+  flexvert_norm = wp.zeros((nworld, nflexvert), dtype=wp.vec3)
   flex_shell = wp.array(mjm.flex_shell, dtype=int)
 
   wp.launch(
     kernel=accumulate_flex_vertex_normals,
-    dim=(d.nworld, m.nflexelemdata // 3),
-    inputs=[m.flex_elem, d.flexvert_xpos],
+    dim=(nworld, nflexelemdata // 3),
+    inputs=[flex_elem, flexvert_xpos],
     outputs=[flexvert_norm],
   )
 
   wp.launch(
     kernel=normalize_vertex_normals,
-    dim=(d.nworld, m.nflexvert),
+    dim=(nworld, nflexvert),
     inputs=[flexvert_norm],
   )
 
@@ -896,10 +904,10 @@ def build_flex_bvh(
     if dim == 2:
       wp.launch(
         kernel=_build_flex_2d_elements,
-        dim=(d.nworld, nelem),
+        dim=(nworld, nelem),
         inputs=[
-          m.flex_elem,
-          d.flexvert_xpos,
+          flex_elem,
+          flexvert_xpos,
           flexvert_norm,
           elem_adr,
           vert_adr,
@@ -912,9 +920,9 @@ def build_flex_bvh(
 
       wp.launch(
         kernel=_build_flex_2d_sides,
-        dim=(d.nworld, nshell),
+        dim=(nworld, nshell),
         inputs=[
-          d.flexvert_xpos,
+          flexvert_xpos,
           flexvert_norm,
           flex_shell,
           shell_adr,
@@ -928,9 +936,9 @@ def build_flex_bvh(
     elif dim == 3:
       wp.launch(
         kernel=_build_flex_3d_shells,
-        dim=(d.nworld, nshell),
+        dim=(nworld, nshell),
         inputs=[
-          d.flexvert_xpos,
+          flexvert_xpos,
           flex_shell,
           shell_adr,
           vert_adr,
@@ -948,10 +956,10 @@ def build_flex_bvh(
     bvh_leaf_size=leaf_size,
   )
 
-  group_root = wp.zeros(d.nworld, dtype=int)
+  group_root = wp.empty(nworld, dtype=int)
   wp.launch(
     kernel=compute_bvh_group_roots,
-    dim=d.nworld,
+    dim=nworld,
     inputs=[flex_mesh.id],
     outputs=[group_root],
   )

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py

@@ -24,6 +24,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import contact_params
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import geom_collision_pair
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import write_contact
+from mujoco.mjx.third_party.mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
 from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
 from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAFACES
@@ -91,7 +92,10 @@ def _hfield_filter(
   r2 = geom_rbound[rbound_id, g2]
 
   # TODO(team): margin?
-  margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
+  if BLEEDING_EDGE_MUJOCO:
+    margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
+  else:
+    margin = wp.max(geom_margin[margin_id, g1], geom_margin[margin_id, g2])
 
   # box-sphere test: horizontal plane
   for i in range(2):

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py

@@ -20,6 +20,7 @@ import warp as wp
 from mujoco.mjx.third_party.mujoco_warp._src.collision_convex import convex_narrowphase
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import primitive_narrowphase
 from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf_narrowphase
+from mujoco.mjx.third_party.mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
 from mujoco.mjx.third_party.mujoco_warp._src.math import upper_tri_index
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseFilter
@@ -101,18 +102,27 @@ def _plane_filter(
   if size1 == 0.0:
     # geom1 is a plane
     dist = wp.dot(xpos2 - xpos1, wp.vec3(xmat1[0, 2], xmat1[1, 2], xmat1[2, 2]))
-    return dist <= size2 + margin1 + margin2
+    if BLEEDING_EDGE_MUJOCO:
+      return dist <= size2 + margin1 + margin2
+    else:
+      return dist <= size2 + wp.max(margin1, margin2)
   elif size2 == 0.0:
     # geom2 is a plane
     dist = wp.dot(xpos1 - xpos2, wp.vec3(xmat2[0, 2], xmat2[1, 2], xmat2[2, 2]))
-    return dist <= size1 + margin1 + margin2
+    if BLEEDING_EDGE_MUJOCO:
+      return dist <= size1 + margin1 + margin2
+    else:
+      return dist <= size1 + wp.max(margin1, margin2)
 
   return True
 
 
 @wp.func
 def _sphere_filter(size1: float, size2: float, margin1: float, margin2: float, xpos1: wp.vec3, xpos2: wp.vec3) -> bool:
-  bound = size1 + size2 + margin1 + margin2
+  if BLEEDING_EDGE_MUJOCO:
+    bound = size1 + size2 + margin1 + margin2
+  else:
+    bound = size1 + size2 + wp.max(margin1, margin2)
   dif = xpos2 - xpos1
   dist_sq = wp.dot(dif, dif)
   return dist_sq <= bound * bound
@@ -141,7 +151,10 @@ def _aabb_filter(
   center1 = xmat1 @ center1 + xpos1
   center2 = xmat2 @ center2 + xpos2
 
-  margin = margin1 + margin2
+  if BLEEDING_EDGE_MUJOCO:
+    margin = margin1 + margin2
+  else:
+    margin = wp.max(margin1, margin2)
 
   max_x1 = -MJ_MAXVAL
   max_y1 = -MJ_MAXVAL
@@ -236,7 +249,10 @@ def _obb_filter(
   xmat2: wp.mat33,
 ) -> bool:
   """Oriented bounding boxes collision (see Gottschalk et al.), see mj_collideOBB."""
-  margin = margin1 + margin2
+  if BLEEDING_EDGE_MUJOCO:
+    margin = margin1 + margin2
+  else:
+    margin = wp.max(margin1, margin2)
 
   xcenter = mat23()
   normal = mat63()

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py

@@ -29,6 +29,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sph
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sphere_capsule
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sphere_cylinder
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sphere_sphere
+from mujoco.mjx.third_party.mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
 from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
 from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
 from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index
@@ -552,8 +553,12 @@ def contact_params(
     solreffriction = wp.vec2(0.0, 0.0)
     solimp = mix * geom_solimp[solimp_id, g1] + (1.0 - mix) * geom_solimp[solimp_id, g2]
     # geom priority is ignored
-    margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
-    gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2]
+    if BLEEDING_EDGE_MUJOCO:
+      margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
+      gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2]
+    else:
+      margin = wp.max(geom_margin[margin_id, g1], geom_margin[margin_id, g2])
+      gap = wp.max(geom_gap[gap_id, g1], geom_gap[gap_id, g2])
 
   friction = vec5(
     wp.max(MJ_MINMU, friction[0]),

mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py

@@ -37,11 +37,14 @@ def _qderiv_actuator_passive_vel(
   actuator_biastype: wp.array(dtype=int),
   actuator_actadr: wp.array(dtype=int),
   actuator_actnum: wp.array(dtype=int),
+  actuator_forcelimited: wp.array(dtype=bool),
   actuator_gainprm: wp.array2d(dtype=vec10f),
   actuator_biasprm: wp.array2d(dtype=vec10f),
+  actuator_forcerange: wp.array2d(dtype=wp.vec2),
   # Data in:
   act_in: wp.array2d(dtype=float),
   ctrl_in: wp.array2d(dtype=float),
+  actuator_force_in: wp.array2d(dtype=float),
   # Out:
   vel_out: wp.array2d(dtype=float),
 ):
@@ -64,6 +67,14 @@ def _qderiv_actuator_passive_vel(
     vel_out[worldid, actid] = 0.0
     return
 
+  # skip if force is clamped by forcerange
+  if actuator_forcelimited[actid]:
+    force = actuator_force_in[worldid, actid]
+    forcerange = actuator_forcerange[worldid % actuator_forcerange.shape[0], actid]
+    if force <= forcerange[0] or force >= forcerange[1]:
+      vel_out[worldid, actid] = 0.0
+      return
+
   vel = float(bias)
   if actuator_dyntype[actid] != DynType.NONE:
     if gain != 0.0:
@@ -77,12 +88,21 @@ def _qderiv_actuator_passive_vel(
   vel_out[worldid, actid] = vel
 
 
+@wp.func
+def _nonzero_mask(x: float) -> float:
+  """Returns 1.0 for non-zero input, 0.0 otherwise."""
+  if x != 0.0:
+    return 1.0
+  return 0.0
+
+
 @cache_kernel
 def _qderiv_actuator_passive_actuation_dense(tile: TileSet, nu: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Data in:
     actuator_moment_in: wp.array3d(dtype=float),
+    qM_in: wp.array3d(dtype=float),
     # In:
     vel_in: wp.array3d(dtype=float),
     adr: wp.array(dtype=int),
@@ -98,6 +118,16 @@ def _qderiv_actuator_passive_actuation_dense(tile: TileSet, nu: int):
     moment_tile = wp.tile_load(actuator_moment_in[worldid], shape=(NU, TILE_SIZE), offset=(0, dofid), bounds_check=False)
     moment_weighted = wp.tile_map(wp.mul, wp.tile_broadcast(vel_tile, shape=(NU, TILE_SIZE)), moment_tile)
     qderiv_tile = wp.tile_matmul(wp.tile_transpose(moment_tile), moment_weighted)
+
+    # Mask out cross-terms for DOF pairs that are structurally zero in M
+    # (e.g., sibling DOFs coupled only through tendons).  Without this,
+    # stale actuation values at sibling positions make A = M - dt*qDeriv
+    # non-positive-definite, causing the tiled Cholesky to produce NaN.
+    # Dropping these terms matches MuJoCo CPU's implicitfast approximation.
+    qM_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid), bounds_check=False)
+    mask_tile = wp.tile_map(_nonzero_mask, qM_tile)
+    qderiv_tile = wp.tile_map(wp.mul, qderiv_tile, mask_tile)
+
     wp.tile_store(qDeriv_out[worldid], qderiv_tile, offset=(dofid, dofid), bounds_check=False)
 
   return kernel
@@ -237,10 +267,13 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
           m.actuator_biastype,
           m.actuator_actadr,
           m.actuator_actnum,
+          m.actuator_forcelimited,
           m.actuator_gainprm,
           m.actuator_biasprm,
+          m.actuator_forcerange,
           d.act,
           d.ctrl,
+          d.actuator_force,
         ],
         outputs=[vel],
       )
@@ -257,7 +290,7 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
           wp.launch_tiled(
             _qderiv_actuator_passive_actuation_dense(tile, m.nu),
             dim=(d.nworld, tile.adr.size),
-            inputs=[d.actuator_moment, vel_3d, tile.adr],
+            inputs=[d.actuator_moment, d.qM, vel_3d, tile.adr],
             outputs=[out],
             block_dim=m.block_dim.qderiv_actuator_dense,
           )

mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py

@@ -2090,6 +2090,13 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any
         if val.upper() not in ("TRUE", "FALSE"):
           raise ValueError(f"Unrecognized value for field: {key}")
         val = val.upper() == "TRUE"
+      elif typ is wp.array and isinstance(val, str):
+        arr = getattr(obj, attr)
+        floats = [float(p) for p in val.strip("[]").split()]
+        val = wp.array([arr.dtype(*floats)], dtype=arr.dtype)
+      elif typ is np.ndarray and isinstance(val, str):
+        arr = getattr(obj, attr)
+        val = np.array([float(p) for p in val.strip("[]").split()], dtype=arr.dtype)
       else:
         val = typ(val)
 
@@ -2156,8 +2163,7 @@ def _build_rays(
 
 def create_render_context(
   mjm: mujoco.MjModel,
-  m: types.Model,
-  d: types.Data,
+  nworld: int = 1,
   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,
@@ -2166,13 +2172,13 @@ def create_render_context(
   enabled_geom_groups: list[int] = [0, 1, 2],
   cam_active: list[bool] | None = None,
   flex_render_smooth: bool = True,
+  use_precomputed_rays: bool = True,
 ) -> types.RenderContext:
   """Creates a render context on device.
 
   Args:
     mjm: The model containing kinematic and dynamic information on host.
-    m: The model on device.
-    d: The data on device.
+    nworld: The number of worlds.
     cam_res: The width and height to render each camera image. If None, uses the
              MuJoCo model values.
     render_rgb: Whether to render RGB images. If None, uses the MuJoCo model values.
@@ -2183,10 +2189,15 @@ def create_render_context(
     cam_active: List of booleans indicating which cameras to include in rendering.
                 If None, all cameras are included.
     flex_render_smooth: Whether to render flex meshes smoothly.
+    use_precomputed_rays: Use precomputed rays instead of computing during rendering.
+                          When using domain randomization for camera intrinsics, set to False.
 
   Returns:
     The render context containing rendering fields and output arrays on device.
   """
+  mjd = mujoco.MjData(mjm)
+  mujoco.mj_forward(mjm, mjd)
+
   # TODO(team): remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
   if use_textures and not hasattr(wp, "Texture2D"):
     warnings.warn("Textures require warp >= 1.12. Disabling textures.")
@@ -2231,7 +2242,7 @@ def create_render_context(
 
   # Flex BVHs
   flex_bvh_id = wp.uint64(0)
-  flex_group_root = wp.zeros(d.nworld, dtype=int)
+  flex_group_root = wp.zeros(nworld, dtype=int)
   flex_mesh = None
   flex_face_point = None
   flex_elemdataadr = None
@@ -2252,7 +2263,7 @@ def create_render_context(
       flex_shell_data,
       flex_faceadr_data,
       flex_nface,
-    ) = bvh.build_flex_bvh(mjm, m, d)
+    ) = bvh.build_flex_bvh(mjm, mjd, nworld)
 
     flex_mesh = fmesh
     flex_bvh_id = fmesh.id
@@ -2352,31 +2363,33 @@ def create_render_context(
 
   znear = mjm.vis.map.znear * mjm.stat.extent
 
-  if m.cam_fovy.shape[0] > 1 or m.cam_intrinsic.shape[0] > 1:
-    ray = None
-  else:
-    ray = wp.zeros(int(total), dtype=wp.vec3)
+  ray = wp.zeros(int(total), dtype=wp.vec3)
 
-    offset = 0
-    for idx, cam_id in enumerate(active_cam_indices):
-      img_w = cam_res_np[idx][0]
-      img_h = cam_res_np[idx][1]
-      wp.launch(
-        kernel=_build_rays,
-        dim=(img_w, img_h),
-        inputs=[
-          offset,
-          img_w,
-          img_h,
-          m.cam_projection.numpy()[cam_id].item(),
-          m.cam_fovy.numpy()[0, cam_id].item(),
-          wp.vec2(m.cam_sensorsize.numpy()[cam_id]),
-          wp.vec4(m.cam_intrinsic.numpy()[0, cam_id]),
-          znear,
-        ],
-        outputs=[ray],
-      )
-      offset += img_w * img_h
+  # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml
+  cam_projection = np.zeros(mjm.ncam, dtype=int)
+  if BLEEDING_EDGE_MUJOCO:
+    cam_projection = mjm.cam_projection
+
+  offset = 0
+  for idx, cam_id in enumerate(active_cam_indices):
+    img_w = cam_res_np[idx][0]
+    img_h = cam_res_np[idx][1]
+    wp.launch(
+      kernel=_build_rays,
+      dim=(img_w, img_h),
+      inputs=[
+        offset,
+        img_w,
+        img_h,
+        int(cam_projection[cam_id]),
+        float(mjm.cam_fovy[cam_id]),
+        wp.vec2(mjm.cam_sensorsize[cam_id]),
+        wp.vec4(mjm.cam_intrinsic[cam_id]),
+        znear,
+      ],
+      outputs=[ray],
+    )
+    offset += img_w * img_h
 
   bvh_ngeom = len(geom_enabled_idx)
 
@@ -2387,6 +2400,7 @@ def create_render_context(
     use_textures=use_textures,
     use_shadows=use_shadows,
     background_color=render_util.pack_rgba_to_uint32(0.1 * 255.0, 0.1 * 255.0, 0.2 * 255.0, 1.0 * 255.0),
+    use_precomputed_rays=use_precomputed_rays,
     bvh_ngeom=bvh_ngeom,
     enabled_geom_ids=wp.array(geom_enabled_idx, dtype=int),
     mesh_registry=mesh_registry,
@@ -2417,14 +2431,14 @@ def create_render_context(
     flex_render_smooth=flex_render_smooth,
     bvh=None,
     bvh_id=None,
-    lower=wp.zeros(d.nworld * bvh_ngeom, dtype=wp.vec3),
-    upper=wp.zeros(d.nworld * bvh_ngeom, dtype=wp.vec3),
-    group=wp.zeros(d.nworld * bvh_ngeom, dtype=int),
-    group_root=wp.zeros(d.nworld, dtype=int),
+    lower=wp.zeros(nworld * bvh_ngeom, dtype=wp.vec3),
+    upper=wp.zeros(nworld * bvh_ngeom, dtype=wp.vec3),
+    group=wp.zeros(nworld * bvh_ngeom, dtype=int),
+    group_root=wp.zeros(nworld, dtype=int),
     ray=ray,
-    rgb_data=wp.zeros((d.nworld, ri), dtype=wp.uint32),
+    rgb_data=wp.zeros((nworld, ri), dtype=wp.uint32),
     rgb_adr=wp.array(rgb_adr, dtype=int),
-    depth_data=wp.zeros((d.nworld, di), dtype=wp.float32),
+    depth_data=wp.zeros((nworld, di), dtype=wp.float32),
     depth_adr=wp.array(depth_adr, dtype=int),
     render_rgb=wp.array(render_rgb, dtype=bool),
     render_depth=wp.array(render_depth, dtype=bool),
@@ -2432,6 +2446,6 @@ def create_render_context(
     total_rays=int(total),
   )
 
-  bvh.build_scene_bvh(m, d, rc)
+  bvh.build_scene_bvh(mjm, mjd, rc, nworld)
 
   return rc

mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py

@@ -495,7 +495,9 @@ def render(m: Model, d: Data, rc: RenderContext):
     # Map active camera index to MuJoCo camera ID
     mujoco_cam_id = cam_id_map[cam_idx]
 
-    if wp.static(rc.ray is None):
+    if wp.static(rc.use_precomputed_rays):
+      ray_dir_local_cam = ray[ray_idx]
+    else:
       img_w = cam_res[cam_idx][0]
       img_h = cam_res[cam_idx][1]
       px = ray_idx_local % img_w
@@ -511,8 +513,6 @@ def render(m: Model, d: Data, rc: RenderContext):
         py,
         wp.static(rc.znear),
       )
-    else:
-      ray_dir_local_cam = ray[ray_idx]
 
     ray_dir_world = cam_xmat_in[world_idx, mujoco_cam_id] @ ray_dir_local_cam
     ray_origin_world = cam_xpos_in[world_idx, mujoco_cam_id]

mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py

@@ -17,6 +17,7 @@ import mujoco
 import warp as wp
 
 from mujoco.mjx.third_party.mujoco_warp._src.types import ProjectionType
+from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
 
 wp.set_module_options({"enable_backward": False})
 
@@ -128,3 +129,80 @@ def compute_ray(
 def pack_rgba_to_uint32(r: float, g: float, b: float, a: float) -> wp.uint32:
   """Pack RGBA values into a single uint32 for efficient memory access."""
   return wp.uint32((int(a) << int(24)) | (int(r) << int(16)) | (int(g) << int(8)) | int(b))
+
+
+@wp.kernel
+def unpack_rgb_kernel(
+  # In:
+  packed: wp.array2d(dtype=wp.uint32),
+  rgb_adr: wp.array(dtype=int),
+  camera_index: int,
+  # Out:
+  rgb_out: wp.array3d(dtype=wp.vec3),
+):
+  """Unpack ABGR uint32 packed pixel data into separate R, G, and B channels."""
+  worldid, pixelid = wp.tid()
+
+  xid = pixelid % rgb_out.shape[2]
+  yid = pixelid // rgb_out.shape[2]
+
+  rgb_adr_offset = rgb_adr[camera_index]
+  val = packed[worldid, rgb_adr_offset + pixelid]
+  b = wp.float32(val & wp.uint32(0xFF)) * wp.static(1.0 / 255.0)
+  g = wp.float32((val >> wp.uint32(8)) & wp.uint32(0xFF)) * wp.static(1.0 / 255.0)
+  r = wp.float32((val >> wp.uint32(16)) & wp.uint32(0xFF)) * wp.static(1.0 / 255.0)
+  rgb_out[worldid, yid, xid] = wp.vec3(r, g, b)
+
+
+@wp.kernel
+def extract_depth_kernel(
+  # In:
+  depth_data: wp.array2d(dtype=float),
+  depth_adr: wp.array(dtype=int),
+  camera_index: int,
+  depth_scale: float,
+  # Out:
+  depth_out: wp.array3d(dtype=float),
+):
+  """Extract the depth data from the render context buffers for a given camera index."""
+  worldid, pixelid = wp.tid()
+  xid = pixelid % depth_out.shape[2]
+  yid = pixelid // depth_out.shape[2]
+
+  depth_adr_offset = depth_adr[camera_index]
+  val = depth_data[worldid, depth_adr_offset + pixelid]
+  depth_out[worldid, yid, xid] = wp.clamp(val / depth_scale, 0.0, 1.0)
+
+
+def get_rgb(rc: RenderContext, camera_index: int, rgb_out: wp.array3d(dtype=wp.vec3)):
+  """Get the RGB data output from the render context buffers for a given camera index.
+
+  Args:
+    rc: The render context on device.
+    camera_index: The index of the camera to get the RGB data for.
+    rgb_out: The output array to store the RGB data in, with shape (nworld, height, width).
+  """
+  wp.launch(
+    unpack_rgb_kernel,
+    dim=(rgb_out.shape[0], rgb_out.shape[1] * rgb_out.shape[2]),
+    inputs=[rc.rgb_data, rc.rgb_adr, camera_index],
+    outputs=[rgb_out],
+  )
+
+
+def get_depth(rc: RenderContext, camera_index: int, depth_scale: float, depth_out: wp.array3d(dtype=float)):
+  """Get the depth data output from the render context buffers for a given camera index.
+
+  Args:
+    rc: The render context on device.
+    camera_index: The index of the camera to get the depth data for.
+    depth_scale: The scale factor to apply to the depth data.
+    depth_out: The output array to store the scaled and clamped depth data in
+      with shape (nworld, height, width).
+  """
+  wp.launch(
+    extract_depth_kernel,
+    dim=(depth_out.shape[0], depth_out.shape[1] * depth_out.shape[2]),
+    inputs=[rc.depth_data, rc.depth_adr, camera_index, depth_scale],
+    outputs=[depth_out],
+  )

mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py

@@ -1781,6 +1781,7 @@ class RenderContext:
     cam_id_map: camera id map
     use_textures: whether to use textures
     use_shadows: whether to use shadows
+    use_precomputed_rays: whether to use precomputed rays
     bvh_ngeom: number of geometries in the BVH
     enabled_geom_ids: enabled geometry ids
     mesh_registry: mesh BVH id to warp mesh mapping
@@ -1834,6 +1835,7 @@ class RenderContext:
   use_textures: bool
   use_shadows: bool
   background_color: wp.uint32
+  use_precomputed_rays: bool
   bvh_ngeom: int
   enabled_geom_ids: array("*", int)
   mesh_registry: dict

mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name="mujoco-warp"
-version = "0.0.2"
+version = "3.5.0"
 # TODO(team): create a distribution list
 authors = [
     {name = "Newton Developers", email = "mujoco@deepmind.com"},

mjx/mujoco/mjx/warp/collision_driver.py

@@ -44,7 +44,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _collision_shim(
     # Model

mjx/mujoco/mjx/warp/forward.py

@@ -44,7 +44,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _forward_shim(
     # Model

mjx/mujoco/mjx/warp/io.py

@@ -30,23 +30,7 @@ def create_render_context(
     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 = mjw.create_render_context(mjm=mjm, nworld=nworld, **kwargs)
   rc.rgb_data_shape = rc.rgb_data.shape
   rc.depth_data_shape = rc.depth_data.shape
   rc.rgb_data = None

mjx/mujoco/mjx/warp/smooth.py

@@ -44,7 +44,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _kinematics_shim(
     # Model