Upload 4 files

init_mujoco.py

@@ -0,0 +1,4 @@
+# Import The Playground (this will clone Mujoco assets)
+
+from mujoco_playground import wrapper
+from mujoco_playground import registry

samples/locomotion.ipynb

@@ -0,0 +1,1758 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "MpkYHwCqk7W-"
+   },
+   "source": [
+    "![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "xBSdkbmGN2K-"
+   },
+   "source": [
+    "### Copyright notice"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "_UbO9uhtBSX5"
+   },
+   "source": [
+    "> <p><small><small>Copyright 2025 DeepMind Technologies Limited.</small></p>\n",
+    "> <p><small><small>Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at <a href=\"http://www.apache.org/licenses/LICENSE-2.0\">http://www.apache.org/licenses/LICENSE-2.0</a>.</small></small></p>\n",
+    "> <p><small><small>Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.</small></small></p>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "dNIJkb_FM2Ux"
+   },
+   "source": [
+    "# Locomotion in The Playground! <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/></a>\n",
+    "\n",
+    "In this notebook, we'll walk through a few locomotion environments available in MuJoCo Playground.\n",
+    "\n",
+    "**A Hugging Face Space with GPU acceleration is required.**\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "cellView": "form",
+    "id": "Xqo7pyX-n72M"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Collecting jax[cuda12]\n",
+      "  Downloading jax-0.6.2-py3-none-any.whl.metadata (13 kB)\n",
+      "Collecting jaxlib<=0.6.2,>=0.6.2 (from jax[cuda12])\n",
+      "  Downloading jaxlib-0.6.2-cp313-cp313-manylinux2014_x86_64.whl.metadata (1.3 kB)\n",
+      "Collecting ml_dtypes>=0.5.0 (from jax[cuda12])\n",
+      "  Downloading ml_dtypes-0.5.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (21 kB)\n",
+      "Collecting numpy>=1.26 (from jax[cuda12])\n",
+      "  Downloading numpy-2.3.1-cp313-cp313-manylinux_2_28_x86_64.whl.metadata (62 kB)\n",
+      "Collecting opt_einsum (from jax[cuda12])\n",
+      "  Downloading opt_einsum-3.4.0-py3-none-any.whl.metadata (6.3 kB)\n",
+      "Collecting scipy>=1.12 (from jax[cuda12])\n",
+      "  Downloading scipy-1.16.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (61 kB)\n",
+      "Collecting jax-cuda12-plugin<=0.6.2,>=0.6.2 (from jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading jax_cuda12_plugin-0.6.2-cp313-cp313-manylinux2014_x86_64.whl.metadata (1.7 kB)\n",
+      "Collecting jax-cuda12-pjrt==0.6.2 (from jax-cuda12-plugin<=0.6.2,>=0.6.2->jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading jax_cuda12_pjrt-0.6.2-py3-none-manylinux2014_x86_64.whl.metadata (579 bytes)\n",
+      "Collecting nvidia-cublas-cu12>=12.1.3.1 (from jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading nvidia_cublas_cu12-12.9.1.4-py3-none-manylinux_2_27_x86_64.whl.metadata (1.7 kB)\n",
+      "Collecting nvidia-cuda-cupti-cu12>=12.1.105 (from jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_x86_64.whl.metadata (1.8 kB)\n",
+      "Collecting nvidia-cuda-nvcc-cu12>=12.6.85 (from jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl.metadata (1.7 kB)\n",
+      "Collecting nvidia-cuda-runtime-cu12>=12.1.105 (from jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.7 kB)\n",
+      "Collecting nvidia-cudnn-cu12<10.0,>=9.8 (from jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl.metadata (1.8 kB)\n",
+      "Collecting nvidia-cufft-cu12>=11.0.2.54 (from jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.8 kB)\n",
+      "Collecting nvidia-cusolver-cu12>=11.4.5.107 (from jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_x86_64.whl.metadata (1.9 kB)\n",
+      "Collecting nvidia-cusparse-cu12>=12.1.0.106 (from jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.8 kB)\n",
+      "Collecting nvidia-nccl-cu12>=2.18.1 (from jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (2.0 kB)\n",
+      "Collecting nvidia-nvjitlink-cu12>=12.1.105 (from jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl.metadata (1.7 kB)\n",
+      "Collecting nvidia-cuda-nvrtc-cu12>=12.1.55 (from jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl.metadata (1.7 kB)\n",
+      "Collecting nvidia-nvshmem-cu12>=3.2.5 (from jax-cuda12-plugin[with-cuda]<=0.6.2,>=0.6.2; extra == \"cuda12\"->jax[cuda12])\n",
+      "  Downloading nvidia_nvshmem_cu12-3.3.9-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (2.1 kB)\n",
+      "Downloading jax-0.6.2-py3-none-any.whl (2.7 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.7/2.7 MB\u001b[0m \u001b[31m161.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading jax_cuda12_plugin-0.6.2-cp313-cp313-manylinux2014_x86_64.whl (15.9 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m15.9/15.9 MB\u001b[0m \u001b[31m167.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading jax_cuda12_pjrt-0.6.2-py3-none-manylinux2014_x86_64.whl (125.3 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m125.3/125.3 MB\u001b[0m \u001b[31m246.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
+      "\u001b[?25hDownloading jaxlib-0.6.2-cp313-cp313-manylinux2014_x86_64.whl (89.9 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m89.9/89.9 MB\u001b[0m \u001b[31m221.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
+      "\u001b[?25hDownloading nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl (706.8 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m706.8/706.8 MB\u001b[0m \u001b[31m71.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
+      "\u001b[?25hDownloading ml_dtypes-0.5.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.7 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.7/4.7 MB\u001b[0m \u001b[31m284.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading numpy-2.3.1-cp313-cp313-manylinux_2_28_x86_64.whl (16.6 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m16.6/16.6 MB\u001b[0m \u001b[31m213.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading nvidia_cublas_cu12-12.9.1.4-py3-none-manylinux_2_27_x86_64.whl (581.2 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m581.2/581.2 MB\u001b[0m \u001b[31m87.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
+      "\u001b[?25hDownloading nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_x86_64.whl (10.8 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m10.8/10.8 MB\u001b[0m \u001b[31m259.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl (40.5 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━���━━━━━━━━━━━━━━\u001b[0m \u001b[32m40.5/40.5 MB\u001b[0m \u001b[31m236.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl (89.6 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m89.6/89.6 MB\u001b[0m \u001b[31m226.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
+      "\u001b[?25hDownloading nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (3.5 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.5/3.5 MB\u001b[0m \u001b[31m220.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (200.9 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m200.9/200.9 MB\u001b[0m \u001b[31m199.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
+      "\u001b[?25hDownloading nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_x86_64.whl (338.1 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m338.1/338.1 MB\u001b[0m \u001b[31m155.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
+      "\u001b[?25hDownloading nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (366.5 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m366.5/366.5 MB\u001b[0m \u001b[31m144.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
+      "\u001b[?25hDownloading nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (322.3 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m322.3/322.3 MB\u001b[0m \u001b[31m169.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
+      "\u001b[?25hDownloading nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl (39.7 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m39.7/39.7 MB\u001b[0m \u001b[31m267.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading nvidia_nvshmem_cu12-3.3.9-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (124.6 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m124.6/124.6 MB\u001b[0m \u001b[31m198.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
+      "\u001b[?25hDownloading scipy-1.16.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (35.1 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m35.1/35.1 MB\u001b[0m \u001b[31m203.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading opt_einsum-3.4.0-py3-none-any.whl (71 kB)\n",
+      "Installing collected packages: jax-cuda12-pjrt, opt_einsum, nvidia-nvshmem-cu12, nvidia-nvjitlink-cu12, nvidia-nccl-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-nvcc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, numpy, jax-cuda12-plugin, scipy, nvidia-cusparse-cu12, nvidia-cufft-cu12, nvidia-cudnn-cu12, ml_dtypes, nvidia-cusolver-cu12, jaxlib, jax\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m20/20\u001b[0m [jax]32m19/20\u001b[0m [jax]ib]cusolver-cu12]2]2]\n",
+      "\u001b[1A\u001b[2KSuccessfully installed jax-0.6.2 jax-cuda12-pjrt-0.6.2 jax-cuda12-plugin-0.6.2 jaxlib-0.6.2 ml_dtypes-0.5.1 numpy-2.3.1 nvidia-cublas-cu12-12.9.1.4 nvidia-cuda-cupti-cu12-12.9.79 nvidia-cuda-nvcc-cu12-12.9.86 nvidia-cuda-nvrtc-cu12-12.9.86 nvidia-cuda-runtime-cu12-12.9.79 nvidia-cudnn-cu12-9.10.2.21 nvidia-cufft-cu12-11.4.1.4 nvidia-cusolver-cu12-11.7.5.82 nvidia-cusparse-cu12-12.5.10.65 nvidia-nccl-cu12-2.27.5 nvidia-nvjitlink-cu12-12.9.86 nvidia-nvshmem-cu12-3.3.9 opt_einsum-3.4.0 scipy-1.16.0\n",
+      "Collecting mujoco\n",
+      "  Downloading mujoco-3.3.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (44 kB)\n",
+      "Collecting absl-py (from mujoco)\n",
+      "  Downloading absl_py-2.3.0-py3-none-any.whl.metadata (2.4 kB)\n",
+      "Collecting etils[epath] (from mujoco)\n",
+      "  Downloading etils-1.12.2-py3-none-any.whl.metadata (6.5 kB)\n",
+      "Collecting glfw (from mujoco)\n",
+      "  Downloading glfw-2.9.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-manylinux_2_28_x86_64.whl.metadata (5.4 kB)\n",
+      "Requirement already satisfied: numpy in /home/user/miniconda/lib/python3.13/site-packages (from mujoco) (2.3.1)\n",
+      "Collecting pyopengl (from mujoco)\n",
+      "  Downloading PyOpenGL-3.1.9-py3-none-any.whl.metadata (3.3 kB)\n",
+      "Collecting fsspec (from etils[epath]->mujoco)\n",
+      "  Downloading fsspec-2025.5.1-py3-none-any.whl.metadata (11 kB)\n",
+      "Collecting importlib_resources (from etils[epath]->mujoco)\n",
+      "  Downloading importlib_resources-6.5.2-py3-none-any.whl.metadata (3.9 kB)\n",
+      "Requirement already satisfied: typing_extensions in /home/user/miniconda/lib/python3.13/site-packages (from etils[epath]->mujoco) (4.12.2)\n",
+      "Collecting zipp (from etils[epath]->mujoco)\n",
+      "  Downloading zipp-3.23.0-py3-none-any.whl.metadata (3.6 kB)\n",
+      "Downloading mujoco-3.3.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (6.6 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.6/6.6 MB\u001b[0m \u001b[31m178.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading absl_py-2.3.0-py3-none-any.whl (135 kB)\n",
+      "Downloading etils-1.12.2-py3-none-any.whl (167 kB)\n",
+      "Downloading fsspec-2025.5.1-py3-none-any.whl (199 kB)\n",
+      "Downloading glfw-2.9.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-manylinux_2_28_x86_64.whl (243 kB)\n",
+      "Downloading importlib_resources-6.5.2-py3-none-any.whl (37 kB)\n",
+      "Downloading PyOpenGL-3.1.9-py3-none-any.whl (3.2 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.2/3.2 MB\u001b[0m \u001b[31m228.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading zipp-3.23.0-py3-none-any.whl (10 kB)\n",
+      "Installing collected packages: pyopengl, glfw, zipp, importlib_resources, fsspec, etils, absl-py, mujoco\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8/8\u001b[0m [mujoco]2m7/8\u001b[0m [mujoco]]\n",
+      "\u001b[1A\u001b[2KSuccessfully installed absl-py-2.3.0 etils-1.12.2 fsspec-2025.5.1 glfw-2.9.0 importlib_resources-6.5.2 mujoco-3.3.3 pyopengl-3.1.9 zipp-3.23.0\n",
+      "Collecting mujoco_mjx\n",
+      "  Downloading mujoco_mjx-3.3.3-py3-none-any.whl.metadata (3.4 kB)\n",
+      "Requirement already satisfied: absl-py in /home/user/miniconda/lib/python3.13/site-packages (from mujoco_mjx) (2.3.0)\n",
+      "Requirement already satisfied: etils[epath] in /home/user/miniconda/lib/python3.13/site-packages (from mujoco_mjx) (1.12.2)\n",
+      "Requirement already satisfied: jax in /home/user/miniconda/lib/python3.13/site-packages (from mujoco_mjx) (0.6.2)\n",
+      "Requirement already satisfied: jaxlib in /home/user/miniconda/lib/python3.13/site-packages (from mujoco_mjx) (0.6.2)\n",
+      "Requirement already satisfied: mujoco>=3.3.3.dev0 in /home/user/miniconda/lib/python3.13/site-packages (from mujoco_mjx) (3.3.3)\n",
+      "Requirement already satisfied: scipy in /home/user/miniconda/lib/python3.13/site-packages (from mujoco_mjx) (1.16.0)\n",
+      "Collecting trimesh (from mujoco_mjx)\n",
+      "  Downloading trimesh-4.6.13-py3-none-any.whl.metadata (18 kB)\n",
+      "Requirement already satisfied: glfw in /home/user/miniconda/lib/python3.13/site-packages (from mujoco>=3.3.3.dev0->mujoco_mjx) (2.9.0)\n",
+      "Requirement already satisfied: numpy in /home/user/miniconda/lib/python3.13/site-packages (from mujoco>=3.3.3.dev0->mujoco_mjx) (2.3.1)\n",
+      "Requirement already satisfied: pyopengl in /home/user/miniconda/lib/python3.13/site-packages (from mujoco>=3.3.3.dev0->mujoco_mjx) (3.1.9)\n",
+      "Requirement already satisfied: fsspec in /home/user/miniconda/lib/python3.13/site-packages (from etils[epath]->mujoco_mjx) (2025.5.1)\n",
+      "Requirement already satisfied: importlib_resources in /home/user/miniconda/lib/python3.13/site-packages (from etils[epath]->mujoco_mjx) (6.5.2)\n",
+      "Requirement already satisfied: typing_extensions in /home/user/miniconda/lib/python3.13/site-packages (from etils[epath]->mujoco_mjx) (4.12.2)\n",
+      "Requirement already satisfied: zipp in /home/user/miniconda/lib/python3.13/site-packages (from etils[epath]->mujoco_mjx) (3.23.0)\n",
+      "Requirement already satisfied: ml_dtypes>=0.5.0 in /home/user/miniconda/lib/python3.13/site-packages (from jax->mujoco_mjx) (0.5.1)\n",
+      "Requirement already satisfied: opt_einsum in /home/user/miniconda/lib/python3.13/site-packages (from jax->mujoco_mjx) (3.4.0)\n",
+      "Downloading mujoco_mjx-3.3.3-py3-none-any.whl (6.7 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.7/6.7 MB\u001b[0m \u001b[31m138.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading trimesh-4.6.13-py3-none-any.whl (712 kB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m712.4/712.4 kB\u001b[0m \u001b[31m124.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hInstalling collected packages: trimesh, mujoco_mjx\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2/2\u001b[0m [mujoco_mjx]2\u001b[0m [mujoco_mjx]\n",
+      "\u001b[1A\u001b[2KSuccessfully installed mujoco_mjx-3.3.3 trimesh-4.6.13\n",
+      "Collecting brax\n",
+      "  Downloading brax-0.12.4-py3-none-any.whl.metadata (20 kB)\n",
+      "Requirement already satisfied: absl-py in /home/user/miniconda/lib/python3.13/site-packages (from brax) (2.3.0)\n",
+      "Requirement already satisfied: etils in /home/user/miniconda/lib/python3.13/site-packages (from brax) (1.12.2)\n",
+      "Collecting flask (from brax)\n",
+      "  Downloading flask-3.1.1-py3-none-any.whl.metadata (3.0 kB)\n",
+      "Collecting flask-cors (from brax)\n",
+      "  Downloading flask_cors-6.0.1-py3-none-any.whl.metadata (5.3 kB)\n",
+      "Collecting flax (from brax)\n",
+      "  Downloading flax-0.10.6-py3-none-any.whl.metadata (11 kB)\n",
+      "Requirement already satisfied: jax>=0.4.6 in /home/user/miniconda/lib/python3.13/site-packages (from brax) (0.6.2)\n",
+      "Requirement already satisfied: jaxlib>=0.4.6 in /home/user/miniconda/lib/python3.13/site-packages (from brax) (0.6.2)\n",
+      "Collecting jaxopt (from brax)\n",
+      "  Downloading jaxopt-0.8.5-py3-none-any.whl.metadata (3.3 kB)\n",
+      "Requirement already satisfied: jinja2 in /home/user/miniconda/lib/python3.13/site-packages (from brax) (3.1.6)\n",
+      "Collecting ml-collections (from brax)\n",
+      "  Downloading ml_collections-1.1.0-py3-none-any.whl.metadata (22 kB)\n",
+      "Requirement already satisfied: mujoco in /home/user/miniconda/lib/python3.13/site-packages (from brax) (3.3.3)\n",
+      "Requirement already satisfied: mujoco-mjx in /home/user/miniconda/lib/python3.13/site-packages (from brax) (3.3.3)\n",
+      "Requirement already satisfied: numpy in /home/user/miniconda/lib/python3.13/site-packages (from brax) (2.3.1)\n",
+      "Collecting optax (from brax)\n",
+      "  Downloading optax-0.2.5-py3-none-any.whl.metadata (7.5 kB)\n",
+      "Collecting orbax-checkpoint (from brax)\n",
+      "  Downloading orbax_checkpoint-0.11.17-py3-none-any.whl.metadata (2.2 kB)\n",
+      "Collecting pillow (from brax)\n",
+      "  Downloading pillow-11.3.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (9.0 kB)\n",
+      "Requirement already satisfied: scipy in /home/user/miniconda/lib/python3.13/site-packages (from brax) (1.16.0)\n",
+      "Collecting tensorboardx (from brax)\n",
+      "  Downloading tensorboardx-2.6.4-py3-none-any.whl.metadata (6.2 kB)\n",
+      "Requirement already satisfied: trimesh in /home/user/miniconda/lib/python3.13/site-packages (from brax) (4.6.13)\n",
+      "Requirement already satisfied: typing-extensions in /home/user/miniconda/lib/python3.13/site-packages (from brax) (4.12.2)\n",
+      "Requirement already satisfied: ml_dtypes>=0.5.0 in /home/user/miniconda/lib/python3.13/site-packages (from jax>=0.4.6->brax) (0.5.1)\n",
+      "Requirement already satisfied: opt_einsum in /home/user/miniconda/lib/python3.13/site-packages (from jax>=0.4.6->brax) (3.4.0)\n",
+      "Collecting blinker>=1.9.0 (from flask->brax)\n",
+      "  Downloading blinker-1.9.0-py3-none-any.whl.metadata (1.6 kB)\n",
+      "Collecting click>=8.1.3 (from flask->brax)\n",
+      "  Downloading click-8.2.1-py3-none-any.whl.metadata (2.5 kB)\n",
+      "Collecting itsdangerous>=2.2.0 (from flask->brax)\n",
+      "  Downloading itsdangerous-2.2.0-py3-none-any.whl.metadata (1.9 kB)\n",
+      "Requirement already satisfied: markupsafe>=2.1.1 in /home/user/miniconda/lib/python3.13/site-packages (from flask->brax) (3.0.2)\n",
+      "Collecting werkzeug>=3.1.0 (from flask->brax)\n",
+      "  Downloading werkzeug-3.1.3-py3-none-any.whl.metadata (3.7 kB)\n",
+      "Collecting msgpack (from flax->brax)\n",
+      "  Downloading msgpack-1.1.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (8.4 kB)\n",
+      "Collecting tensorstore (from flax->brax)\n",
+      "  Downloading tensorstore-0.1.75-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (21 kB)\n",
+      "Requirement already satisfied: rich>=11.1 in /home/user/miniconda/lib/python3.13/site-packages (from flax->brax) (13.9.4)\n",
+      "Requirement already satisfied: PyYAML>=5.4.1 in /home/user/miniconda/lib/python3.13/site-packages (from flax->brax) (6.0.2)\n",
+      "Collecting treescope>=0.1.7 (from flax->brax)\n",
+      "  Downloading treescope-0.1.9-py3-none-any.whl.metadata (6.6 kB)\n",
+      "Requirement already satisfied: markdown-it-py>=2.2.0 in /home/user/miniconda/lib/python3.13/site-packages (from rich>=11.1->flax->brax) (2.2.0)\n",
+      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /home/user/miniconda/lib/python3.13/site-packages (from rich>=11.1->flax->brax) (2.19.1)\n",
+      "Requirement already satisfied: mdurl~=0.1 in /home/user/miniconda/lib/python3.13/site-packages (from markdown-it-py>=2.2.0->rich>=11.1->flax->brax) (0.1.0)\n",
+      "Requirement already satisfied: glfw in /home/user/miniconda/lib/python3.13/site-packages (from mujoco->brax) (2.9.0)\n",
+      "Requirement already satisfied: pyopengl in /home/user/miniconda/lib/python3.13/site-packages (from mujoco->brax) (3.1.9)\n",
+      "Requirement already satisfied: fsspec in /home/user/miniconda/lib/python3.13/site-packages (from etils[epath]->mujoco->brax) (2025.5.1)\n",
+      "Requirement already satisfied: importlib_resources in /home/user/miniconda/lib/python3.13/site-packages (from etils[epath]->mujoco->brax) (6.5.2)\n",
+      "Requirement already satisfied: zipp in /home/user/miniconda/lib/python3.13/site-packages (from etils[epath]->mujoco->brax) (3.23.0)\n",
+      "Collecting chex>=0.1.87 (from optax->brax)\n",
+      "  Downloading chex-0.1.89-py3-none-any.whl.metadata (17 kB)\n",
+      "Requirement already satisfied: setuptools in /home/user/miniconda/lib/python3.13/site-packages (from chex>=0.1.87->optax->brax) (78.1.1)\n",
+      "Collecting toolz>=0.9.0 (from chex>=0.1.87->optax->brax)\n",
+      "  Downloading toolz-1.0.0-py3-none-any.whl.metadata (5.1 kB)\n",
+      "Requirement already satisfied: nest_asyncio in /home/user/miniconda/lib/python3.13/site-packages (from orbax-checkpoint->brax) (1.6.0)\n",
+      "Collecting protobuf (from orbax-checkpoint->brax)\n",
+      "  Downloading protobuf-6.31.1-cp39-abi3-manylinux2014_x86_64.whl.metadata (593 bytes)\n",
+      "Collecting humanize (from orbax-checkpoint->brax)\n",
+      "  Downloading humanize-4.12.3-py3-none-any.whl.metadata (7.8 kB)\n",
+      "Collecting simplejson>=3.16.0 (from orbax-checkpoint->brax)\n",
+      "  Downloading simplejson-3.20.1-cp313-cp313-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (3.3 kB)\n",
+      "Requirement already satisfied: packaging in /home/user/miniconda/lib/python3.13/site-packages (from tensorboardx->brax) (24.2)\n",
+      "Downloading brax-0.12.4-py3-none-any.whl (341 kB)\n",
+      "Downloading flask-3.1.1-py3-none-any.whl (103 kB)\n",
+      "Downloading blinker-1.9.0-py3-none-any.whl (8.5 kB)\n",
+      "Downloading click-8.2.1-py3-none-any.whl (102 kB)\n",
+      "Downloading itsdangerous-2.2.0-py3-none-any.whl (16 kB)\n",
+      "Downloading werkzeug-3.1.3-py3-none-any.whl (224 kB)\n",
+      "Downloading flask_cors-6.0.1-py3-none-any.whl (13 kB)\n",
+      "Downloading flax-0.10.6-py3-none-any.whl (447 kB)\n",
+      "Downloading treescope-0.1.9-py3-none-any.whl (182 kB)\n",
+      "Downloading jaxopt-0.8.5-py3-none-any.whl (172 kB)\n",
+      "Downloading ml_collections-1.1.0-py3-none-any.whl (76 kB)\n",
+      "Downloading msgpack-1.1.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (423 kB)\n",
+      "Downloading optax-0.2.5-py3-none-any.whl (354 kB)\n",
+      "Downloading chex-0.1.89-py3-none-any.whl (99 kB)\n",
+      "Downloading toolz-1.0.0-py3-none-any.whl (56 kB)\n",
+      "Downloading orbax_checkpoint-0.11.17-py3-none-any.whl (479 kB)\n",
+      "Downloading simplejson-3.20.1-cp313-cp313-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (152 kB)\n",
+      "Downloading tensorstore-0.1.75-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (18.8 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m18.8/18.8 MB\u001b[0m \u001b[31m25.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
+      "\u001b[?25hDownloading humanize-4.12.3-py3-none-any.whl (128 kB)\n",
+      "Downloading pillow-11.3.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (6.6 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.6/6.6 MB\u001b[0m \u001b[31m40.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading protobuf-6.31.1-cp39-abi3-manylinux2014_x86_64.whl (321 kB)\n",
+      "Downloading tensorboardx-2.6.4-py3-none-any.whl (87 kB)\n",
+      "Installing collected packages: werkzeug, treescope, toolz, simplejson, protobuf, pillow, msgpack, ml-collections, itsdangerous, humanize, click, blinker, tensorstore, tensorboardx, flask, flask-cors, orbax-checkpoint, jaxopt, chex, optax, flax, brax\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m22/22\u001b[0m [brax]2m21/22\u001b[0m [brax]]]heckpoint]\n",
+      "\u001b[1A\u001b[2KSuccessfully installed blinker-1.9.0 brax-0.12.4 chex-0.1.89 click-8.2.1 flask-3.1.1 flask-cors-6.0.1 flax-0.10.6 humanize-4.12.3 itsdangerous-2.2.0 jaxopt-0.8.5 ml-collections-1.1.0 msgpack-1.1.1 optax-0.2.5 orbax-checkpoint-0.11.17 pillow-11.3.0 protobuf-6.31.1 simplejson-3.20.1 tensorboardx-2.6.4 tensorstore-0.1.75 toolz-1.0.0 treescope-0.1.9 werkzeug-3.1.3\n",
+      "Collecting mediapy\n",
+      "  Downloading mediapy-1.2.4-py3-none-any.whl.metadata (4.8 kB)\n",
+      "Requirement already satisfied: ipython in /home/user/miniconda/lib/python3.13/site-packages (from mediapy) (9.4.0)\n",
+      "Collecting matplotlib (from mediapy)\n",
+      "  Downloading matplotlib-3.10.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (11 kB)\n",
+      "Requirement already satisfied: numpy in /home/user/miniconda/lib/python3.13/site-packages (from mediapy) (2.3.1)\n",
+      "Requirement already satisfied: Pillow in /home/user/miniconda/lib/python3.13/site-packages (from mediapy) (11.3.0)\n",
+      "Requirement already satisfied: decorator in /home/user/miniconda/lib/python3.13/site-packages (from ipython->mediapy) (5.2.1)\n",
+      "Requirement already satisfied: ipython-pygments-lexers in /home/user/miniconda/lib/python3.13/site-packages (from ipython->mediapy) (1.1.1)\n",
+      "Requirement already satisfied: jedi>=0.16 in /home/user/miniconda/lib/python3.13/site-packages (from ipython->mediapy) (0.19.2)\n",
+      "Requirement already satisfied: matplotlib-inline in /home/user/miniconda/lib/python3.13/site-packages (from ipython->mediapy) (0.1.7)\n",
+      "Requirement already satisfied: pexpect>4.3 in /home/user/miniconda/lib/python3.13/site-packages (from ipython->mediapy) (4.9.0)\n",
+      "Requirement already satisfied: prompt_toolkit<3.1.0,>=3.0.41 in /home/user/miniconda/lib/python3.13/site-packages (from ipython->mediapy) (3.0.51)\n",
+      "Requirement already satisfied: pygments>=2.4.0 in /home/user/miniconda/lib/python3.13/site-packages (from ipython->mediapy) (2.19.1)\n",
+      "Requirement already satisfied: stack_data in /home/user/miniconda/lib/python3.13/site-packages (from ipython->mediapy) (0.6.3)\n",
+      "Requirement already satisfied: traitlets>=5.13.0 in /home/user/miniconda/lib/python3.13/site-packages (from ipython->mediapy) (5.14.3)\n",
+      "Requirement already satisfied: wcwidth in /home/user/miniconda/lib/python3.13/site-packages (from prompt_toolkit<3.1.0,>=3.0.41->ipython->mediapy) (0.2.13)\n",
+      "Requirement already satisfied: parso<0.9.0,>=0.8.4 in /home/user/miniconda/lib/python3.13/site-packages (from jedi>=0.16->ipython->mediapy) (0.8.4)\n",
+      "Requirement already satisfied: ptyprocess>=0.5 in /home/user/miniconda/lib/python3.13/site-packages (from pexpect>4.3->ipython->mediapy) (0.7.0)\n",
+      "Collecting contourpy>=1.0.1 (from matplotlib->mediapy)\n",
+      "  Downloading contourpy-1.3.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (5.5 kB)\n",
+      "Collecting cycler>=0.10 (from matplotlib->mediapy)\n",
+      "  Downloading cycler-0.12.1-py3-none-any.whl.metadata (3.8 kB)\n",
+      "Collecting fonttools>=4.22.0 (from matplotlib->mediapy)\n",
+      "  Downloading fonttools-4.58.4-cp313-cp313-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl.metadata (106 kB)\n",
+      "Collecting kiwisolver>=1.3.1 (from matplotlib->mediapy)\n",
+      "  Downloading kiwisolver-1.4.8-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (6.2 kB)\n",
+      "Requirement already satisfied: packaging>=20.0 in /home/user/miniconda/lib/python3.13/site-packages (from matplotlib->mediapy) (24.2)\n",
+      "Collecting pyparsing>=2.3.1 (from matplotlib->mediapy)\n",
+      "  Downloading pyparsing-3.2.3-py3-none-any.whl.metadata (5.0 kB)\n",
+      "Requirement already satisfied: python-dateutil>=2.7 in /home/user/miniconda/lib/python3.13/site-packages (from matplotlib->mediapy) (2.9.0.post0)\n",
+      "Requirement already satisfied: six>=1.5 in /home/user/miniconda/lib/python3.13/site-packages (from python-dateutil>=2.7->matplotlib->mediapy) (1.17.0)\n",
+      "Requirement already satisfied: executing>=1.2.0 in /home/user/miniconda/lib/python3.13/site-packages (from stack_data->ipython->mediapy) (2.2.0)\n",
+      "Requirement already satisfied: asttokens>=2.1.0 in /home/user/miniconda/lib/python3.13/site-packages (from stack_data->ipython->mediapy) (3.0.0)\n",
+      "Requirement already satisfied: pure-eval in /home/user/miniconda/lib/python3.13/site-packages (from stack_data->ipython->mediapy) (0.2.3)\n",
+      "Downloading mediapy-1.2.4-py3-none-any.whl (26 kB)\n",
+      "Downloading matplotlib-3.10.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (8.6 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8.6/8.6 MB\u001b[0m \u001b[31m49.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading contourpy-1.3.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (322 kB)\n",
+      "Downloading cycler-0.12.1-py3-none-any.whl (8.3 kB)\n",
+      "Downloading fonttools-4.58.4-cp313-cp313-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl (4.9 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.9/4.9 MB\u001b[0m \u001b[31m34.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading kiwisolver-1.4.8-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.5 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.5/1.5 MB\u001b[0m \u001b[31m13.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading pyparsing-3.2.3-py3-none-any.whl (111 kB)\n",
+      "Installing collected packages: pyparsing, kiwisolver, fonttools, cycler, contourpy, matplotlib, mediapy\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7/7\u001b[0m [mediapy]m6/7\u001b[0m [mediapy]ib]\n",
+      "\u001b[1A\u001b[2KSuccessfully installed contourpy-1.3.2 cycler-0.12.1 fonttools-4.58.4 kiwisolver-1.4.8 matplotlib-3.10.3 mediapy-1.2.4 pyparsing-3.2.3\n"
+     ]
+    }
+   ],
+   "source": [
+    "#@title Install pre-requisites\n",
+    "!pip install \"jax[cuda12]\"\n",
+    "!pip install mujoco\n",
+    "!pip install mujoco_mjx\n",
+    "!pip install brax\n",
+    "!pip install mediapy"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {
+    "cellView": "form",
+    "id": "IbZxYDxzoz5R"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Tue Jul  1 12:12:35 2025       \n",
+      "+-----------------------------------------------------------------------------------------+\n",
+      "| NVIDIA-SMI 570.158.01             Driver Version: 570.158.01     CUDA Version: 12.8     |\n",
+      "|-----------------------------------------+------------------------+----------------------+\n",
+      "| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |\n",
+      "| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |\n",
+      "|                                         |                        |               MIG M. |\n",
+      "|=========================================+========================+======================|\n",
+      "|   0  NVIDIA L40S                    On  |   00000000:30:00.0 Off |                    0 |\n",
+      "| N/A   38C    P8             36W /  350W |       0MiB /  46068MiB |      0%      Default |\n",
+      "|                                         |                        |                  N/A |\n",
+      "+-----------------------------------------+------------------------+----------------------+\n",
+      "                                                                                         \n",
+      "+-----------------------------------------------------------------------------------------+\n",
+      "| Processes:                                                                              |\n",
+      "|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |\n",
+      "|        ID   ID                                                               Usage      |\n",
+      "|=========================================================================================|\n",
+      "|  No running processes found                                                             |\n",
+      "+-----------------------------------------------------------------------------------------+\n",
+      "Setting environment variable to use GPU rendering:\n",
+      "env: MUJOCO_GL=egl\n",
+      "Checking that the installation succeeded:\n",
+      "Installation successful.\n"
+     ]
+    }
+   ],
+   "source": [
+    "# @title Check if MuJoCo installation was successful\n",
+    "\n",
+    "import distutils.util\n",
+    "import os\n",
+    "import subprocess\n",
+    "\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",
+    "\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",
+    "try:\n",
+    "  print('Checking that the installation succeeded:')\n",
+    "  import mujoco\n",
+    "\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",
+    "\n",
+    "print('Installation successful.')\n",
+    "\n",
+    "# Tell XLA to use Triton GEMM, this improves steps/sec by ~30% on some GPUs\n",
+    "xla_flags = os.environ.get('XLA_FLAGS', '')\n",
+    "xla_flags += ' --xla_gpu_triton_gemm_any=True'\n",
+    "os.environ['XLA_FLAGS'] = xla_flags"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "cellView": "form",
+    "id": "T5f4w3Kq2X14"
+   },
+   "outputs": [],
+   "source": [
+    "# @title Import packages for plotting and creating graphics\n",
+    "import json\n",
+    "import itertools\n",
+    "import time\n",
+    "from typing import Callable, List, NamedTuple, Optional, Union\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)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {
+    "cellView": "form",
+    "id": "ObF1UXrkb0Nd"
+   },
+   "outputs": [],
+   "source": [
+    "# @title Import MuJoCo, MJX, and Brax\n",
+    "from datetime import datetime\n",
+    "import functools\n",
+    "import os\n",
+    "from typing import Any, Dict, Sequence, Tuple, Union\n",
+    "from brax import base\n",
+    "from brax import envs\n",
+    "from brax import math\n",
+    "from brax.base import Base, Motion, Transform\n",
+    "from brax.base import State as PipelineState\n",
+    "from brax.envs.base import Env, PipelineEnv, State\n",
+    "from brax.io import html, mjcf, model\n",
+    "from brax.mjx.base import State as MjxState\n",
+    "from brax.training.agents.ppo import networks as ppo_networks\n",
+    "from brax.training.agents.ppo import train as ppo\n",
+    "from brax.training.agents.sac import networks as sac_networks\n",
+    "from brax.training.agents.sac import train as sac\n",
+    "from etils import epath\n",
+    "from flax import struct\n",
+    "from flax.training import orbax_utils\n",
+    "from IPython.display import HTML, clear_output\n",
+    "import jax\n",
+    "from jax import numpy as jp\n",
+    "from matplotlib import pyplot as plt\n",
+    "import mediapy as media\n",
+    "from ml_collections import config_dict\n",
+    "import mujoco\n",
+    "from mujoco import mjx\n",
+    "import numpy as np\n",
+    "from orbax import checkpoint as ocp"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {
+    "cellView": "form",
+    "id": "UoTLSx4cFRdy"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Collecting playground\n",
+      "  Downloading playground-0.0.5-py3-none-any.whl.metadata (8.7 kB)\n",
+      "Requirement already satisfied: brax>=0.12.1 in /home/user/miniconda/lib/python3.13/site-packages (from playground) (0.12.4)\n",
+      "Requirement already satisfied: etils in /home/user/miniconda/lib/python3.13/site-packages (from playground) (1.12.2)\n",
+      "Requirement already satisfied: flax in /home/user/miniconda/lib/python3.13/site-packages (from playground) (0.10.6)\n",
+      "Requirement already satisfied: jax in /home/user/miniconda/lib/python3.13/site-packages (from playground) (0.6.2)\n",
+      "Collecting lxml (from playground)\n",
+      "  Downloading lxml-6.0.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (6.6 kB)\n",
+      "Requirement already satisfied: ml-collections in /home/user/miniconda/lib/python3.13/site-packages (from playground) (1.1.0)\n",
+      "Requirement already satisfied: mujoco-mjx>=3.2.7 in /home/user/miniconda/lib/python3.13/site-packages (from playground) (3.3.3)\n",
+      "Requirement already satisfied: mujoco>=3.2.7 in /home/user/miniconda/lib/python3.13/site-packages (from playground) (3.3.3)\n",
+      "Requirement already satisfied: tqdm in /home/user/miniconda/lib/python3.13/site-packages (from playground) (4.67.1)\n",
+      "Requirement already satisfied: absl-py in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (2.3.0)\n",
+      "Requirement already satisfied: flask in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (3.1.1)\n",
+      "Requirement already satisfied: flask-cors in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (6.0.1)\n",
+      "Requirement already satisfied: jaxlib>=0.4.6 in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (0.6.2)\n",
+      "Requirement already satisfied: jaxopt in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (0.8.5)\n",
+      "Requirement already satisfied: jinja2 in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (3.1.6)\n",
+      "Requirement already satisfied: numpy in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (2.3.1)\n",
+      "Requirement already satisfied: optax in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (0.2.5)\n",
+      "Requirement already satisfied: orbax-checkpoint in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (0.11.17)\n",
+      "Requirement already satisfied: pillow in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (11.3.0)\n",
+      "Requirement already satisfied: scipy in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (1.16.0)\n",
+      "Requirement already satisfied: tensorboardx in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (2.6.4)\n",
+      "Requirement already satisfied: trimesh in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (4.6.13)\n",
+      "Requirement already satisfied: typing-extensions in /home/user/miniconda/lib/python3.13/site-packages (from brax>=0.12.1->playground) (4.12.2)\n",
+      "Requirement already satisfied: ml_dtypes>=0.5.0 in /home/user/miniconda/lib/python3.13/site-packages (from jax->playground) (0.5.1)\n",
+      "Requirement already satisfied: opt_einsum in /home/user/miniconda/lib/python3.13/site-packages (from jax->playground) (3.4.0)\n",
+      "Requirement already satisfied: glfw in /home/user/miniconda/lib/python3.13/site-packages (from mujoco>=3.2.7->playground) (2.9.0)\n",
+      "Requirement already satisfied: pyopengl in /home/user/miniconda/lib/python3.13/site-packages (from mujoco>=3.2.7->playground) (3.1.9)\n",
+      "Requirement already satisfied: fsspec in /home/user/miniconda/lib/python3.13/site-packages (from etils[epath]->mujoco>=3.2.7->playground) (2025.5.1)\n",
+      "Requirement already satisfied: importlib_resources in /home/user/miniconda/lib/python3.13/site-packages (from etils[epath]->mujoco>=3.2.7->playground) (6.5.2)\n",
+      "Requirement already satisfied: zipp in /home/user/miniconda/lib/python3.13/site-packages (from etils[epath]->mujoco>=3.2.7->playground) (3.23.0)\n",
+      "Requirement already satisfied: blinker>=1.9.0 in /home/user/miniconda/lib/python3.13/site-packages (from flask->brax>=0.12.1->playground) (1.9.0)\n",
+      "Requirement already satisfied: click>=8.1.3 in /home/user/miniconda/lib/python3.13/site-packages (from flask->brax>=0.12.1->playground) (8.2.1)\n",
+      "Requirement already satisfied: itsdangerous>=2.2.0 in /home/user/miniconda/lib/python3.13/site-packages (from flask->brax>=0.12.1->playground) (2.2.0)\n",
+      "Requirement already satisfied: markupsafe>=2.1.1 in /home/user/miniconda/lib/python3.13/site-packages (from flask->brax>=0.12.1->playground) (3.0.2)\n",
+      "Requirement already satisfied: werkzeug>=3.1.0 in /home/user/miniconda/lib/python3.13/site-packages (from flask->brax>=0.12.1->playground) (3.1.3)\n",
+      "Requirement already satisfied: msgpack in /home/user/miniconda/lib/python3.13/site-packages (from flax->playground) (1.1.1)\n",
+      "Requirement already satisfied: tensorstore in /home/user/miniconda/lib/python3.13/site-packages (from flax->playground) (0.1.75)\n",
+      "Requirement already satisfied: rich>=11.1 in /home/user/miniconda/lib/python3.13/site-packages (from flax->playground) (13.9.4)\n",
+      "Requirement already satisfied: PyYAML>=5.4.1 in /home/user/miniconda/lib/python3.13/site-packages (from flax->playground) (6.0.2)\n",
+      "Requirement already satisfied: treescope>=0.1.7 in /home/user/miniconda/lib/python3.13/site-packages (from flax->playground) (0.1.9)\n",
+      "Requirement already satisfied: markdown-it-py>=2.2.0 in /home/user/miniconda/lib/python3.13/site-packages (from rich>=11.1->flax->playground) (2.2.0)\n",
+      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /home/user/miniconda/lib/python3.13/site-packages (from rich>=11.1->flax->playground) (2.19.1)\n",
+      "Requirement already satisfied: mdurl~=0.1 in /home/user/miniconda/lib/python3.13/site-packages (from markdown-it-py>=2.2.0->rich>=11.1->flax->playground) (0.1.0)\n",
+      "Requirement already satisfied: chex>=0.1.87 in /home/user/miniconda/lib/python3.13/site-packages (from optax->brax>=0.12.1->playground) (0.1.89)\n",
+      "Requirement already satisfied: setuptools in /home/user/miniconda/lib/python3.13/site-packages (from chex>=0.1.87->optax->brax>=0.12.1->playground) (78.1.1)\n",
+      "Requirement already satisfied: toolz>=0.9.0 in /home/user/miniconda/lib/python3.13/site-packages (from chex>=0.1.87->optax->brax>=0.12.1->playground) (1.0.0)\n",
+      "Requirement already satisfied: nest_asyncio in /home/user/miniconda/lib/python3.13/site-packages (from orbax-checkpoint->brax>=0.12.1->playground) (1.6.0)\n",
+      "Requirement already satisfied: protobuf in /home/user/miniconda/lib/python3.13/site-packages (from orbax-checkpoint->brax>=0.12.1->playground) (6.31.1)\n",
+      "Requirement already satisfied: humanize in /home/user/miniconda/lib/python3.13/site-packages (from orbax-checkpoint->brax>=0.12.1->playground) (4.12.3)\n",
+      "Requirement already satisfied: simplejson>=3.16.0 in /home/user/miniconda/lib/python3.13/site-packages (from orbax-checkpoint->brax>=0.12.1->playground) (3.20.1)\n",
+      "Requirement already satisfied: packaging in /home/user/miniconda/lib/python3.13/site-packages (from tensorboardx->brax>=0.12.1->playground) (24.2)\n",
+      "Downloading playground-0.0.5-py3-none-any.whl (7.4 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.4/7.4 MB\u001b[0m \u001b[31m88.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hDownloading lxml-6.0.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (5.2 MB)\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.2/5.2 MB\u001b[0m \u001b[31m163.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[?25hInstalling collected packages: lxml, playground\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2/2\u001b[0m [playground]2\u001b[0m [playground]\n",
+      "\u001b[1A\u001b[2KSuccessfully installed lxml-6.0.0 playground-0.0.5\n"
+     ]
+    }
+   ],
+   "source": [
+    "#@title Install MuJoCo Playground\n",
+    "!pip install playground"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {
+    "cellView": "form",
+    "id": "gYm2h7m8w3Nv"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "mujoco_menagerie not found. Downloading...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Cloning mujoco_menagerie: ██████████| 100/100 [00:13<00:00]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Checking out commit 14ceccf557cc47240202f2354d684eca58ff8de4\n",
+      "Successfully downloaded mujoco_menagerie\n"
+     ]
+    }
+   ],
+   "source": [
+    "#@title Import The Playground\n",
+    "\n",
+    "from mujoco_playground import wrapper\n",
+    "from mujoco_playground import registry"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "LcibXbyKt4FI"
+   },
+   "source": [
+    "# Locomotion\n",
+    "\n",
+    "MuJoCo Playground contains a host of quadrupedal and bipedal environments (all listed below after running the command)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {
+    "id": "ox0Gze9Ct5AM"
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "('ApolloJoystickFlatTerrain',\n",
+       " 'BarkourJoystick',\n",
+       " 'BerkeleyHumanoidJoystickFlatTerrain',\n",
+       " 'BerkeleyHumanoidJoystickRoughTerrain',\n",
+       " 'G1JoystickFlatTerrain',\n",
+       " 'G1JoystickRoughTerrain',\n",
+       " 'Go1JoystickFlatTerrain',\n",
+       " 'Go1JoystickRoughTerrain',\n",
+       " 'Go1Getup',\n",
+       " 'Go1Handstand',\n",
+       " 'Go1Footstand',\n",
+       " 'H1InplaceGaitTracking',\n",
+       " 'H1JoystickGaitTracking',\n",
+       " 'Op3Joystick',\n",
+       " 'SpotFlatTerrainJoystick',\n",
+       " 'SpotGetup',\n",
+       " 'SpotJoystickGaitTracking',\n",
+       " 'T1JoystickFlatTerrain',\n",
+       " 'T1JoystickRoughTerrain')"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "registry.locomotion.ALL_ENVS"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "_R01tjWfI-i6"
+   },
+   "source": [
+    "# Quadrupedal\n",
+    "\n",
+    "Let's jump right into quadrupedal locomotion! While we have environments available for the Google Barkour and Boston Dynamics Spot robots, the Unitree Go1 environment contains the most trainable policies that were transferred onto the real robot. We'll go right ahead and show a few policies using the Unitree Go1!\n",
+    "\n",
+    "First, let's train a joystick policy, which tracks linear and yaw velocity commands."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {
+    "id": "kPJeoQeEJBSA"
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "WARNING:2025-07-01 11:54:03,909:jax._src.xla_bridge:794: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n",
+      "WARNING:jax._src.xla_bridge:An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n"
+     ]
+    }
+   ],
+   "source": [
+    "env_name = 'Go1JoystickFlatTerrain'\n",
+    "env = registry.load(env_name)\n",
+    "env_cfg = registry.get_default_config(env_name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "6n9UT9N1wR5K"
+   },
+   "outputs": [],
+   "source": [
+    "env_cfg"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Thm7nZueM4cz"
+   },
+   "source": [
+    "## Joystick\n",
+    "\n",
+    "Let's train the joystick policy and visualize rollouts:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "B9T_UVZYLDdM"
+   },
+   "outputs": [],
+   "source": [
+    "from mujoco_playground.config import locomotion_params\n",
+    "ppo_params = locomotion_params.brax_ppo_config(env_name)\n",
+    "ppo_params"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Aefr2OS01D9g"
+   },
+   "source": [
+    "Domain randomization was used to make the policy robust to sim-to-real transfer. Certain environments in the Playground have domain randomization functions implemented. They're available in the registry and can be passed directly to brax RL algorithms. The [domain randomization](https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/_src/locomotion/go1/randomize.py) function randomizes over friction, armature, center of mass of the torso, and link masses, amongst other simulation parameters."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "UVA4Bn681DZT"
+   },
+   "outputs": [],
+   "source": [
+    "registry.get_domain_randomizer(env_name)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "vBEEQyY6M5OC"
+   },
+   "source": [
+    "### Train\n",
+    "\n",
+    "The policy takes 7 minutes to train on an RTX 4090."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "XKFzyP7wM5OD"
+   },
+   "outputs": [],
+   "source": [
+    "x_data, y_data, y_dataerr = [], [], []\n",
+    "times = [datetime.now()]\n",
+    "\n",
+    "\n",
+    "def progress(num_steps, metrics):\n",
+    "  clear_output(wait=True)\n",
+    "\n",
+    "  times.append(datetime.now())\n",
+    "  x_data.append(num_steps)\n",
+    "  y_data.append(metrics[\"eval/episode_reward\"])\n",
+    "  y_dataerr.append(metrics[\"eval/episode_reward_std\"])\n",
+    "\n",
+    "  plt.xlim([0, ppo_params[\"num_timesteps\"] * 1.25])\n",
+    "  plt.xlabel(\"# environment steps\")\n",
+    "  plt.ylabel(\"reward per episode\")\n",
+    "  plt.title(f\"y={y_data[-1]:.3f}\")\n",
+    "  plt.errorbar(x_data, y_data, yerr=y_dataerr, color=\"blue\")\n",
+    "\n",
+    "  display(plt.gcf())\n",
+    "\n",
+    "randomizer = registry.get_domain_randomizer(env_name)\n",
+    "ppo_training_params = dict(ppo_params)\n",
+    "network_factory = ppo_networks.make_ppo_networks\n",
+    "if \"network_factory\" in ppo_params:\n",
+    "  del ppo_training_params[\"network_factory\"]\n",
+    "  network_factory = functools.partial(\n",
+    "      ppo_networks.make_ppo_networks,\n",
+    "      **ppo_params.network_factory\n",
+    "  )\n",
+    "\n",
+    "train_fn = functools.partial(\n",
+    "    ppo.train, **dict(ppo_training_params),\n",
+    "    network_factory=network_factory,\n",
+    "    randomization_fn=randomizer,\n",
+    "    progress_fn=progress\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "FGrlulWbM5OD"
+   },
+   "outputs": [],
+   "source": [
+    "make_inference_fn, params, metrics = train_fn(\n",
+    "    environment=env,\n",
+    "    eval_env=registry.load(env_name, config=env_cfg),\n",
+    "    wrap_env_fn=wrapper.wrap_for_brax_training,\n",
+    ")\n",
+    "print(f\"time to jit: {times[1] - times[0]}\")\n",
+    "print(f\"time to train: {times[-1] - times[1]}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "AUxSNhq3UqmC"
+   },
+   "source": [
+    "Let's rollout and render the resulting policy!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "RBM89g5A2Yoi"
+   },
+   "outputs": [],
+   "source": [
+    "# Enable perturbation in the eval env.\n",
+    "env_cfg = registry.get_default_config(env_name)\n",
+    "env_cfg.pert_config.enable = True\n",
+    "env_cfg.pert_config.velocity_kick = [3.0, 6.0]\n",
+    "env_cfg.pert_config.kick_wait_times = [5.0, 15.0]\n",
+    "env_cfg.command_config.a = [1.5, 0.8, 2*jp.pi]\n",
+    "eval_env = registry.load(env_name, config=env_cfg)\n",
+    "velocity_kick_range = [0.0, 0.0]  # Disable velocity kick.\n",
+    "kick_duration_range = [0.05, 0.2]\n",
+    "\n",
+    "jit_reset = jax.jit(eval_env.reset)\n",
+    "jit_step = jax.jit(eval_env.step)\n",
+    "jit_inference_fn = jax.jit(make_inference_fn(params, deterministic=True))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "cellView": "form",
+    "id": "C_1CY9xDoUKw"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Rollout and Render\n",
+    "from mujoco_playground._src.gait import draw_joystick_command\n",
+    "\n",
+    "x_vel = 0.0  #@param {type: \"number\"}\n",
+    "y_vel = 0.0  #@param {type: \"number\"}\n",
+    "yaw_vel = 3.14  #@param {type: \"number\"}\n",
+    "\n",
+    "\n",
+    "def sample_pert(rng):\n",
+    "  rng, key1, key2 = jax.random.split(rng, 3)\n",
+    "  pert_mag = jax.random.uniform(\n",
+    "      key1, minval=velocity_kick_range[0], maxval=velocity_kick_range[1]\n",
+    "  )\n",
+    "  duration_seconds = jax.random.uniform(\n",
+    "      key2, minval=kick_duration_range[0], maxval=kick_duration_range[1]\n",
+    "  )\n",
+    "  duration_steps = jp.round(duration_seconds / eval_env.dt).astype(jp.int32)\n",
+    "  state.info[\"pert_mag\"] = pert_mag\n",
+    "  state.info[\"pert_duration\"] = duration_steps\n",
+    "  state.info[\"pert_duration_seconds\"] = duration_seconds\n",
+    "  return rng\n",
+    "\n",
+    "\n",
+    "rng = jax.random.PRNGKey(0)\n",
+    "rollout = []\n",
+    "modify_scene_fns = []\n",
+    "\n",
+    "swing_peak = []\n",
+    "rewards = []\n",
+    "linvel = []\n",
+    "angvel = []\n",
+    "track = []\n",
+    "foot_vel = []\n",
+    "rews = []\n",
+    "contact = []\n",
+    "command = jp.array([x_vel, y_vel, yaw_vel])\n",
+    "\n",
+    "state = jit_reset(rng)\n",
+    "if state.info[\"steps_since_last_pert\"] < state.info[\"steps_until_next_pert\"]:\n",
+    "  rng = sample_pert(rng)\n",
+    "state.info[\"command\"] = command\n",
+    "for i in range(env_cfg.episode_length):\n",
+    "  if state.info[\"steps_since_last_pert\"] < state.info[\"steps_until_next_pert\"]:\n",
+    "    rng = sample_pert(rng)\n",
+    "  act_rng, rng = jax.random.split(rng)\n",
+    "  ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
+    "  state = jit_step(state, ctrl)\n",
+    "  state.info[\"command\"] = command\n",
+    "  rews.append(\n",
+    "      {k: v for k, v in state.metrics.items() if k.startswith(\"reward/\")}\n",
+    "  )\n",
+    "  rollout.append(state)\n",
+    "  swing_peak.append(state.info[\"swing_peak\"])\n",
+    "  rewards.append(\n",
+    "      {k[7:]: v for k, v in state.metrics.items() if k.startswith(\"reward/\")}\n",
+    "  )\n",
+    "  linvel.append(env.get_global_linvel(state.data))\n",
+    "  angvel.append(env.get_gyro(state.data))\n",
+    "  track.append(\n",
+    "      env._reward_tracking_lin_vel(\n",
+    "          state.info[\"command\"], env.get_local_linvel(state.data)\n",
+    "      )\n",
+    "  )\n",
+    "\n",
+    "  feet_vel = state.data.sensordata[env._foot_linvel_sensor_adr]\n",
+    "  vel_xy = feet_vel[..., :2]\n",
+    "  vel_norm = jp.sqrt(jp.linalg.norm(vel_xy, axis=-1))\n",
+    "  foot_vel.append(vel_norm)\n",
+    "\n",
+    "  contact.append(state.info[\"last_contact\"])\n",
+    "\n",
+    "  xyz = np.array(state.data.xpos[env._torso_body_id])\n",
+    "  xyz += np.array([0, 0, 0.2])\n",
+    "  x_axis = state.data.xmat[env._torso_body_id, 0]\n",
+    "  yaw = -np.arctan2(x_axis[1], x_axis[0])\n",
+    "  modify_scene_fns.append(\n",
+    "      functools.partial(\n",
+    "          draw_joystick_command,\n",
+    "          cmd=state.info[\"command\"],\n",
+    "          xyz=xyz,\n",
+    "          theta=yaw,\n",
+    "          scl=abs(state.info[\"command\"][0])\n",
+    "          / env_cfg.command_config.a[0],\n",
+    "      )\n",
+    "  )\n",
+    "\n",
+    "\n",
+    "render_every = 2\n",
+    "fps = 1.0 / eval_env.dt / render_every\n",
+    "traj = rollout[::render_every]\n",
+    "mod_fns = modify_scene_fns[::render_every]\n",
+    "\n",
+    "scene_option = mujoco.MjvOption()\n",
+    "scene_option.geomgroup[2] = True\n",
+    "scene_option.geomgroup[3] = False\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = False\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_PERTFORCE] = True\n",
+    "\n",
+    "frames = eval_env.render(\n",
+    "    traj,\n",
+    "    camera=\"track\",\n",
+    "    scene_option=scene_option,\n",
+    "    width=640,\n",
+    "    height=480,\n",
+    "    modify_scene_fns=mod_fns,\n",
+    ")\n",
+    "media.show_video(frames, fps=fps, loop=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "1QHdoJ2r30En"
+   },
+   "source": [
+    "Let's visualize the feet positions and the positional drift compared to the commanded linear and angular velocity. This is useful for debugging how well the policy follows the commands!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "cellView": "form",
+    "id": "gyyynm3ozEet"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Plot each foot in a 2x2 grid.\n",
+    "\n",
+    "swing_peak = jp.array(swing_peak)\n",
+    "names = [\"FR\", \"FL\", \"RR\", \"RL\"]\n",
+    "colors = [\"r\", \"g\", \"b\", \"y\"]\n",
+    "fig, axs = plt.subplots(2, 2)\n",
+    "for i, ax in enumerate(axs.flat):\n",
+    "  ax.plot(swing_peak[:, i], color=colors[i])\n",
+    "  ax.set_ylim([0, env_cfg.reward_config.max_foot_height * 1.25])\n",
+    "  ax.axhline(env_cfg.reward_config.max_foot_height, color=\"k\", linestyle=\"--\")\n",
+    "  ax.set_title(names[i])\n",
+    "  ax.set_xlabel(\"time\")\n",
+    "  ax.set_ylabel(\"height\")\n",
+    "plt.tight_layout()\n",
+    "plt.show()\n",
+    "\n",
+    "linvel_x = jp.array(linvel)[:, 0]\n",
+    "linvel_y = jp.array(linvel)[:, 1]\n",
+    "angvel_yaw = jp.array(angvel)[:, 2]\n",
+    "\n",
+    "# Plot whether velocity is within the command range.\n",
+    "linvel_x = jp.convolve(linvel_x, jp.ones(10) / 10, mode=\"same\")\n",
+    "linvel_y = jp.convolve(linvel_y, jp.ones(10) / 10, mode=\"same\")\n",
+    "angvel_yaw = jp.convolve(angvel_yaw, jp.ones(10) / 10, mode=\"same\")\n",
+    "\n",
+    "fig, axes = plt.subplots(3, 1, figsize=(10, 10))\n",
+    "axes[0].plot(linvel_x)\n",
+    "axes[1].plot(linvel_y)\n",
+    "axes[2].plot(angvel_yaw)\n",
+    "\n",
+    "axes[0].set_ylim(\n",
+    "    -env_cfg.command_config.a[0], env_cfg.command_config.a[0]\n",
+    ")\n",
+    "axes[1].set_ylim(\n",
+    "    -env_cfg.command_config.a[1], env_cfg.command_config.a[1]\n",
+    ")\n",
+    "axes[2].set_ylim(\n",
+    "    -env_cfg.command_config.a[2], env_cfg.command_config.a[2]\n",
+    ")\n",
+    "\n",
+    "for i, ax in enumerate(axes):\n",
+    "  ax.axhline(state.info[\"command\"][i], color=\"red\", linestyle=\"--\")\n",
+    "\n",
+    "labels = [\"dx\", \"dy\", \"dyaw\"]\n",
+    "for i, ax in enumerate(axes):\n",
+    "  ax.set_ylabel(labels[i])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "t1QAHuYBQBbl"
+   },
+   "source": [
+    "Now let's visualize what it looks like to slowly increase linear velocity commands."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "cellView": "form",
+    "id": "Q0EuQiVlzh5u"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Slowly increase linvel commands\n",
+    "\n",
+    "rng = jax.random.PRNGKey(0)\n",
+    "rollout = []\n",
+    "modify_scene_fns = []\n",
+    "swing_peak = []\n",
+    "linvel = []\n",
+    "angvel = []\n",
+    "\n",
+    "x = -0.25\n",
+    "command = jp.array([x, 0, 0])\n",
+    "state = jit_reset(rng)\n",
+    "for i in range(1_400):\n",
+    "  # Increase the forward velocity by 0.25 m/s every 200 steps.\n",
+    "  if i % 200 == 0:\n",
+    "    x += 0.25\n",
+    "    print(f\"Setting x to {x}\")\n",
+    "    command = jp.array([x, 0, 0])\n",
+    "  state.info[\"command\"] = command\n",
+    "  if state.info[\"steps_since_last_pert\"] < state.info[\"steps_until_next_pert\"]:\n",
+    "    rng = sample_pert(rng)\n",
+    "  act_rng, rng = jax.random.split(rng)\n",
+    "  ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
+    "  state = jit_step(state, ctrl)\n",
+    "  rollout.append(state)\n",
+    "  swing_peak.append(state.info[\"swing_peak\"])\n",
+    "  linvel.append(env.get_global_linvel(state.data))\n",
+    "  angvel.append(env.get_gyro(state.data))\n",
+    "  xyz = np.array(state.data.xpos[env._torso_body_id])\n",
+    "  xyz += np.array([0, 0, 0.2])\n",
+    "  x_axis = state.data.xmat[env._torso_body_id, 0]\n",
+    "  yaw = -np.arctan2(x_axis[1], x_axis[0])\n",
+    "  modify_scene_fns.append(\n",
+    "      functools.partial(\n",
+    "          draw_joystick_command,\n",
+    "          cmd=command,\n",
+    "          xyz=xyz,\n",
+    "          theta=yaw,\n",
+    "          scl=abs(command[0]) / env_cfg.command_config.a[0],\n",
+    "      )\n",
+    "  )\n",
+    "\n",
+    "\n",
+    "# Plot each foot in a 2x2 grid.\n",
+    "swing_peak = jp.array(swing_peak)\n",
+    "names = [\"FR\", \"FL\", \"RR\", \"RL\"]\n",
+    "colors = [\"r\", \"g\", \"b\", \"y\"]\n",
+    "fig, axs = plt.subplots(2, 2)\n",
+    "for i, ax in enumerate(axs.flat):\n",
+    "  ax.plot(swing_peak[:, i], color=colors[i])\n",
+    "  ax.set_ylim([0, env_cfg.reward_config.max_foot_height * 1.25])\n",
+    "  ax.axhline(env_cfg.reward_config.max_foot_height, color=\"k\", linestyle=\"--\")\n",
+    "  ax.set_title(names[i])\n",
+    "  ax.set_xlabel(\"time\")\n",
+    "  ax.set_ylabel(\"height\")\n",
+    "plt.tight_layout()\n",
+    "plt.show()\n",
+    "\n",
+    "linvel_x = jp.array(linvel)[:, 0]\n",
+    "linvel_y = jp.array(linvel)[:, 1]\n",
+    "angvel_yaw = jp.array(angvel)[:, 2]\n",
+    "\n",
+    "# Plot whether velocity is within the command range.\n",
+    "linvel_x = jp.convolve(linvel_x, jp.ones(10) / 10, mode=\"same\")\n",
+    "linvel_y = jp.convolve(linvel_y, jp.ones(10) / 10, mode=\"same\")\n",
+    "angvel_yaw = jp.convolve(angvel_yaw, jp.ones(10) / 10, mode=\"same\")\n",
+    "\n",
+    "fig, axes = plt.subplots(3, 1, figsize=(10, 10))\n",
+    "axes[0].plot(linvel_x)\n",
+    "axes[1].plot(linvel_y)\n",
+    "axes[2].plot(angvel_yaw)\n",
+    "\n",
+    "axes[0].set_ylim(\n",
+    "    -env_cfg.command_config.a[0], env_cfg.command_config.a[0]\n",
+    ")\n",
+    "axes[1].set_ylim(\n",
+    "    -env_cfg.command_config.a[1], env_cfg.command_config.a[1]\n",
+    ")\n",
+    "axes[2].set_ylim(\n",
+    "    -env_cfg.command_config.a[2], env_cfg.command_config.a[2]\n",
+    ")\n",
+    "\n",
+    "for i, ax in enumerate(axes):\n",
+    "  ax.axhline(state.info[\"command\"][i], color=\"red\", linestyle=\"--\")\n",
+    "\n",
+    "labels = [\"dx\", \"dy\", \"dyaw\"]\n",
+    "for i, ax in enumerate(axes):\n",
+    "  ax.set_ylabel(labels[i])\n",
+    "\n",
+    "\n",
+    "render_every = 2\n",
+    "fps = 1.0 / eval_env.dt / render_every\n",
+    "print(f\"fps: {fps}\")\n",
+    "\n",
+    "traj = rollout[::render_every]\n",
+    "mod_fns = modify_scene_fns[::render_every]\n",
+    "assert len(traj) == len(mod_fns)\n",
+    "\n",
+    "scene_option = mujoco.MjvOption()\n",
+    "scene_option.geomgroup[2] = True\n",
+    "scene_option.geomgroup[3] = False\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_PERTFORCE] = True\n",
+    "\n",
+    "frames = eval_env.render(\n",
+    "    traj,\n",
+    "    camera=\"track\",\n",
+    "    height=480,\n",
+    "    width=640,\n",
+    "    modify_scene_fns=mod_fns,\n",
+    "    scene_option=scene_option,\n",
+    ")\n",
+    "media.show_video(frames, fps=fps, loop=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "0RHZvXgmzrEJ"
+   },
+   "source": [
+    "## Handstand\n",
+    "\n",
+    "Additional policies are available for the Unitree Go1 such as fall-recovery, handstand, and footstand policies. We'll use the handstand policy as an opportunity to demonstrate finetuning policies from prior checkpoints. This will allow us to quickly iterate on training curriculums by modifying the enviornment config between runs.\n",
+    "\n",
+    "For the Go1 handstand policy, we'll first train with the default configuration, and then add an energy penalty to make the policy smoother and more likely to transfer onto the robot."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "RYriZOAxzEk_"
+   },
+   "outputs": [],
+   "source": [
+    "from mujoco_playground.config import locomotion_params\n",
+    "\n",
+    "env_name = 'Go1Handstand'\n",
+    "env = registry.load(env_name)\n",
+    "env_cfg = registry.get_default_config(env_name)\n",
+    "ppo_params = locomotion_params.brax_ppo_config(env_name)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "3nB5ugbdS5kk"
+   },
+   "source": [
+    "Let's create a checkpoint directory and then train a policy with checkpointing."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "EyEDpHisS7eO"
+   },
+   "outputs": [],
+   "source": [
+    "ckpt_path = epath.Path(\"checkpoints\").resolve() / env_name\n",
+    "ckpt_path.mkdir(parents=True, exist_ok=True)\n",
+    "print(f\"{ckpt_path}\")\n",
+    "\n",
+    "with open(ckpt_path / \"config.json\", \"w\") as fp:\n",
+    "  json.dump(env_cfg.to_dict(), fp, indent=4)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "lCRUYofXSNGT"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Training fn definition\n",
+    "x_data, y_data, y_dataerr = [], [], []\n",
+    "times = [datetime.now()]\n",
+    "\n",
+    "\n",
+    "def policy_params_fn(current_step, make_policy, params):\n",
+    "  del make_policy  # Unused.\n",
+    "  orbax_checkpointer = ocp.PyTreeCheckpointer()\n",
+    "  save_args = orbax_utils.save_args_from_target(params)\n",
+    "  path = ckpt_path / f\"{current_step}\"\n",
+    "  orbax_checkpointer.save(path, params, force=True, save_args=save_args)\n",
+    "\n",
+    "\n",
+    "def progress(num_steps, metrics):\n",
+    "  clear_output(wait=True)\n",
+    "\n",
+    "  times.append(datetime.now())\n",
+    "  x_data.append(num_steps)\n",
+    "  y_data.append(metrics[\"eval/episode_reward\"])\n",
+    "  y_dataerr.append(metrics[\"eval/episode_reward_std\"])\n",
+    "\n",
+    "  plt.xlim([0, ppo_params[\"num_timesteps\"] * 1.25])\n",
+    "  plt.xlabel(\"# environment steps\")\n",
+    "  plt.ylabel(\"reward per episode\")\n",
+    "  plt.title(f\"y={y_data[-1]:.3f}\")\n",
+    "  plt.errorbar(x_data, y_data, yerr=y_dataerr, color=\"blue\")\n",
+    "\n",
+    "  display(plt.gcf())\n",
+    "\n",
+    "randomizer = registry.get_domain_randomizer(env_name)\n",
+    "ppo_training_params = dict(ppo_params)\n",
+    "network_factory = ppo_networks.make_ppo_networks\n",
+    "if \"network_factory\" in ppo_params:\n",
+    "  del ppo_training_params[\"network_factory\"]\n",
+    "  network_factory = functools.partial(\n",
+    "      ppo_networks.make_ppo_networks,\n",
+    "      **ppo_params.network_factory\n",
+    "  )\n",
+    "\n",
+    "train_fn = functools.partial(\n",
+    "    ppo.train, **dict(ppo_training_params),\n",
+    "    network_factory=network_factory,\n",
+    "    randomization_fn=randomizer,\n",
+    "    progress_fn=progress,\n",
+    "    policy_params_fn=policy_params_fn,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "A1oK80x1anPp"
+   },
+   "source": [
+    "The initial policy takes 8 minutes to train on an RTX 4090."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "MY6P3abhSNGU"
+   },
+   "outputs": [],
+   "source": [
+    "make_inference_fn, params, metrics = train_fn(\n",
+    "    environment=registry.load(env_name, config=env_cfg),\n",
+    "    eval_env=registry.load(env_name, config=env_cfg),\n",
+    "    wrap_env_fn=wrapper.wrap_for_brax_training,\n",
+    ")\n",
+    "print(f\"time to jit: {times[1] - times[0]}\")\n",
+    "print(f\"time to train: {times[-1] - times[1]}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "4s6PkZ4GWV4Z"
+   },
+   "source": [
+    "Let's visualize the current policy."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "cellView": "form",
+    "id": "WiWOtc_6WbcX"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Rollout and Render\n",
+    "inference_fn = make_inference_fn(params, deterministic=True)\n",
+    "jit_inference_fn = jax.jit(inference_fn)\n",
+    "\n",
+    "eval_env = registry.load(env_name, config=env_cfg)\n",
+    "jit_reset = jax.jit(eval_env.reset)\n",
+    "jit_step = jax.jit(eval_env.step)\n",
+    "\n",
+    "rng = jax.random.PRNGKey(12345)\n",
+    "rollout = []\n",
+    "rewards = []\n",
+    "torso_height = []\n",
+    "actions = []\n",
+    "torques = []\n",
+    "power = []\n",
+    "qfrc_constraint = []\n",
+    "qvels = []\n",
+    "power1 = []\n",
+    "power2 = []\n",
+    "for _ in range(10):\n",
+    "  rng, reset_rng = jax.random.split(rng)\n",
+    "  state = jit_reset(reset_rng)\n",
+    "  for i in range(env_cfg.episode_length // 2):\n",
+    "    act_rng, rng = jax.random.split(rng)\n",
+    "    ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
+    "    actions.append(ctrl)\n",
+    "    state = jit_step(state, ctrl)\n",
+    "    rollout.append(state)\n",
+    "    rewards.append(\n",
+    "        {k[7:]: v for k, v in state.metrics.items() if k.startswith(\"reward/\")}\n",
+    "    )\n",
+    "    torso_height.append(state.data.qpos[2])\n",
+    "    torques.append(state.data.actuator_force)\n",
+    "    qvel = state.data.qvel[6:]\n",
+    "    power.append(jp.sum(jp.abs(qvel * state.data.actuator_force)))\n",
+    "    qfrc_constraint.append(jp.linalg.norm(state.data.qfrc_constraint[6:]))\n",
+    "    qvels.append(jp.max(jp.abs(qvel)))\n",
+    "    frc = state.data.actuator_force\n",
+    "    qvel = state.data.qvel[6:]\n",
+    "    power1.append(jp.sum(frc * qvel))\n",
+    "    power2.append(jp.sum(jp.abs(frc * qvel)))\n",
+    "\n",
+    "\n",
+    "render_every = 2\n",
+    "fps = 1.0 / eval_env.dt / render_every\n",
+    "traj = rollout[::render_every]\n",
+    "\n",
+    "scene_option = mujoco.MjvOption()\n",
+    "scene_option.geomgroup[2] = True\n",
+    "scene_option.geomgroup[3] = False\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = False\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = False\n",
+    "\n",
+    "frames = eval_env.render(\n",
+    "    traj, camera=\"side\", scene_option=scene_option, height=480, width=640\n",
+    ")\n",
+    "media.show_video(frames, fps=fps, loop=False)\n",
+    "\n",
+    "power = jp.array(power1)\n",
+    "print(f\"Max power: {jp.max(power)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "v5p0Z3PPSRik"
+   },
+   "source": [
+    "Notice that the above policy looks jittery and unlikely to transfer on the robot. The max power output is also quite high.\n",
+    "\n",
+    "The sim-to-real deployment of the handstand policy was trained using a curriculum on the `energy_termination_threshold`, `energy` and `dof_acc`, which are config values that penalize high torques and high power output. Let's finetune the above policy with a decreased  `energy_termination_threshold`, as well as non-zero values for `energy` and `dof_acc` rewards to get a smoother policy."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "hrjoVL-_WN-r"
+   },
+   "source": [
+    "### Finetune the previous checkpoint"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "jTxAySRSSu96"
+   },
+   "outputs": [],
+   "source": [
+    "env_cfg = registry.get_default_config(env_name)\n",
+    "env_cfg.energy_termination_threshold = 400  # lower energy termination threshold\n",
+    "env_cfg.reward_config.energy = -0.003  # non-zero negative `energy` reward\n",
+    "env_cfg.reward_config.dof_acc = -2.5e-7  # non-zero negative `dof_acc` reward\n",
+    "\n",
+    "FINETUNE_PATH = epath.Path(ckpt_path)\n",
+    "latest_ckpts = list(FINETUNE_PATH.glob(\"*\"))\n",
+    "latest_ckpts = [ckpt for ckpt in latest_ckpts if ckpt.is_dir()]\n",
+    "latest_ckpts.sort(key=lambda x: int(x.name))\n",
+    "latest_ckpt = latest_ckpts[-1]\n",
+    "restore_checkpoint_path = latest_ckpt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "_M5IqOR6z4bV"
+   },
+   "outputs": [],
+   "source": [
+    "x_data, y_data, y_dataerr = [], [], []\n",
+    "times = [datetime.now()]\n",
+    "\n",
+    "make_inference_fn, params, metrics = train_fn(\n",
+    "    environment=registry.load(env_name, config=env_cfg),\n",
+    "    eval_env=registry.load(env_name, config=env_cfg),\n",
+    "    wrap_env_fn=wrapper.wrap_for_brax_training,\n",
+    "    restore_checkpoint_path=restore_checkpoint_path,  # restore from the checkpoint!\n",
+    "    seed=1,\n",
+    ")\n",
+    "print(f\"time to jit: {times[1] - times[0]}\")\n",
+    "print(f\"time to train: {times[-1] - times[1]}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "cellView": "form",
+    "id": "tzG8eY2lz4dk"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Rollout and Render Finetune Policy\n",
+    "inference_fn = make_inference_fn(params, deterministic=True)\n",
+    "jit_inference_fn = jax.jit(inference_fn)\n",
+    "\n",
+    "eval_env = registry.load(env_name, config=env_cfg)\n",
+    "jit_reset = jax.jit(eval_env.reset)\n",
+    "jit_step = jax.jit(eval_env.step)\n",
+    "\n",
+    "rng = jax.random.PRNGKey(12345)\n",
+    "rollout = []\n",
+    "rewards = []\n",
+    "torso_height = []\n",
+    "actions = []\n",
+    "torques = []\n",
+    "power = []\n",
+    "qfrc_constraint = []\n",
+    "qvels = []\n",
+    "power1 = []\n",
+    "power2 = []\n",
+    "for _ in range(10):\n",
+    "  rng, reset_rng = jax.random.split(rng)\n",
+    "  state = jit_reset(reset_rng)\n",
+    "  for i in range(env_cfg.episode_length // 2):\n",
+    "    act_rng, rng = jax.random.split(rng)\n",
+    "    ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
+    "    actions.append(ctrl)\n",
+    "    state = jit_step(state, ctrl)\n",
+    "    rollout.append(state)\n",
+    "    rewards.append(\n",
+    "        {k[7:]: v for k, v in state.metrics.items() if k.startswith(\"reward/\")}\n",
+    "    )\n",
+    "    torso_height.append(state.data.qpos[2])\n",
+    "    torques.append(state.data.actuator_force)\n",
+    "    qvel = state.data.qvel[6:]\n",
+    "    power.append(jp.sum(jp.abs(qvel * state.data.actuator_force)))\n",
+    "    qfrc_constraint.append(jp.linalg.norm(state.data.qfrc_constraint[6:]))\n",
+    "    qvels.append(jp.max(jp.abs(qvel)))\n",
+    "    frc = state.data.actuator_force\n",
+    "    qvel = state.data.qvel[6:]\n",
+    "    power1.append(jp.sum(frc * qvel))\n",
+    "    power2.append(jp.sum(jp.abs(frc * qvel)))\n",
+    "\n",
+    "\n",
+    "render_every = 2\n",
+    "fps = 1.0 / eval_env.dt / render_every\n",
+    "traj = rollout[::render_every]\n",
+    "\n",
+    "scene_option = mujoco.MjvOption()\n",
+    "scene_option.geomgroup[2] = True\n",
+    "scene_option.geomgroup[3] = False\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = False\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = False\n",
+    "\n",
+    "frames = eval_env.render(\n",
+    "    traj, camera=\"side\", scene_option=scene_option, height=480, width=640\n",
+    ")\n",
+    "media.show_video(frames, fps=fps, loop=False)\n",
+    "\n",
+    "power = jp.array(power1)\n",
+    "print(f\"Max power: {jp.max(power)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "yCyibqGMiAca"
+   },
+   "source": [
+    "The final policy should exhibit smoother behavior and have less power output! Feel free to finetune the policy some more using different reward terms to get the best behavior."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "26o77FfWXvVp"
+   },
+   "source": [
+    "# Bipedal\n",
+    "\n",
+    "MuJoCo Playground also comes with a host of bipedal environments, such as the Berkely Humanoid and the Unitree G1/H1. Let's demonstrate a joystick policy on the Berkeley Humanoid. The initial policy takes 17 minutes to train on an RTX 4090."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "ESNd18FUanPt"
+   },
+   "outputs": [],
+   "source": [
+    "env_name = 'BerkeleyHumanoidJoystickFlatTerrain'\n",
+    "env = registry.load(env_name)\n",
+    "env_cfg = registry.get_default_config(env_name)\n",
+    "ppo_params = locomotion_params.brax_ppo_config(env_name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "nibLoRu8anPt"
+   },
+   "outputs": [],
+   "source": [
+    "x_data, y_data, y_dataerr = [], [], []\n",
+    "times = [datetime.now()]\n",
+    "\n",
+    "randomizer = registry.get_domain_randomizer(env_name)\n",
+    "ppo_training_params = dict(ppo_params)\n",
+    "network_factory = ppo_networks.make_ppo_networks\n",
+    "if \"network_factory\" in ppo_params:\n",
+    "  del ppo_training_params[\"network_factory\"]\n",
+    "  network_factory = functools.partial(\n",
+    "      ppo_networks.make_ppo_networks,\n",
+    "      **ppo_params.network_factory\n",
+    "  )\n",
+    "\n",
+    "train_fn = functools.partial(\n",
+    "    ppo.train, **dict(ppo_training_params),\n",
+    "    network_factory=network_factory,\n",
+    "    randomization_fn=randomizer,\n",
+    "    progress_fn=progress\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "16dqomv0anPt"
+   },
+   "outputs": [],
+   "source": [
+    "make_inference_fn, params, metrics = train_fn(\n",
+    "    environment=env,\n",
+    "    eval_env=registry.load(env_name, config=env_cfg),\n",
+    "    wrap_env_fn=wrapper.wrap_for_brax_training,\n",
+    ")\n",
+    "print(f\"time to jit: {times[1] - times[0]}\")\n",
+    "print(f\"time to train: {times[-1] - times[1]}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "cellView": "form",
+    "id": "sBHDF-JFanPt"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Rollout and Render\n",
+    "from mujoco_playground._src.gait import draw_joystick_command\n",
+    "\n",
+    "env = registry.load(env_name)\n",
+    "eval_env = registry.load(env_name)\n",
+    "jit_reset = jax.jit(eval_env.reset)\n",
+    "jit_step = jax.jit(eval_env.step)\n",
+    "jit_inference_fn = jax.jit(make_inference_fn(params, deterministic=True))\n",
+    "\n",
+    "rng = jax.random.PRNGKey(1)\n",
+    "\n",
+    "rollout = []\n",
+    "modify_scene_fns = []\n",
+    "\n",
+    "x_vel = 1.0  #@param {type: \"number\"}\n",
+    "y_vel = 0.0  #@param {type: \"number\"}\n",
+    "yaw_vel = 0.0  #@param {type: \"number\"}\n",
+    "command = jp.array([x_vel, y_vel, yaw_vel])\n",
+    "\n",
+    "phase_dt = 2 * jp.pi * eval_env.dt * 1.5\n",
+    "phase = jp.array([0, jp.pi])\n",
+    "\n",
+    "for j in range(1):\n",
+    "  print(f\"episode {j}\")\n",
+    "  state = jit_reset(rng)\n",
+    "  state.info[\"phase_dt\"] = phase_dt\n",
+    "  state.info[\"phase\"] = phase\n",
+    "  for i in range(env_cfg.episode_length):\n",
+    "    act_rng, rng = jax.random.split(rng)\n",
+    "    ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
+    "    state = jit_step(state, ctrl)\n",
+    "    if state.done:\n",
+    "      break\n",
+    "    state.info[\"command\"] = command\n",
+    "    rollout.append(state)\n",
+    "\n",
+    "    xyz = np.array(state.data.xpos[eval_env.mj_model.body(\"torso\").id])\n",
+    "    xyz += np.array([0, 0.0, 0])\n",
+    "    x_axis = state.data.xmat[eval_env._torso_body_id, 0]\n",
+    "    yaw = -np.arctan2(x_axis[1], x_axis[0])\n",
+    "    modify_scene_fns.append(\n",
+    "        functools.partial(\n",
+    "            draw_joystick_command,\n",
+    "            cmd=state.info[\"command\"],\n",
+    "            xyz=xyz,\n",
+    "            theta=yaw,\n",
+    "            scl=np.linalg.norm(state.info[\"command\"]),\n",
+    "        )\n",
+    "    )\n",
+    "\n",
+    "render_every = 1\n",
+    "fps = 1.0 / eval_env.dt / render_every\n",
+    "print(f\"fps: {fps}\")\n",
+    "traj = rollout[::render_every]\n",
+    "mod_fns = modify_scene_fns[::render_every]\n",
+    "\n",
+    "scene_option = mujoco.MjvOption()\n",
+    "scene_option.geomgroup[2] = True\n",
+    "scene_option.geomgroup[3] = False\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = False\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_PERTFORCE] = False\n",
+    "\n",
+    "frames = eval_env.render(\n",
+    "    traj,\n",
+    "    camera=\"track\",\n",
+    "    scene_option=scene_option,\n",
+    "    width=640*2,\n",
+    "    height=480,\n",
+    "    modify_scene_fns=mod_fns,\n",
+    ")\n",
+    "media.show_video(frames, fps=fps, loop=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "CBtrAqns35sI"
+   },
+   "source": [
+    "🙌 Hasta la vista!"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "A100",
+   "machine_shape": "hm",
+   "private_outputs": true,
+   "provenance": [],
+   "toc_visible": true
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

samples/manipulation.ipynb

@@ -0,0 +1,650 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "MpkYHwCqk7W-"
+      },
+      "source": [
+        "![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "xBSdkbmGN2K-"
+      },
+      "source": [
+        "### Copyright notice"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_UbO9uhtBSX5"
+      },
+      "source": [
+        "> <p><small><small>Copyright 2025 DeepMind Technologies Limited.</small></p>\n",
+        "> <p><small><small>Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at <a href=\"http://www.apache.org/licenses/LICENSE-2.0\">http://www.apache.org/licenses/LICENSE-2.0</a>.</small></small></p>\n",
+        "> <p><small><small>Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.</small></small></p>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dNIJkb_FM2Ux"
+      },
+      "source": [
+        "# Manipulation in The Playground! <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/manipulation.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/></a>\n",
+        "\n",
+        "In this notebook, we'll walk through a couple manipulation environments available in MuJoCo Playground.\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"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Xqo7pyX-n72M",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Install pre-requisites\n",
+        "!pip install mujoco\n",
+        "!pip install mujoco_mjx\n",
+        "!pip install brax"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "IbZxYDxzoz5R"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Check if MuJoCo installation was successful\n",
+        "\n",
+        "import distutils.util\n",
+        "import os\n",
+        "import subprocess\n",
+        "\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",
+        "\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",
+        "try:\n",
+        "  print('Checking that the installation succeeded:')\n",
+        "  import mujoco\n",
+        "\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",
+        "\n",
+        "print('Installation successful.')\n",
+        "\n",
+        "# Tell XLA to use Triton GEMM, this improves steps/sec by ~30% on some GPUs\n",
+        "xla_flags = os.environ.get('XLA_FLAGS', '')\n",
+        "xla_flags += ' --xla_gpu_triton_gemm_any=True'\n",
+        "os.environ['XLA_FLAGS'] = xla_flags"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "T5f4w3Kq2X14",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Import packages for plotting and creating graphics\n",
+        "import json\n",
+        "import itertools\n",
+        "import time\n",
+        "from typing import Callable, List, NamedTuple, Optional, Union\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)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ObF1UXrkb0Nd"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Import MuJoCo, MJX, and Brax\n",
+        "from datetime import datetime\n",
+        "import functools\n",
+        "import os\n",
+        "from typing import Any, Dict, Sequence, Tuple, Union\n",
+        "from brax import base\n",
+        "from brax import envs\n",
+        "from brax import math\n",
+        "from brax.base import Base, Motion, Transform\n",
+        "from brax.base import State as PipelineState\n",
+        "from brax.envs.base import Env, PipelineEnv, State\n",
+        "from brax.io import html, mjcf, model\n",
+        "from brax.mjx.base import State as MjxState\n",
+        "from brax.training.agents.ppo import networks as ppo_networks\n",
+        "from brax.training.agents.ppo import train as ppo\n",
+        "from brax.training.agents.sac import networks as sac_networks\n",
+        "from brax.training.agents.sac import train as sac\n",
+        "from etils import epath\n",
+        "from flax import struct\n",
+        "from flax.training import orbax_utils\n",
+        "from IPython.display import HTML, clear_output\n",
+        "import jax\n",
+        "from jax import numpy as jp\n",
+        "from matplotlib import pyplot as plt\n",
+        "import mediapy as media\n",
+        "from ml_collections import config_dict\n",
+        "import mujoco\n",
+        "from mujoco import mjx\n",
+        "import numpy as np\n",
+        "from orbax import checkpoint as ocp"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title Install MuJoCo Playground\n",
+        "!pip install playground"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "UoTLSx4cFRdy"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title Import The Playground\n",
+        "\n",
+        "from mujoco_playground import wrapper\n",
+        "from mujoco_playground import registry"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "gYm2h7m8w3Nv"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Manipulation\n",
+        "\n",
+        "MuJoCo Playground contains several manipulation environments (all listed below after running the command)."
+      ],
+      "metadata": {
+        "id": "LcibXbyKt4FI"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "registry.manipulation.ALL_ENVS"
+      ],
+      "metadata": {
+        "id": "ox0Gze9Ct5AM"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Franka Emika Panda\n",
+        "\n",
+        "Let's start off with the simplest environment, simply picking up a cube with the Franka Emika Panda."
+      ],
+      "metadata": {
+        "id": "_R01tjWfI-i6"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "env_name = 'PandaPickCubeOrientation'\n",
+        "env = registry.load(env_name)\n",
+        "env_cfg = registry.get_default_config(env_name)"
+      ],
+      "metadata": {
+        "id": "kPJeoQeEJBSA"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "env_cfg"
+      ],
+      "metadata": {
+        "id": "6n9UT9N1wR5K"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Train Policy\n",
+        "\n",
+        "Let's train the pick cube policy and visualize rollouts. The policy takes roughly 3 minutes to train on an RTX 4090."
+      ],
+      "metadata": {
+        "id": "Thm7nZueM4cz"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from mujoco_playground.config import manipulation_params\n",
+        "ppo_params = manipulation_params.brax_ppo_config(env_name)\n",
+        "ppo_params"
+      ],
+      "metadata": {
+        "id": "B9T_UVZYLDdM"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vBEEQyY6M5OC"
+      },
+      "source": [
+        "### PPO"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XKFzyP7wM5OD"
+      },
+      "outputs": [],
+      "source": [
+        "x_data, y_data, y_dataerr = [], [], []\n",
+        "times = [datetime.now()]\n",
+        "\n",
+        "\n",
+        "def progress(num_steps, metrics):\n",
+        "  clear_output(wait=True)\n",
+        "\n",
+        "  times.append(datetime.now())\n",
+        "  x_data.append(num_steps)\n",
+        "  y_data.append(metrics[\"eval/episode_reward\"])\n",
+        "  y_dataerr.append(metrics[\"eval/episode_reward_std\"])\n",
+        "\n",
+        "  plt.xlim([0, ppo_params[\"num_timesteps\"] * 1.25])\n",
+        "  plt.xlabel(\"# environment steps\")\n",
+        "  plt.ylabel(\"reward per episode\")\n",
+        "  plt.title(f\"y={y_data[-1]:.3f}\")\n",
+        "  plt.errorbar(x_data, y_data, yerr=y_dataerr, color=\"blue\")\n",
+        "\n",
+        "  display(plt.gcf())\n",
+        "\n",
+        "ppo_training_params = dict(ppo_params)\n",
+        "network_factory = ppo_networks.make_ppo_networks\n",
+        "if \"network_factory\" in ppo_params:\n",
+        "  del ppo_training_params[\"network_factory\"]\n",
+        "  network_factory = functools.partial(\n",
+        "      ppo_networks.make_ppo_networks,\n",
+        "      **ppo_params.network_factory\n",
+        "  )\n",
+        "\n",
+        "train_fn = functools.partial(\n",
+        "    ppo.train, **dict(ppo_training_params),\n",
+        "    network_factory=network_factory,\n",
+        "    progress_fn=progress,\n",
+        "    seed=1\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "FGrlulWbM5OD"
+      },
+      "outputs": [],
+      "source": [
+        "make_inference_fn, params, metrics = train_fn(\n",
+        "    environment=env,\n",
+        "    wrap_env_fn=wrapper.wrap_for_brax_training,\n",
+        ")\n",
+        "print(f\"time to jit: {times[1] - times[0]}\")\n",
+        "print(f\"time to train: {times[-1] - times[1]}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Visualize Rollouts"
+      ],
+      "metadata": {
+        "id": "mHVmccs-oMSo"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "sG5a2FFXoUKw"
+      },
+      "outputs": [],
+      "source": [
+        "jit_reset = jax.jit(env.reset)\n",
+        "jit_step = jax.jit(env.step)\n",
+        "jit_inference_fn = jax.jit(make_inference_fn(params, deterministic=True))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "C_1CY9xDoUKw"
+      },
+      "outputs": [],
+      "source": [
+        "rng = jax.random.PRNGKey(42)\n",
+        "rollout = []\n",
+        "n_episodes = 1\n",
+        "\n",
+        "for _ in range(n_episodes):\n",
+        "  state = jit_reset(rng)\n",
+        "  rollout.append(state)\n",
+        "  for i in range(env_cfg.episode_length):\n",
+        "    act_rng, rng = jax.random.split(rng)\n",
+        "    ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
+        "    state = jit_step(state, ctrl)\n",
+        "    rollout.append(state)\n",
+        "\n",
+        "render_every = 1\n",
+        "frames = env.render(rollout[::render_every])\n",
+        "rewards = [s.reward for s in rollout]\n",
+        "media.show_video(frames, fps=1.0 / env.dt / render_every)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "While the above policy is very simple, the work was extended using the Madrona batch renderer, and policies were transferred on a real robot. We encourage folks to check out the Madrona-MJX tutorial notebooks ([part 1](https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/training_vision_1.ipynb) and [part 2](https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/training_vision_2.ipynb))!"
+      ],
+      "metadata": {
+        "id": "v5r4FwivlnoG"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Dexterous Manipulation\n",
+        "\n",
+        "Let's now train a policy that was transferred onto a real Leap Hand robot with the `LeapCubeReorient` environment! The environment contains a cube placed in the center of the hand, and the goal is to re-orient the cube in SO(3)."
+      ],
+      "metadata": {
+        "id": "YVQsrEE3mMj8"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "env_name = 'LeapCubeReorient'\n",
+        "env = registry.load(env_name)\n",
+        "env_cfg = registry.get_default_config(env_name)"
+      ],
+      "metadata": {
+        "id": "oPaTdWqVmuPt"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "env_cfg"
+      ],
+      "metadata": {
+        "id": "c0OII08RmuPt"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Train Policy\n",
+        "\n",
+        "Let's train an initial policy and visualize the rollouts. Notice that the PPO parameters contain `policy_obs_key` and `value_obs_key` fields, which allow us to train brax PPO with [asymmetric](https://arxiv.org/abs/1710.06542) observations for the actor and the critic. While the actor recieves proprioceptive state similar in nature to the real-world camera tracking sensors, the critic network recieves privileged state only available in the simulator. This enables more sample efficient learning, and we are able to train an initial policy in 33 minutes on a single RTX 4090.\n",
+        "\n",
+        "Depending on the GPU device and topology, training can be brought down to 10-20 minutes as shown in the MuJoCo Playground technical report."
+      ],
+      "metadata": {
+        "id": "3g335ZYFmuPt"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from mujoco_playground.config import manipulation_params\n",
+        "ppo_params = manipulation_params.brax_ppo_config(env_name)\n",
+        "ppo_params"
+      ],
+      "metadata": {
+        "id": "cc1Ka4hYmuPt"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Ulr1ih6PmuPu"
+      },
+      "source": [
+        "### PPO"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "gzwRjUGLmuPu"
+      },
+      "outputs": [],
+      "source": [
+        "x_data, y_data, y_dataerr = [], [], []\n",
+        "times = [datetime.now()]\n",
+        "\n",
+        "\n",
+        "def progress(num_steps, metrics):\n",
+        "  clear_output(wait=True)\n",
+        "\n",
+        "  times.append(datetime.now())\n",
+        "  x_data.append(num_steps)\n",
+        "  y_data.append(metrics[\"eval/episode_reward\"])\n",
+        "  y_dataerr.append(metrics[\"eval/episode_reward_std\"])\n",
+        "\n",
+        "  plt.xlim([0, ppo_params[\"num_timesteps\"] * 1.25])\n",
+        "  plt.xlabel(\"# environment steps\")\n",
+        "  plt.ylabel(\"reward per episode\")\n",
+        "  plt.title(f\"y={y_data[-1]:.3f}\")\n",
+        "  plt.errorbar(x_data, y_data, yerr=y_dataerr, color=\"blue\")\n",
+        "\n",
+        "  display(plt.gcf())\n",
+        "\n",
+        "ppo_training_params = dict(ppo_params)\n",
+        "network_factory = ppo_networks.make_ppo_networks\n",
+        "if \"network_factory\" in ppo_params:\n",
+        "  del ppo_training_params[\"network_factory\"]\n",
+        "  network_factory = functools.partial(\n",
+        "      ppo_networks.make_ppo_networks,\n",
+        "      **ppo_params.network_factory\n",
+        "  )\n",
+        "\n",
+        "train_fn = functools.partial(\n",
+        "    ppo.train, **dict(ppo_training_params),\n",
+        "    network_factory=network_factory,\n",
+        "    progress_fn=progress,\n",
+        "    seed=1\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YmortADGmuPu"
+      },
+      "outputs": [],
+      "source": [
+        "make_inference_fn, params, metrics = train_fn(\n",
+        "    environment=env,\n",
+        "    wrap_env_fn=wrapper.wrap_for_brax_training,\n",
+        ")\n",
+        "print(f\"time to jit: {times[1] - times[0]}\")\n",
+        "print(f\"time to train: {times[-1] - times[1]}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Visualize Rollouts"
+      ],
+      "metadata": {
+        "id": "xIB_677emuPu"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "xBgGvZpTmuPu"
+      },
+      "outputs": [],
+      "source": [
+        "jit_reset = jax.jit(env.reset)\n",
+        "jit_step = jax.jit(env.step)\n",
+        "jit_inference_fn = jax.jit(make_inference_fn(params, deterministic=True))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ksj6_9PwmuPu"
+      },
+      "outputs": [],
+      "source": [
+        "rng = jax.random.PRNGKey(42)\n",
+        "rollout = []\n",
+        "n_episodes = 1\n",
+        "\n",
+        "for _ in range(n_episodes):\n",
+        "  state = jit_reset(rng)\n",
+        "  rollout.append(state)\n",
+        "  for i in range(env_cfg.episode_length):\n",
+        "    act_rng, rng = jax.random.split(rng)\n",
+        "    ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
+        "    state = jit_step(state, ctrl)\n",
+        "    rollout.append(state)\n",
+        "\n",
+        "render_every = 1\n",
+        "frames = env.render(rollout[::render_every])\n",
+        "rewards = [s.reward for s in rollout]\n",
+        "media.show_video(frames, fps=1.0 / env.dt / render_every)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "The above policy solves the task, but may look a little bit jittery. To get robust sim-to-real transfer,  we retrained from previous checkpoints using a curriculum on the maximum torque to facilitate exploration early on in the curriculum, and to produce smoother actions for the final policy. More details can be found in the MuJoCo Playground technical report!"
+      ],
+      "metadata": {
+        "id": "dWIVTxq5nhs5"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "CBtrAqns35sI"
+      },
+      "source": [
+        "🙌 Thanks for stopping by The Playground!"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "private_outputs": true,
+      "toc_visible": true,
+      "provenance": [],
+      "machine_shape": "hm",
+      "gpuType": "A100"
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

samples/tutorial.ipynb

@@ -0,0 +1,2258 @@
+{
+ "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/python/tutorial.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/></a></center></h1>\n",
+    "\n",
+    "This notebook provides an introductory tutorial for [**MuJoCo** physics](https://github.com/google-deepmind/mujoco#readme), using the native Python bindings.\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": {
+    "id": "YvyGCsgSCxHQ"
+   },
+   "source": [
+    "# All imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "Xqo7pyX-n72M"
+   },
+   "outputs": [],
+   "source": [
+    "!pip install mujoco\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()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "t0CF6Gvkt_Cw"
+   },
+   "source": [
+    "# MuJoCo basics\n",
+    "\n",
+    "We begin by defining and loading a simple model:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "3KJVqak6xdJa"
+   },
+   "outputs": [],
+   "source": [
+    "xml = \"\"\"\n",
+    "<mujoco>\n",
+    "  <worldbody>\n",
+    "    <geom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/>\n",
+    "    <geom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/>\n",
+    "  </worldbody>\n",
+    "</mujoco>\n",
+    "\"\"\"\n",
+    "model = mujoco.MjModel.from_xml_string(xml)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "slhf39lGxvDI"
+   },
+   "source": [
+    "The `xml` string is written in MuJoCo's [MJCF](http://www.mujoco.org/book/modeling.html), which is an [XML](https://en.wikipedia.org/wiki/XML#Key_terminology)-based modeling language.\n",
+    "  - The only required element is `<mujoco>`. The smallest valid MJCF model is `<mujoco/>` which is a completely empty model.\n",
+    "  - All physical elements live inside the `<worldbody>` which is always the top-level body and constitutes the global origin in Cartesian coordinates.\n",
+    "  - We define two geoms in the world named `red_box` and `green_sphere`.\n",
+    "  - **Question:** The `red_box` has no position, the `green_sphere` has no type, why is that?\n",
+    "    - **Answer:** MJCF attributes have *default values*. The default position is `0 0 0`, the default geom type is `sphere`. The MJCF language is described in the documentation's [XML Reference chapter](https://mujoco.readthedocs.io/en/latest/XMLreference.html).\n",
+    "\n",
+    "The `from_xml_string()` method invokes the model compiler, which creates a binary `mjModel` instance."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "gf9h_wi9weet"
+   },
+   "source": [
+    "## mjModel\n",
+    "\n",
+    "MuJoCo's `mjModel`, contains the *model description*, i.e., all quantities which *do not change over time*. The complete description of `mjModel` can be found at the end of the header file [`mjmodel.h`](https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjmodel.h). Note that the header files contain short, useful inline comments, describing each field.\n",
+    "\n",
+    "Examples of quantities that can be found in `mjModel` are `ngeom`, the number of geoms in the scene and `geom_rgba`, their respective colors:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "F40Pe6DY3Q0g"
+   },
+   "outputs": [],
+   "source": [
+    "model.ngeom"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "MOIJG9pzx8cA"
+   },
+   "outputs": [],
+   "source": [
+    "model.geom_rgba"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "bzcLjdY23Kvp"
+   },
+   "source": [
+    "## Named access\n",
+    "\n",
+    "The MuJoCo Python bindings provide convenient [accessors](https://mujoco.readthedocs.io/en/latest/python.html#named-access) using names. Calling the `model.geom()` accessor without a name string generates a convenient error that tells us what the valid names are."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "9AuTwbLFyJxQ"
+   },
+   "outputs": [],
+   "source": [
+    "try:\n",
+    "  model.geom()\n",
+    "except KeyError as e:\n",
+    "  print(e)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "qkfLK3h2zrqr"
+   },
+   "source": [
+    "Calling the named accessor without specifying a property will tell us what all the valid properties are:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "9X95TlWnyEEw"
+   },
+   "outputs": [],
+   "source": [
+    "model.geom('green_sphere')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "mS9qDLevKsJq"
+   },
+   "source": [
+    "Let's read the `green_sphere`'s rgba values:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "xsBlJAV7zpHb"
+   },
+   "outputs": [],
+   "source": [
+    "model.geom('green_sphere').rgba"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "8a8hswjjKyIa"
+   },
+   "source": [
+    "This functionality is a convenience shortcut for MuJoCo's [`mj_name2id`](https://mujoco.readthedocs.io/en/latest/APIreference.html?highlight=mj_name2id#mj-name2id) function:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "Ng92hNUoKnVq"
+   },
+   "outputs": [],
+   "source": [
+    "id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'green_sphere')\n",
+    "model.geom_rgba[id, :]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "5WL_SaJPLl3r"
+   },
+   "source": [
+    "Similarly, the read-only `id` and `name` properties can be used to convert from id to name and back:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "2CbGSmRZeE5p"
+   },
+   "outputs": [],
+   "source": [
+    "print('id of \"green_sphere\": ', model.geom('green_sphere').id)\n",
+    "print('name of geom 1: ', model.geom(1).name)\n",
+    "print('name of body 0: ', model.body(0).name)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "3RIizubaL_du"
+   },
+   "source": [
+    "Note that the 0th body is always the `world`. It cannot be renamed.\n",
+    "\n",
+    "The `id` and `name` attributes are useful in Python comprehensions:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "m3MtIE5F1K7s"
+   },
+   "outputs": [],
+   "source": [
+    "[model.geom(i).name for i in range(model.ngeom)]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "t5hY0fyXFLcf"
+   },
+   "source": [
+    "## `mjData`\n",
+    "`mjData` contains the *state* and quantities that depend on it. The state is made up of time, [generalized](https://en.wikipedia.org/wiki/Generalized_coordinates) positions and generalized velocities. These are respectively `data.time`, `data.qpos` and `data.qvel`. In order to make a new `mjData`, all we need is our `mjModel`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "FV2Hy6m948nr"
+   },
+   "outputs": [],
+   "source": [
+    "data = mujoco.MjData(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "-KmNuvlJ46u0"
+   },
+   "source": [
+    "`mjData` also contains *functions of the state*, for example the Cartesian positions of objects in the world frame. The (x, y, z) positions of our two geoms are in `data.geom_xpos`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "CPwDcAQ0-uUE"
+   },
+   "outputs": [],
+   "source": [
+    "print(data.geom_xpos)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Sjst5xGXX3sr"
+   },
+   "source": [
+    "Wait, why are both of our geoms at the origin? Didn't we offset the green sphere? The answer is that derived quantities in `mjData` need to be explicitly propagated (see [below](#scrollTo=QY1gpms1HXeN)). In our case, the minimal required function is [`mj_kinematics`](https://mujoco.readthedocs.io/en/latest/APIreference.html#mj-kinematics), which computes global Cartesian poses for all objects (excluding cameras and lights)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "tfe0YeZRYNTr"
+   },
+   "outputs": [],
+   "source": [
+    "mujoco.mj_kinematics(model, data)\n",
+    "print('raw access:\\n', data.geom_xpos)\n",
+    "\n",
+    "# MjData also supports named access:\n",
+    "print('\\nnamed access:\\n', data.geom('green_sphere').xpos)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "eU7uWNsTwmcZ"
+   },
+   "source": [
+    "# Basic rendering, simulation, and animation\n",
+    "\n",
+    "In order to render we'll need to instantiate a `Renderer` object and call its `render` method.\n",
+    "\n",
+    "We'll also reload our model to make the colab's sections independent."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "xK3c0-UDxMrN"
+   },
+   "outputs": [],
+   "source": [
+    "xml = \"\"\"\n",
+    "<mujoco>\n",
+    "  <worldbody>\n",
+    "    <geom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/>\n",
+    "    <geom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/>\n",
+    "  </worldbody>\n",
+    "</mujoco>\n",
+    "\"\"\"\n",
+    "# Make model and data\n",
+    "model = mujoco.MjModel.from_xml_string(xml)\n",
+    "data = mujoco.MjData(model)\n",
+    "\n",
+    "# Make renderer, render and show the pixels\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  media.show_image(renderer.render())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "ZkFSHeYGxlT5"
+   },
+   "source": [
+    "Hmmm, why the black pixels?\n",
+    "\n",
+    "**Answer:** For the same reason as above, we first need to propagate the values in `mjData`. This time we'll call [`mj_forward`](https://mujoco.readthedocs.io/en/latest/APIreference/APIfunctions.html#mj-forward), which invokes the entire pipeline up to the computation of accelerations i.e., it computes $\\dot x = f(x)$, where $x$ is the state. This function does more than we actually need, but unless we care about saving computation time, it's good practice to call `mj_forward` since then we know we are not missing anything.\n",
+    "\n",
+    "We also need to update the `mjvScene` which is an object held by the renderer describing the visual scene. We'll later see that the scene can include visual objects which are not part of the physical model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "pvh47r97huS4"
+   },
+   "outputs": [],
+   "source": [
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  mujoco.mj_forward(model, data)\n",
+    "  renderer.update_scene(data)\n",
+    "\n",
+    "  media.show_image(renderer.render())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "6oDW1dOUifw6"
+   },
+   "source": [
+    "This worked, but this image is a bit dark. Let's add a light and re-render."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "iqzJj2NIr_2V"
+   },
+   "outputs": [],
+   "source": [
+    "xml = \"\"\"\n",
+    "<mujoco>\n",
+    "  <worldbody>\n",
+    "    <light name=\"top\" pos=\"0 0 1\"/>\n",
+    "    <geom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/>\n",
+    "    <geom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/>\n",
+    "  </worldbody>\n",
+    "</mujoco>\n",
+    "\"\"\"\n",
+    "model = mujoco.MjModel.from_xml_string(xml)\n",
+    "data = mujoco.MjData(model)\n",
+    "\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  mujoco.mj_forward(model, data)\n",
+    "  renderer.update_scene(data)\n",
+    "\n",
+    "  media.show_image(renderer.render())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "HS4K38Eirww9"
+   },
+   "source": [
+    "Much better!\n",
+    "\n",
+    "Note that all values in the `mjModel` instance are writable. While it's generally not recommended to do this but rather to change the values in the XML, because it's easy to make an invalid model, some values are safe to write into, for example colors:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "GBNcQVYJrt2h"
+   },
+   "outputs": [],
+   "source": [
+    "# Run this cell multiple times for different colors\n",
+    "model.geom('red_box').rgba[:3] = np.random.rand(3)\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  renderer.update_scene(data)\n",
+    "\n",
+    "  media.show_image(renderer.render())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "-P95E-QHizQq"
+   },
+   "source": [
+    "# Simulation\n",
+    "\n",
+    "Now let's simulate and make a video. We'll use MuJoCo's main high level function `mj_step`, which steps the state $x_{t+h} = f(x_t)$.\n",
+    "\n",
+    "Note that in the code block below we are *not* rendering after each call to `mj_step`. This is because the default timestep is 2ms, and we want a 60fps video, not 500fps."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "NdVnHOYisiKl"
+   },
+   "outputs": [],
+   "source": [
+    "duration = 3.8  # (seconds)\n",
+    "framerate = 60  # (Hz)\n",
+    "\n",
+    "# Simulate and display video.\n",
+    "frames = []\n",
+    "mujoco.mj_resetData(model, data)  # Reset state and time.\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  while data.time < duration:\n",
+    "    mujoco.mj_step(model, data)\n",
+    "    if len(frames) < data.time * framerate:\n",
+    "      renderer.update_scene(data)\n",
+    "      pixels = renderer.render()\n",
+    "      frames.append(pixels)\n",
+    "\n",
+    "media.show_video(frames, fps=framerate)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "tYN4sL9RnsCU"
+   },
+   "source": [
+    "Hmmm, the video is playing, but nothing is moving, why is that?\n",
+    "\n",
+    "This is because this model has no [degrees of freedom](https://www.google.com/url?sa=D&q=https%3A%2F%2Fen.wikipedia.org%2Fwiki%2FDegrees_of_freedom_(mechanics)) (DoFs). The things that move (and which have inertia) are called *bodies*. We add DoFs by adding *joints* to bodies, specifying how they can move with respect to their parents. Let's make a new body that contains our geoms, add a hinge joint and re-render, while visualizing the joint axis using the visualization option object `MjvOption`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "LbWf84VYst5m"
+   },
+   "outputs": [],
+   "source": [
+    "xml = \"\"\"\n",
+    "<mujoco>\n",
+    "  <worldbody>\n",
+    "    <light name=\"top\" pos=\"0 0 1\"/>\n",
+    "    <body name=\"box_and_sphere\" euler=\"0 0 -30\">\n",
+    "      <joint name=\"swing\" type=\"hinge\" axis=\"1 -1 0\" pos=\"-.2 -.2 -.2\"/>\n",
+    "      <geom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/>\n",
+    "      <geom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/>\n",
+    "    </body>\n",
+    "  </worldbody>\n",
+    "</mujoco>\n",
+    "\"\"\"\n",
+    "model = mujoco.MjModel.from_xml_string(xml)\n",
+    "data = mujoco.MjData(model)\n",
+    "\n",
+    "# enable joint visualization option:\n",
+    "scene_option = mujoco.MjvOption()\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = True\n",
+    "\n",
+    "duration = 3.8  # (seconds)\n",
+    "framerate = 60  # (Hz)\n",
+    "\n",
+    "# Simulate and display video.\n",
+    "frames = []\n",
+    "mujoco.mj_resetData(model, data)\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  while data.time < duration:\n",
+    "    mujoco.mj_step(model, data)\n",
+    "    if len(frames) < data.time * framerate:\n",
+    "      renderer.update_scene(data, scene_option=scene_option)\n",
+    "      pixels = renderer.render()\n",
+    "      frames.append(pixels)\n",
+    "\n",
+    "media.show_video(frames, fps=framerate)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Ymv-tvWCpl6V"
+   },
+   "source": [
+    "Note that we rotated the `box_and_sphere` body by 30° around the Z (vertical) axis, with the directive `euler=\"0 0 -30\"`. This was made to emphasize that the poses of elements in the [kinematic tree](https://www.google.com/url?sa=D&q=https%3A%2F%2Fen.wikipedia.org%2Fwiki%2FKinematic_chain) are always with respect to their *parent body*, so our two geoms were also rotated by this transformation.\n",
+    "\n",
+    "Physics options live in `mjModel.opt`, for example the timestep:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "5yvAJokcpyX_"
+   },
+   "outputs": [],
+   "source": [
+    "model.opt.timestep"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "SdkwLeGUp9B2"
+   },
+   "source": [
+    "Let's flip gravity and re-render:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "ocjPQG8Dp2F-"
+   },
+   "outputs": [],
+   "source": [
+    "print('default gravity', model.opt.gravity)\n",
+    "model.opt.gravity = (0, 0, 10)\n",
+    "print('flipped gravity', model.opt.gravity)\n",
+    "\n",
+    "# Simulate and display video.\n",
+    "frames = []\n",
+    "mujoco.mj_resetData(model, data)\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  while data.time < duration:\n",
+    "    mujoco.mj_step(model, data)\n",
+    "    if len(frames) < data.time * framerate:\n",
+    "      renderer.update_scene(data, scene_option=scene_option)\n",
+    "      pixels = renderer.render()\n",
+    "      frames.append(pixels)\n",
+    "\n",
+    "media.show_video(frames, fps=60)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "FsxDDgXBqg_J"
+   },
+   "source": [
+    "We could also have done this in XML using the top-level `<option>` element:\n",
+    "```xml\n",
+    "<mujoco>\n",
+    "  <option gravity=\"0 0 10\"/>\n",
+    "  ...\n",
+    "</mujoco>\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "QY1gpms1HXeN"
+   },
+   "source": [
+    "### Understanding Degrees of Freedom\n",
+    "\n",
+    "In the real world, all rigid objects have 6 degrees-of-freedom: 3 translations and 3 rotations. Real-world joints act as constraints, removing relative degrees-of-freedom from bodies connected by joints. Some physics simulation software use this representation which is known as the \"Cartesian\" or \"subtractive\" representation, but it is inefficient. MuJoCo uses a representation known as the \"Lagrangian\", \"generalized\" or \"additive\" representation, whereby objects have no degrees of freedom unless explicitly added using joints.\n",
+    "\n",
+    "Our model, which has a single hinge joint, has one degree of freedom, and the entire state is defined by this joint's angle and angular velocity. These are the system's generalized position and velocity."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "wEdfGEfSKAOC"
+   },
+   "outputs": [],
+   "source": [
+    "print('Total number of DoFs in the model:', model.nv)\n",
+    "print('Generalized positions:', data.qpos)\n",
+    "print('Generalized velocities:', data.qvel)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Z8E-P5xONUSn"
+   },
+   "source": [
+    "MuJoCo's use of generalized coordinates is the reason that calling a function (e.g. [`mj_forward`](https://mujoco.readthedocs.io/en/latest/APIreference.html#mj-forward)) is required before rendering or reading the global poses of objects – Cartesian positions are *derived* from the generalized positions and need to be explicitly computed."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "SHppAOjvSupc"
+   },
+   "source": [
+    "# Example: Simulating free bodies with the self-inverting \"tippe-top\"\n",
+    "\n",
+    "A free body is a body with a [free joint](https://www.google.com/url?sa=D&q=https%3A%2F%2Fmujoco.readthedocs.io%2Fen%2Flatest%2FXMLreference.html%3Fhighlight%3Dfreejoint%23body-freejoint) having 6 DoFs, i.e., 3 translations and 3 rotations. We could give our `box_and_sphere` body a free joint and watch it fall, but let's look at something more interesting. A \"tippe top\" is a spinning toy which flips itself ([video](https://www.youtube.com/watch?v=kbYpVrdcszQ), [Wikipedia](https://en.wikipedia.org/wiki/Tippe_top)). We model it as follows:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "xasXQpVMjIwA"
+   },
+   "outputs": [],
+   "source": [
+    "tippe_top = \"\"\"\n",
+    "<mujoco model=\"tippe top\">\n",
+    "  <option integrator=\"RK4\"/>\n",
+    "\n",
+    "  <asset>\n",
+    "    <texture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
+    "     rgb2=\".2 .3 .4\" width=\"300\" height=\"300\"/>\n",
+    "    <material name=\"grid\" texture=\"grid\" texrepeat=\"8 8\" reflectance=\".2\"/>\n",
+    "  </asset>\n",
+    "\n",
+    "  <worldbody>\n",
+    "    <geom size=\".2 .2 .01\" type=\"plane\" material=\"grid\"/>\n",
+    "    <light pos=\"0 0 .6\"/>\n",
+    "    <camera name=\"closeup\" pos=\"0 -.1 .07\" xyaxes=\"1 0 0 0 1 2\"/>\n",
+    "    <body name=\"top\" pos=\"0 0 .02\">\n",
+    "      <freejoint/>\n",
+    "      <geom name=\"ball\" type=\"sphere\" size=\".02\" />\n",
+    "      <geom name=\"stem\" type=\"cylinder\" pos=\"0 0 .02\" size=\"0.004 .008\"/>\n",
+    "      <geom name=\"ballast\" type=\"box\" size=\".023 .023 0.005\"  pos=\"0 0 -.015\"\n",
+    "       contype=\"0\" conaffinity=\"0\" group=\"3\"/>\n",
+    "    </body>\n",
+    "  </worldbody>\n",
+    "\n",
+    "  <keyframe>\n",
+    "    <key name=\"spinning\" qpos=\"0 0 0.02 1 0 0 0\" qvel=\"0 0 0 0 1 200\" />\n",
+    "  </keyframe>\n",
+    "</mujoco>\n",
+    "\"\"\"\n",
+    "model = mujoco.MjModel.from_xml_string(tippe_top)\n",
+    "data = mujoco.MjData(model)\n",
+    "\n",
+    "mujoco.mj_forward(model, data)\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  renderer.update_scene(data, camera=\"closeup\")\n",
+    "\n",
+    "  media.show_image(renderer.render())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "bvHlr6maJYIG"
+   },
+   "source": [
+    "Note several new features of this model definition:\n",
+    "1. A 6-DoF free joint is added with the `<freejoint/>` clause.\n",
+    "2. We use the `<option/>` clause to set the integrator to the 4th order [Runge Kutta](https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods). Runge-Kutta has a higher rate of convergence than the default Euler integrator, which in many cases increases the accuracy at a given timestep size.\n",
+    "3. We define the floor's grid material inside the `<asset/>` clause and reference it in the `\"floor\"` geom.\n",
+    "4. We use an invisible and non-colliding box geom called `ballast` to move the top's center-of-mass lower. Having a low center of mass is (counter-intuitively) required for the flipping behavior to occur.\n",
+    "5. We save our initial spinning state as a *keyframe*. It has a high rotational velocity around the Z-axis, but is not perfectly oriented with the world, which introduces the symmetry-breaking required for the flipping.\n",
+    "6. We define a `<camera>` in our model, and then render from it using the `camera` argument to `update_scene()`.\n",
+    "Let us examine the state:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "o4S9nYhHOKmb"
+   },
+   "outputs": [],
+   "source": [
+    "print('positions', data.qpos)\n",
+    "print('velocities', data.qvel)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "71UgzBAqWdtZ"
+   },
+   "source": [
+    "The velocities are easy to interpret, 6 zeros, one for each DoF. What about the length 7 positions? We can see the initial 2cm height of the body; the subsequent four numbers are the 3D orientation, defined by a *unit quaternion*. 3D orientations are represented with **4** numbers while angular velocities are **3** numbers. For more information see the Wikipedia article on [quaternions and spatial rotation](https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation).\n",
+    "\n",
+    "Let's make a video:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "5P4HkhKNGQvs"
+   },
+   "outputs": [],
+   "source": [
+    "duration = 7    # (seconds)\n",
+    "framerate = 60  # (Hz)\n",
+    "\n",
+    "# Simulate and display video.\n",
+    "frames = []\n",
+    "mujoco.mj_resetDataKeyframe(model, data, 0)  # Reset the state to keyframe 0\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  while data.time < duration:\n",
+    "    mujoco.mj_step(model, data)\n",
+    "    if len(frames) < data.time * framerate:\n",
+    "      renderer.update_scene(data, \"closeup\")\n",
+    "      pixels = renderer.render()\n",
+    "      frames.append(pixels)\n",
+    "\n",
+    "media.show_video(frames, fps=framerate)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "rRuFKD2ubPgu"
+   },
+   "source": [
+    "### Measuring values from `mjData`\n",
+    "As mentioned above, the `mjData` structure contains the dynamic variables and intermediate results produced by the simulation which are *expected to change* on each timestep. Below we simulate for 2000 timesteps and plot the angular velocity of the top and height of the stem as a function of time."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "1XXB6asJoZ2N"
+   },
+   "outputs": [],
+   "source": [
+    "timevals = []\n",
+    "angular_velocity = []\n",
+    "stem_height = []\n",
+    "\n",
+    "# Simulate and save data\n",
+    "mujoco.mj_resetDataKeyframe(model, data, 0)\n",
+    "while data.time < duration:\n",
+    "  mujoco.mj_step(model, data)\n",
+    "  timevals.append(data.time)\n",
+    "  angular_velocity.append(data.qvel[3:6].copy())\n",
+    "  stem_height.append(data.geom_xpos[2,2]);\n",
+    "\n",
+    "dpi = 120\n",
+    "width = 600\n",
+    "height = 800\n",
+    "figsize = (width / dpi, height / dpi)\n",
+    "_, ax = plt.subplots(2, 1, figsize=figsize, dpi=dpi, sharex=True)\n",
+    "\n",
+    "ax[0].plot(timevals, angular_velocity)\n",
+    "ax[0].set_title('angular velocity')\n",
+    "ax[0].set_ylabel('radians / second')\n",
+    "\n",
+    "ax[1].plot(timevals, stem_height)\n",
+    "ax[1].set_xlabel('time (seconds)')\n",
+    "ax[1].set_ylabel('meters')\n",
+    "_ = ax[1].set_title('stem height')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "u_zN8vATwcGy"
+   },
+   "source": [
+    "# Example: A chaotic pendulum"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "g1MKUEL_eSCM"
+   },
+   "source": [
+    "Below is a model of a chaotic pendulum, similar to [this one](https://www.exploratorium.edu/exhibits/chaotic-pendulum) in the San Francisco Exploratorium."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "3jHYTV-bwfrS"
+   },
+   "outputs": [],
+   "source": [
+    "chaotic_pendulum = \"\"\"\n",
+    "<mujoco>\n",
+    "  <option timestep=\".001\">\n",
+    "    <flag energy=\"enable\" contact=\"disable\"/>\n",
+    "  </option>\n",
+    "\n",
+    "  <default>\n",
+    "    <joint type=\"hinge\" axis=\"0 -1 0\"/>\n",
+    "    <geom type=\"capsule\" size=\".02\"/>\n",
+    "  </default>\n",
+    "\n",
+    "  <worldbody>\n",
+    "    <light pos=\"0 -.4 1\"/>\n",
+    "    <camera name=\"fixed\" pos=\"0 -1 0\" xyaxes=\"1 0 0 0 0 1\"/>\n",
+    "    <body name=\"0\" pos=\"0 0 .2\">\n",
+    "      <joint name=\"root\"/>\n",
+    "      <geom fromto=\"-.2 0 0 .2 0 0\" rgba=\"1 1 0 1\"/>\n",
+    "      <geom fromto=\"0 0 0 0 0 -.25\" rgba=\"1 1 0 1\"/>\n",
+    "      <body name=\"1\" pos=\"-.2 0 0\">\n",
+    "        <joint/>\n",
+    "        <geom fromto=\"0 0 0 0 0 -.2\" rgba=\"1 0 0 1\"/>\n",
+    "      </body>\n",
+    "      <body name=\"2\" pos=\".2 0 0\">\n",
+    "        <joint/>\n",
+    "        <geom fromto=\"0 0 0 0 0 -.2\" rgba=\"0 1 0 1\"/>\n",
+    "      </body>\n",
+    "      <body name=\"3\" pos=\"0 0 -.25\">\n",
+    "        <joint/>\n",
+    "        <geom fromto=\"0 0 0 0 0 -.2\" rgba=\"0 0 1 1\"/>\n",
+    "      </body>\n",
+    "    </body>\n",
+    "  </worldbody>\n",
+    "</mujoco>\n",
+    "\"\"\"\n",
+    "model = mujoco.MjModel.from_xml_string(chaotic_pendulum)\n",
+    "data = mujoco.MjData(model)\n",
+    "height = 480\n",
+    "width = 640\n",
+    "\n",
+    "with mujoco.Renderer(model, height, width) as renderer:\n",
+    "  mujoco.mj_forward(model, data)\n",
+    "  renderer.update_scene(data, camera=\"fixed\")\n",
+    "\n",
+    "  media.show_image(renderer.render())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "EKZrTBSS5f49"
+   },
+   "source": [
+    "## Timing\n",
+    "Let's see a video of it in action while we time the components:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "-kNWvE9dNwYW"
+   },
+   "outputs": [],
+   "source": [
+    "# setup\n",
+    "n_seconds = 6\n",
+    "framerate = 30  # Hz\n",
+    "n_frames = int(n_seconds * framerate)\n",
+    "frames = []\n",
+    "height = 240\n",
+    "width = 320\n",
+    "\n",
+    "# set initial state\n",
+    "mujoco.mj_resetData(model, data)\n",
+    "data.joint('root').qvel = 10\n",
+    "\n",
+    "# simulate and record frames\n",
+    "frame = 0\n",
+    "sim_time = 0\n",
+    "render_time = 0\n",
+    "n_steps = 0\n",
+    "with mujoco.Renderer(model, height, width) as renderer:\n",
+    "  for i in range(n_frames):\n",
+    "    while data.time * framerate < i:\n",
+    "      tic = time.time()\n",
+    "      mujoco.mj_step(model, data)\n",
+    "      sim_time += time.time() - tic\n",
+    "      n_steps += 1\n",
+    "    tic = time.time()\n",
+    "    renderer.update_scene(data, \"fixed\")\n",
+    "    frame = renderer.render()\n",
+    "    render_time += time.time() - tic\n",
+    "    frames.append(frame)\n",
+    "\n",
+    "# print timing and play video\n",
+    "step_time = 1e6*sim_time/n_steps\n",
+    "step_fps = n_steps/sim_time\n",
+    "print(f'simulation: {step_time:5.3g} μs/step  ({step_fps:5.0f}Hz)')\n",
+    "frame_time = 1e6*render_time/n_frames\n",
+    "frame_fps = n_frames/render_time\n",
+    "print(f'rendering:  {frame_time:5.3g} μs/frame ({frame_fps:5.0f}Hz)')\n",
+    "print('\\n')\n",
+    "\n",
+    "# show video\n",
+    "media.show_video(frames, fps=framerate)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Iqi_m8HT-X5k"
+   },
+   "source": [
+    "Note that rendering is **much** slower than the simulated physics.\n",
+    "\n",
+    "## Chaos\n",
+    "This is a [chaotic](https://en.wikipedia.org/wiki/Chaos_theory) system (small pertubations in initial conditions accumulate quickly):"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "Pa_19EfvOzzg"
+   },
+   "outputs": [],
+   "source": [
+    "PERTURBATION = 1e-7\n",
+    "SIM_DURATION = 10 # seconds\n",
+    "NUM_REPEATS = 8\n",
+    "\n",
+    "# preallocate\n",
+    "n_steps = int(SIM_DURATION / model.opt.timestep)\n",
+    "sim_time = np.zeros(n_steps)\n",
+    "angle = np.zeros(n_steps)\n",
+    "energy = np.zeros(n_steps)\n",
+    "\n",
+    "# prepare plotting axes\n",
+    "_, ax = plt.subplots(2, 1, figsize=(8, 6), sharex=True)\n",
+    "\n",
+    "# simulate NUM_REPEATS times with slightly different initial conditions\n",
+    "for _ in range(NUM_REPEATS):\n",
+    "  # initialize\n",
+    "  mujoco.mj_resetData(model, data)\n",
+    "  data.qvel[0] = 10 # root joint velocity\n",
+    "  # perturb initial velocities\n",
+    "  data.qvel[:] += PERTURBATION * np.random.randn(model.nv)\n",
+    "\n",
+    "  # simulate\n",
+    "  for i in range(n_steps):\n",
+    "    mujoco.mj_step(model, data)\n",
+    "    sim_time[i] = data.time\n",
+    "    angle[i] = data.joint('root').qpos\n",
+    "    energy[i] = data.energy[0] + data.energy[1]\n",
+    "\n",
+    "  # plot\n",
+    "  ax[0].plot(sim_time, angle)\n",
+    "  ax[1].plot(sim_time, energy)\n",
+    "\n",
+    "# finalize plot\n",
+    "ax[0].set_title('root angle')\n",
+    "ax[0].set_ylabel('radian')\n",
+    "ax[1].set_title('total energy')\n",
+    "ax[1].set_ylabel('Joule')\n",
+    "ax[1].set_xlabel('second')\n",
+    "plt.tight_layout()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "daSIA_ewFGxV"
+   },
+   "source": [
+    "## Timestep and accuracy\n",
+    "**Question:** Why is the energy varying at all? There is no friction or damping, this system should conserve energy.\n",
+    "\n",
+    "**Answer:** Because of the discretization of time.\n",
+    "\n",
+    "If we decrease the timestep we'll get better accuracy and better energy conservation:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "4z-7KN_fFme-"
+   },
+   "outputs": [],
+   "source": [
+    "SIM_DURATION = 10 # (seconds)\n",
+    "TIMESTEPS = np.power(10, np.linspace(-2, -4, 5))\n",
+    "\n",
+    "# prepare plotting axes\n",
+    "_, ax = plt.subplots(1, 1)\n",
+    "\n",
+    "for dt in TIMESTEPS:\n",
+    "   # set timestep, print\n",
+    "  model.opt.timestep = dt\n",
+    "\n",
+    "  # allocate\n",
+    "  n_steps = int(SIM_DURATION / model.opt.timestep)\n",
+    "  sim_time = np.zeros(n_steps)\n",
+    "  energy = np.zeros(n_steps)\n",
+    "\n",
+    "  # initialize\n",
+    "  mujoco.mj_resetData(model, data)\n",
+    "  data.qvel[0] = 9 # root joint velocity\n",
+    "\n",
+    "  # simulate\n",
+    "  print('{} steps at dt = {:2.2g}ms'.format(n_steps, 1000*dt))\n",
+    "  for i in range(n_steps):\n",
+    "    mujoco.mj_step(model, data)\n",
+    "    sim_time[i] = data.time\n",
+    "    energy[i] = data.energy[0] + data.energy[1]\n",
+    "\n",
+    "  # plot\n",
+    "  ax.plot(sim_time, energy, label='timestep = {:2.2g}ms'.format(1000*dt))\n",
+    "\n",
+    "# finalize plot\n",
+    "ax.set_title('energy')\n",
+    "ax.set_ylabel('Joule')\n",
+    "ax.set_xlabel('second')\n",
+    "ax.legend(frameon=True);\n",
+    "plt.tight_layout()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "jsVkUm7QKb9I"
+   },
+   "source": [
+    "## Timestep and divergence\n",
+    "When we increase the time step, the simulation quickly diverges:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "FbdUA4zDPbDP"
+   },
+   "outputs": [],
+   "source": [
+    "SIM_DURATION = 10 # (seconds)\n",
+    "TIMESTEPS = np.power(10, np.linspace(-2, -1.5, 7))\n",
+    "\n",
+    "# get plotting axes\n",
+    "ax = plt.gca()\n",
+    "\n",
+    "for dt in TIMESTEPS:\n",
+    "  # set timestep\n",
+    "  model.opt.timestep = dt\n",
+    "\n",
+    "  # allocate\n",
+    "  n_steps = int(SIM_DURATION / model.opt.timestep)\n",
+    "  sim_time = np.zeros(n_steps)\n",
+    "  energy = np.zeros(n_steps) * np.nan\n",
+    "  speed = np.zeros(n_steps) * np.nan\n",
+    "\n",
+    "  # initialize\n",
+    "  mujoco.mj_resetData(model, data)\n",
+    "  data.qvel[0] = 11 # set root joint velocity\n",
+    "\n",
+    "  # simulate\n",
+    "  print('simulating {} steps at dt = {:2.2g}ms'.format(n_steps, 1000*dt))\n",
+    "  for i in range(n_steps):\n",
+    "    mujoco.mj_step(model, data)\n",
+    "    if data.warning.number.any():\n",
+    "      warning_index = np.nonzero(data.warning.number)[0][0]\n",
+    "      warning = mujoco.mjtWarning(warning_index).name\n",
+    "      print(f'stopped due to divergence ({warning}) at timestep {i}.\\n')\n",
+    "      break\n",
+    "    sim_time[i] = data.time\n",
+    "    energy[i] = sum(abs(data.qvel))\n",
+    "    speed[i] = np.linalg.norm(data.qvel)\n",
+    "\n",
+    "  # plot\n",
+    "  ax.plot(sim_time, energy, label='timestep = {:2.2g}ms'.format(1000*dt))\n",
+    "  ax.set_yscale('log')\n",
+    "\n",
+    "# finalize plot\n",
+    "ax.set_ybound(1, 1e3)\n",
+    "ax.set_title('energy')\n",
+    "ax.set_ylabel('Joule')\n",
+    "ax.set_xlabel('second')\n",
+    "ax.legend(frameon=True, loc='lower right');\n",
+    "plt.tight_layout()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "FITYfGyy3XPL"
+   },
+   "source": [
+    "# Contacts\n",
+    "\n",
+    "Let's go back to our box and sphere example and give it a free joint:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "2n1VNVv_FkbB"
+   },
+   "outputs": [],
+   "source": [
+    "free_body_MJCF = \"\"\"\n",
+    "<mujoco>\n",
+    "  <asset>\n",
+    "    <texture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
+    "    rgb2=\".2 .3 .4\" width=\"300\" height=\"300\" mark=\"edge\" markrgb=\".2 .3 .4\"/>\n",
+    "    <material name=\"grid\" texture=\"grid\" texrepeat=\"2 2\" texuniform=\"true\"\n",
+    "    reflectance=\".2\"/>\n",
+    "  </asset>\n",
+    "\n",
+    "  <worldbody>\n",
+    "    <light pos=\"0 0 1\" mode=\"trackcom\"/>\n",
+    "    <geom name=\"ground\" type=\"plane\" pos=\"0 0 -.5\" size=\"2 2 .1\" material=\"grid\" solimp=\".99 .99 .01\" solref=\".001 1\"/>\n",
+    "    <body name=\"box_and_sphere\" pos=\"0 0 0\">\n",
+    "      <freejoint/>\n",
+    "      <geom name=\"red_box\" type=\"box\" size=\".1 .1 .1\" rgba=\"1 0 0 1\" solimp=\".99 .99 .01\"  solref=\".001 1\"/>\n",
+    "      <geom name=\"green_sphere\" size=\".06\" pos=\".1 .1 .1\" rgba=\"0 1 0 1\"/>\n",
+    "      <camera name=\"fixed\" pos=\"0 -.6 .3\" xyaxes=\"1 0 0 0 1 2\"/>\n",
+    "      <camera name=\"track\" pos=\"0 -.6 .3\" xyaxes=\"1 0 0 0 1 2\" mode=\"track\"/>\n",
+    "    </body>\n",
+    "  </worldbody>\n",
+    "</mujoco>\n",
+    "\"\"\"\n",
+    "model = mujoco.MjModel.from_xml_string(free_body_MJCF)\n",
+    "data = mujoco.MjData(model)\n",
+    "height = 400\n",
+    "width = 600\n",
+    "\n",
+    "with mujoco.Renderer(model, height, width) as renderer:\n",
+    "  mujoco.mj_forward(model, data)\n",
+    "  renderer.update_scene(data, \"fixed\")\n",
+    "\n",
+    "  media.show_image(renderer.render())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Z2amdQCn8REu"
+   },
+   "source": [
+    "Let render this body rolling on the floor, in slow-motion, while visualizing contact points and forces:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "HlRhFs_d3WLP"
+   },
+   "outputs": [],
+   "source": [
+    "n_frames = 200\n",
+    "height = 240\n",
+    "width = 320\n",
+    "frames = []\n",
+    "\n",
+    "# visualize contact frames and forces, make body transparent\n",
+    "options = mujoco.MjvOption()\n",
+    "mujoco.mjv_defaultOption(options)\n",
+    "options.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True\n",
+    "options.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = True\n",
+    "options.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = True\n",
+    "\n",
+    "# tweak scales of contact visualization elements\n",
+    "model.vis.scale.contactwidth = 0.1\n",
+    "model.vis.scale.contactheight = 0.03\n",
+    "model.vis.scale.forcewidth = 0.05\n",
+    "model.vis.map.force = 0.3\n",
+    "\n",
+    "# random initial rotational velocity:\n",
+    "mujoco.mj_resetData(model, data)\n",
+    "data.qvel[3:6] = 5*np.random.randn(3)\n",
+    "\n",
+    "# Simulate and display video.\n",
+    "with mujoco.Renderer(model, height, width) as renderer:\n",
+    "  for i in range(n_frames):\n",
+    "    while data.time < i/120.0: #1/4x real time\n",
+    "      mujoco.mj_step(model, data)\n",
+    "    renderer.update_scene(data, \"track\", options)\n",
+    "    frame = renderer.render()\n",
+    "    frames.append(frame)\n",
+    "\n",
+    "media.show_video(frames, fps=30)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "_181TbtVSMBl"
+   },
+   "source": [
+    "## Analysis of contact forces\n",
+    "\n",
+    "Let's rerun the above simulation (with a different random initial condition) and\n",
+    "plot some values related to the contacts"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "BMqyWeHki8Eg"
+   },
+   "outputs": [],
+   "source": [
+    "n_steps = 499\n",
+    "\n",
+    "# allocate\n",
+    "sim_time = np.zeros(n_steps)\n",
+    "ncon = np.zeros(n_steps)\n",
+    "force = np.zeros((n_steps,3))\n",
+    "velocity = np.zeros((n_steps, model.nv))\n",
+    "penetration = np.zeros(n_steps)\n",
+    "acceleration = np.zeros((n_steps, model.nv))\n",
+    "forcetorque = np.zeros(6)\n",
+    "\n",
+    "# random initial rotational velocity:\n",
+    "mujoco.mj_resetData(model, data)\n",
+    "data.qvel[3:6] = 2*np.random.randn(3)\n",
+    "\n",
+    "# simulate and save data\n",
+    "for i in range(n_steps):\n",
+    "  mujoco.mj_step(model, data)\n",
+    "  sim_time[i] = data.time\n",
+    "  ncon[i] = data.ncon\n",
+    "  velocity[i] = data.qvel[:]\n",
+    "  acceleration[i] = data.qacc[:]\n",
+    "  # iterate over active contacts, save force and distance\n",
+    "  for j,c in enumerate(data.contact):\n",
+    "    mujoco.mj_contactForce(model, data, j, forcetorque)\n",
+    "    force[i] += forcetorque[0:3]\n",
+    "    penetration[i] = min(penetration[i], c.dist)\n",
+    "  # we could also do\n",
+    "  # force[i] += data.qfrc_constraint[0:3]\n",
+    "  # do you see why?\n",
+    "\n",
+    "# plot\n",
+    "_, ax = plt.subplots(3, 2, sharex=True, figsize=(10, 10))\n",
+    "\n",
+    "lines = ax[0,0].plot(sim_time, force)\n",
+    "ax[0,0].set_title('contact force')\n",
+    "ax[0,0].set_ylabel('Newton')\n",
+    "ax[0,0].legend(lines, ('normal z', 'friction x', 'friction y'));\n",
+    "\n",
+    "ax[1,0].plot(sim_time, acceleration)\n",
+    "ax[1,0].set_title('acceleration')\n",
+    "ax[1,0].set_ylabel('(meter,radian)/s/s')\n",
+    "ax[1,0].legend(['ax', 'ay', 'az', 'αx', 'αy', 'αz'])\n",
+    "\n",
+    "ax[2,0].plot(sim_time, velocity)\n",
+    "ax[2,0].set_title('velocity')\n",
+    "ax[2,0].set_ylabel('(meter,radian)/s')\n",
+    "ax[2,0].set_xlabel('second')\n",
+    "ax[2,0].legend(['vx', 'vy', 'vz', 'ωx', 'ωy', 'ωz'])\n",
+    "\n",
+    "ax[0,1].plot(sim_time, ncon)\n",
+    "ax[0,1].set_title('number of contacts')\n",
+    "ax[0,1].set_yticks(range(6))\n",
+    "\n",
+    "ax[1,1].plot(sim_time, force[:,0])\n",
+    "ax[1,1].set_yscale('log')\n",
+    "ax[1,1].set_title('normal (z) force - log scale')\n",
+    "ax[1,1].set_ylabel('Newton')\n",
+    "z_gravity = -model.opt.gravity[2]\n",
+    "mg = model.body(\"box_and_sphere\").mass[0] * z_gravity\n",
+    "mg_line = ax[1,1].plot(sim_time, np.ones(n_steps)*mg, label='m*g', linewidth=1)\n",
+    "ax[1,1].legend()\n",
+    "\n",
+    "ax[2,1].plot(sim_time, 1000*penetration)\n",
+    "ax[2,1].set_title('penetration depth')\n",
+    "ax[2,1].set_ylabel('millimeter')\n",
+    "ax[2,1].set_xlabel('second')\n",
+    "\n",
+    "plt.tight_layout()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "zV5PkYzFXu42"
+   },
+   "source": [
+    "## Friction\n",
+    "\n",
+    "Let's see the effect of changing friction values"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "2R_gKoYyXwda"
+   },
+   "outputs": [],
+   "source": [
+    "MJCF = \"\"\"\n",
+    "<mujoco>\n",
+    "  <asset>\n",
+    "    <texture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
+    "     rgb2=\".2 .3 .4\" width=\"300\" height=\"300\" mark=\"none\"/>\n",
+    "    <material name=\"grid\" texture=\"grid\" texrepeat=\"6 6\"\n",
+    "     texuniform=\"true\" reflectance=\".2\"/>\n",
+    "     <material name=\"wall\" rgba='.5 .5 .5 1'/>\n",
+    "  </asset>\n",
+    "\n",
+    "  <default>\n",
+    "    <geom type=\"box\" size=\".05 .05 .05\" />\n",
+    "    <joint type=\"free\"/>\n",
+    "  </default>\n",
+    "\n",
+    "  <worldbody>\n",
+    "    <light name=\"light\" pos=\"-.2 0 1\"/>\n",
+    "    <geom name=\"ground\" type=\"plane\" size=\".5 .5 10\" material=\"grid\"\n",
+    "     zaxis=\"-.3 0 1\" friction=\".1\"/>\n",
+    "    <camera name=\"y\" pos=\"-.1 -.6 .3\" xyaxes=\"1 0 0 0 1 2\"/>\n",
+    "    <body pos=\"0 0 .1\">\n",
+    "      <joint/>\n",
+    "      <geom/>\n",
+    "    </body>\n",
+    "    <body pos=\"0 .2 .1\">\n",
+    "      <joint/>\n",
+    "      <geom friction=\".33\"/>\n",
+    "    </body>\n",
+    "  </worldbody>\n",
+    "\n",
+    "</mujoco>\n",
+    "\"\"\"\n",
+    "n_frames = 60\n",
+    "height = 300\n",
+    "width = 300\n",
+    "frames = []\n",
+    "\n",
+    "# load\n",
+    "model = mujoco.MjModel.from_xml_string(MJCF)\n",
+    "data = mujoco.MjData(model)\n",
+    "\n",
+    "# Simulate and display video.\n",
+    "with mujoco.Renderer(model, height, width) as renderer:\n",
+    "  mujoco.mj_resetData(model, data)\n",
+    "  for i in range(n_frames):\n",
+    "    while data.time < i/30.0:\n",
+    "      mujoco.mj_step(model, data)\n",
+    "    renderer.update_scene(data, \"y\")\n",
+    "    frame = renderer.render()\n",
+    "    frames.append(frame)\n",
+    "\n",
+    "media.show_video(frames, fps=30)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "ArmmaPqGP6W7"
+   },
+   "source": [
+    "# Tendons, actuators and sensors"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "VJz84c97c8Df"
+   },
+   "outputs": [],
+   "source": [
+    "MJCF = \"\"\"\n",
+    "<mujoco>\n",
+    "  <asset>\n",
+    "    <texture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
+    "     rgb2=\".2 .3 .4\" width=\"300\" height=\"300\" mark=\"none\"/>\n",
+    "    <material name=\"grid\" texture=\"grid\" texrepeat=\"1 1\"\n",
+    "     texuniform=\"true\" reflectance=\".2\"/>\n",
+    "  </asset>\n",
+    "\n",
+    "  <worldbody>\n",
+    "    <light name=\"light\" pos=\"0 0 1\"/>\n",
+    "    <geom name=\"floor\" type=\"plane\" pos=\"0 0 -.5\" size=\"2 2 .1\" material=\"grid\"/>\n",
+    "    <site name=\"anchor\" pos=\"0 0 .3\" size=\".01\"/>\n",
+    "    <camera name=\"fixed\" pos=\"0 -1.3 .5\" xyaxes=\"1 0 0 0 1 2\"/>\n",
+    "\n",
+    "    <geom name=\"pole\" type=\"cylinder\" fromto=\".3 0 -.5 .3 0 -.1\" size=\".04\"/>\n",
+    "    <body name=\"bat\" pos=\".3 0 -.1\">\n",
+    "      <joint name=\"swing\" type=\"hinge\" damping=\"1\" axis=\"0 0 1\"/>\n",
+    "      <geom name=\"bat\" type=\"capsule\" fromto=\"0 0 .04 0 -.3 .04\"\n",
+    "       size=\".04\" rgba=\"0 0 1 1\"/>\n",
+    "    </body>\n",
+    "\n",
+    "    <body name=\"box_and_sphere\" pos=\"0 0 0\">\n",
+    "      <joint name=\"free\" type=\"free\"/>\n",
+    "      <geom name=\"red_box\" type=\"box\" size=\".1 .1 .1\" rgba=\"1 0 0 1\"/>\n",
+    "      <geom name=\"green_sphere\"  size=\".06\" pos=\".1 .1 .1\" rgba=\"0 1 0 1\"/>\n",
+    "      <site name=\"hook\" pos=\"-.1 -.1 -.1\" size=\".01\"/>\n",
+    "      <site name=\"IMU\"/>\n",
+    "    </body>\n",
+    "  </worldbody>\n",
+    "\n",
+    "  <tendon>\n",
+    "    <spatial name=\"wire\" limited=\"true\" range=\"0 0.35\" width=\"0.003\">\n",
+    "      <site site=\"anchor\"/>\n",
+    "      <site site=\"hook\"/>\n",
+    "    </spatial>\n",
+    "  </tendon>\n",
+    "\n",
+    "  <actuator>\n",
+    "    <motor name=\"my_motor\" joint=\"swing\" gear=\"1\"/>\n",
+    "  </actuator>\n",
+    "\n",
+    "  <sensor>\n",
+    "    <accelerometer name=\"accelerometer\" site=\"IMU\"/>\n",
+    "  </sensor>\n",
+    "</mujoco>\n",
+    "\"\"\"\n",
+    "model = mujoco.MjModel.from_xml_string(MJCF)\n",
+    "data = mujoco.MjData(model)\n",
+    "height = 480\n",
+    "width = 480\n",
+    "\n",
+    "with mujoco.Renderer(model, height, width) as renderer:\n",
+    "  mujoco.mj_forward(model, data)\n",
+    "  renderer.update_scene(data, \"fixed\")\n",
+    "\n",
+    "  media.show_image(renderer.render())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "u8z2vrOr_RVD"
+   },
+   "source": [
+    "actuated bat and passive \"piñata\":"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "z-zoBCuBv2Xi"
+   },
+   "outputs": [],
+   "source": [
+    "n_frames = 180\n",
+    "height = 240\n",
+    "width = 320\n",
+    "frames = []\n",
+    "fps = 60.0\n",
+    "times = []\n",
+    "sensordata = []\n",
+    "\n",
+    "# constant actuator signal\n",
+    "mujoco.mj_resetData(model, data)\n",
+    "data.ctrl = 20\n",
+    "\n",
+    "# Simulate and display video.\n",
+    "with mujoco.Renderer(model, height, width) as renderer:\n",
+    "  for i in range(n_frames):\n",
+    "    while data.time < i/fps:\n",
+    "      mujoco.mj_step(model, data)\n",
+    "      times.append(data.time)\n",
+    "      sensordata.append(data.sensor('accelerometer').data.copy())\n",
+    "    renderer.update_scene(data, \"fixed\")\n",
+    "    frame = renderer.render()\n",
+    "    frames.append(frame)\n",
+    "\n",
+    "media.show_video(frames, fps=fps)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "gwHMy_iRA7Jh"
+   },
+   "source": [
+    "Let's plot the values measured by our accelerometer sensor:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "uy4wSEMAAJgn"
+   },
+   "outputs": [],
+   "source": [
+    "ax = plt.gca()\n",
+    "\n",
+    "ax.plot(np.asarray(times), np.asarray(sensordata), label=[f\"axis {v}\" for v in ['x', 'y', 'z']])\n",
+    "\n",
+    "# finalize plot\n",
+    "ax.set_title('Accelerometer values')\n",
+    "ax.set_ylabel('meter/second^2')\n",
+    "ax.set_xlabel('second')\n",
+    "ax.legend(frameon=True, loc='lower right')\n",
+    "plt.tight_layout()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "0YKSTtJ_BQ7x"
+   },
+   "source": [
+    "Note how the moments when the body is hit by the bat are clearly visible in the accelerometer measurements."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "1kOs1wTc7uCZ"
+   },
+   "source": [
+    "# Advanced rendering\n",
+    "\n",
+    "Like joint visualization, additional rendering options are exposed as parameters to the `render` method.\n",
+    "\n",
+    "Let's bring back our first model:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "mTDgsk2xcgwH"
+   },
+   "outputs": [],
+   "source": [
+    "xml = \"\"\"\n",
+    "<mujoco>\n",
+    "  <worldbody>\n",
+    "    <light name=\"top\" pos=\"0 0 1\"/>\n",
+    "    <body name=\"box_and_sphere\" euler=\"0 0 -30\">\n",
+    "      <joint name=\"swing\" type=\"hinge\" axis=\"1 -1 0\" pos=\"-.2 -.2 -.2\"/>\n",
+    "      <geom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/>\n",
+    "      <geom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/>\n",
+    "    </body>\n",
+    "  </worldbody>\n",
+    "</mujoco>\n",
+    "\"\"\"\n",
+    "model = mujoco.MjModel.from_xml_string(xml)\n",
+    "data = mujoco.MjData(model)\n",
+    "\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  mujoco.mj_forward(model, data)\n",
+    "  renderer.update_scene(data)\n",
+    "  media.show_image(renderer.render())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "VePXamL_6XUc"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Enable transparency and frame visualization {vertical-output: true}\n",
+    "\n",
+    "scene_option.frame = mujoco.mjtFrame.mjFRAME_GEOM\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = True\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  renderer.update_scene(data, scene_option=scene_option)\n",
+    "  frame = renderer.render()\n",
+    "  media.show_image(frame)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "PVcpcvww9lZ8"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Depth rendering {vertical-output: true}\n",
+    "\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  # update renderer to render depth\n",
+    "  renderer.enable_depth_rendering()\n",
+    "\n",
+    "  # reset the scene\n",
+    "  renderer.update_scene(data)\n",
+    "\n",
+    "  # depth is a float array, in meters.\n",
+    "  depth = renderer.render()\n",
+    "\n",
+    "  # Shift nearest values to the origin.\n",
+    "  depth -= depth.min()\n",
+    "  # Scale by 2 mean distances of near rays.\n",
+    "  depth /= 2*depth[depth <= 1].mean()\n",
+    "  # Scale to [0, 255]\n",
+    "  pixels = 255*np.clip(depth, 0, 1)\n",
+    "\n",
+    "  media.show_image(pixels.astype(np.uint8))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "PNwiIrgpx7T8"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Segmentation rendering {vertical-output: true}\n",
+    "\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  renderer.disable_depth_rendering()\n",
+    "\n",
+    "  # update renderer to render segmentation\n",
+    "  renderer.enable_segmentation_rendering()\n",
+    "\n",
+    "  # reset the scene\n",
+    "  renderer.update_scene(data)\n",
+    "\n",
+    "  seg = renderer.render()\n",
+    "\n",
+    "  # Display the contents of the first channel, which contains object\n",
+    "  # IDs. The second channel, seg[:, :, 1], contains object types.\n",
+    "  geom_ids = seg[:, :, 0]\n",
+    "  # Infinity is mapped to -1\n",
+    "  geom_ids = geom_ids.astype(np.float64) + 1\n",
+    "  # Scale to [0, 1]\n",
+    "  geom_ids = geom_ids / geom_ids.max()\n",
+    "  pixels = 255*geom_ids\n",
+    "  media.show_image(pixels.astype(np.uint8))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "wo72mo0mGIXr"
+   },
+   "source": [
+    "## The camera matrix\n",
+    "\n",
+    "For a description of the camera matrix see the article [Camera matrix](https://en.wikipedia.org/wiki/Camera_matrix) on Wikipedia."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "sDYwClpxaxab"
+   },
+   "outputs": [],
+   "source": [
+    "def compute_camera_matrix(renderer, data):\n",
+    "  \"\"\"Returns the 3x4 camera matrix.\"\"\"\n",
+    "  # If the camera is a 'free' camera, we get its position and orientation\n",
+    "  # from the scene data structure. It is a stereo camera, so we average over\n",
+    "  # the left and right channels. Note: we call `self.update()` in order to\n",
+    "  # ensure that the contents of `scene.camera` are correct.\n",
+    "  renderer.update_scene(data)\n",
+    "  pos = np.mean([camera.pos for camera in renderer.scene.camera], axis=0)\n",
+    "  z = -np.mean([camera.forward for camera in renderer.scene.camera], axis=0)\n",
+    "  y = np.mean([camera.up for camera in renderer.scene.camera], axis=0)\n",
+    "  rot = np.vstack((np.cross(y, z), y, z))\n",
+    "  fov = model.vis.global_.fovy\n",
+    "\n",
+    "  # Translation matrix (4x4).\n",
+    "  translation = np.eye(4)\n",
+    "  translation[0:3, 3] = -pos\n",
+    "\n",
+    "  # Rotation matrix (4x4).\n",
+    "  rotation = np.eye(4)\n",
+    "  rotation[0:3, 0:3] = rot\n",
+    "\n",
+    "  # Focal transformation matrix (3x4).\n",
+    "  focal_scaling = (1./np.tan(np.deg2rad(fov)/2)) * renderer.height / 2.0\n",
+    "  focal = np.diag([-focal_scaling, focal_scaling, 1.0, 0])[0:3, :]\n",
+    "\n",
+    "  # Image matrix (3x3).\n",
+    "  image = np.eye(3)\n",
+    "  image[0, 2] = (renderer.width - 1) / 2.0\n",
+    "  image[1, 2] = (renderer.height - 1) / 2.0\n",
+    "  return image @ focal @ rotation @ translation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "My0N4_7PDJ_q"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Project from world to camera coordinates {vertical-output: true}\n",
+    "\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  renderer.disable_segmentation_rendering()\n",
+    "  # reset the scene\n",
+    "  renderer.update_scene(data)\n",
+    "\n",
+    "  # Get the world coordinates of the box corners\n",
+    "  box_pos = data.geom_xpos[model.geom('red_box').id]\n",
+    "  box_mat = data.geom_xmat[model.geom('red_box').id].reshape(3, 3)\n",
+    "  box_size = model.geom_size[model.geom('red_box').id]\n",
+    "  offsets = np.array([-1, 1]) * box_size[:, None]\n",
+    "  xyz_local = np.stack(list(itertools.product(*offsets))).T\n",
+    "  xyz_global = box_pos[:, None] + box_mat @ xyz_local\n",
+    "\n",
+    "  # Camera matrices multiply homogenous [x, y, z, 1] vectors.\n",
+    "  corners_homogeneous = np.ones((4, xyz_global.shape[1]), dtype=float)\n",
+    "  corners_homogeneous[:3, :] = xyz_global\n",
+    "\n",
+    "  # Get the camera matrix.\n",
+    "  m = compute_camera_matrix(renderer, data)\n",
+    "\n",
+    "  # Project world coordinates into pixel space. See:\n",
+    "  # https://en.wikipedia.org/wiki/3D_projection#Mathematical_formula\n",
+    "  xs, ys, s = m @ corners_homogeneous\n",
+    "  # x and y are in the pixel coordinate system.\n",
+    "  x = xs / s\n",
+    "  y = ys / s\n",
+    "\n",
+    "  # Render the camera view and overlay the projected corner coordinates.\n",
+    "  pixels = renderer.render()\n",
+    "  fig, ax = plt.subplots(1, 1)\n",
+    "  ax.imshow(pixels)\n",
+    "  ax.plot(x, y, '+', c='w')\n",
+    "  ax.set_axis_off()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "AGm5-e0sHEAF"
+   },
+   "source": [
+    "## Modifying the scene\n",
+    "\n",
+    "Let's add some arbitrary geometry to the `mjvScene`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "Z6NDYJ8IOVt7"
+   },
+   "outputs": [],
+   "source": [
+    "def get_geom_speed(model, data, geom_name):\n",
+    "  \"\"\"Returns the speed of a geom.\"\"\"\n",
+    "  geom_vel = np.zeros(6)\n",
+    "  geom_type = mujoco.mjtObj.mjOBJ_GEOM\n",
+    "  geom_id = data.geom(geom_name).id\n",
+    "  mujoco.mj_objectVelocity(model, data, geom_type, geom_id, geom_vel, 0)\n",
+    "  return np.linalg.norm(geom_vel)\n",
+    "\n",
+    "def add_visual_capsule(scene, point1, point2, radius, rgba):\n",
+    "  \"\"\"Adds one capsule to an mjvScene.\"\"\"\n",
+    "  if scene.ngeom >= scene.maxgeom:\n",
+    "    return\n",
+    "  scene.ngeom += 1  # increment ngeom\n",
+    "  # initialise a new capsule, add it to the scene using mjv_connector\n",
+    "  mujoco.mjv_initGeom(scene.geoms[scene.ngeom-1],\n",
+    "                      mujoco.mjtGeom.mjGEOM_CAPSULE, np.zeros(3),\n",
+    "                      np.zeros(3), np.zeros(9), rgba.astype(np.float32))\n",
+    "  mujoco.mjv_connector(scene.geoms[scene.ngeom-1],\n",
+    "                       mujoco.mjtGeom.mjGEOM_CAPSULE, radius,\n",
+    "                       point1, point2)\n",
+    "\n",
+    " # traces of time, position and speed\n",
+    "times = []\n",
+    "positions = []\n",
+    "speeds = []\n",
+    "offset = model.jnt_axis[0]/16  # offset along the joint axis\n",
+    "\n",
+    "def modify_scene(scn):\n",
+    "  \"\"\"Draw position trace, speed modifies width and colors.\"\"\"\n",
+    "  if len(positions) > 1:\n",
+    "    for i in range(len(positions)-1):\n",
+    "      rgba=np.array((np.clip(speeds[i]/10, 0, 1),\n",
+    "                     np.clip(1-speeds[i]/10, 0, 1),\n",
+    "                     .5, 1.))\n",
+    "      radius=.003*(1+speeds[i])\n",
+    "      point1 = positions[i] + offset*times[i]\n",
+    "      point2 = positions[i+1] + offset*times[i+1]\n",
+    "      add_visual_capsule(scn, point1, point2, radius, rgba)\n",
+    "\n",
+    "duration = 6    # (seconds)\n",
+    "framerate = 30  # (Hz)\n",
+    "\n",
+    "# Simulate and display video.\n",
+    "frames = []\n",
+    "\n",
+    "# Reset state and time.\n",
+    "mujoco.mj_resetData(model, data)\n",
+    "mujoco.mj_forward(model, data)\n",
+    "\n",
+    "with mujoco.Renderer(model) as renderer:\n",
+    "  while data.time < duration:\n",
+    "    # append data to the traces\n",
+    "    positions.append(data.geom_xpos[data.geom(\"green_sphere\").id].copy())\n",
+    "    times.append(data.time)\n",
+    "    speeds.append(get_geom_speed(model, data, \"green_sphere\"))\n",
+    "    mujoco.mj_step(model, data)\n",
+    "    if len(frames) < data.time * framerate:\n",
+    "      renderer.update_scene(data)\n",
+    "      modify_scene(renderer.scene)\n",
+    "      pixels = renderer.render()\n",
+    "      frames.append(pixels)\n",
+    "\n",
+    "media.show_video(frames, fps=framerate)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "p6wHrjRg1EGF"
+   },
+   "source": [
+    "## Multiple frames in the same scene\n",
+    "\n",
+    "Sometimes one would like to draw the same geometry multiple times, for example when a model is tracking states from motion-capture, it's nice to have the data\n",
+    "visualized next to the model. Unlike `mjv_updateScene` (called by the `Renderer`'s `update_scene` method) which clears the scene at every call, `mjv_addGeoms` will add visual geoms to an existing scene:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "T4b_8n6t1ASk"
+   },
+   "outputs": [],
+   "source": [
+    "# Get MuJoCo's standard humanoid model.\n",
+    "print('Getting MuJoCo humanoid XML description from GitHub:')\n",
+    "!git clone https://github.com/google-deepmind/mujoco\n",
+    "with open('mujoco/model/humanoid/humanoid.xml', 'r') as f:\n",
+    "  xml = f.read()\n",
+    "\n",
+    "# Load the model, make two MjData's.\n",
+    "model = mujoco.MjModel.from_xml_string(xml)\n",
+    "data = mujoco.MjData(model)\n",
+    "data2 = mujoco.MjData(model)\n",
+    "\n",
+    "# Episode parameters.\n",
+    "duration = 3       # (seconds)\n",
+    "framerate = 60     # (Hz)\n",
+    "data.qpos[0:2] = [-.5, -.5]  # Initial x-y position (m)\n",
+    "data.qvel[2] = 4   # Initial vertical velocity (m/s)\n",
+    "ctrl_phase = 2 * np.pi * np.random.rand(model.nu)  # Control phase\n",
+    "ctrl_freq = 1     # Control frequency\n",
+    "\n",
+    "# Visual options for the \"ghost\" model.\n",
+    "vopt2 = mujoco.MjvOption()\n",
+    "vopt2.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = True  # Transparent.\n",
+    "pert = mujoco.MjvPerturb()  # Empty MjvPerturb object\n",
+    "# We only want dynamic objects (the humanoid). Static objects (the floor)\n",
+    "# should not be re-drawn. The mjtCatBit flag lets us do that, though we could\n",
+    "# equivalently use mjtVisFlag.mjVIS_STATIC\n",
+    "catmask = mujoco.mjtCatBit.mjCAT_DYNAMIC\n",
+    "\n",
+    "# Simulate and render.\n",
+    "frames = []\n",
+    "with mujoco.Renderer(model, 480, 640) as renderer:\n",
+    "  while data.time < duration:\n",
+    "    # Sinusoidal control signal.\n",
+    "    data.ctrl = np.sin(ctrl_phase + 2 * np.pi * data.time * ctrl_freq)\n",
+    "    mujoco.mj_step(model, data)\n",
+    "    if len(frames) < data.time * framerate:\n",
+    "      # This draws the regular humanoid from `data`.\n",
+    "      renderer.update_scene(data)\n",
+    "\n",
+    "      # Copy qpos to data2, move the humanoid sideways, call mj_forward.\n",
+    "      data2.qpos = data.qpos\n",
+    "      data2.qpos[0] += 1.5\n",
+    "      data2.qpos[1] += 1\n",
+    "      mujoco.mj_forward(model, data2)\n",
+    "\n",
+    "      # Call mjv_addGeoms to add the ghost humanoid to the scene.\n",
+    "      mujoco.mjv_addGeoms(model, data2, vopt2, pert, catmask, renderer.scene)\n",
+    "\n",
+    "      # Render and add the frame.\n",
+    "      pixels = renderer.render()\n",
+    "      frames.append(pixels)\n",
+    "\n",
+    "# Render video at half real-time.\n",
+    "media.show_video(frames, fps=framerate/2)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Zzzugf-qPExb"
+   },
+   "source": [
+    "## Camera control\n",
+    "\n",
+    "Cameras can be controlled dynamically in order to achieve cinematic effects. Run the three cells below to see the difference between rendering from a static and moving camera.\n",
+    "\n",
+    "The camera-control code smoothly transitions between two trajectories, one orbiting a fixed point, the other tracking a moving object. Parameter values in the code were obtained by iterating quickly on low-res videos."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "cellView": "form",
+    "id": "-SW-K9WuPGrp"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Load the \"dominos\" model\n",
+    "\n",
+    "dominos_xml = \"\"\"\n",
+    "<mujoco>\n",
+    "  <asset>\n",
+    "    <texture type=\"skybox\" builtin=\"gradient\" rgb1=\".3 .5 .7\" rgb2=\"0 0 0\" width=\"32\" height=\"512\"/>\n",
+    "    <texture name=\"grid\" type=\"2d\" builtin=\"checker\" width=\"512\" height=\"512\" rgb1=\".1 .2 .3\" rgb2=\".2 .3 .4\"/>\n",
+    "    <material name=\"grid\" texture=\"grid\" texrepeat=\"2 2\" texuniform=\"true\" reflectance=\".2\"/>\n",
+    "  </asset>\n",
+    "\n",
+    "  <statistic meansize=\".01\"/>\n",
+    "\n",
+    "  <visual>\n",
+    "    <global offheight=\"2160\" offwidth=\"3840\"/>\n",
+    "    <quality offsamples=\"8\"/>\n",
+    "  </visual>\n",
+    "\n",
+    "  <default>\n",
+    "    <geom type=\"box\" solref=\".005 1\"/>\n",
+    "    <default class=\"static\">\n",
+    "      <geom rgba=\".3 .5 .7 1\"/>\n",
+    "    </default>\n",
+    "  </default>\n",
+    "\n",
+    "  <option timestep=\"5e-4\"/>\n",
+    "\n",
+    "  <worldbody>\n",
+    "    <light pos=\".3 -.3 .8\" mode=\"trackcom\" diffuse=\"1 1 1\" specular=\".3 .3 .3\"/>\n",
+    "    <light pos=\"0 -.3 .4\" mode=\"targetbodycom\" target=\"box\" diffuse=\".8 .8 .8\" specular=\".3 .3 .3\"/>\n",
+    "    <geom name=\"floor\" type=\"plane\" size=\"3 3 .01\" pos=\"-0.025 -0.295  0\" material=\"grid\"/>\n",
+    "    <geom name=\"ramp\" pos=\".25 -.45 -.03\" size=\".04 .1 .07\" euler=\"-30 0 0\" class=\"static\"/>\n",
+    "    <camera name=\"top\" pos=\"-0.37 -0.78 0.49\" xyaxes=\"0.78 -0.63 0 0.27 0.33 0.9\"/>\n",
+    "\n",
+    "    <body name=\"ball\" pos=\".25 -.45 .1\">\n",
+    "      <freejoint name=\"ball\"/>\n",
+    "      <geom name=\"ball\" type=\"sphere\" size=\".02\" rgba=\".65 .81 .55 1\"/>\n",
+    "    </body>\n",
+    "\n",
+    "    <body pos=\".26 -.3 .03\" euler=\"0 0 -90.0\">\n",
+    "      <freejoint/>\n",
+    "      <geom size=\".0015 .015 .03\" rgba=\"1 .5 .5 1\"/>\n",
+    "    </body>\n",
+    "\n",
+    "    <body pos=\".26 -.27 .04\" euler=\"0 0 -81.0\">\n",
+    "      <freejoint/>\n",
+    "      <geom size=\".002 .02 .04\" rgba=\"1 1 .5 1\"/>\n",
+    "    </body>\n",
+    "\n",
+    "    <body pos=\".24 -.21 .06\" euler=\"0 0 -63.0\">\n",
+    "      <freejoint/>\n",
+    "      <geom size=\".003 .03 .06\" rgba=\".5 1 .5 1\"/>\n",
+    "    </body>\n",
+    "\n",
+    "    <body pos=\".2 -.16 .08\" euler=\"0 0 -45.0\">\n",
+    "      <freejoint/>\n",
+    "      <geom size=\".004 .04 .08\" rgba=\".5 1 1 1\"/>\n",
+    "    </body>\n",
+    "\n",
+    "    <body pos=\".15 -.12 .1\" euler=\"0 0 -27.0\">\n",
+    "      <freejoint/>\n",
+    "      <geom size=\".005 .05 .1\" rgba=\".5 .5 1 1\"/>\n",
+    "    </body>\n",
+    "\n",
+    "    <body pos=\".09 -.1 .12\" euler=\"0 0 -9.0\">\n",
+    "      <freejoint/>\n",
+    "      <geom size=\".006 .06 .12\" rgba=\"1 .5 1 1\"/>\n",
+    "    </body>\n",
+    "\n",
+    "    <body name=\"seasaw_wrapper\" pos=\"-.23 -.1 0\" euler=\"0 0 30\">\n",
+    "      <geom size=\".01 .01 .015\" pos=\"0 .05 .015\" class=\"static\"/>\n",
+    "      <geom size=\".01 .01 .015\" pos=\"0 -.05 .015\" class=\"static\"/>\n",
+    "      <geom type=\"cylinder\" size=\".01 .0175\" pos=\"-.09 0 .0175\" class=\"static\"/>\n",
+    "      <body name=\"seasaw\" pos=\"0 0 .03\">\n",
+    "        <joint axis=\"0 1 0\"/>\n",
+    "        <geom type=\"cylinder\" size=\".005 .039\" zaxis=\"0 1 0\" rgba=\".84 .15 .33 1\"/>\n",
+    "        <geom size=\".1 .02 .005\" pos=\"0 0 .01\" rgba=\".84 .15 .33 1\"/>\n",
+    "      </body>\n",
+    "    </body>\n",
+    "\n",
+    "    <body name=\"box\" pos=\"-.3 -.14 .05501\" euler=\"0 0 -30\">\n",
+    "      <freejoint name=\"box\"/>\n",
+    "      <geom name=\"box\" size=\".01 .01 .01\" rgba=\".0 .7 .79 1\"/>\n",
+    "    </body>\n",
+    "  </worldbody>\n",
+    "</mujoco>\n",
+    "\"\"\"\n",
+    "model = mujoco.MjModel.from_xml_string(dominos_xml)\n",
+    "data = mujoco.MjData(model)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "cellView": "form",
+    "id": "a2WruafiPhPk"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Render from fixed camera\n",
+    "\n",
+    "duration = 2.5  # (seconds)\n",
+    "framerate = 60  # (Hz)\n",
+    "height = 1024\n",
+    "width = 1440\n",
+    "\n",
+    "# Simulate and display video.\n",
+    "frames = []\n",
+    "mujoco.mj_resetData(model, data)  # Reset state and time.\n",
+    "with mujoco.Renderer(model, height, width) as renderer:\n",
+    "  while data.time < duration:\n",
+    "    mujoco.mj_step(model, data)\n",
+    "    if len(frames) < data.time * framerate:\n",
+    "      renderer.update_scene(data, camera='top')\n",
+    "      pixels = renderer.render()\n",
+    "      frames.append(pixels)\n",
+    "\n",
+    "media.show_video(frames, fps=framerate)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "cellView": "form",
+    "id": "Kie3y-27bQ3J"
+   },
+   "outputs": [],
+   "source": [
+    "#@title Render from moving camera\n",
+    "\n",
+    "duration = 3  # (seconds)\n",
+    "height = 1024\n",
+    "width = 1440\n",
+    "\n",
+    "# find time when box is thrown (speed > 2cm/s)\n",
+    "throw_time = 0.0\n",
+    "mujoco.mj_resetData(model, data)\n",
+    "while data.time < duration and not throw_time:\n",
+    "  mujoco.mj_step(model, data)\n",
+    "  box_speed = np.linalg.norm(data.joint('box').qvel[:3])\n",
+    "  if box_speed > 0.02:\n",
+    "    throw_time = data.time\n",
+    "assert throw_time > 0\n",
+    "\n",
+    "def mix(time, t0=0.0, width=1.0):\n",
+    "  \"\"\"Sigmoidal mixing function.\"\"\"\n",
+    "  t = (time - t0) / width\n",
+    "  s = 1 / (1 + np.exp(-t))\n",
+    "  return 1 - s, s\n",
+    "\n",
+    "def unit_cos(t):\n",
+    "  \"\"\"Unit cosine sigmoid from (0,0) to (1,1).\"\"\"\n",
+    "  return 0.5 - np.cos(np.pi*np.clip(t, 0, 1))/2\n",
+    "\n",
+    "def orbit_motion(t):\n",
+    "  \"\"\"Return orbit trajectory.\"\"\"\n",
+    "  distance = 0.9\n",
+    "  azimuth = 140 + 100 * unit_cos(t)\n",
+    "  elevation = -30\n",
+    "  lookat = data.geom('floor').xpos.copy()\n",
+    "  return distance, azimuth, elevation, lookat\n",
+    "\n",
+    "def track_motion():\n",
+    "  \"\"\"Return box-track trajectory.\"\"\"\n",
+    "  distance = 0.08\n",
+    "  azimuth = 280\n",
+    "  elevation = -10\n",
+    "  lookat = data.geom('box').xpos.copy()\n",
+    "  return distance, azimuth, elevation, lookat\n",
+    "\n",
+    "def cam_motion():\n",
+    "  \"\"\"Return sigmoidally-mixed {orbit, box-track} trajectory.\"\"\"\n",
+    "  d0, a0, e0, l0 = orbit_motion(data.time / throw_time)\n",
+    "  d1, a1, e1, l1 = track_motion()\n",
+    "  mix_time = 0.3\n",
+    "  w0, w1 = mix(data.time, throw_time, mix_time)\n",
+    "  return w0*d0+w1*d1, w0*a0+w1*a1, w0*e0+w1*e1, w0*l0+w1*l1\n",
+    "\n",
+    "# Make a camera.\n",
+    "cam = mujoco.MjvCamera()\n",
+    "mujoco.mjv_defaultCamera(cam)\n",
+    "\n",
+    "# Simulate and display video.\n",
+    "framerate = 60  # (Hz)\n",
+    "slowdown = 4    # 4x slow-down\n",
+    "mujoco.mj_resetData(model, data)\n",
+    "frames = []\n",
+    "with mujoco.Renderer(model, height, width) as renderer:\n",
+    "  while data.time < duration:\n",
+    "    mujoco.mj_step(model, data)\n",
+    "    if len(frames) < data.time * framerate * slowdown:\n",
+    "      cam.distance, cam.azimuth, cam.elevation, cam.lookat = cam_motion()\n",
+    "      renderer.update_scene(data, cam)\n",
+    "      pixels = renderer.render()\n",
+    "      frames.append(pixels)\n",
+    "\n",
+    "media.show_video(frames, fps=framerate)"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "collapsed_sections": [
+    "YvyGCsgSCxHQ"
+   ],
+   "gpuClass": "premium",
+   "private_outputs": true
+  },
+  "gpuClass": "premium",
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}