Taylor Howell
10f4b36fb7
PiperOrigin-RevId: 825958166 Change-Id: I1e90d78d7413956cfa84711336be26a23a18e5aa
259 lines
8.1 KiB
ReStructuredText
259 lines
8.1 KiB
ReStructuredText
.. _MJW:
|
|
|
|
====================
|
|
MuJoCo Warp (MJWarp)
|
|
====================
|
|
|
|
.. toctree::
|
|
:hidden:
|
|
|
|
API <api.rst>
|
|
|
|
MuJoCo Warp (MJWarp) is an implementation of MuJoCo written in `Warp <https://nvidia.github.io/warp/>`__ and optimized
|
|
for `Nvidia <https://nvidia.com>`__ hardware and parallel simulation. MJWarp lives in the
|
|
`google-deepmind/mujoco_warp <https://github.com/google-deepmind/mujoco_warp>`__ GitHub repository and is currently in
|
|
beta.
|
|
|
|
.. TODO: remove after release
|
|
|
|
.. admonition:: Beta software
|
|
:class: attention
|
|
|
|
- MJWarp is beta software and is under active development.
|
|
- MJWarp developers will triage and respond to
|
|
`bug reports and feature requests <https://github.com/google-deepmind/mujoco_warp/issues>`__.
|
|
- MJWarp is mostly feature complete but requires performance optimization, documentation, and testing.
|
|
- The intended audience during Beta are physics engine enthusiasts and learning framework integrators.
|
|
|
|
.. _MJW_install:
|
|
|
|
Installation
|
|
============
|
|
|
|
The beta version of MuJoCo Warp is installed from GitHub. Please note that the beta version of MuJoCo Warp does not
|
|
support all versions of MuJoCo, Warp, CUDA, Nvidia drivers, etc.
|
|
|
|
.. code-block:: shell
|
|
|
|
git clone https://github.com/google-deepmind/mujoco_warp.git
|
|
cd mujoco_warp
|
|
python3 -m venv env
|
|
source env/bin/activate
|
|
pip install --upgrade pip
|
|
pip install uv
|
|
uv pip install -e .[dev,cuda]
|
|
|
|
Test the Installation
|
|
|
|
.. code-block:: shell
|
|
|
|
pytest
|
|
|
|
.. _MJW_Usage:
|
|
|
|
Basic usage
|
|
===========
|
|
|
|
Once installed, the package can be imported via ``import mujoco_warp as mjw``. Structs, functions, and enums are
|
|
available directly from the top-level ``mjw`` module.
|
|
|
|
Structs
|
|
-------
|
|
Before running MJWarp functions on an Nvidia GPU, structs must be copied onto the device via ``mjw.put_model`` and
|
|
``mjw.make_data`` or ``mjw.put_data`` functions. Placing an :ref:`mjModel` on device yields an ``mjw.Model``. Placing
|
|
an :ref:`mjData` on device yields an ``mjw.Data``:
|
|
|
|
.. code-block:: python
|
|
|
|
mjm = mujoco.MjModel.from_xml_string("...")
|
|
mjd = mujoco.MjData(mjm)
|
|
m = mjw.put_model(mjm)
|
|
d = mjw.put_data(mjm, mjd)
|
|
|
|
These MJWarp variants mirror their MuJoCo counterparts but have a few key differences:
|
|
|
|
#. ``mjw.Model`` and ``mjw.Data`` contain Warp arrays that are copied onto device.
|
|
#. Some fields are missing from ``mjw.Model`` and ``mjw.Data`` for features that are unsupported.
|
|
|
|
``nworld``, ``nconmax``, and ``njmax``
|
|
--------------------------------------
|
|
MJWarp is optimized for parallel simulation. A batch of simulations can be specified with three parameters:
|
|
|
|
- ``nworld``: Number of worlds to simulate.
|
|
- ``nconmax``: Expected number of contacts per world. The maximum number of contacts for all worlds is
|
|
``nconmax * nworld``.
|
|
- ``njmax``: Maximum number of constraints per world.
|
|
|
|
.. admonition:: Semantic difference for ``nconmax`` and ``njmax``.
|
|
:class: note
|
|
|
|
It is possible for the number of contacts per world to exceed ``nconmax`` if the total number of contacts for all
|
|
worlds does not exceed ``nworld x nconmax``. However, the number of constraints per world is strictly limited by
|
|
``njmax``.
|
|
|
|
Functions
|
|
_________
|
|
|
|
MuJoCo functions are exposed as MJWarp 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 ``mjw``
|
|
module.
|
|
|
|
Minimal example
|
|
---------------
|
|
|
|
.. code-block:: python
|
|
|
|
# Throw a ball at 100 different velocities.
|
|
|
|
import mujoco
|
|
import mujoco_warp as mjw
|
|
import warp as wp
|
|
|
|
_MJCF=r"""
|
|
<mujoco>
|
|
<worldbody>
|
|
<body>
|
|
<freejoint/>
|
|
<geom size=".15" mass="1" type="sphere"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
mjm = mujoco.MjModel.from_xml_string(_MJCF)
|
|
m = mjw.put_model(mjm)
|
|
d = mjw.make_data(mjm, nworld=100)
|
|
|
|
# initialize velocities
|
|
wp.copy(d.qvel, wp.array([[float(i) / 100, 0, 0, 0, 0, 0] for i in range(100)], dtype=float))
|
|
|
|
# simulate physics
|
|
mjw.step(m, d)
|
|
|
|
print(f'qpos:\n{d.qpos.numpy()}')
|
|
|
|
A call to ``mjw.step`` is comprised of a collection of kernel launches. Warp will launch these kernels individually if
|
|
this function is called directly. To improve performance, especially if the function will be called multiple times, it
|
|
is recommended to capture the operations that comprise the function as a CUDA graph
|
|
|
|
.. code-block:: python
|
|
|
|
with wp.ScopedCapture() as capture:
|
|
mjw.step(m, d)
|
|
|
|
The graph can then be launched or re-launched
|
|
|
|
.. code-block:: python
|
|
|
|
wp.capture_launch(capture.graph)
|
|
|
|
and will typically be significantly faster compared to calling ``mjw.step`` directly. Please see the
|
|
`Warp Graph API reference <https://nvidia.github.io/warp/modules/runtime.html#graph-api-reference>`__ for details.
|
|
|
|
.. _MJW_Cli:
|
|
|
|
Helpful Command Line Scripts
|
|
----------------------------
|
|
|
|
Benchmark an environment with testspeed
|
|
|
|
.. code-block:: shell
|
|
|
|
mjwarp-testspeed benchmark/humanoid/humanoid.xml
|
|
|
|
Interactive environment simulation with MJWarp
|
|
|
|
.. code-block:: shell
|
|
|
|
mjwarp-viewer benchmark/humanoid/humanoid.xml
|
|
|
|
Feature Parity
|
|
==============
|
|
|
|
MJWarp supports most of the main simulation features of MuJoCo, with a few exceptions. MJWarp will raise an exception if
|
|
asked to copy to device an :ref:`mjModel` with field values referencing unsupported features.
|
|
|
|
The following features are **not supported** in MJWarp:
|
|
|
|
.. list-table::
|
|
:width: 90%
|
|
:align: left
|
|
:widths: 2 5
|
|
:header-rows: 1
|
|
|
|
* - Category
|
|
- Feature
|
|
* - :ref:`Equality <mjtEq>`
|
|
- ``FLEX``
|
|
* - :ref:`Integrator <mjtIntegrator>`
|
|
- ``IMPLICIT``, ``IMPLICITFAST`` not supported with fluid drag
|
|
* - :ref:`Solver <mjtSolver>`
|
|
- ``PGS``, ``noslip``, :ref:`islands <soIsland>`
|
|
* - Fluid Model
|
|
- :ref:`flEllipsoid`
|
|
* - :ref:`Sensors <mjtSensor>`
|
|
- ``GEOMDIST``, ``GEOMNORMAL``, ``GEOMFROMTO``
|
|
* - Flex
|
|
- ``VERTCOLLIDE=false``, ``INTERNAL=true``, ``nflex > 1``
|
|
* - Jacobian format
|
|
- ``SPARSE``
|
|
* - Option
|
|
- :ref:`contact override <COverride>`
|
|
* - Plugins
|
|
- ``All`` except ``SDF``
|
|
* - :ref:`User parameters <CUser>`
|
|
- ``All``
|
|
|
|
Batched ``Model`` Fields
|
|
========================
|
|
|
|
To enable batched simulation with different model parameter values, many ``mjw.Model`` fields have a leading batch
|
|
dimension. By default, the leading dimension is 1 (i.e., ``field.shape[0] == 1``) and the same value(s) will be applied
|
|
to all worlds. It is possible to override one of these fields with a ``wp.array`` that has a leading dimension greater
|
|
than one. This field will be indexed with a modulo operation of the world id and batch dimension:
|
|
``field[worldid % field.shape[0]]``. Importantly, the field shape should be overridden prior to graph capture (i.e.,
|
|
``wp.ScopedCapture``)
|
|
|
|
.. code-block:: python
|
|
|
|
# override shape and values
|
|
m.dof_damping = wp.array([[0.1], [0.2]], dtype=float) # nworld=2
|
|
|
|
with wp.ScopedCapture() as capture:
|
|
mjw.step(m, d)
|
|
|
|
It is possible to override the field shape and set the field values after graph capture
|
|
|
|
.. code-block:: python
|
|
|
|
# override shape
|
|
m.dof_damping = wp.empty((2, 1), dtype=float)
|
|
|
|
with wp.ScopedCapture() as capture:
|
|
mjw.step(m, d)
|
|
|
|
# set batched values
|
|
dof_damping_batch = wp.array([[0.1], [0.2]], dtype=float) # nworld=2
|
|
wp.copy(m.dof_damping, dof_damping_batch) # m.dof = dof_damping_batch will not work correctly
|
|
|
|
.. admonition:: Heterogeneous worlds
|
|
:class: note
|
|
|
|
Heterogeneous worlds, for example: per-world meshes or number of degrees of freedom, are not currently available.
|
|
|
|
Parallel Linesearch
|
|
===================
|
|
|
|
In addition to the constraint solver's iterative linesearch, MJWarp provides a parallel linesearch routine that
|
|
evaluates a set of step sizes in parallel and selects the best one. The step sizes are spaced logarithmically from
|
|
``Model.opt.ls_parallel_min_step`` to 1 and the number of step sizes to evaluate is set via ``Model.opt.ls_iterations``.
|
|
|
|
To enable this routine set ``Model.opt.ls_parallel=True`` or add a custom numeric field to the XML
|
|
|
|
.. code-block:: xml
|
|
|
|
<custom>
|
|
<numeric name="ls_parallel" data="1"/>
|
|
</custom>
|