Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 868759523
Change-Id: I83a9f014840e66d7775a3684aa2e59ec5a077191

mjx/mujoco/mjx/_src/types.py

@@ -871,10 +871,10 @@ class Model(PyTreeNode):
   cam_poscom0: jax.Array
   cam_pos0: jax.Array
   cam_mat0: jax.Array
-  cam_fovy: np.ndarray
+  cam_fovy: jax.Array
   cam_resolution: np.ndarray
   cam_sensorsize: np.ndarray
-  cam_intrinsic: np.ndarray
+  cam_intrinsic: jax.Array
   light_mode: np.ndarray
   light_type: jax.Array
   light_castshadow: jax.Array

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

@@ -28,6 +28,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Model as Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data as Data
 # isort: on
 
+from mujoco.mjx.third_party.mujoco_warp._src.bvh import refit_bvh as refit_bvh
 from mujoco.mjx.third_party.mujoco_warp._src.collision_driver import collision as collision
 from mujoco.mjx.third_party.mujoco_warp._src.collision_driver import nxn_broadphase as nxn_broadphase
 from mujoco.mjx.third_party.mujoco_warp._src.collision_driver import sap_broadphase as sap_broadphase
@@ -46,14 +47,19 @@ from mujoco.mjx.third_party.mujoco_warp._src.forward import rungekutta4 as runge
 from mujoco.mjx.third_party.mujoco_warp._src.forward import step1 as step1
 from mujoco.mjx.third_party.mujoco_warp._src.forward import step2 as step2
 from mujoco.mjx.third_party.mujoco_warp._src.inverse import inverse as inverse
+from mujoco.mjx.third_party.mujoco_warp._src.io import create_render_context as create_render_context
 from mujoco.mjx.third_party.mujoco_warp._src.io import get_data_into as get_data_into
 from mujoco.mjx.third_party.mujoco_warp._src.io import make_data as make_data
 from mujoco.mjx.third_party.mujoco_warp._src.io import put_data as put_data
 from mujoco.mjx.third_party.mujoco_warp._src.io import put_model as put_model
 from mujoco.mjx.third_party.mujoco_warp._src.io import reset_data as reset_data
+from mujoco.mjx.third_party.mujoco_warp._src.io import set_const as set_const
+from mujoco.mjx.third_party.mujoco_warp._src.io import set_const_0 as set_const_0
+from mujoco.mjx.third_party.mujoco_warp._src.io import set_const_fixed as set_const_fixed
 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.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
@@ -75,6 +81,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.smooth import transmission as trans
 from mujoco.mjx.third_party.mujoco_warp._src.solver import solve as solve
 from mujoco.mjx.third_party.mujoco_warp._src.support import contact_force as contact_force
 from mujoco.mjx.third_party.mujoco_warp._src.support import get_state as get_state
+from mujoco.mjx.third_party.mujoco_warp._src.support import jac as jac
 from mujoco.mjx.third_party.mujoco_warp._src.support import mul_m as mul_m
 from mujoco.mjx.third_party.mujoco_warp._src.support import set_state as set_state
 from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate as xfrc_accumulate
@@ -92,6 +99,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType as GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType as IntegratorType
 from mujoco.mjx.third_party.mujoco_warp._src.types import JointType as JointType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Option as Option
+from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext as RenderContext
 from mujoco.mjx.third_party.mujoco_warp._src.types import SolverType as SolverType
 from mujoco.mjx.third_party.mujoco_warp._src.types import State as State
 from mujoco.mjx.third_party.mujoco_warp._src.types import Statistic as Statistic

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

@@ -16,7 +16,7 @@
 """Utilities for benchmarking MuJoCo Warp."""
 
 import time
-from typing import Callable, Optional, Tuple
+from typing import Callable, Tuple
 
 import numpy as np
 import warp as wp
@@ -24,6 +24,7 @@ import warp as wp
 from mujoco.mjx.third_party.mujoco_warp._src import warp_util
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
 from mujoco.mjx.third_party.mujoco_warp._src.util_misc import halton
 
 
@@ -87,10 +88,11 @@ def benchmark(
   m: Model,
   d: Data,
   nstep: int,
-  ctrls: Optional[np.ndarray] = None,
+  ctrls: np.ndarray | None = None,
   event_trace: bool = False,
   measure_alloc: bool = False,
   measure_solver_niter: bool = False,
+  render_context: RenderContext | None = None,
 ) -> Tuple[float, float, dict, list, list, list, int]:
   """Benchmark a function of Model and Data.
 
@@ -103,6 +105,7 @@ def benchmark(
     event_trace: If True, time routines decorated with @event_scope.
     measure_alloc: If True, record number of contacts and constraints.
     measure_solver_niter: If True, record the number of solver iterations.
+    render_context: The render context to use for rendering.
 
   Returns:
     - Time to JIT fn.
@@ -120,8 +123,14 @@ def benchmark(
   with warp_util.EventTracer(enabled=event_trace) as tracer:
     # capture the whole function as a CUDA graph
     jit_beg = time.perf_counter()
-    with wp.ScopedCapture() as capture:
+
-      fn(m, d)
+    if render_context is not None:
+      with wp.ScopedCapture() as capture:
+        fn(m, d, render_context)
+    else:
+      with wp.ScopedCapture() as capture:
+        fn(m, d)
+
     jit_end = time.perf_counter()
     jit_duration = jit_end - jit_beg
 

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

@@ -29,6 +29,8 @@ 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
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAHORIZON
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
@@ -38,7 +40,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 # TODO(team): improve compile time to enable backward pass
 wp.set_module_options({"enable_backward": False})
@@ -46,42 +47,6 @@ wp.set_module_options({"enable_backward": False})
 vec_maxconpair = wp.types.vector(length=MJ_MAXCONPAIR, dtype=float)
 mat_maxconpair = wp.types.matrix(shape=(MJ_MAXCONPAIR, 3), dtype=float)
 
-_CONVEX_COLLISION_PAIRS = [
-  (GeomType.HFIELD, GeomType.SPHERE),
-  (GeomType.HFIELD, GeomType.CAPSULE),
-  (GeomType.HFIELD, GeomType.ELLIPSOID),
-  (GeomType.HFIELD, GeomType.CYLINDER),
-  (GeomType.HFIELD, GeomType.BOX),
-  (GeomType.HFIELD, GeomType.MESH),
-  (GeomType.SPHERE, GeomType.ELLIPSOID),
-  (GeomType.SPHERE, GeomType.MESH),
-  (GeomType.CAPSULE, GeomType.ELLIPSOID),
-  (GeomType.CAPSULE, GeomType.CYLINDER),
-  (GeomType.CAPSULE, GeomType.MESH),
-  (GeomType.ELLIPSOID, GeomType.ELLIPSOID),
-  (GeomType.ELLIPSOID, GeomType.CYLINDER),
-  (GeomType.ELLIPSOID, GeomType.BOX),
-  (GeomType.ELLIPSOID, GeomType.MESH),
-  (GeomType.CYLINDER, GeomType.CYLINDER),
-  (GeomType.CYLINDER, GeomType.BOX),
-  (GeomType.CYLINDER, GeomType.MESH),
-  (GeomType.BOX, GeomType.MESH),
-  (GeomType.MESH, GeomType.MESH),
-]
-
-
-def _check_convex_collision_pairs():
-  prev_idx = -1
-  for pair in _CONVEX_COLLISION_PAIRS:
-    idx = upper_trid_index(len(GeomType), pair[0].value, pair[1].value)
-    if pair[1] < pair[0] or idx <= prev_idx:
-      return False
-    prev_idx = idx
-  return True
-
-
-assert _check_convex_collision_pairs(), "_CONVEX_COLLISION_PAIRS is in invalid order."
-
 
 @wp.func
 def _hfield_filter(
@@ -91,6 +56,11 @@ def _hfield_filter(
   geom_size: wp.array2d(dtype=wp.vec3),
   geom_rbound: wp.array2d(dtype=float),
   geom_margin: wp.array2d(dtype=float),
+  mesh_vertadr: wp.array(dtype=int),
+  mesh_vertnum: wp.array(dtype=int),
+  mesh_graphadr: wp.array(dtype=int),
+  mesh_vert: wp.array(dtype=wp.vec3),
+  mesh_graph: wp.array(dtype=int),
   hfield_size: wp.array(dtype=wp.vec4),
   # Data in:
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
@@ -126,14 +96,14 @@ def _hfield_filter(
   # box-sphere test: horizontal plane
   for i in range(2):
     if (size1[i] < pos[i] - r2 - margin) or (-size1[i] > pos[i] + r2 + margin):
-      return True, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf
+      return True, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL
 
   # box-sphere test: vertical direction
   if size1[2] < pos[2] - r2 - margin:  # up
-    return True, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf
+    return True, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL
 
   if -size1[3] > pos[2] + r2 + margin:  # down
-    return True, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf
+    return True, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL
 
   mat2 = geom_xmat_in[worldid, g2]
   mat = mat1T @ mat2
@@ -148,6 +118,15 @@ def _hfield_filter(
 
   geomtype2 = geom_type[g2]
 
+  # load mesh vertex data for support function queries
+  if geomtype2 == GeomType.MESH:
+    dataid = geom_dataid[g2]
+    geom2.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1)
+    geom2.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1)
+    geom2.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1)
+    geom2.vert = mesh_vert
+    geom2.graph = mesh_graph
+
   # use support functions for tight AABB bounds
   xmax = support(geom2, geomtype2, wp.vec3(1.0, 0.0, 0.0)).point[0]
   xmin = support(geom2, geomtype2, wp.vec3(-1.0, 0.0, 0.0)).point[0]
@@ -165,7 +144,7 @@ def _hfield_filter(
     or (zmin - margin > size1[2])
     or (zmax + margin < -size1[3])
   ):
-    return True, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf
+    return True, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL
   else:
     return False, xmin, xmax, ymin, ymax, zmin, zmax
 
@@ -176,11 +155,12 @@ def ccd_hfield_kernel_builder(
   geomtype2: int,
   gjk_iterations: int,
   epa_iterations: int,
+  geomgeomid: int,
 ):
   """Kernel builder for heightfield CCD collisions (no multiccd args)."""
 
   # runs convex collision on a set of geom pairs to recover contact info
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def ccd_hfield_kernel(
     # Model:
     opt_ccd_tolerance: wp.array(dtype=float),
@@ -196,15 +176,10 @@ def ccd_hfield_kernel_builder(
     geom_friction: wp.array2d(dtype=wp.vec3),
     geom_margin: wp.array2d(dtype=float),
     geom_gap: wp.array2d(dtype=float),
-    hfield_adr: wp.array(dtype=int),
-    hfield_nrow: wp.array(dtype=int),
-    hfield_ncol: wp.array(dtype=int),
-    hfield_size: wp.array(dtype=wp.vec4),
-    hfield_data: wp.array(dtype=float),
     mesh_vertadr: wp.array(dtype=int),
     mesh_vertnum: wp.array(dtype=int),
-    mesh_vert: wp.array(dtype=wp.vec3),
     mesh_graphadr: wp.array(dtype=int),
+    mesh_vert: wp.array(dtype=wp.vec3),
     mesh_graph: wp.array(dtype=int),
     mesh_polynum: wp.array(dtype=int),
     mesh_polyadr: wp.array(dtype=int),
@@ -215,6 +190,11 @@ def ccd_hfield_kernel_builder(
     mesh_polymapadr: wp.array(dtype=int),
     mesh_polymapnum: wp.array(dtype=int),
     mesh_polymap: wp.array(dtype=int),
+    hfield_size: wp.array(dtype=wp.vec4),
+    hfield_nrow: wp.array(dtype=int),
+    hfield_ncol: wp.array(dtype=int),
+    hfield_adr: wp.array(dtype=int),
+    hfield_data: wp.array(dtype=float),
     pair_dim: wp.array(dtype=int),
     pair_solref: wp.array2d(dtype=wp.vec2),
     pair_solreffriction: wp.array2d(dtype=wp.vec2),
@@ -223,24 +203,23 @@ def ccd_hfield_kernel_builder(
     pair_gap: wp.array2d(dtype=float),
     pair_friction: wp.array2d(dtype=vec5),
     # Data in:
-    naconmax_in: int,
     geom_xpos_in: wp.array2d(dtype=wp.vec3),
     geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    naconmax_in: int,
+    naccdmax_in: int,
+    ncollision_in: wp.array(dtype=int),
+    # In:
     collision_pair_in: wp.array(dtype=wp.vec2i),
     collision_pairid_in: wp.array(dtype=wp.vec2i),
     collision_worldid_in: wp.array(dtype=int),
-    ncollision_in: wp.array(dtype=int),
+    epa_vert_in: wp.array2d(dtype=wp.vec3),
-    # In:
+    epa_vert_index_in: wp.array2d(dtype=int),
-    epa_vert1_in: wp.array2d(dtype=wp.vec3),
-    epa_vert2_in: wp.array2d(dtype=wp.vec3),
-    epa_vert_index1_in: wp.array2d(dtype=int),
-    epa_vert_index2_in: wp.array2d(dtype=int),
     epa_face_in: wp.array2d(dtype=int),
     epa_pr_in: wp.array2d(dtype=wp.vec3),
     epa_norm2_in: wp.array2d(dtype=float),
     epa_horizon_in: wp.array2d(dtype=int),
+    nccd_in: wp.array(dtype=int),
     # Data out:
-    nacon_out: wp.array(dtype=int),
     contact_dist_out: wp.array(dtype=float),
     contact_pos_out: wp.array(dtype=wp.vec3),
     contact_frame_out: wp.array(dtype=wp.mat33),
@@ -254,27 +233,48 @@ def ccd_hfield_kernel_builder(
     contact_worldid_out: wp.array(dtype=int),
     contact_type_out: wp.array(dtype=int),
     contact_geomcollisionid_out: wp.array(dtype=int),
+    nacon_out: wp.array(dtype=int),
   ):
-    tid = wp.tid()
+    collisionid = wp.tid()
-    if tid >= ncollision_in[0]:
+    if collisionid >= ncollision_in[0]:
       return
 
-    geoms = collision_pair_in[tid]
+    geoms = collision_pair_in[collisionid]
     g1 = geoms[0]
     g2 = geoms[1]
 
     if geom_type[g1] != geomtype1 or geom_type[g2] != geomtype2:
       return
 
-    worldid = collision_worldid_in[tid]
+    worldid = collision_worldid_in[collisionid]
 
     # height field filter
     no_hf_collision, xmin, xmax, ymin, ymax, zmin, zmax = _hfield_filter(
-      geom_type, geom_dataid, geom_size, geom_rbound, geom_margin, hfield_size, geom_xpos_in, geom_xmat_in, worldid, g1, g2
+      geom_type,
+      geom_dataid,
+      geom_size,
+      geom_rbound,
+      geom_margin,
+      mesh_vertadr,
+      mesh_vertnum,
+      mesh_graphadr,
+      mesh_vert,
+      mesh_graph,
+      hfield_size,
+      geom_xpos_in,
+      geom_xmat_in,
+      worldid,
+      g1,
+      g2,
     )
     if no_hf_collision:
       return
 
+    ccdid = wp.atomic_add(nccd_in, wp.static(geomgeomid), 1)
+    if ccdid >= naccdmax_in:
+      wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid)
+      return
+
     _, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params(
       geom_condim,
       geom_priority,
@@ -293,7 +293,7 @@ def ccd_hfield_kernel_builder(
       pair_friction,
       collision_pair_in,
       collision_pairid_in,
-      tid,
+      collisionid,
       worldid,
     )
 
@@ -365,26 +365,23 @@ def ccd_hfield_kernel_builder(
     hfield_contact_dist = vec_maxconpair()
     hfield_contact_pos = mat_maxconpair()
     hfield_contact_normal = mat_maxconpair()
-    min_dist = float(wp.inf)
+    min_dist = float(MJ_MAXVAL)
-    min_normal = wp.vec3(wp.inf, wp.inf, wp.inf)
+    min_normal = wp.vec3(MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL)
-    min_pos = wp.vec3(wp.inf, wp.inf, wp.inf)
+    min_pos = wp.vec3(MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL)
     min_id = int(-1)
 
-    # TODO(team): height field margin?
+    # geom1 margin added to z-axis of hfield prism top points below
-    geom1.margin = margin
     geom2.margin = margin
 
     # EPA memory
-    epa_vert1 = epa_vert1_in[tid]
+    epa_vert = epa_vert_in[ccdid]
-    epa_vert2 = epa_vert2_in[tid]
+    epa_vert_index = epa_vert_index_in[ccdid]
-    epa_vert_index1 = epa_vert_index1_in[tid]
+    epa_face = epa_face_in[ccdid]
-    epa_vert_index2 = epa_vert_index2_in[tid]
+    epa_pr = epa_pr_in[ccdid]
-    epa_face = epa_face_in[tid]
+    epa_norm2 = epa_norm2_in[ccdid]
-    epa_pr = epa_pr_in[tid]
+    epa_horizon = epa_horizon_in[ccdid]
-    epa_norm2 = epa_norm2_in[tid]
-    epa_horizon = epa_horizon_in[tid]
 
-    collision_pairid = collision_pairid_in[tid]
+    collision_pairid = collision_pairid_in[collisionid]
 
     # process all prisms in subgrid
     count = int(0)
@@ -439,11 +436,7 @@ def ccd_hfield_kernel_builder(
           geom1.hfprism = prism
 
           # prism center
-          x1 = geom1.pos
+          x1 = geom1.pos + geom1.rot @ (prism[0] + prism[1] + prism[2] + prism[3] + prism[4] + prism[5]) * wp.static(1.0 / 6.0)
-          x1_ = wp.vec3(0.0, 0.0, 0.0)
-          for i in range(6):
-            x1_ += prism[i]
-          x1 += geom1.rot @ (x1_ / 6.0)
 
           dist, ncontact, w1, w2, idx = ccd(
             opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
@@ -456,10 +449,8 @@ def ccd_hfield_kernel_builder(
             geomtype2,
             x1,
             geom2.pos,
-            epa_vert1,
+            epa_vert,
-            epa_vert2,
+            epa_vert_index,
-            epa_vert_index1,
-            epa_vert_index2,
             epa_face,
             epa_pr,
             epa_norm2,
@@ -529,13 +520,12 @@ def ccd_hfield_kernel_builder(
     )
 
     # TODO(team): routine for select subset of contacts
-    # TODO(team): if use_multiccd?
     if wp.static(True):
       MIN_DIST_TO_NEXT_CONTACT = 1.0e-3
 
       # contact 1: furthest from minimum distance contact
       id1 = int(-1)
-      dist1 = float(-wp.inf)
+      dist1 = float(-MJ_MAXVAL)
       for i in range(count):
         if i == min_id:
           continue
@@ -589,7 +579,7 @@ def ccd_hfield_kernel_builder(
       dist_min1 = wp.cross(min_normal, min_pos - pos1)
 
       id2 = int(-1)
-      dist_12 = float(-wp.inf)
+      dist_12 = float(-MJ_MAXVAL)
       for i in range(count):
         if i == min_id or i == id1:
           continue
@@ -644,7 +634,7 @@ def ccd_hfield_kernel_builder(
       vec_12 = wp.cross(min_normal, pos1 - pos2)
 
       id3 = int(-1)
-      dist3 = float(-wp.inf)
+      dist3 = float(-MJ_MAXVAL)
       for i in range(count):
         if i == min_id or i == id1 or i == id2:
           continue
@@ -704,6 +694,7 @@ def ccd_kernel_builder(
   gjk_iterations: int,
   epa_iterations: int,
   use_multiccd: bool,
+  geomgeomid: int,
 ):
   """Kernel builder for non-heightfield CCD collisions (no hfield args)."""
 
@@ -711,14 +702,11 @@ def ccd_kernel_builder(
   def eval_ccd_write_contact(
     # Model:
     opt_ccd_tolerance: wp.array(dtype=float),
-    geom_type: wp.array(dtype=int),
     # Data in:
     naconmax_in: int,
     # In:
-    epa_vert1_in: wp.array2d(dtype=wp.vec3),
+    epa_vert_in: wp.array2d(dtype=wp.vec3),
-    epa_vert2_in: wp.array2d(dtype=wp.vec3),
+    epa_vert_index_in: wp.array2d(dtype=int),
-    epa_vert_index1_in: wp.array2d(dtype=int),
-    epa_vert_index2_in: wp.array2d(dtype=int),
     epa_face_in: wp.array2d(dtype=int),
     epa_pr_in: wp.array2d(dtype=wp.vec3),
     epa_norm2_in: wp.array2d(dtype=float),
@@ -738,7 +726,7 @@ def ccd_kernel_builder(
     geom2: Geom,
     geoms: wp.vec2i,
     worldid: int,
-    tid: int,
+    ccdid: int,
     margin: float,
     gap: float,
     condim: int,
@@ -748,7 +736,6 @@ def ccd_kernel_builder(
     solimp: vec5,
     x1: wp.vec3,
     x2: wp.vec3,
-    count: int,
     pairid: wp.vec2i,
     # Data out:
     contact_dist_out: wp.array(dtype=float),
@@ -771,12 +758,12 @@ def ccd_kernel_builder(
     witness2 = mat43()
     geom1.margin = margin
     geom2.margin = margin
-    if pairid[1] >= 0:
+    is_collision_sensor = pairid[1] >= 0
-      # if collision sensor, set large cutoff to work with various sensor cutoff values
+    if is_collision_sensor:
       cutoff = 1.0e32
     else:
       cutoff = 0.0
-    dist, ncontact, w1, w2, idx = ccd(
+    dist, ncollision, w1, w2, multiccd_idx = ccd(
       opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
       cutoff,
       gjk_iterations,
@@ -787,14 +774,12 @@ def ccd_kernel_builder(
       geomtype2,
       x1,
       x2,
-      epa_vert1_in[tid],
+      epa_vert_in[ccdid],
-      epa_vert2_in[tid],
+      epa_vert_index_in[ccdid],
-      epa_vert_index1_in[tid],
+      epa_face_in[ccdid],
-      epa_vert_index2_in[tid],
+      epa_pr_in[ccdid],
-      epa_face_in[tid],
+      epa_norm2_in[ccdid],
-      epa_pr_in[tid],
+      epa_horizon_in[ccdid],
-      epa_norm2_in[tid],
-      epa_horizon_in[tid],
     )
 
     if dist >= 0.0 and pairid[1] == -1:
@@ -803,30 +788,33 @@ def ccd_kernel_builder(
     witness1[0] = w1
     witness2[0] = w2
 
-    if wp.static(use_multiccd):
+    if wp.static(use_multiccd or (geomtype1 == GeomType.BOX and geomtype2 == GeomType.BOX)):
-      if (
+      if wp.static(geomtype1 == GeomType.MESH):
-        geom1.margin == 0.0
+        # verify that geom1 mesh data is present for multicontact
-        and geom2.margin == 0.0
+        if geom1.mesh_polyadr < 0:
-        and (geomtype1 == GeomType.BOX or (geomtype1 == GeomType.MESH and geom1.mesh_polyadr > -1))
+          multiccd_idx = -1
-        and (geomtype2 == GeomType.BOX or (geomtype2 == GeomType.MESH and geom2.mesh_polyadr > -1))
+
-      ):
+      if wp.static(geomtype2 == GeomType.MESH):
-        ncontact, witness1, witness2 = multicontact(
+        # verify that geom2 mesh data is present for multicontact
-          multiccd_polygon_in[tid],
+        if geom2.mesh_polyadr < 0:
-          multiccd_clipped_in[tid],
+          multiccd_idx = -1
-          multiccd_pnormal_in[tid],
+
-          multiccd_pdist_in[tid],
+      if multiccd_idx > -1:
-          multiccd_idx1_in[tid],
+        ncollision, witness1, witness2 = multicontact(
-          multiccd_idx2_in[tid],
+          multiccd_polygon_in[ccdid],
-          multiccd_n1_in[tid],
+          multiccd_clipped_in[ccdid],
-          multiccd_n2_in[tid],
+          multiccd_pnormal_in[ccdid],
-          multiccd_endvert_in[tid],
+          multiccd_pdist_in[ccdid],
-          multiccd_face1_in[tid],
+          multiccd_idx1_in[ccdid],
-          multiccd_face2_in[tid],
+          multiccd_idx2_in[ccdid],
-          epa_vert1_in[tid],
+          multiccd_n1_in[ccdid],
-          epa_vert2_in[tid],
+          multiccd_n2_in[ccdid],
-          epa_vert_index1_in[tid],
+          multiccd_endvert_in[ccdid],
-          epa_vert_index2_in[tid],
+          multiccd_face1_in[ccdid],
-          epa_face_in[tid, idx],
+          multiccd_face2_in[ccdid],
+          epa_vert_in[ccdid],
+          epa_vert_index_in[ccdid],
+          epa_face_in[ccdid, multiccd_idx],
           w1,
           w2,
           geom1,
@@ -835,7 +823,7 @@ def ccd_kernel_builder(
           geomtype2,
         )
 
-    for i in range(ncontact):
+    for i in range(ncollision):
       points[i] = 0.5 * (witness1[i] + witness2[i])
     normal = witness1[0] - witness2[0]
     frame = make_frame(normal)
@@ -845,8 +833,9 @@ def ccd_kernel_builder(
       frame *= -1.0
       geoms = wp.vec2i(geoms[1], geoms[0])
 
-    for i in range(ncontact):
+    nactive = int(0)  # number of contacts contributing to the physics
-      write_contact(
+    for i in range(ncollision):
+      active = write_contact(
         naconmax_in,
         i,
         dist,
@@ -877,13 +866,12 @@ def ccd_kernel_builder(
         contact_geomcollisionid_out,
         nacon_out,
       )
-      if count + (i + 1) >= MJ_MAXCONPAIR:
+      nactive += active
-        return i + 1
 
-    return ncontact
+    return nactive
 
   # runs convex collision on a set of geom pairs to recover contact info (non-heightfield)
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def ccd_kernel(
     # Model:
     opt_ccd_tolerance: wp.array(dtype=float),
@@ -900,8 +888,8 @@ def ccd_kernel_builder(
     geom_gap: wp.array2d(dtype=float),
     mesh_vertadr: wp.array(dtype=int),
     mesh_vertnum: wp.array(dtype=int),
-    mesh_vert: wp.array(dtype=wp.vec3),
     mesh_graphadr: wp.array(dtype=int),
+    mesh_vert: wp.array(dtype=wp.vec3),
     mesh_graph: wp.array(dtype=int),
     mesh_polynum: wp.array(dtype=int),
     mesh_polyadr: wp.array(dtype=int),
@@ -920,18 +908,17 @@ def ccd_kernel_builder(
     pair_gap: wp.array2d(dtype=float),
     pair_friction: wp.array2d(dtype=vec5),
     # Data in:
-    naconmax_in: int,
     geom_xpos_in: wp.array2d(dtype=wp.vec3),
     geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    naconmax_in: int,
+    naccdmax_in: int,
+    ncollision_in: wp.array(dtype=int),
+    # In:
     collision_pair_in: wp.array(dtype=wp.vec2i),
     collision_pairid_in: wp.array(dtype=wp.vec2i),
     collision_worldid_in: wp.array(dtype=int),
-    ncollision_in: wp.array(dtype=int),
+    epa_vert_in: wp.array2d(dtype=wp.vec3),
-    # In:
+    epa_vert_index_in: wp.array2d(dtype=int),
-    epa_vert1_in: wp.array2d(dtype=wp.vec3),
-    epa_vert2_in: wp.array2d(dtype=wp.vec3),
-    epa_vert_index1_in: wp.array2d(dtype=int),
-    epa_vert_index2_in: wp.array2d(dtype=int),
     epa_face_in: wp.array2d(dtype=int),
     epa_pr_in: wp.array2d(dtype=wp.vec3),
     epa_norm2_in: wp.array2d(dtype=float),
@@ -947,8 +934,8 @@ def ccd_kernel_builder(
     multiccd_endvert_in: wp.array2d(dtype=wp.vec3),
     multiccd_face1_in: wp.array2d(dtype=wp.vec3),
     multiccd_face2_in: wp.array2d(dtype=wp.vec3),
+    nccd_in: wp.array(dtype=int),
     # Data out:
-    nacon_out: wp.array(dtype=int),
     contact_dist_out: wp.array(dtype=float),
     contact_pos_out: wp.array(dtype=wp.vec3),
     contact_frame_out: wp.array(dtype=wp.mat33),
@@ -962,19 +949,25 @@ def ccd_kernel_builder(
     contact_worldid_out: wp.array(dtype=int),
     contact_type_out: wp.array(dtype=int),
     contact_geomcollisionid_out: wp.array(dtype=int),
+    nacon_out: wp.array(dtype=int),
   ):
-    tid = wp.tid()
+    collisionid = wp.tid()
-    if tid >= ncollision_in[0]:
+    if collisionid >= ncollision_in[0]:
       return
 
-    geoms = collision_pair_in[tid]
+    geoms = collision_pair_in[collisionid]
     g1 = geoms[0]
     g2 = geoms[1]
 
     if geom_type[g1] != geomtype1 or geom_type[g2] != geomtype2:
       return
 
-    worldid = collision_worldid_in[tid]
+    ccdid = wp.atomic_add(nccd_in, wp.static(geomgeomid), 1)
+    if ccdid >= naccdmax_in:
+      wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid)
+      return
+
+    worldid = collision_worldid_in[collisionid]
 
     _, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params(
       geom_condim,
@@ -994,7 +987,7 @@ def ccd_kernel_builder(
       pair_friction,
       collision_pair_in,
       collision_pairid_in,
-      tid,
+      collisionid,
       worldid,
     )
 
@@ -1024,12 +1017,9 @@ def ccd_kernel_builder(
 
     eval_ccd_write_contact(
       opt_ccd_tolerance,
-      geom_type,
       naconmax_in,
-      epa_vert1_in,
+      epa_vert_in,
-      epa_vert2_in,
+      epa_vert_index_in,
-      epa_vert_index1_in,
-      epa_vert_index2_in,
       epa_face_in,
       epa_pr_in,
       epa_norm2_in,
@@ -1049,7 +1039,7 @@ def ccd_kernel_builder(
       geom2,
       geoms,
       worldid,
-      tid,
+      ccdid,
       margin,
       gap,
       condim,
@@ -1059,8 +1049,7 @@ def ccd_kernel_builder(
       solimp,
       geom1.pos,
       geom2.pos,
-      0,
+      collision_pairid_in[collisionid],
-      collision_pairid_in[tid],
       contact_dist_out,
       contact_pos_out,
       contact_frame_out,
@@ -1080,37 +1069,8 @@ def ccd_kernel_builder(
   return ccd_kernel
 
 
-# Heightfield collision pairs handled by ccd_hfield_kernel_builder
-_HFIELD_COLLISION_PAIRS = [
-  (GeomType.HFIELD, GeomType.SPHERE),
-  (GeomType.HFIELD, GeomType.CAPSULE),
-  (GeomType.HFIELD, GeomType.ELLIPSOID),
-  (GeomType.HFIELD, GeomType.CYLINDER),
-  (GeomType.HFIELD, GeomType.BOX),
-  (GeomType.HFIELD, GeomType.MESH),
-]
-
-# Non-heightfield collision pairs handled by ccd_kernel_builder
-_NON_HFIELD_COLLISION_PAIRS = [
-  (GeomType.SPHERE, GeomType.ELLIPSOID),
-  (GeomType.SPHERE, GeomType.MESH),
-  (GeomType.CAPSULE, GeomType.ELLIPSOID),
-  (GeomType.CAPSULE, GeomType.CYLINDER),
-  (GeomType.CAPSULE, GeomType.MESH),
-  (GeomType.ELLIPSOID, GeomType.ELLIPSOID),
-  (GeomType.ELLIPSOID, GeomType.CYLINDER),
-  (GeomType.ELLIPSOID, GeomType.BOX),
-  (GeomType.ELLIPSOID, GeomType.MESH),
-  (GeomType.CYLINDER, GeomType.CYLINDER),
-  (GeomType.CYLINDER, GeomType.BOX),
-  (GeomType.CYLINDER, GeomType.MESH),
-  (GeomType.BOX, GeomType.MESH),
-  (GeomType.MESH, GeomType.MESH),
-]
-
-
 @event_scope
-def convex_narrowphase(m: Model, d: Data):
+def convex_narrowphase(m: Model, d: Data, ctx: CollisionContext, collision_table: list[tuple[GeomType, GeomType]]):
   """Runs narrowphase collision detection for convex geom pairs.
 
   This function handles collision detection for pairs of convex geometries that were
@@ -1126,40 +1086,56 @@ def convex_narrowphase(m: Model, d: Data):
   computations for non-existent pair types.
   """
 
-  def _pair_count(p1: int, p2: int) -> int:
+  def _pair_count(p1: int, p2: int) -> Tuple[int, int]:
-    return m.geom_pair_type_count[upper_trid_index(len(GeomType), p1, p2)]
+    idx = upper_trid_index(len(GeomType), p1, p2)
+    return m.geom_pair_type_count[idx], idx
 
+  ncollision = sum(_pair_count(g[0].value, g[1].value)[0] for g in collision_table)
   # no convex collisions, early return
-  if not any(_pair_count(g[0].value, g[1].value) for g in _CONVEX_COLLISION_PAIRS):
+
+  if ncollision == 0:
     return
 
-  epa_iterations = m.opt.ccd_iterations
+  # compute nmaxpolygon and nmaxmeshdeg given the geom pairs for the model
+  nboxbox, _ = _pair_count(GeomType.BOX.value, GeomType.BOX.value)
+  if (GeomType.BOX, GeomType.BOX) not in collision_table:
+    nboxbox = 0
+  nboxmesh, _ = _pair_count(GeomType.BOX.value, GeomType.MESH.value)
+  nmeshmesh, _ = _pair_count(GeomType.MESH.value, GeomType.MESH.value)
+
+  epa_iterations = 16 if nboxbox == ncollision else m.opt.ccd_iterations
 
   # set to true to enable multiccd
   use_multiccd = m.opt.enableflags & EnableBit.MULTICCD
-  nmaxpolygon = m.nmaxpolygon if use_multiccd else 0
-  nmaxmeshdeg = m.nmaxmeshdeg if use_multiccd else 0
 
-  # epa_vert1: vertices in EPA polytope in geom 1 space
+  # need at least 4 (square sides) if there's a box collision needing multiccd
-  epa_vert1 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=wp.vec3)
+  nmaxpolygon = 4 if nboxbox > 0 else 0
-  # epa_vert2: vertices in EPA polytope in geom 2 space
+  nmaxmeshdeg = 3 if nboxbox > 0 else 0
-  epa_vert2 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=wp.vec3)
+
-  # epa_vert_index1: vertex indices in EPA polytope for geom 1
+  # need to allocate more memory if there's meshes
-  epa_vert_index1 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=int)
+  if use_multiccd and nmeshmesh + nboxmesh > 0:
-  # epa_vert_index2: vertex indices in EPA polytope for geom 2  (naconmax, 5 + CCDiter)
+    minval = 4 if nboxmesh else nmaxpolygon
-  epa_vert_index2 = wp.empty(shape=(d.naconmax, 5 + epa_iterations), dtype=int)
+    nmaxpolygon = max(m.nmaxpolygon, minval)
+    nmaxmeshdeg = max(m.nmaxmeshdeg, 3)
+
+  # ccd collider count
+  nccd = wp.zeros(len(GeomType) * (len(GeomType) + 1) // 2, dtype=int)
+
+  # epa_vert: vertices in EPA polytope
+  epa_vert = wp.empty(shape=(d.naccdmax, 10 + 2 * epa_iterations), dtype=wp.vec3)
+  # epa_vert_index: vertex indices in EPA polytope
+  epa_vert_index = wp.empty(shape=(d.naccdmax, 10 + 2 * epa_iterations), dtype=int)
   # epa_face: faces of polytope represented by three indices
-  epa_face = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int)
+  epa_face = wp.empty(shape=(d.naccdmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=int)
   # epa_pr: projection of origin on polytope faces
-  epa_pr = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=wp.vec3)
+  epa_pr = wp.empty(shape=(d.naccdmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=wp.vec3)
   # epa_norm2: epa_pr * epa_pr
-  epa_norm2 = wp.empty(shape=(d.naconmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=float)
+  epa_norm2 = wp.empty(shape=(d.naccdmax, 6 + MJ_MAX_EPAFACES * epa_iterations), dtype=float)
   # epa_horizon: index pair (i j) of edges on horizon
-  epa_horizon = wp.empty(shape=(d.naconmax, MJ_MAX_EPAHORIZON), dtype=int)
+  epa_horizon = wp.empty(shape=(d.naccdmax, MJ_MAX_EPAHORIZON), dtype=int)
 
   # Contact outputs
   contact_outputs = [
-    d.nacon,
     d.contact.dist,
     d.contact.pos,
     d.contact.frame,
@@ -1173,15 +1149,17 @@ def convex_narrowphase(m: Model, d: Data):
     d.contact.worldid,
     d.contact.type,
     d.contact.geomcollisionid,
+    d.nacon,
   ]
 
   # Launch heightfield collision kernels (no multiccd args, 72 args total)
-  for geom_pair in _HFIELD_COLLISION_PAIRS:
+  for geom_pair in collision_table:
     g1 = geom_pair[0].value
     g2 = geom_pair[1].value
-    if _pair_count(g1, g2):
+    count, geomgeomid = _pair_count(g1, g2)
+    if (g1 == GeomType.HFIELD or g2 == GeomType.HFIELD) and count:
       wp.launch(
-        ccd_hfield_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations),
+        ccd_hfield_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations, geomgeomid),
         dim=d.naconmax,
         inputs=[
           m.opt.ccd_tolerance,
@@ -1197,15 +1175,10 @@ def convex_narrowphase(m: Model, d: Data):
           m.geom_friction,
           m.geom_margin,
           m.geom_gap,
-          m.hfield_adr,
-          m.hfield_nrow,
-          m.hfield_ncol,
-          m.hfield_size,
-          m.hfield_data,
           m.mesh_vertadr,
           m.mesh_vertnum,
-          m.mesh_vert,
           m.mesh_graphadr,
+          m.mesh_vert,
           m.mesh_graph,
           m.mesh_polynum,
           m.mesh_polyadr,
@@ -1216,6 +1189,11 @@ def convex_narrowphase(m: Model, d: Data):
           m.mesh_polymapadr,
           m.mesh_polymapnum,
           m.mesh_polymap,
+          m.hfield_size,
+          m.hfield_nrow,
+          m.hfield_ncol,
+          m.hfield_adr,
+          m.hfield_data,
           m.pair_dim,
           m.pair_solref,
           m.pair_solreffriction,
@@ -1223,56 +1201,57 @@ def convex_narrowphase(m: Model, d: Data):
           m.pair_margin,
           m.pair_gap,
           m.pair_friction,
-          d.naconmax,
           d.geom_xpos,
           d.geom_xmat,
-          d.collision_pair,
+          d.naconmax,
-          d.collision_pairid,
+          d.naccdmax,
-          d.collision_worldid,
           d.ncollision,
-          epa_vert1,
+          ctx.collision_pair,
-          epa_vert2,
+          ctx.collision_pairid,
-          epa_vert_index1,
+          ctx.collision_worldid,
-          epa_vert_index2,
+          epa_vert,
+          epa_vert_index,
           epa_face,
           epa_pr,
           epa_norm2,
           epa_horizon,
+          nccd,
         ],
         outputs=contact_outputs,
       )
 
   # Allocate multiccd arrays only for non-heightfield collisions
   # multiccd_polygon: clipped contact surface
-  multiccd_polygon = wp.empty(shape=(d.naconmax, 2 * nmaxpolygon), dtype=wp.vec3)
+  multiccd_polygon = wp.empty(shape=(d.naccdmax, 2 * nmaxpolygon), dtype=wp.vec3)
   # multiccd_clipped: clipped contact surface (intermediate)
-  multiccd_clipped = wp.empty(shape=(d.naconmax, 2 * nmaxpolygon), dtype=wp.vec3)
+  multiccd_clipped = wp.empty(shape=(d.naccdmax, 2 * nmaxpolygon), dtype=wp.vec3)
   # multiccd_pnormal: plane normal of clipping polygon
-  multiccd_pnormal = wp.empty(shape=(d.naconmax, nmaxpolygon), dtype=wp.vec3)
+  multiccd_pnormal = wp.empty(shape=(d.naccdmax, nmaxpolygon), dtype=wp.vec3)
   # multiccd_pdist: plane distance of clipping polygon
-  multiccd_pdist = wp.empty(shape=(d.naconmax, nmaxpolygon), dtype=float)
+  multiccd_pdist = wp.empty(shape=(d.naccdmax, nmaxpolygon), dtype=float)
   # multiccd_idx1: list of normal index candidates for Geom 1
-  multiccd_idx1 = wp.empty(shape=(d.naconmax, nmaxmeshdeg), dtype=int)
+  multiccd_idx1 = wp.empty(shape=(d.naccdmax, nmaxmeshdeg), dtype=int)
   # multiccd_idx2: list of normal index candidates for Geom 2
-  multiccd_idx2 = wp.empty(shape=(d.naconmax, nmaxmeshdeg), dtype=int)
+  multiccd_idx2 = wp.empty(shape=(d.naccdmax, nmaxmeshdeg), dtype=int)
   # multiccd_n1: list of normal candidates for Geom 1
-  multiccd_n1 = wp.empty(shape=(d.naconmax, nmaxmeshdeg), dtype=wp.vec3)
+  multiccd_n1 = wp.empty(shape=(d.naccdmax, nmaxmeshdeg), dtype=wp.vec3)
   # multiccd_n2: list of normal candidates for Geom 1
-  multiccd_n2 = wp.empty(shape=(d.naconmax, nmaxmeshdeg), dtype=wp.vec3)
+  multiccd_n2 = wp.empty(shape=(d.naccdmax, nmaxmeshdeg), dtype=wp.vec3)
   # multiccd_endvert: list of edge vertices candidates
-  multiccd_endvert = wp.empty(shape=(d.naconmax, nmaxmeshdeg), dtype=wp.vec3)
+  multiccd_endvert = wp.empty(shape=(d.naccdmax, nmaxmeshdeg), dtype=wp.vec3)
   # multiccd_face1: contact face
-  multiccd_face1 = wp.empty(shape=(d.naconmax, nmaxpolygon), dtype=wp.vec3)
+  multiccd_face1 = wp.empty(shape=(d.naccdmax, nmaxpolygon), dtype=wp.vec3)
   # multiccd_face2: contact face
-  multiccd_face2 = wp.empty(shape=(d.naconmax, nmaxpolygon), dtype=wp.vec3)
+  multiccd_face2 = wp.empty(shape=(d.naccdmax, nmaxpolygon), dtype=wp.vec3)
 
   # Launch non-heightfield collision kernels (no hfield args, 78 args total)
-  for geom_pair in _NON_HFIELD_COLLISION_PAIRS:
+  for geom_pair in collision_table:
     g1 = geom_pair[0].value
     g2 = geom_pair[1].value
-    if _pair_count(g1, g2):
+    count, geomgeomid = _pair_count(g1, g2)
+    if g1 != GeomType.HFIELD and g2 != GeomType.HFIELD and count:
       wp.launch(
-        ccd_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations, use_multiccd),
+        ccd_kernel_builder(g1, g2, m.opt.ccd_iterations, epa_iterations, use_multiccd, geomgeomid),
         dim=d.naconmax,
         inputs=[
           m.opt.ccd_tolerance,
@@ -1289,8 +1268,8 @@ def convex_narrowphase(m: Model, d: Data):
           m.geom_gap,
           m.mesh_vertadr,
           m.mesh_vertnum,
-          m.mesh_vert,
           m.mesh_graphadr,
+          m.mesh_vert,
           m.mesh_graph,
           m.mesh_polynum,
           m.mesh_polyadr,
@@ -1308,17 +1287,16 @@ def convex_narrowphase(m: Model, d: Data):
           m.pair_margin,
           m.pair_gap,
           m.pair_friction,
-          d.naconmax,
           d.geom_xpos,
           d.geom_xmat,
-          d.collision_pair,
+          d.naconmax,
-          d.collision_pairid,
+          d.naccdmax,
-          d.collision_worldid,
           d.ncollision,
-          epa_vert1,
+          ctx.collision_pair,
-          epa_vert2,
+          ctx.collision_pairid,
-          epa_vert_index1,
+          ctx.collision_worldid,
-          epa_vert_index2,
+          epa_vert,
+          epa_vert_index,
           epa_face,
           epa_pr,
           epa_norm2,
@@ -1334,6 +1312,7 @@ def convex_narrowphase(m: Model, d: Data):
           multiccd_endvert,
           multiccd_face1,
           multiccd_face2,
+          nccd,
         ],
         outputs=contact_outputs,
       )

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

@@ -24,17 +24,65 @@ 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
 from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType
+from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext
+from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
+from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat23
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 
+# Corresponding table to MuJoCo's mjCOLLISIONFUNC table in engine_collision_driver.c
+MJ_COLLISION_TABLE = {
+  (GeomType.PLANE, GeomType.SPHERE): CollisionType.PRIMITIVE,
+  (GeomType.PLANE, GeomType.CAPSULE): CollisionType.PRIMITIVE,
+  (GeomType.PLANE, GeomType.ELLIPSOID): CollisionType.PRIMITIVE,
+  (GeomType.PLANE, GeomType.CYLINDER): CollisionType.PRIMITIVE,
+  (GeomType.PLANE, GeomType.BOX): CollisionType.PRIMITIVE,
+  (GeomType.PLANE, GeomType.MESH): CollisionType.PRIMITIVE,
+  (GeomType.HFIELD, GeomType.SPHERE): CollisionType.CONVEX,
+  (GeomType.HFIELD, GeomType.CAPSULE): CollisionType.CONVEX,
+  (GeomType.HFIELD, GeomType.ELLIPSOID): CollisionType.CONVEX,
+  (GeomType.HFIELD, GeomType.CYLINDER): CollisionType.CONVEX,
+  (GeomType.HFIELD, GeomType.BOX): CollisionType.CONVEX,
+  (GeomType.HFIELD, GeomType.MESH): CollisionType.CONVEX,
+  (GeomType.SPHERE, GeomType.SPHERE): CollisionType.PRIMITIVE,
+  (GeomType.SPHERE, GeomType.CAPSULE): CollisionType.PRIMITIVE,
+  (GeomType.SPHERE, GeomType.ELLIPSOID): CollisionType.CONVEX,
+  (GeomType.SPHERE, GeomType.CYLINDER): CollisionType.PRIMITIVE,
+  (GeomType.SPHERE, GeomType.BOX): CollisionType.PRIMITIVE,
+  (GeomType.SPHERE, GeomType.MESH): CollisionType.CONVEX,
+  (GeomType.CAPSULE, GeomType.CAPSULE): CollisionType.PRIMITIVE,
+  (GeomType.CAPSULE, GeomType.ELLIPSOID): CollisionType.CONVEX,
+  (GeomType.CAPSULE, GeomType.CYLINDER): CollisionType.CONVEX,
+  (GeomType.CAPSULE, GeomType.BOX): CollisionType.PRIMITIVE,
+  (GeomType.CAPSULE, GeomType.MESH): CollisionType.CONVEX,
+  (GeomType.ELLIPSOID, GeomType.ELLIPSOID): CollisionType.CONVEX,
+  (GeomType.ELLIPSOID, GeomType.CYLINDER): CollisionType.CONVEX,
+  (GeomType.ELLIPSOID, GeomType.BOX): CollisionType.CONVEX,
+  (GeomType.ELLIPSOID, GeomType.MESH): CollisionType.CONVEX,
+  (GeomType.CYLINDER, GeomType.CYLINDER): CollisionType.CONVEX,
+  (GeomType.CYLINDER, GeomType.BOX): CollisionType.CONVEX,
+  (GeomType.CYLINDER, GeomType.MESH): CollisionType.CONVEX,
+  (GeomType.BOX, GeomType.BOX): CollisionType.CONVEX,  # overwritten by NATIVECCD disable flag
+  (GeomType.BOX, GeomType.MESH): CollisionType.CONVEX,
+  (GeomType.MESH, GeomType.MESH): CollisionType.CONVEX,
+}
+
+
+def create_collision_context(naconmax: int) -> CollisionContext:
+  """Create a CollisionContext with allocated arrays."""
+  return CollisionContext(
+    collision_pair=wp.empty(naconmax, dtype=wp.vec2i),
+    collision_pairid=wp.empty(naconmax, dtype=wp.vec2i),
+    collision_worldid=wp.empty(naconmax, dtype=int),
+  )
+
 
 @wp.kernel
 def _zero_nacon_ncollision(
@@ -293,10 +341,11 @@ def _add_geom_pair(
   worldid: int,
   nxnid: int,
   # Data out:
+  ncollision_out: wp.array(dtype=int),
+  # Out:
   collision_pair_out: wp.array(dtype=wp.vec2i),
   collision_pairid_out: wp.array(dtype=wp.vec2i),
   collision_worldid_out: wp.array(dtype=int),
-  ncollision_out: wp.array(dtype=int),
 ):
   pairid = wp.atomic_add(ncollision_out, 0, 1)
 
@@ -329,15 +378,15 @@ def _binary_search(values: wp.array(dtype=Any), value: Any, lower: int, upper: i
 
 
 def _sap_project(opt_broadphase: int):
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def sap_project(
     # Model:
     ngeom: int,
     geom_rbound: wp.array2d(dtype=float),
     geom_margin: wp.array2d(dtype=float),
     # Data in:
-    nworld_in: int,
     geom_xpos_in: wp.array2d(dtype=wp.vec3),
+    nworld_in: int,
     # In:
     direction_in: wp.vec3,
     # Out:
@@ -402,7 +451,7 @@ def _sap_range(
 
 @cache_kernel
 def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Model:
     ngeom: int,
@@ -412,19 +461,20 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
     geom_margin: wp.array2d(dtype=float),
     nxn_pairid: wp.array(dtype=wp.vec2i),
     # Data in:
-    nworld_in: int,
-    naconmax_in: int,
     geom_xpos_in: wp.array2d(dtype=wp.vec3),
     geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    nworld_in: int,
+    naconmax_in: int,
     # In:
     sort_index_in: wp.array2d(dtype=int),
     cumulative_sum_in: wp.array(dtype=int),
     nsweep_in: int,
     # Data out:
+    ncollision_out: wp.array(dtype=int),
+    # Out:
     collision_pair_out: wp.array(dtype=wp.vec2i),
     collision_pairid_out: wp.array(dtype=wp.vec2i),
     collision_worldid_out: wp.array(dtype=int),
-    ncollision_out: wp.array(dtype=int),
   ):
     worldgeomid = wp.tid()
 
@@ -472,17 +522,17 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
           geom2,
           worldid,
           idx,
+          ncollision_out,
           collision_pair_out,
           collision_pairid_out,
           collision_worldid_out,
-          ncollision_out,
         )
 
   return kernel
 
 
 def _segmented_sort(tile_size: int):
-  @wp.kernel
+  @wp.kernel(module="unique")
   def segmented_sort(
     # In:
     projection_lower_in: wp.array2d(dtype=float),
@@ -508,7 +558,7 @@ def _segmented_sort(tile_size: int):
 
 
 @event_scope
-def sap_broadphase(m: Model, d: Data):
+def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
   """Runs broadphase collision detection using a sweep-and-prune (SAP) algorithm.
 
   This method is more efficient than the N-squared approach for large numbers of
@@ -543,7 +593,7 @@ def sap_broadphase(m: Model, d: Data):
   wp.launch(
     kernel=_sap_project(m.opt.broadphase),
     dim=(d.nworld, m.ngeom),
-    inputs=[m.ngeom, m.geom_rbound, m.geom_margin, d.nworld, d.geom_xpos, direction],
+    inputs=[m.ngeom, m.geom_rbound, m.geom_margin, d.geom_xpos, d.nworld, direction],
     outputs=[
       projection_lower.reshape((-1, m.ngeom)),
       projection_upper,
@@ -588,21 +638,21 @@ def sap_broadphase(m: Model, d: Data):
       m.geom_rbound,
       m.geom_margin,
       m.nxn_pairid,
-      d.nworld,
-      d.naconmax,
       d.geom_xpos,
       d.geom_xmat,
+      d.nworld,
+      d.naconmax,
       sort_index.reshape((-1, m.ngeom)),
       cumulative_sum.reshape(-1),
       nsweep,
     ],
-    outputs=[d.collision_pair, d.collision_pairid, d.collision_worldid, d.ncollision],
+    outputs=[d.ncollision, ctx.collision_pair, ctx.collision_pairid, ctx.collision_worldid],
   )
 
 
 @cache_kernel
 def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int):
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Model:
     geom_type: wp.array(dtype=int),
@@ -612,14 +662,15 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
     nxn_geom_pair: wp.array(dtype=wp.vec2i),
     nxn_pairid: wp.array(dtype=wp.vec2i),
     # Data in:
-    naconmax_in: int,
     geom_xpos_in: wp.array2d(dtype=wp.vec3),
     geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    naconmax_in: int,
     # Data out:
+    ncollision_out: wp.array(dtype=int),
+    # Out:
     collision_pair_out: wp.array(dtype=wp.vec2i),
     collision_pairid_out: wp.array(dtype=wp.vec2i),
     collision_worldid_out: wp.array(dtype=int),
-    ncollision_out: wp.array(dtype=int),
   ):
     worldid, elementid = wp.tid()
 
@@ -641,17 +692,17 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
         geom2,
         worldid,
         elementid,
+        ncollision_out,
         collision_pair_out,
         collision_pairid_out,
         collision_worldid_out,
-        ncollision_out,
       )
 
   return kernel
 
 
 @event_scope
-def nxn_broadphase(m: Model, d: Data):
+def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
   """Runs broadphase collision detection using a brute-force N-squared approach.
 
   This function iterates through a pre-filtered list of all possible geometry pairs and
@@ -674,27 +725,34 @@ def nxn_broadphase(m: Model, d: Data):
       m.geom_margin,
       m.nxn_geom_pair_filtered,
       m.nxn_pairid_filtered,
-      d.naconmax,
       d.geom_xpos,
       d.geom_xmat,
+      d.naconmax,
     ],
     outputs=[
-      d.collision_pair,
-      d.collision_pairid,
-      d.collision_worldid,
       d.ncollision,
+      ctx.collision_pair,
+      ctx.collision_pairid,
+      ctx.collision_worldid,
     ],
   )
 
 
-def _narrowphase(m, d):
+def _narrowphase(m: Model, d: Data, ctx: CollisionContext):
+  collision_table = MJ_COLLISION_TABLE
+  if m.opt.disableflags & DisableBit.NATIVECCD:
+    collision_table[(GeomType.BOX, GeomType.BOX)] = CollisionType.PRIMITIVE
+
+  convex_pairs = [key for key, value in collision_table.items() if value == CollisionType.CONVEX]
+  primitive_pairs = [key for key, value in collision_table.items() if value == CollisionType.PRIMITIVE]
+
   # TODO(team): we should reject far-away contacts in the narrowphase instead of constraint
   #             partitioning because we can move some pressure of the atomics
-  convex_narrowphase(m, d)
+  convex_narrowphase(m, d, ctx, convex_pairs)
-  primitive_narrowphase(m, d)
+  primitive_narrowphase(m, d, ctx, primitive_pairs)
 
   if m.has_sdf_geom:
-    sdf_narrowphase(m, d)
+    sdf_narrowphase(m, d, ctx)
 
 
 @event_scope
@@ -715,15 +773,18 @@ def collision(m: Model, d: Data):
   This function will do nothing except zero out arrays if collision detection is disabled
   via `m.opt.disableflags` or if `d.nacon` is 0.
   """
-  # zero contact and collision counters
-  wp.launch(_zero_nacon_ncollision, dim=1, outputs=[d.nacon, d.ncollision])
-
   if d.naconmax == 0 or m.opt.disableflags & (DisableBit.CONSTRAINT | DisableBit.CONTACT):
+    d.nacon.zero_()
     return
 
-  if m.opt.broadphase == BroadphaseType.NXN:
+  ctx = create_collision_context(d.naconmax)
-    nxn_broadphase(m, d)
-  else:
-    sap_broadphase(m, d)
 
-  _narrowphase(m, d)
+  # zero counters
+  wp.launch(_zero_nacon_ncollision, dim=1, outputs=[d.nacon, d.ncollision])
+
+  if m.opt.broadphase == BroadphaseType.NXN:
+    nxn_broadphase(m, d, ctx)
+  else:
+    sap_broadphase(m, d, ctx)
+
+  _narrowphase(m, d, ctx)

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

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+import math
 from typing import Tuple
 
 import warp as wp
@@ -30,9 +31,11 @@ FLOAT_MAX = 1e30
 MINVAL = 1e-15
 MIN_DIST = 1e-10
 
-# TODO(kbayes): write out formulas to derive these constants
+FACE_TOL = wp.static(math.cos(0.0016))
-FACE_TOL = 0.99999872
+EDGE_TOL = wp.static(math.sin(0.0016))
-EDGE_TOL = 0.00159999931
+
+# tolarance used by multicontact for intersecting a plane and a line segment
+INTERSECT_TOL = 0.0000003
 
 # Bit flags for face status in EPA polytope.
 # Defined at module scope to avoid Warp's intermediate type issues with literals.
@@ -59,11 +62,9 @@ class GJKResult:
 class Polytope:
   status: int
 
-  # vertices in polytope
+  # vertices in polytope (packed geom1 followed by geom2)
-  vert1: wp.array(dtype=wp.vec3)
+  vert: wp.array(dtype=wp.vec3)
-  vert2: wp.array(dtype=wp.vec3)
+  vert_index: wp.array(dtype=int)
-  vert_index1: wp.array(dtype=int)
-  vert_index2: wp.array(dtype=int)
   nvert: int
 
   # faces in polytope
@@ -107,9 +108,9 @@ def support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
     tmp = wp.sign(local_dir)
     res = wp.cw_mul(tmp, geom.size)
     sp.point = geom.rot @ res + geom.pos
-    sp.vertex_index = wp.where(tmp[0] > 0, 1, 0)
+    sp.vertex_index = wp.where(tmp[0] > 0.0, 1, 0)
-    sp.vertex_index += wp.where(tmp[1] > 0, 2, 0)
+    sp.vertex_index += wp.where(tmp[1] > 0.0, 2, 0)
-    sp.vertex_index += wp.where(tmp[2] > 0, 4, 0)
+    sp.vertex_index += wp.where(tmp[2] > 0.0, 4, 0)
   elif geomtype == GeomType.CAPSULE:
     res = local_dir * geom.size[0]
     # add cylinder contribution
@@ -177,7 +178,7 @@ def support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
   elif geomtype == GeomType.HFIELD:
     max_dist = float(FLOAT_MIN)
     # TODO(kbayes): Support edge prisms
-    sp.vertex_index = wp.where(dir[2] < 0, -2, -3)
+    sp.vertex_index = wp.where(dir[2] < 0.0, -2, -3)
     for i in range(6):
       vert = geom.hfprism[i]
       dist = wp.dot(vert, dir)
@@ -197,7 +198,10 @@ def _attach_face(pt: Polytope, idx: int, v1: int, v2: int, v3: int) -> float:
     return 0.0
 
   # compute witness point v
-  r, ret = _project_origin_plane(pt.vert1[v3] - pt.vert2[v3], pt.vert1[v2] - pt.vert2[v2], pt.vert1[v1] - pt.vert2[v1])
+  p1 = pt.vert[2 * v1] - pt.vert[2 * v1 + 1]
+  p2 = pt.vert[2 * v2] - pt.vert[2 * v2 + 1]
+  p3 = pt.vert[2 * v3] - pt.vert[2 * v3 + 1]
+  r, ret = _project_origin_plane(p3, p2, p1)
   if ret:
     return 0.0
 
@@ -214,13 +218,13 @@ def _epa_support(
   pt: Polytope, idx: int, geom1: Geom, geom2: Geom, geom1_type: int, geom2_type: int, dir: wp.vec3
 ) -> Tuple[int, int]:
   sp = support(geom1, geom1_type, dir)
-  pt.vert1[idx] = sp.point
+  pt.vert[2 * idx] = sp.point
-  pt.vert_index1[idx] = sp.vertex_index
+  pt.vert_index[2 * idx] = sp.vertex_index
   index1 = sp.cached_index
 
   sp = support(geom2, geom2_type, -dir)
-  pt.vert2[idx] = sp.point
+  pt.vert[2 * idx + 1] = sp.point
-  pt.vert_index2[idx] = sp.vertex_index
+  pt.vert_index[2 * idx + 1] = sp.vertex_index
   index2 = sp.cached_index
 
   return index1, index2
@@ -265,9 +269,9 @@ def _det3(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3) -> float:
 
 @wp.func
 def _same_sign(a: float, b: float) -> int:
-  if a > 0 and b > 0:
+  if a > 0.0 and b > 0.0:
     return 1
-  if a < 0 and b < 0:
+  if a < 0.0 and b < 0.0:
     return -1
   return 0
 
