diff --git a/mjx/training_apg.ipynb b/mjx/training_apg.ipynb
index 26879392..4a7192dc 100644
--- a/mjx/training_apg.ipynb
+++ b/mjx/training_apg.ipynb
@@ -1,39 +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": {
+    "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",
+    "<!-- 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",
+    "-->"
+   ]
+  },
+  {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -106,12 +106,97 @@
     "## Setup: Imports and installations"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install MuJoCo, MJX, and Brax\n",
+    "!pip install mujoco\n",
+    "!pip install mujoco_mjx\n",
+    "!pip install brax"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#@title Check if MuJoCo installation was successful\n",
+    "\n",
+    "# Set up GPU rendering.\n",
+    "from google.colab import files\n",
+    "import distutils.util\n",
+    "import os\n",
+    "import subprocess\n",
+    "if subprocess.run('nvidia-smi').returncode:\n",
+    "  raise RuntimeError(\n",
+    "      'Cannot communicate with GPU. '\n",
+    "      'Make sure you are using a GPU Colab runtime. '\n",
+    "      'Go to the Runtime menu and select Choose runtime type.')\n",
+    "\n",
+    "# Add an ICD config so that glvnd can pick up the Nvidia EGL driver.\n",
+    "# This is usually installed as part of an Nvidia driver package, but the Colab\n",
+    "# kernel doesn't install its driver via APT, and as a result the ICD is missing.\n",
+    "# (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)\n",
+    "NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json'\n",
+    "if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):\n",
+    "  with open(NVIDIA_ICD_CONFIG_PATH, 'w') as f:\n",
+    "    f.write(\"\"\"{\n",
+    "    \"file_format_version\" : \"1.0.0\",\n",
+    "    \"ICD\" : {\n",
+    "        \"library_path\" : \"libEGL_nvidia.so.0\"\n",
+    "    }\n",
+    "}\n",
+    "\"\"\")\n",
+    "\n",
+    "# Configure MuJoCo to use the EGL rendering backend (requires GPU)\n",
+    "print('Setting environment variable to use GPU rendering:')\n",
+    "%env MUJOCO_GL=egl\n",
+    "\n",
+    "# Check if installation was succesful.\n",
+    "try:\n",
+    "  print('Checking that the installation succeeded:')\n",
+    "  import mujoco\n",
+    "  mujoco.MjModel.from_xml_string('<mujoco/>')\n",
+    "except Exception as e:\n",
+    "  raise e from RuntimeError(\n",
+    "      'Something went wrong during installation. Check the shell output above '\n",
+    "      'for more information.\\n'\n",
+    "      'If using a hosted Colab runtime, make sure you enable GPU acceleration '\n",
+    "      'by going to the Runtime menu and selecting \"Choose runtime type\".')\n",
+    "\n",
+    "print('Installation successful.')\n",
+    "\n",
+    "# Other imports and helper functions\n",
+    "import time\n",
+    "import itertools\n",
+    "import numpy as np\n",
+    "\n",
+    "# Graphics and plotting.\n",
+    "print('Installing mediapy:')\n",
+    "!command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n",
+    "!pip install -q mediapy\n",
+    "import mediapy as media\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# More legible printing from numpy.\n",
+    "np.set_printoptions(precision=3, suppress=True, linewidth=100)\n",
+    "\n",
+    "from IPython.display import clear_output\n",
+    "clear_output()"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
+    "#@title Import MuJoCo, MJX, and Brax\n",
+    "\n",
     "import os\n",
     "os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"] = \"0.8\" # 0.9 causes too much lag. \n",
     "from datetime import datetime\n",
@@ -162,7 +247,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 0,
    "metadata": {},
    "outputs": [],
    "source": [