Add MJX and bump version number to 3.0.0.
Co-authored-by: Baruch Tabanpour <btaba@google.com> PiperOrigin-RevId: 574327508 Change-Id: Ia9b62fbc929c6869dfcec87636b2e10d405a1060
This commit is contained in:
Erik Frey
committed by
Saran Tunyasuvunakool
Saran Tunyasuvunakool
parent
3f3d5a3b49
commit
8f9c690c85
@@ -522,7 +522,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr
|
||||
- Maximum number of UI rectangles.
|
||||
Defined in `mjui.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjui.h>`_.
|
||||
* - ``mjVERSION_HEADER``
|
||||
- 238
|
||||
- 300
|
||||
- The version of the MuJoCo headers; changes with every release. This is an integer equal to 100x the software
|
||||
version, so 210 corresponds to version 2.1. Defined in mujoco.h. The API function :ref:`mj_version` returns a
|
||||
number with the same meaning but for the compiled library.
|
||||
|
||||
+56
-45
@@ -2,17 +2,28 @@
|
||||
Changelog
|
||||
=========
|
||||
|
||||
Upcoming version (not yet released)
|
||||
-----------------------------------
|
||||
Version 3.0.0 (October 18, 2023)
|
||||
--------------------------------
|
||||
|
||||
New features
|
||||
^^^^^^^^^^^^
|
||||
|
||||
1. Added simulation on GPU and TPU via the new :doc:`mjx` (MJX) Python module. Python users can now
|
||||
natively run MuJoCo simulations at millions of steps per second on Google TPU or their own accelerator hardware.
|
||||
|
||||
- MJX is designed to work with on-device reinforcement learning algorithms. This Colab notebook demonstrates using
|
||||
MJX along with reinforcement learning to train humanoid and quadruped robots to locomote: |colab|
|
||||
- The MJX API is compatible with MuJoCo but is missing some features in this release. See the outline of
|
||||
:ref:`MJX feature parity <MjxFeatureParity>` for more details.
|
||||
|
||||
.. |colab| image:: https://colab.research.google.com/assets/colab-badge.svg
|
||||
:target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb
|
||||
|
||||
.. youtube:: QewlEqIZi1o
|
||||
:align: right
|
||||
:width: 240px
|
||||
|
||||
1. Added new signed distance field (SDF) collision primitive. SDFs can take any shape and are not constrained to be
|
||||
2. Added new signed distance field (SDF) collision primitive. SDFs can take any shape and are not constrained to be
|
||||
convex. Collision points are found by minimizing the maximum of the two colliding SDFs via gradient descent.
|
||||
|
||||
- Added new SDF plugin for defining implicit geometries. The plugin must define methods computing an SDF and its
|
||||
@@ -22,7 +33,7 @@ New features
|
||||
:align: right
|
||||
:width: 240px
|
||||
|
||||
2. Added new low-level model element called ``flex``, used to define deformable objects. These
|
||||
3. Added new low-level model element called ``flex``, used to define deformable objects. These
|
||||
`simplicial complexes <https://en.wikipedia.org/wiki/Simplicial_complex>`__ can be of dimension 1, 2
|
||||
or 3, corresponding to stretchable lines, triangles or tetrahedra. Two new MJCF elements are used
|
||||
to define flexes. The top-level :ref:`deformable<deformable>` section contains the low-level flex definition.
|
||||
@@ -36,19 +47,19 @@ New features
|
||||
functionality is currently available both via :ref:`deformable<CDeformable>` and :ref:`composite<CComposite>`,
|
||||
and both are modifiable by the first-party
|
||||
`elasticity plugins <https://github.com/google-deepmind/mujoco/tree/main/plugin/elasticity>`__. We expect some of
|
||||
this functionallity to be unified in the future.
|
||||
this functionality to be unified in the future.
|
||||
|
||||
.. youtube:: Vc1tq0fFvQA
|
||||
:align: right
|
||||
:width: 240px
|
||||
|
||||
3. Added constraint island discovery with :ref:`mj_island`. Constraint islands are disjoint sets of constraints
|
||||
4. Added constraint island discovery with :ref:`mj_island`. Constraint islands are disjoint sets of constraints
|
||||
and degrees-of-freedom that do not interact. The only solver which currently supports islands is
|
||||
:ref:`CG<option-solver>`. Island discovery can be activated using a new :ref:`enable flag<option-flag-island>`.
|
||||
If island discovery is enabled, geoms, contacts and tendons will be colored according to the corresponding island,
|
||||
see video. Island discovery is currently disabled for models that have deformable objects (see prevous item).
|
||||
see video. Island discovery is currently disabled for models that have deformable objects (see previous item).
|
||||
|
||||
4. Added :ref:`mjThreadPool` and :ref:`mjTask` which allow for multi-threaded operations within the MuJoCo engine
|
||||
5. Added :ref:`mjThreadPool` and :ref:`mjTask` which allow for multi-threaded operations within the MuJoCo engine
|
||||
pipeline. If engine-internal threading is enabled, the following operations will be multi-threaded:
|
||||
|
||||
- Island constraint resolution, if island discovery is :ref:`enabled<option-flag-island>` and the
|
||||
@@ -60,7 +71,7 @@ New features
|
||||
Engine-internal threading is a work in progress and currently only available in first-party code via the
|
||||
:ref:`testspeed<saTestspeed>` utility, exposed with the ``npoolthread`` flag.
|
||||
|
||||
5. Added capability to initialize :ref:`composite<body-composite>` particles from OBJ files. Fixes :github:issue:`642`
|
||||
6. Added capability to initialize :ref:`composite<body-composite>` particles from OBJ files. Fixes :github:issue:`642`
|
||||
and :github:issue:`674`.
|
||||
|
||||
General
|
||||
@@ -69,32 +80,32 @@ General
|
||||
.. admonition:: Breaking API changes
|
||||
:class: attention
|
||||
|
||||
6. Removed the macros ``mjMARKSTACK`` and ``mjFREESTACK``.
|
||||
7. Removed the macros ``mjMARKSTACK`` and ``mjFREESTACK``.
|
||||
|
||||
**Migration:** These macros have been replaced by new functions :ref:`mj_markStack` and
|
||||
:ref:`mj_freeStack`. These functions manage the :ref:`mjData stack<siStack>` in a fully encapsulated way (i.e.,
|
||||
without introducing a local variable at the call site).
|
||||
|
||||
7. Renamed ``mj_stackAlloc`` to :ref:`mj_stackAllocNum`. The new function :ref:`mj_stackAllocByte` allocates an
|
||||
8. Renamed ``mj_stackAlloc`` to :ref:`mj_stackAllocNum`. The new function :ref:`mj_stackAllocByte` allocates an
|
||||
arbitrary number of bytes and has an additional argument for specifying the alignment of the returned pointer.
|
||||
|
||||
**Migration:** The functionality for allocating ``mjtNum`` arrays is now available via :ref:`mj_stackAllocNum`.
|
||||
|
||||
8. Renamed the ``nstack`` field in :ref:`mjModel` and :ref:`mjData` to ``narena``. Changed ``narena``, ``pstack``,
|
||||
9. Renamed the ``nstack`` field in :ref:`mjModel` and :ref:`mjData` to ``narena``. Changed ``narena``, ``pstack``,
|
||||
and ``maxuse_stack`` to count number of bytes rather than number of :ref:`mjtNum` |-| s.
|
||||
|
||||
9. Changed :ref:`mjData.solver<mjData>`, the array used to collect solver diagnostic information.
|
||||
This array of :ref:`mjSolverStat` structs is now of length ``mjNISLAND * mjNSOLVER``, interpreted as as a matrix.
|
||||
Each row of length ``mjNSOLVER`` contains separate solver statistics for each constraint island.
|
||||
If the solver does not use islands, only row 0 is filled.
|
||||
10. Changed :ref:`mjData.solver<mjData>`, the array used to collect solver diagnostic information.
|
||||
This array of :ref:`mjSolverStat` structs is now of length ``mjNISLAND * mjNSOLVER``, interpreted as as a matrix.
|
||||
Each row of length ``mjNSOLVER`` contains separate solver statistics for each constraint island.
|
||||
If the solver does not use islands, only row 0 is filled.
|
||||
|
||||
- The new constant :ref:`mjNISLAND<glNumeric>` was set to 20.
|
||||
- :ref:`mjNSOLVER<glNumeric>` was reduced from 1000 to 200.
|
||||
- Added :ref:`mjData.solver_nisland<mjData>`: the number of islands for which the solver ran.
|
||||
- Renamed ``mjData.solver_iter`` to ``solver_niter``. Both this member and ``mjData.solver_nnz`` are now integer
|
||||
vectors of length ``mjNISLAND``.
|
||||
- The new constant :ref:`mjNISLAND<glNumeric>` was set to 20.
|
||||
- :ref:`mjNSOLVER<glNumeric>` was reduced from 1000 to 200.
|
||||
- Added :ref:`mjData.solver_nisland<mjData>`: the number of islands for which the solver ran.
|
||||
- Renamed ``mjData.solver_iter`` to ``solver_niter``. Both this member and ``mjData.solver_nnz`` are now integer
|
||||
vectors of length ``mjNISLAND``.
|
||||
|
||||
10. Removed ``mjOption.collision`` and the associated ``option/collision`` attribute.
|
||||
11. Removed ``mjOption.collision`` and the associated ``option/collision`` attribute.
|
||||
|
||||
**Migration:**
|
||||
|
||||
@@ -105,39 +116,39 @@ General
|
||||
:ref:`conaffinity<body-geom-conaffinity>` attributes in the model and then setting them globally to ``0`` using
|
||||
|br| ``<default> <geom contype="0" conaffinity="0"/> </default>``.
|
||||
|
||||
11. Removed the :at:`rope` and :at:`cloth` composite objects.
|
||||
12. Removed the :at:`rope` and :at:`cloth` composite objects.
|
||||
|
||||
**Migration:** Users should use the :at:`cable` and :at:`shell` elasticity plugins.
|
||||
|
||||
12. Added :ref:`mjData.eq_active<mjData>` user input variable, for enabling/disabling the state of equality
|
||||
13. Added :ref:`mjData.eq_active<mjData>` user input variable, for enabling/disabling the state of equality
|
||||
constraints. Renamed ``mjModel.eq_active`` to :ref:`mjModel.eq_active0<mjModel>`, which now has the semantic of
|
||||
"initial value of ``mjData.eq_active``". Fixes :github:issue:`876`.
|
||||
|
||||
**Migration:** Replace uses of ``mjModel.eq_active`` with ``mjData.eq_active``.
|
||||
|
||||
13. Changed the default of :ref:`autolimits<compiler-autolimits>` from "false" to "true". This is a minor breaking
|
||||
14. Changed the default of :ref:`autolimits<compiler-autolimits>` from "false" to "true". This is a minor breaking
|
||||
change. The potential breakage applies to models which have elements with "range" defined and "limited" not set.
|
||||
Such models cannot be loaded since version 2.2.2 (July 2022).
|
||||
|
||||
14. Added a new :ref:`dyntype<actuator-general-dyntype>`, ``filterexact``, which updates first-order filter states with
|
||||
15. Added a new :ref:`dyntype<actuator-general-dyntype>`, ``filterexact``, which updates first-order filter states with
|
||||
the exact formula rather than with Euler integration.
|
||||
15. Added an actuator attribute, :ref:`actearly<actuator-general-actearly>`, which uses semi-implicit integration for
|
||||
16. Added an actuator attribute, :ref:`actearly<actuator-general-actearly>`, which uses semi-implicit integration for
|
||||
actuator forces: using the next step's actuator state to compute the current actuator forces.
|
||||
16. Renamed ``actuatorforcerange`` and ``actuatorforcelimited``, introduced in the previous version to
|
||||
17. Renamed ``actuatorforcerange`` and ``actuatorforcelimited``, introduced in the previous version to
|
||||
:ref:`actuatorfrcrange<body-joint-actuatorfrcrange>` and
|
||||
:ref:`actuatorfrclimited<body-joint-actuatorfrclimited>`, respectively.
|
||||
17. Added the flag :ref:`eulerdamp<option-flag-eulerdamp>`, which disables implicit integration of joint damping in the
|
||||
18. Added the flag :ref:`eulerdamp<option-flag-eulerdamp>`, which disables implicit integration of joint damping in the
|
||||
Euler integrator. See the :ref:`Numerical Integration<geIntegration>` section for more details.
|
||||
18. Added the flag :ref:`invdiscrete<option-flag-invdiscrete>`, which enables discrete-time inverse dynamics for all
|
||||
19. Added the flag :ref:`invdiscrete<option-flag-invdiscrete>`, which enables discrete-time inverse dynamics for all
|
||||
:ref:`integrators<option-integrator>` other than ``RK4``. See the flag documentation for more details.
|
||||
19. Added :ref:`ls_iterations<option-ls_iterations>` and :ref:`ls_tolerance<option-ls_tolerance>` options for adjusting
|
||||
20. Added :ref:`ls_iterations<option-ls_iterations>` and :ref:`ls_tolerance<option-ls_tolerance>` options for adjusting
|
||||
linesearch stopping criteria in CG and Newton solvers. These can be useful for performance tuning.
|
||||
20. Added ``mesh_pos`` and ``mesh_quat`` fields to :ref:`mjModel` to store the normalizing transformation applied to
|
||||
21. Added ``mesh_pos`` and ``mesh_quat`` fields to :ref:`mjModel` to store the normalizing transformation applied to
|
||||
mesh assets. Fixes :github:issue:`409`.
|
||||
21. Added camera :ref:`resolution<body-camera-resolution>` attribute and :ref:`camprojection<sensor-camprojection>`
|
||||
22. Added camera :ref:`resolution<body-camera-resolution>` attribute and :ref:`camprojection<sensor-camprojection>`
|
||||
sensor. If camera resolution is set to positive values, the camera projection sensor will report the location of a
|
||||
target site, projected onto the camera image, in pixel coordinates.
|
||||
22. Added :ref:`camera<body-camera>` calibration attributes:
|
||||
23. Added :ref:`camera<body-camera>` calibration attributes:
|
||||
|
||||
- The new attributes are :ref:`resolution<body-camera-resolution>`, :ref:`focal<body-camera-focal>`,
|
||||
:ref:`focalpixel<body-camera-focalpixel>`, :ref:`principal<body-camera-principal>`,
|
||||
@@ -146,21 +157,21 @@ General
|
||||
attributes are specified. See the following
|
||||
`example model <https://github.com/deepmind/mujoco/blob/main/test/engine/testdata/vis_visualize/frustum.xml>`__.
|
||||
- Note that these attributes only take effect for offline rendering and do not affect interactive visualisation.
|
||||
23. Implemented reversed Z rendering for better depth precision. An enum :ref:`mjtDepthMap` was added with values
|
||||
24. Implemented reversed Z rendering for better depth precision. An enum :ref:`mjtDepthMap` was added with values
|
||||
``mjDEPTH_ZERONEAR`` and ``mjDEPTH_ZEROFAR``, which can be used to set the new ``readDepthMap`` attribute in
|
||||
:ref:`mjrContext` to control how the depth returned by :ref:`mjr_readPixels` is mapped from ``znear`` to ``zfar``.
|
||||
Contribution :github:pull:`978` by `Levi Burner <https://github.com/aftersomemath>`__.
|
||||
24. Deleted the code sample ``testxml``. The functionality provided by this utility is implemented in the
|
||||
25. Deleted the code sample ``testxml``. The functionality provided by this utility is implemented in the
|
||||
`WriteReadCompare <https://github.com/google-deepmind/mujoco/blob/main/test/xml/xml_native_writer_test.cc>`__ test.
|
||||
25. Deleted the code sample ``derivative``. Functionality provided by :ref:`mjd_transitionFD`.
|
||||
26. Deleted the code sample ``derivative``. Functionality provided by :ref:`mjd_transitionFD`.
|
||||
|
||||
Python bindings
|
||||
^^^^^^^^^^^^^^^
|
||||
|
||||
26. Fixed :github:issue:`870` where calling ``update_scene`` with an invalid camera name used the default camera.
|
||||
27. Added ``user_scn`` to the :ref:`passive viewer<PyViewerPassive>` handle, which allows users to add custom
|
||||
27. Fixed :github:issue:`870` where calling ``update_scene`` with an invalid camera name used the default camera.
|
||||
28. Added ``user_scn`` to the :ref:`passive viewer<PyViewerPassive>` handle, which allows users to add custom
|
||||
visualization geoms (:github:issue:`1023`).
|
||||
28. Added optional boolean keyword arguments ``show_left_ui`` and ``show_right_ui`` to the functions ``viewer.launch``
|
||||
29. Added optional boolean keyword arguments ``show_left_ui`` and ``show_right_ui`` to the functions ``viewer.launch``
|
||||
and ``viewer.launch_passive``, which allow users to launch a viewer with UI panels hidden.
|
||||
|
||||
Simulate
|
||||
@@ -170,11 +181,11 @@ Simulate
|
||||
:align: right
|
||||
:width: 240px
|
||||
|
||||
29. Added **state history** mechanism to :ref:`simulate<saSimulate>` and the managed
|
||||
30. Added **state history** mechanism to :ref:`simulate<saSimulate>` and the managed
|
||||
:ref:`Python viewer<PyViewerManaged>`. State history can be viewed by scrubbing the History slider and (more
|
||||
precisely) with the left and right arrow keys. See screen capture:
|
||||
|
||||
30. The ``LOADING...`` label is now shown correctly. Contribution :github:pull:`1070` by
|
||||
31. The ``LOADING...`` label is now shown correctly. Contribution :github:pull:`1070` by
|
||||
`Levi Burner <https://github.com/aftersomemath>`__.
|
||||
|
||||
Documentation
|
||||
@@ -184,17 +195,17 @@ Documentation
|
||||
:align: right
|
||||
:width: 240px
|
||||
|
||||
31. Added :doc:`detailed documentation <computation/fluid>` of fluid force modeling, and an illustrative example model
|
||||
32. Added :doc:`detailed documentation <computation/fluid>` of fluid force modeling, and an illustrative example model
|
||||
showing `tumbling cards <https://github.com/google-deepmind/mujoco/blob/main/model/cards/cards.xml>`__ using the
|
||||
ellipsoid-based fluid model.
|
||||
|
||||
Bug fixes
|
||||
^^^^^^^^^
|
||||
|
||||
32. Fixed a bug that was causing :ref:`geom margin<body-geom-margin>` to be ignored during the construction of
|
||||
33. Fixed a bug that was causing :ref:`geom margin<body-geom-margin>` to be ignored during the construction of
|
||||
midphase collision trees.
|
||||
|
||||
33. Fixed a bug that was generating incorrect values in ``efc_diagApprox`` for weld equality constraints.
|
||||
34. Fixed a bug that was generating incorrect values in ``efc_diagApprox`` for weld equality constraints.
|
||||
|
||||
|
||||
Version 2.3.7 (July 20, 2023)
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 26 KiB |
@@ -14,6 +14,7 @@
|
||||
programming/index.rst
|
||||
APIreference/index.rst
|
||||
python
|
||||
MJX <mjx>
|
||||
unity
|
||||
models
|
||||
changelog
|
||||
|
||||
+312
@@ -0,0 +1,312 @@
|
||||
==========
|
||||
MuJoCo XLA
|
||||
==========
|
||||
|
||||
Starting with version 3.0.0, MuJoCo includes MuJoCo XLA (MJX) under the
|
||||
`mjx <https://github.com/google-deepmind/mujoco/tree/main/mjx>`__ directory. MJX allows MuJoCo to run on compute
|
||||
hardware supported by the `XLA <https://www.tensorflow.org/xla>`__ compiler via the
|
||||
`JAX <https://github.com/google/jax#readme>`__ framework. MJX runs on a
|
||||
`all platforms supported by JAX <https://jax.readthedocs.io/en/latest/installation.html#supported-platforms>`__: Nvidia
|
||||
and AMD GPUs, Apple Silicon, and `Google Cloud TPUs <https://cloud.google.com/tpu>`__.
|
||||
|
||||
The MJX API is consistent with the main simulation functions in the MuJoCo API, although it is currently missing some
|
||||
features. While the :ref:`API documentation <Mainsimulation>` is applicable to both libraries, we indicate features
|
||||
unsupported by MJX in the :ref:`notes <MjxFeatureParity>` below.
|
||||
|
||||
MJX is distributed as a separate package called ``mujoco-mjx`` on `PyPI <https://pypi.org/project/mujoco-mjx>`__.
|
||||
Although it depends on the main ``mujoco`` package for model compilation and visualization, it is a re-implementation of
|
||||
MuJoCo that uses the same algorithms as the MuJoCo implementation. However, in order to properly leverage JAX, MJX
|
||||
deliberately diverges from the MuJoCo API in a few places, see below.
|
||||
|
||||
MJX is a successor to the `generalized physics pipeline <https://github.com/google/brax/tree/main/brax/generalized>`__
|
||||
in Google's `Brax <https://github.com/google/brax>`__ physics and reinforcement learning library. MJX was built
|
||||
by core contributors to both MuJoCo and Brax, who will together continue to support both Brax (for its reinforcement
|
||||
learning algorithms and included environments) and MJX (for its physics algorithms). A future version of Brax will
|
||||
depend on the ``mujoco-mjx`` package, and Brax's existing
|
||||
`generalized pipeline <https://github.com/google/brax/tree/main/brax/generalized>`__ will be deprecated. This change
|
||||
will be largely transparent to users of Brax.
|
||||
|
||||
.. _MjxNotebook:
|
||||
|
||||
Tutorial notebook
|
||||
=================
|
||||
|
||||
The following IPython notebook demonstrates the use of MJX along with reinforcement learning to train humanoid and
|
||||
quadruped robots to locomote: |colab|.
|
||||
|
||||
.. |colab| image:: https://colab.research.google.com/assets/colab-badge.svg
|
||||
:target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb
|
||||
|
||||
.. _MjxInstallation:
|
||||
|
||||
Installation
|
||||
============
|
||||
|
||||
The recommended way to install this package is via `PyPI <https://pypi.org/project/mujoco-mjx/>`__:
|
||||
|
||||
.. code-block:: shell
|
||||
|
||||
pip install mujoco-mjx
|
||||
|
||||
A copy of the MuJoCo library is provided as part of this package's depdendencies and does **not** need to be downloaded
|
||||
or installed separately.
|
||||
|
||||
.. _MjxUsage:
|
||||
|
||||
Basic usage
|
||||
===========
|
||||
|
||||
Once installed, the package can be imported via ``from mujoco import mjx``. Structs, functions, and enums are available
|
||||
directly from the top-level ``mjx`` module.
|
||||
|
||||
.. _MjxStructs:
|
||||
|
||||
Structs
|
||||
-------
|
||||
|
||||
Before running MJX functions on an accelerator device, structs must be copied onto the device via the ``mjx.device_put``
|
||||
function. Placing an :ref:`mjModel` on device yields an ``mjx.Model``. Placing an :ref:`mjData` on device yields
|
||||
an ``mjx.Data``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
model = mujoco.MjModel.from_xml_string("...")
|
||||
data = mujoco.MjData(model)
|
||||
mjx_model = mjx.device_put(model)
|
||||
mjx_data = mjx.device_put(data)
|
||||
|
||||
These MJX variants mirror their MuJoCo counterparts but have three key differences:
|
||||
|
||||
#. Fields in ``mjx.Model`` and ``mjx.Data`` are JAX arrays copied onto device, instead of numpy arrays.
|
||||
#. Some fields are missing from ``mjx.Model`` and ``mjx.Data`` for features that are
|
||||
:ref:`unsupported <mjxFeatureParity>` in MJX.
|
||||
#. Arrays in ``mjx.Model`` and ``mjx.Data`` support adding batch dimensions. Batch dimensions are a natural way to
|
||||
express domain randomization (in the case of ``mjx.Model``) or high-throughput simulation for reinforcement learning
|
||||
(in the case of ``mjx.Data``).
|
||||
|
||||
|
||||
Neither ``mjx.Model`` nor ``mjx.Data`` are meant to be constructed manually. An ``mjx.Data`` may be created by calling
|
||||
``mjx.make_data``, which mirrors the :ref:`mj_makeData` function in MuJoCo:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
model = mujoco.MjModel.from_xml_string("...")
|
||||
mjx_model = mjx.device_put(model)
|
||||
mjx_data = mjx.make_data(model)
|
||||
|
||||
Using ``mx.make_data`` may be preferable when constructing batched ``mjx.Data`` structures inside of a ``vmap``.
|
||||
|
||||
.. _MjxFunctions:
|
||||
|
||||
Functions
|
||||
---------
|
||||
|
||||
MuJoCo functions are exposed as MJX functions of the same name, but following
|
||||
`PEP 8 <https://peps.python.org/pep-0008/>`__-compliant names. Most of the :ref:`main simulation <Mainsimulation>` and
|
||||
some of the :ref:`sub-components <Subcomponents>` for forward simulation are available from the top-level ``mjx`` module.
|
||||
|
||||
MJX functions are not `JIT compiled <https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html>`__ by default -- we
|
||||
leave it to the user to JIT MJX functions, or JIT their own functions that reference MJX functions. See the
|
||||
:ref:`minimal example <MjxExample>` below.
|
||||
|
||||
.. _MjxEnums:
|
||||
|
||||
Enums and constants
|
||||
-------------------
|
||||
|
||||
MJX enums are available as ``mjx.EnumType.ENUM_VALUE``, for example ``mjx.JointType.FREE``. Enums for unsupported MJX
|
||||
features are omitted from the MJX enum declaration. MJX declares no constants but references MuJoCo constants directly.
|
||||
|
||||
.. _MjxExample:
|
||||
|
||||
Minimal example
|
||||
---------------
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Throw a ball at 100 different velocities.
|
||||
|
||||
import jax
|
||||
import mujoco
|
||||
from mujoco import mjx
|
||||
|
||||
XML=r"""
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body>
|
||||
<freejoint/>
|
||||
<geom size=".15" mass="1" type="sphere"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
"""
|
||||
|
||||
model = mujoco.MjModel.from_xml_string(XML)
|
||||
mjx_model = mjx.device_put(model)
|
||||
|
||||
@jax.vmap
|
||||
def batched_step(vel):
|
||||
mjx_data = mjx.make_data(mjx_model)
|
||||
qvel = mjx_data.qvel.at[0].set(vel)
|
||||
mjx_data = mjx_data.replace(qvel=qvel)
|
||||
pos = mjx.step(mjx_model, mjx_data).qpos[0]
|
||||
return pos
|
||||
|
||||
vel = jax.numpy.arange(0.0, 1.0, 0.01)
|
||||
pos = jax.jit(batched_step)(vel)
|
||||
print(pos)
|
||||
|
||||
.. _MjxFeatureParity:
|
||||
|
||||
Feature Parity
|
||||
==============
|
||||
|
||||
MJX supports most of the main simulation features of MuJoCo, with a few exceptions. MJX will raise an exception if
|
||||
asked to copy to device an :ref:`mjModel` with field values referencing unsupported features.
|
||||
|
||||
The following features are **fully supported** in MJX:
|
||||
|
||||
.. list-table::
|
||||
:width: 90%
|
||||
:align: left
|
||||
:widths: 1 5
|
||||
:header-rows: 1
|
||||
|
||||
* - Category
|
||||
- Feature
|
||||
* - Dynamics
|
||||
- :ref:`Forward <mj_forward>`
|
||||
* - :ref:`Joint <mjtJoint>`
|
||||
- ``FREE``, ``BALL``, ``SLIDE``, ``HINGE``
|
||||
* - :ref:`Transmission <mjtTrn>`
|
||||
- ``TRN_JOINT``
|
||||
* - :ref:`Actuation <geactuation>`
|
||||
- ``DYN_NONE``, ``DYN_INTEGRATOR``, ``DYN_FILTER``, ``GAIN_FIXED``, ``GAIN_AFFINE``, ``BIAS_NONE``,
|
||||
``BIAS_AFFINE``
|
||||
* - :ref:`Geom <mjtGeom>`
|
||||
- ``PLANE``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH``
|
||||
* - :ref:`Constraint <mjtConstraint>`
|
||||
- ``EQUALITY``, ``FRICTION_DOF``, ``LIMIT_JOINT``, ``CONTACT_PYRAMIDAL``
|
||||
* - :ref:`Integrator <mjtIntegrator>`
|
||||
- ``EULER``, ``RK4``
|
||||
* - :ref:`Cone <mjtCone>`
|
||||
- ``PYRAMIDAL``
|
||||
* - :ref:`Condim <coContact>`
|
||||
- 3
|
||||
* - :ref:`Solver <mjtSolver>`
|
||||
- ``CG``
|
||||
* - Fluid Model
|
||||
- :ref:`flInertia`
|
||||
|
||||
The following features are **in development** and coming soon:
|
||||
|
||||
.. list-table::
|
||||
:width: 90%
|
||||
:align: left
|
||||
:widths: 1 5
|
||||
:header-rows: 1
|
||||
|
||||
* - Category
|
||||
- Feature
|
||||
* - Dynamics
|
||||
- :ref:`Inverse <mj_inverse>`
|
||||
* - :ref:`Transmission <mjtTrn>`
|
||||
- ``TRN_TENDON``
|
||||
* - :ref:`Geom <mjtGeom>`
|
||||
- ``HFIELD``, ``ELLIPSOID``, ``CYLINDER``, ``SDF``
|
||||
* - :ref:`Integrator <mjtIntegrator>`
|
||||
- ``IMPLICIT``, ``IMPLICITFAST``
|
||||
* - :ref:`Cone <mjtCone>`
|
||||
- ``ELLIPTIC``
|
||||
* - :ref:`Condim <coContact>`
|
||||
- 1, 4, 6
|
||||
* - :ref:`Solver <mjtSolver>`
|
||||
- ``NEWTON``
|
||||
* - Fluid Model
|
||||
- :ref:`flEllipsoid`
|
||||
* - :ref:`Tendons <tendon>`
|
||||
- :ref:`Spatial <tendon-spatial>`, :ref:`Fixed <tendon-fixed>`
|
||||
|
||||
The following features are **unsupported**:
|
||||
|
||||
.. list-table::
|
||||
:width: 90%
|
||||
:align: left
|
||||
:widths: 1 5
|
||||
:header-rows: 1
|
||||
|
||||
* - Category
|
||||
- Feature
|
||||
* - :ref:`Transmission <mjtTrn>`
|
||||
- ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_SITE``, ``TRN_BODY``, ``MUSCLE``
|
||||
* - :ref:`Solver <mjtSolver>`
|
||||
- ``PGS``
|
||||
* - :ref:`Callbacks <glphysics>`
|
||||
- ``mjDYN_USER``, ``mjGAIN_USER``, ``mjBIAS_USER``, ``mjSENS_USER``
|
||||
|
||||
.. _MjxSharpBits:
|
||||
|
||||
🔪 MJX - The Sharp Bits 🔪
|
||||
==========================
|
||||
|
||||
GPUs and TPUs have unique performance tradeoffs that MJX is subject to. MJX specializes in simulating big batches of
|
||||
parallel identical physics scenes using algorithms that can be efficiently vectorized on
|
||||
`SIMD hardware <https://en.wikipedia.org/wiki/Single_instruction,_multiple_data>`__. This specialization is useful
|
||||
for machine learning workloads such as `reinforcement learning <https://en.wikipedia.org/wiki/Reinforcement_learning>`__
|
||||
that require massive data throughput.
|
||||
|
||||
There are certain workflows that MJX is ill-suited for:
|
||||
|
||||
Single scene simulation
|
||||
Simulating a single scene (1 instance of :ref:`mjData`), MJX can be **10x** slower than MuJoCo, which has been
|
||||
carefully optimized for CPU. MJX works best when simulating thousands or tens of thousands of scenes in parallel.
|
||||
|
||||
Large, complex scenes with many contacts
|
||||
Accelerators exhibit poor performance for
|
||||
`branching code <https://aschrein.github.io/jekyll/update/2019/06/13/whatsup-with-my-branches-on-gpu.html#tldr>`__.
|
||||
Branching is used in broad-phase collision detection, when identifying potential collisions between large numbers of
|
||||
bodies in a scene. MJX ships with a simple branchless broad-phase algorithm (see performance tuning) but it is not as
|
||||
powerful as the one in MuJoCo.
|
||||
|
||||
To see how this affects simulation, let us consider a physics scene with increasing numbers of physics bodies. We
|
||||
simulate a scene with a variable number of humanoids (from 1 to 10) and then compare MJX's performance on an Nvidia
|
||||
A100 GPU to MuJoCo on a 12-core workstation:
|
||||
|
||||
.. figure:: images/mjx/mujoco_vs_mjx_large_scene.png
|
||||
:width: 658px
|
||||
:align: center
|
||||
|
||||
Notice that as we increase the number of humanoids (which increases the number of potential contacts in a scene), MJX
|
||||
performance degrades more rapidly than MuJoCo. At the limit, for such a large scene, MuJoCo performance nearly
|
||||
matches MJX.
|
||||
|
||||
Scenes with collisions between meshes with many vertices
|
||||
MJX supports mesh geometries and can determine if two meshes are colliding using branchless versions of
|
||||
`mesh collision algorithms <https://ubm-twvideo01.s3.amazonaws.com/o1/vault/gdc2013/slides/822403Gregorius_Dirk_TheSeparatingAxisTest.pdf>`__.
|
||||
These algorithms work well for smaller meshes (with hundreds of vertices) but suffer with large meshes. With careful
|
||||
tuning, MJX can simulate scenes with mesh collisions well -- see the MJX
|
||||
`shadow hand <https://github.com/google-deepmind/mujoco/tree/main/mjx/benchmark/model/shadow_hand/scene_right.xmld>`__
|
||||
config for an example.
|
||||
|
||||
.. _MjxPerformance:
|
||||
|
||||
Performance tuning
|
||||
==================
|
||||
|
||||
For MJX to perform well, some configuration parameters should be adjusted from their default MuJoCo values:
|
||||
|
||||
:ref:`option` element
|
||||
For now, solver must be set to ``CG`` (but Newton is on its way!). 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 sim2real.
|
||||
|
||||
:ref:`contact-pair` element
|
||||
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 <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.
|
||||
|
||||
:ref:`option-flag` element
|
||||
Disabling ``eulerdamp`` can help performance and is often not needed for stability.
|
||||
+2
-2
@@ -30,14 +30,14 @@ _____
|
||||
|
||||
The MuJoCo app needs to be run at least once before the native library can be used, in order to register the library as
|
||||
a trusted binary. Then, copy the dynamic library file from
|
||||
``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.2.3.8.dylib`` (it can be
|
||||
``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.3.0.0.dylib`` (it can be
|
||||
found by browsing the contents of ``MuJoCo.app``) and rename it as ``mujoco.dylib``.
|
||||
|
||||
Linux
|
||||
_____
|
||||
|
||||
Expand the ``tar.gz`` archive to ``~/.mujoco``. Then copy the dynamic library from
|
||||
``~/.mujoco/mujoco-2.3.8/lib/libmujoco.so.2.3.8`` and rename it as ``libmujoco.so``.
|
||||
``~/.mujoco/mujoco-3.0.0/lib/libmujoco.so.3.0.0`` and rename it as ``libmujoco.so``.
|
||||
|
||||
Windows
|
||||
_______
|
||||
|
||||
Reference in New Issue
Block a user