Loading files
This file is pinned to the selected release. Download the original to open it in your design tool.
b718d79d4This file is pinned to the selected release. Download the original to open it in your design tool.
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "bR2IJtzSilBv"
},
"source": [
"\n",
"\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",
"<!-- 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": "QPdJNe3k62mx"
},
"source": [
"## All imports\n"
]
},
{
"cell_type": "code",
"execution_count": 0,
"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 numpy as np\n",
"from typing import Callable, Optional, Union, List\n",
"import scipy.linalg\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": "J5fL6p-Sx5DB"
},
"source": [
"## Load and render the standard humanoid"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "RecFjafkfX4V"
},
"outputs": [],
"source": [
"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"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5_2cf2qgy0AX"
},
"source": [
"The XML is used to instantiate an `MjModel`. Given the model, we can create an `MjData` which holds the simulation state, and an instance of the `Renderer` class defined above."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "yftlgN0yznRe"
},
"outputs": [],
"source": [
"model = mujoco.MjModel.from_xml_string(xml)\n",
"data = mujoco.MjData(model)\n",
"renderer = mujoco.Renderer(model)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "9IvE5_N5zznN"
},
"source": [
"The state in the `data` object is in the default configuration. Let's invoke the forward dynamics to populate all the derived quantities (like the positions of geoms in the world), update the scene and render it:"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "F6ZhQv2l0OOu"
},
"outputs": [],
"source": [
"mujoco.mj_forward(model, data)\n",
"renderer.update_scene(data)\n",
"media.show_image(renderer.render())"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GzJsMlpW0_8G"
},
"source": [
"The model comes with some built-in \"keyframes\" which are saved simulation states.\n",
"\n",
"`mj_resetDataKeyframe` can be used to load them. Let's see what they look like:"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "wWzbBpCBgAzE"
},
"outputs": [],
"source": [
"for key in range(model.nkey):\n",
" mujoco.mj_resetDataKeyframe(model, data, key)\n",
" mujoco.mj_forward(model, data)\n",
" renderer.update_scene(data)\n",
" media.show_image(renderer.render())"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "MuJp0smS2yEb"
},
"source": [
"Now let's simulate the physics and render to make a video."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "I75-J4DowklB"
},
"outputs": [],
"source": [
"DURATION = 3 # seconds\n",
"FRAMERATE = 60 # Hz\n",
"\n",
"# Initialize to the standing-on-one-leg pose.\n",
"mujoco.mj_resetDataKeyframe(model, data, 1)\n",
"\n",
"frames = []\n",
"while data.time < DURATION:\n",
" # Step the simulation.\n",
" mujoco.mj_step(model, data)\n",
"\n",
" # Render and save frames.\n",
" if len(frames) < data.time * FRAMERATE:\n",
" renderer.update_scene(data)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"\n",
"# Display video.\n",
"media.show_video(frames, fps=FRAMERATE)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Qr9LpiLj4pSV"
},
"source": [
"The model defines built-in torque actuators which we can use to drive the humanoid's joints by setting the `data.ctrl` vector. Let's see what happens if we inject noise into it.\n",
"\n",
"While we're here, let's use a custom camera that will track the humanoid's center of mass."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "kFeaLU7n42Iu"
},
"outputs": [],
"source": [
"DURATION = 3 # seconds\n",
"FRAMERATE = 60 # Hz\n",
"\n",
"# Make a new camera, move it to a closer distance.\n",
"camera = mujoco.MjvCamera()\n",
"mujoco.mjv_defaultFreeCamera(model, camera)\n",
"camera.distance = 2\n",
"\n",
"mujoco.mj_resetDataKeyframe(model, data, 1)\n",
"\n",
"frames = []\n",
"while data.time < DURATION:\n",
" # Set control vector.\n",
" data.ctrl = np.random.randn(model.nu)\n",
"\n",
" # Step the simulation.\n",
" mujoco.mj_step(model, data)\n",
"\n",
" # Render and save frames.\n",
" if len(frames) < data.time * FRAMERATE:\n",
" # Set the lookat point to the humanoid's center of mass.\n",
" camera.lookat = data.body('torso').subtree_com\n",
"\n",
" renderer.update_scene(data, camera)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"\n",
"media.show_video(frames, fps=FRAMERATE)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "9XVqpSg78SH9"
},
"source": [
"## Stable standing on one leg\n",
"\n",
"Clearly this initial pose is not stable. We'll try to find a stabilising control law using a [Linear Quadratic Regulator](https://en.wikipedia.org/wiki/Linear%E2%80%93quadratic_regulator)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "iPZ3TztDDLSg"
},
"source": [
"### Recap of LQR theory\n",
"There are many online resources explaining this theory, developed by Rudolph Kalman in the 1960s, but we'll provide a minimal recap.\n",
"\n",
"Given a dynamical system which is linear in the state $x$ and control $u$,\n",
"$$\n",
"x_{t+h} = A x_t + B u_t\n",
"$$\n",
"if the system fulfills a controllability criterion, it is possible to stabilize it (drive $x$ to 0) in an optimal fashion, as follows. Define a quadratic cost function over states and controls $J(x,u)$ using two Symmetric Positive Definite matrices $Q$ and $R$:\n",
"$$\n",
"J(x,u) = x^T Q x + u^T R u\n",
"$$\n",
"\n",
"The cost-to-go $V^\\pi(x_0)$, also known as the Value function, is the total sum of future costs, letting the state start at $x_0$ and evolve according to the dynamics, while using a control law $u=\\pi(x)$:\n",
"$$\n",
"V^\\pi(x_0) = \\sum_{t=0}^\\infty J(x_t, \\pi(x_t))\n",
"$$\n",
"Kalman's central result can now be stated. The optimal control law which minimizes the cost-to-go (over all possible control laws!) is linear\n",
"$$\n",
"\\pi^*(x) = \\underset{\\pi}{\\text{argmin}}\\; V^\\pi(x)=-Kx\n",
"$$\n",
"and the optimal cost-to-go is quadratic\n",
"$$\n",
"V^*(x) =\\underset{\\pi}{\\min}\\; V^\\pi(x) = x^T P x\n",
"$$\n",
"The matrix $P$ obeys the Riccati equation\n",
"$$\n",
"P = Q + A^T P A - A^T P B (R+B^T P B)^{-1} B^T P A\n",
"$$\n",
"and its relationship to the control gain matrix $K$ is\n",
"$$\n",
"K = (R + B^T P B)^{-1} B^T P A\n",
"$$"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ek1RjwKBNT3C"
},
"source": [
"### Understanding linearization setpoints\n",
"\n",
"Of course our humanoid simulation is anything but linear. But while MuJoCo's `mj_step` function computes some non-linear dynamics $x_{t+h} = f(x_t,u_t)$, we can *linearize* this function around any state-control pair. Using shortcuts for the next state $y=x_{t+h}$, the current state $x=x_t$ and the current control $u=u_t$, and using $\\delta$ to mean \"small change in\", we can write\n",
"$$\n",
"\\delta y = \\frac{\\partial f}{\\partial x}\\delta x+ \\frac{\\partial f}{\\partial u}\\delta u\n",
"$$\n",
"In other words, the partial derivative matrices decribe a linear relationship between perturbations to $x$ and $u$ and changes to $y$. Comparing to the theory above, we can identify the partial derivative (Jacobian) matrices with the transition matrices $A$ and $B$, when considering the linearized dynamical system:\n",
"$$\n",
"A = \\frac{\\partial f}{\\partial x} \\quad\n",
"B = \\frac{\\partial f}{\\partial u}\n",
"$$\n",
"In order to perform the linearization, we need to choose some setpoints $x$ and $u$ around which we will linearize. We already know $x$, this is our initial pose of standing on one leg. But what about $u$? How do we find the \"best\" control around which to linearise?\n",
"\n",
"The answer is inverse dynamics."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "6wPXh8xxWX0y"
},
"source": [
"### Finding the control setpoint using inverse dynamics\n",
"\n",
"MuJoCo's forward dynamics function `mj_forward`, which we used above in order to propagate derived quantities, computes the acceleration given the state and all the forces in the system, some of which are created by the actuators.\n",
"\n",
"The inverse dynamics function takes the acceleration as *input*, and computes the forces required to create the acceleration. Uniquely, MuJoCo's [fast inverse dynamics](https://doi.org/10.1109/ICRA.2014.6907751) takes into account all constraints, including contacts. Let's see how it works.\n",
"\n",
"We'll call the forward dynamics at our desired position setpoint, set the acceleration in `data.qacc` to 0, and call the inverse dynamics:"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "8Q6Ceuf7ZHGQ"
},
"outputs": [],
"source": [
"mujoco.mj_resetDataKeyframe(model, data, 1)\n",
"mujoco.mj_forward(model, data)\n",
"data.qacc = 0 # Assert that there is no the acceleration.\n",
"mujoco.mj_inverse(model, data)\n",
"print(data.qfrc_inverse)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "6tqFzfwNa2i8"
},
"source": [
"Examining the forces found by the inverse dynamics, we see something rather disturbing. There is a very large force applied at the 3rd degree-of-freedom (DoF), the vertical motion DoF of the root joint.\n",
"\n",
"This means that in order to explain our assertion that the acceleration is zero, the inverse dynamics has to invent a \"magic\" force applied directly to the root joint. Let's see how this force varies as we move our humanoid up and down by just 1mm, in increments of 1$\\mu$m:"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "2eN8b1EZ5GO9"
},
"outputs": [],
"source": [
"height_offsets = np.linspace(-0.001, 0.001, 2001)\n",
"vertical_forces = []\n",
"for offset in height_offsets:\n",
" mujoco.mj_resetDataKeyframe(model, data, 1)\n",
" mujoco.mj_forward(model, data)\n",
" data.qacc = 0\n",
" # Offset the height by `offset`.\n",
" data.qpos[2] += offset\n",
" mujoco.mj_inverse(model, data)\n",
" vertical_forces.append(data.qfrc_inverse[2])\n",
"\n",
"# Find the height-offset at which the vertical force is smallest.\n",
"idx = np.argmin(np.abs(vertical_forces))\n",
"best_offset = height_offsets[idx]\n",
"\n",
"# Plot the relationship.\n",
"plt.figure(figsize=(10, 6))\n",
"plt.plot(height_offsets * 1000, vertical_forces, linewidth=3)\n",
"# Red vertical line at offset corresponding to smallest vertical force.\n",
"plt.axvline(x=best_offset*1000, color='red', linestyle='--')\n",
"# Green horizontal line at the humanoid's weight.\n",
"weight = model.body_subtreemass[1]*np.linalg.norm(model.opt.gravity)\n",
"plt.axhline(y=weight, color='green', linestyle='--')\n",
"plt.xlabel('Height offset (mm)')\n",
"plt.ylabel('Vertical force (N)')\n",
"plt.grid(which='major', color='#DDDDDD', linewidth=0.8)\n",
"plt.grid(which='minor', color='#EEEEEE', linestyle=':', linewidth=0.5)\n",
"plt.minorticks_on()\n",
"plt.title(f'Smallest vertical force '\n",
" f'found at offset {best_offset*1000:.4f}mm.')\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "lpldLvBreQj7"
},
"source": [
"In the plot above we can see the strong non-linear relationship due to foot contacts. On the left, as we push the humanoid into the floor, the only way to explain the fact that it is not jumping out of the floor is a large external force pushing it **down**. On the right, as we move the humanoid away from the floor the only way to explain the zero acceleration is a force holding it **up**, and we can clearly see the height at which the foot no longer touches the ground, and the required force is exactly equal to the humanoid's weight (green line), and remains constant as we keep moving up.\n",
"\n",
"Near -0.5mm is the perfect height offset (red line), where the zero vertical acceleration can be entirely explained by internal joint forces, without resorting to \"magical\" external forces. Let's correct the height of our initial pose, save it in `qpos0`, and compute to inverse dynamics forces again:"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "qaw4gxg46h2G"
},
"outputs": [],
"source": [
"mujoco.mj_resetDataKeyframe(model, data, 1)\n",
"mujoco.mj_forward(model, data)\n",
"data.qacc = 0\n",
"data.qpos[2] += best_offset\n",
"qpos0 = data.qpos.copy() # Save the position setpoint.\n",
"mujoco.mj_inverse(model, data)\n",
"qfrc0 = data.qfrc_inverse.copy()\n",
"print('desired forces:', qfrc0)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "y3ofBSh-jdY5"
},
"source": [
"Much better, the forces on the root joint are small. Now that we have forces that can reasonably be produced by the actuators, how do we find the actuator values that will create them? For simple `motor` actuators like the humanoid's, we can simply \"divide\" by the actuation moment arm matrix, i.e. multiply by its pseudo-inverse:"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "a1PF_yHdPvLl"
},
"outputs": [],
"source": [
"actuator_moment = np.zeros((model.nu, model.nv))\n",
"mujoco.mju_sparse2dense(\n",
" actuator_moment,\n",
" data.actuator_moment.reshape(-1),\n",
" data.moment_rownnz,\n",
" data.moment_rowadr,\n",
" data.moment_colind.reshape(-1),\n",
")\n",
"ctrl0 = np.atleast_2d(qfrc0) @ np.linalg.pinv(actuator_moment)\n",
"ctrl0 = ctrl0.flatten() # Save the ctrl setpoint.\n",
"print('control setpoint:', ctrl0)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "h6bLO26Ekvir"
},
"source": [
"More elaborate actuators would require a different method to recover $\\frac{\\partial \\texttt{ qfrc_actuator}}{\\partial \\texttt{ ctrl}}$, and finite-differencing is always an easy option.\n",
"\n",
"Let's apply these controls in the forward dynamics and compare the forces they produce with the desired forces printed above:"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "dDLihz5hk9Wt"
},
"outputs": [],
"source": [
"data.ctrl = ctrl0\n",
"mujoco.mj_forward(model, data)\n",
"print('actuator forces:', data.qfrc_actuator)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "V2XVaRruloKG"
},
"source": [
"Because the humanoid is fully-actuated (apart from the root joint), and the required forces are all within the actuator limits, we can see a perfect match with the desired forces across all internal joints. There is still some mismatch in the root joint, but it's small. Let's see what the simulation looks like when we apply these controls:"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "8cQpEF7MmDPI"
},
"outputs": [],
"source": [
"DURATION = 3 # seconds\n",
"FRAMERATE = 60 # Hz\n",
"\n",
"# Set the state and controls to their setpoints.\n",
"mujoco.mj_resetData(model, data)\n",
"data.qpos = qpos0\n",
"data.ctrl = ctrl0\n",
"\n",
"frames = []\n",
"while data.time < DURATION:\n",
" # Step the simulation.\n",
" mujoco.mj_step(model, data)\n",
"\n",
" # Render and save frames.\n",
" if len(frames) < data.time * FRAMERATE:\n",
" # Set the lookat point to the humanoid's center of mass.\n",
" camera.lookat = data.body('torso').subtree_com\n",
" renderer.update_scene(data, camera)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"\n",
"media.show_video(frames, fps=FRAMERATE)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "rGPdr2_9P2sz"
},
"source": [
"Comparing to the completely passive video we made above, we can see that this is a much better control setpoint. The humanoid still falls down, but it tries to stabilize and succeeds for a short while."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pWqJ4hqzm7Xq"
},
"source": [
"### Choosing the $Q$ and $R$ matrices\n",
"\n",
"In order to obtain the LQR feedback control law, we will need to design the $Q$ and $R$ matrices. Due to the linear structure, the solution is invariant to a scaling of both matrices, so without loss of generality we can choose $R$ to be the identity matrix:"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "SxOeHSf1nspy"
},
"outputs": [],
"source": [
"nu = model.nu # Alias for the number of actuators.\n",
"R = np.eye(nu)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pEluM4mNRghL"
},
"source": [
"Choosing $Q$ is more elaborate. We will construct it as a sum of two terms.\n",
"\n",
"First, a balancing cost that will keep the center of mass (CoM) over the foot. In order to describe it, we will use kinematic Jacobians which map between joint space and global Cartesian positions. MuJoCo computes these analytically."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "9LEK4MCHkH15"
},
"outputs": [],
"source": [
"nv = model.nv # Shortcut for the number of DoFs.\n",
"\n",
"# Get the Jacobian for the root body (torso) CoM.\n",
"mujoco.mj_resetData(model, data)\n",
"data.qpos = qpos0\n",
"mujoco.mj_forward(model, data)\n",
"jac_com = np.zeros((3, nv))\n",
"mujoco.mj_jacSubtreeCom(model, data, jac_com, model.body('torso').id)\n",
"\n",
"# Get the Jacobian for the left foot.\n",
"jac_foot = np.zeros((3, nv))\n",
"mujoco.mj_jacBodyCom(model, data, jac_foot, None, model.body('foot_left').id)\n",
"\n",
"jac_diff = jac_com - jac_foot\n",
"Qbalance = jac_diff.T @ jac_diff"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "daDCbcAskiML"
},
"source": [
"Second, a cost for joints moving away from their initial configuration. We will want different coefficients for different sets of joints:\n",
"- The free joint will get a coefficient of 0, as that is already taken care of by the CoM cost term.\n",
"- The joints required for balancing on the left leg, i.e. the left leg joints and the horizontal abdominal joints, should stay quite close to their initial values.\n",
"- All the other joints should have a smaller coefficient, so that the humanoid will, for example, be able to flail its arms in order to balance.\n",
"\n",
"Let's get the indices of all these joint sets.\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "S731eD3Jm1mJ"
},
"outputs": [],
"source": [
"# Get all joint names.\n",
"joint_names = [model.joint(i).name for i in range(model.njnt)]\n",
"\n",
"# Get indices into relevant sets of joints.\n",
"root_dofs = range(6)\n",
"body_dofs = range(6, nv)\n",
"abdomen_dofs = [\n",
" model.joint(name).dofadr[0]\n",
" for name in joint_names\n",
" if 'abdomen' in name\n",
" and not 'z' in name\n",
"]\n",
"left_leg_dofs = [\n",
" model.joint(name).dofadr[0]\n",
" for name in joint_names\n",
" if 'left' in name\n",
" and ('hip' in name or 'knee' in name or 'ankle' in name)\n",
" and not 'z' in name\n",
"]\n",
"balance_dofs = abdomen_dofs + left_leg_dofs\n",
"other_dofs = np.setdiff1d(body_dofs, balance_dofs)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "OeHYWdQdm-vE"
},
"source": [
"We are now ready to construct the Q matrix. Note that the coefficient of the balancing term is quite high. This is due to 3 separate reasons:\n",
"- It's the thing we care about most. Balancing means keeping the CoM over the foot.\n",
"- We have less control authority over the CoM (relative to body joints).\n",
"- In the balancing context, units of length are \"bigger\". If the knee bends by 0.1 radians (≈6°), we can probably still recover. If the CoM position is 10cm sideways from the foot position, we are likely on our way to the floor."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "reIA8___o3Z4"
},
"outputs": [],
"source": [
"# Cost coefficients.\n",
"BALANCE_COST = 1000 # Balancing.\n",
"BALANCE_JOINT_COST = 3 # Joints required for balancing.\n",
"OTHER_JOINT_COST = .3 # Other joints.\n",
"\n",
"# Construct the Qjoint matrix.\n",
"Qjoint = np.eye(nv)\n",
"Qjoint[root_dofs, root_dofs] *= 0 # Don't penalize free joint directly.\n",
"Qjoint[balance_dofs, balance_dofs] *= BALANCE_JOINT_COST\n",
"Qjoint[other_dofs, other_dofs] *= OTHER_JOINT_COST\n",
"\n",
"# Construct the Q matrix for position DoFs.\n",
"Qpos = BALANCE_COST * Qbalance + Qjoint\n",
"\n",
"# No explicit penalty for velocities.\n",
"Q = np.block([[Qpos, np.zeros((nv, nv))],\n",
" [np.zeros((nv, 2*nv))]])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "U9EEBeIsnJVA"
},
"source": [
"### Computing the LQR gain matrix $K$\n",
"\n",
"Before we solve for the LQR controller, we need the $A$ and $B$ matrices. These are computed by MuJoCo's `mjd_transitionFD` function which computes them using efficient finite-difference derivatives, exploiting the configurable computation pipeline to avoid recomputing quantities which haven't changed."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "NB4ZStYrpx1B"
},
"outputs": [],
"source": [
"# Set the initial state and control.\n",
"mujoco.mj_resetData(model, data)\n",
"data.ctrl = ctrl0\n",
"data.qpos = qpos0\n",
"\n",
"# Allocate the A and B matrices, compute them.\n",
"A = np.zeros((2*nv, 2*nv))\n",
"B = np.zeros((2*nv, nu))\n",
"epsilon = 1e-6\n",
"flg_centered = True\n",
"mujoco.mjd_transitionFD(model, data, epsilon, flg_centered, A, B, None, None)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "wvYn5_PEpsP6"
},
"source": [
"We are now ready to solve for our stabilizing controller. We will use `scipy`'s `solve_discrete_are` to solve the Riccati equation and get the feedback gain matrix using the formula described in the recap."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "azjsWl4v_K16"
},
"outputs": [],
"source": [
"# Solve discrete Riccati equation.\n",
"P = scipy.linalg.solve_discrete_are(A, B, Q, R)\n",
"\n",
"# Compute the feedback gain matrix K.\n",
"K = np.linalg.inv(R + B.T @ P @ B) @ B.T @ P @ A"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "JyvEBxrlXkxH"
},
"source": [
"### Stable standing\n",
"\n",
"We can now try our stabilising controller.\n",
"\n",
"Note that in order to apply our gain matrix $K$, we need to use `mj_differentiatePos` which computes the difference of two positions. This is important because the root orientation is given by a length-4 quaternion, while the difference of two quaternions (in the tangent space) is length-3. In MuJoCo notation, positions (`qpos`) are of size `nq` while a position differences (and velocities) are of size `nv`.\n"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "Z_57VMUDpGrj"
},
"outputs": [],
"source": [
"# Parameters.\n",
"DURATION = 5 # seconds\n",
"FRAMERATE = 60 # Hz\n",
"\n",
"# Reset data, set initial pose.\n",
"mujoco.mj_resetData(model, data)\n",
"data.qpos = qpos0\n",
"\n",
"# Allocate position difference dq.\n",
"dq = np.zeros(model.nv)\n",
"\n",
"frames = []\n",
"while data.time < DURATION:\n",
" # Get state difference dx.\n",
" mujoco.mj_differentiatePos(model, dq, 1, qpos0, data.qpos)\n",
" dx = np.hstack((dq, data.qvel)).T\n",
"\n",
" # LQR control law.\n",
" data.ctrl = ctrl0 - K @ dx\n",
"\n",
" # Step the simulation.\n",
" mujoco.mj_step(model, data)\n",
"\n",
" # Render and save frames.\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": "ONg32ETtrZNl"
},
"source": [
"### Final video\n",
"\n",
"The video above is a bit disappointing, as the humanoid is basically motionless. Let's fix that and also add a few flourishes for our finale:\n",
"- Inject smoothed noise on top of the LQR controller so that the balancing action is more pronounced yet not jerky.\n",
"- Add contact force visualization to the scene.\n",
"- Smoothly orbit the camera around the humanoid.\n",
"- Instantiate a new renderer with higher resolution."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "zJmbOJMurRna"
},
"outputs": [],
"source": [
"# Parameters.\n",
"DURATION = 12 # seconds\n",
"FRAMERATE = 60 # Hz\n",
"TOTAL_ROTATION = 15 # degrees\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",
"mujoco.mjv_defaultFreeCamera(model, camera)\n",
"camera.distance = 2.3\n",
"\n",
"# Enable contact force visualisation.\n",
"scene_option = mujoco.MjvOption()\n",
"scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = True\n",
"\n",
"# Set the scale of visualized contact forces to 1cm/N.\n",
"model.vis.map.force = 0.01\n",
"\n",
"# Define smooth orbiting function.\n",
"def unit_smooth(normalised_time: float) -> float:\n",
" return 1 - np.cos(normalised_time*2*np.pi)\n",
"def azimuth(time: float) -> float:\n",
" return 100 + unit_smooth(data.time/DURATION) * TOTAL_ROTATION\n",
"\n",
"# Precompute some noise.\n",
"np.random.seed(1)\n",
"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",
"kernel /= np.linalg.norm(kernel)\n",
"for i in range(nu):\n",
" perturb[:, i] = np.convolve(perturb[:, i], kernel, mode='same')\n",
"\n",
"# Reset data, set initial pose.\n",
"mujoco.mj_resetData(model, data)\n",
"data.qpos = qpos0\n",
"\n",
"# New renderer instance with higher resolution.\n",
"renderer = mujoco.Renderer(model, width=1280, height=720)\n",
"\n",
"frames = []\n",
"step = 0\n",
"while data.time < DURATION:\n",
" # Get state difference dx.\n",
" mujoco.mj_differentiatePos(model, dq, 1, qpos0, data.qpos)\n",
" dx = np.hstack((dq, data.qvel)).T\n",
"\n",
" # LQR control law.\n",
" data.ctrl = ctrl0 - K @ dx\n",
"\n",
" # Add perturbation, increment step.\n",
" data.ctrl += CTRL_STD * perturb[step]\n",
" step += 1\n",
"\n",
" # Step the simulation.\n",
" mujoco.mj_step(model, data)\n",
"\n",
" # Render and save frames.\n",
" if len(frames) < data.time * FRAMERATE:\n",
" camera.azimuth = azimuth(data.time)\n",
" renderer.update_scene(data, camera, scene_option)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"\n",
"media.show_video(frames, fps=FRAMERATE)"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [
"QPdJNe3k62mx"
],
"private_outputs": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}