From 3d24bdb2171eba447394c30e5f9d3b23bc73cc65 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 4 Mar 2024 02:48:20 -0800 Subject: [PATCH] Small improvements to MJX performance tuning section. PiperOrigin-RevId: 612377295 Change-Id: I542aba7a33372556b62e3cdb4bf0c1c4d97dc98a --- doc/changelog.rst | 2 +- doc/mjx.rst | 36 +++++++++++++++++++----------------- 2 files changed, 20 insertions(+), 18 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 54230557..2eb1e5da 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -37,7 +37,7 @@ MJX Python bindings ^^^^^^^^^^^^^^^ 10. Fixed incorrect data types in the bindings for the ``geom``, ``vert``, ``elem``, and ``flex`` array members - of the ``mjContact`` struct, and all array members of the ``mjrContext`` struct. + of the ``mjContact`` struct, and all array members of the ``mjrContext`` struct. Version 3.1.2 (February 05, 2024) diff --git a/doc/mjx.rst b/doc/mjx.rst index 28fb312d..af935449 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -126,7 +126,7 @@ Minimal example .. code-block:: python - # Throw a ball at 100 different velocities. + # Throw a ball at 100 different velocities. import jax import mujoco @@ -335,33 +335,35 @@ Performance tuning For MJX to perform well, some configuration parameters should be adjusted from their default MuJoCo values: -:ref:`option` element - The ``iterations`` and ``ls_iterations`` attributes---which control solver and linesearch iterations, respectively--- - should be brought down to just low enough that the simulation remains stable. Accurate solver forces are not so - important in reinforcement learning in which domain randomization is often used to add noise to physics for sim-to-real. - The ``NEWTON`` :ref:`Solver ` often delivers reasonable convergence with one solver iteration, and performs - well on GPU. ``CG`` is currently a better choice for TPU. +:ref:`option/iterations` and :ref:`option/ls_iterations` + The :ref:`iterations` and :ref:`ls_iterations` attributes---which control + solver and linesearch iterations, respectively---should be brought down to just low enough that the simulation remains + stable. Accurate solver forces are not so important in reinforcement learning in which domain randomization is often + used to add noise to physics for sim-to-real. The ``NEWTON`` :ref:`Solver ` delivers excellent convergence + with very few (often just one) solver iterations, and performs well on GPU. ``CG`` is currently a better choice for + TPU. -:ref:`contact-pair` element +:ref:`contact/pair` Consider explicitly marking geoms for collision detection to reduce the number of contacts that MJX must consider during each step. Enabling only an explicit list of valid contacts can have a dramatic effect on simulation performance in MJX. Doing this well often requires an understanding of the task -- for example, the `OpenAI Gym Humanoid `__ task resets when the humanoid starts to fall, so full contact with the floor is not needed. -:ref:`option-flag` element - Disabling ``eulerdamp`` can help performance and is often not needed for stability. +:ref:`option/flag/eulerdamp` + Disabling ``eulerdamp`` can help performance and is often not needed for stability. Read the + :ref:`Numerical Integration` section for details regarding the semantics of this flag. -:ref:`option-jacobian` element +:ref:`option/jacobian` Explicitly setting "dense" or "sparse" may speed up simulation depending on your device. Modern TPUs have specialized hardware for rapidly operating over sparse matrices, whereas GPUs tend to be faster with dense matrices as long as - they fit onto the device. As such, the behavior in MJX for the default "auto" setting is sparse if ``nv`` is 60 or - greater, or if MJX detects a TPU as the default backend, otherwise "dense". For TPU, using "sparse" with the - Newton solver can speed up simulation by 2x to 3x. For GPU, choosing "dense" may impart a more modest speedup of 10% - to 20%, as long as the dense matrices can fit on the device. + they fit onto the device. As such, the behavior in MJX for the default "auto" setting is sparse if ``nv >= 60`` (60 or + more degrees of freedom), or if MJX detects a TPU as the default backend, otherwise "dense". For TPU, using "sparse" + with the Newton solver can speed up simulation by 2x to 3x. For GPU, choosing "dense" may impart a more modest speedup + of 10% to 20%, as long as the dense matrices can fit on the device. -GPU performance tuning ----------------------- +GPU performance +--------------- The following environment variables should be set: