MuJoCo Warp documentation: Large scenes

PiperOrigin-RevId: 940590426
Change-Id: I006e245b54cd4165928927bee76f44d730bdd895

doc/mjwarp/index.rst

@@ -64,8 +64,8 @@ Complex scenes
 --------------
 
 MJWarp scales better than MJX for scenes with many geoms or degrees of freedom, but not as well as MuJoCo. There may be
-significant performance degradation in MJWarp for scenes beyond 60 DoFs. Supporting these larger scenes is a high
-priority and progress is tracked in GitHub issues for: sparse Jacobians
+significant performance degradation in MJWarp for scenes beyond 60 degrees of freedom (DoFs). Supporting these larger scenes
+is a high priority and progress is tracked in GitHub issues for: sparse Jacobians
 `#88 <https://github.com/google-deepmind/mujoco_warp/issues/88>`__, block Cholesky factorization and solve
 `#320 <https://github.com/google-deepmind/mujoco_warp/issues/320>`__, constraint islands
 `#886 <https://github.com/google-deepmind/mujoco_warp/issues/886>`__, and sleeping islands
@@ -357,8 +357,97 @@ The :func:`mjw.make_data <mujoco_warp.make_data>` or :func:`mjw.put_data <mujoco
 requirements for CCD. The value for this parameter should be the maximum number of contacts generated by a CCD
 collider, per world or for all worlds, respectively. For example, a batched simulation with 10 worlds that generates 80
 total contacts with per-collider contacts: mesh-mesh: 30 (CCD), ellipsoid-ellipsoid: 10 (CCD), and sphere-sphere: 40
-(primitive) should set `nconmax`_ / `naconmax`_ to at least 8 / 80 (may require more for broadphase) and ``nccdmax`` /
-``naccdmax`` to 3 / 30.
+(primitive) should set `nconmax`_ / `naconmax`_ to at least 8 / 80 (may require more for broadphase) and
+``nccdmax`` / ``naccdmax`` to 3 / 30.
+
+Large scenes
+------------
+
+Simulating scenes with many DoFs (i.e., `nv`) can be computationally
+expensive. However, in many scenarios, a significant portion of the scene may
+be stationary. MJWarp can put stationary objects to *sleep* (see
+:ref:`Sleeping`), excluding them from the working set of many of its
+calculations. Furthermore, MJWarp groups bodies into independent
+:ref:`islands <soIsland>`; if all bodies in an island are stationary, the
+entire island is put to sleep. Currently, both the collision pipeline and the
+constraint solver benefit from sleeping, and more sleeping-aware components
+may be added in the future.
+
+Compact solver
+~~~~~~~~~~~~~~
+
+To optimize performance in scenes with many total DoFs but a relatively
+small number of active DoFs (typically fewer than 64, such as two robot
+arms with grippers (16 DoFs) and 8 active objects (48 DoFs)), MJWarp
+provides a **compact solver** that leverages this sleeping mechanism:
+
+1. Identifies the set of active DOFs for each world, determined from the active islands.
+2. **Compacts** these active DOFs into a single, contiguous dense workspace of a known maximum size (``nvmax``).
+3. Executes the constraint solver (Newton) using GPU-optimized tile operations (such as blocked Cholesky
+   factorization) of fixed size on this compacted space.
+4. Scatters the results back to the global state, freezing the inactive DOFs.
+
+By using a fixed-size compacted workspace, the solver avoids GPU thread divergence and leverages
+high-performance tensor/matrix operations optimized for fixed tile sizes.
+
+.. rubric:: Enabling the compact solver
+
+1. Enable the Newton solver:
+
+   - Via XML:
+
+     .. code-block:: xml
+
+        <option solver="Newton"/>
+
+   - Via Python ``MjSpec``:
+
+     .. code-block:: python
+
+        spec = mujoco.MjSpec()
+        spec.option.solver = mujoco.mjtSolver.mjSOL_NEWTON
+
+2. Enable sleep:
+
+   - Via XML:
+
+     .. code-block:: xml
+
+        <option>
+          <flag sleep="enable"/>
+        </option>
+
+   - Via Python ``MjSpec``:
+
+     .. code-block:: python
+
+        spec = mujoco.MjSpec()
+        spec.option.enableflags |= mujoco.mjtEnableBit.mjENBL_SLEEP
+
+3. Specify the maximum expected active DOFs for any world (``nvmax``) when allocating data. This sizes the
+   compacted workspace.
+
+   - In Python:
+
+     .. code-block:: python
+
+        # Allocate data with a maximum of 64 active DOFs per world
+        d = mjw.make_data(mjm, nworld=2048, nvmax=64)
+
+   - Via the command line:
+
+     .. code-block:: shell
+
+        mjwarp-testspeed scene.xml --nvmax=64
+
+   If ``nvmax`` is not specified, it defaults to the full number of DOFs (``nv``). Sizing ``nvmax``
+   to a tight upper bound of the expected active DOFs significantly reduces GPU memory usage and improves
+   throughput.
+
+.. note::
+   Consider increasing the sleep tolerance setting (e.g., ``sleep_tolerance="0.01"`` in XML options or
+   ``spec.option.sleep_tolerance = 0.01`` in Python) from its default value (0.0001) to more quickly
+   sleep objects.
 
 .. _mjwBatch:
 
@@ -1020,16 +1109,11 @@ exceptions:
 - :ref:`mjDSBL_MIDPHASE <mjtDisablebit>` is not available.
 - :ref:`mjDSBL_AUTORESET <mjtDisablebit>` is not available.
 - :ref:`mjDSBL_NATIVECCD <mjtDisablebit>` changes the default box-box collider from CCD to a primitive collider.
-- :ref:`mjDSBL_ISLAND <mjtDisablebit>` is not currently available. Constraint island discovery is tracked in GitHub issue
-  `#886 <https://github.com/google-deepmind/mujoco_warp/issues/886>`__.
 
 :ref:`enableflags <option-flag>` has the following differences:
 
 - :ref:`mjENBL_OVERRIDE <mjtEnablebit>` is not available.
 - :ref:`mjENBL_FWDINV <mjtEnablebit>` is not available.
-- Constraint island sleeping enabled via :ref:`mjENBL_ISLAND <mjtEnablebit>` is not currently available. This feature is
-  tracked in GitHub issues `#886 <https://github.com/google-deepmind/mujoco_warp/issues/886>`__ and
-  `#887 <https://github.com/google-deepmind/mujoco_warp/issues/887>`__.
 
 Additional MJWarp-only options are available: