diff --git a/doc/mjwarp/index.rst b/doc/mjwarp/index.rst index 4b72f454..add25fb3 100644 --- a/doc/mjwarp/index.rst +++ b/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 `__, block Cholesky factorization and solve `#320 `__, constraint islands `#886 `__, and sleeping islands @@ -357,8 +357,97 @@ The :func:`mjw.make_data ` or :func:`mjw.put_data `; 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 + + + + - 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 ` is not available. - :ref:`mjDSBL_AUTORESET ` is not available. - :ref:`mjDSBL_NATIVECCD ` changes the default box-box collider from CCD to a primitive collider. -- :ref:`mjDSBL_ISLAND ` is not currently available. Constraint island discovery is tracked in GitHub issue - `#886 `__. :ref:`enableflags ` has the following differences: - :ref:`mjENBL_OVERRIDE ` is not available. - :ref:`mjENBL_FWDINV ` is not available. -- Constraint island sleeping enabled via :ref:`mjENBL_ISLAND ` is not currently available. This feature is - tracked in GitHub issues `#886 `__ and - `#887 `__. Additional MJWarp-only options are available: