From 2a4af3246ec354614543248307530cb4d01dad7d Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Fri, 5 Jun 2026 04:56:22 -0700
Subject: [PATCH] Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 927227210
Change-Id: Ic72c58524c8e91e6fff44b4cab7f4a0f7de2bb33
---
 .../mujoco_warp/_src/collision_driver.py      |  175 +-
 .../third_party/mujoco_warp/_src/forward.py   |   99 +-
 .../mjx/third_party/mujoco_warp/_src/io.py    |  130 +-
 .../third_party/mujoco_warp/_src/island.py    |   15 +-
 .../third_party/mujoco_warp/_src/render.py    |  218 ++-
 .../mjx/third_party/mujoco_warp/_src/sleep.py | 1032 +++++++++++
 .../third_party/mujoco_warp/_src/solver.py    | 1617 +++++++++--------
 .../mjx/third_party/mujoco_warp/_src/types.py |  125 +-
 .../third_party/mujoco_warp/_src/warp_util.py |   25 -
 mjx/mujoco/mjx/warp/collision_driver.py       |   11 +
 mjx/mujoco/mjx/warp/forward.py                |  685 ++++---
 mjx/mujoco/mjx/warp/render.py                 |   28 +
 mjx/mujoco/mjx/warp/smooth.py                 |    1 +
 mjx/mujoco/mjx/warp/types.py                  |   65 +-
 14 files changed, 3023 insertions(+), 1203 deletions(-)
 create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/_src/sleep.py

diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
index 7c2fe28e..e3199833 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Any
+from typing import Any, Optional
 
 import warp as wp
 
@@ -30,8 +30,10 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType
 from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
+from mujoco.mjx.third_party.mujoco_warp._src.types import 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 SleepState
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat23
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
@@ -77,14 +79,24 @@ MJ_COLLISION_TABLE = {
 }
 
 
-@wp.kernel
-def _zero_nacon_ncollision(
-  # Data out:
-  nacon_out: wp.array[int],
-  ncollision_out: wp.array[int],
-):
-  ncollision_out[0] = 0
-  nacon_out[0] = 0
+@cache_kernel
+def _zero_nacon_ncollision(enable_sleep: bool = False):
+  @wp.kernel(module="unique", enable_backward=False)
+  def zero_nacon_ncollision(
+    # In:
+    skip_in: wp.array[int],
+    # Data out:
+    nacon_out: wp.array[int],
+    ncollision_out: wp.array[int],
+  ):
+    ncollision_out[0] = 0
+    if wp.static(enable_sleep):
+      if skip_in[0] != 0:
+        nacon_out[0] = 0
+    else:
+      nacon_out[0] = 0
+
+  return zero_nacon_ncollision
 
 
 @wp.func
@@ -376,7 +388,7 @@ def _binary_search(values: wp.array[Any], value: Any, lower: int, upper: int) ->
 
 
 @cache_kernel
-def _sap_project(opt_broadphase: int):
+def _sap_project(opt_broadphase: int, enable_sleep: bool = False):
   @wp.kernel(module="unique", enable_backward=False)
   def sap_project(
     # Model:
@@ -389,6 +401,7 @@ def _sap_project(opt_broadphase: int):
     nworld_in: int,
     # In:
     direction_in: wp.vec3,
+    skip_in: wp.array[int],
     # Out:
     projection_lower_out: wp.array2d[float],
     projection_upper_out: wp.array2d[float],
@@ -397,6 +410,10 @@ def _sap_project(opt_broadphase: int):
   ):
     worldid, geomid = wp.tid()
 
+    if wp.static(enable_sleep):
+      if skip_in[0] == 0:
+        return
+
     xpos = geom_xpos_in[worldid, geomid]
     rbound = geom_rbound[worldid % geom_rbound.shape[0], geomid]
 
@@ -424,38 +441,51 @@ def _sap_project(opt_broadphase: int):
   return sap_project
 
 
-@wp.kernel
-def _sap_range(
-  # Model:
-  ngeom: int,
-  # In:
-  projection_lower_in: wp.array2d[float],
-  projection_upper_in: wp.array2d[float],
-  sort_index_in: wp.array2d[int],
-  # Out:
-  range_out: wp.array2d[int],
-):
-  worldid, geomid = wp.tid()
+@cache_kernel
+def _sap_range(enable_sleep: bool = False):
+  @wp.kernel(module="unique", enable_backward=False)
+  def sap_range(
+    # Model:
+    ngeom: int,
+    # In:
+    projection_lower_in: wp.array2d[float],
+    projection_upper_in: wp.array2d[float],
+    sort_index_in: wp.array2d[int],
+    skip_in: wp.array[int],
+    # Out:
+    range_out: wp.array2d[int],
+  ):
+    worldid, geomid = wp.tid()
 
-  # current bounding geom
-  idx = sort_index_in[worldid, geomid]
+    if wp.static(enable_sleep):
+      if skip_in[0] == 0:
+        range_out[worldid, geomid] = 0
+        return
 
-  upper = projection_upper_in[worldid, idx]
+    # current bounding geom
+    idx = sort_index_in[worldid, geomid]
 
-  limit = _binary_search(projection_lower_in[worldid], upper, geomid + 1, ngeom)
-  limit = wp.min(ngeom - 1, limit)
+    upper = projection_upper_in[worldid, idx]
 
-  # range of geoms for the sweep and prune process
-  range_out[worldid, geomid] = limit - geomid
+    limit = _binary_search(projection_lower_in[worldid], upper, geomid + 1, ngeom)
+    limit = wp.min(ngeom - 1, limit)
+
+    # range of geoms for the sweep and prune process
+    range_out[worldid, geomid] = limit - geomid
+
+  return sap_range
 
 
 @cache_kernel
-def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int):
+def _sap_broadphase(
+  opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int, enable_sleep: bool = False
+):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Model:
     ngeom: int,
     geom_type: wp.array[int],
+    geom_bodyid: wp.array[int],
     geom_aabb: wp.array3d[wp.vec3],
     geom_rbound: wp.array2d[float],
     geom_margin: wp.array2d[float],
@@ -464,12 +494,14 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
     # Data in:
     geom_xpos_in: wp.array2d[wp.vec3],
     geom_xmat_in: wp.array2d[wp.mat33],
+    body_awake_in: wp.array2d[int],
     nworld_in: int,
     naconmax_in: int,
     # In:
     sort_index_in: wp.array2d[int],
     cumulative_sum_in: wp.array[int],
     nsweep_in: int,
+    skip_in: wp.array[int],
     # Data out:
     ncollision_out: wp.array[int],
     # Out:
@@ -479,6 +511,10 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
   ):
     worldgeomid = wp.tid()
 
+    if wp.static(enable_sleep):
+      if skip_in[0] == 0:
+        return
+
     nworldgeom = nworld_in * ngeom
     nworkpackages = cumulative_sum_in[nworldgeom - 1]
 
@@ -509,6 +545,16 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
       if pairid[0] < -1 and pairid[1] < 0:
         continue
 
+      if wp.static(enable_sleep):
+        b1 = geom_bodyid[geom1]
+        b2 = geom_bodyid[geom2]
+        s1 = body_awake_in[worldid, b1]
+        s2 = body_awake_in[worldid, b2]
+        if s1 == SleepState.ASLEEP and s2 == SleepState.ASLEEP:
+          continue
+        if (s1 == SleepState.ASLEEP and s2 == SleepState.STATIC) or (s2 == SleepState.ASLEEP and s1 == SleepState.STATIC):
+          continue
+
       if (
         wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin, ngeom_gap))(
           geom_aabb, geom_rbound, geom_margin, geom_gap, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
@@ -560,7 +606,7 @@ def _segmented_sort(tile_size: int):
 
 
 @event_scope
-def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
+def sap_broadphase(m: Model, d: Data, ctx: CollisionContext, skip: Optional[wp.array] = None):
   """Runs broadphase collision detection using a sweep-and-prune (SAP) algorithm.
 
   This method is more efficient than the N-squared approach for large numbers of
@@ -578,6 +624,8 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
   - `SAP_SEGMENTED`: Uses a segmented sort.
   """
   nworldgeom = d.nworld * m.ngeom
+  skip_in = skip if skip is not None else wp.ones(1, dtype=int)
+  enable_sleep = bool(m.opt.enableflags & EnableBit.SLEEP)
 
   # TODO(team): direction
 
@@ -593,9 +641,9 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
   segmented_index = wp.empty(d.nworld + 1 if m.opt.broadphase == BroadphaseType.SAP_SEGMENTED else 0, dtype=int)
 
   wp.launch(
-    kernel=_sap_project(m.opt.broadphase),
+    kernel=_sap_project(m.opt.broadphase, enable_sleep),
     dim=(d.nworld, m.ngeom),
-    inputs=[m.ngeom, m.geom_rbound, m.geom_margin, m.geom_gap, d.geom_xpos, d.nworld, direction],
+    inputs=[m.ngeom, m.geom_rbound, m.geom_margin, m.geom_gap, d.geom_xpos, d.nworld, direction, skip_in],
     outputs=[
       projection_lower.reshape((-1, m.ngeom)),
       projection_upper,
@@ -618,9 +666,9 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
     )
 
   wp.launch(
-    kernel=_sap_range,
+    kernel=_sap_range(enable_sleep),
     dim=(d.nworld, m.ngeom),
-    inputs=[m.ngeom, projection_lower.reshape((-1, m.ngeom)), projection_upper, sort_index.reshape((-1, m.ngeom))],
+    inputs=[m.ngeom, projection_lower.reshape((-1, m.ngeom)), projection_upper, sort_index.reshape((-1, m.ngeom)), skip_in],
     outputs=[range_],
   )
 
@@ -632,12 +680,18 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
   nsweep = 5 * nworldgeom
   wp.launch(
     kernel=_sap_broadphase(
-      m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
+      m.opt.broadphase_filter,
+      m.geom_aabb.shape[0],
+      m.geom_rbound.shape[0],
+      m.geom_margin.shape[0],
+      m.geom_gap.shape[0],
+      enable_sleep,
     ),
     dim=nsweep,
     inputs=[
       m.ngeom,
       m.geom_type,
+      m.geom_bodyid,
       m.geom_aabb,
       m.geom_rbound,
       m.geom_margin,
@@ -645,22 +699,27 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
       m.nxn_pairid,
       d.geom_xpos,
       d.geom_xmat,
+      d.body_awake,
       d.nworld,
       d.naconmax,
       sort_index.reshape((-1, m.ngeom)),
       cumulative_sum.reshape(-1),
       nsweep,
+      skip_in,
     ],
     outputs=[d.ncollision, ctx.collision_pair, ctx.collision_pairid, ctx.collision_worldid],
   )
 
 
 @cache_kernel
-def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int):
+def _nxn_broadphase(
+  opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int, enable_sleep: bool = False
+):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Model:
     geom_type: wp.array[int],
+    geom_bodyid: wp.array[int],
     geom_aabb: wp.array3d[wp.vec3],
     geom_rbound: wp.array2d[float],
     geom_margin: wp.array2d[float],
@@ -670,7 +729,10 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
     # Data in:
     geom_xpos_in: wp.array2d[wp.vec3],
     geom_xmat_in: wp.array2d[wp.mat33],
+    body_awake_in: wp.array2d[int],
     naconmax_in: int,
+    # In:
+    skip_in: wp.array[int],
     # Data out:
     ncollision_out: wp.array[int],
     # Out:
@@ -680,10 +742,24 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
   ):
     worldid, elementid = wp.tid()
 
+    if wp.static(enable_sleep):
+      if skip_in[0] == 0:
+        return
+
     geom = nxn_geom_pair[elementid]
     geom1 = geom[0]
     geom2 = geom[1]
 
+    if wp.static(enable_sleep):
+      b1 = geom_bodyid[geom1]
+      b2 = geom_bodyid[geom2]
+      s1 = body_awake_in[worldid, b1]
+      s2 = body_awake_in[worldid, b2]
+      if s1 == SleepState.ASLEEP and s2 == SleepState.ASLEEP:
+        return
+      if (s1 == SleepState.ASLEEP and s2 == SleepState.STATIC) or (s2 == SleepState.ASLEEP and s1 == SleepState.STATIC):
+        return
+
     if (
       wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin, ngeom_gap))(
         geom_aabb, geom_rbound, geom_margin, geom_gap, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
@@ -708,7 +784,7 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
 
 
 @event_scope
-def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
+def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext, skip: Optional[wp.array] = None):
   """Runs broadphase collision detection using a brute-force N-squared approach.
 
   This function iterates through a pre-filtered list of all possible geometry pairs and
@@ -721,13 +797,21 @@ def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
   The initial list of pairs is filtered at model creation time to exclude pairs based on
   `contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags.
   """
+  enable_sleep = bool(m.opt.enableflags & EnableBit.SLEEP)
+  skip_in = skip if skip is not None else wp.ones(1, dtype=int)
   wp.launch(
     _nxn_broadphase(
-      m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
+      m.opt.broadphase_filter,
+      m.geom_aabb.shape[0],
+      m.geom_rbound.shape[0],
+      m.geom_margin.shape[0],
+      m.geom_gap.shape[0],
+      enable_sleep,
     ),
     dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]),
     inputs=[
       m.geom_type,
+      m.geom_bodyid,
       m.geom_aabb,
       m.geom_rbound,
       m.geom_margin,
@@ -736,7 +820,9 @@ def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
       m.nxn_pairid_filtered,
       d.geom_xpos,
       d.geom_xmat,
+      d.body_awake,
       d.naconmax,
+      skip_in,
     ],
     outputs=[
       d.ncollision,
@@ -768,7 +854,7 @@ def _narrowphase(m: Model, d: Data, ctx: CollisionContext):
 
 
 @event_scope
-def collision(m: Model, d: Data):
+def collision(m: Model, d: Data, skip: Optional[wp.array] = None):
   """Runs the full collision detection pipeline.
 
   This function orchestrates the broadphase and narrowphase collision detection stages. It
@@ -789,15 +875,18 @@ def collision(m: Model, d: Data):
     d.nacon.zero_()
     return
 
+  # TODO(team): create context outside collision?
   ctx = create_collision_context(d.naconmax)
+  skip_in = skip if skip is not None else wp.ones(1, dtype=int)
+  enable_sleep = bool(m.opt.enableflags & EnableBit.SLEEP)
 
   # zero counters
-  wp.launch(_zero_nacon_ncollision, dim=1, outputs=[d.nacon, d.ncollision])
+  wp.launch(_zero_nacon_ncollision(enable_sleep), dim=1, inputs=[skip_in], outputs=[d.nacon, d.ncollision])
 
   if m.opt.broadphase == BroadphaseType.NXN:
-    nxn_broadphase(m, d, ctx)
+    nxn_broadphase(m, d, ctx, skip)
   else:
-    sap_broadphase(m, d, ctx)
+    sap_broadphase(m, d, ctx, skip)
 
   _narrowphase(m, d, ctx)
 
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
index 00f22815..60375d5d 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
@@ -25,6 +25,7 @@ from mujoco.mjx.third_party.mujoco_warp._src import island
 from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import passive
 from mujoco.mjx.third_party.mujoco_warp._src import sensor
+from mujoco.mjx.third_party.mujoco_warp._src import sleep
 from mujoco.mjx.third_party.mujoco_warp._src import smooth
 from mujoco.mjx.third_party.mujoco_warp._src import solver
 from mujoco.mjx.third_party.mujoco_warp._src import types
@@ -329,6 +330,11 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
 
   wp.copy(d.qacc_warmstart, d.qacc)
 
+  if not (m.opt.disableflags & DisableBit.ISLAND) and (m.opt.enableflags & EnableBit.SLEEP):
+    sleep.sleep(m, d)
+    fwd_velocity(m, d)
+    sleep.update_sleep(m, d)
+
 
 @wp.kernel
 def _compute_damping_deriv(
@@ -658,13 +664,37 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
   smooth.camlight(m, d)
   smooth.flex(m, d)
   smooth.tendon(m, d)
+
+  sleep_enabled = not (m.opt.disableflags & DisableBit.ISLAND) and (m.opt.enableflags & EnableBit.SLEEP)
+
+  if sleep_enabled and m.ntendon > 0:
+    sleep.wake_tendon(m, d)
+    sleep.update_sleep_trees(m, d)
+
   smooth.crb(m, d)
   smooth.tendon_armature(m, d)
   if factorize:
     smooth.factor_m(m, d)
   if m.opt.run_collision_detection:
-    collision_driver.collision(m, d)
+    if sleep_enabled:
+      # pass 1
+      collision_driver.collision(m, d)
+      # check for newly awake
+      skip = wp.zeros(1, dtype=int)
+      sleep.wake_collision(m, d, skip)
+      sleep.update_sleep(m, d)
+      # pass 2: broadphase kernels early-return if skip[0] is 0
+      collision_driver.collision(m, d, skip)
+    else:
+      collision_driver.collision(m, d)
+
   constraint.make_constraint(m, d)
+
+  if sleep_enabled:
+    if m.neq > 0:
+      sleep.wake_equality(m, d)
+    sleep.update_sleep(m, d)
+
   if m.ntree > 1 and not (m.opt.disableflags & types.DisableBit.ISLAND):
     island.island(m, d)
   smooth.transmission(m, d)
@@ -1242,23 +1272,39 @@ def fwd_actuation(m: Model, d: Data):
   )
 
 
-@wp.kernel
-def _qfrc_smooth(
-  # Data in:
-  qfrc_applied_in: wp.array2d[float],
-  qfrc_bias_in: wp.array2d[float],
-  qfrc_passive_in: wp.array2d[float],
-  qfrc_actuator_in: wp.array2d[float],
-  # Data out:
-  qfrc_smooth_out: wp.array2d[float],
-):
-  worldid, dofid = wp.tid()
-  qfrc_smooth_out[worldid, dofid] = (
-    qfrc_passive_in[worldid, dofid]
-    - qfrc_bias_in[worldid, dofid]
-    + qfrc_actuator_in[worldid, dofid]
-    + qfrc_applied_in[worldid, dofid]
-  )
+@cache_kernel
+def _qfrc_smooth(enable_sleep: bool):
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # Model:
+    body_treeid: wp.array[int],
+    dof_bodyid: wp.array[int],
+    # Data in:
+    qfrc_applied_in: wp.array2d[float],
+    tree_awake_in: wp.array2d[int],
+    qfrc_bias_in: wp.array2d[float],
+    qfrc_passive_in: wp.array2d[float],
+    qfrc_actuator_in: wp.array2d[float],
+    # Data out:
+    qfrc_smooth_out: wp.array2d[float],
+  ):
+    worldid, dofid = wp.tid()
+
+    if wp.static(enable_sleep):
+      bodyid = dof_bodyid[dofid]
+      tree = body_treeid[bodyid]
+      if tree >= 0 and tree_awake_in[worldid, tree] == 0:
+        qfrc_smooth_out[worldid, dofid] = 0.0
+        return
+
+    qfrc_smooth_out[worldid, dofid] = (
+      qfrc_passive_in[worldid, dofid]
+      - qfrc_bias_in[worldid, dofid]
+      + qfrc_actuator_in[worldid, dofid]
+      + qfrc_applied_in[worldid, dofid]
+    )
+
+  return kernel
 
 
 @event_scope
@@ -1270,10 +1316,19 @@ def fwd_acceleration(m: Model, d: Data, factorize: bool = False):
     d: The data object containing the current state and output arrays.
     factorize: Flag to factorize inertia matrix.
   """
+  enable_sleep = bool(m.opt.enableflags & EnableBit.SLEEP)
   wp.launch(
-    _qfrc_smooth,
+    _qfrc_smooth(enable_sleep),
     dim=(d.nworld, m.nv),
-    inputs=[d.qfrc_applied, d.qfrc_bias, d.qfrc_passive, d.qfrc_actuator],
+    inputs=[
+      m.body_treeid,
+      m.dof_bodyid,
+      d.qfrc_applied,
+      d.tree_awake,
+      d.qfrc_bias,
+      d.qfrc_passive,
+      d.qfrc_actuator,
+    ],
     outputs=[d.qfrc_smooth],
   )
   xfrc_accumulate(m, d, d.qfrc_smooth)
@@ -1287,6 +1342,10 @@ def fwd_acceleration(m: Model, d: Data, factorize: bool = False):
 @event_scope
 def forward(m: Model, d: Data):
   """Forward dynamics."""
+  if not (m.opt.disableflags & DisableBit.ISLAND) and (m.opt.enableflags & EnableBit.SLEEP):
+    sleep.wake(m, d)
+    sleep.update_sleep(m, d)
+
   energy = m.opt.enableflags & EnableBit.ENERGY
 
   fwd_position(m, d, factorize=False)
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
index 2810bc79..92d0f7a1 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
@@ -154,6 +154,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     (mjm.geom_type, types.GeomType, mujoco.mjtGeom),
     (mjm.sensor_type, types.SensorType, mujoco.mjtSensor),
     (mjm.wrap_type, types.WrapType, mujoco.mjtWrap),
+    (mjm.tree_sleep_policy, types.SleepPolicy, mujoco.mjtSleepPolicy),
   ):
     missing = ~np.isin(field, field_type)
     if missing.any():
@@ -178,6 +179,9 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     if unsupported:
       raise NotImplementedError(f"{mj_type(unsupported).name} is unsupported.")
 
+  if (mjm.opt.enableflags & mujoco.mjtEnableBit.mjENBL_SLEEP) and (mjm.eq_type == mujoco.mjtEq.mjEQ_FLEX).any():
+    raise NotImplementedError("Flex equality constraints are not supported with sleeping enabled.")
+
   if mjm.opt.noslip_iterations > 0:
     raise NotImplementedError(f"noslip solver not implemented.")
 
@@ -297,6 +301,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.nmaxpyramid = np.maximum(1, 2 * (m.nmaxcondim - 1))
   m.has_sdf_geom = (mjm.geom_type == mujoco.mjtGeom.mjGEOM_SDF).any()
   m.block_dim = types.BlockDim()
+  if mjm.nv > 500:
+    m.block_dim.linesearch_iterative = 512
   m.is_sparse = is_sparse(mjm)
   m.has_fluid = mjm.opt.wind.any() or mjm.opt.density > 0 or mjm.opt.viscosity > 0
 
