Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 885715192
Change-Id: I13fcb8312c1515bb1bd2d6cdfd0d256e37cb74de

mjx/mujoco/mjx/_src/io.py

@@ -205,6 +205,9 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any:
     )
   if isinstance(wp_field, mjwp_types.BlockDim):
     return mjxw.types.BlockDim(**wp_field.__dict__)
+  if isinstance(wp_field, mjwp_types.Callback):
+    return wp_field
+
   if isinstance(wp_field, tuple) and is_static(wp_field[0]):
     return wp_field
 

mjx/mujoco/mjx/_src/types.py

@@ -848,9 +848,9 @@ class Model(PyTreeNode):
   actuator_forcerange: jax.Array
   actuator_actrange: jax.Array
   actuator_gear: jax.Array
-  actuator_cranklength: np.ndarray
+  actuator_cranklength: jax.Array
   actuator_acc0: jax.Array
-  actuator_lengthrange: np.ndarray
+  actuator_lengthrange: jax.Array
   sensor_type: np.ndarray
   sensor_datatype: np.ndarray
   sensor_needstage: np.ndarray

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

@@ -56,6 +56,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.io import reset_data as reset_data
 from mujoco.mjx.third_party.mujoco_warp._src.io import set_const as set_const
 from mujoco.mjx.third_party.mujoco_warp._src.io import set_const_0 as set_const_0
 from mujoco.mjx.third_party.mujoco_warp._src.io import set_const_fixed as set_const_fixed
+from mujoco.mjx.third_party.mujoco_warp._src.io import set_length_range as set_length_range
+from mujoco.mjx.third_party.mujoco_warp._src.island import island as island
 from mujoco.mjx.third_party.mujoco_warp._src.passive import passive as passive
 from mujoco.mjx.third_party.mujoco_warp._src.ray import ray as ray
 from mujoco.mjx.third_party.mujoco_warp._src.ray import rays as rays

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

@@ -57,12 +57,12 @@ def ctrl_noise(
   worldid, actid = wp.tid()
 
   # convert rate and scale to discrete time (Ornstein-Uhlenbeck)
-  rate = wp.exp(-opt_timestep[0] / ctrlnoiserate)
+  rate = wp.exp(-opt_timestep[worldid % opt_timestep.shape[0]] / ctrlnoiserate)
   scale = ctrlnoisestd * wp.sqrt(1.0 - rate * rate)
 
   midpoint = 0.0
   halfrange = 1.0
-  ctrlrange = actuator_ctrlrange[0, actid]
+  ctrlrange = actuator_ctrlrange[worldid % actuator_ctrlrange.shape[0], actid]
   is_limited = actuator_ctrllimited[actid]
   if is_limited:
     midpoint = 0.5 * (ctrlrange[1] + ctrlrange[0])

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

@@ -215,7 +215,8 @@ def _compute_bvh_bounds(
     lower_bound, upper_bound = _compute_box_bounds(pos, rot, size)
   elif type == GeomType.HFIELD:
     size = hfield_bounds_size[geom_dataid[geom_id]]
-    lower_bound, upper_bound = _compute_box_bounds(pos, rot, size)
+    hfield_center = pos + rot[:, 2] * size[2]
+    lower_bound, upper_bound = _compute_box_bounds(hfield_center, rot, size)
 
   lower_out[world_id * bvh_ngeom + geom_local_id] = lower_bound
   upper_out[world_id * bvh_ngeom + geom_local_id] = upper_bound

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

@@ -15,32 +15,34 @@
 
 from typing import Tuple
 
-import warp as wp
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext
-
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import ccd
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import multicontact
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import support
-from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import contact_params
+from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import geom_collision_pair
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import write_contact
-from mujoco.mjx.third_party.mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
 from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
 from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index
+from mujoco.mjx.third_party.mujoco_warp._src.types import Data
+from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit
+from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
+from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
+from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAFACES
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAHORIZON
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
-from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext
-from mujoco.mjx.third_party.mujoco_warp._src.types import Data
-from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit
-from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
-from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
-from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
+import warp as wp
 
 # TODO(team): improve compile time to enable backward pass
 wp.set_module_options({"enable_backward": False})
@@ -92,10 +94,7 @@ def _hfield_filter(
   r2 = geom_rbound[rbound_id, g2]
 
   # TODO(team): margin?
-  if BLEEDING_EDGE_MUJOCO:
   margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
-  else:
-    margin = wp.max(geom_margin[margin_id, g1], geom_margin[margin_id, g2])
 
   # box-sphere test: horizontal plane
   for i in range(2):
@@ -789,6 +788,13 @@ def ccd_kernel_builder(
     if dist >= 0.0 and pairid[1] == -1:
       return 0
 
+    # CCD operates on margin-inflated shapes (support() inflates each geom by
+    # 0.5 * margin).  The returned dist is therefore relative to the inflated
+    # geometry.  Correct back to the true surface-to-surface distance so that
+    # the constraint pipeline (pos = dist - includemargin) works consistently
+    # with the primitive narrowphase, which reports un-inflated distances.
+    dist += margin
+
     witness1[0] = w1
     witness2[0] = w2
 

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

@@ -0,0 +1,372 @@
+# 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.
+# ==============================================================================
+
+"""Core collision types and utilities shared across collision modules."""
+
+import dataclasses
+from typing import Tuple
+
+from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
+from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
+from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
+from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
+import warp as wp
+
+wp.set_module_options({"enable_backward": False})
+
+
+@wp.struct
+class Geom:
+  """Geom properties for pairwise collision detection.
+
+  Bundles a geometry's pose, size, surface normal, and mesh topology data into
+  a single struct that can be passed to Warp collision kernels.
+  """
+
+  pos: wp.vec3
+  rot: wp.mat33
+  normal: wp.vec3
+  size: wp.vec3
+  margin: float
+  hfprism: mat63
+  vertadr: int
+  vertnum: int
+  vert: wp.array(dtype=wp.vec3)
+  graphadr: int
+  graph: wp.array(dtype=int)
+  mesh_polynum: int
+  mesh_polyadr: int
+  mesh_polynormal: wp.array(dtype=wp.vec3)
+  mesh_polyvertadr: wp.array(dtype=int)
+  mesh_polyvertnum: wp.array(dtype=int)
+  mesh_polyvert: wp.array(dtype=int)
+  mesh_polymapadr: wp.array(dtype=int)
+  mesh_polymapnum: wp.array(dtype=int)
+  mesh_polymap: wp.array(dtype=int)
+  index: int
+
+
+@wp.func
+def geom_collision_pair(
+    # Model:
+    geom_type: wp.array(dtype=int),
+    geom_dataid: wp.array(dtype=int),
+    geom_size: wp.array2d(dtype=wp.vec3),
+    mesh_vertadr: wp.array(dtype=int),
+    mesh_vertnum: wp.array(dtype=int),
+    mesh_graphadr: wp.array(dtype=int),
+    mesh_vert: wp.array(dtype=wp.vec3),
+    mesh_graph: wp.array(dtype=int),
+    mesh_polynum: wp.array(dtype=int),
+    mesh_polyadr: wp.array(dtype=int),
+    mesh_polynormal: wp.array(dtype=wp.vec3),
+    mesh_polyvertadr: wp.array(dtype=int),
+    mesh_polyvertnum: wp.array(dtype=int),
+    mesh_polyvert: wp.array(dtype=int),
+    mesh_polymapadr: wp.array(dtype=int),
+    mesh_polymapnum: wp.array(dtype=int),
+    mesh_polymap: wp.array(dtype=int),
+    # Data in:
+    geom_xpos_in: wp.array2d(dtype=wp.vec3),
+    geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    # In:
+    geoms: wp.vec2i,
+    worldid: int,
+) -> Tuple[Geom, Geom]:
+  geom1 = Geom()
+  geom2 = Geom()
+
+  g1 = geoms[0]
+  g2 = geoms[1]
+  geom_type1 = geom_type[g1]
+  geom_type2 = geom_type[g2]
+
+  geom1.pos = geom_xpos_in[worldid, g1]
+  geom1.rot = geom_xmat_in[worldid, g1]
+  geom1.size = geom_size[worldid % geom_size.shape[0], g1]
+  # z-axis of the rotation matrix, used as the surface normal for plane collisions
+  geom1.normal = wp.vec3(geom1.rot[0, 2], geom1.rot[1, 2], geom1.rot[2, 2])
+
+  geom2.pos = geom_xpos_in[worldid, g2]
+  geom2.rot = geom_xmat_in[worldid, g2]
+  geom2.size = geom_size[worldid % geom_size.shape[0], g2]
+  # z-axis of the rotation matrix, used as the surface normal for plane collisions
+  geom2.normal = wp.vec3(geom2.rot[0, 2], geom2.rot[1, 2], geom2.rot[2, 2])
+
+  if geom_type1 == GeomType.MESH:
+    dataid = geom_dataid[g1]
+    geom1.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1)
+    geom1.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1)
+    geom1.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1)
+    geom1.mesh_polynum = wp.where(dataid >= 0, mesh_polynum[dataid], -1)
+    geom1.mesh_polyadr = wp.where(dataid >= 0, mesh_polyadr[dataid], -1)
+
+    geom1.vert = mesh_vert
+    geom1.graph = mesh_graph
+    geom1.mesh_polynormal = mesh_polynormal
+    geom1.mesh_polyvertadr = mesh_polyvertadr
+    geom1.mesh_polyvertnum = mesh_polyvertnum
+    geom1.mesh_polyvert = mesh_polyvert
+    geom1.mesh_polymapadr = mesh_polymapadr
+    geom1.mesh_polymapnum = mesh_polymapnum
+    geom1.mesh_polymap = mesh_polymap
+
+  if geom_type2 == GeomType.MESH:
+    dataid = geom_dataid[g2]
+    geom2.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1)
+    geom2.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1)
+    geom2.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1)
+    geom2.mesh_polynum = wp.where(dataid >= 0, mesh_polynum[dataid], -1)
+    geom2.mesh_polyadr = wp.where(dataid >= 0, mesh_polyadr[dataid], -1)
+
+    geom2.vert = mesh_vert
+    geom2.graph = mesh_graph
+    geom2.mesh_polynormal = mesh_polynormal
+    geom2.mesh_polyvertadr = mesh_polyvertadr
+    geom2.mesh_polyvertnum = mesh_polyvertnum
+    geom2.mesh_polyvert = mesh_polyvert
+    geom2.mesh_polymapadr = mesh_polymapadr
+    geom2.mesh_polymapnum = mesh_polymapnum
+    geom2.mesh_polymap = mesh_polymap
+
+  geom1.index = -1
+  geom1.margin = 0.0
+
+  geom2.index = -1
+  geom2.margin = 0.0
+
+  return geom1, geom2
+
+
+@wp.func
+def write_contact(
+    # Data in:
+    naconmax_in: int,
+    # In:
+    id_: int,
+    dist_in: float,
+    pos_in: wp.vec3,
+    frame_in: wp.mat33,
+    margin_in: float,
+    gap_in: float,
+    condim_in: int,
+    friction_in: vec5,
+    solref_in: wp.vec2,
+    solreffriction_in: wp.vec2,
+    solimp_in: vec5,
+    geoms_in: wp.vec2i,
+    pairid_in: wp.vec2i,
+    worldid_in: int,
+    # Data out:
+    contact_dist_out: wp.array(dtype=float),
+    contact_pos_out: wp.array(dtype=wp.vec3),
+    contact_frame_out: wp.array(dtype=wp.mat33),
+    contact_includemargin_out: wp.array(dtype=float),
+    contact_friction_out: wp.array(dtype=vec5),
+    contact_solref_out: wp.array(dtype=wp.vec2),
+    contact_solreffriction_out: wp.array(dtype=wp.vec2),
+    contact_solimp_out: wp.array(dtype=vec5),
+    contact_dim_out: wp.array(dtype=int),
+    contact_geom_out: wp.array(dtype=wp.vec2i),
+    contact_worldid_out: wp.array(dtype=int),
+    contact_type_out: wp.array(dtype=int),
+    contact_geomcollisionid_out: wp.array(dtype=int),
+    nacon_out: wp.array(dtype=int),
+) -> int:
+  """Atomically write a detected contact into the contact output arrays.
+
+  Returns 1 if the contact is active (dist < margin), 0 otherwise.
+  """
+  active = dist_in < margin_in
+
+  # skip contact and no collision sensor
+  if (pairid_in[0] == -2 or not active) and pairid_in[1] == -1:
+    return 0
+
+  contact_type = 0
+
+  if pairid_in[0] >= -1 and active:
+    contact_type |= ContactType.CONSTRAINT
+
+  if pairid_in[1] >= 0:
+    contact_type |= ContactType.SENSOR
+
+  cid = wp.atomic_add(nacon_out, 0, 1)
+  if cid < naconmax_in:
+    contact_dist_out[cid] = dist_in
+    contact_pos_out[cid] = pos_in
+    contact_frame_out[cid] = frame_in
+    contact_geom_out[cid] = geoms_in
+    contact_worldid_out[cid] = worldid_in
+    includemargin = margin_in - gap_in
+    contact_includemargin_out[cid] = includemargin
+    contact_dim_out[cid] = condim_in
+    contact_friction_out[cid] = friction_in
+    contact_solref_out[cid] = solref_in
+    contact_solreffriction_out[cid] = solreffriction_in
+    contact_solimp_out[cid] = solimp_in
+    contact_type_out[cid] = contact_type
+    contact_geomcollisionid_out[cid] = id_
+    return int(active)
+  return 0
+
+
+@wp.func
+def contact_params(
+    # Model:
+    geom_condim: wp.array(dtype=int),
+    geom_priority: wp.array(dtype=int),
+    geom_solmix: wp.array2d(dtype=float),
+    geom_solref: wp.array2d(dtype=wp.vec2),
+    geom_solimp: wp.array2d(dtype=vec5),
+    geom_friction: wp.array2d(dtype=wp.vec3),
+    geom_margin: wp.array2d(dtype=float),
+    geom_gap: wp.array2d(dtype=float),
+    pair_dim: wp.array(dtype=int),
+    pair_solref: wp.array2d(dtype=wp.vec2),
+    pair_solreffriction: wp.array2d(dtype=wp.vec2),
+    pair_solimp: wp.array2d(dtype=vec5),
+    pair_margin: wp.array2d(dtype=float),
+    pair_gap: wp.array2d(dtype=float),
+    pair_friction: wp.array2d(dtype=vec5),
+    # In:
+    collision_pair_in: wp.array(dtype=wp.vec2i),
+    collision_pairid_in: wp.array(dtype=wp.vec2i),
+    cid: int,
+    worldid: int,
+):
+  """Resolve contact parameters for a collision pair.
+
+  Uses explicit pair overrides when available, otherwise mixes geom-level
+  properties by priority and solmix weights.
+  """
+  geoms = collision_pair_in[cid]
+  pairid = collision_pairid_in[cid][0]
+
+  # TODO(team): early return if collision sensor but no contact
+  # (ie, pairid[0] < -1 and pairid[1] < 0)
+
+  if pairid > -1:
+    margin = pair_margin[worldid % pair_margin.shape[0], pairid]
+    gap = pair_gap[worldid % pair_gap.shape[0], pairid]
+    condim = pair_dim[pairid]
+    friction = pair_friction[worldid % pair_friction.shape[0], pairid]
+    solref = pair_solref[worldid % pair_solref.shape[0], pairid]
+    solreffriction = pair_solreffriction[
+        worldid % pair_solreffriction.shape[0], pairid
+    ]
+    solimp = pair_solimp[worldid % pair_solimp.shape[0], pairid]
+  else:
+    g1 = geoms[0]
+    g2 = geoms[1]
+    solmix_id = worldid % geom_solmix.shape[0]
+    friction_id = worldid % geom_friction.shape[0]
+    solref_id = worldid % geom_solref.shape[0]
+    solimp_id = worldid % geom_solimp.shape[0]
+    margin_id = worldid % geom_margin.shape[0]
+    gap_id = worldid % geom_gap.shape[0]
+
+    solmix1 = geom_solmix[solmix_id, g1]
+    solmix2 = geom_solmix[solmix_id, g2]
+
+    condim1 = geom_condim[g1]
+    condim2 = geom_condim[g2]
+
+    # priority
+    p1 = geom_priority[g1]
+    p2 = geom_priority[g2]
+
+    if p1 > p2:
+      mix = 1.0
+      condim = condim1
+      max_geom_friction = geom_friction[friction_id, g1]
+    elif p2 > p1:
+      mix = 0.0
+      condim = condim2
+      max_geom_friction = geom_friction[friction_id, g2]
+    else:
+      mix = safe_div(solmix1, solmix1 + solmix2)
+      mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 < MJ_MINVAL), 0.5, mix)
+      mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 >= MJ_MINVAL), 0.0, mix)
+      mix = wp.where((solmix1 >= MJ_MINVAL) and (solmix2 < MJ_MINVAL), 1.0, mix)
+      condim = wp.max(condim1, condim2)
+      max_geom_friction = wp.max(
+          geom_friction[friction_id, g1], geom_friction[friction_id, g2]
+      )
+
+    friction = vec5(
+        max_geom_friction[0],
+        max_geom_friction[0],
+        max_geom_friction[1],
+        max_geom_friction[2],
+        max_geom_friction[2],
+    )
+
+    if (
+        geom_solref[solref_id, g1][0] > 0.0
+        and geom_solref[solref_id, g2][0] > 0.0
+    ):
+      solref = (
+          mix * geom_solref[solref_id, g1]
+          + (1.0 - mix) * geom_solref[solref_id, g2]
+      )
+    else:
+      solref = wp.min(geom_solref[solref_id, g1], geom_solref[solref_id, g2])
+
+    solreffriction = wp.vec2(0.0, 0.0)
+    solimp = (
+        mix * geom_solimp[solimp_id, g1]
+        + (1.0 - mix) * geom_solimp[solimp_id, g2]
+    )
+    # geom priority is ignored
+    margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
+    gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2]
+
+  friction = vec5(
+      wp.max(MJ_MINMU, friction[0]),
+      wp.max(MJ_MINMU, friction[1]),
+      wp.max(MJ_MINMU, friction[2]),
+      wp.max(MJ_MINMU, friction[3]),
+      wp.max(MJ_MINMU, friction[4]),
+  )
+
+  return geoms, margin, gap, condim, friction, solref, solreffriction, solimp
+
+
+@dataclasses.dataclass
+class CollisionContext:
+  """Collision driver intermediate arrays.
+
+  Attributes:
+    collision_pair: collision pairs from broadphase             (naconmax, 2)
+    collision_pairid: ids from broadphase                       (naconmax, 2)
+    collision_worldid: collision world ids from broadphase      (naconmax,)
+  """
+
+  collision_pair: wp.array
+  collision_pairid: wp.array
+  collision_worldid: wp.array
+
+
+def create_collision_context(naconmax: int) -> CollisionContext:
+  """Create a CollisionContext with allocated arrays."""
+  return CollisionContext(
+      collision_pair=wp.empty(naconmax, dtype=wp.vec2i),
+      collision_pairid=wp.empty(naconmax, dtype=wp.vec2i),
+      collision_worldid=wp.empty(naconmax, dtype=int),
+  )

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

@@ -15,26 +15,25 @@
 
 from typing import Any
 
-import warp as wp
-
 from mujoco.mjx.third_party.mujoco_warp._src.collision_convex import convex_narrowphase
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import create_collision_context
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import primitive_narrowphase
 from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf_narrowphase
-from mujoco.mjx.third_party.mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
 from mujoco.mjx.third_party.mujoco_warp._src.math import upper_tri_index
-from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseFilter
 from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType
-from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext
 from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
-from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat23
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
+import warp as wp
 
 wp.set_module_options({"enable_backward": False})
 
@@ -76,15 +75,6 @@ MJ_COLLISION_TABLE = {
 }
 
 
