Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 892971818
Change-Id: Ibce0c41294a4fcf2825a246a9533e8927fdc7e82

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

@@ -64,6 +64,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.ray import rays as rays
 from mujoco.mjx.third_party.mujoco_warp._src.render import render as render
 from mujoco.mjx.third_party.mujoco_warp._src.render_util import get_depth as get_depth
 from mujoco.mjx.third_party.mujoco_warp._src.render_util import get_rgb as get_rgb
+from mujoco.mjx.third_party.mujoco_warp._src.render_util import get_segmentation as get_segmentation
 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
@@ -92,6 +93,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate as x
 from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType as BiasType
 from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseFilter as BroadphaseFilter
 from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType as BroadphaseType
+from mujoco.mjx.third_party.mujoco_warp._src.types import Callback as Callback
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType as ConeType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Constraint as Constraint
 from mujoco.mjx.third_party.mujoco_warp._src.types import Contact as Contact

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

@@ -189,12 +189,12 @@ def _compute_bvh_bounds(
   upper_out: wp.array(dtype=wp.vec3),
   group_out: wp.array(dtype=int),
 ):
-  world_id, geom_local_id = wp.tid()
+  worldid, geom_local_id = wp.tid()
   geom_id = enabled_geom_ids[geom_local_id]
 
-  pos = geom_xpos_in[world_id, geom_id]
+  pos = geom_xpos_in[worldid, geom_id]
-  rot = geom_xmat_in[world_id, geom_id]
+  rot = geom_xmat_in[worldid, geom_id]
-  size = geom_size[world_id % geom_size.shape[0], geom_id]
+  size = geom_size[worldid % geom_size.shape[0], geom_id]
   type = geom_type[geom_id]
 
   # TODO: Investigate branch elimination with static loop unrolling
@@ -218,9 +218,9 @@ def _compute_bvh_bounds(
     hfield_center = pos + rot[:, 2] * size[2]
     lower_bound, upper_bound = _compute_box_bounds(hfield_center, rot, size)
 
-  lower_out[world_id * bvh_ngeom + geom_local_id] = lower_bound
+  lower_out[worldid * bvh_ngeom + geom_local_id] = lower_bound
-  upper_out[world_id * bvh_ngeom + geom_local_id] = upper_bound
+  upper_out[worldid * bvh_ngeom + geom_local_id] = upper_bound
-  group_out[world_id * bvh_ngeom + geom_local_id] = world_id
+  group_out[worldid * bvh_ngeom + geom_local_id] = worldid
 
 
 @wp.kernel
@@ -235,14 +235,70 @@ def compute_bvh_group_roots(
   group_root_out[tid] = root
 
 
+@wp.kernel
+def _compute_flex_bvh_bounds(
+  # Model:
+  flex_vertadr: wp.array(dtype=int),
+  flex_vertnum: wp.array(dtype=int),
+  flex_edge: wp.array(dtype=wp.vec2i),
+  flex_radius: wp.array(dtype=float),
+  # Data in:
+  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  # In:
+  flex_geom_flexid: wp.array(dtype=int),
+  flex_geom_edgeid: wp.array(dtype=int),
+  bvh_ngeom: int,
+  total_bvh_size: int,
+  # Out:
+  lower_out: wp.array(dtype=wp.vec3),
+  upper_out: wp.array(dtype=wp.vec3),
+  group_out: wp.array(dtype=int),
+):
+  worldid, flexlocalid = wp.tid()
+
+  flex_id = flex_geom_flexid[flexlocalid]
+  edge_id = flex_geom_edgeid[flexlocalid]
+  out_idx = worldid * total_bvh_size + bvh_ngeom + flexlocalid
+  radius = flex_radius[flex_id]
+  inflate = wp.vec3(radius, radius, radius)
+
+  if edge_id >= 0:  # capsule (1D edge)
+    edge = flex_edge[edge_id]
+    vert_adr = flex_vertadr[flex_id]
+    v0 = flexvert_xpos_in[worldid, vert_adr + edge[0]]
+    v1 = flexvert_xpos_in[worldid, vert_adr + edge[1]]
+    lower_out[out_idx] = wp.min(v0, v1) - inflate
+    upper_out[out_idx] = wp.max(v0, v1) + inflate
+  else:  # mesh (2D/3D)
+    vert_adr = flex_vertadr[flex_id]
+    nvert = flex_vertnum[flex_id]
+    min_bound = wp.vec3(MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL)
+    max_bound = wp.vec3(-MJ_MAXVAL, -MJ_MAXVAL, -MJ_MAXVAL)
+    for i in range(nvert):
+      v = flexvert_xpos_in[worldid, vert_adr + i]
+      min_bound = wp.min(min_bound, v)
+      max_bound = wp.max(max_bound, v)
+    lower_out[out_idx] = min_bound - inflate
+    upper_out[out_idx] = max_bound + inflate
+
+  group_out[out_idx] = worldid
+
+
 def build_scene_bvh(mjm: mujoco.MjModel, mjd: mujoco.MjData, rc: RenderContext, nworld: int):
   """Build a global BVH for all geometries in all worlds."""
+  total_bvh_size = rc.bvh_ngeom + rc.bvh_nflexgeom
+
   geom_type = wp.array(mjm.geom_type, dtype=int)
   geom_dataid = wp.array(mjm.geom_dataid, dtype=int)
   geom_size = wp.array(np.tile(mjm.geom_size[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3)
   geom_xpos = wp.array(np.tile(mjd.geom_xpos[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3)
   geom_xmat = wp.array(np.tile(mjd.geom_xmat.reshape(mjm.ngeom, 3, 3)[np.newaxis, :, :, :], (nworld, 1, 1, 1)), dtype=wp.mat33)
 
+  flex_vertadr = wp.array(mjm.flex_vertadr, dtype=int)
+  flex_vertnum = wp.array(mjm.flex_vertnum, dtype=int)
+  flex_edge = wp.array(mjm.flex_edge, dtype=wp.vec2i)
+  flex_radius = wp.array(mjm.flex_radius, dtype=float)
+
   wp.launch(
     kernel=_compute_bvh_bounds,
     dim=(nworld, rc.bvh_ngeom),
@@ -252,7 +308,7 @@ def build_scene_bvh(mjm: mujoco.MjModel, mjd: mujoco.MjData, rc: RenderContext,
       geom_size,
       geom_xpos,
       geom_xmat,
-      rc.bvh_ngeom,
+      total_bvh_size,
       rc.enabled_geom_ids,
       rc.mesh_bounds_size,
       rc.hfield_bounds_size,
@@ -262,6 +318,26 @@ def build_scene_bvh(mjm: mujoco.MjModel, mjd: mujoco.MjData, rc: RenderContext,
     ],
   )
 
+  flexvert_xpos = wp.array(np.tile(mjd.flexvert_xpos[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3)
+  wp.launch(
+    kernel=_compute_flex_bvh_bounds,
+    dim=(nworld, rc.bvh_nflexgeom),
+    inputs=[
+      flex_vertadr,
+      flex_vertnum,
+      flex_edge,
+      flex_radius,
+      flexvert_xpos,
+      rc.flex_geom_flexid,
+      rc.flex_geom_edgeid,
+      rc.bvh_ngeom,
+      total_bvh_size,
+      rc.lower,
+      rc.upper,
+      rc.group,
+    ],
+  )
+
   bvh = wp.Bvh(rc.lower, rc.upper, groups=rc.group, constructor="sah")
 
   # BVH handle must be stored to avoid garbage collection
@@ -277,6 +353,8 @@ def build_scene_bvh(mjm: mujoco.MjModel, mjd: mujoco.MjData, rc: RenderContext,
 
 
 def refit_scene_bvh(m: Model, d: Data, rc: RenderContext):
+  total_bvh_size = rc.bvh_ngeom + rc.bvh_nflexgeom
+
   wp.launch(
     kernel=_compute_bvh_bounds,
     dim=(d.nworld, rc.bvh_ngeom),
@@ -286,7 +364,7 @@ def refit_scene_bvh(m: Model, d: Data, rc: RenderContext):
       m.geom_size,
       d.geom_xpos,
       d.geom_xmat,
-      rc.bvh_ngeom,
+      total_bvh_size,
       rc.enabled_geom_ids,
       rc.mesh_bounds_size,
       rc.hfield_bounds_size,
@@ -296,6 +374,26 @@ def refit_scene_bvh(m: Model, d: Data, rc: RenderContext):
     ],
   )
 
+  if rc.bvh_nflexgeom > 0:
+    wp.launch(
+      kernel=_compute_flex_bvh_bounds,
+      dim=(d.nworld, rc.bvh_nflexgeom),
+      inputs=[
+        m.flex_vertadr,
+        m.flex_vertnum,
+        m.flex_edge,
+        m.flex_radius,
+        d.flexvert_xpos,
+        rc.flex_geom_flexid,
+        rc.flex_geom_edgeid,
+        rc.bvh_ngeom,
+        total_bvh_size,
+        rc.lower,
+        rc.upper,
+        rc.group,
+      ],
+    )
+
   rc.bvh.refit()
 
 
@@ -500,6 +598,12 @@ def build_hfield_bvh(
 @wp.kernel
 def accumulate_flex_vertex_normals(
   # Model:
+  nflex: int,
+  flex_dim: wp.array(dtype=int),
+  flex_vertadr: wp.array(dtype=int),
+  flex_elemadr: wp.array(dtype=int),
+  flex_elemnum: wp.array(dtype=int),
+  flex_elemdataadr: wp.array(dtype=int),
   flex_elem: wp.array(dtype=int),
   # Data in:
   flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
@@ -509,10 +613,22 @@ def accumulate_flex_vertex_normals(
   """Accumulate per-vertex normals by summing adjacent face normals."""
   worldid, elemid = wp.tid()
 
-  elem_base = elemid * 3
+  for i in range(nflex):
-  i0 = flex_elem[elem_base + 0]
+    locid = elemid - flex_elemadr[i]
-  i1 = flex_elem[elem_base + 1]
+    if locid >= 0 and locid < flex_elemnum[i]:
-  i2 = flex_elem[elem_base + 2]
+      f = i
+      break
+
+  if flex_dim[f] == 1 or flex_dim[f] == 3:
+    return
+
+  local_elemid = elemid - flex_elemadr[f]
+  elem_adr = flex_elemdataadr[f]
+  vert_adr = flex_vertadr[f]
+  elem_base = elem_adr + local_elemid * 3
+  i0 = vert_adr + flex_elem[elem_base + 0]
+  i1 = vert_adr + flex_elem[elem_base + 1]
+  i2 = vert_adr + flex_elem[elem_base + 2]
 
   v0 = flexvert_xpos_in[worldid, i0]
   v1 = flexvert_xpos_in[worldid, i1]
@@ -611,11 +727,12 @@ def _build_flex_2d_elements(
 
 @wp.kernel
 def _build_flex_2d_sides(
+  # Model:
+  flex_shell: wp.array(dtype=int),
   # Data in:
   flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
   # In:
   flexvert_norm_in: wp.array2d(dtype=wp.vec3),
-  flex_shell_in: wp.array(dtype=int),
   shell_adr: int,
   vert_adr: int,
   face_offset: int,
@@ -635,8 +752,8 @@ def _build_flex_2d_sides(
   worldid, shellid = wp.tid()
 
   base = shell_adr + 2 * shellid
-  i0 = vert_adr + flex_shell_in[base + 0]
+  i0 = vert_adr + flex_shell[base + 0]
-  i1 = vert_adr + flex_shell_in[base + 1]
+  i1 = vert_adr + flex_shell[base + 1]
 
   v0 = flexvert_xpos_in[worldid, i0]
   v1 = flexvert_xpos_in[worldid, i1]
@@ -672,10 +789,11 @@ def _build_flex_2d_sides(
 
 @wp.kernel
 def _build_flex_3d_shells(
+  # Model:
+  flex_shell: wp.array(dtype=int),
   # Data in:
   flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
   # In:
-  flex_shell_in: wp.array(dtype=int),
   shell_adr: int,
   vert_adr: int,
   face_offset: int,
@@ -693,9 +811,9 @@ def _build_flex_3d_shells(
   worldid, shellid = wp.tid()
 
   base = shell_adr + shellid * 3
-  i0 = vert_adr + flex_shell_in[base + 0]
+  i0 = vert_adr + flex_shell[base + 0]
-  i1 = vert_adr + flex_shell_in[base + 1]
+  i1 = vert_adr + flex_shell[base + 1]
-  i2 = vert_adr + flex_shell_in[base + 2]
+  i2 = vert_adr + flex_shell[base + 2]
 
   face_id = worldid * nface + face_offset + shellid
   base = face_id * 3
@@ -716,53 +834,36 @@ def _build_flex_3d_shells(
 
 
 @wp.kernel
-def _update_flex_face_points(
+def _update_flex_2d_face_points(
   # Model:
-  nflex: int,
-  flex_dim: wp.array(dtype=int),
   flex_vertadr: wp.array(dtype=int),
   flex_elemnum: wp.array(dtype=int),
+  flex_elemdataadr: wp.array(dtype=int),
+  flex_shelldataadr: wp.array(dtype=int),
   flex_elem: wp.array(dtype=int),
+  flex_shell: wp.array(dtype=int),
+  flex_radius: wp.array(dtype=float),
   # Data in:
   flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
   # In:
-  flex_shell_in: wp.array(dtype=int),
   flexvert_norm_in: wp.array2d(dtype=wp.vec3),
-  flex_elemdataadr: wp.array(dtype=int),
+  flex_id: int,
-  flex_shelldataadr: wp.array(dtype=int),
+  nface: int,
-  flex_faceadr: wp.array(dtype=int),
-  flex_radius: wp.array(dtype=float),
-  flex_workadr: wp.array(dtype=int),
-  flex_worknum: wp.array(dtype=int),
-  nfaces: int,
   smooth: bool,
   # Out:
   face_point_out: wp.array(dtype=wp.vec3),
 ):
   worldid, workid = wp.tid()
 
-  # identify which flex this work item belongs to
+  elem_adr = flex_elemdataadr[flex_id]
-  f = int(0)
+  vert_adr = flex_vertadr[flex_id]
-  locid = int(0)
+  radius = flex_radius[flex_id]
-  for i in range(nflex):
+  nelem = flex_elemnum[flex_id]
-    locid = workid - flex_workadr[i]
+  world_face_offset = worldid * nface
-    if locid >= 0 and locid < flex_worknum[i]:
-      f = i
-      break
 
-  dim = flex_dim[f]
+  if workid < nelem:
-  face_offset = flex_faceadr[f]
-  world_face_offset = worldid * nfaces
-  vert_adr = flex_vertadr[f]
-
-  if dim == 2:
-    radius = flex_radius[f]
-    elem_count = flex_elemnum[f]
-
-    if locid < elem_count:
     # 2D element faces
-      elemid = locid
+    elemid = workid
-      elem_adr = flex_elemdataadr[f]
     ebase = elem_adr + elemid * 3
     i0 = vert_adr + flex_elem[ebase + 0]
     i1 = vert_adr + flex_elem[ebase + 1]
@@ -792,24 +893,24 @@ def _update_flex_face_points(
     p1_neg = v1 - radius * n1
     p2_neg = v2 - radius * n2
 
-      face_id0 = world_face_offset + face_offset + (2 * elemid)
+    face_id0 = world_face_offset + (2 * elemid)
     base0 = face_id0 * 3
     face_point_out[base0 + 0] = p0_pos
     face_point_out[base0 + 1] = p1_pos
     face_point_out[base0 + 2] = p2_pos
 
-      face_id1 = world_face_offset + face_offset + (2 * elemid + 1)
+    face_id1 = world_face_offset + (2 * elemid + 1)
     base1 = face_id1 * 3
     face_point_out[base1 + 0] = p0_neg
     face_point_out[base1 + 1] = p1_neg
     face_point_out[base1 + 2] = p2_neg
   else:
     # 2D shell faces
-      shellid = locid - elem_count
+    shell_adr = flex_shelldataadr[flex_id]
-      shell_adr = flex_shelldataadr[f]
+    shellid = workid - nelem
     sbase = shell_adr + 2 * shellid
-      i0 = vert_adr + flex_shell_in[sbase + 0]
+    i0 = vert_adr + flex_shell[sbase + 0]
-      i1 = vert_adr + flex_shell_in[sbase + 1]
+    i1 = vert_adr + flex_shell[sbase + 1]
 
     v0 = flexvert_xpos_in[worldid, i0]
     v1 = flexvert_xpos_in[worldid, i1]
@@ -817,7 +918,7 @@ def _update_flex_face_points(
     n0 = flexvert_norm_in[worldid, i0]
     n1 = flexvert_norm_in[worldid, i1]
 
-      shell_face_offset = face_offset + (2 * elem_count)
+    shell_face_offset = 2 * nelem
     face_id0 = world_face_offset + shell_face_offset + (2 * shellid)
     base0 = face_id0 * 3
     face_point_out[base0 + 0] = v0 + radius * n0
@@ -829,50 +930,67 @@ def _update_flex_face_points(
     face_point_out[base1 + 0] = v1 - radius * n1
     face_point_out[base1 + 1] = v0 + radius * n0
     face_point_out[base1 + 2] = v0 - radius * n0
-  else:
-    # 3D shell faces
-    shellid = locid
-    shell_adr = flex_shelldataadr[f]
-    sbase = shell_adr + shellid * 3
-    i0 = vert_adr + flex_shell_in[sbase + 0]
-    i1 = vert_adr + flex_shell_in[sbase + 1]
-    i2 = vert_adr + flex_shell_in[sbase + 2]
 
-    v0 = flexvert_xpos_in[worldid, i0]
-    v1 = flexvert_xpos_in[worldid, i1]
-    v2 = flexvert_xpos_in[worldid, i2]
 
-    face_id = world_face_offset + face_offset + shellid
+@wp.kernel
+def _update_flex_3d_face_points(
+  # Model:
+  flex_vertadr: wp.array(dtype=int),
+  flex_shelldataadr: wp.array(dtype=int),
+  flex_shell: wp.array(dtype=int),
+  # Data in:
+  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  # In:
+  flex_id: int,
+  nface: int,
+  # Out:
+  face_point_out: wp.array(dtype=wp.vec3),
+):
+  worldid, shellid = wp.tid()
+
+  shell_adr = flex_shelldataadr[flex_id]
+  vert_adr = flex_vertadr[flex_id]
+
+  face_id = worldid * nface + shellid
   fbase = face_id * 3
 
-    face_point_out[fbase + 0] = v0
+  sbase = shell_adr + shellid * 3
-    face_point_out[fbase + 1] = v1
+  i0 = vert_adr + flex_shell[sbase + 0]
-    face_point_out[fbase + 2] = v2
+  i1 = vert_adr + flex_shell[sbase + 1]
+  i2 = vert_adr + flex_shell[sbase + 2]
+
+  face_point_out[fbase + 0] = flexvert_xpos_in[worldid, i0]
+  face_point_out[fbase + 1] = flexvert_xpos_in[worldid, i1]
+  face_point_out[fbase + 2] = flexvert_xpos_in[worldid, i2]
 
 
 def build_flex_bvh(
-  mjm: mujoco.MjModel, mjd: mujoco.MjData, nworld: int, constructor: str = "sah", leaf_size: int = 2
+  mjm: mujoco.MjModel,
-) -> tuple[wp.Mesh, wp.array, wp.array, wp.array, wp.array, wp.array, int]:
+  mjd: mujoco.MjData,
-  """Create a Warp mesh BVH from flex data."""
+  nworld: int,
-  if (mjm.flex_dim == 1).any():
+  flex_id: int,
-    raise ValueError("1D Flex objects are not currently supported.")
+  constructor: str = "sah",
-
+  leaf_size: int = 2,
-  nflex = mjm.nflex
+) -> tuple[wp.Mesh, wp.array, wp.array, wp.array, int]:
+  """Create a Warp mesh BVH for a single 2D or 3D flex."""
   nflexvert = mjm.nflexvert
-  nflexelemdata = len(mjm.flex_elem)
 
+  flex_dim = wp.array(mjm.flex_dim, dtype=int)
+  flex_elemadr = wp.array(mjm.flex_elemadr, dtype=int)
+  flex_elemnum = wp.array(mjm.flex_elemnum, dtype=int)
   flex_elem = wp.array(mjm.flex_elem, dtype=int)
+  flex_elemdataadr = wp.array(mjm.flex_elemdataadr, dtype=int)
+  flex_vertadr = wp.array(mjm.flex_vertadr, dtype=int)
   flexvert_xpos = wp.array(np.tile(mjd.flexvert_xpos[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3)
 
-  flex_faceadr = [0]
+  dim = int(mjm.flex_dim[flex_id])
-  for f in range(nflex):
+  nelem = int(mjm.flex_elemnum[flex_id])
-    if mjm.flex_dim[f] == 2:
+  nshell = int(mjm.flex_shellnum[flex_id])
-      flex_faceadr.append(flex_faceadr[-1] + 2 * mjm.flex_elemnum[f] + 2 * mjm.flex_shellnum[f])
-    elif mjm.flex_dim[f] == 3:
-      flex_faceadr.append(flex_faceadr[-1] + mjm.flex_shellnum[f])
 
-  nface = int(flex_faceadr[-1])
+  if dim == 2:
-  flex_faceadr = flex_faceadr[:-1]
+    nface = 2 * nelem + 2 * nshell
+  else:
+    nface = nshell
 
   face_point = wp.empty(nface * 3 * nworld, dtype=wp.vec3)
   face_index = wp.empty(nface * 3 * nworld, dtype=wp.int32)
@@ -883,8 +1001,8 @@ def build_flex_bvh(
 
   wp.launch(
     kernel=accumulate_flex_vertex_normals,
-    dim=(nworld, nflexelemdata // 3),
+    dim=(nworld, mjm.nflexelem),
-    inputs=[flex_elem, flexvert_xpos],
+    inputs=[mjm.nflex, flex_dim, flex_vertadr, flex_elemadr, flex_elemnum, flex_elemdataadr, flex_elem, flexvert_xpos],
     outputs=[flexvert_norm],
   )
 
@@ -894,13 +1012,9 @@ def build_flex_bvh(
     inputs=[flexvert_norm],
   )
 
-  for f in range(nflex):
+  elem_adr = mjm.flex_elemdataadr[flex_id]
-    dim = mjm.flex_dim[f]
+  shell_adr = mjm.flex_shelldataadr[flex_id]
-    elem_adr = mjm.flex_elemdataadr[f]
+  vert_adr = mjm.flex_vertadr[flex_id]
-    nelem = mjm.flex_elemnum[f]
-    shell_adr = mjm.flex_shelldataadr[f]
-    nshell = mjm.flex_shellnum[f]
-    vert_adr = mjm.flex_vertadr[f]
 
   if dim == 2:
     wp.launch(
@@ -912,8 +1026,8 @@ def build_flex_bvh(
         flexvert_norm,
         elem_adr,
         vert_adr,
-          flex_faceadr[f],
+        0,  # face_offset
-          mjm.flex_radius[f],
+        mjm.flex_radius[flex_id],
         nface,
       ],
       outputs=[face_point, face_index, group],
@@ -923,13 +1037,13 @@ def build_flex_bvh(
       kernel=_build_flex_2d_sides,
       dim=(nworld, nshell),
       inputs=[
+        flex_shell,
         flexvert_xpos,
         flexvert_norm,
-          flex_shell,
         shell_adr,
         vert_adr,
-          flex_faceadr[f] + 2 * nelem,
+        2 * nelem,  # face_offset
-          mjm.flex_radius[f],
+        mjm.flex_radius[flex_id],
         nface,
       ],
       outputs=[face_point, face_index, group],
@@ -939,11 +1053,11 @@ def build_flex_bvh(
       kernel=_build_flex_3d_shells,
       dim=(nworld, nshell),
       inputs=[
-          flexvert_xpos,
         flex_shell,
+        flexvert_xpos,
         shell_adr,
         vert_adr,
-          flex_faceadr[f],
+        0,  # face_offset
         nface,
       ],
       outputs=[face_point, face_index, group],
@@ -965,24 +1079,23 @@ def build_flex_bvh(
     outputs=[group_root],
   )
 
-  return (
+  return flex_mesh, group_root
-    flex_mesh,
-    face_point,
-    group_root,
-    flex_shell,
-    flex_faceadr,
-    nface,
-  )
 
 
 def refit_flex_bvh(m: Model, d: Data, rc: RenderContext):
-  """Refit the flex BVH."""
+  """Refit per-flex BVHs."""
   flexvert_norm = wp.zeros(d.flexvert_xpos.shape, dtype=wp.vec3)
 
   wp.launch(
     kernel=accumulate_flex_vertex_normals,
-    dim=(d.nworld, m.nflexelemdata // 3),
+    dim=(d.nworld, m.nflexelem),
     inputs=[
+      m.nflex,
+      m.flex_dim,
+      m.flex_vertadr,
+      m.flex_elemadr,
+      m.flex_elemnum,
+      m.flex_elemdataadr,
       m.flex_elem,
       d.flexvert_xpos,
     ],
@@ -991,32 +1104,49 @@ def refit_flex_bvh(m: Model, d: Data, rc: RenderContext):
 
   wp.launch(
     kernel=normalize_vertex_normals,
-    dim=(d.nworld, m.nflexvert),
+    dim=(d.nworld, d.flexvert_xpos.shape[1]),
     inputs=[flexvert_norm],
   )
 
+  for i in range(m.nflex):
+    if rc.flex_dim_np[i] == 1:
+      continue
+    mesh = rc.flex_mesh_registry[i]
+    nface = mesh.points.shape[0] // (3 * d.nworld)
+
+    if rc.flex_dim_np[i] == 2:
       wp.launch(
-    kernel=_update_flex_face_points,
+        kernel=_update_flex_2d_face_points,
-    dim=(d.nworld, rc.flex_nwork),
+        dim=(d.nworld, nface // 2),
         inputs=[
-      m.nflex,
-      m.flex_dim,
           m.flex_vertadr,
           m.flex_elemnum,
+          m.flex_elemdataadr,
+          m.flex_shelldataadr,
           m.flex_elem,
+          m.flex_shell,
+          m.flex_radius,
           d.flexvert_xpos,
-      rc.flex_shell,
           flexvert_norm,
-      rc.flex_elemdataadr,
+          i,
-      rc.flex_shelldataadr,
+          nface,
-      rc.flex_faceadr,
-      rc.flex_radius,
-      rc.flex_workadr,
-      rc.flex_worknum,
-      rc.flex_nface,
           rc.flex_render_smooth,
         ],
-    outputs=[rc.flex_face_point],
+        outputs=[mesh.points],
+      )
+    else:
+      wp.launch(
+        kernel=_update_flex_3d_face_points,
+        dim=(d.nworld, nface),
+        inputs=[
+          m.flex_vertadr,
+          m.flex_shelldataadr,
+          m.flex_shell,
+          d.flexvert_xpos,
+          i,
+          nface,
+        ],
+        outputs=[mesh.points],
       )
 
-  rc.flex_mesh.refit()
+    mesh.refit()

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

@@ -15,34 +15,35 @@
 
 from typing import Tuple
 
+import warp as wp
+
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext
-from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import ccd
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import multicontact
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import support
-from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import contact_params
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
+from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import contact_params
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import geom_collision_pair
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import write_contact
 from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
 from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index
-from mujoco.mjx.third_party.mujoco_warp._src.types import Data
-from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit
-from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
-from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
-from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
 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 Data
+from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit
+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 mat43
+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
-import warp as wp
 
 # TODO(team): improve compile time to enable backward pass
 wp.set_module_options({"enable_backward": False})
@@ -233,6 +234,7 @@ def ccd_hfield_kernel_builder(
     contact_solimp_out: wp.array(dtype=vec5),
     contact_dim_out: wp.array(dtype=int),
     contact_geom_out: wp.array(dtype=wp.vec2i),
+    contact_efc_address_out: wp.array2d(dtype=int),
     contact_worldid_out: wp.array(dtype=int),
     contact_type_out: wp.array(dtype=int),
     contact_geomcollisionid_out: wp.array(dtype=int),
@@ -516,6 +518,7 @@ def ccd_hfield_kernel_builder(
       contact_solimp_out,
       contact_dim_out,
       contact_geom_out,
+      contact_efc_address_out,
       contact_worldid_out,
       contact_type_out,
       contact_geomcollisionid_out,
@@ -572,6 +575,7 @@ def ccd_hfield_kernel_builder(
         contact_solimp_out,
         contact_dim_out,
         contact_geom_out,
+        contact_efc_address_out,
         contact_worldid_out,
         contact_type_out,
         contact_geomcollisionid_out,
@@ -626,6 +630,7 @@ def ccd_hfield_kernel_builder(
         contact_solimp_out,
         contact_dim_out,
         contact_geom_out,
+        contact_efc_address_out,
         contact_worldid_out,
         contact_type_out,
         contact_geomcollisionid_out,
@@ -681,6 +686,7 @@ def ccd_hfield_kernel_builder(
         contact_solimp_out,
         contact_dim_out,
         contact_geom_out,
+        contact_efc_address_out,
         contact_worldid_out,
         contact_type_out,
         contact_geomcollisionid_out,
@@ -751,6 +757,7 @@ def ccd_kernel_builder(
     contact_solimp_out: wp.array(dtype=vec5),
     contact_dim_out: wp.array(dtype=int),
     contact_geom_out: wp.array(dtype=wp.vec2i),
+    contact_efc_address_out: wp.array2d(dtype=int),
     contact_worldid_out: wp.array(dtype=int),
     contact_type_out: wp.array(dtype=int),
     contact_geomcollisionid_out: wp.array(dtype=int),
@@ -871,6 +878,7 @@ def ccd_kernel_builder(
         contact_solimp_out,
         contact_dim_out,
         contact_geom_out,
+        contact_efc_address_out,
         contact_worldid_out,
         contact_type_out,
         contact_geomcollisionid_out,
@@ -956,6 +964,7 @@ def ccd_kernel_builder(
     contact_solimp_out: wp.array(dtype=vec5),
     contact_dim_out: wp.array(dtype=int),
     contact_geom_out: wp.array(dtype=wp.vec2i),
+    contact_efc_address_out: wp.array2d(dtype=int),
     contact_worldid_out: wp.array(dtype=int),
     contact_type_out: wp.array(dtype=int),
     contact_geomcollisionid_out: wp.array(dtype=int),
@@ -1070,6 +1079,7 @@ def ccd_kernel_builder(
       contact_solimp_out,
       contact_dim_out,
       contact_geom_out,
+      contact_efc_address_out,
       contact_worldid_out,
       contact_type_out,
       contact_geomcollisionid_out,
@@ -1156,6 +1166,7 @@ def convex_narrowphase(m: Model, d: Data, ctx: CollisionContext, collision_table
     d.contact.solimp,
     d.contact.dim,
     d.contact.geom,
+    d.contact.efc_address,
     d.contact.worldid,
     d.contact.type,
     d.contact.geomcollisionid,

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

@@ -18,14 +18,15 @@
 import dataclasses
 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_MINMU
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
-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 vec5
-import warp as wp
 
 wp.set_module_options({"enable_backward": False})
 
@@ -183,6 +184,7 @@ def write_contact(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -222,6 +224,8 @@ def write_contact(
     contact_solimp_out[cid] = solimp_in
     contact_type_out[cid] = contact_type
     contact_geomcollisionid_out[cid] = id_
+    for i in range(contact_efc_address_out.shape[1]):
+      contact_efc_address_out[cid, i] = -1
     return int(active)
   return 0
 
@@ -267,9 +271,7 @@ def contact_params(
     condim = pair_dim[pairid]
     friction = pair_friction[worldid % pair_friction.shape[0], pairid]
     solref = pair_solref[worldid % pair_solref.shape[0], pairid]
-    solreffriction = pair_solreffriction[
+    solreffriction = pair_solreffriction[worldid % pair_solreffriction.shape[0], pairid]
-        worldid % pair_solreffriction.shape[0], pairid
-    ]
     solimp = pair_solimp[worldid % pair_solimp.shape[0], pairid]
   else:
     g1 = geoms[0]
@@ -305,9 +307,7 @@ def contact_params(
       mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 >= MJ_MINVAL), 0.0, mix)
       mix = wp.where((solmix1 >= MJ_MINVAL) and (solmix2 < MJ_MINVAL), 1.0, mix)
       condim = wp.max(condim1, condim2)
-      max_geom_friction = wp.max(
+      max_geom_friction = wp.max(geom_friction[friction_id, g1], geom_friction[friction_id, g2])
-          geom_friction[friction_id, g1], geom_friction[friction_id, g2]
-      )
 
     friction = vec5(
       max_geom_friction[0],
@@ -317,22 +317,13 @@ def contact_params(
       max_geom_friction[2],
     )
 
-    if (
+    if geom_solref[solref_id, g1][0] > 0.0 and geom_solref[solref_id, g2][0] > 0.0:
-        geom_solref[solref_id, g1][0] > 0.0
+      solref = mix * geom_solref[solref_id, g1] + (1.0 - mix) * geom_solref[solref_id, g2]
-        and geom_solref[solref_id, g2][0] > 0.0
-    ):
-      solref = (
-          mix * geom_solref[solref_id, g1]
-          + (1.0 - mix) * geom_solref[solref_id, g2]
-      )
     else:
       solref = wp.min(geom_solref[solref_id, g1], geom_solref[solref_id, g2])
 
     solreffriction = wp.vec2(0.0, 0.0)
-    solimp = (
+    solimp = mix * geom_solimp[solimp_id, g1] + (1.0 - mix) * geom_solimp[solimp_id, g2]
-        mix * geom_solimp[solimp_id, g1]
-        + (1.0 - mix) * geom_solimp[solimp_id, g2]
-    )
     # geom priority is ignored
     margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
     gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2]

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

@@ -15,25 +15,27 @@
 
 from typing import Any
 
+import warp as wp
+
 from mujoco.mjx.third_party.mujoco_warp._src.collision_convex import convex_narrowphase
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import create_collision_context
+from mujoco.mjx.third_party.mujoco_warp._src.collision_flex import flex_narrowphase
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import primitive_narrowphase
 from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf_narrowphase
 from mujoco.mjx.third_party.mujoco_warp._src.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 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.types import MJ_MAXVAL
-from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 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
-import warp as wp
 
 wp.set_module_options({"enable_backward": False})
 
@@ -290,25 +292,13 @@ def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound
     # 8: obb
 
     aabb_id = worldid % ngeom_aabb if wp.static(ngeom_aabb > 1) else 0
-    center1, center2 = (
+    center1, center2 = geom_aabb[aabb_id, geom1, 0], geom_aabb[aabb_id, geom2, 0]  # kernel_analyzer: ignore
-        geom_aabb[aabb_id, geom1, 0],
+    size1, size2 = geom_aabb[aabb_id, geom1, 1], geom_aabb[aabb_id, geom2, 1]  # kernel_analyzer: ignore
-        geom_aabb[aabb_id, geom2, 0],
-    )  # kernel_analyzer: ignore
-    size1, size2 = (
-        geom_aabb[aabb_id, geom1, 1],
-        geom_aabb[aabb_id, geom2, 1],
-    )  # kernel_analyzer: ignore
 
     rbound_id = worldid % ngeom_rbound if wp.static(ngeom_rbound > 1) else 0
-    rbound1, rbound2 = (
+    rbound1, rbound2 = geom_rbound[rbound_id, geom1], geom_rbound[rbound_id, geom2]  # kernel_analyzer: ignore
-        geom_rbound[rbound_id, geom1],
-        geom_rbound[rbound_id, geom2],
-    )  # kernel_analyzer: ignore
     margin_id = worldid % ngeom_margin if wp.static(ngeom_margin > 1) else 0
-    margin1, margin2 = (
+    margin1, margin2 = geom_margin[margin_id, geom1], geom_margin[margin_id, geom2]  # kernel_analyzer: ignore
-        geom_margin[margin_id, geom1],
-        geom_margin[margin_id, geom2],
-    )  # kernel_analyzer: ignore
     xpos1, xpos2 = geom_xpos_in[worldid, geom1], geom_xpos_in[worldid, geom2]
     xmat1, xmat2 = geom_xmat_in[worldid, geom1], geom_xmat_in[worldid, geom2]
 
@@ -757,6 +747,9 @@ def _narrowphase(m: Model, d: Data, ctx: CollisionContext):
   if m.has_sdf_geom:
     sdf_narrowphase(m, d, ctx)
 
+  if m.nflex > 0:
+    flex_narrowphase(m, d)
+
 
 @event_scope
 def collision(m: Model, d: Data):

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

@@ -0,0 +1,834 @@
+# 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.
+"""Flex collision detection (geom vs flex triangles)."""
+
+import warp as wp
+
+from mujoco.mjx.third_party.mujoco_warp._src import collision_primitive_core
+from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
+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 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 vec5
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
+
+wp.set_module_options({"enable_backward": False})
+
+
+@wp.func
+def _write_flex_contact(
+  # Data in:
+  naconmax_in: int,
+  # In:
+  dist: float,
+  pos: wp.vec3,
+  frame: wp.mat33,
+  margin: float,
+  condim: int,
+  friction: vec5,
+  solref: wp.vec2,
+  solimp: vec5,
+  geom: int,
+  flexid: int,
+  vertid: int,
+  worldid: int,
+  # Data out:
+  contact_dist_out: wp.array(dtype=float),
+  contact_pos_out: wp.array(dtype=wp.vec3),
+  contact_frame_out: wp.array(dtype=wp.mat33),
+  contact_includemargin_out: wp.array(dtype=float),
+  contact_friction_out: wp.array(dtype=vec5),
+  contact_solref_out: wp.array(dtype=wp.vec2),
+  contact_solreffriction_out: wp.array(dtype=wp.vec2),
+  contact_solimp_out: wp.array(dtype=vec5),
+  contact_dim_out: wp.array(dtype=int),
+  contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_flex_out: wp.array(dtype=wp.vec2i),
+  contact_vert_out: wp.array(dtype=wp.vec2i),
+  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),
+):
+  if dist >= margin or dist >= MJ_MAXVAL:
+    return
+
+  id_ = wp.atomic_add(nacon_out, 0, 1)
+  if id_ >= naconmax_in:
+    return
+
+  contact_dist_out[id_] = dist
+  contact_pos_out[id_] = pos
+  contact_frame_out[id_] = frame
+  contact_includemargin_out[id_] = margin
+  contact_friction_out[id_] = friction
+  contact_solref_out[id_] = solref
+  contact_solreffriction_out[id_] = wp.vec2(0.0, 0.0)
+  contact_solimp_out[id_] = solimp
+  contact_dim_out[id_] = condim
+  contact_geom_out[id_] = wp.vec2i(geom, -1)
+  contact_flex_out[id_] = wp.vec2i(-1, flexid)
+  contact_vert_out[id_] = wp.vec2i(-1, vertid)
+  contact_worldid_out[id_] = worldid
+  contact_type_out[id_] = 1
+  contact_geomcollisionid_out[id_] = 0
+
+
+@wp.func
+def _collide_geom_triangle(
+  # Data in:
+  naconmax_in: int,
+  # In:
+  gtype: int,
+  pos: wp.vec3,
+  rot: wp.mat33,
+  size_val: wp.vec3,
+  t1: wp.vec3,
+  t2: wp.vec3,
+  t3: wp.vec3,
+  tri_radius: float,
+  margin: float,
+  condim: int,
+  friction: vec5,
+  solref: wp.vec2,
+  solimp: vec5,
+  geomid: int,
+  flexid: int,
+  vertex_id: int,
+  worldid: int,
+  # Data out:
+  contact_dist_out: wp.array(dtype=float),
+  contact_pos_out: wp.array(dtype=wp.vec3),
+  contact_frame_out: wp.array(dtype=wp.mat33),
+  contact_includemargin_out: wp.array(dtype=float),
+  contact_friction_out: wp.array(dtype=vec5),
+  contact_solref_out: wp.array(dtype=wp.vec2),
+  contact_solreffriction_out: wp.array(dtype=wp.vec2),
+  contact_solimp_out: wp.array(dtype=vec5),
+  contact_dim_out: wp.array(dtype=int),
+  contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_flex_out: wp.array(dtype=wp.vec2i),
+  contact_vert_out: wp.array(dtype=wp.vec2i),
+  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),
+):
+  if gtype == int(GeomType.SPHERE):
+    sphere_radius = size_val[0]
+    dist, contact_pos, nrm = collision_primitive_core.sphere_triangle(pos, sphere_radius, t1, t2, t3, tri_radius)
+    if dist < margin:
+      _write_flex_contact(
+        naconmax_in,
+        dist,
+        contact_pos,
+        make_frame(nrm),
+        margin,
+        condim,
+        friction,
+        solref,
+        solimp,
+        geomid,
+        flexid,
+        vertex_id,
+        worldid,
+        contact_dist_out,
+        contact_pos_out,
+        contact_frame_out,
+        contact_includemargin_out,
+        contact_friction_out,
+        contact_solref_out,
+        contact_solreffriction_out,
+        contact_solimp_out,
+        contact_dim_out,
+        contact_geom_out,
+        contact_flex_out,
+        contact_vert_out,
+        contact_worldid_out,
+        contact_type_out,
+        contact_geomcollisionid_out,
+        nacon_out,
+      )
+    return
+
+  # Capsule, box, cylinder all return up to 2 contacts - compute then share writing code
+  dists = wp.vec2(collision_primitive_core.MJ_MAXVAL, collision_primitive_core.MJ_MAXVAL)
+  poss = collision_primitive_core.mat23f(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
+  nrms = collision_primitive_core.mat23f(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
+
+  if gtype == int(GeomType.CAPSULE):
+    cap_radius = size_val[0]
+    cap_half_len = size_val[1]
+    cap_axis = wp.vec3(rot[0, 2], rot[1, 2], rot[2, 2])
+    dists, poss, nrms = collision_primitive_core.capsule_triangle(
+      pos, cap_axis, cap_radius, cap_half_len, t1, t2, t3, tri_radius
+    )
+  elif gtype == int(GeomType.BOX):
+    dists, poss, nrms = collision_primitive_core.box_triangle(pos, rot, size_val, t1, t2, t3, tri_radius)
+  elif gtype == int(GeomType.CYLINDER):
+    cyl_radius = size_val[0]
+    cyl_half_height = size_val[1]
+    cyl_axis = wp.vec3(rot[0, 2], rot[1, 2], rot[2, 2])
+    dists, poss, nrms = collision_primitive_core.cylinder_triangle(
+      pos, cyl_axis, cyl_radius, cyl_half_height, t1, t2, t3, tri_radius
+    )
+
+  # Write up to 2 contacts (shared code for capsule/box/cylinder)
+  if dists[0] < margin:
+    p1 = wp.vec3(poss[0, 0], poss[0, 1], poss[0, 2])
+    n1 = wp.vec3(nrms[0, 0], nrms[0, 1], nrms[0, 2])
+    _write_flex_contact(
+      naconmax_in,
+      dists[0],
+      p1,
+      make_frame(n1),
+      margin,
+      condim,
+      friction,
+      solref,
+      solimp,
+      geomid,
+      flexid,
+      vertex_id,
+      worldid,
+      contact_dist_out,
+      contact_pos_out,
+      contact_frame_out,
+      contact_includemargin_out,
+      contact_friction_out,
+      contact_solref_out,
+      contact_solreffriction_out,
+      contact_solimp_out,
+      contact_dim_out,
+      contact_geom_out,
+      contact_flex_out,
+      contact_vert_out,
+      contact_worldid_out,
+      contact_type_out,
+      contact_geomcollisionid_out,
+      nacon_out,
+    )
+  if dists[1] < margin:
+    p2 = wp.vec3(poss[1, 0], poss[1, 1], poss[1, 2])
+    n2 = wp.vec3(nrms[1, 0], nrms[1, 1], nrms[1, 2])
+    _write_flex_contact(
+      naconmax_in,
+      dists[1],
+      p2,
+      make_frame(n2),
+      margin,
+      condim,
+      friction,
+      solref,
+      solimp,
+      geomid,
+      flexid,
+      vertex_id,
+      worldid,
+      contact_dist_out,
+      contact_pos_out,
+      contact_frame_out,
+      contact_includemargin_out,
+      contact_friction_out,
+      contact_solref_out,
+      contact_solreffriction_out,
+      contact_solimp_out,
+      contact_dim_out,
+      contact_geom_out,
+      contact_flex_out,
+      contact_vert_out,
+      contact_worldid_out,
+      contact_type_out,
+      contact_geomcollisionid_out,
+      nacon_out,
+    )
+
+
+@wp.kernel
+def _flex_plane_narrowphase(
+  # Model:
+  ngeom: int,
+  nflexvert: int,
+  geom_type: wp.array(dtype=int),
+  geom_condim: wp.array(dtype=int),
+  geom_solref: wp.array2d(dtype=wp.vec2),
+  geom_solimp: wp.array2d(dtype=vec5),
+  geom_friction: wp.array2d(dtype=wp.vec3),
+  geom_margin: wp.array2d(dtype=float),
+  flex_condim: wp.array(dtype=int),
+  flex_friction: wp.array(dtype=wp.vec3),
+  flex_margin: wp.array(dtype=float),
+  flex_vertadr: wp.array(dtype=int),
+  flex_radius: wp.array(dtype=float),
+  flex_vertflexid: wp.array(dtype=int),
+  # Data in:
+  geom_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  nworld_in: int,
+  naconmax_in: int,
+  # Data out:
+  contact_dist_out: wp.array(dtype=float),
+  contact_pos_out: wp.array(dtype=wp.vec3),
+  contact_frame_out: wp.array(dtype=wp.mat33),
+  contact_includemargin_out: wp.array(dtype=float),
+  contact_friction_out: wp.array(dtype=vec5),
+  contact_solref_out: wp.array(dtype=wp.vec2),
+  contact_solreffriction_out: wp.array(dtype=wp.vec2),
+  contact_solimp_out: wp.array(dtype=vec5),
+  contact_dim_out: wp.array(dtype=int),
+  contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_flex_out: wp.array(dtype=wp.vec2i),
+  contact_vert_out: wp.array(dtype=wp.vec2i),
+  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),
+):
+  worldid, vertid = wp.tid()
+
+  flexid = flex_vertflexid[vertid]
+  radius = flex_radius[flexid]
+  flex_margin_val = flex_margin[flexid]
+  flex_condim_val = flex_condim[flexid]
+  flex_fric = flex_friction[flexid]
+  # Convert global vertid to local vertex index within this flex
+  local_vertid = vertid - flex_vertadr[flexid]
+
+  vert = flexvert_xpos_in[worldid, vertid]
+
+  # TODO: Add a broadphase
+  for geomid in range(ngeom):
+    gtype = geom_type[geomid]
+    if gtype != int(GeomType.PLANE):
+      continue
+
+    plane_pos = geom_xpos_in[worldid, geomid]
+    plane_rot = geom_xmat_in[worldid, geomid]
+    plane_normal = wp.vec3(plane_rot[0, 2], plane_rot[1, 2], plane_rot[2, 2])
+
+    margin = geom_margin[worldid % geom_margin.shape[0], geomid] + flex_margin_val
+
+    diff = vert - plane_pos
+    signed_dist = wp.dot(diff, plane_normal)
+    dist = signed_dist - radius
+
+    if dist < margin:
+      geom_condim_val = geom_condim[geomid]
+      condim = wp.max(geom_condim_val, flex_condim_val)
+      solref = geom_solref[worldid % geom_solref.shape[0], geomid]
+      solimp = geom_solimp[worldid % geom_solimp.shape[0], geomid]
+      geom_fric = geom_friction[worldid % geom_friction.shape[0], geomid]
+      fric0 = wp.max(geom_fric[0], flex_fric[0])
+      fric1 = wp.max(geom_fric[1], flex_fric[1])
+      fric2 = wp.max(geom_fric[2], flex_fric[2])
+      friction = vec5(
+        wp.max(MJ_MINMU, fric0),
+        wp.max(MJ_MINMU, fric0),
+        wp.max(MJ_MINMU, fric1),
+        wp.max(MJ_MINMU, fric2),
+        wp.max(MJ_MINMU, fric2),
+      )
+
+      contact_pos = vert - plane_normal * (dist * 0.5 + radius)
+      _write_flex_contact(
+        naconmax_in,
+        dist,
+        contact_pos,
+        make_frame(plane_normal),
+        margin,
+        condim,
+        friction,
+        solref,
+        solimp,
+        geomid,
+        flexid,
+        local_vertid,
+        worldid,
+        contact_dist_out,
+        contact_pos_out,
+        contact_frame_out,
+        contact_includemargin_out,
+        contact_friction_out,
+        contact_solref_out,
+        contact_solreffriction_out,
+        contact_solimp_out,
+        contact_dim_out,
+        contact_geom_out,
+        contact_flex_out,
+        contact_vert_out,
+        contact_worldid_out,
+        contact_type_out,
+        contact_geomcollisionid_out,
+        nacon_out,
+      )
+
+
+@wp.kernel
+def _flex_narrowphase_dim2(
+  # Model:
+  ngeom: int,
+  nflex: int,
+  geom_type: wp.array(dtype=int),
+  geom_contype: wp.array(dtype=int),
+  geom_conaffinity: wp.array(dtype=int),
+  geom_condim: wp.array(dtype=int),
+  geom_solref: wp.array2d(dtype=wp.vec2),
+  geom_solimp: wp.array2d(dtype=vec5),
+  geom_size: wp.array2d(dtype=wp.vec3),
+  geom_friction: wp.array2d(dtype=wp.vec3),
+  geom_margin: wp.array2d(dtype=float),
+  flex_contype: wp.array(dtype=int),
+  flex_conaffinity: wp.array(dtype=int),
+  flex_margin: wp.array(dtype=float),
+  flex_dim: wp.array(dtype=int),
+  flex_vertadr: wp.array(dtype=int),
+  flex_elemadr: wp.array(dtype=int),
+  flex_elemnum: wp.array(dtype=int),
+  flex_elemdataadr: wp.array(dtype=int),
+  flex_elem: wp.array(dtype=int),
+  flex_radius: wp.array(dtype=float),
+  # Data in:
+  geom_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  nworld_in: int,
+  naconmax_in: int,
+  # Data out:
+  contact_dist_out: wp.array(dtype=float),
+  contact_pos_out: wp.array(dtype=wp.vec3),
+  contact_frame_out: wp.array(dtype=wp.mat33),
+  contact_includemargin_out: wp.array(dtype=float),
+  contact_friction_out: wp.array(dtype=vec5),
+  contact_solref_out: wp.array(dtype=wp.vec2),
+  contact_solreffriction_out: wp.array(dtype=wp.vec2),
+  contact_solimp_out: wp.array(dtype=vec5),
+  contact_dim_out: wp.array(dtype=int),
+  contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_flex_out: wp.array(dtype=wp.vec2i),
+  contact_vert_out: wp.array(dtype=wp.vec2i),
+  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),
+):
+  worldid, elemid = wp.tid()
+
+  flexid = int(-1)
+  for i in range(nflex):
+    if flex_dim[i] != 2:
+      continue
+    elem_adr = flex_elemadr[i]
+    elem_num = flex_elemnum[i]
+    if elemid >= elem_adr and elemid < elem_adr + elem_num:
+      flexid = i
+      break
+
+  if flexid < 0:
+    return
+
+  vert_adr = flex_vertadr[flexid]
+  tri_radius = flex_radius[flexid]
+  tri_margin = flex_margin[flexid]
+
+  elem_data_idx = flex_elemdataadr[flexid] + (elemid - flex_elemadr[flexid]) * 3
+  v0_local = flex_elem[elem_data_idx]
+  v1_local = flex_elem[elem_data_idx + 1]
+  v2_local = flex_elem[elem_data_idx + 2]
+
+  t1 = flexvert_xpos_in[worldid, vert_adr + v0_local]
+  t2 = flexvert_xpos_in[worldid, vert_adr + v1_local]
+  t3 = flexvert_xpos_in[worldid, vert_adr + v2_local]
+
+  # TODO: Add a broadphase
+  for geomid in range(ngeom):
+    gtype = geom_type[geomid]
+    if (
+      gtype != int(GeomType.SPHERE)
+      and gtype != int(GeomType.CAPSULE)
+      and gtype != int(GeomType.BOX)
+      and gtype != int(GeomType.CYLINDER)
+    ):
+      continue
+
+    g_contype = geom_contype[geomid]
+    g_conaffinity = geom_conaffinity[geomid]
+    f_contype = flex_contype[flexid]
+    f_conaffinity = flex_conaffinity[flexid]
+    if not ((g_contype & f_conaffinity) or (f_contype & g_conaffinity)):
+      continue
+
+    geom_margin_val = geom_margin[worldid % geom_margin.shape[0], geomid]
+    margin = geom_margin_val + tri_margin
+
+    geom_pos = geom_xpos_in[worldid, geomid]
+    geom_rot = geom_xmat_in[worldid, geomid]
+    geom_size_val = geom_size[worldid % geom_size.shape[0], geomid]
+
+    condim = geom_condim[geomid]
+    gf = geom_friction[worldid % geom_friction.shape[0], geomid]
+    friction = vec5(
+      wp.max(MJ_MINMU, gf[0]),
+      wp.max(MJ_MINMU, gf[0]),
+      wp.max(MJ_MINMU, gf[1]),
+      wp.max(MJ_MINMU, gf[2]),
+      wp.max(MJ_MINMU, gf[2]),
+    )
+    solref = geom_solref[worldid % geom_solref.shape[0], geomid]
+    solimp = geom_solimp[worldid % geom_solimp.shape[0], geomid]
+
+    _collide_geom_triangle(
+      naconmax_in,
+      gtype,
+      geom_pos,
+      geom_rot,
+      geom_size_val,
+      t1,
+      t2,
+      t3,
+      tri_radius,
+      margin,
+      condim,
+      friction,
+      solref,
+      solimp,
+      geomid,
+      flexid,
+      v0_local,
+      worldid,
+      contact_dist_out,
+      contact_pos_out,
+      contact_frame_out,
+      contact_includemargin_out,
+      contact_friction_out,
+      contact_solref_out,
+      contact_solreffriction_out,
+      contact_solimp_out,
+      contact_dim_out,
+      contact_geom_out,
+      contact_flex_out,
+      contact_vert_out,
+      contact_worldid_out,
+      contact_type_out,
+      contact_geomcollisionid_out,
+      nacon_out,
+    )
+
+
+@wp.kernel
+def _flex_narrowphase_dim3(
+  # Model:
+  ngeom: int,
+  nflex: int,
+  geom_type: wp.array(dtype=int),
+  geom_contype: wp.array(dtype=int),
+  geom_conaffinity: wp.array(dtype=int),
+  geom_condim: wp.array(dtype=int),
+  geom_solref: wp.array2d(dtype=wp.vec2),
+  geom_solimp: wp.array2d(dtype=vec5),
+  geom_size: wp.array2d(dtype=wp.vec3),
+  geom_friction: wp.array2d(dtype=wp.vec3),
+  geom_margin: wp.array2d(dtype=float),
+  flex_contype: wp.array(dtype=int),
+  flex_conaffinity: wp.array(dtype=int),
+  flex_margin: wp.array(dtype=float),
+  flex_dim: wp.array(dtype=int),
+  flex_vertadr: wp.array(dtype=int),
+  flex_shellnum: wp.array(dtype=int),
+  flex_shelldataadr: wp.array(dtype=int),
+  flex_shell: wp.array(dtype=int),
+  flex_radius: wp.array(dtype=float),
+  # Data in:
+  geom_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  nworld_in: int,
+  naconmax_in: int,
+  # Data out:
+  contact_dist_out: wp.array(dtype=float),
+  contact_pos_out: wp.array(dtype=wp.vec3),
+  contact_frame_out: wp.array(dtype=wp.mat33),
+  contact_includemargin_out: wp.array(dtype=float),
+  contact_friction_out: wp.array(dtype=vec5),
+  contact_solref_out: wp.array(dtype=wp.vec2),
+  contact_solreffriction_out: wp.array(dtype=wp.vec2),
+  contact_solimp_out: wp.array(dtype=vec5),
+  contact_dim_out: wp.array(dtype=int),
+  contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_flex_out: wp.array(dtype=wp.vec2i),
+  contact_vert_out: wp.array(dtype=wp.vec2i),
+  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),
+):
+  worldid, shellid = wp.tid()
+
+  flexid = int(-1)
+  shell_offset = int(0)
+  for i in range(nflex):
+    if flex_dim[i] != 3:
+      continue
+    shell_num = flex_shellnum[i]
+    if shellid >= shell_offset and shellid < shell_offset + shell_num:
+      flexid = i
+      break
+    shell_offset += shell_num
+
+  if flexid < 0:
+    return
+
+  vert_adr = flex_vertadr[flexid]
+  tri_radius = flex_radius[flexid]
+  tri_margin = flex_margin[flexid]
+
+  shell_adr = flex_shelldataadr[flexid]
+  local_shellid = shellid - shell_offset
+  shell_data_idx = shell_adr + local_shellid * 3
+
+  v0_local = flex_shell[shell_data_idx]
+  v1_local = flex_shell[shell_data_idx + 1]
+  v2_local = flex_shell[shell_data_idx + 2]
+
+  t1 = flexvert_xpos_in[worldid, vert_adr + v0_local]
+  t2 = flexvert_xpos_in[worldid, vert_adr + v1_local]
+  t3 = flexvert_xpos_in[worldid, vert_adr + v2_local]
+
+  # TODO: Add a broadphase
+  for geomid in range(ngeom):
+    gtype = geom_type[geomid]
+    if (
+      gtype != int(GeomType.SPHERE)
+      and gtype != int(GeomType.CAPSULE)
+      and gtype != int(GeomType.BOX)
+      and gtype != int(GeomType.CYLINDER)
+    ):
+      continue
+
+    g_contype = geom_contype[geomid]
+    g_conaffinity = geom_conaffinity[geomid]
+    f_contype = flex_contype[flexid]
+    f_conaffinity = flex_conaffinity[flexid]
+    if not ((g_contype & f_conaffinity) or (f_contype & g_conaffinity)):
+      continue
+
+    geom_margin_val = geom_margin[worldid % geom_margin.shape[0], geomid]
+    margin = geom_margin_val + tri_margin
+
+    geom_pos = geom_xpos_in[worldid, geomid]
+    geom_rot = geom_xmat_in[worldid, geomid]
+    geom_size_val = geom_size[worldid % geom_size.shape[0], geomid]
+
+    condim = geom_condim[geomid]
+    gf = geom_friction[worldid % geom_friction.shape[0], geomid]
+    friction = vec5(
+      wp.max(MJ_MINMU, gf[0]),
+      wp.max(MJ_MINMU, gf[0]),
+      wp.max(MJ_MINMU, gf[1]),
+      wp.max(MJ_MINMU, gf[2]),
+      wp.max(MJ_MINMU, gf[2]),
+    )
+    solref = geom_solref[worldid % geom_solref.shape[0], geomid]
+    solimp = geom_solimp[worldid % geom_solimp.shape[0], geomid]
+
+    _collide_geom_triangle(
+      naconmax_in,
+      gtype,
+      geom_pos,
+      geom_rot,
+      geom_size_val,
+      t1,
+      t2,
+      t3,
+      tri_radius,
+      margin,
+      condim,
+      friction,
+      solref,
+      solimp,
+      geomid,
+      flexid,
+      v0_local,
+      worldid,
+      contact_dist_out,
+      contact_pos_out,
+      contact_frame_out,
+      contact_includemargin_out,
+      contact_friction_out,
+      contact_solref_out,
+      contact_solreffriction_out,
+      contact_solimp_out,
+      contact_dim_out,
+      contact_geom_out,
+      contact_flex_out,
+      contact_vert_out,
+      contact_worldid_out,
+      contact_type_out,
+      contact_geomcollisionid_out,
+      nacon_out,
+    )
+
+
+@event_scope
+def flex_narrowphase(m: Model, d: Data):
+  """Runs collision detection between geoms and flex elements."""
+  if m.nflex == 0:
+    return
+
+  wp.launch(
+    _flex_narrowphase_dim2,
+    dim=(d.nworld, m.nflexelem),
+    inputs=[
+      m.ngeom,
+      m.nflex,
+      m.geom_type,
+      m.geom_contype,
+      m.geom_conaffinity,
+      m.geom_condim,
+      m.geom_solref,
+      m.geom_solimp,
+      m.geom_size,
+      m.geom_friction,
+      m.geom_margin,
+      m.flex_contype,
+      m.flex_conaffinity,
+      m.flex_margin,
+      m.flex_dim,
+      m.flex_vertadr,
+      m.flex_elemadr,
+      m.flex_elemnum,
+      m.flex_elemdataadr,
+      m.flex_elem,
+      m.flex_radius,
+      d.geom_xpos,
+      d.geom_xmat,
+      d.flexvert_xpos,
+      d.nworld,
+      d.naconmax,
+    ],
+    outputs=[
+      d.contact.dist,
+      d.contact.pos,
+      d.contact.frame,
+      d.contact.includemargin,
+      d.contact.friction,
+      d.contact.solref,
+      d.contact.solreffriction,
+      d.contact.solimp,
+      d.contact.dim,
+      d.contact.geom,
+      d.contact.flex,
+      d.contact.vert,
+      d.contact.worldid,
+      d.contact.type,
+      d.contact.geomcollisionid,
+      d.nacon,
+    ],
+  )
+
+  wp.launch(
+    _flex_narrowphase_dim3,
+    dim=(d.nworld, m.nflexshelldata // 3),
+    inputs=[
+      m.ngeom,
+      m.nflex,
+      m.geom_type,
+      m.geom_contype,
+      m.geom_conaffinity,
+      m.geom_condim,
+      m.geom_solref,
+      m.geom_solimp,
+      m.geom_size,
+      m.geom_friction,
+      m.geom_margin,
+      m.flex_contype,
+      m.flex_conaffinity,
+      m.flex_margin,
+      m.flex_dim,
+      m.flex_vertadr,
+      m.flex_shellnum,
+      m.flex_shelldataadr,
+      m.flex_shell,
+      m.flex_radius,
+      d.geom_xpos,
+      d.geom_xmat,
+      d.flexvert_xpos,
+      d.nworld,
+      d.naconmax,
+    ],
+    outputs=[
+      d.contact.dist,
+      d.contact.pos,
+      d.contact.frame,
+      d.contact.includemargin,
+      d.contact.friction,
+      d.contact.solref,
+      d.contact.solreffriction,
+      d.contact.solimp,
+      d.contact.dim,
+      d.contact.geom,
+      d.contact.flex,
+      d.contact.vert,
+      d.contact.worldid,
+      d.contact.type,
+      d.contact.geomcollisionid,
+      d.nacon,
+    ],
+  )
+
+  wp.launch(
+    _flex_plane_narrowphase,
+    dim=(d.nworld, m.nflexvert),
+    inputs=[
+      m.ngeom,
+      m.nflexvert,
+      m.geom_type,
+      m.geom_condim,
+      m.geom_solref,
+      m.geom_solimp,
+      m.geom_friction,
+      m.geom_margin,
+      m.flex_condim,
+      m.flex_friction,
+      m.flex_margin,
+      m.flex_vertadr,
+      m.flex_radius,
+      m.flex_vertflexid,
+      d.geom_xpos,
+      d.geom_xmat,
+      d.flexvert_xpos,
+      d.nworld,
+      d.naconmax,
+    ],
+    outputs=[
+      d.contact.dist,
+      d.contact.pos,
+      d.contact.frame,
+      d.contact.includemargin,
+      d.contact.friction,
+      d.contact.solref,
+      d.contact.solreffriction,
+      d.contact.solimp,
+      d.contact.dim,
+      d.contact.geom,
+      d.contact.flex,
+      d.contact.vert,
+      d.contact.worldid,
+      d.contact.type,
+      d.contact.geomcollisionid,
+      d.nacon,
+    ],
+  )

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

@@ -16,11 +16,12 @@
 import math
 from typing import Tuple
 
+import warp as wp
+
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
-import warp as wp
 
 # TODO(team): improve compile time to enable backward pass
 wp.set_module_options({"enable_backward": False})
@@ -581,16 +582,19 @@ def gjk(
   simplex_index2 = wp.vec4i()
   n = int(0)
   coordinates = wp.vec4()  # barycentric coordinates
-  epsilon = wp.where(is_discrete, 0.0, 0.5 * tolerance * tolerance)
+  tol2 = tolerance * tolerance
+  epsilon = wp.where(is_discrete, 0.0, 0.5 * tol2)
 
   # set initial guess
   x_k = x1_0 - x2_0
+  xnorm_old = FLOAT_MAX
 
-  for k in range(gjk_iterations):
+  for _ in range(gjk_iterations):
     xnorm = wp.dot(x_k, x_k)
     # TODO(kbayes): determine new constant here
-    if xnorm < 1e-12:
+    if xnorm < tol2 or wp.abs(xnorm_old - xnorm) < tol2:
       break
+    xnorm_old = xnorm
     dir_neg = x_k / wp.sqrt(xnorm)
 
     # compute kth support point in geom1
@@ -663,13 +667,6 @@ def gjk(
     if n == 4:
       break
 
-    if k == gjk_iterations - 1:
-      wp.printf(
-          "Warning: opt.ccd_iterations, currently set to %d, needs to be"
-          " increased.\n",
-          gjk_iterations,
-      )
-
   result = GJKResult()
 
   # compute the approximate witness points
@@ -1205,7 +1202,6 @@ def _is_invalid_face(face: int) -> bool:
 def _epa(
   # In:
   tolerance: float,
-  gjk_iterations: int,
   epa_iterations: int,
   pt: Polytope,
   geom1: Geom,
@@ -1226,7 +1222,7 @@ def _epa(
   # so iterations must be cap to limit the number of generated vertices
   # (one new vertex per iteration)
   epa_iterations = wp.min(epa_iterations, 1000)
-  for k in range(epa_iterations):
+  for _ in range(epa_iterations):
     pidx = idx
     idx = int(-1)
     lower2 = float(FLOAT_MAX)
@@ -1325,13 +1321,6 @@ def _epa(
     # clear horizon
     pt.nhorizon = 0
 
-    if k == epa_iterations - 1:
-      wp.printf(
-          "Warning: opt.ccd_iterations, currently set to %d, needs to be"
-          " increased.\n",
-          gjk_iterations,
-      )
-
   # return from valid face
   if idx > -1:
     x1, x2, dist = _epa_witness(pt, geom1, geom2, geomtype1, geomtype2, idx)
@@ -2347,7 +2336,7 @@ def ccd(
   if pt.status:
     return result.dist, 1, result.x1, result.x2, -1
 
-  dist, x1, x2, idx = _epa(tolerance, gjk_iterations, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
+  dist, x1, x2, idx = _epa(tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
   if idx == -1:
     return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1
 

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

@@ -15,9 +15,11 @@
 
 from typing import Tuple
 
+import warp as wp
+
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext
-from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import box_box
@@ -34,15 +36,14 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sph
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sphere_sphere
 from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
 from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_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 mat43
-from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
 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
-import warp as wp
 
 wp.set_module_options({"enable_backward": False})
 
@@ -304,6 +305,7 @@ def plane_sphere_wrapper(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -339,6 +341,7 @@ def plane_sphere_wrapper(
     contact_solimp_out,
     contact_dim_out,
     contact_geom_out,
+    contact_efc_address_out,
     contact_worldid_out,
     contact_type_out,
     contact_geomcollisionid_out,
@@ -374,6 +377,7 @@ def sphere_sphere_wrapper(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -408,6 +412,7 @@ def sphere_sphere_wrapper(
     contact_solimp_out,
     contact_dim_out,
     contact_geom_out,
+    contact_efc_address_out,
     contact_worldid_out,
     contact_type_out,
     contact_geomcollisionid_out,
@@ -443,6 +448,7 @@ def sphere_capsule_wrapper(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -480,6 +486,7 @@ def sphere_capsule_wrapper(
     contact_solimp_out,
     contact_dim_out,
     contact_geom_out,
+    contact_efc_address_out,
     contact_worldid_out,
     contact_type_out,
     contact_geomcollisionid_out,
@@ -515,6 +522,7 @@ def capsule_capsule_wrapper(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -564,6 +572,7 @@ def capsule_capsule_wrapper(
       contact_solimp_out,
       contact_dim_out,
       contact_geom_out,
+      contact_efc_address_out,
       contact_worldid_out,
       contact_type_out,
       contact_geomcollisionid_out,
@@ -599,6 +608,7 @@ def plane_capsule_wrapper(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -644,6 +654,7 @@ def plane_capsule_wrapper(
       contact_solimp_out,
       contact_dim_out,
       contact_geom_out,
+      contact_efc_address_out,
       contact_worldid_out,
       contact_type_out,
       contact_geomcollisionid_out,
@@ -679,6 +690,7 @@ def plane_ellipsoid_wrapper(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -713,6 +725,7 @@ def plane_ellipsoid_wrapper(
     contact_solimp_out,
     contact_dim_out,
     contact_geom_out,
+    contact_efc_address_out,
     contact_worldid_out,
     contact_type_out,
     contact_geomcollisionid_out,
@@ -748,6 +761,7 @@ def plane_box_wrapper(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -784,6 +798,7 @@ def plane_box_wrapper(
       contact_solimp_out,
       contact_dim_out,
       contact_geom_out,
+      contact_efc_address_out,
       contact_worldid_out,
       contact_type_out,
       contact_geomcollisionid_out,
@@ -819,6 +834,7 @@ def plane_convex_wrapper(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -855,6 +871,7 @@ def plane_convex_wrapper(
       contact_solimp_out,
       contact_dim_out,
       contact_geom_out,
+      contact_efc_address_out,
       contact_worldid_out,
       contact_type_out,
       contact_geomcollisionid_out,
@@ -890,6 +907,7 @@ def sphere_cylinder_wrapper(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -934,6 +952,7 @@ def sphere_cylinder_wrapper(
     contact_solimp_out,
     contact_dim_out,
     contact_geom_out,
+    contact_efc_address_out,
     contact_worldid_out,
     contact_type_out,
     contact_geomcollisionid_out,
@@ -969,6 +988,7 @@ def plane_cylinder_wrapper(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -1015,6 +1035,7 @@ def plane_cylinder_wrapper(
       contact_solimp_out,
       contact_dim_out,
       contact_geom_out,
+      contact_efc_address_out,
       contact_worldid_out,
       contact_type_out,
       contact_geomcollisionid_out,
@@ -1050,6 +1071,7 @@ def sphere_box_wrapper(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -1083,6 +1105,7 @@ def sphere_box_wrapper(
     contact_solimp_out,
     contact_dim_out,
     contact_geom_out,
+    contact_efc_address_out,
     contact_worldid_out,
     contact_type_out,
     contact_geomcollisionid_out,
@@ -1118,6 +1141,7 @@ def capsule_box_wrapper(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -1166,6 +1190,7 @@ def capsule_box_wrapper(
       contact_solimp_out,
       contact_dim_out,
       contact_geom_out,
+      contact_efc_address_out,
       contact_worldid_out,
       contact_type_out,
       contact_geomcollisionid_out,
@@ -1201,6 +1226,7 @@ def box_box_wrapper(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -1245,6 +1271,7 @@ def box_box_wrapper(
       contact_solimp_out,
       contact_dim_out,
       contact_geom_out,
+      contact_efc_address_out,
       contact_worldid_out,
       contact_type_out,
       contact_geomcollisionid_out,
@@ -1327,6 +1354,7 @@ def _primitive_narrowphase(primitive_collisions_types, primitive_collisions_func
     contact_solimp_out: wp.array(dtype=vec5),
     contact_dim_out: wp.array(dtype=int),
     contact_geom_out: wp.array(dtype=wp.vec2i),
+    contact_efc_address_out: wp.array2d(dtype=int),
     contact_worldid_out: wp.array(dtype=int),
     contact_type_out: wp.array(dtype=int),
     contact_geomcollisionid_out: wp.array(dtype=int),
@@ -1416,6 +1444,7 @@ def _primitive_narrowphase(primitive_collisions_types, primitive_collisions_func
           contact_solimp_out,
           contact_dim_out,
           contact_geom_out,
+          contact_efc_address_out,
           contact_worldid_out,
           contact_type_out,
           contact_geomcollisionid_out,
@@ -1511,6 +1540,7 @@ def primitive_narrowphase(m: Model, d: Data, ctx: CollisionContext, collision_ta
       d.contact.solimp,
       d.contact.dim,
       d.contact.geom,
+      d.contact.efc_address,
       d.contact.worldid,
       d.contact.type,
       d.contact.geomcollisionid,

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

@@ -1489,3 +1489,507 @@ def capsule_box(
     mat23f(pos1[0], pos1[1], pos1[2], pos2[0], pos2[1], pos2[2]),
     mat23f(normal1[0], normal1[1], normal1[2], normal2[0], normal2[1], normal2[2]),
   )
+
+
+@wp.func
+def _tri_area_sign(p1: wp.vec2, p2: wp.vec2, p3: wp.vec2) -> float:
+  """Sign of (signed) area of planar triangle."""
+  return wp.sign((p1[0] - p3[0]) * (p2[1] - p3[1]) - (p2[0] - p3[0]) * (p1[1] - p3[1]))
+
+
+@wp.func
+def _tri_point_segment(p: wp.vec2, u: wp.vec2, v: wp.vec2) -> wp.vec2:
+  """Find nearest point to p within line segment (u, v)."""
+  uv = v - u
+  up = p - u
+
+  denom = wp.max(MJ_MINVAL, wp.dot(uv, uv))
+  a = wp.dot(uv, up) / denom
+
+  if a <= 0.0:
+    return u
+  elif a >= 1.0:
+    return v
+  else:
+    return u + a * uv
+
+
+@wp.func
+def sphere_triangle(
+  sphere_pos: wp.vec3,
+  sphere_radius: float,
+  t1: wp.vec3,
+  t2: wp.vec3,
+  t3: wp.vec3,
+  tri_radius: float,
+) -> Tuple[float, wp.vec3, wp.vec3]:
+  """Core contact geometry calculation for sphere-triangle collision.
+
+  Port of mjraw_SphereTriangle from engine_collision_primitive.c
+
+  Args:
+    sphere_pos: Center position of the sphere.
+    sphere_radius: Radius of the sphere.
+    t1: Triangle vertex positions.
+    t2: Triangle vertex positions.
+    t3: Triangle vertex positions.
+    tri_radius: Triangle (flex element) radius.
+
+  Returns:
+    - Contact distance (MJ_MAXVAL if no collision).
+    - Contact position.
+    - Contact normal vector.
+  """
+  S = sphere_pos - t1
+  A = t2 - t1
+  B = t3 - t1
+
+  N = wp.normalize(wp.cross(A, B))
+
+  dstS = wp.dot(N, S)
+
+  P = S - dstS * N
+
+  V1 = wp.normalize(A)
+  lenA = wp.length(A)
+  V2 = wp.normalize(wp.cross(N, A))
+
+  o = wp.vec2(0.0, 0.0)
+  a = wp.vec2(lenA, 0.0)
+  b = wp.vec2(wp.dot(V1, B), wp.dot(V2, B))
+  p = wp.vec2(wp.dot(V1, P), wp.dot(V2, P))
+
+  sign1 = _tri_area_sign(p, o, a)
+  sign2 = _tri_area_sign(p, a, b)
+  sign3 = _tri_area_sign(p, b, o)
+
+  X = wp.vec3(0.0)
+  if sign1 == sign2 and sign2 == sign3:
+    X = P
+  else:
+    x0 = _tri_point_segment(p, o, a)
+    x1 = _tri_point_segment(p, a, b)
+    x2 = _tri_point_segment(p, b, o)
+
+    d0 = wp.length(p - x0)
+    d1 = wp.length(p - x1)
+    d2 = wp.length(p - x2)
+
+    if d0 < d1 and d0 < d2:
+      X = x0[0] * V1 + x0[1] * V2
+    elif d1 < d2:
+      X = x1[0] * V1 + x1[1] * V2
+    else:
+      X = x2[0] * V1 + x2[1] * V2
+
+  nrm = X - S
+  dst = wp.length(nrm)
+
+  if dst > MJ_MINVAL:
+    nrm = nrm / dst
+  else:
+    nrm = N
+
+  dist = dst - sphere_radius - tri_radius
+  pos = sphere_pos + nrm * (sphere_radius + 0.5 * dist)
+
+  return dist, pos, nrm
+
+
+@wp.func
+def box_triangle(
+  box_pos: wp.vec3,
+  box_rot: wp.mat33,
+  box_size: wp.vec3,
+  t1: wp.vec3,
+  t2: wp.vec3,
+  t3: wp.vec3,
+  tri_radius: float,
+) -> Tuple[wp.vec2, mat23f, mat23f]:
+  """Core contact geometry calculation for box-triangle collision.
+
+  Port of mjraw_BoxTriangle from engine_collision_primitive.c
+
+  Args:
+    box_pos: Center position of the box.
+    box_rot: Orientation matrix of the box.
+    box_size: Half-sizes of the box.
+    t1: Triangle vertex positions.
+    t2: Triangle vertex positions.
+    t3: Triangle vertex positions.
+    tri_radius: Triangle (flex element) radius.
+
+  Returns:
+    - wp.vec2 of distances for up to 2 contacts (MJ_MAXVAL if no collision).
+    - mat23f of contact positions (2 x vec3).
+    - mat23f of contact normals (2 x vec3).
+  """
+  dist1 = MJ_MAXVAL
+  dist2 = MJ_MAXVAL
+  pos1 = wp.vec3(0.0)
+  pos2 = wp.vec3(0.0)
+  nrm1 = wp.vec3(0.0)
+  nrm2 = wp.vec3(0.0)
+  cnt = 0
+
+  box_rotT = wp.transpose(box_rot)
+
+  for vi in range(3):
+    vert = wp.vec3(0.0)
+    if vi == 0:
+      vert = t1
+    elif vi == 1:
+      vert = t2
+    else:
+      vert = t3
+
+    diff = vert - box_pos
+    local = box_rotT @ diff
+
+    maxaxis = 0
+    maxval = wp.abs(local[0]) - box_size[0]
+    for j in range(1, 3):
+      val = wp.abs(local[j]) - box_size[j]
+      if val > maxval:
+        maxval = val
+        maxaxis = j
+
+    inside = True
+    for j in range(3):
+      if wp.abs(local[j]) > box_size[j] + tri_radius:
+        inside = False
+
+    if inside and cnt < 2:
+      nrm_local = wp.vec3(0.0)
+      if maxaxis == 0:
+        nrm_local = wp.vec3(wp.sign(local[0]), 0.0, 0.0)
+      elif maxaxis == 1:
+        nrm_local = wp.vec3(0.0, wp.sign(local[1]), 0.0)
+      else:
+        nrm_local = wp.vec3(0.0, 0.0, wp.sign(local[2]))
+
+      nrm_global = box_rot @ nrm_local
+      d = maxval - tri_radius
+      offset = tri_radius + d * 0.5
+      p = vert - nrm_global * offset
+
+      if cnt == 0:
+        dist1 = d
+        pos1 = p
+        nrm1 = nrm_global
+      else:
+        dist2 = d
+        pos2 = p
+        nrm2 = nrm_global
+      cnt += 1
+
+  for i in range(8):
+    if cnt >= 2:
+      break
+
+    vec = wp.vec3(
+      wp.where(i & 1, box_size[0], -box_size[0]),
+      wp.where(i & 2, box_size[1], -box_size[1]),
+      wp.where(i & 4, box_size[2], -box_size[2]),
+    )
+    corner = box_rot @ vec + box_pos
+
+    d, p, n = sphere_triangle(corner, 0.0, t1, t2, t3, tri_radius)
+    if d < MJ_MAXVAL:
+      if cnt == 0:
+        dist1 = d
+        pos1 = p
+        nrm1 = n
+      elif cnt == 1:
+        dist2 = d
+        pos2 = p
+        nrm2 = n
+      cnt += 1
+
+  return (
+    wp.vec2(dist1, dist2),
+    mat23f(pos1[0], pos1[1], pos1[2], pos2[0], pos2[1], pos2[2]),
+    mat23f(nrm1[0], nrm1[1], nrm1[2], nrm2[0], nrm2[1], nrm2[2]),
+  )
+
+
+@wp.func
+def capsule_triangle(
+  capsule_pos: wp.vec3,
+  capsule_axis: wp.vec3,
+  capsule_radius: float,
+  capsule_half_length: float,
+  t1: wp.vec3,
+  t2: wp.vec3,
+  t3: wp.vec3,
+  tri_radius: float,
+) -> Tuple[wp.vec2, mat23f, mat23f]:
+  """Core contact geometry calculation for capsule-triangle collision.
+
+  Port of mjraw_CapsuleTriangle from engine_collision_primitive.c
+
+  Args:
+    capsule_pos: Center position of the capsule.
+    capsule_axis: Unit axis direction of the capsule.
+    capsule_radius: Radius of the capsule.
+    capsule_half_length: Half-length of the capsule cylinder.
+    t1: Triangle vertex positions.
+    t2: Triangle vertex positions.
+    t3: Triangle vertex positions.
+    tri_radius: Triangle (flex element) radius.
+
+  Returns:
+    - wp.vec2 of distances for up to 2 contacts (MJ_MAXVAL if no collision).
+    - mat23f of contact positions (2 x vec3).
+    - mat23f of contact normals (2 x vec3).
+  """
+  dist1 = MJ_MAXVAL
+  dist2 = MJ_MAXVAL
+  pos1 = wp.vec3(0.0)
+  pos2 = wp.vec3(0.0)
+  nrm1 = wp.vec3(0.0)
+  nrm2 = wp.vec3(0.0)
+  cnt = 0
+
+  p1 = capsule_pos - capsule_axis * capsule_half_length
+  p2 = capsule_pos + capsule_axis * capsule_half_length
+
+  d, p, n = sphere_triangle(p1, capsule_radius, t1, t2, t3, tri_radius)
+  if d < MJ_MAXVAL:
+    dist1 = d
+    pos1 = p
+    nrm1 = n
+    cnt = 1
+
+  d, p, n = sphere_triangle(p2, capsule_radius, t1, t2, t3, tri_radius)
+  if d < MJ_MAXVAL and cnt < 2:
+    if cnt == 0:
+      dist1 = d
+      pos1 = p
+      nrm1 = n
+    else:
+      dist2 = d
+      pos2 = p
+      nrm2 = n
+    cnt += 1
+
+  ab = p2 - p1
+  ab_len_sq = 4.0 * capsule_half_length * capsule_half_length
+
+  for vi in range(3):
+    if cnt >= 2:
+      break
+
+    vert = wp.vec3(0.0)
+    if vi == 0:
+      vert = t1
+    elif vi == 1:
+      vert = t2
+    else:
+      vert = t3
+
+    vec = vert - p1
+    t_param = wp.dot(vec, ab) / wp.max(MJ_MINVAL, ab_len_sq)
+
+    if t_param > MJ_MINVAL and t_param < 1.0 - MJ_MINVAL:
+      closest = p1 + ab * t_param
+      diff = vert - closest
+      dist_raw = wp.length(diff)
+
+      if dist_raw > MJ_MINVAL:
+        nrm = diff / dist_raw
+        d = dist_raw - capsule_radius - tri_radius
+        p = (closest + vert + nrm * (capsule_radius - tri_radius)) * 0.5
+
+        if cnt == 0:
+          dist1 = d
+          pos1 = p
+          nrm1 = nrm
+        else:
+          dist2 = d
+          pos2 = p
+          nrm2 = nrm
+        cnt += 1
+
+  return (
+    wp.vec2(dist1, dist2),
+    mat23f(pos1[0], pos1[1], pos1[2], pos2[0], pos2[1], pos2[2]),
+    mat23f(nrm1[0], nrm1[1], nrm1[2], nrm2[0], nrm2[1], nrm2[2]),
+  )
+
+
+@wp.func
+def cylinder_triangle(
+  cylinder_pos: wp.vec3,
+  cylinder_axis: wp.vec3,
+  cylinder_radius: float,
+  cylinder_half_height: float,
+  t1: wp.vec3,
+  t2: wp.vec3,
+  t3: wp.vec3,
+  tri_radius: float,
+) -> Tuple[wp.vec2, mat23f, mat23f]:
+  """Core contact geometry calculation for cylinder-triangle collision.
+
+  Args:
+    cylinder_pos: Center position of the cylinder.
+    cylinder_axis: Unit axis direction of the cylinder.
+    cylinder_radius: Radius of the cylinder.
+    cylinder_half_height: Half-height of the cylinder.
+    t1: Triangle vertex positions.
+    t2: Triangle vertex positions.
+    t3: Triangle vertex positions.
+    tri_radius: Triangle (flex element) radius.
+
+  Returns:
+    - wp.vec2 of distances for up to 2 contacts (MJ_MAXVAL if no collision).
+    - mat23f of contact positions (2 x vec3).
+    - mat23f of contact normals (2 x vec3).
+  """
+  dist1 = MJ_MAXVAL
+  dist2 = MJ_MAXVAL
+  pos1 = wp.vec3(0.0)
+  pos2 = wp.vec3(0.0)
+  nrm1 = wp.vec3(0.0)
+  nrm2 = wp.vec3(0.0)
+  cnt = int(0)
+
+  p1 = cylinder_pos - cylinder_axis * cylinder_half_height
+  p2 = cylinder_pos + cylinder_axis * cylinder_half_height
+
+  ab = p2 - p1
+  ab_len_sq = 4.0 * cylinder_half_height * cylinder_half_height
+
+  for vi in range(3):
+    if cnt >= 2:
+      break
+
+    vert = wp.vec3(0.0)
+    if vi == 0:
+      vert = t1
+    elif vi == 1:
+      vert = t2
+    else:
+      vert = t3
+
+    vec = vert - p1
+    t_param = wp.dot(vec, ab) / wp.max(MJ_MINVAL, ab_len_sq)
+
+    if t_param > MJ_MINVAL and t_param < 1.0 - MJ_MINVAL:
+      closest = p1 + ab * t_param
+      diff = vert - closest
+      dist_raw = wp.length(diff)
+
+      if dist_raw < cylinder_radius + tri_radius:
+        if dist_raw > MJ_MINVAL:
+          nrm = diff / dist_raw
+          d = dist_raw - cylinder_radius - tri_radius
+          p = (closest + vert + nrm * (cylinder_radius - tri_radius)) * 0.5
+        else:
+          dist_to_side = cylinder_radius
+          dist_to_p2 = (1.0 - t_param) * wp.sqrt(ab_len_sq)
+          dist_to_p1 = t_param * wp.sqrt(ab_len_sq)
+
+          if dist_to_p2 < dist_to_side and dist_to_p2 < dist_to_p1:
+            nrm = cylinder_axis
+            d = -dist_to_p2 - tri_radius
+            p = vert
+          elif dist_to_p1 < dist_to_side:
+            nrm = -cylinder_axis
+            d = -dist_to_p1 - tri_radius
+            p = vert
+          else:
+            tri_normal = wp.normalize(wp.cross(t2 - t1, t3 - t1))
+            nrm = tri_normal
+            d = -cylinder_radius - tri_radius
+            p = closest
+
+        if cnt == 0:
+          dist1 = d
+          pos1 = p
+          nrm1 = nrm
+        else:
+          dist2 = d
+          pos2 = p
+          nrm2 = nrm
+        cnt += 1
+    elif t_param <= MJ_MINVAL:
+      diff = vert - p1
+      signed_dist = wp.dot(diff, cylinder_axis)
+      perp = diff - cylinder_axis * signed_dist
+      perp_len = wp.length(perp)
+
+      if perp_len < cylinder_radius:
+        d = -signed_dist - tri_radius
+        nrm = -cylinder_axis
+        p = vert - nrm * (tri_radius + d * 0.5)
+        if cnt == 0:
+          dist1 = d
+          pos1 = p
+          nrm1 = nrm
+        else:
+          dist2 = d
+          pos2 = p
+          nrm2 = nrm
+        cnt += 1
+      elif perp_len < cylinder_radius + tri_radius:
+        edge_dir = perp / perp_len
+        edge_point = p1 + edge_dir * cylinder_radius
+        diff_to_edge = vert - edge_point
+        dist_raw = wp.length(diff_to_edge)
+        if dist_raw > MJ_MINVAL:
+          nrm = diff_to_edge / dist_raw
+          d = dist_raw - tri_radius
+          p = vert - nrm * (tri_radius + d * 0.5)
+          if cnt == 0:
+            dist1 = d
+            pos1 = p
+            nrm1 = nrm
+          else:
+            dist2 = d
+            pos2 = p
+            nrm2 = nrm
+          cnt += 1
+    else:
+      diff = vert - p2
+      signed_dist = wp.dot(diff, cylinder_axis)
+      perp = diff - cylinder_axis * signed_dist
+      perp_len = wp.length(perp)
+
+      if perp_len < cylinder_radius:
+        d = signed_dist - tri_radius
+        nrm = cylinder_axis
+        p = vert - nrm * (tri_radius + d * 0.5)
+        if cnt == 0:
+          dist1 = d
+          pos1 = p
+          nrm1 = nrm
+        else:
+          dist2 = d
+          pos2 = p
+          nrm2 = nrm
+        cnt += 1
+      elif perp_len < cylinder_radius + tri_radius:
+        edge_dir = perp / perp_len
+        edge_point = p2 + edge_dir * cylinder_radius
+        diff_to_edge = vert - edge_point
+        dist_raw = wp.length(diff_to_edge)
+        if dist_raw > MJ_MINVAL:
+          nrm = diff_to_edge / dist_raw
+          d = dist_raw - tri_radius
+          p = vert - nrm * (tri_radius + d * 0.5)
+          if cnt == 0:
+            dist1 = d
+            pos1 = p
+            nrm1 = nrm
+          else:
+            dist2 = d
+            pos2 = p
+            nrm2 = nrm
+          cnt += 1
+
+  return (
+    wp.vec2(dist1, dist2),
+    mat23f(pos1[0], pos1[1], pos1[2], pos2[0], pos2[1], pos2[2]),
+    mat23f(nrm1[0], nrm1[1], nrm1[2], nrm2[0], nrm2[1], nrm2[2]),
+  )

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

@@ -15,6 +15,8 @@
 
 from typing import Tuple
 
+import warp as wp
+
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair
@@ -27,9 +29,9 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec8
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i
+from mujoco.mjx.third_party.mujoco_warp._src.types import vec_pluginattr
 from mujoco.mjx.third_party.mujoco_warp._src.util_misc import halton
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-import warp as wp
 
 wp.set_module_options({"enable_backward": False})
 
@@ -38,8 +40,8 @@ wp.set_module_options({"enable_backward": False})
 class OptimizationParams:
   rel_mat: wp.mat33
   rel_pos: wp.vec3
-  attr1: wp.vec3
+  attr1: vec_pluginattr
-  attr2: wp.vec3
+  attr2: vec_pluginattr
 
 
 @wp.struct
@@ -83,14 +85,18 @@ def get_sdf_params(
   oct_coeff: wp.array(dtype=vec8),
   mesh_octadr: wp.array(dtype=int),
   plugin: wp.array(dtype=int),
-    plugin_attr: wp.array(dtype=wp.vec3f),
+  plugin_attr: wp.array(dtype=vec_pluginattr),
   # In:
   g_type: int,
   g_size: wp.vec3,
   plugin_id: int,
   mesh_id: int,
-) -> Tuple[wp.vec3, int, VolumeData, MeshData]:
+) -> Tuple[vec_pluginattr, int, VolumeData, MeshData]:
-  attributes = g_size
+  # default attributes from geom size, first 3 values copied
+  attributes = vec_pluginattr()
+  attributes[0] = g_size[0]
+  attributes[1] = g_size[1]
+  attributes[2] = g_size[2]
   plugin_index = -1
   volume_data = VolumeData()
 
@@ -108,6 +114,16 @@ def get_sdf_params(
     volume_data.oct_coeff = oct_coeff
     volume_data.valid = True
 
+  elif g_type == GeomType.MESH and mesh_id != -1 and mesh_octadr[mesh_id] != -1:
+    octadr = mesh_octadr[mesh_id]
+    volume_data.center = oct_aabb[octadr, 0]
+    volume_data.half_size = oct_aabb[octadr, 1]
+    volume_data.root = octadr
+    volume_data.oct_aabb = oct_aabb
+    volume_data.oct_child = oct_child
+    volume_data.oct_coeff = oct_coeff
+    volume_data.valid = True
+
   return attributes, plugin_index, volume_data, MeshData()
 
 
@@ -215,24 +231,28 @@ def grad_ellipsoid(p: wp.vec3, size: wp.vec3) -> wp.vec3:
 
 
 @wp.func
-def user_sdf(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> float:
+def user_sdf(p: wp.vec3, attr: vec_pluginattr, sdf_type: int) -> float:
+  """User-defined SDF function.
+
+  Access attributes via attr[i] where i is the attribute index (0 to _NPLUGINATTR-1).
+  """
   wp.printf("ERROR: user_sdf function must be implemented by user code\n")
   return 0.0
 
 
 @wp.func
-def user_sdf_grad(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> wp.vec3:
+def user_sdf_grad(p: wp.vec3, attr: vec_pluginattr, sdf_type: int) -> wp.vec3:
+  """User-defined SDF gradient function.
+
+  Access attributes via attr[i] where i is the attribute index (0 to _NPLUGINATTR-1).
+  """
   wp.printf("ERROR: user_sdf_grad function must be implemented by user code\n")
   return wp.vec3(0.0)
 
 
 @wp.func
 def find_oct(
-    oct_child: wp.array(dtype=vec8i),
+  oct_child: wp.array(dtype=vec8i), oct_aabb: wp.array2d(dtype=wp.vec3), p: wp.vec3, grad: bool, root: int
-    oct_aabb: wp.array2d(dtype=wp.vec3),
-    p: wp.vec3,
-    grad: bool,
-    root: int,
 ) -> Tuple[int, Tuple[vec8, vec8, vec8]]:
   stack = root
   niter = int(100)
@@ -342,13 +362,7 @@ def box_project(center: wp.vec3, half_size: wp.vec3, xyz: wp.vec3) -> Tuple[floa
 @wp.func
 def sample_volume_sdf(xyz: wp.vec3, volume_data: VolumeData) -> float:
   dist0, point = box_project(volume_data.center, volume_data.half_size, xyz)
-  node, weights = find_oct(
+  node, weights = find_oct(volume_data.oct_child, volume_data.oct_aabb, point, grad=False, root=volume_data.root)
-      volume_data.oct_child,
-      volume_data.oct_aabb,
-      point,
-      grad=False,
-      root=volume_data.root,
-  )
   return dist0 + wp.dot(weights[0], volume_data.oct_coeff[node])
 
 
@@ -365,13 +379,7 @@ def sample_volume_grad(xyz: wp.vec3, volume_data: VolumeData) -> wp.vec3:
     grad_y = (sample_volume_sdf(xyz + dy, volume_data) - f) / h
     grad_z = (sample_volume_sdf(xyz + dz, volume_data) - f) / h
     return wp.vec3(grad_x, grad_y, grad_z)
-  node, weights = find_oct(
+  node, weights = find_oct(volume_data.oct_child, volume_data.oct_aabb, point, grad=True, root=volume_data.root)
-      volume_data.oct_child,
-      volume_data.oct_aabb,
-      point,
-      grad=True,
-      root=volume_data.root,
-  )
   grad_x = wp.dot(weights[0], volume_data.oct_coeff[node])
   grad_y = wp.dot(weights[1], volume_data.oct_coeff[node])
   grad_z = wp.dot(weights[2], volume_data.oct_coeff[node])
@@ -379,15 +387,17 @@ def sample_volume_grad(xyz: wp.vec3, volume_data: VolumeData) -> wp.vec3:
 
 
 @wp.func
-def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: VolumeData, mesh_data: MeshData) -> float:
+def sdf(type: int, p: wp.vec3, attr: vec_pluginattr, sdf_type: int, volume_data: VolumeData, mesh_data: MeshData) -> float:
+  # extract first 3 elements as vec3 for primitive sdf functions
+  attr_vec3 = wp.vec3(attr[0], attr[1], attr[2])
   if type == GeomType.PLANE:
     return p[2]
   elif type == GeomType.SPHERE:
-    return sphere(p, attr)
+    return sphere(p, attr_vec3)
   elif type == GeomType.BOX:
-    return box(p, attr)
+    return box(p, attr_vec3)
   elif type == GeomType.ELLIPSOID:
-    return ellipsoid(p, attr)
+    return ellipsoid(p, attr_vec3)
   elif type == GeomType.MESH and mesh_data.valid:
     mesh_data.pnt = p
     mesh_data.vec = -wp.normalize(p)
@@ -425,21 +435,27 @@ def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: Volume
       return sample_volume_sdf(p, volume_data)
     else:
       return user_sdf(p, attr, sdf_type)
+  elif type == GeomType.MESH and volume_data.valid:
+    return sample_volume_sdf(p, volume_data)
   wp.printf("ERROR: SDF type not implemented\n")
   return 0.0
 
 
 @wp.func
-def sdf_grad(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: VolumeData, mesh_data: MeshData) -> wp.vec3:
+def sdf_grad(
+  type: int, p: wp.vec3, attr: vec_pluginattr, sdf_type: int, volume_data: VolumeData, mesh_data: MeshData
+) -> wp.vec3:
+  # extract first 3 elements as vec3 for primitive sdf functions
+  attr_vec3 = wp.vec3(attr[0], attr[1], attr[2])
   if type == GeomType.PLANE:
     grad = wp.vec3(0.0, 0.0, 1.0)
     return grad
   elif type == GeomType.SPHERE:
     return grad_sphere(p)
   elif type == GeomType.BOX:
-    return grad_box(p, attr)
+    return grad_box(p, attr_vec3)
   elif type == GeomType.ELLIPSOID:
-    return grad_ellipsoid(p, attr)
+    return grad_ellipsoid(p, attr_vec3)
   elif type == GeomType.MESH and mesh_data.valid:
     mesh_data.pnt = p
     mesh_data.vec = -wp.normalize(p)
@@ -466,6 +482,8 @@ def sdf_grad(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: V
       return sample_volume_grad(p, volume_data)
     else:
       return user_sdf_grad(p, attr, sdf_type)
+  elif type == GeomType.MESH and volume_data.valid:
+    return sample_volume_grad(p, volume_data)
   wp.printf("ERROR: SDF grad type not implemented\n")
   return wp.vec3(0.0)
 
@@ -476,8 +494,8 @@ def clearance(
   type1: int,
   p1: wp.vec3,
   p2: wp.vec3,
-  s1: wp.vec3,
+  s1: vec_pluginattr,
-  s2: wp.vec3,
+  s2: vec_pluginattr,
   sdf_type1: int,
   sdf_type2: int,
   sfd_intersection: bool,
@@ -606,8 +624,8 @@ def gradient_descent(
   # In:
   type1: int,
   x0_initial: wp.vec3,
-  attr1: wp.vec3,
+  attr1: vec_pluginattr,
-  attr2: wp.vec3,
+  attr2: vec_pluginattr,
   pos1: wp.vec3,
   rot1: wp.mat33,
   pos2: wp.vec3,
@@ -687,7 +705,7 @@ def _sdf_narrowphase(
   pair_gap: wp.array2d(dtype=float),
   pair_friction: wp.array2d(dtype=vec5),
   plugin: wp.array(dtype=int),
-    plugin_attr: wp.array(dtype=wp.vec3f),
+  plugin_attr: wp.array(dtype=vec_pluginattr),
   geom_plugin_index: wp.array(dtype=int),
   # Data in:
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
@@ -711,6 +729,7 @@ def _sdf_narrowphase(
   contact_solimp_out: wp.array(dtype=vec5),
   contact_dim_out: wp.array(dtype=int),
   contact_geom_out: wp.array(dtype=wp.vec2i),
+  contact_efc_address_out: wp.array2d(dtype=int),
   contact_worldid_out: wp.array(dtype=int),
   contact_type_out: wp.array(dtype=int),
   contact_geomcollisionid_out: wp.array(dtype=int),
@@ -799,29 +818,11 @@ def _sdf_narrowphase(
   rot1 = geom1.rot
 
   attr1, g1_plugin_id, volume_data1, mesh_data1 = get_sdf_params(
-      oct_child,
+    oct_child, oct_aabb, oct_coeff, mesh_octadr, plugin, plugin_attr, type1, geom1.size, g1_plugin, geom_dataid[g1]
-      oct_aabb,
-      oct_coeff,
-      mesh_octadr,
-      plugin,
-      plugin_attr,
-      type1,
-      geom1.size,
-      g1_plugin,
-      geom_dataid[g1],
   )
 
   attr2, g2_plugin_id, volume_data2, mesh_data2 = get_sdf_params(
-      oct_child,
+    oct_child, oct_aabb, oct_coeff, mesh_octadr, plugin, plugin_attr, type2, geom2.size, g2_plugin, geom_dataid[g2]
-      oct_aabb,
-      oct_coeff,
-      mesh_octadr,
-      plugin,
-      plugin_attr,
-      type2,
-      geom2.size,
-      g2_plugin,
-      geom_dataid[g2],
   )
 
   mesh_data1.nmeshface = nmeshface
@@ -900,6 +901,7 @@ def _sdf_narrowphase(
     contact_solimp_out,
     contact_dim_out,
     contact_geom_out,
+    contact_efc_address_out,
     contact_worldid_out,
     contact_type_out,
     contact_geomcollisionid_out,
@@ -977,6 +979,7 @@ def sdf_narrowphase(m: Model, d: Data, ctx: CollisionContext):
       d.contact.solimp,
       d.contact.dim,
       d.contact.geom,
+      d.contact.efc_address,
       d.contact.worldid,
       d.contact.type,
       d.contact.geomcollisionid,

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

@@ -15,6 +15,7 @@
 
 import warp as wp
 
+from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
 from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
@@ -30,18 +31,24 @@ wp.set_module_options({"enable_backward": False})
 @wp.kernel
 def _qderiv_actuator_passive_vel(
   # Model:
+  opt_timestep: wp.array(dtype=float),
   actuator_dyntype: wp.array(dtype=int),
   actuator_gaintype: wp.array(dtype=int),
   actuator_biastype: wp.array(dtype=int),
   actuator_actadr: wp.array(dtype=int),
   actuator_actnum: wp.array(dtype=int),
   actuator_forcelimited: wp.array(dtype=bool),
+  actuator_actlimited: wp.array(dtype=bool),
+  actuator_dynprm: wp.array2d(dtype=vec10f),
   actuator_gainprm: wp.array2d(dtype=vec10f),
   actuator_biasprm: wp.array2d(dtype=vec10f),
+  actuator_actearly: wp.array(dtype=bool),
   actuator_forcerange: wp.array2d(dtype=wp.vec2),
+  actuator_actrange: wp.array2d(dtype=wp.vec2),
   # Data in:
   act_in: wp.array2d(dtype=float),
   ctrl_in: wp.array2d(dtype=float),
+  act_dot_in: wp.array2d(dtype=float),
   actuator_force_in: wp.array2d(dtype=float),
   # Out:
   vel_out: wp.array2d(dtype=float),
@@ -76,9 +83,24 @@ def _qderiv_actuator_passive_vel(
   vel = float(bias)
   if actuator_dyntype[actid] != DynType.NONE:
     if gain != 0.0:
-      act_first = actuator_actadr[actid]
+      act_adr = actuator_actadr[actid] + actuator_actnum[actid] - 1
-      act_last = act_first + actuator_actnum[actid] - 1
+
-      vel += gain * act_in[worldid, act_last]
+      # use next activation if actearly is set (matching forward pass)
+      if actuator_actearly[actid]:
+        act = next_act(
+          opt_timestep[worldid % opt_timestep.shape[0]],
+          actuator_dyntype[actid],
+          actuator_dynprm[worldid % actuator_dynprm.shape[0], actid],
+          actuator_actrange[worldid % actuator_actrange.shape[0], actid],
+          act_in[worldid, act_adr],
+          act_dot_in[worldid, act_adr],
+          1.0,
+          actuator_actlimited[actid],
+        )
+      else:
+        act = act_in[worldid, act_adr]
+
+      vel += gain * act
   else:
     if gain != 0.0:
       vel += gain * ctrl_in[worldid, actid]
@@ -95,10 +117,9 @@ def _nonzero_mask(x: float) -> float:
 
 
 @wp.kernel
-def _qderiv_actuator_passive_actuation_sparse(
+def _qderiv_actuator_passive_actuation_dense(
   # Model:
   nu: int,
-    is_sparse: bool,
   # Data in:
   moment_rownnz_in: wp.array2d(dtype=int),
   moment_rowadr_in: wp.array2d(dtype=int),
@@ -142,14 +163,65 @@ def _qderiv_actuator_passive_actuation_sparse(
 
     qderiv_contrib += moment_i * moment_j * vel
 
-  if is_sparse:
-    qDeriv_out[worldid, 0, elemid] = qderiv_contrib
-  else:
   qDeriv_out[worldid, dofiid, dofjid] = qderiv_contrib
   if dofiid != dofjid:
     qDeriv_out[worldid, dofjid, dofiid] = qderiv_contrib
 
 
+@wp.kernel
+def _qderiv_actuator_passive_actuation_sparse(
+  # Model:
+  M_rownnz: wp.array(dtype=int),
+  M_rowadr: wp.array(dtype=int),
+  # Data in:
+  moment_rownnz_in: wp.array2d(dtype=int),
+  moment_rowadr_in: wp.array2d(dtype=int),
+  moment_colind_in: wp.array2d(dtype=int),
+  actuator_moment_in: wp.array2d(dtype=float),
+  # In:
+  vel_in: wp.array2d(dtype=float),
+  qMj: wp.array(dtype=int),
+  # Out:
+  qDeriv_out: wp.array3d(dtype=float),
+):
+  worldid, actid = wp.tid()
+
+  vel = vel_in[worldid, actid]
+  if vel == 0.0:
+    return
+
+  rownnz = moment_rownnz_in[worldid, actid]
+  rowadr = moment_rowadr_in[worldid, actid]
+
+  for i in range(rownnz):
+    rowadri = rowadr + i
+    moment_i = actuator_moment_in[worldid, rowadri]
+    if moment_i == 0.0:
+      continue
+    dofi = moment_colind_in[worldid, rowadri]
+
+    for j in range(i + 1):
+      rowadrj = rowadr + j
+      moment_j = actuator_moment_in[worldid, rowadrj]
+      if moment_j == 0.0:
+        continue
+      dofj = moment_colind_in[worldid, rowadrj]
+
+      contrib = moment_i * moment_j * vel
+
+      # Search the corresponding elemid
+      # TODO: This could be precalculated for improved performance
+      row = dofi
+      col = dofj
+      row_startk = M_rowadr[row] - 1
+      row_nnz = M_rownnz[row]
+      for k in range(row_nnz):
+        row_startk += 1
+        if qMj[row_startk] == col:
+          wp.atomic_add(qDeriv_out[worldid, 0], row_startk, contrib)
+          break
+
+
 @wp.kernel
 def _qderiv_actuator_passive(
   # Model:
@@ -176,7 +248,7 @@ def _qderiv_actuator_passive(
   else:
     qderiv = qDeriv_in[worldid, dofiid, dofjid]
 
-  if not opt_disableflags & DisableBit.DAMPER and dofiid == dofjid:
+  if not (opt_disableflags & DisableBit.DAMPER) and dofiid == dofjid:
     qderiv -= dof_damping[worldid % dof_damping.shape[0], dofiid]
 
   qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
@@ -196,10 +268,13 @@ def _qderiv_tendon_damping(
   # Model:
   ntendon: int,
   opt_timestep: wp.array(dtype=float),
+  ten_J_rownnz: wp.array(dtype=int),
+  ten_J_rowadr: wp.array(dtype=int),
+  ten_J_colind: wp.array(dtype=int),
   tendon_damping: wp.array2d(dtype=float),
   is_sparse: bool,
   # Data in:
-  ten_J_in: wp.array3d(dtype=float),
+  ten_J_in: wp.array2d(dtype=float),
   # In:
   qMi: wp.array(dtype=int),
   qMj: wp.array(dtype=int),
@@ -213,7 +288,24 @@ def _qderiv_tendon_damping(
   qderiv = float(0.0)
   tendon_damping_id = worldid % tendon_damping.shape[0]
   for tenid in range(ntendon):
-    qderiv -= ten_J_in[worldid, tenid, dofiid] * ten_J_in[worldid, tenid, dofjid] * tendon_damping[tendon_damping_id, tenid]
+    damping = tendon_damping[tendon_damping_id, tenid]
+    if damping == 0.0:
+      continue
+
+    rownnz = ten_J_rownnz[tenid]
+    rowadr = ten_J_rowadr[tenid]
+    Ji = float(0.0)
+    Jj = float(0.0)
+    for k in range(rownnz):
+      if Ji != 0.0 and Jj != 0.0:
+        break
+      sparseid = rowadr + k
+      colind = ten_J_colind[sparseid]
+      if colind == dofiid:
+        Ji = ten_J_in[worldid, sparseid]
+      if colind == dofjid:
+        Jj = ten_J_in[worldid, sparseid]
+    qderiv -= Ji * Jj * damping
 
   qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
 
@@ -242,41 +334,45 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
   if ~(m.opt.disableflags & (DisableBit.ACTUATION | DisableBit.DAMPER)):
     # TODO(team): only clear elements not set by _qderiv_actuator_passive
     out.zero_()
-    if m.nu > 0 and not m.opt.disableflags & DisableBit.ACTUATION:
+    if m.nu > 0 and not (m.opt.disableflags & DisableBit.ACTUATION):
       vel = wp.empty((d.nworld, m.nu), dtype=float)
       wp.launch(
         _qderiv_actuator_passive_vel,
         dim=(d.nworld, m.nu),
         inputs=[
+          m.opt.timestep,
           m.actuator_dyntype,
           m.actuator_gaintype,
           m.actuator_biastype,
           m.actuator_actadr,
           m.actuator_actnum,
           m.actuator_forcelimited,
+          m.actuator_actlimited,
+          m.actuator_dynprm,
           m.actuator_gainprm,
           m.actuator_biasprm,
+          m.actuator_actearly,
           m.actuator_forcerange,
+          m.actuator_actrange,
           d.act,
           d.ctrl,
+          d.act_dot,
           d.actuator_force,
         ],
         outputs=[vel],
       )
+      if m.is_sparse:
         wp.launch(
           _qderiv_actuator_passive_actuation_sparse,
+          dim=(d.nworld, m.nu),
+          inputs=[m.M_rownnz, m.M_rowadr, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMj],
+          outputs=[out],
+        )
+      else:
+        wp.launch(
+          _qderiv_actuator_passive_actuation_dense,
           dim=(d.nworld, qMi.size),
-          inputs=[
+          inputs=[m.nu, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMi, qMj],
-              m.nu,
-              m.is_sparse,
-              d.moment_rownnz,
-              d.moment_rowadr,
-              d.moment_colind,
-              d.actuator_moment,
-              vel,
-              qMi,
-              qMj,
-          ],
           outputs=[out],
         )
     wp.launch(
@@ -298,11 +394,22 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
     # TODO(team): directly utilize qM for these settings
     wp.copy(out, d.qM)
 
-  if not m.opt.disableflags & DisableBit.DAMPER:
+  if not (m.opt.disableflags & DisableBit.DAMPER):
     wp.launch(
       _qderiv_tendon_damping,
       dim=(d.nworld, qMi.size),
-      inputs=[m.ntendon, m.opt.timestep, m.tendon_damping, m.is_sparse, d.ten_J, qMi, qMj],
+      inputs=[
+        m.ntendon,
+        m.opt.timestep,
+        m.ten_J_rownnz,
+        m.ten_J_rowadr,
+        m.ten_J_colind,
+        m.tendon_damping,
+        m.is_sparse,
+        d.ten_J,
+        qMi,
+        qMj,
+      ],
       outputs=[out],
     )
 

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

@@ -15,6 +15,8 @@
 
 from typing import Optional
 
+import warp as wp
+
 from mujoco.mjx.third_party.mujoco_warp._src import collision_driver
 from mujoco.mjx.third_party.mujoco_warp._src import constraint
 from mujoco.mjx.third_party.mujoco_warp._src import derivative
@@ -25,7 +27,9 @@ from mujoco.mjx.third_party.mujoco_warp._src import sensor
 from mujoco.mjx.third_party.mujoco_warp._src import smooth
 from mujoco.mjx.third_party.mujoco_warp._src import solver
 from mujoco.mjx.third_party.mujoco_warp._src import util_misc
+from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
 from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
@@ -34,14 +38,12 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import GainType
 from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType
 from mujoco.mjx.third_party.mujoco_warp._src.types import JointType
-from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet
 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
-import warp as wp
 
 wp.set_module_options({"enable_backward": False})
 
@@ -127,37 +129,6 @@ def _next_velocity(
   qvel_out[worldid, dofid] = qvel_in[worldid, dofid] + qacc_scale_in * qacc_in[worldid, dofid] * timestep
 
 
-# TODO(team): kernel analyzer array slice?
-@wp.func
-def _next_act(
-  # Model:
-  opt_timestep: float,  # kernel_analyzer: ignore
-  actuator_dyntype: int,  # kernel_analyzer: ignore
-  actuator_dynprm: vec10f,  # kernel_analyzer: ignore
-  actuator_actrange: wp.vec2,  # kernel_analyzer: ignore
-  # Data In:
-  act_in: float,  # kernel_analyzer: ignore
-  act_dot_in: float,  # kernel_analyzer: ignore
-  # In:
-  act_dot_scale: float,
-  clamp: bool,
-) -> float:
-  # advance actuation
-  if actuator_dyntype == DynType.FILTEREXACT:
-    tau = wp.max(MJ_MINVAL, actuator_dynprm[0])
-    act = act_in + act_dot_scale * act_dot_in * tau * (1.0 - wp.exp(-opt_timestep / tau))
-  elif actuator_dyntype == DynType.USER:
-    return act_in
-  else:
-    act = act_in + act_dot_scale * act_dot_in * opt_timestep
-
-  # clamp to actrange
-  if clamp:
-    act = wp.clamp(act, actuator_actrange[0], actuator_actrange[1])
-
-  return act
-
-
 @wp.kernel
 def _next_activation(
   # Model:
@@ -184,7 +155,7 @@ def _next_activation(
   actadr = actuator_actadr[uid]
   actnum = actuator_actnum[uid]
   for j in range(actadr, actadr + actnum):
-    act = _next_act(
+    act = next_act(
       opt_timestep[opt_timestep_id],
       actuator_dyntype[uid],
       actuator_dynprm[actuator_dynprm_id, uid],
@@ -201,12 +172,16 @@ def _next_activation(
 def _next_time(
   # Model:
   opt_timestep: wp.array(dtype=float),
+  is_sparse: bool,
   # Data in:
   nefc_in: wp.array(dtype=int),
   time_in: wp.array(dtype=float),
+  efc_J_rownnz_in: wp.array2d(dtype=int),
+  efc_J_rowadr_in: wp.array2d(dtype=int),
   nworld_in: int,
   naconmax_in: int,
   njmax_in: int,
+  njmax_nnz_in: int,
   nacon_in: wp.array(dtype=int),
   ncollision_in: wp.array(dtype=int),
   # Data out:
@@ -218,6 +193,11 @@ def _next_time(
 
   if nefc > njmax_in:
     wp.printf("nefc overflow - please increase njmax to %u\n", nefc)
+  elif nefc > 0 and is_sparse:
+    efcid = wp.min(nefc, njmax_in) - 1
+    efc_nnz = efc_J_rowadr_in[worldid, efcid] + efc_J_rownnz_in[worldid, efcid]
+    if efc_nnz > njmax_nnz_in:
+      wp.printf("njmax_nnz overflow - please increase njmax_nnz to %u\n", efc_nnz)
 
   if worldid == 0:
     ncollision = ncollision_in[0]
@@ -274,7 +254,20 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
   wp.launch(
     _next_time,
     dim=d.nworld,
-    inputs=[m.opt.timestep, d.nefc, d.time, d.nworld, d.naconmax, d.njmax, d.nacon, d.ncollision],
+    inputs=[
+      m.opt.timestep,
+      m.is_sparse,
+      d.nefc,
+      d.time,
+      d.efc.J_rownnz,
+      d.efc.J_rowadr,
+      d.nworld,
+      d.naconmax,
+      d.njmax,
+      d.njmax_nnz,
+      d.nacon,
+      d.ncollision,
+    ],
     outputs=[d.time],
   )
 
@@ -294,9 +287,7 @@ def _euler_damp_qfrc_sparse(
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
 
   adr = dof_Madr[tid]
-  qM_integration_out[worldid, 0, adr] += (
+  qM_integration_out[worldid, 0, adr] += timestep * dof_damping[worldid % dof_damping.shape[0], tid]
-      timestep * dof_damping[worldid % dof_damping.shape[0], tid]
-  )
 
 
 @cache_kernel
@@ -336,7 +327,7 @@ def _tile_euler_dense(tile: TileSet):
 def euler(m: Model, d: Data):
   """Euler integrator, semi-implicit in velocity."""
   # integrate damping implicitly
-  if not m.opt.disableflags & (DisableBit.EULERDAMP | DisableBit.DAMPER):
+  if not (m.opt.disableflags & (DisableBit.EULERDAMP | DisableBit.DAMPER)):
     qacc = wp.empty((d.nworld, m.nv), dtype=float)
     if m.is_sparse:
       qM = wp.clone(d.qM)
@@ -571,50 +562,49 @@ def _actuator_velocity(
   actuator_velocity_out[worldid, actid] = vel
 
 
-@cache_kernel
+@wp.kernel
-def _tendon_velocity(nv: int):
+def _tendon_velocity(
-  @wp.kernel(module="unique", enable_backward=False)
+  # Model:
-  def tendon_velocity(
+  ten_J_rownnz: wp.array(dtype=int),
+  ten_J_rowadr: wp.array(dtype=int),
+  ten_J_colind: wp.array(dtype=int),
   # Data in:
   qvel_in: wp.array2d(dtype=float),
-    ten_J_in: wp.array3d(dtype=float),
+  ten_J_in: wp.array2d(dtype=float),
   # Data out:
   ten_velocity_out: wp.array2d(dtype=float),
-  ):
+):
   worldid, tenid = wp.tid()
-    ten_J_tile = wp.tile_load(ten_J_in[worldid, tenid], shape=wp.static(nv))
-    qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(nv))
-    ten_J_qvel_tile = wp.tile_map(wp.mul, ten_J_tile, qvel_tile)
-    ten_velocity_tile = wp.tile_reduce(wp.add, ten_J_qvel_tile)
-    ten_velocity_out[worldid, tenid] = ten_velocity_tile[0]
 
-  return tendon_velocity
+  velocity = float(0.0)
+  rownnz = ten_J_rownnz[tenid]
+  rowadr = ten_J_rowadr[tenid]
+  for i in range(rownnz):
+    sparseid = rowadr + i
+    J = ten_J_in[worldid, sparseid]
+    if J != 0.0:
+      colind = ten_J_colind[sparseid]
+      velocity += J * qvel_in[worldid, colind]
+
+  ten_velocity_out[worldid, tenid] = velocity
 
 
 @event_scope
 def fwd_velocity(m: Model, d: Data):
   """Velocity-dependent computations."""
-  wp.launch_tiled(
+  wp.launch(
     _actuator_velocity,
     dim=(d.nworld, m.nu),
-      inputs=[
+    inputs=[d.qvel, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment],
-          d.qvel,
-          d.moment_rownnz,
-          d.moment_rowadr,
-          d.moment_colind,
-          d.actuator_moment,
-      ],
     outputs=[d.actuator_velocity],
     block_dim=m.block_dim.actuator_velocity,
   )
 
-  # TODO(team): sparse version
+  wp.launch(
-  wp.launch_tiled(
+    _tendon_velocity,
-    _tendon_velocity(m.nv),
     dim=(d.nworld, m.ntendon),
-    inputs=[d.qvel, d.ten_J],
+    inputs=[m.ten_J_rownnz, m.ten_J_rowadr, m.ten_J_colind, d.qvel, d.ten_J],
     outputs=[d.ten_velocity],
-    block_dim=m.block_dim.tendon_velocity,
   )
 
   smooth.com_vel(m, d)
@@ -693,7 +683,7 @@ def _actuator_force(
       if dyntype == DynType.INTEGRATOR or dyntype == DynType.NONE:
         act = act_in[worldid, act_last]
 
-      ctrl_act = _next_act(
+      ctrl_act = next_act(
         opt_timestep[worldid % opt_timestep.shape[0]],
         dyntype,
         dynprm,
@@ -720,9 +710,7 @@ def _actuator_force(
     gain = gainprm[0] + gainprm[1] * length + gainprm[2] * velocity
   elif gaintype == GainType.MUSCLE:
     acc0 = actuator_acc0[worldid % actuator_acc0.shape[0], uid]
-    lengthrange = actuator_lengthrange[
+    lengthrange = actuator_lengthrange[worldid % actuator_lengthrange.shape[0], uid]
-        worldid % actuator_lengthrange.shape[0], uid
-    ]
     gain = util_misc.muscle_gain(length, velocity, lengthrange, acc0, gainprm)
   # GainType.USER: gain stays 0, modified by act_gain_callback
 
@@ -735,9 +723,7 @@ def _actuator_force(
     bias = biasprm[0] + biasprm[1] * length + biasprm[2] * velocity
   elif biastype == BiasType.MUSCLE:
     acc0 = actuator_acc0[worldid % actuator_acc0.shape[0], uid]
-    lengthrange = actuator_lengthrange[
+    lengthrange = actuator_lengthrange[worldid % actuator_lengthrange.shape[0], uid]
-        worldid % actuator_lengthrange.shape[0], uid
-    ]
     bias = util_misc.muscle_bias(length, lengthrange, acc0, biasprm)
 
   force = gain * ctrl_act + bias
@@ -812,10 +798,7 @@ def _qfrc_actuator(
   for i in range(rownnz):
     sparseid = rowadr + i
     colind = moment_colind_in[worldid, sparseid]
-    qfrc = (
+    qfrc = actuator_moment_in[worldid, sparseid] * actuator_force_in[worldid, actid]
-        actuator_moment_in[worldid, sparseid]
-        * actuator_force_in[worldid, actid]
-    )
     wp.atomic_add(qfrc_actuator_out[worldid], colind, qfrc)
 
 

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

@@ -14,24 +14,24 @@
 # ==============================================================================
 
 import dataclasses
-from typing import Any, Optional, Sequence
 import warnings
+from typing import Any, Optional, Sequence
 
 import mujoco
+import numpy as np
+import warp as wp
+
 from mujoco.mjx.third_party.mujoco_warp._src import bvh
 from mujoco.mjx.third_party.mujoco_warp._src import math as mjmath
 from mujoco.mjx.third_party.mujoco_warp._src import render_util
 from mujoco.mjx.third_party.mujoco_warp._src import smooth
 from mujoco.mjx.third_party.mujoco_warp._src import types
 from mujoco.mjx.third_party.mujoco_warp._src import warp_util
-from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
-from mujoco.mjx.third_party.mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN
+from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType
 from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
 from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
-import numpy as np
-import warp as wp
 
 
 def _create_array(data: Any, spec: wp.array, sizes: dict[str, int]) -> wp.array | None:
@@ -114,9 +114,6 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     if unsupported:
       raise NotImplementedError(f"{mj_type(unsupported).name} is unsupported.")
 
-  if ((mjm.flex_contype != 0) | (mjm.flex_conaffinity != 0)).any():
-    raise NotImplementedError("Flex collisions are not implemented.")
-
   if mjm.opt.noslip_iterations > 0:
     raise NotImplementedError(f"noslip solver not implemented.")
 
@@ -226,6 +223,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.is_sparse = is_sparse(mjm)
   m.has_fluid = mjm.opt.wind.any() or mjm.opt.density > 0 or mjm.opt.viscosity > 0
 
+  m.max_ten_J_rownnz = int(mjm.ten_J_rownnz.max()) if mjm.ntendon else 0
+
   # body ids grouped by tree level (depth-based traversal)
   bodies, body_depth = {}, np.zeros(mjm.nbody, dtype=int) - 1
   for i in range(mjm.nbody):
@@ -364,6 +363,46 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     )
   )
 
+  # check for unsupported margin + multicontact / box-box CCD combinations
+  use_multiccd = mjm.opt.enableflags & types.EnableBit.MULTICCD
+  nativeccd_disabled = mjm.opt.disableflags & types.DisableBit.NATIVECCD
+  BOX = int(mujoco.mjtGeom.mjGEOM_BOX)
+  MESH = int(mujoco.mjtGeom.mjGEOM_MESH)
+
+  has_boxbox = m.geom_pair_type_count[geom_trid_index(BOX, BOX)] > 0
+  has_multiccd_pairs = has_boxbox or (
+    use_multiccd
+    and (m.geom_pair_type_count[geom_trid_index(BOX, MESH)] > 0 or m.geom_pair_type_count[geom_trid_index(MESH, MESH)] > 0)
+  )
+
+  if has_multiccd_pairs:
+
+    def _check_margin(name, t1, t2, margin):
+      if use_multiccd:
+        raise NotImplementedError(
+          f"{name} has non-zero margin ({margin}) with MULTICCD enabled. Set margin to 0 or disable MULTICCD."
+        )
+      if t1 == BOX and t2 == BOX and not nativeccd_disabled:
+        raise NotImplementedError(
+          f"{name} has non-zero margin ({margin}) with NATIVECCD enabled. Set margin to 0 or disable NATIVECCD."
+        )
+
+    geom_name = lambda g: mujoco.mj_id2name(mjm, mujoco.mjtObj.mjOBJ_GEOM, g) or str(g)
+
+    for idx in np.nonzero(nxn_include & (nxn_pairid_contact == -1))[0]:
+      g1, g2 = int(geom1[idx]), int(geom2[idx])
+      t1, t2 = int(mjm.geom_type[g1]), int(mjm.geom_type[g2])
+      m1, m2 = float(mjm.geom_margin[g1]), float(mjm.geom_margin[g2])
+      if (m1 or m2) and t1 in (BOX, MESH) and t2 in (BOX, MESH):
+        _check_margin(f"geom pair ({geom_name(g1)}, {geom_name(g2)})", t1, t2, (m1, m2))
+
+    for pid in range(mjm.npair):
+      g1, g2 = int(mjm.pair_geom1[pid]), int(mjm.pair_geom2[pid])
+      t1, t2 = int(mjm.geom_type[g1]), int(mjm.geom_type[g2])
+      pm = float(mjm.pair_margin[pid])
+      if pm and t1 in (BOX, MESH) and t2 in (BOX, MESH):
+        _check_margin(f"pair {pid} ({geom_name(g1)}, {geom_name(g2)})", t1, t2, pm)
+
   m.nmaxpolygon = np.append(mjm.mesh_polyvertnum, 0).max()
   m.nmaxmeshdeg = np.append(mjm.mesh_polymapnum, 0).max()
 
@@ -390,9 +429,11 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
           current = []
       else:
         current.append(v)
-    # Pad with zeros if less than 3
+    if len(attr_values) > types._NPLUGINATTR:
-    attr_values += [0.0] * (3 - len(attr_values))
+      raise ValueError(f"Plugin has {len(attr_values)} attributes, which exceeds the maximum of {types._NPLUGINATTR}. ")
-    m.plugin_attr.append(attr_values[:3])
+    # pad with zeros to _NPLUGINATTR
+    attr_values += [0.0] * (types._NPLUGINATTR - len(attr_values))
+    m.plugin_attr.append(attr_values[: types._NPLUGINATTR])
 
   # equality constraint addresses
   m.eq_connect_adr = np.nonzero(mjm.eq_type == types.EqType.CONNECT)[0]
@@ -542,6 +583,15 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
       Madr_ki -= 1
   m.qLD_updates = tuple(wp.array(qLD_updates[i], dtype=wp.vec3i) for i in sorted(qLD_updates))
 
+  # Build concatenated updates for fused kernel
+  all_updates_flat = []
+  level_offsets = [0]
+  for level in sorted(qLD_updates):
+    all_updates_flat.extend(qLD_updates[level])
+    level_offsets.append(len(all_updates_flat))
+  m.qLD_all_updates = all_updates_flat if all_updates_flat else [(0, 0, 0)]
+  m.qLD_level_offsets = level_offsets
+
   # indices for sparse qM_fullm (used in solver)
   m.qM_fullm_i, m.qM_fullm_j = [], []
   for i in range(mjm.nv):
@@ -631,9 +681,168 @@ def _default_njmax(mjm: mujoco.MjModel, mjd: Optional[mujoco.MjData] = None) ->
   return int(valid_sizes[np.searchsorted(valid_sizes, njmax)])
 
 
-def _resolve_batch_size(
+def _body_pair_nnz(mjm: mujoco.MjModel, body1: int, body2: int) -> int:
-    na: int | None, n: int | None, nworld: int, default: int
+  """Returns the number of unique DOFs in the kinematic tree union of two bodies."""
-) -> int:
+  body1 = mjm.body_weldid[body1]
+  body2 = mjm.body_weldid[body2]
+  da1 = mjm.body_dofadr[body1] + mjm.body_dofnum[body1] - 1
+  da2 = mjm.body_dofadr[body2] + mjm.body_dofnum[body2] - 1
+  nnz = 0
+  while da1 >= 0 or da2 >= 0:
+    da = max(da1, da2)
+    if da1 == da:
+      da1 = mjm.dof_parentid[da1]
+    if da2 == da:
+      da2 = mjm.dof_parentid[da2]
+    nnz += 1
+  return nnz
+
+
+def _default_njmax_nnz(mjm: mujoco.MjModel, nconmax: int, njmax: int) -> int:
+  """Returns a heuristic estimate for the number of non-zeros in the sparse constraint Jacobian.
+
+  Assumes all equality, friction, and limit constraints are active and computes
+  their non-zeros. For contacts, assumes njmax contact rows at the maximum
+  body-pair non-zeros from all enabled collision pairs.
+
+  Args:
+    mjm: The model containing kinematic and dynamic information (host).
+    nconmax: Maximum number of contacts per world.
+    njmax: Maximum number of constraint rows per world.
+
+  Returns:
+    Estimated number of non-zeros in the constraint Jacobian.
+  """
+  total_nnz = 0
+
+  def _eq_bodies(i):
+    """Returns body pair for equality constraint i."""
+    obj1id, obj2id = mjm.eq_obj1id[i], mjm.eq_obj2id[i]
+    if mjm.eq_objtype[i] == mujoco.mjtObj.mjOBJ_SITE:
+      return mjm.site_bodyid[obj1id], mjm.site_bodyid[obj2id]
+    return obj1id, obj2id
+
+  # equality constraints (assume all active)
+  for i in range(mjm.neq):
+    eq_type = mjm.eq_type[i]
+
+    if eq_type == mujoco.mjtEq.mjEQ_CONNECT:
+      total_nnz += 3 * _body_pair_nnz(mjm, *_eq_bodies(i))
+
+    elif eq_type == mujoco.mjtEq.mjEQ_WELD:
+      total_nnz += 6 * _body_pair_nnz(mjm, *_eq_bodies(i))
+
+    elif eq_type == mujoco.mjtEq.mjEQ_JOINT:
+      total_nnz += 2 if mjm.eq_obj2id[i] >= 0 else 1
+
+    elif eq_type == mujoco.mjtEq.mjEQ_TENDON:
+      obj1id = mjm.eq_obj1id[i]
+      obj2id = mjm.eq_obj2id[i]
+      rownnz1 = mjm.ten_J_rownnz[obj1id] if obj1id < mjm.ntendon else 0
+      if obj2id >= 0 and obj2id < mjm.ntendon:
+        rowadr1 = mjm.ten_J_rowadr[obj1id]
+        rowadr2 = mjm.ten_J_rowadr[obj2id]
+        rownnz2 = mjm.ten_J_rownnz[obj2id]
+        cols = set()
+        for j in range(rownnz1):
+          cols.add(mjm.ten_J_colind[rowadr1 + j])
+        for j in range(rownnz2):
+          cols.add(mjm.ten_J_colind[rowadr2 + j])
+        total_nnz += len(cols)
+      else:
+        total_nnz += rownnz1
+
+    elif eq_type == mujoco.mjtEq.mjEQ_FLEX:
+      obj1id = mjm.eq_obj1id[i]
+      if obj1id < mjm.nflex:
+        edge_start = mjm.flex_edgeadr[obj1id]
+        edge_count = mjm.flex_edgenum[obj1id]
+        for e in range(edge_count):
+          total_nnz += mjm.flexedge_J_rownnz[edge_start + e]
+
+  # friction constraints
+  total_nnz += (mjm.dof_frictionloss > 0).sum()
+  for i in range(mjm.ntendon):
+    if mjm.tendon_frictionloss[i] > 0:
+      total_nnz += mjm.ten_J_rownnz[i]
+
+  # limit constraints (assume all active)
+  for i in range(mjm.njnt):
+    if mjm.jnt_limited[i]:
+      jnt_type = mjm.jnt_type[i]
+      if jnt_type == mujoco.mjtJoint.mjJNT_BALL:
+        total_nnz += 3
+      elif jnt_type in (mujoco.mjtJoint.mjJNT_SLIDE, mujoco.mjtJoint.mjJNT_HINGE):
+        total_nnz += 1
+  for i in range(mjm.ntendon):
+    if mjm.tendon_limited[i]:
+      total_nnz += mjm.ten_J_rownnz[i]
+
+  # contact constraints: njmax rows at max body-pair non-zeros
+  max_contact_nnz = 0
+
+  # contact pairs
+  for i in range(mjm.npair):
+    g1, g2 = mjm.pair_geom1[i], mjm.pair_geom2[i]
+    b1, b2 = mjm.geom_bodyid[g1], mjm.geom_bodyid[g2]
+    max_contact_nnz = max(max_contact_nnz, _body_pair_nnz(mjm, b1, b2))
+
+  # filter geom-geom pairs (unique body pairs, filtered)
+  body_pair_seen = set()
+  for i in range(mjm.ngeom):
+    bi = mjm.geom_bodyid[i]
+    cti, cai = mjm.geom_contype[i], mjm.geom_conaffinity[i]
+    for j in range(i + 1, mjm.ngeom):
+      bj = mjm.geom_bodyid[j]
+      if bi == bj:
+        continue
+      if mjm.body_weldid[bi] == 0 and mjm.body_weldid[bj] == 0:
+        continue
+      bp = (min(bi, bj), max(bi, bj))
+      if bp in body_pair_seen:
+        continue
+      ctj, caj = mjm.geom_contype[j], mjm.geom_conaffinity[j]
+      if not ((cti & caj) or (ctj & cai)):
+        continue
+      body_pair_seen.add(bp)
+      max_contact_nnz = max(max_contact_nnz, _body_pair_nnz(mjm, bi, bj))
+
+  # flex vertex contacts
+  for fi in range(mjm.nflex):
+    fct = mjm.flex_contype[fi]
+    fca = mjm.flex_conaffinity[fi]
+
+    vert_start = mjm.flex_vertadr[fi]
+    vert_count = mjm.flex_vertnum[fi]
+    flex_bodies = {mjm.flex_vertbodyid[vert_start + v] for v in range(vert_count)}
+
+    geom_bodies = set()
+    for g in range(mjm.ngeom):
+      ct, ca = mjm.geom_contype[g], mjm.geom_conaffinity[g]
+      if (fct & ca) or (ct & fca):
+        geom_bodies.add(mjm.geom_bodyid[g])
+
+    for fb in flex_bodies:
+      for gb in geom_bodies:
+        if fb != gb:
+          max_contact_nnz = max(max_contact_nnz, _body_pair_nnz(mjm, fb, gb))
+
+    # flex self-collision
+    if mjm.flex_selfcollide[fi]:
+      flex_body_list = sorted(flex_bodies)
+      for idx1 in range(len(flex_body_list)):
+        for idx2 in range(idx1 + 1, len(flex_body_list)):
+          max_contact_nnz = max(
+            max_contact_nnz,
+            _body_pair_nnz(mjm, flex_body_list[idx1], flex_body_list[idx2]),
+          )
+
+  total_nnz += njmax * max_contact_nnz
+
+  return int(min(max(total_nnz, 1), njmax * mjm.nv))
+
+
+def _resolve_batch_size(na: int | None, n: int | None, nworld: int, default: int) -> int:
   if na is not None:
     return na
   if n is not None:
@@ -647,6 +856,7 @@ def make_data(
   nconmax: Optional[int] = None,
   nccdmax: Optional[int] = None,
   njmax: Optional[int] = None,
+  njmax_nnz: Optional[int] = None,
   naconmax: Optional[int] = None,
   naccdmax: Optional[int] = None,
 ) -> types.Data:
@@ -660,6 +870,7 @@ def make_data(
     nccdmax: Number of CCD contacts to allocate per world. Same semantics as nconmax.
     njmax: Number of constraints to allocate per world. Constraint arrays are
            batched by world: no world may have more than njmax constraints.
+    njmax_nnz: Number of non-zeros in constraint Jacobian (sparse). Defaults to njmax * nv.
     naconmax: Number of contacts to allocate for all worlds. Overrides nconmax.
     naccdmax: Maximum number of CCD contacts. Defaults to naconmax.
 
@@ -709,20 +920,32 @@ def make_data(
   sizes["naconmax"] = naconmax
   sizes["njmax"] = njmax
 
+  if njmax_nnz is None:
+    if is_sparse(mjm):
+      njmax_nnz = _default_njmax_nnz(mjm, nconmax, njmax)
+    else:
+      njmax_nnz = njmax * mjm.nv
+
   contact = types.Contact(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Contact)})
+  contact.efc_address = wp.array(np.full((naconmax, sizes["nmaxpyramid"]), -1, dtype=int), dtype=int)
   efc = types.Constraint(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Constraint)})
 
-  if SPARSE_CONSTRAINT_JACOBIAN:
+  if is_sparse(mjm):
     efc.J_rownnz = wp.zeros((nworld, njmax), dtype=int)
     efc.J_rowadr = wp.zeros((nworld, njmax), dtype=int)
-    efc.J_colind = wp.zeros((nworld, 1, njmax * mjm.nv), dtype=int)
+    efc.J_colind = wp.zeros((nworld, 1, njmax_nnz), dtype=int)
-    efc.J = wp.zeros((nworld, 1, njmax * mjm.nv), dtype=float)
+    efc.J = wp.zeros((nworld, 1, njmax_nnz), dtype=float)
   else:
     efc.J_rownnz = wp.zeros((nworld, 0), dtype=int)
     efc.J_rowadr = wp.zeros((nworld, 0), dtype=int)
     efc.J_colind = wp.zeros((nworld, 0, 0), dtype=int)
     efc.J = wp.zeros((nworld, sizes["njmax_pad"], sizes["nv_pad"]), dtype=float)
 
+  contact_kwargs = {}
+  for f in dataclasses.fields(types.Contact):
+    contact_kwargs[f.name] = _create_array(None, f.type, sizes)
+  contact = types.Contact(**contact_kwargs)
+
   # world body and static geom (attached to the world) poses are precomputed
   # this speeds up scenes with many static geoms (e.g. terrains)
   # TODO(team): remove this when we introduce dof islands + sleeping
@@ -734,9 +957,7 @@ def make_data(
   mocap_id = mjm.body_mocapid[mocap_body]
 
   d_kwargs = {
-      "qpos": wp.array(
+    "qpos": wp.array(np.tile(mjm.qpos0, nworld), shape=(nworld, mjm.nq), dtype=float),
-          np.tile(mjm.qpos0, nworld), shape=(nworld, mjm.nq), dtype=float
-      ),
     "contact": contact,
     "efc": efc,
     "nworld": nworld,
@@ -744,52 +965,23 @@ def make_data(
     "naccdmax": naccdmax,
     "njmax": njmax,
     "njmax_pad": sizes["njmax_pad"],
+    "njmax_nnz": njmax_nnz,
     "qM": None,
     "qLD": None,
     # world body
-      "xquat": wp.array(
+    "xquat": wp.array(np.tile(mjd.xquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat),
-          np.tile(mjd.xquat, (nworld, 1)),
+    "xmat": wp.array(np.tile(mjd.xmat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.mat33),
-          shape=(nworld, mjm.nbody),
+    "ximat": wp.array(np.tile(mjd.ximat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.mat33),
-          dtype=wp.quat,
-      ),
-      "xmat": wp.array(
-          np.tile(mjd.xmat, (nworld, 1)),
-          shape=(nworld, mjm.nbody),
-          dtype=wp.mat33,
-      ),
-      "ximat": wp.array(
-          np.tile(mjd.ximat, (nworld, 1)),
-          shape=(nworld, mjm.nbody),
-          dtype=wp.mat33,
-      ),
     # static geoms
-      "geom_xpos": wp.array(
+    "geom_xpos": wp.array(np.tile(mjd.geom_xpos, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.vec3),
-          np.tile(mjd.geom_xpos, (nworld, 1)),
+    "geom_xmat": wp.array(np.tile(mjd.geom_xmat, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.mat33),
-          shape=(nworld, mjm.ngeom),
-          dtype=wp.vec3,
-      ),
-      "geom_xmat": wp.array(
-          np.tile(mjd.geom_xmat, (nworld, 1)),
-          shape=(nworld, mjm.ngeom),
-          dtype=wp.mat33,
-      ),
     # mocap
-      "mocap_pos": wp.array(
+    "mocap_pos": wp.array(np.tile(mjm.body_pos[mocap_body[mocap_id]], (nworld, 1)), shape=(nworld, mjm.nmocap), dtype=wp.vec3),
-          np.tile(mjm.body_pos[mocap_body[mocap_id]], (nworld, 1)),
-          shape=(nworld, mjm.nmocap),
-          dtype=wp.vec3,
-      ),
     "mocap_quat": wp.array(
-          np.tile(mjm.body_quat[mocap_body[mocap_id]], (nworld, 1)),
+      np.tile(mjm.body_quat[mocap_body[mocap_id]], (nworld, 1)), shape=(nworld, mjm.nmocap), dtype=wp.quat
-          shape=(nworld, mjm.nmocap),
-          dtype=wp.quat,
     ),
     # equality constraints
-      "eq_active": wp.array(
+    "eq_active": wp.array(np.tile(mjm.eq_active0.astype(bool), (nworld, 1)), shape=(nworld, mjm.neq), dtype=bool),
-          np.tile(mjm.eq_active0.astype(bool), (nworld, 1)),
-          shape=(nworld, mjm.neq),
-          dtype=bool,
-      ),
     # island arrays
     "nisland": None,
     "tree_island": None,
@@ -822,6 +1014,7 @@ def put_data(
   nconmax: Optional[int] = None,
   nccdmax: Optional[int] = None,
   njmax: Optional[int] = None,
+  njmax_nnz: Optional[int] = None,
   naconmax: Optional[int] = None,
   naccdmax: Optional[int] = None,
 ) -> types.Data:
@@ -836,6 +1029,7 @@ def put_data(
     nccdmax: Number of CCD contacts to allocate per world. Same semantics as nconmax.
     njmax: Number of constraints to allocate per world.  Constraint arrays are
            batched by world: no world may have more than njmax constraints.
+    njmax_nnz: Number of non-zeros in constraint Jacobian (sparse). Defaults to njmax * nv.
     naconmax: Number of contacts to allocate for all worlds. Overrides nconmax.
     naccdmax: Maximum number of CCD contacts. Defaults to naconmax.
 
@@ -898,6 +1092,12 @@ def put_data(
   sizes["naconmax"] = naconmax
   sizes["njmax"] = njmax
 
+  if njmax_nnz is None:
+    if is_sparse(mjm):
+      njmax_nnz = _default_njmax_nnz(mjm, nconmax, njmax)
+    else:
+      njmax_nnz = njmax * mjm.nv
+
   # ensure static geom positions are computed
   # TODO: remove once MjData creation semantics are fixed
   mujoco.mj_kinematics(mjm, mjd)
@@ -915,7 +1115,7 @@ def put_data(
 
   contact = types.Contact(**contact_kwargs)
 
-  contact.efc_address = np.zeros((naconmax, sizes["nmaxpyramid"]), dtype=int)
+  contact.efc_address = np.full((naconmax, sizes["nmaxpyramid"]), -1, dtype=int)
   for i in range(mjd.ncon):
     efc_address = mjd.contact.efc_address[i]
     if efc_address == -1:
@@ -945,43 +1145,28 @@ def put_data(
 
   efc = types.Constraint(**efc_kwargs)
 
-  if SPARSE_CONSTRAINT_JACOBIAN:
+  if is_sparse(mjm):
-    # TODO(team): process efc_J sparsity structure for nv row shift
+    J_rownnz = np.zeros(njmax, dtype=np.int32)
-    efc.J_rownnz = wp.array(
+    J_rowadr = np.zeros(njmax, dtype=np.int32)
-        np.full((nworld, njmax), mjm.nv, dtype=int), dtype=int
+    J_colind = np.zeros(njmax_nnz, dtype=np.int32)
-    )
+    J = np.zeros(njmax_nnz, dtype=np.float64)
-    efc.J_rowadr = wp.array(
-        np.tile(
-            np.arange(0, njmax * mjm.nv, mjm.nv)
-            if mjm.nv
-            else np.zeros(njmax, dtype=int),
-            (nworld, 1),
-        ),
-        dtype=int,
-    )
-    efc.J_colind = wp.array(
-        np.tile(np.arange(mjm.nv), (nworld, njmax)).reshape((nworld, 1, -1)),
-        dtype=int,
-    )
-
-    mj_efc_J = np.zeros((mjd.nefc, mjm.nv))
     if mjd.nefc:
       if mujoco.mj_isSparse(mjm):
-        mujoco.mju_sparse2dense(
+        J_rownnz[: mjd.nefc] = mjd.efc_J_rownnz[: mjd.nefc]
-            mj_efc_J,
+        J_rowadr[: mjd.nefc] = mjd.efc_J_rowadr[: mjd.nefc]
-            mjd.efc_J,
+        nnz = int(mjd.efc_J_rownnz[: mjd.nefc].sum())
-            mjd.efc_J_rownnz,
+        J_colind[:nnz] = mjd.efc_J_colind[:nnz]
-            mjd.efc_J_rowadr,
+        J[:nnz] = mjd.efc_J[:nnz]
-            mjd.efc_J_colind,
-        )
       else:
-        mj_efc_J = mjd.efc_J.reshape((mjd.nefc, mjm.nv))
+        dense_J = mjd.efc_J.reshape((-1, mjm.nv))[: mjd.nefc]
-    efc_J = np.zeros((njmax, mjm.nv), dtype=float)
+        mujoco.mju_dense2sparse(
-    efc_J[: mjd.nefc, : mjm.nv] = mj_efc_J
+          J[: mjd.nefc * mjm.nv], dense_J, J_rownnz[: mjd.nefc], J_rowadr[: mjd.nefc], J_colind[: mjd.nefc * mjm.nv]
-    efc.J = wp.array(
-        np.tile(efc_J.reshape(-1), (nworld, 1, 1)).reshape((nworld, 1, -1)),
-        dtype=float,
         )
+
+    efc.J_rownnz = wp.array(np.tile(J_rownnz, (nworld, 1)), dtype=int)
+    efc.J_rowadr = wp.array(np.tile(J_rowadr, (nworld, 1)), dtype=int)
+    efc.J_colind = wp.array(np.tile(J_colind, (nworld, 1)).reshape((nworld, 1, -1)), dtype=int)
+    efc.J = wp.array(np.tile(J, (nworld, 1)).reshape((nworld, 1, -1)), dtype=float)
   else:
     efc.J_rownnz = wp.zeros((nworld, 0), dtype=int)
     efc.J_rowadr = wp.zeros((nworld, 0), dtype=int)
@@ -990,13 +1175,7 @@ def put_data(
     mj_efc_J = np.zeros((mjd.nefc, mjm.nv))
     if mjd.nefc:
       if mujoco.mj_isSparse(mjm):
-        mujoco.mju_sparse2dense(
+        mujoco.mju_sparse2dense(mj_efc_J, mjd.efc_J, mjd.efc_J_rownnz, mjd.efc_J_rowadr, mjd.efc_J_colind)
-            mj_efc_J,
-            mjd.efc_J,
-            mjd.efc_J_rownnz,
-            mjd.efc_J_rowadr,
-            mjd.efc_J_colind,
-        )
       else:
         mj_efc_J = mjd.efc_J.reshape((mjd.nefc, mjm.nv))
     efc_J = np.zeros((nworld, sizes["njmax_pad"], sizes["nv_pad"]), dtype=float)
@@ -1012,11 +1191,11 @@ def put_data(
     "naccdmax": naccdmax,
     "njmax": njmax,
     "njmax_pad": sizes["njmax_pad"],
+    "njmax_nnz": njmax_nnz,
     # fields set after initialization:
     "solver_niter": None,
     "qM": None,
     "qLD": None,
-      "ten_J": None,
     "nacon": None,
     # island arrays
     "nisland": None,
@@ -1050,29 +1229,6 @@ def put_data(
   d.nisland = wp.array(np.full(nworld, mjd.nisland), dtype=int)
   d.tree_island = wp.array(np.tile(mjd.tree_island, (nworld, 1)), dtype=int)
 
-  ten_J = np.zeros((mjm.ntendon, mjm.nv))
-  if mujoco.mj_isSparse(mjm) or check_version("mujoco>=3.5.1.dev872479828"):
-    if mjm.ntendon:
-      if check_version("mujoco>=3.5.1.dev875093374"):
-        mujoco.mju_sparse2dense(
-            ten_J,
-            mjd.ten_J.reshape(-1),
-            mjm.ten_J_rownnz,
-            mjm.ten_J_rowadr,
-            mjm.ten_J_colind.reshape(-1),
-        )
-      else:
-        mujoco.mju_sparse2dense(
-            ten_J,
-            mjd.ten_J.reshape(-1),
-            mjd.ten_J_rownnz,
-            mjd.ten_J_rowadr,
-            mjd.ten_J_colind.reshape(-1),
-        )
-  else:
-    ten_J = mjd.ten_J.reshape((mjm.ntendon, mjm.nv))
-  d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float)
-
   d.nacon = wp.array([mjd.ncon * nworld], dtype=int)
 
   return d
@@ -1233,7 +1389,7 @@ def get_data_into(
     mujoco.mj_factorM(mjm, result)
 
   if nefc > 0:
-    if SPARSE_CONSTRAINT_JACOBIAN:
+    if is_sparse(mjm):
       efc_J = np.zeros((nefc, mjm.nv))
       mujoco.mju_sparse2dense(
         efc_J,
@@ -1276,24 +1432,7 @@ def get_data_into(
 
   # tendon
   result.ten_length[:] = d.ten_length.numpy()[world_id]
-  if check_version("mujoco>=3.5.1.dev869712136"):
+  if mjm.ntendon > 0:
-    ten_J = d.ten_J.numpy()[world_id]
-    if check_version("mujoco>=3.5.1.dev875093374"):
-      ten_J_rownnz = mjm.ten_J_rownnz
-      ten_J_rowadr = mjm.ten_J_rowadr
-      ten_J_colind = mjm.ten_J_colind.reshape(-1)
-    else:
-      ten_J_rownnz = result.ten_J_rownnz
-      ten_J_rowadr = result.ten_J_rowadr
-      ten_J_colind = result.ten_J_colind.reshape(-1)
-    mujoco.mju_dense2sparse(
-        result.ten_J,
-        ten_J,
-        ten_J_rownnz,
-        ten_J_rowadr,
-        ten_J_colind,
-    )
-  else:
     result.ten_J[:] = d.ten_J.numpy()[world_id]
   result.ten_wrapadr[:] = d.ten_wrapadr.numpy()[world_id]
   result.ten_wrapnum[:] = d.ten_wrapnum.numpy()[world_id]
@@ -1428,12 +1567,8 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     mocapid = body_mocapid[bodyid]
 
     if mocapid >= 0:
-      mocap_pos_out[worldid, mocapid] = body_pos[
+      mocap_pos_out[worldid, mocapid] = body_pos[worldid % body_pos.shape[0], bodyid]
-          worldid % body_pos.shape[0], bodyid
+      mocap_quat_out[worldid, mocapid] = body_quat[worldid % body_quat.shape[0], bodyid]
-      ]
-      mocap_quat_out[worldid, mocapid] = body_quat[
-          worldid % body_quat.shape[0], bodyid
-      ]
 
   @wp.kernel(module="unique", enable_backward=False)
   def reset_contact(
@@ -1453,6 +1588,8 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     contact_solimp_out: wp.array(dtype=types.vec5),
     contact_dim_out: wp.array(dtype=int),
     contact_geom_out: wp.array(dtype=wp.vec2i),
+    contact_flex_out: wp.array(dtype=wp.vec2i),
+    contact_vert_out: wp.array(dtype=wp.vec2i),
     contact_efc_address_out: wp.array2d(dtype=int),
     contact_worldid_out: wp.array(dtype=int),
     contact_type_out: wp.array(dtype=int),
@@ -1479,8 +1616,10 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     contact_solimp_out[conid] = types.vec5(0.0, 0.0, 0.0, 0.0, 0.0)
     contact_dim_out[conid] = 0
     contact_geom_out[conid] = wp.vec2i(0, 0)
+    contact_flex_out[conid] = wp.vec2i(0, 0)
+    contact_vert_out[conid] = wp.vec2i(0, 0)
     for i in range(nefcaddress):
-      contact_efc_address_out[conid, i] = 0
+      contact_efc_address_out[conid, i] = -1
     contact_worldid_out[conid] = 0
     contact_type_out[conid] = 0
     contact_geomcollisionid_out[conid] = 0
@@ -1519,6 +1658,8 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
       d.contact.solimp,
       d.contact.dim,
       d.contact.geom,
+      d.contact.flex,
+      d.contact.vert,
       d.contact.efc_address,
       d.contact.worldid,
       d.contact.type,
@@ -1812,28 +1953,45 @@ def _finalize_body_invweight0(
 @wp.kernel
 def _copy_tendon_jacobian(
   tenid_target: int,
-  ten_J_in: wp.array3d(dtype=float),
+  ten_J_rownnz: wp.array(dtype=int),
+  ten_J_rowadr: wp.array(dtype=int),
+  ten_J_colind: wp.array(dtype=int),
+  ten_J_in: wp.array2d(dtype=float),
   ten_J_vec_out: wp.array2d(dtype=float),
 ):
   worldid = wp.tid()
   nv = ten_J_in.shape[2]
-  for i in range(nv):
+  rownnz = ten_J_rownnz[tenid_target]
-    ten_J_vec_out[worldid, i] = ten_J_in[worldid, tenid_target, i]
+  rowadr = ten_J_rowadr[tenid_target]
+  for i in range(rownnz):
+    colind = ten_J_colind[rowadr + i]
+    ten_J_vec_out[worldid, colind] = ten_J_in[worldid, rowadr + i]
 
 
 @wp.kernel
 def _compute_tendon_dot_product(
+  # Model:
+  ten_J_rownnz: wp.array(dtype=int),
+  ten_J_rowadr: wp.array(dtype=int),
+  ten_J_colind: wp.array(dtype=int),
+  # In:
   tenid_target: int,
-  nv: int,
+  ten_J_in: wp.array2d(dtype=float),
-  ten_J_in: wp.array3d(dtype=float),
   result_vec_in: wp.array2d(dtype=float),
+  # Out:
   tendon_invweight0_out: wp.array2d(dtype=float),
 ):
   worldid = wp.tid()
   tendon_invweight0_id = worldid % tendon_invweight0_out.shape[0]
   dot_prod = float(0.0)
-  for i in range(nv):
+
-    dot_prod += ten_J_in[worldid, tenid_target, i] * result_vec_in[worldid, i]
+  rownnz = ten_J_rownnz[tenid_target]
+  rowadr = ten_J_rowadr[tenid_target]
+  for i in range(rownnz):
+    sparseid = rowadr + i
+    colind = ten_J_colind[sparseid]
+    dot_prod += ten_J_in[worldid, sparseid] * result_vec_in[worldid, colind]
+
   tendon_invweight0_out[tendon_invweight0_id, tenid_target] = dot_prod
 
 
@@ -2167,16 +2325,22 @@ def set_const_0(m: types.Model, d: types.Data):
 
   # tendon_invweight0[t] = J_t * inv(M) * J_t'
   if m.ntendon > 0:
-    ten_J_vec = wp.zeros((d.nworld, m.nv), dtype=float)
+    ten_J_vec = wp.empty((d.nworld, m.nv), dtype=float)
-    ten_result_vec = wp.zeros((d.nworld, m.nv), dtype=float)
+    ten_result_vec = wp.empty((d.nworld, m.nv), dtype=float)
 
     for tenid in range(m.ntendon):
-      wp.launch(_copy_tendon_jacobian, dim=d.nworld, inputs=[tenid, d.ten_J], outputs=[ten_J_vec])
+      ten_J_vec.zero_()
+      wp.launch(
+        _copy_tendon_jacobian,
+        dim=d.nworld,
+        inputs=[tenid, m.ten_J_rownnz, m.ten_J_rowadr, m.ten_J_colind, d.ten_J],
+        outputs=[ten_J_vec],
+      )
       smooth.solve_m(m, d, ten_result_vec, ten_J_vec)
       wp.launch(
         _compute_tendon_dot_product,
         dim=d.nworld,
-        inputs=[tenid, m.nv, d.ten_J, ten_result_vec],
+        inputs=[m.ten_J_rownnz, m.ten_J_rowadr, m.ten_J_colind, tenid, d.ten_J, ten_result_vec],
         outputs=[m.tendon_invweight0],
       )
 
@@ -2203,13 +2367,7 @@ def set_const_0(m: types.Model, d: types.Data):
       wp.launch(
         _copy_actuator_moment,
         dim=d.nworld,
-          inputs=[
+        inputs=[actid, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment],
-              actid,
-              d.moment_rownnz,
-              d.moment_rowadr,
-              d.moment_colind,
-              d.actuator_moment,
-          ],
         outputs=[act_moment_vec],
       )
       smooth.solve_m(m, d, act_result_vec, act_moment_vec)
@@ -2255,16 +2413,12 @@ def set_const(m: types.Model, d: types.Data):
     Field                            | Notes
     ---------------------------------|----------------------------------------------
     qpos0, qpos_spring               |
-    body_mass, body_inertia,         | Mass and inertia are usually scaled
+    body_mass, body_inertia,         | Mass and inertia are usually scaled together
-    together
     body_ipos, body_iquat            | since inertia is sum(m * r^2).
-    body_pos, body_quat              | Unsafe for static bodies (invalidates
+    body_pos, body_quat              | Unsafe for static bodies (invalidates BVH).
-    BVH).
+    body_gravcomp                    | If changing from 0 to >0 bodies, required.
-    body_gravcomp                    | If changing from 0 to >0 bodies,
-    required.
     dof_armature                     |
-    eq_data                          | For connect/weld, offsets computed if not
+    eq_data                          | For connect/weld, offsets computed if not set.
-    set.
     hfield_size                      |
     tendon_stiffness, tendon_damping | Only if changing from/to zero.
     actuator_gainprm, actuator_biasprm | For position actuators with dampratio.
@@ -2486,6 +2640,7 @@ def create_render_context(
   cam_res: list[tuple[int, int]] | tuple[int, int] | None = None,
   render_rgb: list[bool] | bool | None = None,
   render_depth: list[bool] | bool | None = None,
+  render_seg: list[bool] | bool | None = None,
   use_textures: bool = True,
   use_shadows: bool = False,
   enabled_geom_groups: list[int] = [0, 1, 2],
@@ -2502,6 +2657,8 @@ def create_render_context(
              MuJoCo model values.
     render_rgb: Whether to render RGB images. If None, uses the MuJoCo model values.
     render_depth: Whether to render depth images. If None, uses the MuJoCo model values.
+    render_seg: Whether to render segmentation (per-pixel geom IDs). If None,
+      uses the MuJoCo model values.
     use_textures: Whether to use textures.
     use_shadows: Whether to use shadows.
     enabled_geom_groups: The geom groups to render.
@@ -2517,10 +2674,13 @@ def create_render_context(
   mjd = mujoco.MjData(mjm)
   mujoco.mj_forward(mjm, mjd)
 
-  # TODO(team): remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
+  constructor = "sah"
-  if use_textures and not hasattr(wp, "Texture2D"):
+  if check_version("warp>=1.13.0.dev20260325"):
-    warnings.warn("Textures require warp >= 1.12. Disabling textures.")
+    # TODO: The cubql constructor and is_cubql_available exist only in
-    use_textures = False
+    # recent Warp 1.13+ builds, modify this after warp is updated to 1.13+.
+    _cubql_avail = getattr(wp, "is_cubql_available", None)
+    if callable(_cubql_avail) and _cubql_avail():
+      constructor = "cubql"
 
   # Mesh BVHs
   nmesh = mjm.nmesh
@@ -2534,7 +2694,7 @@ def create_render_context(
   mesh_bounds_size = [wp.vec3(0.0, 0.0, 0.0) for _ in range(nmesh)]
 
   for mid in used_mesh_id:
-    mesh, half = bvh.build_mesh_bvh(mjm, mid)
+    mesh, half = bvh.build_mesh_bvh(mjm, mid, constructor=constructor)
     mesh_registry[mesh.id] = mesh
     mesh_bvh_id[mid] = mesh.id
     mesh_bounds_size[mid] = half
@@ -2551,7 +2711,7 @@ def create_render_context(
   hfield_bounds_size = [wp.vec3(0.0, 0.0, 0.0) for _ in range(nhfield)]
 
   for hid in used_hfield_id:
-    hmesh, hhalf = bvh.build_hfield_bvh(mjm, hid)
+    hmesh, hhalf = bvh.build_hfield_bvh(mjm, hid, constructor=constructor)
     hfield_registry[hmesh.id] = hmesh
     hfield_bvh_id[hid] = hmesh.id
     hfield_bounds_size[hid] = hhalf
@@ -2560,65 +2720,33 @@ def create_render_context(
   hfield_bounds_size_arr = wp.array(hfield_bounds_size, dtype=wp.vec3)
 
   # Flex BVHs
-  flex_bvh_id = wp.uint64(0)
-  flex_group_root = wp.zeros(nworld, dtype=int)
-  flex_mesh = None
-  flex_face_point = None
-  flex_elemdataadr = None
-  flex_shell = None
-  flex_shelldataadr = None
-  flex_faceadr = None
-  flex_nface = 0
-  flex_radius = None
-  flex_workadr = None
-  flex_worknum = None
-  flex_nwork = 0
-
-  if mjm.nflex > 0:
-    (
-      fmesh,
-      face_point,
-      flex_group_roots,
-      flex_shell_data,
-      flex_faceadr_data,
-      flex_nface,
-    ) = bvh.build_flex_bvh(mjm, mjd, nworld)
-
-    flex_mesh = fmesh
-    flex_bvh_id = fmesh.id
-    flex_face_point = face_point
-    flex_group_root = flex_group_roots
-    flex_elemdataadr = wp.array(mjm.flex_elemdataadr, dtype=int)
-    flex_shell = flex_shell_data
-    flex_shelldataadr = wp.array(mjm.flex_shelldataadr, dtype=int)
-    flex_faceadr = wp.array(flex_faceadr_data, dtype=int)
-    flex_radius = wp.array(mjm.flex_radius, dtype=float)
-
-    # precompute work item layout for unified refit kernel
   nflex = mjm.nflex
-    workadr = np.zeros(nflex, dtype=np.int32)
+  flex_registry = {}
-    worknum = np.zeros(nflex, dtype=np.int32)
+
-    cumsum = 0
+  # Scene BVH flex primitives: 1D → one capsule per edge, 2D/3D → one box per flex
+  flex_geom_flexid = []
+  flex_geom_edgeid = []
+  flex_bvh_id = np.full(nflex, 0, dtype=wp.uint64)
+  flex_group_root = np.zeros((nflex, nworld), dtype=int)
+
   for f in range(nflex):
-      workadr[f] = cumsum
+    if mjm.flex_dim[f] == 1:
-      if mjm.flex_dim[f] == 2:
+      edge_adr = mjm.flex_edgeadr[f]
-        worknum[f] = mjm.flex_elemnum[f] + mjm.flex_shellnum[f]
+      flex_geom_flexid.extend([f] * mjm.flex_edgenum[f])
+      flex_geom_edgeid.extend([edge_adr + e for e in range(mjm.flex_edgenum[f])])
+      flex_group_root[f] = np.zeros(nworld, dtype=int)
     else:
-        worknum[f] = mjm.flex_shellnum[f]
+      flex_geom_flexid.append(f)
-      cumsum += worknum[f]
+      flex_geom_edgeid.append(-1)
-    flex_workadr = wp.array(workadr, dtype=int)
+      fmesh, group_root = bvh.build_flex_bvh(mjm, mjd, nworld, f)
-    flex_worknum = wp.array(worknum, dtype=int)
+      flex_registry[f] = fmesh
-    flex_nwork = int(cumsum)
+      flex_bvh_id[f] = fmesh.id
+      flex_group_root[f] = group_root.numpy()
 
   textures_registry = []
-  # TODO: remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
-  if hasattr(wp, "Texture2D"):
   for i in range(mjm.ntex):
     textures_registry.append(render_util.create_warp_texture(mjm, i))
   textures = wp.array(textures_registry, dtype=wp.Texture2D)
-  else:
-    # Dummy array when texture support isn't available (warp < 1.12)
-    textures = wp.zeros(1, dtype=int)
 
   # Filter active cameras
   if cam_active is not None:
@@ -2642,31 +2770,31 @@ def create_render_context(
   cam_res_arr = wp.array(active_cam_res, dtype=wp.vec2i)
 
   if render_rgb is None:
-    render_rgb = [
+    render_rgb = [mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_RGB for i in active_cam_indices]
-        mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_RGB
-        for i in active_cam_indices
-    ]
   elif isinstance(render_rgb, bool):
     render_rgb = [render_rgb] * ncam
 
   if render_depth is None:
-    render_depth = [
+    render_depth = [mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_DEPTH for i in active_cam_indices]
-        mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_DEPTH
-        for i in active_cam_indices
-    ]
   if isinstance(render_depth, bool):
     render_depth = [render_depth] * ncam
 
-  assert len(render_rgb) == ncam and len(render_depth) == ncam, (
+  if render_seg is None:
-      "render_rgb and render_depth must be a bool or a list of bools with"
+    render_seg = [mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_SEG for i in active_cam_indices]
-      f" length {ncam}"
+  elif isinstance(render_seg, bool):
+    render_seg = [render_seg] * ncam
+
+  assert len(render_rgb) == ncam and len(render_depth) == ncam and len(render_seg) == ncam, (
+    f"render_rgb, render_depth, and render_seg must be a bool or a list of bools with length {ncam}"
   )
 
   rgb_adr = -1 * np.ones(ncam, dtype=int)
   depth_adr = -1 * np.ones(ncam, dtype=int)
+  seg_adr = -1 * np.ones(ncam, dtype=int)
   cam_res_np = cam_res_arr.numpy()
   ri = 0
   di = 0
+  si = 0
   total = 0
 
   for idx in range(ncam):
@@ -2676,6 +2804,9 @@ def create_render_context(
     if render_depth[idx]:
       depth_adr[idx] = di
       di += cam_res_np[idx][0] * cam_res_np[idx][1]
+    if render_seg[idx]:
+      seg_adr[idx] = si
+      si += cam_res_np[idx][0] * cam_res_np[idx][1]
 
     total += cam_res_np[idx][0] * cam_res_np[idx][1]
 
@@ -2729,26 +2860,20 @@ def create_render_context(
     hfield_registry=hfield_registry,
     hfield_bvh_id=hfield_bvh_id_arr,
     hfield_bounds_size=hfield_bounds_size_arr,
-    flex_mesh=flex_mesh,
+    flex_mesh_registry=flex_registry,
     flex_rgba=wp.array(mjm.flex_rgba, dtype=wp.vec4),
-    flex_bvh_id=flex_bvh_id,
+    flex_bvh_id=wp.array(flex_bvh_id, dtype=wp.uint64),
-    flex_face_point=flex_face_point,
+    flex_group_root=wp.array(flex_group_root, dtype=int),
-    flex_faceadr=flex_faceadr,
-    flex_nface=flex_nface,
-    flex_nwork=flex_nwork,
-    flex_group_root=flex_group_root,
-    flex_elemdataadr=flex_elemdataadr,
-    flex_shell=flex_shell,
-    flex_shelldataadr=flex_shelldataadr,
-    flex_radius=flex_radius,
-    flex_workadr=flex_workadr,
-    flex_worknum=flex_worknum,
     flex_render_smooth=flex_render_smooth,
+    bvh_nflexgeom=len(flex_geom_flexid),
+    flex_dim_np=mjm.flex_dim,
+    flex_geom_flexid=wp.array(flex_geom_flexid, dtype=int),
+    flex_geom_edgeid=wp.array(flex_geom_edgeid, dtype=int),
     bvh=None,
     bvh_id=None,
-    lower=wp.zeros(nworld * bvh_ngeom, dtype=wp.vec3),
+    lower=wp.zeros(nworld * (bvh_ngeom + len(flex_geom_flexid)), dtype=wp.vec3),
-    upper=wp.zeros(nworld * bvh_ngeom, dtype=wp.vec3),
+    upper=wp.zeros(nworld * (bvh_ngeom + len(flex_geom_flexid)), dtype=wp.vec3),
-    group=wp.zeros(nworld * bvh_ngeom, dtype=int),
+    group=wp.zeros(nworld * (bvh_ngeom + len(flex_geom_flexid)), dtype=int),
     group_root=wp.zeros(nworld, dtype=int),
     ray=ray,
     rgb_data=wp.zeros((nworld, ri), dtype=wp.uint32),
@@ -2757,6 +2882,9 @@ def create_render_context(
     depth_adr=wp.array(depth_adr, dtype=int),
     render_rgb=wp.array(render_rgb, dtype=bool),
     render_depth=wp.array(render_depth, dtype=bool),
+    seg_data=wp.zeros((nworld, max(si, 1)), dtype=int),
+    seg_adr=wp.array(seg_adr, dtype=int),
+    render_seg=wp.array(render_seg, dtype=bool),
     znear=znear,
     total_rays=int(total),
   )

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

@@ -13,12 +13,13 @@
 # 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
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-import warp as wp
 
 
 @wp.kernel

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

@@ -83,6 +83,35 @@ def quat_to_mat(quat: wp.quat) -> wp.mat33:
   )
 
 
+@wp.func
+def quat_z2vec(vec: wp.vec3) -> wp.quat:
+  """Compute quaternion performing rotation from z-axis to given vector."""
+  quat = wp.quat(0.0, 0.0, 0.0, 1.0)
+
+  # normalize vector; if too small, no rotation
+  norm = wp.length(vec)
+  if norm < types.MJ_MINVAL:
+    return quat
+  vec = vec / norm
+
+  axis = wp.vec3(-vec[1], vec[0], 0.0)
+  a = wp.length(axis)
+
+  # almost parallel
+  if a < types.MJ_MINVAL:
+    # opposite: 180 deg rotation around x axis
+    if vec[2] < 0.0:
+      quat = wp.quat(1.0, 0.0, 0.0, 0.0)
+    return quat
+
+  # make quaternion from angle and axis
+  axis = axis / a
+  angle = wp.atan2(a, vec[2])
+  quat = axis_angle_to_quat(axis, angle)
+
+  return quat
+
+
 @wp.func
 def quat_inv(quat: wp.quat) -> wp.quat:
   return wp.quat(quat[0], -quat[1], -quat[2], -quat[3])

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

@@ -89,8 +89,8 @@ def _spring_damper_dof_passive(
   stiffness = jnt_stiffness[worldid % jnt_stiffness.shape[0], jntid]
   damping = dof_damping[worldid % dof_damping.shape[0], dofid]
 
-  has_stiffness = stiffness != 0.0 and not opt_disableflags & DisableBit.SPRING
+  has_stiffness = stiffness != 0.0 and not (opt_disableflags & DisableBit.SPRING)
-  has_damping = damping != 0.0 and not opt_disableflags & DisableBit.DAMPER
+  has_damping = damping != 0.0 and not (opt_disableflags & DisableBit.DAMPER)
 
   if not has_stiffness:
     qfrc_spring_out[worldid, dofid] = 0.0
@@ -182,11 +182,14 @@ def _spring_damper_dof_passive(
 @wp.kernel
 def _spring_damper_tendon_passive(
   # Model:
+  ten_J_rownnz: wp.array(dtype=int),
+  ten_J_rowadr: wp.array(dtype=int),
+  ten_J_colind: wp.array(dtype=int),
   tendon_stiffness: wp.array2d(dtype=float),
   tendon_damping: wp.array2d(dtype=float),
   tendon_lengthspring: wp.array2d(dtype=wp.vec2),
   # Data in:
-  ten_J_in: wp.array3d(dtype=float),
+  ten_J_in: wp.array2d(dtype=float),
   ten_length_in: wp.array2d(dtype=float),
   ten_velocity_in: wp.array2d(dtype=float),
   # In:
@@ -196,7 +199,7 @@ def _spring_damper_tendon_passive(
   qfrc_spring_out: wp.array2d(dtype=float),
   qfrc_damper_out: wp.array2d(dtype=float),
 ):
-  worldid, tenid, dofid = wp.tid()
+  worldid, tenid, dofid_sparse = wp.tid()
 
   stiffness = tendon_stiffness[worldid % tendon_stiffness.shape[0], tenid]
   damping = tendon_damping[worldid % tendon_damping.shape[0], tenid]
@@ -207,7 +210,13 @@ def _spring_damper_tendon_passive(
   if not has_stiffness and not has_damping:
     return
 
-  J = ten_J_in[worldid, tenid, dofid]
+  rownnz = ten_J_rownnz[tenid]
+  if dofid_sparse >= rownnz:
+    return
+  rowadr = ten_J_rowadr[tenid]
+  sparseid = rowadr + dofid_sparse
+  J = ten_J_in[worldid, sparseid]
+  dofid = ten_J_colind[sparseid]
 
   if has_stiffness:
     # compute spring force along tendon
@@ -565,6 +574,7 @@ def _flex_elasticity(
   flex_edgeadr: wp.array(dtype=int),
   flex_elemadr: wp.array(dtype=int),
   flex_elemnum: wp.array(dtype=int),
+  flex_elemdataadr: wp.array(dtype=int),
   flex_elemedgeadr: wp.array(dtype=int),
   flex_vertbodyid: wp.array(dtype=int),
   flex_elem: wp.array(dtype=int),
@@ -590,32 +600,39 @@ def _flex_elasticity(
       f = i
       break
 
+  local_elemid = elemid - flex_elemadr[f]
   dim = flex_dim[f]
   nvert = dim + 1
   nedge = nvert * (nvert - 1) / 2
   edges = wp.where(
+    dim == 1,
+    wp.matrix(0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, shape=(6, 2), dtype=int),
+    wp.where(
       dim == 3,
       wp.matrix(0, 1, 1, 2, 2, 0, 2, 3, 0, 3, 1, 3, shape=(6, 2), dtype=int),
       wp.matrix(1, 2, 2, 0, 0, 1, 0, 0, 0, 0, 0, 0, shape=(6, 2), dtype=int),
+    ),
   )
   if timestep > 0.0 and not dsbl_damper:
     kD = flex_damping[f] / timestep
   else:
     kD = 0.0
 
+  elem_data_adr = flex_elemdataadr[f] + local_elemid * (dim + 1)
+  vbase = flex_vertadr[f]
   gradient = wp.matrix(0.0, shape=(6, 6))
   for e in range(nedge):
-    vert0 = flex_elem[(dim + 1) * elemid + edges[e, 0]]
+    vert0 = flex_elem[elem_data_adr + edges[e, 0]]
-    vert1 = flex_elem[(dim + 1) * elemid + edges[e, 1]]
+    vert1 = flex_elem[elem_data_adr + edges[e, 1]]
-    xpos0 = flexvert_xpos_in[worldid, vert0]
+    xpos0 = flexvert_xpos_in[worldid, vbase + vert0]
-    xpos1 = flexvert_xpos_in[worldid, vert1]
+    xpos1 = flexvert_xpos_in[worldid, vbase + vert1]
     for i in range(3):
       gradient[e, 0 + i] = xpos0[i] - xpos1[i]
       gradient[e, 3 + i] = xpos1[i] - xpos0[i]
 
   elongation = wp.spatial_vectorf(0.0)
   for e in range(nedge):
-    idx = flex_elemedge[elemid * nedge + e]
+    idx = flex_elemedge[flex_elemedgeadr[f] + local_elemid * nedge + e]
     vel = flexedge_velocity_in[worldid, flex_edgeadr[f] + idx]
     deformed = flexedge_length_in[worldid, flex_edgeadr[f] + idx]
     reference = flexedge_length0[flex_edgeadr[f] + idx]
@@ -638,7 +655,7 @@ def _flex_elasticity(
           force[edges[ed2, i], x] -= elongation[ed1] * gradient[ed2, 3 * i + x] * metric[ed1, ed2]
 
   for v in range(nvert):
-    vert = flex_elem[(dim + 1) * elemid + v]
+    vert = flex_elem[elem_data_adr + v]
     bodyid = flex_vertbodyid[flex_vertadr[f] + vert]
     for x in range(3):
       wp.atomic_add(qfrc_spring_out, worldid, body_dofadr[bodyid] + x, force[v, x])
@@ -742,8 +759,11 @@ def passive(m: Model, d: Data):
   if m.ntendon:
     wp.launch(
       _spring_damper_tendon_passive,
-      dim=(d.nworld, m.ntendon, m.nv),
+      dim=(d.nworld, m.ntendon, m.max_ten_J_rownnz),
       inputs=[
+        m.ten_J_rownnz,
+        m.ten_J_rowadr,
+        m.ten_J_colind,
         m.tendon_stiffness,
         m.tendon_damping,
         m.tendon_lengthspring,
@@ -772,6 +792,7 @@ def passive(m: Model, d: Data):
         m.flex_edgeadr,
         m.flex_elemadr,
         m.flex_elemnum,
+        m.flex_elemdataadr,
         m.flex_elemedgeadr,
         m.flex_vertbodyid,
         m.flex_elem,

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

@@ -752,7 +752,8 @@ def ray_mesh_with_bvh_anyhit(
 @wp.func
 def ray_flex_with_bvh(
   # In:
-  bvh_id: wp.uint64,
+  flex_bvh_id: wp.array(dtype=wp.uint64),
+  flexid: int,
   group_root: int,
   pnt: wp.vec3,
   vec: wp.vec3,
@@ -769,7 +770,7 @@ def ray_flex_with_bvh(
   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)
+  hit = wp.mesh_query_ray(flex_bvh_id[flexid], pnt, vec, max_t, t, u, v, sign, n, f, group_root)
 
   if hit:
     return t, n, u, v, f
@@ -777,6 +778,23 @@ def ray_flex_with_bvh(
   return -1.0, wp.vec3(0.0, 0.0, 0.0), 0.0, 0.0, -1
 
 
+@wp.func
+def ray_flex_with_bvh_anyhit(
+  # In:
+  flex_bvh_id: wp.array(dtype=wp.uint64),
+  flexid: int,
+  group_root: int,
+  pnt: wp.vec3,
+  vec: wp.vec3,
+  max_t: float,
+) -> bool:
+  """Returns True if there is any hit for ray flex intersections.
+
+  Requires wp.Mesh be constructed and their ids to be passed. Flex are already in world space.
+  """
+  return wp.mesh_query_ray_anyhit(flex_bvh_id[flexid], pnt, vec, max_t, group_root)
+
+
 @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.

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

@@ -23,6 +23,7 @@ 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_flex_with_bvh_anyhit
 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
@@ -39,10 +40,6 @@ 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:
@@ -51,7 +48,7 @@ def sample_texture(
   # In:
   geom_id: int,
   tex_repeat: wp.vec2,
-  tex: TEXTURE_DTYPE,
+  tex: wp.Texture2D,
   pos: wp.vec3,
   rot: wp.mat33,
   mesh_facetexcoord: wp.array(dtype=wp.vec3i),
@@ -94,17 +91,26 @@ def cast_ray(
   geom_type: wp.array(dtype=int),
   geom_dataid: wp.array(dtype=int),
   geom_size: wp.array2d(dtype=wp.vec3),
+  flex_vertadr: wp.array(dtype=int),
+  flex_edge: wp.array(dtype=wp.vec2i),
+  flex_radius: wp.array(dtype=float),
   # Data in:
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
   geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
   # In:
   bvh_id: wp.uint64,
   group_root: int,
-  world_id: int,
+  worldid: int,
   bvh_ngeom: int,
+  flex_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),
+  flex_geom_flexid: wp.array(dtype=int),
+  flex_geom_edgeid: wp.array(dtype=int),
+  flex_bvh_id: wp.array(dtype=wp.uint64),
+  flex_group_root: wp.array2d(dtype=int),
   ray_origin_world: wp.vec3,
   ray_dir_world: wp.vec3,
 ) -> Tuple[int, float, wp.vec3, float, float, int, int]:
@@ -118,91 +124,127 @@ def cast_ray(
 
   query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
   bounds_nr = int(0)
+  ngeom = bvh_ngeom + flex_bvh_ngeom
 
   while wp.bvh_query_next(query, bounds_nr, dist):
     gi_global = bounds_nr
-    gi_bvh_local = gi_global - (world_id * bvh_ngeom)
+    local_id = gi_global - (worldid * ngeom)
-    gi = enabled_geom_ids[gi_bvh_local]
 
+    d = float(-1.0)
     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)
+    hit_geom_id = int(-1)
+
+    if local_id < bvh_ngeom:
+      gi = enabled_geom_ids[local_id]
+      gtype = geom_type[gi]
+    else:
+      gi = local_id - bvh_ngeom
+      gtype = GeomType.FLEX
+
+    hit_geom_id = gi
 
     # TODO: Investigate branch elimination with static loop unrolling
-    if geom_type[gi] == GeomType.PLANE:
+    if gtype == GeomType.PLANE:
       d, n = ray_plane(
-        geom_xpos_in[world_id, gi],
+        geom_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
-        geom_size[world_id % geom_size.shape[0], gi],
+        geom_size[worldid % geom_size.shape[0], gi],
         ray_origin_world,
         ray_dir_world,
       )
-    if geom_type[gi] == GeomType.HFIELD:
+    if gtype == 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_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
         ray_origin_world,
         ray_dir_world,
         dist,
       )
-    if geom_type[gi] == GeomType.SPHERE:
+    if gtype == GeomType.SPHERE:
       d, n = ray_sphere(
-        geom_xpos_in[world_id, gi],
+        geom_xpos_in[worldid, gi],
-        geom_size[world_id % geom_size.shape[0], gi][0] * geom_size[world_id % geom_size.shape[0], gi][0],
+        geom_size[worldid % geom_size.shape[0], gi][0] * geom_size[worldid % geom_size.shape[0], gi][0],
         ray_origin_world,
         ray_dir_world,
       )
-    if geom_type[gi] == GeomType.ELLIPSOID:
+    if gtype == GeomType.ELLIPSOID:
       d, n = ray_ellipsoid(
-        geom_xpos_in[world_id, gi],
+        geom_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
-        geom_size[world_id % geom_size.shape[0], gi],
+        geom_size[worldid % geom_size.shape[0], gi],
         ray_origin_world,
         ray_dir_world,
       )
-    if geom_type[gi] == GeomType.CAPSULE:
+    if gtype == GeomType.CAPSULE:
       d, n = ray_capsule(
-        geom_xpos_in[world_id, gi],
+        geom_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
-        geom_size[world_id % geom_size.shape[0], gi],
+        geom_size[worldid % geom_size.shape[0], gi],
         ray_origin_world,
         ray_dir_world,
       )
-    if geom_type[gi] == GeomType.CYLINDER:
+    if gtype == GeomType.CYLINDER:
       d, n = ray_cylinder(
-        geom_xpos_in[world_id, gi],
+        geom_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
-        geom_size[world_id % geom_size.shape[0], gi],
+        geom_size[worldid % geom_size.shape[0], gi],
         ray_origin_world,
         ray_dir_world,
       )
-    if geom_type[gi] == GeomType.BOX:
+    if gtype == GeomType.BOX:
       d, all, n = ray_box(
-        geom_xpos_in[world_id, gi],
+        geom_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
-        geom_size[world_id % geom_size.shape[0], gi],
+        geom_size[worldid % geom_size.shape[0], gi],
         ray_origin_world,
         ray_dir_world,
       )
-    if geom_type[gi] == GeomType.MESH:
+    if gtype == 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_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
         ray_origin_world,
         ray_dir_world,
         dist,
       )
+    if gtype == GeomType.FLEX:
+      hit_geom_id = -2
+      flexid = flex_geom_flexid[gi]
+      edge_id = flex_geom_edgeid[gi]
+
+      if edge_id >= 0:
+        edge = flex_edge[edge_id]
+        vert_adr = flex_vertadr[flexid]
+        v0 = flexvert_xpos_in[worldid, vert_adr + edge[0]]
+        v1 = flexvert_xpos_in[worldid, vert_adr + edge[1]]
+        pos = 0.5 * (v0 + v1)
+        vec = v1 - v0
+
+        length = wp.length(vec)
+        edgeq = math.quat_z2vec(vec)
+        mat = math.quat_to_mat(edgeq)
+        size = wp.vec3(flex_radius[flexid], 0.5 * length, 0.0)
+
+        d, n = ray_capsule(pos, mat, size, ray_origin_world, ray_dir_world)
+        hit_mesh_id = flexid
+      else:
+        flex_gr = flex_group_root[worldid, flexid]
+        d, n, u, v, f = ray_flex_with_bvh(flex_bvh_id, flexid, flex_gr, ray_origin_world, ray_dir_world, dist)
+        if d >= 0.0:
+          hit_mesh_id = flexid
 
     if d >= 0.0 and d < dist:
       dist = d
       normal = n
-      geom_id = gi
+      geom_id = hit_geom_id
       bary_u = u
       bary_v = v
       face_idx = f
@@ -217,17 +259,26 @@ def cast_ray_first_hit(
   geom_type: wp.array(dtype=int),
   geom_dataid: wp.array(dtype=int),
   geom_size: wp.array2d(dtype=wp.vec3),
+  flex_vertadr: wp.array(dtype=int),
+  flex_edge: wp.array(dtype=wp.vec2i),
+  flex_radius: wp.array(dtype=float),
   # Data in:
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
   geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
   # In:
   bvh_id: wp.uint64,
   group_root: int,
-  world_id: int,
+  worldid: int,
   bvh_ngeom: int,
+  bvh_nflexgeom: int,
   enabled_geom_ids: wp.array(dtype=int),
   mesh_bvh_id: wp.array(dtype=wp.uint64),
   hfield_bvh_id: wp.array(dtype=wp.uint64),
+  flex_geom_flexid: wp.array(dtype=int),
+  flex_geom_edgeid: wp.array(dtype=int),
+  flex_bvh_id: wp.array(dtype=wp.uint64),
+  flex_group_root: wp.array2d(dtype=int),
   ray_origin_world: wp.vec3,
   ray_dir_world: wp.vec3,
   max_dist: float,
@@ -235,76 +286,114 @@ def cast_ray_first_hit(
   """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)
+  ngeom = bvh_ngeom + bvh_nflexgeom
 
   while wp.bvh_query_next(query, bounds_nr, max_dist):
     gi_global = bounds_nr
-    gi_bvh_local = gi_global - (world_id * bvh_ngeom)
+    local_id = gi_global - (worldid * ngeom)
-    gi = enabled_geom_ids[gi_bvh_local]
+
+    d = float(-1.0)
+    n = wp.vec3(0.0, 0.0, 0.0)
+
+    if local_id < bvh_ngeom:
+      gi = enabled_geom_ids[local_id]
+      gtype = geom_type[gi]
+    else:
+      gi = local_id - bvh_ngeom
+      gtype = GeomType.FLEX
 
     # TODO: Investigate branch elimination with static loop unrolling
-    if geom_type[gi] == GeomType.PLANE:
+    if gtype == GeomType.PLANE:
       d, n = ray_plane(
-        geom_xpos_in[world_id, gi],
+        geom_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
-        geom_size[world_id % geom_size.shape[0], gi],
+        geom_size[worldid % geom_size.shape[0], gi],
         ray_origin_world,
         ray_dir_world,
       )
-    if geom_type[gi] == GeomType.HFIELD:
+    if gtype == 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_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
         ray_origin_world,
         ray_dir_world,
         max_dist,
       )
-    if geom_type[gi] == GeomType.SPHERE:
+    if gtype == GeomType.SPHERE:
       d, n = ray_sphere(
-        geom_xpos_in[world_id, gi],
+        geom_xpos_in[worldid, gi],
-        geom_size[world_id % geom_size.shape[0], gi][0] * geom_size[world_id % geom_size.shape[0], gi][0],
+        geom_size[worldid % geom_size.shape[0], gi][0] * geom_size[worldid % geom_size.shape[0], gi][0],
         ray_origin_world,
         ray_dir_world,
       )
-    if geom_type[gi] == GeomType.ELLIPSOID:
+    if gtype == GeomType.ELLIPSOID:
       d, n = ray_ellipsoid(
-        geom_xpos_in[world_id, gi],
+        geom_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
-        geom_size[world_id % geom_size.shape[0], gi],
+        geom_size[worldid % geom_size.shape[0], gi],
         ray_origin_world,
         ray_dir_world,
       )
-    if geom_type[gi] == GeomType.CAPSULE:
+    if gtype == GeomType.CAPSULE:
       d, n = ray_capsule(
-        geom_xpos_in[world_id, gi],
+        geom_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
-        geom_size[world_id % geom_size.shape[0], gi],
+        geom_size[worldid % geom_size.shape[0], gi],
         ray_origin_world,
         ray_dir_world,
       )
-    if geom_type[gi] == GeomType.CYLINDER:
+    if gtype == GeomType.CYLINDER:
       d, n = ray_cylinder(
-        geom_xpos_in[world_id, gi],
+        geom_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
-        geom_size[world_id % geom_size.shape[0], gi],
+        geom_size[worldid % geom_size.shape[0], gi],
         ray_origin_world,
         ray_dir_world,
       )
-    if geom_type[gi] == GeomType.BOX:
+    if gtype == GeomType.BOX:
       d, all, n = ray_box(
-        geom_xpos_in[world_id, gi],
+        geom_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
-        geom_size[world_id % geom_size.shape[0], gi],
+        geom_size[worldid % geom_size.shape[0], gi],
         ray_origin_world,
         ray_dir_world,
       )
-    if geom_type[gi] == GeomType.MESH:
+    if gtype == GeomType.MESH:
       hit = ray_mesh_with_bvh_anyhit(
         mesh_bvh_id,
         geom_dataid[gi],
-        geom_xpos_in[world_id, gi],
+        geom_xpos_in[worldid, gi],
-        geom_xmat_in[world_id, gi],
+        geom_xmat_in[worldid, gi],
+        ray_origin_world,
+        ray_dir_world,
+        max_dist,
+      )
+      d = 0.0 if hit else -1.0
+    if gtype == GeomType.FLEX:
+      flexid = flex_geom_flexid[gi]
+      edge_id = flex_geom_edgeid[gi]
+
+      if edge_id >= 0:
+        edge = flex_edge[edge_id]
+        vert_adr = flex_vertadr[flexid]
+        v0 = flexvert_xpos_in[worldid, vert_adr + edge[0]]
+        v1 = flexvert_xpos_in[worldid, vert_adr + edge[1]]
+        pos = 0.5 * (v0 + v1)
+        vec = v1 - v0
+
+        length = wp.length(vec)
+        edgeq = math.quat_z2vec(vec)
+        mat = math.quat_to_mat(edgeq)
+        size = wp.vec3(flex_radius[flexid], 0.5 * length, 0.0)
+
+        d, n = ray_capsule(pos, mat, size, ray_origin_world, ray_dir_world)
+      else:
+        hit = ray_flex_with_bvh_anyhit(
+          flex_bvh_id,
+          flexid,
+          flex_group_root[worldid, flexid],
           ray_origin_world,
           ray_dir_world,
           max_dist,
@@ -323,18 +412,27 @@ def compute_lighting(
   geom_type: wp.array(dtype=int),
   geom_dataid: wp.array(dtype=int),
   geom_size: wp.array2d(dtype=wp.vec3),
+  flex_vertadr: wp.array(dtype=int),
+  flex_edge: wp.array(dtype=wp.vec2i),
+  flex_radius: wp.array(dtype=float),
   # Data in:
   geom_xpos_in: wp.array2d(dtype=wp.vec3),
   geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
   # In:
   use_shadows: bool,
   bvh_id: wp.uint64,
   group_root: int,
   bvh_ngeom: int,
+  bvh_nflexgeom: int,
   enabled_geom_ids: wp.array(dtype=int),
-  world_id: int,
+  worldid: int,
   mesh_bvh_id: wp.array(dtype=wp.uint64),
   hfield_bvh_id: wp.array(dtype=wp.uint64),
+  flex_geom_flexid: wp.array(dtype=int),
+  flex_geom_edgeid: wp.array(dtype=int),
+  flex_bvh_id: wp.array(dtype=wp.uint64),
+  flex_group_root: wp.array2d(dtype=int),
   lightactive: bool,
   lighttype: int,
   lightcastshadow: bool,
@@ -385,15 +483,24 @@ def compute_lighting(
       geom_type,
       geom_dataid,
       geom_size,
+      flex_vertadr,
+      flex_edge,
+      flex_radius,
       geom_xpos_in,
       geom_xmat_in,
+      flexvert_xpos_in,
       bvh_id,
       group_root,
-      world_id,
+      worldid,
       bvh_ngeom,
+      bvh_nflexgeom,
       enabled_geom_ids,
       mesh_bvh_id,
       hfield_bvh_id,
+      flex_geom_flexid,
+      flex_geom_edgeid,
+      flex_bvh_id,
+      flex_group_root,
       shadow_origin,
       L,
       max_t,
@@ -418,6 +525,7 @@ def render(m: Model, d: Data, rc: RenderContext):
   """
   rc.rgb_data.fill_(rc.background_color)
   rc.depth_data.fill_(0.0)
+  rc.seg_data.fill_(-1)
 
   @wp.kernel(module="unique", enable_backward=False)
   def _render_megakernel(
@@ -434,6 +542,9 @@ def render(m: Model, d: Data, rc: RenderContext):
     light_type: wp.array2d(dtype=int),
     light_castshadow: wp.array2d(dtype=bool),
     light_active: wp.array2d(dtype=bool),
+    flex_vertadr: wp.array(dtype=int),
+    flex_edge: wp.array(dtype=wp.vec2i),
+    flex_radius: wp.array(dtype=float),
     mesh_faceadr: wp.array(dtype=int),
     mat_texid: wp.array3d(dtype=int),
     mat_texrepeat: wp.array2d(dtype=wp.vec2),
@@ -445,21 +556,25 @@ def render(m: Model, d: Data, rc: RenderContext):
     cam_xmat_in: wp.array2d(dtype=wp.mat33),
     light_xpos_in: wp.array2d(dtype=wp.vec3),
     light_xdir_in: wp.array2d(dtype=wp.vec3),
+    flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
     # In:
     nrender: int,
     use_shadows: bool,
     bvh_ngeom: int,
+    bvh_nflexgeom: 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),
+    seg_adr: wp.array(dtype=int),
     render_rgb: wp.array(dtype=bool),
     render_depth: wp.array(dtype=bool),
+    render_seg: wp.array(dtype=bool),
     bvh_id: wp.uint64,
     group_root: wp.array(dtype=int),
-    flex_bvh_id: wp.uint64,
+    flex_bvh_id: wp.array(dtype=wp.uint64),
-    flex_group_root: wp.array(dtype=int),
+    flex_group_root: wp.array2d(dtype=int),
     enabled_geom_ids: wp.array(dtype=int),
     mesh_bvh_id: wp.array(dtype=wp.uint64),
     mesh_facetexcoord: wp.array(dtype=wp.vec3i),
@@ -467,46 +582,48 @@ def render(m: Model, d: Data, rc: RenderContext):
     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
+    flex_geom_flexid: wp.array(dtype=int),
-    textures: wp.array(dtype=TEXTURE_DTYPE),
+    flex_geom_edgeid: wp.array(dtype=int),
+    textures: wp.array(dtype=wp.Texture2D),
     # Out:
     rgb_out: wp.array2d(dtype=wp.uint32),
     depth_out: wp.array2d(dtype=float),
+    seg_out: wp.array2d(dtype=int),
   ):
-    world_idx, ray_idx = wp.tid()
+    worldid, rayid = wp.tid()
 
-    # Map global ray_idx -> (cam_idx, ray_idx_local) using cumulative sizes
+    # Map global rayid -> (cam_idx, rayid_local) using cumulative sizes
     cam_idx = int(-1)
-    ray_idx_local = int(-1)
+    rayid_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:
+      if rayid < accum + num_i:
         cam_idx = i
-        ray_idx_local = ray_idx - accum
+        rayid_local = rayid - accum
         break
       accum += num_i
-    if cam_idx == -1 or ray_idx_local < 0:
+    if cam_idx == -1 or rayid_local < 0:
       return
 
-    if not render_rgb[cam_idx] and not render_depth[cam_idx]:
+    if not render_rgb[cam_idx] and not render_depth[cam_idx] and not render_seg[cam_idx]:
       return
 
     # Map active camera index to MuJoCo camera ID
     mujoco_cam_id = cam_id_map[cam_idx]
 
     if wp.static(rc.use_precomputed_rays):
-      ray_dir_local_cam = ray[ray_idx]
+      ray_dir_local_cam = ray[rayid]
     else:
       img_w = cam_res[cam_idx][0]
       img_h = cam_res[cam_idx][1]
-      px = ray_idx_local % img_w
+      px = rayid_local % img_w
-      py = ray_idx_local // img_w
+      py = rayid_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_fovy[worldid % cam_fovy.shape[0], mujoco_cam_id],
         cam_sensorsize[mujoco_cam_id],
-        cam_intrinsic[world_idx % cam_intrinsic.shape[0], mujoco_cam_id],
+        cam_intrinsic[worldid % cam_intrinsic.shape[0], mujoco_cam_id],
         img_w,
         img_h,
         px,
@@ -514,38 +631,37 @@ def render(m: Model, d: Data, rc: RenderContext):
         wp.static(rc.znear),
       )
 
-    ray_dir_world = cam_xmat_in[world_idx, mujoco_cam_id] @ ray_dir_local_cam
+    ray_dir_world = cam_xmat_in[worldid, mujoco_cam_id] @ ray_dir_local_cam
-    ray_origin_world = cam_xpos_in[world_idx, mujoco_cam_id]
+    ray_origin_world = cam_xpos_in[worldid, mujoco_cam_id]
 
     geom_id, dist, normal, u, v, f, mesh_id = cast_ray(
       geom_type,
       geom_dataid,
       geom_size,
+      flex_vertadr,
+      flex_edge,
+      flex_radius,
       geom_xpos_in,
       geom_xmat_in,
+      flexvert_xpos_in,
       bvh_id,
-      group_root[world_idx],
+      group_root[worldid],
-      world_idx,
+      worldid,
       bvh_ngeom,
+      bvh_nflexgeom,
       enabled_geom_ids,
       mesh_bvh_id,
       hfield_bvh_id,
+      flex_geom_flexid,
+      flex_geom_edgeid,
+      flex_bvh_id,
+      flex_group_root,
       ray_origin_world,
       ray_dir_world,
     )
 
-    if wp.static(m.nflex > 0):
+    if render_seg[cam_idx] and geom_id != -1:
-      d, n, u, v, f = ray_flex_with_bvh(
+      seg_out[worldid, seg_adr[cam_idx] + rayid_local] = geom_id
-        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:
@@ -556,9 +672,7 @@ def render(m: Model, d: Data, rc: RenderContext):
       # In camera-local coordinates, the optical axis is -Z. The Z-component of the
       # normalized ray direction is negative, so -ray_dir_local_cam[2] gives cos(θ)
       # between the ray and the optical axis.
-      depth_out[world_idx, depth_adr[cam_idx] + ray_idx_local] = dist * (
+      depth_out[worldid, depth_adr[cam_idx] + rayid_local] = dist * (-ray_dir_local_cam[2])
-          -ray_dir_local_cam[2]
-      )
 
     if not render_rgb[cam_idx]:
       return
@@ -567,31 +681,30 @@ def render(m: Model, d: Data, rc: RenderContext):
     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
+      # We encode flex_id in mesh_id for flex ray hits during cast_ray
-      # is used until further flex support is added.
+      color = flex_rgba[mesh_id]
-      color = flex_rgba[0]
+    elif geom_matid[worldid % geom_matid.shape[0], geom_id] == -1:
-    elif geom_matid[world_idx % geom_matid.shape[0], geom_id] == -1:
+      color = geom_rgba[worldid % geom_rgba.shape[0], geom_id]
-      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]]
+      color = mat_rgba[worldid % mat_rgba.shape[0], geom_matid[worldid % 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]
+        mat_id = geom_matid[worldid % geom_matid.shape[0], geom_id]
         if mat_id >= 0:
-          tex_id = mat_texid[world_idx % mat_texid.shape[0], mat_id, 1]
+          tex_id = mat_texid[worldid % 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],
+              mat_texrepeat[worldid % mat_texrepeat.shape[0], mat_id],
               textures[tex_id],
-              geom_xpos_in[world_idx, geom_id],
+              geom_xpos_in[worldid, geom_id],
-              geom_xmat_in[world_idx, geom_id],
+              geom_xmat_in[worldid, geom_id],
               mesh_facetexcoord,
               mesh_texcoord,
               mesh_texcoord_offsets,
@@ -616,21 +729,30 @@ def render(m: Model, d: Data, rc: RenderContext):
         geom_type,
         geom_dataid,
         geom_size,
+        flex_vertadr,
+        flex_edge,
+        flex_radius,
         geom_xpos_in,
         geom_xmat_in,
+        flexvert_xpos_in,
         use_shadows,
         bvh_id,
-        group_root[world_idx],
+        group_root[worldid],
         bvh_ngeom,
+        bvh_nflexgeom,
         enabled_geom_ids,
-        world_idx,
+        worldid,
         mesh_bvh_id,
         hfield_bvh_id,
-        light_active[world_idx % light_active.shape[0], l],
+        flex_geom_flexid,
-        light_type[world_idx % light_type.shape[0], l],
+        flex_geom_edgeid,
-        light_castshadow[world_idx % light_castshadow.shape[0], l],
+        flex_bvh_id,
-        light_xpos_in[world_idx, l],
+        flex_group_root,
-        light_xdir_in[world_idx, l],
+        light_active[worldid % light_active.shape[0], l],
+        light_type[worldid % light_type.shape[0], l],
+        light_castshadow[worldid % light_castshadow.shape[0], l],
+        light_xpos_in[worldid, l],
+        light_xdir_in[worldid, l],
         normal,
         hit_point,
       )
@@ -639,7 +761,7 @@ def render(m: Model, d: Data, rc: RenderContext):
     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(
+    rgb_out[worldid, rgb_adr[cam_idx] + rayid_local] = pack_rgba_to_uint32(
       hit_color[0] * 255.0,
       hit_color[1] * 255.0,
       hit_color[2] * 255.0,
@@ -662,6 +784,9 @@ def render(m: Model, d: Data, rc: RenderContext):
       m.light_type,
       m.light_castshadow,
       m.light_active,
+      m.flex_vertadr,
+      m.flex_edge,
+      m.flex_radius,
       m.mesh_faceadr,
       m.mat_texid,
       m.mat_texrepeat,
@@ -672,16 +797,20 @@ def render(m: Model, d: Data, rc: RenderContext):
       d.cam_xmat,
       d.light_xpos,
       d.light_xdir,
+      d.flexvert_xpos,
       rc.nrender,
       rc.use_shadows,
       rc.bvh_ngeom,
+      rc.bvh_nflexgeom,
       rc.cam_res,
       rc.cam_id_map,
       rc.ray,
       rc.rgb_adr,
       rc.depth_adr,
+      rc.seg_adr,
       rc.render_rgb,
       rc.render_depth,
+      rc.render_seg,
       rc.bvh_id,
       rc.group_root,
       rc.flex_bvh_id,
@@ -693,10 +822,13 @@ def render(m: Model, d: Data, rc: RenderContext):
       rc.mesh_texcoord_offsets,
       rc.hfield_bvh_id,
       rc.flex_rgba,
+      rc.flex_geom_flexid,
+      rc.flex_geom_edgeid,
       rc.textures,
     ],
     outputs=[
       rc.rgb_data,
       rc.depth_data,
+      rc.seg_data,
     ],
   )

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

@@ -206,3 +206,41 @@ def get_depth(rc: RenderContext, camera_index: int, depth_scale: float, depth_ou
     inputs=[rc.depth_data, rc.depth_adr, camera_index, depth_scale],
     outputs=[depth_out],
   )
+
+
+@wp.kernel
+def _extract_seg_kernel(
+  # In:
+  seg_data: wp.array2d(dtype=int),
+  seg_adr: wp.array(dtype=int),
+  camera_index: int,
+  # Out:
+  seg_out: wp.array3d(dtype=int),
+):
+  """Extract per-pixel geom IDs from the render context buffers for a given camera index."""
+  worldid, pixelid = wp.tid()
+  xid = pixelid % seg_out.shape[2]
+  yid = pixelid // seg_out.shape[2]
+
+  seg_adr_offset = seg_adr[camera_index]
+  seg_out[worldid, yid, xid] = seg_data[worldid, seg_adr_offset + pixelid]
+
+
+def get_segmentation(rc: RenderContext, camera_index: int, seg_out: wp.array3d(dtype=int)):
+  """Get the segmentation data from the render context buffers for a given camera index.
+
+  Each pixel contains the MuJoCo geom ID of the geometry hit by the ray, -1 for
+  background, or -2 for flex bodies.
+
+  Args:
+    rc: The render context on device.
+    camera_index: The index of the camera to get the segmentation data for.
+    seg_out: The output array to store the geom IDs in, with shape
+      (nworld, height, width).
+  """
+  wp.launch(
+    _extract_seg_kernel,
+    dim=(seg_out.shape[0], seg_out.shape[1] * seg_out.shape[2]),
+    inputs=[rc.seg_data, rc.seg_adr, camera_index],
+    outputs=[seg_out],
+  )

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

@@ -15,12 +15,17 @@
 
 from typing import Any, Tuple
 
+import warp as wp
+
 from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import ray
 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_MAXCONPAIR
+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
 from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
@@ -28,9 +33,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import DataType
 from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import JointType
-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 MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
 from mujoco.mjx.third_party.mujoco_warp._src.types import SensorType
@@ -40,10 +42,10 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec6
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec8
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i
+from mujoco.mjx.third_party.mujoco_warp._src.types import vec_pluginattr
 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
-import warp as wp
 
 wp.set_module_options({"enable_backward": False})
 
@@ -2126,6 +2128,7 @@ def _sensor_tactile(
   oct_coeff: wp.array(dtype=vec8),
   geom_type: wp.array(dtype=int),
   geom_bodyid: wp.array(dtype=int),
+  geom_dataid: wp.array(dtype=int),
   geom_size: wp.array2d(dtype=wp.vec3),
   mesh_vertadr: wp.array(dtype=int),
   mesh_vertnum: wp.array(dtype=int),
@@ -2140,7 +2143,7 @@ def _sensor_tactile(
   sensor_dim: wp.array(dtype=int),
   sensor_adr: wp.array(dtype=int),
   plugin: wp.array(dtype=int),
-    plugin_attr: wp.array(dtype=wp.vec3f),
+  plugin_attr: wp.array(dtype=vec_pluginattr),
   geom_plugin_index: wp.array(dtype=int),
   taxel_vertadr: wp.array(dtype=int),
   taxel_sensorid: wp.array(dtype=int),
@@ -2220,31 +2223,16 @@ def _sensor_tactile(
       contact_type,
       geom_size[worldid % geom_size.shape[0], geom],
       plugin_id,
-        mesh_id,
+      geom_dataid[geom],
     )
 
-    depth = wp.min(
+    depth = wp.min(sdf(contact_type, lpos, plugin_attributes, plugin_index, volume_data, mesh_data), 0.0)
-        sdf(
-            contact_type,
-            lpos,
-            plugin_attributes,
-            plugin_index,
-            volume_data,
-            mesh_data,
-        ),
-        0.0,
-    )
     if depth >= 0.0:
       continue
 
-    vel_sensor = _transform_spatial(
+    vel_sensor = _transform_spatial(cvel_in[worldid, parent_weld], xpos - subtree_com_in[worldid, body_rootid[parent_weld]])
-        cvel_in[worldid, parent_weld],
-        xpos - subtree_com_in[worldid, body_rootid[parent_weld]],
-    )
     vel_other = _transform_spatial(
-        cvel_in[worldid, body],
+      cvel_in[worldid, body], geom_xpos_in[worldid, geom] - subtree_com_in[worldid, body_rootid[body]]
-        geom_xpos_in[worldid, geom]
-        - subtree_com_in[worldid, body_rootid[body]],
     )
     vel_rel = vel_sensor - vel_other
 
@@ -2259,24 +2247,9 @@ def _sensor_tactile(
       forceT[2] = wp.abs(wp.dot(vel_rel, tang2))
 
     dim = sensor_dim[sensor_id] // 3
-    wp.atomic_add(
+    wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
-        sensordata_out,
+    wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 1 * dim + vertid, forceT[1])
-        worldid,
+    wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 2 * dim + vertid, forceT[2])
-        sensor_adr[sensor_id] + 0 * dim + vertid,
-        forceT[0],
-    )
-    wp.atomic_add(
-        sensordata_out,
-        worldid,
-        sensor_adr[sensor_id] + 1 * dim + vertid,
-        forceT[1],
-    )
-    wp.atomic_add(
-        sensordata_out,
-        worldid,
-        sensor_adr[sensor_id] + 2 * dim + vertid,
-        forceT[2],
-    )
 
 
 @wp.func
@@ -2533,6 +2506,7 @@ def sensor_acc(m: Model, d: Data):
       m.oct_coeff,
       m.geom_type,
       m.geom_bodyid,
+      m.geom_dataid,
       m.geom_size,
       m.mesh_vertadr,
       m.mesh_vertnum,
@@ -2882,12 +2856,12 @@ def energy_pos(m: Model, d: Data):
   wp.launch(_energy_pos_zero, dim=d.nworld, outputs=[d.energy])
 
   # init potential energy: -sum_i(body_i.mass * dot(gravity, body_i.pos))
-  if not m.opt.disableflags & DisableBit.GRAVITY:
+  if not (m.opt.disableflags & DisableBit.GRAVITY):
     wp.launch(
       _energy_pos_gravity, dim=(d.nworld, m.nbody - 1), inputs=[m.opt.gravity, m.body_mass, d.xipos], outputs=[d.energy]
     )
 
-  if not m.opt.disableflags & DisableBit.SPRING:
+  if not (m.opt.disableflags & DisableBit.SPRING):
     # add joint-level springs
     wp.launch(
       _energy_pos_passive_joint,

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

@@ -17,17 +17,17 @@ import dataclasses
 from math import ceil
 from math import sqrt
 
+import warp as wp
+
 from mujoco.mjx.third_party.mujoco_warp._src import math
 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 import types
 from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_func
 from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_solve_func
-from mujoco.mjx.third_party.mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN
 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 scoped_mathdx_gemm_disabled
-import warp as wp
 
 wp.set_module_options({"enable_backward": False})
 
@@ -139,18 +139,10 @@ def create_solver_context(m: types.Model, d: types.Data) -> SolverContext:
     prev_grad=wp.empty((nworld, nv), dtype=float),
     prev_Mgrad=wp.empty((nworld, nv), dtype=float),
     beta=wp.empty((nworld,), dtype=float),
-      h=wp.zeros((nworld, nv_pad, nv_pad), dtype=float)
+    h=wp.zeros((nworld, nv_pad, nv_pad), dtype=float) if alloc_h else wp.empty((nworld, 0, 0), dtype=float),
-      if alloc_h
+    hfactor=wp.zeros((nworld, nv_pad, nv_pad), dtype=float) if alloc_hfactor else wp.empty((nworld, 0, 0), dtype=float),
-      else wp.empty((nworld, 0, 0), dtype=float),
+    changed_efc_ids=wp.empty((nworld, njmax), dtype=int) if alloc_h else wp.empty((nworld, 0), dtype=int),
-      hfactor=wp.zeros((nworld, nv_pad, nv_pad), dtype=float)
+    changed_efc_count=wp.empty((nworld,), dtype=int) if alloc_h else wp.empty((0,), dtype=int),
-      if alloc_hfactor
-      else wp.empty((nworld, 0, 0), dtype=float),
-      changed_efc_ids=wp.empty((nworld, njmax), dtype=int)
-      if alloc_h
-      else wp.empty((nworld, 0), dtype=int),
-      changed_efc_count=wp.empty((nworld,), dtype=int)
-      if alloc_h
-      else wp.empty((0,), dtype=int),
   )
 
 
@@ -892,18 +884,19 @@ def _compute_efc_eval_pt_3alphas_elliptic(
 
 
 @cache_kernel
-def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv: bool):
+def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv: bool, is_sparse: bool):
   """Factory for iterative linesearch kernel.
 
   Args:
-    block_dim: Number of threads per block for tile reductions.
     ls_iterations: Max linesearch iterations (compile-time constant for loop optimization).
     cone_type: Friction cone type (PYRAMIDAL or ELLIPTIC) for compile-time optimization.
     fuse_jv: Whether to compute jv = J @ search in-kernel (efficient for small nv).
+    is_sparse: Use sparse matrix representation for constraint Jacobian.
   """
   LS_ITERATIONS = ls_iterations
   IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
   FUSE_JV = fuse_jv
+  IS_SPARSE = is_sparse
 
   # Native snippet for CUDA __syncthreads()
   @wp.func_native(snippet="WP_TILE_SYNC();")
@@ -976,15 +969,13 @@ def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv:
     if wp.static(FUSE_JV):
       for efcid in range(tid, nefc, wp.block_dim()):
         jv = float(0.0)
-        if wp.static(SPARSE_CONSTRAINT_JACOBIAN):
+        if wp.static(IS_SPARSE):
           rownnz = efc_J_rownnz_in[worldid, efcid]
           rowadr = efc_J_rowadr_in[worldid, efcid]
           for k in range(rownnz):
             sparseid = rowadr + k
             colind = efc_J_colind_in[worldid, 0, sparseid]
-            jv += (
+            jv += efc_J_in[worldid, 0, sparseid] * ctx_search_in[worldid, colind]
-                efc_J_in[worldid, 0, sparseid] * ctx_search_in[worldid, colind]
-            )
         else:
           for i in range(nv):
             jv += efc_J_in[worldid, efcid, i] * ctx_search_in[worldid, i]
@@ -1360,7 +1351,7 @@ def _linesearch_iterative(m: types.Model, d: types.Data, ctx: SolverContext, fus
     fuse_jv: Whether jv is computed in-kernel (True) or pre-computed (False).
   """
   wp.launch_tiled(
-      linesearch_iterative(m.opt.ls_iterations, m.opt.cone, fuse_jv),
+    linesearch_iterative(m.opt.ls_iterations, m.opt.cone, fuse_jv, m.is_sparse),
     dim=d.nworld,
     inputs=[
       m.nv,
@@ -1420,7 +1411,7 @@ def linesearch_zero_jv(
 
 
 @cache_kernel
-def linesearch_jv_fused(opt_is_sparse: bool, nv: int, dofs_per_thread: int):
+def linesearch_jv_fused(is_sparse: bool, nv: int, dofs_per_thread: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Data in:
@@ -1446,23 +1437,21 @@ def linesearch_jv_fused(opt_is_sparse: bool, nv: int, dofs_per_thread: int):
     jv_out = float(0.0)
 
     if wp.static(dofs_per_thread >= nv):
-      if wp.static(SPARSE_CONSTRAINT_JACOBIAN):
+      if wp.static(is_sparse):
         # Sparse: iterate over non-zero entries in the row
         rownnz = efc_J_rownnz_in[worldid, efcid]
         rowadr = efc_J_rowadr_in[worldid, efcid]
         for k in range(rownnz):
           sparseid = rowadr + k
           colind = efc_J_colind_in[worldid, 0, sparseid]
-          jv_out += (
+          jv_out += efc_J_in[worldid, 0, sparseid] * ctx_search_in[worldid, colind]
-              efc_J_in[worldid, 0, sparseid] * ctx_search_in[worldid, colind]
-          )
       else:
         for i in range(wp.static(min(dofs_per_thread, nv))):
           jv_out += efc_J_in[worldid, efcid, i] * ctx_search_in[worldid, i]
       ctx_jv_out[worldid, efcid] = jv_out
 
     else:
-      if wp.static(SPARSE_CONSTRAINT_JACOBIAN):
+      if wp.static(is_sparse):
         # Sparse: thread 0 handles entire row (sparse entries << nv typically)
         if dofstart == 0:
           rownnz = efc_J_rownnz_in[worldid, efcid]
@@ -1470,9 +1459,7 @@ def linesearch_jv_fused(opt_is_sparse: bool, nv: int, dofs_per_thread: int):
           for k in range(rownnz):
             sparseid = rowadr + k
             colind = efc_J_colind_in[worldid, 0, sparseid]
-            jv_out += (
+            jv_out += efc_J_in[worldid, 0, sparseid] * ctx_search_in[worldid, colind]
-                efc_J_in[worldid, 0, sparseid] * ctx_search_in[worldid, colind]
-            )
           ctx_jv_out[worldid, efcid] = jv_out
       else:
         for i in range(wp.static(dofs_per_thread)):
@@ -1583,10 +1570,7 @@ def linesearch_prepare_quad(
 
     dim = contact_dim_in[conid]
     friction = contact_friction_in[conid]
-    mu = (
+    mu = friction[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
-        friction[0]
-        * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
-    )
 
     u0 = Jaref * mu
     v0 = jv * mu
@@ -1709,15 +1693,7 @@ def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.arra
     wp.launch(
       linesearch_jv_fused(m.is_sparse, m.nv, dofs_per_thread),
       dim=(d.nworld, d.njmax, threads_per_efc),
-        inputs=[
+      inputs=[d.nefc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, ctx.search, ctx.done],
-            d.nefc,
-            d.efc.J_rownnz,
-            d.efc.J_rowadr,
-            d.efc.J_colind,
-            d.efc.J,
-            ctx.search,
-            ctx.done,
-        ],
       outputs=[ctx.jv],
     )
 
@@ -1744,7 +1720,7 @@ def solve_init_efc(
 
 
 @cache_kernel
-def solve_init_jaref(opt_is_sparse: bool, nv: int, dofs_per_thread: int):
+def solve_init_jaref(is_sparse: bool, nv: int, dofs_per_thread: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Data in:
@@ -1764,7 +1740,7 @@ def solve_init_jaref(opt_is_sparse: bool, nv: int, dofs_per_thread: int):
       return
 
     jaref = float(0.0)
-    if wp.static(SPARSE_CONSTRAINT_JACOBIAN):
+    if wp.static(is_sparse):
       rownnz = efc_J_rownnz_in[worldid, efcid]
       rowadr = efc_J_rowadr_in[worldid, efcid]
       for i in range(rownnz):
@@ -1785,9 +1761,7 @@ def solve_init_jaref(opt_is_sparse: bool, nv: int, dofs_per_thread: int):
             jaref += efc_J_in[worldid, efcid, ii] * qacc_in[worldid, ii]
 
         if dofstart == 0:
-          wp.atomic_add(
+          wp.atomic_add(ctx_Jaref_out, worldid, efcid, jaref - efc_aref_in[worldid, efcid])
-              ctx_Jaref_out, worldid, efcid, jaref - efc_aref_in[worldid, efcid]
-          )
         else:
           wp.atomic_add(ctx_Jaref_out, worldid, efcid, jaref)
 
@@ -1869,9 +1843,7 @@ def update_constraint_efc(track_changes: bool):
 
     # Read old QUADRATIC status before overwriting
     if wp.static(TRACK_CHANGES):
-      old_quad = (
+      old_quad = efc_state_out[worldid, efcid] == types.ConstraintState.QUADRATIC.value
-          efc_state_out[worldid, efcid] == types.ConstraintState.QUADRATIC.value
-      )
 
     efc_D = efc_D_in[worldid, efcid]
     Jaref = ctx_Jaref_in[worldid, efcid]
@@ -1919,10 +1891,7 @@ def update_constraint_efc(track_changes: bool):
 
       dim = contact_dim_in[conid]
       friction = contact_friction_in[conid]
-      mu = (
+      mu = friction[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
-          friction[0]
-          * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
-      )
 
       efcid0 = contact_efc_address_in[conid, 0]
       if efcid0 < 0:
@@ -2076,8 +2045,8 @@ def update_constraint_gauss_cost(nv: int, dofs_per_thread: int):
           gauss_cost += (efc_Ma_in[worldid, ii] - qfrc_smooth_in[worldid, ii]) * (
             qacc_in[worldid, ii] - qacc_smooth_in[worldid, ii]
           )
-      wp.atomic_add(ctx_gauss_out, worldid, gauss_cost)
+      wp.atomic_add(ctx_gauss_out, worldid, 0.5 * gauss_cost)
-      wp.atomic_add(ctx_cost_out, worldid, gauss_cost)
+      wp.atomic_add(ctx_cost_out, worldid, 0.5 * gauss_cost)
 
   return kernel
 
@@ -2182,12 +2151,7 @@ def update_gradient_h_incremental_sparse(
         wp.atomic_add(ctx_h_out[worldid, colindj], colindi, h)
 
 
-def _update_constraint(
+def _update_constraint(m: types.Model, d: types.Data, ctx: SolverContext | InverseContext, track_changes: bool = False):
-    m: types.Model,
-    d: types.Data,
-    ctx: SolverContext | InverseContext,
-    track_changes: bool = False,
-):
   """Update constraint arrays after each solve iteration."""
   wp.launch(
     update_constraint_init_cost,
@@ -2217,30 +2181,16 @@ def _update_constraint(
     update_constraint_efc(track_changes),
     dim=(d.nworld, d.njmax),
     inputs=efc_inputs,
-      outputs=[
+    outputs=[d.efc.force, d.efc.state, ctx.cost, ctx.changed_efc_ids, ctx.changed_efc_count],
-          d.efc.force,
-          d.efc.state,
-          ctx.cost,
-          ctx.changed_efc_ids,
-          ctx.changed_efc_count,
-      ],
   )
 
   # qfrc_constraint = efc_J.T @ efc_force
-  if SPARSE_CONSTRAINT_JACOBIAN:
+  if m.is_sparse:
     d.qfrc_constraint.zero_()
     wp.launch(
       update_constraint_init_qfrc_constraint_sparse,
       dim=(d.nworld, d.njmax),
-        inputs=[
+      inputs=[d.nefc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, d.efc.force, ctx.done],
-            d.nefc,
-            d.efc.J_rownnz,
-            d.efc.J_rowadr,
-            d.efc.J_colind,
-            d.efc.J,
-            d.efc.force,
-            ctx.done,
-        ],
       outputs=[d.qfrc_constraint],
     )
   else:
@@ -2441,9 +2391,7 @@ def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int):
 
     nefc = nefc_in[worldid]
 
-    sum_val = wp.tile_load(
+    sum_val = wp.tile_load(qM_in[worldid], shape=(nv_pad, nv_pad), bounds_check=True)
-        qM_in[worldid], shape=(nv_pad, nv_pad), bounds_check=True
-    )
 
     # Each tile processes one output tile by looping over all constraints
     for k in range(0, njmax, TILE_SIZE_K):
@@ -2453,12 +2401,7 @@ def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int):
       # AD: leaving bounds-check disabled here because I'm not entirely sure that
       # everything always hits the fast path. The padding takes care of any
       #  potential OOB accesses.
-      J_kj = wp.tile_load(
+      J_kj = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE_K, nv_pad), offset=(k, 0), bounds_check=False)
-          efc_J_in[worldid],
-          shape=(TILE_SIZE_K, nv_pad),
-          offset=(k, 0),
-          bounds_check=False,
-      )
 
       # state check
       D_k = wp.tile_load(efc_D_in[worldid], shape=TILE_SIZE_K, offset=k, bounds_check=False)
@@ -2473,11 +2416,7 @@ def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int):
       active_tile = wp.tile_map(active_check, tid_tile, threshold_tile)
       D_k = wp.tile_map(wp.mul, active_tile, D_k)
 
-      J_ki = wp.tile_map(
+      J_ki = wp.tile_map(wp.mul, wp.tile_transpose(J_kj), wp.tile_broadcast(D_k, shape=(nv_pad, TILE_SIZE_K)))
-          wp.mul,
-          wp.tile_transpose(J_kj),
-          wp.tile_broadcast(D_k, shape=(nv_pad, TILE_SIZE_K)),
-      )
 
       sum_val += wp.tile_matmul(J_ki, J_kj)
 
@@ -2514,7 +2453,7 @@ def update_gradient_JTCJ_sparse(
   nblocks_perblock: int,
   dim_block: int,
   # Out:
-    h_out: wp.array3d(dtype=float),
+  ctx_h_out: wp.array3d(dtype=float),
 ):
   conid_start, elementid = wp.tid()
 
@@ -2529,20 +2468,37 @@ def update_gradient_JTCJ_sparse(
 
     worldid = contact_worldid_in[conid]
     if ctx_done_in[worldid]:
-      return
+      continue
 
     condim = contact_dim_in[conid]
 
     if condim == 1:
-      return
+      continue
 
     # check contact status
     if contact_dist_in[conid] - contact_includemargin_in[conid] >= 0.0:
-      return
+      continue
 
     efcid0 = contact_efc_address_in[conid, 0]
     if efc_state_in[worldid, efcid0] != types.ConstraintState.CONE:
-      return
+      continue
+
+    # All dims share the same sparsity pattern. Scan colind once to find
+    # the sparse positions of dof1id and dof2id. Skip if either is absent.
+    rownnz = efc_J_rownnz_in[worldid, efcid0]
+    rowadr0 = efc_J_rowadr_in[worldid, efcid0]
+    pos1 = int(-1)
+    pos2 = int(-1)
+    for k in range(rownnz):
+      col = efc_J_colind_in[worldid, 0, rowadr0 + k]
+      if col == dof1id:
+        pos1 = k
+      if col == dof2id:
+        pos2 = k
+      if pos1 >= 0 and pos2 >= 0:
+        break
+    if pos1 < 0 or pos2 < 0:
+      continue
 
     fri = contact_friction_in[conid]
     mu = fri[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
@@ -2551,7 +2507,7 @@ def update_gradient_JTCJ_sparse(
     dm = math.safe_div(efc_D_in[worldid, efcid0], mu2 * (1.0 + mu2))
 
     if dm == 0.0:
-      return
+      continue
 
     n = ctx_Jaref_in[worldid, efcid0] * mu
     u = types.vec6(n, 0.0, 0.0, 0.0, 0.0, 0.0)
@@ -2570,52 +2526,40 @@ def update_gradient_JTCJ_sparse(
     t = wp.max(t, types.MJ_MINVAL)
     ttt = wp.max(t * t * t, types.MJ_MINVAL)
 
+    # Precompute common subexpressions.
+    mu_over_t = math.safe_div(mu, t)
+    mu_n_over_ttt = mu * math.safe_div(n, ttt)
+    mu2_minus_mu_n_over_t = mu2 - mu * math.safe_div(n, t)
+
     h = float(0.0)
 
     for dim1id in range(condim):
       if dim1id == 0:
-        efcid1 = efcid0
+        rowadr1 = rowadr0
+        dm_fri1 = dm * mu
       else:
         efcid1 = contact_efc_address_in[conid, dim1id]
-
-      # TODO(team): improve performance for sparse code path
-      rownnz1 = efc_J_rownnz_in[worldid, efcid1]
         rowadr1 = efc_J_rowadr_in[worldid, efcid1]
+        dm_fri1 = dm * fri[dim1id - 1]
 
-      efc_J11 = float(0.0)
+      # Direct J reads using cached sparse positions.
-      efc_J12 = float(0.0)
+      efc_J11 = efc_J_in[worldid, 0, rowadr1 + pos1]
-      for i1 in range(rownnz1):
+      efc_J12 = efc_J_in[worldid, 0, rowadr1 + pos2]
-        sparseid1 = rowadr1 + i1
-        colind1 = efc_J_colind_in[worldid, 0, sparseid1]
-        if dof1id == colind1:
-          efc_J11 = efc_J_in[worldid, 0, sparseid1]
-        if dof2id == colind1:
-          efc_J12 = efc_J_in[worldid, 0, sparseid1]
-        if efc_J11 != 0.0 and efc_J12 != 0.0:
-          break
 
       ui = u[dim1id]
 
       for dim2id in range(0, dim1id + 1):
         if dim2id == 0:
-          efcid2 = efcid0
+          rowadr2 = rowadr0
+          dm_fri12 = dm_fri1 * mu
         else:
           efcid2 = contact_efc_address_in[conid, dim2id]
-
-        rownnz2 = efc_J_rownnz_in[worldid, efcid2]
           rowadr2 = efc_J_rowadr_in[worldid, efcid2]
+          dm_fri12 = dm_fri1 * fri[dim2id - 1]
 
-        efc_J21 = float(0.0)
+        # Direct J reads using cached sparse positions.
-        efc_J22 = float(0.0)
+        efc_J21 = efc_J_in[worldid, 0, rowadr2 + pos1]
-        for i2 in range(rownnz2):
+        efc_J22 = efc_J_in[worldid, 0, rowadr2 + pos2]
-          sparseid2 = rowadr2 + i2
-          colind2 = efc_J_colind_in[worldid, 0, sparseid2]
-          if dof1id == colind2:
-            efc_J21 = efc_J_in[worldid, 0, sparseid2]
-          if dof2id == colind2:
-            efc_J22 = efc_J_in[worldid, 0, sparseid2]
-          if efc_J21 != 0.0 and efc_J22 != 0.0:
-            break
 
         uj = u[dim2id]
 
@@ -2623,28 +2567,17 @@ def update_gradient_JTCJ_sparse(
         if dim1id == 0 and dim2id == 0:
           hcone = 1.0
         elif dim1id == 0:
-          hcone = -math.safe_div(mu, t) * uj
+          hcone = -mu_over_t * uj
         elif dim2id == 0:
-          hcone = -math.safe_div(mu, t) * ui
+          hcone = -mu_over_t * ui
         else:
-          hcone = mu * math.safe_div(n, ttt) * ui * uj
+          hcone = mu_n_over_ttt * ui * uj
 
           # add to diagonal: mu^2 - mu * n / t
           if dim1id == dim2id:
-            hcone += mu2 - mu * math.safe_div(n, t)
+            hcone += mu2_minus_mu_n_over_t
 
-        # pre and post multiply by diag(mu, friction) scale by dm
+        hcone *= dm_fri12
-        if dim1id == 0:
-          fri1 = mu
-        else:
-          fri1 = fri[dim1id - 1]
-
-        if dim2id == 0:
-          fri2 = mu
-        else:
-          fri2 = fri[dim2id - 1]
-
-        hcone *= dm * fri1 * fri2
 
         if hcone != 0.0:
           h += hcone * efc_J11 * efc_J22
@@ -2652,7 +2585,7 @@ def update_gradient_JTCJ_sparse(
           if dim1id != dim2id:
             h += hcone * efc_J12 * efc_J21
 
-    h_out[worldid, dof1id, dof2id] += h
+    ctx_h_out[worldid, dof1id, dof2id] += h
 
 
 @wp.kernel
@@ -2863,9 +2796,7 @@ def padding_h(nv: int, ctx_done_in: wp.array(dtype=bool), ctx_h_out: wp.array3d(
   ctx_h_out[worldid, dofid, dofid] = 1.0
 
 
-def _cholesky_factorize_solve(
+def _cholesky_factorize_solve(m: types.Model, d: types.Data, ctx: SolverContext):
-    m: types.Model, d: types.Data, ctx: SolverContext
-):
   """Cholesky factorize ctx.h and solve for Mgrad."""
   if m.nv <= _BLOCK_CHOLESKY_DIM:
     wp.launch_tiled(
@@ -2886,36 +2817,14 @@ def _cholesky_factorize_solve(
     wp.launch_tiled(
       update_gradient_cholesky_blocked(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad),
       dim=d.nworld,
-        inputs=[
+      inputs=[ctx.done, ctx.grad.reshape(shape=(d.nworld, ctx.grad.shape[1], 1)), ctx.h, ctx.hfactor],
-            ctx.done,
-            ctx.grad.reshape(shape=(d.nworld, ctx.grad.shape[1], 1)),
-            ctx.h,
-            ctx.hfactor,
-        ],
       outputs=[ctx.Mgrad.reshape(shape=(d.nworld, ctx.Mgrad.shape[1], 1))],
       block_dim=m.block_dim.update_gradient_cholesky_blocked,
     )
 
 
-def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
+@wp.kernel
-  # grad = Ma - qfrc_smooth - qfrc_constraint
+def _JTDAJ_sparse(
-  wp.launch(update_gradient_zero_grad_dot, dim=(d.nworld), inputs=[ctx.done], outputs=[ctx.grad_dot])
-
-  wp.launch(
-    update_gradient_grad,
-    dim=(d.nworld, m.nv),
-    inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, ctx.done],
-    outputs=[ctx.grad, ctx.grad_dot],
-  )
-
-  if m.opt.solver == types.SolverType.CG:
-    smooth.solve_m(m, d, ctx.Mgrad, ctx.grad)
-  elif m.opt.solver == types.SolverType.NEWTON:
-    # h = qM + (efc_J.T * efc_D * active) @ efc_J
-    if SPARSE_CONSTRAINT_JACOBIAN:
-      # TODO(team): improve performance for sparse code path
-      @wp.kernel(module="unique", enable_backward=False)
-      def _JTDAJ_sparse(
   # Data in:
   nefc_in: wp.array(dtype=int),
   efc_J_rownnz_in: wp.array2d(dtype=int),
@@ -2928,7 +2837,7 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
   ctx_done_in: wp.array(dtype=bool),
   # Out:
   h_out: wp.array3d(dtype=float),
-      ):
+):
   worldid, efcid = wp.tid()
 
   if ctx_done_in[worldid]:
@@ -2966,72 +2875,31 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
       if i != j:
         wp.atomic_add(h_out[worldid, colindj], colindi, h)
 
+
+def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
+  # grad = Ma - qfrc_smooth - qfrc_constraint
+  wp.launch(update_gradient_zero_grad_dot, dim=(d.nworld), inputs=[ctx.done], outputs=[ctx.grad_dot])
+
+  wp.launch(
+    update_gradient_grad,
+    dim=(d.nworld, m.nv),
+    inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, ctx.done],
+    outputs=[ctx.grad, ctx.grad_dot],
+  )
+
+  if m.opt.solver == types.SolverType.CG:
+    smooth.solve_m(m, d, ctx.Mgrad, ctx.grad)
+  elif m.opt.solver == types.SolverType.NEWTON:
+    # h = qM + (efc_J.T * efc_D * active) @ efc_J
+    if m.is_sparse:
+      ctx.h.zero_()
       wp.launch(
         _JTDAJ_sparse,
         dim=(d.nworld, d.njmax),
-          inputs=[
+        inputs=[d.nefc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, d.efc.D, d.efc.state, ctx.done],
-              d.nefc,
-              d.efc.J_rownnz,
-              d.efc.J_rowadr,
-              d.efc.J_colind,
-              d.efc.J,
-              d.efc.D,
-              d.efc.state,
-              ctx.done,
-          ],
         outputs=[ctx.h],
       )
 
-      if m.is_sparse:
-        wp.launch(
-            update_gradient_set_h_qM_lower_sparse,
-            dim=(d.nworld, m.qM_fullm_i.size),
-            inputs=[m.qM_fullm_i, m.qM_fullm_j, d.qM, ctx.done],
-            outputs=[ctx.h],
-        )
-      else:
-        # dense M: copy qM directly into h
-        @wp.kernel(module="unique", enable_backward=False)
-        def _set_h_qM_dense(
-            nv: int,
-            qM_in: wp.array3d(dtype=float),
-            ctx_done_in: wp.array(dtype=bool),
-            ctx_h_out: wp.array3d(dtype=float),
-        ):
-          worldid, i, j = wp.tid()
-          if ctx_done_in[worldid]:
-            return
-          if i >= nv or j >= nv:
-            return
-          if i >= j:
-            ctx_h_out[worldid, i, j] += qM_in[worldid, i, j]
-
-        wp.launch(
-            _set_h_qM_dense,
-            dim=(d.nworld, m.nv, m.nv),
-            inputs=[m.nv, d.qM, ctx.done],
-            outputs=[ctx.h],
-        )
-    elif m.is_sparse:
-      num_blocks_ceil = ceil(m.nv / types.TILE_SIZE_JTDAJ_SPARSE)
-      lower_triangle_dim = int(num_blocks_ceil * (num_blocks_ceil + 1) / 2)
-      with scoped_mathdx_gemm_disabled():
-        wp.launch_tiled(
-            update_gradient_JTDAJ_sparse_tiled(
-                types.TILE_SIZE_JTDAJ_SPARSE, d.njmax
-            ),
-            dim=(d.nworld, lower_triangle_dim),
-            inputs=[
-                d.nefc,
-                d.efc.J,
-                d.efc.D,
-                d.efc.state,
-                ctx.done,
-            ],
-            outputs=[ctx.h],
-            block_dim=m.block_dim.update_gradient_JTDAJ_sparse,
-        )
-
       wp.launch(
         update_gradient_set_h_qM_lower_sparse,
         dim=(d.nworld, m.qM_fullm_i.size),
@@ -3041,9 +2909,7 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
     else:
       with scoped_mathdx_gemm_disabled():
         wp.launch_tiled(
-            update_gradient_JTDAJ_dense_tiled(
+          update_gradient_JTDAJ_dense_tiled(m.nv_pad, types.TILE_SIZE_JTDAJ_DENSE, d.njmax),
-                m.nv_pad, types.TILE_SIZE_JTDAJ_DENSE, d.njmax
-            ),
           dim=d.nworld,
           inputs=[
             d.nefc,
@@ -3081,10 +2947,10 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
 
       nblocks_perblock = int((d.naconmax + dim_block - 1) / dim_block)
 
-      if SPARSE_CONSTRAINT_JACOBIAN:
+      if m.is_sparse:
         wp.launch(
           update_gradient_JTCJ_sparse,
-            dim=(d.naconmax, m.dof_tri_row.size),
+          dim=(dim_block, m.dof_tri_row.size),
           inputs=[
             m.opt.impratio_invsqrt,
             m.dof_tri_row,
@@ -3142,20 +3008,13 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
     raise ValueError(f"Unknown solver type: {m.opt.solver}")
 
 
-def _update_gradient_incremental(
+def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverContext):
-    m: types.Model, d: types.Data, ctx: SolverContext
-):
   """Incremental gradient update: update H for changed constraints + re-factorize.
 
   Skips the full J^T*D*J rebuild by applying only the delta from constraints
   that changed QUADRATIC state, then re-factorizes and solves.
   """
-  wp.launch(
+  wp.launch(update_gradient_zero_grad_dot, dim=(d.nworld), inputs=[ctx.done], outputs=[ctx.grad_dot])
-      update_gradient_zero_grad_dot,
-      dim=(d.nworld),
-      inputs=[ctx.done],
-      outputs=[ctx.grad_dot],
-  )
 
   wp.launch(
     update_gradient_grad,
@@ -3165,7 +3024,7 @@ def _update_gradient_incremental(
   )
 
   # Update lower triangle of H with delta from changed constraints
-  if SPARSE_CONSTRAINT_JACOBIAN:
+  if m.is_sparse:
     wp.launch(
       update_gradient_h_incremental_sparse,
       dim=(d.nworld, ctx.changed_efc_ids.shape[1]),
@@ -3347,10 +3206,7 @@ def _solver_iteration(
   # path in update_constraint_efc has early returns that skip state change
   # tracking, and the additional JTCJ Hessian term depends on Jaref which
   # changes every iteration.
-  incremental = (
+  incremental = m.opt.solver == types.SolverType.NEWTON and m.opt.cone != types.ConeType.ELLIPTIC
-      m.opt.solver == types.SolverType.NEWTON
-      and m.opt.cone != types.ConeType.ELLIPTIC
-  )
 
   if incremental:
     # Must complete before update_constraint_efc which atomically increments.
@@ -3424,15 +3280,7 @@ def init_context(m: types.Model, d: types.Data, ctx: SolverContext | InverseCont
   wp.launch(
     solve_init_jaref(m.is_sparse, m.nv, dofs_per_thread),
     dim=(d.nworld, d.njmax, threads_per_efc),
-      inputs=[
+    inputs=[d.nefc, d.qacc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, d.efc.aref],
-          d.nefc,
-          d.qacc,
-          d.efc.J_rownnz,
-          d.efc.J_rowadr,
-          d.efc.J_colind,
-          d.efc.J,
-          d.efc.aref,
-      ],
     outputs=[ctx.Jaref],
   )
 

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

@@ -18,18 +18,52 @@ from typing import Optional, Tuple
 import warp as wp
 
 from mujoco.mjx.third_party.mujoco_warp._src.math import motion_cross
+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 Data
+from mujoco.mjx.third_party.mujoco_warp._src.types import DynType
 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 vec5
+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
 
 wp.set_module_options({"enable_backward": False})
 
 
+# TODO(team): kernel analyzer array slice?
+@wp.func
+def next_act(
+  # Model:
+  opt_timestep: float,  # kernel_analyzer: ignore
+  actuator_dyntype: int,  # kernel_analyzer: ignore
+  actuator_dynprm: vec10f,  # kernel_analyzer: ignore
+  actuator_actrange: wp.vec2,  # kernel_analyzer: ignore
+  # Data In:
+  act_in: float,  # kernel_analyzer: ignore
+  act_dot_in: float,  # kernel_analyzer: ignore
+  # In:
+  act_dot_scale: float,
+  clamp: bool,
+) -> float:
+  # advance actuation
+  if actuator_dyntype == DynType.FILTEREXACT:
+    tau = wp.max(MJ_MINVAL, actuator_dynprm[0])
+    act = act_in + act_dot_scale * act_dot_in * tau * (1.0 - wp.exp(-opt_timestep / tau))
+  elif actuator_dyntype == DynType.USER:
+    return act_in
+  else:
+    act = act_in + act_dot_scale * act_dot_in * opt_timestep
+
+  # clamp to actrange
+  if clamp:
+    act = wp.clamp(act, actuator_actrange[0], actuator_actrange[1])
+
+  return act
+
+
 @cache_kernel
 def mul_m_sparse(check_skip: bool):
   @wp.kernel(module="unique")

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

@@ -33,11 +33,8 @@ MJ_MAX_EPAFACES = 5
 TILE_SIZE_JTDAJ_SPARSE = 16
 TILE_SIZE_JTDAJ_DENSE = 16
 
-# TODO(team): remove after improving performance for sparse constraint jacobian
+# maximum number of plugin attributes
-SPARSE_CONSTRAINT_JACOBIAN = False
+_NPLUGINATTR = 128
-
-# 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
@@ -53,7 +50,6 @@ class BlockDim:
   # forward
   euler_dense: int = 32
   actuator_velocity: int = 32
-  tendon_velocity: int = 32
   # ray
   ray: int = 64
   # sensor
@@ -63,6 +59,7 @@ class BlockDim:
   cholesky_factorize: int = 32
   cholesky_solve: int = 32
   cholesky_factorize_solve: int = 32
+  solve_LD_sparse_fused: int = 64
   # solver
   update_gradient_cholesky: int = 64
   update_gradient_cholesky_blocked: int = 32
@@ -351,6 +348,7 @@ class GeomType(enum.IntEnum):
     BOX: box
     MESH: mesh
     SDF: sdf
+    FLEX: flex
   """
 
   PLANE = mujoco.mjtGeom.mjGEOM_PLANE
@@ -362,6 +360,7 @@ class GeomType(enum.IntEnum):
   BOX = mujoco.mjtGeom.mjGEOM_BOX
   MESH = mujoco.mjtGeom.mjGEOM_MESH
   SDF = mujoco.mjtGeom.mjGEOM_SDF
+  FLEX = mujoco.mjtGeom.mjGEOM_FLEX
   # unsupported: NGEOMTYPES, ARROW*, LINE, SKIN, LABEL, NONE
 
 
@@ -662,6 +661,10 @@ class vec11f(wp.types.vector(length=11, dtype=float)):
   pass
 
 
+class vec_pluginattr(wp.types.vector(length=_NPLUGINATTR, dtype=float)):
+  pass
+
+
 class mat23f(wp.types.matrix(shape=(2, 3), dtype=float)):
   pass
 
@@ -679,6 +682,7 @@ vec6 = vec6f
 vec8 = vec8f
 vec10 = vec10f
 vec11 = vec11f
+vec128 = vec_pluginattr
 mat23 = mat23f
 mat43 = mat43f
 mat63 = mat63f
@@ -841,6 +845,7 @@ class Model:
     nflexelem: number of elements in all flexes
     nflexelemdata: number of element vertex ids in all flexes
     nflexelemedge: number of element edge ids in all flexes
+    nflexshelldata: number of shell fragment vertex ids in all flexes
     nJfe: number of non-zeros in sparse flexedge Jacobian
     nmesh: number of meshes
     nmeshvert: number of vertices for all meshes
@@ -857,6 +862,7 @@ class Model:
     nexclude: number of excluded geom pairs
     neq: number of equality constraints
     ntendon: number of tendons
+    nJten: number of non-zeros in sparse tendon Jacobian
     nwrap: number of wrap objects in all tendon paths
     nsensor: number of sensors
     nmocap: number of mocap bodies
@@ -973,6 +979,11 @@ class Model:
     light_poscom0: global position rel. to sub-com in qpos0  (*, nlight, 3)
     light_pos0: global position rel. to body in qpos0        (*, nlight, 3)
     light_dir0: global direction in qpos0                    (*, nlight, 3)
+    flex_contype: flex contact type                          (nflex,)
+    flex_conaffinity: flex contact affinity                  (nflex,)
+    flex_condim: contact dimensionality (1, 3, 4, 6)         (nflex,)
+    flex_friction: friction for (slide, spin, roll)          (nflex, 3)
+    flex_margin: detect contact if dist<margin               (nflex,)
     flex_dim: 1: lines, 2: triangles, 3: tetrahedra          (nflex,)
     flex_vertadr: first vertex address                       (nflex,)
     flex_vertnum: number of vertices                         (nflex,)
@@ -980,17 +991,24 @@ class Model:
     flex_edgenum: number of edges                            (nflex,)
     flex_elemadr: first element address                      (nflex,)
     flex_elemnum: number of elements                         (nflex,)
-    flex_elemedgeadr: first element address                  (nflex,)
+    flex_elemdataadr: first element vertex id address        (nflex,)
+    flex_elemedgeadr: first element edge id address          (nflex,)
+    flex_shellnum: number of shells                          (nflex,)
+    flex_shelldataadr: first shell data address              (nflex,)
     flex_vertbodyid: vertex body ids                         (nflexvert,)
     flex_edge: edge vertex ids (2 per edge)                  (nflexedge, 2)
     flex_edgeflap: adjacent vertex ids (dim=2 only)          (nflexedge, 2)
     flex_elem: element vertex ids (dim+1 per elem)           (nflexelemdata,)
     flex_elemedge: element edge ids                          (nflexelemedge,)
+    flex_shell: shell fragment vertex ids (dim per frag)     (nflexshelldata,)
+    flex_vert: vertex local positions                        (nflexvert, 3)
     flexedge_length0: edge lengths in qpos0                  (nflexedge,)
     flexedge_invweight0: inv. inertia for the edge           (nflexedge,)
+    flex_radius: radius around primitive element             (nflex,)
     flex_stiffness: finite element stiffness matrix          (nflexelem, 21)
     flex_bending: bending stiffness                          (nflexedge, 17)
     flex_damping: Rayleigh's damping coefficient             (nflex,)
+    flex_centered: flex vertices are centered at body origin (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,)
@@ -1020,8 +1038,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,
+    mat_texid: texture id for rendering                      (*, nmat, mjNTEXROLE)
-      mjNTEXROLE)
     mat_texrepeat: texture repeat for rendering              (*, nmat, 2)
     mat_rgba: rgba                                           (*, nmat, 4)
     pair_dim: contact dimensionality                         (npair,)
@@ -1044,16 +1061,15 @@ class Model:
     eq_data: numeric data for constraint                     (*, neq, mjNEQDATA)
     tendon_adr: address of first object in tendon's path     (ntendon,)
     tendon_num: number of objects in tendon's path           (ntendon,)
+    ten_J_rownnz: number of non-zeros in each tendon row     (ntendon,)
+    ten_J_rowadr: row start address for sparse ten_J         (ntendon,)
+    ten_J_colind: column indices in sparse ten_J             (nJten,)
     tendon_limited: does tendon have length limits           (ntendon,)
     tendon_actfrclimited: does ten have actuator force limit (ntendon,)
-    tendon_solref_lim: constraint solver reference: limit    (*, ntendon,
+    tendon_solref_lim: constraint solver reference: limit    (*, ntendon, mjNREF)
-      mjNREF)
+    tendon_solimp_lim: constraint solver impedance: limit    (*, ntendon, mjNIMP)
-    tendon_solimp_lim: constraint solver impedance: limit    (*, ntendon,
+    tendon_solref_fri: constraint solver reference: friction (*, ntendon, mjNREF)
-      mjNIMP)
+    tendon_solimp_fri: constraint solver impedance: friction (*, ntendon, mjNIMP)
-    tendon_solref_fri: constraint solver reference: friction (*, ntendon,
-      mjNREF)
-    tendon_solimp_fri: constraint solver impedance: friction (*, ntendon,
-      mjNIMP)
     tendon_range: tendon length limits                       (*, ntendon, 2)
     tendon_actfrcrange: range of total actuator force        (*, ntendon, 2)
     tendon_margin: min distance for limit detection          (*, ntendon)
@@ -1099,11 +1115,12 @@ class Model:
     sensor_adr: address in sensor array                      (nsensor,)
     sensor_cutoff: cutoff for real and positive; 0: ignore   (nsensor,)
     plugin: globally registered plugin slot number           (nplugin,)
-    plugin_attr: config attributes of geom plugin            (nplugin, 3)
+    plugin_attr: config attributes of geom plugin            (nplugin, _NPLUGINATTR)
     M_rownnz: number of non-zeros in each row of qM          (nv,)
     M_rowadr: index of each row in qM                        (nv,)
     M_colind: column indices of non-zeros in qM              (nM,)
     mapM2M: index mapping from M (legacy) to M (CSR)         (nC)
+    flex_vertflexid: flex id for each flex vertex            (nflexvert,)
 
   warp only fields:
     callback: custom physics callbacks
@@ -1120,14 +1137,12 @@ class Model:
     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
+    has_fluid: True if wind, density, or viscosity are non-zero at put_model time
-      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 +
+    body_branch_start: start index in body_branches for each branch   (nbranch + 1,)
-      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
@@ -1153,6 +1168,7 @@ class Model:
     tendon_site_pair_adr: site pair tendon address
     tendon_geom_adr: geom tendon address
     tendon_limited_adr: addresses for limited tendons
+    max_ten_J_rownnz: maximum number of non-zeros in a tendon row
     ten_wrapadr_site: wrap object starting address for sites
     ten_wrapnum_site: number of site wrap objects per tendon
     wrap_jnt_adr: addresses for joint tendon wrap object
@@ -1189,6 +1205,8 @@ class Model:
     taxel_sensorid: address for tactile sensors
     qM_tiles: tiling configuration
     qLD_updates: tuple of index triples for sparse factorization
+    qLD_all_updates: tuple of all levels concatenated
+    qLD_level_offsets: tuple of start offsets for each level
     qM_fullm_i: sparse mass matrix addressing
     qM_fullm_j: sparse mass matrix addressing
     qM_mulm_rowadr: sparse matmul row pointers
@@ -1216,6 +1234,7 @@ class Model:
   nflexelem: int
   nflexelemdata: int
   nflexelemedge: int
+  nflexshelldata: int
   nJfe: int
   nmesh: int
   nmeshvert: int
@@ -1232,6 +1251,7 @@ class Model:
   nexclude: int
   neq: int
   ntendon: int
+  nJten: int
   nwrap: int
   nsensor: int
   nmocap: int
@@ -1348,6 +1368,11 @@ class Model:
   light_poscom0: array("*", "nlight", wp.vec3)
   light_pos0: array("*", "nlight", wp.vec3)
   light_dir0: array("*", "nlight", wp.vec3)
+  flex_contype: array("nflex", int)
+  flex_conaffinity: array("nflex", int)
+  flex_condim: array("nflex", int)
+  flex_friction: array("nflex", wp.vec3)
+  flex_margin: array("nflex", float)
   flex_dim: array("nflex", int)
   flex_vertadr: array("nflex", int)
   flex_vertnum: array("nflex", int)
@@ -1355,17 +1380,24 @@ class Model:
   flex_edgenum: array("nflex", int)
   flex_elemadr: array("nflex", int)
   flex_elemnum: array("nflex", int)
+  flex_elemdataadr: array("nflex", int)
   flex_elemedgeadr: array("nflex", int)
+  flex_shellnum: array("nflex", int)
+  flex_shelldataadr: array("nflex", int)
   flex_vertbodyid: array("nflexvert", int)
   flex_edge: array("nflexedge", wp.vec2i)
   flex_edgeflap: array("nflexedge", wp.vec2i)
   flex_elem: array("nflexelemdata", int)
   flex_elemedge: array("nflexelemedge", int)
+  flex_shell: array("nflexshelldata", int)
+  flex_vert: array("nflexvert", wp.vec3)
   flexedge_length0: array("nflexedge", float)
   flexedge_invweight0: array("nflexedge", float)
+  flex_radius: array("nflex", float)
   flex_stiffness: array("nflexelem", 21, float)
   flex_bending: array("nflexedge", 17, float)
   flex_damping: array("nflex", float)
+  flex_centered: array("nflex", bool)
   flexedge_J_rownnz: array("nflexedge", int)
   flexedge_J_rowadr: array("nflexedge", int)
   flexedge_J_colind: array("nJfe", int)
@@ -1418,6 +1450,9 @@ class Model:
   eq_data: array("*", "neq", vec11)
   tendon_adr: array("ntendon", int)
   tendon_num: array("ntendon", int)
+  ten_J_rownnz: array("ntendon", int)
+  ten_J_rowadr: array("ntendon", int)
+  ten_J_colind: array("nJten", int)
   tendon_limited: array("ntendon", int)
   tendon_actfrclimited: array("ntendon", bool)
   tendon_solref_lim: array("*", "ntendon", wp.vec2)
@@ -1469,11 +1504,12 @@ class Model:
   sensor_adr: array("nsensor", int)
   sensor_cutoff: array("nsensor", float)
   plugin: array("nplugin", int)
-  plugin_attr: array("nplugin", wp.vec3f)
+  plugin_attr: array("nplugin", vec_pluginattr)
   M_rownnz: array("nv", int)
   M_rowadr: array("nv", int)
   M_colind: array("nC", int)
   mapM2M: array("nC", int)
+  flex_vertflexid: array("nflexvert", int)
   # warp only fields:
   callback: Callback
   nbranch: int
@@ -1515,6 +1551,7 @@ class Model:
   tendon_site_pair_adr: wp.array(dtype=int)
   tendon_geom_adr: wp.array(dtype=int)
   tendon_limited_adr: wp.array(dtype=int)
+  max_ten_J_rownnz: int
   ten_wrapadr_site: wp.array(dtype=int)
   ten_wrapnum_site: wp.array(dtype=int)
   wrap_jnt_adr: wp.array(dtype=int)
@@ -1546,6 +1583,8 @@ class Model:
   taxel_sensorid: wp.array(dtype=int)
   qM_tiles: tuple[TileSet, ...]
   qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...]
+  qLD_all_updates: wp.array(dtype=wp.vec3i)
+  qLD_level_offsets: wp.array(dtype=int)
   qM_fullm_i: wp.array(dtype=int)
   qM_fullm_j: wp.array(dtype=int)
   # Gather-based sparse mul_m indices (thread per DOF, no atomics)
@@ -1598,6 +1637,8 @@ class Contact:
   solimp: array("naconmax", vec5)
   dim: array("naconmax", int)
   geom: array("naconmax", wp.vec2i)
+  flex: array("naconmax", wp.vec2i)
+  vert: array("naconmax", wp.vec2i)
   efc_address: array("naconmax", "nmaxpyramid", int)
   worldid: array("naconmax", int)
   type: array("naconmax", int)
@@ -1611,14 +1652,14 @@ class Constraint:
   Attributes:
     type: constraint type (ConstraintType)            (nworld, njmax)
     id: id of object of specific type                 (nworld, njmax)
-    J_rownnz: number of non-zeros in J row            (nworld, 0) dense (nworld,
+    J_rownnz: number of non-zeros in J row            (nworld, 0) dense
-      njmax) sparse
+                                                      (nworld, njmax) sparse
-    J_rowadr: row start address in colind array       (nworld, 0) dense (nworld,
+    J_rowadr: row start address in colind array       (nworld, 0) dense
-      njmax) sparse
+                                                      (nworld, njmax) sparse
     J_colind: column indices in J                     (nworld, 0, 0) dense
                                                       (nworld, 1, njmax * nv) sparse
-    J: constraint Jacobian                            (nworld, njmax_pad,
+    J: constraint Jacobian                            (nworld, njmax_pad, nv_pad) dense
-      nv_pad) dense (nworld, 1, njmax * nv) sparse
+                                                      (nworld, 1, njmax * nv) sparse
     pos: constraint position (equality, contact)      (nworld, njmax)
     margin: inclusion margin (contact)                (nworld, njmax)
     D: constraint mass                                (nworld, njmax_pad)
@@ -1627,7 +1668,6 @@ class Constraint:
     frictionloss: frictionloss (friction)             (nworld, njmax)
     force: constraint force in constraint space       (nworld, njmax)
     state: constraint state                           (nworld, njmax_pad)
-
   warp only fields:
     Ma: M*qacc                                        (nworld, nv)
   """
@@ -1668,88 +1708,55 @@ class Data:
     qacc_warmstart: acceleration used for warmstart             (nworld, nv)
     ctrl: control                                               (nworld, nu)
     qfrc_applied: applied generalized force                     (nworld, nv)
-    xfrc_applied: applied Cartesian force/torque                (nworld, nbody,
+    xfrc_applied: applied Cartesian force/torque                (nworld, nbody, 6)
-      6)
     eq_active: enable/disable constraints                       (nworld, neq)
-    mocap_pos: position of mocap bodies                         (nworld, nmocap,
+    mocap_pos: position of mocap bodies                         (nworld, nmocap, 3)
-      3)
+    mocap_quat: orientation of mocap bodies                     (nworld, nmocap, 4)
-    mocap_quat: orientation of mocap bodies                     (nworld, nmocap,
-      4)
     qacc: acceleration                                          (nworld, nv)
     act_dot: time-derivative of actuator activation             (nworld, na)
-    sensordata: sensor data array                               (nworld,
+    sensordata: sensor data array                               (nworld, nsensordata,)
-      nsensordata,)
+    xpos: Cartesian position of body frame                      (nworld, nbody, 3)
-    xpos: Cartesian position of body frame                      (nworld, nbody,
+    xquat: Cartesian orientation of body frame                  (nworld, nbody, 4)
-      3)
+    xmat: Cartesian orientation of body frame                   (nworld, nbody, 3, 3)
-    xquat: Cartesian orientation of body frame                  (nworld, nbody,
+    xipos: Cartesian position of body com                       (nworld, nbody, 3)
-      4)
+    ximat: Cartesian orientation of body inertia                (nworld, nbody, 3, 3)
-    xmat: Cartesian orientation of body frame                   (nworld, nbody,
+    xanchor: Cartesian position of joint anchor                 (nworld, njnt, 3)
-      3, 3)
+    xaxis: Cartesian joint axis                                 (nworld, njnt, 3)
-    xipos: Cartesian position of body com                       (nworld, nbody,
+    geom_xpos: Cartesian geom position                          (nworld, ngeom, 3)
-      3)
+    geom_xmat: Cartesian geom orientation                       (nworld, ngeom, 3, 3)
-    ximat: Cartesian orientation of body inertia                (nworld, nbody,
+    site_xpos: Cartesian site position                          (nworld, nsite, 3)
-      3, 3)
+    site_xmat: Cartesian site orientation                       (nworld, nsite, 3, 3)
-    xanchor: Cartesian position of joint anchor                 (nworld, njnt,
+    cam_xpos: Cartesian camera position                         (nworld, ncam, 3)
-      3)
+    cam_xmat: Cartesian camera orientation                      (nworld, ncam, 3, 3)
-    xaxis: Cartesian joint axis                                 (nworld, njnt,
+    light_xpos: Cartesian light position                        (nworld, nlight, 3)
-      3)
+    light_xdir: Cartesian light direction                       (nworld, nlight, 3)
-    geom_xpos: Cartesian geom position                          (nworld, ngeom,
+    subtree_com: center of mass of each subtree                 (nworld, nbody, 3)
-      3)
-    geom_xmat: Cartesian geom orientation                       (nworld, ngeom,
-      3, 3)
-    site_xpos: Cartesian site position                          (nworld, nsite,
-      3)
-    site_xmat: Cartesian site orientation                       (nworld, nsite,
-      3, 3)
-    cam_xpos: Cartesian camera position                         (nworld, ncam,
-      3)
-    cam_xmat: Cartesian camera orientation                      (nworld, ncam,
-      3, 3)
-    light_xpos: Cartesian light position                        (nworld, nlight,
-      3)
-    light_xdir: Cartesian light direction                       (nworld, nlight,
-      3)
-    subtree_com: center of mass of each subtree                 (nworld, nbody,
-      3)
     cdof: com-based motion axis of each dof (rot:lin)           (nworld, nv, 6)
-    cinert: com-based body inertia and mass                     (nworld, nbody,
+    cinert: com-based body inertia and mass                     (nworld, nbody, 10)
-      10)
+    flexvert_xpos: cartesian flex vertex positions              (nworld, nflexvert, 3)
-    flexvert_xpos: cartesian flex vertex positions              (nworld,
-      nflexvert, 3)
     flexedge_J: edge length Jacobian                            (nworld, nJfe)
-    flexedge_length: flex edge lengths                          (nworld,
+    flexedge_length: flex edge lengths                          (nworld, nflexedge, 1)
-      nflexedge, 1)
+    ten_wrapadr: start address of tendon's path                 (nworld, ntendon)
-    ten_wrapadr: start address of tendon's path                 (nworld,
+    ten_wrapnum: number of wrap points in path                  (nworld, ntendon)
-      ntendon)
+    ten_J: tendon Jacobian                                      (nworld, nJten)
-    ten_wrapnum: number of wrap points in path                  (nworld,
+    ten_length: tendon lengths                                  (nworld, ntendon)
-      ntendon)
+    wrap_obj: geomid; -1: site; -2: pulley                      (nworld, nwrap, 2)
-    ten_J: tendon Jacobian                                      (nworld,
+    wrap_xpos: Cartesian 3D points in all paths                 (nworld, nwrap, 6)
-      ntendon, nv)
-    ten_length: tendon lengths                                  (nworld,
-      ntendon)
-    wrap_obj: geomid; -1: site; -2: pulley                      (nworld, nwrap,
-      2)
-    wrap_xpos: Cartesian 3D points in all paths                 (nworld, nwrap,
-      6)
     actuator_length: actuator lengths                           (nworld, nu)
     moment_rownnz: number of non-zeros in actuator_moment row   (nworld, nu)
     moment_rowadr: row start address in actuator_moment         (nworld, nu)
     moment_colind: column indices in sparse actuator_moment     (nworld, nJmom)
     actuator_moment: actuator moments                           (nworld, nJmom)
-    crb: com-based composite inertia and mass                   (nworld, nbody,
+    crb: com-based composite inertia and mass                   (nworld, nbody, 10)
-      10)
+    qM: total inertia                                           (nworld, nv, nv) if dense
-    qM: total inertia                                           (nworld, nv, nv)
+                                                                (nworld, 1, nM) if sparse
-      if dense (nworld, 1, nM) if sparse
+    qLD: L'*D*L factorization of M                              (nworld, nv, nv) if dense
-    qLD: L'*D*L factorization of M                              (nworld, nv, nv)
+                                                                (nworld, 1, nC) if sparse
-      if dense (nworld, 1, nC) if sparse
     qLDiagInv: 1/diag(D)                                        (nworld, nv)
-    flexedge_velocity: flex edge velocities                     (nworld,
+    flexedge_velocity: flex edge velocities                     (nworld, nflexedge)
-      nflexedge)
+    ten_velocity: tendon velocities                             (nworld, ntendon)
-    ten_velocity: tendon velocities                             (nworld,
-      ntendon)
     actuator_velocity: actuator velocities                      (nworld, nu)
-    cvel: com-based velocity (rot:lin)                          (nworld, nbody,
+    cvel: com-based velocity (rot:lin)                          (nworld, nbody, 6)
-      6)
     cdof_dot: time-derivative of cdof (rot:lin)                 (nworld, nv, 6)
     qfrc_bias: C(qpos,qvel)                                     (nworld, nv)
     qfrc_spring: passive spring force                           (nworld, nv)
@@ -1757,23 +1764,19 @@ class Data:
     qfrc_gravcomp: passive gravity compensation force           (nworld, nv)
     qfrc_fluid: passive fluid force                             (nworld, nv)
     qfrc_passive: total passive force                           (nworld, nv)
-    subtree_linvel: linear velocity of subtree com              (nworld, nbody,
+    subtree_linvel: linear velocity of subtree com              (nworld, nbody, 3)
-      3)
+    subtree_angmom: angular momentum about subtree com          (nworld, nbody, 3)
-    subtree_angmom: angular momentum about subtree com          (nworld, nbody,
-      3)
     actuator_force: actuator force in actuation space           (nworld, nu)
     qfrc_actuator: actuator force                               (nworld, nv)
     qfrc_smooth: net unconstrained force                        (nworld, nv)
     qacc_smooth: unconstrained acceleration                     (nworld, nv)
     qfrc_constraint: constraint force                           (nworld, nv)
     qfrc_inverse: net external force; should equal:             (nworld, nv)
-      qfrc_applied + J.T @ xfrc_applied + qfrc_actuator
+                  qfrc_applied + J.T @ xfrc_applied
-    cacc: com-based acceleration                                (nworld, nbody,
+                  + qfrc_actuator
-      6)
+    cacc: com-based acceleration                                (nworld, nbody, 6)
-    cfrc_int: com-based interaction force with parent           (nworld, nbody,
+    cfrc_int: com-based interaction force with parent           (nworld, nbody, 6)
-      6)
+    cfrc_ext: com-based external force on body                  (nworld, nbody, 6)
-    cfrc_ext: com-based external force on body                  (nworld, nbody,
-      6)
     contact: contact data
     efc: constraint data
     tree_island: island ID per tree (-1 if unconstrained)       (nworld, ntree)
@@ -1784,6 +1787,7 @@ class Data:
     naccdmax: maximum number of contacts for CCD (all worlds)
     njmax: maximum number of constraints per world
     njmax_pad: njmax rounded up to the nearest multiple of TILE_SIZE_JTDAJ
+    njmax_nnz: number of non-zeros in constraint Jacobian
     nacon: number of detected contacts (across all worlds)      (1,)
     ncollision: collision count from broadphase                 (1,)
   """
@@ -1832,7 +1836,7 @@ class Data:
   flexedge_length: array("nworld", "nflexedge", float)
   ten_wrapadr: array("nworld", "ntendon", int)
   ten_wrapnum: array("nworld", "ntendon", int)
-  ten_J: array("nworld", "ntendon", "nv", float)
+  ten_J: array("nworld", "nJten", float)
   ten_length: array("nworld", "ntendon", float)
   wrap_obj: array("nworld", "nwrap", wp.vec2i)
   wrap_xpos: array("nworld", "nwrap", wp.spatial_vector)
@@ -1877,6 +1881,7 @@ class Data:
   naccdmax: int
   njmax: int
   njmax_pad: int
+  njmax_nnz: int
   nacon: array(1, int)
   ncollision: array(1, int)
 
@@ -1905,21 +1910,12 @@ class RenderContext:
     hfield_registry: hfield BVH id to warp mesh mapping
     hfield_bvh_id: hfield BVH ids
     hfield_bounds_size: hfield bounds half-extents
-    flex_mesh: flex mesh
+    flex_mesh_registry: per-flex mesh BVH registry (prevents garbage collection)
     flex_rgba: flex rgba
-    flex_bvh_id: flex BVH id
+    flex_bvh_id: per-flex BVH ids
-    flex_face_point: flex face points
+    flex_group_root: per-flex group roots (nworld x n_flex_bvh)
-    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
+    flex_dim: flex dimension per flex (1D/2D/3D)
     bvh: scene BVH
     bvh_id: scene BVH id
     lower: lower bounds
@@ -1929,12 +1925,13 @@ class RenderContext:
     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
+    seg_data: segmentation data (per-pixel geom IDs)
+    seg_adr: segmentation addresses
+    render_seg: per-camera segmentation render flags
     znear: near plane distance
     total_rays: total number of rays
   """
@@ -1954,27 +1951,20 @@ class RenderContext:
   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("*", wp.Texture2D)
-  textures: array("*", TEXTURE_DTYPE)
+  textures_registry: list[wp.Texture2D]
-  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_mesh_registry: dict
   flex_rgba: array("nflex", wp.vec4)
-  flex_bvh_id: wp.uint64
+  flex_bvh_id: array("*", wp.uint64)
-  flex_face_point: array("*", wp.vec3)
+  flex_group_root: array("nworld", "*", int)
-  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_nflexgeom: int
+  flex_dim_np: array("nflex", int)
+  flex_geom_flexid: array("*", int)
+  flex_geom_edgeid: array("*", int)
   bvh: wp.Bvh
   bvh_id: wp.uint64
   lower: array("*", wp.vec3)
@@ -1988,5 +1978,8 @@ class RenderContext:
   depth_adr: array("ncam", int)
   render_rgb: array("ncam", bool)
   render_depth: array("ncam", bool)
+  seg_data: array("*", int)
+  seg_adr: array("ncam", int)
+  render_seg: array("ncam", bool)
   znear: float
   total_rays: int

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

@@ -37,9 +37,7 @@ def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
   """
   # Split on both '.' and '-'
   parts = re.split(r"[.\-]", version_str)
-  return tuple(
+  return tuple([(0, int(p)) if p.isdigit() else (-1, p) for p in parts] + [(0, 0)])
-      [(0, int(p)) if p.isdigit() else (-1, p) for p in parts] + [(0, 0)]
-  )
 
 
 def check_version(spec: str) -> bool:
@@ -65,9 +63,7 @@ def check_version(spec: str) -> bool:
   """
   match = re.match(r"^([a-zA-Z0-9_\-]+)(>=|<=|>|<|==|!=)(.+)$", spec)
   if not match:
-    raise ValueError(
+    raise ValueError(f"Invalid version spec '{spec}'. Expected format: 'package>=version'")
-        f"Invalid version spec '{spec}'. Expected format: 'package>=version'"
-    )
   package_name, op, version_str = match.groups()
 
   required_version = _parse_version(version_str)

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

@@ -28,7 +28,7 @@ requires-python = ">=3.10"
 dependencies = [
     "absl-py",
     "etils[epath]",
-    "mujoco>=3.5.0",
+    "mujoco>=3.6.0",
     "numpy",
     "warp-lang>=1.12",
 ]
@@ -55,7 +55,7 @@ dev = [
     "ruff",
     "pygls>=1.0.0,<2.0.0",
     "lsprotocol>=2023.0.1,<2024.0.0",
-    "mujoco>=3.5.0.dev0",
+    "mujoco>=3.6.0.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

@@ -18,7 +18,7 @@
 Usage: mjwarp-viewer <mjcf XML path> [flags]
 
 Example:
-  mjwarp-viewer benchmark/humanoid/humanoid.xml -o "opt.solver=cg"
+  mjwarp-viewer benchmarks/humanoid/humanoid.xml -o "opt.solver=cg"
 """
 
 import copy
@@ -56,6 +56,7 @@ _CLEAR_WARP_CACHE = flags.DEFINE_bool("clear_warp_cache", False, "Clear warp cac
 _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.")
+_NJMAX_NNZ = flags.DEFINE_integer("njmax_nnz", None, "Maximum number of non-zeros in constraint Jacobian.")
 _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.")
@@ -149,7 +150,7 @@ def _main(argv: Sequence[str]) -> None:
     override_model(mjm, _OVERRIDE.value)
     m = mjw.put_model(mjm)
     override_model(m, _OVERRIDE.value)
-    d = mjw.put_data(mjm, mjd, nconmax=_NCONMAX.value, njmax=_NJMAX.value, nccdmax=_NCCDMAX.value)
+    d = mjw.put_data(mjm, mjd, nconmax=_NCONMAX.value, njmax=_NJMAX.value, njmax_nnz=_NJMAX_NNZ.value, nccdmax=_NCCDMAX.value)
     graph = _compile_step(m, d) if wp.get_device().is_cuda else None
     if graph is None:
       mjw.step(m, d)  # warmup step

mjx/mujoco/mjx/warp/bvh.py

@@ -48,20 +48,27 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _refit_bvh_shim(
     # Model
     nworld: int,
     flex_dim: wp.array(dtype=int),
+    flex_edge: wp.array(dtype=wp.vec2i),
     flex_elem: wp.array(dtype=int),
+    flex_elemadr: wp.array(dtype=int),
+    flex_elemdataadr: wp.array(dtype=int),
     flex_elemnum: wp.array(dtype=int),
+    flex_radius: wp.array(dtype=float),
+    flex_shell: wp.array(dtype=int),
+    flex_shelldataadr: wp.array(dtype=int),
     flex_vertadr: wp.array(dtype=int),
+    flex_vertnum: wp.array(dtype=int),
     geom_dataid: wp.array(dtype=int),
     geom_size: wp.array2d(dtype=wp.vec3),
     geom_type: wp.array(dtype=int),
     nflex: int,
-    nflexelemdata: int,
+    nflexelem: int,
-    nflexvert: int,
     # Data
     flexvert_xpos: wp.array2d(dtype=wp.vec3),
     geom_xmat: wp.array2d(dtype=wp.mat33),
@@ -77,15 +84,21 @@ def _refit_bvh_shim(
   _d.efc = _e
   _d.contact = _c
   _m.flex_dim = flex_dim
+  _m.flex_edge = flex_edge
   _m.flex_elem = flex_elem
+  _m.flex_elemadr = flex_elemadr
+  _m.flex_elemdataadr = flex_elemdataadr
   _m.flex_elemnum = flex_elemnum
+  _m.flex_radius = flex_radius
+  _m.flex_shell = flex_shell
+  _m.flex_shelldataadr = flex_shelldataadr
   _m.flex_vertadr = flex_vertadr
+  _m.flex_vertnum = flex_vertnum
   _m.geom_dataid = geom_dataid
   _m.geom_size = geom_size
   _m.geom_type = geom_type
   _m.nflex = nflex
-  _m.nflexelemdata = nflexelemdata
+  _m.nflexelem = nflexelem
-  _m.nflexvert = nflexvert
   _d.flexvert_xpos = flexvert_xpos
   _d.geom_xmat = geom_xmat
   _d.geom_xpos = geom_xpos
@@ -113,15 +126,21 @@ def _refit_bvh_jax_impl(
   out = jf(
       d.qpos.shape[0],
       m._impl.flex_dim,
+      m._impl.flex_edge,
       m._impl.flex_elem,
+      m._impl.flex_elemadr,
+      m._impl.flex_elemdataadr,
       m._impl.flex_elemnum,
+      m._impl.flex_radius,
+      m._impl.flex_shell,
+      m._impl.flex_shelldataadr,
       m._impl.flex_vertadr,
+      m._impl.flex_vertnum,
       m.geom_dataid,
       m.geom_size,
       m.geom_type,
       m._impl.nflex,
-      m._impl.nflexelemdata,
+      m._impl.nflexelem,
-      m._impl.nflexvert,
       d._impl.flexvert_xpos,
       d.geom_xmat,
       d.geom_xpos,

mjx/mujoco/mjx/warp/collision_driver.py

@@ -52,8 +52,26 @@ def _collision_shim(
     # Model
     nworld: int,
     block_dim: mjwp_types.BlockDim,
+    flex_conaffinity: wp.array(dtype=int),
+    flex_condim: wp.array(dtype=int),
+    flex_contype: wp.array(dtype=int),
+    flex_dim: wp.array(dtype=int),
+    flex_elem: wp.array(dtype=int),
+    flex_elemadr: wp.array(dtype=int),
+    flex_elemdataadr: wp.array(dtype=int),
+    flex_elemnum: wp.array(dtype=int),
+    flex_friction: wp.array(dtype=wp.vec3),
+    flex_margin: wp.array(dtype=float),
+    flex_radius: wp.array(dtype=float),
+    flex_shell: wp.array(dtype=int),
+    flex_shelldataadr: wp.array(dtype=int),
+    flex_shellnum: wp.array(dtype=int),
+    flex_vertadr: wp.array(dtype=int),
+    flex_vertflexid: wp.array(dtype=int),
     geom_aabb: wp.array3d(dtype=wp.vec3),
+    geom_conaffinity: wp.array(dtype=int),
     geom_condim: wp.array(dtype=int),
+    geom_contype: wp.array(dtype=int),
     geom_dataid: wp.array(dtype=int),
     geom_friction: wp.array2d(dtype=wp.vec3),
     geom_gap: wp.array2d(dtype=float),
@@ -90,6 +108,10 @@ def _collision_shim(
     mesh_vert: wp.array(dtype=wp.vec3),
     mesh_vertadr: wp.array(dtype=int),
     mesh_vertnum: wp.array(dtype=int),
+    nflex: int,
+    nflexelem: int,
+    nflexshelldata: int,
+    nflexvert: int,
     ngeom: int,
     nmaxmeshdeg: int,
     nmaxpolygon: int,
@@ -108,7 +130,7 @@ def _collision_shim(
     pair_solref: wp.array2d(dtype=wp.vec2),
     pair_solreffriction: wp.array2d(dtype=wp.vec2),
     plugin: wp.array(dtype=int),
-    plugin_attr: wp.array(dtype=wp.vec3f),
+    plugin_attr: wp.array(dtype=mjwp_types.vec_pluginattr),
     opt__broadphase: int,
     opt__broadphase_filter: int,
     opt__ccd_iterations: int,
@@ -120,12 +142,15 @@ def _collision_shim(
     # Data
     naccdmax: int,
     naconmax: int,
+    flexvert_xpos: wp.array2d(dtype=wp.vec3),
     geom_xmat: wp.array2d(dtype=wp.mat33),
     geom_xpos: wp.array2d(dtype=wp.vec3),
     nacon: wp.array(dtype=int),
     ncollision: wp.array(dtype=int),
     contact__dim: wp.array(dtype=int),
     contact__dist: wp.array(dtype=float),
+    contact__efc_address: wp.array2d(dtype=int),
+    contact__flex: wp.array(dtype=wp.vec2i),
     contact__frame: wp.array(dtype=wp.mat33),
     contact__friction: wp.array(dtype=mjwp_types.vec5),
     contact__geom: wp.array(dtype=wp.vec2i),
@@ -136,6 +161,7 @@ def _collision_shim(
     contact__solref: wp.array(dtype=wp.vec2),
     contact__solreffriction: wp.array(dtype=wp.vec2),
     contact__type: wp.array(dtype=int),
+    contact__vert: wp.array(dtype=wp.vec2i),
     contact__worldid: wp.array(dtype=int),
 ):
   _m.stat = _s
@@ -144,8 +170,26 @@ def _collision_shim(
   _d.efc = _e
   _d.contact = _c
   _m.block_dim = block_dim
+  _m.flex_conaffinity = flex_conaffinity
+  _m.flex_condim = flex_condim
+  _m.flex_contype = flex_contype
+  _m.flex_dim = flex_dim
+  _m.flex_elem = flex_elem
+  _m.flex_elemadr = flex_elemadr
+  _m.flex_elemdataadr = flex_elemdataadr
+  _m.flex_elemnum = flex_elemnum
+  _m.flex_friction = flex_friction
+  _m.flex_margin = flex_margin
+  _m.flex_radius = flex_radius
+  _m.flex_shell = flex_shell
+  _m.flex_shelldataadr = flex_shelldataadr
+  _m.flex_shellnum = flex_shellnum
+  _m.flex_vertadr = flex_vertadr
+  _m.flex_vertflexid = flex_vertflexid
   _m.geom_aabb = geom_aabb
+  _m.geom_conaffinity = geom_conaffinity
   _m.geom_condim = geom_condim
+  _m.geom_contype = geom_contype
   _m.geom_dataid = geom_dataid
   _m.geom_friction = geom_friction
   _m.geom_gap = geom_gap
@@ -182,6 +226,10 @@ def _collision_shim(
   _m.mesh_vert = mesh_vert
   _m.mesh_vertadr = mesh_vertadr
   _m.mesh_vertnum = mesh_vertnum
+  _m.nflex = nflex
+  _m.nflexelem = nflexelem
+  _m.nflexshelldata = nflexshelldata
+  _m.nflexvert = nflexvert
   _m.ngeom = ngeom
   _m.nmaxmeshdeg = nmaxmeshdeg
   _m.nmaxpolygon = nmaxpolygon
@@ -211,6 +259,8 @@ def _collision_shim(
   _m.plugin_attr = plugin_attr
   _d.contact.dim = contact__dim
   _d.contact.dist = contact__dist
+  _d.contact.efc_address = contact__efc_address
+  _d.contact.flex = contact__flex
   _d.contact.frame = contact__frame
   _d.contact.friction = contact__friction
   _d.contact.geom = contact__geom
@@ -221,7 +271,9 @@ def _collision_shim(
   _d.contact.solref = contact__solref
   _d.contact.solreffriction = contact__solreffriction
   _d.contact.type = contact__type
+  _d.contact.vert = contact__vert
   _d.contact.worldid = contact__worldid
+  _d.flexvert_xpos = flexvert_xpos
   _d.geom_xmat = geom_xmat
   _d.geom_xpos = geom_xpos
   _d.naccdmax = naccdmax
@@ -238,6 +290,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       'ncollision': d._impl.ncollision.shape,
       'contact__dim': d._impl.contact__dim.shape,
       'contact__dist': d._impl.contact__dist.shape,
+      'contact__efc_address': d._impl.contact__efc_address.shape,
+      'contact__flex': d._impl.contact__flex.shape,
       'contact__frame': d._impl.contact__frame.shape,
       'contact__friction': d._impl.contact__friction.shape,
       'contact__geom': d._impl.contact__geom.shape,
@@ -248,11 +302,12 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       'contact__solref': d._impl.contact__solref.shape,
       'contact__solreffriction': d._impl.contact__solreffriction.shape,
       'contact__type': d._impl.contact__type.shape,
+      'contact__vert': d._impl.contact__vert.shape,
       'contact__worldid': d._impl.contact__worldid.shape,
   }
   jf = ffi.jax_callable_variadic_tuple(
       _collision_shim,
-      num_outputs=15,
+      num_outputs=18,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames=set([
@@ -260,6 +315,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
           'ncollision',
           'contact__dim',
           'contact__dist',
+          'contact__efc_address',
+          'contact__flex',
           'contact__frame',
           'contact__friction',
           'contact__geom',
@@ -270,6 +327,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
           'contact__solref',
           'contact__solreffriction',
           'contact__type',
+          'contact__vert',
           'contact__worldid',
       ]),
       stage_in_argnames=set([
@@ -299,8 +357,26 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
   out = jf(
       d.qpos.shape[0],
       m._impl.block_dim,
+      m._impl.flex_conaffinity,
+      m._impl.flex_condim,
+      m._impl.flex_contype,
+      m._impl.flex_dim,
+      m._impl.flex_elem,
+      m._impl.flex_elemadr,
+      m._impl.flex_elemdataadr,
+      m._impl.flex_elemnum,
+      m._impl.flex_friction,
+      m._impl.flex_margin,
+      m._impl.flex_radius,
+      m._impl.flex_shell,
+      m._impl.flex_shelldataadr,
+      m._impl.flex_shellnum,
+      m._impl.flex_vertadr,
+      m._impl.flex_vertflexid,
       m.geom_aabb,
+      m.geom_conaffinity,
       m.geom_condim,
+      m.geom_contype,
       m.geom_dataid,
       m.geom_friction,
       m.geom_gap,
@@ -337,6 +413,10 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m.mesh_vert,
       m.mesh_vertadr,
       m.mesh_vertnum,
+      m._impl.nflex,
+      m._impl.nflexelem,
+      m._impl.nflexshelldata,
+      m._impl.nflexvert,
       m.ngeom,
       m._impl.nmaxmeshdeg,
       m._impl.nmaxpolygon,
@@ -366,12 +446,15 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m.opt._impl.sdf_iterations,
       d._impl.naccdmax,
       d._impl.naconmax,
+      d._impl.flexvert_xpos,
       d.geom_xmat,
       d.geom_xpos,
       d._impl.nacon,
       d._impl.ncollision,
       d._impl.contact__dim,
       d._impl.contact__dist,
+      d._impl.contact__efc_address,
+      d._impl.contact__flex,
       d._impl.contact__frame,
       d._impl.contact__friction,
       d._impl.contact__geom,
@@ -382,6 +465,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       d._impl.contact__solref,
       d._impl.contact__solreffriction,
       d._impl.contact__type,
+      d._impl.contact__vert,
       d._impl.contact__worldid,
   )
   d = d.tree_replace({
@@ -389,17 +473,20 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       '_impl.ncollision': out[1],
       '_impl.contact__dim': out[2],
       '_impl.contact__dist': out[3],
-      '_impl.contact__frame': out[4],
+      '_impl.contact__efc_address': out[4],
-      '_impl.contact__friction': out[5],
+      '_impl.contact__flex': out[5],
-      '_impl.contact__geom': out[6],
+      '_impl.contact__frame': out[6],
-      '_impl.contact__geomcollisionid': out[7],
+      '_impl.contact__friction': out[7],
-      '_impl.contact__includemargin': out[8],
+      '_impl.contact__geom': out[8],
-      '_impl.contact__pos': out[9],
+      '_impl.contact__geomcollisionid': out[9],
-      '_impl.contact__solimp': out[10],
+      '_impl.contact__includemargin': out[10],
-      '_impl.contact__solref': out[11],
+      '_impl.contact__pos': out[11],
-      '_impl.contact__solreffriction': out[12],
+      '_impl.contact__solimp': out[12],
-      '_impl.contact__type': out[13],
+      '_impl.contact__solref': out[13],
-      '_impl.contact__worldid': out[14],
+      '_impl.contact__solreffriction': out[14],
+      '_impl.contact__type': out[15],
+      '_impl.contact__vert': out[16],
+      '_impl.contact__worldid': out[17],
   })
   return d
 

mjx/mujoco/mjx/warp/forward.py

@@ -138,6 +138,10 @@ def _forward_shim(
     eq_type: wp.array(dtype=int),
     eq_wld_adr: wp.array(dtype=int),
     flex_bending: wp.array2d(dtype=float),
+    flex_centered: wp.array(dtype=bool),
+    flex_conaffinity: wp.array(dtype=int),
+    flex_condim: wp.array(dtype=int),
+    flex_contype: wp.array(dtype=int),
     flex_damping: wp.array(dtype=float),
     flex_dim: wp.array(dtype=int),
     flex_edge: wp.array(dtype=wp.vec2i),
@@ -146,12 +150,22 @@ def _forward_shim(
     flex_edgenum: wp.array(dtype=int),
     flex_elem: wp.array(dtype=int),
     flex_elemadr: wp.array(dtype=int),
+    flex_elemdataadr: wp.array(dtype=int),
     flex_elemedge: wp.array(dtype=int),
     flex_elemedgeadr: wp.array(dtype=int),
     flex_elemnum: wp.array(dtype=int),
+    flex_friction: wp.array(dtype=wp.vec3),
+    flex_margin: wp.array(dtype=float),
+    flex_radius: wp.array(dtype=float),
+    flex_shell: wp.array(dtype=int),
+    flex_shelldataadr: wp.array(dtype=int),
+    flex_shellnum: wp.array(dtype=int),
     flex_stiffness: wp.array2d(dtype=float),
+    flex_vert: wp.array(dtype=wp.vec3),
     flex_vertadr: wp.array(dtype=int),
     flex_vertbodyid: wp.array(dtype=int),
+    flex_vertflexid: wp.array(dtype=int),
+    flex_vertnum: wp.array(dtype=int),
     flexedge_J_colind: wp.array(dtype=int),
     flexedge_J_rowadr: wp.array(dtype=int),
     flexedge_J_rownnz: wp.array(dtype=int),
@@ -159,7 +173,9 @@ def _forward_shim(
     flexedge_length0: wp.array(dtype=float),
     geom_aabb: wp.array3d(dtype=wp.vec3),
     geom_bodyid: wp.array(dtype=int),
+    geom_conaffinity: wp.array(dtype=int),
     geom_condim: wp.array(dtype=int),
+    geom_contype: wp.array(dtype=int),
     geom_dataid: wp.array(dtype=int),
     geom_fluid: wp.array2d(dtype=float),
     geom_friction: wp.array2d(dtype=wp.vec3),
@@ -213,6 +229,7 @@ def _forward_shim(
     light_targetbodyid: wp.array(dtype=int),
     mapM2M: wp.array(dtype=int),
     mat_rgba: wp.array2d(dtype=wp.vec4),
+    max_ten_J_rownnz: int,
     mesh_face: wp.array(dtype=wp.vec3i),
     mesh_faceadr: wp.array(dtype=int),
     mesh_graph: wp.array(dtype=int),
@@ -235,6 +252,7 @@ def _forward_shim(
     mesh_vertadr: wp.array(dtype=int),
     mesh_vertnum: wp.array(dtype=int),
     nC: int,
+    nJten: int,
     na: int,
     nacttrnbody: int,
     nbody: int,
@@ -244,6 +262,7 @@ def _forward_shim(
     nflex: int,
     nflexedge: int,
     nflexelem: int,
+    nflexshelldata: int,
     nflexvert: int,
     ngeom: int,
     ngravcomp: int,
@@ -279,7 +298,9 @@ def _forward_shim(
     pair_solref: wp.array2d(dtype=wp.vec2),
     pair_solreffriction: wp.array2d(dtype=wp.vec2),
     plugin: wp.array(dtype=int),
-    plugin_attr: wp.array(dtype=wp.vec3f),
+    plugin_attr: wp.array(dtype=mjwp_types.vec_pluginattr),
+    qLD_all_updates: wp.array(dtype=wp.vec3i),
+    qLD_level_offsets: wp.array(dtype=int),
     qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...],
     qM_fullm_i: wp.array(dtype=int),
     qM_fullm_j: wp.array(dtype=int),
@@ -323,6 +344,9 @@ def _forward_shim(
     site_type: wp.array(dtype=int),
     taxel_sensorid: wp.array(dtype=int),
     taxel_vertadr: wp.array(dtype=int),
+    ten_J_colind: wp.array(dtype=int),
+    ten_J_rowadr: wp.array(dtype=int),
+    ten_J_rownnz: wp.array(dtype=int),
     tendon_actfrclimited: wp.array(dtype=bool),
     tendon_actfrcrange: wp.array2d(dtype=wp.vec2),
     tendon_adr: wp.array(dtype=int),
@@ -382,6 +406,7 @@ def _forward_shim(
     naccdmax: int,
     naconmax: int,
     njmax: int,
+    njmax_nnz: int,
     act: wp.array2d(dtype=float),
     act_dot: wp.array2d(dtype=float),
     actuator_force: wp.array2d(dtype=float),
@@ -446,7 +471,7 @@ def _forward_shim(
     subtree_angmom: wp.array2d(dtype=wp.vec3),
     subtree_com: wp.array2d(dtype=wp.vec3),
     subtree_linvel: wp.array2d(dtype=wp.vec3),
-    ten_J: wp.array3d(dtype=float),
+    ten_J: wp.array2d(dtype=float),
     ten_length: wp.array2d(dtype=float),
     ten_velocity: wp.array2d(dtype=float),
     ten_wrapadr: wp.array2d(dtype=int),
@@ -466,6 +491,7 @@ def _forward_shim(
     contact__dim: wp.array(dtype=int),
     contact__dist: wp.array(dtype=float),
     contact__efc_address: wp.array2d(dtype=int),
+    contact__flex: wp.array(dtype=wp.vec2i),
     contact__frame: wp.array(dtype=wp.mat33),
     contact__friction: wp.array(dtype=mjwp_types.vec5),
     contact__geom: wp.array(dtype=wp.vec2i),
@@ -476,6 +502,7 @@ def _forward_shim(
     contact__solref: wp.array(dtype=wp.vec2),
     contact__solreffriction: wp.array(dtype=wp.vec2),
     contact__type: wp.array(dtype=int),
+    contact__vert: wp.array(dtype=wp.vec2i),
     contact__worldid: wp.array(dtype=int),
     efc__D: wp.array2d(dtype=float),
     efc__J: wp.array3d(dtype=float),
@@ -585,6 +612,10 @@ def _forward_shim(
   _m.eq_type = eq_type
   _m.eq_wld_adr = eq_wld_adr
   _m.flex_bending = flex_bending
+  _m.flex_centered = flex_centered
+  _m.flex_conaffinity = flex_conaffinity
+  _m.flex_condim = flex_condim
+  _m.flex_contype = flex_contype
   _m.flex_damping = flex_damping
   _m.flex_dim = flex_dim
   _m.flex_edge = flex_edge
@@ -593,12 +624,22 @@ def _forward_shim(
   _m.flex_edgenum = flex_edgenum
   _m.flex_elem = flex_elem
   _m.flex_elemadr = flex_elemadr
+  _m.flex_elemdataadr = flex_elemdataadr
   _m.flex_elemedge = flex_elemedge
   _m.flex_elemedgeadr = flex_elemedgeadr
   _m.flex_elemnum = flex_elemnum
+  _m.flex_friction = flex_friction
+  _m.flex_margin = flex_margin
+  _m.flex_radius = flex_radius
+  _m.flex_shell = flex_shell
+  _m.flex_shelldataadr = flex_shelldataadr
+  _m.flex_shellnum = flex_shellnum
   _m.flex_stiffness = flex_stiffness
+  _m.flex_vert = flex_vert
   _m.flex_vertadr = flex_vertadr
   _m.flex_vertbodyid = flex_vertbodyid
+  _m.flex_vertflexid = flex_vertflexid
+  _m.flex_vertnum = flex_vertnum
   _m.flexedge_J_colind = flexedge_J_colind
   _m.flexedge_J_rowadr = flexedge_J_rowadr
   _m.flexedge_J_rownnz = flexedge_J_rownnz
@@ -606,7 +647,9 @@ def _forward_shim(
   _m.flexedge_length0 = flexedge_length0
   _m.geom_aabb = geom_aabb
   _m.geom_bodyid = geom_bodyid
+  _m.geom_conaffinity = geom_conaffinity
   _m.geom_condim = geom_condim
+  _m.geom_contype = geom_contype
   _m.geom_dataid = geom_dataid
   _m.geom_fluid = geom_fluid
   _m.geom_friction = geom_friction
@@ -660,6 +703,7 @@ def _forward_shim(
   _m.light_targetbodyid = light_targetbodyid
   _m.mapM2M = mapM2M
   _m.mat_rgba = mat_rgba
+  _m.max_ten_J_rownnz = max_ten_J_rownnz
   _m.mesh_face = mesh_face
   _m.mesh_faceadr = mesh_faceadr
   _m.mesh_graph = mesh_graph
@@ -682,6 +726,7 @@ def _forward_shim(
   _m.mesh_vertadr = mesh_vertadr
   _m.mesh_vertnum = mesh_vertnum
   _m.nC = nC
+  _m.nJten = nJten
   _m.na = na
   _m.nacttrnbody = nacttrnbody
   _m.nbody = nbody
@@ -691,6 +736,7 @@ def _forward_shim(
   _m.nflex = nflex
   _m.nflexedge = nflexedge
   _m.nflexelem = nflexelem
+  _m.nflexshelldata = nflexshelldata
   _m.nflexvert = nflexvert
   _m.ngeom = ngeom
   _m.ngravcomp = ngravcomp
@@ -753,6 +799,8 @@ def _forward_shim(
   _m.pair_solreffriction = pair_solreffriction
   _m.plugin = plugin
   _m.plugin_attr = plugin_attr
+  _m.qLD_all_updates = qLD_all_updates
+  _m.qLD_level_offsets = qLD_level_offsets
   _m.qLD_updates = qLD_updates
   _m.qM_fullm_i = qM_fullm_i
   _m.qM_fullm_j = qM_fullm_j
@@ -797,6 +845,9 @@ def _forward_shim(
   _m.stat.meaninertia = stat__meaninertia
   _m.taxel_sensorid = taxel_sensorid
   _m.taxel_vertadr = taxel_vertadr
+  _m.ten_J_colind = ten_J_colind
+  _m.ten_J_rowadr = ten_J_rowadr
+  _m.ten_J_rownnz = ten_J_rownnz
   _m.tendon_actfrclimited = tendon_actfrclimited
   _m.tendon_actfrcrange = tendon_actfrcrange
   _m.tendon_adr = tendon_adr
@@ -842,6 +893,7 @@ def _forward_shim(
   _d.contact.dim = contact__dim
   _d.contact.dist = contact__dist
   _d.contact.efc_address = contact__efc_address
+  _d.contact.flex = contact__flex
   _d.contact.frame = contact__frame
   _d.contact.friction = contact__friction
   _d.contact.geom = contact__geom
@@ -852,6 +904,7 @@ def _forward_shim(
   _d.contact.solref = contact__solref
   _d.contact.solreffriction = contact__solreffriction
   _d.contact.type = contact__type
+  _d.contact.vert = contact__vert
   _d.contact.worldid = contact__worldid
   _d.crb = crb
   _d.ctrl = ctrl
@@ -895,6 +948,7 @@ def _forward_shim(
   _d.nf = nf
   _d.nisland = nisland
   _d.njmax = njmax
+  _d.njmax_nnz = njmax_nnz
   _d.nl = nl
   _d.qLD = qLD
   _d.qLDiagInv = qLDiagInv
@@ -1018,6 +1072,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'contact__dim': d._impl.contact__dim.shape,
       'contact__dist': d._impl.contact__dist.shape,
       'contact__efc_address': d._impl.contact__efc_address.shape,
+      'contact__flex': d._impl.contact__flex.shape,
       'contact__frame': d._impl.contact__frame.shape,
       'contact__friction': d._impl.contact__friction.shape,
       'contact__geom': d._impl.contact__geom.shape,
@@ -1028,6 +1083,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'contact__solref': d._impl.contact__solref.shape,
       'contact__solreffriction': d._impl.contact__solreffriction.shape,
       'contact__type': d._impl.contact__type.shape,
+      'contact__vert': d._impl.contact__vert.shape,
       'contact__worldid': d._impl.contact__worldid.shape,
       'efc__D': d._impl.efc__D.shape,
       'efc__J': d._impl.efc__J.shape,
@@ -1047,7 +1103,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _forward_shim,
-      num_outputs=100,
+      num_outputs=102,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames=set([
@@ -1125,6 +1181,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'contact__dim',
           'contact__dist',
           'contact__efc_address',
+          'contact__flex',
           'contact__frame',
           'contact__friction',
           'contact__geom',
@@ -1135,6 +1192,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'contact__solref',
           'contact__solreffriction',
           'contact__type',
+          'contact__vert',
           'contact__worldid',
           'efc__D',
           'efc__J',
@@ -1417,6 +1475,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m.eq_type,
       m._impl.eq_wld_adr,
       m._impl.flex_bending,
+      m._impl.flex_centered,
+      m._impl.flex_conaffinity,
+      m._impl.flex_condim,
+      m._impl.flex_contype,
       m._impl.flex_damping,
       m._impl.flex_dim,
       m._impl.flex_edge,
@@ -1425,12 +1487,22 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_edgenum,
       m._impl.flex_elem,
       m._impl.flex_elemadr,
+      m._impl.flex_elemdataadr,
       m._impl.flex_elemedge,
       m._impl.flex_elemedgeadr,
       m._impl.flex_elemnum,
+      m._impl.flex_friction,
+      m._impl.flex_margin,
+      m._impl.flex_radius,
+      m._impl.flex_shell,
+      m._impl.flex_shelldataadr,
+      m._impl.flex_shellnum,
       m._impl.flex_stiffness,
+      m._impl.flex_vert,
       m._impl.flex_vertadr,
       m._impl.flex_vertbodyid,
+      m._impl.flex_vertflexid,
+      m._impl.flex_vertnum,
       m._impl.flexedge_J_colind,
       m._impl.flexedge_J_rowadr,
       m._impl.flexedge_J_rownnz,
@@ -1438,7 +1510,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.flexedge_length0,
       m.geom_aabb,
       m.geom_bodyid,
+      m.geom_conaffinity,
       m.geom_condim,
+      m.geom_contype,
       m.geom_dataid,
       m.geom_fluid,
       m.geom_friction,
@@ -1492,6 +1566,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.light_targetbodyid,
       m._impl.mapM2M,
       m.mat_rgba,
+      m._impl.max_ten_J_rownnz,
       m.mesh_face,
       m.mesh_faceadr,
       m.mesh_graph,
@@ -1514,6 +1589,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m.mesh_vertadr,
       m.mesh_vertnum,
       m.nC,
+      m.nJten,
       m.na,
       m._impl.nacttrnbody,
       m.nbody,
@@ -1523,6 +1599,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.nflex,
       m._impl.nflexedge,
       m._impl.nflexelem,
+      m._impl.nflexshelldata,
       m._impl.nflexvert,
       m.ngeom,
       m.ngravcomp,
@@ -1559,6 +1636,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m.pair_solreffriction,
       m._impl.plugin,
       m._impl.plugin_attr,
+      m._impl.qLD_all_updates,
+      m._impl.qLD_level_offsets,
       m._impl.qLD_updates,
       m._impl.qM_fullm_i,
       m._impl.qM_fullm_j,
@@ -1602,6 +1681,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m.site_type,
       m._impl.taxel_sensorid,
       m._impl.taxel_vertadr,
+      m._impl.ten_J_colind,
+      m._impl.ten_J_rowadr,
+      m._impl.ten_J_rownnz,
       m.tendon_actfrclimited,
       m.tendon_actfrcrange,
       m.tendon_adr,
@@ -1660,6 +1742,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d._impl.naccdmax,
       d._impl.naconmax,
       d._impl.njmax,
+      d._impl.njmax_nnz,
       d.act,
       d.act_dot,
       d.actuator_force,
@@ -1744,6 +1827,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d._impl.contact__dim,
       d._impl.contact__dist,
       d._impl.contact__efc_address,
+      d._impl.contact__flex,
       d._impl.contact__frame,
       d._impl.contact__friction,
       d._impl.contact__geom,
@@ -1754,6 +1838,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d._impl.contact__solref,
       d._impl.contact__solreffriction,
       d._impl.contact__type,
+      d._impl.contact__vert,
       d._impl.contact__worldid,
       d._impl.efc__D,
       d._impl.efc__J,
@@ -1846,32 +1931,34 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       '_impl.contact__dim': out[71],
       '_impl.contact__dist': out[72],
       '_impl.contact__efc_address': out[73],
-      '_impl.contact__frame': out[74],
+      '_impl.contact__flex': out[74],
-      '_impl.contact__friction': out[75],
+      '_impl.contact__frame': out[75],
-      '_impl.contact__geom': out[76],
+      '_impl.contact__friction': out[76],
-      '_impl.contact__geomcollisionid': out[77],
+      '_impl.contact__geom': out[77],
-      '_impl.contact__includemargin': out[78],
+      '_impl.contact__geomcollisionid': out[78],
-      '_impl.contact__pos': out[79],
+      '_impl.contact__includemargin': out[79],
-      '_impl.contact__solimp': out[80],
+      '_impl.contact__pos': out[80],
-      '_impl.contact__solref': out[81],
+      '_impl.contact__solimp': out[81],
-      '_impl.contact__solreffriction': out[82],
+      '_impl.contact__solref': out[82],
-      '_impl.contact__type': out[83],
+      '_impl.contact__solreffriction': out[83],
-      '_impl.contact__worldid': out[84],
+      '_impl.contact__type': out[84],
-      '_impl.efc__D': out[85],
+      '_impl.contact__vert': out[85],
-      '_impl.efc__J': out[86],
+      '_impl.contact__worldid': out[86],
-      '_impl.efc__J_colind': out[87],
+      '_impl.efc__D': out[87],
-      '_impl.efc__J_rowadr': out[88],
+      '_impl.efc__J': out[88],
-      '_impl.efc__J_rownnz': out[89],
+      '_impl.efc__J_colind': out[89],
-      '_impl.efc__Ma': out[90],
+      '_impl.efc__J_rowadr': out[90],
-      '_impl.efc__aref': out[91],
+      '_impl.efc__J_rownnz': out[91],
-      '_impl.efc__force': out[92],
+      '_impl.efc__Ma': out[92],
-      '_impl.efc__frictionloss': out[93],
+      '_impl.efc__aref': out[93],
-      '_impl.efc__id': out[94],
+      '_impl.efc__force': out[94],
-      '_impl.efc__margin': out[95],
+      '_impl.efc__frictionloss': out[95],
-      '_impl.efc__pos': out[96],
+      '_impl.efc__id': out[96],
-      '_impl.efc__state': out[97],
+      '_impl.efc__margin': out[97],
-      '_impl.efc__type': out[98],
+      '_impl.efc__pos': out[98],
-      '_impl.efc__vel': out[99],
+      '_impl.efc__state': out[99],
+      '_impl.efc__type': out[100],
+      '_impl.efc__vel': out[101],
   })
   return d
 
@@ -1980,6 +2067,10 @@ def _step_shim(
     eq_type: wp.array(dtype=int),
     eq_wld_adr: wp.array(dtype=int),
     flex_bending: wp.array2d(dtype=float),
+    flex_centered: wp.array(dtype=bool),
+    flex_conaffinity: wp.array(dtype=int),
+    flex_condim: wp.array(dtype=int),
+    flex_contype: wp.array(dtype=int),
     flex_damping: wp.array(dtype=float),
     flex_dim: wp.array(dtype=int),
     flex_edge: wp.array(dtype=wp.vec2i),
@@ -1988,12 +2079,22 @@ def _step_shim(
     flex_edgenum: wp.array(dtype=int),
     flex_elem: wp.array(dtype=int),
     flex_elemadr: wp.array(dtype=int),
+    flex_elemdataadr: wp.array(dtype=int),
     flex_elemedge: wp.array(dtype=int),
     flex_elemedgeadr: wp.array(dtype=int),
     flex_elemnum: wp.array(dtype=int),
+    flex_friction: wp.array(dtype=wp.vec3),
+    flex_margin: wp.array(dtype=float),
+    flex_radius: wp.array(dtype=float),
+    flex_shell: wp.array(dtype=int),
+    flex_shelldataadr: wp.array(dtype=int),
+    flex_shellnum: wp.array(dtype=int),
     flex_stiffness: wp.array2d(dtype=float),
+    flex_vert: wp.array(dtype=wp.vec3),
     flex_vertadr: wp.array(dtype=int),
     flex_vertbodyid: wp.array(dtype=int),
+    flex_vertflexid: wp.array(dtype=int),
+    flex_vertnum: wp.array(dtype=int),
     flexedge_J_colind: wp.array(dtype=int),
     flexedge_J_rowadr: wp.array(dtype=int),
     flexedge_J_rownnz: wp.array(dtype=int),
@@ -2001,7 +2102,9 @@ def _step_shim(
     flexedge_length0: wp.array(dtype=float),
     geom_aabb: wp.array3d(dtype=wp.vec3),
     geom_bodyid: wp.array(dtype=int),
+    geom_conaffinity: wp.array(dtype=int),
     geom_condim: wp.array(dtype=int),
+    geom_contype: wp.array(dtype=int),
     geom_dataid: wp.array(dtype=int),
     geom_fluid: wp.array2d(dtype=float),
     geom_friction: wp.array2d(dtype=wp.vec3),
@@ -2055,6 +2158,7 @@ def _step_shim(
     light_targetbodyid: wp.array(dtype=int),
     mapM2M: wp.array(dtype=int),
     mat_rgba: wp.array2d(dtype=wp.vec4),
+    max_ten_J_rownnz: int,
     mesh_face: wp.array(dtype=wp.vec3i),
     mesh_faceadr: wp.array(dtype=int),
     mesh_graph: wp.array(dtype=int),
@@ -2077,6 +2181,7 @@ def _step_shim(
     mesh_vertadr: wp.array(dtype=int),
     mesh_vertnum: wp.array(dtype=int),
     nC: int,
+    nJten: int,
     nM: int,
     na: int,
     nacttrnbody: int,
@@ -2087,6 +2192,7 @@ def _step_shim(
     nflex: int,
     nflexedge: int,
     nflexelem: int,
+    nflexshelldata: int,
     nflexvert: int,
     ngeom: int,
     ngravcomp: int,
@@ -2122,7 +2228,9 @@ def _step_shim(
     pair_solref: wp.array2d(dtype=wp.vec2),
     pair_solreffriction: wp.array2d(dtype=wp.vec2),
     plugin: wp.array(dtype=int),
-    plugin_attr: wp.array(dtype=wp.vec3f),
+    plugin_attr: wp.array(dtype=mjwp_types.vec_pluginattr),
+    qLD_all_updates: wp.array(dtype=wp.vec3i),
+    qLD_level_offsets: wp.array(dtype=int),
     qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...],
     qM_fullm_i: wp.array(dtype=int),
     qM_fullm_j: wp.array(dtype=int),
@@ -2166,6 +2274,9 @@ def _step_shim(
     site_type: wp.array(dtype=int),
     taxel_sensorid: wp.array(dtype=int),
     taxel_vertadr: wp.array(dtype=int),
+    ten_J_colind: wp.array(dtype=int),
+    ten_J_rowadr: wp.array(dtype=int),
+    ten_J_rownnz: wp.array(dtype=int),
     tendon_actfrclimited: wp.array(dtype=bool),
     tendon_actfrcrange: wp.array2d(dtype=wp.vec2),
     tendon_adr: wp.array(dtype=int),
@@ -2226,6 +2337,7 @@ def _step_shim(
     naccdmax: int,
     naconmax: int,
     njmax: int,
+    njmax_nnz: int,
     act: wp.array2d(dtype=float),
     act_dot: wp.array2d(dtype=float),
     actuator_force: wp.array2d(dtype=float),
@@ -2290,7 +2402,7 @@ def _step_shim(
     subtree_angmom: wp.array2d(dtype=wp.vec3),
     subtree_com: wp.array2d(dtype=wp.vec3),
     subtree_linvel: wp.array2d(dtype=wp.vec3),
-    ten_J: wp.array3d(dtype=float),
+    ten_J: wp.array2d(dtype=float),
     ten_length: wp.array2d(dtype=float),
     ten_velocity: wp.array2d(dtype=float),
     ten_wrapadr: wp.array2d(dtype=int),
@@ -2310,6 +2422,7 @@ def _step_shim(
     contact__dim: wp.array(dtype=int),
     contact__dist: wp.array(dtype=float),
     contact__efc_address: wp.array2d(dtype=int),
+    contact__flex: wp.array(dtype=wp.vec2i),
     contact__frame: wp.array(dtype=wp.mat33),
     contact__friction: wp.array(dtype=mjwp_types.vec5),
     contact__geom: wp.array(dtype=wp.vec2i),
@@ -2320,6 +2433,7 @@ def _step_shim(
     contact__solref: wp.array(dtype=wp.vec2),
     contact__solreffriction: wp.array(dtype=wp.vec2),
     contact__type: wp.array(dtype=int),
+    contact__vert: wp.array(dtype=wp.vec2i),
     contact__worldid: wp.array(dtype=int),
     efc__D: wp.array2d(dtype=float),
     efc__J: wp.array3d(dtype=float),
@@ -2429,6 +2543,10 @@ def _step_shim(
   _m.eq_type = eq_type
   _m.eq_wld_adr = eq_wld_adr
   _m.flex_bending = flex_bending
+  _m.flex_centered = flex_centered
+  _m.flex_conaffinity = flex_conaffinity
+  _m.flex_condim = flex_condim
+  _m.flex_contype = flex_contype
   _m.flex_damping = flex_damping
   _m.flex_dim = flex_dim
   _m.flex_edge = flex_edge
@@ -2437,12 +2555,22 @@ def _step_shim(
   _m.flex_edgenum = flex_edgenum
   _m.flex_elem = flex_elem
   _m.flex_elemadr = flex_elemadr
+  _m.flex_elemdataadr = flex_elemdataadr
   _m.flex_elemedge = flex_elemedge
   _m.flex_elemedgeadr = flex_elemedgeadr
   _m.flex_elemnum = flex_elemnum
+  _m.flex_friction = flex_friction
+  _m.flex_margin = flex_margin
+  _m.flex_radius = flex_radius
+  _m.flex_shell = flex_shell
+  _m.flex_shelldataadr = flex_shelldataadr
+  _m.flex_shellnum = flex_shellnum
   _m.flex_stiffness = flex_stiffness
+  _m.flex_vert = flex_vert
   _m.flex_vertadr = flex_vertadr
   _m.flex_vertbodyid = flex_vertbodyid
+  _m.flex_vertflexid = flex_vertflexid
+  _m.flex_vertnum = flex_vertnum
   _m.flexedge_J_colind = flexedge_J_colind
   _m.flexedge_J_rowadr = flexedge_J_rowadr
   _m.flexedge_J_rownnz = flexedge_J_rownnz
@@ -2450,7 +2578,9 @@ def _step_shim(
   _m.flexedge_length0 = flexedge_length0
   _m.geom_aabb = geom_aabb
   _m.geom_bodyid = geom_bodyid
+  _m.geom_conaffinity = geom_conaffinity
   _m.geom_condim = geom_condim
+  _m.geom_contype = geom_contype
   _m.geom_dataid = geom_dataid
   _m.geom_fluid = geom_fluid
   _m.geom_friction = geom_friction
@@ -2504,6 +2634,7 @@ def _step_shim(
   _m.light_targetbodyid = light_targetbodyid
   _m.mapM2M = mapM2M
   _m.mat_rgba = mat_rgba
+  _m.max_ten_J_rownnz = max_ten_J_rownnz
   _m.mesh_face = mesh_face
   _m.mesh_faceadr = mesh_faceadr
   _m.mesh_graph = mesh_graph
@@ -2526,6 +2657,7 @@ def _step_shim(
   _m.mesh_vertadr = mesh_vertadr
   _m.mesh_vertnum = mesh_vertnum
   _m.nC = nC
+  _m.nJten = nJten
   _m.nM = nM
   _m.na = na
   _m.nacttrnbody = nacttrnbody
@@ -2536,6 +2668,7 @@ def _step_shim(
   _m.nflex = nflex
   _m.nflexedge = nflexedge
   _m.nflexelem = nflexelem
+  _m.nflexshelldata = nflexshelldata
   _m.nflexvert = nflexvert
   _m.ngeom = ngeom
   _m.ngravcomp = ngravcomp
@@ -2599,6 +2732,8 @@ def _step_shim(
   _m.pair_solreffriction = pair_solreffriction
   _m.plugin = plugin
   _m.plugin_attr = plugin_attr
+  _m.qLD_all_updates = qLD_all_updates
+  _m.qLD_level_offsets = qLD_level_offsets
   _m.qLD_updates = qLD_updates
   _m.qM_fullm_i = qM_fullm_i
   _m.qM_fullm_j = qM_fullm_j
@@ -2643,6 +2778,9 @@ def _step_shim(
   _m.stat.meaninertia = stat__meaninertia
   _m.taxel_sensorid = taxel_sensorid
   _m.taxel_vertadr = taxel_vertadr
+  _m.ten_J_colind = ten_J_colind
+  _m.ten_J_rowadr = ten_J_rowadr
+  _m.ten_J_rownnz = ten_J_rownnz
   _m.tendon_actfrclimited = tendon_actfrclimited
   _m.tendon_actfrcrange = tendon_actfrcrange
   _m.tendon_adr = tendon_adr
@@ -2688,6 +2826,7 @@ def _step_shim(
   _d.contact.dim = contact__dim
   _d.contact.dist = contact__dist
   _d.contact.efc_address = contact__efc_address
+  _d.contact.flex = contact__flex
   _d.contact.frame = contact__frame
   _d.contact.friction = contact__friction
   _d.contact.geom = contact__geom
@@ -2698,6 +2837,7 @@ def _step_shim(
   _d.contact.solref = contact__solref
   _d.contact.solreffriction = contact__solreffriction
   _d.contact.type = contact__type
+  _d.contact.vert = contact__vert
   _d.contact.worldid = contact__worldid
   _d.crb = crb
   _d.ctrl = ctrl
@@ -2741,6 +2881,7 @@ def _step_shim(
   _d.nf = nf
   _d.nisland = nisland
   _d.njmax = njmax
+  _d.njmax_nnz = njmax_nnz
   _d.nl = nl
   _d.qLD = qLD
   _d.qLDiagInv = qLDiagInv
@@ -2868,6 +3009,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'contact__dim': d._impl.contact__dim.shape,
       'contact__dist': d._impl.contact__dist.shape,
       'contact__efc_address': d._impl.contact__efc_address.shape,
+      'contact__flex': d._impl.contact__flex.shape,
       'contact__frame': d._impl.contact__frame.shape,
       'contact__friction': d._impl.contact__friction.shape,
       'contact__geom': d._impl.contact__geom.shape,
@@ -2878,6 +3020,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'contact__solref': d._impl.contact__solref.shape,
       'contact__solreffriction': d._impl.contact__solreffriction.shape,
       'contact__type': d._impl.contact__type.shape,
+      'contact__vert': d._impl.contact__vert.shape,
       'contact__worldid': d._impl.contact__worldid.shape,
       'efc__D': d._impl.efc__D.shape,
       'efc__J': d._impl.efc__J.shape,
@@ -2897,7 +3040,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _step_shim,
-      num_outputs=104,
+      num_outputs=106,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames=set([
@@ -2979,6 +3122,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'contact__dim',
           'contact__dist',
           'contact__efc_address',
+          'contact__flex',
           'contact__frame',
           'contact__friction',
           'contact__geom',
@@ -2989,6 +3133,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'contact__solref',
           'contact__solreffriction',
           'contact__type',
+          'contact__vert',
           'contact__worldid',
           'efc__D',
           'efc__J',
@@ -3275,6 +3420,10 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.eq_type,
       m._impl.eq_wld_adr,
       m._impl.flex_bending,
+      m._impl.flex_centered,
+      m._impl.flex_conaffinity,
+      m._impl.flex_condim,
+      m._impl.flex_contype,
       m._impl.flex_damping,
       m._impl.flex_dim,
       m._impl.flex_edge,
@@ -3283,12 +3432,22 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_edgenum,
       m._impl.flex_elem,
       m._impl.flex_elemadr,
+      m._impl.flex_elemdataadr,
       m._impl.flex_elemedge,
       m._impl.flex_elemedgeadr,
       m._impl.flex_elemnum,
+      m._impl.flex_friction,
+      m._impl.flex_margin,
+      m._impl.flex_radius,
+      m._impl.flex_shell,
+      m._impl.flex_shelldataadr,
+      m._impl.flex_shellnum,
       m._impl.flex_stiffness,
+      m._impl.flex_vert,
       m._impl.flex_vertadr,
       m._impl.flex_vertbodyid,
+      m._impl.flex_vertflexid,
+      m._impl.flex_vertnum,
       m._impl.flexedge_J_colind,
       m._impl.flexedge_J_rowadr,
       m._impl.flexedge_J_rownnz,
@@ -3296,7 +3455,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.flexedge_length0,
       m.geom_aabb,
       m.geom_bodyid,
+      m.geom_conaffinity,
       m.geom_condim,
+      m.geom_contype,
       m.geom_dataid,
       m.geom_fluid,
       m.geom_friction,
@@ -3350,6 +3511,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.light_targetbodyid,
       m._impl.mapM2M,
       m.mat_rgba,
+      m._impl.max_ten_J_rownnz,
       m.mesh_face,
       m.mesh_faceadr,
       m.mesh_graph,
@@ -3372,6 +3534,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.mesh_vertadr,
       m.mesh_vertnum,
       m.nC,
+      m.nJten,
       m.nM,
       m.na,
       m._impl.nacttrnbody,
@@ -3382,6 +3545,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.nflex,
       m._impl.nflexedge,
       m._impl.nflexelem,
+      m._impl.nflexshelldata,
       m._impl.nflexvert,
       m.ngeom,
       m.ngravcomp,
@@ -3418,6 +3582,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.pair_solreffriction,
       m._impl.plugin,
       m._impl.plugin_attr,
+      m._impl.qLD_all_updates,
+      m._impl.qLD_level_offsets,
       m._impl.qLD_updates,
       m._impl.qM_fullm_i,
       m._impl.qM_fullm_j,
@@ -3461,6 +3627,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.site_type,
       m._impl.taxel_sensorid,
       m._impl.taxel_vertadr,
+      m._impl.ten_J_colind,
+      m._impl.ten_J_rowadr,
+      m._impl.ten_J_rownnz,
       m.tendon_actfrclimited,
       m.tendon_actfrcrange,
       m.tendon_adr,
@@ -3520,6 +3689,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.naccdmax,
       d._impl.naconmax,
       d._impl.njmax,
+      d._impl.njmax_nnz,
       d.act,
       d.act_dot,
       d.actuator_force,
@@ -3604,6 +3774,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.contact__dim,
       d._impl.contact__dist,
       d._impl.contact__efc_address,
+      d._impl.contact__flex,
       d._impl.contact__frame,
       d._impl.contact__friction,
       d._impl.contact__geom,
@@ -3614,6 +3785,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.contact__solref,
       d._impl.contact__solreffriction,
       d._impl.contact__type,
+      d._impl.contact__vert,
       d._impl.contact__worldid,
       d._impl.efc__D,
       d._impl.efc__J,
@@ -3710,32 +3882,34 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       '_impl.contact__dim': out[75],
       '_impl.contact__dist': out[76],
       '_impl.contact__efc_address': out[77],
-      '_impl.contact__frame': out[78],
+      '_impl.contact__flex': out[78],
-      '_impl.contact__friction': out[79],
+      '_impl.contact__frame': out[79],
-      '_impl.contact__geom': out[80],
+      '_impl.contact__friction': out[80],
-      '_impl.contact__geomcollisionid': out[81],
+      '_impl.contact__geom': out[81],
-      '_impl.contact__includemargin': out[82],
+      '_impl.contact__geomcollisionid': out[82],
-      '_impl.contact__pos': out[83],
+      '_impl.contact__includemargin': out[83],
-      '_impl.contact__solimp': out[84],
+      '_impl.contact__pos': out[84],
-      '_impl.contact__solref': out[85],
+      '_impl.contact__solimp': out[85],
-      '_impl.contact__solreffriction': out[86],
+      '_impl.contact__solref': out[86],
-      '_impl.contact__type': out[87],
+      '_impl.contact__solreffriction': out[87],
-      '_impl.contact__worldid': out[88],
+      '_impl.contact__type': out[88],
-      '_impl.efc__D': out[89],
+      '_impl.contact__vert': out[89],
-      '_impl.efc__J': out[90],
+      '_impl.contact__worldid': out[90],
-      '_impl.efc__J_colind': out[91],
+      '_impl.efc__D': out[91],
-      '_impl.efc__J_rowadr': out[92],
+      '_impl.efc__J': out[92],
-      '_impl.efc__J_rownnz': out[93],
+      '_impl.efc__J_colind': out[93],
-      '_impl.efc__Ma': out[94],
+      '_impl.efc__J_rowadr': out[94],
-      '_impl.efc__aref': out[95],
+      '_impl.efc__J_rownnz': out[95],
-      '_impl.efc__force': out[96],
+      '_impl.efc__Ma': out[96],
-      '_impl.efc__frictionloss': out[97],
+      '_impl.efc__aref': out[97],
-      '_impl.efc__id': out[98],
+      '_impl.efc__force': out[98],
-      '_impl.efc__margin': out[99],
+      '_impl.efc__frictionloss': out[99],
-      '_impl.efc__pos': out[100],
+      '_impl.efc__id': out[100],
-      '_impl.efc__state': out[101],
+      '_impl.efc__margin': out[101],
-      '_impl.efc__type': out[102],
+      '_impl.efc__pos': out[102],
-      '_impl.efc__vel': out[103],
+      '_impl.efc__state': out[103],
+      '_impl.efc__type': out[104],
+      '_impl.efc__vel': out[105],
   })
   return d
 

mjx/mujoco/mjx/warp/forward_test.py

@@ -157,7 +157,16 @@ class ForwardTest(parameterized.TestCase):
           m.ten_J_rowadr,
           m.ten_J_colind,
       )
-      tu.assert_eq(dx._impl.ten_J, ten_J, 'ten_J')
+      # convert sparse warp ten_J to dense representation
+      warp_ten_J = np.zeros((m.ntendon, m.nv))
+      mujoco.mju_sparse2dense(
+          warp_ten_J,
+          np.asarray(dx._impl.ten_J),
+          mx._impl.ten_J_rownnz,
+          mx._impl.ten_J_rowadr,
+          mx._impl.ten_J_colind,
+      )
+      tu.assert_eq(warp_ten_J, ten_J, 'ten_J')
       tu.assert_attr_eq(dx._impl, d, 'ten_wrapadr')
       tu.assert_attr_eq(dx._impl, d, 'ten_wrapnum')
       tu.assert_attr_eq(dx._impl, d, 'wrap_xpos')

mjx/mujoco/mjx/warp/render.py

@@ -48,6 +48,7 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _render_shim(
     # Model
@@ -56,6 +57,9 @@ def _render_shim(
     cam_intrinsic: wp.array2d(dtype=wp.vec4),
     cam_projection: wp.array(dtype=int),
     cam_sensorsize: wp.array(dtype=wp.vec2),
+    flex_edge: wp.array(dtype=wp.vec2i),
+    flex_radius: wp.array(dtype=float),
+    flex_vertadr: wp.array(dtype=int),
     geom_dataid: wp.array(dtype=int),
     geom_matid: wp.array2d(dtype=int),
     geom_rgba: wp.array2d(dtype=wp.vec4),
@@ -68,11 +72,11 @@ def _render_shim(
     mat_texid: wp.array3d(dtype=int),
     mat_texrepeat: wp.array2d(dtype=wp.vec2),
     mesh_faceadr: wp.array(dtype=int),
-    nflex: int,
     nlight: int,
     # Data
     cam_xmat: wp.array2d(dtype=wp.mat33),
     cam_xpos: wp.array2d(dtype=wp.vec3),
+    flexvert_xpos: wp.array2d(dtype=wp.vec3),
     geom_xmat: wp.array2d(dtype=wp.mat33),
     geom_xpos: wp.array2d(dtype=wp.vec3),
     light_xdir: wp.array2d(dtype=wp.vec3),
@@ -91,6 +95,9 @@ def _render_shim(
   _m.cam_intrinsic = cam_intrinsic
   _m.cam_projection = cam_projection
   _m.cam_sensorsize = cam_sensorsize
+  _m.flex_edge = flex_edge
+  _m.flex_radius = flex_radius
+  _m.flex_vertadr = flex_vertadr
   _m.geom_dataid = geom_dataid
   _m.geom_matid = geom_matid
   _m.geom_rgba = geom_rgba
@@ -103,10 +110,10 @@ def _render_shim(
   _m.mat_texid = mat_texid
   _m.mat_texrepeat = mat_texrepeat
   _m.mesh_faceadr = mesh_faceadr
-  _m.nflex = nflex
   _m.nlight = nlight
   _d.cam_xmat = cam_xmat
   _d.cam_xpos = cam_xpos
+  _d.flexvert_xpos = flexvert_xpos
   _d.geom_xmat = geom_xmat
   _d.geom_xpos = geom_xpos
   _d.light_xdir = light_xdir
@@ -155,6 +162,9 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
       m.cam_intrinsic,
       m._impl.cam_projection,
       m.cam_sensorsize,
+      m._impl.flex_edge,
+      m._impl.flex_radius,
+      m._impl.flex_vertadr,
       m.geom_dataid,
       m.geom_matid,
       m.geom_rgba,
@@ -167,10 +177,10 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
       m.mat_texid,
       m._impl.mat_texrepeat,
       m.mesh_faceadr,
-      m._impl.nflex,
       m.nlight,
       d.cam_xmat,
       d.cam_xpos,
+      d._impl.flexvert_xpos,
       d.geom_xmat,
       d.geom_xpos,
       d._impl.light_xdir,

mjx/mujoco/mjx/warp/smooth.py

@@ -297,17 +297,20 @@ def kinematics_vmap(
 def _tendon_shim(
     # Model
     nworld: int,
+    body_dofadr: wp.array(dtype=int),
+    body_dofnum: wp.array(dtype=int),
     body_parentid: wp.array(dtype=int),
     body_rootid: wp.array(dtype=int),
-    dof_bodyid: wp.array(dtype=int),
     geom_bodyid: wp.array(dtype=int),
     geom_size: wp.array2d(dtype=wp.vec3),
     jnt_dofadr: wp.array(dtype=int),
     jnt_qposadr: wp.array(dtype=int),
     ntendon: int,
-    nv: int,
     nwrap: int,
     site_bodyid: wp.array(dtype=int),
+    ten_J_colind: wp.array(dtype=int),
+    ten_J_rowadr: wp.array(dtype=int),
+    ten_J_rownnz: wp.array(dtype=int),
     tendon_adr: wp.array(dtype=int),
     tendon_geom_adr: wp.array(dtype=int),
     tendon_jnt_adr: wp.array(dtype=int),
@@ -327,7 +330,7 @@ def _tendon_shim(
     qpos: wp.array2d(dtype=float),
     site_xpos: wp.array2d(dtype=wp.vec3),
     subtree_com: wp.array2d(dtype=wp.vec3),
-    ten_J: wp.array3d(dtype=float),
+    ten_J: wp.array2d(dtype=float),
     ten_length: wp.array2d(dtype=float),
     ten_wrapadr: wp.array2d(dtype=int),
     ten_wrapnum: wp.array2d(dtype=int),
@@ -339,17 +342,20 @@ def _tendon_shim(
   _m.callback = _cb
   _d.efc = _e
   _d.contact = _c
+  _m.body_dofadr = body_dofadr
+  _m.body_dofnum = body_dofnum
   _m.body_parentid = body_parentid
   _m.body_rootid = body_rootid
-  _m.dof_bodyid = dof_bodyid
   _m.geom_bodyid = geom_bodyid
   _m.geom_size = geom_size
   _m.jnt_dofadr = jnt_dofadr
   _m.jnt_qposadr = jnt_qposadr
   _m.ntendon = ntendon
-  _m.nv = nv
   _m.nwrap = nwrap
   _m.site_bodyid = site_bodyid
+  _m.ten_J_colind = ten_J_colind
+  _m.ten_J_rowadr = ten_J_rowadr
+  _m.ten_J_rownnz = ten_J_rownnz
   _m.tendon_adr = tendon_adr
   _m.tendon_geom_adr = tendon_geom_adr
   _m.tendon_jnt_adr = tendon_jnt_adr
@@ -416,17 +422,20 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
   )
   out = jf(
       d.qpos.shape[0],
+      m.body_dofadr,
+      m.body_dofnum,
       m.body_parentid,
       m.body_rootid,
-      m.dof_bodyid,
       m.geom_bodyid,
       m.geom_size,
       m.jnt_dofadr,
       m.jnt_qposadr,
       m.ntendon,
-      m.nv,
       m.nwrap,
       m.site_bodyid,
+      m._impl.ten_J_colind,
+      m._impl.ten_J_rowadr,
+      m._impl.ten_J_rownnz,
       m.tendon_adr,
       m._impl.tendon_geom_adr,
       m._impl.tendon_jnt_adr,

mjx/mujoco/mjx/warp/types.py

@@ -83,7 +83,7 @@ class BlockDim:
   qderiv_actuator_dense: int
   ray: int
   segmented_sort: int
-  tendon_velocity: int
+  solve_LD_sparse_fused: int
   update_gradient_JTDAJ_dense: int
   update_gradient_JTDAJ_sparse: int
   update_gradient_cholesky: int
@@ -141,6 +141,10 @@ class ModelWarp(PyTreeNode):
   eq_ten_adr: np.ndarray
   eq_wld_adr: np.ndarray
   flex_bending: np.ndarray
+  flex_centered: np.ndarray
+  flex_conaffinity: np.ndarray
+  flex_condim: np.ndarray
+  flex_contype: np.ndarray
   flex_damping: np.ndarray
   flex_dim: np.ndarray
   flex_edge: np.ndarray
@@ -149,12 +153,21 @@ class ModelWarp(PyTreeNode):
   flex_edgenum: np.ndarray
   flex_elem: np.ndarray
   flex_elemadr: np.ndarray
+  flex_elemdataadr: np.ndarray
   flex_elemedge: np.ndarray
   flex_elemedgeadr: np.ndarray
   flex_elemnum: np.ndarray
+  flex_friction: np.ndarray
+  flex_margin: np.ndarray
+  flex_radius: np.ndarray
+  flex_shell: np.ndarray
+  flex_shelldataadr: np.ndarray
+  flex_shellnum: np.ndarray
   flex_stiffness: np.ndarray
+  flex_vert: np.ndarray
   flex_vertadr: np.ndarray
   flex_vertbodyid: np.ndarray
+  flex_vertflexid: np.ndarray
   flex_vertnum: np.ndarray
   flexedge_J_colind: np.ndarray
   flexedge_J_rowadr: np.ndarray
@@ -173,6 +186,7 @@ class ModelWarp(PyTreeNode):
   light_targetbodyid: np.ndarray
   mapM2M: np.ndarray
   mat_texrepeat: jax.Array
+  max_ten_J_rownnz: int
   mesh_polyadr: np.ndarray
   mesh_polymap: np.ndarray
   mesh_polymapadr: np.ndarray
@@ -191,6 +205,7 @@ class ModelWarp(PyTreeNode):
   nflexelem: int
   nflexelemdata: int
   nflexelemedge: int
+  nflexshelldata: int
   nflexvert: int
   nmaxcondim: int
   nmaxmeshdeg: int
@@ -213,6 +228,8 @@ class ModelWarp(PyTreeNode):
   oct_coeff: np.ndarray
   plugin: np.ndarray
   plugin_attr: np.ndarray
+  qLD_all_updates: np.ndarray
+  qLD_level_offsets: np.ndarray
   qLD_updates: Tuple[np.ndarray, ...]
   qM_fullm_i: np.ndarray
   qM_fullm_j: np.ndarray
@@ -240,6 +257,9 @@ class ModelWarp(PyTreeNode):
   sensor_vel_adr: np.ndarray
   taxel_sensorid: np.ndarray
   taxel_vertadr: np.ndarray
+  ten_J_colind: np.ndarray
+  ten_J_rowadr: np.ndarray
+  ten_J_rownnz: np.ndarray
   ten_wrapadr_site: np.ndarray
   ten_wrapnum_site: np.ndarray
   tendon_geom_adr: np.ndarray
@@ -266,6 +286,7 @@ class DataWarp(PyTreeNode):
   contact__dim: jax.Array
   contact__dist: jax.Array
   contact__efc_address: jax.Array
+  contact__flex: jax.Array
   contact__frame: jax.Array
   contact__friction: jax.Array
   contact__geom: jax.Array
@@ -276,6 +297,7 @@ class DataWarp(PyTreeNode):
   contact__solref: jax.Array
   contact__solreffriction: jax.Array
   contact__type: jax.Array
+  contact__vert: jax.Array
   contact__worldid: jax.Array
   crb: jax.Array
   efc__D: jax.Array
@@ -312,6 +334,7 @@ class DataWarp(PyTreeNode):
   nf: jax.Array
   nisland: jax.Array
   njmax: int
+  njmax_nnz: int
   njmax_pad: int
   nl: jax.Array
   nworld: int
@@ -335,6 +358,7 @@ DATA_NON_VMAP = {
     'contact__dim',
     'contact__dist',
     'contact__efc_address',
+    'contact__flex',
     'contact__frame',
     'contact__friction',
     'contact__geom',
@@ -345,12 +369,14 @@ DATA_NON_VMAP = {
     'contact__solref',
     'contact__solreffriction',
     'contact__type',
+    'contact__vert',
     'contact__worldid',
     'naccdmax',
     'nacon',
     'naconmax',
     'ncollision',
     'njmax',
+    'njmax_nnz',
     'njmax_pad',
     'nworld',
 }
@@ -398,6 +424,7 @@ _NDIM = {
         'contact__dim': 1,
         'contact__dist': 1,
         'contact__efc_address': 2,
+        'contact__flex': 2,
         'contact__frame': 3,
         'contact__friction': 2,
         'contact__geom': 2,
@@ -408,6 +435,7 @@ _NDIM = {
         'contact__solref': 2,
         'contact__solreffriction': 2,
         'contact__type': 1,
+        'contact__vert': 2,
         'contact__worldid': 1,
         'crb': 3,
         'ctrl': 2,
@@ -451,6 +479,7 @@ _NDIM = {
         'nf': 1,
         'nisland': 1,
         'njmax': 0,
+        'njmax_nnz': 0,
         'njmax_pad': 0,
         'nl': 1,
         'nworld': 0,
@@ -480,7 +509,7 @@ _NDIM = {
         'subtree_angmom': 3,
         'subtree_com': 3,
         'subtree_linvel': 3,
-        'ten_J': 3,
+        'ten_J': 2,
         'ten_length': 2,
         'ten_velocity': 2,
         'ten_wrapadr': 2,
@@ -535,7 +564,7 @@ _NDIM = {
         'block_dim__qderiv_actuator_dense': 0,
         'block_dim__ray': 0,
         'block_dim__segmented_sort': 0,
-        'block_dim__tendon_velocity': 0,
+        'block_dim__solve_LD_sparse_fused': 0,
         'block_dim__update_gradient_JTDAJ_dense': 0,
         'block_dim__update_gradient_JTDAJ_sparse': 0,
         'block_dim__update_gradient_cholesky': 0,
@@ -608,6 +637,10 @@ _NDIM = {
         'eq_wld_adr': 1,
         'exclude_signature': 1,
         'flex_bending': 2,
+        'flex_centered': 1,
+        'flex_conaffinity': 1,
+        'flex_condim': 1,
+        'flex_contype': 1,
         'flex_damping': 1,
         'flex_dim': 1,
         'flex_edge': 2,
@@ -616,12 +649,21 @@ _NDIM = {
         'flex_edgenum': 1,
         'flex_elem': 1,
         'flex_elemadr': 1,
+        'flex_elemdataadr': 1,
         'flex_elemedge': 1,
         'flex_elemedgeadr': 1,
         'flex_elemnum': 1,
+        'flex_friction': 2,
+        'flex_margin': 1,
+        'flex_radius': 1,
+        'flex_shell': 1,
+        'flex_shelldataadr': 1,
+        'flex_shellnum': 1,
         'flex_stiffness': 2,
+        'flex_vert': 2,
         'flex_vertadr': 1,
         'flex_vertbodyid': 1,
+        'flex_vertflexid': 1,
         'flex_vertnum': 1,
         'flexedge_J_colind': 1,
         'flexedge_J_rowadr': 1,
@@ -692,6 +734,7 @@ _NDIM = {
         'mat_rgba': 3,
         'mat_texid': 3,
         'mat_texrepeat': 3,
+        'max_ten_J_rownnz': 0,
         'mesh_face': 2,
         'mesh_faceadr': 1,
         'mesh_graph': 1,
@@ -717,6 +760,7 @@ _NDIM = {
         'nC': 0,
         'nJfe': 0,
         'nJmom': 0,
+        'nJten': 0,
         'nM': 0,
         'na': 0,
         'nacttrnbody': 0,
@@ -730,6 +774,7 @@ _NDIM = {
         'nflexelem': 0,
         'nflexelemdata': 0,
         'nflexelemedge': 0,
+        'nflexshelldata': 0,
         'nflexvert': 0,
         'ngeom': 0,
         'ngravcomp': 0,
@@ -811,6 +856,8 @@ _NDIM = {
         'pair_solreffriction': 3,
         'plugin': 1,
         'plugin_attr': 2,
+        'qLD_all_updates': 2,
+        'qLD_level_offsets': 1,
         'qLD_updates': -1,
         'qM_fullm_i': 1,
         'qM_fullm_j': 1,
@@ -856,6 +903,9 @@ _NDIM = {
         'stat__meaninertia': 1,
         'taxel_sensorid': 1,
         'taxel_vertadr': 1,
+        'ten_J_colind': 1,
+        'ten_J_rowadr': 1,
+        'ten_J_rownnz': 1,
         'ten_wrapadr_site': 1,
         'ten_wrapnum_site': 1,
         'tendon_actfrclimited': 1,
@@ -940,6 +990,7 @@ _BATCH_DIM = {
         'contact__dim': False,
         'contact__dist': False,
         'contact__efc_address': False,
+        'contact__flex': False,
         'contact__frame': False,
         'contact__friction': False,
         'contact__geom': False,
@@ -950,6 +1001,7 @@ _BATCH_DIM = {
         'contact__solref': False,
         'contact__solreffriction': False,
         'contact__type': False,
+        'contact__vert': False,
         'contact__worldid': False,
         'crb': True,
         'ctrl': True,
@@ -993,6 +1045,7 @@ _BATCH_DIM = {
         'nf': True,
         'nisland': True,
         'njmax': False,
+        'njmax_nnz': False,
         'njmax_pad': False,
         'nl': True,
         'nworld': False,
@@ -1077,7 +1130,7 @@ _BATCH_DIM = {
         'block_dim__qderiv_actuator_dense': False,
         'block_dim__ray': False,
         'block_dim__segmented_sort': False,
-        'block_dim__tendon_velocity': False,
+        'block_dim__solve_LD_sparse_fused': False,
         'block_dim__update_gradient_JTDAJ_dense': False,
         'block_dim__update_gradient_JTDAJ_sparse': False,
         'block_dim__update_gradient_cholesky': False,
@@ -1150,6 +1203,10 @@ _BATCH_DIM = {
         'eq_wld_adr': False,
         'exclude_signature': False,
         'flex_bending': False,
+        'flex_centered': False,
+        'flex_conaffinity': False,
+        'flex_condim': False,
+        'flex_contype': False,
         'flex_damping': False,
         'flex_dim': False,
         'flex_edge': False,
@@ -1158,12 +1215,21 @@ _BATCH_DIM = {
         'flex_edgenum': False,
         'flex_elem': False,
         'flex_elemadr': False,
+        'flex_elemdataadr': False,
         'flex_elemedge': False,
         'flex_elemedgeadr': False,
         'flex_elemnum': False,
+        'flex_friction': False,
+        'flex_margin': False,
+        'flex_radius': False,
+        'flex_shell': False,
+        'flex_shelldataadr': False,
+        'flex_shellnum': False,
         'flex_stiffness': False,
+        'flex_vert': False,
         'flex_vertadr': False,
         'flex_vertbodyid': False,
+        'flex_vertflexid': False,
         'flex_vertnum': False,
         'flexedge_J_colind': False,
         'flexedge_J_rowadr': False,
@@ -1234,6 +1300,7 @@ _BATCH_DIM = {
         'mat_rgba': True,
         'mat_texid': True,
         'mat_texrepeat': True,
+        'max_ten_J_rownnz': False,
         'mesh_face': False,
         'mesh_faceadr': False,
         'mesh_graph': False,
@@ -1259,6 +1326,7 @@ _BATCH_DIM = {
         'nC': False,
         'nJfe': False,
         'nJmom': False,
+        'nJten': False,
         'nM': False,
         'na': False,
         'nacttrnbody': False,
@@ -1272,6 +1340,7 @@ _BATCH_DIM = {
         'nflexelem': False,
         'nflexelemdata': False,
         'nflexelemedge': False,
+        'nflexshelldata': False,
         'nflexvert': False,
         'ngeom': False,
         'ngravcomp': False,
@@ -1353,6 +1422,8 @@ _BATCH_DIM = {
         'pair_solreffriction': True,
         'plugin': False,
         'plugin_attr': False,
+        'qLD_all_updates': False,
+        'qLD_level_offsets': False,
         'qLD_updates': False,
         'qM_fullm_i': False,
         'qM_fullm_j': False,
@@ -1398,6 +1469,9 @@ _BATCH_DIM = {
         'stat__meaninertia': True,
         'taxel_sensorid': False,
         'taxel_vertadr': False,
+        'ten_J_colind': False,
+        'ten_J_rowadr': False,
+        'ten_J_rownnz': False,
         'ten_wrapadr_site': False,
         'ten_wrapnum_site': False,
         'tendon_actfrclimited': False,