Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 927227210
Change-Id: Ic72c58524c8e91e6fff44b4cab7f4a0f7de2bb33

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
+@cache_kernel
-def _zero_nacon_ncollision(
+def _zero_nacon_ncollision(enable_sleep: bool = False):
-  # Data out:
+  @wp.kernel(module="unique", enable_backward=False)
-  nacon_out: wp.array[int],
+  def zero_nacon_ncollision(
-  ncollision_out: wp.array[int],
+    # In:
-):
+    skip_in: wp.array[int],
-  ncollision_out[0] = 0
+    # Data out:
-  nacon_out[0] = 0
+    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
+@cache_kernel
-def _sap_range(
+def _sap_range(enable_sleep: bool = False):
-  # Model:
+  @wp.kernel(module="unique", enable_backward=False)
-  ngeom: int,
+  def sap_range(
-  # In:
+    # Model:
-  projection_lower_in: wp.array2d[float],
+    ngeom: int,
-  projection_upper_in: wp.array2d[float],
+    # In:
-  sort_index_in: wp.array2d[int],
+    projection_lower_in: wp.array2d[float],
-  # Out:
+    projection_upper_in: wp.array2d[float],
-  range_out: wp.array2d[int],
+    sort_index_in: wp.array2d[int],
-):
+    skip_in: wp.array[int],
-  worldid, geomid = wp.tid()
+    # Out:
+    range_out: wp.array2d[int],
+  ):
+    worldid, geomid = wp.tid()
 
-  # current bounding geom
+    if wp.static(enable_sleep):
-  idx = sort_index_in[worldid, geomid]
+      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)
+    upper = projection_upper_in[worldid, idx]
-  limit = wp.min(ngeom - 1, limit)
 
-  # range of geoms for the sweep and prune process
+    limit = _binary_search(projection_lower_in[worldid], upper, geomid + 1, ngeom)
-  range_out[worldid, geomid] = limit - geomid
+    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)
 

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
+@cache_kernel
-def _qfrc_smooth(
+def _qfrc_smooth(enable_sleep: bool):
-  # Data in:
+  @wp.kernel(module="unique", enable_backward=False)
-  qfrc_applied_in: wp.array2d[float],
+  def kernel(
-  qfrc_bias_in: wp.array2d[float],
+    # Model:
-  qfrc_passive_in: wp.array2d[float],
+    body_treeid: wp.array[int],
-  qfrc_actuator_in: wp.array2d[float],
+    dof_bodyid: wp.array[int],
-  # Data out:
+    # Data in:
-  qfrc_smooth_out: wp.array2d[float],
+    qfrc_applied_in: wp.array2d[float],
-):
+    tree_awake_in: wp.array2d[int],
-  worldid, dofid = wp.tid()
+    qfrc_bias_in: wp.array2d[float],
-  qfrc_smooth_out[worldid, dofid] = (
+    qfrc_passive_in: wp.array2d[float],
-    qfrc_passive_in[worldid, dofid]
+    qfrc_actuator_in: wp.array2d[float],
-    - qfrc_bias_in[worldid, dofid]
+    # Data out:
-    + qfrc_actuator_in[worldid, dofid]
+    qfrc_smooth_out: wp.array2d[float],
-    + qfrc_applied_in[worldid, dofid]
+  ):
-  )
+    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)

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
+    use_ambient_lighting: Top-level ambient switch. When False, skips all
-      lighting term before applying model lights.
+                          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)

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],
+    inputs=[m.nv, d.dof_island, d.island_idofadr, d.nidof],
-    outputs=[d.island_nv, d.map_dof2idof, d.map_idof2dof, d.dof_islandid, unconstrained_cnt],
+    outputs=[d.island_nv, d.island_dofadr, d.map_dof2idof, d.map_idof2dof, d.dof_islandid, unconstrained_cnt],
   )
 
   # 5. Map Constraints

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:
+    enable_specular: bool,
-    light_contribution = float(0.0)
+    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 not default_attenuation:
-      if lighttype == 0:  # spot light
+        light_attenuation_factor = wp.vec3(1.0, dist_to_light, dist_to_light * dist_to_light)
-        spot_dir = wp.normalize(lightdir)
+        attenuation = math.safe_div(1.0, wp.dot(light_attenuation_factor, lightattenuation))
-        cos_theta = wp.dot(-L, spot_dir)
+      if has_spot:
-        spot_factor = wp.min(1.0, wp.max(0.0, (cos_theta - 0.85) * 10.0))
+        if lighttype == 0:  # spot light
-        attenuation = attenuation * spot_factor
+          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)
+    mat_spec = DEFAULT_MAT_SPECULAR
-    if wp.static(rc.use_ambient_lighting):
+    mat_shin_exp = DEFAULT_MAT_SHININESS_EXPONENT
-      len_n = wp.length(normal)
+    mat_emis = DEFAULT_MAT_EMISSION
-      n = normal if len_n > 0.0 else wp.vec3(0.0, 0.0, 1.0)
+    if wp.static(rc.enable_specular or rc.enable_emission):
-      n = wp.normalize(n)
+      if geom_id != -2:
-      hemispheric = 0.5 * (n[2] + 1.0)
+        mat_id_for_spec = geom_matid[worldid % geom_matid.shape[0], geom_id]
-      ambient_color = wp.vec3(0.4, 0.4, 0.45) * hemispheric + wp.vec3(0.1, 0.1, 0.12) * (1.0 - hemispheric)
+        if mat_id_for_spec >= 0:
-      result = 0.5 * wp.cw_mul(base_color, ambient_color)
+          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
+    result = wp.vec3(0.0)
-    for l in range(wp.static(m.nlight)):
+    if wp.static(rc.enable_emission):
-      light_contribution = compute_lighting(
+      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_active_worldid[light_index],
-        light_type[worldid % light_type.shape[0], l],
+        light_type_worldid[light_index],
-        light_castshadow[worldid % light_castshadow.shape[0], l],
+        light_castshadow_worldid[light_index],
-        light_xpos_in[worldid, l],
+        light_xpos_in_worldid[light_index],
-        light_xdir_in[worldid, l],
+        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,

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
+    use_ambient_lighting: top-level switch for ambient contributions
-    background_color: background color
+    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 = ()

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

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,

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,

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}
 )

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,