-def create_collision_context(naconmax: int) -> CollisionContext:
-  """Create a CollisionContext with allocated arrays."""
-  return CollisionContext(
-    collision_pair=wp.empty(naconmax, dtype=wp.vec2i),
-    collision_pairid=wp.empty(naconmax, dtype=wp.vec2i),
-    collision_worldid=wp.empty(naconmax, dtype=int),
-  )
-
-
 @wp.kernel
 def _zero_nacon_ncollision(
   # Data out:
@@ -102,27 +92,18 @@ def _plane_filter(
   if size1 == 0.0:
     # geom1 is a plane
     dist = wp.dot(xpos2 - xpos1, wp.vec3(xmat1[0, 2], xmat1[1, 2], xmat1[2, 2]))
-    if BLEEDING_EDGE_MUJOCO:
     return dist <= size2 + margin1 + margin2
-    else:
-      return dist <= size2 + wp.max(margin1, margin2)
   elif size2 == 0.0:
     # geom2 is a plane
     dist = wp.dot(xpos1 - xpos2, wp.vec3(xmat2[0, 2], xmat2[1, 2], xmat2[2, 2]))
-    if BLEEDING_EDGE_MUJOCO:
     return dist <= size1 + margin1 + margin2
-    else:
-      return dist <= size1 + wp.max(margin1, margin2)
 
   return True
 
 
 @wp.func
 def _sphere_filter(size1: float, size2: float, margin1: float, margin2: float, xpos1: wp.vec3, xpos2: wp.vec3) -> bool:
-  if BLEEDING_EDGE_MUJOCO:
   bound = size1 + size2 + margin1 + margin2
-  else:
-    bound = size1 + size2 + wp.max(margin1, margin2)
   dif = xpos2 - xpos1
   dist_sq = wp.dot(dif, dif)
   return dist_sq <= bound * bound
@@ -151,10 +132,7 @@ def _aabb_filter(
   center1 = xmat1 @ center1 + xpos1
   center2 = xmat2 @ center2 + xpos2
 
-  if BLEEDING_EDGE_MUJOCO:
   margin = margin1 + margin2
-  else:
-    margin = wp.max(margin1, margin2)
 
   max_x1 = -MJ_MAXVAL
   max_y1 = -MJ_MAXVAL
@@ -249,10 +227,7 @@ def _obb_filter(
   xmat2: wp.mat33,
 ) -> bool:
   """Oriented bounding boxes collision (see Gottschalk et al.), see mj_collideOBB."""
-  if BLEEDING_EDGE_MUJOCO:
   margin = margin1 + margin2
-  else:
-    margin = wp.max(margin1, margin2)
 
   xcenter = mat23()
   normal = mat63()
@@ -315,13 +290,25 @@ def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound
     # 8: obb
 
     aabb_id = worldid % ngeom_aabb if wp.static(ngeom_aabb > 1) else 0
-    center1, center2 = geom_aabb[aabb_id, geom1, 0], geom_aabb[aabb_id, geom2, 0]
+    center1, center2 = (
-    size1, size2 = geom_aabb[aabb_id, geom1, 1], geom_aabb[aabb_id, geom2, 1]
+        geom_aabb[aabb_id, geom1, 0],
+        geom_aabb[aabb_id, geom2, 0],
+    )  # kernel_analyzer: ignore
+    size1, size2 = (
+        geom_aabb[aabb_id, geom1, 1],
+        geom_aabb[aabb_id, geom2, 1],
+    )  # kernel_analyzer: ignore
 
     rbound_id = worldid % ngeom_rbound if wp.static(ngeom_rbound > 1) else 0
-    rbound1, rbound2 = geom_rbound[rbound_id, geom1], geom_rbound[rbound_id, geom2]
+    rbound1, rbound2 = (
+        geom_rbound[rbound_id, geom1],
+        geom_rbound[rbound_id, geom2],
+    )  # kernel_analyzer: ignore
     margin_id = worldid % ngeom_margin if wp.static(ngeom_margin > 1) else 0
-    margin1, margin2 = geom_margin[margin_id, geom1], geom_margin[margin_id, geom2]
+    margin1, margin2 = (
+        geom_margin[margin_id, geom1],
+        geom_margin[margin_id, geom2],
+    )  # kernel_analyzer: ignore
     xpos1, xpos2 = geom_xpos_in[worldid, geom1], geom_xpos_in[worldid, geom2]
     xmat1, xmat2 = geom_xmat_in[worldid, geom1], geom_xmat_in[worldid, geom2]
 
@@ -804,3 +791,6 @@ def collision(m: Model, d: Data):
     sap_broadphase(m, d, ctx)
 
   _narrowphase(m, d, ctx)
+
+  if m.callback.contactfilter:
+    m.callback.contactfilter(m, d)

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

@@ -16,12 +16,11 @@
 import math
 from typing import Tuple
 
-import warp as wp
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom
-
-from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
+import warp as wp
 
 # TODO(team): improve compile time to enable backward pass
 wp.set_module_options({"enable_backward": False})
@@ -581,14 +580,13 @@ def gjk(
   simplex_index1 = wp.vec4i()
   simplex_index2 = wp.vec4i()
   n = int(0)
-  cnt = int(1)
   coordinates = wp.vec4()  # barycentric coordinates
   epsilon = wp.where(is_discrete, 0.0, 0.5 * tolerance * tolerance)
 
   # set initial guess
   x_k = x1_0 - x2_0
 
-  for _ in range(gjk_iterations):
+  for k in range(gjk_iterations):
     xnorm = wp.dot(x_k, x_k)
     # TODO(kbayes): determine new constant here
     if xnorm < 1e-12:
@@ -665,10 +663,12 @@ def gjk(
     if n == 4:
       break
 
-    cnt += 1
+    if k == gjk_iterations - 1:
-
+      wp.printf(
-  if cnt == gjk_iterations:
+          "Warning: opt.ccd_iterations, currently set to %d, needs to be"
-    wp.printf("Warning: opt.ccd_iterations, currently set to %d, needs to be increased.\n", gjk_iterations)
+          " increased.\n",
+          gjk_iterations,
+      )
 
   result = GJKResult()
 
@@ -1220,14 +1220,13 @@ def _epa(
   idx = int(-1)
   pidx = int(-1)
   epsilon = wp.where(is_discrete, 1e-15, tolerance)
-  cnt = int(1)
   nvalid = pt.nface  # number of potential faces for expanding the polytope
 
   # the face vertices are encoded in 10-bits that index the vertex array,
   # so iterations must be cap to limit the number of generated vertices
   # (one new vertex per iteration)
   epa_iterations = wp.min(epa_iterations, 1000)
-  for _ in range(epa_iterations):
+  for k in range(epa_iterations):
     pidx = idx
     idx = int(-1)
     lower2 = float(FLOAT_MAX)
@@ -1325,10 +1324,13 @@ def _epa(
 
     # clear horizon
     pt.nhorizon = 0
-    cnt += 1
 
-  if cnt == epa_iterations:
+    if k == epa_iterations - 1:
-    wp.printf("Warning: opt.ccd_iterations, currently set to %d, needs to be increased.\n", gjk_iterations)
+      wp.printf(
+          "Warning: opt.ccd_iterations, currently set to %d, needs to be"
+          " increased.\n",
+          gjk_iterations,
+      )
 
   # return from valid face
   if idx > -1:

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

@@ -15,8 +15,11 @@
 
 from typing import Tuple
 
-import warp as wp
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext
-
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import Geom
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import box_box
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import capsule_box
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import capsule_capsule
@@ -29,142 +32,21 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sph
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sphere_capsule
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sphere_cylinder
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive_core import sphere_sphere
-from mujoco.mjx.third_party.mujoco_warp._src.io import BLEEDING_EDGE_MUJOCO
 from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
-from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
 from mujoco.mjx.third_party.mujoco_warp._src.math import upper_trid_index
-from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
-from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU
-from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
-from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext
-from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
-from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat43
-from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
+import warp as wp
 
 wp.set_module_options({"enable_backward": False})
 
 
-@wp.struct
-class Geom:
-  pos: wp.vec3
-  rot: wp.mat33
-  normal: wp.vec3
-  size: wp.vec3
-  margin: float
-  hfprism: mat63
-  vertadr: int
-  vertnum: int
-  vert: wp.array(dtype=wp.vec3)
-  graphadr: int
-  graph: wp.array(dtype=int)
-  mesh_polynum: int
-  mesh_polyadr: int
-  mesh_polynormal: wp.array(dtype=wp.vec3)
-  mesh_polyvertadr: wp.array(dtype=int)
-  mesh_polyvertnum: wp.array(dtype=int)
-  mesh_polyvert: wp.array(dtype=int)
-  mesh_polymapadr: wp.array(dtype=int)
-  mesh_polymapnum: wp.array(dtype=int)
-  mesh_polymap: wp.array(dtype=int)
-  index: int
-
-
-@wp.func
-def geom_collision_pair(
-  # Model:
-  geom_type: wp.array(dtype=int),
-  geom_dataid: wp.array(dtype=int),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  mesh_vertadr: wp.array(dtype=int),
-  mesh_vertnum: wp.array(dtype=int),
-  mesh_graphadr: wp.array(dtype=int),
-  mesh_vert: wp.array(dtype=wp.vec3),
-  mesh_graph: wp.array(dtype=int),
-  mesh_polynum: wp.array(dtype=int),
-  mesh_polyadr: wp.array(dtype=int),
-  mesh_polynormal: wp.array(dtype=wp.vec3),
-  mesh_polyvertadr: wp.array(dtype=int),
-  mesh_polyvertnum: wp.array(dtype=int),
-  mesh_polyvert: wp.array(dtype=int),
-  mesh_polymapadr: wp.array(dtype=int),
-  mesh_polymapnum: wp.array(dtype=int),
-  mesh_polymap: wp.array(dtype=int),
-  # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
-  # In:
-  geoms: wp.vec2i,
-  worldid: int,
-) -> Tuple[Geom, Geom]:
-  geom1 = Geom()
-  geom2 = Geom()
-
-  g1 = geoms[0]
-  g2 = geoms[1]
-  geom_type1 = geom_type[g1]
-  geom_type2 = geom_type[g2]
-
-  geom1.pos = geom_xpos_in[worldid, g1]
-  geom1.rot = geom_xmat_in[worldid, g1]
-  geom1.size = geom_size[worldid % geom_size.shape[0], g1]
-  geom1.normal = wp.vec3(geom1.rot[0, 2], geom1.rot[1, 2], geom1.rot[2, 2])  # plane
-
-  geom2.pos = geom_xpos_in[worldid, g2]
-  geom2.rot = geom_xmat_in[worldid, g2]
-  geom2.size = geom_size[worldid % geom_size.shape[0], g2]
-  geom2.normal = wp.vec3(geom2.rot[0, 2], geom2.rot[1, 2], geom2.rot[2, 2])  # plane
-
-  if geom_type1 == GeomType.MESH:
-    dataid = geom_dataid[g1]
-    geom1.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1)
-    geom1.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1)
-    geom1.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1)
-    geom1.mesh_polynum = wp.where(dataid >= 0, mesh_polynum[dataid], -1)
-    geom1.mesh_polyadr = wp.where(dataid >= 0, mesh_polyadr[dataid], -1)
-
-    geom1.vert = mesh_vert
-    geom1.graph = mesh_graph
-    geom1.mesh_polynormal = mesh_polynormal
-    geom1.mesh_polyvertadr = mesh_polyvertadr
-    geom1.mesh_polyvertnum = mesh_polyvertnum
-    geom1.mesh_polyvert = mesh_polyvert
-    geom1.mesh_polymapadr = mesh_polymapadr
-    geom1.mesh_polymapnum = mesh_polymapnum
-    geom1.mesh_polymap = mesh_polymap
-
-  if geom_type2 == GeomType.MESH:
-    dataid = geom_dataid[g2]
-    geom2.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1)
-    geom2.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1)
-    geom2.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1)
-    geom2.mesh_polynum = wp.where(dataid >= 0, mesh_polynum[dataid], -1)
-    geom2.mesh_polyadr = wp.where(dataid >= 0, mesh_polyadr[dataid], -1)
-
-    geom2.vert = mesh_vert
-    geom2.graph = mesh_graph
-    geom2.mesh_polynormal = mesh_polynormal
-    geom2.mesh_polyvertadr = mesh_polyvertadr
-    geom2.mesh_polyvertnum = mesh_polyvertnum
-    geom2.mesh_polyvert = mesh_polyvert
-    geom2.mesh_polymapadr = mesh_polymapadr
-    geom2.mesh_polymapnum = mesh_polymapnum
-    geom2.mesh_polymap = mesh_polymap
-
-  geom1.index = -1
-  geom1.margin = 0.0
-
-  geom2.index = -1
-  geom2.margin = 0.0
-
-  return geom1, geom2
-
-
 @wp.func
 def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tuple[wp.vec4, mat43, wp.vec3]:
   """Core contact geometry calculation for plane-convex collision.
@@ -394,183 +276,6 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup
   return contact_dist, contact_pos, plane_normal
 
 
-@wp.func
-def write_contact(
-  # Data in:
-  naconmax_in: int,
-  # In:
-  id_: int,
-  dist_in: float,
-  pos_in: wp.vec3,
-  frame_in: wp.mat33,
-  margin_in: float,
-  gap_in: float,
-  condim_in: int,
-  friction_in: vec5,
-  solref_in: wp.vec2,
-  solreffriction_in: wp.vec2,
-  solimp_in: vec5,
-  geoms_in: wp.vec2i,
-  pairid_in: wp.vec2i,
-  worldid_in: int,
-  # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
-) -> int:
-  active = dist_in < margin_in
-
-  # skip contact and no collision sensor
-  if (pairid_in[0] == -2 or not active) and pairid_in[1] == -1:
-    return 0
-
-  contact_type = 0
-
-  if pairid_in[0] >= -1 and active:
-    contact_type |= ContactType.CONSTRAINT
-
-  if pairid_in[1] >= 0:
-    contact_type |= ContactType.SENSOR
-
-  cid = wp.atomic_add(nacon_out, 0, 1)
-  if cid < naconmax_in:
-    contact_dist_out[cid] = dist_in
-    contact_pos_out[cid] = pos_in
-    contact_frame_out[cid] = frame_in
-    contact_geom_out[cid] = geoms_in
-    contact_worldid_out[cid] = worldid_in
-    includemargin = margin_in - gap_in
-    contact_includemargin_out[cid] = includemargin
-    contact_dim_out[cid] = condim_in
-    contact_friction_out[cid] = friction_in
-    contact_solref_out[cid] = solref_in
-    contact_solreffriction_out[cid] = solreffriction_in
-    contact_solimp_out[cid] = solimp_in
-    contact_type_out[cid] = contact_type
-    contact_geomcollisionid_out[cid] = id_
-    return int(active)
-  return 0
-
-
-@wp.func
-def contact_params(
-  # Model:
-  geom_condim: wp.array(dtype=int),
-  geom_priority: wp.array(dtype=int),
-  geom_solmix: wp.array2d(dtype=float),
-  geom_solref: wp.array2d(dtype=wp.vec2),
-  geom_solimp: wp.array2d(dtype=vec5),
-  geom_friction: wp.array2d(dtype=wp.vec3),
-  geom_margin: wp.array2d(dtype=float),
-  geom_gap: wp.array2d(dtype=float),
-  pair_dim: wp.array(dtype=int),
-  pair_solref: wp.array2d(dtype=wp.vec2),
-  pair_solreffriction: wp.array2d(dtype=wp.vec2),
-  pair_solimp: wp.array2d(dtype=vec5),
-  pair_margin: wp.array2d(dtype=float),
-  pair_gap: wp.array2d(dtype=float),
-  pair_friction: wp.array2d(dtype=vec5),
-  # In:
-  collision_pair_in: wp.array(dtype=wp.vec2i),
-  collision_pairid_in: wp.array(dtype=wp.vec2i),
-  cid: int,
-  worldid: int,
-):
-  geoms = collision_pair_in[cid]
-  pairid = collision_pairid_in[cid][0]
-
-  # TODO(team): early return if collision sensor but no contact
-  # (ie, pairid[0] < -1 and pairid[1] < 0)
-
-  if pairid > -1:
-    margin = pair_margin[worldid, pairid]
-    gap = pair_gap[worldid, pairid]
-    condim = pair_dim[pairid]
-    friction = pair_friction[worldid, pairid]
-    solref = pair_solref[worldid, pairid]
-    solreffriction = pair_solreffriction[worldid, pairid]
-    solimp = pair_solimp[worldid, pairid]
-  else:
-    g1 = geoms[0]
-    g2 = geoms[1]
-    solmix_id = worldid % geom_solmix.shape[0]
-    friction_id = worldid % geom_friction.shape[0]
-    solref_id = worldid % geom_solref.shape[0]
-    solimp_id = worldid % geom_solimp.shape[0]
-    margin_id = worldid % geom_margin.shape[0]
-    gap_id = worldid % geom_gap.shape[0]
-
-    solmix1 = geom_solmix[solmix_id, g1]
-    solmix2 = geom_solmix[solmix_id, g2]
-
-    condim1 = geom_condim[g1]
-    condim2 = geom_condim[g2]
-
-    # priority
-    p1 = geom_priority[g1]
-    p2 = geom_priority[g2]
-
-    if p1 > p2:
-      mix = 1.0
-      condim = condim1
-      max_geom_friction = geom_friction[friction_id, g1]
-    elif p2 > p1:
-      mix = 0.0
-      condim = condim2
-      max_geom_friction = geom_friction[friction_id, g2]
-    else:
-      mix = safe_div(solmix1, solmix1 + solmix2)
-      mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 < MJ_MINVAL), 0.5, mix)
-      mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 >= MJ_MINVAL), 0.0, mix)
-      mix = wp.where((solmix1 >= MJ_MINVAL) and (solmix2 < MJ_MINVAL), 1.0, mix)
-      condim = wp.max(condim1, condim2)
-      max_geom_friction = wp.max(geom_friction[friction_id, g1], geom_friction[friction_id, g2])
-
-    friction = vec5(
-      max_geom_friction[0],
-      max_geom_friction[0],
-      max_geom_friction[1],
-      max_geom_friction[2],
-      max_geom_friction[2],
-    )
-
-    if geom_solref[solref_id, g1][0] > 0.0 and geom_solref[solref_id, g2][0] > 0.0:
-      solref = mix * geom_solref[solref_id, g1] + (1.0 - mix) * geom_solref[solref_id, g2]
-    else:
-      solref = wp.min(geom_solref[solref_id, g1], geom_solref[solref_id, g2])
-
-    solreffriction = wp.vec2(0.0, 0.0)
-    solimp = mix * geom_solimp[solimp_id, g1] + (1.0 - mix) * geom_solimp[solimp_id, g2]
-    # geom priority is ignored
-    if BLEEDING_EDGE_MUJOCO:
-      margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
-      gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2]
-    else:
-      margin = wp.max(geom_margin[margin_id, g1], geom_margin[margin_id, g2])
-      gap = wp.max(geom_gap[gap_id, g1], geom_gap[gap_id, g2])
-
-  friction = vec5(
-    wp.max(MJ_MINMU, friction[0]),
-    wp.max(MJ_MINMU, friction[1]),
-    wp.max(MJ_MINMU, friction[2]),
-    wp.max(MJ_MINMU, friction[3]),
-    wp.max(MJ_MINMU, friction[4]),
-  )
-
-  return geoms, margin, gap, condim, friction, solref, solreffriction, solimp
-
-
 @wp.func
 def plane_sphere_wrapper(
   # Data in:

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

@@ -15,14 +15,12 @@
 
 from typing import Tuple
 
-import warp as wp
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import CollisionContext
-
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_params
-from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import contact_params
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair
-from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import geom_collision_pair
+from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
-from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import write_contact
 from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
 from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh
-from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionContext
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
@@ -31,6 +29,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec8
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i
 from mujoco.mjx.third_party.mujoco_warp._src.util_misc import halton
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
+import warp as wp
 
 wp.set_module_options({"enable_backward": False})
 
@@ -56,6 +55,7 @@ class VolumeData:
   oct_aabb: wp.array2d(dtype=wp.vec3)
   oct_child: wp.array(dtype=vec8i)
   oct_coeff: wp.array(dtype=vec8)
+  root: int = 0
   valid: bool = False
 
 
@@ -81,6 +81,7 @@ def get_sdf_params(
     oct_child: wp.array(dtype=vec8i),
     oct_aabb: wp.array2d(dtype=wp.vec3),
     oct_coeff: wp.array(dtype=vec8),
+    mesh_octadr: wp.array(dtype=int),
     plugin: wp.array(dtype=int),
     plugin_attr: wp.array(dtype=wp.vec3f),
     # In:
@@ -98,8 +99,10 @@ def get_sdf_params(
     plugin_index = plugin[plugin_id]
 
   elif g_type == GeomType.SDF and mesh_id != -1:
-    volume_data.center = oct_aabb[mesh_id, 0]
+    octadr = mesh_octadr[mesh_id]
-    volume_data.half_size = oct_aabb[mesh_id, 1]
+    volume_data.center = oct_aabb[octadr, 0]
+    volume_data.half_size = oct_aabb[octadr, 1]
+    volume_data.root = octadr
     volume_data.oct_aabb = oct_aabb
     volume_data.oct_child = oct_child
     volume_data.oct_coeff = oct_coeff
@@ -225,9 +228,13 @@ def user_sdf_grad(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> wp.vec3:
 
 @wp.func
 def find_oct(
-  oct_child: wp.array(dtype=vec8i), oct_aabb: wp.array2d(dtype=wp.vec3), p: wp.vec3, grad: bool
+    oct_child: wp.array(dtype=vec8i),
+    oct_aabb: wp.array2d(dtype=wp.vec3),
+    p: wp.vec3,
+    grad: bool,
+    root: int,
 ) -> Tuple[int, Tuple[vec8, vec8, vec8]]:
-  stack = int(0)
+  stack = root
   niter = int(100)
   rx = vec8(0.0)
   ry = vec8(0.0)
@@ -258,8 +265,10 @@ def find_oct(
     coord = wp.cw_div(p - vmin, vmax - vmin)
 
     # check if the node is a leaf
+    # child indices are relative to root (mesh_octadr offset)
+    child0 = oct_child[node][0]
     if (
-      oct_child[node][0] == -1
+        child0 == -1
         and oct_child[node][1] == -1
         and oct_child[node][2] == -1
         and oct_child[node][3] == -1
@@ -282,10 +291,12 @@ def find_oct(
       return node, (rx, ry, rz)
 
     # compute which of 8 children to visit next
+    # child indices are stored relative to mesh_octadr, add root offset
     x = 0 if coord[0] < 0.5 else 1
     y = 0 if coord[1] < 0.5 else 1
     z = 0 if coord[2] < 0.5 else 1
-    stack = oct_child[node][4 * z + 2 * y + x]
+    child = oct_child[node][4 * z + 2 * y + x]
+    stack = child + root if child != -1 else -1
 
   wp.print("ERROR: Node not found\n")
   return -1, (rx, ry, rz)
@@ -331,7 +342,13 @@ def box_project(center: wp.vec3, half_size: wp.vec3, xyz: wp.vec3) -> Tuple[floa
 @wp.func
 def sample_volume_sdf(xyz: wp.vec3, volume_data: VolumeData) -> float:
   dist0, point = box_project(volume_data.center, volume_data.half_size, xyz)
-  node, weights = find_oct(volume_data.oct_child, volume_data.oct_aabb, point, grad=False)
+  node, weights = find_oct(
+      volume_data.oct_child,
+      volume_data.oct_aabb,
+      point,
+      grad=False,
+      root=volume_data.root,
+  )
   return dist0 + wp.dot(weights[0], volume_data.oct_coeff[node])
 
 
@@ -348,7 +365,13 @@ def sample_volume_grad(xyz: wp.vec3, volume_data: VolumeData) -> wp.vec3:
     grad_y = (sample_volume_sdf(xyz + dy, volume_data) - f) / h
     grad_z = (sample_volume_sdf(xyz + dz, volume_data) - f) / h
     return wp.vec3(grad_x, grad_y, grad_z)
-  node, weights = find_oct(volume_data.oct_child, volume_data.oct_aabb, point, grad=True)
+  node, weights = find_oct(
+      volume_data.oct_child,
+      volume_data.oct_aabb,
+      point,
+      grad=True,
+      root=volume_data.root,
+  )
   grad_x = wp.dot(weights[0], volume_data.oct_coeff[node])
   grad_y = wp.dot(weights[1], volume_data.oct_coeff[node])
   grad_z = wp.dot(weights[2], volume_data.oct_coeff[node])
@@ -642,6 +665,7 @@ def _sdf_narrowphase(
     mesh_vertadr: wp.array(dtype=int),
     mesh_vertnum: wp.array(dtype=int),
     mesh_faceadr: wp.array(dtype=int),
+    mesh_octadr: wp.array(dtype=int),
     mesh_graphadr: wp.array(dtype=int),
     mesh_vert: wp.array(dtype=wp.vec3),
     mesh_face: wp.array(dtype=wp.vec3i),
@@ -775,11 +799,29 @@ def _sdf_narrowphase(
   rot1 = geom1.rot
 
   attr1, g1_plugin_id, volume_data1, mesh_data1 = get_sdf_params(
-    oct_child, oct_aabb, oct_coeff, plugin, plugin_attr, type1, geom1.size, g1_plugin, geom_dataid[g1]
+      oct_child,
+      oct_aabb,
+      oct_coeff,
+      mesh_octadr,
+      plugin,
+      plugin_attr,
+      type1,
+      geom1.size,
+      g1_plugin,
+      geom_dataid[g1],
   )
 
   attr2, g2_plugin_id, volume_data2, mesh_data2 = get_sdf_params(
-    oct_child, oct_aabb, oct_coeff, plugin, plugin_attr, type2, geom2.size, g2_plugin, geom_dataid[g2]
+      oct_child,
+      oct_aabb,
+      oct_coeff,
+      mesh_octadr,
+      plugin,
+      plugin_attr,
+      type2,
+      geom2.size,
+      g2_plugin,
+      geom_dataid[g2],
   )
 
   mesh_data1.nmeshface = nmeshface
@@ -890,6 +932,7 @@ def sdf_narrowphase(m: Model, d: Data, ctx: CollisionContext):
           m.mesh_vertadr,
           m.mesh_vertnum,
           m.mesh_faceadr,
+          m.mesh_octadr,
           m.mesh_graphadr,
           m.mesh_vert,
           m.mesh_face,

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

@@ -21,9 +21,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 TileSet
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
 wp.set_module_options({"enable_backward": False})
@@ -96,49 +94,16 @@ def _nonzero_mask(x: float) -> float:
   return 0.0
 
 
-@cache_kernel
-def _qderiv_actuator_passive_actuation_dense(tile: TileSet, nu: int):
-  @wp.kernel(module="unique", enable_backward=False)
-  def kernel(
-    # Data in:
-    actuator_moment_in: wp.array3d(dtype=float),
-    qM_in: wp.array3d(dtype=float),
-    # In:
-    vel_in: wp.array3d(dtype=float),
-    adr: wp.array(dtype=int),
-    # Out:
-    qDeriv_out: wp.array3d(dtype=float),
-  ):
-    worldid, nodeid = wp.tid()
-    TILE_SIZE = wp.static(tile.size)
-    NU = wp.static(nu)
-
-    dofid = adr[nodeid]
-    vel_tile = wp.tile_load(vel_in[worldid], shape=(NU, 1), bounds_check=False)
-    moment_tile = wp.tile_load(actuator_moment_in[worldid], shape=(NU, TILE_SIZE), offset=(0, dofid), bounds_check=False)
-    moment_weighted = wp.tile_map(wp.mul, wp.tile_broadcast(vel_tile, shape=(NU, TILE_SIZE)), moment_tile)
-    qderiv_tile = wp.tile_matmul(wp.tile_transpose(moment_tile), moment_weighted)
-
-    # Mask out cross-terms for DOF pairs that are structurally zero in M
-    # (e.g., sibling DOFs coupled only through tendons).  Without this,
-    # stale actuation values at sibling positions make A = M - dt*qDeriv
-    # non-positive-definite, causing the tiled Cholesky to produce NaN.
-    # Dropping these terms matches MuJoCo CPU's implicitfast approximation.
-    qM_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid), bounds_check=False)
-    mask_tile = wp.tile_map(_nonzero_mask, qM_tile)
-    qderiv_tile = wp.tile_map(wp.mul, qderiv_tile, mask_tile)
-
-    wp.tile_store(qDeriv_out[worldid], qderiv_tile, offset=(dofid, dofid), bounds_check=False)
-
-  return kernel
-
-
 @wp.kernel
 def _qderiv_actuator_passive_actuation_sparse(
     # Model:
     nu: int,
+    is_sparse: bool,
     # Data in:
-  actuator_moment_in: wp.array3d(dtype=float),
+    moment_rownnz_in: wp.array2d(dtype=int),
+    moment_rowadr_in: wp.array2d(dtype=int),
+    moment_colind_in: wp.array2d(dtype=int),
+    actuator_moment_in: wp.array2d(dtype=float),
     # In:
     vel_in: wp.array2d(dtype=float),
     qMi: wp.array(dtype=int),
@@ -156,12 +121,33 @@ def _qderiv_actuator_passive_actuation_sparse(
     if vel == 0.0:
       continue
 
-    moment_i = actuator_moment_in[worldid, actid, dofiid]
+    # TODO(team): restructure sparse version for better parallelism?
-    moment_j = actuator_moment_in[worldid, actid, dofjid]
+    moment_i = float(0.0)
+    moment_j = float(0.0)
+
+    rownnz = moment_rownnz_in[worldid, actid]
+    rowadr = moment_rowadr_in[worldid, actid]
+    for i in range(rownnz):
+      sparseid = rowadr + i
+      colind = moment_colind_in[worldid, sparseid]
+      if colind == dofiid:
+        moment_i = actuator_moment_in[worldid, sparseid]
+      if colind == dofjid:
+        moment_j = actuator_moment_in[worldid, sparseid]
+      if moment_i != 0.0 and moment_j != 0.0:
+        break
+
+    if moment_i == 0 and moment_j == 0:
+      continue
 
     qderiv_contrib += moment_i * moment_j * vel
 
+  if is_sparse:
     qDeriv_out[worldid, 0, elemid] = qderiv_contrib
+  else:
+    qDeriv_out[worldid, dofiid, dofjid] = qderiv_contrib
+    if dofiid != dofjid:
+      qDeriv_out[worldid, dofjid, dofiid] = qderiv_contrib
 
 
 @wp.kernel
@@ -277,23 +263,22 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
         ],
         outputs=[vel],
       )
-      if m.is_sparse:
       wp.launch(
           _qderiv_actuator_passive_actuation_sparse,
           dim=(d.nworld, qMi.size),
-          inputs=[m.nu, d.actuator_moment, vel, qMi, qMj],
+          inputs=[
+              m.nu,
+              m.is_sparse,
+              d.moment_rownnz,
+              d.moment_rowadr,
+              d.moment_colind,
+              d.actuator_moment,
+              vel,
+              qMi,
+              qMj,
+          ],
           outputs=[out],
       )
-      else:
-        vel_3d = vel.reshape(vel.shape + (1,))
-        for tile in m.qM_tiles:
-          wp.launch_tiled(
-            _qderiv_actuator_passive_actuation_dense(tile, m.nu),
-            dim=(d.nworld, tile.adr.size),
-            inputs=[d.actuator_moment, d.qM, vel_3d, tile.adr],
-            outputs=[out],
-            block_dim=m.block_dim.qderiv_actuator_dense,
-          )
     wp.launch(
       _qderiv_actuator_passive,
       dim=(d.nworld, qMi.size),

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

@@ -15,11 +15,10 @@
 
 from typing import Optional
 
-import warp as wp
-
 from mujoco.mjx.third_party.mujoco_warp._src import collision_driver
 from mujoco.mjx.third_party.mujoco_warp._src import constraint
 from mujoco.mjx.third_party.mujoco_warp._src import derivative
+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
@@ -27,7 +26,6 @@ from mujoco.mjx.third_party.mujoco_warp._src import smooth
 from mujoco.mjx.third_party.mujoco_warp._src import solver
 from mujoco.mjx.third_party.mujoco_warp._src import util_misc
 from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate
-from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
@@ -36,12 +34,14 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import GainType
 from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType
 from mujoco.mjx.third_party.mujoco_warp._src.types import JointType
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet
 from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
+import warp as wp
 
 wp.set_module_options({"enable_backward": False})
 
@@ -146,6 +146,8 @@ def _next_act(
   if actuator_dyntype == DynType.FILTEREXACT:
     tau = wp.max(MJ_MINVAL, actuator_dynprm[0])
     act = act_in + act_dot_scale * act_dot_in * tau * (1.0 - wp.exp(-opt_timestep / tau))
+  elif actuator_dyntype == DynType.USER:
+    return act_in
   else:
     act = act_in + act_dot_scale * act_dot_in * opt_timestep
 
@@ -161,6 +163,8 @@ def _next_activation(
     # Model:
     opt_timestep: wp.array(dtype=float),
     actuator_dyntype: wp.array(dtype=int),
+    actuator_actadr: wp.array(dtype=int),
+    actuator_actnum: wp.array(dtype=int),
     actuator_actlimited: wp.array(dtype=bool),
     actuator_dynprm: wp.array2d(dtype=vec10f),
     actuator_actrange: wp.array2d(dtype=wp.vec2),
@@ -173,21 +177,24 @@ def _next_activation(
     # Data out:
     act_out: wp.array2d(dtype=float),
 ):
-  worldid, actid = wp.tid()
+  worldid, uid = wp.tid()
   opt_timestep_id = worldid % opt_timestep.shape[0]
   actuator_dynprm_id = worldid % actuator_dynprm.shape[0]
   actuator_actrange_id = worldid % actuator_actrange.shape[0]
+  actadr = actuator_actadr[uid]
+  actnum = actuator_actnum[uid]
+  for j in range(actadr, actadr + actnum):
     act = _next_act(
         opt_timestep[opt_timestep_id],
-    actuator_dyntype[actid],
+        actuator_dyntype[uid],
-    actuator_dynprm[actuator_dynprm_id, actid],
+        actuator_dynprm[actuator_dynprm_id, uid],
-    actuator_actrange[actuator_actrange_id, actid],
+        actuator_actrange[actuator_actrange_id, uid],
-    act_in[worldid, actid],
+        act_in[worldid, j],
-    act_dot_in[worldid, actid],
+        act_dot_in[worldid, j],
         act_dot_scale,
-    limit and actuator_actlimited[actid],
+        limit and actuator_actlimited[uid],
     )
-  act_out[worldid, actid] = act
+    act_out[worldid, j] = act
 
 
 @wp.kernel
@@ -230,10 +237,12 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
   # advance activations
   wp.launch(
       _next_activation,
-    dim=(d.nworld, m.na),
+      dim=(d.nworld, m.nu),
       inputs=[
           m.opt.timestep,
           m.actuator_dyntype,
+          m.actuator_actadr,
+          m.actuator_actnum,
           m.actuator_actlimited,
           m.actuator_dynprm,
           m.actuator_actrange,
@@ -285,7 +294,9 @@ def _euler_damp_qfrc_sparse(
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
 
   adr = dof_Madr[tid]
-  qM_integration_out[worldid, 0, adr] += timestep * dof_damping[worldid, tid]
+  qM_integration_out[worldid, 0, adr] += (
+      timestep * dof_damping[worldid % dof_damping.shape[0], tid]
+  )
 
 
 @cache_kernel
@@ -380,8 +391,20 @@ def _rk_perturb_state(
   if m.na and act_t0 is not None:
     wp.launch(
         _next_activation,
-      dim=(d.nworld, m.na),
+        dim=(d.nworld, m.nu),
-      inputs=[m.opt.timestep, act_t0, d.act_dot, scale, False],
+        inputs=[
+            m.opt.timestep,
+            m.actuator_dyntype,
+            m.actuator_actadr,
+            m.actuator_actnum,
+            m.actuator_actlimited,
+            m.actuator_dynprm,
+            m.actuator_actrange,
+            act_t0,
+            d.act_dot,
+            scale,
+            False,
+        ],
         outputs=[d.act],
     )
 
@@ -517,28 +540,35 @@ 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 False and not (m.opt.disableflags & DisableBit.ISLAND):
+    island.island(m, d)
   smooth.transmission(m, d)
 
 
-# TODO(team): sparse actuator_moment version
+@wp.kernel
-@cache_kernel
+def _actuator_velocity(
-def _actuator_velocity(nv: int):
-  @wp.kernel(module="unique", enable_backward=False)
-  def actuator_velocity(
     # Data in:
     qvel_in: wp.array2d(dtype=float),
-    actuator_moment_in: wp.array3d(dtype=float),
+    moment_rownnz_in: wp.array2d(dtype=int),
+    moment_rowadr_in: wp.array2d(dtype=int),
+    moment_colind_in: wp.array2d(dtype=int),
+    actuator_moment_in: wp.array2d(dtype=float),
     # Data out:
     actuator_velocity_out: wp.array2d(dtype=float),
-  ):
+):
   worldid, actid = wp.tid()
-    moment_tile = wp.tile_load(actuator_moment_in[worldid, actid], shape=wp.static(nv))
-    qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(nv))
-    moment_qvel_tile = wp.tile_map(wp.mul, moment_tile, qvel_tile)
-    actuator_velocity_tile = wp.tile_reduce(wp.add, moment_qvel_tile)
-    actuator_velocity_out[worldid, actid] = actuator_velocity_tile[0]
 
-  return actuator_velocity
+  rownnz = moment_rownnz_in[worldid, actid]
+  rowadr = moment_rowadr_in[worldid, actid]
+
+  vel = float(0.0)
+  for i in range(rownnz):
+    sparseid = rowadr + i
+    colind = moment_colind_in[worldid, sparseid]
+    vel += actuator_moment_in[worldid, sparseid] * qvel_in[worldid, colind]
+
+  actuator_velocity_out[worldid, actid] = vel
 
 
 @cache_kernel
@@ -565,9 +595,15 @@ def _tendon_velocity(nv: int):
 def fwd_velocity(m: Model, d: Data):
   """Velocity-dependent computations."""
   wp.launch_tiled(
-    _actuator_velocity(m.nv),
+      _actuator_velocity,
       dim=(d.nworld, m.nu),
-    inputs=[d.qvel, d.actuator_moment],
+      inputs=[
+          d.qvel,
+          d.moment_rownnz,
+          d.moment_rowadr,
+          d.moment_colind,
+          d.actuator_moment,
+      ],
       outputs=[d.actuator_velocity],
       block_dim=m.block_dim.actuator_velocity,
   )
@@ -607,8 +643,8 @@ def _actuator_force(
     actuator_ctrlrange: wp.array2d(dtype=wp.vec2),
     actuator_forcerange: wp.array2d(dtype=wp.vec2),
     actuator_actrange: wp.array2d(dtype=wp.vec2),
-  actuator_acc0: wp.array(dtype=float),
+    actuator_acc0: wp.array2d(dtype=float),
-  actuator_lengthrange: wp.array(dtype=wp.vec2),
+    actuator_lengthrange: wp.array2d(dtype=wp.vec2),
     # Data in:
     act_in: wp.array2d(dtype=float),
     ctrl_in: wp.array2d(dtype=float),
@@ -643,9 +679,11 @@ def _actuator_force(
       act = act_in[worldid, act_last]
       act_dot = (ctrl - act) / wp.max(dynprm[0], MJ_MINVAL)
     elif dyntype == DynType.MUSCLE:
-      dynprm = actuator_dynprm[worldid, uid]
+      dynprm = actuator_dynprm[worldid % actuator_dynprm.shape[0], uid]
       act = act_in[worldid, act_last]
       act_dot = util_misc.muscle_dynamics(ctrl, act, dynprm)
+    elif dyntype == DynType.USER:
+      act_dot = 0.0  # set by act_dyn_callback
     else:  # DynType.NONE
       act_dot = 0.0
 
@@ -681,20 +719,25 @@ def _actuator_force(
   elif gaintype == GainType.AFFINE:
     gain = gainprm[0] + gainprm[1] * length + gainprm[2] * velocity
   elif gaintype == GainType.MUSCLE:
-    acc0 = actuator_acc0[uid]
+    acc0 = actuator_acc0[worldid % actuator_acc0.shape[0], uid]
-    lengthrange = actuator_lengthrange[uid]
+    lengthrange = actuator_lengthrange[
+        worldid % actuator_lengthrange.shape[0], uid
+    ]
     gain = util_misc.muscle_gain(length, velocity, lengthrange, acc0, gainprm)
+  # GainType.USER: gain stays 0, modified by act_gain_callback
 
   # bias
   biastype = actuator_biastype[uid]
   biasprm = actuator_biasprm[worldid % actuator_biasprm.shape[0], uid]
 
-  bias = 0.0  # BiasType.NONE
+  bias = 0.0  # BiasType.NONE or BiasType.USER (modified by act_bias_callback)
   if biastype == BiasType.AFFINE:
     bias = biasprm[0] + biasprm[1] * length + biasprm[2] * velocity
   elif biastype == BiasType.MUSCLE:
-    acc0 = actuator_acc0[uid]
+    acc0 = actuator_acc0[worldid % actuator_acc0.shape[0], uid]
-    lengthrange = actuator_lengthrange[uid]
+    lengthrange = actuator_lengthrange[
+        worldid % actuator_lengthrange.shape[0], uid
+    ]
     bias = util_misc.muscle_bias(length, lengthrange, acc0, biasprm)
 
   force = gain * ctrl_act + bias
@@ -752,32 +795,54 @@ def _tendon_actuator_force_clamp(
 
 @wp.kernel
 def _qfrc_actuator(
+    # Data in:
+    moment_rownnz_in: wp.array2d(dtype=int),
+    moment_rowadr_in: wp.array2d(dtype=int),
+    moment_colind_in: wp.array2d(dtype=int),
+    actuator_moment_in: wp.array2d(dtype=float),
+    actuator_force_in: wp.array2d(dtype=float),
+    # Data out:
+    qfrc_actuator_out: wp.array2d(dtype=float),
+):
+  worldid, actid = wp.tid()
+
+  rownnz = moment_rownnz_in[worldid, actid]
+  rowadr = moment_rowadr_in[worldid, actid]
+
+  for i in range(rownnz):
+    sparseid = rowadr + i
+    colind = moment_colind_in[worldid, sparseid]
+    qfrc = (
+        actuator_moment_in[worldid, sparseid]
+        * actuator_force_in[worldid, actid]
+    )
+    wp.atomic_add(qfrc_actuator_out[worldid], colind, qfrc)
+
+
+@wp.kernel
+def _qfrc_actuator_gravcomp_limits(
     # Model:
-  nu: int,
     ngravcomp: int,
     jnt_actfrclimited: wp.array(dtype=bool),
     jnt_actgravcomp: wp.array(dtype=int),
     jnt_actfrcrange: wp.array2d(dtype=wp.vec2),
     dof_jntid: wp.array(dtype=int),
     # Data in:
-  actuator_moment_in: wp.array3d(dtype=float),
     qfrc_gravcomp_in: wp.array2d(dtype=float),
-  actuator_force_in: wp.array2d(dtype=float),
+    qfrc_actuator_in: wp.array2d(dtype=float),
     # Data out:
     qfrc_actuator_out: wp.array2d(dtype=float),
 ):
   worldid, dofid = wp.tid()
-
-  qfrc = float(0.0)
-  for uid in range(nu):
-    qfrc += actuator_moment_in[worldid, uid, dofid] * actuator_force_in[worldid, uid]
-
   jntid = dof_jntid[dofid]
 
+  qfrc = qfrc_actuator_in[worldid, dofid]
+
   # actuator-level gravity compensation, skip if added as passive force
   if ngravcomp and jnt_actgravcomp[jntid]:
     qfrc += qfrc_gravcomp_in[worldid, dofid]
 
+  # limits
   if jnt_actfrclimited[jntid]:
     frcrange = jnt_actfrcrange[worldid % jnt_actfrcrange.shape[0], jntid]
     qfrc = wp.clamp(qfrc, frcrange[0], frcrange[1])
@@ -825,6 +890,13 @@ def fwd_actuation(m: Model, d: Data):
     outputs=[d.act_dot, d.actuator_force],
   )
 
+  if m.callback.act_dyn:
+    m.callback.act_dyn(m, d)
+  if m.callback.act_gain:
+    m.callback.act_gain(m, d)
+  if m.callback.act_bias:
+    m.callback.act_bias(m, d)
+
   if m.ntendon:
     # total actuator force at tendon
     ten_actfrc = wp.zeros((d.nworld, m.ntendon), dtype=float)
@@ -842,19 +914,31 @@ def fwd_actuation(m: Model, d: Data):
       outputs=[d.actuator_force],
     )
 
+  # TODO(team): optimize performance
+  d.qfrc_actuator.zero_()
   wp.launch(
       _qfrc_actuator,
+      dim=(d.nworld, m.nu),
+      inputs=[
+          d.moment_rownnz,
+          d.moment_rowadr,
+          d.moment_colind,
+          d.actuator_moment,
+          d.actuator_force,
+      ],
+      outputs=[d.qfrc_actuator],
+  )
+  wp.launch(
+      _qfrc_actuator_gravcomp_limits,
       dim=(d.nworld, m.nv),
       inputs=[
-      m.nu,
           m.ngravcomp,
           m.jnt_actfrclimited,
           m.jnt_actgravcomp,
           m.jnt_actfrcrange,
           m.dof_jntid,
-      d.actuator_moment,
           d.qfrc_gravcomp,
-      d.actuator_force,
+          d.qfrc_actuator,
       ],
       outputs=[d.qfrc_actuator],
   )
@@ -923,6 +1007,9 @@ def forward(m: Model, d: Data):
     if m.sensor_e_kinetic == 0:  # not computed by sensor
       sensor.energy_vel(m, d)
 
+  if not (m.opt.disableflags & DisableBit.ACTUATION):
+    if m.callback.control:
+      m.callback.control(m, d)
   fwd_actuation(m, d)
   fwd_acceleration(m, d, factorize=True)
 
@@ -971,6 +1058,10 @@ def step1(m: Model, d: Data):
     if m.sensor_e_kinetic == 0:  # not computed by sensor
       sensor.energy_vel(m, d)
 
+  if not (m.opt.disableflags & DisableBit.ACTUATION):
+    if m.callback.control:
+      m.callback.control(m, d)
+
 
 @event_scope
 def step2(m: Model, d: Data):

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

@@ -14,32 +14,24 @@
 # ==============================================================================
 
 import dataclasses
-import importlib.metadata
-import warnings
 from typing import Any, Optional, Sequence
+import warnings
 
 import mujoco
-import numpy as np
-import packaging.version
-import warp as wp
-
 from mujoco.mjx.third_party.mujoco_warp._src import bvh
+from mujoco.mjx.third_party.mujoco_warp._src import math as mjmath
 from mujoco.mjx.third_party.mujoco_warp._src import render_util
 from mujoco.mjx.third_party.mujoco_warp._src import smooth
 from mujoco.mjx.third_party.mujoco_warp._src import types
 from mujoco.mjx.third_party.mujoco_warp._src import warp_util
-
+from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType
-
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
-def _is_mujoco_dev() -> bool:
+from mujoco.mjx.third_party.mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN
-  """Checks if mujoco version is > 3.4.0."""
+from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
-  version_str = getattr(mujoco, "__version__", None)
+from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
-  if not version_str:
+from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
-    version_str = importlib.metadata.version("mujoco")
+import numpy as np
-  version_str = version_str.split("-")[0].split(".dev")[0]
+import warp as wp
-  return packaging.version.parse(version_str) > packaging.version.parse("3.4.0")
-
-
-BLEEDING_EDGE_MUJOCO = _is_mujoco_dev()
 
 
 def _create_array(data: Any, spec: wp.array, sizes: dict[str, int]) -> wp.array | None:
@@ -218,6 +210,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
 
   m.opt = opt
   m.stat = stat
+  m.callback = types.Callback()
 
   m.nv_pad = _get_padded_sizes(
     mjm.nv, 0, is_sparse(mjm), types.TILE_SIZE_JTDAJ_SPARSE if is_sparse(mjm) else types.TILE_SIZE_JTDAJ_DENSE
@@ -585,17 +578,9 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
       m.qM_mulm_madr.append(madr)
     m.qM_mulm_rowadr.append(len(m.qM_mulm_col))
 
-  # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
-  if BLEEDING_EDGE_MUJOCO:
   m.flexedge_J_rownnz = mjm.flexedge_J_rownnz
   m.flexedge_J_rowadr = mjm.flexedge_J_rowadr
   m.flexedge_J_colind = mjm.flexedge_J_colind.reshape(-1)
-  else:
-    mjd = mujoco.MjData(mjm)
-    mujoco.mj_forward(mjm, mjd)
-    m.flexedge_J_rownnz = mjd.flexedge_J_rownnz
-    m.flexedge_J_rowadr = mjd.flexedge_J_rowadr
-    m.flexedge_J_colind = mjd.flexedge_J_colind.reshape(-1)
 
   # place m on device
   sizes = dict({"*": 1}, **{f.name: getattr(m, f.name) for f in dataclasses.fields(types.Model) if f.type is int})
@@ -646,6 +631,16 @@ def _default_njmax(mjm: mujoco.MjModel, mjd: Optional[mujoco.MjData] = None) ->
   return int(valid_sizes[np.searchsorted(valid_sizes, njmax)])
 
 
+def _resolve_batch_size(
+    na: int | None, n: int | None, nworld: int, default: int
+) -> int:
+  if na is not None:
+    return na
+  if n is not None:
+    return n * nworld
+  return default
+
+
 def make_data(
   mjm: mujoco.MjModel,
   nworld: int = 1,
@@ -675,37 +670,36 @@ def make_data(
   if nconmax is None:
     nconmax = _default_nconmax(mjm)
 
-  if nconmax < 0:
-    raise ValueError("nconmax must be >= 0")
-
-  if nccdmax is None:
-    nccdmax = nconmax
-  elif nccdmax < 0:
-    raise ValueError("nccdmax must be >= 0")
-  elif nccdmax > nconmax:
-    raise ValueError(f"nccdmax ({nccdmax}) must be <= nconmax ({nconmax})")
-
   if njmax is None:
     njmax = _default_njmax(mjm)
 
+  if nconmax < 0:
+    raise ValueError("nconmax must be >= 0")
+
   if njmax < 0:
     raise ValueError("njmax must be >= 0")
 
   if nworld < 1:
     raise ValueError(f"nworld must be >= 1")
 
-  if naconmax is None:
+  naconmax = _resolve_batch_size(naconmax, nconmax, nworld, 0)
-    naconmax = nworld * nconmax
+  if naconmax < 0:
-  elif naconmax < 0:
     raise ValueError("naconmax must be >= 0")
 
-  if naccdmax is None:
+  naccdmax = _resolve_batch_size(naccdmax, nccdmax, nworld, naconmax)
-    naccdmax = nworld * nccdmax
+  if naccdmax < 0:
-  elif naccdmax < 0:
     raise ValueError("naccdmax must be >= 0")
   elif naccdmax > naconmax:
     raise ValueError(f"naccdmax ({naccdmax}) must be <= naconmax ({naconmax})")
 
+  if nccdmax is None:
+    nccdmax = nconmax
+  else:
+    if nccdmax < 0:
+      raise ValueError("nccdmax must be >= 0")
+    elif nccdmax > nconmax:
+      raise ValueError(f"nccdmax ({nccdmax}) must be <= nconmax ({nconmax})")
+
   sizes = dict({"*": 1}, **{f.name: getattr(mjm, f.name, None) for f in dataclasses.fields(types.Model) if f.type is int})
   sizes["nmaxcondim"] = np.concatenate(([0], mjm.geom_condim, mjm.pair_dim)).max()
   sizes["nmaxpyramid"] = np.maximum(1, 2 * (sizes["nmaxcondim"] - 1))
@@ -718,6 +712,17 @@ def make_data(
   contact = types.Contact(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Contact)})
   efc = types.Constraint(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Constraint)})
 
+  if SPARSE_CONSTRAINT_JACOBIAN:
+    efc.J_rownnz = wp.zeros((nworld, njmax), dtype=int)
+    efc.J_rowadr = wp.zeros((nworld, njmax), dtype=int)
+    efc.J_colind = wp.zeros((nworld, 1, njmax * mjm.nv), dtype=int)
+    efc.J = wp.zeros((nworld, 1, njmax * mjm.nv), dtype=float)
+  else:
+    efc.J_rownnz = wp.zeros((nworld, 0), dtype=int)
+    efc.J_rowadr = wp.zeros((nworld, 0), dtype=int)
+    efc.J_colind = wp.zeros((nworld, 0, 0), dtype=int)
+    efc.J = wp.zeros((nworld, sizes["njmax_pad"], sizes["nv_pad"]), dtype=float)
+
   # world body and static geom (attached to the world) poses are precomputed
   # this speeds up scenes with many static geoms (e.g. terrains)
   # TODO(team): remove this when we introduce dof islands + sleeping
@@ -729,31 +734,65 @@ def make_data(
   mocap_id = mjm.body_mocapid[mocap_body]
 
   d_kwargs = {
-    "qpos": wp.array(np.tile(mjm.qpos0, nworld), shape=(nworld, mjm.nq), dtype=float),
+      "qpos": wp.array(
+          np.tile(mjm.qpos0, nworld), shape=(nworld, mjm.nq), dtype=float
+      ),
       "contact": contact,
       "efc": efc,
       "nworld": nworld,
       "naconmax": naconmax,
       "naccdmax": naccdmax,
       "njmax": njmax,
+      "njmax_pad": sizes["njmax_pad"],
       "qM": None,
       "qLD": None,
       # world body
-    "xquat": wp.array(np.tile(mjd.xquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat),
+      "xquat": wp.array(
-    "xmat": wp.array(np.tile(mjd.xmat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.mat33),
+          np.tile(mjd.xquat, (nworld, 1)),
-    "ximat": wp.array(np.tile(mjd.ximat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.mat33),
+          shape=(nworld, mjm.nbody),
+          dtype=wp.quat,
+      ),
+      "xmat": wp.array(
+          np.tile(mjd.xmat, (nworld, 1)),
+          shape=(nworld, mjm.nbody),
+          dtype=wp.mat33,
+      ),
+      "ximat": wp.array(
+          np.tile(mjd.ximat, (nworld, 1)),
+          shape=(nworld, mjm.nbody),
+          dtype=wp.mat33,
+      ),
       # static geoms
-    "geom_xpos": wp.array(np.tile(mjd.geom_xpos, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.vec3),
+      "geom_xpos": wp.array(
-    "geom_xmat": wp.array(np.tile(mjd.geom_xmat, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.mat33),
+          np.tile(mjd.geom_xpos, (nworld, 1)),
+          shape=(nworld, mjm.ngeom),
+          dtype=wp.vec3,
+      ),
+      "geom_xmat": wp.array(
+          np.tile(mjd.geom_xmat, (nworld, 1)),
+          shape=(nworld, mjm.ngeom),
+          dtype=wp.mat33,
+      ),
       # mocap
-    "mocap_pos": wp.array(np.tile(mjm.body_pos[mocap_body[mocap_id]], (nworld, 1)), shape=(nworld, mjm.nmocap), dtype=wp.vec3),
+      "mocap_pos": wp.array(
+          np.tile(mjm.body_pos[mocap_body[mocap_id]], (nworld, 1)),
+          shape=(nworld, mjm.nmocap),
+          dtype=wp.vec3,
+      ),
       "mocap_quat": wp.array(
-      np.tile(mjm.body_quat[mocap_body[mocap_id]], (nworld, 1)), shape=(nworld, mjm.nmocap), dtype=wp.quat
+          np.tile(mjm.body_quat[mocap_body[mocap_id]], (nworld, 1)),
+          shape=(nworld, mjm.nmocap),
+          dtype=wp.quat,
       ),
       # equality constraints
-    "eq_active": wp.array(np.tile(mjm.eq_active0.astype(bool), (nworld, 1)), shape=(nworld, mjm.neq), dtype=bool),
+      "eq_active": wp.array(
-    # flexedge
+          np.tile(mjm.eq_active0.astype(bool), (nworld, 1)),
-    "flexedge_J": None,
+          shape=(nworld, mjm.neq),
+          dtype=bool,
+      ),
+      # island arrays
+      "nisland": None,
+      "tree_island": None,
   }
   for f in dataclasses.fields(types.Data):
     if f.name in d_kwargs:
@@ -769,7 +808,9 @@ def make_data(
     d.qM = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float)
     d.qLD = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float)
 
-  d.flexedge_J = wp.zeros((nworld, 1, mjd.flexedge_J.size), dtype=float)
+  # island discovery arrays
+  d.nisland = wp.zeros((nworld,), dtype=int)
+  d.tree_island = wp.zeros((nworld, mjm.ntree), dtype=int)
 
   return d
 
@@ -808,39 +849,43 @@ def put_data(
   if nconmax is None:
     nconmax = _default_nconmax(mjm, mjd)
 
-  if nconmax < 0:
-    raise ValueError("nconmax must be >= 0")
-
-  if nccdmax is None:
-    nccdmax = nconmax
-  elif nccdmax < 0:
-    raise ValueError("nccdmax must be >= 0")
-  elif nccdmax > nconmax:
-    raise ValueError(f"nccdmax ({nccdmax}) must be <= nconmax ({nconmax})")
-
   if njmax is None:
     njmax = _default_njmax(mjm, mjd)
 
+  if nconmax < 0:
+    raise ValueError("nconmax must be >= 0")
+
   if njmax < 0:
     raise ValueError("njmax must be >= 0")
 
   if nworld < 1:
     raise ValueError(f"nworld must be >= 1")
 
-  if naconmax is None:
+  naconmax_is_input = naconmax is not None
-    if mjd.ncon > nconmax:
+  naconmax = _resolve_batch_size(naconmax, nconmax, nworld, 0)
+  if naconmax < 0:
+    raise ValueError("naconmax must be >= 0")
+
+  if not naconmax_is_input and mjd.ncon > nconmax:
     raise ValueError(f"nconmax overflow (nconmax must be >= {mjd.ncon})")
-    naconmax = nworld * nconmax
   elif naconmax < mjd.ncon * nworld:
     raise ValueError(f"naconmax overflow (naconmax must be >= {mjd.ncon * nworld})")
 
-  if naccdmax is None:
+  naccdmax = _resolve_batch_size(naccdmax, nccdmax, nworld, naconmax)
-    naccdmax = nworld * nccdmax
+
-  elif naccdmax < 0:
+  if naccdmax < 0:
     raise ValueError("naccdmax must be >= 0")
   elif naccdmax > naconmax:
     raise ValueError(f"naccdmax ({naccdmax}) must be <= naconmax ({naconmax})")
 
+  if nccdmax is None:
+    nccdmax = nconmax
+  else:
+    if nccdmax < 0:
+      raise ValueError("nccdmax must be >= 0")
+    elif nccdmax > nconmax:
+      raise ValueError(f"nccdmax ({nccdmax}) must be <= nconmax ({nconmax})")
+
   if mjd.nefc > njmax:
     raise ValueError(f"njmax overflow (njmax must be >= {mjd.nefc})")
 
@@ -887,7 +932,7 @@ def put_data(
   contact.geomcollisionid = wp.empty((naconmax,), dtype=int)  # TODO(team): set values
 
   # create efc
-  efc_kwargs = {"J": None}
+  efc_kwargs = {"J_rownnz": None, "J_rowadr": None, "J_colind": None, "J": None}
 
   for f in dataclasses.fields(types.Constraint):
     if f.name in efc_kwargs:
@@ -900,14 +945,63 @@ def put_data(
 
   efc = types.Constraint(**efc_kwargs)
 
+  if SPARSE_CONSTRAINT_JACOBIAN:
+    # TODO(team): process efc_J sparsity structure for nv row shift
+    efc.J_rownnz = wp.array(
+        np.full((nworld, njmax), mjm.nv, dtype=int), dtype=int
+    )
+    efc.J_rowadr = wp.array(
+        np.tile(
+            np.arange(0, njmax * mjm.nv, mjm.nv)
+            if mjm.nv
+            else np.zeros(njmax, dtype=int),
+            (nworld, 1),
+        ),
+        dtype=int,
+    )
+    efc.J_colind = wp.array(
+        np.tile(np.arange(mjm.nv), (nworld, njmax)).reshape((nworld, 1, -1)),
+        dtype=int,
+    )
+
+    mj_efc_J = np.zeros((mjd.nefc, mjm.nv))
+    if mjd.nefc:
       if mujoco.mj_isSparse(mjm):
-    efc_j = np.zeros((mjd.nefc, mjm.nv))
+        mujoco.mju_sparse2dense(
-    mujoco.mju_sparse2dense(efc_j, mjd.efc_J, mjd.efc_J_rownnz, mjd.efc_J_rowadr, mjd.efc_J_colind)
+            mj_efc_J,
+            mjd.efc_J,
+            mjd.efc_J_rownnz,
+            mjd.efc_J_rowadr,
+            mjd.efc_J_colind,
+        )
       else:
-    efc_j = mjd.efc_J.reshape((mjd.nefc, mjm.nv))
+        mj_efc_J = mjd.efc_J.reshape((mjd.nefc, mjm.nv))
-  efc.J = np.zeros((nworld, sizes["njmax_pad"], sizes["nv_pad"]), dtype=f.type.dtype)
+    efc_J = np.zeros((njmax, mjm.nv), dtype=float)
-  efc.J[:, : mjd.nefc, : mjm.nv] = np.tile(efc_j, (nworld, 1, 1))
+    efc_J[: mjd.nefc, : mjm.nv] = mj_efc_J
-  efc.J = wp.array(efc.J, dtype=float)
+    efc.J = wp.array(
+        np.tile(efc_J.reshape(-1), (nworld, 1, 1)).reshape((nworld, 1, -1)),
+        dtype=float,
+    )
+  else:
+    efc.J_rownnz = wp.zeros((nworld, 0), dtype=int)
+    efc.J_rowadr = wp.zeros((nworld, 0), dtype=int)
+    efc.J_colind = wp.zeros((nworld, 0, 0), dtype=int)
+
+    mj_efc_J = np.zeros((mjd.nefc, mjm.nv))
+    if mjd.nefc:
+      if mujoco.mj_isSparse(mjm):
+        mujoco.mju_sparse2dense(
+            mj_efc_J,
+            mjd.efc_J,
+            mjd.efc_J_rownnz,
+            mjd.efc_J_rowadr,
+            mjd.efc_J_colind,
+        )
+      else:
+        mj_efc_J = mjd.efc_J.reshape((mjd.nefc, mjm.nv))
+    efc_J = np.zeros((nworld, sizes["njmax_pad"], sizes["nv_pad"]), dtype=float)
+    efc_J[:, : mjd.nefc, : mjm.nv] = np.tile(mj_efc_J, (nworld, 1, 1))
+    efc.J = wp.array(efc_J, dtype=float)
 
   # create data
   d_kwargs = {
@@ -917,14 +1011,16 @@ def put_data(
       "naconmax": naconmax,
       "naccdmax": naccdmax,
       "njmax": njmax,
+      "njmax_pad": sizes["njmax_pad"],
       # fields set after initialization:
       "solver_niter": None,
       "qM": None,
       "qLD": None,
       "ten_J": None,
-    "actuator_moment": None,
-    "flexedge_J": None,
       "nacon": None,
+      # island arrays
+      "nisland": None,
+      "tree_island": None,
   }
   for f in dataclasses.fields(types.Data):
     if f.name in d_kwargs:
@@ -950,19 +1046,32 @@ def put_data(
     d.qM = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), qM_padded), dtype=float)
     d.qLD = wp.array(np.full((nworld, mjm.nv, mjm.nv), qLD), dtype=float)
 
-  d.flexedge_J = wp.array(np.tile(mjd.flexedge_J.reshape(-1), (nworld, 1)).reshape((nworld, 1, -1)), dtype=float)
+  # island arrays
+  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)
 
-  if mjm.ntendon:
   ten_J = np.zeros((mjm.ntendon, mjm.nv))
-    mujoco.mju_sparse2dense(ten_J, mjd.ten_J.reshape(-1), mjm.ten_J_rownnz, mjm.ten_J_rowadr, mjm.ten_J_colind.reshape(-1))
+  if mujoco.mj_isSparse(mjm) or check_version("mujoco>=3.5.1.dev872479828"):
-    d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float)
+    if mjm.ntendon:
+      if check_version("mujoco>=3.5.1.dev875093374"):
+        mujoco.mju_sparse2dense(
+            ten_J,
+            mjd.ten_J.reshape(-1),
+            mjm.ten_J_rownnz,
+            mjm.ten_J_rowadr,
+            mjm.ten_J_colind.reshape(-1),
+        )
       else:
-    d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), 0.0), dtype=float)
+        mujoco.mju_sparse2dense(
-
+            ten_J,
-  # TODO(taylorhowell): sparse actuator_moment
+            mjd.ten_J.reshape(-1),
-  actuator_moment = np.zeros((mjm.nu, mjm.nv))
+            mjd.ten_J_rownnz,
-  mujoco.mju_sparse2dense(actuator_moment, mjd.actuator_moment, mjd.moment_rownnz, mjd.moment_rowadr, mjd.moment_colind)
+            mjd.ten_J_rowadr,
-  d.actuator_moment = wp.array(np.full((nworld, mjm.nu, mjm.nv), actuator_moment), dtype=float)
+            mjd.ten_J_colind.reshape(-1),
+        )
+  else:
+    ten_J = mjd.ten_J.reshape((mjm.ntendon, mjm.nv))
+  d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float)
 
   d.nacon = wp.array([mjd.ncon * nworld], dtype=int)
 
@@ -1066,28 +1175,15 @@ def get_data_into(
   result.cinert[:] = d.cinert.numpy()[world_id]
   result.flexvert_xpos[:] = d.flexvert_xpos.numpy()[world_id]
   if mjm.nflexedge > 0:
-    # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
-    if not BLEEDING_EDGE_MUJOCO:
-      m = put_model(mjm)
-      result.flexedge_J_rownnz[:] = m.flexedge_J_rownnz.numpy()
-      result.flexedge_J_rowadr[:] = m.flexedge_J_rowadr.numpy()
-      result.flexedge_J_colind[:, :] = m.flexedge_J_colind.numpy().reshape((mjm.nflexedge, mjm.nv))
-      mujoco.mju_sparse2dense(
-        result.flexedge_J,
-        d.flexedge_J.numpy()[world_id].reshape(-1),
-        m.flexedge_J_rownnz.numpy(),
-        m.flexedge_J_rowadr.numpy(),
-        m.flexedge_J_colind.numpy(),
-      )
-    else:
     result.flexedge_J[:] = d.flexedge_J.numpy()[world_id].reshape(-1)
   result.flexedge_length[:] = d.flexedge_length.numpy()[world_id]
   result.flexedge_velocity[:] = d.flexedge_velocity.numpy()[world_id]
   result.actuator_length[:] = d.actuator_length.numpy()[world_id]
-  actuator_moment = d.actuator_moment.numpy()[world_id]
+  result.moment_rownnz[:] = d.moment_rownnz.numpy()[world_id]
-  mujoco.mju_dense2sparse(
+  result.moment_rowadr[:] = d.moment_rowadr.numpy()[world_id]
-    result.actuator_moment, actuator_moment, result.moment_rownnz, result.moment_rowadr, result.moment_colind
+  if mjm.nu:
-  )
+    result.moment_colind[:] = d.moment_colind.numpy()[world_id]
+    result.actuator_moment[:] = d.actuator_moment.numpy()[world_id]
   result.crb[:] = d.crb.numpy()[world_id]
   result.qLDiagInv[:] = d.qLDiagInv.numpy()[world_id]
   result.ten_velocity[:] = d.ten_velocity.numpy()[world_id]
@@ -1137,11 +1233,29 @@ def get_data_into(
     mujoco.mj_factorM(mjm, result)
 
   if nefc > 0:
-    if mujoco.mj_isSparse(mjm):
+    if SPARSE_CONSTRAINT_JACOBIAN:
-      efc_J = d.efc.J.numpy()[world_id, efc_idx, : mjm.nv]
+      efc_J = np.zeros((nefc, mjm.nv))
-      mujoco.mju_dense2sparse(result.efc_J, efc_J, result.efc_J_rownnz, result.efc_J_rowadr, result.efc_J_colind)
+      mujoco.mju_sparse2dense(
+          efc_J,
+          d.efc.J.numpy()[world_id, 0],
+          d.efc.J_rownnz.numpy()[world_id, :nefc],
+          d.efc.J_rowadr.numpy()[world_id, :nefc],
+          d.efc.J_colind.numpy()[world_id, 0],
+      )
     else:
-      result.efc_J[: nefc * mjm.nv] = d.efc.J.numpy()[world_id, :nefc, : mjm.nv].flatten()
+      efc_J = d.efc.J.numpy()[world_id, :nefc, : mjm.nv]
+
+    # write to mujoco result (format depends on mj_isSparse)
+    if mujoco.mj_isSparse(mjm):
+      mujoco.mju_dense2sparse(
+          result.efc_J,
+          efc_J[efc_idx],
+          result.efc_J_rownnz,
+          result.efc_J_rowadr,
+          result.efc_J_colind,
+      )
+    else:
+      result.efc_J[: nefc * mjm.nv] = efc_J[efc_idx].flatten()
 
   # efc
   result.efc_type[:] = d.efc.type.numpy()[world_id, efc_idx]
@@ -1162,15 +1276,22 @@ def get_data_into(
 
   # tendon
   result.ten_length[:] = d.ten_length.numpy()[world_id]
-  # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
+  if check_version("mujoco>=3.5.1.dev869712136"):
-  if BLEEDING_EDGE_MUJOCO:
     ten_J = d.ten_J.numpy()[world_id]
+    if check_version("mujoco>=3.5.1.dev875093374"):
+      ten_J_rownnz = mjm.ten_J_rownnz
+      ten_J_rowadr = mjm.ten_J_rowadr
+      ten_J_colind = mjm.ten_J_colind.reshape(-1)
+    else:
+      ten_J_rownnz = result.ten_J_rownnz
+      ten_J_rowadr = result.ten_J_rowadr
+      ten_J_colind = result.ten_J_colind.reshape(-1)
     mujoco.mju_dense2sparse(
         result.ten_J,
         ten_J,
-        mjm.ten_J_rownnz,
+        ten_J_rownnz,
-        mjm.ten_J_rowadr,
+        ten_J_rowadr,
-        mjm.ten_J_colind,
+        ten_J_colind,
     )
   else:
     result.ten_J[:] = d.ten_J.numpy()[world_id]
@@ -1182,6 +1303,12 @@ def get_data_into(
   # sensors
   result.sensordata[:] = d.sensordata.numpy()[world_id]
 
+  # islands
+  nisland = d.nisland.numpy()[world_id]
+  result.nisland = nisland
+  if nisland:
+    result.tree_island[:] = d.tree_island.numpy()[world_id]
+
 
 def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
   """Clear data, set defaults; optionally by world.
@@ -1301,8 +1428,12 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     mocapid = body_mocapid[bodyid]
 
     if mocapid >= 0:
-      mocap_pos_out[worldid, mocapid] = body_pos[worldid, bodyid]
+      mocap_pos_out[worldid, mocapid] = body_pos[
-      mocap_quat_out[worldid, mocapid] = body_quat[worldid, bodyid]
+          worldid % body_pos.shape[0], bodyid
+      ]
+      mocap_quat_out[worldid, mocapid] = body_quat[
+          worldid % body_quat.shape[0], bodyid
+      ]
 
   @wp.kernel(module="unique", enable_backward=False)
   def reset_contact(
@@ -1761,13 +1892,22 @@ def _compute_light_pos0(
 @wp.kernel
 def _copy_actuator_moment(
     actid_target: int,
-  actuator_moment_in: wp.array3d(dtype=float),
+    moment_rownnz_in: wp.array2d(dtype=int),
+    moment_rowadr_in: wp.array2d(dtype=int),
+    moment_colind_in: wp.array2d(dtype=int),
+    actuator_moment_in: wp.array2d(dtype=float),
     act_moment_vec_out: wp.array2d(dtype=float),
 ):
   worldid = wp.tid()
-  nv = actuator_moment_in.shape[2]
+  nv = act_moment_vec_out.shape[1]
   for i in range(nv):
-    act_moment_vec_out[worldid, i] = actuator_moment_in[worldid, actid_target, i]
+    act_moment_vec_out[worldid, i] = 0.0
+  rownnz = moment_rownnz_in[worldid, actid_target]
+  rowadr = moment_rowadr_in[worldid, actid_target]
+  for i in range(rownnz):
+    sparseid = rowadr + i
+    col = moment_colind_in[worldid, sparseid]
+    act_moment_vec_out[worldid, col] = actuator_moment_in[worldid, sparseid]
 
 
 @wp.kernel
@@ -1775,13 +1915,116 @@ def _compute_actuator_acc0(
     actid_target: int,
     nv: int,
     result_vec_in: wp.array2d(dtype=float),
-  actuator_acc0_out: wp.array(dtype=float),
+    actuator_acc0_out: wp.array2d(dtype=float),
 ):
   worldid = wp.tid()
   norm_sq = float(0.0)
   for i in range(nv):
     norm_sq += result_vec_in[worldid, i] * result_vec_in[worldid, i]
-  actuator_acc0_out[actid_target] = wp.sqrt(norm_sq)
+  actuator_acc0_out[worldid, actid_target] = wp.sqrt(norm_sq)
+
+
+@wp.kernel
+def _compute_dof_M0(
+    dof_bodyid: wp.array(dtype=int),
+    dof_armature: wp.array2d(dtype=float),
+    cdof_in: wp.array2d(dtype=wp.spatial_vector),
+    crb_in: wp.array2d(dtype=vec10),
+    dof_M0_out: wp.array2d(dtype=float),
+):
+  worldid, dofid = wp.tid()
+  bodyid = dof_bodyid[dofid]
+  armature = dof_armature[worldid % dof_armature.shape[0], dofid]
+  buf = mjmath.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid])
+  dof_M0_out[worldid, dofid] = armature + wp.dot(cdof_in[worldid, dofid], buf)
+
+
+@wp.kernel
+def _resolve_dampratio(
+    actuator_biastype: wp.array(dtype=int),
+    actuator_gainprm: wp.array2d(dtype=types.vec10f),
+    moment_rownnz_in: wp.array2d(dtype=int),
+    moment_rowadr_in: wp.array2d(dtype=int),
+    moment_colind_in: wp.array2d(dtype=int),
+    actuator_moment_in: wp.array2d(dtype=float),
+    dof_M0_in: wp.array2d(dtype=float),
+    nv: int,
+    actuator_biasprm: wp.array2d(dtype=types.vec10f),
+):
+  worldid, actid = wp.tid()
+  biastype = actuator_biastype[actid]
+
+  # only affine bias (position actuators)
+  if biastype != BiasType.AFFINE:
+    return
+
+  gainprm_id = worldid % actuator_gainprm.shape[0]
+  biasprm_id = worldid % actuator_biasprm.shape[0]
+  kp = actuator_gainprm[gainprm_id, actid][0]
+
+  biasprm = actuator_biasprm[biasprm_id, actid]
+  # dampratio condition: gainprm[0] == -biasprm[1] and biasprm[2] > 0
+  if wp.abs(kp + biasprm[1]) > MJ_MINVAL:
+    return
+  if biasprm[2] <= 0.0:
+    return
+
+  dampratio = biasprm[2]
+
+  # compute reflected mass: sum(dof_M0[j] / moment[i,j]^2) for active DOFs
+  mass = float(0.0)
+  rownnz = moment_rownnz_in[worldid, actid]
+  rowadr = moment_rowadr_in[worldid, actid]
+  for k in range(rownnz):
+    sparseid = rowadr + k
+    j = moment_colind_in[worldid, sparseid]
+    moment = actuator_moment_in[worldid, sparseid]
+    if wp.abs(moment) > MJ_MINVAL:
+      mass += dof_M0_in[worldid, j] / (moment * moment)
+
+  damping = dampratio * 2.0 * wp.sqrt(kp * mass)
+
+  # write -damping to biasprm[2]
+  new_biasprm = biasprm
+  new_biasprm[2] = -damping
+  actuator_biasprm[biasprm_id, actid] = new_biasprm
+
+
+@wp.kernel
+def _set_length_range(
+    actuator_trntype: wp.array(dtype=int),
+    actuator_trnid: wp.array(dtype=wp.vec2i),
+    actuator_gear: wp.array2d(dtype=wp.spatial_vector),
+    jnt_limited: wp.array(dtype=int),
+    jnt_range: wp.array2d(dtype=wp.vec2),
+    tendon_limited: wp.array(dtype=int),
+    tendon_range: wp.array2d(dtype=wp.vec2),
+    ntendon: int,
+    actuator_lengthrange_out: wp.array2d(dtype=wp.vec2),
+):
+  worldid, actid = wp.tid()
+  trntype = actuator_trntype[actid]
+  id0 = actuator_trnid[actid][0]
+  gear0 = actuator_gear[worldid % actuator_gear.shape[0], actid][0]
+
+  lr = wp.vec2(0.0, 0.0)
+
+  if trntype == TrnType.JOINT or trntype == TrnType.JOINTINPARENT:
+    if jnt_limited[id0]:
+      rng = jnt_range[worldid % jnt_range.shape[0], id0]
+      if gear0 > 0.0:
+        lr = wp.vec2(rng[0] * gear0, rng[1] * gear0)
+      else:
+        lr = wp.vec2(rng[1] * gear0, rng[0] * gear0)
+  elif trntype == TrnType.TENDON:
+    if ntendon > 0 and tendon_limited[id0]:
+      rng = tendon_range[worldid % tendon_range.shape[0], id0]
+      if gear0 > 0.0:
+        lr = wp.vec2(rng[0] * gear0, rng[1] * gear0)
+      else:
+        lr = wp.vec2(rng[1] * gear0, rng[0] * gear0)
+
+  actuator_lengthrange_out[worldid, actid] = lr
 
 
 # kernel_analyzer: on
@@ -1823,6 +2066,9 @@ def set_const_0(m: types.Model, d: types.Data):
     - cam_pos0, cam_poscom0, cam_mat0: camera references
     - light_pos0, light_poscom0, light_dir0: light references
     - actuator_acc0: acceleration from unit actuator force
+    - actuator_biasprm[2] (dampratio resolution): for position actuators where
+      gainprm[0] == -biasprm[1] and biasprm[2] > 0, converts dampratio to
+      damping via biasprm[2] = -dampratio * 2 * sqrt(kp * reflected_mass)
 
   Args:
     m: The model containing kinematic and dynamic information (device).
@@ -1954,10 +2200,46 @@ def set_const_0(m: types.Model, d: types.Data):
     act_result_vec = wp.zeros((d.nworld, m.nv), dtype=float)
 
     for actid in range(m.nu):
-      wp.launch(_copy_actuator_moment, dim=d.nworld, inputs=[actid, d.actuator_moment], outputs=[act_moment_vec])
+      wp.launch(
+          _copy_actuator_moment,
+          dim=d.nworld,
+          inputs=[
+              actid,
+              d.moment_rownnz,
+              d.moment_rowadr,
+              d.moment_colind,
+              d.actuator_moment,
+          ],
+          outputs=[act_moment_vec],
+      )
       smooth.solve_m(m, d, act_result_vec, act_moment_vec)
       wp.launch(_compute_actuator_acc0, dim=d.nworld, inputs=[actid, m.nv, act_result_vec], outputs=[m.actuator_acc0])
 
+  # resolve dampratio: compute dof_M0, then convert dampratio to damping
+  if m.nu > 0 and m.nv > 0:
+    dof_M0 = wp.zeros((d.nworld, m.nv), dtype=float)
+    wp.launch(
+        _compute_dof_M0,
+        dim=(d.nworld, m.nv),
+        inputs=[m.dof_bodyid, m.dof_armature, d.cdof, d.crb],
+        outputs=[dof_M0],
+    )
+    wp.launch(
+        _resolve_dampratio,
+        dim=(d.nworld, m.nu),
+        inputs=[
+            m.actuator_biastype,
+            m.actuator_gainprm,
+            d.moment_rownnz,
+            d.moment_rowadr,
+            d.moment_colind,
+            d.actuator_moment,
+            dof_M0,
+            m.nv,
+        ],
+        outputs=[m.actuator_biasprm],
+    )
+
   wp.copy(d.qpos, qpos_saved)
 
 
@@ -1973,12 +2255,16 @@ def set_const(m: types.Model, d: types.Data):
     Field                            | Notes
     ---------------------------------|----------------------------------------------
     qpos0, qpos_spring               |
-    body_mass, body_inertia,         | Mass and inertia are usually scaled together
+    body_mass, body_inertia,         | Mass and inertia are usually scaled
+    together
     body_ipos, body_iquat            | since inertia is sum(m * r^2).
-    body_pos, body_quat              | Unsafe for static bodies (invalidates BVH).
+    body_pos, body_quat              | Unsafe for static bodies (invalidates
-    body_gravcomp                    | If changing from 0 to >0 bodies, required.
+    BVH).
+    body_gravcomp                    | If changing from 0 to >0 bodies,
+    required.
     dof_armature                     |
-    eq_data                          | For connect/weld, offsets computed if not set.
+    eq_data                          | For connect/weld, offsets computed if not
+    set.
     hfield_size                      |
     tendon_stiffness, tendon_damping | Only if changing from/to zero.
     actuator_gainprm, actuator_biasprm | For position actuators with dampratio.
@@ -2004,8 +2290,9 @@ def set_const(m: types.Model, d: types.Data):
       - cam_pos0, cam_poscom0, cam_mat0: camera references
       - light_pos0, light_poscom0, light_dir0: light references
       - actuator_acc0: acceleration from unit actuator force
+      - actuator_biasprm[2] (dampratio resolution)
 
-  Skips: dof_M0, actuator_length0 (not in mjwarp).
+  Skips: actuator_length0 (not in mjwarp).
 
   Args:
     m: The model containing kinematic and dynamic information (device).
@@ -2015,6 +2302,39 @@ def set_const(m: types.Model, d: types.Data):
   set_const_0(m, d)
 
 
+def set_length_range(m: types.Model, d: types.Data, index: int = -1):
+  """Compute feasible actuator length ranges from joint/tendon limits.
+
+  For joint and tendon transmissions with limits, copies the range directly
+  from jnt_range or tendon_range scaled by gear. Actuators without limits
+  keep (0, 0). This covers the common robotics use case; simulation-based
+  computation for general transmissions is not yet implemented.
+
+  Args:
+    m: The model containing kinematic and dynamic information (device).
+    d: The data object (unused, kept for API compatibility with MuJoCo C).
+    index: Actuator index to compute for, or -1 for all actuators.
+  """
+  if m.nu == 0:
+    return
+
+  wp.launch(
+      _set_length_range,
+      dim=(d.nworld, m.nu),
+      inputs=[
+          m.actuator_trntype,
+          m.actuator_trnid,
+          m.actuator_gear,
+          m.jnt_limited,
+          m.jnt_range,
+          m.tendon_limited,
+          m.tendon_range,
+          m.ntendon,
+      ],
+      outputs=[m.actuator_lengthrange],
+  )
+
+
 def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any] | Sequence[str]):
   """Overrides model parameters.
 
@@ -2321,26 +2641,25 @@ def create_render_context(
 
   cam_res_arr = wp.array(active_cam_res, dtype=wp.vec2i)
 
-  if render_rgb and isinstance(render_rgb, bool):
+  if render_rgb is None:
+    render_rgb = [
+        mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_RGB
+        for i in active_cam_indices
+    ]
+  elif isinstance(render_rgb, bool):
     render_rgb = [render_rgb] * ncam
-  elif render_rgb is None:
-    # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml
-    if BLEEDING_EDGE_MUJOCO:
-      render_rgb = [mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_RGB for i in active_cam_indices]
-    else:
-      render_rgb = [True] * ncam
 
-  if render_depth and isinstance(render_depth, bool):
+  if render_depth is None:
+    render_depth = [
+        mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_DEPTH
+        for i in active_cam_indices
+    ]
+  if isinstance(render_depth, bool):
     render_depth = [render_depth] * ncam
-  elif render_depth is None:
-    # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml
-    if BLEEDING_EDGE_MUJOCO:
-      render_depth = [mjm.cam_output[i] & mujoco.mjtCamOutBit.mjCAMOUT_DEPTH for i in active_cam_indices]
-    else:
-      render_depth = [True] * ncam
 
   assert len(render_rgb) == ncam and len(render_depth) == ncam, (
-    f"Render RGB and depth must be provided for all active cameras (got {len(render_rgb)}, {len(render_depth)}, expected {ncam})"
+      "render_rgb and render_depth must be a bool or a list of bools with"
+      f" length {ncam}"
   )
 
   rgb_adr = -1 * np.ones(ncam, dtype=int)
@@ -2364,9 +2683,6 @@ def create_render_context(
 
   ray = wp.zeros(int(total), dtype=wp.vec3)
 
-  # TODO: remove after mjwarp depends on mujoco >= 3.4.1 in pyproject.toml
-  cam_projection = np.zeros(mjm.ncam, dtype=int)
-  if BLEEDING_EDGE_MUJOCO:
   cam_projection = mjm.cam_projection
 
   offset = 0

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

@@ -13,12 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 
-import warp as wp
-
 from mujoco.mjx.third_party.mujoco_warp._src import types
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
 from mujoco.mjx.third_party.mujoco_warp._src.types import EqType
 from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
+import warp as wp
 
 
 @wp.kernel
@@ -44,7 +44,7 @@ def _tree_edges(
   # Out:
   tree_tree: wp.array3d(dtype=int),  # kernel_analyzer: off
 ):
-  """Find tree edges. Launch: (nworld, njmax)."""
+  """Find tree edges."""
   worldid, efcid = wp.tid()
 
   # skip if beyond active constraints
@@ -176,3 +176,89 @@ def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d(dtype=int)):
     ],
     outputs=[tree_tree],
   )
+
+
+@wp.kernel
+def _flood_fill(
+    # Model:
+    ntree: int,
+    # In:
+    tree_tree_in: wp.array3d(dtype=int),
+    labels_in: wp.array2d(dtype=int),
+    stack_in: wp.array2d(dtype=int),
+    # Data out:
+    nisland_out: wp.array(dtype=int),
+    tree_island_out: wp.array2d(dtype=int),
+    # Out:
+    stack_out: wp.array2d(dtype=int),
+):
+  """DFS flood fill to discover islands using tree_tree matrix."""
+  worldid = wp.tid()
+  nisland = int(0)
+
+  # iterate over trees
+  for i in range(ntree):
+    # already assigned
+    if labels_in[worldid, i] != -1:
+      continue
+
+    # check if tree has any edges
+    has_edge = int(0)
+    for j in range(ntree):
+      if tree_tree_in[worldid, i, j] != 0:
+        has_edge = 1
+        break
+    if has_edge == 0:
+      continue
+
+    # DFS: push i onto stack
+    nstack = int(0)
+    stack_out[worldid, nstack] = i
+    nstack = nstack + 1
+
+    while nstack > 0:
+      # pop v from stack
+      nstack = nstack - 1
+      v = stack_in[worldid, nstack]
+
+      # already assigned
+      if labels_in[worldid, v] != -1:
+        continue
+
+      # assign to current island
+      tree_island_out[worldid, v] = nisland
+
+      # push neighbors
+      for neighbor in range(ntree):
+        if tree_tree_in[worldid, v, neighbor] != 0:
+          if labels_in[worldid, neighbor] == -1:
+            stack_out[worldid, nstack] = neighbor
+            nstack = nstack + 1
+
+    # island filled
+    nisland = nisland + 1
+
+  nisland_out[worldid] = nisland
+
+
+@event_scope
+def island(m: types.Model, d: types.Data):
+  """Discover constraint islands."""
+  if m.ntree == 0:
+    d.nisland.zero_()
+    return
+
+  # Step 1: Find tree edges
+  tree_tree = wp.zeros((d.nworld, m.ntree, m.ntree), dtype=int)
+  tree_edges(m, d, tree_tree)
+
+  # Step 2: DFS flood fill
+  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],
+  )

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

@@ -266,9 +266,9 @@ def _gravity_force(
 @wp.kernel
 def _fluid_force(
     # Model:
+    opt_wind: wp.array(dtype=wp.vec3),
     opt_density: wp.array(dtype=float),
     opt_viscosity: wp.array(dtype=float),
-  opt_wind: wp.array(dtype=wp.vec3),
     body_rootid: wp.array(dtype=int),
     body_geomnum: wp.array(dtype=int),
     body_geomadr: wp.array(dtype=int),
@@ -498,9 +498,9 @@ def _fluid(m: Model, d: Data):
       _fluid_force,
       dim=(d.nworld, m.nbody),
       inputs=[
+          m.opt.wind,
           m.opt.density,
           m.opt.viscosity,
-      m.opt.wind,
           m.body_rootid,
           m.body_geomnum,
           m.body_geomadr,
@@ -846,3 +846,6 @@ def passive(m: Model, d: Data):
       d.qfrc_passive,
     ],
   )
+
+  if m.callback.passive:
+    m.callback.passive(m, d)

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

@@ -66,19 +66,19 @@ def _ray_eliminate(
 ) -> bool:
   """Eliminate ray."""
   bodyid = geom_bodyid[geomid]
-  matid = geom_matid[geomid]
+  matid = geom_matid[geomid]  # kernel_analyzer: ignore
 
   # body exclusion
   if bodyid == bodyexclude:
     return True
 
   # invisible geom exclusion
-  if matid < 0 and geom_rgba[geomid][3] == 0.0:
+  if matid < 0 and geom_rgba[geomid][3] == 0.0:  # kernel_analyzer: ignore
     return True
 
   # invisible material exclusion
   if matid >= 0:
-    if mat_rgba[matid][3] == 0.0:
+    if mat_rgba[matid][3] == 0.0:  # kernel_analyzer: ignore
       return True
 
   # static exclusion

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

@@ -552,7 +552,13 @@ def render(m: Model, d: Data, rc: RenderContext):
       return
 
     if render_depth[cam_idx]:
-      depth_out[world_idx, depth_adr[cam_idx] + ray_idx_local] = dist
+      # Planar depth: project Euclidean distance onto the camera's optical axis.
+      # In camera-local coordinates, the optical axis is -Z. The Z-component of the
+      # normalized ray direction is negative, so -ray_dir_local_cam[2] gives cos(θ)
+      # between the ray and the optical axis.
+      depth_out[world_idx, depth_adr[cam_idx] + ray_idx_local] = dist * (
+          -ray_dir_local_cam[2]
+      )
 
     if not render_rgb[cam_idx]:
       return

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

@@ -15,16 +15,12 @@
 
 from typing import Any, Tuple
 
-import warp as wp
-
 from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import ray
 from mujoco.mjx.third_party.mujoco_warp._src import smooth
 from mujoco.mjx.third_party.mujoco_warp._src import support
 from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import get_sdf_params
 from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf
-from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
-from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
 from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
@@ -32,9 +28,13 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import DataType
 from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import JointType
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
 from mujoco.mjx.third_party.mujoco_warp._src.types import SensorType
+from mujoco.mjx.third_party.mujoco_warp._src.types import Stage
 from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec6
@@ -43,6 +43,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec8i
 from mujoco.mjx.third_party.mujoco_warp._src.util_misc import inside_geom
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
+import warp as wp
 
 wp.set_module_options({"enable_backward": False})
 
@@ -898,6 +899,9 @@ def sensor_pos(m: Model, d: Data):
     ],
   )
 
+  if m.callback.sensor:
+    m.callback.sensor(m, d, Stage.POS)
+
 
 @wp.func
 def _velocimeter(
@@ -1437,6 +1441,9 @@ def sensor_vel(m: Model, d: Data):
     ],
   )
 
+  if m.callback.sensor:
+    m.callback.sensor(m, d, Stage.VEL)
+
 
 @wp.func
 def _accelerometer(
@@ -2072,6 +2079,43 @@ def _transform_spatial(vec: wp.spatial_vector, dif: wp.vec3) -> wp.vec3:
   return wp.spatial_bottom(vec) - wp.cross(dif, wp.spatial_top(vec))
 
 
+@wp.kernel
+def _preprocess_tactile_contacts(
+    # Model:
+    body_weldid: wp.array(dtype=int),
+    geom_bodyid: wp.array(dtype=int),
+    # Data in:
+    contact_geom_in: wp.array(dtype=wp.vec2i),
+    contact_worldid_in: wp.array(dtype=int),
+    nacon_in: wp.array(dtype=int),
+    # Out:
+    weld_geom_count_out: wp.array2d(dtype=int),
+    weld_geom_list_out: wp.array3d(dtype=int),
+):
+  conid = wp.tid()
+  ncon = nacon_in[0]
+  if conid >= ncon:
+    return
+  worldid = contact_worldid_in[conid]
+  contact_geom = contact_geom_in[conid]
+  weld1 = body_weldid[geom_bodyid[contact_geom[0]]]
+  weld2 = body_weldid[geom_bodyid[contact_geom[1]]]
+  geom1 = contact_geom[0]
+  geom2 = contact_geom[1]
+
+  for side in range(2):
+    if side == 0:
+      weld = weld1
+      geom = geom2
+    else:
+      weld = weld2
+      geom = geom1
+
+    idx = wp.atomic_add(weld_geom_count_out[worldid], weld, 1)
+    if idx < MJ_MAXCONPAIR:
+      weld_geom_list_out[worldid, weld, idx] = geom
+
+
 @wp.kernel
 def _sensor_tactile(
     # Model:
@@ -2084,7 +2128,10 @@ def _sensor_tactile(
     geom_bodyid: wp.array(dtype=int),
     geom_size: wp.array2d(dtype=wp.vec3),
     mesh_vertadr: wp.array(dtype=int),
+    mesh_vertnum: wp.array(dtype=int),
+    mesh_octadr: wp.array(dtype=int),
     mesh_normaladr: wp.array(dtype=int),
+    mesh_normalnum: wp.array(dtype=int),
     mesh_vert: wp.array(dtype=wp.vec3),
     mesh_normal: wp.array(dtype=wp.vec3),
     mesh_quat: wp.array(dtype=wp.quat),
@@ -2102,38 +2149,23 @@ def _sensor_tactile(
     geom_xmat_in: wp.array2d(dtype=wp.mat33),
     subtree_com_in: wp.array2d(dtype=wp.vec3),
     cvel_in: wp.array2d(dtype=wp.spatial_vector),
-  contact_geom_in: wp.array(dtype=wp.vec2i),
+    # In:
-  contact_worldid_in: wp.array(dtype=int),
+    weld_geom_count_in: wp.array2d(dtype=int),
-  nacon_in: wp.array(dtype=int),
+    weld_geom_list_in: wp.array3d(dtype=int),
     # Data out:
     sensordata_out: wp.array2d(dtype=float),
 ):
-  conid, taxelid = wp.tid()
+  worldid, taxelid = wp.tid()
 
-  if conid >= nacon_in[0]:
-    return
-
-  worldid = contact_worldid_in[conid]
-
-  # get sensor_id
   sensor_id = taxel_sensorid[taxelid]
-
-  # get parent weld id
   mesh_id = sensor_objid[sensor_id]
   geom_id = sensor_refid[sensor_id]
   parent_body = geom_bodyid[geom_id]
   parent_weld = body_weldid[parent_body]
 
-  # contact geom
+  geom_count = weld_geom_count_in[worldid, parent_weld]
-  body1 = body_weldid[geom_bodyid[contact_geom_in[conid][0]]]
+  if geom_count == 0:
-  body2 = body_weldid[geom_bodyid[contact_geom_in[conid][1]]]
-  if body1 == parent_weld:
-    geom = contact_geom_in[conid][1]
-  elif body2 == parent_weld:
-    geom = contact_geom_in[conid][0]
-  else:
     return
-  body = geom_bodyid[geom]
 
   # vertex local position
   vertid = taxel_vertadr[taxelid] - mesh_vertadr[mesh_id]
@@ -2143,18 +2175,46 @@ def _sensor_tactile(
   xpos = geom_xmat_in[worldid, geom_id] @ pos
   xpos += geom_xpos_in[worldid, geom_id]
 
-  # position in other geom frame
+  has_frame = mesh_normalnum[mesh_id] == 3 * mesh_vertnum[mesh_id]
+  normal_stride = 3 if has_frame else 1
+  offset = mesh_normaladr[mesh_id] + normal_stride * vertid
+  quat = mesh_quat[mesh_id]
+  normal = math.rot_vec_quat(mesh_normal[offset], quat)
+  tang1 = wp.vec3(0.0, 0.0, 0.0)
+  tang2 = wp.vec3(0.0, 0.0, 0.0)
+  if has_frame:
+    tang1 = math.rot_vec_quat(mesh_normal[offset + 1], quat)
+    tang2 = math.rot_vec_quat(mesh_normal[offset + 2], quat)
+
+  for g in range(MJ_MAXCONPAIR):
+    if g >= geom_count:
+      break
+
+    geom = weld_geom_list_in[worldid, parent_weld, g]
+    if geom < 0:
+      continue
+
+    is_dup = int(0)
+    for j in range(g):
+      if weld_geom_list_in[worldid, parent_weld, j] == geom:
+        is_dup = int(1)
+        break
+    if is_dup == int(1):
+      continue
+
+    body = geom_bodyid[geom]
+
     tmp = xpos - geom_xpos_in[worldid, geom]
     lpos = wp.transpose(geom_xmat_in[worldid, geom]) @ tmp
 
     plugin_id = geom_plugin_index[geom]
-
     contact_type = geom_type[geom]
 
     plugin_attributes, plugin_index, volume_data, mesh_data = get_sdf_params(
         oct_child,
         oct_aabb,
         oct_coeff,
+        mesh_octadr,
         plugin,
         plugin_attr,
         contact_type,
@@ -2163,37 +2223,60 @@ def _sensor_tactile(
         mesh_id,
     )
 
-  depth = wp.min(sdf(contact_type, lpos, plugin_attributes, plugin_index, volume_data, mesh_data), 0.0)
+    depth = wp.min(
+        sdf(
+            contact_type,
+            lpos,
+            plugin_attributes,
+            plugin_index,
+            volume_data,
+            mesh_data,
+        ),
+        0.0,
+    )
     if depth >= 0.0:
-    return
+      continue
 
-  # get velocity in global
+    vel_sensor = _transform_spatial(
-  vel_sensor = _transform_spatial(cvel_in[worldid, parent_weld], xpos - subtree_com_in[worldid, body_rootid[parent_weld]])
+        cvel_in[worldid, parent_weld],
+        xpos - subtree_com_in[worldid, body_rootid[parent_weld]],
+    )
     vel_other = _transform_spatial(
-    cvel_in[worldid, body], geom_xpos_in[worldid, geom] - subtree_com_in[worldid, body_rootid[body]]
+        cvel_in[worldid, body],
+        geom_xpos_in[worldid, geom]
+        - subtree_com_in[worldid, body_rootid[body]],
     )
     vel_rel = vel_sensor - vel_other
 
-  # get contact force/torque, rotate into node frame
-  offset = mesh_normaladr[mesh_id] + 3 * vertid
-  normal = math.rot_vec_quat(mesh_normal[offset], mesh_quat[mesh_id])
-  tang1 = math.rot_vec_quat(mesh_normal[offset + 1], mesh_quat[mesh_id])
-  tang2 = math.rot_vec_quat(mesh_normal[offset + 2], mesh_quat[mesh_id])
     kMaxDepth = 0.05
     pressure = depth / wp.max(kMaxDepth - depth, MJ_MINVAL)
     force = wp.mul(normal, pressure)
 
-  # one row of mat^T * force
+    forceT = wp.vec3(0.0, 0.0, 0.0)
-  forceT = wp.vec3()
     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))
 
-  # add to sensor output
+    dim = sensor_dim[sensor_id] // 3
-  dim = sensor_dim[sensor_id] / 3
+    wp.atomic_add(
-  wp.atomic_add(sensordata_out[worldid], sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
+        sensordata_out,
-  wp.atomic_add(sensordata_out[worldid], sensor_adr[sensor_id] + 1 * dim + vertid, forceT[1])
+        worldid,
-  wp.atomic_add(sensordata_out[worldid], sensor_adr[sensor_id] + 2 * dim + vertid, forceT[2])
+        sensor_adr[sensor_id] + 0 * dim + vertid,
+        forceT[0],
+    )
+    wp.atomic_add(
+        sensordata_out,
+        worldid,
+        sensor_adr[sensor_id] + 1 * dim + vertid,
+        forceT[1],
+    )
+    wp.atomic_add(
+        sensordata_out,
+        worldid,
+        sensor_adr[sensor_id] + 2 * dim + vertid,
+        forceT[2],
+    )
 
 
 @wp.func
@@ -2421,9 +2504,27 @@ def sensor_acc(m: Model, d: Data):
     ],
   )
 
+  weld_geom_count = wp.zeros((d.nworld, m.nbody), dtype=int)
+  weld_geom_list = wp.full((d.nworld, m.nbody, MJ_MAXCONPAIR), -1, dtype=int)
+  wp.launch(
+      _preprocess_tactile_contacts,
+      dim=d.naconmax,
+      inputs=[
+          m.body_weldid,
+          m.geom_bodyid,
+          d.contact.geom,
+          d.contact.worldid,
+          d.nacon,
+      ],
+      outputs=[
+          weld_geom_count,
+          weld_geom_list,
+      ],
+  )
+
   wp.launch(
       _sensor_tactile,
-    dim=(d.naconmax, m.nsensortaxel),
+      dim=(d.nworld, m.nsensortaxel),
       inputs=[
           m.body_rootid,
           m.body_weldid,
@@ -2434,7 +2535,10 @@ def sensor_acc(m: Model, d: Data):
           m.geom_bodyid,
           m.geom_size,
           m.mesh_vertadr,
+          m.mesh_vertnum,
+          m.mesh_octadr,
           m.mesh_normaladr,
+          m.mesh_normalnum,
           m.mesh_vert,
           m.mesh_normal,
           m.mesh_quat,
@@ -2451,9 +2555,8 @@ def sensor_acc(m: Model, d: Data):
           d.geom_xmat,
           d.subtree_com,
           d.cvel,
-      d.contact.geom,
+          weld_geom_count,
-      d.contact.worldid,
+          weld_geom_list,
-      d.nacon,
       ],
       outputs=[
           d.sensordata,
@@ -2616,6 +2719,9 @@ def sensor_acc(m: Model, d: Data):
     ],
   )
 
+  if m.callback.sensor:
+    m.callback.sensor(m, d, Stage.ACC)
+
 
 @wp.kernel
 def _energy_pos_zero(

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

@@ -14,29 +14,29 @@
 # ==============================================================================
 
 
-import warp as wp
-
 from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import support
 from mujoco.mjx.third_party.mujoco_warp._src import util_misc
-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 CamLightType
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import EqType
 from mujoco.mjx.third_party.mujoco_warp._src.types import JointType
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
 from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
+from mujoco.mjx.third_party.mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN
 from mujoco.mjx.third_party.mujoco_warp._src.types import TileSet
 from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
-from mujoco.mjx.third_party.mujoco_warp._src.types import WrapType
-from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec11
+from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
+from mujoco.mjx.third_party.mujoco_warp._src.types import WrapType
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
+import warp as wp
 
 wp.set_module_options({"enable_backward": False})
 
@@ -241,10 +241,8 @@ def _flex_vertices(
 def _flex_edges(
     # Model:
     nflex: int,
-  body_parentid: wp.array(dtype=int),
     body_rootid: wp.array(dtype=int),
     body_dofadr: wp.array(dtype=int),
-  dof_bodyid: wp.array(dtype=int),
     flex_vertadr: wp.array(dtype=int),
     flex_edgeadr: wp.array(dtype=int),
     flex_edgenum: wp.array(dtype=int),
@@ -258,7 +256,7 @@ def _flex_edges(
     cdof_in: wp.array2d(dtype=wp.spatial_vector),
     flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
     # Data out:
-  flexedge_J_out: wp.array3d(dtype=float),
+    flexedge_J_out: wp.array2d(dtype=float),
     flexedge_length_out: wp.array2d(dtype=float),
     flexedge_velocity_out: wp.array2d(dtype=float),
 ):
@@ -285,64 +283,40 @@ def _flex_edges(
 
   dofi = body_dofadr[b1]
   dofj = body_dofadr[b2]
-  dofi0 = dofi + 0
-  dofi1 = dofi + 1
-  dofi2 = dofi + 2
-  dofj0 = dofj + 0
-  dofj1 = dofj + 1
-  dofj2 = dofj + 2
 
-  vel1 = wp.vec3(qvel_in[worldid, dofi0], qvel_in[worldid, dofi1], qvel_in[worldid, dofi2])
+  vel1 = wp.vec3(
-  vel2 = wp.vec3(qvel_in[worldid, dofj0], qvel_in[worldid, dofj1], qvel_in[worldid, dofj2])
+      qvel_in[worldid, dofi],
+      qvel_in[worldid, dofi + 1],
+      qvel_in[worldid, dofi + 2],
+  )
+  vel2 = wp.vec3(
+      qvel_in[worldid, dofj],
+      qvel_in[worldid, dofj + 1],
+      qvel_in[worldid, dofj + 2],
+  )
   flexedge_velocity_out[worldid, edgeid] = wp.dot(vel2 - vel1, edge)
 
   rowadr = flexedge_J_rowadr[edgeid]
 
-  sparseid0 = rowadr + 0
+  # compute offsets once per body (avoids 12 redundant tree-ancestry walks in jac_dof)
-  sparseid1 = rowadr + 1
+  offset1 = pos1 - wp.vec3(subtree_com_in[worldid, body_rootid[b1]])
-  sparseid2 = rowadr + 2
+  offset2 = pos2 - wp.vec3(subtree_com_in[worldid, body_rootid[b2]])
-  sparseid3 = rowadr + 3
-  sparseid4 = rowadr + 4
-  sparseid5 = rowadr + 5
 
-  # TODO(team): jacdif
+  # body1 DOFs: b1 is in subtree, b2 is not -> jacdif = 0 - jacp1 = -jacp1
+  for k in range(3):
+    cdof = cdof_in[worldid, dofi + k]
+    cdof_ang = wp.spatial_top(cdof)
+    cdof_lin = wp.spatial_bottom(cdof)
+    jacp1 = cdof_lin + wp.cross(cdof_ang, offset1)
+    flexedge_J_out[worldid, rowadr + k] = wp.dot(-jacp1, edge)
 
-  jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofi0, worldid)
+  # body2 DOFs: b2 is in subtree, b1 is not -> jacdif = jacp2 - 0 = jacp2
-  jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofi0, worldid)
+  for k in range(3):
-  jacdif = jacp2 - jacp1
+    cdof = cdof_in[worldid, dofj + k]
-  Ji0 = wp.dot(jacdif, edge)
+    cdof_ang = wp.spatial_top(cdof)
-
+    cdof_lin = wp.spatial_bottom(cdof)
-  jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofi1, worldid)
+    jacp2 = cdof_lin + wp.cross(cdof_ang, offset2)
-  jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofi1, worldid)
+    flexedge_J_out[worldid, rowadr + 3 + k] = wp.dot(jacp2, edge)
-  jacdif = jacp2 - jacp1
-  Ji1 = wp.dot(jacdif, edge)
-
-  jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofi2, worldid)
-  jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofi2, worldid)
-  jacdif = jacp2 - jacp1
-  Ji2 = wp.dot(jacdif, edge)
-
-  jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofj0, worldid)
-  jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofj0, worldid)
-  jacdif = jacp2 - jacp1
-  Jj0 = wp.dot(jacdif, edge)
-
-  jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofj1, worldid)
-  jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofj1, worldid)
-  jacdif = jacp2 - jacp1
-  Jj1 = wp.dot(jacdif, edge)
-
-  jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofj2, worldid)
-  jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofj2, worldid)
-  jacdif = jacp2 - jacp1
-  Jj2 = wp.dot(jacdif, edge)
-
-  flexedge_J_out[worldid, 0, sparseid0] = Ji0
-  flexedge_J_out[worldid, 0, sparseid1] = Ji1
-  flexedge_J_out[worldid, 0, sparseid2] = Ji2
-  flexedge_J_out[worldid, 0, sparseid3] = Jj0
-  flexedge_J_out[worldid, 0, sparseid4] = Jj1
-  flexedge_J_out[worldid, 0, sparseid5] = Jj2
 
 
 @event_scope
@@ -414,10 +388,8 @@ def flex(m: Model, d: Data):
     dim=(d.nworld, m.nflexedge),
     inputs=[
       m.nflex,
-      m.body_parentid,
       m.body_rootid,
       m.body_dofadr,
-      m.dof_bodyid,
       m.flex_vertadr,
       m.flex_edgeadr,
       m.flex_edgenum,
@@ -818,7 +790,9 @@ def _qM_sparse(
   bodyid = dof_bodyid[dofid]
 
   # init M(i,i) with armature inertia
-  qM_out[worldid, 0, madr_ij] = dof_armature[worldid, dofid]
+  qM_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])
@@ -907,7 +881,7 @@ def _tendon_armature(
   if is_sparse:  # is_sparse is not batched
     madr_ij = dof_Madr[dofid]
 
-  armature = tendon_armature[worldid, tenid]
+  armature = tendon_armature[worldid % tendon_armature.shape[0], tenid]
 
   if armature == 0.0:
     return
@@ -1320,7 +1294,7 @@ def _cfrc_ext_equality(
   )
 
   id = efc_id_in[worldid, efcid]
-  eq_data_ = eq_data[worldid, id]
+  eq_data_ = eq_data[worldid % eq_data.shape[0], id]
   body_semantic = eq_objtype[id] == ObjType.BODY
 
   obj1 = eq_obj1id[id]
@@ -1341,7 +1315,7 @@ def _cfrc_ext_equality(
       else:
         offset = wp.vec3(eq_data_[3], eq_data_[4], eq_data_[5])
     else:
-      offset = site_pos[worldid, obj1]
+      offset = site_pos[worldid % site_pos.shape[0], obj1]
 
     # transform point on body1: local -> global
     pos = xmat_in[worldid, bodyid1] @ offset + xpos_in[worldid, bodyid1]
@@ -1363,7 +1337,7 @@ def _cfrc_ext_equality(
       else:
         offset = wp.vec3(eq_data_[0], eq_data_[1], eq_data_[2])
     else:
-      offset = site_pos[worldid, obj2]
+      offset = site_pos[worldid % site_pos.shape[0], obj2]
 
     # transform point on body2: local -> global
     pos = xmat_in[worldid, bodyid2] @ offset + xpos_in[worldid, bodyid2]
@@ -1558,7 +1532,7 @@ def _tendon_dot(
 ):
   worldid, tenid = wp.tid()
 
-  armature = tendon_armature[worldid, tenid]
+  armature = tendon_armature[worldid % tendon_armature.shape[0], tenid]
   if armature == 0.0:
     return
 
@@ -1716,7 +1690,7 @@ def _tendon_bias_coef(
 ):
   worldid, tenid, dofid = wp.tid()
 
-  armature = tendon_armature[worldid, tenid]
+  armature = tendon_armature[worldid % tendon_armature.shape[0], tenid]
   if armature == 0.0:
     return
 
@@ -1740,7 +1714,7 @@ def _tendon_bias_qfrc(
 ):
   worldid, tenid, dofid = wp.tid()
 
-  armature = tendon_armature[worldid, tenid]
+  armature = tendon_armature[worldid % tendon_armature.shape[0], tenid]
   if armature == 0.0:
     return
 
@@ -1935,7 +1909,7 @@ def _transmission(
     actuator_trntype: wp.array(dtype=int),
     actuator_trnid: wp.array(dtype=wp.vec2i),
     actuator_gear: wp.array2d(dtype=wp.spatial_vector),
-  actuator_cranklength: wp.array(dtype=float),
+    actuator_cranklength: wp.array2d(dtype=float),
     # Data in:
     qpos_in: wp.array2d(dtype=float),
     xquat_in: wp.array2d(dtype=wp.quat),
@@ -1945,9 +1919,14 @@ def _transmission(
     cdof_in: wp.array2d(dtype=wp.spatial_vector),
     ten_J_in: wp.array3d(dtype=float),
     ten_length_in: wp.array2d(dtype=float),
+    # In:
+    moment_nnz: wp.array(dtype=int),
     # Data out:
     actuator_length_out: wp.array2d(dtype=float),
-  actuator_moment_out: wp.array3d(dtype=float),
+    moment_rownnz_out: wp.array2d(dtype=int),
+    moment_rowadr_out: wp.array2d(dtype=int),
+    moment_colind_out: wp.array2d(dtype=int),
+    actuator_moment_out: wp.array2d(dtype=float),
 ):
   worldid, actid = wp.tid()
   trntype = actuator_trntype[actid]
@@ -1960,20 +1939,33 @@ def _transmission(
     qadr = jnt_qposadr[jntid]
     vadr = jnt_dofadr[jntid]
     if jnt_typ == JointType.FREE:
+      moment_rownnz_out[worldid, actid] = 6
+      rowadr = wp.atomic_add(moment_nnz, worldid, 6)
+      moment_rowadr_out[worldid, actid] = rowadr
+      moment_colind_out[worldid, rowadr + 0] = vadr + 0
+      moment_colind_out[worldid, rowadr + 1] = vadr + 1
+      moment_colind_out[worldid, rowadr + 2] = vadr + 2
+      moment_colind_out[worldid, rowadr + 3] = vadr + 3
+      moment_colind_out[worldid, rowadr + 4] = vadr + 4
+      moment_colind_out[worldid, rowadr + 5] = vadr + 5
       actuator_length_out[worldid, actid] = 0.0
       if trntype == TrnType.JOINTINPARENT:
         quat = wp.normalize(wp.quat(qpos[qadr + 3], qpos[qadr + 4], qpos[qadr + 5], qpos[qadr + 6]))
         quat_neg = math.quat_inv(quat)
         gearaxis = math.rot_vec_quat(wp.spatial_bottom(gear), quat_neg)
-        actuator_moment_out[worldid, actid, vadr + 0] = gear[0]
+        actuator_moment_out[worldid, rowadr + 0] = gear[0]
-        actuator_moment_out[worldid, actid, vadr + 1] = gear[1]
+        actuator_moment_out[worldid, rowadr + 1] = gear[1]
-        actuator_moment_out[worldid, actid, vadr + 2] = gear[2]
+        actuator_moment_out[worldid, rowadr + 2] = gear[2]
-        actuator_moment_out[worldid, actid, vadr + 3] = gearaxis[0]
+        actuator_moment_out[worldid, rowadr + 3] = gearaxis[0]
-        actuator_moment_out[worldid, actid, vadr + 4] = gearaxis[1]
+        actuator_moment_out[worldid, rowadr + 4] = gearaxis[1]
-        actuator_moment_out[worldid, actid, vadr + 5] = gearaxis[2]
+        actuator_moment_out[worldid, rowadr + 5] = gearaxis[2]
       else:
-        for i in range(6):
+        actuator_moment_out[worldid, rowadr + 0] = gear[0]
-          actuator_moment_out[worldid, actid, vadr + i] = gear[i]
+        actuator_moment_out[worldid, rowadr + 1] = gear[1]
+        actuator_moment_out[worldid, rowadr + 2] = gear[2]
+        actuator_moment_out[worldid, rowadr + 3] = gear[3]
+        actuator_moment_out[worldid, rowadr + 4] = gear[4]
+        actuator_moment_out[worldid, rowadr + 5] = gear[5]
     elif jnt_typ == JointType.BALL:
       q = wp.quat(qpos[qadr + 0], qpos[qadr + 1], qpos[qadr + 2], qpos[qadr + 3])
       q = wp.normalize(q)
@@ -1983,11 +1975,25 @@ def _transmission(
         quat_neg = math.quat_inv(q)
         gearaxis = math.rot_vec_quat(gearaxis, quat_neg)
       actuator_length_out[worldid, actid] = wp.dot(axis_angle, gearaxis)
+
+      nnz = 3
+      moment_rownnz_out[worldid, actid] = nnz
+      rowadr = wp.atomic_add(moment_nnz, worldid, nnz)
+      moment_rowadr_out[worldid, actid] = rowadr
+
       for i in range(3):
-        actuator_moment_out[worldid, actid, vadr + i] = gearaxis[i]
+        sparseid = rowadr + i
+        moment_colind_out[worldid, sparseid] = vadr + i
+        actuator_moment_out[worldid, sparseid] = gearaxis[i]
     elif jnt_typ == JointType.SLIDE or jnt_typ == JointType.HINGE:
       actuator_length_out[worldid, actid] = qpos[qadr] * gear[0]
-      actuator_moment_out[worldid, actid, vadr] = gear[0]
+
+      nnz = 1
+      moment_rownnz_out[worldid, actid] = nnz
+      rowadr = wp.atomic_add(moment_nnz, worldid, nnz)
+      moment_rowadr_out[worldid, actid] = rowadr
+      moment_colind_out[worldid, rowadr] = vadr
+      actuator_moment_out[worldid, rowadr] = gear[0]
     else:
       wp.printf("unrecognized joint type")
   elif trntype == TrnType.SLIDERCRANK:
@@ -1996,7 +2002,7 @@ def _transmission(
     id = trnid[0]
     idslider = trnid[1]
     gear0 = gear[0]
-    rod = actuator_cranklength[actid]
+    rod = actuator_cranklength[worldid % actuator_cranklength.shape[0], actid]
     site_xmat = site_xmat_in[worldid, idslider]
     axis = wp.vec3(site_xmat[0, 2], site_xmat[1, 2], site_xmat[2, 2])
     site_xpos_id = site_xpos_in[worldid, id]
@@ -2027,26 +2033,77 @@ def _transmission(
       dldv = axis
       dlda = vec
 
-    # apply chain rule
+    # count dofs
-    # TODO(team): parallelize?
+    b1 = body_weldid[site_bodyid[id]]
-    for i in range(nv):
+    b2 = body_weldid[site_bodyid[idslider]]
+    da1_init = int(-1)
+    da2_init = int(-1)
+    if b1 > 0:
+      da1_init = body_dofadr[b1] + body_dofnum[b1] - 1
+    if b2 > 0:
+      da2_init = body_dofadr[b2] + body_dofnum[b2] - 1
+
+    da1 = da1_init
+    da2 = da2_init
+    ndof = int(0)
+    while da1 >= 0 or da2 >= 0:
+      da = wp.max(da1, da2)
+      ndof += 1
+      if da1 == da:
+        da1 = dof_parentid[da1]
+      if da2 == da:
+        da2 = dof_parentid[da2]
+
+    moment_rownnz_out[worldid, actid] = ndof
+    rowadr = wp.atomic_add(moment_nnz, worldid, ndof)
+    moment_rowadr_out[worldid, actid] = rowadr
+
+    # traverse dofs
+    da1 = da1_init
+    da2 = da2_init
+
+    ptr = ndof - 1
+    while da1 >= 0 or da2 >= 0:
+      da = wp.max(da1, da2)
+
       # get Jacobians of axis(jacA) and vec(jac)
-      # mj_jacPointAxis
       jacp, jacr = support.jac_dof(
-        body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, site_xpos_idslider, site_bodyid[idslider], i, worldid
+          body_parentid,
+          body_rootid,
+          dof_bodyid,
+          subtree_com_in,
+          cdof_in,
+          site_xpos_idslider,
+          site_bodyid[idslider],
+          da,
+          worldid,
       )
       jacS = jacp
       jacA = wp.cross(jacr, axis)
-
-      # mj_jacSite
       jac, _ = support.jac_dof(
-        body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, site_xpos_id, site_bodyid[id], i, worldid
+          body_parentid,
+          body_rootid,
+          dof_bodyid,
+          subtree_com_in,
+          cdof_in,
+          site_xpos_id,
+          site_bodyid[id],
+          da,
+          worldid,
       )
       jac -= jacS
 
       # apply the chain rule
       moment = wp.dot(dlda, jacA) + wp.dot(dldv, jac)
-      actuator_moment_out[worldid, actid, i] = moment * gear0
+      sparseid = rowadr + ptr
+      moment_colind_out[worldid, sparseid] = da
+      actuator_moment_out[worldid, sparseid] = moment * gear0
+      ptr -= 1
+
+      if da1 == da:
+        da1 = dof_parentid[da1]
+      if da2 == da:
+        da2 = dof_parentid[da2]
   elif trntype == TrnType.TENDON:
     tenid = actuator_trnid[actid][0]
 
@@ -2057,19 +2114,46 @@ def _transmission(
     adr = tendon_adr[tenid]
     if wrap_type[adr] == WrapType.JOINT:
       ten_num = tendon_num[tenid]
+      rowadr = wp.atomic_add(moment_nnz, worldid, ten_num)
+      moment_rownnz_out[worldid, actid] = ten_num
+      moment_rowadr_out[worldid, actid] = rowadr
+
       for i in range(ten_num):
         dofadr = jnt_dofadr[wrap_objid[adr + i]]
-        actuator_moment_out[worldid, actid, dofadr] = ten_J_in[worldid, tenid, dofadr] * gear0
+        sparseid = rowadr + i
+        moment_colind_out[worldid, sparseid] = dofadr
+        actuator_moment_out[worldid, sparseid] = (
+            ten_J_in[worldid, tenid, dofadr] * gear0
+        )
     else:  # spatial
+      # TODO(team): sparse tendon jacobian
+      ten_nnz = int(0)
       for dofadr in range(nv):
-        actuator_moment_out[worldid, actid, dofadr] = ten_J_in[worldid, tenid, dofadr] * gear0
+        if ten_J_in[worldid, tenid, dofadr] != 0.0:
+          ten_nnz += 1
+      rowadr = wp.atomic_add(moment_nnz, worldid, ten_nnz)
+      moment_rownnz_out[worldid, actid] = ten_nnz
+      moment_rowadr_out[worldid, actid] = rowadr
+      ptr = int(0)
+      for dofadr in range(nv):
+        J = ten_J_in[worldid, tenid, dofadr]
+        if J != 0.0:
+          sparseid = rowadr + ptr
+          moment_colind_out[worldid, sparseid] = dofadr
+          actuator_moment_out[worldid, sparseid] = J * gear0
+          ptr += 1
   elif trntype == TrnType.BODY:
     # cannot compute meaningful length, set to zero
     actuator_length_out[worldid, actid] = 0.0
 
     # initialize moment
+    rowadr = wp.atomic_add(moment_nnz, worldid, nv)
+    moment_rownnz_out[worldid, actid] = nv
+    moment_rowadr_out[worldid, actid] = rowadr
     for i in range(nv):
-      actuator_moment_out[worldid, actid, i] = 0.0
+      sparseid = rowadr + i
+      moment_colind_out[worldid, sparseid] = i
+      actuator_moment_out[worldid, sparseid] = 0.0
 
     # moment computed by _transmission_body_moment and _transmission_body_moment_scale
   elif trntype == TrnType.SITE:
@@ -2077,7 +2161,7 @@ def _transmission(
     siteid = trnid[0]
     refid = trnid[1]
 
-    gear = actuator_gear[worldid, actid]
+    gear = actuator_gear[actuator_gear_id, actid]
     site_quat_id = worldid % site_quat.shape[0]
     gear_translation = wp.spatial_top(gear)
     gear_rotational = wp.spatial_bottom(gear)
@@ -2089,9 +2173,27 @@ def _transmission(
       wrench_translation = site_xmat @ gear_translation
       wrench_rotation = site_xmat @ gear_rotational
 
-      # moment: global Jacobian projected on wrench
+      # count dofs
-      # TODO(team): parallelize
+      b1 = body_weldid[site_bodyid[siteid]]
-      for i in range(nv):
+      da_init = int(-1)
+      if b1 > 0:
+        da_init = body_dofadr[b1] + body_dofnum[b1] - 1
+
+      da = da_init
+      ndof = int(0)
+      while da >= 0:
+        ndof += 1
+        da = dof_parentid[da]
+
+      moment_rownnz_out[worldid, actid] = ndof
+      rowadr = wp.atomic_add(moment_nnz, worldid, ndof)
+      moment_rowadr_out[worldid, actid] = rowadr
+      actuator_length_out[worldid, actid] = 0.0
+
+      # traverse dofs
+      da = da_init
+      ptr = ndof - 1
+      while da >= 0:
         jacp, jacr = support.jac_dof(
             body_parentid,
             body_rootid,
@@ -2100,11 +2202,17 @@ def _transmission(
             cdof_in,
             site_xpos_in[worldid, siteid],
             site_bodyid[siteid],
-          i,
+            da,
             worldid,
         )
-        actuator_length_out[worldid, actid] = 0.0
+        moment = wp.dot(jacp, wrench_translation) + wp.dot(
-        actuator_moment_out[worldid, actid, i] = wp.dot(jacp, wrench_translation) + wp.dot(jacr, wrench_rotation)
+            jacr, wrench_rotation
+        )
+        sparseid = rowadr + ptr
+        moment_colind_out[worldid, sparseid] = da
+        actuator_moment_out[worldid, sparseid] = moment
+        ptr -= 1
+        da = dof_parentid[da]
     # reference site defined
     else:
       # initialize last dof address for each body
@@ -2162,38 +2270,79 @@ def _transmission(
 
       actuator_length_out[worldid, actid] = length
 
-      # TODO(team): parallelize
+      # count dofs
-      for i in range(nv):
+      da1_init = int(-1)
-        jacp, jacr = support.jac_dof(
+      da2_init = int(-1)
-          body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, site_xpos, site_bodyid[siteid], i, worldid
+      if b0 > 0:
-        )
+        da1_init = body_dofadr[b0] + body_dofnum[b0] - 1
+      if b1 > 0:
+        da2_init = body_dofadr[b1] + body_dofnum[b1] - 1
 
-        # jacref: global Jacobian of reference site
+      da1 = da1_init
-        jacpref, jacrref = support.jac_dof(
+      da2 = da2_init
-          body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, ref_xpos, site_bodyid[refid], i, worldid
+      ndof = int(0)
-        )
+      while da1 >= 0 or da2 >= 0:
-
+        da = wp.max(da1, da2)
-        jacpdif = jacp - jacpref
+        if da1 == da and da2 == da:
-        jacrdif = jacr - jacrref
-
-        # if common ancestral dof was found, clear the columns of its parental chain
-        da = dofadr_common
-        while da >= 0:
-          if da == i:
-            jacpdif = wp.vec3(0.0)
-            jacrdif = wp.vec3(0.0)
           break
-          da = dof_parentid[da]
+        ndof += 1
+        if da1 == da:
+          da1 = dof_parentid[da1]
+        if da2 == da:
+          da2 = dof_parentid[da2]
+
+      moment_rownnz_out[worldid, actid] = ndof
+      rowadr = wp.atomic_add(moment_nnz, worldid, ndof)
+      moment_rowadr_out[worldid, actid] = rowadr
+
+      # traverse dofs
+      da1 = da1_init
+      da2 = da2_init
+
+      ptr = ndof - 1
+      while da1 >= 0 or da2 >= 0:
+        da = wp.max(da1, da2)
+        if da1 == da and da2 == da:
+          break
+
+        jacp, jacr = support.jac_dof(
+            body_parentid,
+            body_rootid,
+            dof_bodyid,
+            subtree_com_in,
+            cdof_in,
+            site_xpos,
+            site_bodyid[siteid],
+            da,
+            worldid,
+        )
+        jacpref, jacrref = support.jac_dof(
+            body_parentid,
+            body_rootid,
+            dof_bodyid,
+            subtree_com_in,
+            cdof_in,
+            ref_xpos,
+            site_bodyid[refid],
+            da,
+            worldid,
+        )
 
-        # moment: global Jacobian projected on wrench
         moment = float(0.0)
-
         if translational_transmission:
-          moment += wp.dot(jacpdif, wrench_translation)
+          moment += wp.dot(jacp - jacpref, wrench_translation)
         if rotational_transmission:
-          moment += wp.dot(jacrdif, wrench_rotation)
+          moment += wp.dot(jacr - jacrref, wrench_rotation)
 
-        actuator_moment_out[worldid, actid, i] = moment
+        sparseid = rowadr + ptr
+        moment_colind_out[worldid, sparseid] = da
+        actuator_moment_out[worldid, sparseid] = moment
+        ptr -= 1
+
+        if da1 == da:
+          da1 = dof_parentid[da1]
+        if da2 == da:
+          da2 = dof_parentid[da2]
   else:
     wp.printf("unhandled transmission type %d\n", trntype)
 
@@ -2211,6 +2360,7 @@ def _transmission_body_moment(
     # Data in:
     subtree_com_in: wp.array2d(dtype=wp.vec3),
     cdof_in: wp.array2d(dtype=wp.spatial_vector),
+    moment_rowadr_in: wp.array2d(dtype=int),
     contact_dist_in: wp.array(dtype=float),
     contact_pos_in: wp.array(dtype=wp.vec3),
     contact_frame_in: wp.array(dtype=wp.mat33),
@@ -2219,10 +2369,15 @@ def _transmission_body_moment(
     contact_geom_in: wp.array(dtype=wp.vec2i),
     contact_efc_address_in: wp.array2d(dtype=int),
     contact_worldid_in: wp.array(dtype=int),
+    efc_J_rownnz_in: wp.array2d(dtype=int),
+    efc_J_rowadr_in: wp.array2d(dtype=int),
+    efc_J_colind_in: wp.array3d(dtype=int),
     efc_J_in: wp.array3d(dtype=float),
     nacon_in: wp.array(dtype=int),
+    # In:
+    efc_is_sparse: bool,
     # Data out:
-  actuator_moment_out: wp.array3d(dtype=float),
+    actuator_moment_out: wp.array2d(dtype=float),
     # Out:
     actuator_trntype_body_ncon_out: wp.array2d(dtype=int),
 ):
@@ -2257,23 +2412,61 @@ def _transmission_body_moment(
   if dofid == 0:
     wp.atomic_add(actuator_trntype_body_ncon_out[worldid], trnbodyid, 1)
 
+  rowadr = moment_rowadr_in[worldid, actid]
+
   # mark contact normals in efc_force
   if contact_exclude == 0:
     contact_dim = contact_dim_in[conid]
     contact_efc_address = contact_efc_address_in[conid]
 
     if contact_dim == 1 or opt_cone == ConeType.ELLIPTIC:
-      efc_force = 1.0
       efcid0 = contact_efc_address[0]
-      wp.atomic_add(actuator_moment_out[worldid, actid], dofid, efc_J_in[worldid, efcid0, dofid] * efc_force)
+      if efc_is_sparse:
-
+        rownnz = efc_J_rownnz_in[worldid, efcid0]
+        if dofid < rownnz:
+          efc_rowadr = efc_J_rowadr_in[worldid, efcid0]
+          efc_sparseid = efc_rowadr + dofid
+          colind = efc_J_colind_in[worldid, 0, efc_sparseid]
+          wp.atomic_add(
+              actuator_moment_out[worldid],
+              rowadr + colind,
+              efc_J_in[worldid, 0, efc_sparseid],
+          )
+        else:
+          return
+      else:
+        colind = dofid
+        wp.atomic_add(
+            actuator_moment_out[worldid],
+            rowadr + colind,
+            efc_J_in[worldid, efcid0, dofid],
+        )
     else:
       npyramid = contact_dim - 1  # number of frictional directions
       efc_force = 0.5 / float(npyramid)
 
       for j in range(2 * npyramid):
         efcid = contact_efc_address[j]
-        wp.atomic_add(actuator_moment_out[worldid, actid], dofid, efc_J_in[worldid, efcid, dofid] * efc_force)
+        if efc_is_sparse:
+          rownnz = efc_J_rownnz_in[worldid, efcid]
+          if dofid < rownnz:
+            efc_rowadr = efc_J_rowadr_in[worldid, efcid]
+            efc_sparseid = efc_rowadr + dofid
+            colind = efc_J_colind_in[worldid, 0, efc_sparseid]
+            wp.atomic_add(
+                actuator_moment_out[worldid],
+                rowadr + colind,
+                efc_J_in[worldid, 0, efc_sparseid] * efc_force,
+            )
+          else:
+            return
+        else:
+          colind = dofid
+          wp.atomic_add(
+              actuator_moment_out[worldid],
+              rowadr + colind,
+              efc_J_in[worldid, efcid, dofid] * efc_force,
+          )
 
   # excluded contact in gap: get Jacobian, accumulate
   elif contact_exclude == 1:
@@ -2282,22 +2475,58 @@ def _transmission_body_moment(
     normal = wp.vec3(contact_frame[0, 0], contact_frame[0, 1], contact_frame[0, 2])
 
     # get Jacobian difference
-    jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, contact_pos, b1, dofid, worldid)
+    efcid0 = contact_efc_address_in[conid][0]
-    jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, contact_pos, b2, dofid, worldid)
+    if efc_is_sparse and efcid0 >= 0:
+      # contact has valid efc row: use sparse pattern
+      if dofid >= efc_J_rownnz_in[worldid, efcid0]:
+        return
+      sparseid = efc_J_rowadr_in[worldid, efcid0] + dofid
+      colind = efc_J_colind_in[worldid, 0, sparseid]
+    else:
+      # excluded contact with no efc row or dense: use dofid directly
+      colind = dofid
+
+    jacp1, _ = support.jac_dof(
+        body_parentid,
+        body_rootid,
+        dof_bodyid,
+        subtree_com_in,
+        cdof_in,
+        contact_pos,
+        b1,
+        colind,
+        worldid,
+    )
+    jacp2, _ = support.jac_dof(
+        body_parentid,
+        body_rootid,
+        dof_bodyid,
+        subtree_com_in,
+        cdof_in,
+        contact_pos,
+        b2,
+        colind,
+        worldid,
+    )
+
     jacdif = jacp2 - jacp1
 
     # project Jacobian along the normal of the contact frame
-    wp.atomic_add(actuator_moment_out[worldid, actid], dofid, wp.dot(normal, jacdif))
+    wp.atomic_add(
+        actuator_moment_out[worldid], rowadr + colind, wp.dot(normal, jacdif)
+    )
 
 
 @wp.kernel
 def _transmission_body_moment_scale(
     # Model:
     actuator_trntype_body_adr: wp.array(dtype=int),
+    # Data in:
+    moment_rowadr_in: wp.array2d(dtype=int),
     # In:
     actuator_trntype_body_ncon_in: wp.array2d(dtype=int),
     # Data out:
-  actuator_moment_out: wp.array3d(dtype=float),
+    actuator_moment_out: wp.array2d(dtype=float),
 ):
   worldid, trnbodyid, dofid = wp.tid()
 
@@ -2305,7 +2534,8 @@ def _transmission_body_moment_scale(
 
   if ncon > 0:
     actid = actuator_trntype_body_adr[trnbodyid]
-    actuator_moment_out[worldid, actid, dofid] /= -float(ncon)
+    rowadr = moment_rowadr_in[worldid, actid]
+    actuator_moment_out[worldid, rowadr + dofid] /= -float(ncon)
 
 
 @event_scope
@@ -2315,10 +2545,12 @@ def transmission(m: Model, d: Data):
   Updates the actuator length and moments for all actuators in the model, including joint
   and tendon transmissions.
   """
-  d.actuator_moment.zero_()
+  # TODO(team): investigate pre-computing moment_rownnz, moment_rowadr, moment_colind
+  moment_nnz = wp.zeros((d.nworld,), dtype=int)
+
   wp.launch(
       _transmission,
-    dim=[d.nworld, m.nu],
+      dim=(d.nworld, m.nu),
       inputs=[
           m.nv,
           m.body_parentid,
@@ -2349,13 +2581,21 @@ def transmission(m: Model, d: Data):
           d.cdof,
           d.ten_J,
           d.ten_length,
+          moment_nnz,
+      ],
+      outputs=[
+          d.actuator_length,
+          d.moment_rownnz,
+          d.moment_rowadr,
+          d.moment_colind,
+          d.actuator_moment,
       ],
-    outputs=[d.actuator_length, d.actuator_moment],
   )
 
   if m.nacttrnbody:
     # compute moments
     ncon = wp.zeros((d.nworld, m.nacttrnbody), dtype=int)
+
     wp.launch(
         _transmission_body_moment,
         dim=(m.nacttrnbody, d.naconmax, m.nv),
@@ -2369,6 +2609,7 @@ def transmission(m: Model, d: Data):
             m.actuator_trntype_body_adr,
             d.subtree_com,
             d.cdof,
+            d.moment_rowadr,
             d.contact.dist,
             d.contact.pos,
             d.contact.frame,
@@ -2377,8 +2618,12 @@ def transmission(m: Model, d: Data):
             d.contact.geom,
             d.contact.efc_address,
             d.contact.worldid,
+            d.efc.J_rownnz,
+            d.efc.J_rowadr,
+            d.efc.J_colind,
             d.efc.J,
             d.nacon,
+            SPARSE_CONSTRAINT_JACOBIAN,
         ],
         outputs=[d.actuator_moment, ncon],
     )
