Add MuJoCo Warp sleeping support (#3731)

Co-authored-by: Eric Heiden <eric-heiden@outlook.com>

Author: vreutskyy

CHANGELOG.md

@@ -35,6 +35,7 @@
 - Add opt-in filtering of static-static, static-kinematic, and kinematic-kinematic contacts during broad-phase collision detection. Set `CollisionPipeline(include_static_kinematic_pairs=False)` to enable filtering; the default preserves existing contact generation. `Model.shape_contact_pairs` remains an unfiltered superset for direct consumers such as `SolverKamino` and hydroelastic SDF setup.
 - Add opt-in `body_frame_origin="com"` to `ModelBuilder.add_rod()` and `ModelBuilder.add_rod_graph()` for COM-centered cable capsule body frames.
 - Add `sign_method` argument to `Mesh.build_sdf` and `SDF.create_from_mesh` support for a `"normal"` (angle-weighted pseudo-normal) sign strategy, for selecting the inside/outside sign of the baked SDF (`"auto"`, `"parity"`, `"winding"`, or `"normal"`).
+- Add opt-in MuJoCo Warp sleeping support to `SolverMuJoCo` with MJCF sleep configuration, initial tree policies, `sleep_tolerance`, compact `nvmax` storage, and a launchable `mujoco_sleeping` example. (#3725)
 - Add `forward_depth_image` output support to `SensorTiledCamera.update()` and `SensorTiledCamera.utils.create_forward_depth_image_output()` for native forward-depth rendering without post-processing `depth_image`.
 - Add the `basic_conveyor_forces` example: a multi-belt conveyor circuit that transports rigid boxes with per-belt velocity fields, applying Coulomb-limited tangential body forces from reported per-contact normal forces across `SolverXPBD`, `SolverVBD`, and `SolverMuJoCo`.
 - Add optional `shear_stiffness`/`shear_damping` and `twist_stiffness`/`twist_damping` controls to `ModelBuilder.add_joint_cable()`, `ModelBuilder.add_rod()`, and `ModelBuilder.add_rod_graph()`; omitted shear defaults to stretch and omitted twist defaults to bend for compatibility.
@@ -55,6 +56,7 @@
 - Compile tiled camera render kernels with CUDA fast math by default for faster rendering; set `SensorTiledCamera.render_config.enable_fast_math = False` for bit-exact, IEEE-precise output.
 - Disable `HydroelasticSDF.Config.pre_prune_contacts` when `CollisionPipeline(deterministic=True)` (implied by any `contact_matching` mode other than `"disabled"`). Pre-pruning ranks faces per thread, so it cannot be made reproducible. The face-contact buffer doubles because `contact_buffer_fraction` no longer applies, and the generated contact set differs from the non-deterministic default. Deterministic hydroelastic contacts also cap the buffer at 2^20 faces; lower `buffer_mult_contact` or `buffer_fraction` if construction now raises.
 - Make `CollisionPipeline` the sole owner of rigid-contact geometry for `SolverVBD`: `"latest"` supplies fresh geometry and `"sticky"` supplies replayed geometry. `SolverVBD(rigid_contact_history=True)` uses either mode's match indices only to warm-start its numeric lambda/penalty state.
+- Upgrade `mujoco` and `mujoco-warp` to 3.11.0. (#3725)
 - Optimize raycast/raytrace queries by restructuring ray-shape intersection into local-space primitives and compile specialized depth/shadow variants that skip unused surface-normal work (mesh shadows also use any-hit queries).
 - Change experimental `SolverVBD` cable constraint slots from `[STRETCH=0, BEND=1]` to `[STRETCH=0, SHEAR=1, BEND=2, TWIST=3]`, allowing each stiffness and constraint mode to be configured independently. Existing cable calls using raw `slot=1` or `JointSlot.ANGULAR` now select shear; use `JointSlot.BEND` (now slot 2) to select bending.
 - Map `shape_material_kf` to per-contact MuJoCo `solreffriction` in `SolverMuJoCo` (elliptic friction cones with Newton contacts); resolve `kf` with priority/`solmix`, treat a resolved `kf = 0` as frictionless, and use native MuJoCo contacts or a pyramidal cone to preserve the previous solref-inherited friction.

asv.conf.json

@@ -16,7 +16,7 @@
   "build_command": ["python -m build --wheel -o {build_cache_dir} {build_dir}"],
   "install_command": [
     "in-dir={env_dir} python -m pip install {wheel_file}[dev] --extra-index-url=https://pypi.nvidia.com/",
-    "in-dir={env_dir} python -m pip install {wheel_file}[sim] warp-lang==1.16.0.dev20260716 mujoco==3.10.0 mujoco-warp==3.10.0.3 --extra-index-url=https://pypi.nvidia.com/",
+    "in-dir={env_dir} python -m pip install warp-lang==1.16.0.dev20260716 mujoco==3.11.0 mujoco-warp==3.11.0 --extra-index-url=https://pypi.nvidia.com/",
     "python -m pip install torch==2.10.0+cu130 --index-url https://download.pytorch.org/whl/cu130"
   ]
 }

docs/solvers/mujoco.rst

@@ -592,6 +592,89 @@ See MuJoCo's `solver documentation
 what each parameter does and when to tune it.
 
 
+Sleeping
+--------
+
+MuJoCo Warp can exclude stationary constraint islands from collision
+detection and the compact Newton solver. Newton keeps this optimization
+disabled by default. Enable it when simulating many passive objects that
+spend substantial time at rest. Sleeping adds island bookkeeping and uses the
+compact-solver path even while everything remains awake, so scenes in which
+most degrees of freedom remain awake can be slower than with sleeping disabled.
+Benchmark representative workloads before enabling it::
+
+    solver = SolverMuJoCo(
+        model,
+        enable_sleeping=True,
+        sleep_tolerance=0.01,
+    )
+
+Sleeping is only supported by Newton's MuJoCo Warp GPU path. It requires
+``solver="newton"`` and ``use_mujoco_contacts=True`` so MuJoCo Warp's
+collision pipeline can wake sleeping bodies, and it does not support the RK4
+integrator. Unsupported combinations raise ``ValueError`` during solver
+construction. When an imported MJCF contains
+``<option><flag sleep="enable"/></option>``, leaving ``enable_sleeping`` as
+``None`` honors that setting. An explicit constructor value takes precedence.
+
+``sleep_tolerance`` is the scaled generalized-velocity threshold [m/s] below
+which an island becomes eligible to sleep. It follows the normal
+`Solver options`_ resolution order, including per-world
+``model.mujoco.sleep_tolerance`` values. The default is ``0.001``.
+
+MuJoCo's per-tree ``body/sleep`` policy is available as the body-frequency
+``model.mujoco.sleep_policy`` custom attribute and is imported from MJCF.
+The supported values are :attr:`~newton.solvers.SolverMuJoCo.SleepPolicy.AUTO`,
+:attr:`~newton.solvers.SolverMuJoCo.SleepPolicy.NEVER`,
+:attr:`~newton.solvers.SolverMuJoCo.SleepPolicy.ALLOWED`, and
+:attr:`~newton.solvers.SolverMuJoCo.SleepPolicy.INIT`. The policy must be
+assigned to the moving root body of a MuJoCo kinematic tree. For a
+programmatically built model, register the MuJoCo custom attributes before
+adding the body::
+
+    builder = newton.ModelBuilder()
+    SolverMuJoCo.register_custom_attributes(builder)
+    root = builder.add_link(
+        custom_attributes={
+            "mujoco:sleep_policy": SolverMuJoCo.SleepPolicy.INIT,
+        },
+    )
+
+``INIT`` makes the tree start asleep and is the intended way to construct a
+solver with compact ``nvmax`` storage. The setting applies to the shared model,
+so replicated worlds receive the same initial tree policies. Their default
+joint velocities must also match because MuJoCo Warp uses one shared initial
+sleep state; solver construction rejects differing per-world defaults.
+
+``nvmax`` sizes the compact solver's active-DOF workspace per world. If it is
+``None``, MuJoCo Warp uses the model's full ``nv``. This is the safe starting
+point but does not reduce compact-workspace memory. Tune a smaller value to
+the maximum number of DOFs expected to be awake simultaneously. The capacity
+must include every DOF that is awake in the initial MuJoCo state; Newton rejects
+smaller values before allocating solver data. ``nvmax`` is rejected when
+sleeping is disabled because MuJoCo Warp would otherwise allocate unused
+compact-solver workspace. If more than ``nvmax`` DOFs become active later,
+MuJoCo Warp sets the ``NVMAX`` bit in
+``solver.mjw_data.overflow``; increase ``nvmax`` and recreate the solver.
+
+MuJoCo Warp automatically wakes islands for applied forces and contacts.
+Actuated trees are not allowed to sleep by default. If their sleep policy is
+explicitly changed to allow sleeping, changing actuator controls alone does not
+wake them; apply a force or set a nonzero velocity first. Newton-side
+joint-position edits wake only the affected sleeping trees.
+:meth:`~newton.solvers.SolverMuJoCo.reset` restores the initial sleep state in
+selected worlds, while
+:meth:`~newton.solvers.SolverMuJoCo.notify_model_changed` wakes all worlds
+after model-property updates. The sleeping path supports whole-step CUDA graph
+capture.
+
+See MuJoCo's `sleeping-islands documentation
+<https://mujoco.readthedocs.io/en/stable/programming/simulation.html#sleeping-islands>`_
+and `MuJoCo Warp compact-solver guide
+<https://mujoco.readthedocs.io/en/latest/mjwarp/#compact-solver>`_ for the
+underlying sleep policy, wake triggers, and ``nvmax`` tuning guidance.
+
+
 .. _mujoco-custom-attributes:
 
 MuJoCo-specific parameters in USD and MJCF

newton/_src/solvers/mujoco/kernels.py

@@ -1953,6 +1953,7 @@ def update_solver_options_kernel(
     newton_tolerance: wp.array[float],
     newton_ls_tolerance: wp.array[float],
     newton_ccd_tolerance: wp.array[float],
+    newton_sleep_tolerance: wp.array[float],
     newton_density: wp.array[float],
     newton_viscosity: wp.array[float],
     newton_wind: wp.array[wp.vec3],
@@ -1962,6 +1963,7 @@ def update_solver_options_kernel(
     opt_tolerance: wp.array[float],
     opt_ls_tolerance: wp.array[float],
     opt_ccd_tolerance: wp.array[float],
+    opt_sleep_tolerance: wp.array[float],
     opt_density: wp.array[float],
     opt_viscosity: wp.array[float],
     opt_wind: wp.array[wp.vec3],
@@ -1974,6 +1976,7 @@ def update_solver_options_kernel(
         newton_tolerance: Per-world tolerance values (None if overridden)
         newton_ls_tolerance: Per-world line search tolerance values (None if overridden)
         newton_ccd_tolerance: Per-world CCD tolerance values (None if overridden)
+        newton_sleep_tolerance: Per-world sleep tolerance values (None if overridden)
         newton_density: Per-world medium density values (None if overridden)
         newton_viscosity: Per-world medium viscosity values (None if overridden)
         newton_wind: Per-world wind velocity vectors (None if overridden)
@@ -1982,6 +1985,7 @@ def update_solver_options_kernel(
         opt_tolerance: MuJoCo Warp opt.tolerance array (shape: nworld)
         opt_ls_tolerance: MuJoCo Warp opt.ls_tolerance array (shape: nworld)
         opt_ccd_tolerance: MuJoCo Warp opt.ccd_tolerance array (shape: nworld)
+        opt_sleep_tolerance: MuJoCo Warp opt.sleep_tolerance array (shape: nworld)
         opt_density: MuJoCo Warp opt.density array (shape: nworld)
         opt_viscosity: MuJoCo Warp opt.viscosity array (shape: nworld)
         opt_wind: MuJoCo Warp opt.wind array (shape: nworld)
@@ -2009,6 +2013,9 @@ def update_solver_options_kernel(
     if newton_ccd_tolerance:
         opt_ccd_tolerance[worldid] = newton_ccd_tolerance[worldid]
 
+    if newton_sleep_tolerance:
+        opt_sleep_tolerance[worldid] = newton_sleep_tolerance[worldid]
+
     if newton_density:
         opt_density[worldid] = newton_density[worldid]
 
@@ -3183,6 +3190,159 @@ def reset_world_buffers_kernel(
         xfrc_applied[worldid, i] = wp.spatial_vector(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
 
 
+@wp.kernel(enable_backward=False)
+def reset_sleeping_state_kernel(
+    world_mask: wp.array[wp.bool],
+    clear_overflow: int,
+    nv: int,
+    nbody: int,
+    ntree: int,
+    awake_value: int,
+    sleep_state_static: int,
+    sleep_state_awake: int,
+    body_rootid: wp.array[wp.int32],
+    body_mocapid: wp.array[wp.int32],
+    body_treeid: wp.array[wp.int32],
+    tree_asleep: wp.array2d[wp.int32],
+    tree_awake: wp.array2d[wp.int32],
+    body_awake: wp.array2d[wp.int32],
+    body_awake_ind: wp.array2d[wp.int32],
+    dof_awake_ind: wp.array2d[wp.int32],
+    ntree_awake: wp.array[wp.int32],
+    nbody_awake: wp.array[wp.int32],
+    nv_awake: wp.array[wp.int32],
+    overflow: wp.array[wp.int32],
+):
+    """Wake every tree and rebuild sleep bookkeeping in selected worlds."""
+    worldid, elemid = wp.tid()
+    if world_mask and not world_mask[worldid]:
+        return
+
+    if elemid < ntree:
+        tree_asleep[worldid, elemid] = awake_value
+        tree_awake[worldid, elemid] = 1
+
+    if elemid < nbody:
+        if body_treeid[elemid] < 0:
+            rootid = body_rootid[elemid]
+            if body_mocapid[rootid] < 0:
+                body_awake[worldid, elemid] = sleep_state_static
+            else:
+                body_awake[worldid, elemid] = sleep_state_awake
+        else:
+            body_awake[worldid, elemid] = sleep_state_awake
+        body_awake_ind[worldid, elemid] = elemid
+
+    if elemid < nv:
+        dof_awake_ind[worldid, elemid] = elemid
+
+    if elemid == 0:
+        ntree_awake[worldid] = ntree
+        nbody_awake[worldid] = nbody
+        nv_awake[worldid] = nv
+        if clear_overflow:
+            overflow[worldid] = 0
+
+
+@wp.kernel(enable_backward=False)
+def restore_sleeping_state_kernel(
+    world_mask: wp.array[wp.bool],
+    clear_overflow: int,
+    nv: int,
+    nbody: int,
+    ntree: int,
+    initial_ntree_awake: int,
+    initial_nbody_awake: int,
+    initial_nv_awake: int,
+    initial_tree_asleep: wp.array[wp.int32],
+    initial_tree_awake: wp.array[wp.int32],
+    initial_body_awake: wp.array[wp.int32],
+    initial_body_awake_ind: wp.array[wp.int32],
+    initial_dof_awake_ind: wp.array[wp.int32],
+    tree_asleep: wp.array2d[wp.int32],
+    tree_awake: wp.array2d[wp.int32],
+    body_awake: wp.array2d[wp.int32],
+    body_awake_ind: wp.array2d[wp.int32],
+    dof_awake_ind: wp.array2d[wp.int32],
+    ntree_awake: wp.array[wp.int32],
+    nbody_awake: wp.array[wp.int32],
+    nv_awake: wp.array[wp.int32],
+    overflow: wp.array[wp.int32],
+):
+    """Restore the initial sleep bookkeeping in selected worlds."""
+    worldid, elemid = wp.tid()
+    if world_mask and not world_mask[worldid]:
+        return
+
+    if elemid < ntree:
+        tree_asleep[worldid, elemid] = initial_tree_asleep[elemid]
+        tree_awake[worldid, elemid] = initial_tree_awake[elemid]
+
+    if elemid < nbody:
+        body_awake[worldid, elemid] = initial_body_awake[elemid]
+        body_awake_ind[worldid, elemid] = initial_body_awake_ind[elemid]
+
+    if elemid < nv:
+        dof_awake_ind[worldid, elemid] = initial_dof_awake_ind[elemid]
+
+    if elemid == 0:
+        ntree_awake[worldid] = initial_ntree_awake
+        nbody_awake[worldid] = initial_nbody_awake
+        nv_awake[worldid] = initial_nv_awake
+        if clear_overflow:
+            overflow[worldid] = 0
+
+
+@wp.kernel(enable_backward=False)
+def copy_qpos_and_detect_tree_change_kernel(
+    qpos_new: wp.array2d[wp.float32],
+    tolerance: float,
+    qpos_treeid: wp.array[wp.int32],
+    qpos: wp.array2d[wp.float32],
+    tree_changed: wp.array2d[wp.int32],
+):
+    """Copy converted coordinates and flag trees with external pose edits."""
+    worldid, i = wp.tid()
+    value = qpos_new[worldid, i]
+    if wp.abs(value - qpos[worldid, i]) > tolerance:
+        treeid = qpos_treeid[i]
+        if treeid >= 0:
+            wp.atomic_max(tree_changed, worldid, treeid, 1)
+    qpos[worldid, i] = value
+
+
+@wp.kernel(enable_backward=False)
+def wake_changed_trees_kernel(
+    tree_changed: wp.array2d[wp.int32],
+    ntree: int,
+    awake_value: int,
+    tree_asleep: wp.array2d[wp.int32],
+):
+    """Wake edited trees and every tree in their sleeping-island cycles."""
+    # Negative values mean awake; non-negative values link the next tree in a
+    # sleeping-island cycle. The O(ntree) scan intentionally uses one walker
+    # per world because parallel walkers for overlapping edits could overwrite
+    # a link before a peer reads it.
+    worldid = wp.tid()
+    for treeid in range(ntree):
+        if tree_changed[worldid, treeid] == 0:
+            continue
+
+        asleep_value = tree_asleep[worldid, treeid]
+        if asleep_value < 0:
+            if awake_value < asleep_value:
+                tree_asleep[worldid, treeid] = awake_value
+            continue
+
+        current = treeid
+        for _step in range(ntree + 1):
+            next_tree = tree_asleep[worldid, current]
+            tree_asleep[worldid, current] = awake_value
+            current = next_tree
+            if current == treeid:
+                break
+
+
 @wp.kernel(enable_backward=False)
 def reset_joint_state_kernel(
     world_mask: wp.array[wp.bool],