{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "MpkYHwCqk7W-" }, "source": [ "![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\n", "\n", "#

Tutorial

\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", "" ] }, { "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('')\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", "\n", " \n", " \n", " \n", " \n", "\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 ``. The smallest valid MJCF model is `` which is a completely empty model.\n", " - All physical elements live inside the `` 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", "\n", " \n", " \n", " \n", " \n", "\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", "\n", " \n", " \n", " \n", " \n", " \n", "\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", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\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 `