Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 945593266
Change-Id: Ic08842caabb92bc32974892ffae98e3295827094

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

@@ -19,8 +19,11 @@ 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 Geom
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_margin_gap
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_material_params
 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 geom_collision_pair_from_types
 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 epa_phase
@@ -707,8 +710,8 @@ def ccd_hfield_kernel_builder(
   return ccd_hfield_kernel
 
 
-_CCD_OVERSUBSCRIBE_WAVES = 4
-_CCD_MIN_BLOCKS = 2
+_CCD_OVERSUBSCRIBE_WAVES = 2
+_CCD_MIN_BLOCKS = 8
 
 
 @cache_kernel
@@ -728,6 +731,17 @@ def ccd_kernel_builder(
   def eval_ccd_write_contact(
     # Model:
     opt_ccd_tolerance: wp.array[float],
+    geom_condim: wp.array[int],
+    geom_priority: wp.array[int],
+    geom_solmix: wp.array2d[float],
+    geom_solref: wp.array2d[wp.vec2],
+    geom_solimp: wp.array2d[vec5],
+    geom_friction: wp.array2d[wp.vec3],
+    pair_dim: wp.array[int],
+    pair_solref: wp.array2d[wp.vec2],
+    pair_solreffriction: wp.array2d[wp.vec2],
+    pair_solimp: wp.array2d[vec5],
+    pair_friction: wp.array2d[vec5],
     # Data in:
     naconmax_in: int,
     naccdmax_in: int,
@@ -756,13 +770,6 @@ def ccd_kernel_builder(
     nccd_in: wp.array[int],
     margin: float,
     gap: float,
-    condim: int,
-    friction: vec5,
-    solref: wp.vec2,
-    solreffriction: wp.vec2,
-    solimp: vec5,
-    x1: wp.vec3,
-    x2: wp.vec3,
     pairid: wp.vec2i,
     # Data out:
     contact_dist_out: wp.array[float],
@@ -780,29 +787,26 @@ def ccd_kernel_builder(
     contact_type_out: wp.array[int],
     contact_geomcollisionid_out: wp.array[int],
     nacon_out: wp.array[int],
-    # Data out:
     overflow_out: wp.array[int],
-  ) -> int:
-    points = mat43()
-    witness1 = mat43()
-    witness2 = mat43()
+  ):
     geom1.margin = margin
     geom2.margin = margin
+    tolerance = opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]]
     is_collision_sensor = pairid[1] >= 0
     if is_collision_sensor:
       cutoff = 1.0e32
     else:
       cutoff = gap
     needs_epa, dist, ncollision, w1, w2, gjk_result, geom1, geom2 = gjk_phase(
-      opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
+      tolerance,
       cutoff,
       gjk_iterations,
       geom1,
       geom2,
       geomtype1,
       geomtype2,
-      x1,
-      x2,
+      geom1.pos,
+      geom2.pos,
     )
 
     ccdid = int(-1)
@@ -814,9 +818,9 @@ def ccd_kernel_builder(
         if wp.static(warn_overflow):
           wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid)
         wp.atomic_or(overflow_out, worldid, OverflowType.CCD)
-        return 0
+        return
       dist, ncollision, w1, w2, multiccd_idx = epa_phase(
-        opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
+        tolerance,
         epa_iterations,
         gjk_result,
         geom1,
@@ -831,8 +835,8 @@ def ccd_kernel_builder(
         epa_horizon_in[ccdid],
       )
 
-    if dist >= gap and pairid[1] == -1:
-      return 0
+    if dist >= gap and not is_collision_sensor:
+      return
 
     # CCD operates on margin-inflated shapes (support() inflates each geom by
     # 0.5 * margin).  The returned dist is therefore relative to the inflated
@@ -841,10 +845,16 @@ def ccd_kernel_builder(
     # with the primitive narrowphase, which reports un-inflated distances.
     dist += margin
 
+    witness1 = mat43()
+    witness2 = mat43()
     witness1[0] = w1
     witness2[0] = w2
 
-    if wp.static(use_multiccd or (geomtype1 == GeomType.BOX and geomtype2 == GeomType.BOX)):
+    if wp.static(
+      (use_multiccd or (geomtype1 == GeomType.BOX and geomtype2 == GeomType.BOX))
+      and (geomtype1 == GeomType.BOX or geomtype1 == GeomType.MESH)
+      and (geomtype2 == GeomType.BOX or geomtype2 == GeomType.MESH)
+    ):
       if wp.static(geomtype1 == GeomType.MESH):
         # verify that geom1 mesh data is present for multicontact
         if geom1.mesh_polyadr < 0:
@@ -879,23 +889,34 @@ def ccd_kernel_builder(
           geomtype2,
         )
 
-    for i in range(ncollision):
-      points[i] = 0.5 * (witness1[i] + witness2[i])
-    normal = witness1[0] - witness2[0]
-    frame = make_frame(normal)
+    condim, friction, solref, solreffriction, solimp = contact_material_params(
+      geom_condim,
+      geom_priority,
+      geom_solmix,
+      geom_solref,
+      geom_solimp,
+      geom_friction,
+      pair_dim,
+      pair_solref,
+      pair_solreffriction,
+      pair_solimp,
+      pair_friction,
+      geoms,
+      pairid[0],
+      worldid,
+    )
 
-    # flip if collision sensor
-    if pairid[1] >= 0:
+    frame = make_frame(witness1[0] - witness2[0])
+    if is_collision_sensor:
       frame *= -1.0
       geoms = wp.vec2i(geoms[1], geoms[0])
 
-    nactive = int(0)  # number of contacts contributing to the physics
     for i in range(ncollision):
-      active = write_contact(
+      write_contact(
         naconmax_in,
         i,
         dist,
-        points[i],
+        0.5 * (witness1[i] + witness2[i]),
         frame,
         margin,
         gap,
@@ -923,9 +944,6 @@ def ccd_kernel_builder(
         contact_geomcollisionid_out,
         nacon_out,
       )
-      nactive += active
-
-    return nactive
 
   # runs convex collision on a set of geom pairs to recover contact info (non-heightfield)
   @wp.kernel(module="unique", enable_backward=False, launch_bounds=(block_dim, _CCD_MIN_BLOCKS))
@@ -1022,31 +1040,18 @@ def ccd_kernel_builder(
         continue
 
       worldid = collision_worldid_in[collisionid]
-
-      _, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params(
-        geom_condim,
-        geom_priority,
-        geom_solmix,
-        geom_solref,
-        geom_solimp,
-        geom_friction,
+      pairid = collision_pairid_in[collisionid]
+      margin, gap = contact_margin_gap(
         geom_margin,
         geom_gap,
-        pair_dim,
-        pair_solref,
-        pair_solreffriction,
-        pair_solimp,
         pair_margin,
         pair_gap,
-        pair_friction,
-        collision_pair_in,
-        collision_pairid_in,
-        collisionid,
+        geoms,
+        pairid[0],
         worldid,
       )
 
-      geom1, geom2 = geom_collision_pair(
-        geom_type,
+      geom1, geom2 = geom_collision_pair_from_types(
         geom_dataid,
         geom_size,
         mesh_vertadr,
@@ -1065,12 +1070,25 @@ def ccd_kernel_builder(
         mesh_polymap,
         geom_xpos_in,
         geom_xmat_in,
+        geomtype1,
+        geomtype2,
         geoms,
         worldid,
       )
 
       eval_ccd_write_contact(
         opt_ccd_tolerance,
+        geom_condim,
+        geom_priority,
+        geom_solmix,
+        geom_solref,
+        geom_solimp,
+        geom_friction,
+        pair_dim,
+        pair_solref,
+        pair_solreffriction,
+        pair_solimp,
+        pair_friction,
         naconmax_in,
         naccdmax_in,
         epa_vert_in,
@@ -1097,14 +1115,7 @@ def ccd_kernel_builder(
         nccd_in,
         margin,
         gap,
-        condim,
-        friction,
-        solref,
-        solreffriction,
-        solimp,
-        geom1.pos,
-        geom2.pos,
-        collision_pairid_in[collisionid],
+        pairid,
         contact_dist_out,
         contact_pos_out,
         contact_frame_out,
@@ -1126,11 +1137,15 @@ def ccd_kernel_builder(
   return ccd_kernel
 
 
-def _ccd_grid_size(kernel, naconmax: int) -> int:
+def _ccd_grid_size(kernel, naconmax: int, device) -> int:
   # Grid-stride launch width for the CCD kernel: a few device waves, capped at the contact
   # capacity. The kernel strides over the actual candidate count, so we avoid launching one
   # (mostly idle) thread per naconmax slot.
-  block_size, min_grid_size = wp.get_suggested_block_size(kernel)
+  if device.is_cpu:
+    # Warp forces CPU block_dim to 1 and has no CUDA occupancy information.
+    return naconmax
+
+  block_size, min_grid_size = wp.get_suggested_block_size(kernel, device)
   return max(1, min(naconmax, _CCD_OVERSUBSCRIBE_WAVES * block_size * min_grid_size))
 
 
@@ -1326,7 +1341,7 @@ def convex_narrowphase(m: Model, d: Data, ctx: CollisionContext, collision_table
         m.block_dim.convex_ccd,
         bool(m.opt.warn_overflow),
       )
-      ccd_grid = _ccd_grid_size(ccd_k, d.naconmax)
+      ccd_grid = _ccd_grid_size(ccd_k, d.naconmax, d.ncollision.device)
       wp.launch(
         ccd_k,
         dim=ccd_grid,

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

@@ -16,7 +16,7 @@
 """Core collision types and utilities shared across collision modules."""
 
 import dataclasses
-from typing import Tuple
+from typing import Any, Tuple
 
 import warp as wp
 
@@ -63,9 +63,8 @@ class Geom:
 
 
 @wp.func
-def geom_collision_pair(
+def geom_collision_pair_from_types(
   # Model:
-  geom_type: wp.array[int],
   geom_dataid: wp.array2d[int],
   geom_size: wp.array2d[wp.vec3],
   mesh_vertadr: wp.array[int],
@@ -86,6 +85,8 @@ def geom_collision_pair(
   geom_xpos_in: wp.array2d[wp.vec3],
   geom_xmat_in: wp.array2d[wp.mat33],
   # In:
+  geom_type1: int,
+  geom_type2: int,
   geoms: wp.vec2i,
   worldid: int,
 ) -> Tuple[Geom, Geom]:
@@ -94,8 +95,6 @@ def geom_collision_pair(
 
   g1 = geoms[0]
   g2 = geoms[1]
-  geom_type1 = geom_type[g1]
-  geom_type2 = geom_type[g2]
 
   geom1.pos = geom_xpos_in[worldid, g1]
   geom1.rot = geom_xmat_in[worldid, g1]
@@ -156,6 +155,61 @@ def geom_collision_pair(
   return geom1, geom2
 
 
+@wp.func
+def geom_collision_pair(
+  # Model:
+  geom_type: wp.array[int],
+  geom_dataid: wp.array2d[int],
+  geom_size: wp.array2d[wp.vec3],
+  mesh_vertadr: wp.array[int],
+  mesh_vertnum: wp.array[int],
+  mesh_graphadr: wp.array[int],
+  mesh_vert: wp.array[wp.vec3],
+  mesh_graph: wp.array[int],
+  mesh_polynum: wp.array[int],
+  mesh_polyadr: wp.array[int],
+  mesh_polynormal: wp.array[wp.vec3],
+  mesh_polyvertadr: wp.array[int],
+  mesh_polyvertnum: wp.array[int],
+  mesh_polyvert: wp.array[int],
+  mesh_polymapadr: wp.array[int],
+  mesh_polymapnum: wp.array[int],
+  mesh_polymap: wp.array[int],
+  # Data in:
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
+  # In:
+  geoms: wp.vec2i,
+  worldid: int,
+) -> Tuple[Geom, Geom]:
+  geom_type1 = geom_type[geoms[0]]
+  geom_type2 = geom_type[geoms[1]]
+  return geom_collision_pair_from_types(
+    geom_dataid,
+    geom_size,
+    mesh_vertadr,
+    mesh_vertnum,
+    mesh_graphadr,
+    mesh_vert,
+    mesh_graph,
+    mesh_polynum,
+    mesh_polyadr,
+    mesh_polynormal,
+    mesh_polyvertadr,
+    mesh_polyvertnum,
+    mesh_polyvert,
+    mesh_polymapadr,
+    mesh_polymapnum,
+    mesh_polymap,
+    geom_xpos_in,
+    geom_xmat_in,
+    geom_type1,
+    geom_type2,
+    geoms,
+    worldid,
+  )
+
+
 @wp.func
 def write_contact(
   # Data in:
@@ -234,7 +288,33 @@ def write_contact(
 
 
 @wp.func
-def contact_params(
+def contact_margin_gap(
+  # Model:
+  geom_margin: wp.array2d[float],
+  geom_gap: wp.array2d[float],
+  pair_margin: wp.array2d[float],
+  pair_gap: wp.array2d[float],
+  # In:
+  geoms: wp.vec2i,
+  pairid: int,
+  worldid: int,
+) -> Tuple[float, float]:
+  if pairid > -1:
+    margin = pair_margin[worldid % pair_margin.shape[0], pairid]
+    gap = pair_gap[worldid % pair_gap.shape[0], pairid]
+  else:
+    g1 = geoms[0]
+    g2 = geoms[1]
+    margin_id = worldid % geom_margin.shape[0]
+    gap_id = worldid % geom_gap.shape[0]
+    margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
+    gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2]
+
+  return margin, gap
+
+
+@wp.func
+def contact_material_params(
   # Model:
   geom_condim: wp.array[int],
   geom_priority: wp.array[int],
@@ -242,35 +322,17 @@ def contact_params(
   geom_solref: wp.array2d[wp.vec2],
   geom_solimp: wp.array2d[vec5],
   geom_friction: wp.array2d[wp.vec3],
-  geom_margin: wp.array2d[float],
-  geom_gap: wp.array2d[float],
   pair_dim: wp.array[int],
   pair_solref: wp.array2d[wp.vec2],
   pair_solreffriction: wp.array2d[wp.vec2],
   pair_solimp: wp.array2d[vec5],
-  pair_margin: wp.array2d[float],
-  pair_gap: wp.array2d[float],
   pair_friction: wp.array2d[vec5],
   # In:
-  collision_pair_in: wp.array[wp.vec2i],
-  collision_pairid_in: wp.array[wp.vec2i],
-  cid: int,
+  geoms: wp.vec2i,
+  pairid: int,
   worldid: int,
 ):
-  """Resolve contact parameters for a collision pair.
-
-  Uses explicit pair overrides when available, otherwise mixes geom-level
-  properties by priority and solmix weights.
-  """
-  geoms = collision_pair_in[cid]
-  pairid = collision_pairid_in[cid][0]
-
-  # TODO(team): early return if collision sensor but no contact
-  # (ie, pairid[0] < -1 and pairid[1] < 0)
-
   if pairid > -1:
-    margin = pair_margin[worldid % pair_margin.shape[0], pairid]
-    gap = pair_gap[worldid % pair_gap.shape[0], pairid]
     condim = pair_dim[pairid]
     friction = pair_friction[worldid % pair_friction.shape[0], pairid]
     solref = pair_solref[worldid % pair_solref.shape[0], pairid]
@@ -283,8 +345,6 @@ def contact_params(
     friction_id = worldid % geom_friction.shape[0]
     solref_id = worldid % geom_solref.shape[0]
     solimp_id = worldid % geom_solimp.shape[0]
-    margin_id = worldid % geom_margin.shape[0]
-    gap_id = worldid % geom_gap.shape[0]
 
     solmix1 = geom_solmix[solmix_id, g1]
     solmix2 = geom_solmix[solmix_id, g2]
@@ -327,9 +387,6 @@ def contact_params(
 
     solreffriction = wp.vec2(0.0, 0.0)
     solimp = mix * geom_solimp[solimp_id, g1] + (1.0 - mix) * geom_solimp[solimp_id, g2]
-    # geom priority is ignored
-    margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
-    gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2]
 
   friction = vec5(
     wp.max(MJ_MINMU, friction[0]),
@@ -339,6 +396,62 @@ def contact_params(
     wp.max(MJ_MINMU, friction[4]),
   )
 
+  return condim, friction, solref, solreffriction, solimp
+
+
+@wp.func
+def contact_params(
+  # Model:
+  geom_condim: wp.array[int],
+  geom_priority: wp.array[int],
+  geom_solmix: wp.array2d[float],
+  geom_solref: wp.array2d[wp.vec2],
+  geom_solimp: wp.array2d[vec5],
+  geom_friction: wp.array2d[wp.vec3],
+  geom_margin: wp.array2d[float],
+  geom_gap: wp.array2d[float],
+  pair_dim: wp.array[int],
+  pair_solref: wp.array2d[wp.vec2],
+  pair_solreffriction: wp.array2d[wp.vec2],
+  pair_solimp: wp.array2d[vec5],
+  pair_margin: wp.array2d[float],
+  pair_gap: wp.array2d[float],
+  pair_friction: wp.array2d[vec5],
+  # In:
+  collision_pair_in: wp.array[wp.vec2i],
+  collision_pairid_in: wp.array[wp.vec2i],
+  cid: int,
+  worldid: int,
+):
+  """Resolve contact parameters for a collision pair.
+
+  Uses explicit pair overrides when available, otherwise mixes geom-level
+  properties by priority and solmix weights.
+  """
+  geoms = collision_pair_in[cid]
+  pairid = collision_pairid_in[cid][0]
+
+  # TODO(team): early return if collision sensor but no contact
+  # (ie, pairid[0] < -1 and pairid[1] < 0)
+
+  margin, gap = contact_margin_gap(geom_margin, geom_gap, pair_margin, pair_gap, geoms, pairid, worldid)
+  condim, friction, solref, solreffriction, solimp = contact_material_params(
+    geom_condim,
+    geom_priority,
+    geom_solmix,
+    geom_solref,
+    geom_solimp,
+    geom_friction,
+    pair_dim,
+    pair_solref,
+    pair_solreffriction,
+    pair_solimp,
+    pair_friction,
+    geoms,
+    pairid,
+    worldid,
+  )
+
   return geoms, margin, gap, condim, friction, solref, solreffriction, solimp
 
 
@@ -357,6 +470,111 @@ class CollisionContext:
   collision_worldid: wp.array
 
 
+@wp.func
+def sap_binary_search(values: wp.array[Any], value: Any, lower: int, upper: int) -> int:
+  """Binary search for the first element > value in sorted array."""
+  while lower < upper:
+    mid = (lower + upper) >> 1
+    if values[mid] > value:
+      upper = mid
+    else:
+      lower = mid + 1
+  return upper
+
+
+@wp.kernel
+def sap_range(
+  # In:
+  n: int,
+  lower_in: wp.array2d[float],
+  upper_in: wp.array2d[float],
+  sort_index_in: wp.array2d[int],
+  # Out:
+  range_out: wp.array2d[int],
+):
+  """Compute the sweep range for each sorted element."""
+  worldid, sortedid = wp.tid()
+
+  idx = sort_index_in[worldid, sortedid]
+  upper = upper_in[worldid, idx]
+
+  limit = sap_binary_search(lower_in[worldid], upper, sortedid + 1, n)
+  limit = wp.min(n - 1, limit)
+
+  range_out[worldid, sortedid] = limit - sortedid
+
+
+@wp.kernel
+def sap_sweep(
+  # In:
+  n: int,
+  sort_index_in: wp.array2d[int],
+  cumulative_sum_in: wp.array[int],
+  nsweep_in: int,
+  aabb_lower_in: wp.array2d[wp.vec3],
+  aabb_upper_in: wp.array2d[wp.vec3],
+  max_pairs: int,
+  # Out:
+  npairs_out: wp.array[int],
+  pair_id1_out: wp.array[int],
+  pair_id2_out: wp.array[int],
+  pair_worldid_out: wp.array[int],
+):
+  """Generic SAP sweep: output AABB-overlapping pairs.
+
+  This is the GPU equivalent of MuJoCo's mj_SAP function. It takes
+  axis-aligned bounding boxes and outputs pairs whose AABBs overlap
+  on all 3 axes. Domain-specific filtering is done by the caller.
+  """
+  worldelemid = wp.tid()
+
+  nworldelem = cumulative_sum_in.shape[0]
+  nworkpackages = cumulative_sum_in[nworldelem - 1]
+
+  while worldelemid < nworkpackages:
+    # Binary search to find sortedid (i) and partner sortedid (j)
+    i = sap_binary_search(cumulative_sum_in, worldelemid, 0, nworldelem)
+    j = i + worldelemid + 1
+    if i > 0:
+      j -= cumulative_sum_in[i - 1]
+
+    worldid = i // n
+    i = i % n
+    j = j % n
+
+    # Get actual element indices from sorted order
+    elem1 = sort_index_in[worldid, i]
+    elem2 = sort_index_in[worldid, j]
+
+    # Ensure elem1 < elem2 for consistent ordering
+    if elem1 > elem2:
+      tmp = elem1
+      elem1 = elem2
+      elem2 = tmp
+
+    worldelemid += nsweep_in
+
+    # AABB overlap test on all 3 axes
+    lower1 = aabb_lower_in[worldid, elem1]
+    upper1 = aabb_upper_in[worldid, elem1]
+    lower2 = aabb_lower_in[worldid, elem2]
+    upper2 = aabb_upper_in[worldid, elem2]
+
+    if lower1[0] > upper2[0] or lower2[0] > upper1[0]:
+      continue
+    if lower1[1] > upper2[1] or lower2[1] > upper1[1]:
+      continue
+    if lower1[2] > upper2[2] or lower2[2] > upper1[2]:
+      continue
+
+    # Output this pair
+    idx = wp.atomic_add(npairs_out, 0, 1)
+    if idx < max_pairs:
+      pair_id1_out[idx] = elem1
+      pair_id2_out[idx] = elem2
+      pair_worldid_out[idx] = worldid
+
+
 def create_collision_context(naconmax: int) -> CollisionContext:
   """Create a CollisionContext with allocated arrays."""
   return CollisionContext(

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

@@ -13,13 +13,15 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Any, Optional
+from typing import Optional
 
 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_core import sap_binary_search
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import sap_range
 from mujoco.mjx.third_party.mujoco_warp._src.collision_flex import flex_collision
 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
@@ -81,7 +83,6 @@ MJ_COLLISION_TABLE = {
 
 # TODO(team): Implement narrowphase flex collision support for:
 #             - HFIELD
-#             - ELLIPSOID
 #             - SDF
 MJ_FLEX_COLLISION_TABLE = {
   (GeomType.PLANE, GeomType.FLEX): CollisionType.PRIMITIVE,
@@ -90,6 +91,7 @@ MJ_FLEX_COLLISION_TABLE = {
   (GeomType.BOX, GeomType.FLEX): CollisionType.PRIMITIVE,
   (GeomType.CYLINDER, GeomType.FLEX): CollisionType.PRIMITIVE,
   (GeomType.MESH, GeomType.FLEX): CollisionType.CONVEX,
+  (GeomType.ELLIPSOID, GeomType.FLEX): CollisionType.CONVEX,
 }
 
 
@@ -369,18 +371,6 @@ def _add_geom_pair(
   collision_worldid_out[pairid] = worldid
 
 
-@wp.func
-def _binary_search(values: wp.array[Any], value: Any, lower: int, upper: int) -> int:
-  while lower < upper:
-    mid = (lower + upper) >> 1
-    if values[mid] > value:
-      upper = mid
-    else:
-      lower = mid + 1
-
-  return upper
-
-
 @cache_kernel
 def _sap_project(opt_broadphase: int):
   @wp.kernel(module="unique", enable_backward=False)
@@ -430,31 +420,6 @@ def _sap_project(opt_broadphase: int):
   return sap_project
 
 
-@wp.kernel
-def _sap_range(
-  # Model:
-  ngeom: int,
-  # In:
-  projection_lower_in: wp.array2d[float],
-  projection_upper_in: wp.array2d[float],
-  sort_index_in: wp.array2d[int],
-  # Out:
-  range_out: wp.array2d[int],
-):
-  worldid, geomid = wp.tid()
-
-  # current bounding geom
-  idx = sort_index_in[worldid, geomid]
-
-  upper = projection_upper_in[worldid, idx]
-
-  limit = _binary_search(projection_lower_in[worldid], upper, geomid + 1, ngeom)
-  limit = wp.min(ngeom - 1, limit)
-
-  # range of geoms for the sweep and prune process
-  range_out[worldid, geomid] = limit - geomid
-
-
 @cache_kernel
 def _sap_broadphase(
   opt_broadphase_filter: int,
@@ -501,7 +466,7 @@ def _sap_broadphase(
 
     while worldgeomid < nworkpackages:
       # binary search to find current and next geom pair indices
-      i = _binary_search(cumulative_sum_in, worldgeomid, 0, nworldgeom)
+      i = sap_binary_search(cumulative_sum_in, worldgeomid, 0, nworldgeom)
       j = i + worldgeomid + 1
 
       if i > 0:
@@ -670,7 +635,7 @@ def sap_broadphase(
     )
 
   wp.launch(
-    kernel=_sap_range,
+    kernel=sap_range,
     dim=(d.nworld, m.ngeom),
     inputs=[m.ngeom, projection_lower.reshape((-1, m.ngeom)), projection_upper, sort_index.reshape((-1, m.ngeom))],
     outputs=[range_],

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

@@ -78,6 +78,9 @@ class Polytope:
   vert_index: wp.array[int]
   nvert: int
 
+  # center point of polytope
+  center: wp.vec3
+
   # faces in polytope
   # 10 bits per each vertex index, while the last significant bits are for
   # invalid and deleted face
@@ -229,6 +232,10 @@ def _attach_face(pt: Polytope, idx: int, v1: int, v2: int, v3: int) -> float:
   if ret:
     return 0.0
 
+  # ensure projection points outward from the polytope
+  if wp.dot(r, p1 - pt.center) < 0.0:
+    r = -r
+
   face = v1 + (v2 << 10) + (v3 << 20)
   pt.face[idx] = face
   pt.face_pr[idx] = r
@@ -255,17 +262,14 @@ def _epa_support(
 
 
 @wp.func
-def _linear_combine(n: int, coefs: wp.vec4, mat: mat43) -> wp.vec3:
-  v = wp.vec3(0.0)
+def _linear_combine(n: int, scl: wp.vec4, mat: mat43) -> wp.vec3:
   if n == 1:
-    v = coefs[0] * mat[0]
-  elif n == 2:
-    v = coefs[0] * mat[0] + coefs[1] * mat[1]
-  elif n == 3:
-    v = coefs[0] * mat[0] + coefs[1] * mat[1] + coefs[2] * mat[2]
-  else:
-    v = coefs[0] * mat[0] + coefs[1] * mat[1] + coefs[2] * mat[2] + coefs[3] * mat[3]
-  return v
+    return scl[0] * mat[0]
+  if n == 2:
+    return scl[0] * mat[0] + scl[1] * mat[1]
+  if n == 3:
+    return scl[0] * mat[0] + scl[1] * mat[1] + scl[2] * mat[2]
+  return scl[0] * mat[0] + scl[1] * mat[1] + scl[2] * mat[2] + scl[3] * mat[3]
 
 
 @wp.func
@@ -278,11 +282,11 @@ def _subdistance(n: int, simplex: mat43) -> wp.vec4:
   if n == 4:
     return _S3D(simplex[0], simplex[1], simplex[2], simplex[3])
   if n == 3:
-    coordinates3 = _S2D(simplex[0], simplex[1], simplex[2])
-    return wp.vec4(coordinates3[0], coordinates3[1], coordinates3[2], 0.0)
+    lmbda3 = _S2D(simplex[0], simplex[1], simplex[2])
+    return wp.vec4(lmbda3[0], lmbda3[1], lmbda3[2], 0.0)
   if n == 2:
-    coordinates2 = _S1D(simplex[0], simplex[1])
-    return wp.vec4(coordinates2[0], coordinates2[1], 0.0, 0.0)
+    lmbda2 = _S1D(simplex[0], simplex[1])
+    return wp.vec4(lmbda2[0], lmbda2[1], 0.0, 0.0)
   return wp.vec4(1.0, 0.0, 0.0, 0.0)
 
 
@@ -368,51 +372,51 @@ def _S3D(s1: wp.vec3, s2: wp.vec3, s3: wp.vec3, s4: wp.vec3) -> wp.vec4:
     return wp.vec4(C41 / m_det, C42 / m_det, C43 / m_det, C44 / m_det)
 
   # find the smallest distance, and use the corresponding barycentric coordinates
-  coordinates = wp.vec4(0.0, 0.0, 0.0, 0.0)
+  lmbda = wp.vec4(0.0, 0.0, 0.0, 0.0)
   dmin = FLOAT_MAX
 
   if not comp1:
-    subcoord = _S2D(s2, s3, s4)
-    x = subcoord[0] * s2 + subcoord[1] * s3 + subcoord[2] * s4
+    sublmbda = _S2D(s2, s3, s4)
+    x = sublmbda[0] * s2 + sublmbda[1] * s3 + sublmbda[2] * s4
     d = wp.dot(x, x)
-    coordinates[0] = 0.0
-    coordinates[1] = subcoord[0]
-    coordinates[2] = subcoord[1]
-    coordinates[3] = subcoord[2]
+    lmbda[0] = 0.0
+    lmbda[1] = sublmbda[0]
+    lmbda[2] = sublmbda[1]
+    lmbda[3] = sublmbda[2]
     dmin = d
 
   if not comp2:
-    subcoord = _S2D(s1, s3, s4)
-    x = subcoord[0] * s1 + subcoord[1] * s3 + subcoord[2] * s4
+    sublmbda = _S2D(s1, s3, s4)
+    x = sublmbda[0] * s1 + sublmbda[1] * s3 + sublmbda[2] * s4
     d = wp.dot(x, x)
     if d < dmin:
-      coordinates[0] = subcoord[0]
-      coordinates[1] = 0.0
-      coordinates[2] = subcoord[1]
-      coordinates[3] = subcoord[2]
+      lmbda[0] = sublmbda[0]
+      lmbda[1] = 0.0
+      lmbda[2] = sublmbda[1]
+      lmbda[3] = sublmbda[2]
       dmin = d
 
   if not comp3:
-    subcoord = _S2D(s1, s2, s4)
-    x = subcoord[0] * s1 + subcoord[1] * s2 + subcoord[2] * s4
+    sublmbda = _S2D(s1, s2, s4)
+    x = sublmbda[0] * s1 + sublmbda[1] * s2 + sublmbda[2] * s4
     d = wp.dot(x, x)
     if d < dmin:
-      coordinates[0] = subcoord[0]
-      coordinates[1] = subcoord[1]
-      coordinates[2] = 0.0
-      coordinates[3] = subcoord[2]
+      lmbda[0] = sublmbda[0]
+      lmbda[1] = sublmbda[1]
+      lmbda[2] = 0.0
+      lmbda[3] = sublmbda[2]
       dmin = d
 
   if not comp4:
-    subcoord = _S2D(s1, s2, s3)
-    x = subcoord[0] * s1 + subcoord[1] * s2 + subcoord[2] * s3
+    sublmbda = _S2D(s1, s2, s3)
+    x = sublmbda[0] * s1 + sublmbda[1] * s2 + sublmbda[2] * s3
     d = wp.dot(x, x)
     if d < dmin:
-      coordinates[0] = subcoord[0]
-      coordinates[1] = subcoord[1]
-      coordinates[2] = subcoord[2]
-      coordinates[3] = 0.0
-  return coordinates
+      lmbda[0] = sublmbda[0]
+      lmbda[1] = sublmbda[1]
+      lmbda[2] = sublmbda[2]
+      lmbda[3] = 0.0
+  return lmbda
 
 
 @wp.func
@@ -528,36 +532,36 @@ def _S2D(s1: wp.vec3, s2: wp.vec3, s3: wp.vec3) -> wp.vec3:
 
   # find the smallest distance, and use the corresponding barycentric coordinates
   dmin = FLOAT_MAX
-  coordinates = wp.vec3(0.0, 0.0, 0.0)
+  lmbda = wp.vec3(0.0, 0.0, 0.0)
 
   if not comp1:
-    subcoord = _S1D(s2, s3)
-    x = subcoord[0] * s2 + subcoord[1] * s3
+    sublmbda = _S1D(s2, s3)
+    x = sublmbda[0] * s2 + sublmbda[1] * s3
     d = wp.dot(x, x)
-    coordinates[0] = 0.0
-    coordinates[1] = subcoord[0]
-    coordinates[2] = subcoord[1]
+    lmbda[0] = 0.0
+    lmbda[1] = sublmbda[0]
+    lmbda[2] = sublmbda[1]
     dmin = d
 
   if not comp2:
-    subcoord = _S1D(s1, s3)
-    x = subcoord[0] * s1 + subcoord[1] * s3
+    sublmbda = _S1D(s1, s3)
+    x = sublmbda[0] * s1 + sublmbda[1] * s3
     d = wp.dot(x, x)
     if d < dmin:
-      coordinates[0] = subcoord[0]
-      coordinates[1] = 0.0
-      coordinates[2] = subcoord[1]
+      lmbda[0] = sublmbda[0]
+      lmbda[1] = 0.0
+      lmbda[2] = sublmbda[1]
       dmin = d
 
   if not comp3:
-    subcoord = _S1D(s1, s2)
-    x = subcoord[0] * s1 + subcoord[1] * s2
+    sublmbda = _S1D(s1, s2)
+    x = sublmbda[0] * s1 + sublmbda[1] * s2
     d = wp.dot(x, x)
     if d < dmin:
-      coordinates[0] = subcoord[0]
-      coordinates[1] = subcoord[1]
-      coordinates[2] = 0.0
-  return coordinates
+      lmbda[0] = sublmbda[0]
+      lmbda[1] = sublmbda[1]
+      lmbda[2] = 0.0
+  return lmbda
 
 
 @wp.func
@@ -566,13 +570,18 @@ def _S1D(s1: wp.vec3, s2: wp.vec3) -> wp.vec2:
   p_o = _project_origin_line(s1, s2)
 
   # find the axis with the largest projection "shadow" of the simplex
-  mu_max = 0.0
+  mu_max = s1[0] - s2[0]
   index = 0
-  for i in range(3):
-    mu = s1[i] - s2[i]
-    if wp.abs(mu) >= wp.abs(mu_max):
-      mu_max = mu
-      index = i
+
+  mu = s1[1] - s2[1]
+  if wp.abs(mu) >= wp.abs(mu_max):
+    mu_max = mu
+    index = 1
+
+  mu = s1[2] - s2[2]
+  if wp.abs(mu) >= wp.abs(mu_max):
+    mu_max = mu
+    index = 2
 
   C1 = p_o[index] - s2[index]
   C2 = s1[index] - p_o[index]
@@ -605,21 +614,21 @@ def gjk(
   simplex_index1 = wp.vec4i()
   simplex_index2 = wp.vec4i()
   n = int(0)
-  coordinates = wp.vec4()  # barycentric coordinates
+  lmbda = wp.vec4()  # barycentric coordinates
   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
+  xnorm2_old = FLOAT_MAX
 
   for _ in range(gjk_iterations):
-    xnorm = wp.dot(x_k, x_k)
+    xnorm2 = wp.dot(x_k, x_k)
     # TODO(kbayes): determine new constant here
-    if xnorm < tol2 or wp.abs(xnorm_old - xnorm) < tol2:
+    if xnorm2 < tol2 or wp.abs(xnorm2_old - xnorm2) < tol2:
       break
-    xnorm_old = xnorm
-    dir_neg = x_k / wp.sqrt(xnorm)
+    xnorm2_old = xnorm2
+    dir_neg = x_k / wp.sqrt(xnorm2)
 
     # compute kth support point in geom1
     sp = support(geom1, geomtype1, -dir_neg)
@@ -636,6 +645,11 @@ def gjk(
     # compute the kth support point
     simplex[n] = simplex1[n] - simplex2[n]
 
+    # stopping criteria using the Frank-Wolfe duality gap given by
+    #  |f(x_k) - f(x_min)|^2 <= < grad f(x_k), (x_k - simplex[n]) >
+    if wp.dot(x_k, x_k - simplex[n]) < epsilon:
+      break
+
     if cutoff == 0.0:
       if wp.dot(x_k, simplex[n]) > 0.0:
         result = GJKResult()
@@ -644,25 +658,20 @@ def gjk(
         return result
     elif cutoff < FLOAT_MAX:
       vs = wp.dot(x_k, simplex[n])
-      if wp.dot(x_k, simplex[n]) > 0.0 and (vs * vs / xnorm) >= cutoff2:
+      if wp.dot(x_k, simplex[n]) > 0.0 and (vs * vs / xnorm2) >= cutoff2:
         result = GJKResult()
         result.dim = 0
         result.dist = FLOAT_MAX
         return result
 
-    # stopping criteria using the Frank-Wolfe duality gap given by
-    #  |f(x_k) - f(x_min)|^2 <= < grad f(x_k), (x_k - simplex[n]) >
-    if wp.dot(x_k, x_k - simplex[n]) < epsilon:
-      break
-
     # run the distance subalgorithm to compute the barycentric coordinates
     # of the closest point to the origin in the simplex
-    coordinates = _subdistance(n + 1, simplex)
+    lmbda = _subdistance(n + 1, simplex)
 
     # remove vertices from the simplex no longer needed
     n = int(0)
     for i in range(4):
-      if coordinates[i] == 0.0:
+      if lmbda[i] == 0.0:
         continue
 
       simplex[n] = simplex[i]
@@ -670,7 +679,7 @@ def gjk(
       simplex2[n] = simplex2[i]
       simplex_index1[n] = simplex_index1[i]
       simplex_index2[n] = simplex_index2[i]
-      coordinates[n] = coordinates[i]
+      lmbda[n] = lmbda[i]
       n += int(1)
 
     # SHOULD NOT OCCUR
@@ -678,7 +687,7 @@ def gjk(
       break
 
     # get the next iteration of x_k
-    x_next = _linear_combine(n, coordinates, simplex)
+    x_next = _linear_combine(n, lmbda, simplex)
 
     # x_k has converged to minimum
     if _almost_equal(x_next, x_k):
@@ -696,8 +705,8 @@ def gjk(
   # compute the approximate witness points
   # if n is zero, then there was an immediate return meaning the initial points
   # are the witness points
-  result.x1 = wp.where(n == 0, x1_0, _linear_combine(n, coordinates, simplex1))
-  result.x2 = wp.where(n == 0, x2_0, _linear_combine(n, coordinates, simplex2))
+  result.x1 = wp.where(n == 0, x1_0, _linear_combine(n, lmbda, simplex1))
+  result.x2 = wp.where(n == 0, x2_0, _linear_combine(n, lmbda, simplex2))
   result.dist = wp.norm_l2(x_k)
 
   result.dim = n
@@ -974,6 +983,9 @@ def _polytope2(
   """Create polytope for EPA given a 1-simplex from GJK."""
   diff = simplex[1] - simplex[0]
 
+  # set the polytope center
+  pt.center = 0.5 * (simplex[0] + simplex[1])
+
   # find component with smallest magnitude (so cross product is largest)
   value = FLOAT_MAX
   index = 0
@@ -1063,6 +1075,9 @@ def _polytope3(
   geomtype2: int,
 ) -> Polytope:
   """Create polytope for EPA given a 2-simplex from GJK."""
+  # set the polytope center
+  pt.center = (simplex[0] + simplex[1] + simplex[2]) * wp.static(1.0 / 3.0)
+
   # get normals in both directions
   n = wp.cross(simplex[1] - simplex[0], simplex[2] - simplex[0])
   if wp.norm_l2(n) < MINVAL:
@@ -1146,6 +1161,9 @@ def _polytope4(
   simplex_index2: wp.vec4i,
 ) -> Tuple[Polytope, GJKResult]:
   """Create polytope for EPA given a 3-simplex from GJK."""
+  # set the polytope center
+  pt.center = 0.25 * (simplex[0] + simplex[1] + simplex[2] + simplex[3])
+
   pt.vert[0] = simplex1[0]
   pt.vert[1] = simplex2[0]
   pt.vert[2] = simplex1[1]
@@ -1269,15 +1287,15 @@ def _epa(
     # compute support point w from the closest face's normal
     lower = wp.sqrt(lower2)
     wi = pt.nvert
-    face_pr_normalized = pt.face_pr[idx] / lower
-    i1, i2 = _epa_support(pt, wi, geom1, geom2, geomtype1, geomtype2, face_pr_normalized)
+    face_pr = pt.face_pr[idx]
+    i1, i2 = _epa_support(pt, wi, geom1, geom2, geomtype1, geomtype2, face_pr / lower)
     w = pt.vert[2 * wi] - pt.vert[2 * wi + 1]
     geom1.index = i1
     geom2.index = i2
     pt.nvert += 1
 
-    # upper bound for kth iteration
-    upper_k = wp.dot(face_pr_normalized, w)
+    # upper bound for kth iteration (dot product before normalizing for better precision)
+    upper_k = wp.dot(face_pr, w) / lower
     if upper_k < upper:
       upper = upper_k
       upper2 = upper * upper

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

@@ -143,6 +143,20 @@ def _jtdaj_groups(mjd: mujoco.MjData) -> tuple[np.ndarray, np.ndarray]:
   return adr, nrow
 
 
+def _get_nflexintcell(mjm: mujoco.MjModel) -> int:
+  nflexintcell = 0
+  if mjm.nflex > 0 and hasattr(mjm, "flex_interp"):
+    for fi in range(mjm.nflex):
+      order = abs(int(mjm.flex_interp[fi]))
+      if order == 0:
+        continue
+      if hasattr(mjm, "flex_edgeequality") and mjm.flex_edgeequality[fi] == 3:
+        continue
+      cx, cy, cz = mjm.flex_cellnum[fi]
+      nflexintcell += int(cx) * int(cy) * int(cz)
+  return nflexintcell
+
+
 def is_sparse(mjm: mujoco.MjModel) -> bool:
   if mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_AUTO:
     if mjm.nv > 32:
@@ -330,6 +344,9 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
   if (mjm.opt.enableflags & mujoco.mjtEnableBit.mjENBL_SLEEP) and (mjm.eq_type == mujoco.mjtEq.mjEQ_FLEX).any():
     raise NotImplementedError("Flex equality constraints are not supported with sleeping enabled.")
 
+  if mjm.nflex > 0 and (mjm.flex_interp < 0).any():
+    raise NotImplementedError("Flex interpolation order < 0 (shell/quad elements) is not supported.")
+
   if mjm.opt.noslip_iterations > 0:
     raise NotImplementedError(f"noslip solver not implemented.")
 
@@ -437,10 +454,16 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
   condim_arrays = [np.array([0]), mjm.geom_condim, mjm.pair_dim]
   if mjm.nflex > 0:
     condim_arrays.append(mjm.flex_condim)
+    if (mjm.geom_type == mujoco.mjtGeom.mjGEOM_SDF).any():
+      raise NotImplementedError("Flex-SDF collision is not implemented.")
+    if (mjm.geom_type == mujoco.mjtGeom.mjGEOM_HFIELD).any():
+      raise NotImplementedError("Flex-HField collision is not implemented.")
   m.nmaxcondim = np.concatenate(condim_arrays).max()
   m.nmaxpyramid = np.maximum(1, 2 * (m.nmaxcondim - 1))
   m.has_sdf_geom = (mjm.geom_type == mujoco.mjtGeom.mjGEOM_SDF).any()
+  m.has_ellipsoid_geom = (mjm.geom_type == mujoco.mjtGeom.mjGEOM_ELLIPSOID).any()
   m.has_flex_selfcollide = bool(mjm.nflex > 0 and np.any(mjm.flex_selfcollide != 0))
+  m.has_3d_flex = bool(mjm.nflex > 0 and np.any(mjm.flex_dim == 3))
   m.max_flex_dim = int(np.max(mjm.flex_dim)) if mjm.nflex > 0 else 0
   m.block_dim = types.BlockDim()
   # Derive CG solver block_dim from nv: clamp(round_up_to_32(nv), 32, 256)
@@ -453,6 +476,28 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
     m.block_dim.linesearch_iterative = 512
   m.is_sparse = is_sparse(mjm)
   m.has_fluid = mjm.opt.wind.any() or mjm.opt.density > 0 or mjm.opt.viscosity > 0
+
+  m.nflexintcell = _get_nflexintcell(mjm)
+
+  # Precompute flex_cell_map
+  flex_cell_map = []
+  if mjm.nflex > 0 and hasattr(mjm, "flex_interp"):
+    for fi in range(mjm.nflex):
+      order = abs(int(mjm.flex_interp[fi]))
+      if order == 0:
+        continue
+      if hasattr(mjm, "flex_edgeequality") and mjm.flex_edgeequality[fi] == 3:
+        continue
+      cx, cy, cz = mjm.flex_cellnum[fi]
+      for ci in range(cx):
+        for cj in range(cy):
+          for ck in range(cz):
+            flex_cell_map.append((fi, ci, cj, ck))
+  if not flex_cell_map:
+    m.flex_cell_map = np.zeros((0, 4), dtype=np.int32)
+  else:
+    m.flex_cell_map = np.array(flex_cell_map, dtype=np.int32)
+
   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)
@@ -699,6 +744,49 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
   m.eq_jnt_adr = np.nonzero(mjm.eq_type == types.EqType.JOINT)[0]
   m.eq_ten_adr = np.nonzero(mjm.eq_type == types.EqType.TENDON)[0]
   m.eq_flex_adr = np.nonzero(mjm.eq_type == types.EqType.FLEX)[0]
+  m.eq_flexstrain_adr = np.nonzero(mjm.eq_type == types.EqType.FLEXSTRAIN)[0]
+  m.neq_flexstrain = m.eq_flexstrain_adr.size
+
+  # Precompute flex strain Jacobian sparsity pattern
+  flexstrain_J_rownnz = []
+  flexstrain_J_colind = []
+
+  if m.neq_flexstrain > 0:
+    for eqstrainid, eqid in enumerate(m.eq_flexstrain_adr):
+      f = int(mjm.eq_obj1id[eqid])
+      order = int(mjm.flex_interp[f])
+      ci = int(mjm.eq_data[eqid, 0])
+      cj = int(mjm.eq_data[eqid, 1])
+      ck = int(mjm.eq_data[eqid, 2])
+
+      cellnum = mjm.flex_cellnum[f]
+      cy = cellnum[1]
+      cz = cellnum[2]
+      nstart = mjm.flex_nodeadr[f]
+      ny_g = cy * order + 1
+      nz_g = cz * order + 1
+
+      node_bodies = [
+        mjm.flex_nodebodyid[nstart + (ci * order + li) * ny_g * nz_g + (cj * order + lj) * nz_g + (ck * order + lk)]
+        for li in range(order + 1)
+        for lj in range(order + 1)
+        for lk in range(order + 1)
+      ]
+
+      active_dof_mask = np.any(body_isdofancestor[node_bodies, :] != 0, axis=0)
+      sorted_dofs = np.nonzero(active_dof_mask)[0].tolist()
+      flexstrain_J_rownnz.append(len(sorted_dofs))
+      flexstrain_J_colind.extend(sorted_dofs)
+
+    m.flexstrain_J_rownnz = np.array(flexstrain_J_rownnz, dtype=np.int32)
+    m.flexstrain_J_colind = np.array(flexstrain_J_colind, dtype=np.int32)
+    m.flexstrain_J_rowadr = np.cumsum([0] + flexstrain_J_rownnz[:-1], dtype=np.int32)
+  else:
+    m.flexstrain_J_rownnz = np.zeros((0,), dtype=np.int32)
+    m.flexstrain_J_rowadr = np.zeros((0,), dtype=np.int32)
+    m.flexstrain_J_colind = np.zeros((0,), dtype=np.int32)
+
+  m.nJfs = m.flexstrain_J_colind.size
 
   # fixed tendon
   m.tendon_jnt_adr, m.wrap_jnt_adr = [], []
@@ -952,7 +1040,6 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
 
   # Populate lookup maps and candidate pairs
   flexelem_geom_pairs = []
-  flexshell_geom_pairs = []
   flexvert_geom_pairs = []
 
   flex_elemflexid = np.zeros(mjm.nflexelem, dtype=np.int32)
@@ -996,6 +1083,7 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
           mujoco.mjtGeom.mjGEOM_BOX,
           mujoco.mjtGeom.mjGEOM_CYLINDER,
           mujoco.mjtGeom.mjGEOM_MESH,
+          mujoco.mjtGeom.mjGEOM_ELLIPSOID,
         ],
       )
       is_pl = mjm.geom_type == mujoco.mjtGeom.mjGEOM_PLANE
@@ -1018,19 +1106,6 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
             for g in matching_primitive_geoms[~filtered]:
               flexelem_geom_pairs.append((elemid, g))
 
-      elif fdim == 3:
-        shelldata_start = mjm.flex_shelldataadr[fi]
-        prev_shells_offset = shell_offset - shell_num
-        for s in range(shell_num):
-          v0 = vert_start + mjm.flex_shell[shelldata_start + s * 3]
-          v1 = vert_start + mjm.flex_shell[shelldata_start + s * 3 + 1]
-          v2 = vert_start + mjm.flex_shell[shelldata_start + s * 3 + 2]
-
-          if len(matching_primitive_geoms) > 0:
-            filtered = _filter_tri_geoms(mjm, v0, v1, v2, matching_primitive_geoms, filterparent)
-            for g in matching_primitive_geoms[~filtered]:
-              flexshell_geom_pairs.append((prev_shells_offset + s, g))
-
       # Planes vs Vertices
       if len(matching_plane_geoms) > 0:
         vert_count = mjm.flex_vertnum[fi]
@@ -1057,13 +1132,10 @@ def put_model(mjm: mujoco.MjModel, batch_sizes: dict[str, int] | None = None) ->
 
   if not flexelem_geom_pairs:
     flexelem_geom_pairs = np.zeros((0, 2), dtype=np.int32)
-  if not flexshell_geom_pairs:
-    flexshell_geom_pairs = np.zeros((0, 2), dtype=np.int32)
   if not flexvert_geom_pairs:
     flexvert_geom_pairs = np.zeros((0, 2), dtype=np.int32)
 
   m.flexelem_geom_pair_filtered = np.array(flexelem_geom_pairs, dtype=np.int32)
-  m.flexshell_geom_pair_filtered = np.array(flexshell_geom_pairs, dtype=np.int32)
   m.flexvert_geom_pair_filtered = np.array(flexvert_geom_pairs, dtype=np.int32)
 
   m.flex_elemflexid = flex_elemflexid
@@ -1208,6 +1280,15 @@ def _default_njmax_nnz(mjm: mujoco.MjModel, nconmax: int, njmax: int) -> int:
         for e in range(edge_count):
           total_nnz += mjm.flexedge_J_rownnz[edge_start + e]
 
+    elif eq_type == mujoco.mjtEq.mjEQ_FLEXSTRAIN:
+      # strain constraints: each cell produces neig rows, each dense (nv)
+      obj1id = mjm.eq_obj1id[i]
+      if obj1id < mjm.nflex and hasattr(mjm, "flex_stiffnessadr"):
+        # estimate neig from stiffness data
+        adr = mjm.flex_stiffnessadr[obj1id]
+        neig = int(mjm.flex_stiffness[adr])
+        total_nnz += neig * mjm.nv
+
   # friction constraints
   total_nnz += (mjm.dof_frictionloss > 0).sum()
   for i in range(mjm.ntendon):
@@ -1494,6 +1575,7 @@ def make_data(
   sizes["njmax"] = njmax
   sizes["nvmax"] = nvmax
   sizes["nvmax_pad"] = _nvmax_pad(nvmax)
+  sizes["nflexintcell"] = _get_nflexintcell(mjm)
 
   if njmax_nnz is None:
     if is_sparse(mjm):
@@ -2523,6 +2605,83 @@ def _copy_tendon_length0(
   tendon_length0_out[tendon_length0_id, tenid] = ten_length_in[worldid, tenid]
 
 
+@wp.kernel
+def _compute_eq_data0(
+  # Model:
+  eq_type: wp.array[int],
+  eq_obj1id: wp.array[int],
+  eq_obj2id: wp.array[int],
+  eq_objtype: wp.array[int],
+  # Data in:
+  xpos_in: wp.array2d[wp.vec3],
+  xquat_in: wp.array2d[wp.quat],
+  xmat_in: wp.array2d[wp.mat33],
+  # Out:
+  eq_data_out: wp.array2d[types.vec11],
+):
+  """Compute eq_data for connect/weld constraints.
+
+  Kinematics must have been evaluated at qpos0 so the constraint is satisfied at qpos0.
+  """
+  worldid, eqid = wp.tid()
+  eq_data_id = worldid % eq_data_out.shape[0]
+
+  eqtype = eq_type[eqid]
+  objtype = eq_objtype[eqid]
+  data = eq_data_out[eq_data_id, eqid]
+
+  if eqtype == int(types.EqType.CONNECT.value):
+    if objtype == int(types.ObjType.BODY.value):
+      obj1id = eq_obj1id[eqid]
+      obj2id = eq_obj2id[eqid]
+
+      # data[0:3] = anchor in body1 local frame; map to global frame
+      anchor1 = wp.vec3(data[0], data[1], data[2])
+      pos = xpos_in[worldid, obj1id] + xmat_in[worldid, obj1id] @ anchor1
+
+      # data[3:6] = anchor position in body2 local frame
+      anchor2 = wp.transpose(xmat_in[worldid, obj2id]) @ (pos - xpos_in[worldid, obj2id])
+      data[3] = anchor2[0]
+      data[4] = anchor2[1]
+      data[5] = anchor2[2]
+      eq_data_out[eq_data_id, eqid] = data
+    elif objtype == int(types.ObjType.SITE.value):
+      # site-based connect, eq_data is unused
+      eq_data_out[eq_data_id, eqid] = types.vec11(0.0)
+  elif eqtype == int(types.EqType.WELD.value):
+    if objtype == int(types.ObjType.BODY.value):
+      quat = wp.quat(data[6], data[7], data[8], data[9])
+      if wp.length_sq(quat) > 0.0:
+        # user has set quaternion data: normalize it and keep the remaining data
+        quat = wp.normalize(quat)
+        data[6] = quat[0]
+        data[7] = quat[1]
+        data[8] = quat[2]
+        data[9] = quat[3]
+        eq_data_out[eq_data_id, eqid] = data
+      else:
+        obj1id = eq_obj1id[eqid]
+        obj2id = eq_obj2id[eqid]
+
+        # data[0:3] = anchor in body2 local frame; map to global frame
+        anchor2 = wp.vec3(data[0], data[1], data[2])
+        pos = xpos_in[worldid, obj2id] + xmat_in[worldid, obj2id] @ anchor2
+
+        # data[3:6] = anchor position in body1 local frame
+        anchor1 = wp.transpose(xmat_in[worldid, obj1id]) @ (pos - xpos_in[worldid, obj1id])
+        data[3] = anchor1[0]
+        data[4] = anchor1[1]
+        data[5] = anchor1[2]
+
+        # data[6:10] = neg(xquat1) * xquat2 = "xquat2 - xquat1" in body1 local frame
+        relquat = mjmath.mul_quat(mjmath.quat_inv(xquat_in[worldid, obj1id]), xquat_in[worldid, obj2id])
+        data[6] = relquat[0]
+        data[7] = relquat[1]
+        data[8] = relquat[2]
+        data[9] = relquat[3]
+        eq_data_out[eq_data_id, eqid] = data
+
+
 @wp.kernel
 def _resolve_tendon_lengthspring(
   ten_length_in: wp.array2d[float],
@@ -3005,6 +3164,8 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
 
   Computes:
     - tendon_length0: tendon resting lengths
+    - eq_data: connect/weld anchor data, recomputed so the constraint is
+      satisfied at qpos0
     - dof_invweight0: inverse inertia for DOFs
     - body_invweight0: inverse spatial inertia for bodies
     - tendon_invweight0: inverse weight for tendons
@@ -3044,6 +3205,13 @@ def set_const_0(m: types.Model, d: types.Data, restore: bool = True):
 
   wp.launch(_copy_tendon_length0, dim=(d.nworld, m.ntendon), inputs=[d.ten_length], outputs=[m.tendon_length0])
 
+  wp.launch(
+    _compute_eq_data0,
+    dim=(d.nworld, m.neq),
+    inputs=[m.eq_type, m.eq_obj1id, m.eq_obj2id, m.eq_objtype, d.xpos, d.xquat, d.xmat],
+    outputs=[m.eq_data],
+  )
+
   # dof_invweight0: computed per joint with averaging for multi-DOF joints
   # FREE: 6 DOFs, trans gets mean(A[0:3]), rot gets mean(A[3:6])
   # BALL: 3 DOFs, all get mean(A[0:3])

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

@@ -652,6 +652,11 @@ def _flex_elasticity(
 
   elem_data_adr = flex_elemdataadr[f] + local_elemid * (dim + 1)
   vbase = flex_vertadr[f]
+
+  # skip trilinear/interp elements (vertbodyid == -1, no simplex stiffness)
+  vert0_check = flex_elem[elem_data_adr]
+  if flex_vertbodyid[vbase + vert0_check] < 0:
+    return
   gradient = wp.matrix(0.0, shape=(6, 6))
   for e in range(nedge):
     vert0 = flex_elem[elem_data_adr + edges[e, 0]]
@@ -765,6 +770,180 @@ def _flex_bending(
       wp.atomic_add(qfrc_spring_out, worldid, body_dofadr[bodyid] + x, force[i, x])
 
 
+@wp.kernel
+def _flex_passive_interp(
+  # Model:
+  nflex: int,
+  body_rootid: wp.array[int],
+  body_dofnum: wp.array[int],
+  body_dofadr: wp.array[int],
+  flex_interp: wp.array[int],
+  flex_cellnum: wp.array[wp.vec3i],
+  flex_nodeadr: wp.array[int],
+  flex_stiffnessadr: wp.array[int],
+  flex_nodebodyid: wp.array[int],
+  flex_node: wp.array[wp.vec3],
+  flex_node0: wp.array[wp.vec3],
+  flex_stiffness: wp.array[float],
+  flex_damping: wp.array[float],
+  flex_edgeequality: wp.array[int],
+  flex_centered: wp.array[bool],
+  flex_cell_map: wp.array[wp.vec4i],
+  # Data in:
+  subtree_com_in: wp.array2d[wp.vec3],
+  cvel_in: wp.array2d[wp.spatial_vector],
+  flexnode_xpos_in: wp.array2d[wp.vec3],
+  # In:
+  dsbl_spring: bool,
+  dsbl_damper: bool,
+  # Data out:
+  qfrc_spring_out: wp.array2d[float],
+  qfrc_damper_out: wp.array2d[float],
+  # Out:
+  displ_scratch_out: wp.array3d[wp.vec3],
+  vel_corot_scratch_out: wp.array3d[wp.vec3],
+):
+  """Corotational passive forces for interpolated flex (trilinear/quadratic)."""
+  worldid, cellid = wp.tid()
+
+  mapping = flex_cell_map[cellid]
+  f = mapping[0]
+  ci = mapping[1]
+  cj = mapping[2]
+  ck = mapping[3]
+
+  order = flex_interp[f]
+  if order <= 0:
+    return
+
+  npc = (order + 1) * (order + 1) * (order + 1)
+  ndof_cell = 3 * npc
+
+  cellnum = flex_cellnum[f]
+  cy = cellnum[1]
+  cz = cellnum[2]
+  nstart = flex_nodeadr[f]
+  ny_g = cy * order + 1
+  nz_g = cz * order + 1
+
+  # Cell stiffness matrix address
+  stiffness_adr_base = flex_stiffnessadr[f]
+  if stiffness_adr_base < 0:
+    return
+
+  cell_idx = ci * cy * cz + cj * cz + ck
+  k_base = stiffness_adr_base + cell_idx * ndof_cell * ndof_cell
+
+  # Skip empty cells (zero stiffness)
+  if flex_stiffness[k_base] == 0.0:
+    return
+
+  cell_quat = support.compute_interp_cell_quat(flexnode_xpos_in, order, ci, cj, ck, cy, cz, ny_g, nz_g, nstart, worldid)
+
+  # mju_negQuat: conjugate (R⁻¹) — negate xyz, keep w
+  cell_quat_inv = wp.quat(-cell_quat[0], -cell_quat[1], -cell_quat[2], cell_quat[3])
+
+  # Pre-compute displacements and velocities in corotational frame
+  # (matches C: rotate all positions/velocities once, then K*u)
+  idx_j = int(0)
+  for li_j in range(order + 1):
+    for lj_j in range(order + 1):
+      for lk_j in range(order + 1):
+        if idx_j < npc:
+          gi_j = ci * order + li_j
+          gj_j = cj * order + lj_j
+          gk_j = ck * order + lk_j
+          gidx_j = gi_j * ny_g * nz_g + gj_j * nz_g + gk_j
+
+          xpos_j = flexnode_xpos_in[worldid, nstart + gidx_j]
+
+          if not dsbl_spring:
+            refpos_j = flex_node0[nstart + gidx_j]
+            xrot_j = wp.quat_rotate(cell_quat_inv, xpos_j)
+            displ_scratch_out[worldid, cellid, idx_j] = xrot_j - refpos_j
+
+          if not dsbl_damper:
+            bodyid_j = flex_nodebodyid[nstart + gidx_j]
+            cvel_j = cvel_in[worldid, bodyid_j]
+            omega_j = wp.spatial_top(cvel_j)
+            vcom_j = wp.spatial_bottom(cvel_j)
+            com_j = subtree_com_in[worldid, body_rootid[bodyid_j]]
+            r_j = xpos_j - com_j
+            vel_world_j = vcom_j + wp.cross(omega_j, r_j)
+            vel_corot_scratch_out[worldid, cellid, idx_j] = wp.quat_rotate(cell_quat_inv, vel_world_j)
+
+          idx_j += 1
+
+  # Compute K*displacement and K*velocity per output node, then scatter forces
+  idx_i = int(0)
+  for li_i in range(order + 1):
+    for lj_i in range(order + 1):
+      for lk_i in range(order + 1):
+        if idx_i < npc:
+          gi_i = ci * order + li_i
+          gj_i = cj * order + lj_i
+          gk_i = ck * order + lk_i
+          gidx_i = gi_i * ny_g * nz_g + gj_i * nz_g + gk_i
+          bodyid_i = flex_nodebodyid[nstart + gidx_i]
+
+          frc_spring = wp.vec3(0.0)
+          frc_damper = wp.vec3(0.0)
+
+          for comp_i in range(3):
+            row = idx_i * 3 + comp_i
+            val_spring = float(0.0)
+            val_damper = float(0.0)
+
+            for idx_j in range(npc):
+              for comp_j in range(3):
+                col = idx_j * 3 + comp_j
+                K_ij = flex_stiffness[k_base + row * ndof_cell + col]
+
+                if not dsbl_spring:
+                  val_spring += K_ij * displ_scratch_out[worldid, cellid, idx_j][comp_j]
+
+                if not dsbl_damper:
+                  val_damper += K_ij * vel_corot_scratch_out[worldid, cellid, idx_j][comp_j]
+
+            frc_spring[comp_i] = val_spring
+            frc_damper[comp_i] = val_damper
+
+          # Rotate forces back to world frame (R)
+          frc_spring_world = wp.quat_rotate(cell_quat, frc_spring)
+          frc_damper_world = wp.quat_rotate(cell_quat, frc_damper)
+
+          # Scale damper force by damping coefficient
+          frc_damper_world = frc_damper_world * flex_damping[f]
+
+          # Apply forces to body DOFs (fast path: nodes at body origin)
+          dofnum_i = body_dofnum[bodyid_i]
+          dofadr_i = body_dofadr[bodyid_i]
+          if dofnum_i > 0:
+            centered = flex_centered[f]
+            node_local = flex_node[nstart + gidx_i]
+            at_origin = node_local[0] == 0.0 and node_local[1] == 0.0 and node_local[2] == 0.0
+
+            if centered or at_origin:
+              for x in range(3):
+                if x < dofnum_i:
+                  if not dsbl_spring:
+                    wp.atomic_add(
+                      qfrc_spring_out,
+                      worldid,
+                      dofadr_i + x,
+                      frc_spring_world[x],
+                    )
+                  if not dsbl_damper:
+                    wp.atomic_add(
+                      qfrc_damper_out,
+                      worldid,
+                      dofadr_i + x,
+                      frc_damper_world[x],
+                    )
+
+          idx_i += 1
+
+
 @event_scope
 def passive(m: Model, d: Data):
   """Adds all passive forces."""
@@ -893,6 +1072,44 @@ def passive(m: Model, d: Data):
       outputs=[d.qfrc_gravcomp],
     )
 
+  # Launch passive interp kernel for interpolated flex (trilinear/quadratic)
+  if m.nflex and m.nflexintcell > 0:
+    displ_scratch = wp.empty((d.nworld, m.nflexintcell, 27), dtype=wp.vec3)
+    vel_corot_scratch = wp.empty((d.nworld, m.nflexintcell, 27), dtype=wp.vec3)
+    wp.launch(
+      _flex_passive_interp,
+      dim=(d.nworld, m.nflexintcell),
+      inputs=[
+        m.nflex,
+        m.body_rootid,
+        m.body_dofnum,
+        m.body_dofadr,
+        m.flex_interp,
+        m.flex_cellnum,
+        m.flex_nodeadr,
+        m.flex_stiffnessadr,
+        m.flex_nodebodyid,
+        m.flex_node,
+        m.flex_node0,
+        m.flex_stiffness,
+        m.flex_damping,
+        m.flex_edgeequality,
+        m.flex_centered,
+        m.flex_cell_map,
+        d.subtree_com,
+        d.cvel,
+        d.flexnode_xpos,
+        dsbl_spring,
+        dsbl_damper,
+      ],
+      outputs=[
+        d.qfrc_spring,
+        d.qfrc_damper,
+        displ_scratch,
+        vel_corot_scratch,
+      ],
+    )
+
   if m.has_fluid:
     _fluid(m, d)
 

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

@@ -228,14 +228,19 @@ def _site_local_to_global(
 def _flex_vertices(
   # Model:
   nflex: int,
+  flex_interp: wp.array[int],
+  flex_cellnum: wp.array[wp.vec3i],
+  flex_nodeadr: wp.array[int],
   flex_vertadr: wp.array[int],
   flex_vertnum: wp.array[int],
   flex_vertbodyid: wp.array[int],
   flex_vert: wp.array[wp.vec3],
+  flex_vert0: wp.array[wp.vec3],
   flex_centered: wp.array[bool],
   # Data in:
   xpos_in: wp.array2d[wp.vec3],
   xmat_in: wp.array2d[wp.mat33],
+  flexnode_xpos_in: wp.array2d[wp.vec3],
   # Data out:
   flexvert_xpos_out: wp.array2d[wp.vec3],
 ):
@@ -246,15 +251,93 @@ def _flex_vertices(
     if locid >= 0 and locid < flex_vertnum[f]:
       break
 
-  bodyid = flex_vertbodyid[vertid]
+  if flex_interp[f] != 0:
+    # Interpolated flex: vertex position = weighted sum of node positions
+    coord = flex_vert0[vertid]
+    cn = flex_cellnum[f]
+    cx = cn[0]
+    cy = cn[1]
+    cz = cn[2]
+
+    # Cell lookup: find containing cell
+    ci = wp.min(int(coord[0] * float(cx)), cx - 1)
+    ci = wp.max(ci, 0)
+    cj = wp.min(int(coord[1] * float(cy)), cy - 1)
+    cj = wp.max(cj, 0)
+    ck = wp.min(int(coord[2] * float(cz)), cz - 1)
+    ck = wp.max(ck, 0)
+
+    # Local parametric coordinates within cell
+    local_x = wp.clamp(coord[0] * float(cx) - float(ci), 0.0, 1.0)
+    local_y = wp.clamp(coord[1] * float(cy) - float(cj), 0.0, 1.0)
+    local_z = wp.clamp(coord[2] * float(cz) - float(ck), 0.0, 1.0)
+    local = wp.vec3(local_x, local_y, local_z)
+
+    # Node grid dimensions
+    ny_g = cy + 1
+    nz_g = cz + 1
+    nstart = flex_nodeadr[f]
+
+    # Accumulate weighted node positions
+    result = wp.vec3(0.0, 0.0, 0.0)
+    for li in range(2):
+      for lj in range(2):
+        for lk in range(2):
+          w = support.eval_basis_trilinear(local, li * 4 + lj * 2 + lk)
+          gi = ci + li
+          gj = cj + lj
+          gk = ck + lk
+          node_idx = gi * ny_g * nz_g + gj * nz_g + gk
+          result += w * flexnode_xpos_in[worldid, nstart + node_idx]
+
+    flexvert_xpos_out[worldid, vertid] = result
+  else:
+    # Non-interpolated flex: vertex position from single body
+    bodyid = flex_vertbodyid[vertid]
+    xpos = xpos_in[worldid, bodyid]
+
+    if flex_centered[f]:
+      flexvert_xpos_out[worldid, vertid] = xpos
+    else:
+      xmat = xmat_in[worldid, bodyid]
+      local_pos = flex_vert[vertid]
+      flexvert_xpos_out[worldid, vertid] = xmat @ local_pos + xpos
+
+
+@wp.kernel
+def _flex_nodes(
+  # Model:
+  nflex: int,
+  flex_nodeadr: wp.array[int],
+  flex_nodenum: wp.array[int],
+  flex_nodebodyid: wp.array[int],
+  flex_node: wp.array[wp.vec3],
+  flex_centered: wp.array[bool],
+  # Data in:
+  xpos_in: wp.array2d[wp.vec3],
+  xmat_in: wp.array2d[wp.mat33],
+  # Data out:
+  flexnode_xpos_out: wp.array2d[wp.vec3],
+):
+  worldid, nodeid = wp.tid()
+
+  for f in range(nflex):
+    locid = nodeid - flex_nodeadr[f]
+    if locid >= 0 and locid < flex_nodenum[f]:
+      break
+
+  bodyid = flex_nodebodyid[nodeid]
   xpos = xpos_in[worldid, bodyid]
 
   if flex_centered[f]:
-    flexvert_xpos_out[worldid, vertid] = xpos
+    flexnode_xpos_out[worldid, nodeid] = xpos
   else:
-    xmat = xmat_in[worldid, bodyid]
-    local_pos = flex_vert[vertid]
-    flexvert_xpos_out[worldid, vertid] = xmat @ local_pos + xpos
+    local_pos = flex_node[nodeid]
+    if local_pos[0] == 0.0 and local_pos[1] == 0.0 and local_pos[2] == 0.0:
+      flexnode_xpos_out[worldid, nodeid] = xpos
+    else:
+      xmat = xmat_in[worldid, bodyid]
+      flexnode_xpos_out[worldid, nodeid] = xmat @ local_pos + xpos
 
 
 @wp.kernel
@@ -302,6 +385,11 @@ def _flex_edges(
   b1 = flex_vertbodyid[vbase0]
   b2 = flex_vertbodyid[vbase1]
 
+  # skip Jacobian/velocity for trilinear flex (vertbodyid == -1)
+  if b1 < 0 or b2 < 0:
+    flexedge_velocity_out[worldid, edgeid] = 0.0
+    return
+
   dofnum1 = body_dofnum[b1]
   dofnum2 = body_dofnum[b2]
 
@@ -417,18 +505,40 @@ def kinematics(m: Model, d: Data):
 
 @event_scope
 def flex(m: Model, d: Data):
+  # Compute node positions first (needed for interpolated vertex positions)
+  wp.launch(
+    _flex_nodes,
+    dim=(d.nworld, m.nflexnode),
+    inputs=[
+      m.nflex,
+      m.flex_nodeadr,
+      m.flex_nodenum,
+      m.flex_nodebodyid,
+      m.flex_node,
+      m.flex_centered,
+      d.xpos,
+      d.xmat,
+    ],
+    outputs=[d.flexnode_xpos],
+  )
+
   wp.launch(
     _flex_vertices,
     dim=(d.nworld, m.nflexvert),
     inputs=[
       m.nflex,
+      m.flex_interp,
+      m.flex_cellnum,
+      m.flex_nodeadr,
       m.flex_vertadr,
       m.flex_vertnum,
       m.flex_vertbodyid,
       m.flex_vert,
+      m.flex_vert0,
       m.flex_centered,
       d.xpos,
       d.xmat,
+      d.flexnode_xpos,
     ],
     outputs=[d.flexvert_xpos],
   )

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

@@ -87,6 +87,7 @@ def _create_solver_context(m: types.Model, d: types.Data) -> SolverContext:
     jv=wp.empty((nworld, njmax), dtype=float),
     quad=wp.empty((nworld, njmax), dtype=wp.vec3),
     alpha=wp.empty((nworld,), dtype=float),
+    grad_scale=wp.empty((nworld,), dtype=float),
     improvement=wp.empty((nworld,), dtype=float),
     prev_grad=wp.empty((nworld, nv), dtype=float),
     prev_Mgrad=wp.empty((nworld, nv), dtype=float),
@@ -839,6 +840,7 @@ def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType,
     ctx_jv_out: wp.array2d[float],
     ctx_quad_out: wp.array2d[wp.vec3],
     ctx_improvement_out: wp.array[float],
+    ctx_alpha_out: wp.array[float],
   ):
     worldid, tid = wp.tid()
 
@@ -1202,7 +1204,11 @@ def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType,
         swap_hi = swap_hi_hi_next or swap_hi_mid or swap_hi_lo_next
 
         # check for convergence
-        ls_done = (not swap_lo and not swap_hi) or (lo[1] < 0.0 and lo[1] > -gtol) or (hi[1] > 0.0 and hi[1] < gtol)
+        ls_done = (
+          (not swap_lo and not swap_hi)
+          or (lo[0] < 0.0 and lo[1] < 0.0 and lo[1] > -gtol)
+          or (hi[0] < 0.0 and hi[1] > 0.0 and hi[1] < gtol)
+        )
 
         # update alpha if improved
         improved = lo[0] < 0.0 or hi[0] < 0.0
@@ -1229,6 +1235,7 @@ def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType,
 
     if tid == 0:
       ctx_improvement_out[worldid] = improvement
+      ctx_alpha_out[worldid] = alpha
 
   return kernel
 
@@ -1276,7 +1283,7 @@ def _linesearch_iterative(m: types.Model, d: types.Data, ctx: SolverContext, fus
       ctx.quad,
       ctx.done,
     ],
-    outputs=[d.qacc, d.efc.Ma, ctx.Jaref, ctx.jv, ctx.quad, ctx.improvement],
+    outputs=[d.qacc, d.efc.Ma, ctx.Jaref, ctx.jv, ctx.quad, ctx.improvement, ctx.alpha],
     block_dim=m.block_dim.linesearch_iterative,
   )
 
@@ -1380,8 +1387,13 @@ def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext):
 
   # jv = J @ search (when not fused into iterative kernel)
   if not fuse_jv:
-    dofs_per_thread = 20 if m.nv > 50 else 50
-    threads_per_efc = ceil(m.nv / dofs_per_thread)
+    if m.is_sparse:
+      # Sparse J has few nonzeros per row, one thread handles them all.
+      dofs_per_thread = m.nv
+      threads_per_efc = 1
+    else:
+      dofs_per_thread = 20 if m.nv > 50 else 50
+      threads_per_efc = ceil(m.nv / dofs_per_thread)
 
     if threads_per_efc > 1:
       wp.launch(
@@ -1636,6 +1648,7 @@ def _update_constraint_init_qfrc_constraint_sparse(
   efc_J_in: wp.array3d[float],
   efc_force_in: wp.array2d[float],
   # In:
+  changed_count_in: wp.array[int],
   ctx_done_in: wp.array[bool],
   # Data out:
   qfrc_constraint_out: wp.array2d[float],
@@ -1645,10 +1658,15 @@ def _update_constraint_init_qfrc_constraint_sparse(
   if ctx_done_in[worldid]:
     return
 
+  if changed_count_in[worldid] == 0:
+    return
+
   if efcid >= nefc_in[worldid]:
     return
 
   force = efc_force_in[worldid, efcid]
+  if force == 0.0:
+    return
 
   rownnz = efc_J_rownnz_in[worldid, efcid]
   rowadr = efc_J_rowadr_in[worldid, efcid]
@@ -1660,29 +1678,58 @@ def _update_constraint_init_qfrc_constraint_sparse(
 
 
 @wp.kernel
-def _update_constraint_init_qfrc_constraint_dense(
+def _qfrc_constraint_from_grad(
   # Data in:
-  nefc_in: wp.array[int],
-  efc_J_in: wp.array3d[float],
-  efc_force_in: wp.array2d[float],
-  njmax_in: int,
+  qfrc_smooth_in: wp.array2d[float],
+  efc_Ma_in: wp.array2d[float],
   # In:
-  ctx_done_in: wp.array[bool],
+  ctx_grad_in: wp.array2d[float],
+  ctx_grad_scale_in: wp.array[float],
   # Data out:
   qfrc_constraint_out: wp.array2d[float],
 ):
   worldid, dofid = wp.tid()
 
-  if ctx_done_in[worldid]:
-    return
+  grad = ctx_grad_scale_in[worldid] * ctx_grad_in[worldid, dofid]
+  qfrc_constraint_out[worldid, dofid] = efc_Ma_in[worldid, dofid] - qfrc_smooth_in[worldid, dofid] - grad
 
-  sum_qfrc = float(0.0)
-  for efcid in range(min(njmax_in, nefc_in[worldid])):
-    efc_J = efc_J_in[worldid, efcid, dofid]
-    force = efc_force_in[worldid, efcid]
-    sum_qfrc += efc_J * force
 
-  qfrc_constraint_out[worldid, dofid] = sum_qfrc
+@cache_kernel
+def _update_constraint_init_qfrc_constraint_dense(stable_fast: bool):
+  STABLE_FAST = stable_fast
+
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # Data in:
+    nefc_in: wp.array[int],
+    efc_J_in: wp.array3d[float],
+    efc_force_in: wp.array2d[float],
+    njmax_in: int,
+    # In:
+    changed_count_in: wp.array[int],
+    ctx_done_in: wp.array[bool],
+    # Data out:
+    qfrc_constraint_out: wp.array2d[float],
+  ):
+    worldid, dofid = wp.tid()
+
+    if ctx_done_in[worldid]:
+      return
+
+    # Fast path: stale qfrc_constraint is never read; recovered after the solve.
+    if wp.static(STABLE_FAST):
+      if changed_count_in[worldid] == 0:
+        return
+
+    sum_qfrc = float(0.0)
+    for efcid in range(min(njmax_in, nefc_in[worldid])):
+      efc_J = efc_J_in[worldid, efcid, dofid]
+      force = efc_force_in[worldid, efcid]
+      sum_qfrc += efc_J * force
+
+    qfrc_constraint_out[worldid, dofid] = sum_qfrc
+
+  return kernel
 
 
 @wp.kernel
@@ -1744,48 +1791,51 @@ def _update_gradient_h_incremental_sparse(
   # In:
   changed_ids_in: wp.array2d[int],
   changed_count_in: wp.array[int],
+  slots_per_world: int,
   # Out:
   ctx_h_out: wp.array3d[float],
 ):
-  """Incrementally update upper triangle of H for changed constraints (sparse J)."""
-  worldid, change_idx = wp.tid()
+  """Incrementally update upper triangle of H for changed constraints (sparse J).
+
+  One warp per changed constraint row: the lanes split the row's upper-triangular
+  entries (same sqrt triangular-number decode as _JTDAJ_sparse), replacing the
+  serial nnz^2 loop that dominated this kernel.
+  """
+  worldid, slot, lane = wp.tid()
 
   n_changes = changed_count_in[worldid]
-  if change_idx >= n_changes:
-    return
+  for change_idx in range(slot, n_changes, slots_per_world):
+    efcid = changed_ids_in[worldid, change_idx]
+    D = efc_D_in[worldid, efcid]
+    sign = float(0.0)
+    if efc_state_in[worldid, efcid] == types.ConstraintState.QUADRATIC.value:
+      sign = D
+    else:
+      sign = -D
 
-  efcid = changed_ids_in[worldid, change_idx]
-  D = efc_D_in[worldid, efcid]
-  sign = float(0.0)
-  if efc_state_in[worldid, efcid] == types.ConstraintState.QUADRATIC.value:
-    sign = D
-  else:
-    sign = -D
+    rownnz = efc_J_rownnz_in[worldid, efcid]
+    rowadr = efc_J_rowadr_in[worldid, efcid]
+    n_entries = rownnz * (rownnz + 1) // 2
 
-  rownnz = efc_J_rownnz_in[worldid, efcid]
-  rowadr = efc_J_rowadr_in[worldid, efcid]
-
-  for ii in range(rownnz):
-    sparseidi = rowadr + ii
-    Ji = efc_J_in[worldid, 0, sparseidi]
-    if Ji == 0.0:
-      continue
-    colindi = efc_J_colind_in[worldid, 0, sparseidi]
-    for jj in range(ii + 1):
-      sparseidj = rowadr + jj
-      Jj = efc_J_in[worldid, 0, sparseidj]
-      if Jj == 0.0:
-        continue
-      colindj = efc_J_colind_in[worldid, 0, sparseidj]
+    for entry in range(lane, n_entries, wp.static(_JTDAJ_THREADS_PER_GROUP)):
+      ii = int((wp.sqrt(float(8 * entry + 1)) - 1.0) * 0.5)
+      jj = entry - ii * (ii + 1) // 2
+      Ji = efc_J_in[worldid, 0, rowadr + ii]
+      Jj = efc_J_in[worldid, 0, rowadr + jj]
       h = sign * Ji * Jj
-      # Ensure upper triangle: smaller index first.
-      if colindi <= colindj:
-        wp.atomic_add(ctx_h_out[worldid, colindi], colindj, h)
-      else:
-        wp.atomic_add(ctx_h_out[worldid, colindj], colindi, h)
+      if h != 0.0:
+        colindi = efc_J_colind_in[worldid, 0, rowadr + ii]
+        colindj = efc_J_colind_in[worldid, 0, rowadr + jj]
+        wp.atomic_add(ctx_h_out[worldid, wp.min(colindi, colindj)], wp.max(colindi, colindj), h)
 
 
-def _update_constraint(m: types.Model, d: types.Data, ctx: SolverContext | InverseContext, track_changes: bool = False):
+def _update_constraint(
+  m: types.Model,
+  d: types.Data,
+  ctx: SolverContext | InverseContext,
+  track_changes: bool = False,
+  stable_fast: bool = False,
+):
   """Update constraint arrays after each solve iteration."""
   efc_inputs = [
     m.opt.impratio_invsqrt,
@@ -1811,59 +1861,97 @@ def _update_constraint(m: types.Model, d: types.Data, ctx: SolverContext | Inver
     outputs=[d.efc.force, d.efc.state, ctx.changed_efc_ids, ctx.changed_efc_count],
   )
 
-  # qfrc_constraint = efc_J.T @ efc_force
+  # qfrc_constraint = efc_J.T @ efc_force. Fast-path worlds with no state flips
+  # skip the rebuild; the public value is recovered after the solve.
+  changed = ctx.changed_efc_count if stable_fast else d.nefc
   if m.is_sparse:
     d.qfrc_constraint.zero_()
     wp.launch(
       _update_constraint_init_qfrc_constraint_sparse,
       dim=(d.nworld, d.njmax),
-      inputs=[d.nefc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, d.efc.force, ctx.done],
+      inputs=[d.nefc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, d.efc.force, changed, ctx.done],
       outputs=[d.qfrc_constraint],
     )
   else:
     wp.launch(
-      _update_constraint_init_qfrc_constraint_dense,
+      _update_constraint_init_qfrc_constraint_dense(stable_fast),
       dim=(d.nworld, m.nv),
-      inputs=[d.nefc, d.efc.J, d.efc.force, d.njmax, ctx.done],
+      inputs=[d.nefc, d.efc.J, d.efc.force, d.njmax, changed, ctx.done],
       outputs=[d.qfrc_constraint],
     )
 
 
-@wp.kernel
-def _update_gradient_zero_grad_dot(
-  # In:
-  ctx_done_in: wp.array[bool],
-  # Out:
-  ctx_grad_dot_out: wp.array[float],
-):
-  worldid = wp.tid()
+@cache_kernel
+def _update_gradient_zero_grad_dot(stable_fast: bool):
+  STABLE_FAST = stable_fast
 
-  if ctx_done_in[worldid]:
-    return
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # In:
+    changed_count_in: wp.array[int],
+    ctx_alpha_in: wp.array[float],
+    ctx_done_in: wp.array[bool],
+    # Out:
+    ctx_grad_dot_out: wp.array[float],
+    ctx_grad_scale_out: wp.array[float],
+  ):
+    worldid = wp.tid()
 
-  ctx_grad_dot_out[worldid] = 0.0
+    if ctx_done_in[worldid]:
+      return
+
+    # Fast path: grad stays stale at its last rebuilt value g. The true
+    # gradient is grad_scale * g, and a linesearch step t along the (equally
+    # stale) search direction changes it to (grad_scale - t) * g.
+    if wp.static(STABLE_FAST):
+      if changed_count_in[worldid] == 0:
+        sigma = ctx_grad_scale_out[worldid]
+        new_sigma = sigma - ctx_alpha_in[worldid]
+        ratio = float(0.0)
+        if sigma != 0.0:
+          ratio = new_sigma / sigma
+        ctx_grad_dot_out[worldid] *= ratio * ratio
+        ctx_grad_scale_out[worldid] = new_sigma
+        return
+
+    ctx_grad_dot_out[worldid] = 0.0
+    ctx_grad_scale_out[worldid] = 1.0
+
+  return kernel
 
 
-@wp.kernel
-def _update_gradient_grad(
-  # Data in:
-  qfrc_smooth_in: wp.array2d[float],
-  qfrc_constraint_in: wp.array2d[float],
-  efc_Ma_in: wp.array2d[float],
-  # In:
-  ctx_done_in: wp.array[bool],
-  # Out:
-  ctx_grad_out: wp.array2d[float],
-  ctx_grad_dot_out: wp.array[float],
-):
-  worldid, dofid = wp.tid()
+@cache_kernel
+def _update_gradient_grad(stable_fast: bool):
+  STABLE_FAST = stable_fast
 
-  if ctx_done_in[worldid]:
-    return
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # Data in:
+    qfrc_smooth_in: wp.array2d[float],
+    qfrc_constraint_in: wp.array2d[float],
+    efc_Ma_in: wp.array2d[float],
+    # In:
+    changed_count_in: wp.array[int],
+    ctx_done_in: wp.array[bool],
+    # Out:
+    ctx_grad_out: wp.array2d[float],
+    ctx_grad_dot_out: wp.array[float],
+  ):
+    worldid, dofid = wp.tid()
 
-  grad = efc_Ma_in[worldid, dofid] - qfrc_smooth_in[worldid, dofid] - qfrc_constraint_in[worldid, dofid]
-  ctx_grad_out[worldid, dofid] = grad
-  wp.atomic_add(ctx_grad_dot_out, worldid, grad * grad)
+    if ctx_done_in[worldid]:
+      return
+
+    # Fast path: grad stays stale (see _update_gradient_zero_grad_dot).
+    if wp.static(STABLE_FAST):
+      if changed_count_in[worldid] == 0:
+        return
+
+    grad = efc_Ma_in[worldid, dofid] - qfrc_smooth_in[worldid, dofid] - qfrc_constraint_in[worldid, dofid]
+    ctx_grad_out[worldid, dofid] = grad
+    wp.atomic_add(ctx_grad_dot_out, worldid, grad * grad)
+
+  return kernel
 
 
 @wp.kernel
@@ -2585,12 +2673,15 @@ def _update_gradient_JTCJ_dense(
 
 
 @cache_kernel
-def _update_gradient_cholesky(tile_size: int):
+def _update_gradient_cholesky(tile_size: int, skip_noflip: bool = False):
+  SKIP_NOFLIP = skip_noflip
+
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # In:
     ctx_grad_in: wp.array2d[float],
     h_in: wp.array3d[float],
+    changed_count_in: wp.array[int],
     ctx_done_in: wp.array[bool],
     # Out:
     ctx_Mgrad_out: wp.array2d[float],
@@ -2601,6 +2692,11 @@ def _update_gradient_cholesky(tile_size: int):
     if ctx_done_in[worldid]:
       return
 
+    # Fast path: skip the solve (see the blocked skip_unchanged variant).
+    if wp.static(SKIP_NOFLIP):
+      if changed_count_in[worldid] == 0:
+        return
+
     mat_tile = wp.tile_load(h_in[worldid], shape=(TILE_SIZE, TILE_SIZE))
     wp.tile_cholesky_inplace(mat_tile, fill_mode="upper")
     input_tile = wp.tile_load(ctx_grad_in[worldid], shape=TILE_SIZE)
@@ -2642,8 +2738,9 @@ def _update_gradient_cholesky_blocked(tile_size: int, matrix_size: int, check_sk
 
 
 @cache_kernel
-def _update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size: int):
+def _update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size: int, skip_noflip: bool = False):
   """Blocked Cholesky that skips factorization when no constraints changed."""
+  SKIP_NOFLIP = skip_noflip
 
   @wp.kernel(module="unique", enable_backward=False, module_options={"enable_mathdx_gemm": False})
   def kernel(
@@ -2662,14 +2759,24 @@ def _update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size
     if ctx_done_in[worldid]:
       return
 
-    if changed_count_in[worldid] > 0:
+    # Fast path: skip the solve; Mgrad stays stale on the unchanged ray, and
+    # the linesearch is invariant to the scale of its direction.
+    if wp.static(SKIP_NOFLIP):
+      if changed_count_in[worldid] == 0:
+        return
+
       wp.static(create_blocked_cholesky_factorize_solve_func(TILE_SIZE, matrix_size))(
         ctx_h_in[worldid], ctx_grad_in[worldid], matrix_size, ctx_hfactor[worldid], ctx_Mgrad_out[worldid]
       )
     else:
-      wp.static(create_blocked_cholesky_solve_func(TILE_SIZE, matrix_size))(
-        ctx_hfactor[worldid], ctx_grad_in[worldid], matrix_size, ctx_Mgrad_out[worldid]
-      )
+      if changed_count_in[worldid] > 0:
+        wp.static(create_blocked_cholesky_factorize_solve_func(TILE_SIZE, matrix_size))(
+          ctx_h_in[worldid], ctx_grad_in[worldid], matrix_size, ctx_hfactor[worldid], ctx_Mgrad_out[worldid]
+        )
+      else:
+        wp.static(create_blocked_cholesky_solve_func(TILE_SIZE, matrix_size))(
+          ctx_hfactor[worldid], ctx_grad_in[worldid], matrix_size, ctx_Mgrad_out[worldid]
+        )
 
   return kernel
 
@@ -2685,7 +2792,9 @@ def _padding_h(nv: int, ctx_done_in: wp.array[bool], ctx_h_out: wp.array3d[float
   ctx_h_out[worldid, dofid, dofid] = 1.0
 
 
-def _cholesky_factorize_solve(m: types.Model, d: types.Data, ctx: SolverContext, skip_unchanged: bool = False):
+def _cholesky_factorize_solve(
+  m: types.Model, d: types.Data, ctx: SolverContext, skip_unchanged: bool = False, skip_noflip: bool = False
+):
   """Cholesky factorize ctx.h and solve for Mgrad.
 
   If skip_unchanged is True (blocked path only), worlds where no constraints
@@ -2693,9 +2802,9 @@ def _cholesky_factorize_solve(m: types.Model, d: types.Data, ctx: SolverContext,
   """
   if m.nv <= _BLOCK_CHOLESKY_DIM:
     wp.launch_tiled(
-      _update_gradient_cholesky(m.nv),
+      _update_gradient_cholesky(m.nv, skip_noflip),
       dim=d.nworld,
-      inputs=[ctx.grad, ctx.h, ctx.done],
+      inputs=[ctx.grad, ctx.h, ctx.changed_efc_count if skip_noflip else d.nefc, ctx.done],
       outputs=[ctx.Mgrad],
       block_dim=m.block_dim.update_gradient_cholesky,
     )
@@ -2709,7 +2818,7 @@ def _cholesky_factorize_solve(m: types.Model, d: types.Data, ctx: SolverContext,
 
     if skip_unchanged:
       wp.launch_tiled(
-        _update_gradient_cholesky_blocked_skip_unchanged(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad),
+        _update_gradient_cholesky_blocked_skip_unchanged(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad, skip_noflip),
         dim=d.nworld,
         inputs=[ctx.done, ctx.grad.reshape(shape=(d.nworld, ctx.grad.shape[1], 1)), ctx.h, ctx.changed_efc_count, ctx.hfactor],
         outputs=[ctx.Mgrad.reshape(shape=(d.nworld, ctx.Mgrad.shape[1], 1))],
@@ -2806,11 +2915,16 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext, compact:
       block_dim=m.block_dim.update_gradient_grad,
     )
   else:
-    wp.launch(_update_gradient_zero_grad_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.grad_dot])
     wp.launch(
-      _update_gradient_grad,
+      _update_gradient_zero_grad_dot(False),
+      dim=d.nworld,
+      inputs=[d.nefc, ctx.alpha, ctx.done],
+      outputs=[ctx.grad_dot, ctx.grad_scale],
+    )
+    wp.launch(
+      _update_gradient_grad(False),
       dim=(d.nworld, m.nv),
-      inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, ctx.done],
+      inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, d.nefc, ctx.done],
       outputs=[ctx.grad, ctx.grad_dot],
     )
 
@@ -3002,26 +3116,33 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext, compact:
     raise ValueError(f"Unknown solver type: {m.opt.solver}")
 
 
-def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverContext):
+def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverContext, stable_fast: bool = False):
   """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(_update_gradient_zero_grad_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.grad_dot])
+  changed = ctx.changed_efc_count if stable_fast else d.nefc
+  wp.launch(
+    _update_gradient_zero_grad_dot(stable_fast),
+    dim=d.nworld,
+    inputs=[changed, ctx.alpha, ctx.done],
+    outputs=[ctx.grad_dot, ctx.grad_scale],
+  )
 
   wp.launch(
-    _update_gradient_grad,
+    _update_gradient_grad(stable_fast),
     dim=(d.nworld, m.nv),
-    inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, ctx.done],
+    inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, changed, ctx.done],
     outputs=[ctx.grad, ctx.grad_dot],
   )
 
   # Update upper triangle of H with delta from changed constraints.
   if m.is_sparse:
+    slots = _jtdaj_groups_per_world(d.nworld, ctx.changed_efc_ids.shape[1])
     wp.launch(
       _update_gradient_h_incremental_sparse,
-      dim=(d.nworld, ctx.changed_efc_ids.shape[1]),
+      dim=(d.nworld, slots, _JTDAJ_THREADS_PER_GROUP),
       inputs=[
         d.efc.J_rownnz,
         d.efc.J_rowadr,
@@ -3031,6 +3152,7 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte
         d.efc.state,
         ctx.changed_efc_ids,
         ctx.changed_efc_count,
+        slots,
       ],
       outputs=[ctx.h],
     )
@@ -3049,7 +3171,7 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte
       outputs=[ctx.h],
     )
 
-  _cholesky_factorize_solve(m, d, ctx, skip_unchanged=True)
+  _cholesky_factorize_solve(m, d, ctx, skip_unchanged=True, skip_noflip=stable_fast)
 
 
 @wp.kernel
@@ -3128,46 +3250,70 @@ def _solve_beta_accumulate(
   wp.atomic_add(ctx_beta_den_out, worldid, den)
 
 
-@wp.kernel
-def _solve_zero_search_dot(
-  # In:
-  ctx_done_in: wp.array[bool],
-  # Out:
-  ctx_search_dot_out: wp.array[float],
-):
-  worldid = wp.tid()
+@cache_kernel
+def _solve_zero_search_dot(stable_fast: bool):
+  STABLE_FAST = stable_fast
 
-  if ctx_done_in[worldid]:
-    return
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # In:
+    changed_count_in: wp.array[int],
+    ctx_done_in: wp.array[bool],
+    # Out:
+    ctx_search_dot_out: wp.array[float],
+  ):
+    worldid = wp.tid()
 
-  ctx_search_dot_out[worldid] = 0.0
+    if ctx_done_in[worldid]:
+      return
+
+    # Fast path: search stays on the same ray; keep search_dot consistent with it.
+    if wp.static(STABLE_FAST):
+      if changed_count_in[worldid] == 0:
+        return
+
+    ctx_search_dot_out[worldid] = 0.0
+
+  return kernel
 
 
-@wp.kernel
-def _solve_search_update(
-  # Model:
-  opt_solver: int,
-  # In:
-  ctx_Mgrad_in: wp.array2d[float],
-  ctx_search_in: wp.array2d[float],
-  ctx_beta_in: wp.array[float],
-  ctx_done_in: wp.array[bool],
-  # Out:
-  ctx_search_out: wp.array2d[float],
-  ctx_search_dot_out: wp.array[float],
-):
-  worldid, dofid = wp.tid()
+@cache_kernel
+def _solve_search_update(stable_fast: bool):
+  STABLE_FAST = stable_fast
 
-  if ctx_done_in[worldid]:
-    return
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # Model:
+    opt_solver: int,
+    # In:
+    changed_count_in: wp.array[int],
+    ctx_Mgrad_in: wp.array2d[float],
+    ctx_search_in: wp.array2d[float],
+    ctx_beta_in: wp.array[float],
+    ctx_done_in: wp.array[bool],
+    # Out:
+    ctx_search_out: wp.array2d[float],
+    ctx_search_dot_out: wp.array[float],
+  ):
+    worldid, dofid = wp.tid()
 
-  search = -1.0 * ctx_Mgrad_in[worldid, dofid]
+    if ctx_done_in[worldid]:
+      return
 
-  if opt_solver == types.SolverType.CG:
-    search += ctx_beta_in[worldid] * ctx_search_in[worldid, dofid]
+    # Fast path: search stays on the stale ray; the linesearch absorbs its scale.
+    if wp.static(STABLE_FAST):
+      if changed_count_in[worldid] == 0:
+        return
 
-  ctx_search_out[worldid, dofid] = search
-  wp.atomic_add(ctx_search_dot_out, worldid, search * search)
+    search = -1.0 * ctx_Mgrad_in[worldid, dofid]
+
+    if opt_solver == types.SolverType.CG:
+      search += ctx_beta_in[worldid] * ctx_search_in[worldid, dofid]
+
+    ctx_search_out[worldid, dofid] = search
+    wp.atomic_add(ctx_search_dot_out, worldid, search * search)
+
+  return kernel
 
 
 @wp.kernel
@@ -3291,6 +3437,16 @@ def _solve_done(
     wp.atomic_add(nsolving_out, 0, -1)
 
 
+def _use_incremental(m: types.Model) -> bool:
+  """Whether constraint state changes are tracked for incremental H updates."""
+  return m.opt.solver == types.SolverType.NEWTON and m.opt.cone != types.ConeType.ELLIPTIC
+
+
+def _stable_fast(m: types.Model, compact: bool) -> bool:
+  """Stable-state fast path: needs state-change tracking; compact scatters qfrc itself."""
+  return _use_incremental(m) and not compact
+
+
 @event_scope
 def _solver_iteration(
   m: types.Model,
@@ -3305,16 +3461,21 @@ 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 = m.opt.solver == types.SolverType.NEWTON and m.opt.cone != types.ConeType.ELLIPTIC
+  incremental = _use_incremental(m)
+  # Stable-state fast path: worlds with no state flips this iteration were
+  # exactly quadratic over the step, so grad/Mgrad/search only changed by a
+  # scalar along the same ray. Skip their qfrc/grad/solve/search updates and
+  # track the scalar in ctx.grad_scale.
+  stable_fast = _stable_fast(m, compact)
 
   if incremental:
     # Must complete before _update_constraint_efc which atomically increments.
     ctx.changed_efc_count.zero_()
 
-  _update_constraint(m, d, ctx, track_changes=incremental)
+  _update_constraint(m, d, ctx, track_changes=incremental, stable_fast=stable_fast)
 
   if incremental:
-    _update_gradient_incremental(m, d, ctx)
+    _update_gradient_incremental(m, d, ctx, stable_fast)
   else:
     _update_gradient(m, d, ctx, compact=compact)
 
@@ -3362,12 +3523,13 @@ def _solver_iteration(
     )
 
   else:
-    wp.launch(_solve_zero_search_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.search_dot])
+    changed = ctx.changed_efc_count if stable_fast else d.nefc
+    wp.launch(_solve_zero_search_dot(stable_fast), dim=d.nworld, inputs=[changed, ctx.done], outputs=[ctx.search_dot])
 
     wp.launch(
-      _solve_search_update,
+      _solve_search_update(stable_fast),
       dim=(d.nworld, m.nv),
-      inputs=[m.opt.solver, ctx.Mgrad, ctx.search, ctx.beta, ctx.done],
+      inputs=[m.opt.solver, changed, ctx.Mgrad, ctx.search, ctx.beta, ctx.done],
       outputs=[ctx.search, ctx.search_dot],
     )
 
@@ -3400,12 +3562,16 @@ def init_context(m: types.Model, d: types.Data, ctx: SolverContext | InverseCont
   # if we are only using 1 thread, it makes sense to do more dofs as we can also skip the
   # init kernel. For more than 1 thread, dofs_per_thread is lower for better load balancing.
 
-  if m.nv > 50:
+  if m.is_sparse:
+    # Sparse J has few nonzeros per row, one thread handles them all.
+    dofs_per_thread = m.nv
+    threads_per_efc = 1
+  elif m.nv > 50:
     dofs_per_thread = 20
+    threads_per_efc = ceil(m.nv / dofs_per_thread)
   else:
     dofs_per_thread = 50
-
-  threads_per_efc = ceil(m.nv / dofs_per_thread)
+    threads_per_efc = ceil(m.nv / dofs_per_thread)
   # we need to clear the jaref array if we're doing atomic adds.
   if threads_per_efc > 1:
     ctx.Jaref.zero_()
@@ -3490,6 +3656,16 @@ def _solve(m: types.Model, d: types.Data, ctx: SolverContext, compact: bool = Fa
     for _ in range(m.opt.iterations):
       _solver_iteration(m, d, ctx, nsolving, compact=compact)
 
+  # Recover the public qfrc_constraint: the fast path leaves it stale, and the
+  # per-iteration zeroing wiped it for worlds that converged early.
+  if _stable_fast(m, compact):
+    wp.launch(
+      _qfrc_constraint_from_grad,
+      dim=(d.nworld, m.nv),
+      inputs=[d.qfrc_smooth, d.efc.Ma, ctx.grad, ctx.grad_scale],
+      outputs=[d.qfrc_constraint],
+    )
+
 
 # Active-DOF compaction solve (nvmax < nv).
 #

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

@@ -64,6 +64,115 @@ def next_act(
   return act
 
 
+@wp.func
+def mat33_to_quat_polar(F: wp.mat33) -> wp.quat:
+  cell_quat = wp.quat(0.0, 0.0, 0.0, 1.0)
+  for _iter in range(10):
+    rot = wp.quat_to_matrix(cell_quat)
+    rot_t = wp.transpose(rot)
+    col1_rot = rot_t[0]
+    col2_rot = rot_t[1]
+    col3_rot = rot_t[2]
+    F_t = wp.transpose(F)
+    col1_mat = F_t[0]
+    col2_mat = F_t[1]
+    col3_mat = F_t[2]
+
+    omega = wp.cross(col1_rot, col1_mat) + wp.cross(col2_rot, col2_mat) + wp.cross(col3_rot, col3_mat)
+    denom = wp.abs(wp.dot(col1_rot, col1_mat) + wp.dot(col2_rot, col2_mat) + wp.dot(col3_rot, col3_mat)) + 1.0e-10
+    omega = omega / denom
+
+    w = wp.length(omega)
+    if w < 1.0e-6:
+      break
+
+    axis = omega / w
+    half_w = 0.5 * w
+    qrot = wp.quat(
+      axis[0] * wp.sin(half_w),
+      axis[1] * wp.sin(half_w),
+      axis[2] * wp.sin(half_w),
+      wp.cos(half_w),
+    )
+    cell_quat = wp.normalize(qrot * cell_quat)
+  return cell_quat
+
+
+@wp.func
+def compute_interp_cell_quat(
+  # Data in:
+  flexnode_xpos_in: wp.array2d[wp.vec3],
+  # In:
+  order: int,
+  ci: int,
+  cj: int,
+  ck: int,
+  cy: int,
+  cz: int,
+  ny_g: int,
+  nz_g: int,
+  nstart: int,
+  worldid: int,
+) -> wp.quat:
+  """Computes corotational cell quaternion from deformation gradient at cell center."""
+  npc = (order + 1) * (order + 1) * (order + 1)
+  F = wp.mat33(0.0)
+  idx = int(0)
+  for li in range(order + 1):
+    for lj in range(order + 1):
+      for lk in range(order + 1):
+        if idx < npc:
+          gi = ci * order + li
+          gj = cj * order + lj
+          gk = ck * order + lk
+          gidx = gi * ny_g * nz_g + gj * nz_g + gk
+
+          node_pos = flexnode_xpos_in[worldid, nstart + gidx]
+
+          if order == 1:
+            dphi_x = float(-1) if li == 0 else float(1)
+            dphi_y = float(-1) if lj == 0 else float(1)
+            dphi_z = float(-1) if lk == 0 else float(1)
+            phi_x = float(0.5)
+            phi_y = float(0.5)
+            phi_z = float(0.5)
+          else:
+            if li == 0:
+              dphi_x = -1.0
+            elif li == 1:
+              dphi_x = 0.0
+            else:
+              dphi_x = 1.0
+            if lj == 0:
+              dphi_y = -1.0
+            elif lj == 1:
+              dphi_y = 0.0
+            else:
+              dphi_y = 1.0
+            if lk == 0:
+              dphi_z = -1.0
+            elif lk == 1:
+              dphi_z = 0.0
+            else:
+              dphi_z = 1.0
+            phi_x = 0.5 if li == 0 or li == 2 else 1.0
+            phi_y = 0.5 if lj == 0 or lj == 2 else 1.0
+            phi_z = 0.5 if lk == 0 or lk == 2 else 1.0
+
+          grad_x = dphi_x * phi_y * phi_z
+          grad_y = phi_x * dphi_y * phi_z
+          grad_z = phi_x * phi_y * dphi_z
+
+          for r in range(3):
+            F[r, 0] += node_pos[r] * grad_x
+            F[r, 1] += node_pos[r] * grad_y
+            F[r, 2] += node_pos[r] * grad_z
+
+          idx += 1
+
+  return mat33_to_quat_polar(F)
+
+
 @cache_kernel
 def mul_m_kernel(check_skip: bool):
   @wp.kernel(module="unique")
@@ -893,3 +1002,84 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
       d.userdata,
     ],
   )
+
+
+@wp.func
+def _phi(s: float, i: int) -> float:
+  """1D trilinear basis function (order=1 only).
+
+  phi(s, 0) = 1 - s
+  phi(s, 1) = s
+  """
+  if i == 0:
+    return 1.0 - s
+  return s
+
+
+@wp.func
+def eval_basis_trilinear(local: wp.vec3, node_idx: int) -> float:
+  """Evaluate trilinear basis function for node_idx at local coords [0,1]^3.
+
+  For order=1 (trilinear), node_idx encodes (i,j,k) via bits:
+    k = node_idx & 1, j = (node_idx >> 1) & 1, i = (node_idx >> 2) & 1
+  """
+  k = node_idx & 1
+  j = (node_idx >> 1) & 1
+  i = (node_idx >> 2) & 1
+  return _phi(local[0], i) * _phi(local[1], j) * _phi(local[2], k)
+
+
+@wp.func
+def select_top4_weights(
+  # In:
+  W_mat: wp.mat33,
+  b_mat: wp.mat33,
+) -> tuple[wp.vec4i, wp.vec4]:
+  """Selects top 4 weights and their corresponding body IDs from 8 voxel corners."""
+  selected_b = wp.vec4i(-1, -1, -1, -1)
+  selected_W = wp.vec4(0.0, 0.0, 0.0, 0.0)
+
+  local_W = W_mat
+  for p in range(4):
+    max_w = -1.0
+    max_b = -1
+    max_r = -1
+    max_c = -1
+    for r in range(3):
+      for c in range(3):
+        idx = 3 * r + c
+        if idx < 8:
+          w = local_W[r, c]
+          if w > max_w:
+            max_w = w
+            max_b = int(b_mat[r, c])
+            max_r = r
+            max_c = c
+    # Record top choice for this pass and mark it as visited
+    if max_r >= 0:
+      local_W[max_r, max_c] = -1.0
+
+      if p == 0:
+        selected_b = wp.vec4i(max_b, -1, -1, -1)
+        selected_W = wp.vec4(max_w, 0.0, 0.0, 0.0)
+      elif p == 1:
+        selected_b = wp.vec4i(selected_b[0], max_b, -1, -1)
+        selected_W = wp.vec4(selected_W[0], max_w, 0.0, 0.0)
+      elif p == 2:
+        selected_b = wp.vec4i(selected_b[0], selected_b[1], max_b, -1)
+        selected_W = wp.vec4(selected_W[0], selected_W[1], max_w, 0.0)
+      else:
+        selected_b = wp.vec4i(selected_b[0], selected_b[1], selected_b[2], max_b)
+        selected_W = wp.vec4(selected_W[0], selected_W[1], selected_W[2], max_w)
+
+  # Normalize selected weights
+  sum_W = selected_W[0] + selected_W[1] + selected_W[2] + selected_W[3]
+  if sum_W > 1.0e-5:
+    selected_W = wp.vec4(
+      selected_W[0] / sum_W,
+      selected_W[1] / sum_W,
+      selected_W[2] / sum_W,
+      selected_W[3] / sum_W,
+    )
+
+  return selected_b, selected_W

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

@@ -0,0 +1,351 @@
+# Copyright 2025 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.
+# ==============================================================================
+
+"""Tests for flex."""
+
+import mujoco
+import numpy as np
+from absl.testing import absltest
+from absl.testing import parameterized
+
+import mujoco.mjx.third_party.mujoco_warp as mjw
+from mujoco.mjx.third_party.mujoco_warp import test_data
+
+_TRILINEAR_STRAIN_XML = """
+<mujoco>
+  <option gravity="0 0 -9.81">
+    <flag contact="disable"/>
+  </option>
+  <worldbody>
+    <flexcomp type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
+              pos="0 0 0.5" name="cube" dim="3" mass="1" radius="0.005"
+              dof="trilinear">
+      <edge equality="strain"/>
+      <contact selfcollide="none"/>
+    </flexcomp>
+  </worldbody>
+</mujoco>
+"""
+
+# tolerance for difference between MuJoCo and MJWarp, mostly due to float32
+_TOLERANCE = 5e-4
+
+
+class TrilinearFlexTest(parameterized.TestCase):
+  def test_flexstrain_constraint_at_rest(self):
+    """Test FLEXSTRAIN constraint count and residuals match MuJoCo at rest."""
+    mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
+    mjd = mujoco.MjData(mjm)
+    mujoco.mj_forward(mjm, mjd)
+
+    m = mjw.put_model(mjm)
+    d = mjw.put_data(mjm, mjd)
+
+    mjw.fwd_position(m, d)
+    mjw.make_constraint(m, d)
+
+    # constraint counts should match
+    ne_warp = d.ne.numpy()[0]
+    self.assertEqual(ne_warp, mjd.ne, f"ne mismatch: warp={ne_warp}, mj={mjd.ne}")
+
+    nefc_warp = d.nefc.numpy()[0]
+    self.assertEqual(nefc_warp, mjd.nefc, f"nefc mismatch: warp={nefc_warp}, mj={mjd.nefc}")
+
+    # residuals should match
+    efc_pos = d.efc.pos.numpy()[0, :nefc_warp]
+    efc_pos_mj = mjd.efc_pos[: mjd.nefc]
+    np.testing.assert_allclose(efc_pos, efc_pos_mj, atol=1e-5, err_msg="FLEXSTRAIN residuals should match MuJoCo at rest")
+
+  def test_flexstrain_constraint_perturbed(self):
+    """Test FLEXSTRAIN residuals and Jacobians match MuJoCo under perturbation."""
+    mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
+    mjd = mujoco.MjData(mjm)
+
+    # perturb first node
+    mjd.qpos[0] += 0.01
+    mjd.qpos[3] += 0.005
+    mujoco.mj_forward(mjm, mjd)
+
+    m = mjw.put_model(mjm)
+    d = mjw.put_data(mjm, mjd)
+
+    mjw.fwd_position(m, d)
+    mjw.make_constraint(m, d)
+
+    nefc = d.nefc.numpy()[0]
+    nv = mjm.nv
+    self.assertEqual(nefc, mjd.nefc)
+
+    # residuals
+    efc_pos_warp = d.efc.pos.numpy()[0, :nefc]
+    efc_pos_mj = mjd.efc_pos[: mjd.nefc]
+    np.testing.assert_allclose(
+      efc_pos_warp, efc_pos_mj, atol=_TOLERANCE, err_msg="FLEXSTRAIN residuals don't match MuJoCo under perturbation"
+    )
+
+    # Jacobians
+    if mujoco.mj_isSparse(mjm):
+      mj_efc_J = np.zeros((mjd.nefc, nv))
+      mujoco.mju_sparse2dense(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, nv))
+
+    if m.is_sparse:
+      warp_efc_J = np.zeros((nefc, nv))
+      mujoco.mju_sparse2dense(
+        warp_efc_J,
+        d.efc.J.numpy()[0, 0],
+        d.efc.J_rownnz.numpy()[0, :nefc],
+        d.efc.J_rowadr.numpy()[0, :nefc],
+        d.efc.J_colind.numpy()[0, 0],
+      )
+    else:
+      warp_efc_J = d.efc.J.numpy()[0, :nefc, :nv]
+
+    np.testing.assert_allclose(warp_efc_J, mj_efc_J, atol=0.01, err_msg="FLEXSTRAIN Jacobians don't match MuJoCo")
+
+  def test_flexstrain_constraint_rotated(self):
+    """Test FLEXSTRAIN residuals and Jacobians match MuJoCo under large rotation perturbation."""
+    mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
+    mjd = mujoco.MjData(mjm)
+
+    # Apply a rotation perturbation: rotate all node positions around Y axis by 30 degrees
+    # (0.5235 radians)
+    theta = 0.5235
+    cos_t = np.cos(theta)
+    sin_t = np.sin(theta)
+    for i in range(0, mjm.nq, 3):
+      x = mjd.qpos[i]
+      z = mjd.qpos[i + 2]
+      mjd.qpos[i] = x * cos_t - z * sin_t
+      mjd.qpos[i + 2] = x * sin_t + z * cos_t
+
+    mujoco.mj_forward(mjm, mjd)
+
+    m = mjw.put_model(mjm)
+    d = mjw.put_data(mjm, mjd)
+
+    mjw.fwd_position(m, d)
+    mjw.make_constraint(m, d)
+
+    nefc = d.nefc.numpy()[0]
+    nv = mjm.nv
+    self.assertEqual(nefc, mjd.nefc)
+
+    # residuals
+    efc_pos_warp = d.efc.pos.numpy()[0, :nefc]
+    efc_pos_mj = mjd.efc_pos[: mjd.nefc]
+    np.testing.assert_allclose(
+      efc_pos_warp, efc_pos_mj, atol=_TOLERANCE, err_msg="FLEXSTRAIN residuals don't match MuJoCo under rotation"
+    )
+
+    # Jacobians
+    if mujoco.mj_isSparse(mjm):
+      mj_efc_J = np.zeros((mjd.nefc, nv))
+      mujoco.mju_sparse2dense(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, nv))
+
+    if m.is_sparse:
+      warp_efc_J = np.zeros((nefc, nv))
+      mujoco.mju_sparse2dense(
+        warp_efc_J,
+        d.efc.J.numpy()[0, 0],
+        d.efc.J_rownnz.numpy()[0, :nefc],
+        d.efc.J_rowadr.numpy()[0, :nefc],
+        d.efc.J_colind.numpy()[0, 0],
+      )
+    else:
+      warp_efc_J = d.efc.J.numpy()[0, :nefc, :nv]
+
+    np.testing.assert_allclose(
+      warp_efc_J, mj_efc_J, atol=0.01, err_msg="FLEXSTRAIN Jacobians don't match MuJoCo under rotation"
+    )
+
+  def test_flexstrain_rotational_invariance(self):
+    """Test that FLEXSTRAIN residuals are invariant under rigid translation."""
+    mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
+    mjd = mujoco.MjData(mjm)
+    mujoco.mj_forward(mjm, mjd)
+
+    # Get reference residuals
+    m = mjw.put_model(mjm)
+    d = mjw.put_data(mjm, mjd)
+    mjw.fwd_position(m, d)
+    mjw.make_constraint(m, d)
+    nefc = d.nefc.numpy()[0]
+    efc_pos_rest = d.efc.pos.numpy()[0, :nefc].copy()
+
+    # Apply uniform translation to all nodes (rigid motion)
+    mjd2 = mujoco.MjData(mjm)
+    # All flex nodes have 3 DOFs (slide joints), shift all x by 0.1
+    for i in range(0, mjm.nq, 3):
+      mjd2.qpos[i] += 0.1  # shift x
+    mujoco.mj_forward(mjm, mjd2)
+
+    d2 = mjw.put_data(mjm, mjd2)
+    mjw.fwd_position(m, d2)
+    mjw.make_constraint(m, d2)
+    nefc2 = d2.nefc.numpy()[0]
+    efc_pos_shifted = d2.efc.pos.numpy()[0, :nefc2]
+
+    # Residuals should remain near zero for rigid translation
+    np.testing.assert_allclose(
+      efc_pos_shifted, efc_pos_rest, atol=1e-4, err_msg="FLEXSTRAIN residuals should be invariant under rigid translation"
+    )
+
+  def test_trilinear_gravity_parity(self):
+    """Test that trilinear simulation matches MuJoCo after multiple steps."""
+    mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
+
+    # MuJoCo reference
+    mjd = mujoco.MjData(mjm)
+    for _ in range(10):
+      mujoco.mj_step(mjm, mjd)
+
+    # Warp
+    mjd_warp = mujoco.MjData(mjm)
+    m = mjw.put_model(mjm)
+    d = mjw.put_data(mjm, mjd_warp)
+    z0 = d.qpos.numpy()[0, 2]
+    for _ in range(10):
+      mjw.step(m, d)
+
+    qpos_warp = d.qpos.numpy()[0]
+    qpos_mj = mjd.qpos
+
+    # The cube should have fallen
+    self.assertLess(qpos_warp[2], z0, "Cube should fall under gravity")
+
+    # Allow larger tolerance for accumulated integration error
+    np.testing.assert_allclose(qpos_warp, qpos_mj, atol=0.01, err_msg="Trilinear qpos diverges from MuJoCo after 10 steps")
+
+  def test_trilinear_node_positions(self):
+    """Test that flexnode_xpos are computed correctly from body kinematics."""
+    mjm = mujoco.MjModel.from_xml_string(_TRILINEAR_STRAIN_XML)
+    mjd = mujoco.MjData(mjm)
+    mujoco.mj_forward(mjm, mjd)
+
+    m = mjw.put_model(mjm)
+    d = mjw.put_data(mjm, mjd)
+    mjw.fwd_position(m, d)
+
+    # Compute expected node positions: xpos_n = body_xpos + body_xmat @ flex_node
+    nflexnode = mjm.nflexnode
+    warp_xpos = d.flexnode_xpos.numpy()[0, :nflexnode]
+    nodeadr = mjm.flex_nodeadr[0]
+    nodenum = mjm.flex_nodenum[0]
+    for n in range(nodenum):
+      bodyid = mjm.flex_nodebodyid[nodeadr + n]
+      body_xpos = mjd.xpos[bodyid]
+      body_xmat = mjd.xmat[bodyid].reshape(3, 3)
+      node_local = mjm.flex_node[nodeadr + n]
+      expected = body_xpos + body_xmat @ node_local
+      np.testing.assert_allclose(warp_xpos[n], expected, atol=1e-5, err_msg=f"flexnode_xpos mismatch for node {n}")
+
+  def test_trilinear_passive_forces_parity(self):
+    """Test passive forces (elasticity) match MuJoCo for trilinear flex."""
+    xml = """
+    <mujoco>
+      <worldbody>
+        <flexcomp type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
+                  pos="0 0 0.5" name="cube" dim="3" mass="1" radius="0.005"
+                  dof="trilinear">
+          <elasticity young="1e4" poisson="0.1" damping="0.01"/>
+          <contact selfcollide="none"/>
+        </flexcomp>
+      </worldbody>
+    </mujoco>
+    """
+    mjm = mujoco.MjModel.from_xml_string(xml)
+    mjd = mujoco.MjData(mjm)
+
+    # perturb first node to generate non-zero elasticity forces
+    mjd.qpos[0] += 0.01
+    mjd.qpos[3] += 0.005
+    mujoco.mj_forward(mjm, mjd)
+
+    m = mjw.put_model(mjm)
+    d = mjw.put_data(mjm, mjd)
+
+    mjw.fwd_position(m, d)
+    mjw.passive(m, d)
+
+    qfrc_passive_warp = d.qfrc_passive.numpy()[0]
+    qfrc_passive_mj = mjd.qfrc_passive
+
+    # Verify they match
+    np.testing.assert_allclose(
+      qfrc_passive_warp, qfrc_passive_mj, atol=_TOLERANCE, err_msg="qfrc_passive mismatch for trilinear flex with elasticity"
+    )
+
+  @parameterized.parameters("strain", "true")
+  def test_trilinear_equality_types(self, equality):
+    """Test trilinear with different equality types."""
+    xml = f"""
+    <mujoco>
+      <option gravity="0 0 -9.81">
+        <flag contact="disable"/>
+      </option>
+      <worldbody>
+        <flexcomp type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
+                  pos="0 0 0.5" name="cube" dim="3" mass="1" radius="0.005"
+                  dof="trilinear">
+          <edge equality="{equality}"/>
+          <contact selfcollide="none"/>
+        </flexcomp>
+      </worldbody>
+    </mujoco>
+    """
+    mjm, mjd, m, d = test_data.fixture(xml=xml)
+
+    # Should not crash
+    mjw.forward(m, d)
+
+    # Constraint count should match
+    self.assertEqual(d.nefc.numpy()[0], mjd.nefc)
+
+  def test_trilinear_contact_qfrc_constraint(self):
+    """Test qfrc_constraint parity for trilinear flex with ground contacts."""
+    xml = """
+    <mujoco>
+      <option gravity="0 0 -9.81"/>
+      <worldbody>
+        <geom type="plane" size="1 1 0.1"/>
+        <flexcomp type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
+                  pos="0 0 0.05" name="cube" dim="3" mass="1" radius="0.02"
+                  dof="trilinear">
+          <edge equality="strain"/>
+          <contact selfcollide="none"/>
+        </flexcomp>
+      </worldbody>
+    </mujoco>
+    """
+    mjm, mjd, m, d = test_data.fixture(xml=xml)
+    mjw.forward(m, d)
+
+    # Verify contacts are generated
+    nacon = d.nacon.numpy()[0]
+    self.assertGreater(nacon, 0, "Expected contacts between flex and plane")
+    self.assertEqual(nacon, mjd.ncon)
+
+    # Verify qfrc_constraint parity
+    qfrc_warp = d.qfrc_constraint.numpy()[0]
+    qfrc_mj = mjd.qfrc_constraint
+    np.testing.assert_allclose(qfrc_warp, qfrc_mj, atol=1e-4, err_msg="qfrc_constraint mismatch for trilinear flex contacts")
+
+
+if __name__ == "__main__":
+  absltest.main()

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

@@ -77,7 +77,7 @@ class BlockDim:
   # collision_driver
   segmented_sort: int = 128
   # collision_convex
-  convex_ccd: int = 256
+  convex_ccd: int = 64
   # forward
   actuator_velocity: int = 32
   # ray
@@ -664,6 +664,7 @@ class EqType(enum.IntEnum):
     WELD: fix relative position and orientation of two bodies
     TENDON: couple the lengths of two tendons with cubic
     FLEX: couple the edge lengths of a flex
+    FLEXSTRAIN: strain constraint for interpolated flex
   """
 
   CONNECT = mujoco.mjtEq.mjEQ_CONNECT
@@ -671,6 +672,7 @@ class EqType(enum.IntEnum):
   JOINT = mujoco.mjtEq.mjEQ_JOINT
   TENDON = mujoco.mjtEq.mjEQ_TENDON
   FLEX = mujoco.mjtEq.mjEQ_FLEX
+  FLEXSTRAIN = mujoco.mjtEq.mjEQ_FLEXSTRAIN
   # unsupported: DISTANCE
 
 
@@ -974,6 +976,7 @@ class Model:
     ncam: number of cameras
     nlight: number of lights
     nflex: number of flexes
+    nflexnode: number of nodes in all flexes
     nflexvert: number of vertices in all flexes
     nflexedge: number of edges in all flexes
     nflexelem: number of elements in all flexes
@@ -1140,6 +1143,10 @@ class Model:
     flex_internal: internal collision enabled                (nflex,)
     flex_selfcollide: self-collision mode                    (nflex,)
     flex_dim: 1: lines, 2: triangles, 3: tetrahedra          (nflex,)
+    flex_interp: interpolation order (0: vertex, 1+: nodes)  (nflex,)
+    flex_cellnum: cell count per dimension                   (nflex, 3)
+    flex_nodeadr: first node address                         (nflex,)
+    flex_nodenum: number of nodes                            (nflex,)
     flex_vertadr: first vertex address                       (nflex,)
     flex_vertnum: number of vertices                         (nflex,)
     flex_edgeadr: first edge address                         (nflex,)
@@ -1147,13 +1154,14 @@ class Model:
     flex_elemadr: first element address                      (nflex,)
     flex_elemnum: number of elements                         (nflex,)
     flex_elemdataadr: first element vertex id address        (nflex,)
-    flex_stiffnessadr: stiffness matrix address               (nflex,)
+    flex_stiffnessadr: stiffness matrix address              (nflex,)
     flex_elemedgeadr: first element edge id address          (nflex,)
     flex_bendingadr: first bending data address              (nflex,)
     flex_shellnum: number of shells                          (nflex,)
     flex_shelldataadr: first shell data address              (nflex,)
     flex_evpairadr: first element-vertex pair address        (nflex,)
     flex_evpairnum: number of element-vertex pairs           (nflex,)
+    flex_nodebodyid: node body ids                           (nflexnode,)
     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)
@@ -1162,12 +1170,17 @@ class Model:
     flex_shell: shell fragment vertex ids (dim per frag)     (nflexshelldata,)
     flex_evpair: element-vertex pair indices                 (nflexevpair, 2)
     flex_vert: vertex local positions                        (nflexvert, 3)
+    flex_vert0: reference vertex positions in qpos0          (nflexvert, 3)
+    flex_node: node local positions                          (nflexnode, 3)
+    flex_node0: reference node positions in qpos0            (nflexnode, 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          (nflexstiffness,)
     flex_bending: bending stiffness                          (nflexbending,)
     flex_damping: Rayleigh's damping coefficient             (nflex,)
+
+    flex_edgeequality: edge equality type (0:none,1:edge,2:vert,3:strain) (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,)
@@ -1312,6 +1325,7 @@ class Model:
     nrangefinder: number of rangefinder sensors
     nmaxcondim: maximum condim across geoms, pairs, and flexes
     nmaxpyramid: maximum number of pyramid directions
+    nflexintcell: total interp cells (non-strain) for passive forces
     nmaxpolygon: maximum number of verts per polygon
     nmaxmeshdeg: maximum number of polygons per vert
     is_sparse: constraint Jacobian/Hessian layout (sparse vs dense). Does not affect M, whose
@@ -1327,6 +1341,8 @@ class Model:
     has_fluid: True if wind, density, or viscosity are non-zero at put_model time
     has_sdf_geom: whether the model contains SDF geoms
     has_flex_selfcollide: whether any flex has self-collision enabled
+    has_ellipsoid_geom: whether the model contains ellipsoid geoms
+    has_3d_flex: whether the model contains 3D flexes
     max_flex_dim: maximum flex dimension in the model
     block_dim: block dim options
     body_tree: list of body ids by tree level
@@ -1356,6 +1372,7 @@ class Model:
     eq_jnt_adr: eq_* addresses of type `JOINT`
     eq_ten_adr: eq_* addresses of type `TENDON`
     eq_flex_adr: eq * addresses of type `FLEX
+    eq_flexstrain_adr: eq_* addresses of type `FLEXSTRAIN`
     tendon_jnt_adr: joint tendon address
     tendon_site_pair_adr: site pair tendon address
     tendon_geom_adr: geom tendon address
@@ -1413,13 +1430,18 @@ class Model:
     M_mulm_col: sparse matmul column indices
     M_mulm_madr: sparse matmul matrix addresses
     flexelem_geom_pair_filtered: conaffinity-filtered element vs geom pairs (*, 2)
-    flexshell_geom_pair_filtered: conaffinity-filtered shell vs geom pairs  (*, 2)
     flexvert_geom_pair_filtered: conaffinity-filtered vertex vs geom pairs  (*, 2)
     flex_elemflexid: maps each element index directly to its flexid         (nflexelem,)
     flex_shellflexid: maps each shell index directly to its flexid          (nflexshelldata,)
     flex_evpairflexid: maps each element-vertex pair directly to its flexid (nflexevpair,)
     flex_vertflexid: maps each vertex index directly to its flexid          (nflexvert,)
     flex_shelladr: maps each flex to its start shell index                  (nflex,)
+    flex_cell_map: precomputed flex cell mapping (nflexintcell,)
+    flexstrain_J_rownnz: number of nonzeros in flex strain Jacobian row (neq_flexstrain,)
+    flexstrain_J_rowadr: row start address in colind array (neq_flexstrain,)
+    flexstrain_J_colind: column indices in sparse flex strain Jacobian (nJfs,)
+    neq_flexstrain: number of flex strain equality constraints
+    nJfs: number of non-zeros in sparse flex strain Jacobian
   """
 
   nq: int
@@ -1438,6 +1460,7 @@ class Model:
   ncam: int
   nlight: int
   nflex: int
+  nflexnode: int
   nflexvert: int
   nflexedge: int
   nflexelem: int
@@ -1604,6 +1627,10 @@ class Model:
   flex_internal: array("nflex", int)
   flex_selfcollide: array("nflex", int)
   flex_dim: array("nflex", int)
+  flex_interp: array("nflex", int)
+  flex_cellnum: array("nflex", wp.vec3i)
+  flex_nodeadr: array("nflex", int)
+  flex_nodenum: array("nflex", int)
   flex_vertadr: array("nflex", int)
   flex_vertnum: array("nflex", int)
   flex_edgeadr: array("nflex", int)
@@ -1618,6 +1645,7 @@ class Model:
   flex_shelldataadr: array("nflex", int)
   flex_evpairadr: array("nflex", int)
   flex_evpairnum: array("nflex", int)
+  flex_nodebodyid: array("nflexnode", int)
   flex_vertbodyid: array("nflexvert", int)
   flex_edge: array("nflexedge", wp.vec2i)
   flex_edgeflap: array("nflexedge", wp.vec2i)
@@ -1626,12 +1654,16 @@ class Model:
   flex_shell: array("nflexshelldata", int)
   flex_evpair: array("nflexevpair", wp.vec2i)
   flex_vert: array("nflexvert", wp.vec3)
+  flex_vert0: array("nflexvert", wp.vec3)
+  flex_node: array("nflexnode", wp.vec3)
+  flex_node0: array("nflexnode", wp.vec3)
   flexedge_length0: array("nflexedge", float)
   flexedge_invweight0: array("nflexedge", float)
   flex_radius: array("nflex", float)
   flex_stiffness: array("nflexstiffness", float)
   flex_bending: array("nflexbending", float)
   flex_damping: array("nflex", float)
+  flex_edgeequality: array("nflex", int)
   flex_centered: array("nflex", bool)
   flexedge_J_rownnz: array("nflexedge", int)
   flexedge_J_rowadr: array("nflexedge", int)
@@ -1774,6 +1806,7 @@ class Model:
   nrangefinder: int
   nmaxcondim: int
   nmaxpyramid: int
+  nflexintcell: int
   nmaxpolygon: int
   nmaxmeshdeg: int
   is_sparse: bool
@@ -1788,6 +1821,8 @@ class Model:
   has_fluid: bool
   has_sdf_geom: bool
   has_flex_selfcollide: bool
+  has_ellipsoid_geom: bool
+  has_3d_flex: bool
   max_flex_dim: int
   block_dim: BlockDim
   body_tree: tuple[wp.array[int], ...]
@@ -1813,6 +1848,7 @@ class Model:
   eq_jnt_adr: wp.array[int]
   eq_ten_adr: wp.array[int]
   eq_flex_adr: wp.array[int]
+  eq_flexstrain_adr: wp.array[int]
   tendon_jnt_adr: wp.array[int]
   tendon_site_pair_adr: wp.array[int]
   tendon_geom_adr: wp.array[int]
@@ -1868,13 +1904,18 @@ class Model:
   M_mulm_col: wp.array[int]  # column index to gather from
   M_mulm_madr: wp.array[int]  # matrix address to read
   flexelem_geom_pair_filtered: wp.array[wp.vec2i]
-  flexshell_geom_pair_filtered: wp.array[wp.vec2i]
   flexvert_geom_pair_filtered: wp.array[wp.vec2i]
   flex_elemflexid: array("nflexelem", int)
   flex_shellflexid: array("nflexshelldata", int)
   flex_evpairflexid: array("nflexevpair", int)
   flex_vertflexid: array("nflexvert", int)
   flex_shelladr: array("nflex", int)
+  flex_cell_map: array("nflexintcell", wp.vec4i)
+  flexstrain_J_rownnz: array("neq_flexstrain", int)
+  flexstrain_J_rowadr: array("neq_flexstrain", int)
+  flexstrain_J_colind: array("nJfs", int)
+  neq_flexstrain: int
+  nJfs: int
 
 
 class ContactType(enum.IntFlag):
@@ -2136,6 +2177,7 @@ class Data:
     ncollision: collision count from broadphase                 (1,)
     flex_aabb_min: dynamic flex object bounding box min         (nworld, nflex, 3)
     flex_aabb_max: dynamic flex object bounding box max         (nworld, nflex, 3)
+    flexnode_xpos: cartesian flex node positions                (nworld, nflexnode, 3)
     overflow: overflow bitmask (OverflowType)                   (nworld,)
   """
 
@@ -2279,6 +2321,7 @@ class Data:
   ncollision: array(1, int)
   flex_aabb_min: array("nworld", "nflex", wp.vec3)
   flex_aabb_max: array("nworld", "nflex", wp.vec3)
+  flexnode_xpos: array("nworld", "nflexnode", wp.vec3)
   overflow: array("nworld", int)
 
 
@@ -2308,6 +2351,7 @@ class SolverContext:
   jv: wp.array2d[float]
   quad: wp.array2d[wp.vec3]
   alpha: wp.array[float]
+  grad_scale: wp.array[float]
   improvement: wp.array[float]
   prev_grad: wp.array2d[float]
   prev_Mgrad: wp.array2d[float]

mjx/mujoco/mjx/warp/collision_driver.py

@@ -72,10 +72,6 @@ def _collision_shim(
     flex_priority: wp.array[int],
     flex_radius: wp.array[float],
     flex_selfcollide: wp.array[int],
-    flex_shell: wp.array[int],
-    flex_shelladr: wp.array[int],
-    flex_shelldataadr: wp.array[int],
-    flex_shellflexid: wp.array[int],
     flex_solimp: wp.array[mjwp_types.vec5],
     flex_solmix: wp.array[float],
     flex_solref: wp.array[wp.vec2],
@@ -84,7 +80,6 @@ def _collision_shim(
     flex_vertflexid: wp.array[int],
     flex_vertnum: wp.array[int],
     flexelem_geom_pair_filtered: wp.array[wp.vec2i],
-    flexshell_geom_pair_filtered: wp.array[wp.vec2i],
     flexvert_geom_pair_filtered: wp.array[wp.vec2i],
     geom_aabb: wp.array3d[wp.vec3],
     geom_bodyid: wp.array[int],
@@ -104,6 +99,8 @@ def _collision_shim(
     geom_solmix: wp.array2d[float],
     geom_solref: wp.array2d[wp.vec2],
     geom_type: wp.array[int],
+    has_3d_flex: bool,
+    has_ellipsoid_geom: bool,
     has_flex_selfcollide: bool,
     has_sdf_geom: bool,
     hfield_adr: wp.array[int],
@@ -221,10 +218,6 @@ def _collision_shim(
   _m.flex_priority = flex_priority
   _m.flex_radius = flex_radius
   _m.flex_selfcollide = flex_selfcollide
-  _m.flex_shell = flex_shell
-  _m.flex_shelladr = flex_shelladr
-  _m.flex_shelldataadr = flex_shelldataadr
-  _m.flex_shellflexid = flex_shellflexid
   _m.flex_solimp = flex_solimp
   _m.flex_solmix = flex_solmix
   _m.flex_solref = flex_solref
@@ -233,7 +226,6 @@ def _collision_shim(
   _m.flex_vertflexid = flex_vertflexid
   _m.flex_vertnum = flex_vertnum
   _m.flexelem_geom_pair_filtered = flexelem_geom_pair_filtered
-  _m.flexshell_geom_pair_filtered = flexshell_geom_pair_filtered
   _m.flexvert_geom_pair_filtered = flexvert_geom_pair_filtered
   _m.geom_aabb = geom_aabb
   _m.geom_bodyid = geom_bodyid
@@ -253,6 +245,8 @@ def _collision_shim(
   _m.geom_solmix = geom_solmix
   _m.geom_solref = geom_solref
   _m.geom_type = geom_type
+  _m.has_3d_flex = has_3d_flex
+  _m.has_ellipsoid_geom = has_ellipsoid_geom
   _m.has_flex_selfcollide = has_flex_selfcollide
   _m.has_sdf_geom = has_sdf_geom
   _m.hfield_adr = hfield_adr
@@ -447,10 +441,6 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_priority,
       m._impl.flex_radius,
       m._impl.flex_selfcollide,
-      m._impl.flex_shell,
-      m._impl.flex_shelladr,
-      m._impl.flex_shelldataadr,
-      m._impl.flex_shellflexid,
       m._impl.flex_solimp,
       m._impl.flex_solmix,
       m._impl.flex_solref,
@@ -459,7 +449,6 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_vertflexid,
       m.flex_vertnum,
       m._impl.flexelem_geom_pair_filtered,
-      m._impl.flexshell_geom_pair_filtered,
       m._impl.flexvert_geom_pair_filtered,
       m.geom_aabb,
       m.geom_bodyid,
@@ -479,6 +468,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m.geom_solmix,
       m.geom_solref,
       m.geom_type,
+      m._impl.has_3d_flex,
+      m._impl.has_ellipsoid_geom,
       m._impl.has_flex_selfcollide,
       m._impl.has_sdf_geom,
       m.hfield_adr,

mjx/mujoco/mjx/warp/forward.py

@@ -140,6 +140,7 @@ def _forward_shim(
     eq_connect_adr: wp.array[int],
     eq_data: wp.array2d[mjwp_types.vec11],
     eq_flex_adr: wp.array[int],
+    eq_flexstrain_adr: wp.array[int],
     eq_jnt_adr: wp.array[int],
     eq_obj1id: wp.array[int],
     eq_obj2id: wp.array[int],
@@ -151,6 +152,8 @@ def _forward_shim(
     eq_wld_adr: wp.array[int],
     flex_bending: wp.array[float],
     flex_bendingadr: wp.array[int],
+    flex_cell_map: wp.array[wp.vec4i],
+    flex_cellnum: wp.array[wp.vec3i],
     flex_centered: wp.array[bool],
     flex_conaffinity: wp.array[int],
     flex_condim: wp.array[int],
@@ -159,6 +162,7 @@ def _forward_shim(
     flex_dim: wp.array[int],
     flex_edge: wp.array[wp.vec2i],
     flex_edgeadr: wp.array[int],
+    flex_edgeequality: wp.array[int],
     flex_edgeflap: wp.array[wp.vec2i],
     flex_edgenum: wp.array[int],
     flex_elem: wp.array[int],
@@ -175,20 +179,25 @@ def _forward_shim(
     flex_friction: wp.array[wp.vec3],
     flex_gap: wp.array[float],
     flex_internal: wp.array[int],
+    flex_interp: wp.array[int],
     flex_margin: wp.array[float],
+    flex_node: wp.array[wp.vec3],
+    flex_node0: wp.array[wp.vec3],
+    flex_nodeadr: wp.array[int],
+    flex_nodebodyid: wp.array[int],
+    flex_nodenum: wp.array[int],
     flex_priority: wp.array[int],
     flex_radius: wp.array[float],
     flex_selfcollide: wp.array[int],
     flex_shell: wp.array[int],
-    flex_shelladr: wp.array[int],
     flex_shelldataadr: wp.array[int],
-    flex_shellflexid: wp.array[int],
     flex_solimp: wp.array[mjwp_types.vec5],
     flex_solmix: wp.array[float],
     flex_solref: wp.array[wp.vec2],
     flex_stiffness: wp.array[float],
     flex_stiffnessadr: wp.array[int],
     flex_vert: wp.array[wp.vec3],
+    flex_vert0: wp.array[wp.vec3],
     flex_vertadr: wp.array[int],
     flex_vertbodyid: wp.array[int],
     flex_vertflexid: wp.array[int],
@@ -199,7 +208,9 @@ def _forward_shim(
     flexedge_invweight0: wp.array[float],
     flexedge_length0: wp.array[float],
     flexelem_geom_pair_filtered: wp.array[wp.vec2i],
-    flexshell_geom_pair_filtered: wp.array[wp.vec2i],
+    flexstrain_J_colind: wp.array[int],
+    flexstrain_J_rowadr: wp.array[int],
+    flexstrain_J_rownnz: wp.array[int],
     flexvert_geom_pair_filtered: wp.array[wp.vec2i],
     geom_aabb: wp.array3d[wp.vec3],
     geom_bodyid: wp.array[int],
@@ -225,6 +236,8 @@ def _forward_shim(
     geom_solmix: wp.array2d[float],
     geom_solref: wp.array2d[wp.vec2],
     geom_type: wp.array[int],
+    has_3d_flex: bool,
+    has_ellipsoid_geom: bool,
     has_flex_selfcollide: bool,
     has_fluid: bool,
     has_sdf_geom: bool,
@@ -296,6 +309,8 @@ def _forward_shim(
     nflexedge: int,
     nflexelem: int,
     nflexevpair: int,
+    nflexintcell: int,
+    nflexnode: int,
     nflexvert: int,
     ngeom: int,
     nhistory: int,
@@ -502,6 +517,7 @@ def _forward_shim(
     flexedge_J: wp.array2d[float],
     flexedge_length: wp.array2d[float],
     flexedge_velocity: wp.array2d[float],
+    flexnode_xpos: wp.array2d[wp.vec3],
     flexvert_xpos: wp.array2d[wp.vec3],
     geom_xmat: wp.array2d[wp.mat33],
     geom_xpos: wp.array2d[wp.vec3],
@@ -712,6 +728,7 @@ def _forward_shim(
   _m.eq_connect_adr = eq_connect_adr
   _m.eq_data = eq_data
   _m.eq_flex_adr = eq_flex_adr
+  _m.eq_flexstrain_adr = eq_flexstrain_adr
   _m.eq_jnt_adr = eq_jnt_adr
   _m.eq_obj1id = eq_obj1id
   _m.eq_obj2id = eq_obj2id
@@ -723,6 +740,8 @@ def _forward_shim(
   _m.eq_wld_adr = eq_wld_adr
   _m.flex_bending = flex_bending
   _m.flex_bendingadr = flex_bendingadr
+  _m.flex_cell_map = flex_cell_map
+  _m.flex_cellnum = flex_cellnum
   _m.flex_centered = flex_centered
   _m.flex_conaffinity = flex_conaffinity
   _m.flex_condim = flex_condim
@@ -731,6 +750,7 @@ def _forward_shim(
   _m.flex_dim = flex_dim
   _m.flex_edge = flex_edge
   _m.flex_edgeadr = flex_edgeadr
+  _m.flex_edgeequality = flex_edgeequality
   _m.flex_edgeflap = flex_edgeflap
   _m.flex_edgenum = flex_edgenum
   _m.flex_elem = flex_elem
@@ -747,20 +767,25 @@ def _forward_shim(
   _m.flex_friction = flex_friction
   _m.flex_gap = flex_gap
   _m.flex_internal = flex_internal
+  _m.flex_interp = flex_interp
   _m.flex_margin = flex_margin
+  _m.flex_node = flex_node
+  _m.flex_node0 = flex_node0
+  _m.flex_nodeadr = flex_nodeadr
+  _m.flex_nodebodyid = flex_nodebodyid
+  _m.flex_nodenum = flex_nodenum
   _m.flex_priority = flex_priority
   _m.flex_radius = flex_radius
   _m.flex_selfcollide = flex_selfcollide
   _m.flex_shell = flex_shell
-  _m.flex_shelladr = flex_shelladr
   _m.flex_shelldataadr = flex_shelldataadr
-  _m.flex_shellflexid = flex_shellflexid
   _m.flex_solimp = flex_solimp
   _m.flex_solmix = flex_solmix
   _m.flex_solref = flex_solref
   _m.flex_stiffness = flex_stiffness
   _m.flex_stiffnessadr = flex_stiffnessadr
   _m.flex_vert = flex_vert
+  _m.flex_vert0 = flex_vert0
   _m.flex_vertadr = flex_vertadr
   _m.flex_vertbodyid = flex_vertbodyid
   _m.flex_vertflexid = flex_vertflexid
@@ -771,7 +796,9 @@ def _forward_shim(
   _m.flexedge_invweight0 = flexedge_invweight0
   _m.flexedge_length0 = flexedge_length0
   _m.flexelem_geom_pair_filtered = flexelem_geom_pair_filtered
-  _m.flexshell_geom_pair_filtered = flexshell_geom_pair_filtered
+  _m.flexstrain_J_colind = flexstrain_J_colind
+  _m.flexstrain_J_rowadr = flexstrain_J_rowadr
+  _m.flexstrain_J_rownnz = flexstrain_J_rownnz
   _m.flexvert_geom_pair_filtered = flexvert_geom_pair_filtered
   _m.geom_aabb = geom_aabb
   _m.geom_bodyid = geom_bodyid
@@ -797,6 +824,8 @@ def _forward_shim(
   _m.geom_solmix = geom_solmix
   _m.geom_solref = geom_solref
   _m.geom_type = geom_type
+  _m.has_3d_flex = has_3d_flex
+  _m.has_ellipsoid_geom = has_ellipsoid_geom
   _m.has_flex_selfcollide = has_flex_selfcollide
   _m.has_fluid = has_fluid
   _m.has_sdf_geom = has_sdf_geom
@@ -868,6 +897,8 @@ def _forward_shim(
   _m.nflexedge = nflexedge
   _m.nflexelem = nflexelem
   _m.nflexevpair = nflexevpair
+  _m.nflexintcell = nflexintcell
+  _m.nflexnode = nflexnode
   _m.nflexvert = nflexvert
   _m.ngeom = ngeom
   _m.nhistory = nhistory
@@ -1104,6 +1135,7 @@ def _forward_shim(
   _d.flexedge_J = flexedge_J
   _d.flexedge_length = flexedge_length
   _d.flexedge_velocity = flexedge_velocity
+  _d.flexnode_xpos = flexnode_xpos
   _d.flexvert_xpos = flexvert_xpos
   _d.geom_xmat = geom_xmat
   _d.geom_xpos = geom_xpos
@@ -1231,6 +1263,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'flexedge_J': d._impl.flexedge_J.shape,
       'flexedge_length': d._impl.flexedge_length.shape,
       'flexedge_velocity': d._impl.flexedge_velocity.shape,
+      'flexnode_xpos': d._impl.flexnode_xpos.shape,
       'flexvert_xpos': d._impl.flexvert_xpos.shape,
       'geom_xmat': d.geom_xmat.shape,
       'geom_xpos': d.geom_xpos.shape,
@@ -1342,7 +1375,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _forward_shim,
-      num_outputs=145,
+      num_outputs=146,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames=set([
@@ -1383,6 +1416,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'flexedge_J',
           'flexedge_length',
           'flexedge_velocity',
+          'flexnode_xpos',
           'flexvert_xpos',
           'geom_xmat',
           'geom_xpos',
@@ -1765,6 +1799,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.eq_connect_adr,
       m.eq_data,
       m._impl.eq_flex_adr,
+      m._impl.eq_flexstrain_adr,
       m._impl.eq_jnt_adr,
       m.eq_obj1id,
       m.eq_obj2id,
@@ -1776,6 +1811,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.eq_wld_adr,
       m._impl.flex_bending,
       m._impl.flex_bendingadr,
+      m._impl.flex_cell_map,
+      m._impl.flex_cellnum,
       m._impl.flex_centered,
       m._impl.flex_conaffinity,
       m._impl.flex_condim,
@@ -1784,6 +1821,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_dim,
       m._impl.flex_edge,
       m._impl.flex_edgeadr,
+      m._impl.flex_edgeequality,
       m._impl.flex_edgeflap,
       m._impl.flex_edgenum,
       m._impl.flex_elem,
@@ -1800,20 +1838,25 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_friction,
       m._impl.flex_gap,
       m._impl.flex_internal,
+      m.flex_interp,
       m._impl.flex_margin,
+      m._impl.flex_node,
+      m.flex_node0,
+      m.flex_nodeadr,
+      m.flex_nodebodyid,
+      m.flex_nodenum,
       m._impl.flex_priority,
       m._impl.flex_radius,
       m._impl.flex_selfcollide,
       m._impl.flex_shell,
-      m._impl.flex_shelladr,
       m._impl.flex_shelldataadr,
-      m._impl.flex_shellflexid,
       m._impl.flex_solimp,
       m._impl.flex_solmix,
       m._impl.flex_solref,
       m._impl.flex_stiffness,
       m._impl.flex_stiffnessadr,
       m._impl.flex_vert,
+      m.flex_vert0,
       m.flex_vertadr,
       m._impl.flex_vertbodyid,
       m._impl.flex_vertflexid,
@@ -1824,7 +1867,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.flexedge_invweight0,
       m._impl.flexedge_length0,
       m._impl.flexelem_geom_pair_filtered,
-      m._impl.flexshell_geom_pair_filtered,
+      m._impl.flexstrain_J_colind,
+      m._impl.flexstrain_J_rowadr,
+      m._impl.flexstrain_J_rownnz,
       m._impl.flexvert_geom_pair_filtered,
       m.geom_aabb,
       m.geom_bodyid,
@@ -1850,6 +1895,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m.geom_solmix,
       m.geom_solref,
       m.geom_type,
+      m._impl.has_3d_flex,
+      m._impl.has_ellipsoid_geom,
       m._impl.has_flex_selfcollide,
       m._impl.has_fluid,
       m._impl.has_sdf_geom,
@@ -1921,6 +1968,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.nflexedge,
       m._impl.nflexelem,
       m._impl.nflexevpair,
+      m._impl.nflexintcell,
+      m._impl.nflexnode,
       m._impl.nflexvert,
       m.ngeom,
       m.nhistory,
@@ -2126,6 +2175,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d._impl.flexedge_J,
       d._impl.flexedge_length,
       d._impl.flexedge_velocity,
+      d._impl.flexnode_xpos,
       d._impl.flexvert_xpos,
       d.geom_xmat,
       d.geom_xpos,
@@ -2280,114 +2330,115 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       '_impl.flexedge_J': out[34],
       '_impl.flexedge_length': out[35],
       '_impl.flexedge_velocity': out[36],
-      '_impl.flexvert_xpos': out[37],
-      'geom_xmat': out[38],
-      'geom_xpos': out[39],
-      'history': out[40],
-      '_impl.island_dofadr': out[41],
-      '_impl.island_idofadr': out[42],
-      '_impl.island_iefcadr': out[43],
-      '_impl.island_ne': out[44],
-      '_impl.island_nefc': out[45],
-      '_impl.island_nf': out[46],
-      '_impl.island_nv': out[47],
-      '_impl.light_xdir': out[48],
-      '_impl.light_xpos': out[49],
-      '_impl.map_dof2idof': out[50],
-      '_impl.map_efc2iefc': out[51],
-      '_impl.map_idof2dof': out[52],
-      '_impl.map_iefc2efc': out[53],
-      '_impl.moment_colind': out[54],
-      '_impl.moment_rowadr': out[55],
-      '_impl.moment_rownnz': out[56],
-      '_impl.nacon': out[57],
-      '_impl.nbody_awake': out[58],
-      '_impl.ncdof': out[59],
-      '_impl.ncollision': out[60],
-      '_impl.ne': out[61],
-      '_impl.nefc': out[62],
-      '_impl.nf': out[63],
-      '_impl.nidof': out[64],
-      '_impl.nisland': out[65],
-      '_impl.nl': out[66],
-      '_impl.ntree_awake': out[67],
-      '_impl.nv_awake': out[68],
-      '_impl.overflow': out[69],
-      '_impl.qLD': out[70],
-      '_impl.qLDiagInv': out[71],
-      'qacc': out[72],
-      'qacc_smooth': out[73],
-      'qfrc_actuator': out[74],
-      'qfrc_bias': out[75],
-      'qfrc_constraint': out[76],
-      '_impl.qfrc_damper': out[77],
-      'qfrc_fluid': out[78],
-      'qfrc_gravcomp': out[79],
-      'qfrc_passive': out[80],
-      'qfrc_smooth': out[81],
-      '_impl.qfrc_spring': out[82],
-      'qvel': out[83],
-      'sensordata': out[84],
-      'site_xmat': out[85],
-      'site_xpos': out[86],
-      '_impl.solver_niter': out[87],
-      '_impl.subtree_angmom': out[88],
-      'subtree_com': out[89],
-      '_impl.subtree_linvel': out[90],
-      '_impl.ten_J': out[91],
-      'ten_length': out[92],
-      '_impl.ten_velocity': out[93],
-      '_impl.ten_wrapadr': out[94],
-      '_impl.ten_wrapnum': out[95],
-      '_impl.tree_asleep': out[96],
-      '_impl.tree_awake': out[97],
-      '_impl.tree_island': out[98],
-      '_impl.wrap_obj': out[99],
-      '_impl.wrap_xpos': out[100],
-      'xanchor': out[101],
-      'xaxis': out[102],
-      'ximat': out[103],
-      'xipos': out[104],
-      'xmat': out[105],
-      'xpos': out[106],
-      'xquat': out[107],
-      '_impl.contact__dim': out[108],
-      '_impl.contact__dist': out[109],
-      '_impl.contact__efc_address': out[110],
-      '_impl.contact__elem': out[111],
-      '_impl.contact__flex': out[112],
-      '_impl.contact__frame': out[113],
-      '_impl.contact__friction': out[114],
-      '_impl.contact__geom': out[115],
-      '_impl.contact__geomcollisionid': out[116],
-      '_impl.contact__includemargin': out[117],
-      '_impl.contact__pos': out[118],
-      '_impl.contact__solimp': out[119],
-      '_impl.contact__solref': out[120],
-      '_impl.contact__solreffriction': out[121],
-      '_impl.contact__type': out[122],
-      '_impl.contact__vert': out[123],
-      '_impl.contact__worldid': out[124],
-      '_impl.efc__D': out[125],
-      '_impl.efc__J': out[126],
-      '_impl.efc__J_colind': out[127],
-      '_impl.efc__J_rowadr': out[128],
-      '_impl.efc__J_rownnz': out[129],
-      '_impl.efc__Jqvel': out[130],
-      '_impl.efc__Ma': out[131],
-      '_impl.efc__aref': out[132],
-      '_impl.efc__force': out[133],
-      '_impl.efc__frictionloss': out[134],
-      '_impl.efc__id': out[135],
-      '_impl.efc__island': out[136],
-      '_impl.efc__jtdaj_adr': out[137],
-      '_impl.efc__jtdaj_nblock': out[138],
-      '_impl.efc__jtdaj_nrow': out[139],
-      '_impl.efc__margin': out[140],
-      '_impl.efc__pos': out[141],
-      '_impl.efc__state': out[142],
-      '_impl.efc__type': out[143],
-      '_impl.efc__vel': out[144],
+      '_impl.flexnode_xpos': out[37],
+      '_impl.flexvert_xpos': out[38],
+      'geom_xmat': out[39],
+      'geom_xpos': out[40],
+      'history': out[41],
+      '_impl.island_dofadr': out[42],
+      '_impl.island_idofadr': out[43],
+      '_impl.island_iefcadr': out[44],
+      '_impl.island_ne': out[45],
+      '_impl.island_nefc': out[46],
+      '_impl.island_nf': out[47],
+      '_impl.island_nv': out[48],
+      '_impl.light_xdir': out[49],
+      '_impl.light_xpos': out[50],
+      '_impl.map_dof2idof': out[51],
+      '_impl.map_efc2iefc': out[52],
+      '_impl.map_idof2dof': out[53],
+      '_impl.map_iefc2efc': out[54],
+      '_impl.moment_colind': out[55],
+      '_impl.moment_rowadr': out[56],
+      '_impl.moment_rownnz': out[57],
+      '_impl.nacon': out[58],
+      '_impl.nbody_awake': out[59],
+      '_impl.ncdof': out[60],
+      '_impl.ncollision': out[61],
+      '_impl.ne': out[62],
+      '_impl.nefc': out[63],
+      '_impl.nf': out[64],
+      '_impl.nidof': out[65],
+      '_impl.nisland': out[66],
+      '_impl.nl': out[67],
+      '_impl.ntree_awake': out[68],
+      '_impl.nv_awake': out[69],
+      '_impl.overflow': out[70],
+      '_impl.qLD': out[71],
+      '_impl.qLDiagInv': out[72],
+      'qacc': out[73],
+      'qacc_smooth': out[74],
+      'qfrc_actuator': out[75],
+      'qfrc_bias': out[76],
+      'qfrc_constraint': out[77],
+      '_impl.qfrc_damper': out[78],
+      'qfrc_fluid': out[79],
+      'qfrc_gravcomp': out[80],
+      'qfrc_passive': out[81],
+      'qfrc_smooth': out[82],
+      '_impl.qfrc_spring': out[83],
+      'qvel': out[84],
+      'sensordata': out[85],
+      'site_xmat': out[86],
+      'site_xpos': out[87],
+      '_impl.solver_niter': out[88],
+      '_impl.subtree_angmom': out[89],
+      'subtree_com': out[90],
+      '_impl.subtree_linvel': out[91],
+      '_impl.ten_J': out[92],
+      'ten_length': out[93],
+      '_impl.ten_velocity': out[94],
+      '_impl.ten_wrapadr': out[95],
+      '_impl.ten_wrapnum': out[96],
+      '_impl.tree_asleep': out[97],
+      '_impl.tree_awake': out[98],
+      '_impl.tree_island': out[99],
+      '_impl.wrap_obj': out[100],
+      '_impl.wrap_xpos': out[101],
+      'xanchor': out[102],
+      'xaxis': out[103],
+      'ximat': out[104],
+      'xipos': out[105],
+      'xmat': out[106],
+      'xpos': out[107],
+      'xquat': out[108],
+      '_impl.contact__dim': out[109],
+      '_impl.contact__dist': out[110],
+      '_impl.contact__efc_address': out[111],
+      '_impl.contact__elem': out[112],
+      '_impl.contact__flex': out[113],
+      '_impl.contact__frame': out[114],
+      '_impl.contact__friction': out[115],
+      '_impl.contact__geom': out[116],
+      '_impl.contact__geomcollisionid': out[117],
+      '_impl.contact__includemargin': out[118],
+      '_impl.contact__pos': out[119],
+      '_impl.contact__solimp': out[120],
+      '_impl.contact__solref': out[121],
+      '_impl.contact__solreffriction': out[122],
+      '_impl.contact__type': out[123],
+      '_impl.contact__vert': out[124],
+      '_impl.contact__worldid': out[125],
+      '_impl.efc__D': out[126],
+      '_impl.efc__J': out[127],
+      '_impl.efc__J_colind': out[128],
+      '_impl.efc__J_rowadr': out[129],
+      '_impl.efc__J_rownnz': out[130],
+      '_impl.efc__Jqvel': out[131],
+      '_impl.efc__Ma': out[132],
+      '_impl.efc__aref': out[133],
+      '_impl.efc__force': out[134],
+      '_impl.efc__frictionloss': out[135],
+      '_impl.efc__id': out[136],
+      '_impl.efc__island': out[137],
+      '_impl.efc__jtdaj_adr': out[138],
+      '_impl.efc__jtdaj_nblock': out[139],
+      '_impl.efc__jtdaj_nrow': out[140],
+      '_impl.efc__margin': out[141],
+      '_impl.efc__pos': out[142],
+      '_impl.efc__state': out[143],
+      '_impl.efc__type': out[144],
+      '_impl.efc__vel': out[145],
   })
   return d
 
@@ -2506,6 +2557,7 @@ def _step_shim(
     eq_connect_adr: wp.array[int],
     eq_data: wp.array2d[mjwp_types.vec11],
     eq_flex_adr: wp.array[int],
+    eq_flexstrain_adr: wp.array[int],
     eq_jnt_adr: wp.array[int],
     eq_obj1id: wp.array[int],
     eq_obj2id: wp.array[int],
@@ -2517,6 +2569,8 @@ def _step_shim(
     eq_wld_adr: wp.array[int],
     flex_bending: wp.array[float],
     flex_bendingadr: wp.array[int],
+    flex_cell_map: wp.array[wp.vec4i],
+    flex_cellnum: wp.array[wp.vec3i],
     flex_centered: wp.array[bool],
     flex_conaffinity: wp.array[int],
     flex_condim: wp.array[int],
@@ -2525,6 +2579,7 @@ def _step_shim(
     flex_dim: wp.array[int],
     flex_edge: wp.array[wp.vec2i],
     flex_edgeadr: wp.array[int],
+    flex_edgeequality: wp.array[int],
     flex_edgeflap: wp.array[wp.vec2i],
     flex_edgenum: wp.array[int],
     flex_elem: wp.array[int],
@@ -2541,20 +2596,25 @@ def _step_shim(
     flex_friction: wp.array[wp.vec3],
     flex_gap: wp.array[float],
     flex_internal: wp.array[int],
+    flex_interp: wp.array[int],
     flex_margin: wp.array[float],
+    flex_node: wp.array[wp.vec3],
+    flex_node0: wp.array[wp.vec3],
+    flex_nodeadr: wp.array[int],
+    flex_nodebodyid: wp.array[int],
+    flex_nodenum: wp.array[int],
     flex_priority: wp.array[int],
     flex_radius: wp.array[float],
     flex_selfcollide: wp.array[int],
     flex_shell: wp.array[int],
-    flex_shelladr: wp.array[int],
     flex_shelldataadr: wp.array[int],
-    flex_shellflexid: wp.array[int],
     flex_solimp: wp.array[mjwp_types.vec5],
     flex_solmix: wp.array[float],
     flex_solref: wp.array[wp.vec2],
     flex_stiffness: wp.array[float],
     flex_stiffnessadr: wp.array[int],
     flex_vert: wp.array[wp.vec3],
+    flex_vert0: wp.array[wp.vec3],
     flex_vertadr: wp.array[int],
     flex_vertbodyid: wp.array[int],
     flex_vertflexid: wp.array[int],
@@ -2565,7 +2625,9 @@ def _step_shim(
     flexedge_invweight0: wp.array[float],
     flexedge_length0: wp.array[float],
     flexelem_geom_pair_filtered: wp.array[wp.vec2i],
-    flexshell_geom_pair_filtered: wp.array[wp.vec2i],
+    flexstrain_J_colind: wp.array[int],
+    flexstrain_J_rowadr: wp.array[int],
+    flexstrain_J_rownnz: wp.array[int],
     flexvert_geom_pair_filtered: wp.array[wp.vec2i],
     geom_aabb: wp.array3d[wp.vec3],
     geom_bodyid: wp.array[int],
@@ -2591,6 +2653,8 @@ def _step_shim(
     geom_solmix: wp.array2d[float],
     geom_solref: wp.array2d[wp.vec2],
     geom_type: wp.array[int],
+    has_3d_flex: bool,
+    has_ellipsoid_geom: bool,
     has_flex_selfcollide: bool,
     has_fluid: bool,
     has_sdf_geom: bool,
@@ -2665,6 +2729,8 @@ def _step_shim(
     nflexedge: int,
     nflexelem: int,
     nflexevpair: int,
+    nflexintcell: int,
+    nflexnode: int,
     nflexvert: int,
     ngeom: int,
     nhistory: int,
@@ -2875,6 +2941,7 @@ def _step_shim(
     flexedge_J: wp.array2d[float],
     flexedge_length: wp.array2d[float],
     flexedge_velocity: wp.array2d[float],
+    flexnode_xpos: wp.array2d[wp.vec3],
     flexvert_xpos: wp.array2d[wp.vec3],
     geom_xmat: wp.array2d[wp.mat33],
     geom_xpos: wp.array2d[wp.vec3],
@@ -3094,6 +3161,7 @@ def _step_shim(
   _m.eq_connect_adr = eq_connect_adr
   _m.eq_data = eq_data
   _m.eq_flex_adr = eq_flex_adr
+  _m.eq_flexstrain_adr = eq_flexstrain_adr
   _m.eq_jnt_adr = eq_jnt_adr
   _m.eq_obj1id = eq_obj1id
   _m.eq_obj2id = eq_obj2id
@@ -3105,6 +3173,8 @@ def _step_shim(
   _m.eq_wld_adr = eq_wld_adr
   _m.flex_bending = flex_bending
   _m.flex_bendingadr = flex_bendingadr
+  _m.flex_cell_map = flex_cell_map
+  _m.flex_cellnum = flex_cellnum
   _m.flex_centered = flex_centered
   _m.flex_conaffinity = flex_conaffinity
   _m.flex_condim = flex_condim
@@ -3113,6 +3183,7 @@ def _step_shim(
   _m.flex_dim = flex_dim
   _m.flex_edge = flex_edge
   _m.flex_edgeadr = flex_edgeadr
+  _m.flex_edgeequality = flex_edgeequality
   _m.flex_edgeflap = flex_edgeflap
   _m.flex_edgenum = flex_edgenum
   _m.flex_elem = flex_elem
@@ -3129,20 +3200,25 @@ def _step_shim(
   _m.flex_friction = flex_friction
   _m.flex_gap = flex_gap
   _m.flex_internal = flex_internal
+  _m.flex_interp = flex_interp
   _m.flex_margin = flex_margin
+  _m.flex_node = flex_node
+  _m.flex_node0 = flex_node0
+  _m.flex_nodeadr = flex_nodeadr
+  _m.flex_nodebodyid = flex_nodebodyid
+  _m.flex_nodenum = flex_nodenum
   _m.flex_priority = flex_priority
   _m.flex_radius = flex_radius
   _m.flex_selfcollide = flex_selfcollide
   _m.flex_shell = flex_shell
-  _m.flex_shelladr = flex_shelladr
   _m.flex_shelldataadr = flex_shelldataadr
-  _m.flex_shellflexid = flex_shellflexid
   _m.flex_solimp = flex_solimp
   _m.flex_solmix = flex_solmix
   _m.flex_solref = flex_solref
   _m.flex_stiffness = flex_stiffness
   _m.flex_stiffnessadr = flex_stiffnessadr
   _m.flex_vert = flex_vert
+  _m.flex_vert0 = flex_vert0
   _m.flex_vertadr = flex_vertadr
   _m.flex_vertbodyid = flex_vertbodyid
   _m.flex_vertflexid = flex_vertflexid
@@ -3153,7 +3229,9 @@ def _step_shim(
   _m.flexedge_invweight0 = flexedge_invweight0
   _m.flexedge_length0 = flexedge_length0
   _m.flexelem_geom_pair_filtered = flexelem_geom_pair_filtered
-  _m.flexshell_geom_pair_filtered = flexshell_geom_pair_filtered
+  _m.flexstrain_J_colind = flexstrain_J_colind
+  _m.flexstrain_J_rowadr = flexstrain_J_rowadr
+  _m.flexstrain_J_rownnz = flexstrain_J_rownnz
   _m.flexvert_geom_pair_filtered = flexvert_geom_pair_filtered
   _m.geom_aabb = geom_aabb
   _m.geom_bodyid = geom_bodyid
@@ -3179,6 +3257,8 @@ def _step_shim(
   _m.geom_solmix = geom_solmix
   _m.geom_solref = geom_solref
   _m.geom_type = geom_type
+  _m.has_3d_flex = has_3d_flex
+  _m.has_ellipsoid_geom = has_ellipsoid_geom
   _m.has_flex_selfcollide = has_flex_selfcollide
   _m.has_fluid = has_fluid
   _m.has_sdf_geom = has_sdf_geom
@@ -3253,6 +3333,8 @@ def _step_shim(
   _m.nflexedge = nflexedge
   _m.nflexelem = nflexelem
   _m.nflexevpair = nflexevpair
+  _m.nflexintcell = nflexintcell
+  _m.nflexnode = nflexnode
   _m.nflexvert = nflexvert
   _m.ngeom = ngeom
   _m.nhistory = nhistory
@@ -3493,6 +3575,7 @@ def _step_shim(
   _d.flexedge_J = flexedge_J
   _d.flexedge_length = flexedge_length
   _d.flexedge_velocity = flexedge_velocity
+  _d.flexnode_xpos = flexnode_xpos
   _d.flexvert_xpos = flexvert_xpos
   _d.geom_xmat = geom_xmat
   _d.geom_xpos = geom_xpos
@@ -3622,6 +3705,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'flexedge_J': d._impl.flexedge_J.shape,
       'flexedge_length': d._impl.flexedge_length.shape,
       'flexedge_velocity': d._impl.flexedge_velocity.shape,
+      'flexnode_xpos': d._impl.flexnode_xpos.shape,
       'flexvert_xpos': d._impl.flexvert_xpos.shape,
       'geom_xmat': d.geom_xmat.shape,
       'geom_xpos': d.geom_xpos.shape,
@@ -3737,7 +3821,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _step_shim,
-      num_outputs=150,
+      num_outputs=151,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames=set([
@@ -3779,6 +3863,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'flexedge_J',
           'flexedge_length',
           'flexedge_velocity',
+          'flexnode_xpos',
           'flexvert_xpos',
           'geom_xmat',
           'geom_xpos',
@@ -4177,6 +4262,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.eq_connect_adr,
       m.eq_data,
       m._impl.eq_flex_adr,
+      m._impl.eq_flexstrain_adr,
       m._impl.eq_jnt_adr,
       m.eq_obj1id,
       m.eq_obj2id,
@@ -4188,6 +4274,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.eq_wld_adr,
       m._impl.flex_bending,
       m._impl.flex_bendingadr,
+      m._impl.flex_cell_map,
+      m._impl.flex_cellnum,
       m._impl.flex_centered,
       m._impl.flex_conaffinity,
       m._impl.flex_condim,
@@ -4196,6 +4284,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_dim,
       m._impl.flex_edge,
       m._impl.flex_edgeadr,
+      m._impl.flex_edgeequality,
       m._impl.flex_edgeflap,
       m._impl.flex_edgenum,
       m._impl.flex_elem,
@@ -4212,20 +4301,25 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_friction,
       m._impl.flex_gap,
       m._impl.flex_internal,
+      m.flex_interp,
       m._impl.flex_margin,
+      m._impl.flex_node,
+      m.flex_node0,
+      m.flex_nodeadr,
+      m.flex_nodebodyid,
+      m.flex_nodenum,
       m._impl.flex_priority,
       m._impl.flex_radius,
       m._impl.flex_selfcollide,
       m._impl.flex_shell,
-      m._impl.flex_shelladr,
       m._impl.flex_shelldataadr,
-      m._impl.flex_shellflexid,
       m._impl.flex_solimp,
       m._impl.flex_solmix,
       m._impl.flex_solref,
       m._impl.flex_stiffness,
       m._impl.flex_stiffnessadr,
       m._impl.flex_vert,
+      m.flex_vert0,
       m.flex_vertadr,
       m._impl.flex_vertbodyid,
       m._impl.flex_vertflexid,
@@ -4236,7 +4330,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.flexedge_invweight0,
       m._impl.flexedge_length0,
       m._impl.flexelem_geom_pair_filtered,
-      m._impl.flexshell_geom_pair_filtered,
+      m._impl.flexstrain_J_colind,
+      m._impl.flexstrain_J_rowadr,
+      m._impl.flexstrain_J_rownnz,
       m._impl.flexvert_geom_pair_filtered,
       m.geom_aabb,
       m.geom_bodyid,
@@ -4262,6 +4358,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.geom_solmix,
       m.geom_solref,
       m.geom_type,
+      m._impl.has_3d_flex,
+      m._impl.has_ellipsoid_geom,
       m._impl.has_flex_selfcollide,
       m._impl.has_fluid,
       m._impl.has_sdf_geom,
@@ -4336,6 +4434,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.nflexedge,
       m._impl.nflexelem,
       m._impl.nflexevpair,
+      m._impl.nflexintcell,
+      m._impl.nflexnode,
       m._impl.nflexvert,
       m.ngeom,
       m.nhistory,
@@ -4545,6 +4645,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.flexedge_J,
       d._impl.flexedge_length,
       d._impl.flexedge_velocity,
+      d._impl.flexnode_xpos,
       d._impl.flexvert_xpos,
       d.geom_xmat,
       d.geom_xpos,
@@ -4701,118 +4802,119 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       '_impl.flexedge_J': out[35],
       '_impl.flexedge_length': out[36],
       '_impl.flexedge_velocity': out[37],
-      '_impl.flexvert_xpos': out[38],
-      'geom_xmat': out[39],
-      'geom_xpos': out[40],
-      'history': out[41],
-      '_impl.island_dofadr': out[42],
-      '_impl.island_idofadr': out[43],
-      '_impl.island_iefcadr': out[44],
-      '_impl.island_ne': out[45],
-      '_impl.island_nefc': out[46],
-      '_impl.island_nf': out[47],
-      '_impl.island_nv': out[48],
-      '_impl.light_xdir': out[49],
-      '_impl.light_xpos': out[50],
-      '_impl.map_dof2idof': out[51],
-      '_impl.map_efc2iefc': out[52],
-      '_impl.map_idof2dof': out[53],
-      '_impl.map_iefc2efc': out[54],
-      '_impl.moment_colind': out[55],
-      '_impl.moment_rowadr': out[56],
-      '_impl.moment_rownnz': out[57],
-      '_impl.nacon': out[58],
-      '_impl.nbody_awake': out[59],
-      '_impl.ncdof': out[60],
-      '_impl.ncollision': out[61],
-      '_impl.ne': out[62],
-      '_impl.nefc': out[63],
-      '_impl.nf': out[64],
-      '_impl.nidof': out[65],
-      '_impl.nisland': out[66],
-      '_impl.nl': out[67],
-      '_impl.ntree_awake': out[68],
-      '_impl.nv_awake': out[69],
-      '_impl.overflow': out[70],
-      '_impl.qLD': out[71],
-      '_impl.qLDiagInv': out[72],
-      '_impl.qLU': out[73],
-      'qacc': out[74],
-      'qacc_smooth': out[75],
-      'qacc_warmstart': out[76],
-      'qfrc_actuator': out[77],
-      'qfrc_bias': out[78],
-      'qfrc_constraint': out[79],
-      '_impl.qfrc_damper': out[80],
-      'qfrc_fluid': out[81],
-      'qfrc_gravcomp': out[82],
-      'qfrc_passive': out[83],
-      'qfrc_smooth': out[84],
-      '_impl.qfrc_spring': out[85],
-      'qpos': out[86],
-      'qvel': out[87],
-      'sensordata': out[88],
-      'site_xmat': out[89],
-      'site_xpos': out[90],
-      '_impl.solver_niter': out[91],
-      '_impl.subtree_angmom': out[92],
-      'subtree_com': out[93],
-      '_impl.subtree_linvel': out[94],
-      '_impl.ten_J': out[95],
-      'ten_length': out[96],
-      '_impl.ten_velocity': out[97],
-      '_impl.ten_wrapadr': out[98],
-      '_impl.ten_wrapnum': out[99],
-      'time': out[100],
-      '_impl.tree_asleep': out[101],
-      '_impl.tree_awake': out[102],
-      '_impl.tree_island': out[103],
-      '_impl.wrap_obj': out[104],
-      '_impl.wrap_xpos': out[105],
-      'xanchor': out[106],
-      'xaxis': out[107],
-      'ximat': out[108],
-      'xipos': out[109],
-      'xmat': out[110],
-      'xpos': out[111],
-      'xquat': out[112],
-      '_impl.contact__dim': out[113],
-      '_impl.contact__dist': out[114],
-      '_impl.contact__efc_address': out[115],
-      '_impl.contact__elem': out[116],
-      '_impl.contact__flex': out[117],
-      '_impl.contact__frame': out[118],
-      '_impl.contact__friction': out[119],
-      '_impl.contact__geom': out[120],
-      '_impl.contact__geomcollisionid': out[121],
-      '_impl.contact__includemargin': out[122],
-      '_impl.contact__pos': out[123],
-      '_impl.contact__solimp': out[124],
-      '_impl.contact__solref': out[125],
-      '_impl.contact__solreffriction': out[126],
-      '_impl.contact__type': out[127],
-      '_impl.contact__vert': out[128],
-      '_impl.contact__worldid': out[129],
-      '_impl.efc__D': out[130],
-      '_impl.efc__J': out[131],
-      '_impl.efc__J_colind': out[132],
-      '_impl.efc__J_rowadr': out[133],
-      '_impl.efc__J_rownnz': out[134],
-      '_impl.efc__Jqvel': out[135],
-      '_impl.efc__Ma': out[136],
-      '_impl.efc__aref': out[137],
-      '_impl.efc__force': out[138],
-      '_impl.efc__frictionloss': out[139],
-      '_impl.efc__id': out[140],
-      '_impl.efc__island': out[141],
-      '_impl.efc__jtdaj_adr': out[142],
-      '_impl.efc__jtdaj_nblock': out[143],
-      '_impl.efc__jtdaj_nrow': out[144],
-      '_impl.efc__margin': out[145],
-      '_impl.efc__pos': out[146],
-      '_impl.efc__state': out[147],
-      '_impl.efc__type': out[148],
-      '_impl.efc__vel': out[149],
+      '_impl.flexnode_xpos': out[38],
+      '_impl.flexvert_xpos': out[39],
+      'geom_xmat': out[40],
+      'geom_xpos': out[41],
+      'history': out[42],
+      '_impl.island_dofadr': out[43],
+      '_impl.island_idofadr': out[44],
+      '_impl.island_iefcadr': out[45],
+      '_impl.island_ne': out[46],
+      '_impl.island_nefc': out[47],
+      '_impl.island_nf': out[48],
+      '_impl.island_nv': out[49],
+      '_impl.light_xdir': out[50],
+      '_impl.light_xpos': out[51],
+      '_impl.map_dof2idof': out[52],
+      '_impl.map_efc2iefc': out[53],
+      '_impl.map_idof2dof': out[54],
+      '_impl.map_iefc2efc': out[55],
+      '_impl.moment_colind': out[56],
+      '_impl.moment_rowadr': out[57],
+      '_impl.moment_rownnz': out[58],
+      '_impl.nacon': out[59],
+      '_impl.nbody_awake': out[60],
+      '_impl.ncdof': out[61],
+      '_impl.ncollision': out[62],
+      '_impl.ne': out[63],
+      '_impl.nefc': out[64],
+      '_impl.nf': out[65],
+      '_impl.nidof': out[66],
+      '_impl.nisland': out[67],
+      '_impl.nl': out[68],
+      '_impl.ntree_awake': out[69],
+      '_impl.nv_awake': out[70],
+      '_impl.overflow': out[71],
+      '_impl.qLD': out[72],
+      '_impl.qLDiagInv': out[73],
+      '_impl.qLU': out[74],
+      'qacc': out[75],
+      'qacc_smooth': out[76],
+      'qacc_warmstart': out[77],
+      'qfrc_actuator': out[78],
+      'qfrc_bias': out[79],
+      'qfrc_constraint': out[80],
+      '_impl.qfrc_damper': out[81],
+      'qfrc_fluid': out[82],
+      'qfrc_gravcomp': out[83],
+      'qfrc_passive': out[84],
+      'qfrc_smooth': out[85],
+      '_impl.qfrc_spring': out[86],
+      'qpos': out[87],
+      'qvel': out[88],
+      'sensordata': out[89],
+      'site_xmat': out[90],
+      'site_xpos': out[91],
+      '_impl.solver_niter': out[92],
+      '_impl.subtree_angmom': out[93],
+      'subtree_com': out[94],
+      '_impl.subtree_linvel': out[95],
+      '_impl.ten_J': out[96],
+      'ten_length': out[97],
+      '_impl.ten_velocity': out[98],
+      '_impl.ten_wrapadr': out[99],
+      '_impl.ten_wrapnum': out[100],
+      'time': out[101],
+      '_impl.tree_asleep': out[102],
+      '_impl.tree_awake': out[103],
+      '_impl.tree_island': out[104],
+      '_impl.wrap_obj': out[105],
+      '_impl.wrap_xpos': out[106],
+      'xanchor': out[107],
+      'xaxis': out[108],
+      'ximat': out[109],
+      'xipos': out[110],
+      'xmat': out[111],
+      'xpos': out[112],
+      'xquat': out[113],
+      '_impl.contact__dim': out[114],
+      '_impl.contact__dist': out[115],
+      '_impl.contact__efc_address': out[116],
+      '_impl.contact__elem': out[117],
+      '_impl.contact__flex': out[118],
+      '_impl.contact__frame': out[119],
+      '_impl.contact__friction': out[120],
+      '_impl.contact__geom': out[121],
+      '_impl.contact__geomcollisionid': out[122],
+      '_impl.contact__includemargin': out[123],
+      '_impl.contact__pos': out[124],
+      '_impl.contact__solimp': out[125],
+      '_impl.contact__solref': out[126],
+      '_impl.contact__solreffriction': out[127],
+      '_impl.contact__type': out[128],
+      '_impl.contact__vert': out[129],
+      '_impl.contact__worldid': out[130],
+      '_impl.efc__D': out[131],
+      '_impl.efc__J': out[132],
+      '_impl.efc__J_colind': out[133],
+      '_impl.efc__J_rowadr': out[134],
+      '_impl.efc__J_rownnz': out[135],
+      '_impl.efc__Jqvel': out[136],
+      '_impl.efc__Ma': out[137],
+      '_impl.efc__aref': out[138],
+      '_impl.efc__force': out[139],
+      '_impl.efc__frictionloss': out[140],
+      '_impl.efc__id': out[141],
+      '_impl.efc__island': out[142],
+      '_impl.efc__jtdaj_adr': out[143],
+      '_impl.efc__jtdaj_nblock': out[144],
+      '_impl.efc__jtdaj_nrow': out[145],
+      '_impl.efc__margin': out[146],
+      '_impl.efc__pos': out[147],
+      '_impl.efc__state': out[148],
+      '_impl.efc__type': out[149],
+      '_impl.efc__vel': out[150],
   })
   return d
 

mjx/mujoco/mjx/warp/types.py

@@ -117,7 +117,7 @@ class BlockDim:
     render: render block dimension (render)
   """
   segmented_sort: int = 128
-  convex_ccd: int = 256
+  convex_ccd: int = 64
   actuator_velocity: int = 32
   ray: int = 64
   contact_sort: int = 64
@@ -206,11 +206,14 @@ class ModelWarp(PyTreeNode):
   dof_tri_row: np.ndarray
   eq_connect_adr: np.ndarray
   eq_flex_adr: np.ndarray
+  eq_flexstrain_adr: np.ndarray
   eq_jnt_adr: np.ndarray
   eq_ten_adr: np.ndarray
   eq_wld_adr: np.ndarray
   flex_bending: np.ndarray
   flex_bendingadr: np.ndarray
+  flex_cell_map: np.ndarray
+  flex_cellnum: np.ndarray
   flex_centered: np.ndarray
   flex_conaffinity: np.ndarray
   flex_condim: np.ndarray
@@ -219,6 +222,7 @@ class ModelWarp(PyTreeNode):
   flex_dim: np.ndarray
   flex_edge: np.ndarray
   flex_edgeadr: np.ndarray
+  flex_edgeequality: np.ndarray
   flex_edgeflap: np.ndarray
   flex_edgenum: np.ndarray
   flex_elem: np.ndarray
@@ -236,6 +240,7 @@ class ModelWarp(PyTreeNode):
   flex_gap: np.ndarray
   flex_internal: np.ndarray
   flex_margin: np.ndarray
+  flex_node: np.ndarray
   flex_priority: np.ndarray
   flex_radius: np.ndarray
   flex_selfcollide: np.ndarray
@@ -258,10 +263,14 @@ class ModelWarp(PyTreeNode):
   flexedge_invweight0: np.ndarray
   flexedge_length0: np.ndarray
   flexelem_geom_pair_filtered: np.ndarray
-  flexshell_geom_pair_filtered: np.ndarray
+  flexstrain_J_colind: np.ndarray
+  flexstrain_J_rowadr: np.ndarray
+  flexstrain_J_rownnz: np.ndarray
   flexvert_geom_pair_filtered: np.ndarray
   geom_pair_type_count: Tuple[int, ...]
   geom_plugin_index: np.ndarray
+  has_3d_flex: bool
+  has_ellipsoid_geom: bool
   has_flex_selfcollide: bool
   has_fluid: bool
   has_sdf_geom: bool
@@ -288,14 +297,18 @@ class ModelWarp(PyTreeNode):
   mesh_polyvertnum: np.ndarray
   mocap_bodyid: np.ndarray
   nJfe: int
+  nJfs: int
   nacttrnbody: int
   nbranch: int
+  neq_flexstrain: int
   nflexbending: int
   nflexedge: int
   nflexelem: int
   nflexelemdata: int
   nflexelemedge: int
   nflexevpair: int
+  nflexintcell: int
+  nflexnode: int
   nflexshelldata: int
   nflexstiffness: int
   nflexvert: int
@@ -452,6 +465,7 @@ class DataWarp(PyTreeNode):
   flexedge_J: jax.Array
   flexedge_length: jax.Array
   flexedge_velocity: jax.Array
+  flexnode_xpos: jax.Array
   flexvert_xpos: jax.Array
   island_dofadr: jax.Array
   island_idofadr: jax.Array
@@ -664,6 +678,7 @@ _NDIM = {
         'flexedge_J': 2,
         'flexedge_length': 2,
         'flexedge_velocity': 2,
+        'flexnode_xpos': 3,
         'flexvert_xpos': 3,
         'geom_xmat': 4,
         'geom_xpos': 3,
@@ -882,6 +897,7 @@ _NDIM = {
         'eq_connect_adr': 1,
         'eq_data': 3,
         'eq_flex_adr': 1,
+        'eq_flexstrain_adr': 1,
         'eq_jnt_adr': 1,
         'eq_obj1id': 1,
         'eq_obj2id': 1,
@@ -894,6 +910,8 @@ _NDIM = {
         'exclude_signature': 1,
         'flex_bending': 1,
         'flex_bendingadr': 1,
+        'flex_cell_map': 2,
+        'flex_cellnum': 2,
         'flex_centered': 1,
         'flex_conaffinity': 1,
         'flex_condim': 1,
@@ -902,6 +920,7 @@ _NDIM = {
         'flex_dim': 1,
         'flex_edge': 2,
         'flex_edgeadr': 1,
+        'flex_edgeequality': 1,
         'flex_edgeflap': 2,
         'flex_edgenum': 1,
         'flex_elem': 1,
@@ -918,7 +937,13 @@ _NDIM = {
         'flex_friction': 2,
         'flex_gap': 1,
         'flex_internal': 1,
+        'flex_interp': 1,
         'flex_margin': 1,
+        'flex_node': 2,
+        'flex_node0': 2,
+        'flex_nodeadr': 1,
+        'flex_nodebodyid': 1,
+        'flex_nodenum': 1,
         'flex_priority': 1,
         'flex_radius': 1,
         'flex_selfcollide': 1,
@@ -933,6 +958,7 @@ _NDIM = {
         'flex_stiffness': 1,
         'flex_stiffnessadr': 1,
         'flex_vert': 2,
+        'flex_vert0': 2,
         'flex_vertadr': 1,
         'flex_vertbodyid': 1,
         'flex_vertflexid': 1,
@@ -943,7 +969,9 @@ _NDIM = {
         'flexedge_invweight0': 1,
         'flexedge_length0': 1,
         'flexelem_geom_pair_filtered': 2,
-        'flexshell_geom_pair_filtered': 2,
+        'flexstrain_J_colind': 1,
+        'flexstrain_J_rowadr': 1,
+        'flexstrain_J_rownnz': 1,
         'flexvert_geom_pair_filtered': 2,
         'geom_aabb': 4,
         'geom_bodyid': 1,
@@ -969,6 +997,8 @@ _NDIM = {
         'geom_solmix': 2,
         'geom_solref': 3,
         'geom_type': 1,
+        'has_3d_flex': 0,
+        'has_ellipsoid_geom': 0,
         'has_flex_selfcollide': 0,
         'has_fluid': 0,
         'has_sdf_geom': 0,
@@ -1051,6 +1081,7 @@ _NDIM = {
         'nC': 0,
         'nD': 0,
         'nJfe': 0,
+        'nJfs': 0,
         'nJmom': 0,
         'nJten': 0,
         'nM': 0,
@@ -1060,6 +1091,7 @@ _NDIM = {
         'nbranch': 0,
         'ncam': 0,
         'neq': 0,
+        'neq_flexstrain': 0,
         'nexclude': 0,
         'nflex': 0,
         'nflexbending': 0,
@@ -1068,6 +1100,8 @@ _NDIM = {
         'nflexelemdata': 0,
         'nflexelemedge': 0,
         'nflexevpair': 0,
+        'nflexintcell': 0,
+        'nflexnode': 0,
         'nflexshelldata': 0,
         'nflexstiffness': 0,
         'nflexvert': 0,
@@ -1365,6 +1399,7 @@ _BATCH_DIM = {
         'flexedge_J': True,
         'flexedge_length': True,
         'flexedge_velocity': True,
+        'flexnode_xpos': True,
         'flexvert_xpos': True,
         'geom_xmat': True,
         'geom_xpos': True,
@@ -1583,6 +1618,7 @@ _BATCH_DIM = {
         'eq_connect_adr': False,
         'eq_data': True,
         'eq_flex_adr': False,
+        'eq_flexstrain_adr': False,
         'eq_jnt_adr': False,
         'eq_obj1id': False,
         'eq_obj2id': False,
@@ -1595,6 +1631,8 @@ _BATCH_DIM = {
         'exclude_signature': False,
         'flex_bending': False,
         'flex_bendingadr': False,
+        'flex_cell_map': False,
+        'flex_cellnum': False,
         'flex_centered': False,
         'flex_conaffinity': False,
         'flex_condim': False,
@@ -1603,6 +1641,7 @@ _BATCH_DIM = {
         'flex_dim': False,
         'flex_edge': False,
         'flex_edgeadr': False,
+        'flex_edgeequality': False,
         'flex_edgeflap': False,
         'flex_edgenum': False,
         'flex_elem': False,
@@ -1619,7 +1658,13 @@ _BATCH_DIM = {
         'flex_friction': False,
         'flex_gap': False,
         'flex_internal': False,
+        'flex_interp': False,
         'flex_margin': False,
+        'flex_node': False,
+        'flex_node0': False,
+        'flex_nodeadr': False,
+        'flex_nodebodyid': False,
+        'flex_nodenum': False,
         'flex_priority': False,
         'flex_radius': False,
         'flex_selfcollide': False,
@@ -1634,6 +1679,7 @@ _BATCH_DIM = {
         'flex_stiffness': False,
         'flex_stiffnessadr': False,
         'flex_vert': False,
+        'flex_vert0': False,
         'flex_vertadr': False,
         'flex_vertbodyid': False,
         'flex_vertflexid': False,
@@ -1644,7 +1690,9 @@ _BATCH_DIM = {
         'flexedge_invweight0': False,
         'flexedge_length0': False,
         'flexelem_geom_pair_filtered': False,
-        'flexshell_geom_pair_filtered': False,
+        'flexstrain_J_colind': False,
+        'flexstrain_J_rowadr': False,
+        'flexstrain_J_rownnz': False,
         'flexvert_geom_pair_filtered': False,
         'geom_aabb': True,
         'geom_bodyid': False,
@@ -1670,6 +1718,8 @@ _BATCH_DIM = {
         'geom_solmix': True,
         'geom_solref': True,
         'geom_type': False,
+        'has_3d_flex': False,
+        'has_ellipsoid_geom': False,
         'has_flex_selfcollide': False,
         'has_fluid': False,
         'has_sdf_geom': False,
@@ -1752,6 +1802,7 @@ _BATCH_DIM = {
         'nC': False,
         'nD': False,
         'nJfe': False,
+        'nJfs': False,
         'nJmom': False,
         'nJten': False,
         'nM': False,
@@ -1761,6 +1812,7 @@ _BATCH_DIM = {
         'nbranch': False,
         'ncam': False,
         'neq': False,
+        'neq_flexstrain': False,
         'nexclude': False,
         'nflex': False,
         'nflexbending': False,
@@ -1769,6 +1821,8 @@ _BATCH_DIM = {
         'nflexelemdata': False,
         'nflexelemedge': False,
         'nflexevpair': False,
+        'nflexintcell': False,
+        'nflexnode': False,
         'nflexshelldata': False,
         'nflexstiffness': False,
         'nflexvert': False,