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

Adds support for explicit dense/sparse mass matrices to MJX.

Browse Source 
This increases performance, particularly for the Newton solver on TPU.

PiperOrigin-RevId: 600696483
Change-Id: If69bb9a2e21ba8dad6ca23f093ce7b7ceae644ff
This commit is contained in:
Erik Frey
2024-01-23 00:31:01 -08:00
committed by Copybara-Service
parent aceb52bd09
commit 0a7be1732c
14 changed files with 324 additions and 143 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+10 -8
View File
@@ -26,21 +26,23 @@ MJX
5. Added :at:`site` transmission.
6. Updated MJX colab tutorial with more stable quadruped environment.
7. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, boxes, and meshes.
8. Added ``mjx.is_sparse`` which mirrors :ref:`mj_isSparse` and ``mjx.full_m`` which mirrors :ref:`mj_fullM`.
9. Added support for specifying sparse or dense mass matrices via :ref:`option-jacobian`.
Python bindings
^^^^^^^^^^^^^^^
8. Improved the implmentation of the :ref:`rollout<PySample>` module. Note the changes below are breaking, dependent
code will require modification.
10. Improved the implmentation of the :ref:`rollout<PySample>` module. Note the changes below are breaking, dependent
code will require modification.
- Uses :ref:`mjSTATE_FULLPHYSICS<geFullPhysics>` as state spec, enabling divergence detection by inspecting time.
- Allows user-defined control spec for any combination of :ref:`user input<geInput>` fields as controls.
- Outputs are no longer squeezed and always have dim=3.
- Uses :ref:`mjSTATE_FULLPHYSICS<geFullPhysics>` as state spec, enabling divergence detection by inspecting time.
- Allows user-defined control spec for any combination of :ref:`user input<geInput>` fields as controls.
- Outputs are no longer squeezed and always have dim=3.
Bug fixes
^^^^^^^^^
9. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes
:github:issue:`1270`.
10. Fixed a bug in the :ref:`muscle model<CMuscle>` that led to non-zero values outside the lower
11. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes
:github:issue:`1270`.
12. Fixed a bug in the :ref:`muscle model<CMuscle>` that led to non-zero values outside the lower
bound of the length range. Fixes :github:issue:`1342`.
doc/mjx.rst
+8
View File
@@ -349,3 +349,11 @@ For MJX to perform well, some configuration parameters should be adjusted from t
:ref:`option-flag` element
Disabling ``eulerdamp`` can help performance and is often not needed for stability.
:ref:`option-jacobian` element
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.