@@ -2387,7 +2632,7 @@ def transmission(m: Model, d: Data):
     wp.launch(
         _transmission_body_moment_scale,
         dim=(d.nworld, m.nacttrnbody, m.nv),
-      inputs=[m.actuator_trntype_body_adr, ncon],
+        inputs=[m.actuator_trntype_body_adr, d.moment_rowadr, ncon],
         outputs=[d.actuator_moment],
     )
 
@@ -2888,7 +3133,7 @@ def _spatial_geom_tendon(
   # get geom information
   geom_xpos = geom_xpos_in[worldid, wrap_objid_geom]
   geom_xmat = geom_xmat_in[worldid, wrap_objid_geom]
-  geomsize = geom_size[worldid, wrap_objid_geom][0]
+  geomsize = geom_size[worldid % geom_size.shape[0], wrap_objid_geom][0]
   geom_type = wrap_type[wrap_adr]
 
   # get body ids for site-geom-site instances
@@ -3147,6 +3392,7 @@ def _spatial_tendon_wrap(
     ten_wrapnum_out[worldid, i] = wrapnum
 
 
+@event_scope
 def tendon(m: Model, d: Data):
   """Computes tendon lengths and moments.
 

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

@@ -14,6 +14,7 @@
 # ==============================================================================
 import dataclasses
 import enum
+from typing import Callable
 
 import mujoco
 import warp as wp
@@ -32,6 +33,9 @@ MJ_MAX_EPAFACES = 5
 TILE_SIZE_JTDAJ_SPARSE = 16
 TILE_SIZE_JTDAJ_DENSE = 16
 
+# TODO(team): remove after improving performance for sparse constraint jacobian
+SPARSE_CONSTRAINT_JACOBIAN = False
+
 # TODO(team): remove after mjwarp depends on warp-lang >= 1.12 in pyproject.toml
 TEXTURE_DTYPE = wp.Texture2D if hasattr(wp, "Texture2D") else int
 
@@ -134,6 +138,20 @@ class ProjectionType(enum.IntEnum):
     ORTHOGRAPHIC = 1
 
 
+class Stage(enum.IntEnum):
+  """Computation stage.
+
+  Attributes:
+    POS: position-dependent
+    VEL: velocity-dependent
+    ACC: acceleration/force-dependent
+  """
+
+  POS = mujoco.mjtStage.mjSTAGE_POS
+  VEL = mujoco.mjtStage.mjSTAGE_VEL
+  ACC = mujoco.mjtStage.mjSTAGE_ACC
+
+
 class DataType(enum.IntFlag):
   """Sensor data types.
 
@@ -167,6 +185,7 @@ class DisableBit(enum.IntFlag):
     SENSOR:       sensors
     EULERDAMP:    implicit damping for Euler integration
     NATIVECCD:    native convex collision detection (ignored in MJWarp)
+    ISLAND:       constraint islands
   """
 
   CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT
@@ -185,7 +204,8 @@ class DisableBit(enum.IntFlag):
   SENSOR = mujoco.mjtDisableBit.mjDSBL_SENSOR
   EULERDAMP = mujoco.mjtDisableBit.mjDSBL_EULERDAMP
   NATIVECCD = mujoco.mjtDisableBit.mjDSBL_NATIVECCD
-  # unsupported: MIDPHASE, AUTORESET, ISLAND
+  ISLAND = mujoco.mjtDisableBit.mjDSBL_ISLAND
+  # unsupported: MIDPHASE, AUTORESET
 
 
 class EnableBit(enum.IntFlag):
@@ -232,6 +252,7 @@ class DynType(enum.IntEnum):
     FILTER: linear filter: da/dt = (u-a) / tau
     FILTEREXACT: linear filter: da/dt = (u-a) / tau, with exact integration
     MUSCLE: piece-wise linear filter with two time constants
+    USER: user-defined dynamics via act_dyn_callback
   """
 
   NONE = mujoco.mjtDyn.mjDYN_NONE
@@ -239,7 +260,7 @@ class DynType(enum.IntEnum):
   FILTER = mujoco.mjtDyn.mjDYN_FILTER
   FILTEREXACT = mujoco.mjtDyn.mjDYN_FILTEREXACT
   MUSCLE = mujoco.mjtDyn.mjDYN_MUSCLE
-  # unsupported: USER
+  USER = mujoco.mjtDyn.mjDYN_USER
 
 
 class GainType(enum.IntEnum):
@@ -249,12 +270,13 @@ class GainType(enum.IntEnum):
     FIXED: fixed gain
     AFFINE: const + kp*length + kv*velocity
     MUSCLE: muscle FLV curve computed by muscle_gain
+    USER: user-defined gain via act_gain_callback
   """
 
   FIXED = mujoco.mjtGain.mjGAIN_FIXED
   AFFINE = mujoco.mjtGain.mjGAIN_AFFINE
   MUSCLE = mujoco.mjtGain.mjGAIN_MUSCLE
-  # unsupported: USER
+  USER = mujoco.mjtGain.mjGAIN_USER
 
 
 class BiasType(enum.IntEnum):
@@ -264,12 +286,13 @@ class BiasType(enum.IntEnum):
     NONE: no bias
     AFFINE: const + kp*length + kv*velocity
     MUSCLE: muscle passive force computed by muscle_bias
+    USER: user-defined bias via act_bias_callback
   """
 
   NONE = mujoco.mjtBias.mjBIAS_NONE
   AFFINE = mujoco.mjtBias.mjBIAS_AFFINE
   MUSCLE = mujoco.mjtBias.mjBIAS_MUSCLE
-  # unsupported: USER
+  USER = mujoco.mjtBias.mjBIAS_USER
 
 
 class JointType(enum.IntEnum):
@@ -462,6 +485,7 @@ class SensorType(enum.IntEnum):
     FRAMELINACC: 3D linear acceleration
     FRAMEANGACC: 3D angular acceleration
     TACTILE: tactile sensor
+    USER: user-defined sensor via sensor_callback
   """
 
   MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER
@@ -511,6 +535,7 @@ class SensorType(enum.IntEnum):
   FRAMELINACC = mujoco.mjtSensor.mjSENS_FRAMELINACC
   FRAMEANGACC = mujoco.mjtSensor.mjSENS_FRAMEANGACC
   TACTILE = mujoco.mjtSensor.mjSENS_TACTILE
+  USER = mujoco.mjtSensor.mjSENS_USER
 
 
 class ObjType(enum.IntEnum):
@@ -768,6 +793,29 @@ class TileSet:
   size: int
 
 
+@dataclasses.dataclass
+class Callback:
+  """Callbacks for custom physics behavior.
+
+  Attributes:
+    passive: custom passive forces, writes to ``Data.qfrc_passive``
+    control: custom control laws, writes to ``Data.ctrl``
+    act_dyn: custom actuator dynamics, writes to ``Data.act_dot``
+    act_gain: custom actuator gains, writes to ``Data.actuator_force``
+    act_bias: custom actuator biases, writes to ``Data.actuator_force``
+    sensor: custom sensors, writes to ``Data.sensordata``
+    contactfilter: custom contact filtering, writes to ``Data.contact``
+  """
+
+  passive: Callable | None = None
+  control: Callable | None = None
+  act_dyn: Callable | None = None
+  act_gain: Callable | None = None
+  act_bias: Callable | None = None
+  sensor: Callable | None = None
+  contactfilter: Callable | None = None
+
+
 @dataclasses.dataclass
 class Model:
   """Model definition and parameters.
@@ -793,6 +841,7 @@ class Model:
     nflexelem: number of elements in all flexes
     nflexelemdata: number of element vertex ids in all flexes
     nflexelemedge: number of element edge ids in all flexes
+    nJfe: number of non-zeros in sparse flexedge Jacobian
     nmesh: number of meshes
     nmeshvert: number of vertices for all meshes
     nmeshnormal: number of normals in all meshes
@@ -812,6 +861,7 @@ class Model:
     nsensor: number of sensors
     nmocap: number of mocap bodies
     nplugin: number of plugin instances
+    nJmom: number of non-zeros in actuator_moment
     ngravcomp: number of bodies with nonzero gravcomp
     nsensordata: number of elements in sensor data vector
     opt: physics options
@@ -947,7 +997,9 @@ class Model:
     mesh_vertadr: first vertex address                       (nmesh,)
     mesh_vertnum: number of vertices                         (nmesh,)
     mesh_faceadr: first face address                         (nmesh,)
+    mesh_octadr: octree address for each mesh                (nmesh,)
     mesh_normaladr: first normal address                     (nmesh,)
+    mesh_normalnum: number of normals                        (nmesh,)
     mesh_graphadr: graph data address; -1: no graph          (nmesh,)
     mesh_vert: vertex positions for all meshes               (nmeshvert, 3)
     mesh_normal: normals for all meshes                      (nmeshnormal, 3)
@@ -968,7 +1020,8 @@ class Model:
     hfield_ncol: number of columns in grid                   (nhfield,)
     hfield_adr: start address in hfield_data                 (nhfield,)
     hfield_data: elevation data                              (nhfielddata,)
-    mat_texid: texture id for rendering                      (*, nmat, mjNTEXROLE)
+    mat_texid: texture id for rendering                      (*, nmat,
+      mjNTEXROLE)
     mat_texrepeat: texture repeat for rendering              (*, nmat, 2)
     mat_rgba: rgba                                           (*, nmat, 4)
     pair_dim: contact dimensionality                         (npair,)
@@ -993,10 +1046,14 @@ class Model:
     tendon_num: number of objects in tendon's path           (ntendon,)
     tendon_limited: does tendon have length limits           (ntendon,)
     tendon_actfrclimited: does ten have actuator force limit (ntendon,)
-    tendon_solref_lim: constraint solver reference: limit    (*, ntendon, mjNREF)
+    tendon_solref_lim: constraint solver reference: limit    (*, ntendon,
-    tendon_solimp_lim: constraint solver impedance: limit    (*, ntendon, mjNIMP)
+      mjNREF)
-    tendon_solref_fri: constraint solver reference: friction (*, ntendon, mjNREF)
+    tendon_solimp_lim: constraint solver impedance: limit    (*, ntendon,
-    tendon_solimp_fri: constraint solver impedance: friction (*, ntendon, mjNIMP)
+      mjNIMP)
+    tendon_solref_fri: constraint solver reference: friction (*, ntendon,
+      mjNREF)
+    tendon_solimp_fri: constraint solver impedance: friction (*, ntendon,
+      mjNIMP)
     tendon_range: tendon length limits                       (*, ntendon, 2)
     tendon_actfrcrange: range of total actuator force        (*, ntendon, 2)
     tendon_margin: min distance for limit detection          (*, ntendon)
@@ -1028,9 +1085,9 @@ class Model:
     actuator_forcerange: range of forces                     (*, nu, 2)
     actuator_actrange: range of activations                  (*, nu, 2)
     actuator_gear: scale length and transmitted force        (*, nu, 6)
-    actuator_cranklength: crank length for slider-crank      (nu,)
+    actuator_cranklength: crank length for slider-crank      (*, nu)
-    actuator_acc0: acceleration from unit force in qpos0     (nu,)
+    actuator_acc0: acceleration from unit force in qpos0     (*, nu)
-    actuator_lengthrange: feasible actuator length range     (nu, 2)
+    actuator_lengthrange: feasible actuator length range     (*, nu, 2)
     sensor_type: sensor type (SensorType)                    (nsensor,)
     sensor_datatype: numeric data type (DataType)            (nsensor,)
     sensor_objtype: type of sensorized object (ObjType)      (nsensor,)
@@ -1049,6 +1106,7 @@ class Model:
     mapM2M: index mapping from M (legacy) to M (CSR)         (nC)
 
   warp only fields:
+    callback: custom physics callbacks
     nbranch: number of branches (leaf-to-root paths)
     nv_pad: number of degrees of freedom + padding
     nacttrnbody: number of actuators with body transmission
@@ -1062,12 +1120,14 @@ class Model:
     nmaxpolygon: maximum number of verts per polygon
     nmaxmeshdeg: maximum number of polygons per vert
     is_sparse: whether to use sparse representations
-    has_fluid: True if wind, density, or viscosity are non-zero at put_model time
+    has_fluid: True if wind, density, or viscosity are non-zero at put_model
+      time
     has_sdf_geom: whether the model contains SDF geoms
     block_dim: block dim options
     body_tree: list of body ids by tree level
     body_branches: flattened body ids for all branches
-    body_branch_start: start index in body_branches for each branch   (nbranch + 1,)
+    body_branch_start: start index in body_branches for each branch   (nbranch +
+      1,)
     mocap_bodyid: id of body for mocap                       (nmocap,)
     body_fluid_ellipsoid: does body use ellipsoid fluid      (nbody,)
     jnt_limited_slide_hinge_adr: limited/slide/hinge jntadr
@@ -1156,6 +1216,7 @@ class Model:
   nflexelem: int
   nflexelemdata: int
   nflexelemedge: int
+  nJfe: int
   nmesh: int
   nmeshvert: int
   nmeshnormal: int
@@ -1175,6 +1236,7 @@ class Model:
   nsensor: int
   nmocap: int
   nplugin: int
+  nJmom: int
   ngravcomp: int
   nsensordata: int
   opt: Option
@@ -1306,11 +1368,13 @@ class Model:
   flex_damping: array("nflex", float)
   flexedge_J_rownnz: array("nflexedge", int)
   flexedge_J_rowadr: array("nflexedge", int)
-  flexedge_J_colind: wp.array(dtype=int)
+  flexedge_J_colind: array("nJfe", int)
   mesh_vertadr: array("nmesh", int)
   mesh_vertnum: array("nmesh", int)
   mesh_faceadr: array("nmesh", int)
+  mesh_octadr: array("nmesh", int)
   mesh_normaladr: array("nmesh", int)
+  mesh_normalnum: array("nmesh", int)
   mesh_graphadr: array("nmesh", int)
   mesh_vert: array("nmeshvert", wp.vec3)
   mesh_normal: array("nmeshnormal", wp.vec3)
@@ -1391,9 +1455,9 @@ class Model:
   actuator_forcerange: array("*", "nu", wp.vec2)
   actuator_actrange: array("*", "nu", wp.vec2)
   actuator_gear: array("*", "nu", wp.spatial_vector)
-  actuator_cranklength: array("nu", float)
+  actuator_cranklength: array("*", "nu", float)
-  actuator_acc0: array("nu", float)
+  actuator_acc0: array("*", "nu", float)
-  actuator_lengthrange: array("nu", wp.vec2)
+  actuator_lengthrange: array("*", "nu", wp.vec2)
   sensor_type: array("nsensor", int)
   sensor_datatype: array("nsensor", int)
   sensor_objtype: array("nsensor", int)
@@ -1411,6 +1475,7 @@ class Model:
   M_colind: array("nC", int)
   mapM2M: array("nC", int)
   # warp only fields:
+  callback: Callback
   nbranch: int
   nv_pad: int
   nacttrnbody: int
@@ -1546,7 +1611,14 @@ class Constraint:
   Attributes:
     type: constraint type (ConstraintType)            (nworld, njmax)
     id: id of object of specific type                 (nworld, njmax)
-    J: constraint Jacobian                            (nworld, njmax_pad, nv_pad)
+    J_rownnz: number of non-zeros in J row            (nworld, 0) dense (nworld,
+      njmax) sparse
+    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
+    J: constraint Jacobian                            (nworld, njmax_pad,
+      nv_pad) dense (nworld, 1, njmax * nv) sparse
     pos: constraint position (equality, contact)      (nworld, njmax)
     margin: inclusion margin (contact)                (nworld, njmax)
     D: constraint mass                                (nworld, njmax_pad)
@@ -1555,13 +1627,17 @@ class Constraint:
     frictionloss: frictionloss (friction)             (nworld, njmax)
     force: constraint force in constraint space       (nworld, njmax)
     state: constraint state                           (nworld, njmax_pad)
+
   warp only fields:
     Ma: M*qacc                                        (nworld, nv)
   """
 
   type: array("nworld", "njmax", int)
   id: array("nworld", "njmax", int)
-  J: array("nworld", "njmax_pad", "nv_pad", float)
+  J_rownnz: wp.array2d(dtype=int)
+  J_rowadr: wp.array2d(dtype=int)
+  J_colind: wp.array3d(dtype=int)
+  J: wp.array3d(dtype=float)
   pos: array("nworld", "njmax", float)
   margin: array("nworld", "njmax", float)
   D: array("nworld", "njmax_pad", float)
@@ -1583,6 +1659,7 @@ class Data:
     nf: number of friction constraints                          (nworld,)
     nl: number of limit constraints                             (nworld,)
     nefc: number of constraints                                 (nworld,)
+    nisland: number of constraint islands                       (nworld,)
     time: simulation time                                       (nworld,)
     energy: potential, kinetic energy                           (nworld, 2)
     qpos: position                                              (nworld, nq)
@@ -1591,52 +1668,88 @@ class Data:
     qacc_warmstart: acceleration used for warmstart             (nworld, nv)
     ctrl: control                                               (nworld, nu)
     qfrc_applied: applied generalized force                     (nworld, nv)
-    xfrc_applied: applied Cartesian force/torque                (nworld, nbody, 6)
+    xfrc_applied: applied Cartesian force/torque                (nworld, nbody,
+      6)
     eq_active: enable/disable constraints                       (nworld, neq)
-    mocap_pos: position of mocap bodies                         (nworld, nmocap, 3)
+    mocap_pos: position of mocap bodies                         (nworld, nmocap,
-    mocap_quat: orientation of mocap bodies                     (nworld, nmocap, 4)
+      3)
+    mocap_quat: orientation of mocap bodies                     (nworld, nmocap,
+      4)
     qacc: acceleration                                          (nworld, nv)
     act_dot: time-derivative of actuator activation             (nworld, na)
-    sensordata: sensor data array                               (nworld, nsensordata,)
+    sensordata: sensor data array                               (nworld,
-    xpos: Cartesian position of body frame                      (nworld, nbody, 3)
+      nsensordata,)
-    xquat: Cartesian orientation of body frame                  (nworld, nbody, 4)
+    xpos: Cartesian position of body frame                      (nworld, nbody,
-    xmat: Cartesian orientation of body frame                   (nworld, nbody, 3, 3)
+      3)
-    xipos: Cartesian position of body com                       (nworld, nbody, 3)
+    xquat: Cartesian orientation of body frame                  (nworld, nbody,
-    ximat: Cartesian orientation of body inertia                (nworld, nbody, 3, 3)
+      4)
-    xanchor: Cartesian position of joint anchor                 (nworld, njnt, 3)
+    xmat: Cartesian orientation of body frame                   (nworld, nbody,
-    xaxis: Cartesian joint axis                                 (nworld, njnt, 3)
+      3, 3)
-    geom_xpos: Cartesian geom position                          (nworld, ngeom, 3)
+    xipos: Cartesian position of body com                       (nworld, nbody,
-    geom_xmat: Cartesian geom orientation                       (nworld, ngeom, 3, 3)
+      3)
-    site_xpos: Cartesian site position                          (nworld, nsite, 3)
+    ximat: Cartesian orientation of body inertia                (nworld, nbody,
-    site_xmat: Cartesian site orientation                       (nworld, nsite, 3, 3)
+      3, 3)
-    cam_xpos: Cartesian camera position                         (nworld, ncam, 3)
+    xanchor: Cartesian position of joint anchor                 (nworld, njnt,
-    cam_xmat: Cartesian camera orientation                      (nworld, ncam, 3, 3)
+      3)
-    light_xpos: Cartesian light position                        (nworld, nlight, 3)
+    xaxis: Cartesian joint axis                                 (nworld, njnt,
-    light_xdir: Cartesian light direction                       (nworld, nlight, 3)
+      3)
-    subtree_com: center of mass of each subtree                 (nworld, nbody, 3)
+    geom_xpos: Cartesian geom position                          (nworld, ngeom,
+      3)
+    geom_xmat: Cartesian geom orientation                       (nworld, ngeom,
+      3, 3)
+    site_xpos: Cartesian site position                          (nworld, nsite,
+      3)
+    site_xmat: Cartesian site orientation                       (nworld, nsite,
+      3, 3)
+    cam_xpos: Cartesian camera position                         (nworld, ncam,
+      3)
+    cam_xmat: Cartesian camera orientation                      (nworld, ncam,
+      3, 3)
+    light_xpos: Cartesian light position                        (nworld, nlight,
+      3)
+    light_xdir: Cartesian light direction                       (nworld, nlight,
+      3)
+    subtree_com: center of mass of each subtree                 (nworld, nbody,
+      3)
     cdof: com-based motion axis of each dof (rot:lin)           (nworld, nv, 6)
-    cinert: com-based body inertia and mass                     (nworld, nbody, 10)
+    cinert: com-based body inertia and mass                     (nworld, nbody,
-    flexvert_xpos: cartesian flex vertex positions              (nworld, nflexvert, 3)
+      10)
-    flexedge_J: edge length Jacobian                            (nworld, 1, nflexedge*6)
+    flexvert_xpos: cartesian flex vertex positions              (nworld,
-    flexedge_length: flex edge lengths                          (nworld, nflexedge, 1)
+      nflexvert, 3)
-    ten_wrapadr: start address of tendon's path                 (nworld, ntendon)
+    flexedge_J: edge length Jacobian                            (nworld, nJfe)
-    ten_wrapnum: number of wrap points in path                  (nworld, ntendon)
+    flexedge_length: flex edge lengths                          (nworld,
-    ten_J: tendon Jacobian                                      (nworld, ntendon, nv)
+      nflexedge, 1)
-    ten_length: tendon lengths                                  (nworld, ntendon)
+    ten_wrapadr: start address of tendon's path                 (nworld,
-    wrap_obj: geomid; -1: site; -2: pulley                      (nworld, nwrap, 2)
+      ntendon)
-    wrap_xpos: Cartesian 3D points in all paths                 (nworld, nwrap, 6)
+    ten_wrapnum: number of wrap points in path                  (nworld,
+      ntendon)
+    ten_J: tendon Jacobian                                      (nworld,
+      ntendon, nv)
+    ten_length: tendon lengths                                  (nworld,
+      ntendon)
+    wrap_obj: geomid; -1: site; -2: pulley                      (nworld, nwrap,
+      2)
+    wrap_xpos: Cartesian 3D points in all paths                 (nworld, nwrap,
+      6)
     actuator_length: actuator lengths                           (nworld, nu)
-    actuator_moment: actuator moments                           (nworld, nu, nv)
+    moment_rownnz: number of non-zeros in actuator_moment row   (nworld, nu)
-    crb: com-based composite inertia and mass                   (nworld, nbody, 10)
+    moment_rowadr: row start address in actuator_moment         (nworld, nu)
-    qM: total inertia                                           (nworld, nv, nv) if dense
+    moment_colind: column indices in sparse actuator_moment     (nworld, nJmom)
-                                                                (nworld, 1, nM) if sparse
+    actuator_moment: actuator moments                           (nworld, nJmom)
-    qLD: L'*D*L factorization of M                              (nworld, nv, nv) if dense
+    crb: com-based composite inertia and mass                   (nworld, nbody,
-                                                                (nworld, 1, nC) if sparse
+      10)
+    qM: total inertia                                           (nworld, nv, nv)
+      if dense (nworld, 1, nM) if sparse
+    qLD: L'*D*L factorization of M                              (nworld, nv, nv)
+      if dense (nworld, 1, nC) if sparse
     qLDiagInv: 1/diag(D)                                        (nworld, nv)
-    flexedge_velocity: flex edge velocities                     (nworld, nflexedge)
+    flexedge_velocity: flex edge velocities                     (nworld,
-    ten_velocity: tendon velocities                             (nworld, ntendon)
+      nflexedge)
+    ten_velocity: tendon velocities                             (nworld,
+      ntendon)
     actuator_velocity: actuator velocities                      (nworld, nu)
-    cvel: com-based velocity (rot:lin)                          (nworld, nbody, 6)
+    cvel: com-based velocity (rot:lin)                          (nworld, nbody,
+      6)
     cdof_dot: time-derivative of cdof (rot:lin)                 (nworld, nv, 6)
     qfrc_bias: C(qpos,qvel)                                     (nworld, nv)
     qfrc_spring: passive spring force                           (nworld, nv)
@@ -1644,27 +1757,33 @@ class Data:
     qfrc_gravcomp: passive gravity compensation force           (nworld, nv)
     qfrc_fluid: passive fluid force                             (nworld, nv)
     qfrc_passive: total passive force                           (nworld, nv)
-    subtree_linvel: linear velocity of subtree com              (nworld, nbody, 3)
+    subtree_linvel: linear velocity of subtree com              (nworld, nbody,
-    subtree_angmom: angular momentum about subtree com          (nworld, nbody, 3)
+      3)
+    subtree_angmom: angular momentum about subtree com          (nworld, nbody,
+      3)
     actuator_force: actuator force in actuation space           (nworld, nu)
     qfrc_actuator: actuator force                               (nworld, nv)
     qfrc_smooth: net unconstrained force                        (nworld, nv)
     qacc_smooth: unconstrained acceleration                     (nworld, nv)
     qfrc_constraint: constraint force                           (nworld, nv)
     qfrc_inverse: net external force; should equal:             (nworld, nv)
-                  qfrc_applied + J.T @ xfrc_applied
+      qfrc_applied + J.T @ xfrc_applied + qfrc_actuator
-                  + qfrc_actuator
+    cacc: com-based acceleration                                (nworld, nbody,
-    cacc: com-based acceleration                                (nworld, nbody, 6)
+      6)
-    cfrc_int: com-based interaction force with parent           (nworld, nbody, 6)
+    cfrc_int: com-based interaction force with parent           (nworld, nbody,
-    cfrc_ext: com-based external force on body                  (nworld, nbody, 6)
+      6)
+    cfrc_ext: com-based external force on body                  (nworld, nbody,
+      6)
     contact: contact data
     efc: constraint data
+    tree_island: island ID per tree (-1 if unconstrained)       (nworld, ntree)
 
   warp only fields:
     nworld: number of worlds
     naconmax: maximum number of contacts (shared across all worlds)
     naccdmax: maximum number of contacts for CCD (all worlds)
     njmax: maximum number of constraints per world
+    njmax_pad: njmax rounded up to the nearest multiple of TILE_SIZE_JTDAJ
     nacon: number of detected contacts (across all worlds)      (1,)
     ncollision: collision count from broadphase                 (1,)
   """
@@ -1674,6 +1793,7 @@ class Data:
   nf: array("nworld", int)
   nl: array("nworld", int)
   nefc: array("nworld", int)
+  nisland: array("nworld", int)
   time: array("nworld", float)
   energy: array("nworld", wp.vec2)
   qpos: array("nworld", "nq", float)
@@ -1708,7 +1828,7 @@ class Data:
   cdof: array("nworld", "nv", wp.spatial_vector)
   cinert: array("nworld", "nbody", vec10)
   flexvert_xpos: array("nworld", "nflexvert", wp.vec3)
-  flexedge_J: wp.array3d(dtype=float)
+  flexedge_J: array("nworld", "nJfe", float)
   flexedge_length: array("nworld", "nflexedge", float)
   ten_wrapadr: array("nworld", "ntendon", int)
   ten_wrapnum: array("nworld", "ntendon", int)
@@ -1717,7 +1837,10 @@ class Data:
   wrap_obj: array("nworld", "nwrap", wp.vec2i)
   wrap_xpos: array("nworld", "nwrap", wp.spatial_vector)
   actuator_length: array("nworld", "nu", float)
-  actuator_moment: array("nworld", "nu", "nv", float)
+  moment_rownnz: array("nworld", "nu", int)
+  moment_rowadr: array("nworld", "nu", int)
+  moment_colind: array("nworld", "nJmom", int)
+  actuator_moment: array("nworld", "nJmom", float)
   crb: array("nworld", "nbody", vec10)
   qM: wp.array3d(dtype=float)
   qLD: wp.array3d(dtype=float)
@@ -1746,31 +1869,18 @@ class Data:
   cfrc_ext: array("nworld", "nbody", wp.spatial_vector)
   contact: Contact
   efc: Constraint
+  tree_island: array("nworld", "ntree", int)
 
   # warp only fields:
   nworld: int
   naconmax: int
   naccdmax: int
   njmax: int
+  njmax_pad: int
   nacon: array(1, int)
   ncollision: array(1, int)
 
 
-@dataclasses.dataclass
-class CollisionContext:
-  """Collision driver intermediate arrays.
-
-  Attributes:
-    collision_pair: collision pairs from broadphase             (naconmax, 2)
-    collision_pairid: ids from broadphase                       (naconmax, 2)
-    collision_worldid: collision world ids from broadphase      (naconmax,)
-  """
-
-  collision_pair: wp.array
-  collision_pairid: wp.array
-  collision_worldid: wp.array
-
-
 @dataclasses.dataclass
 class RenderContext:
   """Context for rendering.
@@ -1794,7 +1904,7 @@ class RenderContext:
     textures_registry: texture registry
     hfield_registry: hfield BVH id to warp mesh mapping
     hfield_bvh_id: hfield BVH ids
-    hfield_bounds_size: hfield bounds size
+    hfield_bounds_size: hfield bounds half-extents
     flex_mesh: flex mesh
     flex_rgba: flex rgba
     flex_bvh_id: flex BVH id

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

@@ -0,0 +1,97 @@
+# 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.
+# ==============================================================================
+"""Package version checking utilities."""
+
+import importlib.metadata
+import operator
+import re
+
+
+def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
+  """Parse a version string into comparable components.
+
+  Both '.' and '-' are treated as separators. 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 (e.g., 1.2.3 >= 1.2.3.dev). Non-numeric
+  components are compared lexicographically (e.g., b >= a).
+
+  Args:
+    version_str: Version string like "3.5.0" or "3.5.0.dev869102767".
+
+  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.
+  """
+  # Split on both '.' and '-'
+  parts = re.split(r"[.\-]", version_str)
+  return tuple(
+      [(0, int(p)) if p.isdigit() else (-1, p) for p in parts] + [(0, 0)]
+  )
+
+
+def check_version(spec: str) -> bool:
+  """Check if an installed package satisfies a version requirement.
+
+  Supports operators: >=, <=, >, <, ==, !=
+
+  Version comparison rules:
+  - Both '.' and '-' are treated as separators
+  - Numeric components are compared numerically
+  - Non-numeric components are compared lexicographically
+  - Stable releases are greater than pre-releases (e.g., 1.2.3 >= 1.2.3.dev)
+
+  Args:
+    spec: Version specification like "numpy>=1.20.0".
+
+  Returns:
+    True if the installed version satisfies the requirement.
+
+  Raises:
+    ValueError: If the spec cannot be parsed.
+    importlib.metadata.PackageNotFoundError: If the package is not installed.
+  """
+  match = re.match(r"^([a-zA-Z0-9_\-]+)(>=|<=|>|<|==|!=)(.+)$", spec)
+  if not match:
+    raise ValueError(
+        f"Invalid version spec '{spec}'. Expected format: 'package>=version'"
+    )
+  package_name, op, version_str = match.groups()
+
+  required_version = _parse_version(version_str)
+
+  try:
+    installed_str = importlib.metadata.version(package_name)
+  except importlib.metadata.PackageNotFoundError as e:
+    # Fallback: import the package and read __version__
+    try:
+      import importlib as _importlib  # noqa: F811
+
+      mod = _importlib.import_module(package_name)
+      installed_str = mod.__version__
+    except (ImportError, AttributeError):
+      raise e
+
+  installed_version = _parse_version(installed_str)
+
+  ops = {
+      ">=": operator.ge,
+      "<=": operator.le,
+      ">": operator.gt,
+      "<": operator.lt,
+      "==": operator.eq,
+      "!=": operator.ne,
+  }
+  return ops[op](installed_version, required_version)

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

@@ -15,6 +15,7 @@
 
 import functools
 import inspect
+import warnings
 
 import warp as wp
 
@@ -144,4 +145,37 @@ def check_toolkit_driver():
   wp.init()
   if wp.get_device().is_cuda:
     if not wp.is_conditional_graph_supported():
-      raise RuntimeError("Minimum supported CUDA version: 12.4.")
+      warnings.warn(
+          """
+        CUDA version < 12.4 detected
+        - graph capture may be unreliable for < 12.3
+        - conditional graph nodes are not available for < 12.4
+          Model.opt.graph_conditional should be set to False
+        """,
+          stacklevel=2,
+      )
+
+
+class scoped_mathdx_gemm_disabled:
+  """Temporarily disable Warp MathDx GEMM kernels within this scope."""
+
+  def __init__(self, disable: bool = True):
+    self._disable = disable
+    self._config = None
+    self._prev = None
+
+  def __enter__(self):
+    if not self._disable:
+      return self
+    config = getattr(wp, "config", None)
+    if config is None or not hasattr(config, "enable_mathdx_gemm"):
+      return self
+    self._config = config
+    self._prev = config.enable_mathdx_gemm
+    self._config.enable_mathdx_gemm = False
+    return self
+
+  def __exit__(self, exc_type, exc, tb):
+    if self._config is not None:
+      self._config.enable_mathdx_gemm = self._prev
+    return False

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name="mujoco-warp"
-version = "3.5.0"
+version = "3.6.0"
 # TODO(team): create a distribution list
 authors = [
     {name = "Newton Developers", email = "mujoco@deepmind.com"},
@@ -21,15 +21,16 @@ classifiers = [
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "Topic :: Scientific/Engineering",
 ]
 requires-python = ">=3.10"
 dependencies = [
     "absl-py",
     "etils[epath]",
-    "mujoco>=3.4.0",
+    "mujoco>=3.5.0",
     "numpy",
-    "warp-lang>=1.11.0",
+    "warp-lang>=1.12",
 ]
 
 [[tool.uv.index]]
@@ -54,7 +55,7 @@ dev = [
     "ruff",
     "pygls>=1.0.0,<2.0.0",
     "lsprotocol>=2023.0.1,<2024.0.0",
-    "mujoco>=3.4.1.dev0",
+    "mujoco>=3.5.0.dev0",
     "warp-lang>=1.11.0.dev0",
 ]
 # TODO(team): cpu and cuda JAX optional dependencies are temporary, remove after we land MJX:Warp

mjx/mujoco/mjx/warp/bvh.py

@@ -44,6 +44,9 @@ _c = mjwarp.Contact(
 _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
+_cb = mjwp_types.Callback(
+    **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
+)
 
 @ffi.format_args_for_warp
 def _refit_bvh_shim(
@@ -70,6 +73,7 @@ def _refit_bvh_shim(
 ):
   _m.stat = _s
   _m.opt = _o
+  _m.callback = _cb
   _d.efc = _e
   _d.contact = _c
   _m.flex_dim = flex_dim

mjx/mujoco/mjx/warp/collision_driver.py

@@ -42,6 +42,10 @@ _c = mjwarp.Contact(
 _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
+_cb = mjwp_types.Callback(
+    **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
+)
+
 
 @ffi.format_args_for_warp
 def _collision_shim(
@@ -73,6 +77,7 @@ def _collision_shim(
     mesh_faceadr: wp.array(dtype=int),
     mesh_graph: wp.array(dtype=int),
     mesh_graphadr: wp.array(dtype=int),
+    mesh_octadr: wp.array(dtype=int),
     mesh_polyadr: wp.array(dtype=int),
     mesh_polymap: wp.array(dtype=int),
     mesh_polymapadr: wp.array(dtype=int),
@@ -135,6 +140,7 @@ def _collision_shim(
 ):
   _m.stat = _s
   _m.opt = _o
+  _m.callback = _cb
   _d.efc = _e
   _d.contact = _c
   _m.block_dim = block_dim
@@ -163,6 +169,7 @@ def _collision_shim(
   _m.mesh_faceadr = mesh_faceadr
   _m.mesh_graph = mesh_graph
   _m.mesh_graphadr = mesh_graphadr
+  _m.mesh_octadr = mesh_octadr
   _m.mesh_polyadr = mesh_polyadr
   _m.mesh_polymap = mesh_polymap
   _m.mesh_polymapadr = mesh_polymapadr
@@ -317,6 +324,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m.mesh_faceadr,
       m.mesh_graph,
       m.mesh_graphadr,
+      m.mesh_octadr,
       m._impl.mesh_polyadr,
       m._impl.mesh_polymap,
       m._impl.mesh_polymapadr,

mjx/mujoco/mjx/warp/forward.py

@@ -42,6 +42,10 @@ _c = mjwarp.Contact(
 _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
+_cb = mjwp_types.Callback(
+    **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
+)
+
 
 @ffi.format_args_for_warp
 def _forward_shim(
@@ -49,7 +53,7 @@ def _forward_shim(
     nworld: int,
     M_rowadr: wp.array(dtype=int),
     M_rownnz: wp.array(dtype=int),
-    actuator_acc0: wp.array(dtype=float),
+    actuator_acc0: wp.array2d(dtype=float),
     actuator_actadr: wp.array(dtype=int),
     actuator_actearly: wp.array(dtype=bool),
     actuator_actlimited: wp.array(dtype=bool),
@@ -57,7 +61,7 @@ def _forward_shim(
     actuator_actrange: wp.array2d(dtype=wp.vec2),
     actuator_biasprm: wp.array2d(dtype=mjwp_types.vec10f),
     actuator_biastype: wp.array(dtype=int),
-    actuator_cranklength: wp.array(dtype=float),
+    actuator_cranklength: wp.array2d(dtype=float),
     actuator_ctrllimited: wp.array(dtype=bool),
     actuator_ctrlrange: wp.array2d(dtype=wp.vec2),
     actuator_dynprm: wp.array2d(dtype=mjwp_types.vec10f),
@@ -67,7 +71,7 @@ def _forward_shim(
     actuator_gainprm: wp.array2d(dtype=mjwp_types.vec10f),
     actuator_gaintype: wp.array(dtype=int),
     actuator_gear: wp.array2d(dtype=wp.spatial_vector),
-    actuator_lengthrange: wp.array(dtype=wp.vec2),
+    actuator_lengthrange: wp.array2d(dtype=wp.vec2),
     actuator_trnid: wp.array(dtype=wp.vec2i),
     actuator_trntype: wp.array(dtype=int),
     actuator_trntype_body_adr: wp.array(dtype=int),
@@ -94,6 +98,7 @@ def _forward_shim(
     body_rootid: wp.array(dtype=int),
     body_subtreemass: wp.array2d(dtype=float),
     body_tree: tuple[wp.array(dtype=int), ...],
+    body_treeid: wp.array(dtype=int),
     body_weldid: wp.array(dtype=int),
     cam_bodyid: wp.array(dtype=int),
     cam_fovy: wp.array2d(dtype=float),
@@ -117,6 +122,7 @@ def _forward_shim(
     dof_parentid: wp.array(dtype=int),
     dof_solimp: wp.array2d(dtype=mjwp_types.vec5),
     dof_solref: wp.array2d(dtype=wp.vec2),
+    dof_treeid: wp.array(dtype=int),
     dof_tri_col: wp.array(dtype=int),
     dof_tri_row: wp.array(dtype=int),
     eq_connect_adr: wp.array(dtype=int),
@@ -214,6 +220,7 @@ def _forward_shim(
     mesh_normal: wp.array(dtype=wp.vec3),
     mesh_normaladr: wp.array(dtype=int),
     mesh_normalnum: wp.array(dtype=int),
+    mesh_octadr: wp.array(dtype=int),
     mesh_polyadr: wp.array(dtype=int),
     mesh_polymap: wp.array(dtype=int),
     mesh_polymapadr: wp.array(dtype=int),
@@ -253,6 +260,7 @@ def _forward_shim(
     nsensortaxel: int,
     nsite: int,
     ntendon: int,
+    ntree: int,
     nu: int,
     nv: int,
     nv_pad: int,
@@ -378,7 +386,7 @@ def _forward_shim(
     act_dot: wp.array2d(dtype=float),
     actuator_force: wp.array2d(dtype=float),
     actuator_length: wp.array2d(dtype=float),
-    actuator_moment: wp.array3d(dtype=float),
+    actuator_moment: wp.array2d(dtype=float),
     actuator_velocity: wp.array2d(dtype=float),
     cacc: wp.array2d(dtype=wp.spatial_vector),
     cam_xmat: wp.array2d(dtype=wp.mat33),
@@ -393,7 +401,7 @@ def _forward_shim(
     cvel: wp.array2d(dtype=wp.spatial_vector),
     energy: wp.array(dtype=wp.vec2),
     eq_active: wp.array2d(dtype=bool),
-    flexedge_J: wp.array3d(dtype=float),
+    flexedge_J: wp.array2d(dtype=float),
     flexedge_length: wp.array2d(dtype=float),
     flexedge_velocity: wp.array2d(dtype=float),
     flexvert_xpos: wp.array2d(dtype=wp.vec3),
@@ -403,11 +411,15 @@ def _forward_shim(
     light_xpos: wp.array2d(dtype=wp.vec3),
     mocap_pos: wp.array2d(dtype=wp.vec3),
     mocap_quat: wp.array2d(dtype=wp.quat),
+    moment_colind: wp.array2d(dtype=int),
+    moment_rowadr: wp.array2d(dtype=int),
+    moment_rownnz: wp.array2d(dtype=int),
     nacon: wp.array(dtype=int),
     ncollision: wp.array(dtype=int),
     ne: wp.array(dtype=int),
     nefc: wp.array(dtype=int),
     nf: wp.array(dtype=int),
+    nisland: wp.array(dtype=int),
     nl: wp.array(dtype=int),
     qLD: wp.array3d(dtype=float),
     qLDiagInv: wp.array2d(dtype=float),
@@ -440,6 +452,7 @@ def _forward_shim(
     ten_wrapadr: wp.array2d(dtype=int),
     ten_wrapnum: wp.array2d(dtype=int),
     time: wp.array(dtype=float),
+    tree_island: wp.array2d(dtype=int),
     wrap_obj: wp.array2d(dtype=wp.vec2i),
     wrap_xpos: wp.array2d(dtype=wp.spatial_vector),
     xanchor: wp.array2d(dtype=wp.vec3),
@@ -466,6 +479,9 @@ def _forward_shim(
     contact__worldid: wp.array(dtype=int),
     efc__D: wp.array2d(dtype=float),
     efc__J: wp.array3d(dtype=float),
+    efc__J_colind: wp.array3d(dtype=int),
+    efc__J_rowadr: wp.array2d(dtype=int),
+    efc__J_rownnz: wp.array2d(dtype=int),
     efc__Ma: wp.array2d(dtype=float),
     efc__aref: wp.array2d(dtype=float),
     efc__force: wp.array2d(dtype=float),
@@ -479,6 +495,7 @@ def _forward_shim(
 ):
   _m.stat = _s
   _m.opt = _o
+  _m.callback = _cb
   _d.efc = _e
   _d.contact = _c
   _m.M_rowadr = M_rowadr
@@ -528,6 +545,7 @@ def _forward_shim(
   _m.body_rootid = body_rootid
   _m.body_subtreemass = body_subtreemass
   _m.body_tree = body_tree
+  _m.body_treeid = body_treeid
   _m.body_weldid = body_weldid
   _m.cam_bodyid = cam_bodyid
   _m.cam_fovy = cam_fovy
@@ -551,6 +569,7 @@ def _forward_shim(
   _m.dof_parentid = dof_parentid
   _m.dof_solimp = dof_solimp
   _m.dof_solref = dof_solref
+  _m.dof_treeid = dof_treeid
   _m.dof_tri_col = dof_tri_col
   _m.dof_tri_row = dof_tri_row
   _m.eq_connect_adr = eq_connect_adr
@@ -648,6 +667,7 @@ def _forward_shim(
   _m.mesh_normal = mesh_normal
   _m.mesh_normaladr = mesh_normaladr
   _m.mesh_normalnum = mesh_normalnum
+  _m.mesh_octadr = mesh_octadr
   _m.mesh_polyadr = mesh_polyadr
   _m.mesh_polymap = mesh_polymap
   _m.mesh_polymapadr = mesh_polymapadr
@@ -687,6 +707,7 @@ def _forward_shim(
   _m.nsensortaxel = nsensortaxel
   _m.nsite = nsite
   _m.ntendon = ntendon
+  _m.ntree = ntree
   _m.nu = nu
   _m.nv = nv
   _m.nv_pad = nv_pad
@@ -837,6 +858,9 @@ def _forward_shim(
   _d.cvel = cvel
   _d.efc.D = efc__D
   _d.efc.J = efc__J
+  _d.efc.J_colind = efc__J_colind
+  _d.efc.J_rowadr = efc__J_rowadr
+  _d.efc.J_rownnz = efc__J_rownnz
   _d.efc.Ma = efc__Ma
   _d.efc.aref = efc__aref
   _d.efc.force = efc__force
@@ -859,6 +883,9 @@ def _forward_shim(
   _d.light_xpos = light_xpos
   _d.mocap_pos = mocap_pos
   _d.mocap_quat = mocap_quat
+  _d.moment_colind = moment_colind
+  _d.moment_rowadr = moment_rowadr
+  _d.moment_rownnz = moment_rownnz
   _d.naccdmax = naccdmax
   _d.nacon = nacon
   _d.naconmax = naconmax
@@ -866,6 +893,7 @@ def _forward_shim(
   _d.ne = ne
   _d.nefc = nefc
   _d.nf = nf
+  _d.nisland = nisland
   _d.njmax = njmax
   _d.nl = nl
   _d.qLD = qLD
@@ -899,6 +927,7 @@ def _forward_shim(
   _d.ten_wrapadr = ten_wrapadr
   _d.ten_wrapnum = ten_wrapnum
   _d.time = time
+  _d.tree_island = tree_island
   _d.wrap_obj = wrap_obj
   _d.wrap_xpos = wrap_xpos
   _d.xanchor = xanchor
@@ -939,11 +968,15 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'geom_xpos': d.geom_xpos.shape,
       'light_xdir': d._impl.light_xdir.shape,
       'light_xpos': d._impl.light_xpos.shape,
+      'moment_colind': d._impl.moment_colind.shape,
+      'moment_rowadr': d._impl.moment_rowadr.shape,
+      'moment_rownnz': d._impl.moment_rownnz.shape,
       'nacon': d._impl.nacon.shape,
       'ncollision': d._impl.ncollision.shape,
       'ne': d._impl.ne.shape,
       'nefc': d._impl.nefc.shape,
       'nf': d._impl.nf.shape,
+      'nisland': d._impl.nisland.shape,
       'nl': d._impl.nl.shape,
       'qLD': d._impl.qLD.shape,
       'qLDiagInv': d._impl.qLDiagInv.shape,
@@ -972,6 +1005,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'ten_velocity': d._impl.ten_velocity.shape,
       'ten_wrapadr': d._impl.ten_wrapadr.shape,
       'ten_wrapnum': d._impl.ten_wrapnum.shape,
+      'tree_island': d._impl.tree_island.shape,
       'wrap_obj': d._impl.wrap_obj.shape,
       'wrap_xpos': d._impl.wrap_xpos.shape,
       'xanchor': d.xanchor.shape,
@@ -997,6 +1031,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'contact__worldid': d._impl.contact__worldid.shape,
       'efc__D': d._impl.efc__D.shape,
       'efc__J': d._impl.efc__J.shape,
+      'efc__J_colind': d._impl.efc__J_colind.shape,
+      'efc__J_rowadr': d._impl.efc__J_rowadr.shape,
+      'efc__J_rownnz': d._impl.efc__J_rownnz.shape,
       'efc__Ma': d._impl.efc__Ma.shape,
       'efc__aref': d._impl.efc__aref.shape,
       'efc__force': d._impl.efc__force.shape,
@@ -1010,7 +1047,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _forward_shim,
-      num_outputs=92,
+      num_outputs=100,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames=set([
@@ -1038,11 +1075,15 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'geom_xpos',
           'light_xdir',
           'light_xpos',
+          'moment_colind',
+          'moment_rowadr',
+          'moment_rownnz',
           'nacon',
           'ncollision',
           'ne',
           'nefc',
           'nf',
+          'nisland',
           'nl',
           'qLD',
           'qLDiagInv',
@@ -1071,6 +1112,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'ten_velocity',
           'ten_wrapadr',
           'ten_wrapnum',
+          'tree_island',
           'wrap_obj',
           'wrap_xpos',
           'xanchor',
@@ -1096,6 +1138,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'contact__worldid',
           'efc__D',
           'efc__J',
+          'efc__J_colind',
+          'efc__J_rowadr',
+          'efc__J_rownnz',
           'efc__Ma',
           'efc__aref',
           'efc__force',
@@ -1113,6 +1158,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'actuator_acc0',
           'actuator_actrange',
           'actuator_biasprm',
+          'actuator_cranklength',
           'actuator_ctrlrange',
           'actuator_dynprm',
           'actuator_force',
@@ -1120,6 +1166,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'actuator_gainprm',
           'actuator_gear',
           'actuator_length',
+          'actuator_lengthrange',
           'body_gravcomp',
           'body_inertia',
           'body_invweight0',
@@ -1330,6 +1377,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m.body_rootid,
       m.body_subtreemass,
       m._impl.body_tree,
+      m.body_treeid,
       m.body_weldid,
       m.cam_bodyid,
       m.cam_fovy,
@@ -1353,6 +1401,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m.dof_parentid,
       m.dof_solimp,
       m.dof_solref,
+      m.dof_treeid,
       m._impl.dof_tri_col,
       m._impl.dof_tri_row,
       m._impl.eq_connect_adr,
@@ -1450,6 +1499,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m.mesh_normal,
       m.mesh_normaladr,
       m.mesh_normalnum,
+      m.mesh_octadr,
       m._impl.mesh_polyadr,
       m._impl.mesh_polymap,
       m._impl.mesh_polymapadr,
@@ -1489,6 +1539,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.nsensortaxel,
       m.nsite,
       m.ntendon,
+      m._impl.ntree,
       m.nu,
       m.nv,
       m._impl.nv_pad,
@@ -1638,11 +1689,15 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d._impl.light_xpos,
       d.mocap_pos,
       d.mocap_quat,
+      d._impl.moment_colind,
+      d._impl.moment_rowadr,
+      d._impl.moment_rownnz,
       d._impl.nacon,
       d._impl.ncollision,
       d._impl.ne,
       d._impl.nefc,
       d._impl.nf,
+      d._impl.nisland,
       d._impl.nl,
       d._impl.qLD,
       d._impl.qLDiagInv,
@@ -1675,6 +1730,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d._impl.ten_wrapadr,
       d._impl.ten_wrapnum,
       d.time,
+      d._impl.tree_island,
       d._impl.wrap_obj,
       d._impl.wrap_xpos,
       d.xanchor,
@@ -1701,6 +1757,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d._impl.contact__worldid,
       d._impl.efc__D,
       d._impl.efc__J,
+      d._impl.efc__J_colind,
+      d._impl.efc__J_rowadr,
+      d._impl.efc__J_rownnz,
       d._impl.efc__Ma,
       d._impl.efc__aref,
       d._impl.efc__force,
@@ -1737,74 +1796,82 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'geom_xpos': out[21],
       '_impl.light_xdir': out[22],
       '_impl.light_xpos': out[23],
-      '_impl.nacon': out[24],
+      '_impl.moment_colind': out[24],
-      '_impl.ncollision': out[25],
+      '_impl.moment_rowadr': out[25],
-      '_impl.ne': out[26],
+      '_impl.moment_rownnz': out[26],
-      '_impl.nefc': out[27],
+      '_impl.nacon': out[27],
-      '_impl.nf': out[28],
+      '_impl.ncollision': out[28],
-      '_impl.nl': out[29],
+      '_impl.ne': out[29],
-      '_impl.qLD': out[30],
+      '_impl.nefc': out[30],
-      '_impl.qLDiagInv': out[31],
+      '_impl.nf': out[31],
-      '_impl.qM': out[32],
+      '_impl.nisland': out[32],
-      'qacc': out[33],
+      '_impl.nl': out[33],
-      'qacc_smooth': out[34],
+      '_impl.qLD': out[34],
-      'qfrc_actuator': out[35],
+      '_impl.qLDiagInv': out[35],
-      'qfrc_bias': out[36],
+      '_impl.qM': out[36],
-      'qfrc_constraint': out[37],
+      'qacc': out[37],
-      '_impl.qfrc_damper': out[38],
+      'qacc_smooth': out[38],
-      'qfrc_fluid': out[39],
+      'qfrc_actuator': out[39],
-      'qfrc_gravcomp': out[40],
+      'qfrc_bias': out[40],
-      'qfrc_passive': out[41],
+      'qfrc_constraint': out[41],
-      'qfrc_smooth': out[42],
+      '_impl.qfrc_damper': out[42],
-      '_impl.qfrc_spring': out[43],
+      'qfrc_fluid': out[43],
-      'qvel': out[44],
+      'qfrc_gravcomp': out[44],
-      'sensordata': out[45],
+      'qfrc_passive': out[45],
-      'site_xmat': out[46],
+      'qfrc_smooth': out[46],
-      'site_xpos': out[47],
+      '_impl.qfrc_spring': out[47],
-      '_impl.solver_niter': out[48],
+      'qvel': out[48],
-      '_impl.subtree_angmom': out[49],
+      'sensordata': out[49],
-      'subtree_com': out[50],
+      'site_xmat': out[50],
-      '_impl.subtree_linvel': out[51],
+      'site_xpos': out[51],
-      '_impl.ten_J': out[52],
+      '_impl.solver_niter': out[52],
-      'ten_length': out[53],
+      '_impl.subtree_angmom': out[53],
-      '_impl.ten_velocity': out[54],
+      'subtree_com': out[54],
-      '_impl.ten_wrapadr': out[55],
+      '_impl.subtree_linvel': out[55],
-      '_impl.ten_wrapnum': out[56],
+      '_impl.ten_J': out[56],
-      '_impl.wrap_obj': out[57],
+      'ten_length': out[57],
-      '_impl.wrap_xpos': out[58],
+      '_impl.ten_velocity': out[58],
-      'xanchor': out[59],
+      '_impl.ten_wrapadr': out[59],
-      'xaxis': out[60],
+      '_impl.ten_wrapnum': out[60],
-      'ximat': out[61],
+      '_impl.tree_island': out[61],
-      'xipos': out[62],
+      '_impl.wrap_obj': out[62],
-      'xmat': out[63],
+      '_impl.wrap_xpos': out[63],
-      'xpos': out[64],
+      'xanchor': out[64],
-      'xquat': out[65],
+      'xaxis': out[65],
-      '_impl.contact__dim': out[66],
+      'ximat': out[66],
-      '_impl.contact__dist': out[67],
+      'xipos': out[67],
-      '_impl.contact__efc_address': out[68],
+      'xmat': out[68],
-      '_impl.contact__frame': out[69],
+      'xpos': out[69],
-      '_impl.contact__friction': out[70],
+      'xquat': out[70],
-      '_impl.contact__geom': out[71],
+      '_impl.contact__dim': out[71],
-      '_impl.contact__geomcollisionid': out[72],
+      '_impl.contact__dist': out[72],
-      '_impl.contact__includemargin': out[73],
+      '_impl.contact__efc_address': out[73],
-      '_impl.contact__pos': out[74],
+      '_impl.contact__frame': out[74],
-      '_impl.contact__solimp': out[75],
+      '_impl.contact__friction': out[75],
-      '_impl.contact__solref': out[76],
+      '_impl.contact__geom': out[76],
-      '_impl.contact__solreffriction': out[77],
+      '_impl.contact__geomcollisionid': out[77],
-      '_impl.contact__type': out[78],
+      '_impl.contact__includemargin': out[78],
-      '_impl.contact__worldid': out[79],
+      '_impl.contact__pos': out[79],
-      '_impl.efc__D': out[80],
+      '_impl.contact__solimp': out[80],
-      '_impl.efc__J': out[81],
+      '_impl.contact__solref': out[81],
-      '_impl.efc__Ma': out[82],
+      '_impl.contact__solreffriction': out[82],
-      '_impl.efc__aref': out[83],
+      '_impl.contact__type': out[83],
-      '_impl.efc__force': out[84],
+      '_impl.contact__worldid': out[84],
-      '_impl.efc__frictionloss': out[85],
+      '_impl.efc__D': out[85],
-      '_impl.efc__id': out[86],
+      '_impl.efc__J': out[86],
-      '_impl.efc__margin': out[87],
+      '_impl.efc__J_colind': out[87],
-      '_impl.efc__pos': out[88],
+      '_impl.efc__J_rowadr': out[88],
-      '_impl.efc__state': out[89],
+      '_impl.efc__J_rownnz': out[89],
-      '_impl.efc__type': out[90],
+      '_impl.efc__Ma': out[90],
-      '_impl.efc__vel': out[91],
+      '_impl.efc__aref': out[91],
+      '_impl.efc__force': out[92],
+      '_impl.efc__frictionloss': out[93],
+      '_impl.efc__id': out[94],
+      '_impl.efc__margin': out[95],
+      '_impl.efc__pos': out[96],
+      '_impl.efc__state': out[97],
+      '_impl.efc__type': out[98],
+      '_impl.efc__vel': out[99],
   })
   return d
 
@@ -1828,7 +1895,7 @@ def _step_shim(
     nworld: int,
     M_rowadr: wp.array(dtype=int),
     M_rownnz: wp.array(dtype=int),
-    actuator_acc0: wp.array(dtype=float),
+    actuator_acc0: wp.array2d(dtype=float),
     actuator_actadr: wp.array(dtype=int),
     actuator_actearly: wp.array(dtype=bool),
     actuator_actlimited: wp.array(dtype=bool),
@@ -1836,7 +1903,7 @@ def _step_shim(
     actuator_actrange: wp.array2d(dtype=wp.vec2),
     actuator_biasprm: wp.array2d(dtype=mjwp_types.vec10f),
     actuator_biastype: wp.array(dtype=int),
-    actuator_cranklength: wp.array(dtype=float),
+    actuator_cranklength: wp.array2d(dtype=float),
     actuator_ctrllimited: wp.array(dtype=bool),
     actuator_ctrlrange: wp.array2d(dtype=wp.vec2),
     actuator_dynprm: wp.array2d(dtype=mjwp_types.vec10f),
@@ -1846,7 +1913,7 @@ def _step_shim(
     actuator_gainprm: wp.array2d(dtype=mjwp_types.vec10f),
     actuator_gaintype: wp.array(dtype=int),
     actuator_gear: wp.array2d(dtype=wp.spatial_vector),
-    actuator_lengthrange: wp.array(dtype=wp.vec2),
+    actuator_lengthrange: wp.array2d(dtype=wp.vec2),
     actuator_trnid: wp.array(dtype=wp.vec2i),
     actuator_trntype: wp.array(dtype=int),
     actuator_trntype_body_adr: wp.array(dtype=int),
@@ -1873,6 +1940,7 @@ def _step_shim(
     body_rootid: wp.array(dtype=int),
     body_subtreemass: wp.array2d(dtype=float),
     body_tree: tuple[wp.array(dtype=int), ...],
+    body_treeid: wp.array(dtype=int),
     body_weldid: wp.array(dtype=int),
     cam_bodyid: wp.array(dtype=int),
     cam_fovy: wp.array2d(dtype=float),
@@ -1896,6 +1964,7 @@ def _step_shim(
     dof_parentid: wp.array(dtype=int),
     dof_solimp: wp.array2d(dtype=mjwp_types.vec5),
     dof_solref: wp.array2d(dtype=wp.vec2),
+    dof_treeid: wp.array(dtype=int),
     dof_tri_col: wp.array(dtype=int),
     dof_tri_row: wp.array(dtype=int),
     eq_connect_adr: wp.array(dtype=int),
@@ -1993,6 +2062,7 @@ def _step_shim(
     mesh_normal: wp.array(dtype=wp.vec3),
     mesh_normaladr: wp.array(dtype=int),
     mesh_normalnum: wp.array(dtype=int),
+    mesh_octadr: wp.array(dtype=int),
     mesh_polyadr: wp.array(dtype=int),
     mesh_polymap: wp.array(dtype=int),
     mesh_polymapadr: wp.array(dtype=int),
@@ -2033,6 +2103,7 @@ def _step_shim(
     nsensortaxel: int,
     nsite: int,
     ntendon: int,
+    ntree: int,
     nu: int,
     nv: int,
     nv_pad: int,
@@ -2159,7 +2230,7 @@ def _step_shim(
     act_dot: wp.array2d(dtype=float),
     actuator_force: wp.array2d(dtype=float),
     actuator_length: wp.array2d(dtype=float),
-    actuator_moment: wp.array3d(dtype=float),
+    actuator_moment: wp.array2d(dtype=float),
     actuator_velocity: wp.array2d(dtype=float),
     cacc: wp.array2d(dtype=wp.spatial_vector),
     cam_xmat: wp.array2d(dtype=wp.mat33),
@@ -2174,7 +2245,7 @@ def _step_shim(
     cvel: wp.array2d(dtype=wp.spatial_vector),
     energy: wp.array(dtype=wp.vec2),
     eq_active: wp.array2d(dtype=bool),
-    flexedge_J: wp.array3d(dtype=float),
+    flexedge_J: wp.array2d(dtype=float),
     flexedge_length: wp.array2d(dtype=float),
     flexedge_velocity: wp.array2d(dtype=float),
     flexvert_xpos: wp.array2d(dtype=wp.vec3),
@@ -2184,11 +2255,15 @@ def _step_shim(
     light_xpos: wp.array2d(dtype=wp.vec3),
     mocap_pos: wp.array2d(dtype=wp.vec3),
     mocap_quat: wp.array2d(dtype=wp.quat),
+    moment_colind: wp.array2d(dtype=int),
+    moment_rowadr: wp.array2d(dtype=int),
+    moment_rownnz: wp.array2d(dtype=int),
     nacon: wp.array(dtype=int),
     ncollision: wp.array(dtype=int),
     ne: wp.array(dtype=int),
     nefc: wp.array(dtype=int),
     nf: wp.array(dtype=int),
+    nisland: wp.array(dtype=int),
     nl: wp.array(dtype=int),
     qLD: wp.array3d(dtype=float),
     qLDiagInv: wp.array2d(dtype=float),
@@ -2221,6 +2296,7 @@ def _step_shim(
     ten_wrapadr: wp.array2d(dtype=int),
     ten_wrapnum: wp.array2d(dtype=int),
     time: wp.array(dtype=float),
+    tree_island: wp.array2d(dtype=int),
     wrap_obj: wp.array2d(dtype=wp.vec2i),
     wrap_xpos: wp.array2d(dtype=wp.spatial_vector),
     xanchor: wp.array2d(dtype=wp.vec3),
@@ -2247,6 +2323,9 @@ def _step_shim(
     contact__worldid: wp.array(dtype=int),
     efc__D: wp.array2d(dtype=float),
     efc__J: wp.array3d(dtype=float),
+    efc__J_colind: wp.array3d(dtype=int),
+    efc__J_rowadr: wp.array2d(dtype=int),
+    efc__J_rownnz: wp.array2d(dtype=int),
     efc__Ma: wp.array2d(dtype=float),
     efc__aref: wp.array2d(dtype=float),
     efc__force: wp.array2d(dtype=float),
@@ -2260,6 +2339,7 @@ def _step_shim(
 ):
   _m.stat = _s
   _m.opt = _o
+  _m.callback = _cb
   _d.efc = _e
   _d.contact = _c
   _m.M_rowadr = M_rowadr
@@ -2309,6 +2389,7 @@ def _step_shim(
   _m.body_rootid = body_rootid
   _m.body_subtreemass = body_subtreemass
   _m.body_tree = body_tree
+  _m.body_treeid = body_treeid
   _m.body_weldid = body_weldid
   _m.cam_bodyid = cam_bodyid
   _m.cam_fovy = cam_fovy
@@ -2332,6 +2413,7 @@ def _step_shim(
   _m.dof_parentid = dof_parentid
   _m.dof_solimp = dof_solimp
   _m.dof_solref = dof_solref
+  _m.dof_treeid = dof_treeid
   _m.dof_tri_col = dof_tri_col
   _m.dof_tri_row = dof_tri_row
   _m.eq_connect_adr = eq_connect_adr
@@ -2429,6 +2511,7 @@ def _step_shim(
   _m.mesh_normal = mesh_normal
   _m.mesh_normaladr = mesh_normaladr
   _m.mesh_normalnum = mesh_normalnum
+  _m.mesh_octadr = mesh_octadr
   _m.mesh_polyadr = mesh_polyadr
   _m.mesh_polymap = mesh_polymap
   _m.mesh_polymapadr = mesh_polymapadr
@@ -2469,6 +2552,7 @@ def _step_shim(
   _m.nsensortaxel = nsensortaxel
   _m.nsite = nsite
   _m.ntendon = ntendon
+  _m.ntree = ntree
   _m.nu = nu
   _m.nv = nv
   _m.nv_pad = nv_pad
@@ -2620,6 +2704,9 @@ def _step_shim(
   _d.cvel = cvel
   _d.efc.D = efc__D
   _d.efc.J = efc__J
+  _d.efc.J_colind = efc__J_colind
+  _d.efc.J_rowadr = efc__J_rowadr
+  _d.efc.J_rownnz = efc__J_rownnz
   _d.efc.Ma = efc__Ma
   _d.efc.aref = efc__aref
   _d.efc.force = efc__force
@@ -2642,6 +2729,9 @@ def _step_shim(
   _d.light_xpos = light_xpos
   _d.mocap_pos = mocap_pos
   _d.mocap_quat = mocap_quat
+  _d.moment_colind = moment_colind
+  _d.moment_rowadr = moment_rowadr
+  _d.moment_rownnz = moment_rownnz
   _d.naccdmax = naccdmax
   _d.nacon = nacon
   _d.naconmax = naconmax
@@ -2649,6 +2739,7 @@ def _step_shim(
   _d.ne = ne
   _d.nefc = nefc
   _d.nf = nf
+  _d.nisland = nisland
   _d.njmax = njmax
   _d.nl = nl
   _d.qLD = qLD
@@ -2682,6 +2773,7 @@ def _step_shim(
   _d.ten_wrapadr = ten_wrapadr
   _d.ten_wrapnum = ten_wrapnum
   _d.time = time
+  _d.tree_island = tree_island
   _d.wrap_obj = wrap_obj
   _d.wrap_xpos = wrap_xpos
   _d.xanchor = xanchor
@@ -2723,11 +2815,15 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'geom_xpos': d.geom_xpos.shape,
       'light_xdir': d._impl.light_xdir.shape,
       'light_xpos': d._impl.light_xpos.shape,
+      'moment_colind': d._impl.moment_colind.shape,
+      'moment_rowadr': d._impl.moment_rowadr.shape,
+      'moment_rownnz': d._impl.moment_rownnz.shape,
       'nacon': d._impl.nacon.shape,
       'ncollision': d._impl.ncollision.shape,
       'ne': d._impl.ne.shape,
       'nefc': d._impl.nefc.shape,
       'nf': d._impl.nf.shape,
+      'nisland': d._impl.nisland.shape,
       'nl': d._impl.nl.shape,
       'qLD': d._impl.qLD.shape,
       'qLDiagInv': d._impl.qLDiagInv.shape,
@@ -2759,6 +2855,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'ten_wrapadr': d._impl.ten_wrapadr.shape,
       'ten_wrapnum': d._impl.ten_wrapnum.shape,
       'time': d.time.shape,
+      'tree_island': d._impl.tree_island.shape,
       'wrap_obj': d._impl.wrap_obj.shape,
       'wrap_xpos': d._impl.wrap_xpos.shape,
       'xanchor': d.xanchor.shape,
@@ -2784,6 +2881,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'contact__worldid': d._impl.contact__worldid.shape,
       'efc__D': d._impl.efc__D.shape,
       'efc__J': d._impl.efc__J.shape,
+      'efc__J_colind': d._impl.efc__J_colind.shape,
+      'efc__J_rowadr': d._impl.efc__J_rowadr.shape,
+      'efc__J_rownnz': d._impl.efc__J_rownnz.shape,
       'efc__Ma': d._impl.efc__Ma.shape,
       'efc__aref': d._impl.efc__aref.shape,
       'efc__force': d._impl.efc__force.shape,
@@ -2797,7 +2897,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _step_shim,
-      num_outputs=96,
+      num_outputs=104,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames=set([
@@ -2826,11 +2926,15 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'geom_xpos',
           'light_xdir',
           'light_xpos',
+          'moment_colind',
+          'moment_rowadr',
+          'moment_rownnz',
           'nacon',
           'ncollision',
           'ne',
           'nefc',
           'nf',
+          'nisland',
           'nl',
           'qLD',
           'qLDiagInv',
@@ -2862,6 +2966,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'ten_wrapadr',
           'ten_wrapnum',
           'time',
+          'tree_island',
           'wrap_obj',
           'wrap_xpos',
           'xanchor',
@@ -2887,6 +2992,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'contact__worldid',
           'efc__D',
           'efc__J',
+          'efc__J_colind',
+          'efc__J_rowadr',
+          'efc__J_rownnz',
           'efc__Ma',
           'efc__aref',
           'efc__force',
@@ -2904,6 +3012,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'actuator_acc0',
           'actuator_actrange',
           'actuator_biasprm',
+          'actuator_cranklength',
           'actuator_ctrlrange',
           'actuator_dynprm',
           'actuator_force',
@@ -2911,6 +3020,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'actuator_gainprm',
           'actuator_gear',
           'actuator_length',
+          'actuator_lengthrange',
           'body_gravcomp',
           'body_inertia',
           'body_invweight0',
@@ -3125,6 +3235,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.body_rootid,
       m.body_subtreemass,
       m._impl.body_tree,
+      m.body_treeid,
       m.body_weldid,
       m.cam_bodyid,
       m.cam_fovy,
@@ -3148,6 +3259,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.dof_parentid,
       m.dof_solimp,
       m.dof_solref,
+      m.dof_treeid,
       m._impl.dof_tri_col,
       m._impl.dof_tri_row,
       m._impl.eq_connect_adr,
@@ -3245,6 +3357,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.mesh_normal,
       m.mesh_normaladr,
       m.mesh_normalnum,
+      m.mesh_octadr,
       m._impl.mesh_polyadr,
       m._impl.mesh_polymap,
       m._impl.mesh_polymapadr,
@@ -3285,6 +3398,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.nsensortaxel,
       m.nsite,
       m.ntendon,
+      m._impl.ntree,
       m.nu,
       m.nv,
       m._impl.nv_pad,
@@ -3435,11 +3549,15 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.light_xpos,
       d.mocap_pos,
       d.mocap_quat,
+      d._impl.moment_colind,
+      d._impl.moment_rowadr,
+      d._impl.moment_rownnz,
       d._impl.nacon,
       d._impl.ncollision,
       d._impl.ne,
       d._impl.nefc,
       d._impl.nf,
+      d._impl.nisland,
       d._impl.nl,
       d._impl.qLD,
       d._impl.qLDiagInv,
@@ -3472,6 +3590,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.ten_wrapadr,
       d._impl.ten_wrapnum,
       d.time,
+      d._impl.tree_island,
       d._impl.wrap_obj,
       d._impl.wrap_xpos,
       d.xanchor,
@@ -3498,6 +3617,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.contact__worldid,
       d._impl.efc__D,
       d._impl.efc__J,
+      d._impl.efc__J_colind,
+      d._impl.efc__J_rowadr,
+      d._impl.efc__J_rownnz,
       d._impl.efc__Ma,
       d._impl.efc__aref,
       d._impl.efc__force,
@@ -3535,77 +3657,85 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'geom_xpos': out[22],
       '_impl.light_xdir': out[23],
       '_impl.light_xpos': out[24],
-      '_impl.nacon': out[25],
+      '_impl.moment_colind': out[25],
-      '_impl.ncollision': out[26],
+      '_impl.moment_rowadr': out[26],
-      '_impl.ne': out[27],
+      '_impl.moment_rownnz': out[27],
-      '_impl.nefc': out[28],
+      '_impl.nacon': out[28],
-      '_impl.nf': out[29],
+      '_impl.ncollision': out[29],
-      '_impl.nl': out[30],
+      '_impl.ne': out[30],
-      '_impl.qLD': out[31],
+      '_impl.nefc': out[31],
-      '_impl.qLDiagInv': out[32],
+      '_impl.nf': out[32],
-      '_impl.qM': out[33],
+      '_impl.nisland': out[33],
-      'qacc': out[34],
+      '_impl.nl': out[34],
-      'qacc_smooth': out[35],
+      '_impl.qLD': out[35],
-      'qacc_warmstart': out[36],
+      '_impl.qLDiagInv': out[36],
-      'qfrc_actuator': out[37],
+      '_impl.qM': out[37],
-      'qfrc_bias': out[38],
+      'qacc': out[38],
-      'qfrc_constraint': out[39],
+      'qacc_smooth': out[39],
-      '_impl.qfrc_damper': out[40],
+      'qacc_warmstart': out[40],
-      'qfrc_fluid': out[41],
+      'qfrc_actuator': out[41],
-      'qfrc_gravcomp': out[42],
+      'qfrc_bias': out[42],
-      'qfrc_passive': out[43],
+      'qfrc_constraint': out[43],
-      'qfrc_smooth': out[44],
+      '_impl.qfrc_damper': out[44],
-      '_impl.qfrc_spring': out[45],
+      'qfrc_fluid': out[45],
-      'qpos': out[46],
+      'qfrc_gravcomp': out[46],
-      'qvel': out[47],
+      'qfrc_passive': out[47],
-      'sensordata': out[48],
+      'qfrc_smooth': out[48],
-      'site_xmat': out[49],
+      '_impl.qfrc_spring': out[49],
-      'site_xpos': out[50],
+      'qpos': out[50],
-      '_impl.solver_niter': out[51],
+      'qvel': out[51],
-      '_impl.subtree_angmom': out[52],
+      'sensordata': out[52],
-      'subtree_com': out[53],
+      'site_xmat': out[53],
-      '_impl.subtree_linvel': out[54],
+      'site_xpos': out[54],
-      '_impl.ten_J': out[55],
+      '_impl.solver_niter': out[55],
-      'ten_length': out[56],
+      '_impl.subtree_angmom': out[56],
-      '_impl.ten_velocity': out[57],
+      'subtree_com': out[57],
-      '_impl.ten_wrapadr': out[58],
+      '_impl.subtree_linvel': out[58],
-      '_impl.ten_wrapnum': out[59],
+      '_impl.ten_J': out[59],
-      'time': out[60],
+      'ten_length': out[60],
-      '_impl.wrap_obj': out[61],
+      '_impl.ten_velocity': out[61],
-      '_impl.wrap_xpos': out[62],
+      '_impl.ten_wrapadr': out[62],
-      'xanchor': out[63],
+      '_impl.ten_wrapnum': out[63],
-      'xaxis': out[64],
+      'time': out[64],
-      'ximat': out[65],
+      '_impl.tree_island': out[65],
-      'xipos': out[66],
+      '_impl.wrap_obj': out[66],
-      'xmat': out[67],
+      '_impl.wrap_xpos': out[67],
-      'xpos': out[68],
+      'xanchor': out[68],
-      'xquat': out[69],
+      'xaxis': out[69],
-      '_impl.contact__dim': out[70],
+      'ximat': out[70],
-      '_impl.contact__dist': out[71],
+      'xipos': out[71],
-      '_impl.contact__efc_address': out[72],
+      'xmat': out[72],
-      '_impl.contact__frame': out[73],
+      'xpos': out[73],
-      '_impl.contact__friction': out[74],
+      'xquat': out[74],
-      '_impl.contact__geom': out[75],
+      '_impl.contact__dim': out[75],
-      '_impl.contact__geomcollisionid': out[76],
+      '_impl.contact__dist': out[76],
-      '_impl.contact__includemargin': out[77],
+      '_impl.contact__efc_address': out[77],
-      '_impl.contact__pos': out[78],
+      '_impl.contact__frame': out[78],
-      '_impl.contact__solimp': out[79],
+      '_impl.contact__friction': out[79],
-      '_impl.contact__solref': out[80],
+      '_impl.contact__geom': out[80],
-      '_impl.contact__solreffriction': out[81],
+      '_impl.contact__geomcollisionid': out[81],
-      '_impl.contact__type': out[82],
+      '_impl.contact__includemargin': out[82],
-      '_impl.contact__worldid': out[83],
+      '_impl.contact__pos': out[83],
-      '_impl.efc__D': out[84],
+      '_impl.contact__solimp': out[84],
-      '_impl.efc__J': out[85],
+      '_impl.contact__solref': out[85],
-      '_impl.efc__Ma': out[86],
+      '_impl.contact__solreffriction': out[86],
-      '_impl.efc__aref': out[87],
+      '_impl.contact__type': out[87],
-      '_impl.efc__force': out[88],
+      '_impl.contact__worldid': out[88],
-      '_impl.efc__frictionloss': out[89],
+      '_impl.efc__D': out[89],
-      '_impl.efc__id': out[90],
+      '_impl.efc__J': out[90],
-      '_impl.efc__margin': out[91],
+      '_impl.efc__J_colind': out[91],
-      '_impl.efc__pos': out[92],
+      '_impl.efc__J_rowadr': out[92],
-      '_impl.efc__state': out[93],
+      '_impl.efc__J_rownnz': out[93],
-      '_impl.efc__type': out[94],
+      '_impl.efc__Ma': out[94],
-      '_impl.efc__vel': out[95],
+      '_impl.efc__aref': out[95],
+      '_impl.efc__force': out[96],
+      '_impl.efc__frictionloss': out[97],
+      '_impl.efc__id': out[98],
+      '_impl.efc__margin': out[99],
+      '_impl.efc__pos': out[100],
+      '_impl.efc__state': out[101],
+      '_impl.efc__type': out[102],
+      '_impl.efc__vel': out[103],
   })
   return d
 

mjx/mujoco/mjx/warp/forward_test.py

@@ -181,7 +181,15 @@ class ForwardTest(parameterized.TestCase):
           d.moment_rowadr,
           d.moment_colind,
       )
-      tu.assert_eq(dx._impl.actuator_moment, actuator_moment, 'actuator_moment')
+      warp_actuator_moment = np.zeros((m.nu, m.nv))
+      mujoco.mju_sparse2dense(
+          warp_actuator_moment,
+          np.asarray(dx._impl.actuator_moment),
+          np.asarray(dx._impl.moment_rownnz),
+          np.asarray(dx._impl.moment_rowadr),
+          np.asarray(dx._impl.moment_colind),
+      )
+      tu.assert_eq(warp_actuator_moment, actuator_moment, 'actuator_moment')
 
       # fwd_velocity
       tu.assert_attr_eq(dx._impl, d, 'actuator_velocity')

mjx/mujoco/mjx/warp/render.py

@@ -44,6 +44,9 @@ _c = mjwarp.Contact(
 _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
+_cb = mjwp_types.Callback(
+    **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
+)
 
 @ffi.format_args_for_warp
 def _render_shim(
@@ -81,6 +84,7 @@ def _render_shim(
 ):
   _m.stat = _s
   _m.opt = _o
+  _m.callback = _cb
   _d.efc = _e
   _d.contact = _c
   _m.cam_fovy = cam_fovy

mjx/mujoco/mjx/warp/smooth.py

@@ -42,6 +42,9 @@ _c = mjwarp.Contact(
 _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
+_cb = mjwp_types.Callback(
+    **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
+)
 
 @ffi.format_args_for_warp
 def _kinematics_shim(
@@ -92,6 +95,7 @@ def _kinematics_shim(
 ):
   _m.stat = _s
   _m.opt = _o
+  _m.callback = _cb
   _d.efc = _e
   _d.contact = _c
   _m.body_branch_start = body_branch_start
@@ -332,6 +336,7 @@ def _tendon_shim(
 ):
   _m.stat = _s
   _m.opt = _o
+  _m.callback = _cb
   _d.efc = _e
   _d.contact = _c
   _m.body_parentid = body_parentid

mjx/mujoco/mjx/warp/types.py

@@ -25,11 +25,20 @@ from mujoco.mjx._src import dataclasses as mjx_dataclasses
 import numpy as np
 if typing.TYPE_CHECKING:
   GraphMode = int
+
+  @dataclasses.dataclass
+  class Callback:
+    pass
+
 else:
   try:
     from warp._src.jax_experimental.ffi import GraphMode
+    from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
+
+    Callback = mjwp_types.Callback
   except ImportError:
     GraphMode = int
+    Callback = None
 PyTreeNode = mjx_dataclasses.PyTreeNode
 
 @dataclasses.dataclass(frozen=True)
@@ -43,7 +52,6 @@ class TileSet:
     adr: address of each tile in the set
     size: size of all the tiles in this set
   """
-
   adr: np.ndarray
   size: int
 
@@ -122,6 +130,7 @@ class ModelWarp(PyTreeNode):
   body_branches: np.ndarray
   body_fluid_ellipsoid: np.ndarray
   body_tree: Tuple[np.ndarray, ...]
+  callback: Callback
   cam_projection: np.ndarray
   collision_sensor_adr: np.ndarray
   dof_tri_col: np.ndarray
@@ -174,6 +183,7 @@ class ModelWarp(PyTreeNode):
   mesh_polyvertadr: np.ndarray
   mesh_polyvertnum: np.ndarray
   mocap_bodyid: np.ndarray
+  nJfe: int
   nacttrnbody: int
   nbranch: int
   nflex: int
@@ -270,6 +280,9 @@ class DataWarp(PyTreeNode):
   crb: jax.Array
   efc__D: jax.Array
   efc__J: jax.Array
+  efc__J_colind: jax.Array
+  efc__J_rowadr: jax.Array
+  efc__J_rownnz: jax.Array
   efc__Ma: jax.Array
   efc__aref: jax.Array
   efc__force: jax.Array
@@ -287,6 +300,9 @@ class DataWarp(PyTreeNode):
   flexvert_xpos: jax.Array
   light_xdir: jax.Array
   light_xpos: jax.Array
+  moment_colind: jax.Array
+  moment_rowadr: jax.Array
+  moment_rownnz: jax.Array
   naccdmax: int
   nacon: jax.Array
   naconmax: int
@@ -294,7 +310,9 @@ class DataWarp(PyTreeNode):
   ne: jax.Array
   nefc: jax.Array
   nf: jax.Array
+  nisland: jax.Array
   njmax: int
+  njmax_pad: int
   nl: jax.Array
   nworld: int
   qLD: jax.Array
@@ -309,6 +327,7 @@ class DataWarp(PyTreeNode):
   ten_velocity: jax.Array
   ten_wrapadr: jax.Array
   ten_wrapnum: jax.Array
+  tree_island: jax.Array
   wrap_obj: jax.Array
   wrap_xpos: jax.Array
   shape = property(lambda self: self.cacc.shape)
@@ -332,6 +351,7 @@ DATA_NON_VMAP = {
     'naconmax',
     'ncollision',
     'njmax',
+    'njmax_pad',
     'nworld',
 }
 
@@ -365,7 +385,7 @@ _NDIM = {
         'act_dot': 2,
         'actuator_force': 2,
         'actuator_length': 2,
-        'actuator_moment': 3,
+        'actuator_moment': 2,
         'actuator_velocity': 2,
         'cacc': 3,
         'cam_xmat': 4,
@@ -394,6 +414,9 @@ _NDIM = {
         'cvel': 3,
         'efc__D': 2,
         'efc__J': 3,
+        'efc__J_colind': 3,
+        'efc__J_rowadr': 2,
+        'efc__J_rownnz': 2,
         'efc__Ma': 2,
         'efc__aref': 2,
         'efc__force': 2,
@@ -406,7 +429,7 @@ _NDIM = {
         'efc__vel': 2,
         'energy': 2,
         'eq_active': 2,
-        'flexedge_J': 3,
+        'flexedge_J': 2,
         'flexedge_length': 2,
         'flexedge_velocity': 2,
         'flexvert_xpos': 3,
@@ -416,6 +439,9 @@ _NDIM = {
         'light_xpos': 3,
         'mocap_pos': 3,
         'mocap_quat': 3,
+        'moment_colind': 2,
+        'moment_rowadr': 2,
+        'moment_rownnz': 2,
         'naccdmax': 0,
         'nacon': 1,
         'naconmax': 0,
@@ -423,7 +449,9 @@ _NDIM = {
         'ne': 1,
         'nefc': 1,
         'nf': 1,
+        'nisland': 1,
         'njmax': 0,
+        'njmax_pad': 0,
         'nl': 1,
         'nworld': 0,
         'qLD': 3,
@@ -458,6 +486,7 @@ _NDIM = {
         'ten_wrapadr': 2,
         'ten_wrapnum': 2,
         'time': 1,
+        'tree_island': 2,
         'wrap_obj': 3,
         'wrap_xpos': 3,
         'xanchor': 3,
@@ -473,7 +502,7 @@ _NDIM = {
         'M_colind': 1,
         'M_rowadr': 1,
         'M_rownnz': 1,
-        'actuator_acc0': 1,
+        'actuator_acc0': 2,
         'actuator_actadr': 1,
         'actuator_actearly': 1,
         'actuator_actlimited': 1,
@@ -481,7 +510,7 @@ _NDIM = {
         'actuator_actrange': 3,
         'actuator_biasprm': 3,
         'actuator_biastype': 1,
-        'actuator_cranklength': 1,
+        'actuator_cranklength': 2,
         'actuator_ctrllimited': 1,
         'actuator_ctrlrange': 3,
         'actuator_dynprm': 3,
@@ -491,7 +520,7 @@ _NDIM = {
         'actuator_gainprm': 3,
         'actuator_gaintype': 1,
         'actuator_gear': 3,
-        'actuator_lengthrange': 2,
+        'actuator_lengthrange': 3,
         'actuator_trnid': 2,
         'actuator_trntype': 1,
         'actuator_trntype_body_adr': 1,
@@ -670,6 +699,7 @@ _NDIM = {
         'mesh_normal': 2,
         'mesh_normaladr': 1,
         'mesh_normalnum': 1,
+        'mesh_octadr': 1,
         'mesh_polyadr': 1,
         'mesh_polymap': 1,
         'mesh_polymapadr': 1,
@@ -685,6 +715,8 @@ _NDIM = {
         'mesh_vertnum': 1,
         'mocap_bodyid': 1,
         'nC': 0,
+        'nJfe': 0,
+        'nJmom': 0,
         'nM': 0,
         'na': 0,
         'nacttrnbody': 0,
@@ -924,6 +956,9 @@ _BATCH_DIM = {
         'cvel': True,
         'efc__D': True,
         'efc__J': True,
+        'efc__J_colind': True,
+        'efc__J_rowadr': True,
+        'efc__J_rownnz': True,
         'efc__Ma': True,
         'efc__aref': True,
         'efc__force': True,
@@ -946,6 +981,9 @@ _BATCH_DIM = {
         'light_xpos': True,
         'mocap_pos': True,
         'mocap_quat': True,
+        'moment_colind': True,
+        'moment_rowadr': True,
+        'moment_rownnz': True,
         'naccdmax': False,
         'nacon': False,
         'naconmax': False,
@@ -953,7 +991,9 @@ _BATCH_DIM = {
         'ne': True,
         'nefc': True,
         'nf': True,
+        'nisland': True,
         'njmax': False,
+        'njmax_pad': False,
         'nl': True,
         'nworld': False,
         'qLD': True,
@@ -988,6 +1028,7 @@ _BATCH_DIM = {
         'ten_wrapadr': True,
         'ten_wrapnum': True,
         'time': True,
+        'tree_island': True,
         'wrap_obj': True,
         'wrap_xpos': True,
         'xanchor': True,
@@ -1003,7 +1044,7 @@ _BATCH_DIM = {
         'M_colind': False,
         'M_rowadr': False,
         'M_rownnz': False,
-        'actuator_acc0': False,
+        'actuator_acc0': True,
         'actuator_actadr': False,
         'actuator_actearly': False,
         'actuator_actlimited': False,
@@ -1011,7 +1052,7 @@ _BATCH_DIM = {
         'actuator_actrange': True,
         'actuator_biasprm': True,
         'actuator_biastype': False,
-        'actuator_cranklength': False,
+        'actuator_cranklength': True,
         'actuator_ctrllimited': False,
         'actuator_ctrlrange': True,
         'actuator_dynprm': True,
@@ -1021,7 +1062,7 @@ _BATCH_DIM = {
         'actuator_gainprm': True,
         'actuator_gaintype': False,
         'actuator_gear': True,
-        'actuator_lengthrange': False,
+        'actuator_lengthrange': True,
         'actuator_trnid': False,
         'actuator_trntype': False,
         'actuator_trntype_body_adr': False,
@@ -1200,6 +1241,7 @@ _BATCH_DIM = {
         'mesh_normal': False,
         'mesh_normaladr': False,
         'mesh_normalnum': False,
+        'mesh_octadr': False,
         'mesh_polyadr': False,
         'mesh_polymap': False,
         'mesh_polymapadr': False,
@@ -1215,6 +1257,8 @@ _BATCH_DIM = {
         'mesh_vertnum': False,
         'mocap_bodyid': False,
         'nC': False,
+        'nJfe': False,
+        'nJmom': False,
         'nM': False,
         'na': False,
         'nacttrnbody': False,