Add new differentiable physics tutorial to readme.

PiperOrigin-RevId: 631458983
Change-Id: Ic0ef9e93397a974be50f0616571a4a5fea6f2b21

README.md

@@ -50,16 +50,17 @@ your machine.
 If you are a Python user, you might want to start with our tutorial notebooks
 running on Google Colab:
 
- - The **introductory tutorial** teaches MuJoCo basics:
+ - The **introductory** tutorial teaches MuJoCo basics:
    [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb)
  - The **LQR** tutorial synthesizes a linear-quadratic controller, balancing a humanoid on one leg:
    [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/LQR.ipynb)
  - The **least-squares** tutorial explains how to use the Python-based nonlinear least-squares solver:
    [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/least_squares.ipynb)
  - The **MJX** tutorial provides usage examples of
-   [MuJoCo XLA](https://mujoco.readthedocs.io/en/stable/mjx.html), a branch of MuJoCo written in
-   JAX:
+   [MuJoCo XLA](https://mujoco.readthedocs.io/en/stable/mjx.html), a branch of MuJoCo written in JAX:
    [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb)
+ - The **differentiable physics** tutorial trains locomotion policies with analytical gradients automatically derived from MuJoCo's physics step:
+   [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/training_apg.ipynb)
 
 ## Installation
 

doc/changelog.rst

@@ -52,14 +52,16 @@ MJX
 17. Add support functions for ``id2name`` and ``name2id``, MJX versions of :ref:`mj_id2name` and :ref:`mj_name2id`.
 18. Added support for :ref:`gravcomp<body-gravcomp>` and :ref:`actuatorgravcomp<body-joint-actuatorgravcomp>`.
 19. Fixed a bug in ``mjx.ray`` for sometimes allowed negative distances for ray-mesh tests.
+20. Added a new `differentiable physics tutorial <https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/training_apg.ipynb>`__ that demonstrates training locomotion policies with analytical gradients
+    automatically derived from the MJX physics step.  Contribution by :github:user:`Andrew-Luo1`.
 
 Bug fixes
 ^^^^^^^^^
-20. Defaults of lights were not being saved, now fixed.
-21. Prevent overwriting of frame names by body names when saving an XML. Bug introduced in 3.1.4.
-22. Fixed bug in Python binding of :ref:`mj_saveModel`: ``buffer`` argument was documented as optional but was actually
+21. Defaults of lights were not being saved, now fixed.
+22. Prevent overwriting of frame names by body names when saving an XML. Bug introduced in 3.1.4.
+23. Fixed bug in Python binding of :ref:`mj_saveModel`: ``buffer`` argument was documented as optional but was actually
     not optional.
-23. Fixed bug that prevented memory allocations larger than 2.15 GB.
+24. Fixed bug that prevented memory allocations larger than 2.15 GB.
 
 
 Version 3.1.4 (April 10th, 2024)

mjx/training_apg.ipynb

@@ -1,6 +1,39 @@
 {
  "cells": [
   {
+    "cell_type": "markdown",
+    "metadata": {
+     "id": "MpkYHwCqk7W-"
+    },
+    "source": [
+     "![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\n",
+     "\n",
+     "# <h1><center>Tutorial  <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/training_apg.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/></a></center></h1>\n",
+     "\n",
+     "This notebook provides a tutorial for differentiable physics for policy learning in [**MuJoCo XLA (MJX)**](https://github.com/google-deepmind/mujoco/blob/main/mjx), a JAX-based implementation of MuJoCo.\n",
+     "\n",
+     "**A Colab runtime with GPU acceleration is required.** If you're using a CPU-only runtime, you can switch using the menu \"Runtime > Change runtime type\".\n",
+     "\n",
+     "\n",
+     "This notebook was written by [Jing Yuan Luo](https://github.com/Andrew-Luo1).\n",
+     "\n",
+     "\u003c!-- Copyright 2021 DeepMind Technologies Limited\n",
+     "\n",
+     "     Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+     "     you may not use this file except in compliance with the License.\n",
+     "     You may obtain a copy of the License at\n",
+     "\n",
+     "         http://www.apache.org/licenses/LICENSE-2.0\n",
+     "\n",
+     "     Unless required by applicable law or agreed to in writing, software\n",
+     "     distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+     "     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+     "     See the License for the specific language governing permissions and\n",
+     "     limitations under the License.\n",
+     "--\u003e"
+    ]
+   },
+   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [