From e42f9c649073db37f15bf35bc1cc3768b299164d Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Wed, 13 Mar 2024 04:41:51 -0700
Subject: [PATCH] Tweaks to basic tutorial notebook.

PiperOrigin-RevId: 615368113
Change-Id: I888b252c75a1f6253dc78494e8348d00e8283ed7
---
 python/LQR.ipynb      |  1 +
 python/tutorial.ipynb | 80 +++++++++++--------------------------------
 2 files changed, 21 insertions(+), 60 deletions(-)

diff --git a/python/LQR.ipynb b/python/LQR.ipynb
index 23eaecd8..afc9b04a 100644
--- a/python/LQR.ipynb
+++ b/python/LQR.ipynb
@@ -935,6 +935,7 @@
   "accelerator": "GPU",
   "colab": {
    "collapsed_sections": [
+    "LBAvTJ0xHKy7",
     "QPdJNe3k62mx"
    ],
    "private_outputs": true,
diff --git a/python/tutorial.ipynb b/python/tutorial.ipynb
index ee50e564..2c5e681e 100644
--- a/python/tutorial.ipynb
+++ b/python/tutorial.ipynb
@@ -10,17 +10,7 @@
     "\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",
-    "**A Colab runtime with GPU acceleration is required.** If you're using a CPU-only runtime, you can switch using the menu \"Runtime > Change runtime type\".\n",
-    "\n",
-    "\n",
-    "\n",
-    "\n",
-    "\n",
-    "\n",
-    "\n",
-    "\n"
+    "This notebook provides an introductory tutorial for [**MuJoCo** physics](https://github.com/google-deepmind/mujoco#readme), using the native Python bindings."
    ]
   },
   {
@@ -49,7 +39,7 @@
     "id": "YvyGCsgSCxHQ"
    },
    "source": [
-    "# Install MuJoCo"
+    "# All imports"
    ]
   },
   {
@@ -60,22 +50,10 @@
    },
    "outputs": [],
    "source": [
-    "!pip install mujoco"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 0,
-   "metadata": {
-    "cellView": "form",
-    "id": "IbZxYDxzoz5R"
-   },
-   "outputs": [],
-   "source": [
-    "#@title Set up rendering, check installation\n",
+    "!pip install mujoco\n",
     "\n",
+    "# Set up GPU rendering.\n",
     "from google.colab import files\n",
-    "\n",
     "import distutils.util\n",
     "import os\n",
     "import subprocess\n",
@@ -104,6 +82,7 @@
     "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",
@@ -115,23 +94,12 @@
     "      '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.')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 0,
-   "metadata": {
-    "cellView": "form",
-    "id": "T5f4w3Kq2X14"
-   },
-   "outputs": [],
-   "source": [
-    "#@title Import packages for plotting and creating graphics\n",
+    "print('Installation successful.')\n",
+    "\n",
+    "# Other imports and helper functions\n",
     "import time\n",
     "import itertools\n",
     "import numpy as np\n",
-    "from typing import Callable, NamedTuple, Optional, Union, List\n",
     "\n",
     "# Graphics and plotting.\n",
     "print('Installing mediapy:')\n",
@@ -141,7 +109,10 @@
     "import matplotlib.pyplot as plt\n",
     "\n",
     "# More legible printing from numpy.\n",
-    "np.set_printoptions(precision=3, suppress=True, linewidth=100)"
+    "np.set_printoptions(precision=3, suppress=True, linewidth=100)\n",
+    "\n",
+    "from IPython.display import clear_output\n",
+    "clear_output()\n"
    ]
   },
   {
@@ -1198,7 +1169,7 @@
     "  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]\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",
@@ -1686,7 +1657,7 @@
    },
    "outputs": [],
    "source": [
-    "#@title Enable transparency and frame visualization\n",
+    "#@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",
@@ -1703,7 +1674,7 @@
    },
    "outputs": [],
    "source": [
-    "#@title Depth rendering\n",
+    "#@title Depth rendering {vertical-output: true}\n",
     "\n",
     "# update renderer to render depth\n",
     "renderer.enable_depth_rendering()\n",
@@ -1734,7 +1705,7 @@
    },
    "outputs": [],
    "source": [
-    "#@title Segmentation rendering\n",
+    "#@title Segmentation rendering {vertical-output: true}\n",
     "\n",
     "# update renderer to render segmentation\n",
     "renderer.enable_segmentation_rendering()\n",
@@ -1808,15 +1779,6 @@
     "  return image @ focal @ rotation @ translation"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "Bs89vS0wLoU0"
-   },
-   "source": [
-    "Let's use the camera matrix to project from world to camera coordinates:"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 0,
@@ -1825,6 +1787,7 @@
    },
    "outputs": [],
    "source": [
+    "#@title Project from world to camera coordinates {vertical-output: true}\n",
     "# reset the scene\n",
     "renderer.update_scene(data)\n",
     "\n",
@@ -1904,7 +1867,7 @@
     "times = []\n",
     "positions = []\n",
     "speeds = []\n",
-    "offset = model.jnt_axis[0]/8  # offset along the joint axis\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",
@@ -2158,11 +2121,8 @@
  "metadata": {
   "accelerator": "GPU",
   "colab": {
-   "collapsed_sections": [
-    "-re3Szx-1Ias"
-   ],
-   "private_outputs": true,
-   "toc_visible": true
+   "gpuClass": "premium",
+   "private_outputs": true
   },
   "gpuClass": "premium",
   "kernelspec": {