Adds support for Newton solver to MJX.

PiperOrigin-RevId: 578738033
Change-Id: Id0e68c7b961380cb1ed88b40d99f202c3fa8e228

doc/changelog.rst

@@ -28,12 +28,37 @@ General
 
 MJX
 ^^^
+- Added support for Newton solver (``mjSOL_NEWTON`` in :ref:`mjtSolver`).  The Newton solver significantly speeds up
+  simulation on GPU.  See updated benchmarks on an Nvidia A100 (reported numbers are steps-per-second):
+
+  .. list-table:: Steps-per-second, Conjugate Gradient vs. Newton on A100
+     :header-rows: 1
+     :align: left
+
+     * - Model
+       - CG
+       - Newton
+       - Speedup
+     * - `Humanoid <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/humanoid>`__
+       - 640,000
+       - 1,020,000
+       - **1.6 x**
+     * - `Barkour v0 <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/barkour_v0>`__
+       - 1,290,000
+       - 1,750,000
+       - **1.35 x**
+     * - `Shadow Hand <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/shadow_hand>`__
+       - 215,000
+       - 270,000
+       - **1.25 x**
+
+  Humanoid is the standard MuJoCo humanoid,
+  `Google Barkour <https://blog.research.google/2023/05/barkour-benchmarking-animal-level.html>`__ and the Shadow Hand
+  are both available in the :ref:`MuJoCo Menagerie<Menagerie>`.
 - Added support for joint equality constraints (``mjEQ_JOINT`` in :ref:`mjtEq`).
 - Fixed bug where mixed ``jnt_limited`` joints were not being constrained correctly.
 - Made ``device_put`` type validation more verbose (fixes :github:issue:`1113`).
-- Removed empty EFC rows from `MJX`, for joints with no limits (fixes :github:issue:`1117`).
-- Fixed bug where equality constraints became inactive (fixes :github:issue:`1129`).
-- Added an error when loading a model with tendons (fixes :github:issue:`1149`).
+- Removed empty EFC rows from ``MJX``, for joints with no limits (fixes :github:issue:`1117`).
 
 Python bindings
 ^^^^^^^^^^^^^^^

doc/mjx.rst

@@ -198,7 +198,7 @@ The following features are **fully supported** in MJX:
    * - :ref:`Condim <coContact>`
      - 3
    * - :ref:`Solver <mjtSolver>`
-     - ``CG``
+     - ``CG``, ``NEWTON``
    * - Fluid Model
      - :ref:`flInertia`
 
@@ -226,8 +226,6 @@ The following features are **in development** and coming soon:
      - ``ELLIPTIC``
    * - :ref:`Condim <coContact>`
      - 1, 4, 6
-   * - :ref:`Solver <mjtSolver>`
-     - ``NEWTON``
    * - Fluid Model
      - :ref:`flEllipsoid`
    * - :ref:`Tendons <tendon>`
@@ -306,10 +304,11 @@ 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.
+  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 sim-to-real.
+  The ``NEWTON`` :ref:`Solver <mjtSolver>` often delivers reasonable convergence with one solver iteration, and performs
+  well on GPU.  ``CG`` is currently a better choice for TPU.
 
 :ref:`contact-pair` element
   Consider explicitly marking geoms for collision detection to reduce the number of contacts that MJX must consider