chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Small improvements to MJX performance tuning section.

Browse Source 
PiperOrigin-RevId: 612377295
Change-Id: I542aba7a33372556b62e3cdb4bf0c1c4d97dc98a
This commit is contained in:
Yuval Tassa
2024-03-04 02:48:20 -08:00
committed by Copybara-Service
parent 73f65ec408
commit 3d24bdb217
2 changed files with 20 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+1 -1
View File
@@ -37,7 +37,7 @@ MJX
Python bindings Python bindings
^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^
10. Fixed incorrect data types in the bindings for the ``geom``, ``vert``, ``elem``, and ``flex`` array members 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) Version 3.1.2 (February 05, 2024)
doc/mjx.rst
+19 -17
View File
@@ -126,7 +126,7 @@ Minimal example
.. code-block:: python .. code-block:: python
# Throw a ball at 100 different velocities. # Throw a ball at 100 different velocities.
import jax import jax
import mujoco import mujoco
@@ -335,33 +335,35 @@ Performance tuning
For MJX to perform well, some configuration parameters should be adjusted from their default MuJoCo values: For MJX to perform well, some configuration parameters should be adjusted from their default MuJoCo values:
:ref:`option` element :ref:`option/iterations<option-iterations>` and :ref:`option/ls_iterations<option-ls_iterations>`
The ``iterations`` and ``ls_iterations`` attributes---which control solver and linesearch iterations, respectively--- The :ref:`iterations<option-iterations>` and :ref:`ls_iterations<option-ls_iterations>` attributes---which control
should be brought down to just low enough that the simulation remains stable. Accurate solver forces are not so solver and linesearch iterations, respectively---should be brought down to just low enough that the simulation remains
important in reinforcement learning in which domain randomization is often used to add noise to physics for sim-to-real. stable. Accurate solver forces are not so important in reinforcement learning in which domain randomization is often
The ``NEWTON`` :ref:`Solver <mjtSolver>` often delivers reasonable convergence with one solver iteration, and performs used to add noise to physics for sim-to-real. The ``NEWTON`` :ref:`Solver <mjtSolver>` delivers excellent convergence
well on GPU. ``CG`` is currently a better choice for TPU. 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<contact-pair>`
Consider explicitly marking geoms for collision detection to reduce the number of contacts that MJX must consider 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 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 performance in MJX. Doing this well often requires an understanding of the task -- for example, the
`OpenAI Gym Humanoid <https://github.com/openai/gym/blob/master/gym/envs/mujoco/humanoid_v4.py>`__ task resets when `OpenAI Gym Humanoid <https://github.com/openai/gym/blob/master/gym/envs/mujoco/humanoid_v4.py>`__ task resets when
the humanoid starts to fall, so full contact with the floor is not needed. the humanoid starts to fall, so full contact with the floor is not needed.
:ref:`option-flag` element :ref:`option/flag/eulerdamp<option-flag-eulerdamp>`
Disabling ``eulerdamp`` can help performance and is often not needed for stability. Disabling ``eulerdamp`` can help performance and is often not needed for stability. Read the
:ref:`Numerical Integration<geIntegration>` section for details regarding the semantics of this flag.
:ref:`option-jacobian` element :ref:`option/jacobian<option-jacobian>`
Explicitly setting "dense" or "sparse" may speed up simulation depending on your device. Modern TPUs have specialized 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 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 they fit onto the device. As such, the behavior in MJX for the default "auto" setting is sparse if ``nv >= 60`` (60 or
greater, or if MJX detects a TPU as the default backend, otherwise "dense". For TPU, using "sparse" with the more degrees of freedom), or if MJX detects a TPU as the default backend, otherwise "dense". For TPU, using "sparse"
Newton solver can speed up simulation by 2x to 3x. For GPU, choosing "dense" may impart a more modest speedup of 10% with the Newton solver can speed up simulation by 2x to 3x. For GPU, choosing "dense" may impart a more modest speedup
to 20%, as long as the dense matrices can fit on the device. 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: The following environment variables should be set: