Tweaks to LQR tutorial notebook:

- Unify all imports in one cell.
- Make final motion more interesting by adding more noise to joints which are not critical to balancing.

PiperOrigin-RevId: 615100505
Change-Id: I8302e441b3ca52d5134c317a35f99c5f74bac36b

python/LQR.ipynb

@@ -10,9 +10,7 @@
     "\n",
     "# <h1><center>LQR tutorial  <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/LQR.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 example of an LQR controller using [**MuJoCo** physics](https://github.com/google-deepmind/mujoco#readme).\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"
+    "This notebook provides an example of an LQR controller using [**MuJoCo** physics](https://github.com/google-deepmind/mujoco#readme)."
    ]
   },
   {
@@ -41,7 +39,7 @@
     "id": "QPdJNe3k62mx"
    },
    "source": [
-    "### Install MuJoCo\n"
+    "### All imports\n"
    ]
   },
   {
@@ -52,22 +50,11 @@
    },
    "outputs": [],
    "source": [
-    "!pip install mujoco"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 0,
-   "metadata": {
-    "cellView": "form",
-    "id": "IbZxYDxzoz5R"
-   },
-   "outputs": [],
-   "source": [
-    "#@title Check if installation was successful\n",
     "\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",
@@ -96,6 +83,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",
@@ -107,19 +95,9 @@
     "      '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 Other imports and helper functions\n",
+    "print('Installation successful.')\n",
+    "\n",
+    "# Other imports and helper functions\n",
     "import numpy as np\n",
     "from typing import Callable, Optional, Union, List\n",
     "import scipy.linalg\n",
@@ -132,7 +110,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"
    ]
   },
   {
@@ -141,7 +122,7 @@
     "id": "J5fL6p-Sx5DB"
    },
    "source": [
-    "## Loading and rendering the standard humanoid"
+    "## Load and render the standard humanoid"
    ]
   },
   {
@@ -152,10 +133,11 @@
    },
    "outputs": [],
    "source": [
+    "\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()"
+    "  xml = f.read()\n"
    ]
   },
   {
@@ -875,8 +857,9 @@
     "DURATION = 12         # seconds\n",
     "FRAMERATE = 60        # Hz\n",
     "TOTAL_ROTATION = 15   # degrees\n",
-    "CTRL_STD = 0.05       # actuator units\n",
     "CTRL_RATE = 0.8       # seconds\n",
+    "BALANCE_STD = 0.01    # actuator units\n",
+    "OTHER_STD = 0.08      # actuator units\n",
     "\n",
     "# Make new camera, set distance.\n",
     "camera = mujoco.MjvCamera()\n",
@@ -901,6 +884,13 @@
     "nsteps = int(np.ceil(DURATION/model.opt.timestep))\n",
     "perturb = np.random.randn(nsteps, nu)\n",
     "\n",
+    "# Scaling vector with different STD for \"balance\" and \"other\"\n",
+    "CTRL_STD = np.empty(nu)\n",
+    "for i in range(nu):\n",
+    "  joint = model.actuator(i).trnid[0]\n",
+    "  dof = model.joint(joint).dofadr[0]\n",
+    "  CTRL_STD[i] = BALANCE_STD if dof in balance_dofs else OTHER_STD\n",
+    "\n",
     "# Smooth the noise.\n",
     "width = int(nsteps * CTRL_RATE/DURATION)\n",
     "kernel = np.exp(-0.5*np.linspace(-3, 3, width)**2)\n",
@@ -926,7 +916,7 @@
     "  data.ctrl = ctrl0 - K @ dx\n",
     "\n",
     "  # Add perturbation, increment step.\n",
-    "  data.ctrl += CTRL_STD*perturb[step]\n",
+    "  data.ctrl += CTRL_STD * perturb[step]\n",
     "  step += 1\n",
     "\n",
     "  # Step the simulation.\n",
@@ -947,9 +937,10 @@
   "accelerator": "GPU",
   "colab": {
    "collapsed_sections": [
-    "LBAvTJ0xHKy7"
+    "QPdJNe3k62mx"
    ],
-   "private_outputs": true
+   "private_outputs": true,
+   "toc_visible": true
   },
   "kernelspec": {
    "display_name": "Python 3",