diff --git a/doc/images/mjx/SPS.svg b/doc/images/mjx/SPS.svg
new file mode 100644
index 00000000..a35a2508
--- /dev/null
+++ b/doc/images/mjx/SPS.svg
@@ -0,0 +1 @@
+
diff --git a/doc/images/mjx/mujoco_vs_mjx_large_scene.png b/doc/images/mjx/mujoco_vs_mjx_large_scene.png
deleted file mode 100644
index d9ef4b54..00000000
Binary files a/doc/images/mjx/mujoco_vs_mjx_large_scene.png and /dev/null differ
diff --git a/doc/mjx.rst b/doc/mjx.rst
index 481da5c8..2bdd9a58 100644
--- a/doc/mjx.rst
+++ b/doc/mjx.rst
@@ -171,7 +171,7 @@ The following features are **fully supported** in MJX:
.. list-table::
:width: 90%
:align: left
- :widths: 1 5
+ :widths: 2 5
:header-rows: 1
* - Category
@@ -182,9 +182,12 @@ The following features are **fully supported** in MJX:
- ``FREE``, ``BALL``, ``SLIDE``, ``HINGE``
* - :ref:`Transmission `
- ``TRN_JOINT``
- * - :ref:`Actuation `
- - ``DYN_NONE``, ``DYN_INTEGRATOR``, ``DYN_FILTER``, ``GAIN_FIXED``, ``GAIN_AFFINE``, ``BIAS_NONE``,
- ``BIAS_AFFINE``
+ * - :ref:`Actuator Dynamics `
+ - ``NONE``, ``INTEGRATOR``, ``FILTER``
+ * - :ref:`Actuator Gain `
+ - ``FIXED``, ``AFFINE``
+ * - :ref:`Actuator Bias `
+ - ``NONE``, ``AFFINE``
* - :ref:`Geom `
- ``PLANE``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH``
* - :ref:`Constraint `
@@ -207,7 +210,7 @@ The following features are **in development** and coming soon:
.. list-table::
:width: 90%
:align: left
- :widths: 1 5
+ :widths: 2 5
:header-rows: 1
* - Category
@@ -216,8 +219,16 @@ The following features are **in development** and coming soon:
- :ref:`Inverse `
* - :ref:`Transmission `
- ``TRN_TENDON``
+ * - :ref:`Actuator Dynamics `
+ - ``MUSCLE``
+ * - :ref:`Actuator Gain `
+ - ``MUSCLE``
+ * - :ref:`Actuator Bias `
+ - ``MUSCLE``
+ * - :ref:`Tendon Wrapping `
+ - ``NONE``, ``JOINT``, ``PULLEY``, ``SITE``, ``SPHERE``, ``CYLINDER``
* - :ref:`Geom `
- - ``HFIELD``, ``ELLIPSOID``, ``CYLINDER``, ``SDF``
+ - ``HFIELD``, ``ELLIPSOID``, ``CYLINDER``
* - :ref:`Constraint `
- ``CONTACT_FRICTIONLESS``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF``
* - :ref:`Integrator `
@@ -233,24 +244,32 @@ The following features are **in development** and coming soon:
* - :ref:`Equality `
- ``TENDON``
* - :ref:`Sensors `
- - All except ``mjSENS_PLUGIN``, ``mjSENS_USER``
+ - All except ``PLUGIN``, ``USER``
The following features are **unsupported**:
.. list-table::
:width: 90%
:align: left
- :widths: 1 5
+ :widths: 2 5
:header-rows: 1
* - Category
- Feature
* - :ref:`Transmission `
- - ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_SITE``, ``TRN_BODY``, ``MUSCLE``
+ - ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_SITE``, ``TRN_BODY``
+ * - :ref:`Actuator Dynamics `
+ - ``FILTEREXACT``, ``USER``
+ * - :ref:`Actuator Gain `
+ - ``USER``
+ * - :ref:`Actuator Bias `
+ - ``USER``
* - :ref:`Solver `
- ``PGS``
- * - :ref:`Callbacks `
- - ``mjDYN_USER``, ``mjGAIN_USER``, ``mjBIAS_USER``, ``mjSENS_USER``
+ * - :ref:`Sensors `
+ - ``PLUGIN``, ``USER``
+ * - :ref:`Geom `
+ - ``SDF``
.. _MjxSharpBits:
@@ -276,17 +295,20 @@ Large, complex scenes with many contacts
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:
+ To see how this affects simulation, let us consider a physics scene with increasing numbers of humanoid bodies,
+ varied from 1 to 10. We simulate this scene using CPU MuJoCo on an Apple M1 Pro and a 64-core AMD 3995WX and time
+ it using :ref:`testspeed`, using ``2 x numcore`` threads. We time the MJX simulation on an Nvidia
+ A100 GPU using a batch size of 8192 and an 8-chip
+ `v5 TPU `__
+ machine using a batch size of 16384. Note the vertical scale is logarithmic.
- .. figure:: images/mjx/mujoco_vs_mjx_large_scene.png
- :width: 658px
- :align: center
+ .. figure:: images/mjx/SPS.svg
+ :width: 95%
+ :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.
+ The values for a single humanoid (leftmost datapoints) for the four timed architectures are **320K**, **1.8M**,
+ **950K** and **2.7M** steps per second, respectively. Note that as we increase the number of humanoids (which
+ increases the number of potential contacts in a scene), MJX throughput decreases more rapidly than MuJoCo.
Scenes with collisions between meshes with many vertices
MJX supports mesh geometries and can determine if two meshes are colliding using branchless versions of