@@ -290,28 +294,25 @@ def _project_origin_plane(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3) -> Tuple[wp.vec
   n = wp.cross(diff32, diff21)
   nv = wp.dot(n, v2)
   nn = wp.dot(n, n)
-  if nn == 0:
+  if nn == 0.0:
     return z, 1
-  if nv != 0 and nn > MINVAL:
+  if nv != 0.0 and nn > MINVAL:
-    v = (nv / nn) * n
+    return (nv / nn) * n, 0
-    return v, 0
 
   # n = (v2 - v1) x (v3 - v1)
   n = wp.cross(diff21, diff31)
   nv = wp.dot(n, v1)
   nn = wp.dot(n, n)
-  if nn == 0:
+  if nn == 0.0:
     return z, 1
-  if nv != 0 and nn > MINVAL:
+  if nv != 0.0 and nn > MINVAL:
-    v = (nv / nn) * n
+    return (nv / nn) * n, 0
-    return v, 0
 
   # n = (v1 - v3) x (v2 - v3)
   n = wp.cross(diff31, diff32)
   nv = wp.dot(n, v3)
   nn = wp.dot(n, n)
-  v = (nv / nn) * n
+  return (nv / nn) * n, 0
-  return v, 0
 
 
 @wp.func
@@ -610,14 +611,14 @@ def gjk(
     simplex[n] = simplex1[n] - simplex2[n]
 
     if cutoff == 0.0:
-      if wp.dot(x_k, simplex[n]) > 0:
+      if wp.dot(x_k, simplex[n]) > 0.0:
         result = GJKResult()
         result.dim = 0
         result.dist = FLOAT_MAX
         return result
     elif cutoff < FLOAT_MAX:
       vs = wp.dot(x_k, simplex[n])
-      if wp.dot(x_k, simplex[n]) > 0 and (vs * vs / xnorm) >= cutoff2:
+      if wp.dot(x_k, simplex[n]) > 0.0 and (vs * vs / xnorm) >= cutoff2:
         result = GJKResult()
         result.dim = 0
         result.dist = FLOAT_MAX
@@ -635,7 +636,7 @@ def gjk(
     # remove vertices from the simplex no longer needed
     n = int(0)
     for i in range(4):
-      if coordinates[i] == 0:
+      if coordinates[i] == 0.0:
         continue
 
       simplex[n] = simplex[i]
@@ -692,7 +693,7 @@ def _same_side(p0: wp.vec3, p1: wp.vec3, p2: wp.vec3, p3: wp.vec3) -> bool:
   n = wp.cross(p1 - p0, p2 - p0)
   dot1 = wp.dot(n, p3 - p0)
   dot2 = wp.dot(n, -p0)
-  return (dot1 > 0 and dot2 > 0) or (dot1 < 0 and dot2 < 0)
+  return (dot1 > 0.0 and dot2 > 0.0) or (dot1 < 0.0 and dot2 < 0.0)
 
 
 @wp.func
@@ -750,7 +751,7 @@ def _tri_point_intersect(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3, p: wp.vec3) -> b
   l2 = coordinates[1]
   l3 = coordinates[2]
 
-  if l1 < 0 or l2 < 0 or l3 < 0:
+  if l1 < 0.0 or l2 < 0.0 or l3 < 0.0:
     return False
 
   pr = wp.vec3()
@@ -766,14 +767,14 @@ def _replace_simplex3(pt: Polytope, v1: int, v2: int, v3: int) -> GJKResult:
 
   # reset GJK simplex
   simplex1 = mat43()
-  simplex1[0] = pt.vert1[v1]
+  simplex1[0] = pt.vert[2 * v1]
-  simplex1[1] = pt.vert1[v2]
+  simplex1[1] = pt.vert[2 * v2]
-  simplex1[2] = pt.vert1[v3]
+  simplex1[2] = pt.vert[2 * v3]
 
   simplex2 = mat43()
-  simplex2[0] = pt.vert2[v1]
+  simplex2[0] = pt.vert[2 * v1 + 1]
-  simplex2[1] = pt.vert2[v2]
+  simplex2[1] = pt.vert[2 * v2 + 1]
-  simplex2[2] = pt.vert2[v3]
+  simplex2[2] = pt.vert[2 * v3 + 1]
 
   simplex = mat43()
   simplex[0] = simplex1[0] - simplex2[0]
@@ -781,14 +782,14 @@ def _replace_simplex3(pt: Polytope, v1: int, v2: int, v3: int) -> GJKResult:
   simplex[2] = simplex1[2] - simplex2[2]
 
   simplex_index1 = wp.vec4i()
-  simplex_index1[0] = pt.vert_index1[v1]
+  simplex_index1[0] = pt.vert_index[2 * v1]
-  simplex_index1[1] = pt.vert_index1[v2]
+  simplex_index1[1] = pt.vert_index[2 * v2]
-  simplex_index1[2] = pt.vert_index1[v3]
+  simplex_index1[2] = pt.vert_index[2 * v3]
 
   simplex_index2 = wp.vec4i()
-  simplex_index2[0] = pt.vert_index2[v1]
+  simplex_index2[0] = pt.vert_index[2 * v1 + 1]
-  simplex_index2[1] = pt.vert_index2[v2]
+  simplex_index2[1] = pt.vert_index[2 * v2 + 1]
-  simplex_index2[2] = pt.vert_index2[v3]
+  simplex_index2[2] = pt.vert_index[2 * v3 + 1]
 
   result.simplex = simplex
   result.simplex1 = simplex1
@@ -827,9 +828,9 @@ def _ray_triangle(v1: wp.vec3, v2: wp.vec3, v3: wp.vec3, v4: wp.vec3, v5: wp.vec
   vol2 = _det3(v4 - v1, v5 - v1, v2 - v1)
   vol3 = _det3(v5 - v1, v3 - v1, v2 - v1)
 
-  if vol1 >= 0 and vol2 >= 0 and vol3 >= 0:
+  if vol1 >= 0.0 and vol2 >= 0.0 and vol3 >= 0.0:
     return 1
-  if vol1 <= 0 and vol2 <= 0 and vol3 <= 0:
+  if vol1 <= 0.0 and vol2 <= 0.0 and vol3 <= 0.0:
     return -1
   return 0
 
@@ -867,9 +868,9 @@ def _epa_witness(
 ) -> Tuple[wp.vec3, wp.vec3, float]:
   face = _get_face_verts(pt.face[face_idx])
   # compute affine coordinates for witness points on plane defined by face
-  v1 = pt.vert1[face[0]] - pt.vert2[face[0]]
+  v1 = pt.vert[2 * face[0]] - pt.vert[2 * face[0] + 1]
-  v2 = pt.vert1[face[1]] - pt.vert2[face[1]]
+  v2 = pt.vert[2 * face[1]] - pt.vert[2 * face[1] + 1]
-  v3 = pt.vert1[face[2]] - pt.vert2[face[2]]
+  v3 = pt.vert[2 * face[2]] - pt.vert[2 * face[2] + 1]
 
   coordinates = _tri_affine_coord(v1, v2, v3, pt.face_pr[face_idx])
   l1 = coordinates[0]
@@ -877,18 +878,18 @@ def _epa_witness(
   l3 = coordinates[2]
 
   # face on geom 2
-  v1 = pt.vert2[face[0]]
+  v1 = pt.vert[2 * face[0] + 1]
-  v2 = pt.vert2[face[1]]
+  v2 = pt.vert[2 * face[1] + 1]
-  v3 = pt.vert2[face[2]]
+  v3 = pt.vert[2 * face[2] + 1]
   x2 = wp.vec3()
   x2[0] = v1[0] * l1 + v2[0] * l2 + v3[0] * l3
   x2[1] = v1[1] * l1 + v2[1] * l2 + v3[1] * l3
   x2[2] = v1[2] * l1 + v2[2] * l2 + v3[2] * l3
 
   # correct witness points for hfield geoms
-  i1 = pt.vert_index1[face[0]]
+  i1 = pt.vert_index[2 * face[0]]
-  i2 = pt.vert_index1[face[1]]
+  i2 = pt.vert_index[2 * face[1]]
-  i3 = pt.vert_index1[face[2]]
+  i3 = pt.vert_index[2 * face[2]]
   if geomtype1 == GeomType.HFIELD and (i1 != i2 or i1 != i3):
     # TODO(kbayes): Fix case where geom2 is near bottom of height field or "extreme" prism heights
     n = wp.vec3(0.0, 0.0, 1.0)
@@ -914,7 +915,7 @@ def _epa_witness(
       x2 = sp.point
 
     coordinates2 = _tri_affine_coord(a, b, c, x2)
-    if coordinates2[0] > 0 and coordinates2[1] > 0 and coordinates2[2] > 0:
+    if coordinates2[0] > 0.0 and coordinates2[1] > 0.0 and coordinates2[2] > 0.0:
       x1 = coordinates2[0] * a + coordinates2[1] * b + coordinates2[2] * c
     else:
       p = c
@@ -924,9 +925,9 @@ def _epa_witness(
     return x1, x2, -wp.norm_l2(x1 - x2)
 
   # face on geom 1
-  v1 = pt.vert1[face[0]]
+  v1 = pt.vert[2 * face[0]]
-  v2 = pt.vert1[face[1]]
+  v2 = pt.vert[2 * face[1]]
-  v3 = pt.vert1[face[2]]
+  v3 = pt.vert[2 * face[2]]
   x1 = wp.vec3()
   x1[0] = v1[0] * l1 + v2[0] * l2 + v3[0] * l3
   x1[1] = v1[1] * l1 + v2[1] * l2 + v3[1] * l3
@@ -971,17 +972,15 @@ def _polytope2(
   d3 = R @ d2
 
   # save vertices and get indices for each one
-  pt.vert1[0] = simplex1[0]
+  pt.vert[0] = simplex1[0]
-  pt.vert1[1] = simplex1[1]
+  pt.vert[1] = simplex2[0]
+  pt.vert[2] = simplex1[1]
+  pt.vert[3] = simplex2[1]
 
-  pt.vert_index1[0] = simplex_index1[0]
+  pt.vert_index[0] = simplex_index1[0]
-  pt.vert_index1[1] = simplex_index1[1]
+  pt.vert_index[1] = simplex_index2[0]
-
+  pt.vert_index[2] = simplex_index1[1]
-  pt.vert2[0] = simplex2[0]
+  pt.vert_index[3] = simplex_index2[1]
-  pt.vert2[1] = simplex2[1]
-
-  pt.vert_index2[0] = simplex_index2[0]
-  pt.vert_index2[1] = simplex_index2[1]
 
   _epa_support(pt, 2, geom1, geom2, geomtype1, geomtype2, d1 / wp.norm_l2(d1))
   _epa_support(pt, 3, geom1, geom2, geomtype1, geomtype2, d2 / wp.norm_l2(d2))
@@ -1013,9 +1012,9 @@ def _polytope2(
     return pt, _replace_simplex3(pt, 1, 4, 3)
 
   # check hexahedron is convex
-  v2 = pt.vert1[2] - pt.vert2[2]
+  v2 = pt.vert[4] - pt.vert[5]
-  v3 = pt.vert1[3] - pt.vert2[3]
+  v3 = pt.vert[6] - pt.vert[7]
-  v4 = pt.vert1[4] - pt.vert2[4]
+  v4 = pt.vert[8] - pt.vert[9]
   if not _ray_triangle(simplex[0], simplex[1], v2, v3, v4):
     pt.status = 1
     return pt, GJKResult()
@@ -1049,21 +1048,19 @@ def _polytope3(
     pt.status = 2
     return pt
 
-  pt.vert1[0] = simplex1[0]
+  pt.vert[0] = simplex1[0]
-  pt.vert1[1] = simplex1[1]
+  pt.vert[1] = simplex2[0]
-  pt.vert1[2] = simplex1[2]
+  pt.vert[2] = simplex1[1]
+  pt.vert[3] = simplex2[1]
+  pt.vert[4] = simplex1[2]
+  pt.vert[5] = simplex2[2]
 
-  pt.vert_index1[0] = simplex_index1[0]
+  pt.vert_index[0] = simplex_index1[0]
-  pt.vert_index1[1] = simplex_index1[1]
+  pt.vert_index[1] = simplex_index2[0]
-  pt.vert_index1[2] = simplex_index1[2]
+  pt.vert_index[2] = simplex_index1[1]
-
+  pt.vert_index[3] = simplex_index2[1]
-  pt.vert2[0] = simplex2[0]
+  pt.vert_index[4] = simplex_index1[2]
-  pt.vert2[1] = simplex2[1]
+  pt.vert_index[5] = simplex_index2[2]
-  pt.vert2[2] = simplex2[2]
-
-  pt.vert_index2[0] = simplex_index2[0]
-  pt.vert_index2[1] = simplex_index2[1]
-  pt.vert_index2[2] = simplex_index2[2]
 
   _epa_support(pt, 3, geom1, geom2, geomtype1, geomtype2, -n)
   _epa_support(pt, 4, geom1, geom2, geomtype1, geomtype2, n)
@@ -1071,8 +1068,8 @@ def _polytope3(
   v1 = simplex[0]
   v2 = simplex[1]
   v3 = simplex[2]
-  v4 = pt.vert1[3] - pt.vert2[3]
+  v4 = pt.vert[6] - pt.vert[7]
-  v5 = pt.vert1[4] - pt.vert2[4]
+  v5 = pt.vert[8] - pt.vert[9]
 
   # check that v4 is not contained in the 2-simplex
   if _tri_point_intersect(v1, v2, v3, v4):
@@ -1128,25 +1125,23 @@ def _polytope4(
   simplex_index2: wp.vec4i,
 ) -> Tuple[Polytope, GJKResult]:
   """Create polytope for EPA given a 3-simplex from GJK."""
-  pt.vert1[0] = simplex1[0]
+  pt.vert[0] = simplex1[0]
-  pt.vert1[1] = simplex1[1]
+  pt.vert[1] = simplex2[0]
-  pt.vert1[2] = simplex1[2]
+  pt.vert[2] = simplex1[1]
-  pt.vert1[3] = simplex1[3]
+  pt.vert[3] = simplex2[1]
+  pt.vert[4] = simplex1[2]
+  pt.vert[5] = simplex2[2]
+  pt.vert[6] = simplex1[3]
+  pt.vert[7] = simplex2[3]
 
-  pt.vert_index1[0] = simplex_index1[0]
+  pt.vert_index[0] = simplex_index1[0]
-  pt.vert_index1[1] = simplex_index1[1]
+  pt.vert_index[1] = simplex_index2[0]
-  pt.vert_index1[2] = simplex_index1[2]
+  pt.vert_index[2] = simplex_index1[1]
-  pt.vert_index1[3] = simplex_index1[3]
+  pt.vert_index[3] = simplex_index2[1]
-
+  pt.vert_index[4] = simplex_index1[2]
-  pt.vert2[0] = simplex2[0]
+  pt.vert_index[5] = simplex_index2[2]
-  pt.vert2[1] = simplex2[1]
+  pt.vert_index[6] = simplex_index1[3]
-  pt.vert2[2] = simplex2[2]
+  pt.vert_index[7] = simplex_index2[3]
-  pt.vert2[3] = simplex2[3]
-
-  pt.vert_index2[0] = simplex_index2[0]
-  pt.vert_index2[1] = simplex_index2[1]
-  pt.vert_index2[2] = simplex_index2[2]
-  pt.vert_index2[3] = simplex_index2[3]
 
   # if the origin is on a face, replace the 3-simplex with a 2-simplex
   if _attach_face(pt, 0, 0, 1, 2) < MIN_DIST:
@@ -1249,7 +1244,7 @@ def _epa(
       break
 
     # check if lower bound is 0
-    if lower2 <= 0:
+    if lower2 <= 0.0:
       break
 
     # compute support point w from the closest face's normal
@@ -1257,7 +1252,7 @@ def _epa(
     wi = pt.nvert
     face_pr_normalized = pt.face_pr[idx] / lower
     i1, i2 = _epa_support(pt, wi, geom1, geom2, geomtype1, geomtype2, face_pr_normalized)
-    w = pt.vert1[wi] - pt.vert2[wi]
+    w = pt.vert[2 * wi] - pt.vert[2 * wi + 1]
     geom1.index = i1
     geom2.index = i2
     pt.nvert += 1
@@ -1275,7 +1270,7 @@ def _epa(
     if is_discrete:
       found_repeated = bool(False)
       for i in range(pt.nvert - 1):
-        if pt.vert_index1[i] == pt.vert_index1[wi] and pt.vert_index2[i] == pt.vert_index2[wi]:
+        if pt.vert_index[2 * i] == pt.vert_index[2 * wi] and pt.vert_index[2 * i + 1] == pt.vert_index[2 * wi + 1]:
           found_repeated = True
           break
       if found_repeated:
@@ -1313,7 +1308,7 @@ def _epa(
     for i in range(pt.nhorizon):
       edge = _get_edge(pt.horizon[i])
       dist2 = _attach_face(pt, pt.nface, wi, edge[0], edge[1])
-      if dist2 == 0:
+      if dist2 == 0.0:
         idx = -1
         break
 
@@ -1348,72 +1343,66 @@ def _area4(a: wp.vec3, b: wp.vec3, c: wp.vec3, d: wp.vec3) -> float:
   return 0.5 * wp.norm_l2(wp.cross(a - d, d - b) + wp.cross(b - c, c - a))
 
 
-@wp.func
-def _next(n: int, i: int) -> int:
-  """Returns (i + 1) mod n for 0 <= i <= n - 1."""
-  return wp.where(i == n - 1, 0, i + 1)
-
-
 @wp.func
 def _polygon_quad(polygon: wp.array(dtype=wp.vec3), npolygon: int) -> wp.vec4i:
-  """Returns the indices of a quadrilateral of maximum area in a convex polygon."""
+  """Returns the indices of a quadrilateral of maximum area in a convex polygon (npolygon > 4)."""
-  b = _next(npolygon, 0)
+  b = int(1)
-  c = _next(npolygon, b)
+  c = int(2)
-  d = _next(npolygon, c)
+  d = int(3)
   res = wp.vec4i(0, b, c, d)
   m = _area4(polygon[0], polygon[b], polygon[c], polygon[d])
   for a in range(npolygon):
     while True:
-      m_next = _area4(polygon[a], polygon[b], polygon[c], polygon[_next(npolygon, d)])
+      m_next = _area4(polygon[a], polygon[b], polygon[c], polygon[(d + 1) % npolygon])
       if m_next <= m:
         break
       m = m_next
-      d = _next(npolygon, d)
+      d = (d + 1) % npolygon
       res = wp.vec4i(a, b, c, d)
       while True:
-        m_next = _area4(polygon[a], polygon[b], polygon[_next(npolygon, c)], polygon[d])
+        m_next = _area4(polygon[a], polygon[b], polygon[(c + 1) % npolygon], polygon[d])
         if m_next <= m:
           break
         m = m_next
-        c = _next(npolygon, c)
+        c = (c + 1) % npolygon
         res = wp.vec4i(a, b, c, d)
       while True:
-        m_next = _area4(polygon[a], polygon[_next(npolygon, b)], polygon[c], polygon[d])
+        m_next = _area4(polygon[a], polygon[(b + 1) % npolygon], polygon[c], polygon[d])
         if m_next <= m:
           break
         m = m_next
-        b = _next(npolygon, b)
+        b = (b + 1) % npolygon
         res = wp.vec4i(a, b, c, d)
     if b == a:
-      b = _next(npolygon, b)
+      b = (b + 1) % npolygon
       if c == b:
-        c = _next(npolygon, c)
+        c = (c + 1) % npolygon
         if d == c:
-          d = _next(npolygon, d)
+          d = (d + 1) % npolygon
   return res
 
 
 # return number (1, 2 or 3) of dimensions of a simplex; reorder vertices if necessary
 @wp.func
 def _feature_dim(
-  face: wp.vec3i, vert_index: wp.array(dtype=int), vert: wp.array(dtype=wp.vec3)
+  face: wp.vec3i, vert_index: wp.array(dtype=int), vert: wp.array(dtype=wp.vec3), offset: int
 ) -> Tuple[int, wp.vec3i, wp.mat33]:
-  v1i = vert_index[face[0]]
+  v1i = vert_index[2 * face[0] + offset]
-  v2i = vert_index[face[1]]
+  v2i = vert_index[2 * face[1] + offset]
-  v3i = vert_index[face[2]]
+  v3i = vert_index[2 * face[2] + offset]
 
   feature_index = wp.vec3i(v1i, v2i, v3i)
   feature_vert = wp.mat33()
-  feature_vert[0] = vert[face[0]]
+  feature_vert[0] = vert[2 * face[0] + offset]
-  feature_vert[1] = vert[face[1]]
+  feature_vert[1] = vert[2 * face[1] + offset]
-  feature_vert[2] = vert[face[2]]
+  feature_vert[2] = vert[2 * face[2] + offset]
 
   if v1i != v2i:
     dim = wp.where(v3i == v1i or v3i == v2i, 2, 3)
     return dim, feature_index, feature_vert
 
   feature_index[1] = v3i
-  feature_vert[1] = vert[face[2]]
+  feature_vert[1] = vert[2 * face[2] + offset]
 
   dim = wp.where(v1i != v3i, 2, 1)
   return dim, feature_index, feature_vert
@@ -1675,7 +1664,7 @@ def _box_normals(
     y = float((v1 & 2) and (v2 & 2)) - float(not (v1 & 2) and not (v2 & 2))
     z = float((v1 & 4) and (v2 & 4)) - float(not (v1 & 4) and not (v2 & 4))
     if x != 0.0:
-      normal_out[c] = mat @ wp.vec3(float(x), 0.0, 0.0)
+      normal_out[c] = mat @ wp.vec3(x, 0.0, 0.0)
       index_out[c] = wp.where(x > 0.0, 0, 1)
       c += 1
     if y != 0.0:
@@ -1686,7 +1675,9 @@ def _box_normals(
       normal_out[c] = mat @ wp.vec3(0.0, 0.0, z)
       index_out[c] = wp.where(z > 0.0, 4, 5)
       c += 1
-    if c == 2:
+    # c is 1 if edge is diagonal of a box face
+    # c is 2 if edge is an external edge of box
+    if c == 1 or c == 2:
       return 2
     return _box_normals2(mat, dir, normal_out, index_out)
 
@@ -1824,18 +1815,15 @@ def _halfspace(a: wp.vec3, n: wp.vec3, p: wp.vec3) -> bool:
 
 
 @wp.func
-def _plane_intersect(pn: wp.vec3, pd: float, a: wp.vec3, b: wp.vec3) -> Tuple[float, wp.vec3]:
+def _plane_intersect(pn: wp.vec3, pd: float, a: wp.vec3, b: wp.vec3) -> float:
-  res = wp.vec3()
+  """Returns the parameter t where the line a + t(b - a) intersects the given plane."""
-  ab = b - a
+  dot = wp.dot(pn, b - a)
-  temp = wp.dot(pn, ab)
+
-  if temp == 0.0:
+  # parallel; no intersection
-    return FLOAT_MAX, res  # parallel; no intersection
+  if wp.abs(dot) < 1e-10:
-  t = (pd - wp.dot(pn, a)) / temp
+    return FLOAT_MAX
-  if t >= 0.0 and t <= 1.0:
+
-    res[0] = a[0] + t * ab[0]
+  return (pd - wp.dot(pn, a)) / dot
-    res[1] = a[1] + t * ab[1]
-    res[2] = a[2] + t * ab[2]
-  return t, res
 
 
 # clip a polygon against another polygon
@@ -1901,9 +1889,10 @@ def _polygon_clip(
         continue
 
       # add new vertex to clipped polygon where PQ intersects the clipping edge
-      t, res = _plane_intersect(pn[e], pd[e], P, Q)
+      t = _plane_intersect(pn[e], pd[e], P, Q)
-      if t >= 0.0 and t <= 1.0:
+      if t > -INTERSECT_TOL and t < 1.0 + INTERSECT_TOL:
-        clipped_out[nclipped] = res
+        t = wp.clamp(t, 0.0, 1.0)
+        clipped_out[nclipped] = P + t * (Q - P)
         nclipped += 1
 
       # add Q as PQ is now back inside the clipping edge
@@ -1937,9 +1926,16 @@ def _polygon_clip(
 
 @wp.func
 def _set_edge(
-  vert1: wp.array(dtype=wp.vec3), vert2: wp.array(dtype=wp.vec3), start: int, end: int, face_out: wp.array(dtype=wp.vec3)
+  # In:
+  vert1: wp.array(dtype=wp.vec3),
+  vert2: wp.array(dtype=wp.vec3),
+  start: int,
+  end: int,
+  offset: int,
+  # Out:
+  face_out: wp.array(dtype=wp.vec3),
 ) -> int:
-  face_out[0] = vert1[start]
+  face_out[0] = vert1[2 * start + offset]
   face_out[1] = vert2[end]
   return 2
 
@@ -1959,10 +1955,8 @@ def multicontact(
   endvert: wp.array(dtype=wp.vec3),
   face1: wp.array(dtype=wp.vec3),
   face2: wp.array(dtype=wp.vec3),
-  epa_vert1: wp.array(dtype=wp.vec3),
+  epa_vert: wp.array(dtype=wp.vec3),
-  epa_vert2: wp.array(dtype=wp.vec3),
+  epa_vert_index: wp.array(dtype=int),
-  epa_vert_index1: wp.array(dtype=int),
-  epa_vert_index2: wp.array(dtype=int),
   epa_face: int,
   x1: wp.vec3,
   x2: wp.vec3,
@@ -1998,8 +1992,8 @@ def multicontact(
   face = _get_face_verts(epa_face)
 
   # get dimensions of features of geoms 1 and 2
-  nface1, feature_index1, feature_vertex1 = _feature_dim(face, epa_vert_index1, epa_vert1)
+  nface1, feature_index1, feature_vertex1 = _feature_dim(face, epa_vert_index, epa_vert, 0)
-  nface2, feature_index2, feature_vertex2 = _feature_dim(face, epa_vert_index2, epa_vert2)
+  nface2, feature_index2, feature_vertex2 = _feature_dim(face, epa_vert_index, epa_vert, 1)
 
   dir = x2 - x1
   dir_neg = -dir
@@ -2115,7 +2109,7 @@ def multicontact(
 
   # recover geom1 matching edge or face
   if is_edge_contact_geom1:
-    nface1 = _set_edge(epa_vert1, endvert, face[0], i, face1)
+    nface1 = _set_edge(epa_vert, endvert, face[0], i, 0, face1)
   else:
     ind = wp.where(is_edge_contact_geom2, idx1[j], idx1[i])
     if geomtype1 == GeomType.BOX:
@@ -2136,7 +2130,7 @@ def multicontact(
 
   # recover geom2 matching edge or face
   if is_edge_contact_geom2:
-    nface2 = _set_edge(epa_vert2, endvert, face[0], i, face2)
+    nface2 = _set_edge(epa_vert, endvert, face[0], i, 1, face2)
   else:
     if geomtype2 == GeomType.BOX:
       nface2 = _box_face(geom2.rot, geom2.pos, geom2.size, idx2[j], face2)
@@ -2197,12 +2191,12 @@ def _inflate(
       c = geom1.hfprism[5]
 
       coordinates = _tri_affine_coord(a, b, c, x2)
-      if coordinates[0] > 0 and coordinates[1] > 0 and coordinates[2] > 0:
+      if coordinates[0] > 0.0 and coordinates[1] > 0.0 and coordinates[2] > 0.0:
         x1 = coordinates[0] * a + coordinates[1] * b + coordinates[2] * c
       else:
         p = c
-        p = wp.where(coordinates[1] > 0, b, p)
+        p = wp.where(coordinates[1] > 0.0, b, p)
-        p = wp.where(coordinates[0] > 0, a, p)
+        p = wp.where(coordinates[0] > 0.0, a, p)
         x1 = x2 - wp.dot(x2 - p, n) * n
       dist = -wp.norm_l2(x1 - x2)
       return dist, x1, x2
@@ -2230,10 +2224,8 @@ def ccd(
   geomtype2: int,
   x_1: wp.vec3,
   x_2: wp.vec3,
-  vert1: wp.array(dtype=wp.vec3),
+  vert: wp.array(dtype=wp.vec3),
-  vert2: wp.array(dtype=wp.vec3),
+  vert_index: wp.array(dtype=int),
-  vert_index1: wp.array(dtype=int),
-  vert_index2: wp.array(dtype=int),
   face: wp.array(dtype=int),
   face_pr: wp.array(dtype=wp.vec3),
   face_norm2: wp.array(dtype=float),
@@ -2289,10 +2281,8 @@ def ccd(
   pt.nface = 0
   pt.nvert = 0
   pt.nhorizon = 0
-  pt.vert1 = vert1
+  pt.vert = vert
-  pt.vert2 = vert2
+  pt.vert_index = vert_index
-  pt.vert_index1 = vert_index1
-  pt.vert_index2 = vert_index2
   pt.face = face
   pt.face_pr = face_pr
   pt.face_norm2 = face_norm2
@@ -2358,4 +2348,13 @@ def ccd(
   dist, x1, x2, idx = _epa(tolerance, gjk_iterations, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
   if idx == -1:
     return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1
+
+  # multicontact not supported for margin
+  if geom1.margin != 0.0 or geom2.margin != 0.0:
+    idx = -1
+
+  # multicontact only supported for boxes and meshes
+  if (geomtype1 != GeomType.BOX and geomtype1 != GeomType.MESH) or (geomtype2 != GeomType.BOX and geomtype2 != GeomType.MESH):
+    idx = -1
+
   return dist, 1, x1, x2, idx

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

@@ -32,8 +32,10 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sph
 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
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext
 from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
@@ -43,7 +45,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 
@@ -173,13 +174,13 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup
     convex: Convex geometry object containing position, rotation, and mesh data.
 
   Returns:
-    - Vector of contact distances (wp.inf for unpopulated contacts).
+    - Vector of contact distances (MJ_MAXVAL for unpopulated contacts).
     - Matrix of contact positions (one per row).
     - Matrix of contact normal vectors (one per row).
   """
   _HUGE_VAL = 1e6
 
-  contact_dist = wp.vec4(wp.inf)
+  contact_dist = wp.vec4(MJ_MAXVAL)
   contact_pos = mat43()
   contact_count = int(0)
 
@@ -426,12 +427,12 @@ def write_contact(
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
   nacon_out: wp.array(dtype=int),
-):
+) -> int:
   active = dist_in < margin_in
 
   # skip contact and no collision sensor
   if (pairid_in[0] == -2 or not active) and pairid_in[1] == -1:
-    return
+    return 0
 
   contact_type = 0
 
@@ -457,6 +458,8 @@ def write_contact(
     contact_solimp_out[cid] = solimp_in
     contact_type_out[cid] = contact_type
     contact_geomcollisionid_out[cid] = id_
+    return int(active)
+  return 0
 
 
 @wp.func
@@ -477,10 +480,9 @@ def contact_params(
   pair_margin: wp.array2d(dtype=float),
   pair_gap: wp.array2d(dtype=float),
   pair_friction: wp.array2d(dtype=vec5),
-  # Data in:
+  # In:
   collision_pair_in: wp.array(dtype=wp.vec2i),
   collision_pairid_in: wp.array(dtype=wp.vec2i),
-  # In:
   cid: int,
   worldid: int,
 ):
@@ -1540,6 +1542,7 @@ def box_box_wrapper(
     )
 
 
+# Map of supported primitive collision functions
 _PRIMITIVE_COLLISIONS = {
   (GeomType.PLANE, GeomType.SPHERE): plane_sphere_wrapper,
   (GeomType.PLANE, GeomType.CAPSULE): plane_capsule_wrapper,
@@ -1557,23 +1560,9 @@ _PRIMITIVE_COLLISIONS = {
 }
 
 
-# TODO(team): _check_collisions shared utility
-def _check_primitive_collisions():
-  prev_idx = -1
-  for types in _PRIMITIVE_COLLISIONS.keys():
-    idx = upper_trid_index(len(GeomType), types[0].value, types[1].value)
-    if types[1] < types[0] or idx <= prev_idx:
-      return False
-    prev_idx = idx
-  return True
-
-
-assert _check_primitive_collisions(), "_PRIMITIVE_COLLISIONS is in invalid order"
-
-
 @cache_kernel
 def _primitive_narrowphase(primitive_collisions_types, primitive_collisions_func):
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def primitive_narrowphase(
     # Model:
     geom_type: wp.array(dtype=int),
@@ -1612,10 +1601,11 @@ def _primitive_narrowphase(primitive_collisions_types, primitive_collisions_func
     geom_xpos_in: wp.array2d(dtype=wp.vec3),
     geom_xmat_in: wp.array2d(dtype=wp.mat33),
     naconmax_in: int,
+    ncollision_in: wp.array(dtype=int),
+    # In:
     collision_pair_in: wp.array(dtype=wp.vec2i),
     collision_pairid_in: wp.array(dtype=wp.vec2i),
     collision_worldid_in: wp.array(dtype=int),
-    ncollision_in: wp.array(dtype=int),
     # Data out:
     contact_dist_out: wp.array(dtype=float),
     contact_pos_out: wp.array(dtype=wp.vec3),
@@ -1730,7 +1720,7 @@ _PRIMITIVE_COLLISION_FUNC = []
 
 
 @event_scope
-def primitive_narrowphase(m: Model, d: Data):
+def primitive_narrowphase(m: Model, d: Data, ctx: CollisionContext, collision_table: list[tuple[GeomType, GeomType]]):
   """Runs collision detection on primitive geom pairs discovered during broadphase.
 
   This function processes collision pairs involving primitive shapes that were
@@ -1749,6 +1739,8 @@ def primitive_narrowphase(m: Model, d: Data):
   # for pair types without collisions, as well as updating the launch dimensions.
 
   for types, func in _PRIMITIVE_COLLISIONS.items():
+    if types not in collision_table:
+      continue
     idx = upper_trid_index(len(GeomType), types[0].value, types[1].value)
     if m.geom_pair_type_count[idx] and types not in _PRIMITIVE_COLLISION_TYPES:
       _PRIMITIVE_COLLISION_TYPES.append(types)
@@ -1793,10 +1785,10 @@ def primitive_narrowphase(m: Model, d: Data):
       d.geom_xpos,
       d.geom_xmat,
       d.naconmax,
-      d.collision_pair,
-      d.collision_pairid,
-      d.collision_worldid,
       d.ncollision,
+      ctx.collision_pair,
+      ctx.collision_pairid,
+      ctx.collision_worldid,
     ],
     outputs=[
       d.contact.dist,

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

@@ -18,6 +18,7 @@ from typing import Any, Tuple
 import warp as wp
 
 MJ_MINVAL = 1e-15
+MJ_MAXVAL = 1e10
 
 
 wp.set_module_options({"enable_backward": False})
@@ -411,13 +412,13 @@ def plane_box(
     margin: Collision tolerance.
 
   Returns:
-    - Vector of contact distances (wp.inf for unpopulated contacts).
+    - Vector of contact distances (MJ_MAXVAL for unpopulated contacts).
     - Matrix of contact positions (one per row).
     - Contact normal vector.
   """
   center_dist = wp.dot(box_pos - plane_pos, plane_normal)
 
-  dist = vec8f(wp.inf)
+  dist = vec8f(MJ_MAXVAL)
   pos = mat83f()
 
   # test all corners, pick bottom 4
@@ -540,7 +541,7 @@ def plane_cylinder(
     - Matrix of contact normal vectors (one per row).
   """
   # Initialize output matrices
-  contact_dist = wp.vec4(wp.inf)
+  contact_dist = wp.vec4(MJ_MAXVAL)
   contact_pos = mat43f()
   contact_count = 0
 
@@ -666,14 +667,14 @@ def box_box(
     margin: Collision tolerance.
 
   Returns:
-    - Vector of contact distances (wp.inf for unpopulated contacts).
+    - Vector of contact distances (MJ_MAXVAL for unpopulated contacts).
     - Matrix of contact positions (one per row).
     - Matrix of contact normal vectors (one per row).
   """
   # Initialize output matrices
   contact_dist = vec8f()
   for i in range(8):
-    contact_dist[i] = wp.inf
+    contact_dist[i] = MJ_MAXVAL
   contact_pos = mat83f()
   contact_normals = mat83f()
   contact_count = 0
@@ -1176,7 +1177,7 @@ def capsule_box(
     box_size: Half-extents of the box along each axis.
 
   Returns:
-    - Vector of contact distances (wp.inf for unpopulated contacts).
+    - Vector of contact distances (MJ_MAXVAL for unpopulated contacts).
     - Matrix of contact positions (one per row).
     - Matrix of contact normal vectors (one per row).
   """
@@ -1348,7 +1349,7 @@ def capsule_box(
     c1 = wp.where((ee2 > 0) == w_neg, 1, 2)
 
   if cltype == -4:  # invalid type
-    return wp.vec2(wp.inf), mat23f(), mat23f()
+    return wp.vec2(MJ_MAXVAL), mat23f(), mat23f()
 
   if cltype >= 0 and cltype // 3 != 1:  # closest to a corner of the box
     c1 = axisdir ^ clcorner
@@ -1479,7 +1480,7 @@ def capsule_box(
     # collide with sphere using core function
     dist2, pos2, normal2 = sphere_box(s2_pos_g, capsule_radius, box_pos, box_rot, box_size)
   else:
-    dist2 = wp.inf
+    dist2 = MJ_MAXVAL
     pos2 = wp.vec3()
     normal2 = wp.vec3()
 

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

@@ -22,6 +22,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import geom_col
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import write_contact
 from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
 from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh
+from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
@@ -68,6 +69,7 @@ class MeshData:
   data_id: int
   pos: wp.vec3
   mat: wp.mat33
+  size: wp.vec3
   pnt: wp.vec3
   vec: wp.vec3
   valid: bool = False
@@ -375,6 +377,7 @@ def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: Volume
       mesh_data.data_id,
       mesh_data.pos,
       mesh_data.mat,
+      mesh_data.size,
       mesh_data.pnt,
       mesh_data.vec,
     )
@@ -388,6 +391,7 @@ def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: Volume
         mesh_data.data_id,
         mesh_data.pos,
         mesh_data.mat,
+        mesh_data.size,
         mesh_data.pnt,
         -mesh_data.vec,
       )
@@ -425,6 +429,7 @@ def sdf_grad(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: V
       mesh_data.data_id,
       mesh_data.pos,
       mesh_data.mat,
+      mesh_data.size,
       mesh_data.pnt,
       mesh_data.vec,
     )
@@ -664,11 +669,11 @@ def _sdf_narrowphase(
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
   geom_xmat_in: wp.array2d(dtype=wp.mat33),
   naconmax_in: int,
+  ncollision_in: wp.array(dtype=int),
+  # In:
   collision_pair_in: wp.array(dtype=wp.vec2i),
   collision_pairid_in: wp.array(dtype=wp.vec2i),
   collision_worldid_in: wp.array(dtype=int),
-  ncollision_in: wp.array(dtype=int),
-  # In:
   sdf_initpoints: int,
   sdf_iterations: int,
   # Data out:
@@ -785,6 +790,7 @@ def _sdf_narrowphase(
   mesh_data1.data_id = geom_dataid[g1]
   mesh_data1.pos = geom1.pos
   mesh_data1.mat = geom1.rot
+  mesh_data1.size = geom1.size
   mesh_data1.pnt = wp.vec3(-1.0)
   mesh_data1.vec = wp.vec3(0.0)
   mesh_data1.valid = True
@@ -797,6 +803,7 @@ def _sdf_narrowphase(
   mesh_data2.data_id = geom_dataid[g2]
   mesh_data2.pos = geom2.pos
   mesh_data2.mat = geom2.rot
+  mesh_data2.size = geom2.size
   mesh_data2.pnt = wp.vec3(-1.0)
   mesh_data2.vec = wp.vec3(0.0)
   mesh_data2.valid = True
@@ -859,7 +866,7 @@ def _sdf_narrowphase(
 
 
 @event_scope
-def sdf_narrowphase(m: Model, d: Data):
+def sdf_narrowphase(m: Model, d: Data, ctx: CollisionContext):
   wp.launch(
     _sdf_narrowphase,
     dim=(m.opt.sdf_initpoints, d.naconmax),
@@ -909,10 +916,10 @@ def sdf_narrowphase(m: Model, d: Data):
       d.geom_xpos,
       d.geom_xmat,
       d.naconmax,
-      d.collision_pair,
-      d.collision_pairid,
-      d.collision_worldid,
       d.ncollision,
+      ctx.collision_pair,
+      ctx.collision_pairid,
+      ctx.collision_worldid,
       m.opt.sdf_initpoints,
       m.opt.sdf_iterations,
     ],

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

@@ -34,21 +34,11 @@ def _zero_constraint_counts(
   nf_out: wp.array(dtype=int),
   nl_out: wp.array(dtype=int),
   nefc_out: wp.array(dtype=int),
-  ne_connect_out: wp.array(dtype=int),
-  ne_weld_out: wp.array(dtype=int),
-  ne_jnt_out: wp.array(dtype=int),
-  ne_ten_out: wp.array(dtype=int),
-  ne_flex_out: wp.array(dtype=int),
 ):
   worldid = wp.tid()
 
   # Zero all constraint counters
   ne_out[worldid] = 0
-  ne_connect_out[worldid] = 0
-  ne_weld_out[worldid] = 0
-  ne_jnt_out[worldid] = 0
-  ne_ten_out[worldid] = 0
-  ne_flex_out[worldid] = 0
   nf_out[worldid] = 0
   nl_out[worldid] = 0
   nefc_out[worldid] = 0
@@ -155,6 +145,7 @@ def _efc_equality_connect(
   # In:
   refsafe_in: int,
   # Data out:
+  ne_out: wp.array(dtype=int),
   nefc_out: wp.array(dtype=int),
   efc_type_out: wp.array2d(dtype=int),
   efc_id_out: wp.array2d(dtype=int),
@@ -165,7 +156,6 @@ def _efc_equality_connect(
   efc_vel_out: wp.array2d(dtype=float),
   efc_aref_out: wp.array2d(dtype=float),
   efc_frictionloss_out: wp.array2d(dtype=float),
-  ne_connect_out: wp.array(dtype=int),
 ):
   """Calculates constraint rows for connect equality constraints."""
   worldid, eqconnectid = wp.tid()
@@ -174,7 +164,7 @@ def _efc_equality_connect(
   if not eq_active_in[worldid, eqid]:
     return
 
-  wp.atomic_add(ne_connect_out, worldid, 3)
+  wp.atomic_add(ne_out, worldid, 3)
   efcid = wp.atomic_add(nefc_out, worldid, 3)
 
   if efcid + 3 >= njmax_in:
@@ -205,7 +195,7 @@ def _efc_equality_connect(
   # compute Jacobian difference (opposite of contact: 0 - 1)
   Jqvel = wp.vec3f(0.0, 0.0, 0.0)
   for dofid in range(nv):  # TODO: parallelize
-    jacp1, _ = support.jac(
+    jacp1, _ = support.jac_dof(
       body_parentid,
       body_rootid,
       dof_bodyid,
@@ -216,7 +206,7 @@ def _efc_equality_connect(
       dofid,
       worldid,
     )
-    jacp2, _ = support.jac(
+    jacp2, _ = support.jac_dof(
       body_parentid,
       body_rootid,
       dof_bodyid,
@@ -293,6 +283,7 @@ def _efc_equality_joint(
   # In:
   refsafe_in: int,
   # Data out:
+  ne_out: wp.array(dtype=int),
   nefc_out: wp.array(dtype=int),
   efc_type_out: wp.array2d(dtype=int),
   efc_id_out: wp.array2d(dtype=int),
@@ -303,7 +294,6 @@ def _efc_equality_joint(
   efc_vel_out: wp.array2d(dtype=float),
   efc_aref_out: wp.array2d(dtype=float),
   efc_frictionloss_out: wp.array2d(dtype=float),
-  ne_jnt_out: wp.array(dtype=int),
 ):
   worldid, eqjntid = wp.tid()
   eqid = eq_jnt_adr[eqjntid]
@@ -311,7 +301,7 @@ def _efc_equality_joint(
   if not eq_active_in[worldid, eqid]:
     return
 
-  wp.atomic_add(ne_jnt_out, worldid, 1)
+  wp.atomic_add(ne_out, worldid, 1)
   efcid = wp.atomic_add(nefc_out, worldid, 1)
 
   if efcid >= njmax_in:
@@ -399,6 +389,7 @@ def _efc_equality_tendon(
   # In:
   refsafe_in: int,
   # Data out:
+  ne_out: wp.array(dtype=int),
   nefc_out: wp.array(dtype=int),
   efc_type_out: wp.array2d(dtype=int),
   efc_id_out: wp.array2d(dtype=int),
@@ -409,7 +400,6 @@ def _efc_equality_tendon(
   efc_vel_out: wp.array2d(dtype=float),
   efc_aref_out: wp.array2d(dtype=float),
   efc_frictionloss_out: wp.array2d(dtype=float),
-  ne_ten_out: wp.array(dtype=int),
 ):
   worldid, eqtenid = wp.tid()
   eqid = eq_ten_adr[eqtenid]
@@ -417,7 +407,7 @@ def _efc_equality_tendon(
   if not eq_active_in[worldid, eqid]:
     return
 
-  wp.atomic_add(ne_ten_out, worldid, 1)
+  wp.atomic_add(ne_out, worldid, 1)
   efcid = wp.atomic_add(nefc_out, worldid, 1)
 
   if efcid >= njmax_in:
@@ -494,6 +484,9 @@ def _efc_equality_flex(
   opt_timestep: wp.array(dtype=float),
   flexedge_length0: wp.array(dtype=float),
   flexedge_invweight0: wp.array(dtype=float),
+  flexedge_J_rownnz: wp.array(dtype=int),
+  flexedge_J_rowadr: wp.array(dtype=int),
+  flexedge_J_colind: wp.array(dtype=int),
   eq_solref: wp.array2d(dtype=wp.vec2),
   eq_solimp: wp.array2d(dtype=vec5),
   eq_flex_adr: wp.array(dtype=int),
@@ -505,6 +498,7 @@ def _efc_equality_flex(
   # In:
   refsafe_in: int,
   # Data out:
+  ne_out: wp.array(dtype=int),
   nefc_out: wp.array(dtype=int),
   efc_type_out: wp.array2d(dtype=int),
   efc_id_out: wp.array2d(dtype=int),
@@ -515,12 +509,11 @@ def _efc_equality_flex(
   efc_vel_out: wp.array2d(dtype=float),
   efc_aref_out: wp.array2d(dtype=float),
   efc_frictionloss_out: wp.array2d(dtype=float),
-  ne_flex_out: wp.array(dtype=int),
 ):
   worldid, eqflexid, edgeid = wp.tid()
   eqid = eq_flex_adr[eqflexid]
 
-  wp.atomic_add(ne_flex_out, worldid, 1)
+  wp.atomic_add(ne_out, worldid, 1)
   efcid = wp.atomic_add(nefc_out, worldid, 1)
 
   if efcid >= njmax_in:
@@ -531,10 +524,20 @@ def _efc_equality_flex(
   solimp = eq_solimp[worldid % eq_solimp.shape[0], eqid]
 
   Jqvel = float(0.0)
+
+  # TODO(team): remove once efc.J is sparse
   for i in range(nv):
-    J = flexedge_J_in[worldid, edgeid, i]
+    efc_J_out[worldid, efcid, i] = 0.0
-    efc_J_out[worldid, efcid, i] = J
+
-    Jqvel += J * qvel_in[worldid, i]
+  rownnz = flexedge_J_rownnz[edgeid]
+  rowadr = flexedge_J_rowadr[edgeid]
+  for i in range(rownnz):
+    sparseid = rowadr + i
+    colind = flexedge_J_colind[sparseid]
+    J = flexedge_J_in[worldid, 0, sparseid]
+    # TODO(team): sparse efc.J
+    efc_J_out[worldid, efcid, colind] = J
+    Jqvel += J * qvel_in[worldid, colind]
 
   _update_efc_row(
     worldid,
@@ -748,6 +751,7 @@ def _efc_equality_weld(
   # In:
   refsafe_in: int,
   # Data out:
+  ne_out: wp.array(dtype=int),
   nefc_out: wp.array(dtype=int),
   efc_type_out: wp.array2d(dtype=int),
   efc_id_out: wp.array2d(dtype=int),
@@ -758,7 +762,6 @@ def _efc_equality_weld(
   efc_vel_out: wp.array2d(dtype=float),
   efc_aref_out: wp.array2d(dtype=float),
   efc_frictionloss_out: wp.array2d(dtype=float),
-  ne_weld_out: wp.array(dtype=int),
 ):
   worldid, eqweldid = wp.tid()
   eqid = eq_wld_adr[eqweldid]
@@ -766,7 +769,7 @@ def _efc_equality_weld(
   if not eq_active_in[worldid, eqid]:
     return
 
-  wp.atomic_add(ne_weld_out, worldid, 6)
+  wp.atomic_add(ne_out, worldid, 6)
   efcid = wp.atomic_add(nefc_out, worldid, 6)
 
   if efcid + 6 >= njmax_in:
@@ -808,7 +811,7 @@ def _efc_equality_weld(
   Jqvelr = wp.vec3f(0.0, 0.0, 0.0)
 
   for dofid in range(nv):  # TODO: parallelize
-    jacp1, jacr1 = support.jac(
+    jacp1, jacr1 = support.jac_dof(
       body_parentid,
       body_rootid,
       dof_bodyid,
@@ -819,7 +822,7 @@ def _efc_equality_weld(
       dofid,
       worldid,
     )
-    jacp2, jacr2 = support.jac(
+    jacp2, jacr2 = support.jac_dof(
       body_parentid,
       body_rootid,
       dof_bodyid,
@@ -1215,8 +1218,12 @@ def _efc_contact_pyramidal(
   opt_impratio_invsqrt: wp.array(dtype=float),
   body_parentid: wp.array(dtype=int),
   body_rootid: wp.array(dtype=int),
+  body_weldid: wp.array(dtype=int),
+  body_dofnum: wp.array(dtype=int),
+  body_dofadr: wp.array(dtype=int),
   body_invweight0: wp.array2d(dtype=wp.vec2),
   dof_bodyid: wp.array(dtype=int),
+  dof_parentid: wp.array(dtype=int),
   geom_bodyid: wp.array(dtype=int),
   # Data in:
   qvel_in: wp.array2d(dtype=float),
@@ -1302,49 +1309,73 @@ def _efc_contact_pyramidal(
       invweight = invweight * 2.0 * fri0 * fri0 * impratio_invsqrt * impratio_invsqrt
 
     Jqvel = float(0.0)
-    for i in range(nv):
+
-      J = float(0.0)
+    # skip fixed bodies
-      Ji = float(0.0)
+    body1 = body_weldid[body1]
-      jac1p, jac1r = support.jac(
+    body2 = body_weldid[body2]
-        body_parentid,
+
-        body_rootid,
+    da1 = body_dofadr[body1] + body_dofnum[body1] - 1
-        dof_bodyid,
+    da2 = body_dofadr[body2] + body_dofnum[body2] - 1
-        subtree_com_in,
+    da = wp.max(da1, da2)
-        cdof_in,
+
-        con_pos,
+    for dofid in range(nv - 1, -1, -1):
-        body1,
+      if dofid == da:
-        i,
+        # TODO(team): contact_jacobian
-        worldid,
+        jac1p, jac1r = support.jac_dof(
-      )
+          body_parentid,
-      jac2p, jac2r = support.jac(
+          body_rootid,
-        body_parentid,
+          dof_bodyid,
-        body_rootid,
+          subtree_com_in,
-        dof_bodyid,
+          cdof_in,
-        subtree_com_in,
+          con_pos,
-        cdof_in,
+          body1,
-        con_pos,
+          dofid,
-        body2,
+          worldid,
-        i,
+        )
-        worldid,
+        jac2p, jac2r = support.jac_dof(
-      )
+          body_parentid,
-      jacp_dif = jac2p - jac1p
+          body_rootid,
-      for xyz in range(3):
+          dof_bodyid,
-        J += frame[0, xyz] * jacp_dif[xyz]
+          subtree_com_in,
+          cdof_in,
+          con_pos,
+          body2,
+          dofid,
+          worldid,
+        )
+
+        J = float(0.0)
+        Ji = float(0.0)
+        if condim > 1:
+          dimid2 = dimid / 2 + 1
+
+        for xyz in range(3):
+          jacp_dif = jac2p[xyz] - jac1p[xyz]
+          J += frame[0, xyz] * jacp_dif
+
+          if condim > 1:
+            if dimid2 < 3:
+              Ji += frame[dimid2, xyz] * jacp_dif
+            else:
+              Ji += frame[dimid2 - 3, xyz] * (jac2r[xyz] - jac1r[xyz])
 
         if condim > 1:
-          if dimid2 < 3:
+          if dimid % 2 == 0:
-            Ji += frame[dimid2, xyz] * jacp_dif[xyz]
+            J += Ji * frii
           else:
-            Ji += frame[dimid2 - 3, xyz] * (jac2r[xyz] - jac1r[xyz])
+            J -= Ji * frii
 
-      if condim > 1:
+        efc_J_out[worldid, efcid, dofid] = J
-        if dimid % 2 == 0:
+        Jqvel += J * qvel_in[worldid, dofid]
-          J += Ji * frii
-        else:
-          J -= Ji * frii
 
-      efc_J_out[worldid, efcid, i] = J
+        # Advance tree pointers and recompute da for next iteration
-      Jqvel += J * qvel_in[worldid, i]
+        if da1 == da:
+          da1 = dof_parentid[da1]
+        if da2 == da:
+          da2 = dof_parentid[da2]
+        da = wp.max(da1, da2)
+      else:
+        efc_J_out[worldid, efcid, dofid] = 0.0
 
     if condim == 1:
       efc_type = ConstraintType.CONTACT_FRICTIONLESS
@@ -1385,8 +1416,12 @@ def _efc_contact_elliptic(
   opt_impratio_invsqrt: wp.array(dtype=float),
   body_parentid: wp.array(dtype=int),
   body_rootid: wp.array(dtype=int),
+  body_weldid: wp.array(dtype=int),
+  body_dofnum: wp.array(dtype=int),
+  body_dofadr: wp.array(dtype=int),
   body_invweight0: wp.array2d(dtype=wp.vec2),
   dof_bodyid: wp.array(dtype=int),
+  dof_parentid: wp.array(dtype=int),
   geom_bodyid: wp.array(dtype=int),
   # Data in:
   qvel_in: wp.array2d(dtype=float),
@@ -1450,49 +1485,69 @@ def _efc_contact_elliptic(
     impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
     contact_efc_address_out[conid, dimid] = efcid
 
+    con_pos = pos_in[conid]
+    frame = frame_in[conid]
+
     geom = geom_in[conid]
     body1 = geom_bodyid[geom[0]]
     body2 = geom_bodyid[geom[1]]
 
-    cpos = pos_in[conid]
-    frame = frame_in[conid]
-
-    # TODO(team): parallelize J and Jqvel computation?
     Jqvel = float(0.0)
-    for i in range(nv):
-      J = float(0.0)
-      jac1p, jac1r = support.jac(
-        body_parentid,
-        body_rootid,
-        dof_bodyid,
-        subtree_com_in,
-        cdof_in,
-        cpos,
-        body1,
-        i,
-        worldid,
-      )
-      jac2p, jac2r = support.jac(
-        body_parentid,
-        body_rootid,
-        dof_bodyid,
-        subtree_com_in,
-        cdof_in,
-        cpos,
-        body2,
-        i,
-        worldid,
-      )
-      for xyz in range(3):
-        if dimid < 3:
-          jac_dif = jac2p[xyz] - jac1p[xyz]
-          J += frame[dimid, xyz] * jac_dif
-        else:
-          jac_dif = jac2r[xyz] - jac1r[xyz]
-          J += frame[dimid - 3, xyz] * jac_dif
 
-      efc_J_out[worldid, efcid, i] = J
+    # skip fixed bodies
-      Jqvel += J * qvel_in[worldid, i]
+    body1 = body_weldid[body1]
+    body2 = body_weldid[body2]
+
+    da1 = body_dofadr[body1] + body_dofnum[body1] - 1
+    da2 = body_dofadr[body2] + body_dofnum[body2] - 1
+    da = wp.max(da1, da2)
+
+    for dofid in range(nv - 1, -1, -1):
+      if dofid == da:
+        # TODO(team): contact jacobian
+        jac1p, jac1r = support.jac_dof(
+          body_parentid,
+          body_rootid,
+          dof_bodyid,
+          subtree_com_in,
+          cdof_in,
+          con_pos,
+          body1,
+          dofid,
+          worldid,
+        )
+        jac2p, jac2r = support.jac_dof(
+          body_parentid,
+          body_rootid,
+          dof_bodyid,
+          subtree_com_in,
+          cdof_in,
+          con_pos,
+          body2,
+          dofid,
+          worldid,
+        )
+
+        J = float(0.0)
+        for xyz in range(3):
+          if dimid < 3:
+            jac_dif = jac2p[xyz] - jac1p[xyz]
+            J += frame[dimid, xyz] * jac_dif
+          else:
+            jac_dif = jac2r[xyz] - jac1r[xyz]
+            J += frame[dimid - 3, xyz] * jac_dif
+
+        efc_J_out[worldid, efcid, dofid] = J
+        Jqvel += J * qvel_in[worldid, dofid]
+
+        # Advance tree pointers and recompute da for next iteration
+        if da1 == da:
+          da1 = dof_parentid[da1]
+        if da2 == da:
+          da2 = dof_parentid[da2]
+        da = wp.max(da1, da2)
+      else:
+        efc_J_out[worldid, efcid, dofid] = 0.0
 
     body_invweight0_id = worldid % body_invweight0.shape[0]
     invweight = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0]
@@ -1549,29 +1604,13 @@ def _efc_contact_elliptic(
     )
 
 
-@wp.kernel
-def _num_equality(
-  # Data in:
-  ne_connect_in: wp.array(dtype=int),
-  ne_weld_in: wp.array(dtype=int),
-  ne_jnt_in: wp.array(dtype=int),
-  ne_ten_in: wp.array(dtype=int),
-  ne_flex_in: wp.array(dtype=int),
-  # Data out:
-  ne_out: wp.array(dtype=int),
-):
-  worldid = wp.tid()
-  ne = ne_connect_in[worldid] + ne_weld_in[worldid] + ne_jnt_in[worldid] + ne_ten_in[worldid] + ne_flex_in[worldid]
-  ne_out[worldid] = ne
-
-
 @event_scope
 def make_constraint(m: types.Model, d: types.Data):
   """Creates constraint jacobians and other supporting data."""
   wp.launch(
     _zero_constraint_counts,
     dim=d.nworld,
-    inputs=[d.ne, d.nf, d.nl, d.nefc, d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten, d.ne_flex],
+    inputs=[d.ne, d.nf, d.nl, d.nefc],
   )
 
   if not (m.opt.disableflags & types.DisableBit.CONSTRAINT):
@@ -1608,6 +1647,7 @@ def make_constraint(m: types.Model, d: types.Data):
           refsafe,
         ],
         outputs=[
+          d.ne,
           d.nefc,
           d.efc.type,
           d.efc.id,
@@ -1618,7 +1658,6 @@ def make_constraint(m: types.Model, d: types.Data):
           d.efc.vel,
           d.efc.aref,
           d.efc.frictionloss,
-          d.ne_connect,
         ],
       )
       wp.launch(
@@ -1653,6 +1692,7 @@ def make_constraint(m: types.Model, d: types.Data):
           refsafe,
         ],
         outputs=[
+          d.ne,
           d.nefc,
           d.efc.type,
           d.efc.id,
@@ -1663,7 +1703,6 @@ def make_constraint(m: types.Model, d: types.Data):
           d.efc.vel,
           d.efc.aref,
           d.efc.frictionloss,
-          d.ne_weld,
         ],
       )
       wp.launch(
@@ -1689,6 +1728,7 @@ def make_constraint(m: types.Model, d: types.Data):
           refsafe,
         ],
         outputs=[
+          d.ne,
           d.nefc,
           d.efc.type,
           d.efc.id,
@@ -1699,7 +1739,6 @@ def make_constraint(m: types.Model, d: types.Data):
           d.efc.vel,
           d.efc.aref,
           d.efc.frictionloss,
-          d.ne_jnt,
         ],
       )
       wp.launch(
@@ -1724,6 +1763,7 @@ def make_constraint(m: types.Model, d: types.Data):
           refsafe,
         ],
         outputs=[
+          d.ne,
           d.nefc,
           d.efc.type,
           d.efc.id,
@@ -1734,7 +1774,6 @@ def make_constraint(m: types.Model, d: types.Data):
           d.efc.vel,
           d.efc.aref,
           d.efc.frictionloss,
-          d.ne_ten,
         ],
       )
 
@@ -1746,6 +1785,9 @@ def make_constraint(m: types.Model, d: types.Data):
           m.opt.timestep,
           m.flexedge_length0,
           m.flexedge_invweight0,
+          m.flexedge_J_rownnz,
+          m.flexedge_J_rowadr,
+          m.flexedge_J_colind,
           m.eq_solref,
           m.eq_solimp,
           m.eq_flex_adr,
@@ -1756,6 +1798,7 @@ def make_constraint(m: types.Model, d: types.Data):
           refsafe,
         ],
         outputs=[
+          d.ne,
           d.nefc,
           d.efc.type,
           d.efc.id,
@@ -1766,17 +1809,9 @@ def make_constraint(m: types.Model, d: types.Data):
           d.efc.vel,
           d.efc.aref,
           d.efc.frictionloss,
-          d.ne_flex,
         ],
       )
 
-      wp.launch(
-        _num_equality,
-        dim=d.nworld,
-        inputs=[d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten, d.ne_flex],
-        outputs=[d.ne],
-      )
-
     if not (m.opt.disableflags & types.DisableBit.FRICTIONLOSS):
       wp.launch(
         _efc_friction_dof,
@@ -1957,8 +1992,12 @@ def make_constraint(m: types.Model, d: types.Data):
             m.opt.impratio_invsqrt,
             m.body_parentid,
             m.body_rootid,
+            m.body_weldid,
+            m.body_dofnum,
+            m.body_dofadr,
             m.body_invweight0,
             m.dof_bodyid,
+            m.dof_parentid,
             m.geom_bodyid,
             d.qvel,
             d.subtree_com,
@@ -2002,8 +2041,12 @@ def make_constraint(m: types.Model, d: types.Data):
             m.opt.impratio_invsqrt,
             m.body_parentid,
             m.body_rootid,
+            m.body_weldid,
+            m.body_dofnum,
+            m.body_dofadr,
             m.body_invweight0,
             m.dof_bodyid,
+            m.dof_parentid,
             m.geom_bodyid,
             d.qvel,
             d.subtree_com,

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

@@ -25,7 +25,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 
@@ -80,12 +79,12 @@ def _qderiv_actuator_passive_vel(
 
 @cache_kernel
 def _qderiv_actuator_passive_actuation_dense(tile: TileSet, nu: int):
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Data in:
-    vel_in: wp.array3d(dtype=float),
     actuator_moment_in: wp.array3d(dtype=float),
     # In:
+    vel_in: wp.array3d(dtype=float),
     adr: wp.array(dtype=int),
     # Out:
     qDeriv_out: wp.array3d(dtype=float),
@@ -140,8 +139,8 @@ def _qderiv_actuator_passive(
   # Model:
   opt_timestep: wp.array(dtype=float),
   opt_disableflags: int,
-  opt_is_sparse: bool,
   dof_damping: wp.array2d(dtype=float),
+  is_sparse: bool,
   # Data in:
   qM_in: wp.array3d(dtype=float),
   # In:
@@ -156,7 +155,7 @@ def _qderiv_actuator_passive(
   dofiid = qMi[elemid]
   dofjid = qMj[elemid]
 
-  if opt_is_sparse:
+  if is_sparse:
     qderiv = qDeriv_in[worldid, 0, elemid]
   else:
     qderiv = qDeriv_in[worldid, dofiid, dofjid]
@@ -166,7 +165,7 @@ def _qderiv_actuator_passive(
 
   qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
 
-  if opt_is_sparse:
+  if is_sparse:
     qDeriv_out[worldid, 0, elemid] = qM_in[worldid, 0, elemid] - qderiv
   else:
     qM = qM_in[worldid, dofiid, dofjid] - qderiv
@@ -181,8 +180,8 @@ def _qderiv_tendon_damping(
   # Model:
   ntendon: int,
   opt_timestep: wp.array(dtype=float),
-  opt_is_sparse: bool,
   tendon_damping: wp.array2d(dtype=float),
+  is_sparse: bool,
   # Data in:
   ten_J_in: wp.array3d(dtype=float),
   # In:
@@ -202,7 +201,7 @@ def _qderiv_tendon_damping(
 
   qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
 
-  if opt_is_sparse:
+  if is_sparse:
     qDeriv_out[worldid, 0, elemid] -= qderiv
   else:
     qDeriv_out[worldid, dofiid, dofjid] -= qderiv
@@ -245,7 +244,7 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
         ],
         outputs=[vel],
       )
-      if m.opt.is_sparse:
+      if m.is_sparse:
         wp.launch(
           _qderiv_actuator_passive_actuation_sparse,
           dim=(d.nworld, qMi.size),
@@ -258,9 +257,9 @@ 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=[vel_3d, d.actuator_moment, tile.adr],
+            inputs=[d.actuator_moment, vel_3d, tile.adr],
             outputs=[out],
-            block_dim=m.block_dim.mul_m_dense,
+            block_dim=m.block_dim.qderiv_actuator_dense,
           )
     wp.launch(
       _qderiv_actuator_passive,
@@ -268,8 +267,8 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
       inputs=[
         m.opt.timestep,
         m.opt.disableflags,
-        m.opt.is_sparse,
         m.dof_damping,
+        m.is_sparse,
         d.qM,
         qMi,
         qMj,
@@ -285,7 +284,7 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
     wp.launch(
       _qderiv_tendon_damping,
       dim=(d.nworld, qMi.size),
-      inputs=[m.ntendon, m.opt.timestep, m.opt.is_sparse, m.tendon_damping, d.ten_J, qMi, qMj],
+      inputs=[m.ntendon, m.opt.timestep, m.tendon_damping, m.is_sparse, d.ten_J, qMi, qMj],
       outputs=[out],
     )
 

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

@@ -42,7 +42,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 
@@ -291,17 +290,17 @@ def _euler_damp_qfrc_sparse(
 
 @cache_kernel
 def _tile_euler_dense(tile: TileSet):
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def euler_dense(
     # Model:
-    dof_damping: wp.array2d(dtype=float),
     opt_timestep: wp.array(dtype=float),
+    dof_damping: wp.array2d(dtype=float),
     # Data in:
     qM_in: wp.array3d(dtype=float),
     efc_Ma_in: wp.array2d(dtype=float),
     # In:
     adr_in: wp.array(dtype=int),
-    # Out:
+    # Data out:
     qacc_out: wp.array2d(dtype=float),
   ):
     worldid, nodeid = wp.tid()
@@ -328,7 +327,7 @@ def euler(m: Model, d: Data):
   # integrate damping implicitly
   if not m.opt.disableflags & (DisableBit.EULERDAMP | DisableBit.DAMPER):
     qacc = wp.empty((d.nworld, m.nv), dtype=float)
-    if m.opt.is_sparse:
+    if m.is_sparse:
       qM = wp.clone(d.qM)
       qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
       qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float)
@@ -344,7 +343,7 @@ def euler(m: Model, d: Data):
         wp.launch_tiled(
           _tile_euler_dense(tile),
           dim=(d.nworld, tile.adr.size),
-          inputs=[m.dof_damping, m.opt.timestep, d.qM, d.efc.Ma, tile.adr],
+          inputs=[m.opt.timestep, m.dof_damping, d.qM, d.efc.Ma, tile.adr],
           outputs=[qacc],
           block_dim=m.block_dim.euler_dense,
         )
@@ -482,7 +481,7 @@ def rungekutta4(m: Model, d: Data):
 def implicit(m: Model, d: Data):
   """Integrates fully implicit in velocity."""
   if ~(m.opt.disableflags | ~(DisableBit.ACTUATION | DisableBit.SPRING | DisableBit.DAMPER)):
-    if m.opt.is_sparse:
+    if m.is_sparse:
       qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
       qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
     else:
@@ -524,7 +523,7 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
 # TODO(team): sparse actuator_moment version
 @cache_kernel
 def _actuator_velocity(nv: int):
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def actuator_velocity(
     # Data in:
     qvel_in: wp.array2d(dtype=float),
@@ -544,7 +543,7 @@ def _actuator_velocity(nv: int):
 
 @cache_kernel
 def _tendon_velocity(nv: int):
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def tendon_velocity(
     # Data in:
     qvel_in: wp.array2d(dtype=float),

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

@@ -99,7 +99,7 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d(dtype=float)):
       outputs=[qfrc],
     )
   elif m.opt.integrator == IntegratorType.IMPLICITFAST:
-    if m.opt.is_sparse:
+    if m.is_sparse:
       qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
     else:
       qDeriv = wp.empty((d.nworld, m.nv, m.nv), dtype=float)
@@ -120,10 +120,8 @@ def inv_constraint(m: Model, d: Data):
     d.qfrc_constraint.zero_()
     return
 
-  # update
+  ctx = solver.create_inverse_context(m, d)
-  h = wp.empty((d.nworld, 0, 0), dtype=float)  # not used
+  solver.init_context(m, d, ctx, grad=False)
-  hfactor = wp.empty((d.nworld, 0, 0), dtype=float)  # not used
-  solver.create_context(m, d, h, hfactor, grad=False)
 
 
 def inverse(m: Model, d: Data):

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

@@ -0,0 +1,178 @@
+# Copyright 2026 The Newton Developers
+#
+# 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 warp as wp
+
+from mujoco.mjx.third_party.mujoco_warp._src import types
+from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
+from mujoco.mjx.third_party.mujoco_warp._src.types import EqType
+from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
+
+
+@wp.kernel
+def _tree_edges(
+  # Model:
+  nv: int,
+  body_treeid: wp.array(dtype=int),
+  jnt_dofadr: wp.array(dtype=int),
+  dof_treeid: wp.array(dtype=int),
+  geom_bodyid: wp.array(dtype=int),
+  site_bodyid: wp.array(dtype=int),
+  eq_type: wp.array(dtype=int),
+  eq_obj1id: wp.array(dtype=int),
+  eq_obj2id: wp.array(dtype=int),
+  eq_objtype: wp.array(dtype=int),
+  # Data in:
+  nefc_in: wp.array(dtype=int),
+  contact_geom_in: wp.array(dtype=wp.vec2i),
+  efc_type_in: wp.array2d(dtype=int),
+  efc_id_in: wp.array2d(dtype=int),
+  efc_J_in: wp.array3d(dtype=float),
+  njmax_in: int,
+  # Out:
+  tree_tree: wp.array3d(dtype=int),  # kernel_analyzer: off
+):
+  """Find tree edges. Launch: (nworld, njmax)."""
+  worldid, efcid = wp.tid()
+
+  # skip if beyond active constraints
+  if efcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  efc_type = efc_type_in[worldid, efcid]
+  efc_id = efc_id_in[worldid, efcid]
+
+  tree0 = int(-1)
+  tree1 = int(-1)
+  use_generic = int(0)
+
+  # equality (connect/weld)
+  if efc_type == ConstraintType.EQUALITY:
+    eq_t = eq_type[efc_id]
+
+    if eq_t == EqType.CONNECT or eq_t == EqType.WELD:
+      b1 = eq_obj1id[efc_id]
+      b2 = eq_obj2id[efc_id]
+
+      # site semantics
+      if eq_objtype[efc_id] == ObjType.SITE:
+        b1 = site_bodyid[b1]
+        b2 = site_bodyid[b2]
+
+      tree0 = body_treeid[b1]
+      tree1 = body_treeid[b2]
+    else:
+      # JOINT, TENDON, FLEX
+      use_generic = 1
+
+  # joint friction
+  elif efc_type == ConstraintType.FRICTION_DOF:
+    tree0 = dof_treeid[efc_id]
+
+  # joint limit
+  elif efc_type == ConstraintType.LIMIT_JOINT:
+    tree0 = dof_treeid[jnt_dofadr[efc_id]]
+
+  # contact
+  elif (
+    efc_type == ConstraintType.CONTACT_FRICTIONLESS
+    or efc_type == ConstraintType.CONTACT_PYRAMIDAL
+    or efc_type == ConstraintType.CONTACT_ELLIPTIC
+  ):
+    geom_pair = contact_geom_in[efc_id]
+    g1 = geom_pair[0]
+    g2 = geom_pair[1]
+
+    # flex contacts have negative geom ids
+    if g1 >= 0 and g2 >= 0:
+      tree0 = body_treeid[geom_bodyid[g1]]
+      tree1 = body_treeid[geom_bodyid[g2]]
+    else:
+      use_generic = 1
+
+  # generic
+  else:
+    use_generic = 1
+
+  # handle static bodies
+  if use_generic == 0:
+    # swap so tree0 is non-negative if possible
+    if tree0 < 0 and tree1 >= 0:
+      tree0 = tree1
+      tree1 = -1
+
+    # mark the edge
+    if tree0 >= 0:
+      if tree1 < 0 or tree0 == tree1:
+        # self-edge
+        wp.atomic_max(tree_tree, worldid, tree0, tree0, 1)
+      else:
+        # cross-tree edge
+        t1 = wp.min(tree0, tree1)
+        t2 = wp.max(tree0, tree1)
+        wp.atomic_max(tree_tree, worldid, t1, t2, 1)
+        wp.atomic_max(tree_tree, worldid, t2, t1, 1)
+    return
+
+  # generic: scan Jacobian row
+  first_tree = int(-1)
+  has_cross_edge = int(0)
+
+  for dof in range(nv):
+    # TODO(team): sparse efc_J
+    # TODO(team): tree dof skip
+    J_val = efc_J_in[worldid, efcid, dof]
+    if J_val != 0.0:
+      tree = dof_treeid[dof]
+      if tree < 0:
+        continue
+      if first_tree == -1:
+        first_tree = tree
+      elif tree != first_tree:
+        t1 = wp.min(first_tree, tree)
+        t2 = wp.max(first_tree, tree)
+        wp.atomic_max(tree_tree, worldid, t1, t2, 1)
+        has_cross_edge = 1
+
+  if first_tree >= 0 and has_cross_edge == 0:
+    wp.atomic_max(tree_tree, worldid, first_tree, first_tree, 1)
+
+
+def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d(dtype=int)):
+  """Compute tree-tree adjacency matrix."""
+  tree_tree.zero_()
+  wp.launch(
+    kernel=_tree_edges,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.nv,
+      m.body_treeid,
+      m.jnt_dofadr,
+      m.dof_treeid,
+      m.geom_bodyid,
+      m.site_bodyid,
+      m.eq_type,
+      m.eq_obj1id,
+      m.eq_obj2id,
+      m.eq_objtype,
+      d.nefc,
+      d.contact.geom,
+      d.efc.type,
+      d.efc.id,
+      d.efc.J,
+      d.njmax,
+    ],
+    outputs=[tree_tree],
+  )

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

@@ -258,7 +258,7 @@ def _gravity_force(
   if gravcomp:
     force = -gravity * body_mass[worldid % body_mass.shape[0], bodyid] * gravcomp
     pos = xipos_in[worldid, bodyid]
-    jac, _ = support.jac(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos, bodyid, dofid, worldid)
+    jac, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos, bodyid, dofid, worldid)
 
     wp.atomic_add(qfrc_gravcomp_out[worldid], dofid, wp.dot(jac, force))
 
@@ -526,9 +526,9 @@ def _fluid(m: Model, d: Data):
 @wp.kernel
 def _qfrc_passive(
   # Model:
-  opt_has_fluid: bool,
   jnt_actgravcomp: wp.array(dtype=int),
   dof_jntid: wp.array(dtype=int),
+  has_fluid: bool,
   # Data in:
   qfrc_spring_in: wp.array2d(dtype=float),
   qfrc_damper_in: wp.array2d(dtype=float),
@@ -548,7 +548,7 @@ def _qfrc_passive(
     qfrc_passive += qfrc_gravcomp_in[worldid, dofid]
 
   # add fluid force
-  if opt_has_fluid:
+  if has_fluid:
     qfrc_passive += qfrc_fluid_in[worldid, dofid]
 
   qfrc_passive_out[worldid, dofid] = qfrc_passive
@@ -703,7 +703,7 @@ def _flex_bending(
     force[i, x] -= flex_bending[edgeid, 16] * frc[i, x]
 
   for i in range(nvert):
-    bodyid = flex_vertbodyid[flex_vertadr[f] + v[i]]
+    bodyid = flex_vertbodyid[v[i]]
     for x in range(3):
       wp.atomic_add(qfrc_spring_out, worldid, body_dofadr[bodyid] + x, force[i, x])
 
@@ -826,16 +826,16 @@ def passive(m: Model, d: Data):
       outputs=[d.qfrc_gravcomp],
     )
 
-  if m.opt.has_fluid:
+  if m.has_fluid:
     _fluid(m, d)
 
   wp.launch(
     _qfrc_passive,
     dim=(d.nworld, m.nv),
     inputs=[
-      m.opt.has_fluid,
       m.jnt_actgravcomp,
       m.dof_jntid,
+      m.has_fluid,
       d.qfrc_spring,
       d.qfrc_damper,
       d.qfrc_gravcomp,

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

@@ -13,15 +13,17 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Optional, Tuple
+from typing import Tuple
 
 import warp as wp
 
 from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec6
 
 wp.set_module_options({"enable_backward": False})
@@ -183,7 +185,7 @@ def _ray_triangle(
 
 
 @wp.func
-def _ray_plane(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
+def ray_plane(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
   """Returns the distance and normal at which a ray intersects with a plane."""
   # map to local frame
   lpnt, lvec = _ray_map(pos, mat, pnt, vec)
@@ -207,7 +209,7 @@ def _ray_plane(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp
 
 
 @wp.func
-def _ray_sphere(pos: wp.vec3, dist_sqr: float, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
+def ray_sphere(pos: wp.vec3, dist_sqr: float, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
   """Returns the distance and normal at which a ray intersects with a sphere."""
   dif = pnt - pos
 
@@ -224,11 +226,11 @@ def _ray_sphere(pos: wp.vec3, dist_sqr: float, pnt: wp.vec3, vec: wp.vec3) -> Tu
 
 
 @wp.func
-def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
+def ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
   """Returns the distance and normal at which a ray intersects with a capsule."""
   # bounding sphere test
   ssz = size[0] + size[1]
-  dist_sphere, normal_sphere = _ray_sphere(pos, ssz * ssz, pnt, vec)
+  dist_sphere, normal_sphere = ray_sphere(pos, ssz * ssz, pnt, vec)
   if dist_sphere < 0:
     return -1.0, wp.vec3()
 
@@ -248,8 +250,9 @@ def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
   sol, xx = _ray_quad(a, b, c)
   part = 0  # -1: bottom, 0: cylinder, 1: top
 
-  # make sure round solution is between flat sides
+  # make sure round solution is between flat sides (must use local z component)
-  if sol >= 0.0 and wp.abs(lpnt[2] + sol * vec[2]) <= size[1]:
+  # TODO: We should add a test to catch this case.
+  if sol >= 0.0 and wp.abs(lpnt[2] + sol * lvec[2]) <= size[1]:
     if x < 0.0 or sol < x:
       x = sol
 
@@ -260,7 +263,7 @@ def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
   c = wp.dot(ldif, ldif) - sq_size0
   _, xx = _ray_quad(a, b, c)
 
-  # accept only top half of sphere
+  # accept only top half of sphere (use local z component)
   for i in range(2):
     if xx[i] >= 0.0 and lpnt[2] + xx[i] * lvec[2] >= size[1]:
       if x < 0.0 or xx[i] < x:
@@ -273,7 +276,7 @@ def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
   c = wp.dot(ldif, ldif) - sq_size0
   _, xx = _ray_quad(a, b, c)
 
-  # accept only bottom half of sphere
+  # accept only bottom half of sphere (use local z component)
   for i in range(2):
     if xx[i] >= 0.0 and lpnt[2] + xx[i] * lvec[2] <= -size[1]:
       if x < 0.0 or xx[i] < x:
@@ -297,7 +300,7 @@ def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
 
 
 @wp.func
-def _ray_ellipsoid(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
+def ray_ellipsoid(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
   """Returns the distance and normal at which a ray intersects with an ellipsoid."""
   # map to local frame
   lpnt, lvec = _ray_map(pos, mat, pnt, vec)
@@ -328,11 +331,11 @@ def _ray_ellipsoid(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec
 
 
 @wp.func
-def _ray_cylinder(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
+def ray_cylinder(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
   """Returns the distance and normal at which a ray intersects with a cylinder."""
   # bounding sphere test
   ssz = size[0] * size[0] + size[1] * size[1]
-  dist_sphere, normal_sphere = _ray_sphere(pos, ssz, pnt, vec)
+  dist_sphere, normal_sphere = ray_sphere(pos, ssz, pnt, vec)
   if dist_sphere < 0:
     return -1.0, wp.vec3()
 
@@ -392,13 +395,13 @@ _IFACE = wp.types.matrix((3, 2), dtype=int)(1, 2, 0, 2, 0, 1)
 
 
 @wp.func
-def _ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, vec6, wp.vec3]:
+def ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, vec6, wp.vec3]:
   """Returns distance, per side information, and normal at which a ray intersects with a box."""
   all = vec6(-1.0, -1.0, -1.0, -1.0, -1.0, -1.0)
 
   # bounding sphere test
   ssz = wp.dot(size, size)
-  dist_sphere, _ = _ray_sphere(pos, ssz, pnt, vec)
+  dist_sphere, _ = ray_sphere(pos, ssz, pnt, vec)
   if dist_sphere < 0:
     return -1.0, all, wp.vec3()
 
@@ -446,7 +449,7 @@ def _ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.v
 
 
 @wp.func
-def _ray_hfield(
+def ray_hfield(
   # Model:
   geom_type: wp.array(dtype=int),
   geom_dataid: wp.array(dtype=int),
@@ -487,10 +490,10 @@ def _ray_hfield(
   top_pos = pos + mat_col * top_scale
 
   # init: intersection with base box
-  x, _, normal_base = _ray_box(base_pos, mat, base_size, pnt, vec)
+  x, _, normal_base = ray_box(base_pos, mat, base_size, pnt, vec)
 
   # check top box: done if no intersection
-  top_intersect, all, normal_top = _ray_box(top_pos, mat, top_size, pnt, vec)
+  top_intersect, all, normal_top = ray_box(top_pos, mat, top_size, pnt, vec)
 
   if top_intersect < 0.0:
     return x, normal_base
@@ -627,10 +630,16 @@ def ray_mesh(
   data_id: int,
   pos: wp.vec3,
   mat: wp.mat33,
+  size: wp.vec3,
   pnt: wp.vec3,
   vec: wp.vec3,
 ) -> Tuple[float, wp.vec3]:
   """Returns the distance and normal for ray mesh intersections."""
+  # bounding box test
+  dist_box, _all, _normal = ray_box(pos, mat, size, pnt, vec)
+  if dist_box < 0.0:
+    return -1.0, wp.vec3()
+
   pnt, vec = _ray_map(pos, mat, pnt, vec)
 
   # compute orthogonal basis vectors
@@ -687,6 +696,87 @@ def ray_mesh(
   return x, normal
 
 
+@wp.func
+def ray_mesh_with_bvh(
+  # In:
+  mesh_bvh_id: wp.array(dtype=wp.uint64),
+  mesh_geom_id: int,
+  pos: wp.vec3,
+  mat: wp.mat33,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  max_t: float,
+) -> Tuple[float, wp.vec3, float, float, int, int]:
+  """Returns intersection information for ray mesh intersections.
+
+  Requires wp.Mesh be constructed and their ids to be passed.
+  """
+  t = float(-1.0)
+  u = float(0.0)
+  v = float(0.0)
+  sign = float(0.0)
+  n = wp.vec3(0.0, 0.0, 0.0)
+  f = int(-1)
+
+  lpnt, lvec = _ray_map(pos, mat, pnt, vec)
+  hit = wp.mesh_query_ray(mesh_bvh_id[mesh_geom_id], lpnt, lvec, max_t, t, u, v, sign, n, f)
+
+  if hit and wp.dot(lvec, n) < 0.0:  # Backface culling in local space
+    normal = mat @ n
+    normal = wp.normalize(normal)
+    return t, normal, u, v, f, mesh_geom_id
+
+  return -1.0, wp.vec3(0.0, 0.0, 0.0), 0.0, 0.0, -1, -1
+
+
+@wp.func
+def ray_mesh_with_bvh_anyhit(
+  # In:
+  mesh_bvh_id: wp.array(dtype=wp.uint64),
+  mesh_geom_id: int,
+  pos: wp.vec3,
+  mat: wp.mat33,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  max_t: float,
+) -> bool:
+  """Returns True if there is any hit for ray mesh intersections.
+
+  Requires wp.Mesh be constructed and their ids to be passed. This variant is useful
+  for shadow ray casts where the only goal is if there is any ray hit.
+  """
+  lpnt, lvec = _ray_map(pos, mat, pnt, vec)
+  return wp.mesh_query_ray_anyhit(mesh_bvh_id[mesh_geom_id], lpnt, lvec, max_t)
+
+
+@wp.func
+def ray_flex_with_bvh(
+  # In:
+  bvh_id: wp.uint64,
+  group_root: int,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  max_t: float,
+) -> Tuple[float, wp.vec3, float, float, int]:
+  """Returns intersection information for flex intersections.
+
+  Requires wp.Mesh be constructed and their ids to be passed. Flex are already in world space.
+  """
+  t = float(-1.0)
+  u = float(0.0)
+  v = float(0.0)
+  sign = float(0.0)
+  n = wp.vec3(0.0, 0.0, 0.0)
+  f = int(-1)
+
+  hit = wp.mesh_query_ray(bvh_id, pnt, vec, max_t, t, u, v, sign, n, f, group_root)
+
+  if hit:
+    return t, n, u, v, f
+
+  return -1.0, wp.vec3(0.0, 0.0, 0.0), 0.0, 0.0, -1
+
+
 @wp.func
 def ray_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3, geomtype: int) -> Tuple[float, wp.vec3]:
   """Returns distance along ray to intersection with geom and normal at intersection point.
@@ -695,17 +785,17 @@ def ray_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.v
   """
   # TODO(team): static loop unrolling to remove unnecessary branching
   if geomtype == GeomType.PLANE:
-    return _ray_plane(pos, mat, size, pnt, vec)
+    return ray_plane(pos, mat, size, pnt, vec)
   elif geomtype == GeomType.SPHERE:
-    return _ray_sphere(pos, size[0] * size[0], pnt, vec)
+    return ray_sphere(pos, size[0] * size[0], pnt, vec)
   elif geomtype == GeomType.CAPSULE:
-    return _ray_capsule(pos, mat, size, pnt, vec)
+    return ray_capsule(pos, mat, size, pnt, vec)
   elif geomtype == GeomType.ELLIPSOID:
-    return _ray_ellipsoid(pos, mat, size, pnt, vec)
+    return ray_ellipsoid(pos, mat, size, pnt, vec)
   elif geomtype == GeomType.CYLINDER:
-    return _ray_cylinder(pos, mat, size, pnt, vec)
+    return ray_cylinder(pos, mat, size, pnt, vec)
   elif geomtype == GeomType.BOX:
-    dist, _, normal = _ray_box(pos, mat, size, pnt, vec)
+    dist, _, normal = ray_box(pos, mat, size, pnt, vec)
     return dist, normal
   else:
     return -1.0, wp.vec3()
@@ -771,11 +861,12 @@ def _ray_geom_mesh(
         geom_dataid[geomid],
         pos,
         mat,
+        geom_size[worldid % geom_size.shape[0], geomid],
         pnt,
         vec,
       )
     elif type == GeomType.HFIELD:
-      return _ray_hfield(
+      return ray_hfield(
         geom_type,
         geom_dataid,
         hfield_size,
@@ -836,7 +927,7 @@ def _ray(
 
   num_threads = wp.block_dim()
 
-  min_dist = float(wp.inf)
+  min_dist = float(MJ_MAXVAL)
   min_geomid = int(-1)
   min_normal = wp.vec3()
 
@@ -874,9 +965,9 @@ def _ray(
         geomid,
       )
       if dist < 0:
-        dist = wp.inf
+        dist = MJ_MAXVAL
     else:
-      dist = wp.inf
+      dist = MJ_MAXVAL
       normal = wp.vec3()
 
     tile_dist = wp.tile(dist)
@@ -891,7 +982,162 @@ def _ray(
       min_geomid = tile_geomid[local_min_geomid[0]]
       min_normal = tile_normal[local_min_geomid[0]]
 
-  if wp.isinf(min_dist):
+  if min_dist >= MJ_MAXVAL:
+    dist_out[worldid, rayid] = -1.0
+  else:
+    dist_out[worldid, rayid] = min_dist
+  geomid_out[worldid, rayid] = min_geomid
+  normal_out[worldid, rayid] = min_normal
+
+
+@wp.func
+def _ray_geom_mesh_bvh(
+  # Model:
+  body_weldid: wp.array(dtype=int),
+  geom_type: wp.array(dtype=int),
+  geom_bodyid: wp.array(dtype=int),
+  geom_dataid: wp.array(dtype=int),
+  geom_matid: wp.array2d(dtype=int),
+  geom_group: wp.array(dtype=int),
+  geom_size: wp.array2d(dtype=wp.vec3),
+  geom_rgba: wp.array2d(dtype=wp.vec4),
+  mat_rgba: wp.array2d(dtype=wp.vec4),
+  # Data in:
+  geom_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  # In:
+  worldid: int,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  geomgroup: vec6,
+  flg_static: bool,
+  bodyexclude: int,
+  geomid: int,
+  mesh_bvh_id: wp.array(dtype=wp.uint64),
+  hfield_bvh_id: wp.array(dtype=wp.uint64),
+  min_dist: float,
+) -> Tuple[float, wp.vec3]:
+  if not _ray_eliminate(
+    body_weldid,
+    geom_bodyid,
+    geom_matid[worldid % geom_matid.shape[0]],
+    geom_group,
+    geom_rgba[worldid % geom_rgba.shape[0]],
+    mat_rgba[worldid % mat_rgba.shape[0]],
+    geomid,
+    geomgroup,
+    flg_static,
+    bodyexclude,
+  ):
+    pos = geom_xpos_in[worldid, geomid]
+    mat = geom_xmat_in[worldid, geomid]
+    gtype = geom_type[geomid]
+
+    if gtype == GeomType.MESH or gtype == GeomType.HFIELD:
+      bvh_ids = mesh_bvh_id if gtype == GeomType.MESH else hfield_bvh_id
+      t, n, u, v, f, geom_mesh_id = ray_mesh_with_bvh(
+        bvh_ids,
+        geom_dataid[geomid],
+        pos,
+        mat,
+        pnt,
+        vec,
+        min_dist,
+      )
+      if t >= 0.0 and t < min_dist:
+        return t, n
+    else:
+      return ray_geom(
+        pos,
+        mat,
+        geom_size[worldid % geom_size.shape[0], geomid],
+        pnt,
+        vec,
+        gtype,
+      )
+
+  return -1.0, wp.vec3()
+
+
+@wp.kernel
+def _ray_bvh(
+  # Model:
+  ngeom: int,
+  body_weldid: wp.array(dtype=int),
+  geom_type: wp.array(dtype=int),
+  geom_bodyid: wp.array(dtype=int),
+  geom_dataid: wp.array(dtype=int),
+  geom_matid: wp.array2d(dtype=int),
+  geom_group: wp.array(dtype=int),
+  geom_size: wp.array2d(dtype=wp.vec3),
+  geom_rgba: wp.array2d(dtype=wp.vec4),
+  mat_rgba: wp.array2d(dtype=wp.vec4),
+  # Data in:
+  geom_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  # In:
+  pnt: wp.array2d(dtype=wp.vec3),
+  vec: wp.array2d(dtype=wp.vec3),
+  geomgroup: vec6,
+  flg_static: bool,
+  bodyexclude: wp.array(dtype=int),
+  bvh_id: wp.uint64,
+  group_root: wp.array(dtype=int),
+  enabled_geom_ids: wp.array(dtype=int),
+  mesh_bvh_id: wp.array(dtype=wp.uint64),
+  hfield_bvh_id: wp.array(dtype=wp.uint64),
+  # Out:
+  dist_out: wp.array2d(dtype=float),
+  geomid_out: wp.array2d(dtype=int),
+  normal_out: wp.array2d(dtype=wp.vec3),
+):
+  worldid, rayid = wp.tid()
+
+  ray_origin = pnt[worldid, rayid]
+  ray_dir = vec[worldid, rayid]
+  body_exclude = bodyexclude[rayid]
+
+  min_dist = float(MJ_MAXVAL)
+  min_geomid = int(-1)
+  min_normal = wp.vec3()
+
+  query = wp.bvh_query_ray(bvh_id, ray_origin, ray_dir, group_root[worldid])
+  bounds_nr = int(0)
+
+  while wp.bvh_query_next(query, bounds_nr, min_dist):
+    bvh_local = bounds_nr - (worldid * ngeom)
+    geomid = enabled_geom_ids[bvh_local]
+
+    dist, normal = _ray_geom_mesh_bvh(
+      body_weldid,
+      geom_type,
+      geom_bodyid,
+      geom_dataid,
+      geom_matid,
+      geom_group,
+      geom_size,
+      geom_rgba,
+      mat_rgba,
+      geom_xpos_in,
+      geom_xmat_in,
+      worldid,
+      pnt[worldid, rayid],
+      vec[worldid, rayid],
+      geomgroup,
+      flg_static,
+      body_exclude,
+      geomid,
+      mesh_bvh_id,
+      hfield_bvh_id,
+      min_dist,
+    )
+
+    if dist >= 0.0 and dist < min_dist:
+      min_dist = dist
+      min_geomid = geomid
+      min_normal = normal
+
+  if min_dist >= MJ_MAXVAL:
     dist_out[worldid, rayid] = -1.0
   else:
     dist_out[worldid, rayid] = min_dist
@@ -904,9 +1150,10 @@ def ray(
   d: Data,
   pnt: wp.array2d(dtype=wp.vec3),
   vec: wp.array2d(dtype=wp.vec3),
-  geomgroup: Optional[vec6] = None,
+  geomgroup: vec6 | None = None,
   flg_static: bool = True,
   bodyexclude: int = -1,
+  rc: RenderContext | None = None,
 ) -> Tuple[wp.array, wp.array, wp.array]:
   """Returns the distance at which rays intersect with primitive geoms.
 
@@ -915,9 +1162,11 @@ def ray(
     d: The data object containing the current state and output arrays (device).
     pnt: Ray origin points.
     vec: Ray directions.
-    geomgroup: Group inclusion/exclusion mask. If all are wp.inf, ignore.
+    geomgroup: Group inclusion/exclusion mask.
     flg_static: If True, allows rays to intersect with static geoms.
     bodyexclude: Ignore geoms on specified body id (-1 to disable).
+    rc: Optional Render context containing BVH information for BVH accelerated ray
+      intersections.
 
   Returns:
     Distances from ray origins to geom surfaces, IDs of intersected geoms (-1 if none),
@@ -935,7 +1184,7 @@ def ray(
   ray_geomid = wp.empty((d.nworld, 1), dtype=int)
   ray_normal = wp.empty((d.nworld, 1), dtype=wp.vec3)
 
-  rays(m, d, pnt, vec, geomgroup, flg_static, ray_bodyexclude, ray_dist, ray_geomid, ray_normal)
+  rays(m, d, pnt, vec, geomgroup, flg_static, ray_bodyexclude, ray_dist, ray_geomid, ray_normal, rc)
 
   return ray_dist, ray_geomid, ray_normal
 
@@ -951,6 +1200,7 @@ def rays(
   dist: wp.array2d(dtype=float),
   geomid: wp.array2d(dtype=int),
   normal: wp.array2d(dtype=wp.vec3),
+  rc: RenderContext | None = None,
 ):
   """Ray intersection for multiple worlds and multiple rays.
 
@@ -968,41 +1218,75 @@ def rays(
     geomid: Output array for IDs of intersected geoms, shape (nworld, nray). -1
       indicates no intersection.
     normal: Output array for normals at intersection points, shape (nworld, nray).
+    rc: Optional Render context containing BVH information for BVH accelerated ray
+      intersections.
   """
-  wp.launch_tiled(
+  # TODO: Investigate building rc if none and removing the non-accelerated path
-    _ray,
+  if rc is None:
-    dim=(d.nworld, pnt.shape[1]),
+    wp.launch_tiled(
-    inputs=[
+      _ray,
-      m.ngeom,
+      dim=(d.nworld, pnt.shape[1]),
-      m.nmeshface,
+      inputs=[
-      m.body_weldid,
+        m.ngeom,
-      m.geom_type,
+        m.nmeshface,
-      m.geom_bodyid,
+        m.body_weldid,
-      m.geom_dataid,
+        m.geom_type,
-      m.geom_matid,
+        m.geom_bodyid,
-      m.geom_group,
+        m.geom_dataid,
-      m.geom_size,
+        m.geom_matid,
-      m.geom_rgba,
+        m.geom_group,
-      m.mesh_vertadr,
+        m.geom_size,
-      m.mesh_faceadr,
+        m.geom_rgba,
-      m.mesh_vert,
+        m.mesh_vertadr,
-      m.mesh_face,
+        m.mesh_faceadr,
-      m.hfield_size,
+        m.mesh_vert,
-      m.hfield_nrow,
+        m.mesh_face,
-      m.hfield_ncol,
+        m.hfield_size,
-      m.hfield_adr,
+        m.hfield_nrow,
-      m.hfield_data,
+        m.hfield_ncol,
-      m.mat_rgba,
+        m.hfield_adr,
-      d.geom_xpos,
+        m.hfield_data,
-      d.geom_xmat,
+        m.mat_rgba,
-      pnt,
+        d.geom_xpos,
-      vec,
+        d.geom_xmat,
-      geomgroup,
+        pnt,
-      flg_static,
+        vec,
-      bodyexclude,
+        geomgroup,
-      dist,
+        flg_static,
-      geomid,
+        bodyexclude,
-      normal,
+        dist,
-    ],
+        geomid,
-    block_dim=m.block_dim.ray,
+        normal,
-  )
+      ],
+      block_dim=m.block_dim.ray,
+    )
+  else:
+    wp.launch(
+      _ray_bvh,
+      dim=(d.nworld, pnt.shape[1]),
+      inputs=[
+        rc.bvh_ngeom,
+        m.body_weldid,
+        m.geom_type,
+        m.geom_bodyid,
+        m.geom_dataid,
+        m.geom_matid,
+        m.geom_group,
+        m.geom_size,
+        m.geom_rgba,
+        m.mat_rgba,
+        d.geom_xpos,
+        d.geom_xmat,
+        pnt,
+        vec,
+        geomgroup,
+        flg_static,
+        bodyexclude,
+        rc.bvh_id,
+        rc.group_root,
+        rc.enabled_geom_ids,
+        rc.mesh_bvh_id,
+        rc.hfield_bvh_id,
+      ],
+      outputs=[dist, geomid, normal],
+    )

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

@@ -0,0 +1,696 @@
+# Copyright 2026 The Newton Developers
+#
+# 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.
+# ==============================================================================
+
+from typing import Tuple
+
+import warp as wp
+
+from mujoco.mjx.third_party.mujoco_warp._src import math
+from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_box
+from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_capsule
+from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_cylinder
+from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_ellipsoid
+from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_flex_with_bvh
+from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh_with_bvh
+from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh_with_bvh_anyhit
+from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_plane
+from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_sphere
+from mujoco.mjx.third_party.mujoco_warp._src.render_util import compute_ray
+from mujoco.mjx.third_party.mujoco_warp._src.render_util import pack_rgba_to_uint32
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import Data
+from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
+from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+from mujoco.mjx.third_party.mujoco_warp._src.types import RenderContext
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
+
+wp.set_module_options({"enable_backward": False})
+
+
+# TODO(team): remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
+from mujoco.mjx.third_party.mujoco_warp._src.types import TEXTURE_DTYPE
+
+
+@wp.func
+def sample_texture(
+  # Model:
+  geom_type: wp.array(dtype=int),
+  mesh_faceadr: wp.array(dtype=int),
+  # In:
+  geom_id: int,
+  tex_repeat: wp.vec2,
+  tex: TEXTURE_DTYPE,
+  pos: wp.vec3,
+  rot: wp.mat33,
+  mesh_facetexcoord: wp.array(dtype=wp.vec3i),
+  mesh_texcoord: wp.array(dtype=wp.vec2),
+  mesh_texcoord_offsets: wp.array(dtype=int),
+  hit_point: wp.vec3,
+  bary_u: float,
+  bary_v: float,
+  f: int,
+  mesh_id: int,
+) -> wp.vec3:
+  uv = wp.vec2(0.0, 0.0)
+
+  if geom_type[geom_id] == GeomType.PLANE:
+    local = wp.transpose(rot) @ (hit_point - pos)
+    uv = wp.vec2(local[0], local[1])
+
+  if geom_type[geom_id] == GeomType.MESH:
+    if f < 0 or mesh_id < 0:
+      return wp.vec3(0.0, 0.0, 0.0)
+
+    face_adr = mesh_faceadr[mesh_id] + f
+    uv0 = mesh_texcoord[mesh_texcoord_offsets[mesh_id] + mesh_facetexcoord[face_adr][0]]
+    uv1 = mesh_texcoord[mesh_texcoord_offsets[mesh_id] + mesh_facetexcoord[face_adr][1]]
+    uv2 = mesh_texcoord[mesh_texcoord_offsets[mesh_id] + mesh_facetexcoord[face_adr][2]]
+    uv = uv0 * bary_u + uv1 * bary_v + uv2 * (1.0 - bary_u - bary_v)
+
+  u = uv[0] * tex_repeat[0]
+  v = uv[1] * tex_repeat[1]
+  u = u - wp.floor(u)
+  v = v - wp.floor(v)
+  tex_color = wp.texture_sample(tex, wp.vec2(u, v), dtype=wp.vec4)
+  return wp.vec3(tex_color[0], tex_color[1], tex_color[2])
+
+
+# TODO: Investigate combining cast_ray and cast_ray_first_hit
+@wp.func
+def cast_ray(
+  # Model:
+  geom_type: wp.array(dtype=int),
+  geom_dataid: wp.array(dtype=int),
+  geom_size: wp.array2d(dtype=wp.vec3),
+  # Data in:
+  geom_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  # In:
+  bvh_id: wp.uint64,
+  group_root: int,
+  world_id: int,
+  bvh_ngeom: int,
+  enabled_geom_ids: wp.array(dtype=int),
+  mesh_bvh_id: wp.array(dtype=wp.uint64),
+  hfield_bvh_id: wp.array(dtype=wp.uint64),
+  ray_origin_world: wp.vec3,
+  ray_dir_world: wp.vec3,
+) -> Tuple[int, float, wp.vec3, float, float, int, int]:
+  dist = float(MJ_MAXVAL)
+  normal = wp.vec3(0.0, 0.0, 0.0)
+  geom_id = int(-1)
+  bary_u = float(0.0)
+  bary_v = float(0.0)
+  face_idx = int(-1)
+  geom_mesh_id = int(-1)
+
+  query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
+  bounds_nr = int(0)
+
+  while wp.bvh_query_next(query, bounds_nr, dist):
+    gi_global = bounds_nr
+    gi_bvh_local = gi_global - (world_id * bvh_ngeom)
+    gi = enabled_geom_ids[gi_bvh_local]
+
+    hit_mesh_id = int(-1)
+    u = float(0.0)
+    v = float(0.0)
+    f = int(-1)
+    n = wp.vec3(0.0, 0.0, 0.0)
+
+    # TODO: Investigate branch elimination with static loop unrolling
+    if geom_type[gi] == GeomType.PLANE:
+      d, n = ray_plane(
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        geom_size[world_id % geom_size.shape[0], gi],
+        ray_origin_world,
+        ray_dir_world,
+      )
+    if geom_type[gi] == GeomType.HFIELD:
+      d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh(
+        hfield_bvh_id,
+        geom_dataid[gi],
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        ray_origin_world,
+        ray_dir_world,
+        dist,
+      )
+    if geom_type[gi] == GeomType.SPHERE:
+      d, n = ray_sphere(
+        geom_xpos_in[world_id, gi],
+        geom_size[world_id % geom_size.shape[0], gi][0] * geom_size[world_id % geom_size.shape[0], gi][0],
+        ray_origin_world,
+        ray_dir_world,
+      )
+    if geom_type[gi] == GeomType.ELLIPSOID:
+      d, n = ray_ellipsoid(
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        geom_size[world_id % geom_size.shape[0], gi],
+        ray_origin_world,
+        ray_dir_world,
+      )
+    if geom_type[gi] == GeomType.CAPSULE:
+      d, n = ray_capsule(
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        geom_size[world_id % geom_size.shape[0], gi],
+        ray_origin_world,
+        ray_dir_world,
+      )
+    if geom_type[gi] == GeomType.CYLINDER:
+      d, n = ray_cylinder(
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        geom_size[world_id % geom_size.shape[0], gi],
+        ray_origin_world,
+        ray_dir_world,
+      )
+    if geom_type[gi] == GeomType.BOX:
+      d, all, n = ray_box(
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        geom_size[world_id % geom_size.shape[0], gi],
+        ray_origin_world,
+        ray_dir_world,
+      )
+    if geom_type[gi] == GeomType.MESH:
+      d, n, u, v, f, hit_mesh_id = ray_mesh_with_bvh(
+        mesh_bvh_id,
+        geom_dataid[gi],
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        ray_origin_world,
+        ray_dir_world,
+        dist,
+      )
+
+    if d >= 0.0 and d < dist:
+      dist = d
+      normal = n
+      geom_id = gi
+      bary_u = u
+      bary_v = v
+      face_idx = f
+      geom_mesh_id = hit_mesh_id
+
+  return geom_id, dist, normal, bary_u, bary_v, face_idx, geom_mesh_id
+
+
+@wp.func
+def cast_ray_first_hit(
+  # Model:
+  geom_type: wp.array(dtype=int),
+  geom_dataid: wp.array(dtype=int),
+  geom_size: wp.array2d(dtype=wp.vec3),
+  # Data in:
+  geom_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  # In:
+  bvh_id: wp.uint64,
+  group_root: int,
+  world_id: int,
+  bvh_ngeom: int,
+  enabled_geom_ids: wp.array(dtype=int),
+  mesh_bvh_id: wp.array(dtype=wp.uint64),
+  hfield_bvh_id: wp.array(dtype=wp.uint64),
+  ray_origin_world: wp.vec3,
+  ray_dir_world: wp.vec3,
+  max_dist: float,
+) -> bool:
+  """A simpler version of casting rays that only checks for the first hit."""
+  query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
+  bounds_nr = int(0)
+
+  while wp.bvh_query_next(query, bounds_nr, max_dist):
+    gi_global = bounds_nr
+    gi_bvh_local = gi_global - (world_id * bvh_ngeom)
+    gi = enabled_geom_ids[gi_bvh_local]
+
+    # TODO: Investigate branch elimination with static loop unrolling
+    if geom_type[gi] == GeomType.PLANE:
+      d, n = ray_plane(
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        geom_size[world_id % geom_size.shape[0], gi],
+        ray_origin_world,
+        ray_dir_world,
+      )
+    if geom_type[gi] == GeomType.HFIELD:
+      d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh(
+        hfield_bvh_id,
+        geom_dataid[gi],
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        ray_origin_world,
+        ray_dir_world,
+        max_dist,
+      )
+    if geom_type[gi] == GeomType.SPHERE:
+      d, n = ray_sphere(
+        geom_xpos_in[world_id, gi],
+        geom_size[world_id % geom_size.shape[0], gi][0] * geom_size[world_id % geom_size.shape[0], gi][0],
+        ray_origin_world,
+        ray_dir_world,
+      )
+    if geom_type[gi] == GeomType.ELLIPSOID:
+      d, n = ray_ellipsoid(
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        geom_size[world_id % geom_size.shape[0], gi],
+        ray_origin_world,
+        ray_dir_world,
+      )
+    if geom_type[gi] == GeomType.CAPSULE:
+      d, n = ray_capsule(
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        geom_size[world_id % geom_size.shape[0], gi],
+        ray_origin_world,
+        ray_dir_world,
+      )
+    if geom_type[gi] == GeomType.CYLINDER:
+      d, n = ray_cylinder(
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        geom_size[world_id % geom_size.shape[0], gi],
+        ray_origin_world,
+        ray_dir_world,
+      )
+    if geom_type[gi] == GeomType.BOX:
+      d, all, n = ray_box(
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        geom_size[world_id % geom_size.shape[0], gi],
+        ray_origin_world,
+        ray_dir_world,
+      )
+    if geom_type[gi] == GeomType.MESH:
+      hit = ray_mesh_with_bvh_anyhit(
+        mesh_bvh_id,
+        geom_dataid[gi],
+        geom_xpos_in[world_id, gi],
+        geom_xmat_in[world_id, gi],
+        ray_origin_world,
+        ray_dir_world,
+        max_dist,
+      )
+      d = 0.0 if hit else -1.0
+
+    if d >= 0.0 and d < max_dist:
+      return True
+
+  return False
+
+
+@wp.func
+def compute_lighting(
+  # Model:
+  geom_type: wp.array(dtype=int),
+  geom_dataid: wp.array(dtype=int),
+  geom_size: wp.array2d(dtype=wp.vec3),
+  # Data in:
+  geom_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  # In:
+  use_shadows: bool,
+  bvh_id: wp.uint64,
+  group_root: int,
+  bvh_ngeom: int,
+  enabled_geom_ids: wp.array(dtype=int),
+  world_id: int,
+  mesh_bvh_id: wp.array(dtype=wp.uint64),
+  hfield_bvh_id: wp.array(dtype=wp.uint64),
+  lightactive: bool,
+  lighttype: int,
+  lightcastshadow: bool,
+  lightpos: wp.vec3,
+  lightdir: wp.vec3,
+  normal: wp.vec3,
+  hitpoint: wp.vec3,
+) -> float:
+  light_contribution = float(0.0)
+
+  # TODO: We should probably only be looping over active lights
+  # in the first place with a static loop of enabled light idx?
+  if not lightactive:
+    return light_contribution
+
+  L = wp.vec3(0.0, 0.0, 0.0)
+  dist_to_light = float(MJ_MAXVAL)
+  attenuation = float(1.0)
+
+  if lighttype == 1:  # directional light
+    L = wp.normalize(-lightdir)
+  else:
+    L, dist_to_light = math.normalize_with_norm(lightpos - hitpoint)
+    attenuation = 1.0 / (1.0 + 0.02 * dist_to_light * dist_to_light)
+    if lighttype == 0:  # spot light
+      spot_dir = wp.normalize(lightdir)
+      cos_theta = wp.dot(-L, spot_dir)
+      spot_factor = wp.min(1.0, wp.max(0.0, (cos_theta - 0.85) / (0.95 - 0.85)))
+      attenuation = attenuation * spot_factor
+
+  ndotl = wp.max(0.0, wp.dot(normal, L))
+  if ndotl == 0.0:
+    return light_contribution
+
+  visible = float(1.0)
+
+  if use_shadows and lightcastshadow:
+    # Nudge the origin slightly along the surface normal to avoid
+    # self-intersection when casting shadow rays
+    eps = 1.0e-4
+    shadow_origin = hitpoint + normal * eps
+    # Distance-limited shadows: cap by dist_to_light (for non-directional)
+    max_t = float(dist_to_light - 1.0e-3)
+    if lighttype == 1:  # directional light
+      max_t = float(1.0e8)
+
+    shadow_hit = cast_ray_first_hit(
+      geom_type,
+      geom_dataid,
+      geom_size,
+      geom_xpos_in,
+      geom_xmat_in,
+      bvh_id,
+      group_root,
+      world_id,
+      bvh_ngeom,
+      enabled_geom_ids,
+      mesh_bvh_id,
+      hfield_bvh_id,
+      shadow_origin,
+      L,
+      max_t,
+    )
+
+    if shadow_hit:
+      visible = 0.3
+
+  return ndotl * attenuation * visible
+
+
+@event_scope
+def render(m: Model, d: Data, rc: RenderContext):
+  """Render the current frame.
+
+  Outputs are stored in buffers within the render context.
+
+  Args:
+    m: The model on device.
+    d: The data on device.
+    rc: The render context on device.
+  """
+  rc.rgb_data.fill_(rc.background_color)
+  rc.depth_data.fill_(0.0)
+
+  @wp.kernel(module="unique", enable_backward=False)
+  def _render_megakernel(
+    # Model:
+    geom_type: wp.array(dtype=int),
+    geom_dataid: wp.array(dtype=int),
+    geom_matid: wp.array2d(dtype=int),
+    geom_size: wp.array2d(dtype=wp.vec3),
+    geom_rgba: wp.array2d(dtype=wp.vec4),
+    cam_projection: wp.array(dtype=int),
+    cam_fovy: wp.array2d(dtype=float),
+    cam_sensorsize: wp.array(dtype=wp.vec2),
+    cam_intrinsic: wp.array2d(dtype=wp.vec4),
+    light_type: wp.array2d(dtype=int),
+    light_castshadow: wp.array2d(dtype=bool),
+    light_active: wp.array2d(dtype=bool),
+    mesh_faceadr: wp.array(dtype=int),
+    mat_texid: wp.array3d(dtype=int),
+    mat_texrepeat: wp.array2d(dtype=wp.vec2),
+    mat_rgba: wp.array2d(dtype=wp.vec4),
+    # Data in:
+    geom_xpos_in: wp.array2d(dtype=wp.vec3),
+    geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    cam_xpos_in: wp.array2d(dtype=wp.vec3),
+    cam_xmat_in: wp.array2d(dtype=wp.mat33),
+    light_xpos_in: wp.array2d(dtype=wp.vec3),
+    light_xdir_in: wp.array2d(dtype=wp.vec3),
+    # In:
+    nrender: int,
+    use_shadows: bool,
+    bvh_ngeom: int,
+    cam_res: wp.array(dtype=wp.vec2i),
+    cam_id_map: wp.array(dtype=int),
+    ray: wp.array(dtype=wp.vec3),
+    rgb_adr: wp.array(dtype=int),
+    depth_adr: wp.array(dtype=int),
+    render_rgb: wp.array(dtype=bool),
+    render_depth: wp.array(dtype=bool),
+    bvh_id: wp.uint64,
+    group_root: wp.array(dtype=int),
+    flex_bvh_id: wp.uint64,
+    flex_group_root: wp.array(dtype=int),
+    enabled_geom_ids: wp.array(dtype=int),
+    mesh_bvh_id: wp.array(dtype=wp.uint64),
+    mesh_facetexcoord: wp.array(dtype=wp.vec3i),
+    mesh_texcoord: wp.array(dtype=wp.vec2),
+    mesh_texcoord_offsets: wp.array(dtype=int),
+    hfield_bvh_id: wp.array(dtype=wp.uint64),
+    flex_rgba: wp.array(dtype=wp.vec4),
+    # TODO: remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
+    textures: wp.array(dtype=TEXTURE_DTYPE),
+    # Out:
+    rgb_out: wp.array2d(dtype=wp.uint32),
+    depth_out: wp.array2d(dtype=float),
+  ):
+    world_idx, ray_idx = wp.tid()
+
+    # Map global ray_idx -> (cam_idx, ray_idx_local) using cumulative sizes
+    cam_idx = int(-1)
+    ray_idx_local = int(-1)
+    accum = int(0)
+    for i in range(nrender):
+      num_i = cam_res[i][0] * cam_res[i][1]
+      if ray_idx < accum + num_i:
+        cam_idx = i
+        ray_idx_local = ray_idx - accum
+        break
+      accum += num_i
+    if cam_idx == -1 or ray_idx_local < 0:
+      return
+
+    if not render_rgb[cam_idx] and not render_depth[cam_idx]:
+      return
+
+    # Map active camera index to MuJoCo camera ID
+    mujoco_cam_id = cam_id_map[cam_idx]
+
+    if wp.static(rc.ray is None):
+      img_w = cam_res[cam_idx][0]
+      img_h = cam_res[cam_idx][1]
+      px = ray_idx_local % img_w
+      py = ray_idx_local // img_w
+      ray_dir_local_cam = compute_ray(
+        cam_projection[mujoco_cam_id],
+        cam_fovy[world_idx % cam_fovy.shape[0], mujoco_cam_id],
+        cam_sensorsize[mujoco_cam_id],
+        cam_intrinsic[world_idx % cam_intrinsic.shape[0], mujoco_cam_id],
+        img_w,
+        img_h,
+        px,
+        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]
+
+    geom_id, dist, normal, u, v, f, mesh_id = cast_ray(
+      geom_type,
+      geom_dataid,
+      geom_size,
+      geom_xpos_in,
+      geom_xmat_in,
+      bvh_id,
+      group_root[world_idx],
+      world_idx,
+      bvh_ngeom,
+      enabled_geom_ids,
+      mesh_bvh_id,
+      hfield_bvh_id,
+      ray_origin_world,
+      ray_dir_world,
+    )
+
+    if wp.static(m.nflex > 0):
+      d, n, u, v, f = ray_flex_with_bvh(
+        flex_bvh_id,
+        flex_group_root[world_idx],
+        ray_origin_world,
+        ray_dir_world,
+        dist,
+      )
+      if d >= 0.0 and d < dist:
+        dist = d
+        normal = n
+        geom_id = -2
+
+    # Early Out
+    if geom_id == -1:
+      return
+
+    if render_depth[cam_idx]:
+      depth_out[world_idx, depth_adr[cam_idx] + ray_idx_local] = dist
+
+    if not render_rgb[cam_idx]:
+      return
+
+    # Shade the pixel
+    hit_point = ray_origin_world + ray_dir_world * dist
+
+    if geom_id == -2:
+      # TODO: Currently flex textures are not supported, and only the first rgba value
+      # is used until further flex support is added.
+      color = flex_rgba[0]
+    elif geom_matid[world_idx % geom_matid.shape[0], geom_id] == -1:
+      color = geom_rgba[world_idx % geom_rgba.shape[0], geom_id]
+    else:
+      color = mat_rgba[world_idx % mat_rgba.shape[0], geom_matid[world_idx % geom_matid.shape[0], geom_id]]
+
+    base_color = wp.vec3(color[0], color[1], color[2])
+    hit_color = base_color
+
+    if wp.static(rc.use_textures):
+      if geom_id != -2:
+        mat_id = geom_matid[world_idx % geom_matid.shape[0], geom_id]
+        if mat_id >= 0:
+          tex_id = mat_texid[world_idx % mat_texid.shape[0], mat_id, 1]
+          if tex_id >= 0:
+            tex_color = sample_texture(
+              geom_type,
+              mesh_faceadr,
+              geom_id,
+              mat_texrepeat[world_idx % mat_texrepeat.shape[0], mat_id],
+              textures[tex_id],
+              geom_xpos_in[world_idx, geom_id],
+              geom_xmat_in[world_idx, geom_id],
+              mesh_facetexcoord,
+              mesh_texcoord,
+              mesh_texcoord_offsets,
+              hit_point,
+              u,
+              v,
+              f,
+              mesh_id,
+            )
+            base_color = wp.cw_mul(base_color, tex_color)
+
+    len_n = wp.length(normal)
+    n = normal if len_n > 0.0 else wp.vec3(0.0, 0.0, 1.0)
+    n = wp.normalize(n)
+    hemispheric = 0.5 * (n[2] + 1.0)
+    ambient_color = wp.vec3(0.4, 0.4, 0.45) * hemispheric + wp.vec3(0.1, 0.1, 0.12) * (1.0 - hemispheric)
+    result = 0.5 * wp.cw_mul(base_color, ambient_color)
+
+    # Apply lighting and shadows
+    for l in range(wp.static(m.nlight)):
+      light_contribution = compute_lighting(
+        geom_type,
+        geom_dataid,
+        geom_size,
+        geom_xpos_in,
+        geom_xmat_in,
+        use_shadows,
+        bvh_id,
+        group_root[world_idx],
+        bvh_ngeom,
+        enabled_geom_ids,
+        world_idx,
+        mesh_bvh_id,
+        hfield_bvh_id,
+        light_active[world_idx % light_active.shape[0], l],
+        light_type[world_idx % light_type.shape[0], l],
+        light_castshadow[world_idx % light_castshadow.shape[0], l],
+        light_xpos_in[world_idx, l],
+        light_xdir_in[world_idx, l],
+        normal,
+        hit_point,
+      )
+      result = result + base_color * light_contribution
+
+    hit_color = wp.min(result, wp.vec3(1.0, 1.0, 1.0))
+    hit_color = wp.max(hit_color, wp.vec3(0.0, 0.0, 0.0))
+
+    rgb_out[world_idx, rgb_adr[cam_idx] + ray_idx_local] = pack_rgba_to_uint32(
+      hit_color[0] * 255.0,
+      hit_color[1] * 255.0,
+      hit_color[2] * 255.0,
+      255.0,
+    )
+
+  wp.launch(
+    kernel=_render_megakernel,
+    dim=(d.nworld, rc.total_rays),
+    inputs=[
+      m.geom_type,
+      m.geom_dataid,
+      m.geom_matid,
+      m.geom_size,
+      m.geom_rgba,
+      m.cam_projection,
+      m.cam_fovy,
+      m.cam_sensorsize,
+      m.cam_intrinsic,
+      m.light_type,
+      m.light_castshadow,
+      m.light_active,
+      m.mesh_faceadr,
+      m.mat_texid,
+      m.mat_texrepeat,
+      m.mat_rgba,
+      d.geom_xpos,
+      d.geom_xmat,
+      d.cam_xpos,
+      d.cam_xmat,
+      d.light_xpos,
+      d.light_xdir,
+      rc.nrender,
+      rc.use_shadows,
+      rc.bvh_ngeom,
+      rc.cam_res,
+      rc.cam_id_map,
+      rc.ray,
+      rc.rgb_adr,
+      rc.depth_adr,
+      rc.render_rgb,
+      rc.render_depth,
+      rc.bvh_id,
+      rc.group_root,
+      rc.flex_bvh_id,
+      rc.flex_group_root,
+      rc.enabled_geom_ids,
+      rc.mesh_bvh_id,
+      rc.mesh_facetexcoord,
+      rc.mesh_texcoord,
+      rc.mesh_texcoord_offsets,
+      rc.hfield_bvh_id,
+      rc.flex_rgba,
+      rc.textures,
+    ],
+    outputs=[
+      rc.rgb_data,
+      rc.depth_data,
+    ],
+  )

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

@@ -0,0 +1,130 @@
+# Copyright 2026 The Newton Developers
+#
+# 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 mujoco
+import warp as wp
+
+from mujoco.mjx.third_party.mujoco_warp._src.types import ProjectionType
+
+wp.set_module_options({"enable_backward": False})
+
+
+@wp.kernel
+def _convert_texture_data(
+  # In:
+  width: int,
+  adr: int,
+  nc: int,
+  tex_data_in: wp.array(dtype=wp.uint8),
+  # Out:
+  tex_data_out: wp.array3d(dtype=float),
+):
+  """Convert uint8 texture data to vec4 format for efficient sampling."""
+  x, y = wp.tid()
+  offset = adr + (y * width + x) * nc
+  r = tex_data_in[offset + 0] if nc > 0 else wp.uint8(0)
+  g = tex_data_in[offset + 1] if nc > 1 else wp.uint8(0)
+  b = tex_data_in[offset + 2] if nc > 2 else wp.uint8(0)
+  a = wp.uint8(255)
+
+  tex_data_out[y, x, 0] = float(r) * wp.static(1.0 / 255.0)
+  tex_data_out[y, x, 1] = float(g) * wp.static(1.0 / 255.0)
+  tex_data_out[y, x, 2] = float(b) * wp.static(1.0 / 255.0)
+  tex_data_out[y, x, 3] = float(a) * wp.static(1.0 / 255.0)
+
+
+def create_warp_texture(mjm: mujoco.MjModel, tex_id: int) -> wp.array:
+  """Create a Warp texture from a MuJoCo model texture data."""
+  tex_adr = mjm.tex_adr[tex_id]
+  tex_width = mjm.tex_width[tex_id]
+  tex_height = mjm.tex_height[tex_id]
+  nchannel = mjm.tex_nchannel[tex_id]
+  tex_data = wp.zeros((tex_height, tex_width, 4), dtype=float)
+
+  wp.launch(
+    _convert_texture_data,
+    dim=(tex_width, tex_height),
+    inputs=[tex_width, tex_adr, nchannel, wp.array(mjm.tex_data, dtype=wp.uint8)],
+    outputs=[tex_data],
+  )
+  return wp.Texture2D(tex_data, filter_mode=wp.TextureFilterMode.LINEAR)
+
+
+@wp.func
+def compute_ray(
+  # In:
+  projection: int,
+  fovy: float,
+  sensorsize: wp.vec2,
+  intrinsic: wp.vec4,
+  img_w: int,
+  img_h: int,
+  px: int,
+  py: int,
+  znear: float,
+) -> wp.vec3:
+  """Compute ray direction for a pixel with per-world camera parameters.
+
+  This combines _camera_frustum_bounds and build_primary_rays logic for use
+  inside a kernel when camera parameters are batched/randomized across worlds.
+  """
+  if projection == ProjectionType.ORTHOGRAPHIC:
+    return wp.vec3(0.0, 0.0, -1.0)
+
+  aspect = float(img_w) / float(img_h)
+  sensor_h = sensorsize[1]
+
+  # Check if we have intrinsics (sensorsize[1] != 0)
+  if sensor_h != 0.0:
+    fx = intrinsic[0]
+    fy = intrinsic[1]
+    cx = intrinsic[2]
+    cy = intrinsic[3]
+    sensor_w = sensorsize[0]
+
+    target_aspect = float(img_w) / float(img_h)
+    sensor_aspect = sensor_w / sensor_h
+    if target_aspect > sensor_aspect:
+      sensor_h = sensor_w / target_aspect
+    elif target_aspect < sensor_aspect:
+      sensor_w = sensor_h * target_aspect
+
+    inv_fx_znear = znear / fx
+    inv_fy_znear = znear / fy
+    left = -inv_fx_znear * (sensor_w * 0.5 - cx)
+    right = inv_fx_znear * (sensor_w * 0.5 + cx)
+    top = inv_fy_znear * (sensor_h * 0.5 - cy)
+    bottom = -inv_fy_znear * (sensor_h * 0.5 + cy)
+  else:
+    fovy_rad = fovy * wp.static(wp.pi / 180.0)
+    half_height = znear * wp.tan(0.5 * fovy_rad)
+    half_width = half_height * aspect
+    left = -half_width
+    right = half_width
+    top = half_height
+    bottom = -half_height
+
+  u = (float(px) + 0.5) / float(img_w)
+  v = (float(py) + 0.5) / float(img_h)
+  x = left + (right - left) * u
+  y = top + (bottom - top) * v
+
+  return wp.normalize(wp.vec3(x, y, -znear))
+
+
+@wp.func
+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))

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

@@ -23,6 +23,7 @@ from mujoco.mjx.third_party.mujoco_warp._src import smooth
 from mujoco.mjx.third_party.mujoco_warp._src import support
 from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import get_sdf_params
 from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
@@ -42,7 +43,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i
 from mujoco.mjx.third_party.mujoco_warp._src.util_misc import inside_geom
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 
@@ -124,10 +124,10 @@ def _magnetometer(
 @wp.func
 def _cam_projection(
   # Model:
-  cam_fovy: wp.array(dtype=float),
+  cam_fovy: wp.array2d(dtype=float),
   cam_resolution: wp.array(dtype=wp.vec2i),
   cam_sensorsize: wp.array(dtype=wp.vec2),
-  cam_intrinsic: wp.array(dtype=wp.vec4),
+  cam_intrinsic: wp.array2d(dtype=wp.vec4),
   # Data in:
   site_xpos_in: wp.array2d(dtype=wp.vec3),
   cam_xpos_in: wp.array2d(dtype=wp.vec3),
@@ -138,8 +138,8 @@ def _cam_projection(
   refid: int,
 ) -> wp.vec2:
   sensorsize = cam_sensorsize[refid]
-  intrinsic = cam_intrinsic[refid]
+  intrinsic = cam_intrinsic[worldid % cam_intrinsic.shape[0], refid]
-  fovy = cam_fovy[refid]
+  fovy = cam_fovy[worldid % cam_fovy.shape[0], refid]
   res = cam_resolution[refid]
 
   target_xpos = site_xpos_in[worldid, objid]
@@ -470,10 +470,10 @@ def _sensor_pos(
   site_quat: wp.array2d(dtype=wp.quat),
   cam_bodyid: wp.array(dtype=int),
   cam_quat: wp.array2d(dtype=wp.quat),
-  cam_fovy: wp.array(dtype=float),
+  cam_fovy: wp.array2d(dtype=float),
   cam_resolution: wp.array(dtype=wp.vec2i),
   cam_sensorsize: wp.array(dtype=wp.vec2),
-  cam_intrinsic: wp.array(dtype=wp.vec4),
+  cam_intrinsic: wp.array2d(dtype=wp.vec4),
   sensor_type: wp.array(dtype=int),
   sensor_datatype: wp.array(dtype=int),
   sensor_objtype: wp.array(dtype=int),
@@ -782,7 +782,7 @@ def sensor_pos(m: Model, d: Data):
       d,
       rangefinder_pnt,
       rangefinder_vec,
-      vec6(wp.inf, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf),
+      vec6(MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL),
       True,
       m.sensor_rangefinder_bodyid,
       rangefinder_dist,
@@ -797,8 +797,7 @@ def sensor_pos(m: Model, d: Data):
     energy_vel(m, d)
 
   # collision sensors (distance, normal, fromto)
-  sensor_collision = wp.empty((d.nworld, m.nsensorcollision, 8, 7), dtype=float)
+  sensor_collision = wp.full((d.nworld, m.nsensorcollision, 8, 7), 1.0e32, dtype=float)
-  sensor_collision.fill_(1.0e32)
   if m.nsensorcollision:
     wp.launch(
       _sensor_collision,
@@ -1787,38 +1786,12 @@ def _sensor_acc(
     nmatch = sensor_contact_nmatch_in[worldid, contactsensorid]
 
     if reduce == 3:  # netforce
-      # compute point: force-weighted centroid of contact position
+      # Single-pass computation: first compute centroid, then wrench about centroid
+      # Pass 1: compute force-weighted centroid of contact positions
       net_pos = wp.vec3(0.0)
-      total_force_magnitude = float(0.0)
-
-      for i in range(nmatch):
-        cid = sensor_contact_matchid_in[worldid, contactsensorid, i]
-
-        contact_forcetorque = support.contact_force_fn(
-          opt_cone,
-          contact_frame_in,
-          contact_friction_in,
-          contact_dim_in,
-          contact_efc_address_in,
-          efc_force_in,
-          njmax_in,
-          nacon_in,
-          worldid,
-          cid,
-          False,
-        )
-
-        weight = wp.norm_l2(wp.spatial_top(contact_forcetorque))
-        net_pos += weight * contact_pos_in[cid]
-        total_force_magnitude += weight
-
-      net_pos /= wp.max(total_force_magnitude, MJ_MINVAL)
-
-      # TODO(team): iterate over matches once
-
-      # compute total wrench about point, in the global frame
       net_force = wp.vec3(0.0)
       net_torque = wp.vec3(0.0)
+      total_force_magnitude = float(0.0)
 
       for i in range(nmatch):
         cid = sensor_contact_matchid_in[worldid, contactsensorid, i]
@@ -1837,8 +1810,15 @@ def _sensor_acc(
           cid,
           False,
         )
-        contact_forcetorque *= dir
 
+        # Accumulate for centroid computation (unsigned force magnitude)
+        weight = wp.norm_l2(wp.spatial_top(contact_forcetorque))
+        contact_pos = contact_pos_in[cid]
+        net_pos += weight * contact_pos
+        total_force_magnitude += weight
+
+        # Apply direction and transform to global frame
+        contact_forcetorque *= dir
         force_local = wp.spatial_top(contact_forcetorque)
         torque_local = wp.spatial_bottom(contact_forcetorque)
 
@@ -1848,12 +1828,18 @@ def _sensor_acc(
         force_global = frameT @ force_local
         torque_global = frameT @ torque_local
 
-        # add to total force, torque
+        # Accumulate force and torque (about origin for now)
         net_force += force_global
         net_torque += torque_global
+        # Accumulate moment contribution: will adjust after centroid is computed
+        net_torque += wp.cross(contact_pos, force_global)
 
-        # add induced moment: torque += (pos - point) x force
+      # Finalize centroid
-        net_torque += wp.cross(contact_pos_in[cid] - net_pos, force_global)
+      net_pos /= wp.max(total_force_magnitude, MJ_MINVAL)
+
+      # Adjust torque: subtract moment from centroid (since we accumulated about origin)
+      # torque_about_centroid = torque_about_origin - centroid x total_force
+      net_torque -= wp.cross(net_pos, net_force)
 
       adr_slot = adr
 
@@ -1888,7 +1874,8 @@ def _sensor_acc(
         out[adr_slot + 1] = 1.0
         out[adr_slot + 2] = 0.0
     else:
-      for i in range(wp.min(nmatch, num)):
+      nslots = wp.min(nmatch, num)
+      for i in range(nslots):
         # sorted contact id
         cid = sensor_contact_matchid_in[worldid, contactsensorid, i]
 
@@ -2165,7 +2152,15 @@ def _sensor_tactile(
   contact_type = geom_type[geom]
 
   plugin_attributes, plugin_index, volume_data, mesh_data = get_sdf_params(
-    oct_child, oct_aabb, oct_coeff, plugin, plugin_attr, contact_type, geom_size[worldid, geom], plugin_id, mesh_id
+    oct_child,
+    oct_aabb,
+    oct_coeff,
+    plugin,
+    plugin_attr,
+    contact_type,
+    geom_size[worldid % geom_size.shape[0], geom],
+    plugin_id,
+    mesh_id,
   )
 
   depth = wp.min(sdf(contact_type, lpos, plugin_attributes, plugin_index, volume_data, mesh_data), 0.0)
@@ -2357,16 +2352,16 @@ def _contact_match(
 
 @cache_kernel
 def _contact_sort(maxmatch: int):
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def contact_sort(
     # Model:
     sensor_intprm: wp.array2d(dtype=int),
     sensor_contact_adr: wp.array(dtype=int),
-    # Data in:
+    # In:
     sensor_contact_nmatch_in: wp.array2d(dtype=int),
     sensor_contact_matchid_in: wp.array3d(dtype=int),
     sensor_contact_criteria_in: wp.array3d(dtype=float),
-    # Data out:
+    # Out:
     sensor_contact_matchid_out: wp.array3d(dtype=int),
   ):
     worldid, contactsensorid = wp.tid()
@@ -2819,13 +2814,13 @@ def energy_pos(m: Model, d: Data):
 
 @cache_kernel
 def _energy_vel_kinetic(nv: int):
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def energy_vel_kinetic(
     # Data in:
     qvel_in: wp.array2d(dtype=float),
     # In:
     Mqvel: wp.array2d(dtype=float),
-    # Out:
+    # Data out:
     energy_out: wp.array(dtype=wp.vec2),
   ):
     worldid = wp.tid()
@@ -2849,12 +2844,13 @@ def energy_vel(m: Model, d: Data):
   # kinetic energy: 0.5 * qvel.T @ M @ qvel
 
   # M @ qvel
-  support.mul_m(m, d, d.efc.mv, d.qvel)
+  mv = wp.zeros((d.nworld, m.nv), dtype=float)
+  support.mul_m(m, d, mv, d.qvel)
 
   wp.launch_tiled(
     _energy_vel_kinetic(m.nv),
     dim=d.nworld,
-    inputs=[d.qvel, d.efc.mv],
+    inputs=[d.qvel, mv],
     outputs=[d.energy],
     block_dim=m.block_dim.energy_vel_kinetic,
   )

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

@@ -23,21 +23,21 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import JointType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import State
-from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import nested_kernel
 
 wp.set_module_options({"enable_backward": False})
 
 
 @cache_kernel
-def mul_m_sparse_diag(check_skip: bool):
+def mul_m_sparse(check_skip: bool):
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique")
-  def _mul_m_sparse_diag(
+  def _mul_m_sparse(
     # Model:
-    dof_Madr: wp.array(dtype=int),
+    qM_mulm_rowadr: wp.array(dtype=int),
+    qM_mulm_col: wp.array(dtype=int),
+    qM_mulm_madr: wp.array(dtype=int),
     # Data in:
     qM_in: wp.array3d(dtype=float),
     # In:
@@ -46,26 +46,33 @@ def mul_m_sparse_diag(check_skip: bool):
     # Out:
     res: wp.array2d(dtype=float),
   ):
-    """Diagonal update for sparse matmul."""
+    """Sparse matmul: one thread per DOF, gather-based (no atomics)."""
     worldid, dofid = wp.tid()
 
     if wp.static(check_skip):
       if skip[worldid]:
         return
 
-    res[worldid, dofid] = qM_in[worldid, 0, dof_Madr[dofid]] * vec[worldid, dofid]
+    # Gather all contributions (diagonal + off-diagonal)
+    acc = float(0.0)
+    start = qM_mulm_rowadr[dofid]
+    end = qM_mulm_rowadr[dofid + 1]
+    for k in range(start, end):
+      col = qM_mulm_col[k]
+      madr = qM_mulm_madr[k]
+      acc += qM_in[worldid, 0, madr] * vec[worldid, col]
 
-  return _mul_m_sparse_diag
+    res[worldid, dofid] = acc
+
+  return _mul_m_sparse
 
 
 @cache_kernel
-def mul_m_sparse_ij(check_skip: bool):
+def mul_m_dense(nv: int, check_skip: bool):
-  @nested_kernel(module="unique", enable_backward=False)
+  """Simple SIMT dense matmul: one thread per output element."""
-  def _mul_m_sparse_ij(
+
-    # Model:
+  @wp.kernel(module="unique")
-    qM_mulm_i: wp.array(dtype=int),
+  def _mul_m_dense(
-    qM_mulm_j: wp.array(dtype=int),
-    qM_madr_ij: wp.array(dtype=int),
     # Data in:
     qM_in: wp.array3d(dtype=float),
     # In:
@@ -74,52 +81,16 @@ def mul_m_sparse_ij(check_skip: bool):
     # Out:
     res: wp.array2d(dtype=float),
   ):
-    """Off-diagonal update for sparse matmul."""
+    worldid, i = wp.tid()
-    worldid, elementid = wp.tid()
 
     if wp.static(check_skip):
       if skip[worldid]:
         return
 
-    i = qM_mulm_i[elementid]
+    acc = float(0.0)
-    j = qM_mulm_j[elementid]
+    for j in range(wp.static(nv)):
-    madr_ij = qM_madr_ij[elementid]
+      acc += qM_in[worldid, i, j] * vec[worldid, j]
-
+    res[worldid, i] = acc
-    qM_ij = qM_in[worldid, 0, madr_ij]
-
-    wp.atomic_add(res[worldid], i, qM_ij * vec[worldid, j])
-    wp.atomic_add(res[worldid], j, qM_ij * vec[worldid, i])
-
-  return _mul_m_sparse_ij
-
-
-@cache_kernel
-def mul_m_dense(tile: TileSet, check_skip: bool):
-  """Returns a matmul kernel for some tile size."""
-
-  @nested_kernel(module="unique", enable_backward=False)
-  def _mul_m_dense(
-    # Data In:
-    qM_in: wp.array3d(dtype=float),
-    # In:
-    adr: wp.array(dtype=int),
-    vec: wp.array3d(dtype=float),
-    skip: wp.array(dtype=bool),
-    # Out:
-    res: wp.array3d(dtype=float),
-  ):
-    worldid, nodeid = wp.tid()
-    TILE_SIZE = wp.static(tile.size)
-
-    if wp.static(check_skip):
-      if skip[worldid]:
-        return
-
-    dofid = adr[nodeid]
-    qM_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid), bounds_check=False)
-    vec_tile = wp.tile_load(vec[worldid], shape=(TILE_SIZE, 1), offset=(dofid, 0), bounds_check=False)
-    res_tile = wp.tile_matmul(qM_tile, vec_tile)
-    wp.tile_store(res[worldid], res_tile, offset=(dofid, 0), bounds_check=False)
 
   return _mul_m_dense
 
@@ -149,36 +120,21 @@ def mul_m(
   if M is None:
     M = d.qM
 
-  if m.opt.is_sparse:
+  if m.is_sparse:
     wp.launch(
-      mul_m_sparse_diag(check_skip),
+      mul_m_sparse(check_skip),
       dim=(d.nworld, m.nv),
-      inputs=[m.dof_Madr, M, vec, skip],
+      inputs=[m.qM_mulm_rowadr, m.qM_mulm_col, m.qM_mulm_madr, M, vec, skip],
-      outputs=[res],
-    )
-
-    wp.launch(
-      mul_m_sparse_ij(check_skip),
-      dim=(d.nworld, m.qM_madr_ij.size),
-      inputs=[m.qM_mulm_i, m.qM_mulm_j, m.qM_madr_ij, M, vec, skip],
       outputs=[res],
     )
 
   else:
-    for tile in m.qM_tiles:
+    wp.launch(
-      wp.launch_tiled(
+      mul_m_dense(m.nv, check_skip),
-        mul_m_dense(tile, check_skip),
+      dim=(d.nworld, m.nv),
-        dim=(d.nworld, tile.adr.size),
+      inputs=[M, vec, skip],
-        inputs=[
+      outputs=[res],
-          M,
+    )
-          tile.adr,
-          # note reshape: tile_matmul expects 2d input
-          vec.reshape(vec.shape + (1,)),
-          skip,
-        ],
-        outputs=[res.reshape(res.shape + (1,))],
-        block_dim=m.block_dim.mul_m_dense,
-      )
 
 
 @wp.kernel
@@ -408,7 +364,7 @@ def transform_force(frc: wp.spatial_vector, offset: wp.vec3) -> wp.spatial_vecto
 
 
 @wp.func
-def jac(
+def jac_dof(
   # Model:
   body_parentid: wp.array(dtype=int),
   body_rootid: wp.array(dtype=int),
@@ -446,8 +402,77 @@ def jac(
   return jacp, jacr
 
 
+@cache_kernel
+def _make_jac_kernel(has_jacp: bool, has_jacr: bool):
+  @wp.kernel(module="unique", enable_backward=False)
+  def _jac(
+    # Model:
+    body_parentid: wp.array(dtype=int),
+    body_rootid: wp.array(dtype=int),
+    dof_bodyid: wp.array(dtype=int),
+    # Data in:
+    subtree_com_in: wp.array2d(dtype=wp.vec3),
+    cdof_in: wp.array2d(dtype=wp.spatial_vector),
+    # In:
+    point_in: wp.array(dtype=wp.vec3),
+    bodyid_in: wp.array(dtype=int),
+    # Out:
+    jacp_out: wp.array3d(dtype=float),
+    jacr_out: wp.array3d(dtype=float),
+  ):
+    worldid, dofid = wp.tid()
+
+    jacp_val, jacr_val = jac_dof(
+      body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, point_in[worldid], bodyid_in[worldid], dofid, worldid
+    )
+
+    if wp.static(has_jacp):
+      jacp_out[worldid, 0, dofid] = jacp_val[0]
+      jacp_out[worldid, 1, dofid] = jacp_val[1]
+      jacp_out[worldid, 2, dofid] = jacp_val[2]
+
+    if wp.static(has_jacr):
+      jacr_out[worldid, 0, dofid] = jacr_val[0]
+      jacr_out[worldid, 1, dofid] = jacr_val[1]
+      jacr_out[worldid, 2, dofid] = jacr_val[2]
+
+  return _jac
+
+
+@event_scope
+def jac(
+  m: Model,
+  d: Data,
+  jacp: wp.array | None,  # wp.array3d(dtype=float)
+  jacr: wp.array | None,  # wp.array3d(dtype=float)
+  point: wp.array(dtype=wp.vec3),
+  body: wp.array(dtype=int),
+):
+  """Compute translational and rotational Jacobian for point on body.
+
+  Args:
+    m: The model containing kinematic and dynamic information (device).
+    d: The data object containing the current state (device).
+    jacp: Output translational Jacobian (optional).
+    jacr: Output rotational Jacobian (optional).
+    point: 3D point in global coordinates.
+    body: Body ID for each world.
+  """
+  kernel = _make_jac_kernel(jacp is not None, jacr is not None)
+
+  jacp_arr = jacp or wp.empty((0, 0, 0), dtype=float)
+  jacr_arr = jacr or wp.empty((0, 0, 0), dtype=float)
+
+  wp.launch(
+    kernel,
+    dim=(d.nworld, m.nv),
+    inputs=[m.body_parentid, m.body_rootid, m.dof_bodyid, d.subtree_com, d.cdof, point, body],
+    outputs=[jacp_arr, jacr_arr],
+  )
+
+
 @wp.func
-def jac_dot(
+def jac_dot_dof(
   # Model:
   body_parentid: wp.array(dtype=int),
   body_rootid: wp.array(dtype=int),
@@ -529,7 +554,7 @@ def get_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, activ
   if sig >= (1 << State.NSTATE):
     raise ValueError(f"invalid state signature {sig} >= 2^mjNSTATE")
 
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def _get_state(
     # Model:
     nq: int,
@@ -668,7 +693,7 @@ def set_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, activ
   if sig >= (1 << State.NSTATE):
     raise ValueError(f"invalid state signature {sig} >= 2^mjNSTATE")
 
-  @nested_kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False)
   def _set_state(
     # Model:
     nq: int,

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

@@ -32,6 +32,9 @@ MJ_MAX_EPAFACES = 5
 TILE_SIZE_JTDAJ_SPARSE = 16
 TILE_SIZE_JTDAJ_DENSE = 16
 
+# TODO(team): remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
+TEXTURE_DTYPE = wp.Texture2D if hasattr(wp, "Texture2D") else int
+
 
 # TODO(team): add check that all wp.launch_tiled 'block_dim' settings are configurable
 @dataclasses.dataclass
@@ -61,9 +64,9 @@ class BlockDim:
   update_gradient_cholesky_blocked: int = 32
   update_gradient_JTDAJ_sparse: int = 64
   update_gradient_JTDAJ_dense: int = 96
-  linesearch_iterative: int = 64
+  linesearch_iterative: int = 32
-  # support
+  # derivative
-  mul_m_dense: int = 32
+  qderiv_actuator_dense: int = 32
 
 
 class BroadphaseType(enum.IntEnum):
@@ -114,6 +117,23 @@ class CamLightType(enum.IntEnum):
   TARGETBODYCOM = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM
 
 
+class ProjectionType(enum.IntEnum):
+  """Type of camera projection.
+
+  Attributes:
+    PERSPECTIVE: perspective projection
+    ORTHOGRAPHIC: orthographic projection
+  """
+
+  # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
+  if hasattr(mujoco, "mjtProjection"):
+    PERSPECTIVE = mujoco.mjtProjection.mjPROJ_PERSPECTIVE
+    ORTHOGRAPHIC = mujoco.mjtProjection.mjPROJ_ORTHOGRAPHIC
+  else:
+    PERSPECTIVE = 0
+    ORTHOGRAPHIC = 1
+
+
 class DataType(enum.IntFlag):
   """Sensor data types.
 
@@ -164,7 +184,8 @@ class DisableBit(enum.IntFlag):
   REFSAFE = mujoco.mjtDisableBit.mjDSBL_REFSAFE
   SENSOR = mujoco.mjtDisableBit.mjDSBL_SENSOR
   EULERDAMP = mujoco.mjtDisableBit.mjDSBL_EULERDAMP
-  # unsupported: MIDPHASE, AUTORESET, NATIVECCD, ISLAND
+  NATIVECCD = mujoco.mjtDisableBit.mjDSBL_NATIVECCD
+  # unsupported: MIDPHASE, AUTORESET, ISLAND
 
 
 class EnableBit(enum.IntFlag):
@@ -321,6 +342,20 @@ class GeomType(enum.IntEnum):
   # unsupported: NGEOMTYPES, ARROW*, LINE, SKIN, LABEL, NONE
 
 
+class CollisionType(enum.IntEnum):
+  """Type of narrowphase collision.
+
+  Attributes:
+    PRIMITIVE: primitive collision
+    CONVEX: convex collision (CCD)
+    SDF: sdf collision
+  """
+
+  PRIMITIVE = 0
+  CONVEX = 1
+  SDF = 2
+
+
 class SolverType(enum.IntEnum):
   """Constraint solver algorithm.
 
@@ -665,10 +700,8 @@ class Option:
 
   warp only fields:
     impratio_invsqrt: ratio of friction-to-normal contact impedance (stored as inverse square root)
-    is_sparse: whether to use sparse representations
     ls_parallel: evaluate engine solver step sizes in parallel
     ls_parallel_min_step: minimum step size for solver linesearch
-    has_fluid: True if wind, density, or viscosity are non-zero at put_model time
     broadphase: broadphase type (BroadphaseType)
     broadphase_filter: broadphase filter bitflag (BroadphaseFilter)
     graph_conditional: flag to use cuda graph conditional
@@ -700,10 +733,8 @@ class Option:
   sdf_iterations: int
   # warp only fields:
   impratio_invsqrt: array("*", float)
-  is_sparse: bool
   ls_parallel: bool
   ls_parallel_min_step: float
-  has_fluid: bool
   broadphase: BroadphaseType
   broadphase_filter: BroadphaseFilter
   graph_conditional: bool
@@ -716,10 +747,10 @@ class Statistic:
   """Model statistics (in qpos0).
 
   Attributes:
-    meaninertia: mean diagonal inertia
+    meaninertia: mean diagonal inertia (per-world)
   """
 
-  meaninertia: float
+  meaninertia: array("*", float)
 
 
 @dataclasses.dataclass
@@ -749,6 +780,7 @@ class Model:
     nbody: number of bodies
     noct: number of total octree cells in all meshes
     njnt: number of joints
+    ntree: number of kinematic trees
     nM: number of non-zeros in sparse inertia matrix
     nC: number of non-zeros in sparse body-dof matrix
     ngeom: number of geoms
@@ -794,6 +826,7 @@ class Model:
     body_jntadr: start addr of joints; -1: no joints         (nbody,)
     body_dofnum: number of motion degrees of freedom         (nbody,)
     body_dofadr: start addr of dofs; -1: no dofs             (nbody,)
+    body_treeid: id of body's tree; -1: static               (nbody,)
     body_geomnum: number of geoms                            (nbody,)
     body_geomadr: start addr of geoms; -1: no geoms          (nbody,)
     body_pos: position offset rel. to parent body            (*, nbody, 3)
@@ -828,6 +861,7 @@ class Model:
     dof_bodyid: id of dof's body                             (nv,)
     dof_jntid: id of dof's joint                             (nv,)
     dof_parentid: id of dof's parent; -1: none               (nv,)
+    dof_treeid: id of dof's tree                             (nv,)
     dof_Madr: dof address in M-diagonal                      (nv,)
     dof_solref: constraint solver reference: frictionloss    (*, nv, NREF)
     dof_solimp: constraint solver impedance: frictionloss    (*, nv, NIMP)
@@ -835,6 +869,9 @@ class Model:
     dof_armature: dof armature inertia/mass                  (*, nv)
     dof_damping: damping coefficient                         (*, nv)
     dof_invweight0: diag. inverse inertia in qpos0           (*, nv)
+    tree_bodynum: number of bodies in tree (incl. root)      (ntree,)
+    tree_dofadr: start address of tree's dofs                (ntree,)
+    tree_dofnum: number of dofs in tree                      (ntree,)
     geom_type: geometric type (GeomType)                     (ngeom,)
     geom_contype: geom contact type                          (ngeom,)
     geom_conaffinity: geom contact affinity                  (ngeom,)
@@ -870,10 +907,11 @@ class Model:
     cam_poscom0: global position rel. to sub-com in qpos0    (*, ncam, 3)
     cam_pos0: global position rel. to body in qpos0          (*, ncam, 3)
     cam_mat0: global orientation in qpos0                    (*, ncam, 3, 3)
-    cam_fovy: y field-of-view (ortho ? len : deg)            (ncam,)
+    cam_projection: projection type (ProjectionType)         (ncam,)
+    cam_fovy: y field-of-view (ortho ? len : deg)            (*, ncam)
     cam_resolution: resolution: pixels [width, height]       (ncam, 2)
     cam_sensorsize: sensor size: length [width, height]      (ncam, 2)
-    cam_intrinsic: [focal length; principal point]           (ncam, 4)
+    cam_intrinsic: [focal length; principal point]           (*, ncam, 4)
     light_mode: light tracking mode (CamLightType)           (nlight,)
     light_bodyid: id of light's body                         (nlight,)
     light_targetbodyid: id of targeted body; -1: none        (nlight,)
@@ -903,6 +941,9 @@ class Model:
     flex_stiffness: finite element stiffness matrix          (nflexelem, 21)
     flex_bending: bending stiffness                          (nflexedge, 17)
     flex_damping: Rayleigh's damping coefficient             (nflex,)
+    flexedge_J_rownnz: number of nonzeros in Jacobian row    (nflexedge,)
+    flexedge_J_rowadr: row start address in colind array     (nflexedge,)
+    flexedge_J_colind: column indices in sparse Jacobian     (nJfe,)
     mesh_vertadr: first vertex address                       (nmesh,)
     mesh_vertnum: number of vertices                         (nmesh,)
     mesh_faceadr: first face address                         (nmesh,)
@@ -927,7 +968,7 @@ class Model:
     hfield_ncol: number of columns in grid                   (nhfield,)
     hfield_adr: start address in hfield_data                 (nhfield,)
     hfield_data: elevation data                              (nhfielddata,)
-    mat_texid: texture id for rendering                      (nmat, mjNTEXROLE)
+    mat_texid: texture id for rendering                      (*, nmat, mjNTEXROLE)
     mat_texrepeat: texture repeat for rendering              (*, nmat, 2)
     mat_rgba: rgba                                           (*, nmat, 4)
     pair_dim: contact dimensionality                         (npair,)
@@ -1008,6 +1049,7 @@ class Model:
     mapM2M: index mapping from M (legacy) to M (CSR)         (nC)
 
   warp only fields:
+    nbranch: number of branches (leaf-to-root paths)
     nv_pad: number of degrees of freedom + padding
     nacttrnbody: number of actuators with body transmission
     nsensorcollision: number of unique collisions for
@@ -1019,9 +1061,13 @@ class Model:
     nmaxpyramid: maximum number of pyramid directions
     nmaxpolygon: maximum number of verts per polygon
     nmaxmeshdeg: maximum number of polygons per vert
+    is_sparse: whether to use sparse representations
+    has_fluid: True if wind, density, or viscosity are non-zero at put_model time
     has_sdf_geom: whether the model contains SDF geoms
     block_dim: block dim options
     body_tree: list of body ids by tree level
+    body_branches: flattened body ids for all branches
+    body_branch_start: start index in body_branches for each branch   (nbranch + 1,)
     mocap_bodyid: id of body for mocap                       (nmocap,)
     body_fluid_ellipsoid: does body use ellipsoid fluid      (nbody,)
     jnt_limited_slide_hinge_adr: limited/slide/hinge jntadr
@@ -1085,9 +1131,9 @@ class Model:
     qLD_updates: tuple of index triples for sparse factorization
     qM_fullm_i: sparse mass matrix addressing
     qM_fullm_j: sparse mass matrix addressing
-    qM_mulm_i: sparse matmul addressing
+    qM_mulm_rowadr: sparse matmul row pointers
-    qM_mulm_j: sparse matmul addressing
+    qM_mulm_col: sparse matmul column indices
-    qM_madr_ij: sparse matmul addressing
+    qM_mulm_madr: sparse matmul matrix addresses
   """
 
   nq: int
@@ -1097,6 +1143,7 @@ class Model:
   nbody: int
   noct: int
   njnt: int
+  ntree: int
   nM: int
   nC: int
   ngeom: int
@@ -1142,6 +1189,7 @@ class Model:
   body_jntadr: array("nbody", int)
   body_dofnum: array("nbody", int)
   body_dofadr: array("nbody", int)
+  body_treeid: array("nbody", int)
   body_geomnum: array("nbody", int)
   body_geomadr: array("nbody", int)
   body_pos: array("*", "nbody", wp.vec3)
@@ -1176,6 +1224,7 @@ class Model:
   dof_bodyid: array("nv", int)
   dof_jntid: array("nv", int)
   dof_parentid: array("nv", int)
+  dof_treeid: array("nv", int)
   dof_Madr: array("nv", int)
   dof_solref: array("*", "nv", wp.vec2)
   dof_solimp: array("*", "nv", vec5)
@@ -1183,6 +1232,9 @@ class Model:
   dof_armature: array("*", "nv", float)
   dof_damping: array("*", "nv", float)
   dof_invweight0: array("*", "nv", float)
+  tree_bodynum: array("ntree", int)
+  tree_dofadr: array("ntree", int)
+  tree_dofnum: array("ntree", int)
   geom_type: array("ngeom", int)
   geom_contype: array("ngeom", int)
   geom_conaffinity: array("ngeom", int)
@@ -1218,10 +1270,11 @@ class Model:
   cam_poscom0: array("*", "ncam", wp.vec3)
   cam_pos0: array("*", "ncam", wp.vec3)
   cam_mat0: array("*", "ncam", wp.mat33)
-  cam_fovy: array("ncam", float)
+  cam_projection: array("ncam", int)
+  cam_fovy: array("*", "ncam", float)
   cam_resolution: array("ncam", wp.vec2i)
   cam_sensorsize: array("ncam", wp.vec2)
-  cam_intrinsic: array("ncam", wp.vec4)
+  cam_intrinsic: array("*", "ncam", wp.vec4)
   light_mode: array("nlight", int)
   light_bodyid: array("nlight", int)
   light_targetbodyid: array("nlight", int)
@@ -1251,6 +1304,9 @@ class Model:
   flex_stiffness: array("nflexelem", 21, float)
   flex_bending: array("nflexedge", 17, float)
   flex_damping: array("nflex", float)
+  flexedge_J_rownnz: array("nflexedge", int)
+  flexedge_J_rowadr: array("nflexedge", int)
+  flexedge_J_colind: wp.array(dtype=int)
   mesh_vertadr: array("nmesh", int)
   mesh_vertnum: array("nmesh", int)
   mesh_faceadr: array("nmesh", int)
@@ -1275,7 +1331,7 @@ class Model:
   hfield_ncol: array("nhfield", int)
   hfield_adr: array("nhfield", int)
   hfield_data: array("nhfielddata", float)
-  mat_texid: array("nmat", 10, int)
+  mat_texid: array("*", "nmat", 10, int)
   mat_texrepeat: array("*", "nmat", wp.vec2)
   mat_rgba: array("*", "nmat", wp.vec4)
   pair_dim: array("npair", int)
@@ -1355,6 +1411,7 @@ class Model:
   M_colind: array("nC", int)
   mapM2M: array("nC", int)
   # warp only fields:
+  nbranch: int
   nv_pad: int
   nacttrnbody: int
   nsensorcollision: int
@@ -1365,9 +1422,13 @@ class Model:
   nmaxpyramid: int
   nmaxpolygon: int
   nmaxmeshdeg: int
+  is_sparse: bool
+  has_fluid: bool
   has_sdf_geom: bool
   block_dim: BlockDim
   body_tree: tuple[wp.array(dtype=int), ...]
+  body_branches: wp.array(dtype=int)
+  body_branch_start: wp.array(dtype=int)
   mocap_bodyid: array("nmocap", int)
   body_fluid_ellipsoid: array("nbody", bool)
   jnt_limited_slide_hinge_adr: wp.array(dtype=int)
@@ -1422,9 +1483,10 @@ class Model:
   qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...]
   qM_fullm_i: wp.array(dtype=int)
   qM_fullm_j: wp.array(dtype=int)
-  qM_mulm_i: wp.array(dtype=int)
+  # Gather-based sparse mul_m indices (thread per DOF, no atomics)
-  qM_mulm_j: wp.array(dtype=int)
+  qM_mulm_rowadr: wp.array(dtype=int)  # start address for each row [nv+1]
-  qM_madr_ij: wp.array(dtype=int)
+  qM_mulm_col: wp.array(dtype=int)  # column index to gather from
+  qM_mulm_madr: wp.array(dtype=int)  # matrix address to read
 
 
 class ContactType(enum.IntFlag):
@@ -1492,26 +1554,9 @@ class Constraint:
     aref: reference pseudo-acceleration               (nworld, njmax)
     frictionloss: frictionloss (friction)             (nworld, njmax)
     force: constraint force in constraint space       (nworld, njmax)
-    Jaref: Jac*qacc - aref                            (nworld, njmax)
-    Ma: M*qacc                                        (nworld, nv)
-    grad: gradient of master cost                     (nworld, nv_pad)
-    grad_dot: dot(grad, grad)                         (nworld,)
-    Mgrad: M / grad                                   (nworld, nv_pad)
-    search: linesearch vector                         (nworld, nv)
-    search_dot: dot(search, search)                   (nworld,)
-    gauss: Gauss Cost                                 (nworld,)
-    cost: constraint + Gauss cost                     (nworld,)
-    prev_cost: cost from previous iter                (nworld,)
     state: constraint state                           (nworld, njmax_pad)
-    mv: qM @ search                                   (nworld, nv)
+  warp only fields:
-    jv: efc_J @ search                                (nworld, njmax)
+    Ma: M*qacc                                        (nworld, nv)
-    quad: quadratic cost coefficients                 (nworld, njmax, 3)
-    quad_gauss: quadratic cost Gauss coefficients     (nworld, 3)
-    alpha: line search step size                      (nworld,)
-    prev_grad: previous grad                          (nworld, nv)
-    prev_Mgrad: previous Mgrad                        (nworld, nv)
-    beta: Polak-Ribiere beta                          (nworld,)
-    done: solver done                                 (nworld,)
   """
 
   type: array("nworld", "njmax", int)
@@ -1524,26 +1569,8 @@ class Constraint:
   aref: array("nworld", "njmax", float)
   frictionloss: array("nworld", "njmax", float)
   force: array("nworld", "njmax", float)
-  Jaref: array("nworld", "njmax", float)
-  Ma: array("nworld", "nv", float)
-  grad: array("nworld", "nv_pad", float)
-  grad_dot: array("nworld", float)
-  Mgrad: array("nworld", "nv_pad", float)
-  search: array("nworld", "nv", float)
-  search_dot: array("nworld", float)
-  gauss: array("nworld", float)
-  cost: array("nworld", float)
-  prev_cost: array("nworld", float)
   state: array("nworld", "njmax_pad", int)
-  mv: array("nworld", "nv", float)
+  Ma: array("nworld", "nv", float)
-  jv: array("nworld", "njmax", float)
-  quad: array("nworld", "njmax", wp.vec3)
-  quad_gauss: array("nworld", wp.vec3)
-  alpha: array("nworld", float)
-  prev_grad: array("nworld", "nv", float)
-  prev_Mgrad: array("nworld", "nv", float)
-  beta: array("nworld", float)
-  done: array("nworld", bool)
 
 
 @dataclasses.dataclass
@@ -1590,7 +1617,7 @@ class Data:
     cdof: com-based motion axis of each dof (rot:lin)           (nworld, nv, 6)
     cinert: com-based body inertia and mass                     (nworld, nbody, 10)
     flexvert_xpos: cartesian flex vertex positions              (nworld, nflexvert, 3)
-    flexedge_J: edge length Jacobian                            (nworld, nflexedge, nv)
+    flexedge_J: edge length Jacobian                            (nworld, 1, nflexedge*6)
     flexedge_length: flex edge lengths                          (nworld, nflexedge, 1)
     ten_wrapadr: start address of tendon's path                 (nworld, ntendon)
     ten_wrapnum: number of wrap points in path                  (nworld, ntendon)
@@ -1606,7 +1633,7 @@ class Data:
     qLD: L'*D*L factorization of M                              (nworld, nv, nv) if dense
                                                                 (nworld, 1, nC) if sparse
     qLDiagInv: 1/diag(D)                                        (nworld, nv)
-    flexedge_velocity: flex edge velocities                     (nworld, nflexedge,)
+    flexedge_velocity: flex edge velocities                     (nworld, nflexedge)
     ten_velocity: tendon velocities                             (nworld, ntendon)
     actuator_velocity: actuator velocities                      (nworld, nu)
     cvel: com-based velocity (rot:lin)                          (nworld, nbody, 6)
@@ -1636,18 +1663,9 @@ class Data:
   warp only fields:
     nworld: number of worlds
     naconmax: maximum number of contacts (shared across all worlds)
+    naccdmax: maximum number of contacts for CCD (all worlds)
     njmax: maximum number of constraints per world
     nacon: number of detected contacts (across all worlds)      (1,)
-    ne_connect: number of equality connect constraints          (nworld,)
-    ne_weld: number of equality weld constraints                (nworld,)
-    ne_jnt: number of equality joint constraints                (nworld,)
-    ne_ten: number of equality tendon constraints               (nworld,)
-    ne_flex: number of flex edge equality constraints           (nworld,)
-    nsolving: number of unconverged worlds                      (1,)
-    subtree_bodyvel: subtree body velocity (ang, vel)           (nworld, nbody, 6)
-    collision_pair: collision pairs from broadphase             (naconmax, 2)
-    collision_pairid: ids from broadphase                       (naconmax, 2)
-    collision_worldid: collision world ids from broadphase      (naconmax,)
     ncollision: collision count from broadphase                 (1,)
   """
 
@@ -1690,7 +1708,7 @@ class Data:
   cdof: array("nworld", "nv", wp.spatial_vector)
   cinert: array("nworld", "nbody", vec10)
   flexvert_xpos: array("nworld", "nflexvert", wp.vec3)
-  flexedge_J: array("nworld", "nflexedge", "nv", float)
+  flexedge_J: wp.array3d(dtype=float)
   flexedge_length: array("nworld", "nflexedge", float)
   ten_wrapadr: array("nworld", "ntendon", int)
   ten_wrapnum: array("nworld", "ntendon", int)
@@ -1732,18 +1750,131 @@ class Data:
   # warp only fields:
   nworld: int
   naconmax: int
+  naccdmax: int
   njmax: int
   nacon: array(1, int)
-  ne_connect: array("nworld", int)
-  ne_weld: array("nworld", int)
-  ne_jnt: array("nworld", int)
-  ne_ten: array("nworld", int)
-  ne_flex: array("nworld", int)
-  nsolving: array(1, int)
-  subtree_bodyvel: array("nworld", "nbody", wp.spatial_vector)
-
-  # warp only: collision driver
-  collision_pair: array("naconmax", wp.vec2i)
-  collision_pairid: array("naconmax", wp.vec2i)
-  collision_worldid: array("naconmax", int)
   ncollision: array(1, int)
+
+
+@dataclasses.dataclass
+class CollisionContext:
+  """Collision driver intermediate arrays.
+
+  Attributes:
+    collision_pair: collision pairs from broadphase             (naconmax, 2)
+    collision_pairid: ids from broadphase                       (naconmax, 2)
+    collision_worldid: collision world ids from broadphase      (naconmax,)
+  """
+
+  collision_pair: wp.array
+  collision_pairid: wp.array
+  collision_worldid: wp.array
+
+
+@dataclasses.dataclass
+class RenderContext:
+  """Context for rendering.
+
+  Attributes:
+    nrender: number of actively rendering cameras
+    cam_res: camera resolution for actively rendering cameras
+    cam_id_map: camera id map
+    use_textures: whether to use textures
+    use_shadows: whether to use shadows
+    bvh_ngeom: number of geometries in the BVH
+    enabled_geom_ids: enabled geometry ids
+    mesh_registry: mesh BVH id to warp mesh mapping
+    mesh_bvh_id: mesh BVH ids
+    mesh_bounds_size: mesh bounds size
+    mesh_texcoord: mesh texture coordinates
+    mesh_texcoord_offsets: mesh texture coordinate offsets
+    mesh_facetexcoord: mesh face texture coordinates
+    textures: textures
+    textures_registry: texture registry
+    hfield_registry: hfield BVH id to warp mesh mapping
+    hfield_bvh_id: hfield BVH ids
+    hfield_bounds_size: hfield bounds size
+    flex_mesh: flex mesh
+    flex_rgba: flex rgba
+    flex_bvh_id: flex BVH id
+    flex_face_point: flex face points
+    flex_faceadr: flex face addresses
+    flex_nface: number of flex faces
+    flex_nwork: total flex work items for refit
+    flex_group_root: flex group roots
+    flex_elemdataadr: flex element data addresses
+    flex_shell: flex shell data
+    flex_shelldataadr: flex shell data addresses
+    flex_radius: flex radius
+    flex_workadr: flex work item addresses for refit
+    flex_worknum: flex work item counts for refit
+    flex_render_smooth: whether to render flex meshes smoothly
+    bvh: scene BVH
+    bvh_id: scene BVH id
+    lower: lower bounds
+    upper: upper bounds
+    group: groups
+    group_root: group roots
+    ray: rays
+    rgb_data: RGB data
+    rgb_adr: RGB addresses
+    rgb_size: per-camera RGB buffer sizes
+    depth_data: depth data
+    depth_adr: depth addresses
+    depth_size: per-camera depth buffer sizes
+    render_rgb: per-camera RGB render flags
+    render_depth: per-camera depth render flags
+    znear: near plane distance
+    total_rays: total number of rays
+  """
+
+  nrender: int
+  cam_res: array("ncam", wp.vec2i)
+  cam_id_map: array("ncam", int)
+  use_textures: bool
+  use_shadows: bool
+  background_color: wp.uint32
+  bvh_ngeom: int
+  enabled_geom_ids: array("*", int)
+  mesh_registry: dict
+  mesh_bvh_id: array("nmesh", wp.uint64)
+  mesh_bounds_size: array("nmesh", wp.vec3)
+  mesh_texcoord: array("*", wp.vec2)
+  mesh_texcoord_offsets: array("nmesh", int)
+  mesh_facetexcoord: array("nmeshface", wp.vec3i)
+  # TODO(team): remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
+  textures: array("*", TEXTURE_DTYPE)
+  textures_registry: list[TEXTURE_DTYPE]
+  hfield_registry: dict
+  hfield_bvh_id: array("nhfield", wp.uint64)
+  hfield_bounds_size: array("nhfield", wp.vec3)
+  flex_mesh: wp.Mesh
+  flex_rgba: array("nflex", wp.vec4)
+  flex_bvh_id: wp.uint64
+  flex_face_point: array("*", wp.vec3)
+  flex_faceadr: array("nflex", int)
+  flex_nface: int
+  flex_nwork: int
+  flex_group_root: array("nworld", int)
+  flex_elemdataadr: array("nflex", int)
+  flex_shell: array("*", int)
+  flex_shelldataadr: array("nflex", int)
+  flex_radius: array("nflex", float)
+  flex_workadr: array("nflex", int)
+  flex_worknum: array("nflex", int)
+  flex_render_smooth: bool
+  bvh: wp.Bvh
+  bvh_id: wp.uint64
+  lower: array("*", wp.vec3)
+  upper: array("*", wp.vec3)
+  group: array("*", int)
+  group_root: array("*", int)
+  ray: array("*", wp.vec3)
+  rgb_data: array("*", wp.uint32)
+  rgb_adr: array("ncam", int)
+  depth_data: array("*", wp.float32)
+  depth_adr: array("ncam", int)
+  render_rgb: array("ncam", bool)
+  render_depth: array("ncam", bool)
+  znear: float
+  total_rays: int

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

@@ -20,6 +20,7 @@ from typing import Tuple
 import warp as wp
 
 from mujoco.mjx.third_party.mujoco_warp._src import math
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import WrapType
@@ -105,13 +106,13 @@ def wrap_circle(end: wp.vec4, side: wp.vec2, radius: float) -> Tuple[float, wp.v
 
   Args:
     end: Two 2D points.
-    side: Optional 2D side point, no side point: wp.vec2(wp.inf).
+    side: Optional 2D side point, no side point: wp.vec2(MJ_MAXVAL).
     radius: Circle radius.
 
   Returns:
     Length of circular wrap or -1.0 if no wrap, pair of 2D wrap points.
   """
-  valid_side = wp.norm_l2(side) < wp.inf
+  valid_side = wp.norm_l2(side) < MJ_MAXVAL
 
   end0 = wp.vec2(end[0], end[1])
   end1 = wp.vec2(end[2], end[3])
@@ -122,13 +123,13 @@ def wrap_circle(end: wp.vec4, side: wp.vec2, radius: float) -> Tuple[float, wp.v
 
   # either point inside circle or circle too small: no wrap
   if (sqlen0 < sqrad) or (sqlen1 < sqrad) or (radius < MJ_MINVAL):
-    return -1.0, wp.vec2(wp.inf), wp.vec2(wp.inf)
+    return -1.0, wp.vec2(MJ_MAXVAL), wp.vec2(MJ_MAXVAL)
 
   # points too close: no wrap
   dif = end1 - end0
   dd = wp.dot(dif, dif)
   if dd < MJ_MINVAL:
-    return -1.0, wp.vec2(wp.inf), wp.vec2(wp.inf)
+    return -1.0, wp.vec2(MJ_MAXVAL), wp.vec2(MJ_MAXVAL)
 
   # find nearest point on line segment to origin: a * dif + d0
   a = -wp.dot(dif, end0) / dd
@@ -137,7 +138,7 @@ def wrap_circle(end: wp.vec4, side: wp.vec2, radius: float) -> Tuple[float, wp.v
   # check for intersection and side
   tmp = a * dif + end0
   if (wp.dot(tmp, tmp) > sqrad) and (not valid_side or wp.dot(side, tmp) >= 0.0):
-    return -1.0, wp.vec2(wp.inf), wp.vec2(wp.inf)
+    return -1.0, wp.vec2(MJ_MAXVAL), wp.vec2(MJ_MAXVAL)
 
   sqrt0 = wp.sqrt(sqlen0 - sqrad)
   sqrt1 = wp.sqrt(sqlen1 - sqrad)
@@ -191,7 +192,7 @@ def wrap_circle(end: wp.vec4, side: wp.vec2, radius: float) -> Tuple[float, wp.v
 
   # check for intersection
   if is_intersect(end0, pnt0, end1, pnt1):
-    return -1.0, wp.vec2(wp.inf), wp.vec2(wp.inf)
+    return -1.0, wp.vec2(MJ_MAXVAL), wp.vec2(MJ_MAXVAL)
 
   # return curve length
   return length_circle(pnt0, pnt1, ind, radius), pnt0, pnt1
@@ -230,7 +231,7 @@ def wrap_inside(
 
   # either point inside circle or circle too small: no wrap
   if (len0 <= radius) or (len1 <= radius) or (radius < MJ_MINVAL) or (len0 < MJ_MINVAL) or (len1 < MJ_MINVAL):
-    return -1.0, wp.vec2(wp.inf), wp.vec2(wp.inf)
+    return -1.0, wp.vec2(MJ_MAXVAL), wp.vec2(MJ_MAXVAL)
 
   # segment-circle intersection: no wrap
   if dd > MJ_MINVAL:
@@ -241,7 +242,7 @@ def wrap_inside(
     if (a > 0.0) and (a < 1.0):
       tmp = end0 + a * dif
       if wp.norm_l2(tmp) <= radius:
-        return -1.0, wp.vec2(wp.inf), wp.vec2(wp.inf)
+        return -1.0, wp.vec2(MJ_MAXVAL), wp.vec2(MJ_MAXVAL)
 
   # prepare default in case of numerical failure: average
   pnt = 0.5 * (end0 + end1)
@@ -335,14 +336,14 @@ def wrap(
     mat: Orientation of geom.
     radius: Geom radius.
     geomtype: Wrap type (mjtWrap).
-    side: 3D position for sidesite, no side point: wp.vec3(wp.inf).
+    side: 3D position for sidesite, no side point: wp.vec3(MJ_MAXVAL).
 
   Returns:
     Length of circular wrap else -1.0 if no wrap, pair of 3D wrap points.
   """
   # check object type
   if geomtype != WrapType.SPHERE and geomtype != WrapType.CYLINDER:
-    return wp.inf, wp.vec3(wp.inf), wp.vec3(wp.inf)
+    return MJ_MAXVAL, wp.vec3(MJ_MAXVAL), wp.vec3(MJ_MAXVAL)
 
   # map sites to wrap object's local frame
   matT = wp.transpose(mat)
@@ -351,7 +352,7 @@ def wrap(
 
   # too close to origin: return
   if (wp.norm_l2(p0) < MJ_MINVAL) or (wp.norm_l2(p1) < MJ_MINVAL):
-    return -1.0, wp.vec3(wp.inf), wp.vec3(wp.inf)
+    return -1.0, wp.vec3(MJ_MAXVAL), wp.vec3(MJ_MAXVAL)
 
   # construct 2D frame for circle wrap
   if geomtype == WrapType.SPHERE:
@@ -399,7 +400,7 @@ def wrap(
   )
 
   # handle sidesite
-  valid_side = wp.norm_l2(side) < wp.inf
+  valid_side = wp.norm_l2(side) < MJ_MAXVAL
 
   if valid_side:
     # side point: apply same projection as x0, x1
@@ -414,7 +415,7 @@ def wrap(
     sidepnt_proj, _ = math.normalize_with_norm(sidepnt_proj)
     sidepnt_proj *= radius
   else:
-    sidepnt_proj = wp.vec2(wp.inf)
+    sidepnt_proj = wp.vec2(MJ_MAXVAL)
 
   # apply inside wrap
   if valid_side and wp.norm_l2(sidepnt) < radius:
@@ -424,7 +425,7 @@ def wrap(
 
   # no wrap: return
   if wlen < 0.0:
-    return -1.0, wp.vec3(wp.inf), wp.vec3(wp.inf)
+    return -1.0, wp.vec3(MJ_MAXVAL), wp.vec3(MJ_MAXVAL)
 
   # reconstruct 3D points in local frame: res
   res0 = axis0 * pnt0[0] + axis1 * pnt0[1]

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

@@ -15,7 +15,6 @@
 
 import functools
 import inspect
-from typing import Callable, Optional
 
 import warp as wp
 
@@ -119,82 +118,6 @@ def event_scope(fn, name: str = ""):
   return wrapper
 
 
-# @nested_kernel decorator to automatically set up modules based on nested
-# function names
-def nested_kernel(
-  f: Optional[Callable] = None,
-  *,
-  enable_backward: Optional[bool] = None,
-  module: Optional[wp.Module] = None,
-):
-  """Decorator to register a Warp kernel from a Python function.
-
-  The function must be defined with type annotations for all arguments.
-  The function must not return anything.
-
-  Example::
-
-      @nested_kernel
-      def my_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float)):
-        tid = wp.tid()
-        b[tid] = a[tid] + 1.0
-
-
-      @nested_kernel(enable_backward=False)
-      def my_kernel_no_backward(a: wp.array(dtype=float, ndim=2), x: float):
-        # the backward pass will not be generated
-        i, j = wp.tid()
-        a[i, j] = x
-
-
-      @nested_kernel(module="unique")
-      def my_kernel_unique_module(a: wp.array(dtype=float), b: wp.array(dtype=float)):
-        # the kernel will be registered in new unique module created just for this
-        # kernel and its dependent functions and structs
-        tid = wp.tid()
-        b[tid] = a[tid] + 1.0
-
-
-      @neste_kernel(enable_backward=False, module=None)
-      def my_kernel_with_args(a: wp.array(dtype=float), b: wp.array(dtype=float)):
-        # can now use arguments even when module=None
-        tid = wp.tid()
-        b[tid] = a[tid] + 1.0
-
-  Args:
-      f: The function to be registered as a kernel.
-      enable_backward: If False, the backward pass will not be generated.
-      module: The :class:`warp.Module` to which the kernel belongs. Alternatively,
-              if a string `"unique"` is provided, the kernel is assigned to a new module
-              named after the kernel name and hash. If None, the module is inferred from
-              the function's module.
-
-  Returns:
-      The registered kernel.
-  """
-
-  def decorator(func):
-    if module is None:
-      # create a module name based on the name of the nested function
-      # get the qualified name, e.g. "main.<locals>.nested_kernel"
-      qualname = func.__qualname__
-      parts = [part for part in qualname.split(".") if part != "<locals>"]
-      outer_functions = parts[:-1]
-      module_name = wp.get_module(".".join([func.__module__] + outer_functions))
-    else:
-      module_name = module
-
-    return wp.kernel(func, enable_backward=enable_backward, module=module_name)
-
-  # Handle both @kernel and @kernel(...) usage patterns
-  if f is None:
-    # Called with arguments: @kernel(enable_backward=False)
-    return decorator
-  else:
-    # Called without arguments: @kernel
-    return decorator(f)
-
-
 _KERNEL_CACHE = {}
 
 
@@ -221,4 +144,4 @@ def check_toolkit_driver():
   wp.init()
   if wp.get_device().is_cuda:
     if not wp.is_conditional_graph_supported():
-      RuntimeError("Minimum supported CUDA version: 12.4.")
+      raise RuntimeError("Minimum supported CUDA version: 12.4.")

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

@@ -54,7 +54,7 @@ dev = [
     "ruff",
     "pygls>=1.0.0,<2.0.0",
     "lsprotocol>=2023.0.1,<2024.0.0",
-    "mujoco>=3.3.7.dev0",
+    "mujoco>=3.4.1.dev0",
     "warp-lang>=1.11.0.dev0",
 ]
 # TODO(team): cpu and cuda JAX optional dependencies are temporary, remove after we land MJX:Warp

mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py

@@ -24,6 +24,7 @@ Example:
 import copy
 import enum
 import logging
+import shutil
 import sys
 import time
 from typing import Sequence
@@ -51,10 +52,11 @@ class EngineOptions(enum.IntEnum):
   C = 1
 
 
-_CLEAR_KERNEL_CACHE = flags.DEFINE_bool("clear_kernel_cache", False, "Clear kernel cache (to calculate full JIT time)")
+_CLEAR_WARP_CACHE = flags.DEFINE_bool("clear_warp_cache", False, "Clear warp caches (kernel, LTO, CUDA compute)")
 _ENGINE = flags.DEFINE_enum_class("engine", EngineOptions.WARP, EngineOptions, "Simulation engine")
 _NCONMAX = flags.DEFINE_integer("nconmax", None, "Maximum number of contacts.")
 _NJMAX = flags.DEFINE_integer("njmax", None, "Maximum number of constraints per world.")
+_NCCDMAX = flags.DEFINE_integer("nccdmax", None, "Maximum number of CCD contacts per world.")
 _OVERRIDE = flags.DEFINE_multi_string("override", [], "Model overrides (notation: foo.bar = baz)", short_name="o")
 _KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulation.")
 _DEVICE = flags.DEFINE_string("device", None, "override the default Warp device")
@@ -134,27 +136,37 @@ def _main(argv: Sequence[str]) -> None:
   else:
     wp.config.quiet = flags.FLAGS["verbosity"].value < 1
     wp.init()
-    if _CLEAR_KERNEL_CACHE.value:
+    wp.set_device(_DEVICE.value)
+    if _CLEAR_WARP_CACHE.value:
       wp.clear_kernel_cache()
+      wp.clear_lto_cache()
+      # Clear CUDA compute cache for truly cold start JIT
+      compute_cache = epath.Path("~/.nv/ComputeCache").expanduser()
+      if compute_cache.exists():
+        shutil.rmtree(compute_cache)
+        compute_cache.mkdir()
 
-    with wp.ScopedDevice(_DEVICE.value):
+    override_model(mjm, _OVERRIDE.value)
-      m = mjw.put_model(mjm)
+    m = mjw.put_model(mjm)
-      override_model(m, _OVERRIDE.value)
+    override_model(m, _OVERRIDE.value)
-      broadphase, filter = mjw.BroadphaseType(m.opt.broadphase).name, mjw.BroadphaseFilter(m.opt.broadphase_filter).name
+    d = mjw.put_data(mjm, mjd, nconmax=_NCONMAX.value, njmax=_NJMAX.value, nccdmax=_NCCDMAX.value)
-      solver, cone = mjw.SolverType(m.opt.solver).name, mjw.ConeType(m.opt.cone).name
+    graph = _compile_step(m, d) if wp.get_device().is_cuda else None
-      integrator = mjw.IntegratorType(m.opt.integrator).name
+    if graph is None:
-      iterations, ls_iterations = m.opt.iterations, m.opt.ls_iterations
+      mjw.step(m, d)  # warmup step
-      ls_str = f"{'parallel' if m.opt.ls_parallel else 'iterative'} linesearch iterations: {ls_iterations}"
+      print("Running Warp unoptimized on CPU.")
-      print(
+    broadphase, filter = mjw.BroadphaseType(m.opt.broadphase).name, mjw.BroadphaseFilter(m.opt.broadphase_filter).name
-        f"  nbody: {m.nbody} nv: {m.nv} ngeom: {m.ngeom} nu: {m.nu} is_sparse: {m.opt.is_sparse}\n"
+    solver, cone = mjw.SolverType(m.opt.solver).name, mjw.ConeType(m.opt.cone).name
-        f"  broadphase: {broadphase} broadphase_filter: {filter}\n"
+    integrator = mjw.IntegratorType(m.opt.integrator).name
-        f"  solver: {solver} cone: {cone} iterations: {iterations} {ls_str}\n"
+    iterations, ls_iterations = m.opt.iterations, m.opt.ls_iterations
-        f"  integrator: {integrator} graph_conditional: {m.opt.graph_conditional}"
+    ls_str = f"{'parallel' if m.opt.ls_parallel else 'iterative'} linesearch iterations: {ls_iterations}"
-      )
+    print(
-      d = mjw.put_data(mjm, mjd, nconmax=_NCONMAX.value, njmax=_NJMAX.value)
+      f"  nbody: {m.nbody} nv: {m.nv} ngeom: {m.ngeom} nu: {m.nu} is_sparse: {m.is_sparse}\n"
-      print(f"Data\n  nworld: {d.nworld} nconmax: {d.naconmax / d.nworld} njmax: {d.njmax}\n")
+      f"  broadphase: {broadphase} broadphase_filter: {filter}\n"
-      graph = _compile_step(m, d)
+      f"  solver: {solver} cone: {cone} iterations: {iterations} {ls_str}\n"
-      print(f"MuJoCo Warp simulating with dt = {m.opt.timestep.numpy()[0]:.3f}...")
+      f"  integrator: {integrator} graph_conditional: {m.opt.graph_conditional}"
+    )
+    print(f"Data\n  nworld: {d.nworld} nconmax: {int(d.naconmax / d.nworld)} njmax: {d.njmax}\n")
+    print(f"MuJoCo Warp simulating with dt = {m.opt.timestep.numpy()[0]:.3f}...")
 
   with mujoco.viewer.launch_passive(mjm, mjd, key_callback=key_callback) as viewer:
     opt = copy.copy(mjm.opt)
@@ -175,22 +187,19 @@ def _main(argv: Sequence[str]) -> None:
         wp.copy(d.qpos, wp.array([mjd.qpos.astype(np.float32)]))
         wp.copy(d.qvel, wp.array([mjd.qvel.astype(np.float32)]))
         wp.copy(d.time, wp.array([mjd.time], dtype=wp.float32))
-
         # if the user changed an option in the MuJoCo Simulate UI, go ahead and recompile the step
         # TODO: update memory tied to option max iterations
         if mjm.opt != opt:
           opt = copy.copy(mjm.opt)
           m = mjw.put_model(mjm)
-          graph = _compile_step(m, d)
+          graph = _compile_step(m, d) if wp.get_device().is_cuda else None
-
+        if _VIEWER_GLOBAL_STATE["running"] or _VIEWER_GLOBAL_STATE["step_once"]:
-        if _VIEWER_GLOBAL_STATE["running"]:
-          wp.capture_launch(graph)
-          wp.synchronize()
-        elif _VIEWER_GLOBAL_STATE["step_once"]:
           _VIEWER_GLOBAL_STATE["step_once"] = False
-          wp.capture_launch(graph)
+          if graph is None:
-          wp.synchronize()
+            mjw.step(m, d)
-
+          else:
+            wp.capture_launch(graph)
+            wp.synchronize()
         mjw.get_data_into(mjd, mjm, d)
 
       viewer.sync()

mjx/mujoco/mjx/warp/collision_driver.py

@@ -42,6 +42,7 @@ _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
@@ -112,10 +113,8 @@ def _collision_shim(
     opt__sdf_initpoints: int,
     opt__sdf_iterations: int,
     # Data
+    naccdmax: int,
     naconmax: int,
-    collision_pair: wp.array(dtype=wp.vec2i),
-    collision_pairid: wp.array(dtype=wp.vec2i),
-    collision_worldid: wp.array(dtype=int),
     geom_xmat: wp.array2d(dtype=wp.mat33),
     geom_xpos: wp.array2d(dtype=wp.vec3),
     nacon: wp.array(dtype=int),
@@ -203,9 +202,6 @@ def _collision_shim(
   _m.pair_solreffriction = pair_solreffriction
   _m.plugin = plugin
   _m.plugin_attr = plugin_attr
-  _d.collision_pair = collision_pair
-  _d.collision_pairid = collision_pairid
-  _d.collision_worldid = collision_worldid
   _d.contact.dim = contact__dim
   _d.contact.dist = contact__dist
   _d.contact.frame = contact__frame
@@ -221,6 +217,7 @@ def _collision_shim(
   _d.contact.worldid = contact__worldid
   _d.geom_xmat = geom_xmat
   _d.geom_xpos = geom_xpos
+  _d.naccdmax = naccdmax
   _d.nacon = nacon
   _d.naconmax = naconmax
   _d.ncollision = ncollision
@@ -230,9 +227,6 @@ def _collision_shim(
 
 def _collision_jax_impl(m: types.Model, d: types.Data):
   output_dims = {
-      'collision_pair': d._impl.collision_pair.shape,
-      'collision_pairid': d._impl.collision_pairid.shape,
-      'collision_worldid': d._impl.collision_worldid.shape,
       'nacon': d._impl.nacon.shape,
       'ncollision': d._impl.ncollision.shape,
       'contact__dim': d._impl.contact__dim.shape,
@@ -251,13 +245,10 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _collision_shim,
-      num_outputs=18,
+      num_outputs=15,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames={
-          'collision_pair',
-          'collision_pairid',
-          'collision_worldid',
           'nacon',
           'ncollision',
           'contact__dim',
@@ -364,10 +355,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m.opt.enableflags,
       m.opt._impl.sdf_initpoints,
       m.opt._impl.sdf_iterations,
+      d._impl.naccdmax,
       d._impl.naconmax,
-      d._impl.collision_pair,
-      d._impl.collision_pairid,
-      d._impl.collision_worldid,
       d.geom_xmat,
       d.geom_xpos,
       d._impl.nacon,
@@ -387,24 +376,21 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       d._impl.contact__worldid,
   )
   d = d.tree_replace({
-      '_impl.collision_pair': out[0],
+      '_impl.nacon': out[0],
-      '_impl.collision_pairid': out[1],
+      '_impl.ncollision': out[1],
-      '_impl.collision_worldid': out[2],
+      '_impl.contact__dim': out[2],
-      '_impl.nacon': out[3],
+      '_impl.contact__dist': out[3],
-      '_impl.ncollision': out[4],
+      '_impl.contact__frame': out[4],
-      '_impl.contact__dim': out[5],
+      '_impl.contact__friction': out[5],
-      '_impl.contact__dist': out[6],
+      '_impl.contact__geom': out[6],
-      '_impl.contact__frame': out[7],
+      '_impl.contact__geomcollisionid': out[7],
-      '_impl.contact__friction': out[8],
+      '_impl.contact__includemargin': out[8],
-      '_impl.contact__geom': out[9],
+      '_impl.contact__pos': out[9],
-      '_impl.contact__geomcollisionid': out[10],
+      '_impl.contact__solimp': out[10],
-      '_impl.contact__includemargin': out[11],
+      '_impl.contact__solref': out[11],
-      '_impl.contact__pos': out[12],
+      '_impl.contact__solreffriction': out[12],
-      '_impl.contact__solimp': out[13],
+      '_impl.contact__type': out[13],
-      '_impl.contact__solref': out[14],
+      '_impl.contact__worldid': out[14],
-      '_impl.contact__solreffriction': out[15],
-      '_impl.contact__type': out[16],
-      '_impl.contact__worldid': out[17],
   })
   return d
 

mjx/mujoco/mjx/warp/smooth.py

@@ -42,10 +42,13 @@ _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
     nworld: int,
+    body_branch_start: wp.array(dtype=int),
+    body_branches: wp.array(dtype=int),
     body_ipos: wp.array2d(dtype=wp.vec3),
     body_iquat: wp.array2d(dtype=wp.quat),
     body_jntadr: wp.array(dtype=int),
@@ -55,7 +58,6 @@ def _kinematics_shim(
     body_pos: wp.array2d(dtype=wp.vec3),
     body_quat: wp.array2d(dtype=wp.quat),
     body_rootid: wp.array(dtype=int),
-    body_tree: tuple[wp.array(dtype=int), ...],
     body_weldid: wp.array(dtype=int),
     geom_bodyid: wp.array(dtype=int),
     geom_pos: wp.array2d(dtype=wp.vec3),
@@ -64,6 +66,8 @@ def _kinematics_shim(
     jnt_pos: wp.array2d(dtype=wp.vec3),
     jnt_qposadr: wp.array(dtype=int),
     jnt_type: wp.array(dtype=int),
+    nbody: int,
+    nbranch: int,
     ngeom: int,
     nsite: int,
     qpos0: wp.array2d(dtype=float),
@@ -90,6 +94,8 @@ def _kinematics_shim(
   _m.opt = _o
   _d.efc = _e
   _d.contact = _c
+  _m.body_branch_start = body_branch_start
+  _m.body_branches = body_branches
   _m.body_ipos = body_ipos
   _m.body_iquat = body_iquat
   _m.body_jntadr = body_jntadr
@@ -99,7 +105,6 @@ def _kinematics_shim(
   _m.body_pos = body_pos
   _m.body_quat = body_quat
   _m.body_rootid = body_rootid
-  _m.body_tree = body_tree
   _m.body_weldid = body_weldid
   _m.geom_bodyid = geom_bodyid
   _m.geom_pos = geom_pos
@@ -108,6 +113,8 @@ def _kinematics_shim(
   _m.jnt_pos = jnt_pos
   _m.jnt_qposadr = jnt_qposadr
   _m.jnt_type = jnt_type
+  _m.nbody = nbody
+  _m.nbranch = nbranch
   _m.ngeom = ngeom
   _m.nsite = nsite
   _m.qpos0 = qpos0
@@ -208,6 +215,8 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
   )
   out = jf(
       d.qpos.shape[0],
+      m._impl.body_branch_start,
+      m._impl.body_branches,
       m.body_ipos,
       m.body_iquat,
       m.body_jntadr,
@@ -217,7 +226,6 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
       m.body_pos,
       m.body_quat,
       m.body_rootid,
-      m._impl.body_tree,
       m.body_weldid,
       m.geom_bodyid,
       m.geom_pos,
@@ -226,6 +234,8 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
       m.jnt_pos,
       m.jnt_qposadr,
       m.jnt_type,
+      m.nbody,
+      m._impl.nbranch,
       m.ngeom,
       m.nsite,
       m.qpos0,

mjx/mujoco/mjx/warp/types.py

@@ -72,7 +72,7 @@ class BlockDim:
   energy_vel_kinetic: int
   euler_dense: int
   linesearch_iterative: int
-  mul_m_dense: int
+  qderiv_actuator_dense: int
   ray: int
   segmented_sort: int
   tendon_velocity: int
@@ -93,7 +93,7 @@ class BlockDim:
 
 class StatisticWarp(PyTreeNode):
   """Derived fields from Statistic."""
-  meaninertia: float
+  meaninertia: jax.Array
 
 class OptionWarp(PyTreeNode):
   """Derived fields from Option."""
@@ -104,9 +104,7 @@ class OptionWarp(PyTreeNode):
   contact_sensor_maxmatch: int
   graph_conditional: bool
   graph_mode: GraphMode
-  has_fluid: bool
   impratio_invsqrt: jax.Array
-  is_sparse: bool
   ls_parallel: bool
   ls_parallel_min_step: float
   run_collision_detection: bool
@@ -120,8 +118,11 @@ class ModelWarp(PyTreeNode):
   M_rownnz: np.ndarray
   actuator_trntype_body_adr: np.ndarray
   block_dim: BlockDim
+  body_branch_start: np.ndarray
+  body_branches: np.ndarray
   body_fluid_ellipsoid: np.ndarray
   body_tree: Tuple[np.ndarray, ...]
+  cam_projection: np.ndarray
   collision_sensor_adr: np.ndarray
   dof_tri_col: np.ndarray
   dof_tri_row: np.ndarray
@@ -146,11 +147,16 @@ class ModelWarp(PyTreeNode):
   flex_vertadr: np.ndarray
   flex_vertbodyid: np.ndarray
   flex_vertnum: np.ndarray
+  flexedge_J_colind: np.ndarray
+  flexedge_J_rowadr: np.ndarray
+  flexedge_J_rownnz: np.ndarray
   flexedge_invweight0: np.ndarray
   flexedge_length0: np.ndarray
   geom_pair_type_count: Tuple[int, ...]
   geom_plugin_index: np.ndarray
+  has_fluid: bool
   has_sdf_geom: bool
+  is_sparse: bool
   jnt_limited_ball_adr: np.ndarray
   jnt_limited_slide_hinge_adr: np.ndarray
   light_active: jax.Array
@@ -169,6 +175,7 @@ class ModelWarp(PyTreeNode):
   mesh_polyvertnum: np.ndarray
   mocap_bodyid: np.ndarray
   nacttrnbody: int
+  nbranch: int
   nflex: int
   nflexedge: int
   nflexelem: int
@@ -185,6 +192,7 @@ class ModelWarp(PyTreeNode):
   nsensorcollision: int
   nsensorcontact: int
   nsensortaxel: int
+  ntree: int
   nv_pad: int
   nxn_geom_pair: np.ndarray
   nxn_geom_pair_filtered: np.ndarray
@@ -198,9 +206,9 @@ class ModelWarp(PyTreeNode):
   qLD_updates: Tuple[np.ndarray, ...]
   qM_fullm_i: np.ndarray
   qM_fullm_j: np.ndarray
-  qM_madr_ij: np.ndarray
+  qM_mulm_col: np.ndarray
-  qM_mulm_i: np.ndarray
+  qM_mulm_madr: np.ndarray
-  qM_mulm_j: np.ndarray
+  qM_mulm_rowadr: np.ndarray
   qM_tiles: Tuple[TileSet, ...]
   rangefinder_sensor_adr: np.ndarray
   sensor_acc_adr: np.ndarray
@@ -228,6 +236,9 @@ class ModelWarp(PyTreeNode):
   tendon_jnt_adr: np.ndarray
   tendon_limited_adr: np.ndarray
   tendon_site_pair_adr: np.ndarray
+  tree_bodynum: np.ndarray
+  tree_dofadr: np.ndarray
+  tree_dofnum: np.ndarray
   wrap_geom_adr: np.ndarray
   wrap_jnt_adr: np.ndarray
   wrap_pulley_scale: np.ndarray
@@ -242,9 +253,6 @@ class DataWarp(PyTreeNode):
   cfrc_ext: jax.Array
   cfrc_int: jax.Array
   cinert: jax.Array
-  collision_pair: jax.Array
-  collision_pairid: jax.Array
-  collision_worldid: jax.Array
   contact__dim: jax.Array
   contact__dist: jax.Array
   contact__efc_address: jax.Array
@@ -262,31 +270,13 @@ class DataWarp(PyTreeNode):
   crb: jax.Array
   efc__D: jax.Array
   efc__J: jax.Array
-  efc__Jaref: jax.Array
   efc__Ma: jax.Array
-  efc__Mgrad: jax.Array
-  efc__alpha: jax.Array
   efc__aref: jax.Array
-  efc__beta: jax.Array
-  efc__cost: jax.Array
-  efc__done: jax.Array
   efc__force: jax.Array
   efc__frictionloss: jax.Array
-  efc__gauss: jax.Array
-  efc__grad: jax.Array
-  efc__grad_dot: jax.Array
   efc__id: jax.Array
-  efc__jv: jax.Array
   efc__margin: jax.Array
-  efc__mv: jax.Array
   efc__pos: jax.Array
-  efc__prev_Mgrad: jax.Array
-  efc__prev_cost: jax.Array
-  efc__prev_grad: jax.Array
-  efc__quad: jax.Array
-  efc__quad_gauss: jax.Array
-  efc__search: jax.Array
-  efc__search_dot: jax.Array
   efc__state: jax.Array
   efc__type: jax.Array
   efc__vel: jax.Array
@@ -297,20 +287,15 @@ class DataWarp(PyTreeNode):
   flexvert_xpos: jax.Array
   light_xdir: jax.Array
   light_xpos: jax.Array
+  naccdmax: int
   nacon: jax.Array
   naconmax: int
   ncollision: jax.Array
   ne: jax.Array
-  ne_connect: jax.Array
-  ne_flex: jax.Array
-  ne_jnt: jax.Array
-  ne_ten: jax.Array
-  ne_weld: jax.Array
   nefc: jax.Array
   nf: jax.Array
   njmax: int
   nl: jax.Array
-  nsolving: jax.Array
   nworld: int
   qLD: jax.Array
   qLDiagInv: jax.Array
@@ -319,7 +304,6 @@ class DataWarp(PyTreeNode):
   qfrc_spring: jax.Array
   solver_niter: jax.Array
   subtree_angmom: jax.Array
-  subtree_bodyvel: jax.Array
   subtree_linvel: jax.Array
   ten_J: jax.Array
   ten_velocity: jax.Array
@@ -329,9 +313,6 @@ class DataWarp(PyTreeNode):
   wrap_xpos: jax.Array
   shape = property(lambda self: self.cacc.shape)
 DATA_NON_VMAP = {
-    'collision_pair',
-    'collision_pairid',
-    'collision_worldid',
     'contact__dim',
     'contact__dist',
     'contact__efc_address',
@@ -346,11 +327,11 @@ DATA_NON_VMAP = {
     'contact__solreffriction',
     'contact__type',
     'contact__worldid',
+    'naccdmax',
     'nacon',
     'naconmax',
     'ncollision',
     'njmax',
-    'nsolving',
     'nworld',
 }
 
@@ -394,9 +375,6 @@ _NDIM = {
         'cfrc_ext': 3,
         'cfrc_int': 3,
         'cinert': 3,
-        'collision_pair': 2,
-        'collision_pairid': 2,
-        'collision_worldid': 1,
         'contact__dim': 1,
         'contact__dist': 1,
         'contact__efc_address': 2,
@@ -416,31 +394,13 @@ _NDIM = {
         'cvel': 3,
         'efc__D': 2,
         'efc__J': 3,
-        'efc__Jaref': 2,
         'efc__Ma': 2,
-        'efc__Mgrad': 2,
-        'efc__alpha': 1,
         'efc__aref': 2,
-        'efc__beta': 1,
-        'efc__cost': 1,
-        'efc__done': 1,
         'efc__force': 2,
         'efc__frictionloss': 2,
-        'efc__gauss': 1,
-        'efc__grad': 2,
-        'efc__grad_dot': 1,
         'efc__id': 2,
-        'efc__jv': 2,
         'efc__margin': 2,
-        'efc__mv': 2,
         'efc__pos': 2,
-        'efc__prev_Mgrad': 2,
-        'efc__prev_cost': 1,
-        'efc__prev_grad': 2,
-        'efc__quad': 3,
-        'efc__quad_gauss': 2,
-        'efc__search': 2,
-        'efc__search_dot': 1,
         'efc__state': 2,
         'efc__type': 2,
         'efc__vel': 2,
@@ -456,20 +416,15 @@ _NDIM = {
         'light_xpos': 3,
         'mocap_pos': 3,
         'mocap_quat': 3,
+        'naccdmax': 0,
         'nacon': 1,
         'naconmax': 0,
         'ncollision': 1,
         'ne': 1,
-        'ne_connect': 1,
-        'ne_flex': 1,
-        'ne_jnt': 1,
-        'ne_ten': 1,
-        'ne_weld': 1,
         'nefc': 1,
         'nf': 1,
         'njmax': 0,
         'nl': 1,
-        'nsolving': 1,
         'nworld': 0,
         'qLD': 3,
         'qLDiagInv': 2,
@@ -495,7 +450,6 @@ _NDIM = {
         'site_xpos': 3,
         'solver_niter': 1,
         'subtree_angmom': 3,
-        'subtree_bodyvel': 3,
         'subtree_com': 3,
         'subtree_linvel': 3,
         'ten_J': 3,
@@ -549,7 +503,7 @@ _NDIM = {
         'block_dim__energy_vel_kinetic': 0,
         'block_dim__euler_dense': 0,
         'block_dim__linesearch_iterative': 0,
-        'block_dim__mul_m_dense': 0,
+        'block_dim__qderiv_actuator_dense': 0,
         'block_dim__ray': 0,
         'block_dim__segmented_sort': 0,
         'block_dim__tendon_velocity': 0,
@@ -557,6 +511,8 @@ _NDIM = {
         'block_dim__update_gradient_JTDAJ_sparse': 0,
         'block_dim__update_gradient_cholesky': 0,
         'block_dim__update_gradient_cholesky_blocked': 0,
+        'body_branch_start': 1,
+        'body_branches': 1,
         'body_conaffinity': 1,
         'body_contype': 1,
         'body_dofadr': 1,
@@ -579,15 +535,17 @@ _NDIM = {
         'body_rootid': 1,
         'body_subtreemass': 2,
         'body_tree': -1,
+        'body_treeid': 1,
         'body_weldid': 1,
         'cam_bodyid': 1,
-        'cam_fovy': 1,
+        'cam_fovy': 2,
-        'cam_intrinsic': 2,
+        'cam_intrinsic': 3,
         'cam_mat0': 4,
         'cam_mode': 1,
         'cam_pos': 3,
         'cam_pos0': 3,
         'cam_poscom0': 3,
+        'cam_projection': 1,
         'cam_quat': 3,
         'cam_resolution': 2,
         'cam_sensorsize': 2,
@@ -603,6 +561,7 @@ _NDIM = {
         'dof_parentid': 1,
         'dof_solimp': 3,
         'dof_solref': 3,
+        'dof_treeid': 1,
         'dof_tri_col': 1,
         'dof_tri_row': 1,
         'eq_active0': 1,
@@ -635,6 +594,9 @@ _NDIM = {
         'flex_vertadr': 1,
         'flex_vertbodyid': 1,
         'flex_vertnum': 1,
+        'flexedge_J_colind': 1,
+        'flexedge_J_rowadr': 1,
+        'flexedge_J_rownnz': 1,
         'flexedge_invweight0': 1,
         'flexedge_length0': 1,
         'geom_aabb': 4,
@@ -661,12 +623,14 @@ _NDIM = {
         'geom_solmix': 2,
         'geom_solref': 3,
         'geom_type': 1,
+        'has_fluid': 0,
         'has_sdf_geom': 0,
         'hfield_adr': 1,
         'hfield_data': 1,
         'hfield_ncol': 1,
         'hfield_nrow': 1,
         'hfield_size': 2,
+        'is_sparse': 0,
         'jnt_actfrclimited': 1,
         'jnt_actfrcrange': 3,
         'jnt_actgravcomp': 1,
@@ -697,7 +661,7 @@ _NDIM = {
         'light_type': 2,
         'mapM2M': 1,
         'mat_rgba': 3,
-        'mat_texid': 2,
+        'mat_texid': 3,
         'mat_texrepeat': 3,
         'mesh_face': 2,
         'mesh_faceadr': 1,
@@ -724,6 +688,7 @@ _NDIM = {
         'na': 0,
         'nacttrnbody': 0,
         'nbody': 0,
+        'nbranch': 0,
         'ncam': 0,
         'neq': 0,
         'nexclude': 0,
@@ -765,6 +730,7 @@ _NDIM = {
         'nsensortaxel': 0,
         'nsite': 0,
         'ntendon': 0,
+        'ntree': 0,
         'nu': 0,
         'nv': 0,
         'nv_pad': 0,
@@ -785,10 +751,8 @@ _NDIM = {
         'opt__enableflags': 0,
         'opt__graph_conditional': 0,
         'opt__gravity': 2,
-        'opt__has_fluid': 0,
         'opt__impratio_invsqrt': 1,
         'opt__integrator': 0,
-        'opt__is_sparse': 0,
         'opt__iterations': 0,
         'opt__ls_iterations': 0,
         'opt__ls_parallel': 0,
@@ -817,9 +781,9 @@ _NDIM = {
         'qLD_updates': -1,
         'qM_fullm_i': 1,
         'qM_fullm_j': 1,
-        'qM_madr_ij': 1,
+        'qM_mulm_col': 1,
-        'qM_mulm_i': 1,
+        'qM_mulm_madr': 1,
-        'qM_mulm_j': 1,
+        'qM_mulm_rowadr': 1,
         'qM_tiles': -1,
         'qpos0': 2,
         'qpos_spring': 2,
@@ -856,7 +820,7 @@ _NDIM = {
         'site_quat': 3,
         'site_size': 2,
         'site_type': 1,
-        'stat__meaninertia': 0,
+        'stat__meaninertia': 1,
         'taxel_sensorid': 1,
         'taxel_vertadr': 1,
         'ten_wrapadr_site': 1,
@@ -883,6 +847,9 @@ _NDIM = {
         'tendon_solref_fri': 3,
         'tendon_solref_lim': 3,
         'tendon_stiffness': 2,
+        'tree_bodynum': 1,
+        'tree_dofadr': 1,
+        'tree_dofnum': 1,
         'wrap_geom_adr': 1,
         'wrap_jnt_adr': 1,
         'wrap_objid': 1,
@@ -902,10 +869,8 @@ _NDIM = {
         'enableflags': 0,
         'graph_conditional': 0,
         'gravity': 2,
-        'has_fluid': 0,
         'impratio_invsqrt': 1,
         'integrator': 0,
-        'is_sparse': 0,
         'iterations': 0,
         'ls_iterations': 0,
         'ls_parallel': 0,
@@ -921,7 +886,7 @@ _NDIM = {
         'viscosity': 1,
         'wind': 2,
     },
-    'Statistic': {'meaninertia': 0},
+    'Statistic': {'meaninertia': 1},
 }
 _BATCH_DIM = {
     'Data': {
@@ -939,9 +904,6 @@ _BATCH_DIM = {
         'cfrc_ext': True,
         'cfrc_int': True,
         'cinert': True,
-        'collision_pair': False,
-        'collision_pairid': False,
-        'collision_worldid': False,
         'contact__dim': False,
         'contact__dist': False,
         'contact__efc_address': False,
@@ -961,31 +923,13 @@ _BATCH_DIM = {
         'cvel': True,
         'efc__D': True,
         'efc__J': True,
-        'efc__Jaref': True,
         'efc__Ma': True,
-        'efc__Mgrad': True,
-        'efc__alpha': True,
         'efc__aref': True,
-        'efc__beta': True,
-        'efc__cost': True,
-        'efc__done': True,
         'efc__force': True,
         'efc__frictionloss': True,
-        'efc__gauss': True,
-        'efc__grad': True,
-        'efc__grad_dot': True,
         'efc__id': True,
-        'efc__jv': True,
         'efc__margin': True,
-        'efc__mv': True,
         'efc__pos': True,
-        'efc__prev_Mgrad': True,
-        'efc__prev_cost': True,
-        'efc__prev_grad': True,
-        'efc__quad': True,
-        'efc__quad_gauss': True,
-        'efc__search': True,
-        'efc__search_dot': True,
         'efc__state': True,
         'efc__type': True,
         'efc__vel': True,
@@ -1001,20 +945,15 @@ _BATCH_DIM = {
         'light_xpos': True,
         'mocap_pos': True,
         'mocap_quat': True,
+        'naccdmax': False,
         'nacon': False,
         'naconmax': False,
         'ncollision': False,
         'ne': True,
-        'ne_connect': True,
-        'ne_flex': True,
-        'ne_jnt': True,
-        'ne_ten': True,
-        'ne_weld': True,
         'nefc': True,
         'nf': True,
         'njmax': False,
         'nl': True,
-        'nsolving': False,
         'nworld': False,
         'qLD': True,
         'qLDiagInv': True,
@@ -1040,7 +979,6 @@ _BATCH_DIM = {
         'site_xpos': True,
         'solver_niter': True,
         'subtree_angmom': True,
-        'subtree_bodyvel': True,
         'subtree_com': True,
         'subtree_linvel': True,
         'ten_J': True,
@@ -1094,7 +1032,7 @@ _BATCH_DIM = {
         'block_dim__energy_vel_kinetic': False,
         'block_dim__euler_dense': False,
         'block_dim__linesearch_iterative': False,
-        'block_dim__mul_m_dense': False,
+        'block_dim__qderiv_actuator_dense': False,
         'block_dim__ray': False,
         'block_dim__segmented_sort': False,
         'block_dim__tendon_velocity': False,
@@ -1102,6 +1040,8 @@ _BATCH_DIM = {
         'block_dim__update_gradient_JTDAJ_sparse': False,
         'block_dim__update_gradient_cholesky': False,
         'block_dim__update_gradient_cholesky_blocked': False,
+        'body_branch_start': False,
+        'body_branches': False,
         'body_conaffinity': False,
         'body_contype': False,
         'body_dofadr': False,
@@ -1124,15 +1064,17 @@ _BATCH_DIM = {
         'body_rootid': False,
         'body_subtreemass': True,
         'body_tree': False,
+        'body_treeid': False,
         'body_weldid': False,
         'cam_bodyid': False,
-        'cam_fovy': False,
+        'cam_fovy': True,
-        'cam_intrinsic': False,
+        'cam_intrinsic': True,
         'cam_mat0': True,
         'cam_mode': False,
         'cam_pos': True,
         'cam_pos0': True,
         'cam_poscom0': True,
+        'cam_projection': False,
         'cam_quat': True,
         'cam_resolution': False,
         'cam_sensorsize': False,
@@ -1148,6 +1090,7 @@ _BATCH_DIM = {
         'dof_parentid': False,
         'dof_solimp': True,
         'dof_solref': True,
+        'dof_treeid': False,
         'dof_tri_col': False,
         'dof_tri_row': False,
         'eq_active0': False,
@@ -1180,6 +1123,9 @@ _BATCH_DIM = {
         'flex_vertadr': False,
         'flex_vertbodyid': False,
         'flex_vertnum': False,
+        'flexedge_J_colind': False,
+        'flexedge_J_rowadr': False,
+        'flexedge_J_rownnz': False,
         'flexedge_invweight0': False,
         'flexedge_length0': False,
         'geom_aabb': True,
@@ -1206,12 +1152,14 @@ _BATCH_DIM = {
         'geom_solmix': True,
         'geom_solref': True,
         'geom_type': False,
+        'has_fluid': False,
         'has_sdf_geom': False,
         'hfield_adr': False,
         'hfield_data': False,
         'hfield_ncol': False,
         'hfield_nrow': False,
         'hfield_size': False,
+        'is_sparse': False,
         'jnt_actfrclimited': False,
         'jnt_actfrcrange': True,
         'jnt_actgravcomp': False,
@@ -1242,7 +1190,7 @@ _BATCH_DIM = {
         'light_type': True,
         'mapM2M': False,
         'mat_rgba': True,
-        'mat_texid': False,
+        'mat_texid': True,
         'mat_texrepeat': True,
         'mesh_face': False,
         'mesh_faceadr': False,
@@ -1269,6 +1217,7 @@ _BATCH_DIM = {
         'na': False,
         'nacttrnbody': False,
         'nbody': False,
+        'nbranch': False,
         'ncam': False,
         'neq': False,
         'nexclude': False,
@@ -1310,6 +1259,7 @@ _BATCH_DIM = {
         'nsensortaxel': False,
         'nsite': False,
         'ntendon': False,
+        'ntree': False,
         'nu': False,
         'nv': False,
         'nv_pad': False,
@@ -1330,10 +1280,8 @@ _BATCH_DIM = {
         'opt__enableflags': False,
         'opt__graph_conditional': False,
         'opt__gravity': True,
-        'opt__has_fluid': False,
         'opt__impratio_invsqrt': True,
         'opt__integrator': False,
-        'opt__is_sparse': False,
         'opt__iterations': False,
         'opt__ls_iterations': False,
         'opt__ls_parallel': False,
@@ -1362,9 +1310,9 @@ _BATCH_DIM = {
         'qLD_updates': False,
         'qM_fullm_i': False,
         'qM_fullm_j': False,
-        'qM_madr_ij': False,
+        'qM_mulm_col': False,
-        'qM_mulm_i': False,
+        'qM_mulm_madr': False,
-        'qM_mulm_j': False,
+        'qM_mulm_rowadr': False,
         'qM_tiles': False,
         'qpos0': True,
         'qpos_spring': True,
@@ -1401,7 +1349,7 @@ _BATCH_DIM = {
         'site_quat': True,
         'site_size': False,
         'site_type': False,
-        'stat__meaninertia': False,
+        'stat__meaninertia': True,
         'taxel_sensorid': False,
         'taxel_vertadr': False,
         'ten_wrapadr_site': False,
@@ -1428,6 +1376,9 @@ _BATCH_DIM = {
         'tendon_solref_fri': True,
         'tendon_solref_lim': True,
         'tendon_stiffness': True,
+        'tree_bodynum': False,
+        'tree_dofadr': False,
+        'tree_dofnum': False,
         'wrap_geom_adr': False,
         'wrap_jnt_adr': False,
         'wrap_objid': False,
@@ -1447,10 +1398,8 @@ _BATCH_DIM = {
         'enableflags': False,
         'graph_conditional': False,
         'gravity': True,
-        'has_fluid': False,
         'impratio_invsqrt': True,
         'integrator': False,
-        'is_sparse': False,
         'iterations': False,
         'ls_iterations': False,
         'ls_parallel': False,
@@ -1466,5 +1415,5 @@ _BATCH_DIM = {
         'viscosity': True,
         'wind': True,
     },
-    'Statistic': {'meaninertia': False},
+    'Statistic': {'meaninertia': True},
 }