@@ -995,6 +1001,7 @@ def _allocate_island_arrays(
   d.dof_island = wp.array(np.tile(mjd.dof_island, (nworld, 1 if island_enabled else 0)), dtype=int)
 
   d.island_dofadr = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_idofadr = wp.empty((nworld, ntree_size), dtype=int)
   d.island_nv = wp.empty((nworld, ntree_size), dtype=int)
   d.island_nefc = wp.empty((nworld, ntree_size), dtype=int)
   d.island_ne = wp.empty((nworld, ntree_size), dtype=int)
@@ -1155,6 +1162,7 @@ def make_data(
     "tree_island": None,
     "dof_island": None,
     "island_dofadr": None,
+    "island_idofadr": None,
     "island_nv": None,
     "island_nefc": None,
     "island_ne": None,
@@ -1171,6 +1179,10 @@ def make_data(
     "iqacc_smooth": None,
     "iqfrc_smooth": None,
     "iqfrc_constraint": None,
+    # sleep state: all trees start fully awake
+    "tree_asleep": wp.array(np.full((nworld, mjm.ntree), -(1 + types.MJ_MINAWAKE)), dtype=int),
+    "tree_awake": wp.array(np.ones((nworld, mjm.ntree)), dtype=int),
+    "body_awake": wp.array(np.ones((nworld, mjm.nbody)), dtype=int),
   }
   for f in dataclasses.fields(types.Data):
     if f.name in d_kwargs:
@@ -1380,6 +1392,7 @@ def put_data(
     "tree_island": None,
     "dof_island": None,
     "island_dofadr": None,
+    "island_idofadr": None,
     "island_nv": None,
     "island_nefc": None,
     "island_ne": None,
@@ -1657,12 +1670,18 @@ def get_data_into(
   # sensors
   result.sensordata[:] = d.sensordata.numpy()[world_id]
 
+  # sleep
+  result.tree_asleep[:] = d.tree_asleep.numpy()[world_id]
+  result.tree_awake[:] = d.tree_awake.numpy()[world_id]
+  result.body_awake[:] = d.body_awake.numpy()[world_id]
+
   # islands
   nisland = d.nisland.numpy()[world_id]
   result.nisland = nisland
   if d.tree_island.shape[1] > 0 and nisland:
     result.tree_island[:] = d.tree_island.numpy()[world_id]
     result.dof_island[:] = d.dof_island.numpy()[world_id]
+    result.island_idofadr[:nisland] = d.island_idofadr.numpy()[world_id, :nisland]
     result.island_dofadr[:nisland] = d.island_dofadr.numpy()[world_id, :nisland]
     result.island_nv[:nisland] = d.island_nv.numpy()[world_id, :nisland]
     result.island_nefc[:nisland] = d.island_nefc.numpy()[world_id, :nisland]
@@ -1743,6 +1762,8 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     nv: int,
     nu: int,
     na: int,
+    nbody: int,
+    ntree: int,
     neq: int,
     nsensordata: int,
     qpos0: wp.array2d[float],
@@ -1757,6 +1778,9 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     nf_out: wp.array[int],
     nl_out: wp.array[int],
     nefc_out: wp.array[int],
+    ntree_awake_out: wp.array[int],
+    nbody_awake_out: wp.array[int],
+    nv_awake_out: wp.array[int],
     time_out: wp.array[float],
     energy_out: wp.array[wp.vec2],
     qpos_out: wp.array2d[float],
@@ -1786,6 +1810,9 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     nefc_out[worldid] = 0
     time_out[worldid] = 0.0
     energy_out[worldid] = wp.vec2(0.0, 0.0)
+    ntree_awake_out[worldid] = ntree
+    nbody_awake_out[worldid] = nbody
+    nv_awake_out[worldid] = nv
     qpos0_id = worldid % qpos0.shape[0]
     for i in range(nq):
       qpos_out[worldid, i] = qpos0[qpos0_id, i]
@@ -1882,6 +1909,47 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     contact_type_out[conid] = 0
     contact_geomcollisionid_out[conid] = 0
 
+  @wp.kernel(module="unique", enable_backward=False)
+  def reset_sleep(
+    # Model:
+    nv: int,
+    nbody: int,
+    ntree: int,
+    body_mocapid: wp.array[int],
+    body_treeid: wp.array[int],
+    # In:
+    mj_minawake: int,
+    reset_in: wp.array[bool],
+    # Data out:
+    tree_asleep_out: wp.array2d[int],
+    tree_awake_out: wp.array2d[int],
+    body_awake_out: wp.array2d[int],
+    body_awake_ind_out: wp.array2d[int],
+    dof_awake_ind_out: wp.array2d[int],
+  ):
+    worldid, elemid = wp.tid()
+
+    if wp.static(reset is not None):
+      if not reset_in[worldid]:
+        return
+
+    if elemid < ntree:
+      tree_asleep_out[worldid, elemid] = -(1 + mj_minawake)
+      tree_awake_out[worldid, elemid] = 1
+
+    if elemid < nbody:
+      if body_treeid[elemid] < 0:
+        if body_mocapid[elemid] >= 0:
+          body_awake_out[worldid, elemid] = int(types.SleepState.AWAKE)
+        else:
+          body_awake_out[worldid, elemid] = int(types.SleepState.STATIC)
+      else:
+        body_awake_out[worldid, elemid] = int(types.SleepState.AWAKE)
+      body_awake_ind_out[worldid, elemid] = elemid
+
+    if elemid < nv:
+      dof_awake_ind_out[worldid, elemid] = elemid
+
   reset_input = reset or wp.ones(d.nworld, dtype=bool)
 
   wp.launch(reset_xfrc_applied, dim=(d.nworld, m.nbody, 6), inputs=[reset_input], outputs=[d.xfrc_applied])
@@ -1925,16 +1993,32 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     ],
   )
 
+  wp.launch(
+    reset_sleep,
+    dim=(d.nworld, max(m.ntree, m.nbody, m.nv)),
+    inputs=[m.nv, m.nbody, m.ntree, m.body_mocapid, m.body_treeid, types.MJ_MINAWAKE, reset_input],
+    outputs=[
+      d.tree_asleep,
+      d.tree_awake,
+      d.body_awake,
+      d.body_awake_ind,
+      d.dof_awake_ind,
+    ],
+  )
+
   wp.launch(
     reset_nworld,
     dim=d.nworld,
-    inputs=[m.nq, m.nv, m.nu, m.na, m.neq, m.nsensordata, m.qpos0, m.eq_active0, d.nworld, reset_input],
+    inputs=[m.nq, m.nv, m.nu, m.na, m.nbody, m.ntree, m.neq, m.nsensordata, m.qpos0, m.eq_active0, d.nworld, reset_input],
     outputs=[
       d.solver_niter,
       d.ne,
       d.nf,
       d.nl,
       d.nefc,
+      d.ntree_awake,
+      d.nbody_awake,
+      d.nv_awake,
       d.time,
       d.energy,
       d.qpos,
@@ -2785,7 +2869,7 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any
     "opt.graph_conditional",
     "opt.contact_sensor_maxmatch",
   }
-  mj_only_fields = {"opt.jacobian"}
+  mj_only_fields = {"opt.jacobian", "vis.quality.offsamples"}
 
   if not isinstance(overrides, dict):
     overrides_dict = {}
@@ -2951,10 +3035,14 @@ def create_render_context(
   use_ambient_lighting: bool = True,
   enabled_geom_groups: list[int] = [0, 1, 2],
   cam_active: list[bool] | None = None,
+  background_color: tuple[float, float, float, float] = (0.1, 0.1, 0.2, 1.0),
   flex_render_smooth: bool = True,
   use_precomputed_rays: bool = True,
   render_skybox: bool = False,
   enable_backface_culling: bool = True,
+  enable_specular: bool = True,
+  enable_emission: bool = True,
+  enable_per_light_ambient: bool = True,
 ) -> types.RenderContext:
   """Creates a render context on device.
 
@@ -2969,8 +3057,9 @@ def create_render_context(
       If None, uses the MuJoCo model values.
     use_textures: Whether to use textures.
     use_shadows: Whether to use shadows.
-    use_ambient_lighting: Whether to add the renderer's hemispheric ambient
-      lighting term before applying model lights.
+    use_ambient_lighting: Top-level ambient switch. When False, skips all
+                          ambient contributions, including headlight ambient,
+                          the no-light fallback, and per-light ambient.
     enabled_geom_groups: The geom groups to render.
     cam_active: List of booleans indicating which cameras to include in rendering.
                 If None, all cameras are included.
@@ -2983,6 +3072,19 @@ def create_render_context(
                              the ray (ray origin inside the geom). Matches MuJoCo's
                              mesh-ray rule. Default True. Disable for a small
                              performance gain when no camera is ever inside a geom.
+    background_color: The color to use for background pixels when no skybox is rendered.
+    enable_specular: Evaluate specular highlights per light. When False the
+                     half-vector normalize and shininess `pow` are dropped at
+                     compile time. Disable for performance when no specular is present.
+    enable_emission: Add `mat_emission * base_color` per shaded pixel. When
+                     False the term is dropped at compile time. Disable for performance
+                     when no emission is present.
+    enable_per_light_ambient: When ambient lighting is enabled, sum each
+                              light's `ambient` color into shaded pixels
+                              even outside its cone or in shadow. When False
+                              the per-light ambient pass is removed at compile
+                              time. Disable for performance when model lights
+                              do not use ambient colors.
 
   Returns:
     The render context containing rendering fields and output arrays on device.
@@ -3163,6 +3265,13 @@ def create_render_context(
   if len(flex_geom_flexid) > 0:
     geom_ray_types.add(int(types.GeomType.FLEX))
   geom_ray_types = tuple(sorted(geom_ray_types))
+  if mjm.nlight == 0:
+    light_attenuation_is_default = True
+    has_spot_lights = False
+  else:
+    atten = np.asarray(mjm.light_attenuation, dtype=np.float32).reshape(-1, 3)
+    light_attenuation_is_default = bool(np.allclose(atten, np.array([1.0, 0.0, 0.0], dtype=np.float32)))
+    has_spot_lights = bool((np.asarray(mjm.light_type) == int(mujoco.mjtLightType.mjLIGHT_SPOT)).any())
 
   rc = types.RenderContext(
     nrender=ncam,
@@ -3171,11 +3280,17 @@ def create_render_context(
     use_textures=use_textures,
     use_shadows=use_shadows,
     use_ambient_lighting=use_ambient_lighting,
-    background_color=render_util.pack_rgba_to_uint32(0.1 * 255.0, 0.1 * 255.0, 0.2 * 255.0, 1.0 * 255.0),
+    background_color=render_util.pack_rgba_to_uint32(
+      background_color[0] * 255.0, background_color[1] * 255.0, background_color[2] * 255.0, background_color[3] * 255.0
+    ),
     use_precomputed_rays=use_precomputed_rays,
     render_skybox=render_skybox,
     skybox_tex_id=skybox_tex_id,
     skybox_face_width=skybox_face_width,
+    headlight_active=bool(mjm.vis.headlight.active),
+    headlight_ambient=wp.vec3(mjm.vis.headlight.ambient),
+    headlight_diffuse=wp.vec3(mjm.vis.headlight.diffuse),
+    headlight_specular=wp.vec3(mjm.vis.headlight.specular),
     bvh_ngeom=bvh_ngeom,
     enabled_geom_ids=wp.array(geom_enabled_idx, dtype=int),
     mesh_registry=mesh_registry,
@@ -3218,6 +3333,11 @@ def create_render_context(
     total_rays=int(total),
     enable_backface_culling=enable_backface_culling,
     geom_ray_types=geom_ray_types,
+    enable_specular=enable_specular,
+    enable_emission=enable_emission,
+    enable_per_light_ambient=enable_per_light_ambient,
+    light_attenuation_is_default=light_attenuation_is_default,
+    has_spot_lights=has_spot_lights,
   )
 
   bvh.build_scene_bvh(mjm, mjd, rc, nworld)
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
index 0e92b034..806f71d3 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
@@ -346,6 +346,7 @@ def _island_map_dofs(
   island_idofadr_in: wp.array2d[int],
   nidof_in: wp.array[int],
   island_nv_inout: wp.array2d[int],
+  island_dofadr_out: wp.array2d[int],
   map_dof2idof_out: wp.array2d[int],
   map_idof2dof_out: wp.array2d[int],
   idof_islandid_out: wp.array2d[int],
@@ -360,6 +361,7 @@ def _island_map_dofs(
     local_idx = wp.atomic_add(island_nv_inout, worldid, island_id, 1)
     idof = island_idofadr_in[worldid, island_id] + local_idx
     idof_islandid_out[worldid, idof] = island_id
+    wp.atomic_min(island_dofadr_out, worldid, island_id, dofid)
   else:
     cnt = wp.atomic_add(unconstrained_cnt_inout, worldid, 0, 1)
     idof = nidof + cnt
@@ -804,7 +806,7 @@ def compute_island_mapping(m: types.Model, d: types.Data, ctx: IslandSolverConte
   Populates island solver context arrays via ctx: nv, nefc, ne, nf,
   iefcadr, nidof, map_dof2idof, map_idof2dof, dof_islandid, map_efc2iefc,
   map_iefc2efc, efc_islandid. Also populates d.dof_island, d.efc.island,
-  and d.island_dofadr.
+  d.island_idofadr, and d.island_dofadr.
 
   Args:
     m: Model.
@@ -816,6 +818,8 @@ def compute_island_mapping(m: types.Model, d: types.Data, ctx: IslandSolverConte
     d.dof_islandid = wp.empty((d.nworld, m.nv), dtype=int)
   if d.efc_islandid.shape[1] != d.njmax:
     d.efc_islandid = wp.empty((d.nworld, d.njmax), dtype=int)
+  if d.island_idofadr.shape[1] != m.ntree:
+    d.island_idofadr = wp.empty((d.nworld, m.ntree), dtype=int)
 
   # Ensure island-local DOF arrays are allocated at the right shape
   if d.iqacc.shape[1] != m.nv:
@@ -829,7 +833,7 @@ def compute_island_mapping(m: types.Model, d: types.Data, ctx: IslandSolverConte
     dim=(d.nworld, m.ntree),
     inputs=[],
     outputs=[
-      d.island_dofadr,
+      d.island_idofadr,
       d.island_nv,
       d.island_nefc,
       d.island_ne,
@@ -901,16 +905,17 @@ def compute_island_mapping(m: types.Model, d: types.Data, ctx: IslandSolverConte
     _island_scan_sizes,
     dim=d.nworld,
     inputs=[d.nisland],
-    outputs=[d.island_dofadr, d.island_nv, d.island_nefc, d.island_efcadr, d.nidof],
+    outputs=[d.island_idofadr, d.island_nv, d.island_nefc, d.island_efcadr, d.nidof],
   )
 
   # 4. Map DOFs
   unconstrained_cnt = wp.zeros((d.nworld, 1), dtype=int)
+  d.island_dofadr.fill_(m.nv)
   wp.launch(
     _island_map_dofs,
     dim=(d.nworld, m.nv),
-    inputs=[m.nv, d.dof_island, d.island_dofadr, d.nidof],
-    outputs=[d.island_nv, d.map_dof2idof, d.map_idof2dof, d.dof_islandid, unconstrained_cnt],
+    inputs=[m.nv, d.dof_island, d.island_idofadr, d.nidof],
+    outputs=[d.island_nv, d.island_dofadr, d.map_dof2idof, d.map_idof2dof, d.dof_islandid, unconstrained_cnt],
   )
 
   # 5. Map Constraints
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
index 05780e5b..5859c6ee 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
@@ -40,6 +40,21 @@ from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
 wp.set_module_options({"enable_backward": False})
 
+# Default value for mat_shininess in MuJoCo is 0.5
+# With an 8 bit image format, the maximum value is 255.0
+# So max shininess value for the Phong lighting model is 128.0
+MAX_SHININESS = 128.0
+# The exponent value for mat_shininess is 0.5 times the max shininess value
+DEFAULT_MAT_SHININESS_EXPONENT = 0.5 * MAX_SHININESS
+
+# Default value for mat_specular in MuJoCo is 0.5
+DEFAULT_MAT_SPECULAR = 0.5
+
+# Default value for mat_emission in MuJoCo is 0.0
+DEFAULT_MAT_EMISSION = 0.0
+
+NO_LIGHT_AMBIENT_FALLBACK = 0.3
+
 
 @wp.func
 def sample_texture(
@@ -416,37 +431,54 @@ def _make_compute_lighting(cast_ray_first_hit: wp.Function) -> wp.Function:
     lightcastshadow: bool,
     lightpos: wp.vec3,
     lightdir: wp.vec3,
+    lightattenuation: wp.vec3,
+    lightcutoff_rad: float,
+    lightexp: float,
+    lightdiff: wp.vec3,
+    lightspec: wp.vec3,
     normal: wp.vec3,
     hitpoint: wp.vec3,
+    view_dir: wp.vec3,
+    mat_spec: float,
+    mat_shin_exp: float,
     cull_backfaces: bool,
-  ) -> float:
-    light_contribution = float(0.0)
+    enable_specular: bool,
+    default_attenuation: bool,
+    has_spot: bool,
+  ) -> Tuple[wp.vec3, wp.vec3]:
+    diff_rgb = wp.vec3(0.0)
+    spec_rgb = wp.vec3(0.0)
 
     # TODO: We should probably only be looping over active lights
     # in the first place with a static loop of enabled light idx?
     if not lightactive:
-      return light_contribution
+      return diff_rgb, spec_rgb
 
-    L = wp.vec3(0.0, 0.0, 0.0)
+    L = wp.vec3(0.0)
     dist_to_light = float(MJ_MAXVAL)
-    attenuation = float(1.0)
+    attenuation = 1.0
 
     if lighttype == 1:  # directional light
-      L = wp.normalize(-lightdir)
+      # MuJoCo guarantees `lightdir` is unit length.
+      L = -lightdir
     else:
       L, dist_to_light = math.normalize_with_norm(lightpos - hitpoint)
-      attenuation = 1.0 / (1.0 + 0.02 * dist_to_light * dist_to_light)
-      if lighttype == 0:  # spot light
-        spot_dir = wp.normalize(lightdir)
-        cos_theta = wp.dot(-L, spot_dir)
-        spot_factor = wp.min(1.0, wp.max(0.0, (cos_theta - 0.85) * 10.0))
-        attenuation = attenuation * spot_factor
+      if not default_attenuation:
+        light_attenuation_factor = wp.vec3(1.0, dist_to_light, dist_to_light * dist_to_light)
+        attenuation = math.safe_div(1.0, wp.dot(light_attenuation_factor, lightattenuation))
+      if has_spot:
+        if lighttype == 0:  # spot light
+          cos_theta = wp.dot(-L, lightdir)
+          cos_cutoff = wp.cos(lightcutoff_rad)
+          if cos_theta < cos_cutoff:
+            return diff_rgb, spec_rgb
+          attenuation = attenuation * wp.pow(wp.max(cos_theta, 0.0), lightexp)
 
     ndotl = wp.max(0.0, wp.dot(normal, L))
     if ndotl == 0.0:
-      return light_contribution
+      return diff_rgb, spec_rgb
 
-    visible = float(1.0)
+    visible = 1.0
 
     if use_shadows and lightcastshadow:
       # Nudge the origin slightly along the surface normal to avoid
@@ -486,9 +518,17 @@ def _make_compute_lighting(cast_ray_first_hit: wp.Function) -> wp.Function:
       )
 
       if shadow_geom_id != -1:
-        visible = 0.3
+        visible = NO_LIGHT_AMBIENT_FALLBACK
 
-    return ndotl * attenuation * visible
+    weight = attenuation * visible
+    diff_rgb = lightdiff * (ndotl * weight)
+    if enable_specular:
+      if mat_spec > 0.0 and mat_shin_exp > 0.0:
+        H = wp.normalize(L + view_dir)
+        ndoth = wp.max(0.0, wp.dot(normal, H))
+        spec_rgb = lightspec * (mat_spec * wp.pow(ndoth, mat_shin_exp) * weight)
+
+    return diff_rgb, spec_rgb
 
   return compute_lighting
 
@@ -530,12 +570,21 @@ def render(m: Model, d: Data, rc: RenderContext):
     light_type: wp.array2d[int],
     light_castshadow: wp.array2d[bool],
     light_active: wp.array2d[bool],
+    light_attenuation: wp.array2d[wp.vec3],
+    light_cutoff: wp.array2d[float],
+    light_exponent: wp.array2d[float],
+    light_ambient: wp.array2d[wp.vec3],
+    light_diffuse: wp.array2d[wp.vec3],
+    light_specular: wp.array2d[wp.vec3],
     flex_vertadr: wp.array[int],
     flex_edge: wp.array[wp.vec2i],
     flex_radius: wp.array[float],
     mesh_faceadr: wp.array[int],
     mat_texid: wp.array3d[int],
     mat_texrepeat: wp.array2d[wp.vec2],
+    mat_emission: wp.array2d[float],
+    mat_specular: wp.array2d[float],
+    mat_shininess: wp.array2d[float],
     mat_rgba: wp.array2d[wp.vec4],
     # Data in:
     geom_xpos_in: wp.array2d[wp.vec3],
@@ -619,8 +668,9 @@ def render(m: Model, d: Data, rc: RenderContext):
         wp.static(rc.znear),
       )
 
-    ray_dir_world = cam_xmat_in[worldid, mujoco_cam_id] @ ray_dir_local_cam
     ray_origin_world = cam_xpos_in[worldid, mujoco_cam_id]
+    cam_mat_world = cam_xmat_in[worldid, mujoco_cam_id]
+    ray_dir_world = cam_mat_world @ ray_dir_local_cam
 
     geom_id, dist, normal, u, v, f, mesh_id = cast_ray(
       geom_type,
@@ -694,7 +744,6 @@ def render(m: Model, d: Data, rc: RenderContext):
       color = mat_rgba[worldid % mat_rgba.shape[0], geom_matid[worldid % geom_matid.shape[0], geom_id]]
 
     base_color = wp.vec3(color[0], color[1], color[2])
-    hit_color = base_color
 
     if wp.static(rc.use_textures):
       if geom_id != -2:
@@ -721,18 +770,48 @@ def render(m: Model, d: Data, rc: RenderContext):
             )
             base_color = wp.cw_mul(base_color, tex_color)
 
-    result = wp.vec3(0.0, 0.0, 0.0)
-    if wp.static(rc.use_ambient_lighting):
-      len_n = wp.length(normal)
-      n = normal if len_n > 0.0 else wp.vec3(0.0, 0.0, 1.0)
-      n = wp.normalize(n)
-      hemispheric = 0.5 * (n[2] + 1.0)
-      ambient_color = wp.vec3(0.4, 0.4, 0.45) * hemispheric + wp.vec3(0.1, 0.1, 0.12) * (1.0 - hemispheric)
-      result = 0.5 * wp.cw_mul(base_color, ambient_color)
+    mat_spec = DEFAULT_MAT_SPECULAR
+    mat_shin_exp = DEFAULT_MAT_SHININESS_EXPONENT
+    mat_emis = DEFAULT_MAT_EMISSION
+    if wp.static(rc.enable_specular or rc.enable_emission):
+      if geom_id != -2:
+        mat_id_for_spec = geom_matid[worldid % geom_matid.shape[0], geom_id]
+        if mat_id_for_spec >= 0:
+          if wp.static(rc.enable_specular):
+            mat_spec = mat_specular[worldid % mat_specular.shape[0], mat_id_for_spec]
+            mat_shin_exp = mat_shininess[worldid % mat_shininess.shape[0], mat_id_for_spec] * MAX_SHININESS
+          if wp.static(rc.enable_emission):
+            mat_emis = mat_emission[worldid % mat_emission.shape[0], mat_id_for_spec]
 
-    # Apply lighting and shadows
-    for l in range(wp.static(m.nlight)):
-      light_contribution = compute_lighting(
+    result = wp.vec3(0.0)
+    if wp.static(rc.enable_emission):
+      result = base_color * mat_emis
+
+    if wp.static(rc.use_ambient_lighting):
+      if wp.static(rc.headlight_active):
+        result = result + wp.cw_mul(base_color, wp.static(rc.headlight_ambient))
+      elif wp.static(m.nlight == 0):
+        result = result + base_color * NO_LIGHT_AMBIENT_FALLBACK
+      if wp.static(rc.enable_per_light_ambient):
+        for light_index in range(wp.static(m.nlight)):
+          if light_active[worldid % light_active.shape[0], light_index]:
+            result = result + wp.cw_mul(base_color, light_ambient[worldid % light_ambient.shape[0], light_index])
+
+    view_dir = -ray_dir_world
+
+    light_cutoff_worldid = light_cutoff[worldid % light_cutoff.shape[0]]
+    light_active_worldid = light_active[worldid % light_active.shape[0]]
+    light_type_worldid = light_type[worldid % light_type.shape[0]]
+    light_castshadow_worldid = light_castshadow[worldid % light_castshadow.shape[0]]
+    light_xpos_in_worldid = light_xpos_in[worldid]
+    light_xdir_in_worldid = light_xdir_in[worldid]
+    light_attenuation_worldid = light_attenuation[worldid % light_attenuation.shape[0]]
+    light_exponent_worldid = light_exponent[worldid % light_exponent.shape[0]]
+    light_diffuse_worldid = light_diffuse[worldid % light_diffuse.shape[0]]
+    light_specular_worldid = light_specular[worldid % light_specular.shape[0]]
+    # Apply Lighting for each light
+    for light_index in range(wp.static(m.nlight)):
+      diff_rgb, spec_rgb = compute_lighting(
         geom_type,
         geom_dataid,
         geom_size,
@@ -755,16 +834,76 @@ def render(m: Model, d: Data, rc: RenderContext):
         flex_geom_edgeid,
         flex_bvh_id,
         flex_group_root,
-        light_active[worldid % light_active.shape[0], l],
-        light_type[worldid % light_type.shape[0], l],
-        light_castshadow[worldid % light_castshadow.shape[0], l],
-        light_xpos_in[worldid, l],
-        light_xdir_in[worldid, l],
+        light_active_worldid[light_index],
+        light_type_worldid[light_index],
+        light_castshadow_worldid[light_index],
+        light_xpos_in_worldid[light_index],
+        light_xdir_in_worldid[light_index],
+        light_attenuation_worldid[light_index],
+        light_cutoff_worldid[light_index] * wp.static(wp.pi / 180.0),
+        light_exponent_worldid[light_index],
+        light_diffuse_worldid[light_index],
+        light_specular_worldid[light_index],
         normal,
         hit_point,
+        view_dir,
+        mat_spec,
+        mat_shin_exp,
         wp.static(rc.enable_backface_culling),
+        wp.static(rc.enable_specular),
+        wp.static(rc.light_attenuation_is_default),
+        wp.static(rc.has_spot_lights),
       )
-      result = result + base_color * light_contribution
+      result = result + wp.cw_mul(base_color, diff_rgb) + spec_rgb
+
+    # Apply Headlight
+    if wp.static(rc.headlight_active):
+      cam_pos = ray_origin_world
+      cam_fwd = -cam_mat_world[:, 2]
+      hl_diff, hl_spec = compute_lighting(
+        geom_type,
+        geom_dataid,
+        geom_size,
+        flex_vertadr,
+        flex_edge,
+        flex_radius,
+        geom_xpos_in,
+        geom_xmat_in,
+        flexvert_xpos_in,
+        use_shadows,
+        bvh_id,
+        group_root[worldid],
+        bvh_ngeom,
+        bvh_nflexgeom,
+        enabled_geom_ids,
+        worldid,
+        mesh_bvh_id,
+        hfield_bvh_id,
+        flex_geom_flexid,
+        flex_geom_edgeid,
+        flex_bvh_id,
+        flex_group_root,
+        True,
+        1,
+        False,
+        cam_pos,
+        cam_fwd,
+        wp.vec3(1.0, 0.0, 0.0),
+        0.0,
+        0.0,
+        wp.static(rc.headlight_diffuse),
+        wp.static(rc.headlight_specular),
+        normal,
+        hit_point,
+        view_dir,
+        mat_spec,
+        mat_shin_exp,
+        wp.static(rc.enable_backface_culling),
+        wp.static(rc.enable_specular),
+        True,
+        False,
+      )
+      result = result + wp.cw_mul(base_color, hl_diff) + hl_spec
 
     hit_color = wp.min(result, wp.vec3(1.0, 1.0, 1.0))
     hit_color = wp.max(hit_color, wp.vec3(0.0, 0.0, 0.0))
@@ -792,12 +931,21 @@ def render(m: Model, d: Data, rc: RenderContext):
       m.light_type,
       m.light_castshadow,
       m.light_active,
+      m.light_attenuation,
+      m.light_cutoff,
+      m.light_exponent,
+      m.light_ambient,
+      m.light_diffuse,
+      m.light_specular,
       m.flex_vertadr,
       m.flex_edge,
       m.flex_radius,
       m.mesh_faceadr,
       m.mat_texid,
       m.mat_texrepeat,
+      m.mat_emission,
+      m.mat_specular,
+      m.mat_shininess,
       m.mat_rgba,
       d.geom_xpos,
       d.geom_xmat,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sleep.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sleep.py
new file mode 100644
index 00000000..724212ad
--- /dev/null
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sleep.py
@@ -0,0 +1,1032 @@
+# Copyright 2026 The Newton Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import Optional
+
+import warp as wp
+
+from mujoco.mjx.third_party.mujoco_warp._src import types
+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.types import SleepPolicy
+from mujoco.mjx.third_party.mujoco_warp._src.types import SleepState
+from mujoco.mjx.third_party.mujoco_warp._src.types import WrapType
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
+
+wp.set_module_options({"enable_backward": False})
+
+# tree_asleep value for fully awake tree
+K_AWAKE_VAL = -(1 + types.MJ_MINAWAKE)
+
+
+@wp.func
+def _sleep_cycle(tree_asleep: wp.array2d[int], ntree: int, worldid: int, treeid: int) -> int:  # kernel_analyzer: ignore
+  if treeid < 0 or treeid >= ntree:
+    return -1
+
+  smallest = int(treeid)
+  current = int(treeid)
+  count = int(0)
+
+  for step in range(ntree + 1):  # safe upper bound
+    next_tree = tree_asleep[worldid, current]
+    if next_tree < 0 or next_tree >= ntree:
+      return -1
+
+    if next_tree < smallest:
+      smallest = next_tree
+
+    current = next_tree
+    count += 1
+    if current == treeid:
+      break
+
+  return smallest
+
+
+@wp.func
+def _tendon_limit_active(
+  tendon_limited: wp.array[int],  # kernel_analyzer: ignore
+  tendon_range: wp.array2d[wp.vec2],  # kernel_analyzer: ignore
+  tendon_margin: wp.array2d[float],  # kernel_analyzer: ignore
+  ten_length: wp.array2d[float],  # kernel_analyzer: ignore
+  worldid: int,  # kernel_analyzer: ignore
+  tenid: int,
+) -> bool:
+  if tendon_limited[tenid] == 0:
+    return False
+
+  length = ten_length[worldid, tenid]
+  margin = tendon_margin[worldid % tendon_margin.shape[0], tenid]
+  r = tendon_range[worldid % tendon_range.shape[0], tenid]
+
+  # low limit
+  dist_low = length - r[0]
+  if dist_low < margin:
+    return True
+
+  # high limit
+  dist_high = r[1] - length
+  if dist_high < margin:
+    return True
+
+  return False
+
+
+@wp.kernel
+def _zero_sleep_counters(
+  # Data out:
+  ntree_awake_out: wp.array[int],
+  nbody_awake_out: wp.array[int],
+  nv_awake_out: wp.array[int],
+):
+  worldid = wp.tid()
+  ntree_awake_out[worldid] = 0
+  nbody_awake_out[worldid] = 0
+  nv_awake_out[worldid] = 0
+
+
+@wp.kernel
+def _update_sleep_trees(
+  # Data in:
+  tree_asleep_in: wp.array2d[int],
+  # Data out:
+  ntree_awake_out: wp.array[int],
+  tree_awake_out: wp.array2d[int],
+):
+  worldid, treeid = wp.tid()
+  is_awake = int(tree_asleep_in[worldid, treeid] < 0)
+  tree_awake_out[worldid, treeid] = is_awake
+  if is_awake == 1:
+    wp.atomic_add(ntree_awake_out, worldid, 1)
+
+
+@wp.kernel
+def _update_sleep_bodies(
+  # Model:
+  body_parentid: wp.array[int],
+  body_rootid: wp.array[int],
+  body_mocapid: wp.array[int],
+  body_treeid: wp.array[int],
+  # Data in:
+  tree_awake_in: wp.array2d[int],
+  # In:
+  flg_staticawake: int,
+  # Data out:
+  nbody_awake_out: wp.array[int],
+  body_awake_out: wp.array2d[int],
+  body_awake_ind_out: wp.array2d[int],
+):
+  worldid, bodyid = wp.tid()
+
+  tree = body_treeid[bodyid]
+
+  # check static
+  if tree < 0:
+    root = body_rootid[bodyid]
+    mocap = body_mocapid[root]
+    if mocap >= 0:
+      state = SleepState.AWAKE
+    else:
+      state = SleepState.AWAKE if flg_staticawake != 0 else SleepState.STATIC
+  else:
+    state = SleepState.AWAKE if tree_awake_in[worldid, tree] == 1 else SleepState.ASLEEP
+
+  body_awake_out[worldid, bodyid] = state
+
+  if state != SleepState.ASLEEP:
+    idx = wp.atomic_add(nbody_awake_out, worldid, 1)
+    body_awake_ind_out[worldid, idx] = bodyid
+
+
+@wp.kernel
+def _update_sleep_dofs(
+  # Model:
+  body_treeid: wp.array[int],
+  dof_bodyid: wp.array[int],
+  # Data in:
+  body_awake_in: wp.array2d[int],
+  # Data out:
+  nv_awake_out: wp.array[int],
+  dof_awake_ind_out: wp.array2d[int],
+):
+  worldid, dofid = wp.tid()
+  bodyid = dof_bodyid[dofid]
+  if body_treeid[bodyid] >= 0 and body_awake_in[worldid, bodyid] == SleepState.AWAKE:
+    idx = wp.atomic_add(nv_awake_out, worldid, 1)
+    dof_awake_ind_out[worldid, idx] = dofid
+
+
+@event_scope
+def update_sleep(m: types.Model, d: types.Data, flg_staticawake: int = 0):
+  """Computes sleeping arrays from tree_asleep."""
+  wp.launch(
+    _zero_sleep_counters,
+    dim=d.nworld,
+    inputs=[],
+    outputs=[d.ntree_awake, d.nbody_awake, d.nv_awake],
+  )
+
+  wp.launch(
+    _update_sleep_trees,
+    dim=(d.nworld, m.ntree),
+    inputs=[d.tree_asleep],
+    outputs=[d.ntree_awake, d.tree_awake],
+  )
+
+  wp.launch(
+    _update_sleep_bodies,
+    dim=(d.nworld, m.nbody),
+    inputs=[
+      m.body_parentid,
+      m.body_rootid,
+      m.body_mocapid,
+      m.body_treeid,
+      d.tree_awake,
+      flg_staticawake,
+    ],
+    outputs=[
+      d.nbody_awake,
+      d.body_awake,
+      d.body_awake_ind,
+    ],
+  )
+
+  wp.launch(
+    _update_sleep_dofs,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      m.body_treeid,
+      m.dof_bodyid,
+      d.body_awake,
+    ],
+    outputs=[d.nv_awake, d.dof_awake_ind],
+  )
+
+
+@event_scope
+def update_sleep_trees(m: types.Model, d: types.Data):
+  """Lightweight update of tree_awake array only, avoiding body/dof kernel launches."""
+  wp.launch(
+    _zero_sleep_counters,
+    dim=d.nworld,
+    inputs=[],
+    outputs=[d.ntree_awake, d.nbody_awake, d.nv_awake],
+  )
+
+  wp.launch(
+    _update_sleep_trees,
+    dim=(d.nworld, m.ntree),
+    inputs=[d.tree_asleep],
+    outputs=[d.ntree_awake, d.tree_awake],
+  )
+
+
+@wp.func
+def _wake_tree(
+  # Model:
+  ntree: int,
+  # In:
+  worldid: int,
+  treeid: int,
+  wakeval: int,
+  # Data out:
+  tree_asleep_out: wp.array2d[int],
+) -> int:
+  """Wakes tree treeid and its associated cycle, returning number of woke trees."""
+  asleep_val = tree_asleep_out[worldid, treeid]
+  if asleep_val < 0:
+    if wakeval < asleep_val:
+      tree_asleep_out[worldid, treeid] = wakeval
+    return 0
+
+  nwoke = int(0)
+  current = int(treeid)
+  for step in range(ntree + 1):  # safe upper bound
+    next_tree = tree_asleep_out[worldid, current]
+    if next_tree < 0 or next_tree >= ntree:
+      break
+
+    tree_asleep_out[worldid, current] = wakeval
+    nwoke += 1
+    current = next_tree
+    if current == treeid:
+      break
+
+  return nwoke
+
+
+@wp.func
+def _tree_can_sleep(
+  # Model:
+  nbody: int,
+  body_treeid: wp.array[int],
+  dof_length: wp.array[float],
+  tree_dofadr: wp.array[int],
+  tree_dofnum: wp.array[int],
+  tree_sleep_policy: wp.array[int],
+  # Data in:
+  qvel_in: wp.array2d[float],
+  qfrc_applied_in: wp.array2d[float],
+  xfrc_applied_in: wp.array2d[wp.spatial_vector],
+  # In:
+  worldid: int,
+  treeid: int,
+  sleep_tolerance: float,
+) -> bool:
+  policy = tree_sleep_policy[treeid]
+  if policy == SleepPolicy.AUTO_NEVER:
+    return False
+
+  # check xfrc_applied
+  for b in range(nbody):
+    if body_treeid[b] == treeid:
+      xfrc = xfrc_applied_in[worldid, b]
+      for i in range(6):
+        if xfrc[i] != 0.0:
+          return False
+
+  # check qfrc_applied
+  dofadr = tree_dofadr[treeid]
+  dofnum = tree_dofnum[treeid]
+  for d in range(dofnum):
+    if qfrc_applied_in[worldid, dofadr + d] != 0.0:
+      return False
+
+  # check qvel
+  for d in range(dofnum):
+    dof_idx = dofadr + d
+    v = qvel_in[worldid, dof_idx]
+    weight = dof_length[dof_idx]
+    if sleep_tolerance > 0.0:
+      if wp.abs(weight * v) >= sleep_tolerance:
+        return False
+    else:
+      if v != 0.0:
+        return False
+
+  return True
+
+
+@wp.kernel
+def _wake_kernel(
+  # Model:
+  nbody: int,
+  ntree: int,
+  body_treeid: wp.array[int],
+  dof_length: wp.array[float],
+  tree_dofadr: wp.array[int],
+  tree_dofnum: wp.array[int],
+  tree_sleep_policy: wp.array[int],
+  # Data in:
+  qvel_in: wp.array2d[float],
+  qfrc_applied_in: wp.array2d[float],
+  xfrc_applied_in: wp.array2d[wp.spatial_vector],
+  tree_awake_in: wp.array2d[int],
+  # Out:
+  tree_asleep_out: wp.array2d[int],  # kernel_analyzer: ignore
+  nwoke_out: wp.array[int],  # kernel_analyzer: ignore
+):
+  worldid, treeid = wp.tid()
+
+  asleep = int(tree_asleep_out[worldid, treeid] >= 0)
+  if asleep == 0:
+    return
+
+  # if tree_awake mismatch or cannot sleep: wake up
+  if tree_awake_in[worldid, treeid] == 1 or not _tree_can_sleep(
+    nbody,
+    body_treeid,
+    dof_length,
+    tree_dofadr,
+    tree_dofnum,
+    tree_sleep_policy,
+    qvel_in,
+    qfrc_applied_in,
+    xfrc_applied_in,
+    worldid,
+    treeid,
+    0.0,  # zero tolerance
+  ):
+    woke = _wake_tree(ntree, worldid, treeid, K_AWAKE_VAL, tree_asleep_out)
+    if woke > 0:
+      wp.atomic_add(nwoke_out, worldid, woke)
+
+
+@wp.kernel
+def _wake_collision_kernel(
+  # Model:
+  ntree: int,
+  body_treeid: wp.array[int],
+  geom_bodyid: wp.array[int],
+  # Data in:
+  tree_awake_in: wp.array2d[int],
+  contact_geom_in: wp.array[wp.vec2i],
+  contact_worldid_in: wp.array[int],
+  nacon_in: wp.array[int],
+  # Out:
+  tree_asleep_out: wp.array2d[int],  # kernel_analyzer: ignore
+  skip_out: wp.array[int],  # kernel_analyzer: ignore
+):
+  conid = wp.tid()
+  if conid >= nacon_in[0]:
+    return
+
+  geom_pair = contact_geom_in[conid]
+  g1 = geom_pair[0]
+  g2 = geom_pair[1]
+
+  if g1 < 0 or g2 < 0:
+    return
+
+  b1 = geom_bodyid[g1]
+  b2 = geom_bodyid[g2]
+  tree1 = body_treeid[b1]
+  tree2 = body_treeid[b2]
+
+  if tree1 < 0 or tree2 < 0:
+    return
+
+  worldid = contact_worldid_in[conid]
+  awake1 = tree_awake_in[worldid, tree1]
+  awake2 = tree_awake_in[worldid, tree2]
+
+  if awake1 == 1 and awake2 == 1:
+    return
+
+  if awake1 == 0 and awake2 == 0:
+    return
+
+  # wake sleeping tree
+  sleeping_tree = tree2 if awake1 == 1 else tree1
+  wakeval = tree_asleep_out[worldid, tree1] if awake1 == 1 else tree_asleep_out[worldid, tree2]
+
+  woke = _wake_tree(ntree, worldid, sleeping_tree, wakeval, tree_asleep_out)
+  if woke > 0:
+    wp.atomic_add(skip_out, 0, woke)
+
+
+@wp.kernel
+def _wake_tendon_kernel(
+  # Model:
+  ntree: int,
+  ntendon: int,
+  body_treeid: wp.array[int],
+  jnt_bodyid: wp.array[int],
+  geom_bodyid: wp.array[int],
+  site_bodyid: wp.array[int],
+  tendon_adr: wp.array[int],
+  tendon_num: wp.array[int],
+  tendon_limited: wp.array[int],
+  tendon_range: wp.array2d[wp.vec2],
+  tendon_margin: wp.array2d[float],
+  wrap_type: wp.array[int],
+  wrap_objid: wp.array[int],
+  # Data in:
+  ten_length_in: wp.array2d[float],
+  tree_awake_in: wp.array2d[int],
+  # Data out:
+  tree_asleep_out: wp.array2d[int],
+  # Out:
+  nwoke_out: wp.array[int],
+):
+  worldid, tenid = wp.tid()
+
+  adr = tendon_adr[tenid]
+  num = tendon_num[tenid]
+
+  # Pass 1: Check if any tree involved in the tendon is awake
+  any_awake = int(0)
+  wakeval = int(K_AWAKE_VAL)
+
+  for i in range(num):
+    idx = adr + i
+    t_type = wrap_type[idx]
+    t_objid = wrap_objid[idx]
+    t = int(-1)
+    if t_type == WrapType.JOINT:
+      t = body_treeid[jnt_bodyid[t_objid]]
+    elif t_type == WrapType.SITE:
+      t = body_treeid[site_bodyid[t_objid]]
+    elif t_type == WrapType.SPHERE or t_type == WrapType.CYLINDER:
+      t = body_treeid[geom_bodyid[t_objid]]
+
+    if t >= 0:
+      if tree_awake_in[worldid, t] == 1:
+        any_awake = 1
+        val = tree_asleep_out[worldid, t]
+        if val < wakeval:
+          wakeval = val
+
+  # Pass 2: If at least one tree is awake and the limit is active, wake up all sleeping trees
+  if any_awake == 1:
+    if _tendon_limit_active(tendon_limited, tendon_range, tendon_margin, ten_length_in, worldid, tenid):
+      for i in range(num):
+        idx = adr + i
+        t_type = wrap_type[idx]
+        t_objid = wrap_objid[idx]
+        t = int(-1)
+        if t_type == WrapType.JOINT:
+          t = body_treeid[jnt_bodyid[t_objid]]
+        elif t_type == WrapType.SITE:
+          t = body_treeid[site_bodyid[t_objid]]
+        elif t_type == WrapType.SPHERE or t_type == WrapType.CYLINDER:
+          t = body_treeid[geom_bodyid[t_objid]]
+
+        if t >= 0:
+          if tree_awake_in[worldid, t] == 0:
+            woke = _wake_tree(ntree, worldid, t, wakeval, tree_asleep_out)
+            if woke > 0:
+              wp.atomic_add(nwoke_out, worldid, woke)
+
+
+@wp.func
+def _tendon_wake_val(
+  # Model:
+  body_treeid: wp.array[int],
+  jnt_bodyid: wp.array[int],
+  geom_bodyid: wp.array[int],
+  site_bodyid: wp.array[int],
+  tendon_adr: wp.array[int],
+  tendon_num: wp.array[int],
+  wrap_type: wp.array[int],
+  wrap_objid: wp.array[int],
+  # Data in:
+  tree_awake_in: wp.array2d[int],
+  # In:
+  worldid: int,
+  tenid: int,
+  # Data out:
+  tree_asleep_out: wp.array2d[int],
+) -> int:
+  """Returns the minimum (most-awake) tree_asleep value if any tree in tendon is awake, else 0."""
+  if tenid < 0:
+    return 0
+
+  adr = tendon_adr[tenid]
+  num = tendon_num[tenid]
+  wakeval = int(0)
+
+  for i in range(num):
+    idx = adr + i
+    t_type = wrap_type[idx]
+    t_objid = wrap_objid[idx]
+    t = int(-1)
+    if t_type == WrapType.JOINT:
+      t = body_treeid[jnt_bodyid[t_objid]]
+    elif t_type == WrapType.SITE:
+      t = body_treeid[site_bodyid[t_objid]]
+    elif t_type == WrapType.SPHERE or t_type == WrapType.CYLINDER:
+      t = body_treeid[geom_bodyid[t_objid]]
+
+    if t >= 0:
+      if tree_awake_in[worldid, t] == 1:
+        val = tree_asleep_out[worldid, t]
+        if wakeval == 0 or val < wakeval:
+          wakeval = val
+
+  return wakeval
+
+
+@wp.func
+def _wake_tendon_trees(
+  # Model:
+  ntree: int,
+  body_treeid: wp.array[int],
+  jnt_bodyid: wp.array[int],
+  geom_bodyid: wp.array[int],
+  site_bodyid: wp.array[int],
+  tendon_adr: wp.array[int],
+  tendon_num: wp.array[int],
+  wrap_type: wp.array[int],
+  wrap_objid: wp.array[int],
+  # Data in:
+  tree_awake_in: wp.array2d[int],
+  # In:
+  worldid: int,
+  tenid: int,
+  wakeval: int,
+  # Data out:
+  tree_asleep_out: wp.array2d[int],
+  # Out:
+  nwoke_out: wp.array[int],
+):
+  """Wakes up all sleeping trees associated with a tendon."""
+  if tenid < 0:
+    return
+
+  adr = tendon_adr[tenid]
+  num = tendon_num[tenid]
+  for i in range(num):
+    idx = adr + i
+    t_type = wrap_type[idx]
+    t_objid = wrap_objid[idx]
+    t = int(-1)
+    if t_type == WrapType.JOINT:
+      t = body_treeid[jnt_bodyid[t_objid]]
+    elif t_type == WrapType.SITE:
+      t = body_treeid[site_bodyid[t_objid]]
+    elif t_type == WrapType.SPHERE or t_type == WrapType.CYLINDER:
+      t = body_treeid[geom_bodyid[t_objid]]
+
+    if t >= 0:
+      if tree_awake_in[worldid, t] == 0:
+        woke = _wake_tree(ntree, worldid, t, wakeval, tree_asleep_out)
+        if woke > 0:
+          wp.atomic_add(nwoke_out, worldid, woke)
+
+
+@wp.kernel
+def _wake_equality_kernel(
+  # Model:
+  ntree: int,
+  neq: int,
+  body_treeid: wp.array[int],
+  jnt_bodyid: wp.array[int],
+  geom_bodyid: wp.array[int],
+  site_bodyid: wp.array[int],
+  eq_type: wp.array[int],
+  eq_obj1id: wp.array[int],
+  eq_obj2id: wp.array[int],
+  eq_objtype: wp.array[int],
+  tendon_adr: wp.array[int],
+  tendon_num: wp.array[int],
+  wrap_type: wp.array[int],
+  wrap_objid: wp.array[int],
+  # Data in:
+  eq_active_in: wp.array2d[bool],  # kernel_analyzer: ignore
+  tree_awake_in: wp.array2d[int],
+  # Data out:
+  tree_asleep_out: wp.array2d[int],  # kernel_analyzer: ignore
+  # Out:
+  nwoke_out: wp.array[int],  # kernel_analyzer: ignore
+):
+  worldid, eqid = wp.tid()
+
+  if not eq_active_in[worldid, eqid]:
+    return
+
+  eqtype = eq_type[eqid]
+  id1 = eq_obj1id[eqid]
+  id2 = eq_obj2id[eqid]
+
+  if eqtype == EqType.CONNECT or eqtype == EqType.WELD or eqtype == EqType.JOINT:
+    tree1 = int(-1)
+    tree2 = int(-1)
+
+    if eqtype == EqType.CONNECT or eqtype == EqType.WELD:
+      if eq_objtype[eqid] == ObjType.BODY:
+        tree1 = body_treeid[id1]
+        tree2 = body_treeid[id2]
+      else:
+        tree1 = body_treeid[site_bodyid[id1]]
+        tree2 = body_treeid[site_bodyid[id2]]
+    elif eqtype == EqType.JOINT:
+      tree1 = body_treeid[jnt_bodyid[id1]] if id1 >= 0 else -1
+      tree2 = body_treeid[jnt_bodyid[id2]] if id2 >= 0 else -1
+
+    s1 = tree_awake_in[worldid, tree1] if tree1 >= 0 else SleepState.STATIC
+    s2 = tree_awake_in[worldid, tree2] if tree2 >= 0 else SleepState.STATIC
+
+    if s1 != SleepState.ASLEEP and s2 != SleepState.ASLEEP:
+      return
+    if s1 == SleepState.STATIC or s2 == SleepState.STATIC:
+      return
+    if tree1 == tree2:
+      return
+
+    if s1 == SleepState.ASLEEP and s2 == SleepState.ASLEEP:
+      cycle1 = _sleep_cycle(tree_asleep_out, ntree, worldid, tree1)
+      cycle2 = _sleep_cycle(tree_asleep_out, ntree, worldid, tree2)
+      if cycle1 != cycle2:
+        w1 = _wake_tree(ntree, worldid, tree1, K_AWAKE_VAL, tree_asleep_out)
+        w2 = _wake_tree(ntree, worldid, tree2, K_AWAKE_VAL, tree_asleep_out)
+        if w1 + w2 > 0:
+          wp.atomic_add(nwoke_out, worldid, w1 + w2)
+    else:
+      sleeping_tree = tree1 if s1 == SleepState.ASLEEP else tree2
+      woke = _wake_tree(ntree, worldid, sleeping_tree, K_AWAKE_VAL, tree_asleep_out)
+      if woke > 0:
+        wp.atomic_add(nwoke_out, worldid, woke)
+
+  elif eqtype == EqType.TENDON:
+    ten1 = id1
+    ten2 = id2
+    w1 = _tendon_wake_val(
+      body_treeid,
+      jnt_bodyid,
+      geom_bodyid,
+      site_bodyid,
+      tendon_adr,
+      tendon_num,
+      wrap_type,
+      wrap_objid,
+      tree_awake_in,
+      worldid,
+      ten1,
+      tree_asleep_out,
+    )
+    w2 = _tendon_wake_val(
+      body_treeid,
+      jnt_bodyid,
+      geom_bodyid,
+      site_bodyid,
+      tendon_adr,
+      tendon_num,
+      wrap_type,
+      wrap_objid,
+      tree_awake_in,
+      worldid,
+      ten2,
+      tree_asleep_out,
+    )
+
+    if w1 < 0 or w2 < 0:
+      wakeval = int(K_AWAKE_VAL)
+      if w1 < 0 and w1 < wakeval:
+        wakeval = w1
+      if w2 < 0 and w2 < wakeval:
+        wakeval = w2
+
+      _wake_tendon_trees(
+        ntree,
+        body_treeid,
+        jnt_bodyid,
+        geom_bodyid,
+        site_bodyid,
+        tendon_adr,
+        tendon_num,
+        wrap_type,
+        wrap_objid,
+        tree_awake_in,
+        worldid,
+        ten1,
+        wakeval,
+        tree_asleep_out,
+        nwoke_out,
+      )
+      _wake_tendon_trees(
+        ntree,
+        body_treeid,
+        jnt_bodyid,
+        geom_bodyid,
+        site_bodyid,
+        tendon_adr,
+        tendon_num,
+        wrap_type,
+        wrap_objid,
+        tree_awake_in,
+        worldid,
+        ten2,
+        wakeval,
+        tree_asleep_out,
+        nwoke_out,
+      )
+
+  # TODO(team): Implement waking for EqType.FLEX constraints.
+
+
+@event_scope
+def wake(m: types.Model, d: types.Data):
+  """Wakes sleeping trees due to user changes/perturbations."""
+  nwoke = wp.zeros((d.nworld,), dtype=int)
+  wp.launch(
+    _wake_kernel,
+    dim=(d.nworld, m.ntree),
+    inputs=[
+      m.nbody,
+      m.ntree,
+      m.body_treeid,
+      m.dof_length,
+      m.tree_dofadr,
+      m.tree_dofnum,
+      m.tree_sleep_policy,
+      d.qvel,
+      d.qfrc_applied,
+      d.xfrc_applied,
+      d.tree_awake,
+      d.tree_asleep,
+    ],
+    outputs=[nwoke],
+  )
+
+
+@event_scope
+def wake_collision(m: types.Model, d: types.Data, skip: Optional[wp.array] = None):
+  """Wakes sleeping trees that touch awake trees."""
+  skip_out = skip if skip is not None else wp.zeros(1, dtype=int)
+  wp.launch(
+    _wake_collision_kernel,
+    dim=d.naconmax,
+    inputs=[
+      m.ntree,
+      m.body_treeid,
+      m.geom_bodyid,
+      d.tree_awake,
+      d.contact.geom,
+      d.contact.worldid,
+      d.nacon,
+      d.tree_asleep,
+    ],
+    outputs=[skip_out],
+  )
+
+
+@event_scope
+def wake_tendon(m: types.Model, d: types.Data):
+  """Wakes sleeping trees with constrained tendons."""
+  if m.ntendon == 0:
+    return
+
+  nwoke = wp.zeros((d.nworld,), dtype=int)
+  wp.launch(
+    _wake_tendon_kernel,
+    dim=(d.nworld, m.ntendon),
+    inputs=[
+      m.ntree,
+      m.ntendon,
+      m.body_treeid,
+      m.jnt_bodyid,
+      m.geom_bodyid,
+      m.site_bodyid,
+      m.tendon_adr,
+      m.tendon_num,
+      m.tendon_limited,
+      m.tendon_range,
+      m.tendon_margin,
+      m.wrap_type,
+      m.wrap_objid,
+      d.ten_length,
+      d.tree_awake,
+    ],
+    outputs=[d.tree_asleep, nwoke],
+  )
+
+
+@event_scope
+def wake_equality(m: types.Model, d: types.Data):
+  """Wakes sleeping trees connected by equality constraints."""
+  if m.neq == 0:
+    return
+
+  nwoke = wp.zeros((d.nworld,), dtype=int)
+  wp.launch(
+    _wake_equality_kernel,
+    dim=(d.nworld, m.neq),
+    inputs=[
+      m.ntree,
+      m.neq,
+      m.body_treeid,
+      m.jnt_bodyid,
+      m.geom_bodyid,
+      m.site_bodyid,
+      m.eq_type,
+      m.eq_obj1id,
+      m.eq_obj2id,
+      m.eq_objtype,
+      m.tendon_adr,
+      m.tendon_num,
+      m.wrap_type,
+      m.wrap_objid,
+      d.eq_active,
+      d.tree_awake,
+      d.tree_asleep,
+    ],
+    outputs=[nwoke],
+  )
+
+
+@wp.kernel
+def _sweep_awake_trees(  # kernel_analyzer: ignore
+  # Model:
+  nbody: int,
+  body_treeid: wp.array[int],
+  dof_length: wp.array[float],
+  tree_dofadr: wp.array[int],
+  tree_dofnum: wp.array[int],
+  tree_sleep_policy: wp.array[int],
+  # Data in:
+  qvel_in: wp.array2d[float],
+  qfrc_applied_in: wp.array2d[float],
+  xfrc_applied_in: wp.array2d[wp.spatial_vector],
+  # In:
+  opt_sleep_tolerance: wp.array[float],
+  # Out:
+  tree_asleep_out: wp.array2d[int],  # kernel_analyzer: ignore
+):
+  worldid, treeid = wp.tid()
+  sleep_tolerance = opt_sleep_tolerance[worldid % opt_sleep_tolerance.shape[0]]
+  as_val = tree_asleep_out[worldid, treeid]
+  if as_val >= 0:
+    return
+
+  if _tree_can_sleep(
+    nbody,
+    body_treeid,
+    dof_length,
+    tree_dofadr,
+    tree_dofnum,
+    tree_sleep_policy,
+    qvel_in,
+    qfrc_applied_in,
+    xfrc_applied_in,
+    worldid,
+    treeid,
+    sleep_tolerance,
+  ):
+    if as_val < -1:
+      tree_asleep_out[worldid, treeid] = as_val + 1
+  else:
+    tree_asleep_out[worldid, treeid] = K_AWAKE_VAL
+
+
+@wp.kernel
+def _check_island_can_sleep(
+  # Model:
+  ntree: int,
+  # Data in:
+  nisland_in: wp.array[int],
+  tree_asleep_in: wp.array2d[int],
+  tree_island_in: wp.array2d[int],
+  # Out:
+  island_can_sleep_out: wp.array2d[int],
+):
+  worldid, treeid = wp.tid()
+  nisland = nisland_in[worldid]
+  island_id = tree_island_in[worldid, treeid]
+  if island_id >= 0 and island_id < nisland:
+    as_val = tree_asleep_in[worldid, treeid]
+    if as_val < -1:
+      # Not ready to sleep yet
+      wp.atomic_min(island_can_sleep_out, worldid, island_id, 0)
+
+
+@wp.kernel
+def _build_cycles(  # kernel_analyzer: ignore
+  # Model:
+  ntree: int,
+  tree_dofadr: wp.array[int],
+  tree_dofnum: wp.array[int],
+  # Data in:
+  nisland_in: wp.array[int],
+  tree_island_in: wp.array2d[int],
+  # In:
+  island_can_sleep_in: wp.array2d[int],
+  # Out:
+  tree_asleep_out: wp.array2d[int],  # kernel_analyzer: ignore
+  qvel_out: wp.array2d[float],  # kernel_analyzer: ignore
+  qacc_out: wp.array2d[float],  # kernel_analyzer: ignore
+  nslept_out: wp.array[int],  # kernel_analyzer: ignore
+):
+  worldid = wp.tid()
+
+  num_islands = nisland_in[worldid]
+  for island_id in range(num_islands):
+    if island_can_sleep_in[worldid, island_id] == 1:
+      first_tree = int(-1)
+      prev_tree = int(-1)
+      n = int(0)
+      for t in range(ntree):
+        if tree_island_in[worldid, t] == island_id:
+          if first_tree == -1:
+            first_tree = t
+          if prev_tree != -1:
+            tree_asleep_out[worldid, prev_tree] = t
+          prev_tree = t
+          n += 1
+
+          # Zero velocities and accelerations
+          dofadr = tree_dofadr[t]
+          dofnum = tree_dofnum[t]
+          for d in range(dofnum):
+            qvel_out[worldid, dofadr + d] = 0.0
+            qacc_out[worldid, dofadr + d] = 0.0
+
+      if first_tree != -1:
+        tree_asleep_out[worldid, prev_tree] = first_tree
+        wp.atomic_add(nslept_out, worldid, n)
+
+  # Sleep unconstrained trees
+  for t in range(ntree):
+    island_id = tree_island_in[worldid, t]
+    if island_id < 0 or island_id >= num_islands:
+      if tree_asleep_out[worldid, t] == -1:
+        tree_asleep_out[worldid, t] = t  # self-cycle
+        wp.atomic_add(nslept_out, worldid, 1)
+
+      # Ensure sleeping tree dof velocity and acceleration remain exactly zero
+      if tree_asleep_out[worldid, t] >= 0:
+        dofadr = tree_dofadr[t]
+        dofnum = tree_dofnum[t]
+        for d in range(dofnum):
+          # TODO(team): shouldn't be necessary to zero if island is already asleep
+          qvel_out[worldid, dofadr + d] = 0.0
+          qacc_out[worldid, dofadr + d] = 0.0
+
+
+@event_scope
+def sleep(m: types.Model, d: types.Data):
+  """Puts trees to sleep according to velocity tolerance."""
+  # 1. Sweep over awake trees and increment counter if they can sleep
+  wp.launch(
+    _sweep_awake_trees,
+    dim=(d.nworld, m.ntree),
+    inputs=[
+      m.nbody,
+      m.body_treeid,
+      m.dof_length,
+      m.tree_dofadr,
+      m.tree_dofnum,
+      m.tree_sleep_policy,
+      d.qvel,
+      d.qfrc_applied,
+      d.xfrc_applied,
+      m.opt.sleep_tolerance,
+    ],
+    outputs=[d.tree_asleep],
+  )
+
+  # 2. Check which constraint islands can sleep (all trees in island must be asleep)
+  island_can_sleep = wp.ones((d.nworld, m.ntree), dtype=int)
+  wp.launch(
+    _check_island_can_sleep,
+    dim=(d.nworld, m.ntree),
+    inputs=[
+      m.ntree,
+      d.nisland,
+      d.tree_asleep,
+      d.tree_island,
+    ],
+    outputs=[island_can_sleep],
+  )
+
+  # 3. Build sleep cycles for sleeping islands and sleep unconstrained trees
+  nslept = wp.zeros((d.nworld,), dtype=int)
+  wp.launch(
+    _build_cycles,
+    dim=d.nworld,
+    inputs=[
+      m.ntree,
+      m.tree_dofadr,
+      m.tree_dofnum,
+      d.nisland,
+      d.tree_island,
+      island_can_sleep,
+      d.tree_asleep,
+      d.qvel,
+      d.qacc,
+    ],
+    outputs=[nslept],
+  )
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
index 1a8ab32b..0329924a 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
@@ -14,7 +14,6 @@
 # ==============================================================================
 
 from math import ceil
-from math import sqrt
 
 import warp as wp
 
@@ -30,7 +29,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import IslandSolverContext
 from mujoco.mjx.third_party.mujoco_warp._src.types import SolverContext
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import scoped_mathdx_gemm_disabled
 
 wp.set_module_options({"enable_backward": False})
 
@@ -53,16 +51,13 @@ def create_inverse_context(m: types.Model, d: types.Data) -> InverseContext:
   return InverseContext(
     Jaref=wp.empty((nworld, njmax), dtype=float),
     search_dot=wp.empty((nworld,), dtype=float),
-    gauss=wp.empty((nworld,), dtype=float),
-    cost=wp.empty((nworld,), dtype=float),
-    prev_cost=wp.empty((nworld,), dtype=float),
     done=wp.empty((nworld,), dtype=bool),
     changed_efc_ids=wp.empty((nworld, 0), dtype=int),
     changed_efc_count=wp.empty((0,), dtype=int),
   )
 
 
-def create_island_solver_context(m: types.Model, d: types.Data) -> IslandSolverContext:
+def _create_island_solver_context(m: types.Model, d: types.Data) -> IslandSolverContext:
   """Create an IslandSolverContext with allocated workspace arrays.
 
   Args:
@@ -100,12 +95,13 @@ def create_island_solver_context(m: types.Model, d: types.Data) -> IslandSolverC
     done=wp.empty((nworld, ntree), dtype=bool),
     solver_niter=wp.empty((nworld, ntree), dtype=int),
     beta=wp.empty((nworld, ntree), dtype=float) if alloc_island_cg else wp.empty((nworld, 0), dtype=float),
+    beta_den=wp.empty((nworld, ntree), dtype=float) if alloc_island_cg else wp.empty((nworld, 0), dtype=float),
     alpha=wp.empty((nworld, ntree), dtype=float),
     Ma=wp.empty((nworld, nv), dtype=float),
   )
 
 
-def create_solver_context(m: types.Model, d: types.Data) -> SolverContext:
+def _create_solver_context(m: types.Model, d: types.Data) -> SolverContext:
   """Create a SolverContext with allocated workspace arrays.
 
   Args:
@@ -126,24 +122,23 @@ def create_solver_context(m: types.Model, d: types.Data) -> SolverContext:
   return SolverContext(
     Jaref=wp.empty((nworld, njmax), dtype=float),
     search_dot=wp.empty((nworld,), dtype=float),
-    gauss=wp.empty((nworld,), dtype=float),
-    cost=wp.empty((nworld,), dtype=float),
-    prev_cost=wp.empty((nworld,), dtype=float),
     done=wp.empty((nworld,), dtype=bool),
     grad=wp.zeros((nworld, nv_pad), dtype=float),
     grad_dot=wp.empty((nworld,), dtype=float),
-    Mgrad=wp.zeros((nworld, nv_pad), dtype=float),
+    Mgrad=wp.empty((nworld, nv_pad), dtype=float),
     search=wp.empty((nworld, nv), dtype=float),
     mv=wp.empty((nworld, nv), dtype=float),
     jv=wp.empty((nworld, njmax), dtype=float),
     quad=wp.empty((nworld, njmax), dtype=wp.vec3),
     quad_gauss=wp.empty((nworld,), dtype=wp.vec3),
     alpha=wp.empty((nworld,), dtype=float),
+    improvement=wp.empty((nworld,), dtype=float),
     prev_grad=wp.empty((nworld, nv), dtype=float),
     prev_Mgrad=wp.empty((nworld, nv), dtype=float),
     beta=wp.empty((nworld,), dtype=float),
-    h=wp.zeros((nworld, nv_pad, nv_pad), dtype=float) if alloc_h else wp.empty((nworld, 0, 0), dtype=float),
-    hfactor=wp.zeros((nworld, nv_pad, nv_pad), dtype=float) if alloc_hfactor else wp.empty((nworld, 0, 0), dtype=float),
+    beta_den=wp.empty((nworld,), dtype=float),
+    h=wp.empty((nworld, nv_pad, nv_pad), dtype=float) if alloc_h else wp.empty((nworld, 0, 0), dtype=float),
+    hfactor=wp.empty((nworld, nv_pad, nv_pad), dtype=float) if alloc_hfactor else wp.empty((nworld, 0, 0), dtype=float),
     changed_efc_ids=wp.empty((nworld, njmax), dtype=int) if alloc_h else wp.empty((nworld, 0), dtype=int),
     changed_efc_count=wp.empty((nworld,), dtype=int) if alloc_h else wp.empty((0,), dtype=int),
   )
@@ -159,19 +154,33 @@ def _in_bracket(x: wp.vec3, y: wp.vec3) -> bool:
   return (x[1] < y[1] and y[1] < 0.0) or (x[1] > y[1] and y[1] > 0.0)
 
 
+@wp.func
+def _eval_pt_direct_alpha_zero(jaref: float, jv: float, d: float) -> wp.vec3:
+  """Eval quadratic constraint at alpha=0."""
+  jvD = jv * d
+  return wp.vec3(0.5 * d * jaref * jaref, jvD * jaref, jv * jvD)
+
+
 @wp.func
 def _eval_pt_direct(jaref: float, jv: float, d: float, alpha: float) -> wp.vec3:
-  """Eval quadratic constraint, return (cost, grad, hessian)."""
+  """Eval quadratic constraint."""
   x = jaref + alpha * jv
   jvD = jv * d
   return wp.vec3(0.5 * d * x * x, jvD * x, jv * jvD)
 
 
 @wp.func
-def _eval_pt_direct_alpha_zero(jaref: float, jv: float, d: float) -> wp.vec3:
-  """Eval quadratic constraint at alpha=0."""
+def _eval_pt_direct_cost_alpha_zero(jaref: float, d: float) -> float:
+  return 0.5 * d * jaref * jaref
+
+
+@wp.func
+def _eval_pt_direct_shifted(jaref: float, jv: float, d: float, alpha: float, offset: float) -> wp.vec3:
+  """Eval quadratic constraint shifted by alpha=0, plus a constant cost offset."""
   jvD = jv * d
-  return wp.vec3(0.5 * d * jaref * jaref, jvD * jaref, jv * jvD)
+  hessian = jv * jvD
+  alpha_h = alpha * hessian
+  return wp.vec3(alpha * (jvD * jaref + 0.5 * alpha_h) + offset, jvD * jaref + alpha_h, hessian)
 
 
 @wp.func
@@ -192,6 +201,24 @@ def _eval_pt_direct_3alphas(
   )
 
 
+@wp.func
+def _eval_pt_direct_shifted_3alphas(
+  jaref: float, jv: float, d: float, lo_alpha: float, hi_alpha: float, mid_alpha: float, offset: float
+) -> tuple[wp.vec3, wp.vec3, wp.vec3]:
+  """Eval shifted quadratic constraint for 3 alphas, plus a constant cost offset."""
+  jvD = jv * d
+  grad0 = jvD * jaref
+  hessian = jv * jvD
+  lo_ah = lo_alpha * hessian
+  hi_ah = hi_alpha * hessian
+  mid_ah = mid_alpha * hessian
+  return (
+    wp.vec3(lo_alpha * (grad0 + 0.5 * lo_ah) + offset, grad0 + lo_ah, hessian),
+    wp.vec3(hi_alpha * (grad0 + 0.5 * hi_ah) + offset, grad0 + hi_ah, hessian),
+    wp.vec3(mid_alpha * (grad0 + 0.5 * mid_ah) + offset, grad0 + mid_ah, hessian),
+  )
+
+
 @wp.func
 def _eval_cost(quad: wp.vec3, alpha: float) -> float:
   return alpha * alpha * quad[2] + alpha * quad[1] + quad[0]
@@ -223,6 +250,11 @@ def _eval_pt_3alphas(quad: wp.vec3, lo_alpha: float, hi_alpha: float, mid_alpha:
   )
 
 
+@wp.func
+def _shift_cost(pt: wp.vec3, cost0: float) -> wp.vec3:
+  return wp.vec3(pt[0] - cost0, pt[1], pt[2])
+
+
 @wp.func
 def _eval_frictionloss_pt(x: float, f: float, rf: float, jv: float, d: float) -> wp.vec3:
   """Eval frictionloss and return (cost, grad, hessian). x = Jaref + alpha * jv."""
@@ -235,6 +267,15 @@ def _eval_frictionloss_pt(x: float, f: float, rf: float, jv: float, d: float) ->
     return wp.vec3(f * (-0.5 * rf + x), f * jv, 0.0)
 
 
+@wp.func
+def _eval_frictionloss_cost(x: float, f: float, rf: float, d: float) -> float:
+  if (-rf < x) and (x < rf):
+    return 0.5 * d * x * x
+  elif x <= -rf:
+    return f * (-0.5 * rf - x)
+  return f * (-0.5 * rf + x)
+
+
 @wp.func
 def _eval_frictionloss_pt_one(x: float, f: float, rf: float, half_d: float, jvD: float, hessian: float, f_jv: float) -> wp.vec3:
   """Eval frictionloss with precomputed shared values."""
@@ -265,15 +306,12 @@ def _eval_frictionloss_pt_3alphas(
 @wp.func
 def _eval_elliptic(
   # In:
-  impratio_invsqrt: float,
-  friction: types.vec5,
+  mu: float,
   quad: wp.vec3,
   quad1: wp.vec3,
   quad2: wp.vec3,
   alpha: float,
 ) -> wp.vec3:
-  mu = friction[0] * impratio_invsqrt
-
   u0 = quad1[0]
   v0 = quad1[1]
   uu = quad1[2]
@@ -324,6 +362,103 @@ def _eval_elliptic(
   return wp.vec3(0.0, 0.0, 0.0)
 
 
+@wp.func
+def _eval_elliptic_cost(
+  # In:
+  mu: float,
+  quad: wp.vec3,
+  quad1: wp.vec3,
+  quad2: wp.vec3,
+  alpha: float,
+) -> float:
+  u0 = quad1[0]
+  v0 = quad1[1]
+  uu = quad1[2]
+  uv = quad2[0]
+  vv = quad2[1]
+  dm = quad2[2]
+
+  N = u0 + alpha * v0
+  Tsqr = uu + alpha * (2.0 * uv + alpha * vv)
+
+  if Tsqr <= 0.0:
+    if N < 0.0:
+      return _eval_cost(quad, alpha)
+  else:
+    T = wp.sqrt(Tsqr)
+    if N >= mu * T:
+      pass
+    elif mu * N + T <= 0.0:
+      return _eval_cost(quad, alpha)
+    else:
+      return 0.5 * dm * (N - mu * T) * (N - mu * T)
+
+  return 0.0
+
+
+@wp.func
+def _eval_constraint(
+  # In:
+  is_equality: bool,
+  is_friction: bool,
+  is_elliptic: bool,
+  jaref: float,
+  D: float,
+  frictionloss: float,
+  efcid: int,
+  efcid0: int,
+  jaref0: float,
+  D0: float,
+  mu: float,
+  ufrictionj: float,
+  TT: float,
+) -> wp.vec3:
+  if is_equality:
+    force = -D * jaref
+    cost = 0.5 * D * jaref * jaref
+    return wp.vec3(force, float(types.ConstraintState.QUADRATIC.value), cost)
+
+  if is_friction:
+    rf = math.safe_div(frictionloss, D)
+    if jaref <= -rf:
+      return wp.vec3(frictionloss, float(types.ConstraintState.LINEARNEG.value), -frictionloss * (0.5 * rf + jaref))
+    elif jaref >= rf:
+      return wp.vec3(-frictionloss, float(types.ConstraintState.LINEARPOS.value), -frictionloss * (0.5 * rf - jaref))
+    else:
+      return wp.vec3(-D * jaref, float(types.ConstraintState.QUADRATIC.value), 0.5 * D * jaref * jaref)
+
+  if is_elliptic:
+    N = jaref0 * mu
+    if TT <= 0.0:
+      T = 0.0
+    else:
+      T = wp.sqrt(TT)
+
+    # Top zone
+    if (N >= mu * T) or ((T <= 0.0) and (N >= 0.0)):
+      return wp.vec3(0.0, float(types.ConstraintState.SATISFIED.value), 0.0)
+    # Bottom zone
+    elif (mu * N + T <= 0.0) or ((T <= 0.0) and (N < 0.0)):
+      return wp.vec3(-D * jaref, float(types.ConstraintState.QUADRATIC.value), 0.5 * D * jaref * jaref)
+    # Middle zone
+    else:
+      dm = math.safe_div(D0, mu * mu * (1.0 + mu * mu))
+      nmt = N - mu * T
+      force_normal = -dm * nmt * mu
+
+      if efcid == efcid0:
+        return wp.vec3(force_normal, float(types.ConstraintState.CONE.value), 0.5 * dm * nmt * nmt)
+      else:
+        force_tangent = -math.safe_div(force_normal, T) * ufrictionj
+
+      return wp.vec3(force_tangent, float(types.ConstraintState.CONE.value), 0.0)
+
+  if jaref >= 0.0:
+    return wp.vec3(0.0, float(types.ConstraintState.SATISFIED.value), 0.0)
+  else:
+    return wp.vec3(-D * jaref, float(types.ConstraintState.QUADRATIC.value), 0.5 * D * jaref * jaref)
+
+
 @wp.func
 def _log_scale(min_value: float, max_value: float, num_values: int, i: int) -> float:
   step = (wp.log(max_value) - wp.log(min_value)) / wp.max(1.0, float(num_values - 1))
@@ -331,7 +466,7 @@ def _log_scale(min_value: float, max_value: float, num_values: int, i: int) -> f
 
 
 @wp.kernel
-def linesearch_parallel_fused(
+def _linesearch_parallel_fused(
   # Model:
   opt_ls_iterations: int,
   opt_impratio_invsqrt: wp.array[float],
@@ -364,16 +499,17 @@ def linesearch_parallel_fused(
 
   alpha = _log_scale(opt_ls_parallel_min_step, 1.0, opt_ls_iterations, alphaid)
 
-  out = _eval_cost(ctx_quad_gauss_in[worldid], alpha)
+  quad_gauss = ctx_quad_gauss_in[worldid]
+  out = alpha * alpha * quad_gauss[2] + alpha * quad_gauss[1]
 
   ne = ne_in[worldid]
   nf = nf_in[worldid]
 
-  # TODO(team): _eval with option to only compute cost
   for efcid in range(min(njmax_in, nefc_in[worldid])):
     # equality
     if efcid < ne:
-      out += _eval_cost(ctx_quad_in[worldid, efcid], alpha)
+      quad = ctx_quad_in[worldid, efcid]
+      out += alpha * alpha * quad[2] + alpha * quad[1]
     # friction
     elif efcid < ne + nf:
       # search point, friction loss, bound (rf)
@@ -381,19 +517,10 @@ def linesearch_parallel_fused(
       dir = ctx_jv_in[worldid, efcid]
       x = start + alpha * dir
       f = efc_frictionloss_in[worldid, efcid]
-      rf = math.safe_div(f, efc_D_in[worldid, efcid])
+      efc_D = efc_D_in[worldid, efcid]
+      rf = math.safe_div(f, efc_D)
 
-      # -bound < x < bound : quadratic
-      if (-rf < x) and (x < rf):
-        quad = ctx_quad_in[worldid, efcid]
-      # x < -bound: linear negative
-      elif x <= -rf:
-        quad = wp.vec3(f * (-0.5 * rf - start), -f * dir, 0.0)
-      # bound < x : linear positive
-      else:
-        quad = wp.vec3(f * (-0.5 * rf + start), f * dir, 0.0)
-
-      out += _eval_cost(quad, alpha)
+      out += _eval_frictionloss_cost(x, f, rf, efc_D) - _eval_frictionloss_cost(start, f, rf, efc_D)
     # limit and contact
     elif efc_type_in[worldid, efcid] == types.ConstraintType.CONTACT_ELLIPTIC:
       # extract contact info
@@ -406,56 +533,36 @@ def linesearch_parallel_fused(
       if efcid != efcid0:
         continue
 
-      friction = contact_friction_in[conid]
-      mu = friction[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
-
       # unpack quad
       efcid1 = contact_efc_address_in[conid, 1]
       efcid2 = contact_efc_address_in[conid, 2]
-      u0 = ctx_quad_in[worldid, efcid1][0]
-      v0 = ctx_quad_in[worldid, efcid1][1]
-      uu = ctx_quad_in[worldid, efcid1][2]
-      uv = ctx_quad_in[worldid, efcid2][0]
-      vv = ctx_quad_in[worldid, efcid2][1]
-      dm = ctx_quad_in[worldid, efcid2][2]
 
-      # compute N, Tsqr
-      N = u0 + alpha * v0
-      Tsqr = uu + alpha * (2.0 * uv + alpha * vv)
+      impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
+      friction = contact_friction_in[conid]
+      quad = ctx_quad_in[worldid, efcid]
+      quad1 = ctx_quad_in[worldid, efcid1]
+      quad2 = ctx_quad_in[worldid, efcid2]
 
-      # no tangential force: top or bottom zone
-      if Tsqr <= 0.0:
-        # bottom zone: quadratic cost
-        if N < 0.0:
-          out += _eval_cost(ctx_quad_in[worldid, efcid], alpha)
-      # otherwise regular processing
-      else:
-        # tangential force
-        T = wp.sqrt(Tsqr)
-
-        # N >= mu * T : top zone
-        if N >= mu * T:
-          # nothing to do
-          pass
-        # mu * N + T <= 0 : bottom zone
-        elif mu * N + T <= 0.0:
-          out += _eval_cost(ctx_quad_in[worldid, efcid], alpha)
-        # otherwise middle zone
-        else:
-          out += 0.5 * dm * (N - mu * T) * (N - mu * T)
+      mu = friction[0] * impratio_invsqrt
+      out += _eval_elliptic_cost(mu, quad, quad1, quad2, alpha)
+      out -= _eval_elliptic_cost(mu, quad, quad1, quad2, 0.0)
     else:
       # search point
-      x = ctx_Jaref_in[worldid, efcid] + alpha * ctx_jv_in[worldid, efcid]
+      start = ctx_Jaref_in[worldid, efcid]
+      x = start + alpha * ctx_jv_in[worldid, efcid]
+      cost0 = wp.where(start < 0.0, ctx_quad_in[worldid, efcid][0], 0.0)
 
       # active
       if x < 0.0:
-        out += _eval_cost(ctx_quad_in[worldid, efcid], alpha)
+        out += _eval_cost(ctx_quad_in[worldid, efcid], alpha) - cost0
+      else:
+        out -= cost0
 
   cost_out[worldid, alphaid] = out
 
 
 @wp.kernel
-def linesearch_parallel_best_alpha(
+def _linesearch_parallel_best_alpha(
   # Model:
   opt_ls_iterations: int,
   opt_ls_parallel_min_step: float,
@@ -464,6 +571,7 @@ def linesearch_parallel_best_alpha(
   cost_in: wp.array2d[float],
   # Out:
   ctx_alpha_out: wp.array[float],
+  ctx_improvement_out: wp.array[float],
 ):
   worldid = wp.tid()
 
@@ -471,14 +579,17 @@ def linesearch_parallel_best_alpha(
     return
 
   bestid = int(0)
-  best_cost = float(types.MJ_MAXVAL)
+  best_cost = float(0.0)
+  improved = bool(False)
   for i in range(opt_ls_iterations):
     cost = cost_in[worldid, i]
     if cost < best_cost:
       best_cost = cost
       bestid = i
+      improved = True
 
-  ctx_alpha_out[worldid] = _log_scale(opt_ls_parallel_min_step, 1.0, opt_ls_iterations, bestid)
+  ctx_alpha_out[worldid] = wp.where(improved, _log_scale(opt_ls_parallel_min_step, 1.0, opt_ls_iterations, bestid), 0.0)
+  ctx_improvement_out[worldid] = wp.where(improved, -best_cost, 0.0)
 
 
 def _linesearch_parallel(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.array2d[float]):
@@ -486,21 +597,21 @@ def _linesearch_parallel(m: types.Model, d: types.Data, ctx: SolverContext, cost
   dofs_per_thread = 20 if m.nv > 50 else 50
   threads_per_efc = ceil(m.nv / dofs_per_thread)
 
-  # quad_gauss = [gauss, search.T @ Ma - search.T @ qfrc_smooth, 0.5 * search.T @ mv]
+  # quad_gauss = [0, search.T @ Ma - search.T @ qfrc_smooth, 0.5 * search.T @ mv]
   if threads_per_efc > 1:
     ctx.quad_gauss.zero_()
 
   wp.launch(
-    linesearch_prepare_gauss(m.nv, dofs_per_thread),
+    _linesearch_prepare_gauss(m.nv, dofs_per_thread),
     dim=(d.nworld, threads_per_efc),
-    inputs=[d.qfrc_smooth, d.efc.Ma, ctx.search, ctx.gauss, ctx.mv, ctx.done],
+    inputs=[d.qfrc_smooth, d.efc.Ma, ctx.search, ctx.mv, ctx.done],
     outputs=[ctx.quad_gauss],
   )
 
   # quad = [0.5 * Jaref * Jaref * efc_D, jv * Jaref * efc_D, 0.5 * jv * jv * efc_D]
 
   wp.launch(
-    linesearch_prepare_quad,
+    _linesearch_prepare_quad,
     dim=(d.nworld, d.njmax),
     inputs=[
       m.opt.impratio_invsqrt,
@@ -520,7 +631,7 @@ def _linesearch_parallel(m: types.Model, d: types.Data, ctx: SolverContext, cost
   )
 
   wp.launch(
-    linesearch_parallel_fused,
+    _linesearch_parallel_fused,
     dim=(d.nworld, m.opt.ls_iterations),
     inputs=[
       m.opt.ls_iterations,
@@ -547,22 +658,22 @@ def _linesearch_parallel(m: types.Model, d: types.Data, ctx: SolverContext, cost
   )
 
   wp.launch(
-    linesearch_parallel_best_alpha,
-    dim=(d.nworld),
+    _linesearch_parallel_best_alpha,
+    dim=d.nworld,
     inputs=[m.opt.ls_iterations, m.opt.ls_parallel_min_step, ctx.done, cost],
-    outputs=[ctx.alpha],
+    outputs=[ctx.alpha, ctx.improvement],
   )
 
   # Teardown: update qacc, Ma, Jaref
   wp.launch(
-    linesearch_qacc_ma,
+    _linesearch_qacc_ma,
     dim=(d.nworld, m.nv),
     inputs=[ctx.search, ctx.mv, ctx.alpha, ctx.done],
     outputs=[d.qacc, d.efc.Ma],
   )
 
   wp.launch(
-    linesearch_jaref,
+    _linesearch_jaref,
     dim=(d.nworld, d.njmax),
     inputs=[d.nefc, ctx.jv, ctx.alpha, ctx.done],
     outputs=[ctx.Jaref],
@@ -582,23 +693,31 @@ def _compute_efc_eval_pt_pyramidal(
   ctx_Jaref: float,
   ctx_jv: float,
 ) -> wp.vec3:
-  """Compute for pyramidal cones (no elliptic contact data needed)."""
+  """Compute shifted cost, gradient, and hessian for pyramidal cones.
+
+  Returns (cost(alpha) - cost(0), grad(alpha), hessian(alpha)) summed across the row.
+  """
   # Limit/other constraint
   if efcid >= ne + nf:
     x = ctx_Jaref + alpha * ctx_jv
+    quad0 = _eval_pt_direct_cost_alpha_zero(ctx_Jaref, efc_D)
+    cost0 = wp.where(ctx_Jaref < 0.0, quad0, 0.0)
+    # _eval_pt_direct_shifted returns quad(alpha) - quad(0); add back quad(0) when the
+    # constraint was inactive at alpha=0 (i.e. cost(0) = 0) so we get quad(alpha) - 0.
+    offset = quad0 - cost0
     if x < 0.0:
-      return _eval_pt_direct(ctx_Jaref, ctx_jv, efc_D, alpha)
-    return wp.vec3(0.0)
+      return _eval_pt_direct_shifted(ctx_Jaref, ctx_jv, efc_D, alpha, offset)
+    return wp.vec3(-cost0, 0.0, 0.0)
 
   # Friction constraint - needs quad for frictionloss computation
   if efcid >= ne:
     f = efc_frictionloss[efcid]
     x = ctx_Jaref + alpha * ctx_jv
     rf = math.safe_div(f, efc_D)
-    return _eval_frictionloss_pt(x, f, rf, ctx_jv, efc_D)
+    return _shift_cost(_eval_frictionloss_pt(x, f, rf, ctx_jv, efc_D), _eval_frictionloss_cost(ctx_Jaref, f, rf, efc_D))
 
   # Equality constraint
-  return _eval_pt_direct(ctx_Jaref, ctx_jv, efc_D, alpha)
+  return _eval_pt_direct_shifted(ctx_Jaref, ctx_jv, efc_D, alpha, 0.0)
 
 
 @wp.func
@@ -621,20 +740,29 @@ def _compute_efc_eval_pt_elliptic(
   quad1: wp.vec3,
   quad2: wp.vec3,
 ) -> wp.vec3:
-  """Compute for elliptic cones (includes elliptic contact data)."""
+  """Compute shifted cost, gradient, and hessian for elliptic cones.
+
+  Returns (cost(alpha) - cost(0), grad(alpha), hessian(alpha)) summed across the row.
+  """
   # Contact/limit/other constraints
   if efcid >= ne + nf:
-    # Contact elliptic
     if efc_type == types.ConstraintType.CONTACT_ELLIPTIC:
       if efcid != efc_address0:  # Not primary row
         return wp.vec3(0.0)
-      return _eval_elliptic(impratio_invsqrt, contact_friction, ctx_quad, quad1, quad2, alpha)
+      mu = contact_friction[0] * impratio_invsqrt
+      cost0 = _eval_elliptic_cost(mu, ctx_quad, quad1, quad2, 0.0)
+      return _shift_cost(_eval_elliptic(mu, ctx_quad, quad1, quad2, alpha), cost0)
 
     # Limit/other constraint — direct eval (no quad read)
     x = ctx_Jaref + alpha * ctx_jv
+    efc_D = efc_D_in[efcid]
+    quad0 = _eval_pt_direct_cost_alpha_zero(ctx_Jaref, efc_D)
+    cost0 = wp.where(ctx_Jaref < 0.0, quad0, 0.0)
+    # See _compute_efc_eval_pt_pyramidal for the offset rationale.
+    offset = quad0 - cost0
     if x < 0.0:
-      return _eval_pt_direct(ctx_Jaref, ctx_jv, efc_D_in[efcid], alpha)
-    return wp.vec3(0.0)
+      return _eval_pt_direct_shifted(ctx_Jaref, ctx_jv, efc_D, alpha, offset)
+    return wp.vec3(-cost0, 0.0, 0.0)
 
   # Friction constraint - load D and frictionloss only here
   if efcid >= ne:
@@ -642,10 +770,11 @@ def _compute_efc_eval_pt_elliptic(
     f = efc_frictionloss[efcid]
     x = ctx_Jaref + alpha * ctx_jv
     rf = math.safe_div(f, efc_D)
-    return _eval_frictionloss_pt(x, f, rf, ctx_jv, efc_D)
+    return _shift_cost(_eval_frictionloss_pt(x, f, rf, ctx_jv, efc_D), _eval_frictionloss_cost(ctx_Jaref, f, rf, efc_D))
 
   # Equality constraint — direct eval (no quad read)
-  return _eval_pt_direct(ctx_Jaref, ctx_jv, efc_D_in[efcid], alpha)
+  efc_D = efc_D_in[efcid]
+  return _eval_pt_direct_shifted(ctx_Jaref, ctx_jv, efc_D, alpha, 0.0)
 
 
 @wp.func
@@ -698,11 +827,11 @@ def _compute_efc_eval_pt_alpha_zero_elliptic(
   """Optimized version for alpha=0.0, elliptic cones."""
   # Contact/limit/other constraints
   if efcid >= ne + nf:
-    # Contact elliptic
     if efc_type == types.ConstraintType.CONTACT_ELLIPTIC:
       if efcid != efc_address0:  # Not primary row
         return wp.vec3(0.0)
-      return _eval_elliptic(impratio_invsqrt, contact_friction, ctx_quad, quad1, quad2, 0.0)
+      mu = contact_friction[0] * impratio_invsqrt
+      return _eval_elliptic(mu, ctx_quad, quad1, quad2, 0.0)
 
     # Limit/other constraint — direct eval (no quad read)
     if ctx_Jaref < 0.0:
@@ -734,7 +863,7 @@ def _compute_efc_eval_pt_3alphas_pyramidal(
   ctx_Jaref: float,
   ctx_jv: float,
 ) -> tuple[wp.vec3, wp.vec3, wp.vec3]:
-  """Compute (cost, gradient, hessian) for 3 alphas, pyramidal cones.
+  """Compute shifted cost, gradient, and hessian for 3 alphas, pyramidal cones.
 
   Returns a tuple of 3 vec3s for (lo_alpha, hi_alpha, mid_alpha).
   Constraint types checked in order: limit/other -> friction -> equality.
@@ -744,11 +873,17 @@ def _compute_efc_eval_pt_3alphas_pyramidal(
     x_lo = ctx_Jaref + lo_alpha * ctx_jv
     x_hi = ctx_Jaref + hi_alpha * ctx_jv
     x_mid = ctx_Jaref + mid_alpha * ctx_jv
-    pt_lo, pt_hi, pt_mid = _eval_pt_direct_3alphas(ctx_Jaref, ctx_jv, efc_D, lo_alpha, hi_alpha, mid_alpha)
-    r_lo = wp.where(x_lo < 0.0, pt_lo, wp.vec3(0.0))
-    r_hi = wp.where(x_hi < 0.0, pt_hi, wp.vec3(0.0))
-    r_mid = wp.where(x_mid < 0.0, pt_mid, wp.vec3(0.0))
-    return (r_lo, r_hi, r_mid)
+    quad0 = _eval_pt_direct_cost_alpha_zero(ctx_Jaref, efc_D)
+    cost0 = wp.where(ctx_Jaref < 0.0, quad0, 0.0)
+    # See _compute_efc_eval_pt_pyramidal for the offset rationale.
+    offset = quad0 - cost0
+    pt_lo, pt_hi, pt_mid = _eval_pt_direct_shifted_3alphas(ctx_Jaref, ctx_jv, efc_D, lo_alpha, hi_alpha, mid_alpha, offset)
+    inactive = wp.vec3(-cost0, 0.0, 0.0)
+    return (
+      wp.where(x_lo < 0.0, pt_lo, inactive),
+      wp.where(x_hi < 0.0, pt_hi, inactive),
+      wp.where(x_mid < 0.0, pt_mid, inactive),
+    )
 
   # Friction constraint - needs quad for frictionloss computation
   if efcid >= ne:
@@ -757,10 +892,12 @@ def _compute_efc_eval_pt_3alphas_pyramidal(
     x_mid = ctx_Jaref + mid_alpha * ctx_jv
     f = efc_frictionloss[efcid]
     rf = math.safe_div(f, efc_D)
-    return _eval_frictionloss_pt_3alphas(x_lo, x_hi, x_mid, f, rf, ctx_jv, efc_D)
+    cost0 = _eval_frictionloss_cost(ctx_Jaref, f, rf, efc_D)
+    lo, hi, mid = _eval_frictionloss_pt_3alphas(x_lo, x_hi, x_mid, f, rf, ctx_jv, efc_D)
+    return (_shift_cost(lo, cost0), _shift_cost(hi, cost0), _shift_cost(mid, cost0))
 
   # Equality constraint: always active
-  return _eval_pt_direct_3alphas(ctx_Jaref, ctx_jv, efc_D, lo_alpha, hi_alpha, mid_alpha)
+  return _eval_pt_direct_shifted_3alphas(ctx_Jaref, ctx_jv, efc_D, lo_alpha, hi_alpha, mid_alpha, 0.0)
 
 
 @wp.func
@@ -785,7 +922,7 @@ def _compute_efc_eval_pt_3alphas_elliptic(
   quad1: wp.vec3,
   quad2: wp.vec3,
 ) -> tuple[wp.vec3, wp.vec3, wp.vec3]:
-  """Compute (cost, gradient, hessian) for 3 alphas, elliptic cones.
+  """Compute shifted cost, gradient, and hessian for 3 alphas, elliptic cones.
 
   Returns a tuple of 3 vec3s for (lo_alpha, hi_alpha, mid_alpha).
   Constraint types checked in order: contact elliptic/limit/other -> friction -> equality.
@@ -801,29 +938,39 @@ def _compute_efc_eval_pt_3alphas_elliptic(
     if efc_type == types.ConstraintType.CONTACT_ELLIPTIC:
       if efcid != efc_address0:  # secondary rows contribute nothing
         return (wp.vec3(0.0), wp.vec3(0.0), wp.vec3(0.0))
-      return (
-        _eval_elliptic(impratio_invsqrt, contact_friction, ctx_quad, quad1, quad2, lo_alpha),
-        _eval_elliptic(impratio_invsqrt, contact_friction, ctx_quad, quad1, quad2, hi_alpha),
-        _eval_elliptic(impratio_invsqrt, contact_friction, ctx_quad, quad1, quad2, mid_alpha),
-      )
+      mu = contact_friction[0] * impratio_invsqrt
+      cost0 = _eval_elliptic_cost(mu, ctx_quad, quad1, quad2, 0.0)
+      lo = _eval_elliptic(mu, ctx_quad, quad1, quad2, lo_alpha)
+      hi = _eval_elliptic(mu, ctx_quad, quad1, quad2, hi_alpha)
+      mid = _eval_elliptic(mu, ctx_quad, quad1, quad2, mid_alpha)
+      return (_shift_cost(lo, cost0), _shift_cost(hi, cost0), _shift_cost(mid, cost0))
 
     # Limit/other constraints — direct eval (no quad read)
     efc_D = efc_D_in[efcid]
-    pt_lo, pt_hi, pt_mid = _eval_pt_direct_3alphas(ctx_Jaref, ctx_jv, efc_D, lo_alpha, hi_alpha, mid_alpha)
-    r_lo = wp.where(x_lo < 0.0, pt_lo, wp.vec3(0.0))
-    r_hi = wp.where(x_hi < 0.0, pt_hi, wp.vec3(0.0))
-    r_mid = wp.where(x_mid < 0.0, pt_mid, wp.vec3(0.0))
-    return (r_lo, r_hi, r_mid)
+    quad0 = _eval_pt_direct_cost_alpha_zero(ctx_Jaref, efc_D)
+    cost0 = wp.where(ctx_Jaref < 0.0, quad0, 0.0)
+    # See _compute_efc_eval_pt_pyramidal for the offset rationale.
+    offset = quad0 - cost0
+    pt_lo, pt_hi, pt_mid = _eval_pt_direct_shifted_3alphas(ctx_Jaref, ctx_jv, efc_D, lo_alpha, hi_alpha, mid_alpha, offset)
+    inactive = wp.vec3(-cost0, 0.0, 0.0)
+    return (
+      wp.where(x_lo < 0.0, pt_lo, inactive),
+      wp.where(x_hi < 0.0, pt_hi, inactive),
+      wp.where(x_mid < 0.0, pt_mid, inactive),
+    )
 
   # Friction constraint - load D and frictionloss only here
   if efcid >= ne:
     efc_D = efc_D_in[efcid]
     f = efc_frictionloss[efcid]
     rf = math.safe_div(f, efc_D)
-    return _eval_frictionloss_pt_3alphas(x_lo, x_hi, x_mid, f, rf, ctx_jv, efc_D)
+    cost0 = _eval_frictionloss_cost(ctx_Jaref, f, rf, efc_D)
+    lo, hi, mid = _eval_frictionloss_pt_3alphas(x_lo, x_hi, x_mid, f, rf, ctx_jv, efc_D)
+    return (_shift_cost(lo, cost0), _shift_cost(hi, cost0), _shift_cost(mid, cost0))
 
   # Equality constraint — direct eval (no quad read)
-  return _eval_pt_direct_3alphas(ctx_Jaref, ctx_jv, efc_D_in[efcid], lo_alpha, hi_alpha, mid_alpha)
+  efc_D = efc_D_in[efcid]
+  return _eval_pt_direct_shifted_3alphas(ctx_Jaref, ctx_jv, efc_D, lo_alpha, hi_alpha, mid_alpha, 0.0)
 
 
 # kernel_analyzer: on
@@ -887,7 +1034,7 @@ def _compute_efc_eval_pt_3alphas_elliptic(
 
 
 @cache_kernel
-def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv: bool, is_sparse: bool):
+def _linesearch_iterative_kernel(ls_iterations: int, cone_type: types.ConeType, fuse_jv: bool, is_sparse: bool):
   """Factory for iterative linesearch kernel.
 
   Args:
@@ -946,7 +1093,6 @@ def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv:
     ctx_Jaref_in: wp.array2d[float],
     ctx_search_in: wp.array2d[float],
     ctx_search_dot_in: wp.array[float],
-    ctx_gauss_in: wp.array[float],
     ctx_mv_in: wp.array2d[float],
     ctx_jv_in: wp.array2d[float],
     ctx_quad_in: wp.array2d[wp.vec3],
@@ -958,6 +1104,7 @@ def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv:
     ctx_Jaref_out: wp.array2d[float],
     ctx_jv_out: wp.array2d[float],
     ctx_quad_out: wp.array2d[wp.vec3],
+    ctx_improvement_out: wp.array[float],
   ):
     worldid, tid = wp.tid()
 
@@ -1113,8 +1260,8 @@ def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv:
     p0_tile = wp.tile(local_p0, preserve_type=True)
     p0_sum = wp.tile_reduce(wp.add, p0_tile)
 
-    # quad_gauss = [gauss, search.T @ Ma - search.T @ qfrc_smooth, 0.5 * search.T @ mv]
-    local_gauss = wp.vec2(0.0)  # vec2 since component 0 is constant (ctx_gauss_in)
+    # quad_gauss = [0, search.T @ Ma - search.T @ qfrc_smooth, 0.5 * search.T @ mv]
+    local_gauss = wp.vec2(0.0)
     for dofid in range(tid, nv, wp.block_dim()):
       search = ctx_search_in[worldid, dofid]
       local_gauss += wp.vec2(
@@ -1125,10 +1272,11 @@ def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv:
     gauss_tile = wp.tile(local_gauss, preserve_type=True)
     gauss_sum = wp.tile_reduce(wp.add, gauss_tile)
     gauss_reduced = gauss_sum[0]
-    ctx_quad_gauss = wp.vec3(ctx_gauss_in[worldid], gauss_reduced[0], gauss_reduced[1])
+    ctx_quad_gauss = wp.vec3(0.0, gauss_reduced[0], gauss_reduced[1])
 
     # add quad_gauss contribution to p0
     p0 = wp.vec3(ctx_quad_gauss[0], ctx_quad_gauss[1], 2.0 * ctx_quad_gauss[2]) + p0_sum[0]
+    p0_delta = wp.vec3(0.0, p0[1], p0[2])
 
     # lo_in at lo_alpha_in = -p0[1] / p0[2]
     lo_alpha_in = -math.safe_div(p0[1], p0[2])
@@ -1189,17 +1337,18 @@ def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv:
     lo_in = _eval_pt(ctx_quad_gauss, lo_alpha_in) + lo_in_sum[0]
 
     # accept Newton step if derivative is small and cost improved
-    initial_converged = wp.abs(lo_in[1]) < gtol and lo_in[0] < p0[0]
+    initial_converged = wp.abs(lo_in[1]) < gtol and lo_in[0] < 0.0
 
     # main iterative loop - skip if already converged
     if not initial_converged:
       alpha = float(0.0)
+      improvement = float(0.0)
 
       # initialize bounds
       lo_less = lo_in[1] < p0[1]
-      lo = wp.where(lo_less, lo_in, p0)
+      lo = wp.where(lo_less, lo_in, p0_delta)
       lo_alpha = wp.where(lo_less, lo_alpha_in, 0.0)
-      hi = wp.where(lo_less, p0, lo_in)
+      hi = wp.where(lo_less, p0_delta, lo_in)
       hi_alpha = wp.where(lo_less, 0.0, lo_alpha_in)
 
       for _ in range(LS_ITERATIONS):
@@ -1322,15 +1471,18 @@ def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv:
         ls_done = (not swap_lo and not swap_hi) or (lo[1] < 0.0 and lo[1] > -gtol) or (hi[1] > 0.0 and hi[1] < gtol)
 
         # update alpha if improved
-        improved = lo[0] < p0[0] or hi[0] < p0[0]
+        improved = lo[0] < 0.0 or hi[0] < 0.0
         lo_better = lo[0] < hi[0]
-        alpha = wp.where(improved and lo_better, lo_alpha, alpha)
-        alpha = wp.where(improved and not lo_better, hi_alpha, alpha)
+        best_alpha = wp.where(lo_better, lo_alpha, hi_alpha)
+        best_delta = wp.where(lo_better, lo[0], hi[0])
+        alpha = wp.where(improved, best_alpha, alpha)
+        improvement = wp.where(improved, -best_delta, improvement)
 
         if ls_done:
           break
     else:
       alpha = lo_alpha_in
+      improvement = -lo_in[0]
 
     # qacc and Ma update
     for dofid in range(tid, nv, wp.block_dim()):
@@ -1341,6 +1493,9 @@ def linesearch_iterative(ls_iterations: int, cone_type: types.ConeType, fuse_jv:
     for efcid in range(tid, nefc, wp.block_dim()):
       ctx_Jaref_out[worldid, efcid] += alpha * ctx_jv_in[worldid, efcid]
 
+    if tid == 0:
+      ctx_improvement_out[worldid] = improvement
+
   return kernel
 
 
@@ -1354,7 +1509,7 @@ def _linesearch_iterative(m: types.Model, d: types.Data, ctx: SolverContext, fus
     fuse_jv: Whether jv is computed in-kernel (True) or pre-computed (False).
   """
   wp.launch_tiled(
-    linesearch_iterative(m.opt.ls_iterations, m.opt.cone, fuse_jv, m.is_sparse),
+    _linesearch_iterative_kernel(m.opt.ls_iterations, m.opt.cone, fuse_jv, m.is_sparse),
     dim=d.nworld,
     inputs=[
       m.nv,
@@ -1382,19 +1537,18 @@ def _linesearch_iterative(m: types.Model, d: types.Data, ctx: SolverContext, fus
       ctx.Jaref,
       ctx.search,
       ctx.search_dot,
-      ctx.gauss,
       ctx.mv,
       ctx.jv,
       ctx.quad,
       ctx.done,
     ],
-    outputs=[d.qacc, d.efc.Ma, ctx.Jaref, ctx.jv, ctx.quad],
+    outputs=[d.qacc, d.efc.Ma, ctx.Jaref, ctx.jv, ctx.quad, ctx.improvement],
     block_dim=m.block_dim.linesearch_iterative,
   )
 
 
 @wp.kernel
-def linesearch_zero_jv(
+def _linesearch_zero_jv(
   # Data in:
   nefc_in: wp.array[int],
   # In:
@@ -1414,7 +1568,7 @@ def linesearch_zero_jv(
 
 
 @cache_kernel
-def linesearch_jv_fused(is_sparse: bool, nv: int, dofs_per_thread: int):
+def _linesearch_jv_fused_kernel(is_sparse: bool, nv: int, dofs_per_thread: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Data in:
@@ -1475,7 +1629,7 @@ def linesearch_jv_fused(is_sparse: bool, nv: int, dofs_per_thread: int):
 
 
 @cache_kernel
-def linesearch_prepare_gauss(nv: int, dofs_per_thread: int):
+def _linesearch_prepare_gauss(nv: int, dofs_per_thread: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Data in:
@@ -1483,7 +1637,6 @@ def linesearch_prepare_gauss(nv: int, dofs_per_thread: int):
     efc_Ma_in: wp.array2d[float],
     # In:
     ctx_search_in: wp.array2d[float],
-    ctx_gauss_in: wp.array[float],
     ctx_mv_in: wp.array2d[float],
     ctx_done_in: wp.array[bool],
     # Out:
@@ -1503,8 +1656,7 @@ def linesearch_prepare_gauss(nv: int, dofs_per_thread: int):
         quad_gauss_1 += search * (efc_Ma_in[worldid, i] - qfrc_smooth_in[worldid, i])
         quad_gauss_2 += 0.5 * search * ctx_mv_in[worldid, i]
 
-      quad_gauss_0 = ctx_gauss_in[worldid]
-      ctx_quad_gauss_out[worldid] = wp.vec3(quad_gauss_0, quad_gauss_1, quad_gauss_2)
+      ctx_quad_gauss_out[worldid] = wp.vec3(0.0, quad_gauss_1, quad_gauss_2)
 
     else:
       for i in range(wp.static(dofs_per_thread)):
@@ -1514,17 +1666,13 @@ def linesearch_prepare_gauss(nv: int, dofs_per_thread: int):
           quad_gauss_1 += search * (efc_Ma_in[worldid, ii] - qfrc_smooth_in[worldid, ii])
           quad_gauss_2 += 0.5 * search * ctx_mv_in[worldid, ii]
 
-      if dofstart == 0:
-        quad_gauss_0 = ctx_gauss_in[worldid]
-        wp.atomic_add(ctx_quad_gauss_out, worldid, wp.vec3(quad_gauss_0, quad_gauss_1, quad_gauss_2))
-      else:
-        wp.atomic_add(ctx_quad_gauss_out, worldid, wp.vec3(0.0, quad_gauss_1, quad_gauss_2))
+      wp.atomic_add(ctx_quad_gauss_out, worldid, wp.vec3(0.0, quad_gauss_1, quad_gauss_2))
 
   return kernel
 
 
 @wp.kernel
-def linesearch_prepare_quad(
+def _linesearch_prepare_quad(
   # Model:
   opt_impratio_invsqrt: wp.array[float],
   # Data in:
@@ -1620,7 +1768,7 @@ def linesearch_prepare_quad(
 
 
 @wp.kernel
-def linesearch_qacc_ma(
+def _linesearch_qacc_ma(
   # In:
   ctx_search_in: wp.array2d[float],
   ctx_mv_in: wp.array2d[float],
@@ -1641,7 +1789,7 @@ def linesearch_qacc_ma(
 
 
 @wp.kernel
-def linesearch_jaref(
+def _linesearch_jaref(
   # Data in:
   nefc_in: wp.array[int],
   # In:
@@ -1687,14 +1835,14 @@ def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.arra
 
     if threads_per_efc > 1:
       wp.launch(
-        linesearch_zero_jv,
+        _linesearch_zero_jv,
         dim=(d.nworld, d.njmax),
         inputs=[d.nefc, ctx.done],
         outputs=[ctx.jv],
       )
 
     wp.launch(
-      linesearch_jv_fused(m.is_sparse, m.nv, dofs_per_thread),
+      _linesearch_jv_fused_kernel(m.is_sparse, m.nv, dofs_per_thread),
       dim=(d.nworld, d.njmax, threads_per_efc),
       inputs=[d.nefc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, ctx.search, ctx.done],
       outputs=[ctx.jv],
@@ -1707,23 +1855,21 @@ def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.arra
 
 
 @wp.kernel
-def solve_init_efc(
+def _solve_init_efc(
   # Data out:
   solver_niter_out: wp.array[int],
   # Out:
   ctx_search_dot_out: wp.array[float],
-  ctx_cost_out: wp.array[float],
   ctx_done_out: wp.array[bool],
 ):
   worldid = wp.tid()
-  ctx_cost_out[worldid] = types.MJ_MAXVAL
   solver_niter_out[worldid] = 0
   ctx_done_out[worldid] = False
   ctx_search_dot_out[worldid] = 0.0
 
 
 @cache_kernel
-def solve_init_jaref(is_sparse: bool, nv: int, dofs_per_thread: int):
+def _solve_init_jaref_kernel(is_sparse: bool, nv: int, dofs_per_thread: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Data in:
@@ -1772,7 +1918,7 @@ def solve_init_jaref(is_sparse: bool, nv: int, dofs_per_thread: int):
 
 
 @wp.kernel
-def solve_init_search(
+def _solve_init_search(
   # In:
   ctx_Mgrad_in: wp.array2d[float],
   # Out:
@@ -1785,28 +1931,8 @@ def solve_init_search(
   wp.atomic_add(ctx_search_dot_out, worldid, search * search)
 
 
-@wp.kernel
-def update_constraint_init_cost(
-  # In:
-  ctx_cost_in: wp.array[float],
-  ctx_done_in: wp.array[bool],
-  # Out:
-  ctx_gauss_out: wp.array[float],
-  ctx_cost_out: wp.array[float],
-  ctx_prev_cost_out: wp.array[float],
-):
-  worldid = wp.tid()
-
-  if ctx_done_in[worldid]:
-    return
-
-  ctx_gauss_out[worldid] = 0.0
-  ctx_prev_cost_out[worldid] = ctx_cost_in[worldid]
-  ctx_cost_out[worldid] = 0.0
-
-
 @cache_kernel
-def update_constraint_efc(track_changes: bool):
+def _update_constraint_efc(track_changes: bool):
   TRACK_CHANGES = track_changes
 
   @wp.kernel(module="unique", enable_backward=False)
@@ -1832,7 +1958,6 @@ def update_constraint_efc(track_changes: bool):
     efc_force_out: wp.array2d[float],
     efc_state_out: wp.array2d[int],
     # Out:
-    ctx_cost_out: wp.array[float],
     changed_ids_out: wp.array2d[int],
     changed_count_out: wp.array[int],
   ):
@@ -1854,56 +1979,33 @@ def update_constraint_efc(track_changes: bool):
     ne = ne_in[worldid]
     nf = nf_in[worldid]
 
-    new_state = types.ConstraintState.SATISFIED.value
+    is_equality = efcid < ne
+    is_friction = (not is_equality) and (efcid < ne + nf)
+    is_elliptic = efc_type_in[worldid, efcid] == types.ConstraintType.CONTACT_ELLIPTIC
 
-    if efcid < ne:
-      # equality
-      efc_force_out[worldid, efcid] = -efc_D * Jaref
-      new_state = types.ConstraintState.QUADRATIC.value
-      wp.atomic_add(ctx_cost_out, worldid, 0.5 * efc_D * Jaref * Jaref)
-    elif efcid < ne + nf:
-      # friction
-      f = efc_frictionloss_in[worldid, efcid]
-      rf = math.safe_div(f, efc_D)
-      if Jaref <= -rf:
-        efc_force_out[worldid, efcid] = f
-        new_state = types.ConstraintState.LINEARNEG.value
-        wp.atomic_add(ctx_cost_out, worldid, -f * (0.5 * rf + Jaref))
-      elif Jaref >= rf:
-        efc_force_out[worldid, efcid] = -f
-        new_state = types.ConstraintState.LINEARPOS.value
-        wp.atomic_add(ctx_cost_out, worldid, -f * (0.5 * rf - Jaref))
-      else:
-        efc_force_out[worldid, efcid] = -efc_D * Jaref
-        new_state = types.ConstraintState.QUADRATIC.value
-        wp.atomic_add(ctx_cost_out, worldid, 0.5 * efc_D * Jaref * Jaref)
-    elif efc_type_in[worldid, efcid] != types.ConstraintType.CONTACT_ELLIPTIC:
-      # limit, frictionless contact, pyramidal friction cone contact
-      if Jaref >= 0.0:
-        efc_force_out[worldid, efcid] = 0.0
-        new_state = types.ConstraintState.SATISFIED.value
-      else:
-        efc_force_out[worldid, efcid] = -efc_D * Jaref
-        new_state = types.ConstraintState.QUADRATIC.value
-        wp.atomic_add(ctx_cost_out, worldid, 0.5 * efc_D * Jaref * Jaref)
-    else:  # elliptic friction cone contact
+    frictionloss = efc_frictionloss_in[worldid, efcid] if is_friction else 0.0
+
+    efcid0 = -1
+    jaref0 = float(0.0)
+    D0 = float(0.0)
+    mu = float(0.0)
+    ufrictionj = float(0.0)
+    TT = float(0.0)
+
+    if is_elliptic:
       conid = efc_id_in[worldid, efcid]
-
       if conid >= nacon_in[0]:
         return
+      efcid0 = contact_efc_address_in[conid, 0]
+      if efcid0 < 0:
+        return
 
       dim = contact_dim_in[conid]
       friction = contact_friction_in[conid]
       mu = friction[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
+      jaref0 = ctx_Jaref_in[worldid, efcid0]
+      D0 = efc_D_in[worldid, efcid0]
 
-      efcid0 = contact_efc_address_in[conid, 0]
-      if efcid0 < 0:
-        return
-
-      N = ctx_Jaref_in[worldid, efcid0] * mu
-
-      ufrictionj = float(0.0)
-      TT = float(0.0)
       for j in range(1, dim):
         efcidj = contact_efc_address_in[conid, j]
         if efcidj < 0:
@@ -1914,35 +2016,24 @@ def update_constraint_efc(track_changes: bool):
         if efcid == efcidj:
           ufrictionj = uj * frictionj
 
-      if TT <= 0.0:
-        T = 0.0
-      else:
-        T = wp.sqrt(TT)
-
-      # top zone
-      if (N >= mu * T) or ((T <= 0.0) and (N >= 0.0)):
-        efc_force_out[worldid, efcid] = 0.0
-        new_state = types.ConstraintState.SATISFIED.value
-      # bottom zone
-      elif (mu * N + T <= 0.0) or ((T <= 0.0) and (N < 0.0)):
-        efc_force_out[worldid, efcid] = -efc_D * Jaref
-        new_state = types.ConstraintState.QUADRATIC.value
-        wp.atomic_add(ctx_cost_out, worldid, 0.5 * efc_D * Jaref * Jaref)
-      # middle zone
-      else:
-        dm = math.safe_div(efc_D_in[worldid, efcid0], mu * mu * (1.0 + mu * mu))
-        nmt = N - mu * T
-
-        force = -dm * nmt * mu
-
-        if efcid == efcid0:
-          efc_force_out[worldid, efcid] = force
-          wp.atomic_add(ctx_cost_out, worldid, 0.5 * dm * nmt * nmt)
-        else:
-          efc_force_out[worldid, efcid] = -math.safe_div(force, T) * ufrictionj
-
-        new_state = types.ConstraintState.CONE.value
+    res = _eval_constraint(
+      is_equality,
+      is_friction,
+      is_elliptic,
+      Jaref,
+      efc_D,
+      frictionloss,
+      efcid,
+      efcid0,
+      jaref0,
+      D0,
+      mu,
+      ufrictionj,
+      TT,
+    )
 
+    new_state = int(res[1])
+    efc_force_out[worldid, efcid] = res[0]
     efc_state_out[worldid, efcid] = new_state
 
     if wp.static(TRACK_CHANGES):
@@ -1955,7 +2046,7 @@ def update_constraint_efc(track_changes: bool):
 
 
 @wp.kernel
-def update_constraint_init_qfrc_constraint_sparse(
+def _update_constraint_init_qfrc_constraint_sparse(
   # Data in:
   nefc_in: wp.array[int],
   efc_J_rownnz_in: wp.array2d[int],
@@ -1988,7 +2079,7 @@ def update_constraint_init_qfrc_constraint_sparse(
 
 
 @wp.kernel
-def update_constraint_init_qfrc_constraint_dense(
+def _update_constraint_init_qfrc_constraint_dense(
   # Data in:
   nefc_in: wp.array[int],
   efc_J_in: wp.array3d[float],
@@ -2013,49 +2104,8 @@ def update_constraint_init_qfrc_constraint_dense(
   qfrc_constraint_out[worldid, dofid] = sum_qfrc
 
 
-@cache_kernel
-def update_constraint_gauss_cost(nv: int, dofs_per_thread: int):
-  @wp.kernel(module="unique", enable_backward=False)
-  def kernel(
-    # Data in:
-    qacc_in: wp.array2d[float],
-    qfrc_smooth_in: wp.array2d[float],
-    qacc_smooth_in: wp.array2d[float],
-    efc_Ma_in: wp.array2d[float],
-    # In:
-    ctx_done_in: wp.array[bool],
-    # Out:
-    ctx_gauss_out: wp.array[float],
-    ctx_cost_out: wp.array[float],
-  ):
-    worldid, dofstart = wp.tid()
-
-    if ctx_done_in[worldid]:
-      return
-
-    gauss_cost = float(0.0)
-
-    if wp.static(dofs_per_thread >= nv):
-      for i in range(wp.static(min(dofs_per_thread, nv))):
-        gauss_cost += (efc_Ma_in[worldid, i] - qfrc_smooth_in[worldid, i]) * (qacc_in[worldid, i] - qacc_smooth_in[worldid, i])
-      ctx_gauss_out[worldid] += 0.5 * gauss_cost
-      ctx_cost_out[worldid] += 0.5 * gauss_cost
-
-    else:
-      for i in range(wp.static(dofs_per_thread)):
-        ii = dofstart * wp.static(dofs_per_thread) + i
-        if ii < nv:
-          gauss_cost += (efc_Ma_in[worldid, ii] - qfrc_smooth_in[worldid, ii]) * (
-            qacc_in[worldid, ii] - qacc_smooth_in[worldid, ii]
-          )
-      wp.atomic_add(ctx_gauss_out, worldid, 0.5 * gauss_cost)
-      wp.atomic_add(ctx_cost_out, worldid, 0.5 * gauss_cost)
-
-  return kernel
-
-
 @wp.kernel
-def update_gradient_h_incremental(
+def _update_gradient_h_incremental(
   # Data in:
   efc_J_in: wp.array3d[float],
   efc_D_in: wp.array2d[float],
@@ -2102,7 +2152,7 @@ def update_gradient_h_incremental(
 
 
 @wp.kernel
-def update_gradient_h_incremental_sparse(
+def _update_gradient_h_incremental_sparse(
   # Data in:
   efc_J_rownnz_in: wp.array2d[int],
   efc_J_rowadr_in: wp.array2d[int],
@@ -2156,13 +2206,6 @@ def update_gradient_h_incremental_sparse(
 
 def _update_constraint(m: types.Model, d: types.Data, ctx: SolverContext | InverseContext, track_changes: bool = False):
   """Update constraint arrays after each solve iteration."""
-  wp.launch(
-    update_constraint_init_cost,
-    dim=(d.nworld),
-    inputs=[ctx.cost, ctx.done],
-    outputs=[ctx.gauss, ctx.cost, ctx.prev_cost],
-  )
-
   efc_inputs = [
     m.opt.impratio_invsqrt,
     d.ne,
@@ -2181,49 +2224,32 @@ def _update_constraint(m: types.Model, d: types.Data, ctx: SolverContext | Inver
   ]
 
   wp.launch(
-    update_constraint_efc(track_changes),
+    _update_constraint_efc(track_changes),
     dim=(d.nworld, d.njmax),
     inputs=efc_inputs,
-    outputs=[d.efc.force, d.efc.state, ctx.cost, ctx.changed_efc_ids, ctx.changed_efc_count],
+    outputs=[d.efc.force, d.efc.state, ctx.changed_efc_ids, ctx.changed_efc_count],
   )
 
   # qfrc_constraint = efc_J.T @ efc_force
   if m.is_sparse:
     d.qfrc_constraint.zero_()
     wp.launch(
-      update_constraint_init_qfrc_constraint_sparse,
+      _update_constraint_init_qfrc_constraint_sparse,
       dim=(d.nworld, d.njmax),
       inputs=[d.nefc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, d.efc.force, ctx.done],
       outputs=[d.qfrc_constraint],
     )
   else:
     wp.launch(
-      update_constraint_init_qfrc_constraint_dense,
+      _update_constraint_init_qfrc_constraint_dense,
       dim=(d.nworld, m.nv),
       inputs=[d.nefc, d.efc.J, d.efc.force, d.njmax, ctx.done],
       outputs=[d.qfrc_constraint],
     )
 
-  # if we are only using 1 thread, it makes sense to do more dofs and skip the atomics.
-  # For more than 1 thread, dofs_per_thread is lower for better load balancing.
-  if m.nv > 50:
-    dofs_per_thread = 20
-  else:
-    dofs_per_thread = 50
-
-  threads_per_efc = ceil(m.nv / dofs_per_thread)
-
-  # gauss = 0.5 * (Ma - qfrc_smooth).T @ (qacc - qacc_smooth)
-  wp.launch(
-    update_constraint_gauss_cost(m.nv, dofs_per_thread),
-    dim=(d.nworld, threads_per_efc),
-    inputs=[d.qacc, d.qfrc_smooth, d.qacc_smooth, d.efc.Ma, ctx.done],
-    outputs=[ctx.gauss, ctx.cost],
-  )
-
 
 @wp.kernel
-def update_gradient_zero_grad_dot(
+def _update_gradient_zero_grad_dot(
   # In:
   ctx_done_in: wp.array[bool],
   # Out:
@@ -2238,7 +2264,7 @@ def update_gradient_zero_grad_dot(
 
 
 @wp.kernel
-def update_gradient_grad(
+def _update_gradient_grad(
   # Data in:
   qfrc_smooth_in: wp.array2d[float],
   qfrc_constraint_in: wp.array2d[float],
@@ -2260,11 +2286,10 @@ def update_gradient_grad(
 
 
 @wp.kernel
-def update_gradient_set_h_M_upper_sparse(
+def _update_gradient_init_h_sparse(
   # Model:
-  M_fullm_upper_i: wp.array[int],
-  M_fullm_upper_j: wp.array[int],
-  M_fullm_upper_elemid: wp.array[int],
+  nv: int,
+  M_elemid: wp.array2d[int],
   # Data in:
   M_in: wp.array3d[float],
   # In:
@@ -2272,19 +2297,29 @@ def update_gradient_set_h_M_upper_sparse(
   # Out:
   ctx_h_out: wp.array3d[float],
 ):
-  worldid, elementid = wp.tid()
+  worldid, i, j = wp.tid()
 
   if ctx_done_in[worldid]:
     return
 
-  i = M_fullm_upper_i[elementid]
-  j = M_fullm_upper_j[elementid]
-  madr = M_fullm_upper_elemid[elementid]
-  ctx_h_out[worldid, i, j] += M_in[worldid, 0, madr]
+  # only write the upper triangle; Cholesky reads the upper triangle only
+  if j < i:
+    return
+
+  if i >= nv or j >= nv:
+    ctx_h_out[worldid, i, j] = 0.0
+    return
+
+  # M is stored in the lower triangle, so transpose the lookup for the upper
+  elemid = M_elemid[j, i]
+  if elemid >= 0:
+    ctx_h_out[worldid, i, j] = M_in[worldid, 0, elemid]
+  else:
+    ctx_h_out[worldid, i, j] = 0.0
 
 
 @wp.func
-def state_check(D: float, state: int) -> float:
+def _state_check(D: float, state: int) -> float:
   if state == types.ConstraintState.QUADRATIC.value:
     return D
   else:
@@ -2292,7 +2327,7 @@ def state_check(D: float, state: int) -> float:
 
 
 @wp.func
-def active_check(tid: int, threshold: int) -> float:
+def _active_check(tid: int, threshold: int) -> float:
   if tid >= threshold:
     return 0.0
   else:
@@ -2300,83 +2335,13 @@ def active_check(tid: int, threshold: int) -> float:
 
 
 @cache_kernel
-def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int):
-  TILE_SIZE = tile_size
-
-  @wp.kernel(module="unique", enable_backward=False)
-  def kernel(
-    # Data in:
-    nefc_in: wp.array[int],
-    efc_J_in: wp.array3d[float],
-    efc_D_in: wp.array2d[float],
-    efc_state_in: wp.array2d[int],
-    # In:
-    ctx_done_in: wp.array[bool],
-    # Out:
-    ctx_h_out: wp.array3d[float],
-  ):
-    worldid, elementid = wp.tid()
-
-    if ctx_done_in[worldid]:
-      return
-
-    nefc = nefc_in[worldid]
-
-    # Upper-triangle tile index: elementid -> (row, col) where row <= col.
-    col = (int(sqrt(float(1 + 8 * elementid))) - 1) // 2
-    row = elementid - (col * (col + 1)) // 2
-
-    offset_row = row * TILE_SIZE
-    offset_col = col * TILE_SIZE
-
-    sum_val = wp.tile_zeros(shape=(TILE_SIZE, TILE_SIZE), dtype=wp.float32)
-
-    # Each tile processes looping over all constraints, producing 1 output tile
-    for k in range(0, njmax, TILE_SIZE):
-      if k >= nefc:
-        break
-
-      # AD: leaving bounds-check disabled here because I'm not entirely sure that
-      # everything always hits the fast path. The padding takes care of any
-      # potential OOB accesses.
-      J_krow = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(k, offset_row), bounds_check=False)
-
-      if offset_row != offset_col:
-        J_kcol = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(k, offset_col), bounds_check=False)
-      else:
-        wp.tile_assign(J_kcol, J_krow, (0, 0))
-
-      D_k = wp.tile_load(efc_D_in[worldid], shape=TILE_SIZE, offset=k, bounds_check=False)
-      state = wp.tile_load(efc_state_in[worldid], shape=TILE_SIZE, offset=k, bounds_check=False)
-
-      D_k = wp.tile_map(state_check, D_k, state)
-
-      # force unused elements to be zero
-      tid_tile = wp.tile_arange(TILE_SIZE, dtype=int)
-      threshold_tile = wp.tile_ones(shape=TILE_SIZE, dtype=int) * (nefc - k)
-
-      active_tile = wp.tile_map(active_check, tid_tile, threshold_tile)
-      D_k = wp.tile_map(wp.mul, active_tile, D_k)
-
-      J_krow = wp.tile_map(wp.mul, wp.tile_transpose(J_krow), wp.tile_broadcast(D_k, shape=(TILE_SIZE, TILE_SIZE)))
-
-      sum_val += wp.tile_matmul(J_krow, J_kcol)
-
-    # AD: setting bounds_check to True explicitly here because for some reason it was
-    # slower to disable it.
-    wp.tile_store(ctx_h_out[worldid], sum_val, offset=(offset_row, offset_col), bounds_check=True)
-
-  return kernel
-
-
-@cache_kernel
-def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int):
+def _update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int):
   if njmax < tile_size:
     tile_size = njmax
 
   TILE_SIZE_K = tile_size
 
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, module_options={"enable_mathdx_gemm": False})
   def kernel(
     # Data in:
     nefc_in: wp.array[int],
@@ -2412,13 +2377,13 @@ def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int):
       D_k = wp.tile_load(efc_D_in[worldid], shape=TILE_SIZE_K, offset=k, bounds_check=False)
       state = wp.tile_load(efc_state_in[worldid], shape=TILE_SIZE_K, offset=k, bounds_check=False)
 
-      D_k = wp.tile_map(state_check, D_k, state)
+      D_k = wp.tile_map(_state_check, D_k, state)
 
       # force unused elements to be zero
       tid_tile = wp.tile_arange(TILE_SIZE_K, dtype=int)
       threshold_tile = wp.tile_ones(shape=TILE_SIZE_K, dtype=int) * (nefc - k)
 
-      active_tile = wp.tile_map(active_check, tid_tile, threshold_tile)
+      active_tile = wp.tile_map(_active_check, tid_tile, threshold_tile)
       D_k = wp.tile_map(wp.mul, active_tile, D_k)
 
       J_ki = wp.tile_map(wp.mul, wp.tile_transpose(J_kj), wp.tile_broadcast(D_k, shape=(nv_pad, TILE_SIZE_K)))
@@ -2432,7 +2397,7 @@ def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int):
 
 # TODO(thowell): combine with JTDAJ ?
 @wp.kernel
-def update_gradient_JTCJ_sparse(
+def _update_gradient_JTCJ_sparse(
   # Model:
   opt_impratio_invsqrt: wp.array[float],
   dof_tri_row: wp.array[int],
@@ -2594,7 +2559,7 @@ def update_gradient_JTCJ_sparse(
 
 
 @wp.kernel
-def update_gradient_JTCJ_dense(
+def _update_gradient_JTCJ_dense(
   # Model:
   opt_impratio_invsqrt: wp.array[float],
   dof_tri_row: wp.array[int],
@@ -2734,7 +2699,7 @@ def update_gradient_JTCJ_dense(
 
 
 @cache_kernel
-def update_gradient_cholesky(tile_size: int):
+def _update_gradient_cholesky(tile_size: int):
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # In:
@@ -2760,8 +2725,8 @@ def update_gradient_cholesky(tile_size: int):
 
 
 @cache_kernel
-def update_gradient_cholesky_blocked(tile_size: int, matrix_size: int):
-  @wp.kernel(module="unique", enable_backward=False)
+def _update_gradient_cholesky_blocked(tile_size: int, matrix_size: int):
+  @wp.kernel(module="unique", enable_backward=False, module_options={"enable_mathdx_gemm": False})
   def kernel(
     # In:
     ctx_done_in: wp.array[bool],
@@ -2790,10 +2755,10 @@ def update_gradient_cholesky_blocked(tile_size: int, matrix_size: int):
 
 
 @cache_kernel
-def update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size: int):
+def _update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size: int):
   """Blocked Cholesky that skips factorization when no constraints changed."""
 
-  @wp.kernel(module="unique", enable_backward=False)
+  @wp.kernel(module="unique", enable_backward=False, module_options={"enable_mathdx_gemm": False})
   def kernel(
     # In:
     ctx_done_in: wp.array[bool],
@@ -2823,7 +2788,7 @@ def update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size:
 
 
 @wp.kernel
-def padding_h(nv: int, ctx_done_in: wp.array[bool], ctx_h_out: wp.array3d[float]):
+def _padding_h(nv: int, ctx_done_in: wp.array[bool], ctx_h_out: wp.array3d[float]):
   worldid, elementid = wp.tid()
 
   if ctx_done_in[worldid]:
@@ -2841,7 +2806,7 @@ def _cholesky_factorize_solve(m: types.Model, d: types.Data, ctx: SolverContext,
   """
   if m.nv <= _BLOCK_CHOLESKY_DIM:
     wp.launch_tiled(
-      update_gradient_cholesky(m.nv),
+      _update_gradient_cholesky(m.nv),
       dim=d.nworld,
       inputs=[ctx.grad, ctx.h, ctx.done],
       outputs=[ctx.Mgrad],
@@ -2849,7 +2814,7 @@ def _cholesky_factorize_solve(m: types.Model, d: types.Data, ctx: SolverContext,
     )
   else:
     wp.launch(
-      padding_h,
+      _padding_h,
       dim=(d.nworld, m.nv_pad - m.nv),
       inputs=[m.nv, ctx.done],
       outputs=[ctx.h],
@@ -2857,7 +2822,7 @@ def _cholesky_factorize_solve(m: types.Model, d: types.Data, ctx: SolverContext,
 
     if skip_unchanged:
       wp.launch_tiled(
-        update_gradient_cholesky_blocked_skip_unchanged(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad),
+        _update_gradient_cholesky_blocked_skip_unchanged(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad),
         dim=d.nworld,
         inputs=[ctx.done, ctx.grad.reshape(shape=(d.nworld, ctx.grad.shape[1], 1)), ctx.h, ctx.changed_efc_count, ctx.hfactor],
         outputs=[ctx.Mgrad.reshape(shape=(d.nworld, ctx.Mgrad.shape[1], 1))],
@@ -2865,7 +2830,7 @@ def _cholesky_factorize_solve(m: types.Model, d: types.Data, ctx: SolverContext,
       )
     else:
       wp.launch_tiled(
-        update_gradient_cholesky_blocked(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad),
+        _update_gradient_cholesky_blocked(types.TILE_SIZE_JTDAJ_DENSE, m.nv_pad),
         dim=d.nworld,
         inputs=[ctx.done, ctx.grad.reshape(shape=(d.nworld, ctx.grad.shape[1], 1)), ctx.h, ctx.hfactor],
         outputs=[ctx.Mgrad.reshape(shape=(d.nworld, ctx.Mgrad.shape[1], 1))],
@@ -2899,7 +2864,7 @@ def _JTDAJ_sparse(
   efc_D = efc_D_in[worldid, efcid]
   efc_state = efc_state_in[worldid, efcid]
 
-  if state_check(efc_D, efc_state) == 0.0:
+  if _state_check(efc_D, efc_state) == 0.0:
     return
 
   rownnz = efc_J_rownnz_in[worldid, efcid]
@@ -2928,10 +2893,10 @@ def _JTDAJ_sparse(
 
 def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
   # grad = Ma - qfrc_smooth - qfrc_constraint
-  wp.launch(update_gradient_zero_grad_dot, dim=(d.nworld), inputs=[ctx.done], outputs=[ctx.grad_dot])
+  wp.launch(_update_gradient_zero_grad_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.grad_dot])
 
   wp.launch(
-    update_gradient_grad,
+    _update_gradient_grad,
     dim=(d.nworld, m.nv),
     inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, ctx.done],
     outputs=[ctx.grad, ctx.grad_dot],
@@ -2942,36 +2907,34 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
   elif m.opt.solver == types.SolverType.NEWTON:
     # h = M + (efc_J.T * efc_D * active) @ efc_J
     if m.is_sparse:
-      ctx.h.zero_()
+      wp.launch(
+        _update_gradient_init_h_sparse,
+        dim=(d.nworld, m.nv_pad, m.nv_pad),
+        inputs=[m.nv, m.M_elemid, d.M, ctx.done],
+        outputs=[ctx.h],
+      )
+
       wp.launch(
         _JTDAJ_sparse,
         dim=(d.nworld, d.njmax),
         inputs=[d.nefc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, d.efc.D, d.efc.state, ctx.done],
         outputs=[ctx.h],
       )
-
-      wp.launch(
-        update_gradient_set_h_M_upper_sparse,
-        dim=(d.nworld, m.M_fullm_upper_i.size),
-        inputs=[m.M_fullm_upper_i, m.M_fullm_upper_j, m.M_fullm_upper_elemid, d.M, ctx.done],
-        outputs=[ctx.h],
-      )
     else:
-      with scoped_mathdx_gemm_disabled():
-        wp.launch_tiled(
-          update_gradient_JTDAJ_dense_tiled(m.nv_pad, types.TILE_SIZE_JTDAJ_DENSE, d.njmax),
-          dim=d.nworld,
-          inputs=[
-            d.nefc,
-            d.M,
-            d.efc.J,
-            d.efc.D,
-            d.efc.state,
-            ctx.done,
-          ],
-          outputs=[ctx.h],
-          block_dim=m.block_dim.update_gradient_JTDAJ_dense,
-        )
+      wp.launch_tiled(
+        _update_gradient_JTDAJ_dense_tiled(m.nv_pad, types.TILE_SIZE_JTDAJ_DENSE, d.njmax),
+        dim=d.nworld,
+        inputs=[
+          d.nefc,
+          d.M,
+          d.efc.J,
+          d.efc.D,
+          d.efc.state,
+          ctx.done,
+        ],
+        outputs=[ctx.h],
+        block_dim=m.block_dim.update_gradient_JTDAJ_dense,
+      )
 
     if m.opt.cone == types.ConeType.ELLIPTIC:
       # Optimization: launching update_gradient_JTCJ with limited number of blocks on a GPU.
@@ -2999,7 +2962,7 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
 
       if m.is_sparse:
         wp.launch(
-          update_gradient_JTCJ_sparse,
+          _update_gradient_JTCJ_sparse,
           dim=(dim_block, m.dof_tri_row.size),
           inputs=[
             m.opt.impratio_invsqrt,
@@ -3028,7 +2991,7 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
         )
       else:
         wp.launch(
-          update_gradient_JTCJ_dense,
+          _update_gradient_JTCJ_dense,
           dim=(dim_block, m.dof_tri_row.size),
           inputs=[
             m.opt.impratio_invsqrt,
@@ -3064,10 +3027,10 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte
   Skips the full J^T*D*J rebuild by applying only the delta from constraints
   that changed QUADRATIC state, then re-factorizes and solves.
   """
-  wp.launch(update_gradient_zero_grad_dot, dim=(d.nworld), inputs=[ctx.done], outputs=[ctx.grad_dot])
+  wp.launch(_update_gradient_zero_grad_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.grad_dot])
 
   wp.launch(
-    update_gradient_grad,
+    _update_gradient_grad,
     dim=(d.nworld, m.nv),
     inputs=[d.qfrc_smooth, d.qfrc_constraint, d.efc.Ma, ctx.done],
     outputs=[ctx.grad, ctx.grad_dot],
@@ -3076,7 +3039,7 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte
   # Update upper triangle of H with delta from changed constraints.
   if m.is_sparse:
     wp.launch(
-      update_gradient_h_incremental_sparse,
+      _update_gradient_h_incremental_sparse,
       dim=(d.nworld, ctx.changed_efc_ids.shape[1]),
       inputs=[
         d.efc.J_rownnz,
@@ -3093,7 +3056,7 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte
   else:
     tri_dim = m.nv * (m.nv + 1) // 2
     wp.launch(
-      update_gradient_h_incremental,
+      _update_gradient_h_incremental,
       dim=(d.nworld, tri_dim),
       inputs=[
         d.efc.J,
@@ -3109,7 +3072,7 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte
 
 
 @wp.kernel
-def solve_prev_grad_Mgrad(
+def _solve_prev_grad_Mgrad(
   # In:
   ctx_grad_in: wp.array2d[float],
   ctx_Mgrad_in: wp.array2d[float],
@@ -3128,9 +3091,18 @@ def solve_prev_grad_Mgrad(
 
 
 @wp.kernel
-def solve_beta(
-  # Model:
-  nv: int,
+def _solve_beta_zero(
+  # Out:
+  ctx_beta_num_out: wp.array[float],
+  ctx_beta_den_out: wp.array[float],
+):
+  worldid = wp.tid()
+  ctx_beta_num_out[worldid] = 0.0
+  ctx_beta_den_out[worldid] = 0.0
+
+
+@wp.kernel
+def _solve_beta_accumulate(
   # In:
   ctx_grad_in: wp.array2d[float],
   ctx_Mgrad_in: wp.array2d[float],
@@ -3138,6 +3110,28 @@ def solve_beta(
   ctx_prev_Mgrad_in: wp.array2d[float],
   ctx_done_in: wp.array[bool],
   # Out:
+  ctx_beta_num_out: wp.array[float],
+  ctx_beta_den_out: wp.array[float],
+):
+  worldid, dofid = wp.tid()
+
+  if ctx_done_in[worldid]:
+    return
+
+  prev_Mgrad = ctx_prev_Mgrad_in[worldid, dofid]
+  num = ctx_grad_in[worldid, dofid] * (ctx_Mgrad_in[worldid, dofid] - prev_Mgrad)
+  den = ctx_prev_grad_in[worldid, dofid] * prev_Mgrad
+  wp.atomic_add(ctx_beta_num_out, worldid, num)
+  wp.atomic_add(ctx_beta_den_out, worldid, den)
+
+
+@wp.kernel
+def _solve_beta_finalize(
+  # In:
+  ctx_beta_num_in: wp.array[float],
+  ctx_beta_den_in: wp.array[float],
+  ctx_done_in: wp.array[bool],
+  # Out:
   ctx_beta_out: wp.array[float],
 ):
   worldid = wp.tid()
@@ -3145,18 +3139,11 @@ def solve_beta(
   if ctx_done_in[worldid]:
     return
 
-  beta_num = float(0.0)
-  beta_den = float(0.0)
-  for dofid in range(nv):
-    prev_Mgrad = ctx_prev_Mgrad_in[worldid][dofid]
-    beta_num += ctx_grad_in[worldid, dofid] * (ctx_Mgrad_in[worldid, dofid] - prev_Mgrad)
-    beta_den += ctx_prev_grad_in[worldid, dofid] * prev_Mgrad
-
-  ctx_beta_out[worldid] = wp.max(0.0, beta_num / wp.max(types.MJ_MINVAL, beta_den))
+  ctx_beta_out[worldid] = wp.max(0.0, ctx_beta_num_in[worldid] / wp.max(types.MJ_MINVAL, ctx_beta_den_in[worldid]))
 
 
 @wp.kernel
-def solve_zero_search_dot(
+def _solve_zero_search_dot(
   # In:
   ctx_done_in: wp.array[bool],
   # Out:
@@ -3171,7 +3158,7 @@ def solve_zero_search_dot(
 
 
 @wp.kernel
-def solve_search_update(
+def _solve_search_update(
   # Model:
   opt_solver: int,
   # In:
@@ -3198,7 +3185,7 @@ def solve_search_update(
 
 
 @wp.kernel
-def solve_done(
+def _solve_done(
   # Model:
   nv: int,
   opt_tolerance: wp.array[float],
@@ -3206,8 +3193,7 @@ def solve_done(
   stat_meaninertia: wp.array[float],
   # In:
   ctx_grad_dot_in: wp.array[float],
-  ctx_cost_in: wp.array[float],
-  ctx_prev_cost_in: wp.array[float],
+  ctx_improvement_in: wp.array[float],
   ctx_done_in: wp.array[bool],
   # Data out:
   solver_niter_out: wp.array[int],
@@ -3224,7 +3210,7 @@ def solve_done(
   tolerance = opt_tolerance[worldid % opt_tolerance.shape[0]]
   meaninertia = stat_meaninertia[worldid % stat_meaninertia.shape[0]]
 
-  improvement = _rescale(nv, meaninertia, ctx_prev_cost_in[worldid] - ctx_cost_in[worldid])
+  improvement = _rescale(nv, meaninertia, ctx_improvement_in[worldid])
   gradient = _rescale(nv, meaninertia, wp.sqrt(ctx_grad_dot_in[worldid]))
   done = (improvement < tolerance) or (gradient < tolerance)
   if done or solver_niter_out[worldid] == opt_iterations:
@@ -3246,20 +3232,20 @@ def _solver_iteration(
 
   if m.opt.solver == types.SolverType.CG:
     wp.launch(
-      solve_prev_grad_Mgrad,
+      _solve_prev_grad_Mgrad,
       dim=(d.nworld, m.nv),
       inputs=[ctx.grad, ctx.Mgrad, ctx.done],
       outputs=[ctx.prev_grad, ctx.prev_Mgrad],
     )
 
   # Incremental H is only valid for non-elliptic cones. The elliptic cone
-  # path in update_constraint_efc has early returns that skip state change
+  # path in _update_constraint_efc has early returns that skip state change
   # tracking, and the additional JTCJ Hessian term depends on Jaref which
   # changes every iteration.
   incremental = m.opt.solver == types.SolverType.NEWTON and m.opt.cone != types.ConeType.ELLIPTIC
 
   if incremental:
-    # Must complete before update_constraint_efc which atomically increments.
+    # Must complete before _update_constraint_efc which atomically increments.
     ctx.changed_efc_count.zero_()
 
   _update_constraint(m, d, ctx, track_changes=incremental)
@@ -3272,23 +3258,34 @@ def _solver_iteration(
   # polak-ribiere
   if m.opt.solver == types.SolverType.CG:
     wp.launch(
-      solve_beta,
+      _solve_beta_zero,
       dim=d.nworld,
-      inputs=[m.nv, ctx.grad, ctx.Mgrad, ctx.prev_grad, ctx.prev_Mgrad, ctx.done],
+      outputs=[ctx.beta, ctx.beta_den],
+    )
+    wp.launch(
+      _solve_beta_accumulate,
+      dim=(d.nworld, m.nv),
+      inputs=[ctx.grad, ctx.Mgrad, ctx.prev_grad, ctx.prev_Mgrad, ctx.done],
+      outputs=[ctx.beta, ctx.beta_den],
+    )
+    wp.launch(
+      _solve_beta_finalize,
+      dim=d.nworld,
+      inputs=[ctx.beta, ctx.beta_den, ctx.done],
       outputs=[ctx.beta],
     )
 
-  wp.launch(solve_zero_search_dot, dim=(d.nworld), inputs=[ctx.done], outputs=[ctx.search_dot])
+  wp.launch(_solve_zero_search_dot, dim=d.nworld, inputs=[ctx.done], outputs=[ctx.search_dot])
 
   wp.launch(
-    solve_search_update,
+    _solve_search_update,
     dim=(d.nworld, m.nv),
     inputs=[m.opt.solver, ctx.Mgrad, ctx.search, ctx.beta, ctx.done],
     outputs=[ctx.search, ctx.search_dot],
   )
 
   wp.launch(
-    solve_done,
+    _solve_done,
     dim=d.nworld,
     inputs=[
       m.nv,
@@ -3296,8 +3293,7 @@ def _solver_iteration(
       m.opt.iterations,
       m.stat.meaninertia,
       ctx.grad_dot,
-      ctx.cost,
-      ctx.prev_cost,
+      ctx.improvement,
       ctx.done,
     ],
     outputs=[d.solver_niter, nsolving, ctx.done],
@@ -3307,9 +3303,9 @@ def _solver_iteration(
 def init_context(m: types.Model, d: types.Data, ctx: SolverContext | InverseContext, grad: bool = True):
   # initialize some efc arrays
   wp.launch(
-    solve_init_efc,
-    dim=(d.nworld),
-    outputs=[d.solver_niter, ctx.search_dot, ctx.cost, ctx.done],
+    _solve_init_efc,
+    dim=d.nworld,
+    outputs=[d.solver_niter, ctx.search_dot, ctx.done],
   )
 
   # jaref = d.efc_J @ d.qacc - d.efc_aref
@@ -3328,7 +3324,7 @@ def init_context(m: types.Model, d: types.Data, ctx: SolverContext | InverseCont
     ctx.Jaref.zero_()
 
   wp.launch(
-    solve_init_jaref(m.is_sparse, m.nv, dofs_per_thread),
+    _solve_init_jaref_kernel(m.is_sparse, m.nv, dofs_per_thread),
     dim=(d.nworld, d.njmax, threads_per_efc),
     inputs=[d.nefc, d.qacc, d.efc.J_rownnz, d.efc.J_rowadr, d.efc.J_colind, d.efc.J, d.efc.aref],
     outputs=[ctx.Jaref],
@@ -3350,15 +3346,15 @@ def solve(m: types.Model, d: types.Data):
     d.solver_niter.fill_(0)
   else:
     if m.ntree > 1 and not (m.opt.disableflags & types.DisableBit.ISLAND):
-      ctx = create_island_solver_context(m, d)
+      ctx = _create_island_solver_context(m, d)
       island.compute_island_mapping(m, d, ctx)
       island.gather_island_inputs(m, d, ctx)
-      _solve_islands(m, d, ctx)
+      _solve_island(m, d, ctx)
       # Ma is needed by Euler/implicit integrators for implicit damping
       scatter_Ma = m.opt.integrator != types.IntegratorType.RK4
       island.scatter_island_results(m, d, ctx, scatter_Ma=scatter_Ma)
     else:
-      ctx = create_solver_context(m, d)
+      ctx = _create_solver_context(m, d)
       _solve(m, d, ctx)
 
 
@@ -3374,7 +3370,7 @@ def _solve(m: types.Model, d: types.Data, ctx: SolverContext):
 
   # search = -Mgrad
   wp.launch(
-    solve_init_search,
+    _solve_init_search,
     dim=(d.nworld, m.nv),
     inputs=[ctx.Mgrad],
     outputs=[ctx.search, ctx.search_dot],
@@ -3404,7 +3400,7 @@ def _solve(m: types.Model, d: types.Data, ctx: SolverContext):
 
 # TODO(team): Consolidate monolithic and island solver code where possible
 @event_scope
-def _solve_islands(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+def _solve_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
   """Solve constraints for all islands in parallel.
 
   All islands are processed simultaneously. Island-local arrays in ctx
@@ -3414,7 +3410,7 @@ def _solve_islands(m: types.Model, d: types.Data, ctx: IslandSolverContext):
   # Initialize iacc from warmstart or smooth
   if not (m.opt.disableflags & types.DisableBit.WARMSTART):
     wp.launch(
-      gather_warmstart_island,
+      _gather_warmstart_island,
       dim=(d.nworld, m.nv),
       inputs=[d.nidof, d.qacc_warmstart, d.map_idof2dof],
       outputs=[d.iqacc],
@@ -3422,26 +3418,20 @@ def _solve_islands(m: types.Model, d: types.Data, ctx: IslandSolverContext):
   else:
     wp.copy(d.iqacc, d.iqacc_smooth)
 
+  # nsolving tracks how many active islands still have unconverged globally
+  nsolving = wp.zeros((1,), dtype=int)
+
   # Initialize island context
-  init_context_island(m, d, ctx)
+  _init_context_island(m, d, ctx, nsolving)
 
   # search = -Mgrad
   wp.launch(
-    solve_init_search_island,
+    _solve_init_search_island,
     dim=(d.nworld, m.nv),
     inputs=[d.nidof, ctx.Mgrad, d.dof_islandid, ctx.done],
     outputs=[ctx.search, ctx.search_dot],
   )
 
-  # nsolving tracks how many active islands still have unconverged globally
-  nsolving = wp.zeros((1,), dtype=int)
-  wp.launch(
-    solve_init_nsolving_island,
-    dim=d.nworld,
-    inputs=[d.nisland],
-    outputs=[nsolving],
-  )
-
   if m.opt.iterations != 0 and m.opt.graph_conditional:
     wp.capture_while(
       nsolving,
@@ -3457,7 +3447,7 @@ def _solve_islands(m: types.Model, d: types.Data, ctx: IslandSolverContext):
 
 
 @wp.kernel
-def gather_warmstart_island(
+def _gather_warmstart_island(
   # Data in:
   nidof_in: wp.array[int],
   qacc_warmstart_in: wp.array2d[float],
@@ -3475,38 +3465,60 @@ def gather_warmstart_island(
   iacc_out[worldid, idofid] = qacc_warmstart_in[worldid, dof]
 
 
-@wp.kernel
-def solve_init_efc_island(
-  # Data in:
-  nisland_in: wp.array[int],
-  # Out:
-  island_cost_out: wp.array2d[float],
-  island_search_dot_out: wp.array2d[float],
-  island_done_out: wp.array2d[bool],
-  island_solver_niter_out: wp.array2d[int],
-):
-  """Initialize per-island solver scalars."""
-  worldid, islandid = wp.tid()
+@cache_kernel
+def _solve_init_efc_island(enable_sleep: bool):
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # Model:
+    ntree: int,
+    # Data in:
+    nisland_in: wp.array[int],
+    tree_awake_in: wp.array2d[int],
+    tree_island_in: wp.array2d[int],
+    # Out:
+    island_cost_out: wp.array2d[float],
+    island_search_dot_out: wp.array2d[float],
+    island_done_out: wp.array2d[bool],
+    island_solver_niter_out: wp.array2d[int],
+    nsolving_out: wp.array[int],
+  ):
+    """Initialize per-island solver scalars."""
+    worldid, islandid = wp.tid()
 
-  if islandid >= nisland_in[worldid]:
-    return
+    if islandid >= nisland_in[worldid]:
+      return
 
-  island_cost_out[worldid, islandid] = 0.0
-  island_search_dot_out[worldid, islandid] = 0.0
-  island_done_out[worldid, islandid] = False
-  island_solver_niter_out[worldid, islandid] = 0
+    is_asleep_flag = int(0)
+    if wp.static(enable_sleep):
+      has_awake_tree = int(0)
+      for t in range(ntree):
+        if tree_island_in[worldid, t] == islandid:
+          if tree_awake_in[worldid, t] == 1:
+            has_awake_tree = int(1)
+            break
+      if has_awake_tree == 0:
+        is_asleep_flag = int(1)
+
+    island_cost_out[worldid, islandid] = 0.0
+    island_search_dot_out[worldid, islandid] = 0.0
+    island_done_out[worldid, islandid] = is_asleep_flag == 1
+    island_solver_niter_out[worldid, islandid] = 0
+    if is_asleep_flag == 0:
+      wp.atomic_add(nsolving_out, 0, 1)
+
+  return kernel
 
 
 @wp.kernel
-def solve_init_jaref_island(
+def _solve_init_jaref_island(
   # Model:
   is_sparse: bool,
   # Data in:
   nefc_in: wp.array[int],
+  island_idofadr_in: wp.array2d[int],
   island_nv_in: wp.array2d[int],
   njmax_in: int,
   # In:
-  island_idofadr_in: wp.array2d[int],
   iefc_J_rownnz_in: wp.array2d[int],
   iefc_J_rowadr_in: wp.array2d[int],
   iefc_J_colind_in: wp.array3d[int],
@@ -3550,7 +3562,7 @@ def solve_init_jaref_island(
 
 
 @wp.kernel
-def solve_init_search_island(
+def _solve_init_search_island(
   # Data in:
   nidof_in: wp.array[int],
   # In:
@@ -3578,33 +3590,28 @@ def solve_init_search_island(
   wp.atomic_add(island_search_dot_out, worldid, islandid, s * s)
 
 
+# TODO(team): remove after updating island solver done criteria to use delta cost
 @wp.kernel
-def update_constraint_init_cost_island(
-  # Data in:
-  nisland_in: wp.array[int],
+def _update_constraint_init_cost(
   # In:
-  island_cost_in: wp.array2d[float],
-  island_done_in: wp.array2d[bool],
+  cost_in: wp.array[float],
+  done_in: wp.array[bool],
   # Out:
-  island_gauss_out: wp.array2d[float],
-  island_cost_out: wp.array2d[float],
-  island_prev_cost_out: wp.array2d[float],
+  gauss_out: wp.array[float],
+  cost_out: wp.array[float],
+  prev_cost_out: wp.array[float],
 ):
-  """Save prev_cost and zero cost/gauss per island."""
-  worldid, islandid = wp.tid()
-
-  if islandid >= nisland_in[worldid]:
-    return
-  if island_done_in[worldid, islandid]:
+  tid = wp.tid()
+  if done_in[tid]:
     return
 
-  island_prev_cost_out[worldid, islandid] = island_cost_in[worldid, islandid]
-  island_cost_out[worldid, islandid] = 0.0
-  island_gauss_out[worldid, islandid] = 0.0
+  prev_cost_out[tid] = cost_in[tid]
+  cost_out[tid] = 0.0
+  gauss_out[tid] = 0.0
 
 
 @wp.kernel
-def update_constraint_efc_island(
+def _update_constraint_efc_island(
   # Model:
   opt_impratio_invsqrt: wp.array[float],
   # Data in:
@@ -3653,63 +3660,34 @@ def update_constraint_efc_island(
   jaref = Jaref_in[worldid, iefcid]
   D = iefc_D_in[worldid, iefcid]
 
-  force = float(0.0)
-  state = types.ConstraintState.SATISFIED.value
-  cost = float(0.0)
+  is_equality = local_iefcid < ine
+  is_friction = (not is_equality) and (local_iefcid < ine + inf)
+  is_elliptic = iefc_type_in[worldid, iefcid] == types.ConstraintType.CONTACT_ELLIPTIC
 
-  # Determine constraint category by position within island
-  if local_iefcid < ine:
-    # Equality constraint: always quadratic
-    force = -D * jaref
-    state = types.ConstraintState.QUADRATIC.value
-    cost = 0.5 * D * jaref * jaref
-  elif local_iefcid < ine + inf:
-    # Friction constraint
-    f = iefc_frictionloss_in[worldid, iefcid]
-    rf = math.safe_div(f, D)
-    if jaref <= -rf:
-      force = f
-      state = types.ConstraintState.LINEARNEG.value
-      cost = -f * (0.5 * rf + jaref)
-    elif jaref >= rf:
-      force = -f
-      state = types.ConstraintState.LINEARPOS.value
-      cost = -f * (0.5 * rf - jaref)
-    else:
-      force = -D * jaref
-      state = types.ConstraintState.QUADRATIC.value
-      cost = 0.5 * D * jaref * jaref
-  elif iefc_type_in[worldid, iefcid] != types.ConstraintType.CONTACT_ELLIPTIC:
-    # Limit, frictionless contact, pyramidal friction cone contact
-    if jaref >= 0.0:
-      force = 0.0
-      state = types.ConstraintState.SATISFIED.value
-      cost = 0.0
-    else:
-      force = -D * jaref
-      state = types.ConstraintState.QUADRATIC.value
-      cost = 0.5 * D * jaref * jaref
-  else:
-    # Elliptic friction cone contact
+  frictionloss = iefc_frictionloss_in[worldid, iefcid] if is_friction else 0.0
+
+  ic0 = int(-1)
+  jaref0 = float(0.0)
+  D0 = float(0.0)
+  mu = float(0.0)
+  ufrictionj = float(0.0)
+  TT = float(0.0)
+
+  if is_elliptic:
     conid = iefc_id_in[worldid, iefcid]
-
     if conid >= nacon_in[0]:
       return
-
-    dim = contact_dim_in[conid]
-    friction = contact_friction_in[conid]
-    mu = friction[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
-
-    # Get island-local index for lead row
     efcid0_global = contact_efc_address_in[conid, 0]
     if efcid0_global < 0:
       return
     ic0 = map_efc2iefc_in[worldid, efcid0_global]
 
-    N = Jaref_in[worldid, ic0] * mu
+    dim = contact_dim_in[conid]
+    friction = contact_friction_in[conid]
+    mu = friction[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
+    jaref0 = Jaref_in[worldid, ic0]
+    D0 = iefc_D_in[worldid, ic0]
 
-    ufrictionj = float(0.0)
-    TT = float(0.0)
     for j in range(1, dim):
       efcidj_global = contact_efc_address_in[conid, j]
       if efcidj_global < 0:
@@ -3721,42 +3699,31 @@ def update_constraint_efc_island(
       if iefcid == icj:
         ufrictionj = uj * frictionj
 
-    if TT <= 0.0:
-      T = 0.0
-    else:
-      T = wp.sqrt(TT)
+  res = _eval_constraint(
+    is_equality,
+    is_friction,
+    is_elliptic,
+    jaref,
+    D,
+    frictionloss,
+    iefcid,
+    ic0,
+    jaref0,
+    D0,
+    mu,
+    ufrictionj,
+    TT,
+  )
 
-    # top zone
-    if (N >= mu * T) or ((T <= 0.0) and (N >= 0.0)):
-      force = 0.0
-      state = types.ConstraintState.SATISFIED.value
-    # bottom zone
-    elif (mu * N + T <= 0.0) or ((T <= 0.0) and (N < 0.0)):
-      force = -D * jaref
-      state = types.ConstraintState.QUADRATIC.value
-      cost = 0.5 * D * jaref * jaref
-    # middle zone
-    else:
-      dm = math.safe_div(iefc_D_in[worldid, ic0], mu * mu * (1.0 + mu * mu))
-      nmt = N - mu * T
-
-      force_normal = -dm * nmt * mu
-
-      if iefcid == ic0:
-        force = force_normal
-        cost = 0.5 * dm * nmt * nmt
-      else:
-        force = -math.safe_div(force_normal, T) * ufrictionj
-
-      state = types.ConstraintState.CONE.value
-
-  iefc_force_out[worldid, iefcid] = force
-  iefc_state_out[worldid, iefcid] = state
-  wp.atomic_add(island_cost_out, worldid, islandid, cost)
+  iefc_force_out[worldid, iefcid] = res[0]
+  iefc_state_out[worldid, iefcid] = int(res[1])
+  cost = res[2]
+  if cost != 0.0:
+    wp.atomic_add(island_cost_out, worldid, islandid, cost)
 
 
 @wp.kernel
-def update_constraint_init_qfrc_constraint_dense_island(
+def _update_constraint_init_qfrc_constraint_dense_island(
   # Data in:
   nefc_in: wp.array[int],
   nidof_in: wp.array[int],
@@ -3795,7 +3762,7 @@ def update_constraint_init_qfrc_constraint_dense_island(
 
 
 @wp.kernel
-def update_constraint_init_qfrc_constraint_sparse_island(
+def _update_constraint_init_qfrc_constraint_sparse_island(
   # Data in:
   nefc_in: wp.array[int],
   njmax_in: int,
@@ -3832,7 +3799,7 @@ def update_constraint_init_qfrc_constraint_sparse_island(
 
 
 @wp.kernel
-def update_constraint_gauss_cost_island(
+def _update_constraint_gauss_cost_island(
   # Data in:
   nidof_in: wp.array[int],
   # In:
@@ -3867,7 +3834,7 @@ def update_constraint_gauss_cost_island(
 
 
 @wp.kernel
-def update_gradient_grad_island(
+def _update_gradient_grad_island(
   # Data in:
   nidof_in: wp.array[int],
   # In:
@@ -3898,16 +3865,16 @@ def update_gradient_grad_island(
 
 
 @wp.kernel
-def linesearch_jv_island(
+def _linesearch_jv_island(
   # Model:
   is_sparse: bool,
   # Data in:
   nefc_in: wp.array[int],
   nidof_in: wp.array[int],
+  island_idofadr_in: wp.array2d[int],
   island_nv_in: wp.array2d[int],
   njmax_in: int,
   # In:
-  island_idofadr_in: wp.array2d[int],
   iefc_J_rownnz_in: wp.array2d[int],
   iefc_J_rowadr_in: wp.array2d[int],
   iefc_J_colind_in: wp.array3d[int],
@@ -4003,39 +3970,15 @@ def _eval_elliptic_cost_island(
   mu2 = mu * mu
   dm = math.safe_div(D0, mu2 * (1.0 + mu2))
 
-  # Compute N, Tsqr
-  N = u0 + alpha * v0
-  Tsqr = uu + alpha * (2.0 * uv + alpha * vv)
+  quad1 = wp.vec3(u0, v0, uu)
+  quad2 = wp.vec3(uv, vv, dm)
 
-  if Tsqr <= 0.0:
-    if N < 0.0:
-      return _eval_pt(quad, alpha)
-    return wp.vec3(0.0)
-
-  T = wp.sqrt(Tsqr)
-
-  # top zone
-  if N >= mu * T:
-    return wp.vec3(0.0)
-  # bottom zone
-  if mu * N + T <= 0.0:
-    return _eval_pt(quad, alpha)
-
-  # middle zone
-  N1 = v0
-  T1 = (uv + alpha * vv) / T
-  T2 = vv / T - (uv + alpha * vv) * T1 / (T * T)
-  nmt = N - mu * T
-  return wp.vec3(
-    0.5 * dm * nmt * nmt,
-    dm * nmt * (N1 - mu * T1),
-    dm * ((N1 - mu * T1) * (N1 - mu * T1) + nmt * (-mu * T2)),
-  )
+  return _eval_elliptic(mu, quad, quad1, quad2, alpha)
 
 
-# TODO(team): refactor linesearch_island
+# TODO(team): refactor _linesearch_kernel_island
 @wp.kernel
-def linesearch_island(
+def _linesearch_kernel_island(
   # Model:
   opt_tolerance: wp.array[float],
   opt_ls_tolerance: wp.array[float],
@@ -4108,7 +4051,9 @@ def linesearch_island(
     s = search_in[worldid, idof]
     quad_gauss_1 += s * (iMa_in[worldid, idof] - ifrc_smooth_in[worldid, idof])
     quad_gauss_2 += 0.5 * s * mv_in[worldid, idof]
-  quad_gauss = wp.vec3(island_gauss_in[worldid, islandid], quad_gauss_1, quad_gauss_2)
+  # Costs are evaluated as deltas from alpha=0 to keep float32 precision on large
+  # absolute costs, so the constant gauss cost (island_gauss_in) is dropped here.
+  quad_gauss = wp.vec3(0.0, quad_gauss_1, quad_gauss_2)
 
   # gtol
   snorm = wp.sqrt(island_search_dot_in[worldid, islandid])
@@ -4157,6 +4102,9 @@ def linesearch_island(
         jvD = jv_val * D
         p0 += wp.vec3(0.5 * D * ja * ja, jvD * ja, jv_val * jvD)
 
+  # Zero the cost component: alpha=0 is the delta reference (grad/hessian are kept).
+  p0 = wp.vec3(0.0, p0[1], p0[2])
+
   # Newton step: lo_alpha_in = -p0[1] / p0[2]
   lo_alpha_in = -math.safe_div(p0[1], p0[2])
 
@@ -4170,37 +4118,58 @@ def linesearch_island(
     ja = Jaref_in[worldid, iefcid]
     jv_val = jv_in[worldid, iefcid]
     if local_iefcid < ine:
-      lo_in += _eval_pt_direct(ja, jv_val, D, lo_alpha_in)
+      lo_in += _eval_pt_direct_shifted(ja, jv_val, D, lo_alpha_in, 0.0)
     elif local_iefcid < ine + inf:
       f = iefc_frictionloss_in[worldid, iefcid]
       rf = math.safe_div(f, D)
       x_a = ja + lo_alpha_in * jv_val
-      lo_in += _eval_frictionloss_pt(x_a, f, rf, jv_val, D)
+      lo_in += _shift_cost(_eval_frictionloss_pt(x_a, f, rf, jv_val, D), _eval_frictionloss_cost(ja, f, rf, D))
     elif iefc_type_in[worldid, iefcid] == types.ConstraintType.CONTACT_ELLIPTIC:
       conid = iefc_id_in[worldid, iefcid]
       if conid < nacon_in[0]:
         ic0 = map_efc2iefc_in[worldid, contact_efc_address_in[conid, 0]]
         if iefcid == ic0:
-          lo_in += _eval_elliptic_cost_island(
+          cost0 = _eval_elliptic_cost_island(
             impratio_invsqrt,
             contact_friction_in,
             contact_dim_in,
             contact_efc_address_in,
             map_efc2iefc_in,
-            lo_alpha_in,
+            0.0,
             conid,
             iefc_D_in,
             Jaref_in,
             jv_in,
             worldid,
+          )[0]
+          lo_in += _shift_cost(
+            _eval_elliptic_cost_island(
+              impratio_invsqrt,
+              contact_friction_in,
+              contact_dim_in,
+              contact_efc_address_in,
+              map_efc2iefc_in,
+              lo_alpha_in,
+              conid,
+              iefc_D_in,
+              Jaref_in,
+              jv_in,
+              worldid,
+            ),
+            cost0,
           )
     else:
+      # Inequality
       x_a = ja + lo_alpha_in * jv_val
+      quad0 = _eval_pt_direct_cost_alpha_zero(ja, D)
+      cost0 = wp.where(ja < 0.0, quad0, 0.0)
       if x_a < 0.0:
-        lo_in += _eval_pt_direct(ja, jv_val, D, lo_alpha_in)
+        lo_in += _eval_pt_direct_shifted(ja, jv_val, D, lo_alpha_in, quad0 - cost0)
+      else:
+        lo_in += wp.vec3(-cost0, 0.0, 0.0)
 
   # Accept Newton step if derivative is small and cost improved
-  initial_converged = wp.abs(lo_in[1]) < gtol and lo_in[0] < p0[0]
+  initial_converged = wp.abs(lo_in[1]) < gtol and lo_in[0] < 0.0
 
   if initial_converged:
     alpha = lo_alpha_in
@@ -4240,16 +4209,17 @@ def linesearch_island(
         ja = Jaref_in[worldid, iefcid]
         jv_val = jv_in[worldid, iefcid]
         if local_iefcid < ine:
-          r_lo, r_hi, r_mid = _eval_pt_direct_3alphas(ja, jv_val, D, lo_next_alpha, hi_next_alpha, mid_alpha)
+          r_lo, r_hi, r_mid = _eval_pt_direct_shifted_3alphas(ja, jv_val, D, lo_next_alpha, hi_next_alpha, mid_alpha, 0.0)
         elif local_iefcid < ine + inf:
           f = iefc_frictionloss_in[worldid, iefcid]
           rf = math.safe_div(f, D)
+          cost0 = _eval_frictionloss_cost(ja, f, rf, D)
           x_lo = ja + lo_next_alpha * jv_val
           x_hi = ja + hi_next_alpha * jv_val
           x_mid = ja + mid_alpha * jv_val
-          r_lo = _eval_frictionloss_pt(x_lo, f, rf, jv_val, D)
-          r_hi = _eval_frictionloss_pt(x_hi, f, rf, jv_val, D)
-          r_mid = _eval_frictionloss_pt(x_mid, f, rf, jv_val, D)
+          r_lo = _shift_cost(_eval_frictionloss_pt(x_lo, f, rf, jv_val, D), cost0)
+          r_hi = _shift_cost(_eval_frictionloss_pt(x_hi, f, rf, jv_val, D), cost0)
+          r_mid = _shift_cost(_eval_frictionloss_pt(x_mid, f, rf, jv_val, D), cost0)
         elif iefc_type_in[worldid, iefcid] == types.ConstraintType.CONTACT_ELLIPTIC:
           conid = iefc_id_in[worldid, iefcid]
           r_lo = wp.vec3(0.0)
@@ -4258,58 +4228,85 @@ def linesearch_island(
           if conid < nacon_in[0]:
             ic0 = map_efc2iefc_in[worldid, contact_efc_address_in[conid, 0]]
             if iefcid == ic0:
-              r_lo = _eval_elliptic_cost_island(
+              cost0 = _eval_elliptic_cost_island(
                 impratio_invsqrt,
                 contact_friction_in,
                 contact_dim_in,
                 contact_efc_address_in,
                 map_efc2iefc_in,
-                lo_next_alpha,
+                0.0,
                 conid,
                 iefc_D_in,
                 Jaref_in,
                 jv_in,
                 worldid,
+              )[0]
+              r_lo = _shift_cost(
+                _eval_elliptic_cost_island(
+                  impratio_invsqrt,
+                  contact_friction_in,
+                  contact_dim_in,
+                  contact_efc_address_in,
+                  map_efc2iefc_in,
+                  lo_next_alpha,
+                  conid,
+                  iefc_D_in,
+                  Jaref_in,
+                  jv_in,
+                  worldid,
+                ),
+                cost0,
               )
-              r_hi = _eval_elliptic_cost_island(
-                impratio_invsqrt,
-                contact_friction_in,
-                contact_dim_in,
-                contact_efc_address_in,
-                map_efc2iefc_in,
-                hi_next_alpha,
-                conid,
-                iefc_D_in,
-                Jaref_in,
-                jv_in,
-                worldid,
+              r_hi = _shift_cost(
+                _eval_elliptic_cost_island(
+                  impratio_invsqrt,
+                  contact_friction_in,
+                  contact_dim_in,
+                  contact_efc_address_in,
+                  map_efc2iefc_in,
+                  hi_next_alpha,
+                  conid,
+                  iefc_D_in,
+                  Jaref_in,
+                  jv_in,
+                  worldid,
+                ),
+                cost0,
               )
-              r_mid = _eval_elliptic_cost_island(
-                impratio_invsqrt,
-                contact_friction_in,
-                contact_dim_in,
-                contact_efc_address_in,
-                map_efc2iefc_in,
-                mid_alpha,
-                conid,
-                iefc_D_in,
-                Jaref_in,
-                jv_in,
-                worldid,
+              r_mid = _shift_cost(
+                _eval_elliptic_cost_island(
+                  impratio_invsqrt,
+                  contact_friction_in,
+                  contact_dim_in,
+                  contact_efc_address_in,
+                  map_efc2iefc_in,
+                  mid_alpha,
+                  conid,
+                  iefc_D_in,
+                  Jaref_in,
+                  jv_in,
+                  worldid,
+                ),
+                cost0,
               )
         else:
+          # Inequality
           x_lo = ja + lo_next_alpha * jv_val
           x_hi = ja + hi_next_alpha * jv_val
           x_mid = ja + mid_alpha * jv_val
-          r_lo = wp.vec3(0.0)
-          r_hi = wp.vec3(0.0)
-          r_mid = wp.vec3(0.0)
+          quad0 = _eval_pt_direct_cost_alpha_zero(ja, D)
+          cost0 = wp.where(ja < 0.0, quad0, 0.0)
+          offset = quad0 - cost0
+          neg_cost0 = wp.vec3(-cost0, 0.0, 0.0)
+          r_lo = neg_cost0
+          r_hi = neg_cost0
+          r_mid = neg_cost0
           if x_lo < 0.0:
-            r_lo = _eval_pt_direct(ja, jv_val, D, lo_next_alpha)
+            r_lo = _eval_pt_direct_shifted(ja, jv_val, D, lo_next_alpha, offset)
           if x_hi < 0.0:
-            r_hi = _eval_pt_direct(ja, jv_val, D, hi_next_alpha)
+            r_hi = _eval_pt_direct_shifted(ja, jv_val, D, hi_next_alpha, offset)
           if x_mid < 0.0:
-            r_mid = _eval_pt_direct(ja, jv_val, D, mid_alpha)
+            r_mid = _eval_pt_direct_shifted(ja, jv_val, D, mid_alpha, offset)
         lo_next += r_lo
         hi_next += r_hi
         mid += r_mid
@@ -4347,7 +4344,7 @@ def linesearch_island(
       ls_done = (swap_lo == 0 and swap_hi == 0) or (lo[1] < 0.0 and lo[1] > -gtol) or (hi[1] > 0.0 and hi[1] < gtol)
 
       # Update alpha if improved
-      if lo[0] < p0[0] or hi[0] < p0[0]:
+      if lo[0] < 0.0 or hi[0] < 0.0:
         if lo[0] < hi[0]:
           alpha = lo_alpha
         else:
@@ -4360,7 +4357,7 @@ def linesearch_island(
 
 
 @wp.kernel
-def linesearch_qacc_ma_island(
+def _linesearch_qacc_ma_island(
   # Data in:
   nidof_in: wp.array[int],
   # In:
@@ -4387,7 +4384,7 @@ def linesearch_qacc_ma_island(
 
 
 @wp.kernel
-def linesearch_jaref_island(
+def _linesearch_jaref_island(
   # Data in:
   nefc_in: wp.array[int],
   njmax_in: int,
@@ -4416,7 +4413,7 @@ def linesearch_jaref_island(
 
 
 @wp.kernel
-def solve_prev_grad_Mgrad_island(
+def _solve_prev_grad_Mgrad_island(
   # Data in:
   nidof_in: wp.array[int],
   # In:
@@ -4445,46 +4442,81 @@ def solve_prev_grad_Mgrad_island(
 
 
 @wp.kernel
-def solve_beta_island(
-  # Data in:
-  nisland_in: wp.array[int],
-  island_nv_in: wp.array2d[int],
+def _solve_beta_island_zero(
   # In:
-  island_idofadr_in: wp.array2d[int],
+  nisland_in: wp.array[int],
+  # Out:
+  island_beta_num_out: wp.array2d[float],
+  island_beta_den_out: wp.array2d[float],
+):
+  """Zero Polak-Ribière numerator and denominator per island."""
+  worldid, islandid = wp.tid()
+
+  if islandid >= nisland_in[worldid]:
+    return
+
+  island_beta_num_out[worldid, islandid] = 0.0
+  island_beta_den_out[worldid, islandid] = 0.0
+
+
+@wp.kernel
+def _solve_beta_island_accumulate(
+  # Data in:
+  nidof_in: wp.array[int],
+  # In:
+  idof_islandid_in: wp.array2d[int],
   grad_in: wp.array2d[float],
   Mgrad_in: wp.array2d[float],
   prev_grad_in: wp.array2d[float],
   prev_Mgrad_in: wp.array2d[float],
   island_done_in: wp.array2d[bool],
   # Out:
+  island_beta_num_out: wp.array2d[float],
+  island_beta_den_out: wp.array2d[float],
+):
+  """Parallel Polak-Ribière beta accumulation per island DOF."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  pMg = prev_Mgrad_in[worldid, idofid]
+  num = grad_in[worldid, idofid] * (Mgrad_in[worldid, idofid] - pMg)
+  den = prev_grad_in[worldid, idofid] * pMg
+  wp.atomic_add(island_beta_num_out, worldid, islandid, num)
+  wp.atomic_add(island_beta_den_out, worldid, islandid, den)
+
+
+@wp.kernel
+def _solve_beta_island_finalize(
+  # Data in:
+  nisland_in: wp.array[int],
+  # In:
+  island_beta_num_in: wp.array2d[float],
+  island_beta_den_in: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  # Out:
   island_beta_out: wp.array2d[float],
 ):
-  """Polak-Ribière beta per island."""
+  """Finalize Polak-Ribière beta per island."""
   worldid, islandid = wp.tid()
 
   if islandid >= nisland_in[worldid]:
     return
 
-  if island_done_in[worldid, islandid]:
-    island_beta_out[worldid, islandid] = 0.0
-    return
-
-  idofadr = island_idofadr_in[worldid, islandid]
-  inv = island_nv_in[worldid, islandid]
-
-  beta_num = float(0.0)
-  beta_den = float(0.0)
-  for i in range(inv):
-    idof = idofadr + i
-    pMg = prev_Mgrad_in[worldid, idof]
-    beta_num += grad_in[worldid, idof] * (Mgrad_in[worldid, idof] - pMg)
-    beta_den += prev_grad_in[worldid, idof] * pMg
-
-  island_beta_out[worldid, islandid] = wp.max(0.0, beta_num / wp.max(types.MJ_MINVAL, beta_den))
+  island_beta_out[worldid, islandid] = wp.max(
+    0.0, island_beta_num_in[worldid, islandid] / wp.max(types.MJ_MINVAL, island_beta_den_in[worldid, islandid])
+  )
 
 
 @wp.kernel
-def solve_search_update_island(
+def _solve_search_update_island(
   # Model:
   opt_solver: int,
   # Data in:
@@ -4520,17 +4552,7 @@ def solve_search_update_island(
 
 
 @wp.kernel
-def solve_init_nsolving_island(
-  nisland_in: wp.array[int],
-  nsolving_out: wp.array[int],
-):
-  """Initialize active island count nsolving on device without CPU stalls."""
-  worldid = wp.tid()
-  wp.atomic_add(nsolving_out, 0, nisland_in[worldid])
-
-
-@wp.kernel
-def solve_done_island(
+def _solve_done_island(
   # Model:
   opt_tolerance: wp.array[float],
   opt_iterations: int,
@@ -4578,7 +4600,7 @@ def solve_done_island(
 
 
 @wp.kernel
-def update_gradient_JTDAJ_island(
+def _update_gradient_JTDAJ_island(
   # Model:
   is_sparse: bool,
   # Data in:
@@ -4652,7 +4674,7 @@ def update_gradient_JTDAJ_island(
 
 
 @wp.kernel
-def update_gradient_set_h_M_sparse_island(
+def _update_gradient_set_h_M_sparse_island(
   # Model:
   M_fullm_i: wp.array[int],
   M_fullm_j: wp.array[int],
@@ -4702,7 +4724,7 @@ def update_gradient_set_h_M_sparse_island(
 
 
 @wp.kernel
-def update_gradient_set_h_M_dense_island(
+def _update_gradient_set_h_M_dense_island(
   # Model:
   nv: int,
   # Data in:
@@ -4737,7 +4759,7 @@ def update_gradient_set_h_M_dense_island(
 
 
 @wp.kernel
-def update_gradient_JTCJ_island(
+def _update_gradient_JTCJ_island(
   # Model:
   opt_impratio_invsqrt: wp.array[float],
   is_sparse: bool,
@@ -4928,14 +4950,23 @@ def update_gradient_JTCJ_island(
             if i != jj:
               wp.atomic_add(ih_out[worldid, idofadr + jj], idofadr + i, val)
 
-            if dim1id != dim2id:
+        if dim1id != dim2id:
+          # Swap-pair contribution: hcone * J[ic2, i] * J[ic1, j].
+          # Together with the loop above this gives the full
+          # hcone * (J[ic1, i] * J[ic2, j] + J[ic2, i] * J[ic1, j])
+          # contribution to cell (i, j).
+          for i in range(inv):
+            J2i = iefc_J_in[worldid, ic2, idofadr + i]
+            if J2i == 0.0:
+              continue
+            for jj in range(i + 1):
               J1j = iefc_J_in[worldid, ic1, idofadr + jj]
-              J2i = iefc_J_in[worldid, ic2, idofadr + i]
-              if J1j != 0.0 and J2i != 0.0:
-                val2 = hcone * J1j * J2i
-                wp.atomic_add(ih_out[worldid, idofadr + i], idofadr + jj, val2)
-                if i != jj:
-                  wp.atomic_add(ih_out[worldid, idofadr + jj], idofadr + i, val2)
+              if J1j == 0.0:
+                continue
+              val = hcone * J2i * J1j
+              wp.atomic_add(ih_out[worldid, idofadr + i], idofadr + jj, val)
+              if i != jj:
+                wp.atomic_add(ih_out[worldid, idofadr + jj], idofadr + i, val)
 
 
 @wp.kernel
@@ -4999,27 +5030,33 @@ def _cholesky_factorize_solve_island(
     Mgrad_out[worldid, adr + i] = s / wp.max(types.MJ_MINVAL, ih_in[worldid, adr + i, adr + i])
 
 
-def init_context_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+def _init_context_island(m: types.Model, d: types.Data, ctx: IslandSolverContext, nsolving: wp.array):
   """Initialize island solver context."""
   # Init per-island scalars
   d.solver_niter.zero_()
+  enable_sleep = bool(m.opt.enableflags & types.EnableBit.SLEEP)
   wp.launch(
-    solve_init_efc_island,
+    _solve_init_efc_island(enable_sleep),
     dim=(d.nworld, m.ntree),
-    inputs=[d.nisland],
-    outputs=[ctx.cost, ctx.search_dot, ctx.done, ctx.solver_niter],
+    inputs=[
+      m.ntree,
+      d.nisland,
+      d.tree_awake,
+      d.tree_island,
+    ],
+    outputs=[ctx.cost, ctx.search_dot, ctx.done, ctx.solver_niter, nsolving],
   )
 
   # Jaref = iefc_J @ iacc - iefc_aref
   wp.launch(
-    solve_init_jaref_island,
+    _solve_init_jaref_island,
     dim=(d.nworld, d.njmax),
     inputs=[
       m.is_sparse,
       d.nefc,
+      d.island_idofadr,
       d.island_nv,
       d.njmax,
-      d.island_dofadr,
       d.efc.iJ_rownnz,
       d.efc.iJ_rowadr,
       d.efc.iJ_colind,
@@ -5059,15 +5096,19 @@ def _update_constraint_island(m: types.Model, d: types.Data, ctx: IslandSolverCo
   """Update constraint arrays for island solver."""
   # Save prev cost, zero cost/gauss
   wp.launch(
-    update_constraint_init_cost_island,
-    dim=(d.nworld, m.ntree),
-    inputs=[d.nisland, ctx.cost, ctx.done],
-    outputs=[ctx.gauss, ctx.cost, ctx.prev_cost],
+    _update_constraint_init_cost,
+    dim=d.nworld * m.ntree,
+    inputs=[ctx.cost.reshape(-1), ctx.done.reshape(-1)],
+    outputs=[
+      ctx.gauss.reshape(-1),
+      ctx.cost.reshape(-1),
+      ctx.prev_cost.reshape(-1),
+    ],
   )
 
   # Compute force, state, cost per EFC
   wp.launch(
-    update_constraint_efc_island,
+    _update_constraint_efc_island,
     dim=(d.nworld, d.njmax),
     inputs=[
       m.opt.impratio_invsqrt,
@@ -5097,7 +5138,7 @@ def _update_constraint_island(m: types.Model, d: types.Data, ctx: IslandSolverCo
   if m.is_sparse:
     d.iqfrc_constraint.zero_()
     wp.launch(
-      update_constraint_init_qfrc_constraint_sparse_island,
+      _update_constraint_init_qfrc_constraint_sparse_island,
       dim=(d.nworld, d.njmax),
       inputs=[
         d.nefc,
@@ -5114,7 +5155,7 @@ def _update_constraint_island(m: types.Model, d: types.Data, ctx: IslandSolverCo
     )
   else:
     wp.launch(
-      update_constraint_init_qfrc_constraint_dense_island,
+      _update_constraint_init_qfrc_constraint_dense_island,
       dim=(d.nworld, m.nv),
       inputs=[
         d.nefc,
@@ -5132,7 +5173,7 @@ def _update_constraint_island(m: types.Model, d: types.Data, ctx: IslandSolverCo
 
   # Gauss cost
   wp.launch(
-    update_constraint_gauss_cost_island,
+    _update_constraint_gauss_cost_island,
     dim=(d.nworld, m.nv),
     inputs=[
       d.nidof,
@@ -5155,7 +5196,7 @@ def _update_gradient_island(m: types.Model, d: types.Data, ctx: IslandSolverCont
 
   # grad = Ma - frc_smooth - frc_constraint, accumulate grad_dot
   wp.launch(
-    update_gradient_grad_island,
+    _update_gradient_grad_island,
     dim=(d.nworld, m.nv),
     inputs=[
       d.nidof,
@@ -5187,7 +5228,7 @@ def _update_gradient_incremental_island(m: types.Model, d: types.Data, ctx: Isla
 
   # grad = Ma - frc_smooth - frc_constraint, accumulate grad_dot
   wp.launch(
-    update_gradient_grad_island,
+    _update_gradient_grad_island,
     dim=(d.nworld, m.nv),
     inputs=[
       d.nidof,
@@ -5205,13 +5246,13 @@ def _update_gradient_incremental_island(m: types.Model, d: types.Data, ctx: Isla
 
   # JTDAJ
   wp.launch(
-    update_gradient_JTDAJ_island,
+    _update_gradient_JTDAJ_island,
     dim=(d.nworld, d.njmax),
     inputs=[
       m.is_sparse,
       d.nefc,
       d.njmax,
-      d.island_dofadr,
+      d.island_idofadr,
       d.island_nv,
       d.efc.iJ_rownnz,
       d.efc.iJ_rowadr,
@@ -5228,7 +5269,7 @@ def _update_gradient_incremental_island(m: types.Model, d: types.Data, ctx: Isla
   # Add mass matrix
   if m.is_sparse:
     wp.launch(
-      update_gradient_set_h_M_sparse_island,
+      _update_gradient_set_h_M_sparse_island,
       dim=(d.nworld, m.M_fullm_i.shape[0]),
       inputs=[
         m.M_fullm_i,
@@ -5244,7 +5285,7 @@ def _update_gradient_incremental_island(m: types.Model, d: types.Data, ctx: Isla
     )
   else:
     wp.launch(
-      update_gradient_set_h_M_dense_island,
+      _update_gradient_set_h_M_dense_island,
       dim=(d.nworld, m.nv),
       inputs=[
         m.nv,
@@ -5260,7 +5301,7 @@ def _update_gradient_incremental_island(m: types.Model, d: types.Data, ctx: Isla
   # Elliptic cone correction: JTCJ
   if m.opt.cone == types.ConeType.ELLIPTIC and d.naconmax > 0:
     wp.launch(
-      update_gradient_JTCJ_island,
+      _update_gradient_JTCJ_island,
       dim=d.naconmax,
       inputs=[
         m.opt.impratio_invsqrt,
@@ -5270,7 +5311,7 @@ def _update_gradient_incremental_island(m: types.Model, d: types.Data, ctx: Isla
         d.contact.dim,
         d.contact.efc_address,
         d.contact.worldid,
-        d.island_dofadr,
+        d.island_idofadr,
         d.naconmax,
         d.nidof,
         d.map_efc2iefc,
@@ -5294,7 +5335,7 @@ def _update_gradient_incremental_island(m: types.Model, d: types.Data, ctx: Isla
     dim=(d.nworld, m.ntree),
     inputs=[
       d.nisland,
-      d.island_dofadr,
+      d.island_idofadr,
       d.island_nv,
       ctx.grad,
       ctx.h,
@@ -5322,15 +5363,15 @@ def _linesearch_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
 
   # jv = J @ search
   wp.launch(
-    linesearch_jv_island,
+    _linesearch_jv_island,
     dim=(d.nworld, d.njmax),
     inputs=[
       m.is_sparse,
       d.nefc,
       d.nidof,
+      d.island_idofadr,
       d.island_nv,
       d.njmax,
-      d.island_dofadr,
       d.efc.iJ_rownnz,
       d.efc.iJ_rowadr,
       d.efc.iJ_colind,
@@ -5344,7 +5385,7 @@ def _linesearch_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
 
   # linesearch
   wp.launch(
-    linesearch_island,
+    _linesearch_kernel_island,
     dim=(d.nworld, m.ntree),
     inputs=[
       m.opt.tolerance,
@@ -5366,7 +5407,7 @@ def _linesearch_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
       d.map_efc2iefc,
       d.njmax,
       d.nacon,
-      d.island_dofadr,
+      d.island_idofadr,
       d.efc.itype,
       d.efc.iid,
       d.efc.iD,
@@ -5386,7 +5427,7 @@ def _linesearch_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
 
   # Update iacc, iMa
   wp.launch(
-    linesearch_qacc_ma_island,
+    _linesearch_qacc_ma_island,
     dim=(d.nworld, m.nv),
     inputs=[
       d.nidof,
@@ -5401,7 +5442,7 @@ def _linesearch_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
 
   # Update Jaref
   wp.launch(
-    linesearch_jaref_island,
+    _linesearch_jaref_island,
     dim=(d.nworld, d.njmax),
     inputs=[
       d.nefc,
@@ -5431,7 +5472,7 @@ def _solver_iteration_island(
   # Save prev_grad, prev_Mgrad for CG
   if is_cg:
     wp.launch(
-      solve_prev_grad_Mgrad_island,
+      _solve_prev_grad_Mgrad_island,
       dim=(d.nworld, m.nv),
       inputs=[d.nidof, ctx.grad, ctx.Mgrad, d.dof_islandid, ctx.done],
       outputs=[ctx.prev_grad, ctx.prev_Mgrad],
@@ -5449,27 +5490,43 @@ def _solver_iteration_island(
   # Polak-Ribière beta (CG only)
   if is_cg:
     wp.launch(
-      solve_beta_island,
+      _solve_beta_island_zero,
       dim=(d.nworld, m.ntree),
+      inputs=[d.nisland],
+      outputs=[ctx.beta, ctx.beta_den],
+    )
+    wp.launch(
+      _solve_beta_island_accumulate,
+      dim=(d.nworld, m.nv),
       inputs=[
-        d.nisland,
-        d.island_nv,
-        d.island_dofadr,
+        d.nidof,
+        d.dof_islandid,
         ctx.grad,
         ctx.Mgrad,
         ctx.prev_grad,
         ctx.prev_Mgrad,
         ctx.done,
       ],
+      outputs=[ctx.beta, ctx.beta_den],
+    )
+    wp.launch(
+      _solve_beta_island_finalize,
+      dim=(d.nworld, m.ntree),
+      inputs=[d.nisland, ctx.beta, ctx.beta_den, ctx.done],
       outputs=[ctx.beta],
     )
 
   # Zero search_dot
-  ctx.search_dot.zero_()
+  wp.launch(
+    _solve_zero_search_dot,
+    dim=d.nworld * m.ntree,
+    inputs=[ctx.done.reshape(-1)],
+    outputs=[ctx.search_dot.reshape(-1)],
+  )
 
   # Search update
   wp.launch(
-    solve_search_update_island,
+    _solve_search_update_island,
     dim=(d.nworld, m.nv),
     inputs=[
       m.opt.solver,
@@ -5486,7 +5543,7 @@ def _solver_iteration_island(
   # Convergence check
   d.solver_niter.zero_()
   wp.launch(
-    solve_done_island,
+    _solve_done_island,
     dim=(d.nworld, m.ntree),
     inputs=[
       m.opt.tolerance,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
index 4ce5a31c..2fe4607d 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
@@ -28,6 +28,7 @@ MJ_MINIMP = mujoco.mjMINIMP  # minimum constraint impedance
 MJ_MAXIMP = mujoco.mjMAXIMP  # maximum constraint impedance
 MJ_MAXCONPAIR = mujoco.mjMAXCONPAIR
 MJ_MINMU = mujoco.mjMINMU  # minimum friction
+MJ_MINAWAKE = mujoco.mjMINAWAKE  # minimum number of timesteps before sleeping
 NEW_GAP_SEMANTICS = check_version("mujoco>=3.9.0.dev914519929")
 TACTILE_DEPTH_SEMANTICS = check_version("mujoco>=3.9.0.dev921980899")
 # maximum size (by number of edges) of an horizon in EPA algorithm
@@ -89,7 +90,7 @@ class BlockDim:
   update_gradient_cholesky: int = 64
   update_gradient_cholesky_blocked: int = 32
   update_gradient_JTDAJ_sparse: int = 64
-  update_gradient_JTDAJ_dense: int = 96
+  update_gradient_JTDAJ_dense: int = 128
   linesearch_iterative: int = 32
   contact_jac_tiled: int = 32
   # derivative
@@ -236,13 +237,44 @@ class EnableBit(enum.IntFlag):
   Attributes:
     ENERGY: energy computation
     INVDISCRETE: discrete-time inverse dynamics
+    SLEEP: sleeping
   """
 
   ENERGY = mujoco.mjtEnableBit.mjENBL_ENERGY
   INVDISCRETE = mujoco.mjtEnableBit.mjENBL_INVDISCRETE
+  SLEEP = mujoco.mjtEnableBit.mjENBL_SLEEP
   # unsupported: OVERRIDE, FWDINV, ISLAND
 
 
+class SleepPolicy(enum.IntEnum):
+  """Per-tree sleep policy.
+
+  Attributes:
+    AUTO: compiler chooses sleep policy
+    AUTO_NEVER: compiler sleep policy: never
+    AUTO_ALLOWED: compiler sleep policy: allowed
+  """
+
+  AUTO = mujoco.mjtSleepPolicy.mjSLEEP_AUTO
+  AUTO_NEVER = mujoco.mjtSleepPolicy.mjSLEEP_AUTO_NEVER
+  AUTO_ALLOWED = mujoco.mjtSleepPolicy.mjSLEEP_AUTO_ALLOWED
+  # unsupported: NEVER, ALLOWED, INIT
+
+
+class SleepState(enum.IntEnum):
+  """Sleep state for bodies.
+
+  Attributes:
+    STATIC: body is static (world body or mocap)
+    ASLEEP: body is asleep
+    AWAKE: body is awake
+  """
+
+  STATIC = mujoco.mjtSleepState.mjS_STATIC
+  ASLEEP = mujoco.mjtSleepState.mjS_ASLEEP
+  AWAKE = mujoco.mjtSleepState.mjS_AWAKE
+
+
 class TrnType(enum.IntEnum):
   """Type of actuator transmission.
 
@@ -749,6 +781,7 @@ class Option:
     tolerance: main solver tolerance
     ls_tolerance: CG/Newton linesearch tolerance
     ccd_tolerance: convex collision detection tolerance
+    sleep_tolerance: sleep velocity tolerance
     gravity: gravitational acceleration
     wind: wind (for lift, drag, and viscosity)
     magnetic: global magnetic flux
@@ -783,6 +816,7 @@ class Option:
   tolerance: array("*", float)
   ls_tolerance: array("*", float)
   ccd_tolerance: array("*", float)
+  sleep_tolerance: array("*", float)
   gravity: array("*", wp.vec3)
   wind: array("*", wp.vec3)
   magnetic: array("*", wp.vec3)
@@ -978,9 +1012,11 @@ class Model:
     dof_damping: damping coefficient                         (*, nv)
     dof_dampingpoly: high-order damping coefficients         (*, nv, 2)
     dof_invweight0: diag. inverse inertia in qpos0           (*, nv)
+    dof_length: dof length for weighting velocity norm       (nv,)
     tree_bodynum: number of bodies in tree (incl. root)      (ntree,)
     tree_dofadr: start address of tree's dofs                (ntree,)
     tree_dofnum: number of dofs in tree                      (ntree,)
+    tree_sleep_policy: tree sleep policy (SleepPolicy)       (ntree,)
     geom_type: geometric type (GeomType)                     (ngeom,)
     geom_contype: geom contact type                          (ngeom,)
     geom_conaffinity: geom contact affinity                  (ngeom,)
@@ -1032,6 +1068,12 @@ class Model:
     light_poscom0: global position rel. to sub-com in qpos0  (*, nlight, 3)
     light_pos0: global position rel. to body in qpos0        (*, nlight, 3)
     light_dir0: global direction in qpos0                    (*, nlight, 3)
+    light_attenuation: OpenGL constant/linear/quadratic      (*, nlight, 3)
+    light_cutoff: spotlight half-cone angle in degrees       (*, nlight)
+    light_exponent: spotlight angular falloff exponent       (*, nlight)
+    light_ambient: ambient RGB                               (*, nlight, 3)
+    light_diffuse: diffuse RGB                               (*, nlight, 3)
+    light_specular: specular RGB                             (*, nlight, 3)
     flex_contype: flex contact type                          (nflex,)
     flex_conaffinity: flex contact affinity                  (nflex,)
     flex_condim: contact dimensionality (1, 3, 4, 6)         (nflex,)
@@ -1100,6 +1142,9 @@ class Model:
     hfield_data: elevation data                              (nhfielddata,)
     mat_texid: texture id for rendering                      (*, nmat, mjNTEXROLE)
     mat_texrepeat: texture repeat for rendering              (*, nmat, 2)
+    mat_emission: emission scalar (self-illumination)        (*, nmat)
+    mat_specular: specular reflection scalar                 (*, nmat)
+    mat_shininess: shininess in [0, 1], mapped to GL [0, 128](*, nmat)
     mat_rgba: rgba                                           (*, nmat, 4)
     pair_dim: contact dimensionality                         (npair,)
     pair_geom1: id of geom1                                  (npair,)
@@ -1402,9 +1447,11 @@ class Model:
   dof_damping: array("*", "nv", float)
   dof_dampingpoly: array("*", "nv", wp.vec2)
   dof_invweight0: array("*", "nv", float)
+  dof_length: array("nv", float)
   tree_bodynum: array("ntree", int)
   tree_dofadr: array("ntree", int)
   tree_dofnum: array("ntree", int)
+  tree_sleep_policy: array("ntree", int)
   geom_type: array("ngeom", int)
   geom_contype: array("ngeom", int)
   geom_conaffinity: array("ngeom", int)
@@ -1456,6 +1503,12 @@ class Model:
   light_poscom0: array("*", "nlight", wp.vec3)
   light_pos0: array("*", "nlight", wp.vec3)
   light_dir0: array("*", "nlight", wp.vec3)
+  light_attenuation: array("*", "nlight", wp.vec3)
+  light_cutoff: array("*", "nlight", float)
+  light_exponent: array("*", "nlight", float)
+  light_ambient: array("*", "nlight", wp.vec3)
+  light_diffuse: array("*", "nlight", wp.vec3)
+  light_specular: array("*", "nlight", wp.vec3)
   flex_contype: array("nflex", int)
   flex_conaffinity: array("nflex", int)
   flex_condim: array("nflex", int)
@@ -1524,6 +1577,9 @@ class Model:
   hfield_data: array("nhfielddata", float)
   mat_texid: array("*", "nmat", 10, int)
   mat_texrepeat: array("*", "nmat", wp.vec2)
+  mat_emission: array("*", "nmat", float)
+  mat_specular: array("*", "nmat", float)
+  mat_shininess: array("*", "nmat", float)
   mat_rgba: array("*", "nmat", wp.vec4)
   pair_dim: array("npair", int)
   pair_geom1: array("npair", int)
@@ -1852,6 +1908,9 @@ class Data:
     nefc: number of constraints                                 (nworld,)
     nisland: number of constraint islands                       (nworld,)
     nidof: total DOFs in islands                                (nworld,)
+    ntree_awake: number of awake trees                          (nworld,)
+    nbody_awake: number of awake bodies                         (nworld,)
+    nv_awake: number of awake dofs                              (nworld,)
     time: simulation time                                       (nworld,)
     energy: potential, kinetic energy                           (nworld, 2)
     qpos: position                                              (nworld, nq)
@@ -1868,6 +1927,7 @@ class Data:
     qacc: acceleration                                          (nworld, nv)
     act_dot: time-derivative of actuator activation             (nworld, na)
     sensordata: sensor data array                               (nworld, nsensordata,)
+    tree_asleep: tree asleep counter; >=0: asleep cycle         (nworld, ntree)
     xpos: Cartesian position of body frame                      (nworld, nbody, 3)
     xquat: Cartesian orientation of body frame                  (nworld, nbody, 4)
     xmat: Cartesian orientation of body frame                   (nworld, nbody, 3, 3)
@@ -1906,6 +1966,10 @@ class Data:
     qLD: upper Cholesky factorization                           (nworld, nv, nv) if dense
          L'*D*L factorization of M                              (nworld, 1, nC) if sparse
     qLDiagInv: 1/diag(D)                                        (nworld, nv)
+    tree_awake: is tree awake; 0: asleep; 1: awake              (nworld, ntree)
+    body_awake: body sleep state (SleepState)                   (nworld, nbody)
+    body_awake_ind: indices of awake/static bodies              (nworld, nbody)
+    dof_awake_ind: indices of awake dofs                        (nworld, nv)
     flexedge_velocity: flex edge velocities                     (nworld, nflexedge)
     ten_velocity: tendon velocities                             (nworld, ntendon)
     actuator_velocity: actuator velocities                      (nworld, nu)
@@ -1936,6 +2000,7 @@ class Data:
     tree_island: island ID per tree (-1 if unconstrained)       (nworld, ntree)
     dof_island: island ID per DOF (-1 if unconstrained)         (nworld, nv)
     island_dofadr: island start address in dof vector           (nworld, ntree)
+    island_idofadr: island start address in idof vector          (nworld, ntree)
     island_nv: DOFs per island                                  (nworld, ntree)
     island_nefc: constraints per island                         (nworld, ntree)
     island_ne: equality constraints per island                  (nworld, ntree)
@@ -1970,6 +2035,9 @@ class Data:
   nefc: array("nworld", int)
   nisland: array("nworld", int)
   nidof: array("nworld", int)
+  ntree_awake: array("nworld", int)
+  nbody_awake: array("nworld", int)
+  nv_awake: array("nworld", int)
   time: array("nworld", float)
   energy: array("nworld", wp.vec2)
   qpos: array("nworld", "nq", float)
@@ -1986,6 +2054,7 @@ class Data:
   qacc: array("nworld", "nv", float)
   act_dot: array("nworld", "na", float)
   sensordata: array("nworld", "nsensordata", float)
+  tree_asleep: array("nworld", "ntree", int)
   xpos: array("nworld", "nbody", wp.vec3)
   xquat: array("nworld", "nbody", wp.quat)
   xmat: array("nworld", "nbody", wp.mat33)
@@ -2022,6 +2091,10 @@ class Data:
   M: wp.array3d[float]
   qLD: wp.array3d[float]
   qLDiagInv: array("nworld", "nv", float)
+  tree_awake: array("nworld", "ntree", int)
+  body_awake: array("nworld", "nbody", int)
+  body_awake_ind: array("nworld", "nbody", int)
+  dof_awake_ind: array("nworld", "nv", int)
   flexedge_velocity: array("nworld", "nflexedge", float)
   ten_velocity: array("nworld", "ntendon", float)
   actuator_velocity: array("nworld", "nu", float)
@@ -2050,6 +2123,7 @@ class Data:
   tree_island: array("nworld", "ntree", int)
   dof_island: array("nworld", "nv", int)
   island_dofadr: array("nworld", "ntree", int)
+  island_idofadr: array("nworld", "ntree", int)
   island_nv: array("nworld", "ntree", int)
   island_nefc: array("nworld", "ntree", int)
   island_ne: array("nworld", "ntree", int)
@@ -2083,9 +2157,6 @@ class InverseContext:
 
   Jaref: wp.array2d[float]
   search_dot: wp.array[float]
-  gauss: wp.array[float]
-  cost: wp.array[float]
-  prev_cost: wp.array[float]
   done: wp.array[bool]
   changed_efc_ids: wp.array2d[int]
   changed_efc_count: wp.array[int]
@@ -2115,6 +2186,7 @@ class IslandSolverContext:
   done: wp.array2d[bool]  # per-island convergence
   solver_niter: wp.array2d[int]  # iterations per island
   beta: wp.array2d[float]
+  beta_den: wp.array2d[float]
   alpha: wp.array2d[float]
   Ma: wp.array2d[float]  # island-local Ma (nworld, nv)
 
@@ -2125,9 +2197,6 @@ class SolverContext:
 
   Jaref: wp.array2d[float]
   search_dot: wp.array[float]
-  gauss: wp.array[float]
-  cost: wp.array[float]
-  prev_cost: wp.array[float]
   done: wp.array[bool]
   grad: wp.array2d[float]
   grad_dot: wp.array[float]
@@ -2138,9 +2207,11 @@ class SolverContext:
   quad: wp.array2d[wp.vec3]
   quad_gauss: wp.array[wp.vec3]
   alpha: wp.array[float]
+  improvement: wp.array[float]
   prev_grad: wp.array2d[float]
   prev_Mgrad: wp.array2d[float]
   beta: wp.array[float]
+  beta_den: wp.array[float]
   h: wp.array3d[float]
   hfactor: wp.array3d[float]
   # Incremental Hessian update (Newton only)
@@ -2158,8 +2229,8 @@ class RenderContext:
     cam_id_map: camera id map
     use_textures: whether to use textures
     use_shadows: whether to use shadows
-    use_ambient_lighting: whether to use ambient lighting
-    background_color: background color
+    use_ambient_lighting: top-level switch for ambient contributions
+    background_color: color used for missed rays when no skybox is rendered
     use_precomputed_rays: whether to use precomputed rays
     bvh_ngeom: number of geometries in the BVH
     enabled_geom_ids: enabled geometry ids
@@ -2204,11 +2275,38 @@ class RenderContext:
     render_skybox: whether to shade missed rays with the MuJoCo skybox texture
     skybox_tex_id: index into textures of the skybox (MuJoCo tex_type == SKYBOX), -1 if none
     skybox_face_width: pixel width of one skybox cube face (0 if no skybox)
+    headlight_active: whether to inject MuJoCo's vis.headlight as a synthetic
+      directional light at the active camera. Read from `mjm.vis.headlight.active`
+      at context creation; users disable the headlight by configuring it on the
+      MuJoCo model (e.g. `<visual><headlight active="0"/></visual>` in XML).
+    headlight_ambient: RGB ambient color of the headlight (from vis.headlight).
+    headlight_diffuse: RGB diffuse color of the headlight.
+    headlight_specular: RGB specular color of the headlight.
     enable_backface_culling: drop primitive ray hits whose normal faces away
       from the ray (i.e. the ray origin is inside the geom). Matches MuJoCo's
       mesh-ray rule. When False, the renderer reports inner-surface hits, which
       is faster but causes a camera placed inside a geom to render that geom's
       back wall.
+    light_attenuation_is_default: True iff every light in the model has the
+      MuJoCo default `attenuation = (1, 0, 0)`. Computed once at context
+      creation; when True the kernel skips the per-light polynomial
+      attenuation evaluation (a divide + 3 multiplies + an add per
+      non-directional light per pixel) via `wp.static`.
+    has_spot_lights: True iff any light in the model has `type == SPOT`.
+      When False, the kernel skips the spot-cone branch (cos cutoff +
+      pow exponent) per non-directional light per pixel via `wp.static`.
+    enable_specular: when True, evaluate the Phong specular highlight per
+      light per pixel (uses `mat_specular` / `mat_shininess`). When False,
+      the entire specular branch is removed at compile time. Useful for
+      depth/segmentation-only workflows or when materials are matte.
+    enable_emission: when True, add `mat_emission * base_color` to each
+      shaded pixel. When False the term is dropped at compile time.
+    enable_per_light_ambient: when True and `use_ambient_lighting` is also
+      True, sum the per-light `light_ambient` colors into each shaded pixel
+      even when the surface normal is perpendicular to the light direction
+      or the pixel is shadowed. When False the second per-light loop for
+      ambient is removed at compile time. Headlight ambient and the no-light
+      fallback are controlled by `use_ambient_lighting`.
     geom_ray_types: tuple of GeomType int values present in the scene, used to
       statically eliminate unused intersection branches in the ray-cast kernels.
   """
@@ -2224,6 +2322,10 @@ class RenderContext:
   render_skybox: bool
   skybox_tex_id: int
   skybox_face_width: int
+  headlight_active: bool
+  headlight_ambient: wp.vec3
+  headlight_diffuse: wp.vec3
+  headlight_specular: wp.vec3
   bvh_ngeom: int
   enabled_geom_ids: array("*", int)
   mesh_registry: dict
@@ -2265,4 +2367,9 @@ class RenderContext:
   znear: float
   total_rays: int
   enable_backface_culling: bool
+  enable_specular: bool
+  enable_emission: bool
+  enable_per_light_ambient: bool
+  light_attenuation_is_default: bool
+  has_spot_lights: bool
   geom_ray_types: tuple = ()
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py
index 960e4266..55fe2cfe 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py
@@ -154,28 +154,3 @@ def check_toolkit_driver():
         """,
         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
diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py
index d3cfd111..cc67809d 100644
--- a/mjx/mujoco/mjx/warp/collision_driver.py
+++ b/mjx/mujoco/mjx/warp/collision_driver.py
@@ -23,6 +23,7 @@ import mujoco.mjx.third_party.mujoco_warp as mjwarp
 from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 import warp as wp
 
+
 _m = mjwarp.Model(
     **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
 )
@@ -45,6 +46,7 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _collision_shim(
     # Model
@@ -72,6 +74,7 @@ def _collision_shim(
     flex_vertadr: wp.array[int],
     flex_vertflexid: wp.array[int],
     geom_aabb: wp.array3d[wp.vec3],
+    geom_bodyid: wp.array[int],
     geom_conaffinity: wp.array[int],
     geom_condim: wp.array[int],
     geom_contype: wp.array[int],
@@ -139,11 +142,13 @@ def _collision_shim(
     opt__ccd_iterations: int,
     opt__ccd_tolerance: wp.array[float],
     opt__disableflags: int,
+    opt__enableflags: int,
     opt__sdf_initpoints: int,
     opt__sdf_iterations: int,
     # Data
     naccdmax: int,
     naconmax: int,
+    body_awake: wp.array2d[int],
     flexvert_xpos: wp.array2d[wp.vec3],
     geom_xmat: wp.array2d[wp.mat33],
     geom_xpos: wp.array2d[wp.vec3],
@@ -194,6 +199,7 @@ def _collision_shim(
   _m.flex_vertadr = flex_vertadr
   _m.flex_vertflexid = flex_vertflexid
   _m.geom_aabb = geom_aabb
+  _m.geom_bodyid = geom_bodyid
   _m.geom_conaffinity = geom_conaffinity
   _m.geom_condim = geom_condim
   _m.geom_contype = geom_contype
@@ -252,6 +258,7 @@ def _collision_shim(
   _m.opt.ccd_iterations = opt__ccd_iterations
   _m.opt.ccd_tolerance = opt__ccd_tolerance
   _m.opt.disableflags = opt__disableflags
+  _m.opt.enableflags = opt__enableflags
   _m.opt.sdf_initpoints = opt__sdf_initpoints
   _m.opt.sdf_iterations = opt__sdf_iterations
   _m.pair_dim = pair_dim
@@ -263,6 +270,7 @@ def _collision_shim(
   _m.pair_solreffriction = pair_solreffriction
   _m.plugin = plugin
   _m.plugin_attr = plugin_attr
+  _d.body_awake = body_awake
   _d.contact.dim = contact__dim
   _d.contact.dist = contact__dist
   _d.contact.efc_address = contact__efc_address
@@ -385,6 +393,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m.flex_vertadr,
       m._impl.flex_vertflexid,
       m.geom_aabb,
+      m.geom_bodyid,
       m.geom_conaffinity,
       m.geom_condim,
       m.geom_contype,
@@ -452,10 +461,12 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m.opt._impl.ccd_iterations,
       m.opt._impl.ccd_tolerance,
       m.opt.disableflags,
+      m.opt.enableflags,
       m.opt._impl.sdf_initpoints,
       m.opt._impl.sdf_iterations,
       d._impl.naccdmax,
       d._impl.naconmax,
+      d._impl.body_awake,
       d._impl.flexvert_xpos,
       d.geom_xmat,
       d.geom_xpos,
diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py
index 5242732d..e5ae1977 100644
--- a/mjx/mujoco/mjx/warp/forward.py
+++ b/mjx/mujoco/mjx/warp/forward.py
@@ -23,6 +23,7 @@ import mujoco.mjx.third_party.mujoco_warp as mjwarp
 from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 import warp as wp
 
+
 _m = mjwarp.Model(
     **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
 )
@@ -45,6 +46,7 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _forward_shim(
     # Model
@@ -52,9 +54,6 @@ def _forward_shim(
     M_elemid: wp.array2d[int],
     M_fullm_i: wp.array[int],
     M_fullm_j: wp.array[int],
-    M_fullm_upper_elemid: wp.array[int],
-    M_fullm_upper_i: wp.array[int],
-    M_fullm_upper_j: wp.array[int],
     M_mulm_col: wp.array[int],
     M_mulm_madr: wp.array[int],
     M_mulm_rowadr: wp.array[int],
@@ -131,6 +130,7 @@ def _forward_shim(
     dof_frictionloss: wp.array2d[float],
     dof_invweight0: wp.array2d[float],
     dof_jntid: wp.array[int],
+    dof_length: wp.array[float],
     dof_parentid: wp.array[int],
     dof_solimp: wp.array2d[mjwp_types.vec5],
     dof_solref: wp.array2d[wp.vec2],
@@ -376,6 +376,7 @@ def _forward_shim(
     tendon_jnt_adr: wp.array[int],
     tendon_length0: wp.array2d[float],
     tendon_lengthspring: wp.array2d[wp.vec2],
+    tendon_limited: wp.array[int],
     tendon_limited_adr: wp.array[int],
     tendon_margin: wp.array2d[float],
     tendon_num: wp.array[int],
@@ -387,6 +388,9 @@ def _forward_shim(
     tendon_solref_lim: wp.array2d[wp.vec2],
     tendon_stiffness: wp.array2d[float],
     tendon_stiffnesspoly: wp.array2d[wp.vec2],
+    tree_dofadr: wp.array[int],
+    tree_dofnum: wp.array[int],
+    tree_sleep_policy: wp.array[int],
     wrap_geom_adr: wp.array[int],
     wrap_jnt_adr: wp.array[int],
     wrap_objid: wp.array[int],
@@ -434,6 +438,8 @@ def _forward_shim(
     actuator_length: wp.array2d[float],
     actuator_moment: wp.array2d[float],
     actuator_velocity: wp.array2d[float],
+    body_awake: wp.array2d[int],
+    body_awake_ind: wp.array2d[int],
     cacc: wp.array2d[wp.spatial_vector],
     cam_xmat: wp.array2d[wp.mat33],
     cam_xpos: wp.array2d[wp.vec3],
@@ -445,6 +451,7 @@ def _forward_shim(
     crb: wp.array2d[mjwp_types.vec10],
     ctrl: wp.array2d[float],
     cvel: wp.array2d[wp.spatial_vector],
+    dof_awake_ind: wp.array2d[int],
     dof_island: wp.array2d[int],
     dof_islandid: wp.array2d[int],
     efc_islandid: wp.array2d[int],
@@ -463,6 +470,7 @@ def _forward_shim(
     iqfrc_smooth: wp.array2d[float],
     island_dofadr: wp.array2d[int],
     island_efcadr: wp.array2d[int],
+    island_idofadr: wp.array2d[int],
     island_ne: wp.array2d[int],
     island_nefc: wp.array2d[int],
     island_nf: wp.array2d[int],
@@ -479,6 +487,7 @@ def _forward_shim(
     moment_rowadr: wp.array2d[int],
     moment_rownnz: wp.array2d[int],
     nacon: wp.array[int],
+    nbody_awake: wp.array[int],
     ncollision: wp.array[int],
     ne: wp.array[int],
     nefc: wp.array[int],
@@ -486,6 +495,8 @@ def _forward_shim(
     nidof: wp.array[int],
     nisland: wp.array[int],
     nl: wp.array[int],
+    ntree_awake: wp.array[int],
+    nv_awake: wp.array[int],
     qLD: wp.array3d[float],
     qLDiagInv: wp.array2d[float],
     qacc: wp.array2d[float],
@@ -516,6 +527,8 @@ def _forward_shim(
     ten_wrapadr: wp.array2d[int],
     ten_wrapnum: wp.array2d[int],
     time: wp.array[float],
+    tree_asleep: wp.array2d[int],
+    tree_awake: wp.array2d[int],
     tree_island: wp.array2d[int],
     wrap_obj: wp.array2d[wp.vec2i],
     wrap_xpos: wp.array2d[wp.spatial_vector],
@@ -580,9 +593,6 @@ def _forward_shim(
   _m.M_elemid = M_elemid
   _m.M_fullm_i = M_fullm_i
   _m.M_fullm_j = M_fullm_j
-  _m.M_fullm_upper_elemid = M_fullm_upper_elemid
-  _m.M_fullm_upper_i = M_fullm_upper_i
-  _m.M_fullm_upper_j = M_fullm_upper_j
   _m.M_mulm_col = M_mulm_col
   _m.M_mulm_madr = M_mulm_madr
   _m.M_mulm_rowadr = M_mulm_rowadr
@@ -659,6 +669,7 @@ def _forward_shim(
   _m.dof_frictionloss = dof_frictionloss
   _m.dof_invweight0 = dof_invweight0
   _m.dof_jntid = dof_jntid
+  _m.dof_length = dof_length
   _m.dof_parentid = dof_parentid
   _m.dof_solimp = dof_solimp
   _m.dof_solref = dof_solref
@@ -932,6 +943,7 @@ def _forward_shim(
   _m.tendon_jnt_adr = tendon_jnt_adr
   _m.tendon_length0 = tendon_length0
   _m.tendon_lengthspring = tendon_lengthspring
+  _m.tendon_limited = tendon_limited
   _m.tendon_limited_adr = tendon_limited_adr
   _m.tendon_margin = tendon_margin
   _m.tendon_num = tendon_num
@@ -943,6 +955,9 @@ def _forward_shim(
   _m.tendon_solref_lim = tendon_solref_lim
   _m.tendon_stiffness = tendon_stiffness
   _m.tendon_stiffnesspoly = tendon_stiffnesspoly
+  _m.tree_dofadr = tree_dofadr
+  _m.tree_dofnum = tree_dofnum
+  _m.tree_sleep_policy = tree_sleep_policy
   _m.wrap_geom_adr = wrap_geom_adr
   _m.wrap_jnt_adr = wrap_jnt_adr
   _m.wrap_objid = wrap_objid
@@ -957,6 +972,8 @@ def _forward_shim(
   _d.actuator_length = actuator_length
   _d.actuator_moment = actuator_moment
   _d.actuator_velocity = actuator_velocity
+  _d.body_awake = body_awake
+  _d.body_awake_ind = body_awake_ind
   _d.cacc = cacc
   _d.cam_xmat = cam_xmat
   _d.cam_xpos = cam_xpos
@@ -984,6 +1001,7 @@ def _forward_shim(
   _d.crb = crb
   _d.ctrl = ctrl
   _d.cvel = cvel
+  _d.dof_awake_ind = dof_awake_ind
   _d.dof_island = dof_island
   _d.dof_islandid = dof_islandid
   _d.efc.D = efc__D
@@ -1030,6 +1048,7 @@ def _forward_shim(
   _d.iqfrc_smooth = iqfrc_smooth
   _d.island_dofadr = island_dofadr
   _d.island_efcadr = island_efcadr
+  _d.island_idofadr = island_idofadr
   _d.island_ne = island_ne
   _d.island_nefc = island_nefc
   _d.island_nf = island_nf
@@ -1048,6 +1067,7 @@ def _forward_shim(
   _d.naccdmax = naccdmax
   _d.nacon = nacon
   _d.naconmax = naconmax
+  _d.nbody_awake = nbody_awake
   _d.ncollision = ncollision
   _d.ne = ne
   _d.nefc = nefc
@@ -1057,6 +1077,8 @@ def _forward_shim(
   _d.njmax = njmax
   _d.njmax_nnz = njmax_nnz
   _d.nl = nl
+  _d.ntree_awake = ntree_awake
+  _d.nv_awake = nv_awake
   _d.qLD = qLD
   _d.qLDiagInv = qLDiagInv
   _d.qacc = qacc
@@ -1087,6 +1109,8 @@ def _forward_shim(
   _d.ten_wrapadr = ten_wrapadr
   _d.ten_wrapnum = ten_wrapnum
   _d.time = time
+  _d.tree_asleep = tree_asleep
+  _d.tree_awake = tree_awake
   _d.tree_island = tree_island
   _d.wrap_obj = wrap_obj
   _d.wrap_xpos = wrap_xpos
@@ -1110,6 +1134,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'actuator_length': d.actuator_length.shape,
       'actuator_moment': d._impl.actuator_moment.shape,
       'actuator_velocity': d._impl.actuator_velocity.shape,
+      'body_awake': d._impl.body_awake.shape,
+      'body_awake_ind': d._impl.body_awake_ind.shape,
       'cacc': d._impl.cacc.shape,
       'cam_xmat': d.cam_xmat.shape,
       'cam_xpos': d.cam_xpos.shape,
@@ -1120,6 +1146,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'cinert': d._impl.cinert.shape,
       'crb': d._impl.crb.shape,
       'cvel': d.cvel.shape,
+      'dof_awake_ind': d._impl.dof_awake_ind.shape,
       'dof_island': d._impl.dof_island.shape,
       'dof_islandid': d._impl.dof_islandid.shape,
       'efc_islandid': d._impl.efc_islandid.shape,
@@ -1137,6 +1164,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'iqfrc_smooth': d._impl.iqfrc_smooth.shape,
       'island_dofadr': d._impl.island_dofadr.shape,
       'island_efcadr': d._impl.island_efcadr.shape,
+      'island_idofadr': d._impl.island_idofadr.shape,
       'island_ne': d._impl.island_ne.shape,
       'island_nefc': d._impl.island_nefc.shape,
       'island_nf': d._impl.island_nf.shape,
@@ -1151,6 +1179,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'moment_rowadr': d._impl.moment_rowadr.shape,
       'moment_rownnz': d._impl.moment_rownnz.shape,
       'nacon': d._impl.nacon.shape,
+      'nbody_awake': d._impl.nbody_awake.shape,
       'ncollision': d._impl.ncollision.shape,
       'ne': d._impl.ne.shape,
       'nefc': d._impl.nefc.shape,
@@ -1158,6 +1187,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'nidof': d._impl.nidof.shape,
       'nisland': d._impl.nisland.shape,
       'nl': d._impl.nl.shape,
+      'ntree_awake': d._impl.ntree_awake.shape,
+      'nv_awake': d._impl.nv_awake.shape,
       'qLD': d._impl.qLD.shape,
       'qLDiagInv': d._impl.qLDiagInv.shape,
       'qacc': d.qacc.shape,
@@ -1184,6 +1215,8 @@ 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_asleep': d._impl.tree_asleep.shape,
+      'tree_awake': d._impl.tree_awake.shape,
       'tree_island': d._impl.tree_island.shape,
       'wrap_obj': d._impl.wrap_obj.shape,
       'wrap_xpos': d._impl.wrap_xpos.shape,
@@ -1241,7 +1274,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _forward_shim,
-      num_outputs=134,
+      num_outputs=143,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames=set([
@@ -1251,6 +1284,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'actuator_length',
           'actuator_moment',
           'actuator_velocity',
+          'body_awake',
+          'body_awake_ind',
           'cacc',
           'cam_xmat',
           'cam_xpos',
@@ -1261,6 +1296,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'cinert',
           'crb',
           'cvel',
+          'dof_awake_ind',
           'dof_island',
           'dof_islandid',
           'efc_islandid',
@@ -1278,6 +1314,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'iqfrc_smooth',
           'island_dofadr',
           'island_efcadr',
+          'island_idofadr',
           'island_ne',
           'island_nefc',
           'island_nf',
@@ -1292,6 +1329,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'moment_rowadr',
           'moment_rownnz',
           'nacon',
+          'nbody_awake',
           'ncollision',
           'ne',
           'nefc',
@@ -1299,6 +1337,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'nidof',
           'nisland',
           'nl',
+          'ntree_awake',
+          'nv_awake',
           'qLD',
           'qLDiagInv',
           'qacc',
@@ -1325,6 +1365,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'ten_velocity',
           'ten_wrapadr',
           'ten_wrapnum',
+          'tree_asleep',
+          'tree_awake',
           'tree_island',
           'wrap_obj',
           'wrap_xpos',
@@ -1567,9 +1609,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.M_elemid,
       m._impl.M_fullm_i,
       m._impl.M_fullm_j,
-      m._impl.M_fullm_upper_elemid,
-      m._impl.M_fullm_upper_i,
-      m._impl.M_fullm_upper_j,
       m._impl.M_mulm_col,
       m._impl.M_mulm_madr,
       m._impl.M_mulm_rowadr,
@@ -1646,6 +1685,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m.dof_frictionloss,
       m.dof_invweight0,
       m.dof_jntid,
+      m._impl.dof_length,
       m.dof_parentid,
       m.dof_solimp,
       m.dof_solref,
@@ -1891,6 +1931,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.tendon_jnt_adr,
       m.tendon_length0,
       m.tendon_lengthspring,
+      m.tendon_limited,
       m._impl.tendon_limited_adr,
       m.tendon_margin,
       m.tendon_num,
@@ -1902,6 +1943,9 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m.tendon_solref_lim,
       m.tendon_stiffness,
       m.tendon_stiffnesspoly,
+      m._impl.tree_dofadr,
+      m._impl.tree_dofnum,
+      m._impl.tree_sleep_policy,
       m._impl.wrap_geom_adr,
       m._impl.wrap_jnt_adr,
       m.wrap_objid,
@@ -1948,6 +1992,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d.actuator_length,
       d._impl.actuator_moment,
       d._impl.actuator_velocity,
+      d._impl.body_awake,
+      d._impl.body_awake_ind,
       d._impl.cacc,
       d.cam_xmat,
       d.cam_xpos,
@@ -1959,6 +2005,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d._impl.crb,
       d.ctrl,
       d.cvel,
+      d._impl.dof_awake_ind,
       d._impl.dof_island,
       d._impl.dof_islandid,
       d._impl.efc_islandid,
@@ -1977,6 +2024,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d._impl.iqfrc_smooth,
       d._impl.island_dofadr,
       d._impl.island_efcadr,
+      d._impl.island_idofadr,
       d._impl.island_ne,
       d._impl.island_nefc,
       d._impl.island_nf,
@@ -1993,6 +2041,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d._impl.moment_rowadr,
       d._impl.moment_rownnz,
       d._impl.nacon,
+      d._impl.nbody_awake,
       d._impl.ncollision,
       d._impl.ne,
       d._impl.nefc,
@@ -2000,6 +2049,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d._impl.nidof,
       d._impl.nisland,
       d._impl.nl,
+      d._impl.ntree_awake,
+      d._impl.nv_awake,
       d._impl.qLD,
       d._impl.qLDiagInv,
       d.qacc,
@@ -2030,6 +2081,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d._impl.ten_wrapadr,
       d._impl.ten_wrapnum,
       d.time,
+      d._impl.tree_asleep,
+      d._impl.tree_awake,
       d._impl.tree_island,
       d._impl.wrap_obj,
       d._impl.wrap_xpos,
@@ -2093,134 +2146,143 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'actuator_length': out[3],
       '_impl.actuator_moment': out[4],
       '_impl.actuator_velocity': out[5],
-      '_impl.cacc': out[6],
-      'cam_xmat': out[7],
-      'cam_xpos': out[8],
-      'cdof': out[9],
-      'cdof_dot': out[10],
-      '_impl.cfrc_ext': out[11],
-      '_impl.cfrc_int': out[12],
-      '_impl.cinert': out[13],
-      '_impl.crb': out[14],
-      'cvel': out[15],
-      '_impl.dof_island': out[16],
-      '_impl.dof_islandid': out[17],
-      '_impl.efc_islandid': out[18],
-      '_impl.energy': out[19],
-      '_impl.flexedge_J': out[20],
-      '_impl.flexedge_length': out[21],
-      '_impl.flexedge_velocity': out[22],
-      '_impl.flexvert_xpos': out[23],
-      'geom_xmat': out[24],
-      'geom_xpos': out[25],
-      'history': out[26],
-      '_impl.iqacc': out[27],
-      '_impl.iqacc_smooth': out[28],
-      '_impl.iqfrc_constraint': out[29],
-      '_impl.iqfrc_smooth': out[30],
-      '_impl.island_dofadr': out[31],
-      '_impl.island_efcadr': out[32],
-      '_impl.island_ne': out[33],
-      '_impl.island_nefc': out[34],
-      '_impl.island_nf': out[35],
-      '_impl.island_nv': out[36],
-      '_impl.light_xdir': out[37],
-      '_impl.light_xpos': out[38],
-      '_impl.map_dof2idof': out[39],
-      '_impl.map_efc2iefc': out[40],
-      '_impl.map_idof2dof': out[41],
-      '_impl.map_iefc2efc': out[42],
-      '_impl.moment_colind': out[43],
-      '_impl.moment_rowadr': out[44],
-      '_impl.moment_rownnz': out[45],
-      '_impl.nacon': out[46],
-      '_impl.ncollision': out[47],
-      '_impl.ne': out[48],
-      '_impl.nefc': out[49],
-      '_impl.nf': out[50],
-      '_impl.nidof': out[51],
-      '_impl.nisland': out[52],
-      '_impl.nl': out[53],
-      '_impl.qLD': out[54],
-      '_impl.qLDiagInv': out[55],
-      'qacc': out[56],
-      'qacc_smooth': out[57],
-      'qfrc_actuator': out[58],
-      'qfrc_bias': out[59],
-      'qfrc_constraint': out[60],
-      '_impl.qfrc_damper': out[61],
-      'qfrc_fluid': out[62],
-      'qfrc_gravcomp': out[63],
-      'qfrc_passive': out[64],
-      'qfrc_smooth': out[65],
-      '_impl.qfrc_spring': out[66],
-      'qvel': out[67],
-      'sensordata': out[68],
-      'site_xmat': out[69],
-      'site_xpos': out[70],
-      '_impl.solver_niter': out[71],
-      '_impl.subtree_angmom': out[72],
-      'subtree_com': out[73],
-      '_impl.subtree_linvel': out[74],
-      '_impl.ten_J': out[75],
-      'ten_length': out[76],
-      '_impl.ten_velocity': out[77],
-      '_impl.ten_wrapadr': out[78],
-      '_impl.ten_wrapnum': out[79],
-      '_impl.tree_island': out[80],
-      '_impl.wrap_obj': out[81],
-      '_impl.wrap_xpos': out[82],
-      'xanchor': out[83],
-      'xaxis': out[84],
-      'ximat': out[85],
-      'xipos': out[86],
-      'xmat': out[87],
-      'xpos': out[88],
-      'xquat': out[89],
-      '_impl.contact__dim': out[90],
-      '_impl.contact__dist': out[91],
-      '_impl.contact__efc_address': out[92],
-      '_impl.contact__flex': out[93],
-      '_impl.contact__frame': out[94],
-      '_impl.contact__friction': out[95],
-      '_impl.contact__geom': out[96],
-      '_impl.contact__geomcollisionid': out[97],
-      '_impl.contact__includemargin': out[98],
-      '_impl.contact__pos': out[99],
-      '_impl.contact__solimp': out[100],
-      '_impl.contact__solref': out[101],
-      '_impl.contact__solreffriction': out[102],
-      '_impl.contact__type': out[103],
-      '_impl.contact__vert': out[104],
-      '_impl.contact__worldid': out[105],
-      '_impl.efc__D': out[106],
-      '_impl.efc__J': out[107],
-      '_impl.efc__J_colind': out[108],
-      '_impl.efc__J_rowadr': out[109],
-      '_impl.efc__J_rownnz': out[110],
-      '_impl.efc__Jqvel': out[111],
-      '_impl.efc__Ma': out[112],
-      '_impl.efc__aref': out[113],
-      '_impl.efc__force': out[114],
-      '_impl.efc__frictionloss': out[115],
-      '_impl.efc__iD': out[116],
-      '_impl.efc__iJ': out[117],
-      '_impl.efc__iJ_colind': out[118],
-      '_impl.efc__iJ_rowadr': out[119],
-      '_impl.efc__iJ_rownnz': out[120],
-      '_impl.efc__iaref': out[121],
-      '_impl.efc__id': out[122],
-      '_impl.efc__iforce': out[123],
-      '_impl.efc__ifrictionloss': out[124],
-      '_impl.efc__iid': out[125],
-      '_impl.efc__island': out[126],
-      '_impl.efc__istate': out[127],
-      '_impl.efc__itype': out[128],
-      '_impl.efc__margin': out[129],
-      '_impl.efc__pos': out[130],
-      '_impl.efc__state': out[131],
-      '_impl.efc__type': out[132],
-      '_impl.efc__vel': out[133],
+      '_impl.body_awake': out[6],
+      '_impl.body_awake_ind': out[7],
+      '_impl.cacc': out[8],
+      'cam_xmat': out[9],
+      'cam_xpos': out[10],
+      'cdof': out[11],
+      'cdof_dot': out[12],
+      '_impl.cfrc_ext': out[13],
+      '_impl.cfrc_int': out[14],
+      '_impl.cinert': out[15],
+      '_impl.crb': out[16],
+      'cvel': out[17],
+      '_impl.dof_awake_ind': out[18],
+      '_impl.dof_island': out[19],
+      '_impl.dof_islandid': out[20],
+      '_impl.efc_islandid': out[21],
+      '_impl.energy': out[22],
+      '_impl.flexedge_J': out[23],
+      '_impl.flexedge_length': out[24],
+      '_impl.flexedge_velocity': out[25],
+      '_impl.flexvert_xpos': out[26],
+      'geom_xmat': out[27],
+      'geom_xpos': out[28],
+      'history': out[29],
+      '_impl.iqacc': out[30],
+      '_impl.iqacc_smooth': out[31],
+      '_impl.iqfrc_constraint': out[32],
+      '_impl.iqfrc_smooth': out[33],
+      '_impl.island_dofadr': out[34],
+      '_impl.island_efcadr': out[35],
+      '_impl.island_idofadr': out[36],
+      '_impl.island_ne': out[37],
+      '_impl.island_nefc': out[38],
+      '_impl.island_nf': out[39],
+      '_impl.island_nv': out[40],
+      '_impl.light_xdir': out[41],
+      '_impl.light_xpos': out[42],
+      '_impl.map_dof2idof': out[43],
+      '_impl.map_efc2iefc': out[44],
+      '_impl.map_idof2dof': out[45],
+      '_impl.map_iefc2efc': out[46],
+      '_impl.moment_colind': out[47],
+      '_impl.moment_rowadr': out[48],
+      '_impl.moment_rownnz': out[49],
+      '_impl.nacon': out[50],
+      '_impl.nbody_awake': out[51],
+      '_impl.ncollision': out[52],
+      '_impl.ne': out[53],
+      '_impl.nefc': out[54],
+      '_impl.nf': out[55],
+      '_impl.nidof': out[56],
+      '_impl.nisland': out[57],
+      '_impl.nl': out[58],
+      '_impl.ntree_awake': out[59],
+      '_impl.nv_awake': out[60],
+      '_impl.qLD': out[61],
+      '_impl.qLDiagInv': out[62],
+      'qacc': out[63],
+      'qacc_smooth': out[64],
+      'qfrc_actuator': out[65],
+      'qfrc_bias': out[66],
+      'qfrc_constraint': out[67],
+      '_impl.qfrc_damper': out[68],
+      'qfrc_fluid': out[69],
+      'qfrc_gravcomp': out[70],
+      'qfrc_passive': out[71],
+      'qfrc_smooth': out[72],
+      '_impl.qfrc_spring': out[73],
+      'qvel': out[74],
+      'sensordata': out[75],
+      'site_xmat': out[76],
+      'site_xpos': out[77],
+      '_impl.solver_niter': out[78],
+      '_impl.subtree_angmom': out[79],
+      'subtree_com': out[80],
+      '_impl.subtree_linvel': out[81],
+      '_impl.ten_J': out[82],
+      'ten_length': out[83],
+      '_impl.ten_velocity': out[84],
+      '_impl.ten_wrapadr': out[85],
+      '_impl.ten_wrapnum': out[86],
+      '_impl.tree_asleep': out[87],
+      '_impl.tree_awake': out[88],
+      '_impl.tree_island': out[89],
+      '_impl.wrap_obj': out[90],
+      '_impl.wrap_xpos': out[91],
+      'xanchor': out[92],
+      'xaxis': out[93],
+      'ximat': out[94],
+      'xipos': out[95],
+      'xmat': out[96],
+      'xpos': out[97],
+      'xquat': out[98],
+      '_impl.contact__dim': out[99],
+      '_impl.contact__dist': out[100],
+      '_impl.contact__efc_address': out[101],
+      '_impl.contact__flex': out[102],
+      '_impl.contact__frame': out[103],
+      '_impl.contact__friction': out[104],
+      '_impl.contact__geom': out[105],
+      '_impl.contact__geomcollisionid': out[106],
+      '_impl.contact__includemargin': out[107],
+      '_impl.contact__pos': out[108],
+      '_impl.contact__solimp': out[109],
+      '_impl.contact__solref': out[110],
+      '_impl.contact__solreffriction': out[111],
+      '_impl.contact__type': out[112],
+      '_impl.contact__vert': out[113],
+      '_impl.contact__worldid': out[114],
+      '_impl.efc__D': out[115],
+      '_impl.efc__J': out[116],
+      '_impl.efc__J_colind': out[117],
+      '_impl.efc__J_rowadr': out[118],
+      '_impl.efc__J_rownnz': out[119],
+      '_impl.efc__Jqvel': out[120],
+      '_impl.efc__Ma': out[121],
+      '_impl.efc__aref': out[122],
+      '_impl.efc__force': out[123],
+      '_impl.efc__frictionloss': out[124],
+      '_impl.efc__iD': out[125],
+      '_impl.efc__iJ': out[126],
+      '_impl.efc__iJ_colind': out[127],
+      '_impl.efc__iJ_rowadr': out[128],
+      '_impl.efc__iJ_rownnz': out[129],
+      '_impl.efc__iaref': out[130],
+      '_impl.efc__id': out[131],
+      '_impl.efc__iforce': out[132],
+      '_impl.efc__ifrictionloss': out[133],
+      '_impl.efc__iid': out[134],
+      '_impl.efc__island': out[135],
+      '_impl.efc__istate': out[136],
+      '_impl.efc__itype': out[137],
+      '_impl.efc__margin': out[138],
+      '_impl.efc__pos': out[139],
+      '_impl.efc__state': out[140],
+      '_impl.efc__type': out[141],
+      '_impl.efc__vel': out[142],
   })
   return d
 
@@ -2249,9 +2311,6 @@ def _step_shim(
     M_elemid: wp.array2d[int],
     M_fullm_i: wp.array[int],
     M_fullm_j: wp.array[int],
-    M_fullm_upper_elemid: wp.array[int],
-    M_fullm_upper_i: wp.array[int],
-    M_fullm_upper_j: wp.array[int],
     M_mulm_col: wp.array[int],
     M_mulm_madr: wp.array[int],
     M_mulm_rowadr: wp.array[int],
@@ -2328,6 +2387,7 @@ def _step_shim(
     dof_frictionloss: wp.array2d[float],
     dof_invweight0: wp.array2d[float],
     dof_jntid: wp.array[int],
+    dof_length: wp.array[float],
     dof_parentid: wp.array[int],
     dof_solimp: wp.array2d[mjwp_types.vec5],
     dof_solref: wp.array2d[wp.vec2],
@@ -2578,6 +2638,7 @@ def _step_shim(
     tendon_jnt_adr: wp.array[int],
     tendon_length0: wp.array2d[float],
     tendon_lengthspring: wp.array2d[wp.vec2],
+    tendon_limited: wp.array[int],
     tendon_limited_adr: wp.array[int],
     tendon_margin: wp.array2d[float],
     tendon_num: wp.array[int],
@@ -2589,6 +2650,9 @@ def _step_shim(
     tendon_solref_lim: wp.array2d[wp.vec2],
     tendon_stiffness: wp.array2d[float],
     tendon_stiffnesspoly: wp.array2d[wp.vec2],
+    tree_dofadr: wp.array[int],
+    tree_dofnum: wp.array[int],
+    tree_sleep_policy: wp.array[int],
     wrap_geom_adr: wp.array[int],
     wrap_jnt_adr: wp.array[int],
     wrap_objid: wp.array[int],
@@ -2618,6 +2682,7 @@ def _step_shim(
     opt__run_collision_detection: bool,
     opt__sdf_initpoints: int,
     opt__sdf_iterations: int,
+    opt__sleep_tolerance: wp.array[float],
     opt__solver: int,
     opt__timestep: wp.array[float],
     opt__tolerance: wp.array[float],
@@ -2636,6 +2701,8 @@ def _step_shim(
     actuator_length: wp.array2d[float],
     actuator_moment: wp.array2d[float],
     actuator_velocity: wp.array2d[float],
+    body_awake: wp.array2d[int],
+    body_awake_ind: wp.array2d[int],
     cacc: wp.array2d[wp.spatial_vector],
     cam_xmat: wp.array2d[wp.mat33],
     cam_xpos: wp.array2d[wp.vec3],
@@ -2647,6 +2714,7 @@ def _step_shim(
     crb: wp.array2d[mjwp_types.vec10],
     ctrl: wp.array2d[float],
     cvel: wp.array2d[wp.spatial_vector],
+    dof_awake_ind: wp.array2d[int],
     dof_island: wp.array2d[int],
     dof_islandid: wp.array2d[int],
     efc_islandid: wp.array2d[int],
@@ -2665,6 +2733,7 @@ def _step_shim(
     iqfrc_smooth: wp.array2d[float],
     island_dofadr: wp.array2d[int],
     island_efcadr: wp.array2d[int],
+    island_idofadr: wp.array2d[int],
     island_ne: wp.array2d[int],
     island_nefc: wp.array2d[int],
     island_nf: wp.array2d[int],
@@ -2681,6 +2750,7 @@ def _step_shim(
     moment_rowadr: wp.array2d[int],
     moment_rownnz: wp.array2d[int],
     nacon: wp.array[int],
+    nbody_awake: wp.array[int],
     ncollision: wp.array[int],
     ne: wp.array[int],
     nefc: wp.array[int],
@@ -2688,6 +2758,8 @@ def _step_shim(
     nidof: wp.array[int],
     nisland: wp.array[int],
     nl: wp.array[int],
+    ntree_awake: wp.array[int],
+    nv_awake: wp.array[int],
     qLD: wp.array3d[float],
     qLDiagInv: wp.array2d[float],
     qLU: wp.array3d[float],
@@ -2719,6 +2791,8 @@ def _step_shim(
     ten_wrapadr: wp.array2d[int],
     ten_wrapnum: wp.array2d[int],
     time: wp.array[float],
+    tree_asleep: wp.array2d[int],
+    tree_awake: wp.array2d[int],
     tree_island: wp.array2d[int],
     wrap_obj: wp.array2d[wp.vec2i],
     wrap_xpos: wp.array2d[wp.spatial_vector],
@@ -2787,9 +2861,6 @@ def _step_shim(
   _m.M_elemid = M_elemid
   _m.M_fullm_i = M_fullm_i
   _m.M_fullm_j = M_fullm_j
-  _m.M_fullm_upper_elemid = M_fullm_upper_elemid
-  _m.M_fullm_upper_i = M_fullm_upper_i
-  _m.M_fullm_upper_j = M_fullm_upper_j
   _m.M_mulm_col = M_mulm_col
   _m.M_mulm_madr = M_mulm_madr
   _m.M_mulm_rowadr = M_mulm_rowadr
@@ -2866,6 +2937,7 @@ def _step_shim(
   _m.dof_frictionloss = dof_frictionloss
   _m.dof_invweight0 = dof_invweight0
   _m.dof_jntid = dof_jntid
+  _m.dof_length = dof_length
   _m.dof_parentid = dof_parentid
   _m.dof_solimp = dof_solimp
   _m.dof_solref = dof_solref
@@ -3069,6 +3141,7 @@ def _step_shim(
   _m.opt.run_collision_detection = opt__run_collision_detection
   _m.opt.sdf_initpoints = opt__sdf_initpoints
   _m.opt.sdf_iterations = opt__sdf_iterations
+  _m.opt.sleep_tolerance = opt__sleep_tolerance
   _m.opt.solver = opt__solver
   _m.opt.timestep = opt__timestep
   _m.opt.tolerance = opt__tolerance
@@ -3144,6 +3217,7 @@ def _step_shim(
   _m.tendon_jnt_adr = tendon_jnt_adr
   _m.tendon_length0 = tendon_length0
   _m.tendon_lengthspring = tendon_lengthspring
+  _m.tendon_limited = tendon_limited
   _m.tendon_limited_adr = tendon_limited_adr
   _m.tendon_margin = tendon_margin
   _m.tendon_num = tendon_num
@@ -3155,6 +3229,9 @@ def _step_shim(
   _m.tendon_solref_lim = tendon_solref_lim
   _m.tendon_stiffness = tendon_stiffness
   _m.tendon_stiffnesspoly = tendon_stiffnesspoly
+  _m.tree_dofadr = tree_dofadr
+  _m.tree_dofnum = tree_dofnum
+  _m.tree_sleep_policy = tree_sleep_policy
   _m.wrap_geom_adr = wrap_geom_adr
   _m.wrap_jnt_adr = wrap_jnt_adr
   _m.wrap_objid = wrap_objid
@@ -3169,6 +3246,8 @@ def _step_shim(
   _d.actuator_length = actuator_length
   _d.actuator_moment = actuator_moment
   _d.actuator_velocity = actuator_velocity
+  _d.body_awake = body_awake
+  _d.body_awake_ind = body_awake_ind
   _d.cacc = cacc
   _d.cam_xmat = cam_xmat
   _d.cam_xpos = cam_xpos
@@ -3196,6 +3275,7 @@ def _step_shim(
   _d.crb = crb
   _d.ctrl = ctrl
   _d.cvel = cvel
+  _d.dof_awake_ind = dof_awake_ind
   _d.dof_island = dof_island
   _d.dof_islandid = dof_islandid
   _d.efc.D = efc__D
@@ -3242,6 +3322,7 @@ def _step_shim(
   _d.iqfrc_smooth = iqfrc_smooth
   _d.island_dofadr = island_dofadr
   _d.island_efcadr = island_efcadr
+  _d.island_idofadr = island_idofadr
   _d.island_ne = island_ne
   _d.island_nefc = island_nefc
   _d.island_nf = island_nf
@@ -3260,6 +3341,7 @@ def _step_shim(
   _d.naccdmax = naccdmax
   _d.nacon = nacon
   _d.naconmax = naconmax
+  _d.nbody_awake = nbody_awake
   _d.ncollision = ncollision
   _d.ne = ne
   _d.nefc = nefc
@@ -3269,6 +3351,8 @@ def _step_shim(
   _d.njmax = njmax
   _d.njmax_nnz = njmax_nnz
   _d.nl = nl
+  _d.ntree_awake = ntree_awake
+  _d.nv_awake = nv_awake
   _d.qLD = qLD
   _d.qLDiagInv = qLDiagInv
   _d.qLU = qLU
@@ -3300,6 +3384,8 @@ def _step_shim(
   _d.ten_wrapadr = ten_wrapadr
   _d.ten_wrapnum = ten_wrapnum
   _d.time = time
+  _d.tree_asleep = tree_asleep
+  _d.tree_awake = tree_awake
   _d.tree_island = tree_island
   _d.wrap_obj = wrap_obj
   _d.wrap_xpos = wrap_xpos
@@ -3324,6 +3410,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'actuator_length': d.actuator_length.shape,
       'actuator_moment': d._impl.actuator_moment.shape,
       'actuator_velocity': d._impl.actuator_velocity.shape,
+      'body_awake': d._impl.body_awake.shape,
+      'body_awake_ind': d._impl.body_awake_ind.shape,
       'cacc': d._impl.cacc.shape,
       'cam_xmat': d.cam_xmat.shape,
       'cam_xpos': d.cam_xpos.shape,
@@ -3334,6 +3422,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'cinert': d._impl.cinert.shape,
       'crb': d._impl.crb.shape,
       'cvel': d.cvel.shape,
+      'dof_awake_ind': d._impl.dof_awake_ind.shape,
       'dof_island': d._impl.dof_island.shape,
       'dof_islandid': d._impl.dof_islandid.shape,
       'efc_islandid': d._impl.efc_islandid.shape,
@@ -3351,6 +3440,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'iqfrc_smooth': d._impl.iqfrc_smooth.shape,
       'island_dofadr': d._impl.island_dofadr.shape,
       'island_efcadr': d._impl.island_efcadr.shape,
+      'island_idofadr': d._impl.island_idofadr.shape,
       'island_ne': d._impl.island_ne.shape,
       'island_nefc': d._impl.island_nefc.shape,
       'island_nf': d._impl.island_nf.shape,
@@ -3365,6 +3455,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'moment_rowadr': d._impl.moment_rowadr.shape,
       'moment_rownnz': d._impl.moment_rownnz.shape,
       'nacon': d._impl.nacon.shape,
+      'nbody_awake': d._impl.nbody_awake.shape,
       'ncollision': d._impl.ncollision.shape,
       'ne': d._impl.ne.shape,
       'nefc': d._impl.nefc.shape,
@@ -3372,6 +3463,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'nidof': d._impl.nidof.shape,
       'nisland': d._impl.nisland.shape,
       'nl': d._impl.nl.shape,
+      'ntree_awake': d._impl.ntree_awake.shape,
+      'nv_awake': d._impl.nv_awake.shape,
       'qLD': d._impl.qLD.shape,
       'qLDiagInv': d._impl.qLDiagInv.shape,
       'qLU': d._impl.qLU.shape,
@@ -3402,6 +3495,8 @@ 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_asleep': d._impl.tree_asleep.shape,
+      'tree_awake': d._impl.tree_awake.shape,
       'tree_island': d._impl.tree_island.shape,
       'wrap_obj': d._impl.wrap_obj.shape,
       'wrap_xpos': d._impl.wrap_xpos.shape,
@@ -3459,7 +3554,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _step_shim,
-      num_outputs=139,
+      num_outputs=148,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames=set([
@@ -3470,6 +3565,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'actuator_length',
           'actuator_moment',
           'actuator_velocity',
+          'body_awake',
+          'body_awake_ind',
           'cacc',
           'cam_xmat',
           'cam_xpos',
@@ -3480,6 +3577,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'cinert',
           'crb',
           'cvel',
+          'dof_awake_ind',
           'dof_island',
           'dof_islandid',
           'efc_islandid',
@@ -3497,6 +3595,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'iqfrc_smooth',
           'island_dofadr',
           'island_efcadr',
+          'island_idofadr',
           'island_ne',
           'island_nefc',
           'island_nf',
@@ -3511,6 +3610,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'moment_rowadr',
           'moment_rownnz',
           'nacon',
+          'nbody_awake',
           'ncollision',
           'ne',
           'nefc',
@@ -3518,6 +3618,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'nidof',
           'nisland',
           'nl',
+          'ntree_awake',
+          'nv_awake',
           'qLD',
           'qLDiagInv',
           'qLU',
@@ -3548,6 +3650,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'ten_wrapadr',
           'ten_wrapnum',
           'time',
+          'tree_asleep',
+          'tree_awake',
           'tree_island',
           'wrap_obj',
           'wrap_xpos',
@@ -3798,9 +3902,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.M_elemid,
       m._impl.M_fullm_i,
       m._impl.M_fullm_j,
-      m._impl.M_fullm_upper_elemid,
-      m._impl.M_fullm_upper_i,
-      m._impl.M_fullm_upper_j,
       m._impl.M_mulm_col,
       m._impl.M_mulm_madr,
       m._impl.M_mulm_rowadr,
@@ -3877,6 +3978,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.dof_frictionloss,
       m.dof_invweight0,
       m.dof_jntid,
+      m._impl.dof_length,
       m.dof_parentid,
       m.dof_solimp,
       m.dof_solref,
@@ -4127,6 +4229,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.tendon_jnt_adr,
       m.tendon_length0,
       m.tendon_lengthspring,
+      m.tendon_limited,
       m._impl.tendon_limited_adr,
       m.tendon_margin,
       m.tendon_num,
@@ -4138,6 +4241,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.tendon_solref_lim,
       m.tendon_stiffness,
       m.tendon_stiffnesspoly,
+      m._impl.tree_dofadr,
+      m._impl.tree_dofnum,
+      m._impl.tree_sleep_policy,
       m._impl.wrap_geom_adr,
       m._impl.wrap_jnt_adr,
       m.wrap_objid,
@@ -4167,6 +4273,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.opt._impl.run_collision_detection,
       m.opt._impl.sdf_initpoints,
       m.opt._impl.sdf_iterations,
+      m.opt._impl.sleep_tolerance,
       m.opt.solver,
       m.opt.timestep,
       m.opt.tolerance,
@@ -4184,6 +4291,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d.actuator_length,
       d._impl.actuator_moment,
       d._impl.actuator_velocity,
+      d._impl.body_awake,
+      d._impl.body_awake_ind,
       d._impl.cacc,
       d.cam_xmat,
       d.cam_xpos,
@@ -4195,6 +4304,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.crb,
       d.ctrl,
       d.cvel,
+      d._impl.dof_awake_ind,
       d._impl.dof_island,
       d._impl.dof_islandid,
       d._impl.efc_islandid,
@@ -4213,6 +4323,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.iqfrc_smooth,
       d._impl.island_dofadr,
       d._impl.island_efcadr,
+      d._impl.island_idofadr,
       d._impl.island_ne,
       d._impl.island_nefc,
       d._impl.island_nf,
@@ -4229,6 +4340,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.moment_rowadr,
       d._impl.moment_rownnz,
       d._impl.nacon,
+      d._impl.nbody_awake,
       d._impl.ncollision,
       d._impl.ne,
       d._impl.nefc,
@@ -4236,6 +4348,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.nidof,
       d._impl.nisland,
       d._impl.nl,
+      d._impl.ntree_awake,
+      d._impl.nv_awake,
       d._impl.qLD,
       d._impl.qLDiagInv,
       d._impl.qLU,
@@ -4267,6 +4381,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.ten_wrapadr,
       d._impl.ten_wrapnum,
       d.time,
+      d._impl.tree_asleep,
+      d._impl.tree_awake,
       d._impl.tree_island,
       d._impl.wrap_obj,
       d._impl.wrap_xpos,
@@ -4331,138 +4447,147 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'actuator_length': out[4],
       '_impl.actuator_moment': out[5],
       '_impl.actuator_velocity': out[6],
-      '_impl.cacc': out[7],
-      'cam_xmat': out[8],
-      'cam_xpos': out[9],
-      'cdof': out[10],
-      'cdof_dot': out[11],
-      '_impl.cfrc_ext': out[12],
-      '_impl.cfrc_int': out[13],
-      '_impl.cinert': out[14],
-      '_impl.crb': out[15],
-      'cvel': out[16],
-      '_impl.dof_island': out[17],
-      '_impl.dof_islandid': out[18],
-      '_impl.efc_islandid': out[19],
-      '_impl.energy': out[20],
-      '_impl.flexedge_J': out[21],
-      '_impl.flexedge_length': out[22],
-      '_impl.flexedge_velocity': out[23],
-      '_impl.flexvert_xpos': out[24],
-      'geom_xmat': out[25],
-      'geom_xpos': out[26],
-      'history': out[27],
-      '_impl.iqacc': out[28],
-      '_impl.iqacc_smooth': out[29],
-      '_impl.iqfrc_constraint': out[30],
-      '_impl.iqfrc_smooth': out[31],
-      '_impl.island_dofadr': out[32],
-      '_impl.island_efcadr': out[33],
-      '_impl.island_ne': out[34],
-      '_impl.island_nefc': out[35],
-      '_impl.island_nf': out[36],
-      '_impl.island_nv': out[37],
-      '_impl.light_xdir': out[38],
-      '_impl.light_xpos': out[39],
-      '_impl.map_dof2idof': out[40],
-      '_impl.map_efc2iefc': out[41],
-      '_impl.map_idof2dof': out[42],
-      '_impl.map_iefc2efc': out[43],
-      '_impl.moment_colind': out[44],
-      '_impl.moment_rowadr': out[45],
-      '_impl.moment_rownnz': out[46],
-      '_impl.nacon': out[47],
-      '_impl.ncollision': out[48],
-      '_impl.ne': out[49],
-      '_impl.nefc': out[50],
-      '_impl.nf': out[51],
-      '_impl.nidof': out[52],
-      '_impl.nisland': out[53],
-      '_impl.nl': out[54],
-      '_impl.qLD': out[55],
-      '_impl.qLDiagInv': out[56],
-      '_impl.qLU': out[57],
-      'qacc': out[58],
-      'qacc_smooth': out[59],
-      'qacc_warmstart': out[60],
-      'qfrc_actuator': out[61],
-      'qfrc_bias': out[62],
-      'qfrc_constraint': out[63],
-      '_impl.qfrc_damper': out[64],
-      'qfrc_fluid': out[65],
-      'qfrc_gravcomp': out[66],
-      'qfrc_passive': out[67],
-      'qfrc_smooth': out[68],
-      '_impl.qfrc_spring': out[69],
-      'qpos': out[70],
-      'qvel': out[71],
-      'sensordata': out[72],
-      'site_xmat': out[73],
-      'site_xpos': out[74],
-      '_impl.solver_niter': out[75],
-      '_impl.subtree_angmom': out[76],
-      'subtree_com': out[77],
-      '_impl.subtree_linvel': out[78],
-      '_impl.ten_J': out[79],
-      'ten_length': out[80],
-      '_impl.ten_velocity': out[81],
-      '_impl.ten_wrapadr': out[82],
-      '_impl.ten_wrapnum': out[83],
-      'time': out[84],
-      '_impl.tree_island': out[85],
-      '_impl.wrap_obj': out[86],
-      '_impl.wrap_xpos': out[87],
-      'xanchor': out[88],
-      'xaxis': out[89],
-      'ximat': out[90],
-      'xipos': out[91],
-      'xmat': out[92],
-      'xpos': out[93],
-      'xquat': out[94],
-      '_impl.contact__dim': out[95],
-      '_impl.contact__dist': out[96],
-      '_impl.contact__efc_address': out[97],
-      '_impl.contact__flex': out[98],
-      '_impl.contact__frame': out[99],
-      '_impl.contact__friction': out[100],
-      '_impl.contact__geom': out[101],
-      '_impl.contact__geomcollisionid': out[102],
-      '_impl.contact__includemargin': out[103],
-      '_impl.contact__pos': out[104],
-      '_impl.contact__solimp': out[105],
-      '_impl.contact__solref': out[106],
-      '_impl.contact__solreffriction': out[107],
-      '_impl.contact__type': out[108],
-      '_impl.contact__vert': out[109],
-      '_impl.contact__worldid': out[110],
-      '_impl.efc__D': out[111],
-      '_impl.efc__J': out[112],
-      '_impl.efc__J_colind': out[113],
-      '_impl.efc__J_rowadr': out[114],
-      '_impl.efc__J_rownnz': out[115],
-      '_impl.efc__Jqvel': out[116],
-      '_impl.efc__Ma': out[117],
-      '_impl.efc__aref': out[118],
-      '_impl.efc__force': out[119],
-      '_impl.efc__frictionloss': out[120],
-      '_impl.efc__iD': out[121],
-      '_impl.efc__iJ': out[122],
-      '_impl.efc__iJ_colind': out[123],
-      '_impl.efc__iJ_rowadr': out[124],
-      '_impl.efc__iJ_rownnz': out[125],
-      '_impl.efc__iaref': out[126],
-      '_impl.efc__id': out[127],
-      '_impl.efc__iforce': out[128],
-      '_impl.efc__ifrictionloss': out[129],
-      '_impl.efc__iid': out[130],
-      '_impl.efc__island': out[131],
-      '_impl.efc__istate': out[132],
-      '_impl.efc__itype': out[133],
-      '_impl.efc__margin': out[134],
-      '_impl.efc__pos': out[135],
-      '_impl.efc__state': out[136],
-      '_impl.efc__type': out[137],
-      '_impl.efc__vel': out[138],
+      '_impl.body_awake': out[7],
+      '_impl.body_awake_ind': out[8],
+      '_impl.cacc': out[9],
+      'cam_xmat': out[10],
+      'cam_xpos': out[11],
+      'cdof': out[12],
+      'cdof_dot': out[13],
+      '_impl.cfrc_ext': out[14],
+      '_impl.cfrc_int': out[15],
+      '_impl.cinert': out[16],
+      '_impl.crb': out[17],
+      'cvel': out[18],
+      '_impl.dof_awake_ind': out[19],
+      '_impl.dof_island': out[20],
+      '_impl.dof_islandid': out[21],
+      '_impl.efc_islandid': out[22],
+      '_impl.energy': out[23],
+      '_impl.flexedge_J': out[24],
+      '_impl.flexedge_length': out[25],
+      '_impl.flexedge_velocity': out[26],
+      '_impl.flexvert_xpos': out[27],
+      'geom_xmat': out[28],
+      'geom_xpos': out[29],
+      'history': out[30],
+      '_impl.iqacc': out[31],
+      '_impl.iqacc_smooth': out[32],
+      '_impl.iqfrc_constraint': out[33],
+      '_impl.iqfrc_smooth': out[34],
+      '_impl.island_dofadr': out[35],
+      '_impl.island_efcadr': out[36],
+      '_impl.island_idofadr': out[37],
+      '_impl.island_ne': out[38],
+      '_impl.island_nefc': out[39],
+      '_impl.island_nf': out[40],
+      '_impl.island_nv': out[41],
+      '_impl.light_xdir': out[42],
+      '_impl.light_xpos': out[43],
+      '_impl.map_dof2idof': out[44],
+      '_impl.map_efc2iefc': out[45],
+      '_impl.map_idof2dof': out[46],
+      '_impl.map_iefc2efc': out[47],
+      '_impl.moment_colind': out[48],
+      '_impl.moment_rowadr': out[49],
+      '_impl.moment_rownnz': out[50],
+      '_impl.nacon': out[51],
+      '_impl.nbody_awake': out[52],
+      '_impl.ncollision': out[53],
+      '_impl.ne': out[54],
+      '_impl.nefc': out[55],
+      '_impl.nf': out[56],
+      '_impl.nidof': out[57],
+      '_impl.nisland': out[58],
+      '_impl.nl': out[59],
+      '_impl.ntree_awake': out[60],
+      '_impl.nv_awake': out[61],
+      '_impl.qLD': out[62],
+      '_impl.qLDiagInv': out[63],
+      '_impl.qLU': out[64],
+      'qacc': out[65],
+      'qacc_smooth': out[66],
+      'qacc_warmstart': out[67],
+      'qfrc_actuator': out[68],
+      'qfrc_bias': out[69],
+      'qfrc_constraint': out[70],
+      '_impl.qfrc_damper': out[71],
+      'qfrc_fluid': out[72],
+      'qfrc_gravcomp': out[73],
+      'qfrc_passive': out[74],
+      'qfrc_smooth': out[75],
+      '_impl.qfrc_spring': out[76],
+      'qpos': out[77],
+      'qvel': out[78],
+      'sensordata': out[79],
+      'site_xmat': out[80],
+      'site_xpos': out[81],
+      '_impl.solver_niter': out[82],
+      '_impl.subtree_angmom': out[83],
+      'subtree_com': out[84],
+      '_impl.subtree_linvel': out[85],
+      '_impl.ten_J': out[86],
+      'ten_length': out[87],
+      '_impl.ten_velocity': out[88],
+      '_impl.ten_wrapadr': out[89],
+      '_impl.ten_wrapnum': out[90],
+      'time': out[91],
+      '_impl.tree_asleep': out[92],
+      '_impl.tree_awake': out[93],
+      '_impl.tree_island': out[94],
+      '_impl.wrap_obj': out[95],
+      '_impl.wrap_xpos': out[96],
+      'xanchor': out[97],
+      'xaxis': out[98],
+      'ximat': out[99],
+      'xipos': out[100],
+      'xmat': out[101],
+      'xpos': out[102],
+      'xquat': out[103],
+      '_impl.contact__dim': out[104],
+      '_impl.contact__dist': out[105],
+      '_impl.contact__efc_address': out[106],
+      '_impl.contact__flex': out[107],
+      '_impl.contact__frame': out[108],
+      '_impl.contact__friction': out[109],
+      '_impl.contact__geom': out[110],
+      '_impl.contact__geomcollisionid': out[111],
+      '_impl.contact__includemargin': out[112],
+      '_impl.contact__pos': out[113],
+      '_impl.contact__solimp': out[114],
+      '_impl.contact__solref': out[115],
+      '_impl.contact__solreffriction': out[116],
+      '_impl.contact__type': out[117],
+      '_impl.contact__vert': out[118],
+      '_impl.contact__worldid': out[119],
+      '_impl.efc__D': out[120],
+      '_impl.efc__J': out[121],
+      '_impl.efc__J_colind': out[122],
+      '_impl.efc__J_rowadr': out[123],
+      '_impl.efc__J_rownnz': out[124],
+      '_impl.efc__Jqvel': out[125],
+      '_impl.efc__Ma': out[126],
+      '_impl.efc__aref': out[127],
+      '_impl.efc__force': out[128],
+      '_impl.efc__frictionloss': out[129],
+      '_impl.efc__iD': out[130],
+      '_impl.efc__iJ': out[131],
+      '_impl.efc__iJ_colind': out[132],
+      '_impl.efc__iJ_rowadr': out[133],
+      '_impl.efc__iJ_rownnz': out[134],
+      '_impl.efc__iaref': out[135],
+      '_impl.efc__id': out[136],
+      '_impl.efc__iforce': out[137],
+      '_impl.efc__ifrictionloss': out[138],
+      '_impl.efc__iid': out[139],
+      '_impl.efc__island': out[140],
+      '_impl.efc__istate': out[141],
+      '_impl.efc__itype': out[142],
+      '_impl.efc__margin': out[143],
+      '_impl.efc__pos': out[144],
+      '_impl.efc__state': out[145],
+      '_impl.efc__type': out[146],
+      '_impl.efc__vel': out[147],
   })
   return d
 
diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py
index ae82028f..b0de6c5c 100644
--- a/mjx/mujoco/mjx/warp/render.py
+++ b/mjx/mujoco/mjx/warp/render.py
@@ -67,9 +67,18 @@ def _render_shim(
     geom_size: wp.array2d[wp.vec3],
     geom_type: wp.array[int],
     light_active: wp.array2d[bool],
+    light_ambient: wp.array2d[wp.vec3],
+    light_attenuation: wp.array2d[wp.vec3],
     light_castshadow: wp.array2d[bool],
+    light_cutoff: wp.array2d[float],
+    light_diffuse: wp.array2d[wp.vec3],
+    light_exponent: wp.array2d[float],
+    light_specular: wp.array2d[wp.vec3],
     light_type: wp.array2d[int],
+    mat_emission: wp.array2d[float],
     mat_rgba: wp.array2d[wp.vec4],
+    mat_shininess: wp.array2d[float],
+    mat_specular: wp.array2d[float],
     mat_texid: wp.array3d[int],
     mat_texrepeat: wp.array2d[wp.vec2],
     mesh_faceadr: wp.array[int],
@@ -107,9 +116,18 @@ def _render_shim(
   _m.geom_size = geom_size
   _m.geom_type = geom_type
   _m.light_active = light_active
+  _m.light_ambient = light_ambient
+  _m.light_attenuation = light_attenuation
   _m.light_castshadow = light_castshadow
+  _m.light_cutoff = light_cutoff
+  _m.light_diffuse = light_diffuse
+  _m.light_exponent = light_exponent
+  _m.light_specular = light_specular
   _m.light_type = light_type
+  _m.mat_emission = mat_emission
   _m.mat_rgba = mat_rgba
+  _m.mat_shininess = mat_shininess
+  _m.mat_specular = mat_specular
   _m.mat_texid = mat_texid
   _m.mat_texrepeat = mat_texrepeat
   _m.mesh_faceadr = mesh_faceadr
@@ -154,6 +172,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
           'geom_xpos',
           'light_active',
           'light_castshadow',
+          'light_cutoff',
           'light_type',
           'mat_rgba',
           'mat_texid',
@@ -178,9 +197,18 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
       m.geom_size,
       m.geom_type,
       m.light_active,
+      m._impl.light_ambient,
+      m._impl.light_attenuation,
       m.light_castshadow,
+      m.light_cutoff,
+      m._impl.light_diffuse,
+      m._impl.light_exponent,
+      m._impl.light_specular,
       m.light_type,
+      m._impl.mat_emission,
       m.mat_rgba,
+      m._impl.mat_shininess,
+      m._impl.mat_specular,
       m.mat_texid,
       m._impl.mat_texrepeat,
       m.mesh_faceadr,
diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py
index bf26a123..ce765d8b 100644
--- a/mjx/mujoco/mjx/warp/smooth.py
+++ b/mjx/mujoco/mjx/warp/smooth.py
@@ -23,6 +23,7 @@ import mujoco.mjx.third_party.mujoco_warp as mjwarp
 from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 import warp as wp
 
+
 _m = mjwarp.Model(
     **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
 )
diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py
index cc8f089d..7130b49b 100644
--- a/mjx/mujoco/mjx/warp/types.py
+++ b/mjx/mujoco/mjx/warp/types.py
@@ -23,7 +23,6 @@ from jax import tree_util
 from jax.interpreters import batching
 from mujoco.mjx._src import dataclasses as mjx_dataclasses
 import numpy as np
-
 if typing.TYPE_CHECKING:
   GraphMode = int
 
@@ -145,6 +144,7 @@ class OptionWarp(PyTreeNode):
   run_collision_detection: bool
   sdf_initpoints: int
   sdf_iterations: int
+  sleep_tolerance: jax.Array
 
 class ModelWarp(PyTreeNode):
   """Derived fields from Model."""
@@ -178,6 +178,7 @@ class ModelWarp(PyTreeNode):
   callback: Callback
   cam_projection: np.ndarray
   collision_sensor_adr: np.ndarray
+  dof_length: np.ndarray
   dof_tri_col: np.ndarray
   dof_tri_row: np.ndarray
   eq_connect_adr: np.ndarray
@@ -231,11 +232,19 @@ class ModelWarp(PyTreeNode):
   is_sparse: bool
   jnt_limited_ball_adr: np.ndarray
   jnt_limited_slide_hinge_adr: np.ndarray
+  light_ambient: jax.Array
+  light_attenuation: jax.Array
   light_bodyid: np.ndarray
+  light_diffuse: jax.Array
+  light_exponent: jax.Array
+  light_specular: jax.Array
   light_targetbodyid: np.ndarray
   mapD2M: np.ndarray
   mapM2D: np.ndarray
   mapM2M: np.ndarray
+  mat_emission: jax.Array
+  mat_shininess: jax.Array
+  mat_specular: jax.Array
   mat_texrepeat: jax.Array
   max_ten_J_rownnz: int
   mesh_polyadr: np.ndarray
@@ -321,6 +330,7 @@ class ModelWarp(PyTreeNode):
   tree_bodynum: np.ndarray
   tree_dofadr: np.ndarray
   tree_dofnum: np.ndarray
+  tree_sleep_policy: np.ndarray
   wrap_geom_adr: np.ndarray
   wrap_jnt_adr: np.ndarray
   wrap_pulley_scale: np.ndarray
@@ -332,6 +342,8 @@ class DataWarp(PyTreeNode):
   M: jax.Array
   actuator_moment: jax.Array
   actuator_velocity: jax.Array
+  body_awake: jax.Array
+  body_awake_ind: jax.Array
   cacc: jax.Array
   cfrc_ext: jax.Array
   cfrc_int: jax.Array
@@ -353,6 +365,7 @@ class DataWarp(PyTreeNode):
   contact__vert: jax.Array
   contact__worldid: jax.Array
   crb: jax.Array
+  dof_awake_ind: jax.Array
   dof_island: jax.Array
   dof_islandid: jax.Array
   efc__D: jax.Array
@@ -395,6 +408,7 @@ class DataWarp(PyTreeNode):
   iqfrc_smooth: jax.Array
   island_dofadr: jax.Array
   island_efcadr: jax.Array
+  island_idofadr: jax.Array
   island_ne: jax.Array
   island_nefc: jax.Array
   island_nf: jax.Array
@@ -411,6 +425,7 @@ class DataWarp(PyTreeNode):
   naccdmax: int
   nacon: jax.Array
   naconmax: int
+  nbody_awake: jax.Array
   ncollision: jax.Array
   ne: jax.Array
   nefc: jax.Array
@@ -421,6 +436,8 @@ class DataWarp(PyTreeNode):
   njmax_nnz: int
   njmax_pad: int
   nl: jax.Array
+  ntree_awake: jax.Array
+  nv_awake: jax.Array
   nworld: int
   qLD: jax.Array
   qLDiagInv: jax.Array
@@ -434,6 +451,8 @@ class DataWarp(PyTreeNode):
   ten_velocity: jax.Array
   ten_wrapadr: jax.Array
   ten_wrapnum: jax.Array
+  tree_asleep: jax.Array
+  tree_awake: jax.Array
   tree_island: jax.Array
   wrap_obj: jax.Array
   wrap_xpos: jax.Array
@@ -499,6 +518,8 @@ _NDIM = {
         'actuator_length': 2,
         'actuator_moment': 2,
         'actuator_velocity': 2,
+        'body_awake': 2,
+        'body_awake_ind': 2,
         'cacc': 3,
         'cam_xmat': 4,
         'cam_xpos': 3,
@@ -526,6 +547,7 @@ _NDIM = {
         'crb': 3,
         'ctrl': 2,
         'cvel': 3,
+        'dof_awake_ind': 2,
         'dof_island': 2,
         'dof_islandid': 2,
         'efc__D': 2,
@@ -572,6 +594,7 @@ _NDIM = {
         'iqfrc_smooth': 2,
         'island_dofadr': 2,
         'island_efcadr': 2,
+        'island_idofadr': 2,
         'island_ne': 2,
         'island_nefc': 2,
         'island_nf': 2,
@@ -590,6 +613,7 @@ _NDIM = {
         'naccdmax': 0,
         'nacon': 1,
         'naconmax': 0,
+        'nbody_awake': 1,
         'ncollision': 1,
         'ne': 1,
         'nefc': 1,
@@ -600,6 +624,8 @@ _NDIM = {
         'njmax_nnz': 0,
         'njmax_pad': 0,
         'nl': 1,
+        'ntree_awake': 1,
+        'nv_awake': 1,
         'nworld': 0,
         'qLD': 3,
         'qLDiagInv': 2,
@@ -633,6 +659,8 @@ _NDIM = {
         'ten_wrapadr': 2,
         'ten_wrapnum': 2,
         'time': 1,
+        'tree_asleep': 2,
+        'tree_awake': 2,
         'tree_island': 2,
         'wrap_obj': 3,
         'wrap_xpos': 3,
@@ -755,6 +783,7 @@ _NDIM = {
         'dof_frictionloss': 2,
         'dof_invweight0': 2,
         'dof_jntid': 1,
+        'dof_length': 1,
         'dof_parentid': 1,
         'dof_solimp': 3,
         'dof_solref': 3,
@@ -867,20 +896,29 @@ _NDIM = {
         'jnt_stiffnesspoly': 3,
         'jnt_type': 1,
         'light_active': 2,
+        'light_ambient': 3,
+        'light_attenuation': 3,
         'light_bodyid': 1,
         'light_castshadow': 2,
+        'light_cutoff': 2,
+        'light_diffuse': 3,
         'light_dir': 3,
         'light_dir0': 3,
+        'light_exponent': 2,
         'light_mode': 1,
         'light_pos': 3,
         'light_pos0': 3,
         'light_poscom0': 3,
+        'light_specular': 3,
         'light_targetbodyid': 1,
         'light_type': 2,
         'mapD2M': 1,
         'mapM2D': 1,
         'mapM2M': 1,
+        'mat_emission': 2,
         'mat_rgba': 3,
+        'mat_shininess': 2,
+        'mat_specular': 2,
         'mat_texid': 3,
         'mat_texrepeat': 3,
         'max_ten_J_rownnz': 0,
@@ -993,6 +1031,7 @@ _NDIM = {
         'opt__run_collision_detection': 0,
         'opt__sdf_initpoints': 0,
         'opt__sdf_iterations': 0,
+        'opt__sleep_tolerance': 1,
         'opt__solver': 0,
         'opt__timestep': 1,
         'opt__tolerance': 1,
@@ -1088,6 +1127,7 @@ _NDIM = {
         'tree_bodynum': 1,
         'tree_dofadr': 1,
         'tree_dofnum': 1,
+        'tree_sleep_policy': 1,
         'wrap_geom_adr': 1,
         'wrap_jnt_adr': 1,
         'wrap_objid': 1,
@@ -1118,6 +1158,7 @@ _NDIM = {
         'run_collision_detection': 0,
         'sdf_initpoints': 0,
         'sdf_iterations': 0,
+        'sleep_tolerance': 1,
         'solver': 0,
         'timestep': 1,
         'tolerance': 1,
@@ -1135,6 +1176,8 @@ _BATCH_DIM = {
         'actuator_length': True,
         'actuator_moment': True,
         'actuator_velocity': True,
+        'body_awake': True,
+        'body_awake_ind': True,
         'cacc': True,
         'cam_xmat': True,
         'cam_xpos': True,
@@ -1162,6 +1205,7 @@ _BATCH_DIM = {
         'crb': True,
         'ctrl': True,
         'cvel': True,
+        'dof_awake_ind': True,
         'dof_island': True,
         'dof_islandid': True,
         'efc__D': True,
@@ -1208,6 +1252,7 @@ _BATCH_DIM = {
         'iqfrc_smooth': True,
         'island_dofadr': True,
         'island_efcadr': True,
+        'island_idofadr': True,
         'island_ne': True,
         'island_nefc': True,
         'island_nf': True,
@@ -1226,6 +1271,7 @@ _BATCH_DIM = {
         'naccdmax': False,
         'nacon': False,
         'naconmax': False,
+        'nbody_awake': True,
         'ncollision': False,
         'ne': True,
         'nefc': True,
@@ -1236,6 +1282,8 @@ _BATCH_DIM = {
         'njmax_nnz': False,
         'njmax_pad': False,
         'nl': True,
+        'ntree_awake': True,
+        'nv_awake': True,
         'nworld': False,
         'qLD': True,
         'qLDiagInv': True,
@@ -1269,6 +1317,8 @@ _BATCH_DIM = {
         'ten_wrapadr': True,
         'ten_wrapnum': True,
         'time': True,
+        'tree_asleep': True,
+        'tree_awake': True,
         'tree_island': True,
         'wrap_obj': True,
         'wrap_xpos': True,
@@ -1391,6 +1441,7 @@ _BATCH_DIM = {
         'dof_frictionloss': True,
         'dof_invweight0': True,
         'dof_jntid': False,
+        'dof_length': False,
         'dof_parentid': False,
         'dof_solimp': True,
         'dof_solref': True,
@@ -1503,20 +1554,29 @@ _BATCH_DIM = {
         'jnt_stiffnesspoly': True,
         'jnt_type': False,
         'light_active': True,
+        'light_ambient': True,
+        'light_attenuation': True,
         'light_bodyid': False,
         'light_castshadow': True,
+        'light_cutoff': True,
+        'light_diffuse': True,
         'light_dir': True,
         'light_dir0': True,
+        'light_exponent': True,
         'light_mode': False,
         'light_pos': True,
         'light_pos0': True,
         'light_poscom0': True,
+        'light_specular': True,
         'light_targetbodyid': False,
         'light_type': True,
         'mapD2M': False,
         'mapM2D': False,
         'mapM2M': False,
+        'mat_emission': True,
         'mat_rgba': True,
+        'mat_shininess': True,
+        'mat_specular': True,
         'mat_texid': True,
         'mat_texrepeat': True,
         'max_ten_J_rownnz': False,
@@ -1629,6 +1689,7 @@ _BATCH_DIM = {
         'opt__run_collision_detection': False,
         'opt__sdf_initpoints': False,
         'opt__sdf_iterations': False,
+        'opt__sleep_tolerance': True,
         'opt__solver': False,
         'opt__timestep': True,
         'opt__tolerance': True,
@@ -1724,6 +1785,7 @@ _BATCH_DIM = {
         'tree_bodynum': False,
         'tree_dofadr': False,
         'tree_dofnum': False,
+        'tree_sleep_policy': False,
         'wrap_geom_adr': False,
         'wrap_jnt_adr': False,
         'wrap_objid': False,
@@ -1754,6 +1816,7 @@ _BATCH_DIM = {
         'run_collision_detection': False,
         'sdf_initpoints': False,
         'sdf_iterations': False,
+        'sleep_tolerance': True,
         'solver': False,
         'timestep': True,
         'tolerance': True,