Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 923922783
Change-Id: I0b407347a70a0e2024aae7ed5bab1592dc40a074

mjx/mujoco/mjx/_src/io.py

@@ -272,6 +272,10 @@ def _put_option(
 
 
   if impl == types.Impl.WARP:
+    if not mjxw.mjwp_io.ENABLE_ISLANDS:
+      fields['disableflags'] = types.DisableBit(
+          fields['disableflags'] | mjwp_types.DisableBit.ISLAND
+      )
     impl_fields = {k: _wp_to_np_type(v) for k, v in impl_fields.items()}
     return types.Option(**fields, _impl=mjxw.types.OptionWarp(**impl_fields))
 
@@ -1272,8 +1276,11 @@ def _get_data_into_warp(
       else:
         value = getattr(d_i, field.name)
 
-      if field.name in ('ne', 'nl', 'nf'):
+      if field.name in ('ne', 'nl', 'nf', 'nisland', 'nidof'):
-        pass
+        if isinstance(value, np.ndarray) and value.size == 0:
+          value = 0
+        else:
+          value = int(value)
       elif field.name in ('nefc', 'ncon'):
         value = {'nefc': nefc, 'ncon': ncon}[field.name]
       elif field.name.endswith('xmat') or field.name == 'ximat':
@@ -1287,9 +1294,11 @@ def _get_data_into_warp(
           'contact',
           'qM',
           'qLD',
+          'qLU',
           'qLDiagInv',
           'ten_J',
           'flexedge_J',
+          'M',
       ):
         continue
       if field.name.startswith('efc_'):

mjx/mujoco/mjx/_src/io_test.py

@@ -530,7 +530,7 @@ class DataIOTest(parameterized.TestCase):
     elif impl == 'warp':
       qm = np.zeros((m.nv, m.nv), dtype=np.float64)
       mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
-      np.testing.assert_allclose(dx._impl.qM, qm)
+      np.testing.assert_allclose(dx._impl.M[:m.nv, :m.nv], qm)
       # TODO(taylorhowell): test efc__J
       np.testing.assert_allclose(dx._impl.efc__aref[:3], d.efc_aref[:3])
 

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

@@ -47,6 +47,10 @@ from mujoco.mjx.third_party.mujoco_warp._src.forward import implicit as implicit
 from mujoco.mjx.third_party.mujoco_warp._src.forward import rungekutta4 as rungekutta4
 from mujoco.mjx.third_party.mujoco_warp._src.forward import step1 as step1
 from mujoco.mjx.third_party.mujoco_warp._src.forward import step2 as step2
+from mujoco.mjx.third_party.mujoco_warp._src.history import init_ctrl_history as init_ctrl_history
+from mujoco.mjx.third_party.mujoco_warp._src.history import init_sensor_history as init_sensor_history
+from mujoco.mjx.third_party.mujoco_warp._src.history import read_ctrl as read_ctrl
+from mujoco.mjx.third_party.mujoco_warp._src.history import read_sensor as read_sensor
 from mujoco.mjx.third_party.mujoco_warp._src.inverse import inverse as inverse
 from mujoco.mjx.third_party.mujoco_warp._src.io import create_render_context as create_render_context
 from mujoco.mjx.third_party.mujoco_warp._src.io import get_data_into as get_data_into

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

@@ -19,48 +19,77 @@ import warp as wp
 
 
 @lru_cache(maxsize=None)
-def create_blocked_cholesky_func(block_size: int):
+def create_blocked_cholesky_factorize_solve_func(block_size: int, matrix_size_static: int):
   @wp.func
-  def blocked_cholesky_func(
+  def blocked_cholesky_factorize_solve_func(
     # In:
     A: wp.array2d[float],
+    b: wp.array2d[float],
     matrix_size: int,
     # Out:
-    L: wp.array2d[float],
+    U: wp.array2d[float],
+    x: wp.array2d[float],
   ):
-    """Computes the Cholesky factorization of a symmetric positive definite matrix A in blocks.
+    """Block Cholesky factorization and solve while keeping the forward RHS live."""
+    rhs_tile = wp.tile_load(b, shape=(matrix_size_static, 1), offset=(0, 0), storage="shared", bounds_check=False)
 
-    It returns a lower-triangular matrix L such that A = L L^T.
-    """
-    # Process the matrix in blocks along its leading dimension.
     for k in range(0, matrix_size, block_size):
       end = k + block_size
+      rhs_view = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(k, 0))
 
-      # Load current diagonal block A[k:end, k:end]
+      A_kk_tile = wp.tile_load(
-      # and update with contributions from previously computed blocks.
+        A, shape=(block_size, block_size), offset=(k, k), storage="shared", bounds_check=False, aligned=True
-      A_kk_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(k, k), storage="shared")
+      )
 
       for j in range(0, k, block_size):
-        L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
+        U_block = wp.tile_load(
-        wp.tile_matmul(L_block, wp.tile_transpose(L_block), A_kk_tile, alpha=-1.0)
+          U, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
+        )
+        wp.tile_matmul(wp.tile_transpose(U_block), U_block, A_kk_tile, alpha=-1.0)
 
-      # Compute the Cholesky factorization for the block
+        y_block = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
-      wp.tile_cholesky_inplace(A_kk_tile)
+        wp.tile_matmul(wp.tile_transpose(U_block), y_block, rhs_view, alpha=-1.0)
-      wp.tile_store(L, A_kk_tile, offset=(k, k))
+
+      wp.tile_cholesky_inplace(A_kk_tile, fill_mode="upper")
+      wp.tile_store(U, A_kk_tile, offset=(k, k), bounds_check=False, aligned=True)
+
+      wp.tile_lower_solve_inplace(wp.tile_transpose(A_kk_tile), rhs_view)
 
-      # Process the blocks below the current block
       for i in range(end, matrix_size, block_size):
-        A_ik_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(i, k), storage="shared")
+        A_ki_tile = wp.tile_load(
+          A, shape=(block_size, block_size), offset=(k, i), storage="shared", bounds_check=False, aligned=True
+        )
 
         for j in range(0, k, block_size):
-          L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
+          U_jk_tile = wp.tile_load(
-          L_2_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
+            U, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
-          wp.tile_matmul(L_tile, wp.tile_transpose(L_2_tile), A_ik_tile, alpha=-1.0)
+          )
+          U_ji_tile = wp.tile_load(
+            U, shape=(block_size, block_size), offset=(j, i), storage="shared", bounds_check=False, aligned=True
+          )
+          wp.tile_matmul(wp.tile_transpose(U_jk_tile), U_ji_tile, A_ki_tile, alpha=-1.0)
 
-        wp.tile_lower_solve_inplace(A_kk_tile, wp.tile_transpose(A_ik_tile))
+        wp.tile_lower_solve_inplace(wp.tile_transpose(A_kk_tile), A_ki_tile)
-        wp.tile_store(L, A_ik_tile, offset=(i, k))
+        wp.tile_store(U, A_ki_tile, offset=(k, i), bounds_check=False, aligned=True)
 
-  return blocked_cholesky_func
+    for i in range(matrix_size - block_size, -1, -block_size):
+      i_end = i + block_size
+      tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
+      for j in range(i_end, matrix_size, block_size):
+        U_tile = wp.tile_load(
+          U, shape=(block_size, block_size), offset=(i, j), storage="shared", bounds_check=False, aligned=True
+        )
+        x_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
+        wp.tile_matmul(U_tile, x_tile, tmp_tile, alpha=-1.0)
+
+      U_tile = wp.tile_load(
+        U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
+      )
+      wp.tile_upper_solve_inplace(U_tile, tmp_tile)
+
+    wp.tile_store(x, rhs_tile, offset=(0, 0), bounds_check=False)
+
+  return blocked_cholesky_factorize_solve_func
 
 
 @lru_cache(maxsize=None)
@@ -68,41 +97,49 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
   @wp.func
   def blocked_cholesky_solve_func(
     # In:
-    L: wp.array2d[float],
+    U: wp.array2d[float],
     b: wp.array2d[float],
     matrix_size: int,
     # Out:
     x: wp.array2d[float],
   ):
-    """Block Cholesky factorization and solve.
+    """Block Cholesky solve.
 
-    Solves A x = b given the Cholesky factor L (A = L L^T) using blocked forward and backward
+    Solves A x = b given the Cholesky factor U (A = U^T U) using blocked forward and backward
     substitution.
     """
     rhs_tile = wp.tile_load(b, shape=(matrix_size_static, 1), offset=(0, 0), storage="shared", bounds_check=False)
 
-    # Forward substitution: solve L y = b
+    # Forward substitution: solve U^T y = b
     for i in range(0, matrix_size, block_size):
       rhs_view = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
       for j in range(0, i, block_size):
-        L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
+        U_block = wp.tile_load(
+          U, shape=(block_size, block_size), offset=(j, i), storage="shared", bounds_check=False, aligned=True
+        )
         y_block = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
-        wp.tile_matmul(L_block, y_block, rhs_view, alpha=-1.0)
+        wp.tile_matmul(wp.tile_transpose(U_block), y_block, rhs_view, alpha=-1.0)
 
-      L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
+      U_tile = wp.tile_load(
-      wp.tile_lower_solve_inplace(L_tile, rhs_view)
+        U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
+      )
+      wp.tile_lower_solve_inplace(wp.tile_transpose(U_tile), rhs_view)
 
-    # Backward substitution: solve L^T x = y
+    # Backward substitution: solve U x = y
     for i in range(matrix_size - block_size, -1, -block_size):
       i_end = i + block_size
       tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
       for j in range(i_end, matrix_size, block_size):
-        L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(j, i), storage="shared")
+        U_tile = wp.tile_load(
+          U, shape=(block_size, block_size), offset=(i, j), storage="shared", bounds_check=False, aligned=True
+        )
         x_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
-        wp.tile_matmul(wp.tile_transpose(L_tile), x_tile, tmp_tile, alpha=-1.0)
+        wp.tile_matmul(U_tile, x_tile, tmp_tile, alpha=-1.0)
-      L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
+      U_tile = wp.tile_load(
+        U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
+      )
 
-      wp.tile_upper_solve_inplace(wp.tile_transpose(L_tile), tmp_tile)
+      wp.tile_upper_solve_inplace(U_tile, tmp_tile)
 
     wp.tile_store(x, rhs_tile, offset=(0, 0), bounds_check=False)
 

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

@@ -25,6 +25,7 @@ from absl import flags
 from etils import epath
 
 import mujoco.mjx.third_party.mujoco_warp as mjw
+from mujoco.mjx.third_party.mujoco_warp._src import io
 from mujoco.mjx.third_party.mujoco_warp._src import warp_util
 from mujoco.mjx.third_party.mujoco_warp._src.io import load_trajectory
 from mujoco.mjx.third_party.mujoco_warp._src.io import override_model
@@ -42,6 +43,12 @@ KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulatio
 EVENT_TRACE = flags.DEFINE_bool("event_trace", False, "print an event trace report")
 NOISE_STD = flags.DEFINE_float("noise_std", 0.01, "add noise to ctrl signal (standard deviation)")
 NOISE_RATE = flags.DEFINE_float("noise_rate", 0.1, "add noise to ctrl signal (noise rate)")
+ENABLE_ISLANDS = flags.DEFINE_bool(
+  "enable_islands",
+  False,
+  "Enable constraint islands solver",
+)
+
 
 DEVICE = flags.DEFINE_string("device", None, "override the default Warp device")
 REPLAY = flags.DEFINE_string("replay", None, "NPZ file with ctrl sequence to replay")
@@ -52,6 +59,12 @@ RENDER_RGB = flags.DEFINE_bool("render_rgb", True, "render RGB image")
 RENDER_DEPTH = flags.DEFINE_bool("render_depth", True, "render depth image")
 RENDER_TEXTURES = flags.DEFINE_bool("render_textures", True, "use textures")
 RENDER_SHADOWS = flags.DEFINE_bool("render_shadows", False, "use shadows")
+RENDER_BACKFACE_CULLING = flags.DEFINE_bool(
+  "render_backface_culling",
+  True,
+  "enable renderer backface culling (RenderContext.enable_backface_culling)",
+)
+RENDER_SKYBOX = flags.DEFINE_bool("render_skybox", True, "render skybox texture if available")
 
 
 def load_model(path: epath.Path) -> mujoco.MjModel:
@@ -128,6 +141,8 @@ def init_structs(
   fn: Callable[..., None], mjm: mujoco.MjModel
 ) -> Tuple[mjw.Model, mjw.Data, mjw.RenderContext | None, list[np.ndarray] | None]:
   """Initialize device structs."""
+  io.ENABLE_ISLANDS = ENABLE_ISLANDS.value
+
   mjd = mujoco.MjData(mjm)
   ctrls = None
   if REPLAY.value:
@@ -152,12 +167,14 @@ def init_structs(
 
     rc = mjw.create_render_context(
       mjm,
-      NWORLD.value,
+      nworld=NWORLD.value,
-      (RENDER_WIDTH.value, RENDER_HEIGHT.value),
+      cam_res=(RENDER_WIDTH.value, RENDER_HEIGHT.value),
-      RENDER_RGB.value,
+      render_rgb=RENDER_RGB.value,
-      RENDER_DEPTH.value,
+      render_depth=RENDER_DEPTH.value,
-      RENDER_TEXTURES.value,
+      use_textures=RENDER_TEXTURES.value,
-      RENDER_SHADOWS.value,
+      use_shadows=RENDER_SHADOWS.value,
+      enable_backface_culling=RENDER_BACKFACE_CULLING.value,
+      render_skybox=RENDER_SKYBOX.value,
     )
 
     return m, d, rc, ctrls

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

@@ -23,6 +23,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_param
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import ccd
+from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import epa_phase
+from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import gjk_phase
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import multicontact
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import support
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
@@ -35,7 +37,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAFACES
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAHORIZON
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
-from mujoco.mjx.third_party.mujoco_warp._src.types import _NEW_GAP_SEMANTICS
+from mujoco.mjx.third_party.mujoco_warp._src.types import NEW_GAP_SEMANTICS
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
@@ -715,6 +717,7 @@ def ccd_kernel_builder(
     opt_ccd_tolerance: wp.array[float],
     # Data in:
     naconmax_in: int,
+    naccdmax_in: int,
     # In:
     epa_vert_in: wp.array2d[wp.vec3],
     epa_vert_index_in: wp.array2d[int],
@@ -737,7 +740,7 @@ def ccd_kernel_builder(
     geom2: Geom,
     geoms: wp.vec2i,
     worldid: int,
-    ccdid: int,
+    nccd_in: wp.array[int],
     margin: float,
     gap: float,
     condim: int,
@@ -774,21 +777,38 @@ def ccd_kernel_builder(
     if is_collision_sensor:
       cutoff = 1.0e32
     else:
-      if wp.static(_NEW_GAP_SEMANTICS):
+      if wp.static(NEW_GAP_SEMANTICS):
         cutoff = gap
       else:
         cutoff = 0.0
-    dist, ncollision, w1, w2, multiccd_idx = ccd(
+    needs_epa, dist, ncollision, w1, w2, gjk_result, geom1, geom2 = gjk_phase(
       opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
       cutoff,
       gjk_iterations,
-      epa_iterations,
       geom1,
       geom2,
       geomtype1,
       geomtype2,
       x1,
       x2,
+    )
+
+    ccdid = int(-1)
+    multiccd_idx = int(-1)
+
+    if needs_epa:
+      ccdid = wp.atomic_add(nccd_in, geomgeomid, 1)
+      if ccdid >= naccdmax_in:
+        wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid)
+        return 0
+      dist, ncollision, w1, w2, multiccd_idx = epa_phase(
+        opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
+        epa_iterations,
+        gjk_result,
+        geom1,
+        geom2,
+        geomtype1,
+        geomtype2,
         epa_vert_in[ccdid],
         epa_vert_index_in[ccdid],
         epa_face_in[ccdid],
@@ -797,7 +817,7 @@ def ccd_kernel_builder(
         epa_horizon_in[ccdid],
       )
 
-    if wp.static(_NEW_GAP_SEMANTICS):
+    if wp.static(NEW_GAP_SEMANTICS):
       if dist >= gap and pairid[1] == -1:
         return 0
     else:
@@ -990,11 +1010,6 @@ def ccd_kernel_builder(
     if geom_type[g1] != geomtype1 or geom_type[g2] != geomtype2:
       return
 
-    ccdid = wp.atomic_add(nccd_in, wp.static(geomgeomid), 1)
-    if ccdid >= naccdmax_in:
-      wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid)
-      return
-
     worldid = collision_worldid_in[collisionid]
 
     _, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params(
@@ -1046,6 +1061,7 @@ def ccd_kernel_builder(
     eval_ccd_write_contact(
       opt_ccd_tolerance,
       naconmax_in,
+      naccdmax_in,
       epa_vert_in,
       epa_vert_index_in,
       epa_face_in,
@@ -1067,7 +1083,7 @@ def ccd_kernel_builder(
       geom2,
       geoms,
       worldid,
-      ccdid,
+      nccd_in,
       margin,
       gap,
       condim,

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

@@ -22,8 +22,8 @@ import warp as wp
 
 from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU
-from mujoco.mjx.third_party.mujoco_warp._src.types import _NEW_GAP_SEMANTICS
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import NEW_GAP_SEMANTICS
 from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
@@ -219,7 +219,7 @@ def write_contact(
     contact_frame_out[cid] = frame_in
     contact_geom_out[cid] = geoms_in
     contact_worldid_out[cid] = worldid_in
-    if wp.static(_NEW_GAP_SEMANTICS):
+    if wp.static(NEW_GAP_SEMANTICS):
       includemargin = margin_in
     else:
       includemargin = margin_in - gap_in

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

@@ -375,6 +375,7 @@ def _binary_search(values: wp.array[Any], value: Any, lower: int, upper: int) ->
   return upper
 
 
+@cache_kernel
 def _sap_project(opt_broadphase: int):
   @wp.kernel(module="unique", enable_backward=False)
   def sap_project(
@@ -531,6 +532,7 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
   return kernel
 
 
+@cache_kernel
 def _segmented_sort(tile_size: int):
   @wp.kernel(module="unique")
   def segmented_sort(
@@ -629,7 +631,9 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
   # assumes each geom has 5 other geoms (batched over all worlds)
   nsweep = 5 * nworldgeom
   wp.launch(
-    kernel=_sap_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]),
+    kernel=_sap_broadphase(
+      m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
+    ),
     dim=nsweep,
     inputs=[
       m.ngeom,
@@ -718,7 +722,9 @@ def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
   `contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags.
   """
   wp.launch(
-    _nxn_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]),
+    _nxn_broadphase(
+      m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
+    ),
     dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]),
     inputs=[
       m.geom_type,

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

@@ -2135,8 +2135,12 @@ def multicontact(
 
   # face1 is an edge; clip face1 against face2
   if is_edge_contact_geom1:
-    approx_dir = wp.norm_l2(dir) * n2[j]
+    approx_dir = -wp.norm_l2(dir) * n2[j]
-    return _polygon_clip(plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped)
+    nclipped, clipped1, clipped2 = _polygon_clip(
+      plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped
+    )
+    # the faces were flipped in calling _polygon_clip so we need to flip them back
+    return nclipped, clipped2, clipped1
 
   # face2 is an edge; clip face2 against face1
   if is_edge_contact_geom2:
@@ -2197,34 +2201,29 @@ def _inflate(
 
 
 @wp.func
-def ccd(
+def gjk_phase(
   # In:
   tolerance: float,
   cutoff: float,
   gjk_iterations: int,
-  epa_iterations: int,
   geom1: Geom,
   geom2: Geom,
   geomtype1: int,
   geomtype2: int,
   x_1: wp.vec3,
   x_2: wp.vec3,
-  vert: wp.array[wp.vec3],
+) -> Tuple[bool, float, int, wp.vec3, wp.vec3, GJKResult, Geom, Geom]:
-  vert_index: wp.array[int],
+  """Run GJK phase of CCD."""
-  face: wp.array[int],
-  face_pr: wp.array[wp.vec3],
-  face_norm2: wp.array[float],
-  horizon: wp.array[int],
-) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
-  """General convex collision detection via GJK/EPA."""
   full_margin1 = 0.0
   full_margin2 = 0.0
   size1 = 0.0
   size2 = 0.0
+  empty = GJKResult()
 
   # determine if the geoms being tested are discrete
   is_discrete = _discrete_geoms(geomtype1, geomtype2) and (geom1.margin == 0.0 and geom2.margin == 0.0)
 
+  # special handling for sphere and capsule (shrink to point and line respectively)
   if geomtype1 == GeomType.SPHERE or geomtype1 == GeomType.CAPSULE:
     size1 = geom1.size[0]
     full_margin1 = size1 + 0.5 * geom1.margin
@@ -2237,7 +2236,6 @@ def ccd(
     geom2.margin = 0.0
     geom2.size = wp.vec3(0.0, geom2.size[1], geom2.size[2])
 
-  # special handling for sphere and capsule (shrink to point and line respectively)
   if size1 + size2 > 0.0:
     cutoff += full_margin1 + full_margin2
     result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
@@ -2245,11 +2243,11 @@ def ccd(
     # shallow penetration, inflate contact
     if result.dist > tolerance:
       if result.dist == FLOAT_MAX:
-        return result.dist, 1, result.x1, result.x2, -1
+        return False, result.dist, 1, result.x1, result.x2, empty, geom1, geom2
       dist, x1, x2 = _inflate(result, geom1, geom2, geomtype1, geomtype2, full_margin1, full_margin2)
-      return dist, 1, x1, x2, -1
+      return False, dist, 1, x1, x2, empty, geom1, geom2
 
-    # deep penetration, reset initial conditions and rerun GJK + EPA
+    # deep penetration: reset initial conditions and rerun GJK + EPA
     geom1.margin = full_margin1 - size1
     geom1.size = wp.vec3(size1, geom1.size[1], geom1.size[2])
     geom2.margin = full_margin2 - size2
@@ -2260,8 +2258,29 @@ def ccd(
 
   # no penetration depth to recover
   if result.dist > tolerance or result.dim < 2:
-    return result.dist, 1, result.x1, result.x2, -1
+    return False, result.dist, 1, result.x1, result.x2, empty, geom1, geom2
 
+  return True, result.dist, 1, result.x1, result.x2, result, geom1, geom2
+
+
+@wp.func
+def epa_phase(
+  # In:
+  tolerance: float,
+  epa_iterations: int,
+  result: GJKResult,
+  geom1: Geom,
+  geom2: Geom,
+  geomtype1: int,
+  geomtype2: int,
+  vert: wp.array[wp.vec3],
+  vert_index: wp.array[int],
+  face: wp.array[int],
+  face_pr: wp.array[wp.vec3],
+  face_norm2: wp.array[float],
+  horizon: wp.array[int],
+) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
+  """Run EPA given GJK result. Returns (dist, ncontact, x1, x2, multiccd_idx)."""
   pt = Polytope()
   pt.nface = 0
   pt.nvert = 0
@@ -2330,6 +2349,7 @@ def ccd(
   if pt.status:
     return result.dist, 1, result.x1, result.x2, -1
 
+  is_discrete = _discrete_geoms(geomtype1, geomtype2) and (geom1.margin == 0.0 and geom2.margin == 0.0)
   dist, x1, x2, idx = _epa(tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
   if idx == -1:
     return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1
@@ -2343,3 +2363,46 @@ def ccd(
     idx = -1
 
   return dist, 1, x1, x2, idx
+
+
+@wp.func
+def ccd(
+  # In:
+  tolerance: float,
+  cutoff: float,
+  gjk_iterations: int,
+  epa_iterations: int,
+  geom1: Geom,
+  geom2: Geom,
+  geomtype1: int,
+  geomtype2: int,
+  x_1: wp.vec3,
+  x_2: wp.vec3,
+  vert: wp.array[wp.vec3],
+  vert_index: wp.array[int],
+  face: wp.array[int],
+  face_pr: wp.array[wp.vec3],
+  face_norm2: wp.array[float],
+  horizon: wp.array[int],
+) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
+  """General convex collision detection via GJK/EPA."""
+  needs_epa, dist, ncontact, x1, x2, result, geom1, geom2 = gjk_phase(
+    tolerance, cutoff, gjk_iterations, geom1, geom2, geomtype1, geomtype2, x_1, x_2
+  )
+  if not needs_epa:
+    return dist, ncontact, x1, x2, -1
+  return epa_phase(
+    tolerance,
+    epa_iterations,
+    result,
+    geom1,
+    geom2,
+    geomtype1,
+    geomtype2,
+    vert,
+    vert_index,
+    face,
+    face_pr,
+    face_norm2,
+    horizon,
+  )

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

@@ -23,6 +23,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collisio
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
 from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
 from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
@@ -183,6 +184,24 @@ def ellipsoid(p: wp.vec3, size: wp.vec3) -> float:
   return k0 * (k0 - 1.0) / denom
 
 
+@wp.func
+def capsule(p: wp.vec3, size: wp.vec3) -> float:
+  r = size[0]
+  h = size[1]
+  pz_clamped = wp.clamp(p[2], -h, h)
+  diff = wp.vec3(p[0], p[1], p[2] - pz_clamped)
+  return wp.length(diff) - r
+
+
+@wp.func
+def cylinder(p: wp.vec3, size: wp.vec3) -> float:
+  r = size[0]
+  h = size[1]
+  dx = wp.length(wp.vec2(p[0], p[1])) - r
+  dy = wp.abs(p[2]) - h
+  return wp.min(wp.max(dx, dy), 0.0) + wp.length(wp.vec2(wp.max(dx, 0.0), wp.max(dy, 0.0)))
+
+
 @wp.func
 def grad_sphere(p: wp.vec3) -> wp.vec3:
   c = wp.length(p)
@@ -230,6 +249,52 @@ def grad_ellipsoid(p: wp.vec3, size: wp.vec3) -> wp.vec3:
   return raw_grad / wp.length(raw_grad)
 
 
+@wp.func
+def grad_capsule(p: wp.vec3, size: wp.vec3) -> wp.vec3:
+  h = size[1]
+  pz_clamped = wp.clamp(p[2], -h, h)
+  diff = wp.vec3(p[0], p[1], p[2] - pz_clamped)
+  c = wp.length(diff)
+  if c > MJ_MINVAL:
+    return diff / c
+  else:
+    return wp.vec3(0.0)
+
+
+@wp.func
+def grad_cylinder(p: wp.vec3, size: wp.vec3) -> wp.vec3:
+  r = size[0]
+  h = size[1]
+
+  radial_dist = wp.length(wp.vec2(p[0], p[1]))
+  if radial_dist > MJ_MINVAL:
+    u = wp.vec3(p[0] / radial_dist, p[1] / radial_dist, 0.0)
+  else:
+    u = wp.vec3(0.0)
+
+  w = wp.vec3(0.0, 0.0, wp.sign(p[2]))
+
+  dx = radial_dist - r
+  dy = wp.abs(p[2]) - h
+
+  if dx > 0.0 and dy > 0.0:
+    v = wp.vec2(dx, dy)
+    len_v = wp.length(v)
+    if len_v > MJ_MINVAL:
+      return u * (dx / len_v) + w * (dy / len_v)
+    else:
+      return wp.vec3(0.0)
+  elif dx > 0.0:
+    return u
+  elif dy > 0.0:
+    return w
+  else:
+    if dx > dy:
+      return u
+    else:
+      return w
+
+
 @wp.func
 def user_sdf(p: wp.vec3, attr: vec_pluginattr, sdf_type: int) -> float:
   """User-defined SDF function.
@@ -287,16 +352,17 @@ def find_oct(
     # check if the node is a leaf
     # child indices are relative to root (mesh_octadr offset)
     child0 = oct_child[node][0]
+    # Evaluate this hot leaf predicate eagerly to avoid branch-heavy codegen.
     if (
-      child0 == -1
+      int(child0 == -1)
-      and oct_child[node][1] == -1
+      & int(oct_child[node][1] == -1)
-      and oct_child[node][2] == -1
+      & int(oct_child[node][2] == -1)
-      and oct_child[node][3] == -1
+      & int(oct_child[node][3] == -1)
-      and oct_child[node][4] == -1
+      & int(oct_child[node][4] == -1)
-      and oct_child[node][5] == -1
+      & int(oct_child[node][5] == -1)
-      and oct_child[node][6] == -1
+      & int(oct_child[node][6] == -1)
-      and oct_child[node][7] == -1
+      & int(oct_child[node][7] == -1)
-    ):
+    ) != 0:
       for j in range(8):
         if not grad:
           rx[j] = (
@@ -394,6 +460,10 @@ def sdf(type: int, p: wp.vec3, attr: vec_pluginattr, sdf_type: int, volume_data:
     return p[2]
   elif type == GeomType.SPHERE:
     return sphere(p, attr_vec3)
+  elif type == GeomType.CAPSULE:
+    return capsule(p, attr_vec3)
+  elif type == GeomType.CYLINDER:
+    return cylinder(p, attr_vec3)
   elif type == GeomType.BOX:
     return box(p, attr_vec3)
   elif type == GeomType.ELLIPSOID:
@@ -452,6 +522,10 @@ def sdf_grad(
     return grad
   elif type == GeomType.SPHERE:
     return grad_sphere(p)
+  elif type == GeomType.CAPSULE:
+    return grad_capsule(p, attr_vec3)
+  elif type == GeomType.CYLINDER:
+    return grad_cylinder(p, attr_vec3)
   elif type == GeomType.BOX:
     return grad_box(p, attr_vec3)
   elif type == GeomType.ELLIPSOID:

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

@@ -134,7 +134,10 @@ def _equality_connect(
   body_dofnum: wp.array[int],
   body_dofadr: wp.array[int],
   body_invweight0: wp.array2d[wp.vec2],
+  jnt_type: wp.array[int],
+  jnt_dofadr: wp.array[int],
   dof_bodyid: wp.array[int],
+  dof_jntid: wp.array[int],
   dof_parentid: wp.array[int],
   site_bodyid: wp.array[int],
   eq_obj1id: wp.array[int],
@@ -154,6 +157,9 @@ def _equality_connect(
   site_xpos_in: wp.array2d[wp.vec3],
   subtree_com_in: wp.array2d[wp.vec3],
   cdof_in: wp.array2d[wp.spatial_vector],
+  cvel_in: wp.array2d[wp.spatial_vector],
+  cdof_dot_in: wp.array2d[wp.spatial_vector],
+  subtree_linvel_in: wp.array2d[wp.vec3],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
@@ -214,6 +220,7 @@ def _equality_connect(
 
   # compute Jacobian difference (opposite of contact: 0 - 1)
   Jqvel = wp.vec3f(0.0, 0.0, 0.0)
+  Jdotv = wp.vec3f(0.0, 0.0, 0.0)
 
   if is_sparse:
     # TODO(team): pre-compute number of non-zeros
@@ -282,6 +289,42 @@ def _equality_connect(
       )
       j1mj2 = jacp1 - jacp2
 
+      jacp1_dot, _ = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos1,
+        body1,
+        da,
+        worldid,
+      )
+      jacp2_dot, _ = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos2,
+        body2,
+        da,
+        worldid,
+      )
+      j1mj2_dot = jacp1_dot - jacp2_dot
+
       sparseid0 = rowadr + nnz
       sparseid1 = rowadr + rownnz + nnz
       sparseid2 = rowadr + 2 * rownnz + nnz
@@ -294,7 +337,9 @@ def _equality_connect(
       efc_J_out[worldid, 0, sparseid1] = j1mj2[1]
       efc_J_out[worldid, 0, sparseid2] = j1mj2[2]
 
-      Jqvel += j1mj2 * qvel_in[worldid, da]
+      qvel = qvel_in[worldid, da]
+      Jqvel += j1mj2 * qvel
+      Jdotv += j1mj2_dot * qvel
 
       nnz += 1
   else:
@@ -326,11 +371,49 @@ def _equality_connect(
       )
       j1mj2 = jacp1 - jacp2
 
+      jacp1_dot, _ = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos1,
+        body1,
+        dofid,
+        worldid,
+      )
+      jacp2_dot, _ = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos2,
+        body2,
+        dofid,
+        worldid,
+      )
+      j1mj2_dot = jacp1_dot - jacp2_dot
+
       efc_J_out[worldid, efcid0, dofid] = j1mj2[0]
       efc_J_out[worldid, efcid1, dofid] = j1mj2[1]
       efc_J_out[worldid, efcid2, dofid] = j1mj2[2]
 
-      Jqvel += j1mj2 * qvel_in[worldid, dofid]
+      qvel = qvel_in[worldid, dofid]
+      Jqvel += j1mj2 * qvel
+      Jdotv += j1mj2_dot * qvel
 
   body_invweight0_id = worldid % body_invweight0.shape[0]
   invweight = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0]
@@ -368,6 +451,8 @@ def _equality_connect(
       efc_frictionloss_out,
     )
 
+    efc_aref_out[worldid, efcidi] -= Jdotv[i]
+
 
 @wp.kernel
 def _equality_joint(
@@ -812,7 +897,10 @@ def _equality_weld(
   body_dofnum: wp.array[int],
   body_dofadr: wp.array[int],
   body_invweight0: wp.array2d[wp.vec2],
+  jnt_type: wp.array[int],
+  jnt_dofadr: wp.array[int],
   dof_bodyid: wp.array[int],
+  dof_jntid: wp.array[int],
   dof_parentid: wp.array[int],
   site_bodyid: wp.array[int],
   site_quat: wp.array2d[wp.quat],
@@ -834,6 +922,9 @@ def _equality_weld(
   site_xpos_in: wp.array2d[wp.vec3],
   subtree_com_in: wp.array2d[wp.vec3],
   cdof_in: wp.array2d[wp.spatial_vector],
+  cvel_in: wp.array2d[wp.spatial_vector],
+  cdof_dot_in: wp.array2d[wp.spatial_vector],
+  subtree_linvel_in: wp.array2d[wp.vec3],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
@@ -903,9 +994,43 @@ def _equality_weld(
     quat = math.mul_quat(xquat_in[worldid, body1], relpose)
     quat1 = math.quat_inv(xquat_in[worldid, body2])
 
+    # quat1 = quat_inv(xquat_in[worldid, body2])
+    q2 = xquat_in[worldid, body2]
+    quat1 = wp.quat(q2[0], -q2[1], -q2[2], -q2[3])
+
+  # compute rotational Jdotv helper quaternions
+  omega1 = wp.spatial_top(cvel_in[worldid, body1])
+  omega2 = wp.spatial_top(cvel_in[worldid, body2])
+  domega = omega1 - omega2
+
+  omega1_q = wp.quat(0.0, omega1[0], omega1[1], omega1[2])
+  omega2_q = wp.quat(0.0, omega2[0], omega2[1], omega2[2])
+  domega_q = wp.quat(0.0, domega[0], domega[1], domega[2])
+
+  if is_site:
+    qdot0r = math.mul_quat(omega1_q, quat) * 0.5
+    qfull1 = math.mul_quat(xquat_in[worldid, body2], site_quat[site_quat_id, obj2id])
+    qdot1 = math.mul_quat(omega2_q, qfull1) * 0.5
+
+    negqdot1 = wp.quat(-qdot1[0], -qdot1[1], -qdot1[2], -qdot1[3])
+    negq1 = wp.quat(qfull1[0], -qfull1[1], -qfull1[2], -qfull1[3])
+
+  else:
+    # qdot0 = mul_quat(xquat_in[worldid, body1], omega1_q) * 0.5
+    u7 = xquat_in[worldid, body1]
+    qdot0 = math.mul_quat(omega1_q, xquat_in[worldid, body1]) * 0.5
+    qdot0r = math.mul_quat(qdot0, relpose)
+    q1_non_site = xquat_in[worldid, body2]
+    qdot1 = math.mul_quat(omega2_q, q1_non_site) * 0.5
+
+    negqdot1 = wp.quat(-qdot1[0], -qdot1[1], -qdot1[2], -qdot1[3])
+    negq1 = wp.quat(q1_non_site[0], -q1_non_site[1], -q1_non_site[2], -q1_non_site[3])
+
   # compute Jacobian difference (opposite of contact: 0 - 1)
   Jqvelp = wp.vec3f(0.0, 0.0, 0.0)
   Jqvelr = wp.vec3f(0.0, 0.0, 0.0)
+  Jdotv_p = wp.vec3f(0.0, 0.0, 0.0)
+  Jdotv_r0 = wp.vec3f(0.0, 0.0, 0.0)
 
   if is_sparse:
     # TODO(team): pre-compute number of non-zeros
@@ -985,6 +1110,44 @@ def _equality_weld(
       jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat)
       jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3])
 
+      jacp1_dot, jacr1_dot = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos1,
+        body1,
+        da,
+        worldid,
+      )
+      jacp2_dot, jacr2_dot = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos2,
+        body2,
+        da,
+        worldid,
+      )
+
+      jacdifp_dot = jacp1_dot - jacp2_dot
+      jacdifr_dot = jacr1_dot - jacr2_dot
+
       sparseid0 = rowadr + nnz
       sparseid1 = rowadr + rownnz + nnz
       sparseid2 = rowadr + 2 * rownnz + nnz
@@ -1008,6 +1171,8 @@ def _equality_weld(
 
       Jqvelp += jacdifp * qvel_in[worldid, da]
       Jqvelr += jacdifr * qvel_in[worldid, da]
+      Jdotv_p += jacdifp_dot * qvel_in[worldid, da]
+      Jdotv_r0 += jacdifr_dot * qvel_in[worldid, da]
 
       nnz += 1
   else:
@@ -1047,12 +1212,52 @@ def _equality_weld(
       jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat)
       jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3])
 
+      jacp1_dot, jacr1_dot = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos1,
+        body1,
+        dofid,
+        worldid,
+      )
+      jacp2_dot, jacr2_dot = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos2,
+        body2,
+        dofid,
+        worldid,
+      )
+
+      jacdifp_dot = jacp1_dot - jacp2_dot
+      jacdifr_dot = jacr1_dot - jacr2_dot
+
       efc_J_out[worldid, efcid3, dofid] = jacdifr[0]
       efc_J_out[worldid, efcid4, dofid] = jacdifr[1]
       efc_J_out[worldid, efcid5, dofid] = jacdifr[2]
 
       Jqvelp += jacdifp * qvel_in[worldid, dofid]
       Jqvelr += jacdifr * qvel_in[worldid, dofid]
+      Jdotv_p += jacdifp_dot * qvel_in[worldid, dofid]
+      Jdotv_r0 += jacdifr_dot * qvel_in[worldid, dofid]
 
   # error is difference in global position and orientation
   cpos = pos1 - pos2
@@ -1070,6 +1275,22 @@ def _equality_weld(
 
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
 
+  djrdv_q = wp.quat(0.0, Jdotv_r0[0], Jdotv_r0[1], Jdotv_r0[2])
+
+  # Term 1: negqdot1 * domega * q0r
+  t1a = math.mul_quat(negqdot1, domega_q)
+  t1 = math.mul_quat(t1a, quat)
+
+  # Term 2: negq1 * djrdv * q0r
+  t2a = math.mul_quat(negq1, djrdv_q)
+  t2 = math.mul_quat(t2a, quat)
+
+  # Term 3: negq1 * domega * qdot0r
+  t3a = math.mul_quat(negq1, domega_q)
+  t3 = math.mul_quat(t3a, qdot0r)
+
+  Jdotv_r = wp.vec3(t1[1] + t2[1] + t3[1], t1[2] + t2[2] + t3[2], t1[3] + t2[3] + t3[3]) * 0.5 * torquescale
+
   for i in range(3):
     _efc_row(
       opt_disableflags,
@@ -1096,6 +1317,8 @@ def _equality_weld(
       efc_frictionloss_out,
     )
 
+    efc_aref_out[worldid, efcid + i] -= Jdotv_p[i]
+
   invweight_r = body_invweight0[body_invweight0_id, body1][1] + body_invweight0[body_invweight0_id, body2][1]
 
   for i in range(3):
@@ -1124,6 +1347,8 @@ def _equality_weld(
       efc_frictionloss_out,
     )
 
+    efc_aref_out[worldid, efcid + 3 + i] -= Jdotv_r[i]
+
 
 @wp.kernel
 def _friction_dof(
@@ -1680,6 +1905,7 @@ def _limit_tendon(
     )
 
 
+@cache_kernel
 def _efc_contact_init(cone_type: types.ConeType, is_sparse: bool):
   IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
   IS_SPARSE = is_sparse
@@ -1803,6 +2029,7 @@ def _efc_contact_init(cone_type: types.ConeType, is_sparse: bool):
   return kernel
 
 
+@cache_kernel
 def _efc_contact_jac_sparse(cone_type: types.ConeType):
   IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
 
@@ -1973,6 +2200,7 @@ def _efc_contact_jac_sparse(cone_type: types.ConeType):
   return kernel
 
 
+@cache_kernel
 def _efc_contact_jac_dense(tile_size: int, cone_type: types.ConeType):
   TILE_SIZE = tile_size
   IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
@@ -2117,6 +2345,7 @@ def _efc_contact_jac_dense(tile_size: int, cone_type: types.ConeType):
   return kernel
 
 
+@cache_kernel
 def _efc_contact_update(cone_type: types.ConeType):
   IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
 
@@ -2297,7 +2526,10 @@ def make_constraint(m: types.Model, d: types.Data):
           m.body_dofnum,
           m.body_dofadr,
           m.body_invweight0,
+          m.jnt_type,
+          m.jnt_dofadr,
           m.dof_bodyid,
+          m.dof_jntid,
           m.dof_parentid,
           m.site_bodyid,
           m.eq_obj1id,
@@ -2316,6 +2548,9 @@ def make_constraint(m: types.Model, d: types.Data):
           d.site_xpos,
           d.subtree_com,
           d.cdof,
+          d.cvel,
+          d.cdof_dot,
+          d.subtree_linvel,
           d.njmax,
           d.njmax_nnz,
         ],
@@ -2351,7 +2586,10 @@ def make_constraint(m: types.Model, d: types.Data):
           m.body_dofnum,
           m.body_dofadr,
           m.body_invweight0,
+          m.jnt_type,
+          m.jnt_dofadr,
           m.dof_bodyid,
+          m.dof_jntid,
           m.dof_parentid,
           m.site_bodyid,
           m.site_quat,
@@ -2372,6 +2610,9 @@ def make_constraint(m: types.Model, d: types.Data):
           d.site_xpos,
           d.subtree_com,
           d.cdof,
+          d.cvel,
+          d.cdof_dot,
+          d.subtree_linvel,
           d.njmax,
           d.njmax_nnz,
         ],

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

@@ -15,6 +15,7 @@
 
 import warp as wp
 
+from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import util_misc
 from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
@@ -24,6 +25,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import DynType
 from mujoco.mjx.third_party.mujoco_warp._src.types import GainType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
@@ -60,13 +62,47 @@ def _qderiv_actuator_passive_vel(
   actuator_gainprm_id = worldid % actuator_gainprm.shape[0]
   actuator_biasprm_id = worldid % actuator_biasprm.shape[0]
 
+  bias = float(0.0)
+
   if actuator_gaintype[actid] == GainType.AFFINE:
     gain = actuator_gainprm[actuator_gainprm_id, actid][2]
+  elif actuator_gaintype[actid] == GainType.DCMOTOR:
+    gain = 0.0
+    dynprm = actuator_dynprm[worldid % actuator_dynprm.shape[0], actid]
+    gainprm = actuator_gainprm[actuator_gainprm_id, actid]
+    te = dynprm[0]
+
+    # controller velocity derivative: dV/dω
+    input_mode = int(gainprm[8])
+    dVdw = 0.0
+    if input_mode == 1:
+      dVdw = -gainprm[6]  # position: -kd
+    elif input_mode == 2:
+      dVdw = -gainprm[4]  # velocity: -kp
+
+    if te > 0.0:
+      # stateful current with actearly: d(K*next_act)/dω
+      # includes both back-EMF (-K) and controller (dVdw) through act_dot
+      R = wp.max(MJ_MINVAL, gainprm[0])
+      K = gainprm[1]
+      s = 1.0 - wp.exp(-opt_timestep[worldid % opt_timestep.shape[0]] / te)
+      bias += K * (dVdw - K) * s / R
+    elif dVdw != 0.0:
+      # stateless: controller terms only (back-EMF handled in bias block)
+      R = wp.max(MJ_MINVAL, gainprm[0])
+      K = gainprm[1]
+      bias += K * dVdw / R
+
+    # LuGre: force includes -sigma1*z_dot, z_dot = a*z + v
+    # d(sigma1*z_dot)/dv = sigma1*(da/dv*z + 1), ignoring higher-order da/dv*z
+    sigma1 = dynprm[6]
+    if sigma1 > 0.0:
+      bias -= sigma1
   else:
     gain = 0.0
 
   if actuator_biastype[actid] == BiasType.AFFINE:
-    bias = actuator_biasprm[actuator_biasprm_id, actid][2]
+    bias += actuator_biasprm[actuator_biasprm_id, actid][2]
   elif actuator_biastype[actid] == BiasType.DCMOTOR:
     dynprm = actuator_dynprm[worldid % actuator_dynprm.shape[0], actid]
     te = dynprm[0]
@@ -86,11 +122,7 @@ def _qderiv_actuator_passive_vel(
         Ta = dynprm[4]
         R *= 1.0 + alpha * (T + Ta - T0)
 
-      bias = -K * K / wp.max(MJ_MINVAL, R)
+      bias += -K * K / wp.max(MJ_MINVAL, R)
-    else:
-      bias = 0.0
-  else:
-    bias = 0.0
 
   if bias == 0.0 and gain == 0.0:
     vel_out[worldid, actid] = 0.0
@@ -151,15 +183,15 @@ def _qderiv_actuator_passive_actuation_dense(
   actuator_moment_in: wp.array2d[float],
   # In:
   vel_in: wp.array2d[float],
-  qMi: wp.array[int],
+  Mi: wp.array[int],
-  qMj: wp.array[int],
+  Mj: wp.array[int],
   # Out:
   qDeriv_out: wp.array3d[float],
 ):
   worldid, elemid = wp.tid()
 
-  dofiid = qMi[elemid]
+  dofiid = Mi[elemid]
-  dofjid = qMj[elemid]
+  dofjid = Mj[elemid]
   qderiv_contrib = float(0.0)
   for actid in range(nu):
     vel = vel_in[worldid, actid]
@@ -195,8 +227,7 @@ def _qderiv_actuator_passive_actuation_dense(
 @wp.kernel
 def _qderiv_actuator_passive_actuation_sparse(
   # Model:
-  M_rownnz: wp.array[int],
+  M_elemid: wp.array2d[int],
-  M_rowadr: wp.array[int],
   # Data in:
   moment_rownnz_in: wp.array2d[int],
   moment_rowadr_in: wp.array2d[int],
@@ -204,7 +235,6 @@ def _qderiv_actuator_passive_actuation_sparse(
   actuator_moment_in: wp.array2d[float],
   # In:
   vel_in: wp.array2d[float],
-  qMj: wp.array[int],
   # Out:
   qDeriv_out: wp.array3d[float],
 ):
@@ -231,19 +261,10 @@ def _qderiv_actuator_passive_actuation_sparse(
         continue
       dofj = moment_colind_in[worldid, rowadrj]
 
+      elemid = M_elemid[dofi, dofj]
+      if elemid >= 0:
         contrib = moment_i * moment_j * vel
-
+        wp.atomic_add(qDeriv_out[worldid, 0], elemid, contrib)
-      # Search the corresponding elemid
-      # TODO: This could be precalculated for improved performance
-      row = dofi
-      col = dofj
-      row_startk = M_rowadr[row] - 1
-      row_nnz = M_rownnz[row]
-      for k in range(row_nnz):
-        row_startk += 1
-        if qMj[row_startk] == col:
-          wp.atomic_add(qDeriv_out[worldid, 0], row_startk, contrib)
-          break
 
 
 @wp.kernel
@@ -254,23 +275,29 @@ def _qderiv_actuator_passive(
   dof_damping: wp.array2d[float],
   dof_dampingpoly: wp.array2d[wp.vec2],
   is_sparse: bool,
+  M_elemid: wp.array2d[int],
   # Data in:
   qvel_in: wp.array2d[float],
-  qM_in: wp.array3d[float],
+  M_in: wp.array3d[float],
   # In:
-  qMi: wp.array[int],
+  Mi: wp.array[int],
-  qMj: wp.array[int],
+  Mj: wp.array[int],
   qDeriv_in: wp.array3d[float],
   # Out:
   qDeriv_out: wp.array3d[float],
 ):
   worldid, elemid = wp.tid()
 
-  dofiid = qMi[elemid]
+  dofiid = Mi[elemid]
-  dofjid = qMj[elemid]
+  dofjid = Mj[elemid]
+
+  madr = M_elemid[dofiid, dofjid]
 
   if is_sparse:
-    qderiv = qDeriv_in[worldid, 0, elemid]
+    if madr >= 0:
+      qderiv = qDeriv_in[worldid, 0, madr]
+    else:
+      qderiv = 0.0
   else:
     qderiv = qDeriv_in[worldid, dofiid, dofjid]
 
@@ -283,12 +310,13 @@ def _qderiv_actuator_passive(
   qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
 
   if is_sparse:
-    qDeriv_out[worldid, 0, elemid] = qM_in[worldid, 0, elemid] - qderiv
+    if madr >= 0:
+      qDeriv_out[worldid, 0, madr] = M_in[worldid, 0, madr] - qderiv
   else:
-    qM = qM_in[worldid, dofiid, dofjid] - qderiv
+    M = M_in[worldid, dofiid, dofjid] - qderiv
-    qDeriv_out[worldid, dofiid, dofjid] = qM
+    qDeriv_out[worldid, dofiid, dofjid] = M
     if dofiid != dofjid:
-      qDeriv_out[worldid, dofjid, dofiid] = qM
+      qDeriv_out[worldid, dofjid, dofiid] = M
 
 
 # TODO(team): improve performance with tile operations?
@@ -303,18 +331,19 @@ def _qderiv_tendon_damping(
   tendon_damping: wp.array2d[float],
   tendon_dampingpoly: wp.array2d[wp.vec2],
   is_sparse: bool,
+  M_elemid: wp.array2d[int],
   # Data in:
   ten_J_in: wp.array2d[float],
   ten_velocity_in: wp.array2d[float],
   # In:
-  qMi: wp.array[int],
+  Mi: wp.array[int],
-  qMj: wp.array[int],
+  Mj: wp.array[int],
   # Out:
   qDeriv_out: wp.array3d[float],
 ):
   worldid, elemid = wp.tid()
-  dofiid = qMi[elemid]
+  dofiid = Mi[elemid]
-  dofjid = qMj[elemid]
+  dofjid = Mj[elemid]
 
   qderiv = float(0.0)
   tendon_damping_id = worldid % tendon_damping.shape[0]
@@ -343,14 +372,283 @@ def _qderiv_tendon_damping(
 
   qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
 
+  madr = M_elemid[dofiid, dofjid]
+
   if is_sparse:
-    qDeriv_out[worldid, 0, elemid] -= qderiv
+    if madr >= 0:
+      qDeriv_out[worldid, 0, madr] -= qderiv
   else:
     qDeriv_out[worldid, dofiid, dofjid] -= qderiv
     if dofiid != dofjid:
       qDeriv_out[worldid, dofjid, dofiid] -= qderiv
 
 
+@wp.kernel
+def deriv_rne_cvel_cdof_dot(
+  # Model:
+  body_parentid: wp.array[int],
+  body_jntnum: wp.array[int],
+  body_jntadr: wp.array[int],
+  body_dofadr: wp.array[int],
+  jnt_type: wp.array[int],
+  # Data in:
+  cdof_in: wp.array2d[wp.spatial_vector],
+  # In:
+  body_tree_: wp.array[int],
+  # Out:
+  Dcvel_out: wp.array3d[wp.spatial_vector],
+  Dcdof_dot_out: wp.array3d[wp.spatial_vector],
+):
+  """Forward pass: compute d(cvel)/d(qvel_k) and d(cdof_dot)/d(qvel_k).
+
+  Mirrors the accumulation order of comvel for each joint type.
+
+  Dcdof_dot for rotation DOFs of free joints (dofid+0..2) is zero because the
+  forward pass sets cdof_dot[dofid+0..2] = 0.  The Dcdof_dot array is
+  zero-initialized so no explicit write is needed.
+  """
+  worldid, nodeid, dofid = wp.tid()
+  bodyid = body_tree_[nodeid]
+  dofadr = body_dofadr[bodyid]
+  jntid = body_jntadr[bodyid]
+  jntnum = body_jntnum[bodyid]
+  pid = body_parentid[bodyid]
+
+  cdof = cdof_in[worldid]
+
+  # Initialize from parent
+  cvel_k = Dcvel_out[worldid, pid, dofid]
+
+  if jntnum == 0:
+    Dcvel_out[worldid, bodyid, dofid] = cvel_k
+    return
+
+  dof_i = dofadr
+
+  for j in range(jntid, jntid + jntnum):
+    jnttype = jnt_type[j]
+
+    if jnttype == 0:  # FREE
+      # rotation DOFs (dof_i+0..2) contribute to cvel
+      if dofid >= dof_i and dofid < dof_i + 3:
+        cvel_k += cdof[dofid]
+
+      # cdof_dot for rotation DOFs is zero (set in forward kinematics),
+      # so Dcdof_dot for rotation DOFs is zero (from wp.zeros init)
+
+      # derivative of cdof_dot for translation DOFs 3,4,5
+      Dcdof_dot_out[worldid, dof_i + 3, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 3])
+      Dcdof_dot_out[worldid, dof_i + 4, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 4])
+      Dcdof_dot_out[worldid, dof_i + 5, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 5])
+
+      # translation DOFs (dof_i+3..5) contribute to cvel
+      if dofid >= dof_i + 3 and dofid < dof_i + 6:
+        cvel_k += cdof[dofid]
+
+      dof_i += 6
+
+    elif jnttype == 1:  # BALL
+      Dcdof_dot_out[worldid, dof_i + 0, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 0])
+      Dcdof_dot_out[worldid, dof_i + 1, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 1])
+      Dcdof_dot_out[worldid, dof_i + 2, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 2])
+
+      if dofid >= dof_i and dofid < dof_i + 3:
+        cvel_k += cdof[dofid]
+
+      dof_i += 3
+    else:  # HINGE or SLIDE
+      Dcdof_dot_out[worldid, dof_i, dofid] = math.motion_cross(cvel_k, cdof[dof_i])
+
+      if dofid == dof_i:
+        cvel_k += cdof[dof_i]
+
+      dof_i += 1
+
+  Dcvel_out[worldid, bodyid, dofid] = cvel_k
+
+
+@wp.kernel
+def deriv_rne_cacc_cfrcbody_forward(
+  # Model:
+  body_parentid: wp.array[int],
+  body_dofnum: wp.array[int],
+  body_dofadr: wp.array[int],
+  # Data in:
+  qvel_in: wp.array2d[float],
+  cinert_in: wp.array2d[vec10],
+  cvel_in: wp.array2d[wp.spatial_vector],
+  cdof_dot_in: wp.array2d[wp.spatial_vector],
+  # In:
+  body_tree_: wp.array[int],
+  Dcvel_in: wp.array3d[wp.spatial_vector],
+  Dcdof_dot_in: wp.array3d[wp.spatial_vector],
+  # Out:
+  Dcacc_out: wp.array3d[wp.spatial_vector],
+  Dcfrcbody_out: wp.array3d[wp.spatial_vector],
+):
+  """Forward pass: compute d(cacc)/d(qvel_k) and d(cfrc_body)/d(qvel_k)."""
+  worldid, nodeid, dofid = wp.tid()
+  bodyid = body_tree_[nodeid]
+  dofadr = body_dofadr[bodyid]
+  dofnum = body_dofnum[bodyid]
+  pid = body_parentid[bodyid]
+
+  qvel = qvel_in[worldid]
+
+  dcacc = Dcacc_out[worldid, pid, dofid]
+
+  for j in range(dofadr, dofadr + dofnum):
+    # Term 1: d(cdof_dot * qvel)/d(qvel_k) when j == dofid
+    if j == dofid:
+      dcacc += cdof_dot_in[worldid, j]
+
+    # Term 2: cdof_dot depends on cvel which depends on qvel_k
+    dcdofdot = Dcdof_dot_in[worldid, j, dofid]
+    dcacc += dcdofdot * qvel[j]
+
+  Dcacc_out[worldid, bodyid, dofid] = dcacc
+
+  # d(cfrc_body)/d(qvel_k)
+  cinert = cinert_in[worldid, bodyid]
+  cvel = cvel_in[worldid, bodyid]
+  dcvel = Dcvel_in[worldid, bodyid, dofid]
+
+  # term1 = cinert * d(cacc)/d(qvel_k)
+  term1 = math.inert_vec(cinert, dcacc)
+
+  # term2 = d(cvel x* (cinert * cvel))/d(qvel_k)
+  cinert_cvel = math.inert_vec(cinert, cvel)
+  cinert_dcvel = math.inert_vec(cinert, dcvel)
+  term2 = math.motion_cross_force(dcvel, cinert_cvel) + math.motion_cross_force(cvel, cinert_dcvel)
+
+  Dcfrcbody_out[worldid, bodyid, dofid] = term1 + term2
+
+
+@wp.kernel
+def deriv_rne_cfrcbody_backward(
+  # Model:
+  body_parentid: wp.array[int],
+  # In:
+  body_tree_: wp.array[int],
+  # Out:
+  Dcfrcbody_out: wp.array3d[wp.spatial_vector],
+):
+  """Backward pass: accumulate d(cfrc_body) from children to parents."""
+  worldid, nodeid, dofid = wp.tid()
+  bodyid = body_tree_[nodeid]
+  pid = body_parentid[bodyid]
+
+  # body_tree never contains bodyid=0 (worldbody), so pid >= 0 is always valid.
+  # Siblings at the same level may share a parent; atomic_add handles this.
+  val = Dcfrcbody_out[worldid, bodyid, dofid]
+  wp.atomic_add(Dcfrcbody_out[worldid, pid], dofid, val)
+
+
+@wp.kernel
+def deriv_rne_body2jnt_sparse(
+  # Model:
+  dof_bodyid: wp.array[int],
+  # Data in:
+  cdof_in: wp.array2d[wp.spatial_vector],
+  # In:
+  timestep: wp.array[float],
+  Di: wp.array[int],
+  Dj: wp.array[int],
+  Dcfrcbody_in: wp.array3d[wp.spatial_vector],
+  flg_subtract: bool,
+  # Out:
+  qDeriv_out: wp.array3d[float],
+):
+  """Project body-space RNE derivatives into joint-space qDeriv (sparse)."""
+  worldid, elemid = wp.tid()
+  dt = timestep[worldid % timestep.shape[0]]
+
+  i = Di[elemid]
+  j = Dj[elemid]
+
+  body_i = dof_bodyid[i]
+  dcfrc = Dcfrcbody_in[worldid, body_i, j]
+  term = wp.dot(cdof_in[worldid, i], dcfrc)
+
+  if flg_subtract:
+    wp.atomic_sub(qDeriv_out[worldid, 0], elemid, dt * term)
+  else:
+    wp.atomic_add(qDeriv_out[worldid, 0], elemid, dt * term)
+
+
+def deriv_rne_vel(m: Model, d: Data, out: wp.array3d[float], flg_subtract: bool = False):
+  """Compute RNE velocity derivatives and add/subtract from the output.
+
+  Implements the analytical derivative of inverse-dynamics Coriolis/centrifugal
+  forces with respect to joint velocities.
+
+  Args:
+    m: The model (device).
+    d: The data (device).
+    out: D-structure output array (nworld, 1, nD) to accumulate RNE terms into.
+    flg_subtract: If True, subtract the RNE derivatives from output instead of adding them.
+  """
+  # TODO(team): consider caching these allocations
+  Dcvel = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
+  Dcdof_dot = wp.zeros((d.nworld, m.nv, m.nv), dtype=wp.spatial_vector)
+  Dcacc = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
+  Dcfrcbody = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
+
+  # Forward pass 1: compute Dcvel and Dcdof_dot
+  for body_tree in m.body_tree:
+    wp.launch(
+      deriv_rne_cvel_cdof_dot,
+      dim=(d.nworld, body_tree.size, m.nv),
+      inputs=[
+        m.body_parentid,
+        m.body_jntnum,
+        m.body_jntadr,
+        m.body_dofadr,
+        m.jnt_type,
+        d.cdof,
+        body_tree,
+      ],
+      outputs=[Dcvel, Dcdof_dot],
+    )
+
+  # Forward pass 2: compute Dcacc and Dcfrcbody
+  for body_tree in m.body_tree:
+    wp.launch(
+      deriv_rne_cacc_cfrcbody_forward,
+      dim=(d.nworld, body_tree.size, m.nv),
+      inputs=[
+        m.body_parentid,
+        m.body_dofnum,
+        m.body_dofadr,
+        d.qvel,
+        d.cinert,
+        d.cvel,
+        d.cdof_dot,
+        body_tree,
+        Dcvel,
+        Dcdof_dot,
+      ],
+      outputs=[Dcacc, Dcfrcbody],
+    )
+
+  # Backward pass: accumulate Dcfrcbody from children to parents
+  for body_tree in reversed(m.body_tree):
+    wp.launch(
+      deriv_rne_cfrcbody_backward,
+      dim=(d.nworld, body_tree.size, m.nv),
+      inputs=[m.body_parentid, body_tree],
+      outputs=[Dcfrcbody],
+    )
+
+  # Project body-space derivatives into joint-space qDeriv (always sparse D-structure)
+  wp.launch(
+    deriv_rne_body2jnt_sparse,
+    dim=(d.nworld, m.qD_fullm_i.size),
+    inputs=[m.dof_bodyid, d.cdof, m.opt.timestep, m.qD_fullm_i, m.qD_fullm_j, Dcfrcbody, flg_subtract],
+    outputs=[out],
+  )
+
+
 @event_scope
 def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
   """Analytical derivative of smooth forces w.r.t. velocities.
@@ -358,12 +656,10 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
   Args:
     m: The model containing kinematic and dynamic information (device).
     d: The data object containing the current state and output arrays (device).
-    out: qM - dt * qDeriv (derivatives of smooth forces w.r.t velocities).
+    out: M - dt * qDeriv (derivatives of smooth forces w.r.t velocities).
   """
-  qMi = m.qM_fullm_i
+  Mi = m.M_fullm_i
-  qMj = m.qM_fullm_j
+  Mj = m.M_fullm_j
-
-  # TODO(team): implicit requires different sparsity structure
 
   if ~(m.opt.disableflags & (DisableBit.ACTUATION | DisableBit.DAMPER)):
     # TODO(team): only clear elements not set by _qderiv_actuator_passive
@@ -399,41 +695,49 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
         wp.launch(
           _qderiv_actuator_passive_actuation_sparse,
           dim=(d.nworld, m.nu),
-          inputs=[m.M_rownnz, m.M_rowadr, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMj],
+          inputs=[
+            m.M_elemid,
+            d.moment_rownnz,
+            d.moment_rowadr,
+            d.moment_colind,
+            d.actuator_moment,
+            vel,
+          ],
           outputs=[out],
         )
       else:
         wp.launch(
           _qderiv_actuator_passive_actuation_dense,
-          dim=(d.nworld, qMi.size),
+          dim=(d.nworld, Mi.size),
-          inputs=[m.nu, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMi, qMj],
+          inputs=[m.nu, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, Mi, Mj],
           outputs=[out],
         )
     wp.launch(
       _qderiv_actuator_passive,
-      dim=(d.nworld, qMi.size),
+      dim=(d.nworld, Mi.size),
       inputs=[
         m.opt.timestep,
         m.opt.disableflags,
         m.dof_damping,
         m.dof_dampingpoly,
         m.is_sparse,
+        m.M_elemid,
         d.qvel,
-        d.qM,
+        d.M,
-        qMi,
+        Mi,
-        qMj,
+        Mj,
         out,
       ],
       outputs=[out],
     )
   else:
-    # TODO(team): directly utilize qM for these settings
+    # TODO(team): directly utilize M for these settings
-    wp.copy(out, d.qM)
+    wp.copy(out, d.M)
 
   if not (m.opt.disableflags & DisableBit.DAMPER):
     wp.launch(
       _qderiv_tendon_damping,
-      dim=(d.nworld, qMi.size),
+      dim=(d.nworld, Mi.size),
       inputs=[
         m.ntendon,
         m.opt.timestep,
@@ -443,12 +747,11 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
         m.tendon_damping,
         m.tendon_dampingpoly,
         m.is_sparse,
+        m.M_elemid,
         d.ten_J,
         d.ten_velocity,
-        qMi,
+        Mi,
-        qMj,
+        Mj,
       ],
       outputs=[out],
     )
-
-  # TODO(team): rne derivative

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

@@ -20,12 +20,14 @@ import warp as wp
 from mujoco.mjx.third_party.mujoco_warp._src import collision_driver
 from mujoco.mjx.third_party.mujoco_warp._src import constraint
 from mujoco.mjx.third_party.mujoco_warp._src import derivative
+from mujoco.mjx.third_party.mujoco_warp._src import history
 from mujoco.mjx.third_party.mujoco_warp._src import island
 from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import passive
 from mujoco.mjx.third_party.mujoco_warp._src import sensor
 from mujoco.mjx.third_party.mujoco_warp._src import smooth
 from mujoco.mjx.third_party.mujoco_warp._src import solver
+from mujoco.mjx.third_party.mujoco_warp._src import types
 from mujoco.mjx.third_party.mujoco_warp._src import util_misc
 from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
 from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate
@@ -302,6 +304,9 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
     outputs=[d.qpos],
   )
 
+  # advance history buffers before time advance
+  history.insert_ctrl_history(m, d)
+
   wp.launch(
     _next_time,
     dim=d.nworld,
@@ -346,17 +351,18 @@ def _compute_damping_deriv(
 def _euler_damp_qfrc_sparse(
   # Model:
   opt_timestep: wp.array[float],
-  dof_Madr: wp.array[int],
+  M_rownnz: wp.array[int],
+  M_rowadr: wp.array[int],
   # In:
   damp_deriv: wp.array2d[float],
   # Out:
-  qM_integration_out: wp.array3d[float],
+  M_integration_out: wp.array3d[float],
 ):
   worldid, tid = wp.tid()
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
 
-  adr = dof_Madr[tid]
+  adr = M_rowadr[tid] + M_rownnz[tid] - 1
-  qM_integration_out[worldid, 0, adr] += timestep * damp_deriv[worldid, tid]
+  M_integration_out[worldid, 0, adr] += timestep * damp_deriv[worldid, tid]
 
 
 @cache_kernel
@@ -366,7 +372,7 @@ def _tile_euler_dense(tile: TileSet):
     # Model:
     opt_timestep: wp.array[float],
     # Data in:
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     efc_Ma_in: wp.array2d[float],
     # In:
     damp_deriv: wp.array2d[float],
@@ -379,14 +385,14 @@ def _tile_euler_dense(tile: TileSet):
     TILE_SIZE = wp.static(tile.size)
 
     dofid = adr_in[nodeid]
-    M_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
+    M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
     damping_tile = wp.tile_load(damp_deriv[worldid], shape=(TILE_SIZE,), offset=(dofid,))
     damping_scaled = damping_tile * timestep
     qm_integration_tile = wp.tile_diag_add(M_tile, damping_scaled)
 
     Ma_tile = wp.tile_load(efc_Ma_in[worldid], shape=(TILE_SIZE,), offset=(dofid,))
-    L_tile = wp.tile_cholesky(qm_integration_tile)
+    L_tile = wp.tile_cholesky(qm_integration_tile, fill_mode="upper")
-    qacc_tile = wp.tile_cholesky_solve(L_tile, Ma_tile)
+    qacc_tile = wp.tile_cholesky_solve(L_tile, Ma_tile, fill_mode="upper")
     wp.tile_store(qacc_out[worldid], qacc_tile, offset=(dofid))
 
   return euler_dense
@@ -409,22 +415,22 @@ def euler(m: Model, d: Data):
     )
 
     if m.is_sparse:
-      qM = wp.clone(d.qM)
+      M = wp.clone(d.M)
       qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
       qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float)
       wp.launch(
         _euler_damp_qfrc_sparse,
         dim=(d.nworld, m.nv),
-        inputs=[m.opt.timestep, m.dof_Madr, damp_deriv],
+        inputs=[m.opt.timestep, m.M_rownnz, m.M_rowadr, damp_deriv],
-        outputs=[qM],
+        outputs=[M],
       )
-      smooth.factor_solve_i(m, d, qM, qLD, qLDiagInv, qacc, d.efc.Ma)
+      smooth.factor_solve_i(m, d, M, qLD, qLDiagInv, qacc, d.efc.Ma)
     else:
-      for tile in m.qM_tiles:
+      for tile in m.M_tiles:
         wp.launch_tiled(
           _tile_euler_dense(tile),
           dim=(d.nworld, tile.adr.size),
-          inputs=[m.opt.timestep, d.qM, d.efc.Ma, damp_deriv, tile.adr],
+          inputs=[m.opt.timestep, d.M, d.efc.Ma, damp_deriv, tile.adr],
           outputs=[qacc],
           block_dim=m.block_dim.euler_dense,
         )
@@ -573,16 +579,62 @@ def rungekutta4(m: Model, d: Data):
   _advance(m, d, qacc_rk, qvel_rk)
 
 
+@wp.kernel
+def _map_m2d(
+  # Model:
+  mapM2D: wp.array[int],
+  is_sparse: bool,
+  # In:
+  qDi: wp.array[int],
+  qDj: wp.array[int],
+  qH_M: wp.array3d[float],
+  # Data out:
+  qLU_out: wp.array3d[float],
+):
+  worldid, elemid = wp.tid()
+  if is_sparse:
+    m_idx = mapM2D[elemid]
+    if m_idx >= 0:
+      qLU_out[worldid, 0, elemid] = qH_M[worldid, 0, m_idx]
+    else:
+      qLU_out[worldid, 0, elemid] = 0.0
+  else:
+    i = qDi[elemid]
+    j = qDj[elemid]
+    qLU_out[worldid, 0, elemid] = qH_M[worldid, i, j]
+
+
 @event_scope
 def implicit(m: Model, d: Data):
   """Integrates fully implicit in velocity."""
-  if ~(m.opt.disableflags | ~(DisableBit.ACTUATION | DisableBit.SPRING | DisableBit.DAMPER)):
+  if m.opt.integrator == IntegratorType.IMPLICIT:
+    qH_M = wp.empty(d.M.shape, dtype=float)
+
+    # 1. Compute M - dt * qDeriv_smooth in M-structure
+    derivative.deriv_smooth_vel(m, d, qH_M)
+
+    # 2. Map M-structure to D-structure
+    wp.launch(
+      _map_m2d,
+      dim=(d.nworld, m.nD),
+      inputs=[m.mapM2D, m.is_sparse, m.qD_fullm_i, m.qD_fullm_j, qH_M],
+      outputs=[d.qLU],
+    )
+
+    # 3. Compute RNE derivatives, scale by timestep, and subtract in-place from qLU
+    derivative.deriv_rne_vel(m, d, d.qLU, flg_subtract=True)
+
+    # 4. Factorize and solve: qacc = qLU \ Ma
+    qacc = wp.empty((d.nworld, m.nv), dtype=float)
+    smooth.factor_solve_lu(m, d, d.qLU, qacc, d.efc.Ma)
+    _advance(m, d, qacc)
+  elif ~(m.opt.disableflags | ~(DisableBit.ACTUATION | DisableBit.SPRING | DisableBit.DAMPER)):
     if m.is_sparse:
-      qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
+      qDeriv = wp.empty((d.nworld, 1, m.nC), dtype=float)
       qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
     else:
-      qDeriv = wp.empty(d.qM.shape, dtype=float)
+      qDeriv = wp.empty(d.M.shape, dtype=float)
-      qLD = wp.empty(d.qM.shape, dtype=float)
+      qLD = wp.empty(d.M.shape, dtype=float)
     qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float)
     derivative.deriv_smooth_vel(m, d, qDeriv)
     qacc = wp.empty((d.nworld, m.nv), dtype=float)
@@ -613,8 +665,7 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
   if m.opt.run_collision_detection:
     collision_driver.collision(m, d)
   constraint.make_constraint(m, d)
-  # TODO(team): remove False after island features are more complete
+  if m.ntree > 1 and not (m.opt.disableflags & types.DisableBit.ISLAND):
-  if False and not (m.opt.disableflags & DisableBit.ISLAND):
     island.island(m, d)
   smooth.transmission(m, d)
 
@@ -1098,6 +1149,13 @@ def fwd_actuation(m: Model, d: Data):
     d.actuator_force.zero_()
     return
 
+  # read delayed ctrl (or direct copy if no delay)
+  if m.nhistory > 0:
+    ctrl = wp.empty((d.nworld, m.nu), dtype=float)
+    history.read_ctrl_delayed(m, d, ctrl)
+  else:
+    ctrl = d.ctrl
+
   wp.launch(
     _actuator_force,
     dim=(d.nworld, m.nu),
@@ -1122,7 +1180,7 @@ def fwd_actuation(m: Model, d: Data):
       m.actuator_acc0,
       m.actuator_lengthrange,
       d.act,
-      d.ctrl,
+      ctrl,
       d.actuator_length,
       d.actuator_velocity,
       m.opt.disableflags & DisableBit.CLAMPCTRL,
@@ -1221,7 +1279,7 @@ def fwd_acceleration(m: Model, d: Data, factorize: bool = False):
   xfrc_accumulate(m, d, d.qfrc_smooth)
 
   if factorize:
-    smooth.factor_solve_i(m, d, d.qM, d.qLD, d.qLDiagInv, d.qacc_smooth, d.qfrc_smooth)
+    smooth.factor_solve_i(m, d, d.M, d.qLD, d.qLDiagInv, d.qacc_smooth, d.qfrc_smooth)
   else:
     smooth.solve_m(m, d, d.qacc_smooth, d.qfrc_smooth)
 
@@ -1266,7 +1324,7 @@ def step(m: Model, d: Data):
     euler(m, d)
   elif m.opt.integrator == IntegratorType.RK4:
     rungekutta4(m, d)
-  elif m.opt.integrator == IntegratorType.IMPLICITFAST:
+  elif m.opt.integrator in (IntegratorType.IMPLICITFAST, IntegratorType.IMPLICIT):
     implicit(m, d)
   else:
     raise NotImplementedError(f"integrator {m.opt.integrator} not implemented.")
@@ -1307,8 +1365,7 @@ def step2(m: Model, d: Data):
   sensor.sensor_acc(m, d)
 
   # integrate with Euler or implicitfast
-  # TODO(team): implicit
+  if m.opt.integrator in (IntegratorType.IMPLICITFAST, IntegratorType.IMPLICIT):
-  if m.opt.integrator == IntegratorType.IMPLICITFAST:
     implicit(m, d)
   else:
     # note: RK4 defaults to Euler

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

@@ -0,0 +1,925 @@
+# Copyright 2026 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import warp as wp
+
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import Data
+from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+
+wp.set_module_options({"enable_backward": False})
+
+
+@wp.func
+def _history_physical_index(cursor: int, n: int, logical: int) -> int:
+  """Convert logical index (0=oldest, n-1=newest) to physical index."""
+  return (cursor + 1 + logical) % n
+
+
+@wp.func
+def _history_find_index(
+  # In:
+  buf: wp.array2d[float],
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  cursor: int,
+  t: float,
+) -> int:
+  """Find logical index i such that times[i-1] < t <= times[i].
+
+  Returns 0 if t <= times[oldest], n if t > times[newest].
+  Uses circular binary search matching MuJoCo C historyFindIndex.
+  """
+  times_offset = buf_offset + 2
+
+  oldest_phys = _history_physical_index(cursor, n, 0)
+  newest_phys = _history_physical_index(cursor, n, n - 1)
+  t_oldest = buf[worldid, times_offset + oldest_phys]
+  t_newest = buf[worldid, times_offset + newest_phys]
+
+  # before or at first element
+  if t <= t_oldest:
+    return 0
+
+  # after last element
+  if t > t_newest:
+    return n
+
+  # circular binary search: find smallest logical i such that times[phys(i)] >= t
+  lo = int(0)
+  hi = int(n - 1)
+  while hi - lo > 1:
+    mid = int((lo + hi) >> 1)
+    mid_phys = _history_physical_index(cursor, n, mid)
+    if buf[worldid, times_offset + mid_phys] < t:
+      lo = mid
+    else:
+      hi = mid
+
+  return hi
+
+
+@wp.func
+def _history_read_scalar(
+  # In:
+  buf: wp.array2d[float],
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  t: float,
+  interp: int,
+) -> float:
+  """Read a scalar value from history buffer at time t.
+
+  interp: 0=zero-order-hold, 1=linear, 2=cubic (Catmull-Rom spline)
+  """
+  cursor = int(buf[worldid, buf_offset + 1])
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + n
+
+  oldest_phys = _history_physical_index(cursor, n, 0)
+  newest_phys = _history_physical_index(cursor, n, n - 1)
+  t_oldest = buf[worldid, times_offset + oldest_phys]
+  t_newest = buf[worldid, times_offset + newest_phys]
+
+  # extrapolate before oldest
+  if t <= t_oldest + 1e-6:
+    return buf[worldid, values_offset + oldest_phys]
+
+  # extrapolate after newest
+  if t >= t_newest - 1e-6:
+    return buf[worldid, values_offset + newest_phys]
+
+  # find bracketing index
+  i = _history_find_index(buf, worldid, buf_offset, n, cursor, t)
+  phys_i = _history_physical_index(cursor, n, i)
+
+  # exact match
+  if wp.abs(t - buf[worldid, times_offset + phys_i]) < 1e-6:
+    return buf[worldid, values_offset + phys_i]
+
+  phys_lo = _history_physical_index(cursor, n, i - 1)
+  phys_hi = phys_i
+
+  # zero-order hold
+  if interp == 0:
+    return buf[worldid, values_offset + phys_lo]
+
+  dt = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo]
+  alpha = (t - buf[worldid, times_offset + phys_lo]) / dt
+  v_lo = buf[worldid, values_offset + phys_lo]
+  v_hi = buf[worldid, values_offset + phys_hi]
+
+  # linear interpolation
+  if interp == 1:
+    return v_lo + alpha * (v_hi - v_lo)
+
+  # cubic spline interpolation (Catmull-Rom)
+  alpha2 = alpha * alpha
+  alpha3 = alpha2 * alpha
+  h00 = 2.0 * alpha3 - 3.0 * alpha2 + 1.0
+  h10 = alpha3 - 2.0 * alpha2 + alpha
+  h01 = -2.0 * alpha3 + 3.0 * alpha2
+  h11 = alpha3 - alpha2
+
+  # finite-differenced Catmull-Rom slopes, 0 at endpoints
+  m_lo = 0.0
+  if i > 1:
+    phys_lo_prev = _history_physical_index(cursor, n, i - 2)
+    dt_lo = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo_prev]
+    m_lo = (v_hi - buf[worldid, values_offset + phys_lo_prev]) / dt_lo
+
+  m_hi = 0.0
+  if i < n - 1:
+    phys_hi_next = _history_physical_index(cursor, n, i + 1)
+    dt_hi = buf[worldid, times_offset + phys_hi_next] - buf[worldid, times_offset + phys_lo]
+    m_hi = (buf[worldid, values_offset + phys_hi_next] - v_lo) / dt_hi
+
+  return h00 * v_lo + h10 * dt * m_lo + h01 * v_hi + h11 * dt * m_hi
+
+
+@wp.func
+def _history_read_vector(
+  # In:
+  adr: int,
+  buf: wp.array2d[float],
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  dim: int,
+  t: float,
+  interp: int,
+  # Data out:
+  sensordata_out: wp.array2d[float],
+) -> int:
+  """Read a vector value from history buffer at time t into sensordata.
+
+  Returns 1 on success (value written to sensordata).
+  interp: 0=zero-order-hold, 1=linear, 2=cubic (Catmull-Rom spline)
+  """
+  cursor = int(buf[worldid, buf_offset + 1])
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + n
+
+  oldest_phys = _history_physical_index(cursor, n, 0)
+  newest_phys = _history_physical_index(cursor, n, n - 1)
+  t_oldest = buf[worldid, times_offset + oldest_phys]
+  t_newest = buf[worldid, times_offset + newest_phys]
+
+  # extrapolate before oldest: copy oldest
+  if t <= t_oldest + 1e-6:
+    for d in range(dim):
+      sensordata_out[worldid, adr + d] = buf[worldid, values_offset + oldest_phys * dim + d]
+    return 1
+
+  # extrapolate after newest: copy newest
+  if t >= t_newest - 1e-6:
+    for d in range(dim):
+      sensordata_out[worldid, adr + d] = buf[worldid, values_offset + newest_phys * dim + d]
+    return 1
+
+  # find bracketing index
+  i = _history_find_index(buf, worldid, buf_offset, n, cursor, t)
+  phys_i = _history_physical_index(cursor, n, i)
+
+  # exact match
+  if wp.abs(t - buf[worldid, times_offset + phys_i]) < 1e-6:
+    for d in range(dim):
+      sensordata_out[worldid, adr + d] = buf[worldid, values_offset + phys_i * dim + d]
+    return 1
+
+  phys_lo = _history_physical_index(cursor, n, i - 1)
+  phys_hi = phys_i
+
+  # zero-order hold
+  if interp == 0:
+    for d in range(dim):
+      sensordata_out[worldid, adr + d] = buf[worldid, values_offset + phys_lo * dim + d]
+    return 1
+
+  dt = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo]
+  alpha = (t - buf[worldid, times_offset + phys_lo]) / dt
+
+  # linear interpolation
+  if interp == 1:
+    for d in range(dim):
+      v_lo = buf[worldid, values_offset + phys_lo * dim + d]
+      v_hi = buf[worldid, values_offset + phys_hi * dim + d]
+      sensordata_out[worldid, adr + d] = v_lo + alpha * (v_hi - v_lo)
+    return 1
+
+  # cubic spline interpolation (Catmull-Rom)
+  alpha2 = alpha * alpha
+  alpha3 = alpha2 * alpha
+  h00 = 2.0 * alpha3 - 3.0 * alpha2 + 1.0
+  h10 = alpha3 - 2.0 * alpha2 + alpha
+  h01 = -2.0 * alpha3 + 3.0 * alpha2
+  h11 = alpha3 - alpha2
+
+  for d in range(dim):
+    v_lo = buf[worldid, values_offset + phys_lo * dim + d]
+    v_hi = buf[worldid, values_offset + phys_hi * dim + d]
+
+    # finite-differenced Catmull-Rom slopes, 0 at endpoints
+    m_lo = 0.0
+    if i > 1:
+      phys_lo_prev = _history_physical_index(cursor, n, i - 2)
+      dt_lo = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo_prev]
+      m_lo = (v_hi - buf[worldid, values_offset + phys_lo_prev * dim + d]) / dt_lo
+
+    m_hi = 0.0
+    if i < n - 1:
+      phys_hi_next = _history_physical_index(cursor, n, i + 1)
+      dt_hi = buf[worldid, times_offset + phys_hi_next] - buf[worldid, times_offset + phys_lo]
+      m_hi = (buf[worldid, values_offset + phys_hi_next * dim + d] - v_lo) / dt_hi
+
+    sensordata_out[worldid, adr + d] = h00 * v_lo + h10 * dt * m_lo + h01 * v_hi + h11 * dt * m_hi
+  return 1
+
+
+@wp.func
+def _history_insert_scalar(
+  # In:
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  t: float,
+  value: float,
+  # Out:
+  buf_out: wp.array2d[float],
+):
+  """Insert a scalar value into history buffer at time t."""
+  cursor = int(buf_out[worldid, buf_offset + 1])
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + n
+
+  i = _history_find_index(buf_out, worldid, buf_offset, n, cursor, t)
+
+  # exact match
+  if i < n:
+    phys_i = _history_physical_index(cursor, n, i)
+    if wp.abs(t - buf_out[worldid, times_offset + phys_i]) < 1e-6:
+      buf_out[worldid, values_offset + phys_i] = value
+      return
+
+  # older than oldest: replace oldest
+  if i == 0:
+    oldest_phys = _history_physical_index(cursor, n, 0)
+    buf_out[worldid, times_offset + oldest_phys] = t
+    buf_out[worldid, values_offset + oldest_phys] = value
+    return
+
+  # newer than newest: advance cursor
+  if i == n:
+    cursor = (cursor + 1) % n
+    buf_out[worldid, buf_offset + 1] = float(cursor)
+    buf_out[worldid, times_offset + cursor] = t
+    buf_out[worldid, values_offset + cursor] = value
+    return
+
+  # out-of-order: shift [1, i-1] left, insert at i-1
+  for j in range(i - 1):
+    src_phys = _history_physical_index(cursor, n, j + 1)
+    dst_phys = _history_physical_index(cursor, n, j)
+    buf_out[worldid, times_offset + dst_phys] = buf_out[worldid, times_offset + src_phys]
+    buf_out[worldid, values_offset + dst_phys] = buf_out[worldid, values_offset + src_phys]
+  insert_phys = _history_physical_index(cursor, n, i - 1)
+  buf_out[worldid, times_offset + insert_phys] = t
+  buf_out[worldid, values_offset + insert_phys] = value
+
+
+@wp.func
+def _history_insert_vector(
+  # In:
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  dim: int,
+  t: float,
+  src: wp.array2d[float],
+  src_adr: int,
+  # Out:
+  buf_out: wp.array2d[float],
+):
+  """Insert a vector value from src[worldid, src_adr:src_adr+dim] into history buffer at time t."""
+  cursor = int(buf_out[worldid, buf_offset + 1])
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + n
+
+  i = _history_find_index(buf_out, worldid, buf_offset, n, cursor, t)
+
+  slot_phys = -1
+
+  # exact match
+  if i < n:
+    phys_i = _history_physical_index(cursor, n, i)
+    if wp.abs(t - buf_out[worldid, times_offset + phys_i]) < 1e-6:
+      slot_phys = phys_i
+
+  if slot_phys < 0:
+    if i == 0:
+      # older than oldest: replace oldest
+      slot_phys = _history_physical_index(cursor, n, 0)
+      buf_out[worldid, times_offset + slot_phys] = t
+    elif i == n:
+      # newer than newest: advance cursor
+      cursor = (cursor + 1) % n
+      buf_out[worldid, buf_offset + 1] = float(cursor)
+      slot_phys = cursor
+      buf_out[worldid, times_offset + slot_phys] = t
+    else:
+      # out-of-order: shift [1, i-1] left, insert at i-1
+      for j in range(i - 1):
+        src_phys = _history_physical_index(cursor, n, j + 1)
+        dst_phys = _history_physical_index(cursor, n, j)
+        buf_out[worldid, times_offset + dst_phys] = buf_out[worldid, times_offset + src_phys]
+        for d in range(dim):
+          buf_out[worldid, values_offset + dst_phys * dim + d] = buf_out[worldid, values_offset + src_phys * dim + d]
+      slot_phys = _history_physical_index(cursor, n, i - 1)
+      buf_out[worldid, times_offset + slot_phys] = t
+
+  # copy values
+  for d in range(dim):
+    buf_out[worldid, values_offset + slot_phys * dim + d] = src[worldid, src_adr + d]
+
+
+@wp.kernel
+def _read_ctrl_delayed_kernel(
+  # Model:
+  actuator_history: wp.array[wp.vec2i],
+  actuator_historyadr: wp.array[int],
+  actuator_delay: wp.array[float],
+  # Data in:
+  time_in: wp.array[float],
+  history_in: wp.array2d[float],
+  ctrl_in: wp.array2d[float],
+  # Data out:
+  ctrl_out: wp.array2d[float],
+):
+  """Read delayed ctrl for each actuator."""
+  worldid, uid = wp.tid()
+
+  hist = actuator_history[uid]
+  nsample = hist[0]
+  delay = actuator_delay[uid]
+
+  if nsample == 0 or delay == 0.0:
+    # no delay: direct copy
+    ctrl_out[worldid, uid] = ctrl_in[worldid, uid]
+  else:
+    interp = hist[1]
+    buf_offset = actuator_historyadr[uid]
+    t = time_in[worldid] - delay
+    ctrl_out[worldid, uid] = _history_read_scalar(history_in, worldid, buf_offset, nsample, t, interp)
+
+
+@wp.kernel
+def _insert_ctrl_history_kernel(
+  # Model:
+  actuator_history: wp.array[wp.vec2i],
+  actuator_historyadr: wp.array[int],
+  # Data in:
+  time_in: wp.array[float],
+  ctrl_in: wp.array2d[float],
+  # Data out:
+  history_out: wp.array2d[float],
+):
+  """Insert current ctrl into history buffers."""
+  worldid, uid = wp.tid()
+
+  hist = actuator_history[uid]
+  nsample = hist[0]
+  if nsample == 0:
+    return
+
+  buf_offset = actuator_historyadr[uid]
+  t = time_in[worldid]
+  value = ctrl_in[worldid, uid]
+  _history_insert_scalar(worldid, buf_offset, nsample, t, value, history_out)
+
+
+@wp.kernel
+def _insert_sensor_history_stage(
+  # Model:
+  sensor_dim: wp.array[int],
+  sensor_adr: wp.array[int],
+  sensor_history: wp.array[wp.vec2i],
+  sensor_historyadr: wp.array[int],
+  sensor_delay: wp.array[float],
+  sensor_interval: wp.array[wp.vec2],
+  # Data in:
+  time_in: wp.array[float],
+  sensordata_in: wp.array2d[float],
+  # In:
+  sensor_ids: wp.array[int],
+  # Data out:
+  history_out: wp.array2d[float],
+):
+  """Insert current sensor values into history buffers for specific sensor IDs."""
+  worldid, idx = wp.tid()
+  sid = sensor_ids[idx]
+
+  hist = sensor_history[sid]
+  nsample = hist[0]
+  if nsample == 0:
+    return
+
+  buf_offset = sensor_historyadr[sid]
+  dim = sensor_dim[sid]
+  interval_val = sensor_interval[sid]
+  period = interval_val[0]
+  t = time_in[worldid]
+
+  if period > 0.0:
+    # interval mode: check if condition is satisfied
+    time_prev = history_out[worldid, buf_offset]  # user slot stores time_prev
+    if time_prev + period <= t:
+      # advance time_prev by exact period
+      history_out[worldid, buf_offset] = time_prev + period
+      # insert sensor value
+      _history_insert_vector(worldid, buf_offset, nsample, dim, t, sensordata_in, sensor_adr[sid], history_out)
+  else:
+    _history_insert_vector(worldid, buf_offset, nsample, dim, t, sensordata_in, sensor_adr[sid], history_out)
+
+
+@wp.kernel
+def _apply_sensor_delay_kernel(
+  # Model:
+  sensor_dim: wp.array[int],
+  sensor_adr: wp.array[int],
+  sensor_history: wp.array[wp.vec2i],
+  sensor_historyadr: wp.array[int],
+  sensor_delay: wp.array[float],
+  sensor_interval: wp.array[wp.vec2],
+  # Data in:
+  time_in: wp.array[float],
+  history_in: wp.array2d[float],
+  # In:
+  sensor_ids: wp.array[int],
+  # Data out:
+  sensordata_out: wp.array2d[float],
+):
+  """Apply delay/interval logic for sensors after computation.
+
+  TODO(team): Revisit always-compute decision for computationally expensive sensors
+  with interval/period (e.g., raytracers)
+  """
+  worldid, idx = wp.tid()
+  sid = sensor_ids[idx]
+
+  hist = sensor_history[sid]
+  nsample = hist[0]
+  if nsample <= 0:
+    return
+
+  delay = sensor_delay[sid]
+  dim = sensor_dim[sid]
+  interp = hist[1]
+  buf_offset = sensor_historyadr[sid]
+  t = time_in[worldid]
+
+  if delay > 0.0:
+    # delay > 0: read delayed value from buffer
+    _history_read_vector(sensor_adr[sid], history_in, worldid, buf_offset, nsample, dim, t - delay, interp, sensordata_out)
+  else:
+    # interval-only (delay == 0, interval > 0): check interval condition
+    interval_val = sensor_interval[sid]
+    period = interval_val[0]
+    if period > 0.0:
+      time_prev = history_in[worldid, buf_offset]  # user slot
+      if time_prev + period > t:
+        # interval condition not satisfied: read from buffer
+        _history_read_vector(sensor_adr[sid], history_in, worldid, buf_offset, nsample, dim, t, interp, sensordata_out)
+      # else: interval condition satisfied, keep computed value
+
+
+def read_ctrl_delayed(m: Model, d: Data, ctrl: wp.array2d[float]):
+  """Read delayed ctrl values for all actuators."""
+  if m.nhistory == 0:
+    wp.copy(ctrl, d.ctrl)
+    return
+
+  wp.launch(
+    _read_ctrl_delayed_kernel,
+    dim=(d.nworld, m.nu),
+    inputs=[
+      m.actuator_history,
+      m.actuator_historyadr,
+      m.actuator_delay,
+      d.time,
+      d.history,
+      d.ctrl,
+    ],
+    outputs=[ctrl],
+  )
+
+
+def insert_ctrl_history(m: Model, d: Data):
+  """Insert current ctrl values into history buffers."""
+  if m.nhistory == 0 or m.nu == 0:
+    return
+
+  wp.launch(
+    _insert_ctrl_history_kernel,
+    dim=(d.nworld, m.nu),
+    inputs=[
+      m.actuator_history,
+      m.actuator_historyadr,
+      d.time,
+      d.ctrl,
+    ],
+    outputs=[d.history],
+  )
+
+
+def apply_sensor_delay(m: Model, d: Data, sensorid: wp.array[int]):
+  """Apply delay/interval logic for given sensors after computation.
+
+  Matches MuJoCo C architecture where the delayed read (mj_sensorPos) occurs
+  before the fresh value insert (mj_advance). We save fresh sensordata,
+  overwrite with delayed values, then insert the saved fresh values.
+  """
+  if m.nhistory == 0 or sensorid.shape[0] == 0:
+    return
+
+  # Save fresh sensordata before delay overwrite
+  fresh_sensordata = wp.empty_like(d.sensordata)
+  wp.copy(fresh_sensordata, d.sensordata)
+
+  # Read delayed values from buffer → overwrite sensordata
+  wp.launch(
+    _apply_sensor_delay_kernel,
+    dim=(d.nworld, sensorid.shape[0]),
+    inputs=[
+      m.sensor_dim,
+      m.sensor_adr,
+      m.sensor_history,
+      m.sensor_historyadr,
+      m.sensor_delay,
+      m.sensor_interval,
+      d.time,
+      d.history,
+      sensorid,
+    ],
+    outputs=[d.sensordata],
+  )
+
+  # Insert saved fresh sensor values into history buffers
+  wp.launch(
+    _insert_sensor_history_stage,
+    dim=(d.nworld, sensorid.shape[0]),
+    inputs=[
+      m.sensor_dim,
+      m.sensor_adr,
+      m.sensor_history,
+      m.sensor_historyadr,
+      m.sensor_delay,
+      m.sensor_interval,
+      d.time,
+      fresh_sensordata,
+      sensorid,
+    ],
+    outputs=[d.history],
+  )
+
+
+@wp.kernel
+def _read_ctrl_kernel(
+  # Model:
+  actuator_history: wp.array[wp.vec2i],
+  actuator_historyadr: wp.array[int],
+  actuator_delay: wp.array[float],
+  # Data in:
+  time_in: wp.array[float],
+  history_in: wp.array2d[float],
+  ctrl_in: wp.array2d[float],
+  # In:
+  uid: int,
+  interp: int,
+  # Out:
+  result_out: wp.array[float],
+):
+  """Read delayed ctrl for 1 actuator across all worlds."""
+  worldid = wp.tid()
+
+  hist = actuator_history[uid]
+  nsample = hist[0]
+
+  if nsample == 0:
+    result_out[worldid] = ctrl_in[worldid, uid]
+  else:
+    interp_val = interp
+    if interp_val < 0:
+      interp_val = hist[1]
+    delay = actuator_delay[uid]
+    buf_offset = actuator_historyadr[uid]
+    t = time_in[worldid] - delay
+    result_out[worldid] = _history_read_scalar(history_in, worldid, buf_offset, nsample, t, interp_val)
+
+
+def read_ctrl(
+  m: Model,
+  d: Data,
+  ctrlid: int,
+  time: wp.array[float],
+  interp: int,
+  result: wp.array2d[float],
+):
+  """Read delayed ctrl for 1 actuator across all worlds.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    ctrlid: actuator index.
+    time: query time per world (nworld,).
+    interp: interpolation mode (-1=model default, 0=ZOH, 1=linear, 2=cubic).
+    result: output buffer (nworld,).
+  """
+  wp.launch(
+    _read_ctrl_kernel,
+    dim=(d.nworld,),
+    inputs=[
+      m.actuator_history,
+      m.actuator_historyadr,
+      m.actuator_delay,
+      time,
+      d.history,
+      d.ctrl,
+      ctrlid,
+      interp,
+    ],
+    outputs=[result],
+  )
+
+
+@wp.kernel
+def _read_sensor_kernel(
+  # Model:
+  sensor_dim: wp.array[int],
+  sensor_adr: wp.array[int],
+  sensor_history: wp.array[wp.vec2i],
+  sensor_historyadr: wp.array[int],
+  sensor_delay: wp.array[float],
+  # Data in:
+  time_in: wp.array[float],
+  history_in: wp.array2d[float],
+  sensordata_in: wp.array2d[float],
+  # In:
+  sid: int,
+  interp: int,
+  # Out:
+  result_out: wp.array2d[float],
+):
+  """Read delayed sensor for 1 sensor across all worlds."""
+  worldid = wp.tid()
+
+  hist = sensor_history[sid]
+  nsample = hist[0]
+  dim = sensor_dim[sid]
+  adr = sensor_adr[sid]
+
+  if nsample == 0:
+    for i in range(dim):
+      result_out[worldid, i] = sensordata_in[worldid, adr + i]
+  else:
+    interp_val = interp
+    if interp_val < 0:
+      interp_val = hist[1]
+    delay = sensor_delay[sid]
+    buf_offset = sensor_historyadr[sid]
+    t = time_in[worldid] - delay
+    _history_read_vector(
+      0,  # write to result_out starting at index 0 (not global sensor adr)
+      history_in,
+      worldid,
+      buf_offset,
+      nsample,
+      dim,
+      t,
+      interp_val,
+      result_out,
+    )
+
+
+def read_sensor(
+  m: Model,
+  d: Data,
+  sensorid: int,
+  time: wp.array[float],
+  interp: int,
+  result: wp.array2d[float],
+):
+  """Read delayed sensor for 1 sensor across all worlds.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    sensorid: sensor index.
+    time: query time per world (nworld,).
+    interp: interpolation mode (-1=model default, 0=ZOH, 1=linear, 2=cubic).
+    result: output buffer (nworld, dim).
+  """
+  wp.launch(
+    _read_sensor_kernel,
+    dim=(d.nworld,),
+    inputs=[
+      m.sensor_dim,
+      m.sensor_adr,
+      m.sensor_history,
+      m.sensor_historyadr,
+      m.sensor_delay,
+      time,
+      d.history,
+      d.sensordata,
+      sensorid,
+      interp,
+    ],
+    outputs=[result],
+  )
+
+
+@wp.kernel
+def _init_ctrl_history_kernel(
+  # kernel_analyzer: off
+  # Model:
+  actuator_history: wp.array[wp.vec2i],
+  actuator_historyadr: wp.array[int],
+  # In:
+  ctrlid: int,
+  times: wp.array[float],
+  values: wp.array2d[float],
+  has_times: int,
+  # Data out:
+  history_out: wp.array2d[float],
+  # kernel_analyzer: on
+):
+  """Initialize history buffer for 1 actuator across all worlds."""
+  worldid = wp.tid()
+
+  nsample = actuator_history[ctrlid][0]
+  buf_offset = actuator_historyadr[ctrlid]
+
+  # preserve user slot
+  user = history_out[worldid, buf_offset]
+
+  # cursor = 0 (samples in order, newest at index nsample-1)
+  history_out[worldid, buf_offset + 1] = float(nsample - 1)
+
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + nsample
+
+  for i in range(nsample):
+    if has_times != 0:
+      history_out[worldid, times_offset + i] = times[i]
+    else:
+      history_out[worldid, times_offset + i] = -MJ_MAXVAL
+    history_out[worldid, values_offset + i] = values[worldid, i]
+
+  # restore user slot
+  history_out[worldid, buf_offset] = user
+
+
+def init_ctrl_history(
+  m: Model,
+  d: Data,
+  ctrlid: int,
+  times: wp.array[float],
+  values: wp.array2d[float],
+):
+  """Initialize history buffer for 1 actuator across all worlds.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    ctrlid: actuator index.
+    times: timestamps or None (nsample,).
+    values: ctrl values (nworld, nsample).
+
+  Raises:
+    ValueError: If times are not strictly increasing.
+  """
+  has_times = 0 if times is None else 1
+  if times is not None:
+    t_np = times.numpy()
+    for i in range(len(t_np) - 1):
+      if t_np[i + 1] - t_np[i] < MJ_MINVAL:
+        raise ValueError(f"times must be strictly increasing, got times[{i}]={t_np[i]} >= times[{i + 1}]={t_np[i + 1]}")
+  if times is None:
+    times = wp.empty(0, dtype=float)
+
+  wp.launch(
+    _init_ctrl_history_kernel,
+    dim=(d.nworld,),
+    inputs=[
+      m.actuator_history,
+      m.actuator_historyadr,
+      ctrlid,
+      times,
+      values,
+      has_times,
+    ],
+    outputs=[d.history],
+  )
+
+
+# kernel_analyzer: off
+@wp.kernel
+def _init_sensor_history_kernel(
+  # Model:
+  sensor_history: wp.array[wp.vec2i],
+  sensor_historyadr: wp.array[int],
+  sensor_dim_arr: wp.array[int],
+  # In:
+  sensorid: int,
+  times: wp.array[float],
+  values: wp.array2d[float],
+  phase: wp.array[float],
+  has_times: int,
+  # Data out:
+  history_out: wp.array2d[float],
+):
+  # kernel_analyzer: on
+  """Initialize history buffer for 1 sensor across all worlds."""
+  worldid = wp.tid()
+
+  nsample = sensor_history[sensorid][0]
+  dim = sensor_dim_arr[sensorid]
+  buf_offset = sensor_historyadr[sensorid]
+
+  # set user slot (phase = last computation time for interval sensors)
+  history_out[worldid, buf_offset] = phase[worldid]
+
+  # cursor = 0 (samples in order, newest at index nsample-1)
+  history_out[worldid, buf_offset + 1] = float(nsample - 1)
+
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + nsample
+
+  for i in range(nsample):
+    if has_times != 0:
+      history_out[worldid, times_offset + i] = times[i]
+    else:
+      history_out[worldid, times_offset + i] = -MJ_MAXVAL
+    for j in range(dim):
+      history_out[worldid, values_offset + i * dim + j] = values[worldid, i * dim + j]
+
+
+def init_sensor_history(
+  m: Model,
+  d: Data,
+  sensorid: int,
+  times: wp.array[float],
+  values: wp.array2d[float],
+  phase: wp.array[float],
+):
+  """Initialize history buffer for 1 sensor across all worlds.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    sensorid: sensor index.
+    times: timestamps or None (nsample,).
+    values: sensor values (nworld, nsample * dim).
+    phase: user slot value per world (nworld,).
+
+  Raises:
+    ValueError: If times are not strictly increasing.
+  """
+  has_times = 0 if times is None else 1
+  if times is not None:
+    t_np = times.numpy()
+    for i in range(len(t_np) - 1):
+      if t_np[i + 1] - t_np[i] < MJ_MINVAL:
+        raise ValueError(f"times must be strictly increasing, got times[{i}]={t_np[i]} >= times[{i + 1}]={t_np[i + 1]}")
+  if times is None:
+    times = wp.empty(0, dtype=float)
+
+  wp.launch(
+    _init_sensor_history_kernel,
+    dim=(d.nworld,),
+    inputs=[
+      m.sensor_history,
+      m.sensor_historyadr,
+      m.sensor_dim,
+      sensorid,
+      times,
+      values,
+      phase,
+      has_times,
+    ],
+    outputs=[d.history],
+  )

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

@@ -95,9 +95,9 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
 
     # TODO(team): qacc = d.qacc if (m.dof_damping == 0.0).all()
 
-    # set qfrc = (d.qM + m.opt.timestep * diag(m.dof_damping)) * d.qacc
+    # set qfrc = (d.M + m.opt.timestep * diag(m.dof_damping)) * d.qacc
 
-    # d.qM @ d.qacc
+    # d.M @ d.qacc
     support.mul_m(m, d, qfrc, d.qacc)
 
     # qfrc += m.opt.timestep * damp_deriv * d.qacc
@@ -109,16 +109,16 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
     )
   elif m.opt.integrator == IntegratorType.IMPLICITFAST:
     if m.is_sparse:
-      qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
+      qDeriv = wp.empty((d.nworld, 1, m.nC), dtype=float)
     else:
       qDeriv = wp.empty((d.nworld, m.nv, m.nv), dtype=float)
     derivative.deriv_smooth_vel(m, d, qDeriv)
     mul_m(m, d, qfrc, d.qacc, M=qDeriv)
-    smooth.factor_solve_i(m, d, d.qM, d.qLD, d.qLDiagInv, qacc, qfrc)
+    smooth.factor_solve_i(m, d, d.M, d.qLD, d.qLDiagInv, qacc, qfrc)
   else:
     raise NotImplementedError(f"integrator {m.opt.integrator} not implemented.")
 
-  # solve for qacc: qfrc = d.qM @ d.qacc
+  # solve for qacc: qfrc = d.M @ d.qacc
   smooth.solve_m(m, d, qacc, qfrc)
 
 

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

@@ -33,6 +33,9 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
 from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
 
+# TODO(team): remove after improving island solver performance
+ENABLE_ISLANDS = False
+
 
 def _is_array_spec(typ) -> bool:
   """Check if a type annotation is an array spec (wp.array instance or bracket annotation)."""
@@ -64,6 +67,61 @@ def _create_array(data: Any, spec, sizes: dict[str, int]) -> wp.array | None:
   return array
 
 
+def _create_constraint(
+  mjm,
+  nworld: int,
+  njmax: int,
+  njmax_nnz: int,
+  sizes: dict,
+  island_enabled: bool,
+  mjd=None,
+) -> types.Constraint:
+  """Construct a types.Constraint with standard and island local fields allocated properly."""
+  efc_kwargs = {"J_rownnz": None, "J_rowadr": None, "J_colind": None, "J": None}
+  sparse = is_sparse(mjm)
+
+  for f in dataclasses.fields(types.Constraint):
+    if f.name == "itype":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
+    elif f.name == "iid":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
+    elif f.name == "iJ_rownnz":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0) if sparse else (nworld, 0), dtype=int)
+    elif f.name == "iJ_rowadr":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0) if sparse else (nworld, 0), dtype=int)
+    elif f.name == "iJ_colind":
+      efc_kwargs[f.name] = wp.empty((nworld, 1, njmax_nnz if island_enabled else 0) if sparse else (nworld, 0, 0), dtype=int)
+    elif f.name == "iJ":
+      efc_kwargs[f.name] = wp.empty(
+        (nworld, 1, njmax_nnz if island_enabled else 0) if sparse else (nworld, njmax if island_enabled else 0, mjm.nv),
+        dtype=float,
+      )
+    elif f.name == "iD":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
+    elif f.name == "iaref":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
+    elif f.name == "ifrictionloss":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
+    elif f.name == "iforce":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
+    elif f.name == "istate":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
+    else:
+      if f.name in efc_kwargs:
+        continue
+
+      if mjd is not None:
+        shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in f.type.shape)
+        val = np.zeros(shape, dtype=f.type.dtype)
+        if f.name in ("type", "id", "pos", "margin", "D", "vel", "aref", "frictionloss", "force"):
+          val[:, : mjd.nefc] = np.tile(getattr(mjd, "efc_" + f.name), (nworld, 1))
+        efc_kwargs[f.name] = wp.array(val, dtype=f.type.dtype)
+      else:
+        efc_kwargs[f.name] = _create_array(None, f.type, sizes)
+
+  return types.Constraint(**efc_kwargs)
+
+
 def is_sparse(mjm: mujoco.MjModel) -> bool:
   if mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_AUTO:
     if mjm.nv > 32:
@@ -183,6 +241,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     opt_kwargs["impratio_invsqrt"] = 1.0 / np.sqrt(np.maximum(mjm.opt.impratio, mujoco.mjMINVAL))
   opt = types.Option(**opt_kwargs)
 
+  # islands are disabled by default while performance is being improved
+  # override by setting io.ENABLE_ISLANDS = True
+  # TODO(team): remove after improving island solver performance
+  if not ENABLE_ISLANDS:
+    opt.disableflags |= types.DisableBit.ISLAND
+
   # C MuJoCo tolerance was chosen for float64 architecture, but we default to float32 on GPU
   # adjust the tolerance for lower precision, to avoid the solver spending iterations needlessly
   # bouncing around the optimal solution
@@ -271,7 +335,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   jnt_limited_slide_hinge = mjm.jnt_limited & np.isin(mjm.jnt_type, (mujoco.mjtJoint.mjJNT_SLIDE, mujoco.mjtJoint.mjJNT_HINGE))
   m.jnt_limited_slide_hinge_adr = np.nonzero(jnt_limited_slide_hinge)[0]
   m.jnt_limited_ball_adr = np.nonzero(mjm.jnt_limited & (mjm.jnt_type == mujoco.mjtJoint.mjJNT_BALL))[0]
-  m.dof_tri_row, m.dof_tri_col = np.tril_indices(mjm.nv)
+  m.dof_tri_row, m.dof_tri_col = np.triu_indices(mjm.nv)
 
   # precompute body_isdofancestor: which DOFs affect each body
   # TODO: Investigate alternative approach such as bitmap
@@ -585,14 +649,14 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     for j in range(mjm.mesh_vertnum[mjm.sensor_objid[i]])
   ]
 
-  # qM_tiles records the block diagonal structure of qM
+  # M_tiles records the block diagonal structure of M
   tile_corners = [i for i in range(mjm.nv) if mjm.dof_parentid[i] == -1]
   tiles = {}
   for i in range(len(tile_corners)):
     tile_beg = tile_corners[i]
     tile_end = mjm.nv if i == len(tile_corners) - 1 else tile_corners[i + 1]
     tiles.setdefault(tile_end - tile_beg, []).append(tile_beg)
-  m.qM_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys()))
+  m.M_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys()))
 
   # qLD_updates has dof tree ordering of qLD updates for sparse factor m
   qLD_updates, dof_depth = {}, np.zeros(mjm.nv, dtype=int) - 1
@@ -620,41 +684,69 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.qLD_all_updates = all_updates_flat if all_updates_flat else [(0, 0, 0)]
   m.qLD_level_offsets = level_offsets
 
-  # indices for sparse qM_fullm (used in solver)
+  # Indices for sparse M_fullm (used in solver). M_fullm_i/j are built by
-  m.qM_fullm_i, m.qM_fullm_j = [], []
+  # walking dof_parentid for each dof, so for joint types whose internal block
+  # MuJoCo stores diagonal-only in the compact (M_rownnz, M_rowadr) layout
+  # (e.g. free joints), the chain-aware layout here has more entries per row
+  # than the compact layout.
+  m.M_fullm_i, m.M_fullm_j = [], []
   for i in range(mjm.nv):
     j = i
     while j > -1:
-      m.qM_fullm_i.append(i)
+      m.M_fullm_i.append(i)
-      m.qM_fullm_j.append(j)
+      m.M_fullm_j.append(j)
       j = mjm.dof_parentid[j]
+  # M_elemid maps (row, col) -> madr index in the native CSR M layout
+  M_elemid = np.full((mjm.nv, mjm.nv), -1, dtype=np.int32)
+  for i in range(mjm.nv):
+    rowadr = mjm.M_rowadr[i]
+    rownnz = mjm.M_rownnz[i]
+    for k in range(rownnz):
+      madr = rowadr + k
+      col = int(mjm.M_colind[madr])
+      M_elemid[i, col] = madr
+  m.M_elemid = M_elemid
+
+  upper_j, upper_i = np.triu_indices(mjm.nv)
+  upper_elemid = M_elemid[upper_i, upper_j]
+  valid_mask = upper_elemid != -1
+  m.M_fullm_upper_i = upper_j[valid_mask].tolist()
+  m.M_fullm_upper_j = upper_i[valid_mask].tolist()
+  m.M_fullm_upper_elemid = upper_elemid[valid_mask].tolist()
+
+  # indices for sparse qD_fullm (used in RNE derivatives)
+  # D-structure is the full square sparsity pattern (both upper and lower triangle)
+  m.qD_fullm_i, m.qD_fullm_j = [], []
+  for i in range(mjm.nv):
+    rowadr = mjm.D_rowadr[i]
+    rownnz = mjm.D_rownnz[i]
+    for k in range(rownnz):
+      m.qD_fullm_i.append(i)
+      m.qD_fullm_j.append(int(mjm.D_colind[rowadr + k]))
+  m.nD = mjm.nD
 
   # Gather-based sparse mul_m: for each row, all (col, madr) including diagonal
   row_elements = [[] for _ in range(mjm.nv)]
 
-  # Add diagonal
   for i in range(mjm.nv):
-    row_elements[i].append((i, mjm.dof_Madr[i]))
+    rowadr = mjm.M_rowadr[i]
-
+    rownnz = mjm.M_rownnz[i]
-  # Add off-diagonals: ancestors (lower) and descendants (upper)
+    for k in range(rownnz):
-  for i in range(mjm.nv):
+      madr = rowadr + k
-    madr_ij, j = mjm.dof_Madr[i], i
+      col = int(mjm.M_colind[madr])
-    while True:
+      row_elements[i].append((col, madr))  # row i gathers M[i,col] * vec[col]
-      madr_ij, j = madr_ij + 1, mjm.dof_parentid[j]
+      if i != col:
-      if j == -1:
+        row_elements[col].append((i, madr))  # row col gathers M[i,col] * vec[i]
-        break
-      row_elements[i].append((j, madr_ij))  # row i gathers M[i,j] * vec[j]
-      row_elements[j].append((i, madr_ij))  # row j gathers M[j,i] * vec[i]
 
   # Flatten into CSR-like arrays
-  m.qM_mulm_rowadr = [0]
+  m.M_mulm_rowadr = [0]
-  m.qM_mulm_col = []
+  m.M_mulm_col = []
-  m.qM_mulm_madr = []
+  m.M_mulm_madr = []
   for i in range(mjm.nv):
     for col, madr in row_elements[i]:
-      m.qM_mulm_col.append(col)
+      m.M_mulm_col.append(col)
-      m.qM_mulm_madr.append(madr)
+      m.M_mulm_madr.append(madr)
-    m.qM_mulm_rowadr.append(len(m.qM_mulm_col))
+    m.M_mulm_rowadr.append(len(m.M_mulm_col))
 
   m.flexedge_J_rownnz = mjm.flexedge_J_rownnz
   m.flexedge_J_rowadr = mjm.flexedge_J_rowadr
@@ -886,6 +978,42 @@ def _resolve_batch_size(na: int | None, n: int | None, nworld: int, default: int
   return default
 
 
+def _allocate_island_arrays(
+  mjm: mujoco.MjModel,
+  d: types.Data,
+  nworld: int,
+  njmax: int,
+  island_enabled: bool,
+  mjd: mujoco.MjData,
+):
+  ntree_size = mjm.ntree if island_enabled else 0
+  nv_size = mjm.nv if island_enabled else 0
+  njmax_size = njmax if island_enabled else 0
+
+  d.nisland = wp.array(np.full(nworld, mjd.nisland), dtype=int)
+  d.tree_island = wp.array(np.tile(mjd.tree_island, (nworld, 1 if island_enabled else 0)), dtype=int)
+  d.dof_island = wp.array(np.tile(mjd.dof_island, (nworld, 1 if island_enabled else 0)), dtype=int)
+
+  d.island_dofadr = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_nv = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_nefc = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_ne = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_nf = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_efcadr = wp.empty((nworld, ntree_size), dtype=int)
+  d.nidof = wp.empty((nworld if island_enabled else 0,), dtype=int)
+  d.map_dof2idof = wp.empty((nworld, nv_size), dtype=int)
+  d.map_idof2dof = wp.empty((nworld, nv_size), dtype=int)
+  d.map_efc2iefc = wp.empty((nworld, njmax_size), dtype=int)
+  d.map_iefc2efc = wp.empty((nworld, njmax_size), dtype=int)
+
+  d.dof_islandid = wp.empty((nworld, nv_size), dtype=int)
+  d.efc_islandid = wp.empty((nworld, njmax_size), dtype=int)
+  d.iqacc = wp.empty((nworld, nv_size), dtype=float)
+  d.iqacc_smooth = wp.empty((nworld, nv_size), dtype=float)
+  d.iqfrc_smooth = wp.empty((nworld, nv_size), dtype=float)
+  d.iqfrc_constraint = wp.empty((nworld, nv_size), dtype=float)
+
+
 def make_data(
   mjm: mujoco.MjModel,
   nworld: int = 1,
@@ -967,7 +1095,8 @@ def make_data(
 
   contact = types.Contact(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Contact)})
   contact.efc_address = wp.array(np.full((naconmax, sizes["nmaxpyramid"]), -1, dtype=int), dtype=int)
-  efc = types.Constraint(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Constraint)})
+
+  efc = _create_constraint(mjm, nworld, njmax, njmax_nnz, sizes, ENABLE_ISLANDS)
 
   if is_sparse(mjm):
     efc.J_rownnz = wp.zeros((nworld, njmax), dtype=int)
@@ -1005,7 +1134,7 @@ def make_data(
     "njmax": njmax,
     "njmax_pad": sizes["njmax_pad"],
     "njmax_nnz": njmax_nnz,
-    "qM": None,
+    "M": None,
     "qLD": None,
     # world body
     "xquat": wp.array(np.tile(mjd.xquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat),
@@ -1024,6 +1153,24 @@ def make_data(
     # island arrays
     "nisland": None,
     "tree_island": None,
+    "dof_island": None,
+    "island_dofadr": None,
+    "island_nv": None,
+    "island_nefc": None,
+    "island_ne": None,
+    "island_nf": None,
+    "island_efcadr": None,
+    "nidof": None,
+    "map_dof2idof": None,
+    "map_idof2dof": None,
+    "map_efc2iefc": None,
+    "map_iefc2efc": None,
+    "dof_islandid": None,
+    "efc_islandid": None,
+    "iqacc": None,
+    "iqacc_smooth": None,
+    "iqfrc_smooth": None,
+    "iqfrc_constraint": None,
   }
   for f in dataclasses.fields(types.Data):
     if f.name in d_kwargs:
@@ -1033,15 +1180,13 @@ def make_data(
   d = types.Data(**d_kwargs)
 
   if is_sparse(mjm):
-    d.qM = wp.zeros((nworld, 1, mjm.nM), dtype=float)
+    d.M = wp.zeros((nworld, 1, mjm.nC), dtype=float)
     d.qLD = wp.zeros((nworld, 1, mjm.nC), dtype=float)
   else:
-    d.qM = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float)
+    d.M = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float)
     d.qLD = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float)
 
-  # island discovery arrays
+  _allocate_island_arrays(mjm, d, nworld, njmax, ENABLE_ISLANDS, mjd)
-  d.nisland = wp.zeros((nworld,), dtype=int)
-  d.tree_island = wp.zeros((nworld, mjm.ntree), dtype=int)
 
   return d
 
@@ -1176,16 +1321,7 @@ def put_data(
   # create efc
   efc_kwargs = {"J_rownnz": None, "J_rowadr": None, "J_colind": None, "J": None}
 
-  for f in dataclasses.fields(types.Constraint):
+  efc = _create_constraint(mjm, nworld, njmax, njmax_nnz, sizes, ENABLE_ISLANDS, mjd)
-    if f.name in efc_kwargs:
-      continue
-    shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in f.type.shape)
-    val = np.zeros(shape, dtype=f.type.dtype)
-    if f.name in ("type", "id", "pos", "margin", "D", "vel", "aref", "frictionloss", "force"):
-      val[:, : mjd.nefc] = np.tile(getattr(mjd, "efc_" + f.name), (nworld, 1))
-    efc_kwargs[f.name] = wp.array(val, dtype=f.type.dtype)
-
-  efc = types.Constraint(**efc_kwargs)
 
   if is_sparse(mjm):
     J_rownnz = np.zeros(njmax, dtype=np.int32)
@@ -1236,12 +1372,30 @@ def put_data(
     "njmax_nnz": njmax_nnz,
     # fields set after initialization:
     "solver_niter": None,
-    "qM": None,
+    "M": None,
     "qLD": None,
     "nacon": None,
     # island arrays
     "nisland": None,
     "tree_island": None,
+    "dof_island": None,
+    "island_dofadr": None,
+    "island_nv": None,
+    "island_nefc": None,
+    "island_ne": None,
+    "island_nf": None,
+    "island_efcadr": None,
+    "nidof": None,
+    "map_dof2idof": None,
+    "map_idof2dof": None,
+    "map_efc2iefc": None,
+    "map_iefc2efc": None,
+    "dof_islandid": None,
+    "efc_islandid": None,
+    "iqacc": None,
+    "iqacc_smooth": None,
+    "iqfrc_smooth": None,
+    "iqfrc_constraint": None,
   }
   for f in dataclasses.fields(types.Data):
     if f.name in d_kwargs:
@@ -1256,20 +1410,25 @@ def put_data(
   d.solver_niter = wp.full((nworld,), mjd.solver_niter[0], dtype=int)
 
   if is_sparse(mjm):
-    d.qM = wp.array(np.full((nworld, 1, mjm.nM), mjd.qM), dtype=float)
+    if check_version("mujoco>=3.8.1.dev910242375"):
+      d.M = wp.array(np.full((nworld, 1, mjm.nC), mjd.M), dtype=float)
+    else:
+      d.M = wp.array(np.full((nworld, 1, mjm.nC), mjd.qM[mjm.mapM2M]), dtype=float)
     d.qLD = wp.array(np.full((nworld, 1, mjm.nC), mjd.qLD), dtype=float)
   else:
-    qM = np.zeros((mjm.nv, mjm.nv))
+    M = np.zeros((mjm.nv, mjm.nv))
-    mujoco.mj_fullM(mjm, qM, mjd.qM)
+    if check_version("mujoco>=3.8.1.dev910242375"):
-    qLD = np.linalg.cholesky(qM) if (mjd.qM != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
+      mujoco.mju_sym2dense(M, mjd.M, mjm.M_rownnz, mjm.M_rowadr, mjm.M_colind)
+      qLD = np.linalg.cholesky(M).T if (mjd.M != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
+    else:
+      mujoco.mj_fullM(mjm, M, mjd.qM)
+      qLD = np.linalg.cholesky(M).T if (mjd.qM != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
     padding = sizes["nv_pad"] - mjm.nv
-    qM_padded = np.pad(qM, ((0, padding), (0, padding)), mode="constant", constant_values=0.0)
+    M_padded = np.pad(M, ((0, padding), (0, padding)), mode="constant", constant_values=0.0)
-    d.qM = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), qM_padded), dtype=float)
+    d.M = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), M_padded), dtype=float)
     d.qLD = wp.array(np.full((nworld, mjm.nv, mjm.nv), qLD), dtype=float)
 
-  # island arrays
+  _allocate_island_arrays(mjm, d, nworld, njmax, ENABLE_ISLANDS, mjd)
-  d.nisland = wp.array(np.full(nworld, mjd.nisland), dtype=int)
-  d.tree_island = wp.array(np.tile(mjd.tree_island, (nworld, 1)), dtype=int)
 
   d.nacon = wp.array([mjd.ncon * nworld], dtype=int)
 
@@ -1403,6 +1562,9 @@ def get_data_into(
   result.qfrc_constraint[:] = d.qfrc_constraint.numpy()[world_id]
   result.qfrc_inverse[:] = d.qfrc_inverse.numpy()[world_id]
 
+  if mjm.nhistory > 0:
+    result.history[:] = d.history.numpy()[world_id]
+
   # contact
   result.contact.dist[:ncon] = d.contact.dist.numpy()[ncon_filter]
   result.contact.pos[:ncon] = d.contact.pos.numpy()[ncon_filter]
@@ -1417,15 +1579,25 @@ def get_data_into(
   result.contact.efc_address[:ncon] = contact_efc_address_ordered[:ncon]
 
   if is_sparse(mjm):
-    result.qM[:] = d.qM.numpy()[world_id, 0]
+    if check_version("mujoco>=3.8.1.dev910242375"):
+      result.M[:] = d.M.numpy()[world_id, 0]
+    else:
+      result.qM[mjm.mapM2M] = d.M.numpy()[world_id, 0]
     result.qLD[:] = d.qLD.numpy()[world_id, 0]
   else:
-    qM = d.qM.numpy()[world_id]
+    M = d.M.numpy()[world_id]
+    if check_version("mujoco>=3.8.1.dev910242375"):
+      for i in range(mjm.nv):
+        adr = mjm.M_rowadr[i]
+        for k in range(mjm.M_rownnz[i]):
+          col = mjm.M_colind[adr + k]
+          result.M[adr + k] = M[i, col]
+    else:
       adr = 0
       for i in range(mjm.nv):
         j = i
         while j >= 0:
-        result.qM[adr] = qM[i, j]
+          result.qM[adr] = M[i, j]
           j = mjm.dof_parentid[j]
           adr += 1
     mujoco.mj_factorM(mjm, result)
@@ -1466,6 +1638,7 @@ def get_data_into(
   result.efc_frictionloss[:] = d.efc.frictionloss.numpy()[world_id, efc_idx]
   result.efc_state[:] = d.efc.state.numpy()[world_id, efc_idx]
   result.efc_force[:] = d.efc.force.numpy()[world_id, efc_idx]
+  result.efc_island[:] = d.efc.island.numpy()[world_id, efc_idx]
 
   # rne_postconstraint
   result.cacc[:] = d.cacc.numpy()[world_id]
@@ -1487,8 +1660,51 @@ def get_data_into(
   # islands
   nisland = d.nisland.numpy()[world_id]
   result.nisland = nisland
-  if nisland:
+  if d.tree_island.shape[1] > 0 and nisland:
     result.tree_island[:] = d.tree_island.numpy()[world_id]
+    result.dof_island[:] = d.dof_island.numpy()[world_id]
+    result.island_dofadr[:nisland] = d.island_dofadr.numpy()[world_id, :nisland]
+    result.island_nv[:nisland] = d.island_nv.numpy()[world_id, :nisland]
+    result.island_nefc[:nisland] = d.island_nefc.numpy()[world_id, :nisland]
+    result.island_ne[:nisland] = d.island_ne.numpy()[world_id, :nisland]
+    result.island_nf[:nisland] = d.island_nf.numpy()[world_id, :nisland]
+    result.island_iefcadr[:nisland] = d.island_efcadr.numpy()[world_id, :nisland]
+    nv = mjm.nv
+    result.map_dof2idof[:nv] = d.map_dof2idof.numpy()[world_id, :nv]
+    result.map_idof2dof[:nv] = d.map_idof2dof.numpy()[world_id, :nv]
+    result.map_efc2iefc[:nefc] = d.map_efc2iefc.numpy()[world_id, :nefc]
+    result.map_iefc2efc[:nefc] = d.map_iefc2efc.numpy()[world_id, :nefc]
+
+    result.iefc_type[:nefc] = d.efc.itype.numpy()[world_id, :nefc]
+    result.iefc_id[:nefc] = d.efc.iid.numpy()[world_id, :nefc]
+    result.iefc_D[:nefc] = d.efc.iD.numpy()[world_id, :nefc]
+    result.iefc_aref[:nefc] = d.efc.iaref.numpy()[world_id, :nefc]
+    result.iefc_frictionloss[:nefc] = d.efc.ifrictionloss.numpy()[world_id, :nefc]
+    result.iefc_state[:nefc] = d.efc.istate.numpy()[world_id, :nefc]
+    result.iefc_force[:nefc] = d.efc.iforce.numpy()[world_id, :nefc]
+
+    if is_sparse(mjm):
+      iefc_J = np.zeros((nefc, mjm.nv))
+      mujoco.mju_sparse2dense(
+        iefc_J,
+        d.efc.iJ.numpy()[world_id, 0],
+        d.efc.iJ_rownnz.numpy()[world_id, :nefc],
+        d.efc.iJ_rowadr.numpy()[world_id, :nefc],
+        d.efc.iJ_colind.numpy()[world_id, 0],
+      )
+    else:
+      iefc_J = d.efc.iJ.numpy()[world_id, :nefc, : mjm.nv]
+
+    if mujoco.mj_isSparse(mjm):
+      mujoco.mju_dense2sparse(
+        result.iefc_J,
+        iefc_J,
+        result.iefc_J_rownnz,
+        result.iefc_J_rowadr,
+        result.iefc_J_colind,
+      )
+    else:
+      result.iefc_J[: nefc * mjm.nv] = iefc_J.flatten()
 
 
 def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
@@ -1511,14 +1727,14 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     xfrc_applied_out[worldid, bodyid][elemid] = 0.0
 
   @wp.kernel(module="unique", enable_backward=False)
-  def reset_qM(reset_in: wp.array[bool], qM_out: wp.array3d[float]):
+  def reset_M(reset_in: wp.array[bool], M_out: wp.array3d[float]):
     worldid, elemid1, elemid2 = wp.tid()
 
     if wp.static(reset is not None):
       if not reset_in[worldid]:
         return
 
-    qM_out[worldid, elemid1, elemid2] = 0.0
+    M_out[worldid, elemid1, elemid2] = 0.0
 
   @wp.kernel(module="unique", enable_backward=False)
   def reset_nworld(
@@ -1670,10 +1886,10 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
 
   wp.launch(reset_xfrc_applied, dim=(d.nworld, m.nbody, 6), inputs=[reset_input], outputs=[d.xfrc_applied])
   wp.launch(
-    reset_qM,
+    reset_M,
-    dim=(d.nworld, d.qM.shape[1], d.qM.shape[2]),
+    dim=(d.nworld, d.M.shape[1], d.M.shape[2]),
     inputs=[reset_input],
-    outputs=[d.qM],
+    outputs=[d.M],
   )
 
   # set mocap_pos/quat = body_pos/quat for mocap bodies
@@ -1786,11 +2002,12 @@ def _copy_tendon_length0(
 def _compute_meaninertia(
   nv: int,
   is_sparse: bool,
-  dof_Madr_in: wp.array[int],
+  M_rownnz_in: wp.array[int],
-  qM_in: wp.array3d[float],
+  M_rowadr_in: wp.array[int],
+  M_in: wp.array3d[float],
   meaninertia_out: wp.array[float],
 ):
-  """Compute mean diagonal inertia from qM at qpos0."""
+  """Compute mean diagonal inertia from M at qpos0."""
   worldid = wp.tid()
 
   if nv == 0:
@@ -1800,12 +2017,12 @@ def _compute_meaninertia(
   total = float(0.0)
   for i in range(nv):
     if is_sparse:
-      # Sparse: qM is flattened lower triangular, diagonal at dof_Madr[i]
+      # Sparse: M is in CSR format, diagonal at M_rowadr_in[i] + M_rownnz_in[i] - 1
-      madr = dof_Madr_in[i]
+      madr = M_rowadr_in[i] + M_rownnz_in[i] - 1
-      total += qM_in[worldid, 0, madr]
+      total += M_in[worldid, 0, madr]
     else:
-      # Dense: qM is 2D matrix, diagonal at [i,i]
+      # Dense: M is 2D matrix, diagonal at [i,i]
-      total += qM_in[worldid, i, i]
+      total += M_in[worldid, i, i]
 
   meaninertia_out[worldid % meaninertia_out.shape[0]] = total / float(nv)
 
@@ -2288,11 +2505,11 @@ def set_const_0(m: types.Model, d: types.Data):
   smooth.factor_m(m, d)
   smooth.transmission(m, d)
 
-  # Compute meaninertia from qM diagonal at qpos0
+  # Compute meaninertia from M diagonal at qpos0
   wp.launch(
     _compute_meaninertia,
     dim=d.nworld,
-    inputs=[m.nv, m.is_sparse, m.dof_Madr, d.qM],
+    inputs=[m.nv, m.is_sparse, m.M_rownnz, m.M_rowadr, d.M],
     outputs=[m.stat.meaninertia],
   )
 
@@ -2452,27 +2669,31 @@ def set_const(m: types.Model, d: types.Data):
 
   Model fields that can be modified safely with set_const:
 
-    Field                            | Notes
+  ==================================  ==============================================
-    ---------------------------------|----------------------------------------------
+  Field                               Notes
-    qpos0, qpos_spring               |
+  ==================================  ==============================================
-    body_mass, body_inertia,         | Mass and inertia are usually scaled together
+  qpos0, qpos_spring
-    body_ipos, body_iquat            | since inertia is sum(m * r^2).
+  body_mass, body_inertia,            Mass and inertia are usually scaled together
-    body_pos, body_quat              | Unsafe for static bodies (invalidates BVH).
+  body_ipos, body_iquat               since inertia is sum(m * r^2).
-    body_gravcomp                    | If changing from 0 to >0 bodies, required.
+  body_pos, body_quat                 Unsafe for static bodies (invalidates BVH).
-    dof_armature                     |
+  body_gravcomp                       If changing from 0 to >0 bodies, required.
-    eq_data                          | For connect/weld, offsets computed if not set.
+  dof_armature
-    hfield_size                      |
+  eq_data                             For connect/weld, offsets computed if not set.
-    tendon_stiffness, tendon_damping | Only if changing from/to zero.
+  hfield_size
-    actuator_gainprm, actuator_biasprm | For position actuators with dampratio.
+  tendon_stiffness, tendon_damping    Only if changing from/to zero.
+  actuator_gainprm, actuator_biasprm  For position actuators with dampratio.
+  ==================================  ==============================================
 
   For selective updates, use the sub-functions directly based on what changed:
 
-    Modified Field  | Call
+  ==============  ===============
-    ----------------|------------------
+  Modified Field  Call
-    body_mass       | set_const
+  ==============  ===============
-    body_gravcomp   | set_const_fixed
+  body_mass       set_const
-    body_inertia    | set_const_0
+  body_gravcomp   set_const_fixed
-    qpos0           | set_const_0
+  body_inertia    set_const_0
+  qpos0           set_const_0
+  ==============  ===============
 
   Computes:
     - Fixed quantities (via set_const_fixed):
@@ -2621,6 +2842,9 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any
       else:
         val = typ(val)
 
+      if attr == "disableflags" and isinstance(obj, types.Option) and not ENABLE_ISLANDS:
+        val = int(val) | types.DisableBit.ISLAND
+
       setattr(obj, attr, val)
 
 
@@ -2724,11 +2948,13 @@ def create_render_context(
   render_seg: list[bool] | bool | None = None,
   use_textures: bool = True,
   use_shadows: bool = False,
+  use_ambient_lighting: bool = True,
   enabled_geom_groups: list[int] = [0, 1, 2],
   cam_active: list[bool] | None = None,
   flex_render_smooth: bool = True,
   use_precomputed_rays: bool = True,
   render_skybox: bool = False,
+  enable_backface_culling: bool = True,
 ) -> types.RenderContext:
   """Creates a render context on device.
 
@@ -2743,6 +2969,8 @@ def create_render_context(
       If None, uses the MuJoCo model values.
     use_textures: Whether to use textures.
     use_shadows: Whether to use shadows.
+    use_ambient_lighting: Whether to add the renderer's hemispheric ambient
+      lighting term before applying model lights.
     enabled_geom_groups: The geom groups to render.
     cam_active: List of booleans indicating which cameras to include in rendering.
                 If None, all cameras are included.
@@ -2751,6 +2979,10 @@ def create_render_context(
                           When using domain randomization for camera intrinsics, set to False.
     render_skybox: Whether to shade missed rays with the MuJoCo skybox texture.
                    Requires the model to contain a texture with type `mjTEXTURE_SKYBOX`.
+    enable_backface_culling: Drop primitive-ray hits whose normal faces away from
+                             the ray (ray origin inside the geom). Matches MuJoCo's
+                             mesh-ray rule. Default True. Disable for a small
+                             performance gain when no camera is ever inside a geom.
 
   Returns:
     The render context containing rendering fields and output arrays on device.
@@ -2758,12 +2990,6 @@ def create_render_context(
   mjd = mujoco.MjData(mjm)
   mujoco.mj_forward(mjm, mjd)
 
-  constructor = "sah"
-  if check_version("warp>=1.13.0.dev20260325"):
-    # TODO: The cubql constructor and is_cubql_available exist only in
-    # recent Warp 1.13+ builds, modify this after warp is updated to 1.13+.
-    _cubql_avail = getattr(wp, "is_cubql_available", None)
-    if callable(_cubql_avail) and _cubql_avail():
   constructor = "cubql"
 
   # Mesh BVHs – build for all meshes so per-world variants are available
@@ -2937,6 +3163,7 @@ def create_render_context(
     cam_id_map=wp.array(active_cam_indices, dtype=int),
     use_textures=use_textures,
     use_shadows=use_shadows,
+    use_ambient_lighting=use_ambient_lighting,
     background_color=render_util.pack_rgba_to_uint32(0.1 * 255.0, 0.1 * 255.0, 0.2 * 255.0, 1.0 * 255.0),
     use_precomputed_rays=use_precomputed_rays,
     render_skybox=render_skybox,
@@ -2982,6 +3209,7 @@ def create_render_context(
     render_seg=wp.array(render_seg, dtype=bool),
     znear=znear,
     total_rays=int(total),
+    enable_backface_culling=enable_backface_culling,
   )
 
   bvh.build_scene_bvh(mjm, mjd, rc, nworld)

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

@@ -18,6 +18,7 @@ import warp as wp
 from mujoco.mjx.third_party.mujoco_warp._src import types
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
 from mujoco.mjx.third_party.mujoco_warp._src.types import EqType
+from mujoco.mjx.third_party.mujoco_warp._src.types import IslandSolverContext
 from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
@@ -35,11 +36,15 @@ def _tree_edges(
   eq_obj1id: wp.array[int],
   eq_obj2id: wp.array[int],
   eq_objtype: wp.array[int],
+  is_sparse: bool,
   # Data in:
   nefc_in: wp.array[int],
   contact_geom_in: wp.array[wp.vec2i],
   efc_type_in: wp.array2d[int],
   efc_id_in: wp.array2d[int],
+  efc_J_rownnz_in: wp.array2d[int],
+  efc_J_rowadr_in: wp.array2d[int],
+  efc_J_colind_in: wp.array3d[int],
   efc_J_in: wp.array3d[float],
   njmax_in: int,
   # Out:
@@ -131,11 +136,22 @@ def _tree_edges(
   first_tree = int(-1)
   has_cross_edge = int(0)
 
-  for dof in range(nv):
+  count = nv
-    # TODO(team): sparse efc_J
+  rowadr = 0
-    # TODO(team): tree dof skip
+  if is_sparse:
+    count = efc_J_rownnz_in[worldid, efcid]
+    rowadr = efc_J_rowadr_in[worldid, efcid]
+
+  for i in range(count):
+    dof = i
+    if is_sparse:
+      sparseid = rowadr + i
+      dof = efc_J_colind_in[worldid, 0, sparseid]
+    else:
       J_val = efc_J_in[worldid, efcid, dof]
-    if J_val != 0.0:
+      if J_val == 0.0:
+        continue
+
     tree = dof_treeid[dof]
     if tree < 0:
       continue
@@ -145,12 +161,14 @@ def _tree_edges(
       t1 = wp.min(first_tree, tree)
       t2 = wp.max(first_tree, tree)
       wp.atomic_max(tree_tree, worldid, t1, t2, 1)
+      wp.atomic_max(tree_tree, worldid, t2, t1, 1)
       has_cross_edge = 1
 
   if first_tree >= 0 and has_cross_edge == 0:
     wp.atomic_max(tree_tree, worldid, first_tree, first_tree, 1)
 
 
+@event_scope
 def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
   """Compute tree-tree adjacency matrix."""
   tree_tree.zero_()
@@ -168,10 +186,14 @@ def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
       m.eq_obj1id,
       m.eq_obj2id,
       m.eq_objtype,
+      m.is_sparse,
       d.nefc,
       d.contact.geom,
       d.efc.type,
       d.efc.id,
+      d.efc.J_rownnz,
+      d.efc.J_rowadr,
+      d.efc.J_colind,
       d.efc.J,
       d.njmax,
     ],
@@ -242,6 +264,19 @@ def _flood_fill(
   nisland_out[worldid] = nisland
 
 
+@event_scope
+def flood_fill(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
+  d.tree_island.fill_(-1)
+  stack_scratch = wp.empty((d.nworld, m.ntree * m.ntree), dtype=int)
+
+  wp.launch(
+    _flood_fill,
+    dim=d.nworld,
+    inputs=[m.ntree, tree_tree, d.tree_island, stack_scratch],
+    outputs=[d.nisland, d.tree_island, stack_scratch],
+  )
+
+
 @event_scope
 def island(m: types.Model, d: types.Data):
   """Discover constraint islands."""
@@ -254,12 +289,778 @@ def island(m: types.Model, d: types.Data):
   tree_edges(m, d, tree_tree)
 
   # Step 2: DFS flood fill
-  d.tree_island.fill_(-1)
+  flood_fill(m, d, tree_tree)
-  stack_scratch = wp.empty((d.nworld, m.ntree * m.ntree), dtype=int)
+
+
+@wp.kernel
+def _island_count_dofs(
+  dof_treeid: wp.array[int],
+  tree_island_in: wp.array2d[int],
+  dof_island_out: wp.array2d[int],
+  island_nv_out: wp.array2d[int],
+):
+  worldid, dofid = wp.tid()
+
+  island_id = tree_island_in[worldid, dof_treeid[dofid]]
+  dof_island_out[worldid, dofid] = island_id
+  if island_id >= 0:
+    wp.atomic_add(island_nv_out, worldid, island_id, 1)
+
+
+@wp.kernel
+def _island_scan_sizes(
+  nisland_in: wp.array[int],
+  island_idofadr_out: wp.array2d[int],
+  island_nv_inout: wp.array2d[int],
+  island_nefc_inout: wp.array2d[int],
+  island_iefcadr_out: wp.array2d[int],
+  nidof_out: wp.array[int],
+):
+  worldid = wp.tid()
+
+  nisland = nisland_in[worldid]
+  if nisland == 0:
+    nidof_out[worldid] = 0
+    return
+
+  # Scan DOFs and Constraints
+  island_idofadr_out[worldid, 0] = 0
+  island_iefcadr_out[worldid, 0] = 0
+  for i in range(1, nisland):
+    island_idofadr_out[worldid, i] = island_idofadr_out[worldid, i - 1] + island_nv_inout[worldid, i - 1]
+    island_iefcadr_out[worldid, i] = island_iefcadr_out[worldid, i - 1] + island_nefc_inout[worldid, i - 1]
+
+  nidof = island_idofadr_out[worldid, nisland - 1] + island_nv_inout[worldid, nisland - 1]
+  nidof_out[worldid] = nidof
+
+  # Reset for recount
+  for i in range(nisland):
+    island_nv_inout[worldid, i] = 0
+    island_nefc_inout[worldid, i] = 0
+
+
+@wp.kernel
+def _island_map_dofs(
+  nv: int,
+  dof_island_in: wp.array2d[int],
+  island_idofadr_in: wp.array2d[int],
+  nidof_in: wp.array[int],
+  island_nv_inout: wp.array2d[int],
+  map_dof2idof_out: wp.array2d[int],
+  map_idof2dof_out: wp.array2d[int],
+  idof_islandid_out: wp.array2d[int],
+  unconstrained_cnt_inout: wp.array2d[int],
+):
+  worldid, dofid = wp.tid()
+
+  nidof = nidof_in[worldid]
+  island_id = dof_island_in[worldid, dofid]
+
+  if island_id >= 0:
+    local_idx = wp.atomic_add(island_nv_inout, worldid, island_id, 1)
+    idof = island_idofadr_in[worldid, island_id] + local_idx
+    idof_islandid_out[worldid, idof] = island_id
+  else:
+    cnt = wp.atomic_add(unconstrained_cnt_inout, worldid, 0, 1)
+    idof = nidof + cnt
+
+  map_dof2idof_out[worldid, dofid] = idof
+  map_idof2dof_out[worldid, idof] = dofid
+
+
+@wp.kernel
+def _island_count_constraints(
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  efc_tree_in: wp.array2d[int],
+  tree_island_in: wp.array2d[int],
+  efc_type_in: wp.array2d[int],
+  efc_island_out: wp.array2d[int],
+  island_nefc_out: wp.array2d[int],
+  island_ne_out: wp.array2d[int],
+  island_nf_out: wp.array2d[int],
+):
+  worldid, efcid = wp.tid()
+  if efcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  efc_tree = efc_tree_in[worldid, efcid]
+  if efc_tree < 0:
+    efc_island_out[worldid, efcid] = -1
+    return
+  island_id = tree_island_in[worldid, efc_tree]
+  efc_island_out[worldid, efcid] = island_id
+
+  if island_id >= 0:
+    wp.atomic_add(island_nefc_out, worldid, island_id, 1)
+
+    efc_type = efc_type_in[worldid, efcid]
+    if efc_type == ConstraintType.EQUALITY:
+      wp.atomic_add(island_ne_out, worldid, island_id, 1)
+    elif efc_type == ConstraintType.FRICTION_DOF or efc_type == ConstraintType.FRICTION_TENDON:
+      wp.atomic_add(island_nf_out, worldid, island_id, 1)
+
+
+@wp.kernel
+def _island_map_constraints(
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  efc_island_in: wp.array2d[int],
+  island_iefcadr_in: wp.array2d[int],
+  island_ne_in: wp.array2d[int],
+  island_nf_in: wp.array2d[int],
+  efc_type_in: wp.array2d[int],
+  # Counters (inout):
+  island_ne_mapped_inout: wp.array2d[int],
+  island_nf_mapped_inout: wp.array2d[int],
+  island_nother_mapped_inout: wp.array2d[int],
+  island_nefc_inout: wp.array2d[int],
+  # Out:
+  map_efc2iefc_out: wp.array2d[int],
+  map_iefc2efc_out: wp.array2d[int],
+  iefc_islandid_out: wp.array2d[int],
+):
+  worldid, efcid = wp.tid()
+  if efcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  island_id = efc_island_in[worldid, efcid]
+  if island_id >= 0:
+    efc_type = efc_type_in[worldid, efcid]
+
+    # 1. Determine local index and absolute index ic based on category
+    if efc_type == ConstraintType.EQUALITY:
+      local_idx = wp.atomic_add(island_ne_mapped_inout, worldid, island_id, 1)
+      ic = island_iefcadr_in[worldid, island_id] + local_idx
+    elif efc_type == ConstraintType.FRICTION_DOF or efc_type == ConstraintType.FRICTION_TENDON:
+      local_idx = wp.atomic_add(island_nf_mapped_inout, worldid, island_id, 1)
+      ic = island_iefcadr_in[worldid, island_id] + island_ne_in[worldid, island_id] + local_idx
+    else:
+      local_idx = wp.atomic_add(island_nother_mapped_inout, worldid, island_id, 1)
+      ic = (
+        island_iefcadr_in[worldid, island_id] + island_ne_in[worldid, island_id] + island_nf_in[worldid, island_id] + local_idx
+      )
+
+    # 2. Increment overall island_nefc counter to reconstruct d.island_nefc
+    wp.atomic_add(island_nefc_inout, worldid, island_id, 1)
+
+    # 3. Store mappings
+    map_efc2iefc_out[worldid, efcid] = ic
+    map_iefc2efc_out[worldid, ic] = efcid
+    iefc_islandid_out[worldid, ic] = island_id
+
+
+@wp.kernel
+def _island_scan_sparse_rows(
+  nisland_in: wp.array[int],
+  efc_J_rownnz_in: wp.array2d[int],
+  island_nefc_in: wp.array2d[int],
+  island_iefcadr_in: wp.array2d[int],
+  map_iefc2efc_in: wp.array2d[int],
+  iefc_J_rownnz_out: wp.array2d[int],
+  iefc_J_rowadr_out: wp.array2d[int],
+):
+  worldid = wp.tid()
+
+  nisland = nisland_in[worldid]
+  if nisland == 0:
+    return
+
+  total_gathered_efc = island_iefcadr_in[worldid, nisland - 1] + island_nefc_in[worldid, nisland - 1]
+
+  running_rowadr = int(0)
+  for ic in range(total_gathered_efc):
+    c = map_iefc2efc_in[worldid, ic]
+    rownnz = efc_J_rownnz_in[worldid, c]
+    iefc_J_rowadr_out[worldid, ic] = running_rowadr
+    iefc_J_rownnz_out[worldid, ic] = rownnz
+    running_rowadr += rownnz
+
+
+@wp.kernel
+def _compute_efc_tree(
+  # Model:
+  nv: int,
+  body_treeid: wp.array[int],
+  jnt_dofadr: wp.array[int],
+  dof_treeid: wp.array[int],
+  geom_bodyid: wp.array[int],
+  site_bodyid: wp.array[int],
+  eq_type: wp.array[int],
+  eq_obj1id: wp.array[int],
+  eq_obj2id: wp.array[int],
+  eq_objtype: wp.array[int],
+  is_sparse: bool,
+  # Data in:
+  nefc_in: wp.array[int],
+  contact_geom_in: wp.array[wp.vec2i],
+  efc_type_in: wp.array2d[int],
+  efc_id_in: wp.array2d[int],
+  efc_J_in: wp.array3d[float],
+  efc_J_rownnz_in: wp.array2d[int],
+  efc_J_rowadr_in: wp.array2d[int],
+  efc_J_colind_in: wp.array3d[int],
+  njmax_in: int,
+  # Out:
+  efc_tree_out: wp.array2d[int],
+):
+  """Compute the first non-negative tree for each constraint."""
+  worldid, efcid = wp.tid()
+
+  if efcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  efc_type = efc_type_in[worldid, efcid]
+  efc_id = efc_id_in[worldid, efcid]
+
+  tree = int(-1)
+  use_generic = int(0)
+
+  # equality (connect/weld)
+  if efc_type == ConstraintType.EQUALITY:
+    eq_t = eq_type[efc_id]
+
+    if eq_t == EqType.CONNECT or eq_t == EqType.WELD:
+      b1 = eq_obj1id[efc_id]
+      b2 = eq_obj2id[efc_id]
+
+      if eq_objtype[efc_id] == ObjType.SITE:
+        b1 = site_bodyid[b1]
+        b2 = site_bodyid[b2]
+
+      t1 = body_treeid[b1]
+      t2 = body_treeid[b2]
+      if t1 >= 0:
+        tree = t1
+      else:
+        tree = t2
+    else:
+      # JOINT, TENDON, FLEX: generic scan
+      use_generic = 1
+
+  # joint friction
+  elif efc_type == ConstraintType.FRICTION_DOF:
+    tree = dof_treeid[efc_id]
+
+  # joint limit
+  elif efc_type == ConstraintType.LIMIT_JOINT:
+    tree = dof_treeid[jnt_dofadr[efc_id]]
+
+  # contact
+  elif (
+    efc_type == ConstraintType.CONTACT_FRICTIONLESS
+    or efc_type == ConstraintType.CONTACT_PYRAMIDAL
+    or efc_type == ConstraintType.CONTACT_ELLIPTIC
+  ):
+    geom_pair = contact_geom_in[efc_id]
+    g1 = geom_pair[0]
+    g2 = geom_pair[1]
+
+    if g1 >= 0 and g2 >= 0:
+      t1 = body_treeid[geom_bodyid[g1]]
+      t2 = body_treeid[geom_bodyid[g2]]
+      if t1 >= 0:
+        tree = t1
+      else:
+        tree = t2
+    else:
+      # flex contacts: generic scan
+      use_generic = 1
+
+  else:
+    # generic: scan Jacobian row
+    use_generic = 1
+
+  if use_generic:
+    count = nv
+    rowadr = 0
+    if is_sparse:
+      count = efc_J_rownnz_in[worldid, efcid]
+      rowadr = efc_J_rowadr_in[worldid, efcid]
+
+    for i in range(count):
+      dof = i
+      if is_sparse:
+        sparseid = rowadr + i
+        dof = efc_J_colind_in[worldid, 0, sparseid]
+      else:
+        J_val = efc_J_in[worldid, efcid, dof]
+        if J_val == 0.0:
+          continue
+
+      t = dof_treeid[dof]
+      if t >= 0:
+        tree = t
+        break
+
+  efc_tree_out[worldid, efcid] = tree
+
+
+@wp.kernel
+def _gather_efc_and_jacobian(
+  # Model:
+  is_sparse: bool,
+  # Data in:
+  nefc_in: wp.array[int],
+  efc_D_in: wp.array2d[float],
+  efc_type_in: wp.array2d[int],
+  efc_id_in: wp.array2d[int],
+  efc_frictionloss_in: wp.array2d[float],
+  efc_aref_in: wp.array2d[float],
+  efc_J_in: wp.array3d[float],
+  # Sparse Jacobian arrays:
+  efc_J_rownnz_in: wp.array2d[int],
+  efc_J_rowadr_in: wp.array2d[int],
+  efc_J_colind_in: wp.array3d[int],
+  iefc_J_rowadr_in: wp.array2d[int],
+  # In:
+  njmax_in: int,
+  map_iefc2efc_in: wp.array2d[int],
+  map_idof2dof_in: wp.array2d[int],
+  map_dof2idof_in: wp.array2d[int],
+  nidof_in: wp.array[int],
+  # Out:
+  iefc_D_out: wp.array2d[float],
+  iefc_type_out: wp.array2d[int],
+  iefc_id_out: wp.array2d[int],
+  iefc_frictionloss_out: wp.array2d[float],
+  iefc_aref_out: wp.array2d[float],
+  iefc_J_out: wp.array3d[float],
+  iefc_J_colind_out: wp.array3d[int],
+):
+  """Gather constraint arrays and Jacobian into island-local dense order."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  c = map_iefc2efc_in[worldid, iefcid]
+
+  # Scalar constraint fields
+  iefc_D_out[worldid, iefcid] = efc_D_in[worldid, c]
+  iefc_type_out[worldid, iefcid] = efc_type_in[worldid, c]
+  iefc_id_out[worldid, iefcid] = efc_id_in[worldid, c]
+  iefc_frictionloss_out[worldid, iefcid] = efc_frictionloss_in[worldid, c]
+  iefc_aref_out[worldid, iefcid] = efc_aref_in[worldid, c]
+
+  # Jacobian: convert to island-local dense format
+  nid = nidof_in[worldid]
+
+  if is_sparse:
+    rownnz = efc_J_rownnz_in[worldid, c]
+    rowadr_in = efc_J_rowadr_in[worldid, c]
+    rowadr_out = iefc_J_rowadr_in[worldid, iefcid]
+    for i in range(rownnz):
+      sparseid_in = rowadr_in + i
+      sparseid_out = rowadr_out + i
+      dof = efc_J_colind_in[worldid, 0, sparseid_in]
+      idof = map_dof2idof_in[worldid, dof]
+      # Store in sparse iefc_J_out and iefc_J_colind_out
+      iefc_J_out[worldid, 0, sparseid_out] = efc_J_in[worldid, 0, sparseid_in]
+      iefc_J_colind_out[worldid, 0, sparseid_out] = idof
+  else:
+    # Dense path: reorder rows by iefc, columns by idof
+    for idof in range(nid):
+      dof = map_idof2dof_in[worldid, idof]
+      iefc_J_out[worldid, iefcid, idof] = efc_J_in[worldid, c, dof]
+
+
+@wp.kernel
+def _gather_dof_arrays(
+  # Data in:
+  qacc_in: wp.array2d[float],
+  qacc_smooth_in: wp.array2d[float],
+  qfrc_smooth_in: wp.array2d[float],
+  # In:
+  nidof_in: wp.array[int],
+  map_idof2dof_in: wp.array2d[int],
+  # Out:
+  iacc_out: wp.array2d[float],
+  iacc_smooth_out: wp.array2d[float],
+  ifrc_smooth_out: wp.array2d[float],
+):
+  """Gather DOF arrays into island-local order."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  dof = map_idof2dof_in[worldid, idofid]
+  iacc_out[worldid, idofid] = qacc_in[worldid, dof]
+  iacc_smooth_out[worldid, idofid] = qacc_smooth_in[worldid, dof]
+  ifrc_smooth_out[worldid, idofid] = qfrc_smooth_in[worldid, dof]
+
+
+@wp.kernel
+def _scatter_dof_arrays(
+  # In:
+  qacc_smooth_in: wp.array2d[float],
+  qfrc_smooth_in: wp.array2d[float],
+  dof_island_in: wp.array2d[int],
+  iacc_in: wp.array2d[float],
+  ifrc_constraint_in: wp.array2d[float],
+  iMa_in: wp.array2d[float],
+  map_dof2idof_in: wp.array2d[int],
+  scatter_Ma: bool,
+  # Out:
+  qacc_out: wp.array2d[float],
+  qfrc_constraint_out: wp.array2d[float],
+  Ma_out: wp.array2d[float],
+):
+  """Scatter island results to global arrays, copy qacc_smooth for non-island DOFs."""
+  worldid, dofid = wp.tid()
+
+  if dof_island_in[worldid, dofid] < 0:
+    qacc_out[worldid, dofid] = qacc_smooth_in[worldid, dofid]
+    qfrc_constraint_out[worldid, dofid] = 0.0
+    if scatter_Ma:
+      Ma_out[worldid, dofid] = qfrc_smooth_in[worldid, dofid]
+  else:
+    idof = map_dof2idof_in[worldid, dofid]
+    qacc_out[worldid, dofid] = iacc_in[worldid, idof]
+    qfrc_constraint_out[worldid, dofid] = ifrc_constraint_in[worldid, idof]
+    if scatter_Ma:
+      Ma_out[worldid, dofid] = iMa_in[worldid, idof]
+
+
+@wp.kernel
+def _scatter_efc_arrays(
+  # In:
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  map_iefc2efc_in: wp.array2d[int],
+  iefc_force_in: wp.array2d[float],
+  iefc_state_in: wp.array2d[int],
+  # Out:
+  efc_force_out: wp.array2d[float],
+  efc_state_out: wp.array2d[int],
+):
+  """Scatter island-local constraint results back to global order."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  c = map_iefc2efc_in[worldid, iefcid]
+  efc_force_out[worldid, c] = iefc_force_in[worldid, iefcid]
+  efc_state_out[worldid, c] = iefc_state_in[worldid, iefcid]
+
+
+@wp.kernel
+def _init_island_arrays(
+  island_idofadr_out: wp.array2d[int],
+  island_nv_out: wp.array2d[int],
+  island_nefc_out: wp.array2d[int],
+  island_ne_out: wp.array2d[int],
+  island_nf_out: wp.array2d[int],
+  island_iefcadr_out: wp.array2d[int],
+  nidof_out: wp.array[int],
+):
+  worldid, islandid = wp.tid()
+  island_nv_out[worldid, islandid] = 0
+  island_nefc_out[worldid, islandid] = 0
+  island_ne_out[worldid, islandid] = 0
+  island_nf_out[worldid, islandid] = 0
+  island_idofadr_out[worldid, islandid] = 0
+  island_iefcadr_out[worldid, islandid] = 0
+  if islandid == 0:
+    nidof_out[worldid] = 0
+
+
+@wp.kernel
+def _init_dof_arrays(
+  dof_island_out: wp.array2d[int],
+  map_dof2idof_out: wp.array2d[int],
+  map_idof2dof_out: wp.array2d[int],
+  idof_islandid_out: wp.array2d[int],
+):
+  worldid, dofid = wp.tid()
+  dof_island_out[worldid, dofid] = -1
+  map_dof2idof_out[worldid, dofid] = 0
+  map_idof2dof_out[worldid, dofid] = 0
+  idof_islandid_out[worldid, dofid] = -1
+
+
+@wp.kernel
+def _init_efc_arrays(
+  efc_island_out: wp.array2d[int],
+  map_efc2iefc_out: wp.array2d[int],
+  map_iefc2efc_out: wp.array2d[int],
+  iefc_islandid_out: wp.array2d[int],
+  efc_tree_out: wp.array2d[int],
+):
+  worldid, efcid = wp.tid()
+  efc_island_out[worldid, efcid] = -1
+  map_efc2iefc_out[worldid, efcid] = 0
+  map_iefc2efc_out[worldid, efcid] = 0
+  iefc_islandid_out[worldid, efcid] = -1
+  efc_tree_out[worldid, efcid] = -1
+
+
+@event_scope
+def compute_island_mapping(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Compute DOF/constraint island mappings after island discovery.
+
+  Populates island solver context arrays via ctx: nv, nefc, ne, nf,
+  iefcadr, nidof, map_dof2idof, map_idof2dof, dof_islandid, map_efc2iefc,
+  map_iefc2efc, efc_islandid. Also populates d.dof_island, d.efc.island,
+  and d.island_dofadr.
+
+  Args:
+    m: Model.
+    d: Data.
+    ctx: IslandSolverContext.
+  """
+  # Ensure dof_islandid / efc_islandid are allocated at the right shape
+  if d.dof_islandid.shape[1] != m.nv:
+    d.dof_islandid = wp.empty((d.nworld, m.nv), dtype=int)
+  if d.efc_islandid.shape[1] != d.njmax:
+    d.efc_islandid = wp.empty((d.nworld, d.njmax), dtype=int)
+
+  # Ensure island-local DOF arrays are allocated at the right shape
+  if d.iqacc.shape[1] != m.nv:
+    nw = d.nworld
+    d.iqacc = wp.empty((nw, m.nv), dtype=float)
+    d.iqacc_smooth = wp.empty((nw, m.nv), dtype=float)
+    d.iqfrc_smooth = wp.empty((nw, m.nv), dtype=float)
+    d.iqfrc_constraint = wp.empty((nw, m.nv), dtype=float)
+  wp.launch(
+    _init_island_arrays,
+    dim=(d.nworld, m.ntree),
+    inputs=[],
+    outputs=[
+      d.island_dofadr,
+      d.island_nv,
+      d.island_nefc,
+      d.island_ne,
+      d.island_nf,
+      d.island_efcadr,
+      d.nidof,
+    ],
+  )
+  wp.launch(
+    _init_dof_arrays,
+    dim=(d.nworld, m.nv),
+    inputs=[],
+    outputs=[d.dof_island, d.map_dof2idof, d.map_idof2dof, d.dof_islandid],
+  )
+  efc_tree = wp.empty((d.nworld, d.njmax), dtype=int)
+  wp.launch(
+    _init_efc_arrays,
+    dim=(d.nworld, d.njmax),
+    inputs=[],
+    outputs=[d.efc.island, d.map_efc2iefc, d.map_iefc2efc, d.efc_islandid, efc_tree],
+  )
 
   wp.launch(
-    _flood_fill,
+    _compute_efc_tree,
-    dim=d.nworld,
+    dim=(d.nworld, d.njmax),
-    inputs=[m.ntree, tree_tree, d.tree_island, stack_scratch],
+    inputs=[
-    outputs=[d.nisland, d.tree_island, stack_scratch],
+      m.nv,
+      m.body_treeid,
+      m.jnt_dofadr,
+      m.dof_treeid,
+      m.geom_bodyid,
+      m.site_bodyid,
+      m.eq_type,
+      m.eq_obj1id,
+      m.eq_obj2id,
+      m.eq_objtype,
+      m.is_sparse,
+      d.nefc,
+      d.contact.geom,
+      d.efc.type,
+      d.efc.id,
+      d.efc.J,
+      d.efc.J_rownnz,
+      d.efc.J_rowadr,
+      d.efc.J_colind,
+      d.njmax,
+    ],
+    outputs=[efc_tree],
+  )
+
+  # 1. Count DOFs per island
+  wp.launch(
+    _island_count_dofs,
+    dim=(d.nworld, m.nv),
+    inputs=[m.dof_treeid, d.tree_island],
+    outputs=[d.dof_island, d.island_nv],
+  )
+
+  # 2. Count Constraints per island
+  wp.launch(
+    _island_count_constraints,
+    dim=(d.nworld, d.njmax),
+    inputs=[d.nefc, d.njmax, efc_tree, d.tree_island, d.efc.type],
+    outputs=[d.efc.island, d.island_nefc, d.island_ne, d.island_nf],
+  )
+
+  # 3. Scan sizes and reset counters for mapping
+  wp.launch(
+    _island_scan_sizes,
+    dim=d.nworld,
+    inputs=[d.nisland],
+    outputs=[d.island_dofadr, d.island_nv, d.island_nefc, d.island_efcadr, d.nidof],
+  )
+
+  # 4. Map DOFs
+  unconstrained_cnt = wp.zeros((d.nworld, 1), dtype=int)
+  wp.launch(
+    _island_map_dofs,
+    dim=(d.nworld, m.nv),
+    inputs=[m.nv, d.dof_island, d.island_dofadr, d.nidof],
+    outputs=[d.island_nv, d.map_dof2idof, d.map_idof2dof, d.dof_islandid, unconstrained_cnt],
+  )
+
+  # 5. Map Constraints
+  ne_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
+  nf_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
+  nother_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
+
+  wp.launch(
+    _island_map_constraints,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      d.nefc,
+      d.njmax,
+      d.efc.island,
+      d.island_efcadr,
+      d.island_ne,
+      d.island_nf,
+      d.efc.type,
+      ne_mapped,
+      nf_mapped,
+      nother_mapped,
+    ],
+    outputs=[d.island_nefc, d.map_efc2iefc, d.map_iefc2efc, d.efc_islandid],
+  )
+
+  # 6. Scan Sparse Rows (if sparse)
+  if m.is_sparse:
+    wp.launch(
+      _island_scan_sparse_rows,
+      dim=d.nworld,
+      inputs=[d.nisland, d.efc.J_rownnz, d.island_nefc, d.island_efcadr, d.map_iefc2efc],
+      outputs=[d.efc.iJ_rownnz, d.efc.iJ_rowadr],
+    )
+
+
+@event_scope
+def gather_island_inputs(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Gather constraint and DOF arrays into island-local order.
+
+  Populates d.iefc (D, type, id, frictionloss, aref, J, J_colind) and
+  d.iqacc, d.iqacc_smooth, d.iqfrc_smooth.
+
+  Must be called after compute_island_mapping() and before per-island solving.
+
+  Args:
+    m: Model.
+    d: Data.
+    ctx: IslandSolverContext whose arrays are populated.
+  """
+  # Gather constraint arrays and dense Jacobian (fused)
+  wp.launch(
+    _gather_efc_and_jacobian,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.is_sparse,
+      d.nefc,
+      d.efc.D,
+      d.efc.type,
+      d.efc.id,
+      d.efc.frictionloss,
+      d.efc.aref,
+      d.efc.J,
+      d.efc.J_rownnz,
+      d.efc.J_rowadr,
+      d.efc.J_colind,
+      d.efc.iJ_rowadr,
+      d.njmax,
+      d.map_iefc2efc,
+      d.map_idof2dof,
+      d.map_dof2idof,
+      d.nidof,
+    ],
+    outputs=[
+      d.efc.iD,
+      d.efc.itype,
+      d.efc.iid,
+      d.efc.ifrictionloss,
+      d.efc.iaref,
+      d.efc.iJ,
+      d.efc.iJ_colind,
+    ],
+  )
+
+  # Gather DOF arrays
+  wp.launch(
+    _gather_dof_arrays,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.qacc,
+      d.qacc_smooth,
+      d.qfrc_smooth,
+      d.nidof,
+      d.map_idof2dof,
+    ],
+    outputs=[
+      d.iqacc,
+      d.iqacc_smooth,
+      d.iqfrc_smooth,
+    ],
+  )
+
+
+@event_scope
+def scatter_island_results(m: types.Model, d: types.Data, ctx: IslandSolverContext, scatter_Ma: bool):
+  """Scatter island-local solver results back to global arrays.
+
+  Reads ctx qacc, qfrc_constraint, Ma, d.efc iforce, istate and
+  writes them back to d.qacc, d.qfrc_constraint, d.efc.Ma, d.efc.force,
+  d.efc.state. Unconstrained DOFs receive qacc_smooth and zero qfrc.
+
+  Args:
+    m: Model.
+    d: Data.
+    ctx: IslandSolverContext which contains results.
+    scatter_Ma: Whether to scatter Ma for Euler/implicit integrators.
+  """
+  # Scatter DOF results (and optionally Ma for Euler/implicit integrators)
+  wp.launch(
+    _scatter_dof_arrays,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.qacc_smooth,
+      d.qfrc_smooth,
+      d.dof_island,
+      d.iqacc,
+      d.iqfrc_constraint,
+      ctx.Ma,
+      d.map_dof2idof,
+      scatter_Ma,
+    ],
+    outputs=[
+      d.qacc,
+      d.qfrc_constraint,
+      d.efc.Ma,
+    ],
+  )
+
+  # Scatter constraint results (force and state from island-local arrays)
+  wp.launch(
+    _scatter_efc_arrays,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      d.nefc,
+      d.njmax,
+      d.map_iefc2efc,
+      d.efc.iforce,
+      d.efc.istate,
+    ],
+    outputs=[
+      d.efc.force,
+      d.efc.state,
+    ],
   )

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

@@ -707,10 +707,14 @@ def ray_mesh_with_bvh(
   pnt: wp.vec3,
   vec: wp.vec3,
   max_t: float,
+  cull_backfaces: bool,
 ) -> Tuple[float, wp.vec3, float, float, int, int]:
   """Returns intersection information for ray mesh intersections.
 
   Requires wp.Mesh be constructed and their ids to be passed.
+
+  When ``cull_backfaces`` is True, the function rejects exit-face hits in the
+  local-space frame. This matches MuJoCo OpenGL rendering's backface culling rule.
   """
   t = float(-1.0)
   u = float(0.0)
@@ -722,7 +726,12 @@ def ray_mesh_with_bvh(
   lpnt, lvec = _ray_map(pos, mat, pnt, vec)
   hit = wp.mesh_query_ray(mesh_bvh_id[mesh_geom_id], lpnt, lvec, max_t, t, u, v, sign, n, f)
 
-  if hit and wp.dot(lvec, n) < 0.0:  # Backface culling in local space
+  if not hit:
+    return -1.0, wp.vec3(0.0, 0.0, 0.0), 0.0, 0.0, -1, -1
+
+  # Front-face hit, or back-face hit when cull is disabled: rotate the
+  # local-space normal into world space and return the hit.
+  if (not cull_backfaces) or wp.dot(lvec, n) < 0.0:
     normal = mat @ n
     normal = wp.normalize(normal)
     return t, normal, u, v, f, mesh_geom_id
@@ -1055,6 +1064,8 @@ def _ray_geom_mesh_bvh(
 
     if gtype == GeomType.MESH or gtype == GeomType.HFIELD:
       bvh_ids = mesh_bvh_id if gtype == GeomType.MESH else hfield_bvh_id
+      # Public ray API (mjw.ray / mjw.rays) preserves MuJoCo's mj_ray cull rule:
+      # rangefinder sensors and user-facing ray casts always cull back-faces.
       t, n, u, v, f, geom_mesh_id = ray_mesh_with_bvh(
         bvh_ids,
         geom_dataid[worldid % geom_dataid.shape[0], geomid],
@@ -1063,6 +1074,7 @@ def _ray_geom_mesh_bvh(
         pnt,
         vec,
         min_dist,
+        True,
       )
       if t >= 0.0 and t < min_dist:
         return t, n

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

@@ -180,6 +180,7 @@ def cast_ray(
   flex_group_root: wp.array2d[int],
   ray_origin_world: wp.vec3,
   ray_dir_world: wp.vec3,
+  cull_backfaces: bool,
 ) -> Tuple[int, float, wp.vec3, float, float, int, int]:
   dist = float(MJ_MAXVAL)
   normal = wp.vec3(0.0, 0.0, 0.0)
@@ -232,6 +233,7 @@ def cast_ray(
         ray_origin_world,
         ray_dir_world,
         dist,
+        cull_backfaces,
       )
     if gtype == GeomType.SPHERE:
       d, n = ray_sphere(
@@ -281,6 +283,7 @@ def cast_ray(
         ray_origin_world,
         ray_dir_world,
         dist,
+        cull_backfaces,
       )
     if gtype == GeomType.FLEX:
       hit_geom_id = -2
@@ -308,6 +311,11 @@ def cast_ray(
         if d >= 0.0:
           hit_mesh_id = flexid
 
+    # Backface cull: drop exit-face hits when the ray origin is inside the geom,
+    # matching ray_mesh_with_bvh's `dot(lvec, n) < 0` rule.
+    if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0:
+      d = -1.0
+
     if d >= 0.0 and d < dist:
       dist = d
       normal = n
@@ -349,6 +357,7 @@ def cast_ray_first_hit(
   ray_origin_world: wp.vec3,
   ray_dir_world: wp.vec3,
   max_dist: float,
+  cull_backfaces: bool,
 ) -> bool:
   """A simpler version of casting rays that only checks for the first hit."""
   query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
@@ -387,6 +396,7 @@ def cast_ray_first_hit(
         ray_origin_world,
         ray_dir_world,
         max_dist,
+        cull_backfaces,
       )
     if gtype == GeomType.SPHERE:
       d, n = ray_sphere(
@@ -467,6 +477,11 @@ def cast_ray_first_hit(
         )
         d = 0.0 if hit else -1.0
 
+    # Backface cull: see cast_ray for rationale. Strict `> 0` keeps tangent
+    # hits and skips branches with a zero-vector normal (mesh/flex anyhit).
+    if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0:
+      d = -1.0
+
     if d >= 0.0 and d < max_dist:
       return True
 
@@ -507,6 +522,7 @@ def compute_lighting(
   lightdir: wp.vec3,
   normal: wp.vec3,
   hitpoint: wp.vec3,
+  cull_backfaces: bool,
 ) -> float:
   light_contribution = float(0.0)
 
@@ -571,6 +587,7 @@ def compute_lighting(
       shadow_origin,
       L,
       max_t,
+      cull_backfaces,
     )
 
     if shadow_hit:
@@ -725,6 +742,7 @@ def render(m: Model, d: Data, rc: RenderContext):
       flex_group_root,
       ray_origin_world,
       ray_dir_world,
+      wp.static(rc.enable_backface_culling),
     )
 
     if render_seg[cam_idx] and geom_id != -1:
@@ -798,6 +816,8 @@ def render(m: Model, d: Data, rc: RenderContext):
             )
             base_color = wp.cw_mul(base_color, tex_color)
 
+    result = wp.vec3(0.0, 0.0, 0.0)
+    if wp.static(rc.use_ambient_lighting):
       len_n = wp.length(normal)
       n = normal if len_n > 0.0 else wp.vec3(0.0, 0.0, 1.0)
       n = wp.normalize(n)
@@ -837,6 +857,7 @@ def render(m: Model, d: Data, rc: RenderContext):
         light_xdir_in[worldid, l],
         normal,
         hit_point,
+        wp.static(rc.enable_backface_culling),
       )
       result = result + base_color * light_contribution
 

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

@@ -17,6 +17,7 @@ from typing import Any, Tuple
 
 import warp as wp
 
+from mujoco.mjx.third_party.mujoco_warp._src import history
 from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import ray
 from mujoco.mjx.third_party.mujoco_warp._src import smooth
@@ -26,6 +27,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import TACTILE_DEPTH_SEMANTICS
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
 from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
@@ -917,6 +919,10 @@ def sensor_pos(m: Model, d: Data):
     ],
   )
 
+  # apply sensor delay/interval for position sensors
+  history.apply_sensor_delay(m, d, m.sensor_pos_adr)
+  history.apply_sensor_delay(m, d, m.sensor_limitpos_adr)
+
   if m.callback.sensor:
     m.callback.sensor(m, d, Stage.POS)
 
@@ -1459,6 +1465,10 @@ def sensor_vel(m: Model, d: Data):
     ],
   )
 
+  # apply sensor delay/interval for velocity sensors
+  history.apply_sensor_delay(m, d, m.sensor_vel_adr)
+  history.apply_sensor_delay(m, d, m.sensor_limitvel_adr)
+
   if m.callback.sensor:
     m.callback.sensor(m, d, Stage.VEL)
 
@@ -2253,13 +2263,23 @@ def _sensor_tactile(
     vel_rel = vel_sensor - vel_other
 
     forceT = wp.vec3(0.0, 0.0, 0.0)
+    if wp.static(TACTILE_DEPTH_SEMANTICS):
       forceT[0] = -depth
+    else:
+      kMaxDepth = 0.05
+      pressure = depth / wp.max(kMaxDepth - depth, MJ_MINVAL)
+      force = wp.mul(normal, pressure)
+      forceT[0] = wp.dot(force, normal)
+
     if has_frame:
       forceT[1] = wp.abs(wp.dot(vel_rel, tang1))
       forceT[2] = wp.abs(wp.dot(vel_rel, tang2))
 
     dim = sensor_dim[sensor_id] // 3
+    if wp.static(TACTILE_DEPTH_SEMANTICS):
       wp.atomic_max(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
+    else:
+      wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
     wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 1 * dim + vertid, forceT[1])
     wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 2 * dim + vertid, forceT[2])
 
@@ -2705,6 +2725,10 @@ def sensor_acc(m: Model, d: Data):
     ],
   )
 
+  # apply sensor delay/interval for acceleration sensors
+  history.apply_sensor_delay(m, d, m.sensor_acc_adr)
+  history.apply_sensor_delay(m, d, m.sensor_limitfrc_adr)
+
   if m.callback.sensor:
     m.callback.sensor(m, d, Stage.ACC)
 

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

@@ -823,37 +823,38 @@ def _crb_accumulate(
 
 
 @wp.kernel
-def _qM_sparse(
+def _M_sparse(
   # Model:
   dof_bodyid: wp.array[int],
   dof_parentid: wp.array[int],
-  dof_Madr: wp.array[int],
   dof_armature: wp.array2d[float],
+  M_rownnz: wp.array[int],
+  M_rowadr: wp.array[int],
   # Data in:
   cdof_in: wp.array2d[wp.spatial_vector],
   crb_in: wp.array2d[vec10],
   # Data out:
-  qM_out: wp.array3d[float],
+  M_out: wp.array3d[float],
 ):
   worldid, dofid = wp.tid()
-  madr_ij = dof_Madr[dofid]  # dof_Madr is not batched
   bodyid = dof_bodyid[dofid]
+  madr_ij = M_rowadr[dofid] + M_rownnz[dofid] - 1
 
   # init M(i,i) with armature inertia
-  qM_out[worldid, 0, madr_ij] = dof_armature[worldid % dof_armature.shape[0], dofid]
+  M_out[worldid, 0, madr_ij] = dof_armature[worldid % dof_armature.shape[0], dofid]
 
   # precompute buf = crb_body_i * cdof_i
   buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid])
 
   # sparse backward pass over ancestors
   while dofid >= 0:
-    qM_out[worldid, 0, madr_ij] += wp.dot(cdof_in[worldid, dofid], buf)
+    M_out[worldid, 0, madr_ij] += wp.dot(cdof_in[worldid, dofid], buf)
-    madr_ij += 1
+    madr_ij -= 1
     dofid = dof_parentid[dofid]
 
 
 @wp.kernel
-def _qM_dense(
+def _M_dense(
   # Model:
   dof_bodyid: wp.array[int],
   dof_parentid: wp.array[int],
@@ -862,7 +863,7 @@ def _qM_dense(
   cdof_in: wp.array2d[wp.spatial_vector],
   crb_in: wp.array2d[vec10],
   # Data out:
-  qM_out: wp.array3d[float],
+  M_out: wp.array3d[float],
 ):
   worldid, dofid = wp.tid()
   bodyid = dof_bodyid[dofid]
@@ -873,15 +874,15 @@ def _qM_dense(
   buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid])
   M += wp.dot(cdof_in[worldid, dofid], buf)
 
-  qM_out[worldid, dofid, dofid] = M
+  M_out[worldid, dofid, dofid] = M
 
   # sparse backward pass over ancestors
   dofidi = dofid
   dofid = dof_parentid[dofid]
   while dofid >= 0:
-    qMij = wp.dot(cdof_in[worldid, dofid], buf)
+    Mij = wp.dot(cdof_in[worldid, dofid], buf)
-    qM_out[worldid, dofidi, dofid] += qMij
+    M_out[worldid, dofidi, dofid] += Mij
-    qM_out[worldid, dofid, dofidi] += qMij
+    M_out[worldid, dofid, dofidi] += Mij
     dofid = dof_parentid[dofid]
 
 
@@ -898,17 +899,17 @@ def crb(m: Model, d: Data):
     body_tree = m.body_tree[i]
     wp.launch(_crb_accumulate, dim=(d.nworld, body_tree.size), inputs=[m.body_parentid, d.crb, body_tree], outputs=[d.crb])
 
-  d.qM.zero_()
+  d.M.zero_()
   if m.is_sparse:
     wp.launch(
-      _qM_sparse,
+      _M_sparse,
       dim=(d.nworld, m.nv),
-      inputs=[m.dof_bodyid, m.dof_parentid, m.dof_Madr, m.dof_armature, d.cdof, d.crb],
+      inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, m.M_rownnz, m.M_rowadr, d.cdof, d.crb],
-      outputs=[d.qM],
+      outputs=[d.M],
     )
   else:
     wp.launch(
-      _qM_dense, dim=(d.nworld, m.nv), inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, d.cdof, d.crb], outputs=[d.qM]
+      _M_dense, dim=(d.nworld, m.nv), inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, d.cdof, d.crb], outputs=[d.M]
     )
 
 
@@ -916,16 +917,17 @@ def crb(m: Model, d: Data):
 def _tendon_armature(
   # Model:
   dof_parentid: wp.array[int],
-  dof_Madr: wp.array[int],
   ten_J_rownnz: wp.array[int],
   ten_J_rowadr: wp.array[int],
   ten_J_colind: wp.array[int],
   tendon_armature: wp.array2d[float],
+  M_rownnz: wp.array[int],
+  M_rowadr: wp.array[int],
   is_sparse: bool,
   # Data in:
   ten_J_in: wp.array2d[float],
   # Data out:
-  qM_out: wp.array3d[float],
+  M_out: wp.array3d[float],
 ):
   worldid, tenid, dofid = wp.tid()
 
@@ -947,7 +949,7 @@ def _tendon_armature(
     return
 
   if is_sparse:
-    madr_ij = dof_Madr[dofid]
+    madr_ij = M_rowadr[dofid] + M_rownnz[dofid] - 1
 
   # sparse backward pass over ancestors
   dofidi = dofid
@@ -967,52 +969,40 @@ def _tendon_armature(
       else:
         ten_Jj = float(0.0)
 
-    qMij = armature * ten_Jj * ten_Ji
+    Mij = armature * ten_Jj * ten_Ji
 
     if is_sparse:
-      wp.atomic_add(qM_out[worldid, 0], madr_ij, qMij)
+      wp.atomic_add(M_out[worldid, 0], madr_ij, Mij)
-      madr_ij += 1
+      madr_ij -= 1
     else:
-      wp.atomic_add(qM_out[worldid, dofidi], dofid, qMij)
+      wp.atomic_add(M_out[worldid, dofidi], dofid, Mij)
       if dofidi != dofid:
-        wp.atomic_add(qM_out[worldid, dofid], dofidi, qMij)
+        wp.atomic_add(M_out[worldid, dofid], dofidi, Mij)
 
     dofid = dof_parentid[dofid]
 
 
 @event_scope
 def tendon_armature(m: Model, d: Data):
-  """Add tendon armature to qM."""
+  """Add tendon armature to M."""
   wp.launch(
     _tendon_armature,
     dim=(d.nworld, m.ntendon, m.max_ten_J_rownnz),
     inputs=[
       m.dof_parentid,
-      m.dof_Madr,
       m.ten_J_rownnz,
       m.ten_J_rowadr,
       m.ten_J_colind,
       m.tendon_armature,
+      m.M_rownnz,
+      m.M_rowadr,
       m.is_sparse,
       d.ten_J,
     ],
-    outputs=[d.qM],
+    outputs=[d.M],
   )
 
 
-@wp.kernel
-def _copy_CSR(
-  # Model:
-  mapM2M: wp.array[int],
-  # In:
-  M_in: wp.array3d[float],
-  # Out:
-  L_out: wp.array3d[float],
-):
-  worldid, ind = wp.tid()
-  L_out[worldid, 0, ind] = M_in[worldid, 0, mapM2M[ind]]
-
-
 @wp.kernel
 def _qLD_acc(
   # Model:
@@ -1055,7 +1045,7 @@ def _qLDiag_div(
 
 def _factor_i_sparse(m: Model, d: Data, M: wp.array3d[float], L: wp.array3d[float], D: wp.array2d[float]):
   """Sparse L'*D*L factorization of inertia-like matrix M, assumed spd."""
-  wp.launch(_copy_CSR, dim=(d.nworld, m.nC), inputs=[m.mapM2M, M], outputs=[L])
+  wp.copy(L, M)
 
   for i in reversed(range(len(m.qLD_updates))):
     qLD_updates = m.qLD_updates[i]
@@ -1071,7 +1061,7 @@ def _tile_cholesky_factorize(tile: TileSet):
   @wp.kernel(module="unique", enable_backward=False)
   def cholesky_factorize(
     # Data in:
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     # In:
     adr: wp.array[int],
     # Out:
@@ -1081,8 +1071,8 @@ def _tile_cholesky_factorize(tile: TileSet):
     TILE_SIZE = wp.static(tile.size)
 
     dofid = adr[nodeid]
-    M_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
+    M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
-    L_tile = wp.tile_cholesky(M_tile)
+    L_tile = wp.tile_cholesky(M_tile, fill_mode="upper")
     wp.tile_store(L_out[worldid], L_tile, offset=(dofid, dofid))
 
   return cholesky_factorize
@@ -1090,7 +1080,7 @@ def _tile_cholesky_factorize(tile: TileSet):
 
 def _factor_i_dense(m: Model, d: Data, M: wp.array, L: wp.array):
   """Dense Cholesky factorization of inertia-like matrix M, assumed spd."""
-  for tile in m.qM_tiles:
+  for tile in m.M_tiles:
     wp.launch_tiled(
       _tile_cholesky_factorize(tile),
       dim=(d.nworld, tile.adr.size),
@@ -1104,9 +1094,9 @@ def _factor_i_dense(m: Model, d: Data, M: wp.array, L: wp.array):
 def factor_m(m: Model, d: Data):
   """Factorization of inertia-like matrix M, assumed spd."""
   if m.is_sparse:
-    _factor_i_sparse(m, d, d.qM, d.qLD, d.qLDiagInv)
+    _factor_i_sparse(m, d, d.M, d.qLD, d.qLDiagInv)
   else:
-    _factor_i_dense(m, d, d.qM, d.qLD)
+    _factor_i_dense(m, d, d.M, d.qLD)
 
 
 @wp.kernel
@@ -2842,15 +2832,15 @@ def _tile_cholesky_solve(tile: TileSet):
     dofid = adr[nodeid]
     y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(dofid,))
     L_tile = wp.tile_load(L[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
-    x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
+    x_slice = wp.tile_cholesky_solve(L_tile, y_slice, fill_mode="upper")
     wp.tile_store(x[worldid], x_slice, offset=(dofid,))
 
   return cholesky_solve
 
 
 def _solve_LD_dense(m: Model, d: Data, L: wp.array3d[float], x: wp.array2d[float], y: wp.array2d[float]):
-  """Computes dense backsubstitution: x = inv(L'*L)*y."""
+  """Computes dense backsubstitution: x = inv(U.T @ U) * y."""
-  for tile in m.qM_tiles:
+  for tile in m.M_tiles:
     wp.launch_tiled(
       _tile_cholesky_solve(tile),
       dim=(d.nworld, tile.adr.size),
@@ -2868,9 +2858,9 @@ def solve_LD(
   x: wp.array2d[float],
   y: wp.array2d[float],
 ):
-  """Computes backsubstitution to solve a linear system of the form x = inv(L'*D*L) * y.
+  """Computes backsubstitution for the inertia factorization.
 
-  L and D are the factors from the Cholesky factorization of the inertia matrix.
+  Sparse models use MuJoCo's L'*D*L factors; dense models use an upper Cholesky factor U.
 
   This function dispatches to either a sparse or dense solver depending on Model options.
 
@@ -2907,8 +2897,9 @@ def _tile_cholesky_factorize_solve(tile: TileSet):
 
   @wp.kernel(module="unique", enable_backward=False)
   def cholesky_factorize_solve(
+    # Data in:
+    M_in: wp.array3d[float],
     # In:
-    M: wp.array3d[float],
     y: wp.array2d[float],
     adr: wp.array[int],
     # Out:
@@ -2919,12 +2910,12 @@ def _tile_cholesky_factorize_solve(tile: TileSet):
     TILE_SIZE = wp.static(tile.size)
 
     dofid = adr[nodeid]
-    M_tile = wp.tile_load(M[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
+    M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
     y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(dofid,))
 
-    L_tile = wp.tile_cholesky(M_tile)
+    L_tile = wp.tile_cholesky(M_tile, fill_mode="upper")
     wp.tile_store(L[worldid], L_tile, offset=(dofid, dofid))
-    x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
+    x_slice = wp.tile_cholesky_solve(L_tile, y_slice, fill_mode="upper")
     wp.tile_store(x[worldid], x_slice, offset=(dofid,))
 
   return cholesky_factorize_solve
@@ -2938,7 +2929,7 @@ def _factor_solve_i_dense(
   y: wp.array2d[float],
   L: wp.array3d[float],
 ):
-  for tile in m.qM_tiles:
+  for tile in m.M_tiles:
     wp.launch_tiled(
       _tile_cholesky_factorize_solve(tile),
       dim=(d.nworld, tile.adr.size),
@@ -2949,9 +2940,9 @@ def _factor_solve_i_dense(
 
 
 def factor_solve_i(m, d, M, L, D, x, y):
-  """Factorizes and solves the linear system: x = inv(L'*D*L) * y or x = inv(L'*L) * y.
+  """Factorizes and solves the inertia-like linear system.
 
-  M is an inertia-like matrix and L, D are its Cholesky-like factors.
+  Sparse models use MuJoCo's L'*D*L factors; dense models use an upper Cholesky factor U.
 
   This function first factorizes the matrix M (sparse or dense depending on model options),
   then solves the system for x given right-hand side y.
@@ -2960,7 +2951,7 @@ def factor_solve_i(m, d, M, L, D, x, y):
     m: The model containing factorization and sparsity information.
     d: The data object containing workspace and factorization results.
     M: The inertia-like matrix to factorize.
-    L: Output lower-triangular factor from the factorization (sparse or dense).
+    L: Output sparse factor or dense upper Cholesky factor.
     D: Output diagonal factor from the factorization (only used for sparse).
     x: Output array for the solution.
     y: Input right-hand side array.
@@ -2972,6 +2963,132 @@ def factor_solve_i(m, d, M, L, D, x, y):
     _factor_solve_i_dense(m, d, M, x, y, L)
 
 
+@cache_kernel
+def _factor_solve_lu_sparse_fused(nv: int):
+  """Fused sparse LU factorization and solve in a single kernel."""
+
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # Model:
+    D_rownnz: wp.array[int],
+    D_rowadr: wp.array[int],
+    D_diag: wp.array[int],
+    D_colind: wp.array[int],
+    # In:
+    qfrc: wp.array2d[float],
+    # Data out:
+    qacc_out: wp.array2d[float],
+    qLU_out: wp.array3d[float],
+  ):
+    worldid = wp.tid()
+    NV = wp.static(nv)
+
+    # Phase 1: LU factorization (in-place on qLU_out)
+    for i in range(NV):
+      qacc_out[worldid, i] = float(D_rownnz[i])
+
+    # process diagonal elements from n-1 down to 0
+    for r_rev in range(NV):
+      i = NV - 1 - r_rev
+
+      rem_i = int(qacc_out[worldid, i])
+      rowadr_i = D_rowadr[i]
+      ii = rowadr_i + rem_i - 1
+      qacc_out[worldid, i] = float(rem_i - 1)
+
+      # cache diagonal element for row i
+      LUii = qLU_out[worldid, 0, ii]
+
+      # rows j above i (j < i), processed from i-1 down to 0
+      for c in range(i):
+        j = i - 1 - c
+
+        # get address of last remaining element of row j
+        rem_j = int(qacc_out[worldid, j])
+        rowadr_j = D_rowadr[j]
+        ji = rowadr_j + rem_j - 1
+
+        # process row j if (j,i) is non-zero
+        if D_colind[ji] == i:
+          # adjust remaining counter
+          rem_j = rem_j - 1
+          qacc_out[worldid, j] = float(rem_j)
+
+          # (j,i) = (j,i) / (i,i)
+          LUji = qLU_out[worldid, 0, ji] / LUii
+          qLU_out[worldid, 0, ji] = LUji
+
+          # (j,k) = (j,k) - (i,k) * (j,i) for k < i
+          icnt = rowadr_i
+          jcnt = rowadr_j
+          jend = rowadr_j + rem_j
+          while jcnt < jend:
+            col_i = D_colind[icnt]
+            col_j = D_colind[jcnt]
+            if col_i == col_j:
+              qLU_out[worldid, 0, jcnt] = qLU_out[worldid, 0, jcnt] - qLU_out[worldid, 0, icnt] * LUji
+              icnt = icnt + 1
+              jcnt = jcnt + 1
+            elif col_i > col_j:
+              jcnt = jcnt + 1
+            else:
+              icnt = icnt + 1
+
+    # Phase 2: LU solve (backward + forward substitution)
+
+    # Backward substitution: solve (U+I)*qacc = qfrc
+    for k_rev in range(NV):
+      i = NV - 1 - k_rev
+
+      diag_i = D_diag[i]
+      rowadr_i = D_rowadr[i]
+      d1 = diag_i + 1
+      nnz_upper = D_rownnz[i] - d1
+
+      acc = qfrc[worldid, i]
+      for j in range(nnz_upper):
+        adr_j = rowadr_i + d1 + j
+        col = D_colind[adr_j]
+        acc = acc - qLU_out[worldid, 0, adr_j] * qacc_out[worldid, col]
+      qacc_out[worldid, i] = acc
+
+    # Forward substitution: solve L*qacc = qacc
+    for i in range(NV):
+      diag_i = D_diag[i]
+      rowadr_i = D_rowadr[i]
+
+      acc = qacc_out[worldid, i]
+      for j in range(diag_i):
+        adr_j = rowadr_i + j
+        col = D_colind[adr_j]
+        acc = acc - qLU_out[worldid, 0, adr_j] * qacc_out[worldid, col]
+
+      qacc_out[worldid, i] = acc / qLU_out[worldid, 0, rowadr_i + diag_i]
+
+  return kernel
+
+
+@event_scope
+def factor_solve_lu(m: Model, d: Data, qLU: wp.array3d[float], qacc: wp.array2d[float], qfrc: wp.array2d[float]):
+  r"""Factorize and solve non-symmetric implicit system: qacc = A \\ qfrc.
+
+  qLU is overwritten in-place with the LU factors, then used to solve for qacc.
+
+  Args:
+    m: The model containing D-structure sparsity information.
+    d: The data object.
+    qLU: array containing the system matrix, overwritten with LU factors.
+    qacc: output array for the solution.
+    qfrc: input right-hand side.
+  """
+  wp.launch(
+    _factor_solve_lu_sparse_fused(m.nv),
+    dim=(d.nworld,),
+    inputs=[m.D_rownnz, m.D_rowadr, m.D_diag, m.D_colind, qfrc],
+    outputs=[qacc, qLU],
+  )
+
+
 @wp.kernel
 def _subtree_vel_forward(
   # Model:

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

@@ -69,11 +69,11 @@ def mul_m_sparse(check_skip: bool):
   @wp.kernel(module="unique")
   def _mul_m_sparse(
     # Model:
-    qM_mulm_rowadr: wp.array[int],
+    M_mulm_rowadr: wp.array[int],
-    qM_mulm_col: wp.array[int],
+    M_mulm_col: wp.array[int],
-    qM_mulm_madr: wp.array[int],
+    M_mulm_madr: wp.array[int],
     # Data in:
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     # In:
     vec: wp.array2d[float],
     skip: wp.array[bool],
@@ -89,12 +89,12 @@ def mul_m_sparse(check_skip: bool):
 
     # Gather all contributions (diagonal + off-diagonal)
     acc = float(0.0)
-    start = qM_mulm_rowadr[dofid]
+    start = M_mulm_rowadr[dofid]
-    end = qM_mulm_rowadr[dofid + 1]
+    end = M_mulm_rowadr[dofid + 1]
     for k in range(start, end):
-      col = qM_mulm_col[k]
+      col = M_mulm_col[k]
-      madr = qM_mulm_madr[k]
+      madr = M_mulm_madr[k]
-      acc += qM_in[worldid, 0, madr] * vec[worldid, col]
+      acc += M_in[worldid, 0, madr] * vec[worldid, col]
 
     res[worldid, dofid] = acc
 
@@ -108,7 +108,7 @@ def mul_m_dense(nv: int, check_skip: bool):
   @wp.kernel(module="unique")
   def _mul_m_dense(
     # Data in:
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     # In:
     vec: wp.array2d[float],
     skip: wp.array[bool],
@@ -123,7 +123,7 @@ def mul_m_dense(nv: int, check_skip: bool):
 
     acc = float(0.0)
     for j in range(wp.static(nv)):
-      acc += qM_in[worldid, i, j] * vec[worldid, j]
+      acc += M_in[worldid, i, j] * vec[worldid, j]
     res[worldid, i] = acc
 
   return _mul_m_dense
@@ -143,8 +143,8 @@ def mul_m(
   Args:
     m: The model containing kinematic and dynamic information (device).
     d: The data object containing the current state and output arrays (device).
-    res: Result: qM @ vec.
+    res: Result: M @ vec.
-    vec: Input vector to multiply by qM.
+    vec: Input vector to multiply by M.
     skip: Per-world bitmask to skip computing output.
     M: Input matrix: M @ vec.
   """
@@ -152,13 +152,13 @@ def mul_m(
   skip = skip or wp.empty(0, dtype=bool)
 
   if M is None:
-    M = d.qM
+    M = d.M
 
   if m.is_sparse:
     wp.launch(
       mul_m_sparse(check_skip),
       dim=(d.nworld, m.nv),
-      inputs=[m.qM_mulm_rowadr, m.qM_mulm_col, m.qM_mulm_madr, M, vec, skip],
+      inputs=[m.M_mulm_rowadr, m.M_mulm_col, m.M_mulm_madr, M, vec, skip],
       outputs=[res],
     )
 
@@ -171,6 +171,340 @@ def mul_m(
     )
 
 
+@wp.kernel
+def _mul_m_island_sparse(
+  # Model:
+  M_mulm_rowadr: wp.array[int],
+  M_mulm_col: wp.array[int],
+  M_mulm_madr: wp.array[int],
+  # Data in:
+  nidof_in: wp.array[int],
+  M_in: wp.array3d[float],
+  map_dof2idof_in: wp.array2d[int],
+  map_idof2dof_in: wp.array2d[int],
+  # In:
+  idof_islandid_in: wp.array2d[int],
+  vec: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  check_skip: int,
+  # Out:
+  res: wp.array2d[float],
+):
+  """Sparse island mul_m for ALL islands in parallel."""
+  worldid, idofid = wp.tid()
+
+  nidof = nidof_in[worldid]
+  if idofid >= nidof:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+
+  if check_skip:
+    if island_done_in[worldid, islandid]:
+      return
+
+  dof = map_idof2dof_in[worldid, idofid]
+
+  acc = float(0.0)
+  start = M_mulm_rowadr[dof]
+  end = M_mulm_rowadr[dof + 1]
+  for k in range(start, end):
+    col = M_mulm_col[k]
+    madr = M_mulm_madr[k]
+    col_idof = map_dof2idof_in[worldid, col]
+    # skip unconstrained DOFs
+    if col_idof < nidof:
+      acc += M_in[worldid, 0, madr] * vec[worldid, col_idof]
+
+  res[worldid, idofid] = acc
+
+
+@wp.kernel
+def _mul_m_island_dense(
+  # Model:
+  nv: int,
+  # Data in:
+  nidof_in: wp.array[int],
+  M_in: wp.array3d[float],
+  map_dof2idof_in: wp.array2d[int],
+  map_idof2dof_in: wp.array2d[int],
+  # In:
+  idof_islandid_in: wp.array2d[int],
+  vec: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  check_skip: int,
+  # Out:
+  res: wp.array2d[float],
+):
+  """Dense island mul_m for ALL islands in parallel."""
+  worldid, idofid = wp.tid()
+
+  nidof = nidof_in[worldid]
+  if idofid >= nidof:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+
+  if check_skip:
+    if island_done_in[worldid, islandid]:
+      return
+
+  dof = map_idof2dof_in[worldid, idofid]
+
+  acc = float(0.0)
+  for j in range(nv):
+    col_idof = map_dof2idof_in[worldid, j]
+    # skip unconstrained DOFs
+    if col_idof < nidof:
+      acc += M_in[worldid, dof, j] * vec[worldid, col_idof]
+
+  res[worldid, idofid] = acc
+
+
+@event_scope
+def mul_m_island(
+  m: Model,
+  d: Data,
+  res: wp.array2d[float],
+  vec: wp.array2d[float],
+  nidof: wp.array[int],
+  map_idof2dof: wp.array2d[int],
+  map_dof2idof: wp.array2d[int],
+  idof_islandid: wp.array2d[int],
+  island_done: Optional[wp.array] = None,
+  M: Optional[wp.array] = None,
+):
+  """Multiply island-local vectors by inertia matrix for all islands in parallel.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    res: Result: qM @ vec (island-local DOF order).
+    vec: Input vector (island-local DOF order).
+    nidof: Number of island DOFs per world.
+    map_idof2dof: Island-local DOF → global DOF map.
+    map_dof2idof: Global DOF → island-local DOF map.
+    idof_islandid: Island ID per island-local DOF.
+    island_done: Per-island done flags (nworld, ntree).
+    M: Optional mass matrix override.
+  """
+  check_skip = int(island_done is not None)
+  island_done = island_done or wp.empty((0, 0), dtype=bool)
+
+  if M is None:
+    M = d.M
+
+  if m.is_sparse:
+    wp.launch(
+      _mul_m_island_sparse,
+      dim=(d.nworld, m.nv),
+      inputs=[
+        m.M_mulm_rowadr,
+        m.M_mulm_col,
+        m.M_mulm_madr,
+        nidof,
+        M,
+        map_dof2idof,
+        map_idof2dof,
+        idof_islandid,
+        vec,
+        island_done,
+        check_skip,
+      ],
+      outputs=[res],
+    )
+  else:
+    wp.launch(
+      _mul_m_island_dense,
+      dim=(d.nworld, m.nv),
+      inputs=[
+        m.nv,
+        nidof,
+        M,
+        map_dof2idof,
+        map_idof2dof,
+        idof_islandid,
+        vec,
+        island_done,
+        check_skip,
+      ],
+      outputs=[res],
+    )
+
+
+@cache_kernel
+def _solve_LD_sparse_island(nv: int, nlevels: int):
+  """Sparse backsubstitution with island-local index remapping.
+
+  Same algorithm as _solve_LD_sparse_fused, but reads/writes x/y in
+  island-local DOF order. The L/D factorization stays in global DOF order.
+  """
+
+  @wp.func_native(snippet="WP_TILE_SYNC();")
+  def _syncthreads():
+    pass
+
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # Data in:
+    nidof_in: wp.array[int],
+    map_dof2idof_in: wp.array2d[int],
+    map_idof2dof_in: wp.array2d[int],
+    # In:
+    L: wp.array3d[float],
+    D: wp.array2d[float],
+    all_updates: wp.array[wp.vec3i],
+    level_offsets: wp.array[int],
+    y: wp.array2d[float],
+    # Out:
+    x_out: wp.array2d[float],
+  ):
+    worldid, tid = wp.tid()
+    NLEVELS = wp.static(nlevels)
+    BLOCK_DIM = wp.block_dim()
+    nid = nidof_in[worldid]
+
+    # Copy y to x_out (island-local, only iterate up to nid)
+    for idof in range(tid, nid, BLOCK_DIM):
+      x_out[worldid, idof] = y[worldid, idof]
+    _syncthreads()
+
+    # Forward substitution
+    for level in range(NLEVELS):
+      level_idx = NLEVELS - 1 - level
+      level_offset = level_offsets[level_idx]
+      level_size = level_offsets[level_idx + 1] - level_offset
+
+      for u in range(tid, level_size, BLOCK_DIM):
+        update = all_updates[level_offset + u]
+        i, k, Madr_ki = update[0], update[1], update[2]
+        idof_i = map_dof2idof_in[worldid, i]
+        if idof_i < nid:
+          idof_k = map_dof2idof_in[worldid, k]
+          wp.atomic_sub(x_out[worldid], idof_i, L[worldid, 0, Madr_ki] * x_out[worldid, idof_k])
+      _syncthreads()
+
+    # Diagonal multiply (only iterate up to nid)
+    for idof in range(tid, nid, BLOCK_DIM):
+      dofid = map_idof2dof_in[worldid, idof]
+      x_out[worldid, idof] *= D[worldid, dofid]
+    _syncthreads()
+
+    # Backward substitution
+    for level in range(NLEVELS):
+      level_idx = level
+      level_offset = level_offsets[level_idx]
+      level_size = level_offsets[level_idx + 1] - level_offset
+
+      for u in range(tid, level_size, BLOCK_DIM):
+        update = all_updates[level_offset + u]
+        i, k, Madr_ki = update[0], update[1], update[2]
+        idof_k = map_dof2idof_in[worldid, k]
+        if idof_k < nid:
+          idof_i = map_dof2idof_in[worldid, i]
+          wp.atomic_sub(x_out[worldid], idof_k, L[worldid, 0, Madr_ki] * x_out[worldid, idof_i])
+      _syncthreads()
+
+  return kernel
+
+
+@cache_kernel
+def _tile_cholesky_solve_island(tile):
+  """Dense Cholesky backsubstitution with island-local index remapping.
+
+  L is loaded from global DOF offsets (factorization unchanged).
+  y/x are loaded/stored at island-local DOF offsets via map_dof2idof.
+  """
+
+  @wp.kernel(module="unique", enable_backward=False)
+  def cholesky_solve(
+    # Data in:
+    nidof_in: wp.array[int],
+    map_dof2idof_in: wp.array2d[int],
+    # In:
+    L: wp.array3d[float],
+    y: wp.array2d[float],
+    adr: wp.array[int],
+    # Out:
+    x: wp.array2d[float],
+  ):
+    worldid, nodeid = wp.tid()
+    TILE_SIZE = wp.static(tile.size)
+
+    dofid = adr[nodeid]
+    idofid = map_dof2idof_in[worldid, dofid]
+
+    # Skip unconstrained trees (uniform branch — all threads in block agree)
+    if idofid >= nidof_in[worldid]:
+      return
+
+    # L stays in global order
+    L_tile = wp.tile_load(L[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
+    # y and x use island-local offsets
+    y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(idofid,))
+    x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
+    wp.tile_store(x[worldid], x_slice, offset=(idofid,))
+
+  return cholesky_solve
+
+
+@event_scope
+def solve_m_island(
+  m: Model,
+  d: Data,
+  res: wp.array2d[float],
+  vec: wp.array2d[float],
+  nidof: wp.array[int],
+  map_idof2dof: wp.array2d[int],
+):
+  """Compute res = M^{-1} @ vec for island-local DOFs.
+
+  Args:
+    m: Model.
+    d: Data.
+    res: Output in island-local DOF order.
+    vec: Input in island-local DOF order.
+    nidof: Number of island DOFs per world.
+    map_idof2dof: Island-local DOF -> global DOF map.
+  """
+  if m.is_sparse:
+    nlevels = len(m.qLD_updates)
+    if wp.get_device().is_cuda:
+      dim_block = m.block_dim.solve_LD_sparse_fused
+    else:
+      dim_block = 1
+
+    wp.launch(
+      _solve_LD_sparse_island(m.nv, nlevels),
+      dim=(d.nworld, dim_block),
+      inputs=[
+        d.nidof,
+        d.map_dof2idof,
+        map_idof2dof,
+        d.qLD,
+        d.qLDiagInv,
+        m.qLD_all_updates,
+        m.qLD_level_offsets,
+        vec,
+      ],
+      outputs=[res],
+      block_dim=dim_block,
+    )
+  else:
+    for tile in m.M_tiles:
+      wp.launch_tiled(
+        _tile_cholesky_solve_island(tile),
+        dim=(d.nworld, tile.adr.size),
+        inputs=[d.nidof, d.map_dof2idof, d.qLD, vec, tile.adr],
+        outputs=[res],
+        block_dim=m.block_dim.cholesky_solve,
+      )
+
+
 @wp.kernel
 def _apply_ft(
   # Model:
@@ -604,11 +938,13 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
     nbody: int,
     neq: int,
     nmocap: int,
+    nhistory: int,
     # Data in:
     time_in: wp.array[float],
     qpos_in: wp.array2d[float],
     qvel_in: wp.array2d[float],
     act_in: wp.array2d[float],
+    history_in: wp.array2d[float],
     qacc_warmstart_in: wp.array2d[float],
     ctrl_in: wp.array2d[float],
     qfrc_applied_in: wp.array2d[float],
@@ -647,6 +983,10 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
           for j in range(na):
             state_out[worldid, adr + j] = act_in[worldid, j]
           adr += na
+        elif element == State.HISTORY:
+          for j in range(nhistory):
+            state_out[worldid, adr + j] = history_in[worldid, j]
+          adr += nhistory
         elif element == State.WARMSTART:
           for j in range(nv):
             state_out[worldid, adr + j] = qacc_warmstart_in[worldid, j]
@@ -700,10 +1040,12 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
       m.nbody,
       m.neq,
       m.nmocap,
+      m.nhistory,
       d.time,
       d.qpos,
       d.qvel,
       d.act,
+      d.history,
       d.qacc_warmstart,
       d.ctrl,
       d.qfrc_applied,
@@ -743,6 +1085,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
     nbody: int,
     neq: int,
     nmocap: int,
+    nhistory: int,
     # In:
     sig_in: int,
     active_in: wp.array[bool],
@@ -752,6 +1095,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
     qpos_out: wp.array2d[float],
     qvel_out: wp.array2d[float],
     act_out: wp.array2d[float],
+    history_out: wp.array2d[float],
     qacc_warmstart_out: wp.array2d[float],
     ctrl_out: wp.array2d[float],
     qfrc_applied_out: wp.array2d[float],
@@ -785,6 +1129,10 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
           for j in range(na):
             act_out[worldid, j] = state_in[worldid, adr + j]
           adr += na
+        elif element == State.HISTORY:
+          for j in range(nhistory):
+            history_out[worldid, j] = state_in[worldid, adr + j]
+          adr += nhistory
         elif element == State.WARMSTART:
           for j in range(nv):
             qacc_warmstart_out[worldid, j] = state_in[worldid, adr + j]
@@ -844,6 +1192,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
       m.nbody,
       m.neq,
       m.nmocap,
+      m.nhistory,
       int(sig),
       active or wp.ones(d.nworld, dtype=bool),
       state,
@@ -853,6 +1202,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
       d.qpos,
       d.qvel,
       d.act,
+      d.history,
       d.qacc_warmstart,
       d.ctrl,
       d.qfrc_applied,

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

@@ -17,18 +17,19 @@ import enum
 from typing import Callable
 
 import mujoco
-from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
 import numpy as np
 import warp as wp
 
+from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
+
 MJ_MINVAL = mujoco.mjMINVAL
 MJ_MAXVAL = mujoco.mjMAXVAL
 MJ_MINIMP = mujoco.mjMINIMP  # minimum constraint impedance
 MJ_MAXIMP = mujoco.mjMAXIMP  # maximum constraint impedance
 MJ_MAXCONPAIR = mujoco.mjMAXCONPAIR
 MJ_MINMU = mujoco.mjMINMU  # minimum friction
-# True if MuJoCo >= 3.9.0 (new margin/gap semantics: includemargin = margin)
+NEW_GAP_SEMANTICS = check_version("mujoco>=3.9.0.dev914519929")
-_NEW_GAP_SEMANTICS = check_version("mujoco>=3.9.0")
+TACTILE_DEPTH_SEMANTICS = check_version("mujoco>=3.9.0.dev921980899")
 # maximum size (by number of edges) of an horizon in EPA algorithm
 MJ_MAX_EPAHORIZON = 24
 # maximum average number of trianglarfaces EPA can insert at each iteration
@@ -47,6 +48,25 @@ class BlockDim:
   """Block dimension 'block_dim' settings for wp.launch_tiled.
 
   TODO(team): experimental and may be removed
+
+  Attributes:
+    segmented_sort: segmented sort block dimension (collision_driver)
+    euler_dense: Euler dense block dimension (forward)
+    actuator_velocity: actuator velocity block dimension (forward)
+    ray: ray block dimension (ray)
+    contact_sort: contact sort block dimension (sensor)
+    energy_vel_kinetic: energy velocity kinetic block dimension (sensor)
+    cholesky_factorize: Cholesky factorize block dimension (smooth)
+    cholesky_solve: Cholesky solve block dimension (smooth)
+    cholesky_factorize_solve: Cholesky factorize and solve block dimension (smooth)
+    solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth)
+    update_gradient_cholesky: update gradient Cholesky block dimension (solver)
+    update_gradient_cholesky_blocked: update gradient Cholesky blocked block dimension (solver)
+    update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension (solver)
+    update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension (solver)
+    linesearch_iterative: linesearch iterative block dimension (solver)
+    contact_jac_tiled: contact Jacobian tiled block dimension (solver)
+    qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
   """
 
   # collision_driver
@@ -131,13 +151,8 @@ class ProjectionType(enum.IntEnum):
     ORTHOGRAPHIC: orthographic projection
   """
 
-  # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
-  if hasattr(mujoco, "mjtProjection"):
   PERSPECTIVE = mujoco.mjtProjection.mjPROJ_PERSPECTIVE
   ORTHOGRAPHIC = mujoco.mjtProjection.mjPROJ_ORTHOGRAPHIC
-  else:
-    PERSPECTIVE = 0
-    ORTHOGRAPHIC = 1
 
 
 class Stage(enum.IntEnum):
@@ -188,6 +203,7 @@ class DisableBit(enum.IntFlag):
     EULERDAMP:    implicit damping for Euler integration
     NATIVECCD:    native convex collision detection (ignored in MJWarp)
     ISLAND:       constraint islands
+    MULTICCD:     disable multiple contacts with CCD
   """
 
   CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT
@@ -217,7 +233,6 @@ class EnableBit(enum.IntFlag):
   Attributes:
     ENERGY: energy computation
     INVDISCRETE: discrete-time inverse dynamics
-    MULTICCD: multiple contacts with CCD
   """
 
   ENERGY = mujoco.mjtEnableBit.mjENBL_ENERGY
@@ -338,12 +353,13 @@ class IntegratorType(enum.IntEnum):
     EULER: semi-implicit Euler
     RK4: 4th-order Runge Kutta
     IMPLICITFAST: implicit in velocity, no rne derivative
+    IMPLICIT: implicit in velocity, with rne derivative
   """
 
   EULER = mujoco.mjtIntegrator.mjINT_EULER
   RK4 = mujoco.mjtIntegrator.mjINT_RK4
   IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
-  # unsupported: IMPLICIT
+  IMPLICIT = mujoco.mjtIntegrator.mjINT_IMPLICIT
 
 
 class GeomType(enum.IntEnum):
@@ -617,6 +633,7 @@ class State(enum.IntEnum):
     QPOS: position
     QVEL: velocity
     ACT: actuator activation
+    HISTORY: delay/interval history buffers
     WARMSTART: acceleration used for warmstart
     CTRL: control
     QFRC_APPLIED: applied generalized force
@@ -625,7 +642,7 @@ class State(enum.IntEnum):
     MOCAP_POS: positions of mocap bodies
     MOCAP_QUAT: orientations of mocap bodies
     NSTATE: number of state elements
-    PHYSICS: QPOS | QVEL | ACT
+    PHYSICS: TIME | QPOS | QVEL | ACT | HISTORY
     FULLPHYSICS: TIME | PHYSICS | PLUGIN
     USER: CTRL | QFRC_APPLIED | XFRC_APPLIED | EQ_ACTIVE | MOCAP_POS | MOCAP_QUAT | USERDATA
     INTEGRATION: FULLPHYSICS | USER | WARMSTART
@@ -635,6 +652,7 @@ class State(enum.IntEnum):
   QPOS = mujoco.mjtState.mjSTATE_QPOS
   QVEL = mujoco.mjtState.mjSTATE_QVEL
   ACT = mujoco.mjtState.mjSTATE_ACT
+  HISTORY = mujoco.mjtState.mjSTATE_HISTORY
   WARMSTART = mujoco.mjtState.mjSTATE_WARMSTART
   CTRL = mujoco.mjtState.mjSTATE_CTRL
   QFRC_APPLIED = mujoco.mjtState.mjSTATE_QFRC_APPLIED
@@ -643,7 +661,7 @@ class State(enum.IntEnum):
   MOCAP_POS = mujoco.mjtState.mjSTATE_MOCAP_POS
   MOCAP_QUAT = mujoco.mjtState.mjSTATE_MOCAP_QUAT
   NSTATE = mujoco.mjtState.mjNSTATE
-  PHYSICS = mujoco.mjtState.mjSTATE_PHYSICS
+  PHYSICS = mujoco.mjtState.mjSTATE_PHYSICS  # includes HISTORY
   FULLPHYSICS = mujoco.mjtState.mjSTATE_FULLPHYSICS
   USER = mujoco.mjtState.mjSTATE_USER
   INTEGRATION = mujoco.mjtState.mjSTATE_INTEGRATION
@@ -861,6 +879,7 @@ class Model:
     ntree: number of kinematic trees
     nM: number of non-zeros in sparse inertia matrix
     nC: number of non-zeros in sparse body-dof matrix
+    nD: number of non-zeros in sparse derivative matrix
     ngeom: number of geoms
     nsite: number of sites
     ncam: number of cameras
@@ -898,6 +917,7 @@ class Model:
     nJmom: number of non-zeros in actuator_moment
     ngravcomp: number of bodies with nonzero gravcomp
     nsensordata: number of elements in sensor data vector
+    nhistory: number of history buffer entries
     opt: physics options
     stat: model statistics
     qpos0: qpos values at default pose                       (*, nq)
@@ -1129,6 +1149,9 @@ class Model:
     actuator_trnid: transmission id: joint, tendon, site     (nu, 2)
     actuator_actadr: first activation address; -1: stateless (nu,)
     actuator_actnum: number of activation variables          (nu,)
+    actuator_history: history buffer sizes                   (nu, 2)
+    actuator_historyadr: history buffer address              (nu,)
+    actuator_delay: delay in seconds                         (nu,)
     actuator_ctrllimited: is control limited                 (nu,)
     actuator_forcelimited: is force limited                  (nu,)
     actuator_actlimited: is activation limited               (nu,)
@@ -1153,12 +1176,22 @@ class Model:
     sensor_dim: number of scalar outputs                     (nsensor,)
     sensor_adr: address in sensor array                      (nsensor,)
     sensor_cutoff: cutoff for real and positive; 0: ignore   (nsensor,)
+    sensor_history: history buffer sizes                     (nsensor, 2)
+    sensor_historyadr: history buffer address                (nsensor,)
+    sensor_delay: delay in seconds                           (nsensor,)
+    sensor_interval: sensor interval and phase               (nsensor, 2)
     plugin: globally registered plugin slot number           (nplugin,)
     plugin_attr: config attributes of geom plugin            (nplugin, _NPLUGINATTR)
-    M_rownnz: number of non-zeros in each row of qM          (nv,)
+    M_rownnz: number of non-zeros in each row of M           (nv,)
-    M_rowadr: index of each row in qM                        (nv,)
+    M_rowadr: index of each row in M                         (nv,)
-    M_colind: column indices of non-zeros in qM              (nM,)
+    M_colind: column indices of non-zeros in M               (nC,)
     mapM2M: index mapping from M (legacy) to M (CSR)         (nC)
+    D_rownnz: non-zeros per row in D-structure               (nv,)
+    D_rowadr: row start addresses in D-structure             (nv,)
+    D_diag: diagonal element index within each row           (nv,)
+    D_colind: column indices in D-structure                  (nD,)
+    mapM2D: index mapping from M to D                        (nD,)
+    mapD2M: index mapping from D to M                        (nC,)
     flex_vertflexid: flex id for each flex vertex            (nflexvert,)
 
   warp only fields:
@@ -1187,8 +1220,8 @@ class Model:
     jnt_limited_slide_hinge_adr: limited/slide/hinge jntadr
     jnt_limited_ball_adr: limited/ball jntadr
     body_isdofancestor: precomputed mask of which DOFs affect each body
-    dof_tri_row: dof lower triangle row (used in solver)
+    dof_tri_row: dof upper triangle row (used in solver)
-    dof_tri_col: dof lower triangle col (used in solver)
+    dof_tri_col: dof upper triangle col (used in solver)
     nxn_geom_pair: collision pair geom ids [-2, ngeom-1]
     nxn_geom_pair_filtered: valid collision pair geom ids
                             [-1, ngeom - 1]
@@ -1243,15 +1276,21 @@ class Model:
     sensor_rangefinder_bodyid: bodyid for rangefinder        (nrangefinder,)
     taxel_vertadr: tactile sensor vertex address             (nsensortaxel,)
     taxel_sensorid: address for tactile sensors
-    qM_tiles: tiling configuration
+    M_tiles: tiling configuration
     qLD_updates: tuple of index triples for sparse factorization
     qLD_all_updates: tuple of all levels concatenated
     qLD_level_offsets: tuple of start offsets for each level
-    qM_fullm_i: sparse mass matrix addressing
+    M_fullm_i: sparse mass matrix addressing
-    qM_fullm_j: sparse mass matrix addressing
+    M_fullm_j: sparse mass matrix addressing
-    qM_mulm_rowadr: sparse matmul row pointers
+    M_elemid: (row, col) -> CSR madr addresses; -1 if not a chain ancestor
-    qM_mulm_col: sparse matmul column indices
+    M_fullm_upper_i: upper-triangle row indices for solver h seeding
-    qM_mulm_madr: sparse matmul matrix addresses
+    M_fullm_upper_j: upper-triangle column indices for solver h seeding
+    M_fullm_upper_elemid: source elemid into M_fullm_i/M_fullm_j
+    qD_fullm_i: D-structure row indices for RNE derivatives
+    qD_fullm_j: D-structure column indices for RNE derivatives
+    M_mulm_rowadr: sparse matmul row pointers
+    M_mulm_col: sparse matmul column indices
+    M_mulm_madr: sparse matmul matrix addresses
   """
 
   nq: int
@@ -1264,6 +1303,7 @@ class Model:
   ntree: int
   nM: int
   nC: int
+  nD: int
   ngeom: int
   nsite: int
   ncam: int
@@ -1301,6 +1341,7 @@ class Model:
   nJmom: int
   ngravcomp: int
   nsensordata: int
+  nhistory: int
   opt: Option
   stat: Statistic
   qpos0: array("*", "nq", float)
@@ -1532,6 +1573,9 @@ class Model:
   actuator_trnid: array("nu", wp.vec2i)
   actuator_actadr: array("nu", int)
   actuator_actnum: array("nu", int)
+  actuator_history: array("nu", wp.vec2i)
+  actuator_historyadr: array("nu", int)
+  actuator_delay: array("nu", float)
   actuator_ctrllimited: array("nu", bool)
   actuator_forcelimited: array("nu", bool)
   actuator_actlimited: array("nu", bool)
@@ -1556,12 +1600,22 @@ class Model:
   sensor_dim: array("nsensor", int)
   sensor_adr: array("nsensor", int)
   sensor_cutoff: array("nsensor", float)
+  sensor_history: array("nsensor", wp.vec2i)
+  sensor_historyadr: array("nsensor", int)
+  sensor_delay: array("nsensor", float)
+  sensor_interval: array("nsensor", wp.vec2)
   plugin: array("nplugin", int)
   plugin_attr: array("nplugin", vec_pluginattr)
   M_rownnz: array("nv", int)
   M_rowadr: array("nv", int)
   M_colind: array("nC", int)
   mapM2M: array("nC", int)
+  D_rownnz: array("nv", int)
+  D_rowadr: array("nv", int)
+  D_diag: array("nv", int)
+  D_colind: array("nD", int)
+  mapM2D: array("nD", int)
+  mapD2M: array("nC", int)
   flex_vertflexid: array("nflexvert", int)
   # warp only fields:
   callback: Callback
@@ -1635,23 +1689,32 @@ class Model:
   sensor_rangefinder_bodyid: array("nrangefinder", int)
   taxel_vertadr: array("nsensortaxel", int)
   taxel_sensorid: wp.array[int]
-  qM_tiles: tuple[TileSet, ...]
+  M_tiles: tuple[TileSet, ...]
   qLD_updates: tuple[wp.array[wp.vec3i], ...]
   qLD_all_updates: wp.array[wp.vec3i]
   qLD_level_offsets: wp.array[int]
-  qM_fullm_i: wp.array[int]
+  # TODO(team): Remove M_fullm_i/j and M_elemid by iterating the M CSR layout
-  qM_fullm_j: wp.array[int]
+  # directly in the solver/derivative kernels
+  M_fullm_i: wp.array[int]
+  M_fullm_j: wp.array[int]
+  M_elemid: wp.array2d[int]  # (row, col) -> CSR madr address; -1 if col is not a chain ancestor of row
+  M_fullm_upper_i: wp.array[int]
+  M_fullm_upper_j: wp.array[int]
+  M_fullm_upper_elemid: wp.array[int]
+  qD_fullm_i: wp.array[int]  # D-structure (full square) row indices for RNE derivatives
+  qD_fullm_j: wp.array[int]  # D-structure (full square) column indices for RNE derivatives
   # Gather-based sparse mul_m indices (thread per DOF, no atomics)
-  qM_mulm_rowadr: wp.array[int]  # start address for each row [nv+1]
+  M_mulm_rowadr: wp.array[int]  # start address for each row [nv+1]
-  qM_mulm_col: wp.array[int]  # column index to gather from
+  M_mulm_col: wp.array[int]  # column index to gather from
-  qM_mulm_madr: wp.array[int]  # matrix address to read
+  M_mulm_madr: wp.array[int]  # matrix address to read
 
 
 class ContactType(enum.IntFlag):
   """Type of contact.
 
-  CONSTRAINT: contact for constraint solver.
+  Attributes:
-  SENSOR: contact for collision sensor (GEOMDIST, GEOMNORMAL, GEOMFROMTO).
+    CONSTRAINT: contact for constraint solver
+    SENSOR: contact for collision sensor (GEOMDIST, GEOMNORMAL, GEOMFROMTO)
   """
 
   CONSTRAINT = 1
@@ -1673,6 +1736,8 @@ class Contact:
     solimp: constraint solver impedance                              (naconmax, 5)
     dim: contact space dimensionality: 1, 3, 4 or 6                  (naconmax,)
     geom: geom ids; -1 for flex                                      (naconmax, 2)
+    flex: flex ids; -1 for geom                                      (naconmax, 2)
+    vert: vertex ids for flex/mesh contact                           (naconmax, 2)
     efc_address: address in efc; -1: not included                    (naconmax, nmaxpyramid)
     worldid: world id                                                (naconmax,)
     type: ContactType                                                (naconmax,)
@@ -1711,9 +1776,9 @@ class Constraint:
     J_rowadr: row start address in colind array       (nworld, 0) dense
                                                       (nworld, njmax) sparse
     J_colind: column indices in J                     (nworld, 0, 0) dense
-                                                      (nworld, 1, njmax * nv) sparse
+                                                      (nworld, 1, njmax_nnz) sparse
     J: constraint Jacobian                            (nworld, njmax_pad, nv_pad) dense
-                                                      (nworld, 1, njmax * nv) sparse
+                                                      (nworld, 1, njmax_nnz) sparse
     pos: constraint position (equality, contact)      (nworld, njmax)
     margin: inclusion margin (contact)                (nworld, njmax)
     D: constraint mass                                (nworld, njmax_pad)
@@ -1722,6 +1787,20 @@ class Constraint:
     frictionloss: frictionloss (friction)             (nworld, njmax)
     force: constraint force in constraint space       (nworld, njmax)
     state: constraint state                           (nworld, njmax_pad)
+    island: island ID per constraint                  (nworld, njmax)
+    itype: island constraint type                     (nworld, njmax)
+    iid: island constraint id                         (nworld, njmax)
+    iJ_rownnz: island J_rownnz                        (nworld, njmax)
+    iJ_rowadr: island J_rowadr                        (nworld, njmax)
+    iJ_colind: island J_colind                        (nworld, 0, 0) dense
+                                                      (nworld, 1, njmax_nnz) sparse
+    iJ: island J                                      (nworld, njmax, nv) dense
+                                                      (nworld, 1, njmax_nnz) sparse
+    iD: island constraint mass                        (nworld, njmax_pad)
+    iaref: island aref                                (nworld, njmax)
+    ifrictionloss: island frictionloss                (nworld, njmax)
+    iforce: island force                              (nworld, njmax)
+    istate: island state                              (nworld, njmax_pad)
   warp only fields:
     Ma: M*qacc                                        (nworld, nv)
     Jqvel: J*qvel                                     (nworld, njmax)
@@ -1729,8 +1808,8 @@ class Constraint:
 
   type: array("nworld", "njmax", int)
   id: array("nworld", "njmax", int)
-  J_rownnz: wp.array2d[int]
+  J_rownnz: array("nworld", "njmax", int)
-  J_rowadr: wp.array2d[int]
+  J_rowadr: array("nworld", "njmax", int)
   J_colind: wp.array3d[int]
   J: wp.array3d[float]
   pos: array("nworld", "njmax", float)
@@ -1741,9 +1820,22 @@ class Constraint:
   frictionloss: array("nworld", "njmax", float)
   force: array("nworld", "njmax", float)
   state: array("nworld", "njmax_pad", int)
+  island: array("nworld", "njmax", int)
   Ma: array("nworld", "nv", float)
   Jqvel: array("nworld", "njmax", float)
 
+  itype: array("nworld", "njmax", int)
+  iid: array("nworld", "njmax", int)
+  iJ_rownnz: array("nworld", "njmax", int)
+  iJ_rowadr: array("nworld", "njmax", int)
+  iJ_colind: wp.array3d[int]
+  iJ: wp.array3d[float]
+  iD: array("nworld", "njmax_pad", float)
+  iaref: array("nworld", "njmax", float)
+  ifrictionloss: array("nworld", "njmax", float)
+  iforce: array("nworld", "njmax", float)
+  istate: array("nworld", "njmax_pad", int)
+
 
 @dataclasses.dataclass
 class Data:
@@ -1756,11 +1848,13 @@ class Data:
     nl: number of limit constraints                             (nworld,)
     nefc: number of constraints                                 (nworld,)
     nisland: number of constraint islands                       (nworld,)
+    nidof: total DOFs in islands                                (nworld,)
     time: simulation time                                       (nworld,)
     energy: potential, kinetic energy                           (nworld, 2)
     qpos: position                                              (nworld, nq)
     qvel: velocity                                              (nworld, nv)
     act: actuator activation                                    (nworld, na)
+    history: history buffer for delays                          (nworld, nhistory)
     qacc_warmstart: acceleration used for warmstart             (nworld, nv)
     ctrl: control                                               (nworld, nu)
     qfrc_applied: applied generalized force                     (nworld, nv)
@@ -1791,7 +1885,7 @@ class Data:
     cinert: com-based body inertia and mass                     (nworld, nbody, 10)
     flexvert_xpos: cartesian flex vertex positions              (nworld, nflexvert, 3)
     flexedge_J: edge length Jacobian                            (nworld, nJfe)
-    flexedge_length: flex edge lengths                          (nworld, nflexedge, 1)
+    flexedge_length: flex edge lengths                          (nworld, nflexedge)
     ten_wrapadr: start address of tendon's path                 (nworld, ntendon)
     ten_wrapnum: number of wrap points in path                  (nworld, ntendon)
     ten_J: tendon Jacobian                                      (nworld, nJten)
@@ -1804,10 +1898,10 @@ class Data:
     moment_colind: column indices in sparse actuator_moment     (nworld, nJmom)
     actuator_moment: actuator moments                           (nworld, nJmom)
     crb: com-based composite inertia and mass                   (nworld, nbody, 10)
-    qM: total inertia                                           (nworld, nv, nv) if dense
+    M: total inertia                                            (nworld, nv_pad, nv_pad) if dense
-                                                                (nworld, 1, nM) if sparse
-    qLD: L'*D*L factorization of M                              (nworld, nv, nv) if dense
                                                                 (nworld, 1, nC) if sparse
+    qLD: upper Cholesky factorization                           (nworld, nv, nv) if dense
+         L'*D*L factorization of M                              (nworld, 1, nC) if sparse
     qLDiagInv: 1/diag(D)                                        (nworld, nv)
     flexedge_velocity: flex edge velocities                     (nworld, nflexedge)
     ten_velocity: tendon velocities                             (nworld, ntendon)
@@ -1822,6 +1916,7 @@ class Data:
     qfrc_passive: total passive force                           (nworld, nv)
     subtree_linvel: linear velocity of subtree com              (nworld, nbody, 3)
     subtree_angmom: angular momentum about subtree com          (nworld, nbody, 3)
+    qLU: sparse LU factorization of (M - dt*qDeriv)             (nworld, 1, nD)
     actuator_force: actuator force in actuation space           (nworld, nu)
     qfrc_actuator: actuator force                               (nworld, nv)
     qfrc_smooth: net unconstrained force                        (nworld, nv)
@@ -1836,6 +1931,23 @@ class Data:
     contact: contact data
     efc: constraint data
     tree_island: island ID per tree (-1 if unconstrained)       (nworld, ntree)
+    dof_island: island ID per DOF (-1 if unconstrained)         (nworld, nv)
+    island_dofadr: island start address in dof vector           (nworld, ntree)
+    island_nv: DOFs per island                                  (nworld, ntree)
+    island_nefc: constraints per island                         (nworld, ntree)
+    island_ne: equality constraints per island                  (nworld, ntree)
+    island_nf: friction constraints per island                  (nworld, ntree)
+    island_efcadr: island start address in efc vector           (nworld, ntree)
+    map_dof2idof: global DOF -> island-local DOF                (nworld, nv)
+    map_idof2dof: island-local DOF -> global DOF                (nworld, nv)
+    map_efc2iefc: global EFC -> island-local EFC                (nworld, njmax)
+    map_iefc2efc: island-local EFC -> global EFC                (nworld, njmax)
+    dof_islandid: island ID per island-DOF                      (nworld, nv)
+    efc_islandid: island ID per island-EFC                      (nworld, njmax)
+    iqacc: island-local qacc                                    (nworld, nv)
+    iqacc_smooth: island-local qacc_smooth                      (nworld, nv)
+    iqfrc_smooth: island-local qfrc_smooth                      (nworld, nv)
+    iqfrc_constraint: island-local qfrc_constraint              (nworld, nv)
 
   warp only fields:
     nworld: number of worlds
@@ -1854,11 +1966,13 @@ class Data:
   nl: array("nworld", int)
   nefc: array("nworld", int)
   nisland: array("nworld", int)
+  nidof: array("nworld", int)
   time: array("nworld", float)
   energy: array("nworld", wp.vec2)
   qpos: array("nworld", "nq", float)
   qvel: array("nworld", "nv", float)
   act: array("nworld", "na", float)
+  history: array("nworld", "nhistory", float)
   qacc_warmstart: array("nworld", "nv", float)
   ctrl: array("nworld", "nu", float)
   qfrc_applied: array("nworld", "nv", float)
@@ -1902,7 +2016,7 @@ class Data:
   moment_colind: array("nworld", "nJmom", int)
   actuator_moment: array("nworld", "nJmom", float)
   crb: array("nworld", "nbody", vec10)
-  qM: wp.array3d[float]
+  M: wp.array3d[float]
   qLD: wp.array3d[float]
   qLDiagInv: array("nworld", "nv", float)
   flexedge_velocity: array("nworld", "nflexedge", float)
@@ -1918,6 +2032,7 @@ class Data:
   qfrc_passive: array("nworld", "nv", float)
   subtree_linvel: array("nworld", "nbody", wp.vec3)
   subtree_angmom: array("nworld", "nbody", wp.vec3)
+  qLU: array("nworld", 1, "nD", float)
   actuator_force: array("nworld", "nu", float)
   qfrc_actuator: array("nworld", "nv", float)
   qfrc_smooth: array("nworld", "nv", float)
@@ -1930,6 +2045,23 @@ class Data:
   contact: Contact
   efc: Constraint
   tree_island: array("nworld", "ntree", int)
+  dof_island: array("nworld", "nv", int)
+  island_dofadr: array("nworld", "ntree", int)
+  island_nv: array("nworld", "ntree", int)
+  island_nefc: array("nworld", "ntree", int)
+  island_ne: array("nworld", "ntree", int)
+  island_nf: array("nworld", "ntree", int)
+  island_efcadr: array("nworld", "ntree", int)
+  map_dof2idof: array("nworld", "nv", int)
+  map_idof2dof: array("nworld", "nv", int)
+  map_efc2iefc: array("nworld", "njmax", int)
+  map_iefc2efc: array("nworld", "njmax", int)
+  dof_islandid: array("nworld", "nv", int)
+  efc_islandid: array("nworld", "njmax", int)
+  iqacc: wp.array2d[float]
+  iqacc_smooth: wp.array2d[float]
+  iqfrc_smooth: wp.array2d[float]
+  iqfrc_constraint: wp.array2d[float]
 
   # warp only fields:
   nworld: int
@@ -1942,6 +2074,77 @@ class Data:
   ncollision: array(1, int)
 
 
+@dataclasses.dataclass
+class InverseContext:
+  """Workspace arrays for inverse dynamics."""
+
+  Jaref: wp.array2d[float]
+  search_dot: wp.array[float]
+  gauss: wp.array[float]
+  cost: wp.array[float]
+  prev_cost: wp.array[float]
+  done: wp.array[bool]
+  changed_efc_ids: wp.array2d[int]
+  changed_efc_count: wp.array[int]
+
+
+@dataclasses.dataclass
+class IslandSolverContext:
+  """Workspace arrays for island constraint solver."""
+
+  # Re-ordered workspace arrays (sized per-dof / per-constraint)
+  Jaref: wp.array2d[float]
+  jv: wp.array2d[float]
+  search: wp.array2d[float]
+  mv: wp.array2d[float]
+  grad: wp.array2d[float]
+  Mgrad: wp.array2d[float]
+  prev_grad: wp.array2d[float]
+  prev_Mgrad: wp.array2d[float]
+  h: wp.array3d[float]
+
+  # Per-island solver scalars (nworld, ntree)
+  cost: wp.array2d[float]
+  prev_cost: wp.array2d[float]
+  gauss: wp.array2d[float]
+  search_dot: wp.array2d[float]
+  grad_dot: wp.array2d[float]
+  done: wp.array2d[bool]  # per-island convergence
+  solver_niter: wp.array2d[int]  # iterations per island
+  beta: wp.array2d[float]
+  alpha: wp.array2d[float]
+  Ma: wp.array2d[float]  # island-local Ma (nworld, nv)
+
+
+@dataclasses.dataclass
+class SolverContext:
+  """Workspace arrays for constraint solver."""
+
+  Jaref: wp.array2d[float]
+  search_dot: wp.array[float]
+  gauss: wp.array[float]
+  cost: wp.array[float]
+  prev_cost: wp.array[float]
+  done: wp.array[bool]
+  grad: wp.array2d[float]
+  grad_dot: wp.array[float]
+  Mgrad: wp.array2d[float]
+  search: wp.array2d[float]
+  mv: wp.array2d[float]
+  jv: wp.array2d[float]
+  quad: wp.array2d[wp.vec3]
+  quad_gauss: wp.array[wp.vec3]
+  alpha: wp.array[float]
+  prev_grad: wp.array2d[float]
+  prev_Mgrad: wp.array2d[float]
+  beta: wp.array[float]
+  h: wp.array3d[float]
+  hfactor: wp.array3d[float]
+  # Incremental Hessian update (Newton only)
+  changed_efc_ids: wp.array2d[int]
+  changed_efc_count: wp.array[int]
+
+
 @dataclasses.dataclass
 class RenderContext:
   """Context for rendering.
@@ -1952,6 +2155,8 @@ class RenderContext:
     cam_id_map: camera id map
     use_textures: whether to use textures
     use_shadows: whether to use shadows
+    use_ambient_lighting: whether to use ambient lighting
+    background_color: background color
     use_precomputed_rays: whether to use precomputed rays
     bvh_ngeom: number of geometries in the BVH
     enabled_geom_ids: enabled geometry ids
@@ -1971,7 +2176,10 @@ class RenderContext:
     flex_bvh_id: per-flex BVH ids
     flex_group_root: per-flex group roots (nworld x n_flex_bvh)
     flex_render_smooth: whether to render flex meshes smoothly
-    flex_dim: flex dimension per flex (1D/2D/3D)
+    bvh_nflexgeom: number of flex geometries in the BVH
+    flex_dim_np: flex dimension per flex (1D/2D/3D)
+    flex_geom_flexid: map from flex geom ID to flex ID
+    flex_geom_edgeid: map from flex geom ID to flex edge ID
     bvh: scene BVH
     bvh_id: scene BVH id
     lower: lower bounds
@@ -1993,6 +2201,11 @@ class RenderContext:
     render_skybox: whether to shade missed rays with the MuJoCo skybox texture
     skybox_tex_id: index into textures of the skybox (MuJoCo tex_type == SKYBOX), -1 if none
     skybox_face_width: pixel width of one skybox cube face (0 if no skybox)
+    enable_backface_culling: drop primitive ray hits whose normal faces away
+      from the ray (i.e. the ray origin is inside the geom). Matches MuJoCo's
+      mesh-ray rule. When False, the renderer reports inner-surface hits, which
+      is faster but causes a camera placed inside a geom to render that geom's
+      back wall.
   """
 
   nrender: int
@@ -2000,6 +2213,7 @@ class RenderContext:
   cam_id_map: array("ncam", int)
   use_textures: bool
   use_shadows: bool
+  use_ambient_lighting: bool
   background_color: wp.uint32
   use_precomputed_rays: bool
   render_skybox: bool
@@ -2045,3 +2259,4 @@ class RenderContext:
   render_seg: array("ncam", bool)
   znear: float
   total_rays: int
+  enable_backface_culling: bool

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

@@ -23,7 +23,7 @@ def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
   """Parse a version string into comparable components.
 
   Dot-separated components form the version. Hyphen-separated suffixes (e.g.,
-  "-foo3") are treated as local build identifiers and stripped before
+  "-newton") are treated as local build identifiers and stripped before
   parsing. Each component is wrapped in a tuple: (0, int) for numeric parts,
   (-1, str) for non-numeric. A (0, 0) sentinel is appended so that stable
   releases sort above pre-release suffixes during Python tuple comparison
@@ -32,13 +32,13 @@ def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
 
   Args:
     version_str: Version string like "3.5.0", "3.5.0.dev869102767", or
-      "3.9.0-foo3".
+      "3.9.0-newton".
 
   Returns:
     Tuple of (type_order, value) pairs for comparison, where type_order is 0
     for integers and -1 for strings, followed by a (0, 0) sentinel.
   """
-  # Strip local build identifier (e.g., "3.9.0-foo3" -> "3.9.0")
+  # Strip local build identifier (e.g., "3.9.0-newton" -> "3.9.0")
   version_str = version_str.split("-", 1)[0]
   parts = version_str.split(".")
   return tuple([(0, int(p)) if p.isdigit() else (-1, p) for p in parts] + [(0, 0)])

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name="mujoco-warp"
-version = "3.8.0.1"
+version = "3.9.0.1"
 # TODO(team): create a distribution list
 authors = [
     {name = "Newton Developers", email = "mujoco@deepmind.com"},
@@ -31,7 +31,7 @@ dependencies = [
     "etils[epath]",
     "mujoco>=3.8.0",
     "numpy",
-    "warp-lang>=1.12",
+    "warp-lang>=1.13",
 ]
 
 [[tool.uv.index]]

mjx/mujoco/mjx/warp/__init__.py

@@ -23,6 +23,7 @@ if not typing.TYPE_CHECKING:
   warp: Any = None
   mujoco_warp: Any = None
   mjwp_types: Any = None
+  mjwp_io: Any = None
   WARP_INSTALLED: bool = False
 
   # pylint: disable=g-import-not-at-top
@@ -35,6 +36,7 @@ if not typing.TYPE_CHECKING:
   try:
     from mujoco.mjx.third_party import mujoco_warp
     from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
+    from mujoco.mjx.third_party.mujoco_warp._src import io as mjwp_io
   except (ImportError, RuntimeError) as e:
     print('Failed to import mujoco_warp:', e)
     pass
@@ -67,7 +69,11 @@ else:
     def BlockDim(self, *args, **kwargs):  # pylint: disable=invalid-name
       pass
 
+  class _MjwpIoStub:
+    ENABLE_ISLANDS = True
+
   WARP_INSTALLED: bool = True
   warp: Any = _WpStub()
   mujoco_warp: Any = _MjwpStub()
   mjwp_types: Any = _MjwpTypesStub()
+  mjwp_io: Any = _MjwpIoStub()

mjx/mujoco/mjx/warp/forward_test.py

@@ -176,8 +176,8 @@ class ForwardTest(parameterized.TestCase):
 
       qm = np.zeros((m.nv, m.nv))
       mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
-      # mjwarp adds padding to qM
+      # mjwarp adds padding to M
-      tu.assert_eq(qm, dx._impl.qM[: m.nv, : m.nv], 'qM')
+      tu.assert_eq(qm, dx._impl.M[: m.nv, : m.nv], 'M')
       # qLD is fused in a cholesky factorize and solve, and not written to.
 
       tu.assert_contact_eq(d, dx, worldid=i)

mjx/mujoco/mjx/warp/smooth_test.py

@@ -112,7 +112,7 @@ class SmoothTest(parameterized.TestCase):
 
     d = mujoco.MjData(m)
     mx = mjx.put_model(m, impl='warp')
-    mx = mx.replace(_impl=mx._impl.replace(qM_tiles=()))
+    mx = mx.replace(_impl=mx._impl.replace(M_tiles=()))
 
     rng = jax.random.PRNGKey(0)
     dx = mjx.make_data(m, impl='warp')
@@ -188,7 +188,7 @@ class SmoothTest(parameterized.TestCase):
     batch_size = 7
     d = mujoco.MjData(m)
     mx = mjx.put_model(m, impl='warp')
-    mx = mx.replace(_impl=mx._impl.replace(qM_tiles=()))
+    mx = mx.replace(_impl=mx._impl.replace(M_tiles=()))
 
     worldids = jp.arange(batch_size)
     dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids)

mjx/mujoco/mjx/warp/types.py

@@ -78,6 +78,29 @@ class BlockDim:
   """Block dimension 'block_dim' settings for wp.launch_tiled.
 
   TODO(team): experimental and may be removed
+
+  Attributes:
+    segmented_sort: segmented sort block dimension (collision_driver)
+    euler_dense: Euler dense block dimension (forward)
+    actuator_velocity: actuator velocity block dimension (forward)
+    ray: ray block dimension (ray)
+    contact_sort: contact sort block dimension (sensor)
+    energy_vel_kinetic: energy velocity kinetic block dimension (sensor)
+    cholesky_factorize: Cholesky factorize block dimension (smooth)
+    cholesky_solve: Cholesky solve block dimension (smooth)
+    cholesky_factorize_solve: Cholesky factorize and solve block dimension
+      (smooth)
+    solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth)
+    update_gradient_cholesky: update gradient Cholesky block dimension (solver)
+    update_gradient_cholesky_blocked: update gradient Cholesky blocked block
+      dimension (solver)
+    update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension
+      (solver)
+    update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension
+      (solver)
+    linesearch_iterative: linesearch iterative block dimension (solver)
+    contact_jac_tiled: contact Jacobian tiled block dimension (solver)
+    qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
   """
 
   actuator_velocity: int
@@ -134,10 +157,26 @@ class OptionWarp(PyTreeNode):
 
 class ModelWarp(PyTreeNode):
   """Derived fields from Model."""
-
+  D_colind: np.ndarray
+  D_diag: np.ndarray
+  D_rowadr: np.ndarray
+  D_rownnz: np.ndarray
   M_colind: np.ndarray
+  M_elemid: np.ndarray
+  M_fullm_i: np.ndarray
+  M_fullm_j: np.ndarray
+  M_fullm_upper_elemid: np.ndarray
+  M_fullm_upper_i: np.ndarray
+  M_fullm_upper_j: np.ndarray
+  M_mulm_col: np.ndarray
+  M_mulm_madr: np.ndarray
+  M_mulm_rowadr: np.ndarray
   M_rowadr: np.ndarray
   M_rownnz: np.ndarray
+  M_tiles: Tuple[TileSet, ...]
+  actuator_delay: np.ndarray
+  actuator_history: np.ndarray
+  actuator_historyadr: np.ndarray
   actuator_trntype_body_adr: np.ndarray
   block_dim: BlockDim
   body_branch_start: np.ndarray
@@ -204,6 +243,8 @@ class ModelWarp(PyTreeNode):
   light_active: jax.Array
   light_bodyid: np.ndarray
   light_targetbodyid: np.ndarray
+  mapD2M: np.ndarray
+  mapM2D: np.ndarray
   mapM2M: np.ndarray
   mat_texrepeat: jax.Array
   max_ten_J_rownnz: int
@@ -249,22 +290,22 @@ class ModelWarp(PyTreeNode):
   oct_coeff: np.ndarray
   plugin: np.ndarray
   plugin_attr: np.ndarray
+  qD_fullm_i: np.ndarray
+  qD_fullm_j: np.ndarray
   qLD_all_updates: np.ndarray
   qLD_level_offsets: np.ndarray
   qLD_updates: Tuple[np.ndarray, ...]
-  qM_fullm_i: np.ndarray
-  qM_fullm_j: np.ndarray
-  qM_mulm_col: np.ndarray
-  qM_mulm_madr: np.ndarray
-  qM_mulm_rowadr: np.ndarray
-  qM_tiles: Tuple[TileSet, ...]
   rangefinder_sensor_adr: np.ndarray
   sensor_acc_adr: np.ndarray
   sensor_adr_to_contact_adr: np.ndarray
   sensor_collision_start_adr: np.ndarray
   sensor_contact_adr: np.ndarray
+  sensor_delay: np.ndarray
   sensor_e_kinetic: bool
   sensor_e_potential: bool
+  sensor_history: np.ndarray
+  sensor_historyadr: np.ndarray
+  sensor_interval: np.ndarray
   sensor_limitfrc_adr: np.ndarray
   sensor_limitpos_adr: np.ndarray
   sensor_limitvel_adr: np.ndarray
@@ -299,7 +340,7 @@ class ModelWarp(PyTreeNode):
 
 class DataWarp(PyTreeNode):
   """Derived fields from Data."""
-
+  M: jax.Array
   actuator_moment: jax.Array
   actuator_velocity: jax.Array
   cacc: jax.Array
@@ -323,6 +364,8 @@ class DataWarp(PyTreeNode):
   contact__vert: jax.Array
   contact__worldid: jax.Array
   crb: jax.Array
+  dof_island: jax.Array
+  dof_islandid: jax.Array
   efc__D: jax.Array
   efc__J: jax.Array
   efc__J_colind: jax.Array
@@ -333,19 +376,46 @@ class DataWarp(PyTreeNode):
   efc__aref: jax.Array
   efc__force: jax.Array
   efc__frictionloss: jax.Array
+  efc__iD: jax.Array
+  efc__iJ: jax.Array
+  efc__iJ_colind: jax.Array
+  efc__iJ_rowadr: jax.Array
+  efc__iJ_rownnz: jax.Array
+  efc__iaref: jax.Array
   efc__id: jax.Array
+  efc__iforce: jax.Array
+  efc__ifrictionloss: jax.Array
+  efc__iid: jax.Array
+  efc__island: jax.Array
+  efc__istate: jax.Array
+  efc__itype: jax.Array
   efc__margin: jax.Array
   efc__pos: jax.Array
   efc__state: jax.Array
   efc__type: jax.Array
   efc__vel: jax.Array
+  efc_islandid: jax.Array
   energy: jax.Array
   flexedge_J: jax.Array
   flexedge_length: jax.Array
   flexedge_velocity: jax.Array
   flexvert_xpos: jax.Array
+  iqacc: jax.Array
+  iqacc_smooth: jax.Array
+  iqfrc_constraint: jax.Array
+  iqfrc_smooth: jax.Array
+  island_dofadr: jax.Array
+  island_efcadr: jax.Array
+  island_ne: jax.Array
+  island_nefc: jax.Array
+  island_nf: jax.Array
+  island_nv: jax.Array
   light_xdir: jax.Array
   light_xpos: jax.Array
+  map_dof2idof: jax.Array
+  map_efc2iefc: jax.Array
+  map_idof2dof: jax.Array
+  map_iefc2efc: jax.Array
   moment_colind: jax.Array
   moment_rowadr: jax.Array
   moment_rownnz: jax.Array
@@ -356,6 +426,7 @@ class DataWarp(PyTreeNode):
   ne: jax.Array
   nefc: jax.Array
   nf: jax.Array
+  nidof: jax.Array
   nisland: jax.Array
   njmax: int
   njmax_nnz: int
@@ -364,7 +435,7 @@ class DataWarp(PyTreeNode):
   nworld: int
   qLD: jax.Array
   qLDiagInv: jax.Array
-  qM: jax.Array
+  qLU: jax.Array
   qfrc_damper: jax.Array
   qfrc_spring: jax.Array
   solver_niter: jax.Array
@@ -401,6 +472,7 @@ DATA_NON_VMAP = {
     'nacon',
     'naconmax',
     'ncollision',
+    'nidof',
     'njmax',
     'njmax_nnz',
     'njmax_pad',
@@ -434,6 +506,7 @@ batching.register_vmappable(DataWarp, int, int, _to_elt, _from_elt, None)
 
 _NDIM = {
     'Data': {
+        'M': 3,
         'act': 2,
         'act_dot': 2,
         'actuator_force': 2,
@@ -467,6 +540,8 @@ _NDIM = {
         'crb': 3,
         'ctrl': 2,
         'cvel': 3,
+        'dof_island': 2,
+        'dof_islandid': 2,
         'efc__D': 2,
         'efc__J': 3,
         'efc__J_colind': 3,
@@ -477,12 +552,25 @@ _NDIM = {
         'efc__aref': 2,
         'efc__force': 2,
         'efc__frictionloss': 2,
+        'efc__iD': 2,
+        'efc__iJ': 3,
+        'efc__iJ_colind': 3,
+        'efc__iJ_rowadr': 2,
+        'efc__iJ_rownnz': 2,
+        'efc__iaref': 2,
         'efc__id': 2,
+        'efc__iforce': 2,
+        'efc__ifrictionloss': 2,
+        'efc__iid': 2,
+        'efc__island': 2,
+        'efc__istate': 2,
+        'efc__itype': 2,
         'efc__margin': 2,
         'efc__pos': 2,
         'efc__state': 2,
         'efc__type': 2,
         'efc__vel': 2,
+        'efc_islandid': 2,
         'energy': 2,
         'eq_active': 2,
         'flexedge_J': 2,
@@ -491,8 +579,23 @@ _NDIM = {
         'flexvert_xpos': 3,
         'geom_xmat': 4,
         'geom_xpos': 3,
+        'history': 2,
+        'iqacc': 2,
+        'iqacc_smooth': 2,
+        'iqfrc_constraint': 2,
+        'iqfrc_smooth': 2,
+        'island_dofadr': 2,
+        'island_efcadr': 2,
+        'island_ne': 2,
+        'island_nefc': 2,
+        'island_nf': 2,
+        'island_nv': 2,
         'light_xdir': 3,
         'light_xpos': 3,
+        'map_dof2idof': 2,
+        'map_efc2iefc': 2,
+        'map_idof2dof': 2,
+        'map_iefc2efc': 2,
         'mocap_pos': 3,
         'mocap_quat': 3,
         'moment_colind': 2,
@@ -505,6 +608,7 @@ _NDIM = {
         'ne': 1,
         'nefc': 1,
         'nf': 1,
+        'nidof': 1,
         'nisland': 1,
         'njmax': 0,
         'njmax_nnz': 0,
@@ -513,7 +617,7 @@ _NDIM = {
         'nworld': 0,
         'qLD': 3,
         'qLDiagInv': 2,
-        'qM': 3,
+        'qLU': 3,
         'qacc': 2,
         'qacc_smooth': 2,
         'qacc_warmstart': 2,
@@ -556,9 +660,23 @@ _NDIM = {
         'xquat': 3,
     },
     'Model': {
+        'D_colind': 1,
+        'D_diag': 1,
+        'D_rowadr': 1,
+        'D_rownnz': 1,
         'M_colind': 1,
+        'M_elemid': 2,
+        'M_fullm_i': 1,
+        'M_fullm_j': 1,
+        'M_fullm_upper_elemid': 1,
+        'M_fullm_upper_i': 1,
+        'M_fullm_upper_j': 1,
+        'M_mulm_col': 1,
+        'M_mulm_madr': 1,
+        'M_mulm_rowadr': 1,
         'M_rowadr': 1,
         'M_rownnz': 1,
+        'M_tiles': -1,
         'actuator_acc0': 2,
         'actuator_actadr': 1,
         'actuator_actearly': 1,
@@ -570,6 +688,7 @@ _NDIM = {
         'actuator_cranklength': 2,
         'actuator_ctrllimited': 1,
         'actuator_ctrlrange': 3,
+        'actuator_delay': 1,
         'actuator_dynprm': 3,
         'actuator_dyntype': 1,
         'actuator_forcelimited': 1,
@@ -577,6 +696,8 @@ _NDIM = {
         'actuator_gainprm': 3,
         'actuator_gaintype': 1,
         'actuator_gear': 3,
+        'actuator_history': 2,
+        'actuator_historyadr': 1,
         'actuator_lengthrange': 3,
         'actuator_trnid': 2,
         'actuator_trntype': 1,
@@ -769,6 +890,8 @@ _NDIM = {
         'light_poscom0': 3,
         'light_targetbodyid': 1,
         'light_type': 2,
+        'mapD2M': 1,
+        'mapM2D': 1,
         'mapM2M': 1,
         'mat_rgba': 3,
         'mat_texid': 3,
@@ -797,6 +920,7 @@ _NDIM = {
         'mesh_vertnum': 1,
         'mocap_bodyid': 1,
         'nC': 0,
+        'nD': 0,
         'nJfe': 0,
         'nJmom': 0,
         'nJten': 0,
@@ -821,6 +945,7 @@ _NDIM = {
         'ngravcomp': 0,
         'nhfield': 0,
         'nhfielddata': 0,
+        'nhistory': 0,
         'njnt': 0,
         'nlight': 0,
         'nmat': 0,
@@ -897,15 +1022,11 @@ _NDIM = {
         'pair_solreffriction': 3,
         'plugin': 1,
         'plugin_attr': 2,
+        'qD_fullm_i': 1,
+        'qD_fullm_j': 1,
         'qLD_all_updates': 2,
         'qLD_level_offsets': 1,
         'qLD_updates': -1,
-        'qM_fullm_i': 1,
-        'qM_fullm_j': 1,
-        'qM_mulm_col': 1,
-        'qM_mulm_madr': 1,
-        'qM_mulm_rowadr': 1,
-        'qM_tiles': -1,
         'qpos0': 2,
         'qpos_spring': 2,
         'rangefinder_sensor_adr': 1,
@@ -916,9 +1037,13 @@ _NDIM = {
         'sensor_contact_adr': 1,
         'sensor_cutoff': 1,
         'sensor_datatype': 1,
+        'sensor_delay': 1,
         'sensor_dim': 1,
         'sensor_e_kinetic': 0,
         'sensor_e_potential': 0,
+        'sensor_history': 2,
+        'sensor_historyadr': 1,
+        'sensor_interval': 2,
         'sensor_intprm': 2,
         'sensor_limitfrc_adr': 1,
         'sensor_limitpos_adr': 1,
@@ -1016,6 +1141,7 @@ _NDIM = {
 }
 _BATCH_DIM = {
     'Data': {
+        'M': True,
         'act': True,
         'act_dot': True,
         'actuator_force': True,
@@ -1049,6 +1175,8 @@ _BATCH_DIM = {
         'crb': True,
         'ctrl': True,
         'cvel': True,
+        'dof_island': True,
+        'dof_islandid': True,
         'efc__D': True,
         'efc__J': True,
         'efc__J_colind': True,
@@ -1059,12 +1187,25 @@ _BATCH_DIM = {
         'efc__aref': True,
         'efc__force': True,
         'efc__frictionloss': True,
+        'efc__iD': True,
+        'efc__iJ': True,
+        'efc__iJ_colind': True,
+        'efc__iJ_rowadr': True,
+        'efc__iJ_rownnz': True,
+        'efc__iaref': True,
         'efc__id': True,
+        'efc__iforce': True,
+        'efc__ifrictionloss': True,
+        'efc__iid': True,
+        'efc__island': True,
+        'efc__istate': True,
+        'efc__itype': True,
         'efc__margin': True,
         'efc__pos': True,
         'efc__state': True,
         'efc__type': True,
         'efc__vel': True,
+        'efc_islandid': True,
         'energy': True,
         'eq_active': True,
         'flexedge_J': True,
@@ -1073,8 +1214,23 @@ _BATCH_DIM = {
         'flexvert_xpos': True,
         'geom_xmat': True,
         'geom_xpos': True,
+        'history': True,
+        'iqacc': True,
+        'iqacc_smooth': True,
+        'iqfrc_constraint': True,
+        'iqfrc_smooth': True,
+        'island_dofadr': True,
+        'island_efcadr': True,
+        'island_ne': True,
+        'island_nefc': True,
+        'island_nf': True,
+        'island_nv': True,
         'light_xdir': True,
         'light_xpos': True,
+        'map_dof2idof': True,
+        'map_efc2iefc': True,
+        'map_idof2dof': True,
+        'map_iefc2efc': True,
         'mocap_pos': True,
         'mocap_quat': True,
         'moment_colind': True,
@@ -1087,6 +1243,7 @@ _BATCH_DIM = {
         'ne': True,
         'nefc': True,
         'nf': True,
+        'nidof': False,
         'nisland': True,
         'njmax': False,
         'njmax_nnz': False,
@@ -1095,7 +1252,7 @@ _BATCH_DIM = {
         'nworld': False,
         'qLD': True,
         'qLDiagInv': True,
-        'qM': True,
+        'qLU': True,
         'qacc': True,
         'qacc_smooth': True,
         'qacc_warmstart': True,
@@ -1138,9 +1295,23 @@ _BATCH_DIM = {
         'xquat': True,
     },
     'Model': {
+        'D_colind': False,
+        'D_diag': False,
+        'D_rowadr': False,
+        'D_rownnz': False,
         'M_colind': False,
+        'M_elemid': False,
+        'M_fullm_i': False,
+        'M_fullm_j': False,
+        'M_fullm_upper_elemid': False,
+        'M_fullm_upper_i': False,
+        'M_fullm_upper_j': False,
+        'M_mulm_col': False,
+        'M_mulm_madr': False,
+        'M_mulm_rowadr': False,
         'M_rowadr': False,
         'M_rownnz': False,
+        'M_tiles': False,
         'actuator_acc0': True,
         'actuator_actadr': False,
         'actuator_actearly': False,
@@ -1152,6 +1323,7 @@ _BATCH_DIM = {
         'actuator_cranklength': True,
         'actuator_ctrllimited': False,
         'actuator_ctrlrange': True,
+        'actuator_delay': False,
         'actuator_dynprm': True,
         'actuator_dyntype': False,
         'actuator_forcelimited': False,
@@ -1159,6 +1331,8 @@ _BATCH_DIM = {
         'actuator_gainprm': True,
         'actuator_gaintype': False,
         'actuator_gear': True,
+        'actuator_history': False,
+        'actuator_historyadr': False,
         'actuator_lengthrange': True,
         'actuator_trnid': False,
         'actuator_trntype': False,
@@ -1351,6 +1525,8 @@ _BATCH_DIM = {
         'light_poscom0': True,
         'light_targetbodyid': False,
         'light_type': True,
+        'mapD2M': False,
+        'mapM2D': False,
         'mapM2M': False,
         'mat_rgba': True,
         'mat_texid': True,
@@ -1379,6 +1555,7 @@ _BATCH_DIM = {
         'mesh_vertnum': False,
         'mocap_bodyid': False,
         'nC': False,
+        'nD': False,
         'nJfe': False,
         'nJmom': False,
         'nJten': False,
@@ -1403,6 +1580,7 @@ _BATCH_DIM = {
         'ngravcomp': False,
         'nhfield': False,
         'nhfielddata': False,
+        'nhistory': False,
         'njnt': False,
         'nlight': False,
         'nmat': False,
@@ -1479,15 +1657,11 @@ _BATCH_DIM = {
         'pair_solreffriction': True,
         'plugin': False,
         'plugin_attr': False,
+        'qD_fullm_i': False,
+        'qD_fullm_j': False,
         'qLD_all_updates': False,
         'qLD_level_offsets': False,
         'qLD_updates': False,
-        'qM_fullm_i': False,
-        'qM_fullm_j': False,
-        'qM_mulm_col': False,
-        'qM_mulm_madr': False,
-        'qM_mulm_rowadr': False,
-        'qM_tiles': False,
         'qpos0': True,
         'qpos_spring': True,
         'rangefinder_sensor_adr': False,
@@ -1498,9 +1672,13 @@ _BATCH_DIM = {
         'sensor_contact_adr': False,
         'sensor_cutoff': False,
         'sensor_datatype': False,
+        'sensor_delay': False,
         'sensor_dim': False,
         'sensor_e_kinetic': False,
         'sensor_e_potential': False,
+        'sensor_history': False,
+        'sensor_historyadr': False,
+        'sensor_interval': False,
         'sensor_intprm': False,
         'sensor_limitfrc_adr': False,
         'sensor_limitpos_adr': False,