diff --git a/mjx/tutorial.ipynb b/mjx/tutorial.ipynb
index befc6b3b..65579cad 100644
--- a/mjx/tutorial.ipynb
+++ b/mjx/tutorial.ipynb
@@ -1,1477 +1,1471 @@
{
- "cells": [
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "MpkYHwCqk7W-"
- },
- "source": [
- "\n",
- "\n",
- "# \u003ch1\u003e\u003ccenter\u003eTutorial \u003ca href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/\u003e\u003c/a\u003e\u003c/center\u003e\u003c/h1\u003e\n",
- "\n",
- "This notebook provides an introductory tutorial for [**MuJoCo XLA (MJX)**](https://github.com/google-deepmind/mujoco/blob/main/mjx), a JAX-based implementation of MuJoCo useful for RL training workloads.\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 \u003e Change runtime type\".\n",
- "\n",
- "\n",
- "\n",
- "\n",
- "\n",
- "\n",
- "\n",
- "\n"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "xBSdkbmGN2K-"
- },
- "source": [
- "### Copyright notice"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "_UbO9uhtBSX5"
- },
- "source": [
- "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eCopyright 2023 DeepMind Technologies Limited.\u003c/small\u003e\u003c/p\u003e\n",
- "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eLicensed 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 \u003ca href=\"http://www.apache.org/licenses/LICENSE-2.0\"\u003ehttp://www.apache.org/licenses/LICENSE-2.0\u003c/a\u003e.\u003c/small\u003e\u003c/small\u003e\u003c/p\u003e\n",
- "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eUnless 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.\u003c/small\u003e\u003c/small\u003e\u003c/p\u003e"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "YvyGCsgSCxHQ"
- },
- "source": [
- "# Install MuJoCo, MJX, and Brax"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "Xqo7pyX-n72M"
- },
- "outputs": [],
- "source": [
- "!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",
- "from google.colab import files\n",
- "\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",
- "# 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\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",
- " mujoco.MjModel.from_xml_string('\u003cmujoco/\u003e')\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.')"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "T5f4w3Kq2X14"
- },
- "outputs": [],
- "source": [
- "#@title Import packages for plotting and creating graphics\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",
- "!command -v ffmpeg \u003e/dev/null || (apt update \u0026\u0026 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": {
- "id": "ObF1UXrkb0Nd"
- },
- "outputs": [],
- "source": [
- "#@title Import MuJoCo, MJX, and Brax\n",
- "\n",
- "\n",
- "from datetime import datetime\n",
- "import functools\n",
- "from IPython.display import HTML\n",
- "import jax\n",
- "from jax import numpy as jp\n",
- "import numpy as np\n",
- "from typing import Any, Dict, Sequence, Tuple, Union\n",
- "\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.envs.base import Env, PipelineEnv, State\n",
- "from brax.mjx.base import State as MjxState\n",
- "from brax.training.agents.ppo import train as ppo\n",
- "from brax.training.agents.ppo import networks as ppo_networks\n",
- "from brax.io import html, mjcf, model\n",
- "\n",
- "from etils import epath\n",
- "from flax import struct\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"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "Nj4-Xmx4DFaq"
- },
- "source": [
- "# Introduction to MJX\n",
- "\n",
- "MJX is an implementation of MuJoCo written in [JAX](https://jax.readthedocs.io/en/latest/index.html), enabling large batch training on GPU/TPU. In this notebook, we will demonstrate how to train RL policies with MJX.\n",
- "\n",
- "Before we get into hefty RL workloads, let's get started with a simpler example! The entrypoint into MJX is through MuJoCo, so first we load a MuJoCo model:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "bNus3mbbDz6a"
- },
- "outputs": [],
- "source": [
- "xml = \"\"\"\n",
- "\u003cmujoco\u003e\n",
- " \u003cworldbody\u003e\n",
- " \u003clight name=\"top\" pos=\"0 0 1\"/\u003e\n",
- " \u003cbody name=\"box_and_sphere\" euler=\"0 0 -30\"\u003e\n",
- " \u003cjoint name=\"swing\" type=\"hinge\" axis=\"1 -1 0\" pos=\"-.2 -.2 -.2\"/\u003e\n",
- " \u003cgeom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/\u003e\n",
- " \u003cgeom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- " \u003c/worldbody\u003e\n",
- "\u003c/mujoco\u003e\n",
- "\"\"\"\n",
- "\n",
- "# Make model, data, and renderer\n",
- "mj_model = mujoco.MjModel.from_xml_string(xml)\n",
- "mj_data = mujoco.MjData(mj_model)\n",
- "renderer = mujoco.Renderer(mj_model)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "Po5oykJbFQbj"
- },
- "source": [
- "Next we take the MuJoCo model and data, and place them on the GPU device using MJX."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "TSpoOWqeEC3P"
- },
- "outputs": [],
- "source": [
- "mjx_model = mjx.put_model(mj_model)\n",
- "mjx_data = mjx.put_data(mj_model, mj_data)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "6rxMMSs4OJJf"
- },
- "source": [
- "Below, we print the `qpos` from MuJoCo and MJX. Notice that the `qpos` for the mjData is a numpy array living on the CPU, while the `qpos` for `mjx.Data` is a JAX Array living on the GPU device."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "ZOD582pfOLP-"
- },
- "outputs": [],
- "source": [
- "print(mj_data.qpos, type(mj_data.qpos))\n",
- "print(mjx_data.qpos, type(mjx_data.qpos), mjx_data.qpos.devices())"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "ZShF9-o_JLm3"
- },
- "source": [
- "Let's run the simulation in MuJoCo and render the trajectory. This example is taken from the [MuJoCo tutorial](https://colab.sandbox.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb)."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "HDlPlX05I3m-"
- },
- "outputs": [],
- "source": [
- "# 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",
- "frames = []\n",
- "mujoco.mj_resetData(mj_model, mj_data)\n",
- "while mj_data.time \u003c duration:\n",
- " mujoco.mj_step(mj_model, mj_data)\n",
- " if len(frames) \u003c mj_data.time * framerate:\n",
- " renderer.update_scene(mj_data, scene_option=scene_option)\n",
- " pixels = renderer.render()\n",
- " frames.append(pixels)\n",
- "\n",
- "# Simulate and display video.\n",
- "media.show_video(frames, fps=framerate)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "m70b_RxBJOyd"
- },
- "source": [
- "Now let's run the same exact simulation on the GPU device using MJX!\n",
- "\n",
- "In the example below, we use `mjx.step` instead of `mujoco.mj_step`, and we also [`jax.jit`](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html) the `mjx.step` so that it runs efficiently on the GPU. After each step, we convert the `mjx.Data` back to `mjData` so that we can use the MuJoCo renderer.\n"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "Pr29xq0-JRQv"
- },
- "outputs": [],
- "source": [
- "\n",
- "jit_step = jax.jit(mjx.step)\n",
- "\n",
- "frames = []\n",
- "mujoco.mj_resetData(mj_model, mj_data)\n",
- "mjx_data = mjx.put_data(mj_model, mj_data)\n",
- "while mjx_data.time \u003c duration:\n",
- " mjx_data = jit_step(mjx_model, mjx_data)\n",
- " if len(frames) \u003c mjx_data.time * framerate:\n",
- " mj_data = mjx.get_data(mj_model, mjx_data)\n",
- " renderer.update_scene(mj_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": "wXsQ4qO2KO3Q"
- },
- "source": [
- "Running single threaded physics simulation on the GPU is not very [efficient](https://mujoco.readthedocs.io/en/stable/mjx.html#mjx-the-sharp-bits). The advantage with MJX is that we can run environments in parallel on a hardware accelerated device. Let's try it out!\n",
- "\n",
- "In the example below, we create 4096 copies of the `mjx.Data` and we run the `mjx.step` over the batched data. Since MJX is implemented in JAX, we take advantage of [`jax.vmap`](https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html) to run the `mjx.step` in parallel over all `mjx.Data`."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "rrdrcKRVK6w9"
- },
- "outputs": [],
- "source": [
- "rng = jax.random.PRNGKey(0)\n",
- "rng = jax.random.split(rng, 4096)\n",
- "batch = jax.vmap(lambda rng: mjx_data.replace(qpos=jax.random.uniform(rng, (1,))))(rng)\n",
- "\n",
- "jit_step = jax.vmap(mjx.step, in_axes=(None, 0))\n",
- "batch = jit_step(mjx_model, batch)\n",
- "\n",
- "print(batch.qpos)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "x4lL220cOj0q"
- },
- "source": [
- "We can copy the batched `mjx.Data` back to MuJoCo like we did before:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "Jtz7j1PDOnw5"
- },
- "outputs": [],
- "source": [
- "batched_mj_data = mjx.get_data(mj_model, batch)\n",
- "print([d.qpos for d in batched_mj_data])"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "RAv6WUVUm78k"
- },
- "source": [
- "# Training a Policy with MJX\n",
- "\n",
- "Running large batch physics simulation is useful for training RL policies. Here we demonstrate training RL policies with MJX using the RL library from [Brax](https://github.com/google/brax).\n",
- "\n",
- "Below, we implement the classic Humanoid environment using MJX and Brax. We inherit from the `MjxEnv` implementation in Brax so that we can step the physics with MJX while training with Brax RL implementations.\n"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "mtGMYNLE3QJN"
- },
- "outputs": [],
- "source": [
- "#@title Humanoid Env\n",
- "\n",
- "class Humanoid(PipelineEnv):\n",
- "\n",
- " def __init__(\n",
- " self,\n",
- " forward_reward_weight=1.25,\n",
- " ctrl_cost_weight=0.1,\n",
- " healthy_reward=5.0,\n",
- " terminate_when_unhealthy=True,\n",
- " healthy_z_range=(1.0, 2.0),\n",
- " reset_noise_scale=1e-2,\n",
- " exclude_current_positions_from_observation=True,\n",
- " **kwargs,\n",
- " ):\n",
- " path = epath.Path(epath.resource_path('mujoco')) / (\n",
- " 'mjx/test_data/humanoid'\n",
- " )\n",
- " mj_model = mujoco.MjModel.from_xml_path(\n",
- " (path / 'humanoid.xml').as_posix())\n",
- " mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG\n",
- " mj_model.opt.iterations = 6\n",
- " mj_model.opt.ls_iterations = 6\n",
- "\n",
- " sys = mjcf.load_model(mj_model)\n",
- "\n",
- " physics_steps_per_control_step = 5\n",
- " kwargs['n_frames'] = kwargs.get(\n",
- " 'n_frames', physics_steps_per_control_step)\n",
- " kwargs['backend'] = 'mjx'\n",
- "\n",
- " super().__init__(sys, **kwargs)\n",
- "\n",
- " self._forward_reward_weight = forward_reward_weight\n",
- " self._ctrl_cost_weight = ctrl_cost_weight\n",
- " self._healthy_reward = healthy_reward\n",
- " self._terminate_when_unhealthy = terminate_when_unhealthy\n",
- " self._healthy_z_range = healthy_z_range\n",
- " self._reset_noise_scale = reset_noise_scale\n",
- " self._exclude_current_positions_from_observation = (\n",
- " exclude_current_positions_from_observation\n",
- " )\n",
- "\n",
- " def reset(self, rng: jp.ndarray) -\u003e State:\n",
- " \"\"\"Resets the environment to an initial state.\"\"\"\n",
- " rng, rng1, rng2 = jax.random.split(rng, 3)\n",
- "\n",
- " low, hi = -self._reset_noise_scale, self._reset_noise_scale\n",
- " qpos = self.sys.qpos0 + jax.random.uniform(\n",
- " rng1, (self.sys.nq,), minval=low, maxval=hi\n",
- " )\n",
- " qvel = jax.random.uniform(\n",
- " rng2, (self.sys.nv,), minval=low, maxval=hi\n",
- " )\n",
- "\n",
- " data = self.pipeline_init(qpos, qvel)\n",
- "\n",
- " obs = self._get_obs(data, jp.zeros(self.sys.nu))\n",
- " reward, done, zero = jp.zeros(3)\n",
- " metrics = {\n",
- " 'forward_reward': zero,\n",
- " 'reward_linvel': zero,\n",
- " 'reward_quadctrl': zero,\n",
- " 'reward_alive': zero,\n",
- " 'x_position': zero,\n",
- " 'y_position': zero,\n",
- " 'distance_from_origin': zero,\n",
- " 'x_velocity': zero,\n",
- " 'y_velocity': zero,\n",
- " }\n",
- " return State(data, obs, reward, done, metrics)\n",
- "\n",
- " def step(self, state: State, action: jp.ndarray) -\u003e State:\n",
- " \"\"\"Runs one timestep of the environment's dynamics.\"\"\"\n",
- " data0 = state.pipeline_state\n",
- " data = self.pipeline_step(data0, action)\n",
- "\n",
- " com_before = data0.subtree_com[1]\n",
- " com_after = data.subtree_com[1]\n",
- " velocity = (com_after - com_before) / self.dt\n",
- " forward_reward = self._forward_reward_weight * velocity[0]\n",
- "\n",
- " min_z, max_z = self._healthy_z_range\n",
- " is_healthy = jp.where(data.q[2] \u003c min_z, 0.0, 1.0)\n",
- " is_healthy = jp.where(data.q[2] \u003e max_z, 0.0, is_healthy)\n",
- " if self._terminate_when_unhealthy:\n",
- " healthy_reward = self._healthy_reward\n",
- " else:\n",
- " healthy_reward = self._healthy_reward * is_healthy\n",
- "\n",
- " ctrl_cost = self._ctrl_cost_weight * jp.sum(jp.square(action))\n",
- "\n",
- " obs = self._get_obs(data, action)\n",
- " reward = forward_reward + healthy_reward - ctrl_cost\n",
- " done = 1.0 - is_healthy if self._terminate_when_unhealthy else 0.0\n",
- " state.metrics.update(\n",
- " forward_reward=forward_reward,\n",
- " reward_linvel=forward_reward,\n",
- " reward_quadctrl=-ctrl_cost,\n",
- " reward_alive=healthy_reward,\n",
- " x_position=com_after[0],\n",
- " y_position=com_after[1],\n",
- " distance_from_origin=jp.linalg.norm(com_after),\n",
- " x_velocity=velocity[0],\n",
- " y_velocity=velocity[1],\n",
- " )\n",
- "\n",
- " return state.replace(\n",
- " pipeline_state=data, obs=obs, reward=reward, done=done\n",
- " )\n",
- "\n",
- " def _get_obs(\n",
- " self, data: mjx.Data, action: jp.ndarray\n",
- " ) -\u003e jp.ndarray:\n",
- " \"\"\"Observes humanoid body position, velocities, and angles.\"\"\"\n",
- " position = data.qpos\n",
- " if self._exclude_current_positions_from_observation:\n",
- " position = position[2:]\n",
- "\n",
- " # external_contact_forces are excluded\n",
- " return jp.concatenate([\n",
- " position,\n",
- " data.qvel,\n",
- " data.cinert[1:].ravel(),\n",
- " data.cvel[1:].ravel(),\n",
- " data.qfrc_actuator,\n",
- " ])\n",
- "\n",
- "\n",
- "envs.register_environment('humanoid', Humanoid)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "P1K6IznI2y83"
- },
- "source": [
- "## Visualize a Rollout\n",
- "\n",
- "Let's instantiate the environment and visualize a short rollout.\n",
- "\n",
- "NOTE: Since episodes terminates early if the torso is below the healthy z-range, the only relevant contacts for this task are between the feet and the plane. We turn off other contacts."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "EhKLFK54C1CH"
- },
- "outputs": [],
- "source": [
- "# instantiate the environment\n",
- "env_name = 'humanoid'\n",
- "env = envs.get_environment(env_name)\n",
- "\n",
- "# define the jit reset/step functions\n",
- "jit_reset = jax.jit(env.reset)\n",
- "jit_step = jax.jit(env.step)\n"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "Ph8u-v2Q2xLS"
- },
- "outputs": [],
- "source": [
- "# initialize the state\n",
- "state = jit_reset(jax.random.PRNGKey(0))\n",
- "rollout = [state.pipeline_state]\n",
- "\n",
- "# grab a trajectory\n",
- "for i in range(10):\n",
- " ctrl = -0.1 * jp.ones(env.sys.nu)\n",
- " state = jit_step(state, ctrl)\n",
- " rollout.append(state.pipeline_state)\n",
- "\n",
- "media.show_video(env.render(rollout, camera='side'), fps=1.0 / env.dt)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "BQDG6NQ1CbZD"
- },
- "source": [
- "## Train Humanoid Policy\n",
- "\n",
- "Let's now train a policy with PPO to make the Humanoid run forwards. Training takes about 6 minutes on a Tesla A100 GPU."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "xLiddQYPApBw"
- },
- "outputs": [],
- "source": [
- "train_fn = functools.partial(\n",
- " ppo.train, num_timesteps=30_000_000, num_evals=5, reward_scaling=0.1,\n",
- " episode_length=1000, normalize_observations=True, action_repeat=1,\n",
- " unroll_length=10, num_minibatches=32, num_updates_per_batch=8,\n",
- " discounting=0.97, learning_rate=3e-4, entropy_cost=1e-3, num_envs=2048,\n",
- " batch_size=1024, seed=0)\n",
- "\n",
- "\n",
- "x_data = []\n",
- "y_data = []\n",
- "ydataerr = []\n",
- "times = [datetime.now()]\n",
- "\n",
- "max_y, min_y = 13000, 0\n",
- "def progress(num_steps, metrics):\n",
- " times.append(datetime.now())\n",
- " x_data.append(num_steps)\n",
- " y_data.append(metrics['eval/episode_reward'])\n",
- " ydataerr.append(metrics['eval/episode_reward_std'])\n",
- "\n",
- " plt.xlim([0, train_fn.keywords['num_timesteps'] * 1.25])\n",
- " plt.ylim([min_y, max_y])\n",
- "\n",
- " plt.xlabel('# environment steps')\n",
- " plt.ylabel('reward per episode')\n",
- " plt.title(f'y={y_data[-1]:.3f}')\n",
- "\n",
- " plt.errorbar(\n",
- " x_data, y_data, yerr=ydataerr)\n",
- " plt.show()\n",
- "\n",
- "make_inference_fn, params, _= train_fn(environment=env, progress_fn=progress)\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": "YYIch0HEApBx"
- },
- "source": [
- "\u003c!-- ## Save and Load Policy --\u003e\n",
- "\n",
- "We can save and load the policy using the brax model API."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "Z8gI6qH6ApBx"
- },
- "outputs": [],
- "source": [
- "#@title Save Model\n",
- "model_path = '/tmp/mjx_brax_policy'\n",
- "model.save_params(model_path, params)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "h4reaWgxApBx"
- },
- "outputs": [],
- "source": [
- "#@title Load Model and Define Inference Function\n",
- "params = model.load_params(model_path)\n",
- "\n",
- "inference_fn = make_inference_fn(params)\n",
- "jit_inference_fn = jax.jit(inference_fn)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "0G357XIfApBy"
- },
- "source": [
- "## Visualize Policy\n",
- "\n",
- "Finally we can visualize the policy."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "osYasMw4ApBy"
- },
- "outputs": [],
- "source": [
- "eval_env = envs.get_environment(env_name)\n",
- "\n",
- "jit_reset = jax.jit(eval_env.reset)\n",
- "jit_step = jax.jit(eval_env.step)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "d-UhypudApBy"
- },
- "outputs": [],
- "source": [
- "# initialize the state\n",
- "rng = jax.random.PRNGKey(0)\n",
- "state = jit_reset(rng)\n",
- "rollout = [state.pipeline_state]\n",
- "\n",
- "# grab a trajectory\n",
- "n_steps = 500\n",
- "render_every = 2\n",
- "\n",
- "for i in range(n_steps):\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.pipeline_state)\n",
- "\n",
- " if state.done:\n",
- " break\n",
- "\n",
- "media.show_video(env.render(rollout[::render_every], camera='side'), fps=1.0 / env.dt / render_every)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "zR-heox6LARK"
- },
- "source": [
- "# MJX Policy in MuJoCo\n",
- "\n",
- "We can also perform the physics step using the original MuJoCo python bindings to show that the policy trained in MJX works in MuJoCo."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "w6ixFi4dApBy"
- },
- "outputs": [],
- "source": [
- "mj_model = eval_env.sys.mj_model\n",
- "mj_data = mujoco.MjData(mj_model)\n",
- "\n",
- "renderer = mujoco.Renderer(mj_model)\n",
- "ctrl = jp.zeros(mj_model.nu)\n",
- "\n",
- "images = []\n",
- "for i in range(n_steps):\n",
- " act_rng, rng = jax.random.split(rng)\n",
- "\n",
- " obs = eval_env._get_obs(mjx.put_data(mj_model, mj_data), ctrl)\n",
- " ctrl, _ = jit_inference_fn(obs, act_rng)\n",
- "\n",
- " mj_data.ctrl = ctrl\n",
- " for _ in range(eval_env._n_frames):\n",
- " mujoco.mj_step(mj_model, mj_data) # Physics step using MuJoCo mj_step.\n",
- "\n",
- " if i % render_every == 0:\n",
- " renderer.update_scene(mj_data, camera='side')\n",
- " images.append(renderer.render())\n",
- "\n",
- "media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "65mIPj6DQNNa"
- },
- "source": [
- "# Training a Policy with Domain Randomization\n",
- "\n",
- "We might also want to include randomization over certain `mjModel` parameters while training a policy. In MJX, we can easily create a batch of environments with randomized values populated in `mjx.Model`. Below, we show a function that randomizes friction and actuator gain/bias."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "h8mhzKjHQuoL"
- },
- "outputs": [],
- "source": [
- "def domain_randomize(sys, rng):\n",
- " \"\"\"Randomizes the mjx.Model.\"\"\"\n",
- " @jax.vmap\n",
- " def rand(rng):\n",
- " _, key = jax.random.split(rng, 2)\n",
- " # friction\n",
- " friction = jax.random.uniform(key, (1,), minval=0.6, maxval=1.4)\n",
- " friction = sys.geom_friction.at[:, 0].set(friction)\n",
- " # actuator\n",
- " _, key = jax.random.split(key, 2)\n",
- " gain_range = (-5, 5)\n",
- " param = jax.random.uniform(\n",
- " key, (1,), minval=gain_range[0], maxval=gain_range[1]\n",
- " ) + sys.actuator_gainprm[:, 0]\n",
- " gain = sys.actuator_gainprm.at[:, 0].set(param)\n",
- " bias = sys.actuator_biasprm.at[:, 1].set(-param)\n",
- " return friction, gain, bias\n",
- "\n",
- " friction, gain, bias = rand(rng)\n",
- "\n",
- " in_axes = jax.tree_map(lambda x: None, sys)\n",
- " in_axes = in_axes.tree_replace({\n",
- " 'geom_friction': 0,\n",
- " 'actuator_gainprm': 0,\n",
- " 'actuator_biasprm': 0,\n",
- " })\n",
- "\n",
- " sys = sys.tree_replace({\n",
- " 'geom_friction': friction,\n",
- " 'actuator_gainprm': gain,\n",
- " 'actuator_biasprm': bias,\n",
- " })\n",
- "\n",
- " return sys, in_axes"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "gnsZo-GWSYYj"
- },
- "source": [
- "If we wanted 10 environments with randomized friction and actuator params, we can call `domain_randomize`, which returns a batched `mjx.Model` along with a dictionary specifying the axes that are batched."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "1K45Kp2ASV9s"
- },
- "outputs": [],
- "source": [
- "rng = jax.random.PRNGKey(0)\n",
- "rng = jax.random.split(rng, 10)\n",
- "batched_sys, _ = domain_randomize(env.sys, rng)\n",
- "\n",
- "print('Single env friction shape: ', env.sys.geom_friction.shape)\n",
- "print('Batched env friction shape: ', batched_sys.geom_friction.shape)\n",
- "\n",
- "print('Friction on geom 0: ', env.sys.geom_friction[0, 0])\n",
- "print('Random frictions on geom 0: ', batched_sys.geom_friction[:, 0, 0])"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "efnxNOnpQFuC"
- },
- "source": [
- "## Quadruped Env\n",
- "\n",
- "Let's define a quadruped environment that takes advantage of the domain randomization function. Here we use the [Barkour vb Quadruped](https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_vb) from [MuJoCo Menagerie](https://github.com/google-deepmind/mujoco_menagerie). We implement an environment that trains a joystick policy with Brax."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "VfyK73gtRXid"
- },
- "outputs": [],
- "source": [
- "!git clone https://github.com/google-deepmind/mujoco_menagerie"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "y79PoJOCIl-O"
- },
- "outputs": [],
- "source": [
- "#@title Barkour vb Quadruped Env\n",
- "\n",
- "def get_config():\n",
- " \"\"\"Returns reward config for barkour quadruped environment.\"\"\"\n",
- "\n",
- " def get_default_rewards_config():\n",
- " default_config = config_dict.ConfigDict(\n",
- " dict(\n",
- " # The coefficients for all reward terms used for training. All\n",
- " # physical quantities are in SI units, if no otherwise specified,\n",
- " # i.e. joint positions are in rad, positions are measured in meters,\n",
- " # torques in Nm, and time in seconds, and forces in Newtons.\n",
- " scales=config_dict.ConfigDict(\n",
- " dict(\n",
- " # Tracking rewards are computed using exp(-delta^2/sigma)\n",
- " # sigma can be a hyperparameters to tune.\n",
- " # Track the base x-y velocity (no z-velocity tracking.)\n",
- " tracking_lin_vel=1.5,\n",
- " # Track the angular velocity along z-axis, i.e. yaw rate.\n",
- " tracking_ang_vel=0.8,\n",
- " # Below are regularization terms, we roughly divide the\n",
- " # terms to base state regularizations, joint\n",
- " # regularizations, and other behavior regularizations.\n",
- " # Penalize the base velocity in z direction, L2 penalty.\n",
- " lin_vel_z=-2.0,\n",
- " # Penalize the base roll and pitch rate. L2 penalty.\n",
- " ang_vel_xy=-0.05,\n",
- " # Penalize non-zero roll and pitch angles. L2 penalty.\n",
- " orientation=-5.0,\n",
- " # L2 regularization of joint torques, |tau|^2.\n",
- " torques=-0.0002,\n",
- " # Penalize the change in the action and encourage smooth\n",
- " # actions. L2 regularization |action - last_action|^2\n",
- " action_rate=-0.01,\n",
- " # Encourage long swing steps. However, it does not\n",
- " # encourage high clearances.\n",
- " feet_air_time=0.2,\n",
- " # Encourage no motion at zero command, L2 regularization\n",
- " # |q - q_default|^2.\n",
- " stand_still=-0.5,\n",
- " # Early termination penalty.\n",
- " termination=-1.0,\n",
- " # Penalizing foot slipping on the ground.\n",
- " foot_slip=-0.1,\n",
- " )\n",
- " ),\n",
- " # Tracking reward = exp(-error^2/sigma).\n",
- " tracking_sigma=0.25,\n",
- " )\n",
- " )\n",
- " return default_config\n",
- "\n",
- " default_config = config_dict.ConfigDict(\n",
- " dict(\n",
- " rewards=get_default_rewards_config(),\n",
- " )\n",
- " )\n",
- "\n",
- " return default_config\n",
- "\n",
- "\n",
- "class BarkourEnv(PipelineEnv):\n",
- " \"\"\"Environment for training the barkour quadruped joystick policy in MJX.\"\"\"\n",
- "\n",
- " def __init__(\n",
- " self,\n",
- " obs_noise: float = 0.05,\n",
- " action_scale: float = 0.3,\n",
- " kick_vel: float = 0.05,\n",
- " **kwargs,\n",
- " ):\n",
- " path = epath.Path('mujoco_menagerie/google_barkour_vb/scene_mjx.xml')\n",
- " sys = mjcf.load(path.as_posix())\n",
- " self._dt = 0.02 # this environment is 50 fps\n",
- " sys = sys.tree_replace({'opt.timestep': 0.004, 'dt': 0.004})\n",
- "\n",
- " # override menagerie params for smoother policy\n",
- " sys = sys.replace(\n",
- " dof_damping=sys.dof_damping.at[6:].set(0.5239),\n",
- " actuator_gainprm=sys.actuator_gainprm.at[:, 0].set(35.0),\n",
- " actuator_biasprm=sys.actuator_biasprm.at[:, 1].set(-35.0),\n",
- " )\n",
- "\n",
- " n_frames = kwargs.pop('n_frames', int(self._dt / sys.opt.timestep))\n",
- " super().__init__(sys, backend='mjx', n_frames=n_frames)\n",
- "\n",
- " self.reward_config = get_config()\n",
- " # set custom from kwargs\n",
- " for k, v in kwargs.items():\n",
- " if k.endswith('_scale'):\n",
- " self.reward_config.rewards.scales[k[:-6]] = v\n",
- "\n",
- " self._torso_idx = mujoco.mj_name2id(\n",
- " sys.mj_model, mujoco.mjtObj.mjOBJ_BODY.value, 'torso'\n",
- " )\n",
- " self._action_scale = action_scale\n",
- " self._obs_noise = obs_noise\n",
- " self._kick_vel = kick_vel\n",
- " self._init_q = jp.array(sys.mj_model.keyframe('home').qpos)\n",
- " self._default_pose = sys.mj_model.keyframe('home').qpos[7:]\n",
- " self.lowers = jp.array([-0.7, -1.0, 0.05] * 4)\n",
- " self.uppers = jp.array([0.52, 2.1, 2.1] * 4)\n",
- " feet_site = [\n",
- " 'foot_front_left',\n",
- " 'foot_hind_left',\n",
- " 'foot_front_right',\n",
- " 'foot_hind_right',\n",
- " ]\n",
- " feet_site_id = [\n",
- " mujoco.mj_name2id(sys.mj_model, mujoco.mjtObj.mjOBJ_SITE.value, f)\n",
- " for f in feet_site\n",
- " ]\n",
- " assert not any(id_ == -1 for id_ in feet_site_id), 'Site not found.'\n",
- " self._feet_site_id = np.array(feet_site_id)\n",
- " lower_leg_body = [\n",
- " 'lower_leg_front_left',\n",
- " 'lower_leg_hind_left',\n",
- " 'lower_leg_front_right',\n",
- " 'lower_leg_hind_right',\n",
- " ]\n",
- " lower_leg_body_id = [\n",
- " mujoco.mj_name2id(sys.mj_model, mujoco.mjtObj.mjOBJ_BODY.value, l)\n",
- " for l in lower_leg_body\n",
- " ]\n",
- " assert not any(id_ == -1 for id_ in lower_leg_body_id), 'Body not found.'\n",
- " self._lower_leg_body_id = np.array(lower_leg_body_id)\n",
- " self._foot_radius = 0.0175\n",
- " self._nv = sys.nv\n",
- "\n",
- " def sample_command(self, rng: jax.Array) -\u003e jax.Array:\n",
- " lin_vel_x = [-0.6, 1.5] # min max [m/s]\n",
- " lin_vel_y = [-0.8, 0.8] # min max [m/s]\n",
- " ang_vel_yaw = [-0.7, 0.7] # min max [rad/s]\n",
- "\n",
- " _, key1, key2, key3 = jax.random.split(rng, 4)\n",
- " lin_vel_x = jax.random.uniform(\n",
- " key1, (1,), minval=lin_vel_x[0], maxval=lin_vel_x[1]\n",
- " )\n",
- " lin_vel_y = jax.random.uniform(\n",
- " key2, (1,), minval=lin_vel_y[0], maxval=lin_vel_y[1]\n",
- " )\n",
- " ang_vel_yaw = jax.random.uniform(\n",
- " key3, (1,), minval=ang_vel_yaw[0], maxval=ang_vel_yaw[1]\n",
- " )\n",
- " new_cmd = jp.array([lin_vel_x[0], lin_vel_y[0], ang_vel_yaw[0]])\n",
- " return new_cmd\n",
- "\n",
- " def reset(self, rng: jax.Array) -\u003e State: # pytype: disable=signature-mismatch\n",
- " rng, key = jax.random.split(rng)\n",
- "\n",
- " pipeline_state = self.pipeline_init(self._init_q, jp.zeros(self._nv))\n",
- "\n",
- " state_info = {\n",
- " 'rng': rng,\n",
- " 'last_act': jp.zeros(12),\n",
- " 'last_vel': jp.zeros(12),\n",
- " 'command': self.sample_command(key),\n",
- " 'last_contact': jp.zeros(4, dtype=bool),\n",
- " 'feet_air_time': jp.zeros(4),\n",
- " 'rewards': {k: 0.0 for k in self.reward_config.rewards.scales.keys()},\n",
- " 'kick': jp.array([0.0, 0.0]),\n",
- " 'step': 0,\n",
- " }\n",
- "\n",
- " obs_history = jp.zeros(15 * 31) # store 15 steps of history\n",
- " obs = self._get_obs(pipeline_state, state_info, obs_history)\n",
- " reward, done = jp.zeros(2)\n",
- " metrics = {'total_dist': 0.0}\n",
- " for k in state_info['rewards']:\n",
- " metrics[k] = state_info['rewards'][k]\n",
- " state = State(pipeline_state, obs, reward, done, metrics, state_info) # pytype: disable=wrong-arg-types\n",
- " return state\n",
- "\n",
- " def step(self, state: State, action: jax.Array) -\u003e State: # pytype: disable=signature-mismatch\n",
- " rng, cmd_rng, kick_noise_2 = jax.random.split(state.info['rng'], 3)\n",
- "\n",
- " # kick\n",
- " push_interval = 10\n",
- " kick_theta = jax.random.uniform(kick_noise_2, maxval=2 * jp.pi)\n",
- " kick = jp.array([jp.cos(kick_theta), jp.sin(kick_theta)])\n",
- " kick *= jp.mod(state.info['step'], push_interval) == 0\n",
- " qvel = state.pipeline_state.qvel # pytype: disable=attribute-error\n",
- " qvel = qvel.at[:2].set(kick * self._kick_vel + qvel[:2])\n",
- " state = state.tree_replace({'pipeline_state.qvel': qvel})\n",
- "\n",
- " # physics step\n",
- " motor_targets = self._default_pose + action * self._action_scale\n",
- " motor_targets = jp.clip(motor_targets, self.lowers, self.uppers)\n",
- " pipeline_state = self.pipeline_step(state.pipeline_state, motor_targets)\n",
- " x, xd = pipeline_state.x, pipeline_state.xd\n",
- "\n",
- " # observation data\n",
- " obs = self._get_obs(pipeline_state, state.info, state.obs)\n",
- " joint_angles = pipeline_state.q[7:]\n",
- " joint_vel = pipeline_state.qd[6:]\n",
- "\n",
- " # foot contact data based on z-position\n",
- " foot_pos = pipeline_state.site_xpos[self._feet_site_id] # pytype: disable=attribute-error\n",
- " foot_contact_z = foot_pos[:, 2] - self._foot_radius\n",
- " contact = foot_contact_z \u003c 1e-3 # a mm or less off the floor\n",
- " contact_filt_mm = contact | state.info['last_contact']\n",
- " contact_filt_cm = (foot_contact_z \u003c 3e-2) | state.info['last_contact']\n",
- " first_contact = (state.info['feet_air_time'] \u003e 0) * contact_filt_mm\n",
- " state.info['feet_air_time'] += self.dt\n",
- "\n",
- " # done if joint limits are reached or robot is falling\n",
- " up = jp.array([0.0, 0.0, 1.0])\n",
- " done = jp.dot(math.rotate(up, x.rot[self._torso_idx - 1]), up) \u003c 0\n",
- " done |= jp.any(joint_angles \u003c self.lowers)\n",
- " done |= jp.any(joint_angles \u003e self.uppers)\n",
- " done |= pipeline_state.x.pos[self._torso_idx - 1, 2] \u003c 0.18\n",
- "\n",
- " # reward\n",
- " rewards = {\n",
- " 'tracking_lin_vel': (\n",
- " self._reward_tracking_lin_vel(state.info['command'], x, xd)\n",
- " ),\n",
- " 'tracking_ang_vel': (\n",
- " self._reward_tracking_ang_vel(state.info['command'], x, xd)\n",
- " ),\n",
- " 'lin_vel_z': self._reward_lin_vel_z(xd),\n",
- " 'ang_vel_xy': self._reward_ang_vel_xy(xd),\n",
- " 'orientation': self._reward_orientation(x),\n",
- " 'torques': self._reward_torques(pipeline_state.qfrc_actuator), # pytype: disable=attribute-error\n",
- " 'action_rate': self._reward_action_rate(action, state.info['last_act']),\n",
- " 'stand_still': self._reward_stand_still(\n",
- " state.info['command'], joint_angles,\n",
- " ),\n",
- " 'feet_air_time': self._reward_feet_air_time(\n",
- " state.info['feet_air_time'],\n",
- " first_contact,\n",
- " state.info['command'],\n",
- " ),\n",
- " 'foot_slip': self._reward_foot_slip(pipeline_state, contact_filt_cm),\n",
- " 'termination': self._reward_termination(done, state.info['step']),\n",
- " }\n",
- " rewards = {\n",
- " k: v * self.reward_config.rewards.scales[k] for k, v in rewards.items()\n",
- " }\n",
- " reward = jp.clip(sum(rewards.values()) * self.dt, 0.0, 10000.0)\n",
- "\n",
- " # state management\n",
- " state.info['kick'] = kick\n",
- " state.info['last_act'] = action\n",
- " state.info['last_vel'] = joint_vel\n",
- " state.info['feet_air_time'] *= ~contact_filt_mm\n",
- " state.info['last_contact'] = contact\n",
- " state.info['rewards'] = rewards\n",
- " state.info['step'] += 1\n",
- " state.info['rng'] = rng\n",
- "\n",
- " # sample new command if more than 500 timesteps achieved\n",
- " state.info['command'] = jp.where(\n",
- " state.info['step'] \u003e 500,\n",
- " self.sample_command(cmd_rng),\n",
- " state.info['command'],\n",
- " )\n",
- " # reset the step counter when done\n",
- " state.info['step'] = jp.where(\n",
- " done | (state.info['step'] \u003e 500), 0, state.info['step']\n",
- " )\n",
- "\n",
- " # log total displacement as a proxy metric\n",
- " state.metrics['total_dist'] = math.normalize(x.pos[self._torso_idx - 1])[1]\n",
- " state.metrics.update(state.info['rewards'])\n",
- "\n",
- " done = jp.float32(done)\n",
- " state = state.replace(\n",
- " pipeline_state=pipeline_state, obs=obs, reward=reward, done=done\n",
- " )\n",
- " return state\n",
- "\n",
- " def _get_obs(\n",
- " self,\n",
- " pipeline_state: base.State,\n",
- " state_info: dict[str, Any],\n",
- " obs_history: jax.Array,\n",
- " ) -\u003e jax.Array:\n",
- " inv_torso_rot = math.quat_inv(pipeline_state.x.rot[0])\n",
- " local_rpyrate = math.rotate(pipeline_state.xd.ang[0], inv_torso_rot)\n",
- "\n",
- " obs = jp.concatenate([\n",
- " jp.array([local_rpyrate[2]]) * 0.25, # yaw rate\n",
- " math.rotate(jp.array([0, 0, -1]), inv_torso_rot), # projected gravity\n",
- " state_info['command'] * jp.array([2.0, 2.0, 0.25]), # command\n",
- " pipeline_state.q[7:] - self._default_pose, # motor angles\n",
- " state_info['last_act'], # last action\n",
- " ])\n",
- "\n",
- " # clip, noise\n",
- " obs = jp.clip(obs, -100.0, 100.0) + self._obs_noise * jax.random.uniform(\n",
- " state_info['rng'], obs.shape, minval=-1, maxval=1\n",
- " )\n",
- " # stack observations through time\n",
- " obs = jp.roll(obs_history, obs.size).at[:obs.size].set(obs)\n",
- "\n",
- " return obs\n",
- "\n",
- " # ------------ reward functions----------------\n",
- " def _reward_lin_vel_z(self, xd: Motion) -\u003e jax.Array:\n",
- " # Penalize z axis base linear velocity\n",
- " return jp.square(xd.vel[0, 2])\n",
- "\n",
- " def _reward_ang_vel_xy(self, xd: Motion) -\u003e jax.Array:\n",
- " # Penalize xy axes base angular velocity\n",
- " return jp.sum(jp.square(xd.ang[0, :2]))\n",
- "\n",
- " def _reward_orientation(self, x: Transform) -\u003e jax.Array:\n",
- " # Penalize non flat base orientation\n",
- " up = jp.array([0.0, 0.0, 1.0])\n",
- " rot_up = math.rotate(up, x.rot[0])\n",
- " return jp.sum(jp.square(rot_up[:2]))\n",
- "\n",
- " def _reward_torques(self, torques: jax.Array) -\u003e jax.Array:\n",
- " # Penalize torques\n",
- " return jp.sqrt(jp.sum(jp.square(torques))) + jp.sum(jp.abs(torques))\n",
- "\n",
- " def _reward_action_rate(\n",
- " self, act: jax.Array, last_act: jax.Array\n",
- " ) -\u003e jax.Array:\n",
- " # Penalize changes in actions\n",
- " return jp.sum(jp.square(act - last_act))\n",
- "\n",
- " def _reward_tracking_lin_vel(\n",
- " self, commands: jax.Array, x: Transform, xd: Motion\n",
- " ) -\u003e jax.Array:\n",
- " # Tracking of linear velocity commands (xy axes)\n",
- " local_vel = math.rotate(xd.vel[0], math.quat_inv(x.rot[0]))\n",
- " lin_vel_error = jp.sum(jp.square(commands[:2] - local_vel[:2]))\n",
- " lin_vel_reward = jp.exp(\n",
- " -lin_vel_error / self.reward_config.rewards.tracking_sigma\n",
- " )\n",
- " return lin_vel_reward\n",
- "\n",
- " def _reward_tracking_ang_vel(\n",
- " self, commands: jax.Array, x: Transform, xd: Motion\n",
- " ) -\u003e jax.Array:\n",
- " # Tracking of angular velocity commands (yaw)\n",
- " base_ang_vel = math.rotate(xd.ang[0], math.quat_inv(x.rot[0]))\n",
- " ang_vel_error = jp.square(commands[2] - base_ang_vel[2])\n",
- " return jp.exp(-ang_vel_error / self.reward_config.rewards.tracking_sigma)\n",
- "\n",
- " def _reward_feet_air_time(\n",
- " self, air_time: jax.Array, first_contact: jax.Array, commands: jax.Array\n",
- " ) -\u003e jax.Array:\n",
- " # Reward air time.\n",
- " rew_air_time = jp.sum((air_time - 0.1) * first_contact)\n",
- " rew_air_time *= (\n",
- " math.normalize(commands[:2])[1] \u003e 0.05\n",
- " ) # no reward for zero command\n",
- " return rew_air_time\n",
- "\n",
- " def _reward_stand_still(\n",
- " self,\n",
- " commands: jax.Array,\n",
- " joint_angles: jax.Array,\n",
- " ) -\u003e jax.Array:\n",
- " # Penalize motion at zero commands\n",
- " return jp.sum(jp.abs(joint_angles - self._default_pose)) * (\n",
- " math.normalize(commands[:2])[1] \u003c 0.1\n",
- " )\n",
- "\n",
- " def _reward_foot_slip(\n",
- " self, pipeline_state: base.State, contact_filt: jax.Array\n",
- " ) -\u003e jax.Array:\n",
- " # get velocities at feet which are offset from lower legs\n",
- " # pytype: disable=attribute-error\n",
- " pos = pipeline_state.site_xpos[self._feet_site_id] # feet position\n",
- " feet_offset = pos - pipeline_state.xpos[self._lower_leg_body_id]\n",
- " # pytype: enable=attribute-error\n",
- " offset = base.Transform.create(pos=feet_offset)\n",
- " foot_indices = self._lower_leg_body_id - 1 # we got rid of the world body\n",
- " foot_vel = offset.vmap().do(pipeline_state.xd.take(foot_indices)).vel\n",
- "\n",
- " # Penalize large feet velocity for feet that are in contact with the ground.\n",
- " return jp.sum(jp.square(foot_vel[:, :2]) * contact_filt.reshape((-1, 1)))\n",
- "\n",
- " def _reward_termination(self, done: jax.Array, step: jax.Array) -\u003e jax.Array:\n",
- " return done \u0026 (step \u003c 500)\n",
- "\n",
- " def render(\n",
- " self, trajectory: List[base.State], camera: str | None = None\n",
- " ) -\u003e Sequence[np.ndarray]:\n",
- " camera = camera or 'track'\n",
- " return super().render(trajectory, camera=camera)\n",
- "\n",
- "envs.register_environment('barkour', BarkourEnv)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "pi_yrcz-Qp3W"
- },
- "outputs": [],
- "source": [
- "env_name = 'barkour'\n",
- "env = envs.get_environment(env_name)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "nxaNFP9mA23H"
- },
- "source": [
- "## Train Policy\n",
- "\n",
- "To train a policy with domain randomization, we pass in the domain randomization function into the brax train function; brax will call the domain randomization function when rolling out episodes. Training the quadruped takes 6 minutes on a Tesla A100 GPU."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "cHJCbESGA7Rk"
- },
- "outputs": [],
- "source": [
- "make_networks_factory = functools.partial(\n",
- " ppo_networks.make_ppo_networks,\n",
- " policy_hidden_layer_sizes=(128, 128, 128, 128))\n",
- "train_fn = functools.partial(\n",
- " ppo.train, num_timesteps=100_000_000, num_evals=10,\n",
- " reward_scaling=1, episode_length=1000, normalize_observations=True,\n",
- " action_repeat=1, unroll_length=20, num_minibatches=32,\n",
- " num_updates_per_batch=4, discounting=0.97, learning_rate=3.0e-4,\n",
- " entropy_cost=1e-2, num_envs=8192, batch_size=256,\n",
- " network_factory=make_networks_factory,\n",
- " randomization_fn=domain_randomize, seed=0)\n",
- "\n",
- "x_data = []\n",
- "y_data = []\n",
- "ydataerr = []\n",
- "times = [datetime.now()]\n",
- "max_y, min_y = 40, 0\n",
- "\n",
- "# Reset environments since internals may be overwritten by tracers from the\n",
- "# domain randomization function.\n",
- "env = envs.get_environment(env_name)\n",
- "eval_env = envs.get_environment(env_name)\n",
- "make_inference_fn, params, _= train_fn(environment=env,\n",
- " progress_fn=progress,\n",
- " eval_env=eval_env)\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": {
- "id": "8Yge-CGP5JoO"
- },
- "outputs": [],
- "source": [
- "# Save and reload params.\n",
- "model_path = '/tmp/mjx_brax_quadruped_policy'\n",
- "model.save_params(model_path, params)\n",
- "params = model.load_params(model_path)\n",
- "\n",
- "inference_fn = make_inference_fn(params)\n",
- "jit_inference_fn = jax.jit(inference_fn)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "L01IrN4oCIkC"
- },
- "source": [
- "## Visualize Policy\n",
- "\n",
- "For the Barkour Quadruped, the joystick commands can be set through `x_vel`, `y_vel`, and `ang_vel`. `x_vel` and `y_vel` define the linear forward and sideways velocities with respect to the quadruped torso. `ang_vel` defines the angular velocity of the torso in the z direction."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "VTbpEtXnEecd"
- },
- "outputs": [],
- "source": [
- "eval_env = envs.get_environment(env_name)\n",
- "\n",
- "jit_reset = jax.jit(eval_env.reset)\n",
- "jit_step = jax.jit(eval_env.step)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "HRRN-8L-BivZ"
- },
- "outputs": [],
- "source": [
- "# @markdown Commands **only used for Barkour Env**:\n",
- "x_vel = 1.0 #@param {type: \"number\"}\n",
- "y_vel = 0.0 #@param {type: \"number\"}\n",
- "ang_vel = -0.5 #@param {type: \"number\"}\n",
- "\n",
- "the_command = jp.array([x_vel, y_vel, ang_vel])\n",
- "\n",
- "# initialize the state\n",
- "rng = jax.random.PRNGKey(0)\n",
- "state = jit_reset(rng)\n",
- "state.info['command'] = the_command\n",
- "rollout = [state.pipeline_state]\n",
- "\n",
- "# grab a trajectory\n",
- "n_steps = 500\n",
- "render_every = 2\n",
- "\n",
- "for i in range(n_steps):\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.pipeline_state)\n",
- "\n",
- "media.show_video(\n",
- " eval_env.render(rollout[::render_every], camera='track'),\n",
- " fps=1.0 / eval_env.dt / render_every)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "aD6H6WD0915X"
- },
- "source": [
- "We can also render the rollout using the Brax renderer."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "V7jqv08X95u4"
- },
- "outputs": [],
- "source": [
- "HTML(html.render(eval_env.sys.replace(dt=eval_env.dt), rollout))"
- ]
- }
- ],
- "metadata": {
- "accelerator": "GPU",
- "colab": {
- "gpuClass": "premium",
- "gpuType": "V100",
- "machine_shape": "hm",
- "private_outputs": true,
- "provenance": [
- {
- "file_id": "1A58SK07tnOzix53E68D0TQ2ePCTZA61f",
- "timestamp": 1707342610876
- }
- ],
- "toc_visible": true
- },
- "kernelspec": {
- "display_name": "Python 3",
- "name": "python3"
- }
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "MpkYHwCqk7W-"
+ },
+ "source": [
+ "\n",
+ "\n",
+ "#
Tutorial
\n",
+ "\n",
+ "This notebook provides an introductory tutorial for [**MuJoCo XLA (MJX)**](https://github.com/google-deepmind/mujoco/blob/main/mjx), a JAX-based implementation of MuJoCo useful for RL training workloads.\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"
+ ]
},
- "nbformat": 4,
- "nbformat_minor": 0
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "xBSdkbmGN2K-"
+ },
+ "source": [
+ "### Copyright notice"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "_UbO9uhtBSX5"
+ },
+ "source": [
+ ">
Copyright 2023 DeepMind Technologies Limited.
\n",
+ ">
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 http://www.apache.org/licenses/LICENSE-2.0.
\n",
+ ">
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.
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "YvyGCsgSCxHQ"
+ },
+ "source": [
+ "# Install MuJoCo, MJX, and Brax"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "Xqo7pyX-n72M"
+ },
+ "outputs": [],
+ "source": [
+ "!pip install mujoco\n",
+ "!pip install mujoco_mjx\n",
+ "!pip install brax"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "cellView": "form",
+ "id": "IbZxYDxzoz5R"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Check if MuJoCo installation was successful\n",
+ "\n",
+ "from google.colab import files\n",
+ "\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",
+ "# 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\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",
+ " 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.')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "cellView": "form",
+ "id": "T5f4w3Kq2X14"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Import packages for plotting and creating graphics\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",
+ "!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": 0,
+ "metadata": {
+ "id": "ObF1UXrkb0Nd"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Import MuJoCo, MJX, and Brax\n",
+ "\n",
+ "\n",
+ "from datetime import datetime\n",
+ "import functools\n",
+ "from IPython.display import HTML\n",
+ "import jax\n",
+ "from jax import numpy as jp\n",
+ "import numpy as np\n",
+ "from typing import Any, Dict, Sequence, Tuple, Union\n",
+ "\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.envs.base import Env, PipelineEnv, State\n",
+ "from brax.mjx.base import State as MjxState\n",
+ "from brax.training.agents.ppo import train as ppo\n",
+ "from brax.training.agents.ppo import networks as ppo_networks\n",
+ "from brax.io import html, mjcf, model\n",
+ "\n",
+ "from etils import epath\n",
+ "from flax import struct\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"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Nj4-Xmx4DFaq"
+ },
+ "source": [
+ "# Introduction to MJX\n",
+ "\n",
+ "MJX is an implementation of MuJoCo written in [JAX](https://jax.readthedocs.io/en/latest/index.html), enabling large batch training on GPU/TPU. In this notebook, we will demonstrate how to train RL policies with MJX.\n",
+ "\n",
+ "Before we get into hefty RL workloads, let's get started with a simpler example! The entrypoint into MJX is through MuJoCo, so first we load a MuJoCo model:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "bNus3mbbDz6a"
+ },
+ "outputs": [],
+ "source": [
+ "xml = \"\"\"\n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ "\"\"\"\n",
+ "\n",
+ "# Make model, data, and renderer\n",
+ "mj_model = mujoco.MjModel.from_xml_string(xml)\n",
+ "mj_data = mujoco.MjData(mj_model)\n",
+ "renderer = mujoco.Renderer(mj_model)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Po5oykJbFQbj"
+ },
+ "source": [
+ "Next we take the MuJoCo model and data, and place them on the GPU device using MJX."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "TSpoOWqeEC3P"
+ },
+ "outputs": [],
+ "source": [
+ "mjx_model = mjx.put_model(mj_model)\n",
+ "mjx_data = mjx.put_data(mj_model, mj_data)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "6rxMMSs4OJJf"
+ },
+ "source": [
+ "Below, we print the `qpos` from MuJoCo and MJX. Notice that the `qpos` for the mjData is a numpy array living on the CPU, while the `qpos` for `mjx.Data` is a JAX Array living on the GPU device."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "ZOD582pfOLP-"
+ },
+ "outputs": [],
+ "source": [
+ "print(mj_data.qpos, type(mj_data.qpos))\n",
+ "print(mjx_data.qpos, type(mjx_data.qpos), mjx_data.qpos.devices())"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "ZShF9-o_JLm3"
+ },
+ "source": [
+ "Let's run the simulation in MuJoCo and render the trajectory. This example is taken from the [MuJoCo tutorial](https://colab.sandbox.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "HDlPlX05I3m-"
+ },
+ "outputs": [],
+ "source": [
+ "# 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",
+ "frames = []\n",
+ "mujoco.mj_resetData(mj_model, mj_data)\n",
+ "while mj_data.time < duration:\n",
+ " mujoco.mj_step(mj_model, mj_data)\n",
+ " if len(frames) < mj_data.time * framerate:\n",
+ " renderer.update_scene(mj_data, scene_option=scene_option)\n",
+ " pixels = renderer.render()\n",
+ " frames.append(pixels)\n",
+ "\n",
+ "# Simulate and display video.\n",
+ "media.show_video(frames, fps=framerate)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "m70b_RxBJOyd"
+ },
+ "source": [
+ "Now let's run the same exact simulation on the GPU device using MJX!\n",
+ "\n",
+ "In the example below, we use `mjx.step` instead of `mujoco.mj_step`, and we also [`jax.jit`](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html) the `mjx.step` so that it runs efficiently on the GPU. After each step, we convert the `mjx.Data` back to `mjData` so that we can use the MuJoCo renderer.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "Pr29xq0-JRQv"
+ },
+ "outputs": [],
+ "source": [
+ "\n",
+ "jit_step = jax.jit(mjx.step)\n",
+ "\n",
+ "frames = []\n",
+ "mujoco.mj_resetData(mj_model, mj_data)\n",
+ "mjx_data = mjx.put_data(mj_model, mj_data)\n",
+ "while mjx_data.time < duration:\n",
+ " mjx_data = jit_step(mjx_model, mjx_data)\n",
+ " if len(frames) < mjx_data.time * framerate:\n",
+ " mj_data = mjx.get_data(mj_model, mjx_data)\n",
+ " renderer.update_scene(mj_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": "wXsQ4qO2KO3Q"
+ },
+ "source": [
+ "Running single threaded physics simulation on the GPU is not very [efficient](https://mujoco.readthedocs.io/en/stable/mjx.html#mjx-the-sharp-bits). The advantage with MJX is that we can run environments in parallel on a hardware accelerated device. Let's try it out!\n",
+ "\n",
+ "In the example below, we create 4096 copies of the `mjx.Data` and we run the `mjx.step` over the batched data. Since MJX is implemented in JAX, we take advantage of [`jax.vmap`](https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html) to run the `mjx.step` in parallel over all `mjx.Data`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "rrdrcKRVK6w9"
+ },
+ "outputs": [],
+ "source": [
+ "rng = jax.random.PRNGKey(0)\n",
+ "rng = jax.random.split(rng, 4096)\n",
+ "batch = jax.vmap(lambda rng: mjx_data.replace(qpos=jax.random.uniform(rng, (1,))))(rng)\n",
+ "\n",
+ "jit_step = jax.vmap(mjx.step, in_axes=(None, 0))\n",
+ "batch = jit_step(mjx_model, batch)\n",
+ "\n",
+ "print(batch.qpos)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "x4lL220cOj0q"
+ },
+ "source": [
+ "We can copy the batched `mjx.Data` back to MuJoCo like we did before:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "Jtz7j1PDOnw5"
+ },
+ "outputs": [],
+ "source": [
+ "batched_mj_data = mjx.get_data(mj_model, batch)\n",
+ "print([d.qpos for d in batched_mj_data])"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "RAv6WUVUm78k"
+ },
+ "source": [
+ "# Training a Policy with MJX\n",
+ "\n",
+ "Running large batch physics simulation is useful for training RL policies. Here we demonstrate training RL policies with MJX using the RL library from [Brax](https://github.com/google/brax).\n",
+ "\n",
+ "Below, we implement the classic Humanoid environment using MJX and Brax. We inherit from the `MjxEnv` implementation in Brax so that we can step the physics with MJX while training with Brax RL implementations.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "mtGMYNLE3QJN"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Humanoid Env\n",
+ "\n",
+ "class Humanoid(PipelineEnv):\n",
+ "\n",
+ " def __init__(\n",
+ " self,\n",
+ " forward_reward_weight=1.25,\n",
+ " ctrl_cost_weight=0.1,\n",
+ " healthy_reward=5.0,\n",
+ " terminate_when_unhealthy=True,\n",
+ " healthy_z_range=(1.0, 2.0),\n",
+ " reset_noise_scale=1e-2,\n",
+ " exclude_current_positions_from_observation=True,\n",
+ " **kwargs,\n",
+ " ):\n",
+ " path = epath.Path(epath.resource_path('mujoco')) / (\n",
+ " 'mjx/test_data/humanoid'\n",
+ " )\n",
+ " mj_model = mujoco.MjModel.from_xml_path(\n",
+ " (path / 'humanoid.xml').as_posix())\n",
+ " mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG\n",
+ " mj_model.opt.iterations = 6\n",
+ " mj_model.opt.ls_iterations = 6\n",
+ "\n",
+ " sys = mjcf.load_model(mj_model)\n",
+ "\n",
+ " physics_steps_per_control_step = 5\n",
+ " kwargs['n_frames'] = kwargs.get(\n",
+ " 'n_frames', physics_steps_per_control_step)\n",
+ " kwargs['backend'] = 'mjx'\n",
+ "\n",
+ " super().__init__(sys, **kwargs)\n",
+ "\n",
+ " self._forward_reward_weight = forward_reward_weight\n",
+ " self._ctrl_cost_weight = ctrl_cost_weight\n",
+ " self._healthy_reward = healthy_reward\n",
+ " self._terminate_when_unhealthy = terminate_when_unhealthy\n",
+ " self._healthy_z_range = healthy_z_range\n",
+ " self._reset_noise_scale = reset_noise_scale\n",
+ " self._exclude_current_positions_from_observation = (\n",
+ " exclude_current_positions_from_observation\n",
+ " )\n",
+ "\n",
+ " def reset(self, rng: jp.ndarray) -> State:\n",
+ " \"\"\"Resets the environment to an initial state.\"\"\"\n",
+ " rng, rng1, rng2 = jax.random.split(rng, 3)\n",
+ "\n",
+ " low, hi = -self._reset_noise_scale, self._reset_noise_scale\n",
+ " qpos = self.sys.qpos0 + jax.random.uniform(\n",
+ " rng1, (self.sys.nq,), minval=low, maxval=hi\n",
+ " )\n",
+ " qvel = jax.random.uniform(\n",
+ " rng2, (self.sys.nv,), minval=low, maxval=hi\n",
+ " )\n",
+ "\n",
+ " data = self.pipeline_init(qpos, qvel)\n",
+ "\n",
+ " obs = self._get_obs(data, jp.zeros(self.sys.nu))\n",
+ " reward, done, zero = jp.zeros(3)\n",
+ " metrics = {\n",
+ " 'forward_reward': zero,\n",
+ " 'reward_linvel': zero,\n",
+ " 'reward_quadctrl': zero,\n",
+ " 'reward_alive': zero,\n",
+ " 'x_position': zero,\n",
+ " 'y_position': zero,\n",
+ " 'distance_from_origin': zero,\n",
+ " 'x_velocity': zero,\n",
+ " 'y_velocity': zero,\n",
+ " }\n",
+ " return State(data, obs, reward, done, metrics)\n",
+ "\n",
+ " def step(self, state: State, action: jp.ndarray) -> State:\n",
+ " \"\"\"Runs one timestep of the environment's dynamics.\"\"\"\n",
+ " data0 = state.pipeline_state\n",
+ " data = self.pipeline_step(data0, action)\n",
+ "\n",
+ " com_before = data0.subtree_com[1]\n",
+ " com_after = data.subtree_com[1]\n",
+ " velocity = (com_after - com_before) / self.dt\n",
+ " forward_reward = self._forward_reward_weight * velocity[0]\n",
+ "\n",
+ " min_z, max_z = self._healthy_z_range\n",
+ " is_healthy = jp.where(data.q[2] < min_z, 0.0, 1.0)\n",
+ " is_healthy = jp.where(data.q[2] > max_z, 0.0, is_healthy)\n",
+ " if self._terminate_when_unhealthy:\n",
+ " healthy_reward = self._healthy_reward\n",
+ " else:\n",
+ " healthy_reward = self._healthy_reward * is_healthy\n",
+ "\n",
+ " ctrl_cost = self._ctrl_cost_weight * jp.sum(jp.square(action))\n",
+ "\n",
+ " obs = self._get_obs(data, action)\n",
+ " reward = forward_reward + healthy_reward - ctrl_cost\n",
+ " done = 1.0 - is_healthy if self._terminate_when_unhealthy else 0.0\n",
+ " state.metrics.update(\n",
+ " forward_reward=forward_reward,\n",
+ " reward_linvel=forward_reward,\n",
+ " reward_quadctrl=-ctrl_cost,\n",
+ " reward_alive=healthy_reward,\n",
+ " x_position=com_after[0],\n",
+ " y_position=com_after[1],\n",
+ " distance_from_origin=jp.linalg.norm(com_after),\n",
+ " x_velocity=velocity[0],\n",
+ " y_velocity=velocity[1],\n",
+ " )\n",
+ "\n",
+ " return state.replace(\n",
+ " pipeline_state=data, obs=obs, reward=reward, done=done\n",
+ " )\n",
+ "\n",
+ " def _get_obs(\n",
+ " self, data: mjx.Data, action: jp.ndarray\n",
+ " ) -> jp.ndarray:\n",
+ " \"\"\"Observes humanoid body position, velocities, and angles.\"\"\"\n",
+ " position = data.qpos\n",
+ " if self._exclude_current_positions_from_observation:\n",
+ " position = position[2:]\n",
+ "\n",
+ " # external_contact_forces are excluded\n",
+ " return jp.concatenate([\n",
+ " position,\n",
+ " data.qvel,\n",
+ " data.cinert[1:].ravel(),\n",
+ " data.cvel[1:].ravel(),\n",
+ " data.qfrc_actuator,\n",
+ " ])\n",
+ "\n",
+ "\n",
+ "envs.register_environment('humanoid', Humanoid)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "P1K6IznI2y83"
+ },
+ "source": [
+ "## Visualize a Rollout\n",
+ "\n",
+ "Let's instantiate the environment and visualize a short rollout.\n",
+ "\n",
+ "NOTE: Since episodes terminates early if the torso is below the healthy z-range, the only relevant contacts for this task are between the feet and the plane. We turn off other contacts."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "EhKLFK54C1CH"
+ },
+ "outputs": [],
+ "source": [
+ "# instantiate the environment\n",
+ "env_name = 'humanoid'\n",
+ "env = envs.get_environment(env_name)\n",
+ "\n",
+ "# define the jit reset/step functions\n",
+ "jit_reset = jax.jit(env.reset)\n",
+ "jit_step = jax.jit(env.step)\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "Ph8u-v2Q2xLS"
+ },
+ "outputs": [],
+ "source": [
+ "# initialize the state\n",
+ "state = jit_reset(jax.random.PRNGKey(0))\n",
+ "rollout = [state.pipeline_state]\n",
+ "\n",
+ "# grab a trajectory\n",
+ "for i in range(10):\n",
+ " ctrl = -0.1 * jp.ones(env.sys.nu)\n",
+ " state = jit_step(state, ctrl)\n",
+ " rollout.append(state.pipeline_state)\n",
+ "\n",
+ "media.show_video(env.render(rollout, camera='side'), fps=1.0 / env.dt)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "BQDG6NQ1CbZD"
+ },
+ "source": [
+ "## Train Humanoid Policy\n",
+ "\n",
+ "Let's now train a policy with PPO to make the Humanoid run forwards. Training takes about 6 minutes on a Tesla A100 GPU."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "xLiddQYPApBw"
+ },
+ "outputs": [],
+ "source": [
+ "train_fn = functools.partial(\n",
+ " ppo.train, num_timesteps=30_000_000, num_evals=5, reward_scaling=0.1,\n",
+ " episode_length=1000, normalize_observations=True, action_repeat=1,\n",
+ " unroll_length=10, num_minibatches=32, num_updates_per_batch=8,\n",
+ " discounting=0.97, learning_rate=3e-4, entropy_cost=1e-3, num_envs=2048,\n",
+ " batch_size=1024, seed=0)\n",
+ "\n",
+ "\n",
+ "x_data = []\n",
+ "y_data = []\n",
+ "ydataerr = []\n",
+ "times = [datetime.now()]\n",
+ "\n",
+ "max_y, min_y = 13000, 0\n",
+ "def progress(num_steps, metrics):\n",
+ " times.append(datetime.now())\n",
+ " x_data.append(num_steps)\n",
+ " y_data.append(metrics['eval/episode_reward'])\n",
+ " ydataerr.append(metrics['eval/episode_reward_std'])\n",
+ "\n",
+ " plt.xlim([0, train_fn.keywords['num_timesteps'] * 1.25])\n",
+ " plt.ylim([min_y, max_y])\n",
+ "\n",
+ " plt.xlabel('# environment steps')\n",
+ " plt.ylabel('reward per episode')\n",
+ " plt.title(f'y={y_data[-1]:.3f}')\n",
+ "\n",
+ " plt.errorbar(\n",
+ " x_data, y_data, yerr=ydataerr)\n",
+ " plt.show()\n",
+ "\n",
+ "make_inference_fn, params, _= train_fn(environment=env, progress_fn=progress)\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": "YYIch0HEApBx"
+ },
+ "source": [
+ "\n",
+ "\n",
+ "We can save and load the policy using the brax model API."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "Z8gI6qH6ApBx"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Save Model\n",
+ "model_path = '/tmp/mjx_brax_policy'\n",
+ "model.save_params(model_path, params)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "h4reaWgxApBx"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Load Model and Define Inference Function\n",
+ "params = model.load_params(model_path)\n",
+ "\n",
+ "inference_fn = make_inference_fn(params)\n",
+ "jit_inference_fn = jax.jit(inference_fn)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "0G357XIfApBy"
+ },
+ "source": [
+ "## Visualize Policy\n",
+ "\n",
+ "Finally we can visualize the policy."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "osYasMw4ApBy"
+ },
+ "outputs": [],
+ "source": [
+ "eval_env = envs.get_environment(env_name)\n",
+ "\n",
+ "jit_reset = jax.jit(eval_env.reset)\n",
+ "jit_step = jax.jit(eval_env.step)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "d-UhypudApBy"
+ },
+ "outputs": [],
+ "source": [
+ "# initialize the state\n",
+ "rng = jax.random.PRNGKey(0)\n",
+ "state = jit_reset(rng)\n",
+ "rollout = [state.pipeline_state]\n",
+ "\n",
+ "# grab a trajectory\n",
+ "n_steps = 500\n",
+ "render_every = 2\n",
+ "\n",
+ "for i in range(n_steps):\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.pipeline_state)\n",
+ "\n",
+ " if state.done:\n",
+ " break\n",
+ "\n",
+ "media.show_video(env.render(rollout[::render_every], camera='side'), fps=1.0 / env.dt / render_every)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "zR-heox6LARK"
+ },
+ "source": [
+ "# MJX Policy in MuJoCo\n",
+ "\n",
+ "We can also perform the physics step using the original MuJoCo python bindings to show that the policy trained in MJX works in MuJoCo."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "w6ixFi4dApBy"
+ },
+ "outputs": [],
+ "source": [
+ "mj_model = eval_env.sys.mj_model\n",
+ "mj_data = mujoco.MjData(mj_model)\n",
+ "\n",
+ "renderer = mujoco.Renderer(mj_model)\n",
+ "ctrl = jp.zeros(mj_model.nu)\n",
+ "\n",
+ "images = []\n",
+ "for i in range(n_steps):\n",
+ " act_rng, rng = jax.random.split(rng)\n",
+ "\n",
+ " obs = eval_env._get_obs(mjx.put_data(mj_model, mj_data), ctrl)\n",
+ " ctrl, _ = jit_inference_fn(obs, act_rng)\n",
+ "\n",
+ " mj_data.ctrl = ctrl\n",
+ " for _ in range(eval_env._n_frames):\n",
+ " mujoco.mj_step(mj_model, mj_data) # Physics step using MuJoCo mj_step.\n",
+ "\n",
+ " if i % render_every == 0:\n",
+ " renderer.update_scene(mj_data, camera='side')\n",
+ " images.append(renderer.render())\n",
+ "\n",
+ "media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "65mIPj6DQNNa"
+ },
+ "source": [
+ "# Training a Policy with Domain Randomization\n",
+ "\n",
+ "We might also want to include randomization over certain `mjModel` parameters while training a policy. In MJX, we can easily create a batch of environments with randomized values populated in `mjx.Model`. Below, we show a function that randomizes friction and actuator gain/bias."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "h8mhzKjHQuoL"
+ },
+ "outputs": [],
+ "source": [
+ "def domain_randomize(sys, rng):\n",
+ " \"\"\"Randomizes the mjx.Model.\"\"\"\n",
+ " @jax.vmap\n",
+ " def rand(rng):\n",
+ " _, key = jax.random.split(rng, 2)\n",
+ " # friction\n",
+ " friction = jax.random.uniform(key, (1,), minval=0.6, maxval=1.4)\n",
+ " friction = sys.geom_friction.at[:, 0].set(friction)\n",
+ " # actuator\n",
+ " _, key = jax.random.split(key, 2)\n",
+ " gain_range = (-5, 5)\n",
+ " param = jax.random.uniform(\n",
+ " key, (1,), minval=gain_range[0], maxval=gain_range[1]\n",
+ " ) + sys.actuator_gainprm[:, 0]\n",
+ " gain = sys.actuator_gainprm.at[:, 0].set(param)\n",
+ " bias = sys.actuator_biasprm.at[:, 1].set(-param)\n",
+ " return friction, gain, bias\n",
+ "\n",
+ " friction, gain, bias = rand(rng)\n",
+ "\n",
+ " in_axes = jax.tree_map(lambda x: None, sys)\n",
+ " in_axes = in_axes.tree_replace({\n",
+ " 'geom_friction': 0,\n",
+ " 'actuator_gainprm': 0,\n",
+ " 'actuator_biasprm': 0,\n",
+ " })\n",
+ "\n",
+ " sys = sys.tree_replace({\n",
+ " 'geom_friction': friction,\n",
+ " 'actuator_gainprm': gain,\n",
+ " 'actuator_biasprm': bias,\n",
+ " })\n",
+ "\n",
+ " return sys, in_axes"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "gnsZo-GWSYYj"
+ },
+ "source": [
+ "If we wanted 10 environments with randomized friction and actuator params, we can call `domain_randomize`, which returns a batched `mjx.Model` along with a dictionary specifying the axes that are batched."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "1K45Kp2ASV9s"
+ },
+ "outputs": [],
+ "source": [
+ "rng = jax.random.PRNGKey(0)\n",
+ "rng = jax.random.split(rng, 10)\n",
+ "batched_sys, _ = domain_randomize(env.sys, rng)\n",
+ "\n",
+ "print('Single env friction shape: ', env.sys.geom_friction.shape)\n",
+ "print('Batched env friction shape: ', batched_sys.geom_friction.shape)\n",
+ "\n",
+ "print('Friction on geom 0: ', env.sys.geom_friction[0, 0])\n",
+ "print('Random frictions on geom 0: ', batched_sys.geom_friction[:, 0, 0])"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "efnxNOnpQFuC"
+ },
+ "source": [
+ "## Quadruped Env\n",
+ "\n",
+ "Let's define a quadruped environment that takes advantage of the domain randomization function. Here we use the [Barkour vb Quadruped](https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_vb) from [MuJoCo Menagerie](https://github.com/google-deepmind/mujoco_menagerie). We implement an environment that trains a joystick policy with Brax."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "VfyK73gtRXid"
+ },
+ "outputs": [],
+ "source": [
+ "!git clone https://github.com/google-deepmind/mujoco_menagerie"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "y79PoJOCIl-O"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Barkour vb Quadruped Env\n",
+ "\n",
+ "def get_config():\n",
+ " \"\"\"Returns reward config for barkour quadruped environment.\"\"\"\n",
+ "\n",
+ " def get_default_rewards_config():\n",
+ " default_config = config_dict.ConfigDict(\n",
+ " dict(\n",
+ " # The coefficients for all reward terms used for training. All\n",
+ " # physical quantities are in SI units, if no otherwise specified,\n",
+ " # i.e. joint positions are in rad, positions are measured in meters,\n",
+ " # torques in Nm, and time in seconds, and forces in Newtons.\n",
+ " scales=config_dict.ConfigDict(\n",
+ " dict(\n",
+ " # Tracking rewards are computed using exp(-delta^2/sigma)\n",
+ " # sigma can be a hyperparameters to tune.\n",
+ " # Track the base x-y velocity (no z-velocity tracking.)\n",
+ " tracking_lin_vel=1.5,\n",
+ " # Track the angular velocity along z-axis, i.e. yaw rate.\n",
+ " tracking_ang_vel=0.8,\n",
+ " # Below are regularization terms, we roughly divide the\n",
+ " # terms to base state regularizations, joint\n",
+ " # regularizations, and other behavior regularizations.\n",
+ " # Penalize the base velocity in z direction, L2 penalty.\n",
+ " lin_vel_z=-2.0,\n",
+ " # Penalize the base roll and pitch rate. L2 penalty.\n",
+ " ang_vel_xy=-0.05,\n",
+ " # Penalize non-zero roll and pitch angles. L2 penalty.\n",
+ " orientation=-5.0,\n",
+ " # L2 regularization of joint torques, |tau|^2.\n",
+ " torques=-0.0002,\n",
+ " # Penalize the change in the action and encourage smooth\n",
+ " # actions. L2 regularization |action - last_action|^2\n",
+ " action_rate=-0.01,\n",
+ " # Encourage long swing steps. However, it does not\n",
+ " # encourage high clearances.\n",
+ " feet_air_time=0.2,\n",
+ " # Encourage no motion at zero command, L2 regularization\n",
+ " # |q - q_default|^2.\n",
+ " stand_still=-0.5,\n",
+ " # Early termination penalty.\n",
+ " termination=-1.0,\n",
+ " # Penalizing foot slipping on the ground.\n",
+ " foot_slip=-0.1,\n",
+ " )\n",
+ " ),\n",
+ " # Tracking reward = exp(-error^2/sigma).\n",
+ " tracking_sigma=0.25,\n",
+ " )\n",
+ " )\n",
+ " return default_config\n",
+ "\n",
+ " default_config = config_dict.ConfigDict(\n",
+ " dict(\n",
+ " rewards=get_default_rewards_config(),\n",
+ " )\n",
+ " )\n",
+ "\n",
+ " return default_config\n",
+ "\n",
+ "\n",
+ "class BarkourEnv(PipelineEnv):\n",
+ " \"\"\"Environment for training the barkour quadruped joystick policy in MJX.\"\"\"\n",
+ "\n",
+ " def __init__(\n",
+ " self,\n",
+ " obs_noise: float = 0.05,\n",
+ " action_scale: float = 0.3,\n",
+ " kick_vel: float = 0.05,\n",
+ " **kwargs,\n",
+ " ):\n",
+ " path = epath.Path('mujoco_menagerie/google_barkour_vb/scene_mjx.xml')\n",
+ " sys = mjcf.load(path.as_posix())\n",
+ " self._dt = 0.02 # this environment is 50 fps\n",
+ " sys = sys.tree_replace({'opt.timestep': 0.004, 'dt': 0.004})\n",
+ "\n",
+ " # override menagerie params for smoother policy\n",
+ " sys = sys.replace(\n",
+ " dof_damping=sys.dof_damping.at[6:].set(0.5239),\n",
+ " actuator_gainprm=sys.actuator_gainprm.at[:, 0].set(35.0),\n",
+ " actuator_biasprm=sys.actuator_biasprm.at[:, 1].set(-35.0),\n",
+ " )\n",
+ "\n",
+ " n_frames = kwargs.pop('n_frames', int(self._dt / sys.opt.timestep))\n",
+ " super().__init__(sys, backend='mjx', n_frames=n_frames)\n",
+ "\n",
+ " self.reward_config = get_config()\n",
+ " # set custom from kwargs\n",
+ " for k, v in kwargs.items():\n",
+ " if k.endswith('_scale'):\n",
+ " self.reward_config.rewards.scales[k[:-6]] = v\n",
+ "\n",
+ " self._torso_idx = mujoco.mj_name2id(\n",
+ " sys.mj_model, mujoco.mjtObj.mjOBJ_BODY.value, 'torso'\n",
+ " )\n",
+ " self._action_scale = action_scale\n",
+ " self._obs_noise = obs_noise\n",
+ " self._kick_vel = kick_vel\n",
+ " self._init_q = jp.array(sys.mj_model.keyframe('home').qpos)\n",
+ " self._default_pose = sys.mj_model.keyframe('home').qpos[7:]\n",
+ " self.lowers = jp.array([-0.7, -1.0, 0.05] * 4)\n",
+ " self.uppers = jp.array([0.52, 2.1, 2.1] * 4)\n",
+ " feet_site = [\n",
+ " 'foot_front_left',\n",
+ " 'foot_hind_left',\n",
+ " 'foot_front_right',\n",
+ " 'foot_hind_right',\n",
+ " ]\n",
+ " feet_site_id = [\n",
+ " mujoco.mj_name2id(sys.mj_model, mujoco.mjtObj.mjOBJ_SITE.value, f)\n",
+ " for f in feet_site\n",
+ " ]\n",
+ " assert not any(id_ == -1 for id_ in feet_site_id), 'Site not found.'\n",
+ " self._feet_site_id = np.array(feet_site_id)\n",
+ " lower_leg_body = [\n",
+ " 'lower_leg_front_left',\n",
+ " 'lower_leg_hind_left',\n",
+ " 'lower_leg_front_right',\n",
+ " 'lower_leg_hind_right',\n",
+ " ]\n",
+ " lower_leg_body_id = [\n",
+ " mujoco.mj_name2id(sys.mj_model, mujoco.mjtObj.mjOBJ_BODY.value, l)\n",
+ " for l in lower_leg_body\n",
+ " ]\n",
+ " assert not any(id_ == -1 for id_ in lower_leg_body_id), 'Body not found.'\n",
+ " self._lower_leg_body_id = np.array(lower_leg_body_id)\n",
+ " self._foot_radius = 0.0175\n",
+ " self._nv = sys.nv\n",
+ "\n",
+ " def sample_command(self, rng: jax.Array) -> jax.Array:\n",
+ " lin_vel_x = [-0.6, 1.5] # min max [m/s]\n",
+ " lin_vel_y = [-0.8, 0.8] # min max [m/s]\n",
+ " ang_vel_yaw = [-0.7, 0.7] # min max [rad/s]\n",
+ "\n",
+ " _, key1, key2, key3 = jax.random.split(rng, 4)\n",
+ " lin_vel_x = jax.random.uniform(\n",
+ " key1, (1,), minval=lin_vel_x[0], maxval=lin_vel_x[1]\n",
+ " )\n",
+ " lin_vel_y = jax.random.uniform(\n",
+ " key2, (1,), minval=lin_vel_y[0], maxval=lin_vel_y[1]\n",
+ " )\n",
+ " ang_vel_yaw = jax.random.uniform(\n",
+ " key3, (1,), minval=ang_vel_yaw[0], maxval=ang_vel_yaw[1]\n",
+ " )\n",
+ " new_cmd = jp.array([lin_vel_x[0], lin_vel_y[0], ang_vel_yaw[0]])\n",
+ " return new_cmd\n",
+ "\n",
+ " def reset(self, rng: jax.Array) -> State: # pytype: disable=signature-mismatch\n",
+ " rng, key = jax.random.split(rng)\n",
+ "\n",
+ " pipeline_state = self.pipeline_init(self._init_q, jp.zeros(self._nv))\n",
+ "\n",
+ " state_info = {\n",
+ " 'rng': rng,\n",
+ " 'last_act': jp.zeros(12),\n",
+ " 'last_vel': jp.zeros(12),\n",
+ " 'command': self.sample_command(key),\n",
+ " 'last_contact': jp.zeros(4, dtype=bool),\n",
+ " 'feet_air_time': jp.zeros(4),\n",
+ " 'rewards': {k: 0.0 for k in self.reward_config.rewards.scales.keys()},\n",
+ " 'kick': jp.array([0.0, 0.0]),\n",
+ " 'step': 0,\n",
+ " }\n",
+ "\n",
+ " obs_history = jp.zeros(15 * 31) # store 15 steps of history\n",
+ " obs = self._get_obs(pipeline_state, state_info, obs_history)\n",
+ " reward, done = jp.zeros(2)\n",
+ " metrics = {'total_dist': 0.0}\n",
+ " for k in state_info['rewards']:\n",
+ " metrics[k] = state_info['rewards'][k]\n",
+ " state = State(pipeline_state, obs, reward, done, metrics, state_info) # pytype: disable=wrong-arg-types\n",
+ " return state\n",
+ "\n",
+ " def step(self, state: State, action: jax.Array) -> State: # pytype: disable=signature-mismatch\n",
+ " rng, cmd_rng, kick_noise_2 = jax.random.split(state.info['rng'], 3)\n",
+ "\n",
+ " # kick\n",
+ " push_interval = 10\n",
+ " kick_theta = jax.random.uniform(kick_noise_2, maxval=2 * jp.pi)\n",
+ " kick = jp.array([jp.cos(kick_theta), jp.sin(kick_theta)])\n",
+ " kick *= jp.mod(state.info['step'], push_interval) == 0\n",
+ " qvel = state.pipeline_state.qvel # pytype: disable=attribute-error\n",
+ " qvel = qvel.at[:2].set(kick * self._kick_vel + qvel[:2])\n",
+ " state = state.tree_replace({'pipeline_state.qvel': qvel})\n",
+ "\n",
+ " # physics step\n",
+ " motor_targets = self._default_pose + action * self._action_scale\n",
+ " motor_targets = jp.clip(motor_targets, self.lowers, self.uppers)\n",
+ " pipeline_state = self.pipeline_step(state.pipeline_state, motor_targets)\n",
+ " x, xd = pipeline_state.x, pipeline_state.xd\n",
+ "\n",
+ " # observation data\n",
+ " obs = self._get_obs(pipeline_state, state.info, state.obs)\n",
+ " joint_angles = pipeline_state.q[7:]\n",
+ " joint_vel = pipeline_state.qd[6:]\n",
+ "\n",
+ " # foot contact data based on z-position\n",
+ " foot_pos = pipeline_state.site_xpos[self._feet_site_id] # pytype: disable=attribute-error\n",
+ " foot_contact_z = foot_pos[:, 2] - self._foot_radius\n",
+ " contact = foot_contact_z < 1e-3 # a mm or less off the floor\n",
+ " contact_filt_mm = contact | state.info['last_contact']\n",
+ " contact_filt_cm = (foot_contact_z < 3e-2) | state.info['last_contact']\n",
+ " first_contact = (state.info['feet_air_time'] > 0) * contact_filt_mm\n",
+ " state.info['feet_air_time'] += self.dt\n",
+ "\n",
+ " # done if joint limits are reached or robot is falling\n",
+ " up = jp.array([0.0, 0.0, 1.0])\n",
+ " done = jp.dot(math.rotate(up, x.rot[self._torso_idx - 1]), up) < 0\n",
+ " done |= jp.any(joint_angles < self.lowers)\n",
+ " done |= jp.any(joint_angles > self.uppers)\n",
+ " done |= pipeline_state.x.pos[self._torso_idx - 1, 2] < 0.18\n",
+ "\n",
+ " # reward\n",
+ " rewards = {\n",
+ " 'tracking_lin_vel': (\n",
+ " self._reward_tracking_lin_vel(state.info['command'], x, xd)\n",
+ " ),\n",
+ " 'tracking_ang_vel': (\n",
+ " self._reward_tracking_ang_vel(state.info['command'], x, xd)\n",
+ " ),\n",
+ " 'lin_vel_z': self._reward_lin_vel_z(xd),\n",
+ " 'ang_vel_xy': self._reward_ang_vel_xy(xd),\n",
+ " 'orientation': self._reward_orientation(x),\n",
+ " 'torques': self._reward_torques(pipeline_state.qfrc_actuator), # pytype: disable=attribute-error\n",
+ " 'action_rate': self._reward_action_rate(action, state.info['last_act']),\n",
+ " 'stand_still': self._reward_stand_still(\n",
+ " state.info['command'], joint_angles,\n",
+ " ),\n",
+ " 'feet_air_time': self._reward_feet_air_time(\n",
+ " state.info['feet_air_time'],\n",
+ " first_contact,\n",
+ " state.info['command'],\n",
+ " ),\n",
+ " 'foot_slip': self._reward_foot_slip(pipeline_state, contact_filt_cm),\n",
+ " 'termination': self._reward_termination(done, state.info['step']),\n",
+ " }\n",
+ " rewards = {\n",
+ " k: v * self.reward_config.rewards.scales[k] for k, v in rewards.items()\n",
+ " }\n",
+ " reward = jp.clip(sum(rewards.values()) * self.dt, 0.0, 10000.0)\n",
+ "\n",
+ " # state management\n",
+ " state.info['kick'] = kick\n",
+ " state.info['last_act'] = action\n",
+ " state.info['last_vel'] = joint_vel\n",
+ " state.info['feet_air_time'] *= ~contact_filt_mm\n",
+ " state.info['last_contact'] = contact\n",
+ " state.info['rewards'] = rewards\n",
+ " state.info['step'] += 1\n",
+ " state.info['rng'] = rng\n",
+ "\n",
+ " # sample new command if more than 500 timesteps achieved\n",
+ " state.info['command'] = jp.where(\n",
+ " state.info['step'] > 500,\n",
+ " self.sample_command(cmd_rng),\n",
+ " state.info['command'],\n",
+ " )\n",
+ " # reset the step counter when done\n",
+ " state.info['step'] = jp.where(\n",
+ " done | (state.info['step'] > 500), 0, state.info['step']\n",
+ " )\n",
+ "\n",
+ " # log total displacement as a proxy metric\n",
+ " state.metrics['total_dist'] = math.normalize(x.pos[self._torso_idx - 1])[1]\n",
+ " state.metrics.update(state.info['rewards'])\n",
+ "\n",
+ " done = jp.float32(done)\n",
+ " state = state.replace(\n",
+ " pipeline_state=pipeline_state, obs=obs, reward=reward, done=done\n",
+ " )\n",
+ " return state\n",
+ "\n",
+ " def _get_obs(\n",
+ " self,\n",
+ " pipeline_state: base.State,\n",
+ " state_info: dict[str, Any],\n",
+ " obs_history: jax.Array,\n",
+ " ) -> jax.Array:\n",
+ " inv_torso_rot = math.quat_inv(pipeline_state.x.rot[0])\n",
+ " local_rpyrate = math.rotate(pipeline_state.xd.ang[0], inv_torso_rot)\n",
+ "\n",
+ " obs = jp.concatenate([\n",
+ " jp.array([local_rpyrate[2]]) * 0.25, # yaw rate\n",
+ " math.rotate(jp.array([0, 0, -1]), inv_torso_rot), # projected gravity\n",
+ " state_info['command'] * jp.array([2.0, 2.0, 0.25]), # command\n",
+ " pipeline_state.q[7:] - self._default_pose, # motor angles\n",
+ " state_info['last_act'], # last action\n",
+ " ])\n",
+ "\n",
+ " # clip, noise\n",
+ " obs = jp.clip(obs, -100.0, 100.0) + self._obs_noise * jax.random.uniform(\n",
+ " state_info['rng'], obs.shape, minval=-1, maxval=1\n",
+ " )\n",
+ " # stack observations through time\n",
+ " obs = jp.roll(obs_history, obs.size).at[:obs.size].set(obs)\n",
+ "\n",
+ " return obs\n",
+ "\n",
+ " # ------------ reward functions----------------\n",
+ " def _reward_lin_vel_z(self, xd: Motion) -> jax.Array:\n",
+ " # Penalize z axis base linear velocity\n",
+ " return jp.square(xd.vel[0, 2])\n",
+ "\n",
+ " def _reward_ang_vel_xy(self, xd: Motion) -> jax.Array:\n",
+ " # Penalize xy axes base angular velocity\n",
+ " return jp.sum(jp.square(xd.ang[0, :2]))\n",
+ "\n",
+ " def _reward_orientation(self, x: Transform) -> jax.Array:\n",
+ " # Penalize non flat base orientation\n",
+ " up = jp.array([0.0, 0.0, 1.0])\n",
+ " rot_up = math.rotate(up, x.rot[0])\n",
+ " return jp.sum(jp.square(rot_up[:2]))\n",
+ "\n",
+ " def _reward_torques(self, torques: jax.Array) -> jax.Array:\n",
+ " # Penalize torques\n",
+ " return jp.sqrt(jp.sum(jp.square(torques))) + jp.sum(jp.abs(torques))\n",
+ "\n",
+ " def _reward_action_rate(\n",
+ " self, act: jax.Array, last_act: jax.Array\n",
+ " ) -> jax.Array:\n",
+ " # Penalize changes in actions\n",
+ " return jp.sum(jp.square(act - last_act))\n",
+ "\n",
+ " def _reward_tracking_lin_vel(\n",
+ " self, commands: jax.Array, x: Transform, xd: Motion\n",
+ " ) -> jax.Array:\n",
+ " # Tracking of linear velocity commands (xy axes)\n",
+ " local_vel = math.rotate(xd.vel[0], math.quat_inv(x.rot[0]))\n",
+ " lin_vel_error = jp.sum(jp.square(commands[:2] - local_vel[:2]))\n",
+ " lin_vel_reward = jp.exp(\n",
+ " -lin_vel_error / self.reward_config.rewards.tracking_sigma\n",
+ " )\n",
+ " return lin_vel_reward\n",
+ "\n",
+ " def _reward_tracking_ang_vel(\n",
+ " self, commands: jax.Array, x: Transform, xd: Motion\n",
+ " ) -> jax.Array:\n",
+ " # Tracking of angular velocity commands (yaw)\n",
+ " base_ang_vel = math.rotate(xd.ang[0], math.quat_inv(x.rot[0]))\n",
+ " ang_vel_error = jp.square(commands[2] - base_ang_vel[2])\n",
+ " return jp.exp(-ang_vel_error / self.reward_config.rewards.tracking_sigma)\n",
+ "\n",
+ " def _reward_feet_air_time(\n",
+ " self, air_time: jax.Array, first_contact: jax.Array, commands: jax.Array\n",
+ " ) -> jax.Array:\n",
+ " # Reward air time.\n",
+ " rew_air_time = jp.sum((air_time - 0.1) * first_contact)\n",
+ " rew_air_time *= (\n",
+ " math.normalize(commands[:2])[1] > 0.05\n",
+ " ) # no reward for zero command\n",
+ " return rew_air_time\n",
+ "\n",
+ " def _reward_stand_still(\n",
+ " self,\n",
+ " commands: jax.Array,\n",
+ " joint_angles: jax.Array,\n",
+ " ) -> jax.Array:\n",
+ " # Penalize motion at zero commands\n",
+ " return jp.sum(jp.abs(joint_angles - self._default_pose)) * (\n",
+ " math.normalize(commands[:2])[1] < 0.1\n",
+ " )\n",
+ "\n",
+ " def _reward_foot_slip(\n",
+ " self, pipeline_state: base.State, contact_filt: jax.Array\n",
+ " ) -> jax.Array:\n",
+ " # get velocities at feet which are offset from lower legs\n",
+ " # pytype: disable=attribute-error\n",
+ " pos = pipeline_state.site_xpos[self._feet_site_id] # feet position\n",
+ " feet_offset = pos - pipeline_state.xpos[self._lower_leg_body_id]\n",
+ " # pytype: enable=attribute-error\n",
+ " offset = base.Transform.create(pos=feet_offset)\n",
+ " foot_indices = self._lower_leg_body_id - 1 # we got rid of the world body\n",
+ " foot_vel = offset.vmap().do(pipeline_state.xd.take(foot_indices)).vel\n",
+ "\n",
+ " # Penalize large feet velocity for feet that are in contact with the ground.\n",
+ " return jp.sum(jp.square(foot_vel[:, :2]) * contact_filt.reshape((-1, 1)))\n",
+ "\n",
+ " def _reward_termination(self, done: jax.Array, step: jax.Array) -> jax.Array:\n",
+ " return done & (step < 500)\n",
+ "\n",
+ " def render(\n",
+ " self, trajectory: List[base.State], camera: str | None = None\n",
+ " ) -> Sequence[np.ndarray]:\n",
+ " camera = camera or 'track'\n",
+ " return super().render(trajectory, camera=camera)\n",
+ "\n",
+ "envs.register_environment('barkour', BarkourEnv)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "pi_yrcz-Qp3W"
+ },
+ "outputs": [],
+ "source": [
+ "env_name = 'barkour'\n",
+ "env = envs.get_environment(env_name)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "nxaNFP9mA23H"
+ },
+ "source": [
+ "## Train Policy\n",
+ "\n",
+ "To train a policy with domain randomization, we pass in the domain randomization function into the brax train function; brax will call the domain randomization function when rolling out episodes. Training the quadruped takes 6 minutes on a Tesla A100 GPU."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "cHJCbESGA7Rk"
+ },
+ "outputs": [],
+ "source": [
+ "make_networks_factory = functools.partial(\n",
+ " ppo_networks.make_ppo_networks,\n",
+ " policy_hidden_layer_sizes=(128, 128, 128, 128))\n",
+ "train_fn = functools.partial(\n",
+ " ppo.train, num_timesteps=100_000_000, num_evals=10,\n",
+ " reward_scaling=1, episode_length=1000, normalize_observations=True,\n",
+ " action_repeat=1, unroll_length=20, num_minibatches=32,\n",
+ " num_updates_per_batch=4, discounting=0.97, learning_rate=3.0e-4,\n",
+ " entropy_cost=1e-2, num_envs=8192, batch_size=256,\n",
+ " network_factory=make_networks_factory,\n",
+ " randomization_fn=domain_randomize, seed=0)\n",
+ "\n",
+ "x_data = []\n",
+ "y_data = []\n",
+ "ydataerr = []\n",
+ "times = [datetime.now()]\n",
+ "max_y, min_y = 40, 0\n",
+ "\n",
+ "# Reset environments since internals may be overwritten by tracers from the\n",
+ "# domain randomization function.\n",
+ "env = envs.get_environment(env_name)\n",
+ "eval_env = envs.get_environment(env_name)\n",
+ "make_inference_fn, params, _= train_fn(environment=env,\n",
+ " progress_fn=progress,\n",
+ " eval_env=eval_env)\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": 0,
+ "metadata": {
+ "id": "8Yge-CGP5JoO"
+ },
+ "outputs": [],
+ "source": [
+ "# Save and reload params.\n",
+ "model_path = '/tmp/mjx_brax_quadruped_policy'\n",
+ "model.save_params(model_path, params)\n",
+ "params = model.load_params(model_path)\n",
+ "\n",
+ "inference_fn = make_inference_fn(params)\n",
+ "jit_inference_fn = jax.jit(inference_fn)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "L01IrN4oCIkC"
+ },
+ "source": [
+ "## Visualize Policy\n",
+ "\n",
+ "For the Barkour Quadruped, the joystick commands can be set through `x_vel`, `y_vel`, and `ang_vel`. `x_vel` and `y_vel` define the linear forward and sideways velocities with respect to the quadruped torso. `ang_vel` defines the angular velocity of the torso in the z direction."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "VTbpEtXnEecd"
+ },
+ "outputs": [],
+ "source": [
+ "eval_env = envs.get_environment(env_name)\n",
+ "\n",
+ "jit_reset = jax.jit(eval_env.reset)\n",
+ "jit_step = jax.jit(eval_env.step)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "HRRN-8L-BivZ"
+ },
+ "outputs": [],
+ "source": [
+ "# @markdown Commands **only used for Barkour Env**:\n",
+ "x_vel = 1.0 #@param {type: \"number\"}\n",
+ "y_vel = 0.0 #@param {type: \"number\"}\n",
+ "ang_vel = -0.5 #@param {type: \"number\"}\n",
+ "\n",
+ "the_command = jp.array([x_vel, y_vel, ang_vel])\n",
+ "\n",
+ "# initialize the state\n",
+ "rng = jax.random.PRNGKey(0)\n",
+ "state = jit_reset(rng)\n",
+ "state.info['command'] = the_command\n",
+ "rollout = [state.pipeline_state]\n",
+ "\n",
+ "# grab a trajectory\n",
+ "n_steps = 500\n",
+ "render_every = 2\n",
+ "\n",
+ "for i in range(n_steps):\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.pipeline_state)\n",
+ "\n",
+ "media.show_video(\n",
+ " eval_env.render(rollout[::render_every], camera='track'),\n",
+ " fps=1.0 / eval_env.dt / render_every)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "aD6H6WD0915X"
+ },
+ "source": [
+ "We can also render the rollout using the Brax renderer."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "V7jqv08X95u4"
+ },
+ "outputs": [],
+ "source": [
+ "HTML(html.render(eval_env.sys.replace(dt=eval_env.dt), rollout))"
+ ]
+ }
+ ],
+ "metadata": {
+ "accelerator": "GPU",
+ "colab": {
+ "gpuClass": "premium",
+ "gpuType": "V100",
+ "machine_shape": "hm",
+ "private_outputs": true,
+ "toc_visible": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
}
diff --git a/python/LQR.ipynb b/python/LQR.ipynb
index 98ec50f8..3c6020ed 100644
--- a/python/LQR.ipynb
+++ b/python/LQR.ipynb
@@ -1,966 +1,961 @@
{
- "cells": [
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "bR2IJtzSilBv"
- },
- "source": [
- "\n",
- "\n",
- "# \u003ch1\u003e\u003ccenter\u003eLQR tutorial \u003ca href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/LQR.ipynb\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/\u003e\u003c/a\u003e\u003c/center\u003e\u003c/h1\u003e\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 \u003e Change runtime type\".\n"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "LBAvTJ0xHKy7"
- },
- "source": [
- "### Copyright notice"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "_UbO9uhtBSX5"
- },
- "source": [
- "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eCopyright 2022 DeepMind Technologies Limited\u003c/small\u003e\u003c/p\u003e\n",
- "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eLicensed 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 \u003ca href=\"http://www.apache.org/licenses/LICENSE-2.0\"\u003ehttp://www.apache.org/licenses/LICENSE-2.0\u003c/a\u003e.\u003c/small\u003e\u003c/small\u003e\u003c/p\u003e\n",
- "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eUnless 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.\u003c/small\u003e\u003c/small\u003e\u003c/p\u003e"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "QPdJNe3k62mx"
- },
- "source": [
- "### Install MuJoCo\n"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "Xqo7pyX-n72M"
- },
- "outputs": [],
- "source": [
- "!pip install mujoco"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "IbZxYDxzoz5R"
- },
- "outputs": [],
- "source": [
- "#@title Check if installation was successful\n",
- "\n",
- "from google.colab import files\n",
- "\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",
- "try:\n",
- " print('Checking that the installation succeeded:')\n",
- " import mujoco\n",
- " mujoco.MjModel.from_xml_string('\u003cmujoco/\u003e')\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.')"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "T5f4w3Kq2X14"
- },
- "outputs": [],
- "source": [
- "#@title 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 \u003e/dev/null || (apt update \u0026\u0026 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": "markdown",
- "metadata": {
- "id": "J5fL6p-Sx5DB"
- },
- "source": [
- "## Loading and rendering the standard humanoid"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "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()"
- ]
- },
- {
- "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": null,
- "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": null,
- "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": null,
- "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": null,
- "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 \u003c DURATION:\n",
- " # Step the simulation.\n",
- " mujoco.mj_step(model, data)\n",
- "\n",
- " # Render and save frames.\n",
- " if len(frames) \u003c 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": null,
- "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 \u003c 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) \u003c 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": null,
- "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": null,
- "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": null,
- "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": null,
- "metadata": {
- "id": "a1PF_yHdPvLl"
- },
- "outputs": [],
- "source": [
- "ctrl0 = np.atleast_2d(qfrc0) @ np.linalg.pinv(data.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": null,
- "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": null,
- "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 \u003c DURATION:\n",
- " # Step the simulation.\n",
- " mujoco.mj_step(model, data)\n",
- "\n",
- " # Render and save frames.\n",
- " if len(frames) \u003c 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": null,
- "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": null,
- "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": null,
- "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 seperate 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": null,
- "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": null,
- "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": null,
- "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": null,
- "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 \u003c 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) \u003c 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": null,
- "metadata": {
- "id": "zJmbOJMurRna"
- },
- "outputs": [],
- "source": [
- "# Parameters.\n",
- "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",
- "\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) -\u003e float:\n",
- " return 1 - np.cos(normalised_time*2*np.pi)\n",
- "def azimuth(time: float) -\u003e 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",
- "# 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 \u003c 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) \u003c 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": [
- "LBAvTJ0xHKy7"
- ],
- "last_runtime": {
- "build_target": "",
- "kind": "local"
- },
- "private_outputs": true,
- "provenance": []
- },
- "kernelspec": {
- "display_name": "Python 3",
- "name": "python3"
- }
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "bR2IJtzSilBv"
+ },
+ "source": [
+ "\n",
+ "\n",
+ "#
LQR tutorial
\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"
+ ]
},
- "nbformat": 4,
- "nbformat_minor": 0
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "LBAvTJ0xHKy7"
+ },
+ "source": [
+ "### Copyright notice"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "_UbO9uhtBSX5"
+ },
+ "source": [
+ ">
Copyright 2022 DeepMind Technologies Limited
\n",
+ ">
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 http://www.apache.org/licenses/LICENSE-2.0.
\n",
+ ">
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.
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "QPdJNe3k62mx"
+ },
+ "source": [
+ "### Install MuJoCo\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "Xqo7pyX-n72M"
+ },
+ "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",
+ "from google.colab import files\n",
+ "\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",
+ "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.')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "cellView": "form",
+ "id": "T5f4w3Kq2X14"
+ },
+ "outputs": [],
+ "source": [
+ "#@title 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)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "J5fL6p-Sx5DB"
+ },
+ "source": [
+ "## Loading and rendering 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()"
+ ]
+ },
+ {
+ "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": [
+ "ctrl0 = np.atleast_2d(qfrc0) @ np.linalg.pinv(data.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_STD = 0.05 # actuator units\n",
+ "CTRL_RATE = 0.8 # seconds\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",
+ "# 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": [
+ "LBAvTJ0xHKy7"
+ ],
+ "private_outputs": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
}
diff --git a/python/tutorial.ipynb b/python/tutorial.ipynb
index c12aadeb..ee50e564 100644
--- a/python/tutorial.ipynb
+++ b/python/tutorial.ipynb
@@ -1,2176 +1,2175 @@
{
- "cells": [
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "MpkYHwCqk7W-"
- },
- "source": [
- "\n",
- "\n",
- "# \u003ch1\u003e\u003ccenter\u003eTutorial \u003ca href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/\u003e\u003c/a\u003e\u003c/center\u003e\u003c/h1\u003e\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 \u003e Change runtime type\".\n",
- "\n",
- "\n",
- "\n",
- "\n",
- "\n",
- "\n",
- "\n",
- "\n"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "xBSdkbmGN2K-"
- },
- "source": [
- "### Copyright notice"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "_UbO9uhtBSX5"
- },
- "source": [
- "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eCopyright 2022 DeepMind Technologies Limited.\u003c/small\u003e\u003c/p\u003e\n",
- "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eLicensed 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 \u003ca href=\"http://www.apache.org/licenses/LICENSE-2.0\"\u003ehttp://www.apache.org/licenses/LICENSE-2.0\u003c/a\u003e.\u003c/small\u003e\u003c/small\u003e\u003c/p\u003e\n",
- "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eUnless 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.\u003c/small\u003e\u003c/small\u003e\u003c/p\u003e"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "YvyGCsgSCxHQ"
- },
- "source": [
- "# Install MuJoCo"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "Xqo7pyX-n72M"
- },
- "outputs": [],
- "source": [
- "!pip install mujoco"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "IbZxYDxzoz5R"
- },
- "outputs": [],
- "source": [
- "#@title Set up rendering, check installation\n",
- "\n",
- "from google.colab import files\n",
- "\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",
- "try:\n",
- " print('Checking that the installation succeeded:')\n",
- " import mujoco\n",
- " mujoco.MjModel.from_xml_string('\u003cmujoco/\u003e')\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.')"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "T5f4w3Kq2X14"
- },
- "outputs": [],
- "source": [
- "#@title Import packages for plotting and creating graphics\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",
- "!command -v ffmpeg \u003e/dev/null || (apt update \u0026\u0026 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": "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",
- "\u003cmujoco\u003e\n",
- " \u003cworldbody\u003e\n",
- " \u003cgeom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/\u003e\n",
- " \u003cgeom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/\u003e\n",
- " \u003c/worldbody\u003e\n",
- "\u003c/mujoco\u003e\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 `\u003cmujoco\u003e`. The smallest valid MJCF model is `\u003cmujoco/\u003e` which is a completely empty model.\n",
- " - All physical elements live inside the `\u003cworldbody\u003e` 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",
- "\u003cmujoco\u003e\n",
- " \u003cworldbody\u003e\n",
- " \u003cgeom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/\u003e\n",
- " \u003cgeom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/\u003e\n",
- " \u003c/worldbody\u003e\n",
- "\u003c/mujoco\u003e\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",
- "renderer = mujoco.Renderer(model)\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": [
- "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",
- "\u003cmujoco\u003e\n",
- " \u003cworldbody\u003e\n",
- " \u003clight name=\"top\" pos=\"0 0 1\"/\u003e\n",
- " \u003cgeom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/\u003e\n",
- " \u003cgeom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/\u003e\n",
- " \u003c/worldbody\u003e\n",
- "\u003c/mujoco\u003e\n",
- "\"\"\"\n",
- "model = mujoco.MjModel.from_xml_string(xml)\n",
- "data = mujoco.MjData(model)\n",
- "renderer = mujoco.Renderer(model)\n",
- "\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",
- "renderer.update_scene(data)\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",
- "while data.time \u003c duration:\n",
- " mujoco.mj_step(model, data)\n",
- " if len(frames) \u003c data.time * framerate:\n",
- " renderer.update_scene(data)\n",
- " pixels = renderer.render()\n",
- " frames.append(pixels)\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\u0026q=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",
- "\u003cmujoco\u003e\n",
- " \u003cworldbody\u003e\n",
- " \u003clight name=\"top\" pos=\"0 0 1\"/\u003e\n",
- " \u003cbody name=\"box_and_sphere\" euler=\"0 0 -30\"\u003e\n",
- " \u003cjoint name=\"swing\" type=\"hinge\" axis=\"1 -1 0\" pos=\"-.2 -.2 -.2\"/\u003e\n",
- " \u003cgeom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/\u003e\n",
- " \u003cgeom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- " \u003c/worldbody\u003e\n",
- "\u003c/mujoco\u003e\n",
- "\"\"\"\n",
- "model = mujoco.MjModel.from_xml_string(xml)\n",
- "data = mujoco.MjData(model)\n",
- "renderer = mujoco.Renderer(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",
- "frames = []\n",
- "mujoco.mj_resetData(model, data)\n",
- "while data.time \u003c duration:\n",
- " mujoco.mj_step(model, data)\n",
- " if len(frames) \u003c data.time * framerate:\n",
- " renderer.update_scene(data, scene_option=scene_option)\n",
- " pixels = renderer.render()\n",
- " frames.append(pixels)\n",
- "\n",
- "# Simulate and display video.\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\u0026q=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",
- "frames = []\n",
- "mujoco.mj_resetData(model, data)\n",
- "while data.time \u003c duration:\n",
- " mujoco.mj_step(model, data)\n",
- " if len(frames) \u003c 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 `\u003coption\u003e` element:\n",
- "```xml\n",
- "\u003cmujoco\u003e\n",
- " \u003coption gravity=\"0 0 10\"/\u003e\n",
- " ...\n",
- "\u003c/mujoco\u003e\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\u0026q=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",
- "\u003cmujoco model=\"tippe top\"\u003e\n",
- " \u003coption integrator=\"RK4\"/\u003e\n",
- "\n",
- " \u003casset\u003e\n",
- " \u003ctexture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
- " rgb2=\".2 .3 .4\" width=\"300\" height=\"300\"/\u003e\n",
- " \u003cmaterial name=\"grid\" texture=\"grid\" texrepeat=\"8 8\" reflectance=\".2\"/\u003e\n",
- " \u003c/asset\u003e\n",
- "\n",
- " \u003cworldbody\u003e\n",
- " \u003cgeom size=\".2 .2 .01\" type=\"plane\" material=\"grid\"/\u003e\n",
- " \u003clight pos=\"0 0 .6\"/\u003e\n",
- " \u003ccamera name=\"closeup\" pos=\"0 -.1 .07\" xyaxes=\"1 0 0 0 1 2\"/\u003e\n",
- " \u003cbody name=\"top\" pos=\"0 0 .02\"\u003e\n",
- " \u003cfreejoint/\u003e\n",
- " \u003cgeom name=\"ball\" type=\"sphere\" size=\".02\" /\u003e\n",
- " \u003cgeom name=\"stem\" type=\"cylinder\" pos=\"0 0 .02\" size=\"0.004 .008\"/\u003e\n",
- " \u003cgeom name=\"ballast\" type=\"box\" size=\".023 .023 0.005\" pos=\"0 0 -.015\"\n",
- " contype=\"0\" conaffinity=\"0\" group=\"3\"/\u003e\n",
- " \u003c/body\u003e\n",
- " \u003c/worldbody\u003e\n",
- "\n",
- " \u003ckeyframe\u003e\n",
- " \u003ckey name=\"spinning\" qpos=\"0 0 0.02 1 0 0 0\" qvel=\"0 0 0 0 1 200\" /\u003e\n",
- " \u003c/keyframe\u003e\n",
- "\u003c/mujoco\u003e\n",
- "\"\"\"\n",
- "model = mujoco.MjModel.from_xml_string(tippe_top)\n",
- "renderer = mujoco.Renderer(model)\n",
- "data = mujoco.MjData(model)\n",
- "mujoco.mj_forward(model, data)\n",
- "renderer.update_scene(data, camera=\"closeup\")\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 `\u003cfreejoint/\u003e` clause.\n",
- "2. We use the `\u003coption/\u003e` 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 `\u003casset/\u003e` 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 `\u003ccamera\u003e` 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",
- "while data.time \u003c duration:\n",
- " mujoco.mj_step(model, data)\n",
- " if len(frames) \u003c 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 \u003c 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",
- "\u003cmujoco\u003e\n",
- " \u003coption timestep=\".001\"\u003e\n",
- " \u003cflag energy=\"enable\" contact=\"disable\"/\u003e\n",
- " \u003c/option\u003e\n",
- "\n",
- " \u003cdefault\u003e\n",
- " \u003cjoint type=\"hinge\" axis=\"0 -1 0\"/\u003e\n",
- " \u003cgeom type=\"capsule\" size=\".02\"/\u003e\n",
- " \u003c/default\u003e\n",
- "\n",
- " \u003cworldbody\u003e\n",
- " \u003clight pos=\"0 -.4 1\"/\u003e\n",
- " \u003ccamera name=\"fixed\" pos=\"0 -1 0\" xyaxes=\"1 0 0 0 0 1\"/\u003e\n",
- " \u003cbody name=\"0\" pos=\"0 0 .2\"\u003e\n",
- " \u003cjoint name=\"root\"/\u003e\n",
- " \u003cgeom fromto=\"-.2 0 0 .2 0 0\" rgba=\"1 1 0 1\"/\u003e\n",
- " \u003cgeom fromto=\"0 0 0 0 0 -.25\" rgba=\"1 1 0 1\"/\u003e\n",
- " \u003cbody name=\"1\" pos=\"-.2 0 0\"\u003e\n",
- " \u003cjoint/\u003e\n",
- " \u003cgeom fromto=\"0 0 0 0 0 -.2\" rgba=\"1 0 0 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- " \u003cbody name=\"2\" pos=\".2 0 0\"\u003e\n",
- " \u003cjoint/\u003e\n",
- " \u003cgeom fromto=\"0 0 0 0 0 -.2\" rgba=\"0 1 0 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- " \u003cbody name=\"3\" pos=\"0 0 -.25\"\u003e\n",
- " \u003cjoint/\u003e\n",
- " \u003cgeom fromto=\"0 0 0 0 0 -.2\" rgba=\"0 0 1 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- " \u003c/body\u003e\n",
- " \u003c/worldbody\u003e\n",
- "\u003c/mujoco\u003e\n",
- "\"\"\"\n",
- "model = mujoco.MjModel.from_xml_string(chaotic_pendulum)\n",
- "renderer = mujoco.Renderer(model, 480, 640)\n",
- "data = mujoco.MjData(model)\n",
- "mujoco.mj_forward(model, data)\n",
- "renderer.update_scene(data, camera=\"fixed\")\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",
- "renderer = mujoco.Renderer(model, 240, 320)\n",
- "\n",
- "\n",
- "# set initial state\n",
- "mujoco.mj_resetData(model, data)\n",
- "data.joint('root').qvel = 10\n",
- "\n",
- "\n",
- "# simulate and record frames\n",
- "frame = 0\n",
- "sim_time = 0\n",
- "render_time = 0\n",
- "n_steps = 0\n",
- "for i in range(n_frames):\n",
- " while data.time * framerate \u003c 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]\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",
- "\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()\n"
- ]
- },
- {
- "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",
- "\u003cmujoco\u003e\n",
- " \u003casset\u003e\n",
- " \u003ctexture 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\"/\u003e\n",
- " \u003cmaterial name=\"grid\" texture=\"grid\" texrepeat=\"2 2\" texuniform=\"true\"\n",
- " reflectance=\".2\"/\u003e\n",
- " \u003c/asset\u003e\n",
- "\n",
- " \u003cworldbody\u003e\n",
- " \u003clight pos=\"0 0 1\" mode=\"trackcom\"/\u003e\n",
- " \u003cgeom name=\"ground\" type=\"plane\" pos=\"0 0 -.5\" size=\"2 2 .1\" material=\"grid\" solimp=\".99 .99 .01\" solref=\".001 1\"/\u003e\n",
- " \u003cbody name=\"box_and_sphere\" pos=\"0 0 0\"\u003e\n",
- " \u003cfreejoint/\u003e\n",
- " \u003cgeom name=\"red_box\" type=\"box\" size=\".1 .1 .1\" rgba=\"1 0 0 1\" solimp=\".99 .99 .01\" solref=\".001 1\"/\u003e\n",
- " \u003cgeom name=\"green_sphere\" size=\".06\" pos=\".1 .1 .1\" rgba=\"0 1 0 1\"/\u003e\n",
- " \u003ccamera name=\"fixed\" pos=\"0 -.6 .3\" xyaxes=\"1 0 0 0 1 2\"/\u003e\n",
- " \u003ccamera name=\"track\" pos=\"0 -.6 .3\" xyaxes=\"1 0 0 0 1 2\" mode=\"track\"/\u003e\n",
- " \u003c/body\u003e\n",
- " \u003c/worldbody\u003e\n",
- "\u003c/mujoco\u003e\n",
- "\"\"\"\n",
- "model = mujoco.MjModel.from_xml_string(free_body_MJCF)\n",
- "renderer = mujoco.Renderer(model, 400, 600)\n",
- "data = mujoco.MjData(model)\n",
- "mujoco.mj_forward(model, data)\n",
- "renderer.update_scene(data, \"fixed\")\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",
- "renderer = mujoco.Renderer(model, height, width)\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 render\n",
- "for i in range(n_frames):\n",
- " while data.time \u003c 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",
- "# show video\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(iter(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",
- "\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",
- "\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",
- "\u003cmujoco\u003e\n",
- " \u003casset\u003e\n",
- " \u003ctexture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
- " rgb2=\".2 .3 .4\" width=\"300\" height=\"300\" mark=\"none\"/\u003e\n",
- " \u003cmaterial name=\"grid\" texture=\"grid\" texrepeat=\"6 6\"\n",
- " texuniform=\"true\" reflectance=\".2\"/\u003e\n",
- " \u003cmaterial name=\"wall\" rgba='.5 .5 .5 1'/\u003e\n",
- " \u003c/asset\u003e\n",
- "\n",
- " \u003cdefault\u003e\n",
- " \u003cgeom type=\"box\" size=\".05 .05 .05\" /\u003e\n",
- " \u003cjoint type=\"free\"/\u003e\n",
- " \u003c/default\u003e\n",
- "\n",
- " \u003cworldbody\u003e\n",
- " \u003clight name=\"light\" pos=\"-.2 0 1\"/\u003e\n",
- " \u003cgeom name=\"ground\" type=\"plane\" size=\".5 .5 10\" material=\"grid\"\n",
- " zaxis=\"-.3 0 1\" friction=\".1\"/\u003e\n",
- " \u003ccamera name=\"y\" pos=\"-.1 -.6 .3\" xyaxes=\"1 0 0 0 1 2\"/\u003e\n",
- " \u003cbody pos=\"0 0 .1\"\u003e\n",
- " \u003cjoint/\u003e\n",
- " \u003cgeom/\u003e\n",
- " \u003c/body\u003e\n",
- " \u003cbody pos=\"0 .2 .1\"\u003e\n",
- " \u003cjoint/\u003e\n",
- " \u003cgeom friction=\".33\"/\u003e\n",
- " \u003c/body\u003e\n",
- " \u003c/worldbody\u003e\n",
- "\n",
- "\u003c/mujoco\u003e\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",
- "renderer = mujoco.Renderer(model, height, width)\n",
- "\n",
- "# simulate and render\n",
- "mujoco.mj_resetData(model, data)\n",
- "for i in range(n_frames):\n",
- " while data.time \u003c i/30.0:\n",
- " mujoco.mj_step(model, data)\n",
- " renderer.update_scene(data, \"y\")\n",
- " frame = renderer.render()\n",
- " frames.append(frame)\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",
- "\u003cmujoco\u003e\n",
- " \u003casset\u003e\n",
- " \u003ctexture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
- " rgb2=\".2 .3 .4\" width=\"300\" height=\"300\" mark=\"none\"/\u003e\n",
- " \u003cmaterial name=\"grid\" texture=\"grid\" texrepeat=\"1 1\"\n",
- " texuniform=\"true\" reflectance=\".2\"/\u003e\n",
- " \u003c/asset\u003e\n",
- "\n",
- " \u003cworldbody\u003e\n",
- " \u003clight name=\"light\" pos=\"0 0 1\"/\u003e\n",
- " \u003cgeom name=\"floor\" type=\"plane\" pos=\"0 0 -.5\" size=\"2 2 .1\" material=\"grid\"/\u003e\n",
- " \u003csite name=\"anchor\" pos=\"0 0 .3\" size=\".01\"/\u003e\n",
- " \u003ccamera name=\"fixed\" pos=\"0 -1.3 .5\" xyaxes=\"1 0 0 0 1 2\"/\u003e\n",
- "\n",
- " \u003cgeom name=\"pole\" type=\"cylinder\" fromto=\".3 0 -.5 .3 0 -.1\" size=\".04\"/\u003e\n",
- " \u003cbody name=\"bat\" pos=\".3 0 -.1\"\u003e\n",
- " \u003cjoint name=\"swing\" type=\"hinge\" damping=\"1\" axis=\"0 0 1\"/\u003e\n",
- " \u003cgeom name=\"bat\" type=\"capsule\" fromto=\"0 0 .04 0 -.3 .04\"\n",
- " size=\".04\" rgba=\"0 0 1 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- "\n",
- " \u003cbody name=\"box_and_sphere\" pos=\"0 0 0\"\u003e\n",
- " \u003cjoint name=\"free\" type=\"free\"/\u003e\n",
- " \u003cgeom name=\"red_box\" type=\"box\" size=\".1 .1 .1\" rgba=\"1 0 0 1\"/\u003e\n",
- " \u003cgeom name=\"green_sphere\" size=\".06\" pos=\".1 .1 .1\" rgba=\"0 1 0 1\"/\u003e\n",
- " \u003csite name=\"hook\" pos=\"-.1 -.1 -.1\" size=\".01\"/\u003e\n",
- " \u003csite name=\"IMU\"/\u003e\n",
- " \u003c/body\u003e\n",
- " \u003c/worldbody\u003e\n",
- "\n",
- " \u003ctendon\u003e\n",
- " \u003cspatial name=\"wire\" limited=\"true\" range=\"0 0.35\" width=\"0.003\"\u003e\n",
- " \u003csite site=\"anchor\"/\u003e\n",
- " \u003csite site=\"hook\"/\u003e\n",
- " \u003c/spatial\u003e\n",
- " \u003c/tendon\u003e\n",
- "\n",
- " \u003cactuator\u003e\n",
- " \u003cmotor name=\"my_motor\" joint=\"swing\" gear=\"1\"/\u003e\n",
- " \u003c/actuator\u003e\n",
- "\n",
- " \u003csensor\u003e\n",
- " \u003caccelerometer name=\"accelerometer\" site=\"IMU\"/\u003e\n",
- " \u003c/sensor\u003e\n",
- "\u003c/mujoco\u003e\n",
- "\"\"\"\n",
- "model = mujoco.MjModel.from_xml_string(MJCF)\n",
- "renderer = mujoco.Renderer(model, 480, 480)\n",
- "data = mujoco.MjData(model)\n",
- "mujoco.mj_forward(model, data)\n",
- "renderer.update_scene(data, \"fixed\")\n",
- "media.show_image(renderer.render())\n"
- ]
- },
- {
- "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",
- "renderer = mujoco.Renderer(model, height, width)\n",
- "\n",
- "# constant actuator signal\n",
- "mujoco.mj_resetData(model, data)\n",
- "data.ctrl = 20\n",
- "\n",
- "# simulate and render\n",
- "for i in range(n_frames):\n",
- " while data.time \u003c 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='timestep = {:2.2g}ms'.format(1000*dt))\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",
- "\u003cmujoco\u003e\n",
- " \u003cworldbody\u003e\n",
- " \u003clight name=\"top\" pos=\"0 0 1\"/\u003e\n",
- " \u003cbody name=\"box_and_sphere\" euler=\"0 0 -30\"\u003e\n",
- " \u003cjoint name=\"swing\" type=\"hinge\" axis=\"1 -1 0\" pos=\"-.2 -.2 -.2\"/\u003e\n",
- " \u003cgeom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/\u003e\n",
- " \u003cgeom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- " \u003c/worldbody\u003e\n",
- "\u003c/mujoco\u003e\n",
- "\"\"\"\n",
- "model = mujoco.MjModel.from_xml_string(xml)\n",
- "renderer = mujoco.Renderer(model)\n",
- "data = mujoco.MjData(model)\n",
- "\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\n",
- "\n",
- "scene_option.frame = mujoco.mjtFrame.mjFRAME_GEOM\n",
- "scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = True\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\n",
- "\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 \u003c= 1].mean()\n",
- "# Scale to [0, 255]\n",
- "pixels = 255*np.clip(depth, 0, 1)\n",
- "\n",
- "media.show_image(pixels.astype(np.uint8))\n",
- "\n",
- "renderer.disable_depth_rendering()"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "PNwiIrgpx7T8"
- },
- "outputs": [],
- "source": [
- "#@title Segmentation 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))\n",
- "\n",
- "renderer.disable_segmentation_rendering()"
- ]
- },
- {
- "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": "markdown",
- "metadata": {
- "id": "Bs89vS0wLoU0"
- },
- "source": [
- "Let's use the camera matrix to project from world to camera coordinates:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "My0N4_7PDJ_q"
- },
- "outputs": [],
- "source": [
- "# reset the scene\n",
- "renderer.update_scene(data)\n",
- "\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 \u003e= scene.maxgeom:\n",
- " return\n",
- " scene.ngeom += 1 # increment ngeom\n",
- " # initialise a new capsule, add it to the scene using mjv_makeConnector\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_makeConnector(scene.geoms[scene.ngeom-1],\n",
- " mujoco.mjtGeom.mjGEOM_CAPSULE, radius,\n",
- " point1[0], point1[1], point1[2],\n",
- " point2[0], point2[1], point2[2])\n",
- "\n",
- " # traces of time, position and speed\n",
- "times = []\n",
- "positions = []\n",
- "speeds = []\n",
- "offset = model.jnt_axis[0]/8 # offset along the joint axis\n",
- "\n",
- "def modify_scene(scn):\n",
- " \"\"\"Draw position trace, speed modifies width and colors.\"\"\"\n",
- " if len(positions) \u003e 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",
- "while data.time \u003c 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) \u003c data.time * framerate:\n",
- " renderer.update_scene(data)\n",
- " modify_scene(renderer.scene)\n",
- " pixels = renderer.render()\n",
- " frames.append(pixels)\n",
- "media.show_video(frames, fps=framerate)"
- ]
- },
- {
- "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",
- "dominos_xml = \"\"\"\n",
- "\u003cmujoco\u003e\n",
- " \u003casset\u003e\n",
- " \u003ctexture type=\"skybox\" builtin=\"gradient\" rgb1=\".3 .5 .7\" rgb2=\"0 0 0\" width=\"32\" height=\"512\"/\u003e\n",
- " \u003ctexture name=\"grid\" type=\"2d\" builtin=\"checker\" width=\"512\" height=\"512\" rgb1=\".1 .2 .3\" rgb2=\".2 .3 .4\"/\u003e\n",
- " \u003cmaterial name=\"grid\" texture=\"grid\" texrepeat=\"2 2\" texuniform=\"true\" reflectance=\".2\"/\u003e\n",
- " \u003c/asset\u003e\n",
- "\n",
- " \u003cstatistic meansize=\".01\"/\u003e\n",
- "\n",
- " \u003cvisual\u003e\n",
- " \u003cglobal offheight=\"2160\" offwidth=\"3840\"/\u003e\n",
- " \u003cquality offsamples=\"8\"/\u003e\n",
- " \u003c/visual\u003e\n",
- "\n",
- " \u003cdefault\u003e\n",
- " \u003cgeom type=\"box\" solref=\".005 1\"/\u003e\n",
- " \u003cdefault class=\"static\"\u003e\n",
- " \u003cgeom rgba=\".3 .5 .7 1\"/\u003e\n",
- " \u003c/default\u003e\n",
- " \u003c/default\u003e\n",
- "\n",
- " \u003coption timestep=\"5e-4\"/\u003e\n",
- "\n",
- " \u003cworldbody\u003e\n",
- " \u003clight pos=\".3 -.3 .8\" mode=\"trackcom\" diffuse=\"1 1 1\" specular=\".3 .3 .3\"/\u003e\n",
- " \u003clight pos=\"0 -.3 .4\" mode=\"targetbodycom\" target=\"box\" diffuse=\".8 .8 .8\" specular=\".3 .3 .3\"/\u003e\n",
- " \u003cgeom name=\"floor\" type=\"plane\" size=\"3 3 .01\" pos=\"-0.025 -0.295 0\" material=\"grid\"/\u003e\n",
- " \u003cgeom name=\"ramp\" pos=\".25 -.45 -.03\" size=\".04 .1 .07\" euler=\"-30 0 0\" class=\"static\"/\u003e\n",
- " \u003ccamera name=\"top\" pos=\"-0.37 -0.78 0.49\" xyaxes=\"0.78 -0.63 0 0.27 0.33 0.9\"/\u003e\n",
- "\n",
- " \u003cbody name=\"ball\" pos=\".25 -.45 .1\"\u003e\n",
- " \u003cfreejoint name=\"ball\"/\u003e\n",
- " \u003cgeom name=\"ball\" type=\"sphere\" size=\".02\" rgba=\".65 .81 .55 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- "\n",
- " \u003cbody pos=\".26 -.3 .03\" euler=\"0 0 -90.0\"\u003e\n",
- " \u003cfreejoint/\u003e\n",
- " \u003cgeom size=\".0015 .015 .03\" rgba=\"1 .5 .5 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- "\n",
- " \u003cbody pos=\".26 -.27 .04\" euler=\"0 0 -81.0\"\u003e\n",
- " \u003cfreejoint/\u003e\n",
- " \u003cgeom size=\".002 .02 .04\" rgba=\"1 1 .5 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- "\n",
- " \u003cbody pos=\".24 -.21 .06\" euler=\"0 0 -63.0\"\u003e\n",
- " \u003cfreejoint/\u003e\n",
- " \u003cgeom size=\".003 .03 .06\" rgba=\".5 1 .5 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- "\n",
- " \u003cbody pos=\".2 -.16 .08\" euler=\"0 0 -45.0\"\u003e\n",
- " \u003cfreejoint/\u003e\n",
- " \u003cgeom size=\".004 .04 .08\" rgba=\".5 1 1 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- "\n",
- " \u003cbody pos=\".15 -.12 .1\" euler=\"0 0 -27.0\"\u003e\n",
- " \u003cfreejoint/\u003e\n",
- " \u003cgeom size=\".005 .05 .1\" rgba=\".5 .5 1 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- "\n",
- " \u003cbody pos=\".09 -.1 .12\" euler=\"0 0 -9.0\"\u003e\n",
- " \u003cfreejoint/\u003e\n",
- " \u003cgeom size=\".006 .06 .12\" rgba=\"1 .5 1 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- "\n",
- " \u003cbody name=\"seasaw_wrapper\" pos=\"-.23 -.1 0\" euler=\"0 0 30\"\u003e\n",
- " \u003cgeom size=\".01 .01 .015\" pos=\"0 .05 .015\" class=\"static\"/\u003e\n",
- " \u003cgeom size=\".01 .01 .015\" pos=\"0 -.05 .015\" class=\"static\"/\u003e\n",
- " \u003cgeom type=\"cylinder\" size=\".01 .0175\" pos=\"-.09 0 .0175\" class=\"static\"/\u003e\n",
- " \u003cbody name=\"seasaw\" pos=\"0 0 .03\"\u003e\n",
- " \u003cjoint axis=\"0 1 0\"/\u003e\n",
- " \u003cgeom type=\"cylinder\" size=\".005 .039\" zaxis=\"0 1 0\" rgba=\".84 .15 .33 1\"/\u003e\n",
- " \u003cgeom size=\".1 .02 .005\" pos=\"0 0 .01\" rgba=\".84 .15 .33 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- " \u003c/body\u003e\n",
- "\n",
- " \u003cbody name=\"box\" pos=\"-.3 -.14 .05501\" euler=\"0 0 -30\"\u003e\n",
- " \u003cfreejoint name=\"box\"/\u003e\n",
- " \u003cgeom name=\"box\" size=\".01 .01 .01\" rgba=\".0 .7 .79 1\"/\u003e\n",
- " \u003c/body\u003e\n",
- " \u003c/worldbody\u003e\n",
- "\u003c/mujoco\u003e\n",
- "\"\"\"\n",
- "model = mujoco.MjModel.from_xml_string(dominos_xml)\n",
- "data = mujoco.MjData(model)\n",
- "renderer = mujoco.Renderer(model, height=1024, width=1440)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "a2WruafiPhPk"
- },
- "outputs": [],
- "source": [
- "#@title Render from fixed camera\n",
- "duration = 2.5 # (seconds)\n",
- "framerate = 60 # (Hz)\n",
- "\n",
- "# Simulate and display video.\n",
- "frames = []\n",
- "mujoco.mj_resetData(model, data) # Reset state and time.\n",
- "while data.time \u003c duration:\n",
- " mujoco.mj_step(model, data)\n",
- " if len(frames) \u003c data.time * framerate:\n",
- " renderer.update_scene(data, camera='top')\n",
- " pixels = renderer.render()\n",
- " frames.append(pixels)\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",
- "duration = 3 # (seconds)\n",
- "\n",
- "# find time when box is thrown (speed \u003e 2cm/s)\n",
- "throw_time = 0.0\n",
- "mujoco.mj_resetData(model, data)\n",
- "while data.time \u003c 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 \u003e 0.02:\n",
- " throw_time = data.time\n",
- "assert throw_time \u003e 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",
- "while data.time \u003c duration:\n",
- " mujoco.mj_step(model, data)\n",
- " if len(frames) \u003c 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",
- "media.show_video(frames, fps=framerate)"
- ]
- }
- ],
- "metadata": {
- "accelerator": "GPU",
- "colab": {
- "collapsed_sections": [
- "-re3Szx-1Ias"
- ],
- "private_outputs": true,
- "provenance": [],
- "toc_visible": true
- },
- "gpuClass": "premium",
- "kernelspec": {
- "display_name": "Python 3",
- "name": "python3"
- }
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "MpkYHwCqk7W-"
+ },
+ "source": [
+ "\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",
+ "**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"
+ ]
},
- "nbformat": 4,
- "nbformat_minor": 0
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "xBSdkbmGN2K-"
+ },
+ "source": [
+ "### Copyright notice"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "_UbO9uhtBSX5"
+ },
+ "source": [
+ ">
Copyright 2022 DeepMind Technologies Limited.
\n",
+ ">
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 http://www.apache.org/licenses/LICENSE-2.0.
\n",
+ ">
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.
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "YvyGCsgSCxHQ"
+ },
+ "source": [
+ "# Install MuJoCo"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "Xqo7pyX-n72M"
+ },
+ "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",
+ "\n",
+ "from google.colab import files\n",
+ "\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",
+ "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.')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "cellView": "form",
+ "id": "T5f4w3Kq2X14"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Import packages for plotting and creating graphics\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",
+ "!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": "markdown",
+ "metadata": {
+ "id": "t0CF6Gvkt_Cw"
+ },
+ "source": [
+ "# MuJoCo basics\n",
+ "\n",
+ "We begin by defining and loading a simple model:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "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": 0,
+ "metadata": {
+ "id": "F40Pe6DY3Q0g"
+ },
+ "outputs": [],
+ "source": [
+ "model.ngeom"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "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": 0,
+ "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": 0,
+ "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": 0,
+ "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": 0,
+ "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": 0,
+ "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": 0,
+ "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": 0,
+ "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": 0,
+ "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": 0,
+ "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": 0,
+ "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",
+ "renderer = mujoco.Renderer(model)\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": 0,
+ "metadata": {
+ "id": "pvh47r97huS4"
+ },
+ "outputs": [],
+ "source": [
+ "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": 0,
+ "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",
+ "renderer = mujoco.Renderer(model)\n",
+ "\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": 0,
+ "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",
+ "renderer.update_scene(data)\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": 0,
+ "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",
+ "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",
+ "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": 0,
+ "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",
+ "renderer = mujoco.Renderer(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",
+ "frames = []\n",
+ "mujoco.mj_resetData(model, data)\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",
+ "# Simulate and display video.\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": 0,
+ "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": 0,
+ "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",
+ "frames = []\n",
+ "mujoco.mj_resetData(model, data)\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 `` 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 `` 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 `` 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": 0,
+ "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": 0,
+ "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",
+ "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": 0,
+ "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": 0,
+ "metadata": {
+ "id": "3jHYTV-bwfrS"
+ },
+ "outputs": [],
+ "source": [
+ "chaotic_pendulum = \"\"\"\n",
+ "\n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ "\"\"\"\n",
+ "model = mujoco.MjModel.from_xml_string(chaotic_pendulum)\n",
+ "renderer = mujoco.Renderer(model, 480, 640)\n",
+ "data = mujoco.MjData(model)\n",
+ "mujoco.mj_forward(model, data)\n",
+ "renderer.update_scene(data, camera=\"fixed\")\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": 0,
+ "metadata": {
+ "id": "-kNWvE9dNwYW"
+ },
+ "outputs": [],
+ "source": [
+ "# setup\n",
+ "n_seconds = 6\n",
+ "framerate = 30 # Hz\n",
+ "n_frames = int(n_seconds * framerate)\n",
+ "frames = []\n",
+ "renderer = mujoco.Renderer(model, 240, 320)\n",
+ "\n",
+ "\n",
+ "# set initial state\n",
+ "mujoco.mj_resetData(model, data)\n",
+ "data.joint('root').qvel = 10\n",
+ "\n",
+ "\n",
+ "# simulate and record frames\n",
+ "frame = 0\n",
+ "sim_time = 0\n",
+ "render_time = 0\n",
+ "n_steps = 0\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": 0,
+ "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": 0,
+ "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": 0,
+ "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]\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",
+ "\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()\n"
+ ]
+ },
+ {
+ "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": 0,
+ "metadata": {
+ "id": "2n1VNVv_FkbB"
+ },
+ "outputs": [],
+ "source": [
+ "free_body_MJCF = \"\"\"\n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ "\"\"\"\n",
+ "model = mujoco.MjModel.from_xml_string(free_body_MJCF)\n",
+ "renderer = mujoco.Renderer(model, 400, 600)\n",
+ "data = mujoco.MjData(model)\n",
+ "mujoco.mj_forward(model, data)\n",
+ "renderer.update_scene(data, \"fixed\")\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": 0,
+ "metadata": {
+ "id": "HlRhFs_d3WLP"
+ },
+ "outputs": [],
+ "source": [
+ "n_frames = 200\n",
+ "height = 240\n",
+ "width = 320\n",
+ "frames = []\n",
+ "renderer = mujoco.Renderer(model, height, width)\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 render\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",
+ "# show video\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": 0,
+ "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(iter(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",
+ "\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",
+ "\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": 0,
+ "metadata": {
+ "id": "2R_gKoYyXwda"
+ },
+ "outputs": [],
+ "source": [
+ "MJCF = \"\"\"\n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ "\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",
+ "renderer = mujoco.Renderer(model, height, width)\n",
+ "\n",
+ "# simulate and render\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",
+ "media.show_video(frames, fps=30)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "ArmmaPqGP6W7"
+ },
+ "source": [
+ "# Tendons, actuators and sensors"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "VJz84c97c8Df"
+ },
+ "outputs": [],
+ "source": [
+ "MJCF = \"\"\"\n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ "\"\"\"\n",
+ "model = mujoco.MjModel.from_xml_string(MJCF)\n",
+ "renderer = mujoco.Renderer(model, 480, 480)\n",
+ "data = mujoco.MjData(model)\n",
+ "mujoco.mj_forward(model, data)\n",
+ "renderer.update_scene(data, \"fixed\")\n",
+ "media.show_image(renderer.render())\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "u8z2vrOr_RVD"
+ },
+ "source": [
+ "actuated bat and passive \"piñata\":"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "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",
+ "renderer = mujoco.Renderer(model, height, width)\n",
+ "\n",
+ "# constant actuator signal\n",
+ "mujoco.mj_resetData(model, data)\n",
+ "data.ctrl = 20\n",
+ "\n",
+ "# simulate and render\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": 0,
+ "metadata": {
+ "id": "uy4wSEMAAJgn"
+ },
+ "outputs": [],
+ "source": [
+ "ax = plt.gca()\n",
+ "\n",
+ "ax.plot(np.asarray(times), np.asarray(sensordata), label='timestep = {:2.2g}ms'.format(1000*dt))\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": 0,
+ "metadata": {
+ "id": "mTDgsk2xcgwH"
+ },
+ "outputs": [],
+ "source": [
+ "xml = \"\"\"\n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ "\"\"\"\n",
+ "model = mujoco.MjModel.from_xml_string(xml)\n",
+ "renderer = mujoco.Renderer(model)\n",
+ "data = mujoco.MjData(model)\n",
+ "\n",
+ "mujoco.mj_forward(model, data)\n",
+ "renderer.update_scene(data)\n",
+ "media.show_image(renderer.render())"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "VePXamL_6XUc"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Enable transparency and frame visualization\n",
+ "\n",
+ "scene_option.frame = mujoco.mjtFrame.mjFRAME_GEOM\n",
+ "scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = True\n",
+ "renderer.update_scene(data, scene_option=scene_option)\n",
+ "frame = renderer.render()\n",
+ "media.show_image(frame)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "PVcpcvww9lZ8"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Depth rendering\n",
+ "\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))\n",
+ "\n",
+ "renderer.disable_depth_rendering()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "PNwiIrgpx7T8"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Segmentation 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))\n",
+ "\n",
+ "renderer.disable_segmentation_rendering()"
+ ]
+ },
+ {
+ "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": 0,
+ "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": "markdown",
+ "metadata": {
+ "id": "Bs89vS0wLoU0"
+ },
+ "source": [
+ "Let's use the camera matrix to project from world to camera coordinates:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "id": "My0N4_7PDJ_q"
+ },
+ "outputs": [],
+ "source": [
+ "# reset the scene\n",
+ "renderer.update_scene(data)\n",
+ "\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": 0,
+ "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_makeConnector\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_makeConnector(scene.geoms[scene.ngeom-1],\n",
+ " mujoco.mjtGeom.mjGEOM_CAPSULE, radius,\n",
+ " point1[0], point1[1], point1[2],\n",
+ " point2[0], point2[1], point2[2])\n",
+ "\n",
+ " # traces of time, position and speed\n",
+ "times = []\n",
+ "positions = []\n",
+ "speeds = []\n",
+ "offset = model.jnt_axis[0]/8 # 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",
+ "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",
+ "media.show_video(frames, fps=framerate)"
+ ]
+ },
+ {
+ "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": 0,
+ "metadata": {
+ "cellView": "form",
+ "id": "-SW-K9WuPGrp"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Load the \"dominos\" model\n",
+ "dominos_xml = \"\"\"\n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ "\"\"\"\n",
+ "model = mujoco.MjModel.from_xml_string(dominos_xml)\n",
+ "data = mujoco.MjData(model)\n",
+ "renderer = mujoco.Renderer(model, height=1024, width=1440)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "cellView": "form",
+ "id": "a2WruafiPhPk"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Render from fixed camera\n",
+ "duration = 2.5 # (seconds)\n",
+ "framerate = 60 # (Hz)\n",
+ "\n",
+ "# Simulate and display video.\n",
+ "frames = []\n",
+ "mujoco.mj_resetData(model, data) # Reset state and time.\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",
+ "media.show_video(frames, fps=framerate)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 0,
+ "metadata": {
+ "cellView": "form",
+ "id": "Kie3y-27bQ3J"
+ },
+ "outputs": [],
+ "source": [
+ "#@title Render from moving camera\n",
+ "duration = 3 # (seconds)\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",
+ "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",
+ "media.show_video(frames, fps=framerate)"
+ ]
+ }
+ ],
+ "metadata": {
+ "accelerator": "GPU",
+ "colab": {
+ "collapsed_sections": [
+ "-re3Szx-1Ias"
+ ],
+ "private_outputs": true,
+ "toc_visible": true
+ },
+ "gpuClass": "premium",
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
}
![](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}