chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
41a873b016f1d3c253fb6c059d29af6c62dc9c23
Mujoco_WASM/python/rollout.ipynb
T
Levi Burner 41a873b016 Notebook showcasing the Python rollout module
2025-02-18 10:57:03 -05:00

2072 lines
94 KiB
Plaintext
Raw Blame History

{
"cells": [
{
"cell_type": "markdown",
"id": "6adc68e0-a943-44ab-9af5-4bc62cc19f34",
"metadata": {
"editable": true,
"id": "6adc68e0-a943-44ab-9af5-4bc62cc19f34",
"tags": []
},
"source": [
"![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\n",
"\n",
"# <h1><center>Rollout Tutorial <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/rollout.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/></a></center></h1>\n",
"\n",
"This notebook provides a tutorial for [**MuJoCo** physics](https://github.com/google-deepmind/mujoco#readme), using the native Python bindings.\n",
"\n",
"This notebook describes the `rollout` module included in the MuJoCo Python library. It performs simulation \"rollouts\" with an underlying C++ function. The rollouts can be multithreaded.\n",
"\n",
"Below, the usage of each argument is explained with examples. An example of using `rollout` with minimize is also given. Then `rollout` is benchmarked against pure python and MJX. Finally, some examples for advanced use cases are provided.\n",
"\n",
"Note the benchmarks were designed to run on a AMD 5800X3D and an RTX 4090. They do not run in a reasonable amount of time on a typical free colab runtime.\n",
"\n",
"<!-- Copyright 2025 DeepMind Technologies Limited\n",
"\n",
" Licensed under the Apache License, Version 2.0 (the \"License\");\n",
" you may not use this file except in compliance with the License.\n",
" You may obtain a copy of the License at\n",
"\n",
" http://www.apache.org/licenses/LICENSE-2.0\n",
"\n",
" Unless required by applicable law or agreed to in writing, software\n",
" distributed under the License is distributed on an \"AS IS\" BASIS,\n",
" WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
" See the License for the specific language governing permissions and\n",
" limitations under the License.\n",
"-->"
]
},
{
"cell_type": "markdown",
"id": "5d8a6604-0948-4a42-a48d-249c7f0c462b",
"metadata": {
"editable": true,
"id": "5d8a6604-0948-4a42-a48d-249c7f0c462b",
"tags": []
},
"source": [
"# All Imports"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0f9fbad1-59d0-40ac-b2b6-99f37313670f",
"metadata": {
"cellView": "form",
"colab": {
"base_uri": "https://localhost:8080/"
},
"editable": true,
"id": "0f9fbad1-59d0-40ac-b2b6-99f37313670f",
"outputId": "5f1cfc9d-a955-486f-8ca7-7f08bd40a837",
"tags": [
"hide-input"
]
},
"outputs": [],
"source": [
"#@title All imports\n",
"\n",
"!pip install mujoco\n",
"!pip install mujoco_mjx\n",
"!pip install brax\n",
"\n",
"# Set up GPU rendering.\n",
"#from google.colab import files\n",
"import distutils.util\n",
"import os\n",
"import subprocess\n",
"if subprocess.run('nvidia-smi').returncode:\n",
" raise RuntimeError(\n",
" 'Cannot communicate with GPU. '\n",
" 'Make sure you are using a GPU Colab runtime. '\n",
" 'Go to the Runtime menu and select Choose runtime type.')\n",
"\n",
"# Add an ICD config so that glvnd can pick up the Nvidia EGL driver.\n",
"# This is usually installed as part of an Nvidia driver package, but the Colab\n",
"# kernel doesn't install its driver via APT, and as a result the ICD is missing.\n",
"# (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)\n",
"NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json'\n",
"if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):\n",
" with open(NVIDIA_ICD_CONFIG_PATH, 'w') as f:\n",
" f.write(\"\"\"{\n",
" \"file_format_version\" : \"1.0.0\",\n",
" \"ICD\" : {\n",
" \"library_path\" : \"libEGL_nvidia.so.0\"\n",
" }\n",
"}\n",
"\"\"\")\n",
"\n",
"# Configure MuJoCo to use the EGL rendering backend (requires GPU)\n",
"print('Setting environment variable to use GPU rendering:')\n",
"%env MUJOCO_GL=egl\n",
"\n",
"# Check if installation was succesful.\n",
"try:\n",
" print('Checking that the installation succeeded:')\n",
" import mujoco\n",
" from mujoco import minimize\n",
" from mujoco import rollout\n",
" from mujoco import mjx\n",
" mujoco.MjModel.from_xml_string('<mujoco/>')\n",
"except Exception as e:\n",
" raise e from RuntimeError(\n",
" 'Something went wrong during installation. Check the shell output above '\n",
" 'for more information.\\n'\n",
" 'If using a hosted Colab runtime, make sure you enable GPU acceleration '\n",
" 'by going to the Runtime menu and selecting \"Choose runtime type\".')\n",
"\n",
"print('Installation successful.')\n",
"\n",
"# 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",
"# Other imports and helper functions\n",
"import copy\n",
"import time\n",
"from multiprocessing import cpu_count\n",
"import threading\n",
"import itertools\n",
"import numpy as np\n",
"import jax\n",
"import jax.numpy as jp\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\n",
"import matplotlib.pyplot as plt\n",
"\n",
"# More legible printing from numpy.\n",
"np.set_printoptions(precision=3, suppress=True, linewidth=100)\n",
"\n",
"from IPython.display import clear_output\n",
"clear_output()\n",
"\n",
"# Set the number of threads to the number of cpu's that the multiprocessing module reports\n",
"nthread = cpu_count()\n",
"\n",
"# Get MuJoCo's standard humanoid model.\n",
"print('Getting MuJoCo humanoid XML description from GitHub:')\n",
"!git clone https://github.com/google-deepmind/mujoco"
]
},
{
"cell_type": "markdown",
"id": "fc69d0f4",
"metadata": {
"id": "fc69d0f4"
},
"source": [
"# Helper Functions"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "082482c7",
"metadata": {
"cellView": "form",
"editable": true,
"id": "082482c7",
"tags": [
"hide-input"
]
},
"outputs": [],
"source": [
"#@title helper functions\n",
"\n",
"def get_state(model, data, nbatch=1):\n",
" state = np.zeros((mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS),))\n",
" mujoco.mj_getState(model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS)\n",
" np.tile(state, (nbatch, 1))\n",
" return state\n",
"\n",
"def benchmark(f, x_list=[None], ntiming=1):\n",
" x_times_list = []\n",
" for x in x_list:\n",
" times = [time.perf_counter()]\n",
" for i in range(ntiming):\n",
" f(x)\n",
" times.append(time.perf_counter())\n",
" x_times_list.append(np.mean(np.diff(times)))\n",
" return np.array(x_times_list)\n",
"\n",
"def render_many(model, data, state, framerate, camera=-1, shift_joint=None, ncols=10, spacing=(1., 1.), shape=(480, 640), transparent=True):\n",
" nbatch = state.shape[0]\n",
"\n",
" if not isinstance(model, mujoco.MjModel):\n",
" model = list(model)\n",
"\n",
" if isinstance(model, list) and len(model) == 1:\n",
" model = model * nbatch\n",
" elif isinstance(model, list):\n",
" assert len(model) == nbatch\n",
" else:\n",
" model = [model] * nbatch\n",
"\n",
" if shift_joint is not None:\n",
" data = copy.copy(data)\n",
"\n",
" # Visual options\n",
" vopt = mujoco.MjvOption()\n",
" vopt.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = transparent # Transparent.\n",
" pert = mujoco.MjvPerturb() # Empty MjvPerturb object\n",
" catmask = mujoco.mjtCatBit.mjCAT_DYNAMIC\n",
"\n",
" # Simulate and render.\n",
" frames = []\n",
" with mujoco.Renderer(model[0], *shape) as renderer:\n",
" for i in range(state.shape[1]):\n",
" if len(frames) < i * model[0].opt.timestep * framerate:\n",
" for j in range(state.shape[0]):\n",
" mujoco.mj_setState(model[j], data, state[j, i, :], mujoco.mjtState.mjSTATE_FULLPHYSICS)\n",
" mujoco.mj_forward(model[j], data)\n",
"\n",
" if shift_joint is not None:\n",
" grid_x = j % ncols\n",
" grid_y = j // ncols\n",
" #print(grid_x, grid_y)\n",
" data.joint(shift_joint).qpos[:3] = data.joint(shift_joint).qpos[:3] + (grid_x * spacing[0], grid_y * spacing[1], 0)\n",
" mujoco.mj_forward(model[j], data)\n",
"\n",
" # Add the first top to the scene\n",
" if j == 0:\n",
" renderer.update_scene(data, camera, scene_option=vopt)\n",
" else:\n",
" mujoco.mjv_addGeoms(model[j], data, vopt, pert, catmask, renderer.scene)\n",
" # Render and add the frame.\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
" return frames"
]
},
{
"cell_type": "markdown",
"id": "c0570c2c",
"metadata": {},
"source": [
"# Using `rollout`\n",
"\n",
"The `rollout.rollout` function in the `mujoco` Python library runs batches of simulations for a fixed number steps. It can run in single or multi-threaded modes. The speedup over pure Python is significant because `rollout` users can easily enable the usage of a lightweight threadpool.\n",
"\n",
"Below we load the \"tippe top\", \"humanoid\", and \"humanoid100\" models which will be used in the following usage examples and benchmarks.\n",
"\n",
"The tippe top is copied from the [tutorial notebook](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb). The humanoid and humanoid100 models are distributed with MuJoCo."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "849b93e5",
"metadata": {},
"outputs": [],
"source": [
"#@title Benchmarked models\n",
"tippe_top = \"\"\"\n",
"<mujoco model=\"tippe top\">\n",
" <option integrator=\"RK4\"/>\n",
"\n",
" <asset>\n",
" <texture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
" rgb2=\".2 .3 .4\" width=\"300\" height=\"300\"/>\n",
" <material name=\"grid\" texture=\"grid\" texrepeat=\"40 40\" reflectance=\".2\"/>\n",
" </asset>\n",
"\n",
" <worldbody>\n",
" <geom size=\"1 1 .01\" type=\"plane\" material=\"grid\"/>\n",
" <light pos=\"0 0 .6\"/>\n",
" <camera name=\"closeup\" pos=\"0 -.1 .07\" xyaxes=\"1 0 0 0 1 2\"/>\n",
" <camera name=\"distant\" pos=\"0 -.4 .4\" xyaxes=\"1 0 0 0 1 1\"/>\n",
" <body name=\"top\" pos=\"0 0 .02\">\n",
" <freejoint name=\"top\"/>\n",
" <site name=\"top\" pos=\"0 0 0\"/>\n",
" <geom name=\"ball\" type=\"sphere\" size=\".02\" />\n",
" <geom name=\"stem\" type=\"cylinder\" pos=\"0 0 .02\" size=\"0.004 .008\"/>\n",
" <geom name=\"ballast\" type=\"box\" size=\".023 .023 0.005\" pos=\"0 0 -.015\"\n",
" contype=\"0\" conaffinity=\"0\" group=\"3\"/>\n",
" </body>\n",
" </worldbody>\n",
"\n",
" <sensor>\n",
" <gyro name=\"gyro\" site=\"top\"/>\n",
" </sensor>\n",
"\n",
" <keyframe>\n",
" <key name=\"spinning\" qpos=\"0 0 0.02 1 0 0 0\" qvel=\"0 0 0 0 1 200\" />\n",
" </keyframe>\n",
"</mujoco>\n",
"\"\"\"\n",
"\n",
"# Create and initialize top model\n",
"top_model = mujoco.MjModel.from_xml_string(tippe_top)\n",
"def init_top(model):\n",
" data = mujoco.MjData(model)\n",
" mujoco.mj_resetDataKeyframe(model, data, 0) # Set to the state to a spinning upside down top\n",
" return data\n",
"top_data = init_top(top_model)\n",
"\n",
"# Create and initialize humanoid model\n",
"humanoid_xml_path = 'mujoco/model/humanoid/humanoid.xml'\n",
"humanoid_model = mujoco.MjModel.from_xml_path(humanoid_xml_path)\n",
"def init_humanoid(model):\n",
" data = mujoco.MjData(model)\n",
" data.qvel[2] = 4 # Make the humanoid jump\n",
" return data\n",
"humanoid_data = init_humanoid(humanoid_model)\n",
"\n",
"# Create and initialize humanoid100 model\n",
"humanoid100_xml_path = 'mujoco/model/humanoid/humanoid100.xml'\n",
"humanoid100_model = mujoco.MjModel.from_xml_path(humanoid100_xml_path)\n",
"def init_humanoid100(model):\n",
" data = mujoco.MjData(model)\n",
" return data\n",
"humanoid100_data = init_humanoid100(humanoid100_model)\n",
"\n",
"start = time.time()\n",
"top_nstep = int(6 / top_model.opt.timestep)\n",
"top_state, _ = rollout.rollout(top_model, top_data, initial_state=get_state(top_model, top_data), nstep=top_nstep)\n",
"\n",
"humanoid_nstep = int(3 / humanoid_model.opt.timestep)\n",
"humanoid_state, _ = rollout.rollout(humanoid_model, humanoid_data, initial_state=get_state(humanoid_model, humanoid_data), nstep=humanoid_nstep)\n",
"\n",
"humanoid100_nstep = int(3 / humanoid100_model.opt.timestep)\n",
"humanoid100_state, _ = rollout.rollout(humanoid100_model, humanoid100_data, initial_state=get_state(humanoid100_model, humanoid100_data), nstep=humanoid100_nstep)\n",
"end = time.time()\n",
"\n",
"start_render = time.time()\n",
"top_frames = render_many(top_model, top_data, top_state, framerate=60, shape=(240, 320), transparent=False)\n",
"humanoid_frames = render_many(humanoid_model, humanoid_data, humanoid_state, framerate=120, shape=(240, 320), transparent=False)\n",
"humanoid100_frames = render_many(humanoid100_model, humanoid100_data, humanoid100_state, framerate=120, shape=(240, 320), transparent=False)\n",
"\n",
"media.show_video(np.concatenate((top_frames, humanoid_frames, humanoid100_frames), axis=2), fps=60) # humanoid and humanoid100 are shown at half speed\n",
"end_render = time.time()\n",
"\n",
"print(f'Rollout took {end-start:.1f} seconds')\n",
"print(f'Rendering took {end_render-start_render:.1f} seconds')"
]
},
{
"cell_type": "markdown",
"id": "55d171f7-541b-4441-aa18-da86d6716410",
"metadata": {
"id": "55d171f7-541b-4441-aa18-da86d6716410"
},
"source": [
"## Detailed Usage\n",
"\n",
"It is helpful to read `rollout`'s docstring before beginning. The main takeaways are that `rollout` runs nbatch rollouts for nstep steps. The MjModel's can be different but should be the same up to parameter values. Passing multiple MjData enables multithreading, one thread per MjData.\n",
"\n",
"Next we give usage examples of the most common arguments. The more advanced arguments are discussed in the \"Advanced Usage\" section."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9cd2f94a-11df-4247-986c-5a56af69a1f5",
"metadata": {
"id": "9cd2f94a-11df-4247-986c-5a56af69a1f5",
"outputId": "e6bfc150-4e04-42f6-98af-58477c686a9b"
},
"outputs": [],
"source": [
"print(rollout.rollout.__doc__)"
]
},
{
"cell_type": "markdown",
"id": "b6f7a094-8352-4b07-99ee-5278e3036cd5",
"metadata": {
"id": "b6f7a094-8352-4b07-99ee-5278e3036cd5",
"tags": []
},
"source": [
"### Example: different initial states\n",
"`rollout` is designed to run nbatch rollouts in parallel for nstep steps. Lets simulate 100 tippe tops with different initial rotation speeds."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe",
"metadata": {
"id": "849af5f2-9de1-4cb9-bc3a-c9b7acf0e3fe",
"outputId": "2fd64fd0-8583-4a09-cda2-5ca828114e88"
},
"outputs": [],
"source": [
"nbatch = 100 # Simulate this many tops\n",
"\n",
"# Get nbatch initial states and scale the initial speed of the tippe top using the batch index\n",
"top_data = init_top(top_model)\n",
"initial_state = get_state(top_model, top_data)\n",
"initial_states = np.tile(initial_state, (nbatch, 1))\n",
"initial_states[:, -1] *= np.linspace(0.5, 1.5, num=nbatch)\n",
"\n",
"# Run the rollout\n",
"start = time.time()\n",
"state, sensordata = rollout.rollout(top_model, [copy.copy(top_data) for _ in range(nthread)], # Create one MjData per thread\n",
" initial_states, nstep=int(top_nstep*1.5))\n",
"end = time.time()\n",
"\n",
"# Use state to render all the tops at once\n",
"start_render = time.time()\n",
"framerate = 60\n",
"media.show_video(render_many(top_model, top_data, state, framerate), fps=framerate)\n",
"end_render = time.time()\n",
"\n",
"print(f'Rollout took {end-start:.1f} seconds')\n",
"print(f'Rendering took {end_render-start_render:.1f} seconds')"
]
},
{
"cell_type": "markdown",
"id": "aa2cf151-bf9a-4a23-b7fe-6a766979d93f",
"metadata": {
"id": "aa2cf151-bf9a-4a23-b7fe-6a766979d93f"
},
"source": [
"Our model has an angular velocity sensor the middle of the top. Let's plot the response using the `sensordata` array that rollout returns."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "957b8566-da31-410b-b385-e78241c5247a",
"metadata": {
"id": "957b8566-da31-410b-b385-e78241c5247a",
"outputId": "313aa855-1516-420e-a92e-90c4cfc71969"
},
"outputs": [],
"source": [
"plt.subplot(3,1,1)\n",
"for i in range(nbatch): plt.plot(sensordata[i, :, 0])\n",
"plt.subplot(3,1,2)\n",
"for i in range(nbatch): plt.plot(sensordata[i, :, 1])\n",
"plt.subplot(3,1,3)\n",
"for i in range(nbatch): plt.plot(sensordata[i, :, 2])"
]
},
{
"cell_type": "markdown",
"id": "58044bc1-f98c-4bbf-a703-40ba075552a0",
"metadata": {
"id": "58044bc1-f98c-4bbf-a703-40ba075552a0"
},
"source": [
"### Example: different models\n",
"100 gray tops is kind of boring. It would be better if they were colorful and different sizes!\n",
"\n",
"`rollout` supports using different models for each rollout, so long as they are of compatibile dimensions. Let's simulate 100 tippe tops with the same initial condition, but different sizes and colors.\n",
"\n",
"**Note:** Strictly speaking, the models must have the same number of states, controls, degrees of freedom, and sensor outputs. The most common use case is multiple models of the same thing up to parameter values."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2",
"metadata": {
"id": "7c39e79e-8942-4fea-b306-ea0cb3c826e2",
"outputId": "d240e9d0-0bbe-497f-bf39-98994a97dcf6"
},
"outputs": [],
"source": [
"# Make 100 tippe tops with different colors and sizes\n",
"nbatch = 100\n",
"spec = mujoco.MjSpec.from_string(tippe_top)\n",
"models = []\n",
"for i in range(nbatch):\n",
" for geom in spec.geoms:\n",
" if geom.name in ['ball', 'stem', 'ballast']:\n",
" geom.rgba[:3] = np.random.rand(3)\n",
" if geom.name == 'stem':\n",
" stem_geom = geom\n",
" if geom.name == 'ball':\n",
" ball_geom = geom\n",
"\n",
" # Save original geom size\n",
" stem_geom_size = np.copy(stem_geom.size)\n",
" ball_geom_size = np.copy(ball_geom.size)\n",
"\n",
" # Scale geoms and compile model\n",
" size_scale = 0.75*np.random.rand(1) + 0.5\n",
" stem_geom.size *= size_scale\n",
" ball_geom.size *= size_scale\n",
" models.append(spec.compile())\n",
"\n",
" # Restore original geom size\n",
" stem_geom.size = stem_geom_size\n",
" ball_geom.size = ball_geom_size\n",
"\n",
"# Reset the intial state\n",
"top_data = init_top(top_model)\n",
"\n",
"# Run the rollout\n",
"start = time.time()\n",
"state, sensordata = rollout.rollout(models, [copy.copy(top_data) for _ in range(nthread)], # Create one MjData per thread\n",
" get_state(top_model, top_data), nstep=int(1.5*top_nstep))\n",
"end = time.time()\n",
"\n",
"# Render video\n",
"start_render = time.time()\n",
"framerate = 60\n",
"cam = mujoco.MjvCamera()\n",
"mujoco.mjv_defaultCamera(cam)\n",
"cam.distance = 0.2\n",
"cam.azimuth = 135\n",
"cam.elevation = -25\n",
"cam.lookat = [0, 0, 0.07]\n",
"models[0].vis.global_.fovy = 60\n",
"frames = render_many(models, top_data, state, framerate, shift_joint='top', spacing=[-0.05, 0.05], camera=cam)\n",
"media.show_video(frames, fps=framerate)\n",
"end_render = time.time()\n",
"\n",
"print(f'Rollout took {end-start:.1f} seconds')\n",
"print(f'Rendering took {end_render-start_render:.1f} seconds')"
]
},
{
"cell_type": "markdown",
"id": "cf485c08-72be-4169-89b6-9d93df8ebbe3",
"metadata": {
"id": "cf485c08-72be-4169-89b6-9d93df8ebbe3"
},
"source": [
"Because the models are now different, the measurements of the gyro sensor are not consistent even though the initial state for each rollout was the same."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03",
"metadata": {
"id": "b8a5d3d4-24e7-41a1-b3bd-7b63c1812b03",
"outputId": "016af14d-caf0-4ac8-8ad6-1a184440a206"
},
"outputs": [],
"source": [
"plt.subplot(3,1,1)\n",
"for i in range(nbatch): plt.plot(sensordata[i, :, 0])\n",
"plt.subplot(3,1,2)\n",
"for i in range(nbatch): plt.plot(sensordata[i, :, 1])\n",
"plt.subplot(3,1,3)\n",
"for i in range(nbatch): plt.plot(sensordata[i, :, 2])"
]
},
{
"cell_type": "markdown",
"id": "3841a669-6cd1-427e-a629-20a10a6e3a34",
"metadata": {
"id": "3841a669-6cd1-427e-a629-20a10a6e3a34"
},
"source": [
"### Example: control inputs\n",
"Open loop controls can be passed to `rollout` via the `control` argument. If passed, `nstep` no longer needs to be specified as it can be inferred from the size of `control`.\n",
"\n",
"Below we simulate 100 of the flailing humanoids from the [tutorial notebook](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb). Each humanoid uses a different control signal."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2a184873-8d24-45da-b444-8d21f5dcd733",
"metadata": {
"id": "2a184873-8d24-45da-b444-8d21f5dcd733",
"outputId": "8663c8ae-ed11-4858-9d85-f00b24d2ccc7"
},
"outputs": [],
"source": [
"# Episode parameters.\n",
"duration = 3 # (seconds)\n",
"framerate = 120 # (Hz)\n",
"humanoid_data.qvel[2] = 4 # Initial vertical velocity (m/s)\n",
"ctrl_phase = 2 * np.pi * np.random.rand(humanoid_model.nu) # Control phase\n",
"ctrl_freq = 1 # Control frequency\n",
"\n",
"# Generate 100 different controls\n",
"nbatch = 100\n",
"nstep = int(duration / humanoid_model.opt.timestep)\n",
"times = np.linspace(0.0, duration, nstep)\n",
"times = np.arange(0.0, duration, humanoid_model.opt.timestep)\n",
"control = np.sin((2 * np.pi * times * ctrl_freq).reshape(nstep, 1) + ctrl_phase.reshape(1, humanoid_model.nu))\n",
"control = np.stack([control]*nbatch, axis=0)\n",
"control += np.random.normal(size=control.shape)\n",
"\n",
"# Initialize the model\n",
"humanoid_data = init_humanoid(humanoid_model)\n",
"\n",
"# Run the rollout\n",
"start = time.time()\n",
"state, _ = rollout.rollout(humanoid_model, [copy.copy(humanoid_data) for _ in range(nthread)],\n",
" get_state(humanoid_model, humanoid_data), control)\n",
"end = time.time()\n",
"\n",
"# Render the rollout\n",
"start_render = time.time()\n",
"framerate=120\n",
"cam = mujoco.MjvCamera()\n",
"mujoco.mjv_defaultCamera(cam)\n",
"cam.distance = 3.5\n",
"cam.azimuth = 132.5\n",
"cam.elevation = -45\n",
"cam.lookat = [0, 0, 3.0]\n",
"humanoid_model.vis.global_.fovy = 60\n",
"frames = render_many(humanoid_model, humanoid_data, state, framerate, shift_joint='root', spacing=[-1.0, 1.0], camera=cam)\n",
"media.show_video(frames, fps=framerate/2) # Show the video at half speed\n",
"end_render = time.time()\n",
"\n",
"print(f'Rollout took {end-start:.1f} seconds')\n",
"print(f'Rendering took {end_render-start_render:.1f} seconds')"
]
},
{
"cell_type": "markdown",
"id": "4d89e4fb-7711-4d23-8ff3-eb5030fa8bf7",
"metadata": {
"id": "4d89e4fb-7711-4d23-8ff3-eb5030fa8bf7"
},
"source": [
"`rollout`'s `control_spec` argument can be used to indicate `control` contains values for actuators, generalized forces, cartesian forces, mocap poses, and/or the activation/deactivation of equality constraints. Internally, this is managed through [mj_setState](https://mujoco.readthedocs.io/en/stable/APIreference/APIfunctions.html#mj-setstate) and `control_spec` corresponds to `mj_setState`'s `spec` argument.\n",
"\n",
"Let's try applying cartesian forces in addition to the control inputs. This will make the humanoids look like they are being dragged while waving their limbs."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5",
"metadata": {
"id": "4b02bb61-912d-47de-a956-aadfcd4c5cd5",
"outputId": "8848acdc-fb51-4f34-a06d-ce5162076e1c"
},
"outputs": [],
"source": [
"xfrc = np.zeros((control.shape[0], control.shape[1], mujoco.mj_stateSize(humanoid_model, mujoco.mjtState.mjSTATE_XFRC_APPLIED)))\n",
"head_id = humanoid_model.body('head').id\n",
"\n",
"# Apply a constant but different force to each model\n",
"xfrc[:, :, 3*head_id:3*head_id+2] = np.random.normal(scale=150.0, size=(control.shape[0], 1, 2))\n",
"\n",
"humanoid_data = init_humanoid(humanoid_model)\n",
"\n",
"control_xfrc = np.concatenate((control, xfrc), axis=2)\n",
"control_spec = mujoco.mjtState.mjSTATE_CTRL.value + mujoco.mjtState.mjSTATE_XFRC_APPLIED.value\n",
"\n",
"start = time.time()\n",
"state, _ = rollout.rollout(humanoid_model, [copy.copy(humanoid_data) for _ in range(nthread)],\n",
" get_state(humanoid_model, humanoid_data), control_xfrc, control_spec=control_spec)\n",
"end = time.time()\n",
"\n",
"start_render = time.time()\n",
"frames = render_many(humanoid_model, humanoid_data, state, framerate, shift_joint='root', spacing=[-1.0, 1.0], camera=cam)\n",
"media.show_video(frames, fps=framerate/2) # Show the video at half speed\n",
"end_render = time.time()\n",
"\n",
"print(f'Rollout took {end-start:.1f} seconds')\n",
"print(f'Rendering took {end_render-start_render:.1f} seconds')"
]
},
{
"cell_type": "markdown",
"id": "32de0cf5-4ac9-49bd-907b-f350c056dee4",
"metadata": {
"id": "32de0cf5-4ac9-49bd-907b-f350c056dee4",
"tags": []
},
"source": [
"# Application: `rollout` + `minimize.least_squares`"
]
},
{
"cell_type": "markdown",
"id": "1ba95d9f-d6d4-4076-9fdf-ba2b604fd4ff",
"metadata": {
"id": "1ba95d9f-d6d4-4076-9fdf-ba2b604fd4ff"
},
"source": [
"`rollout` can be easily used with MuJoCo's nonlinear least squares utility, `minimize.least_squares`. Because `minimize` uses finite-differencing to estimate jacobians, it benefits greatly from multi-threaded rollouts.\n",
"\n",
"As an example let's consider the \"reach\" sample from the [least squares notebook](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/least_squares.ipynb). The code is copied here with a small modification that allows multithreading.\n",
"\n",
"The goal is for the humanoid to reach a target with one of its hands. By default, the humanoid does a jump, but does not reach its hand to the target."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d0d33759-b049-4fc7-bfc3-abaebc7c17cb",
"metadata": {
"cellView": "form",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 281
},
"id": "d0d33759-b049-4fc7-bfc3-abaebc7c17cb",
"outputId": "196452d4-8237-43dd-d77d-0e1dbee6f7e5",
"tags": []
},
"outputs": [],
"source": [
"#@title Humanoid Reaching XML and code\n",
"xml = \"\"\"\n",
"<mujoco model=\"Humanoid\">\n",
" <option timestep=\"0.005\"/>\n",
"\n",
" <visual>\n",
" <map force=\"0.1\" zfar=\"30\"/>\n",
" <rgba haze=\"0.15 0.25 0.35 1\"/>\n",
" <global offwidth=\"2560\" offheight=\"1440\" elevation=\"-20\" azimuth=\"120\"/>\n",
" </visual>\n",
"\n",
" <statistic center=\"0 0 0.7\"/>\n",
"\n",
" <asset>\n",
" <texture type=\"skybox\" builtin=\"gradient\" rgb1=\".3 .5 .7\" rgb2=\"0 0 0\" width=\"32\" height=\"512\"/>\n",
" <texture name=\"body\" type=\"cube\" builtin=\"flat\" mark=\"cross\" width=\"128\" height=\"128\" rgb1=\"0.8 0.6 0.4\" rgb2=\"0.8 0.6 0.4\" markrgb=\"1 1 1\" random=\"0.01\"/>\n",
" <material name=\"body\" texture=\"body\" texuniform=\"true\" rgba=\"0.8 0.6 .4 1\"/>\n",
" <texture name=\"grid\" type=\"2d\" builtin=\"checker\" width=\"512\" height=\"512\" rgb1=\".1 .2 .3\" rgb2=\".2 .3 .4\"/>\n",
" <material name=\"grid\" texture=\"grid\" texrepeat=\"1 1\" texuniform=\"true\" reflectance=\".2\"/>\n",
" </asset>\n",
"\n",
" <default>\n",
" <position inheritrange=\"0.95\"/>\n",
" <default class=\"body\">\n",
"\n",
" <!-- geoms -->\n",
" <geom type=\"capsule\" condim=\"1\" friction=\".7\" solimp=\"0.9 .99 .003\" solref=\".015 1\" material=\"body\" group=\"1\"/>\n",
" <default class=\"thigh\">\n",
" <geom size=\".06\"/>\n",
" </default>\n",
" <default class=\"shin\">\n",
" <geom fromto=\"0 0 0 0 0 -.3\" size=\".049\"/>\n",
" </default>\n",
" <default class=\"foot\">\n",
" <geom size=\".027\"/>\n",
" <default class=\"foot1\">\n",
" <geom fromto=\"-.07 -.01 0 .14 -.03 0\"/>\n",
" </default>\n",
" <default class=\"foot2\">\n",
" <geom fromto=\"-.07 .01 0 .14 .03 0\"/>\n",
" </default>\n",
" </default>\n",
" <default class=\"arm_upper\">\n",
" <geom size=\".04\"/>\n",
" </default>\n",
" <default class=\"arm_lower\">\n",
" <geom size=\".031\"/>\n",
" </default>\n",
" <default class=\"hand\">\n",
" <geom type=\"sphere\" size=\".04\"/>\n",
" </default>\n",
"\n",
" <!-- joints -->\n",
" <joint type=\"hinge\" damping=\".2\" stiffness=\"1\" armature=\".01\" limited=\"true\" solimplimit=\"0 .99 .01\"/>\n",
" <default class=\"joint_big\">\n",
" <joint damping=\"5\" stiffness=\"10\"/>\n",
" <default class=\"hip_x\">\n",
" <joint range=\"-30 10\"/>\n",
" </default>\n",
" <default class=\"hip_z\">\n",
" <joint range=\"-60 35\"/>\n",
" </default>\n",
" <default class=\"hip_y\">\n",
" <joint axis=\"0 1 0\" range=\"-150 20\"/>\n",
" </default>\n",
" <default class=\"joint_big_stiff\">\n",
" <joint stiffness=\"20\"/>\n",
" </default>\n",
" </default>\n",
" <default class=\"knee\">\n",
" <joint pos=\"0 0 .02\" axis=\"0 -1 0\" range=\"-160 2\"/>\n",
" </default>\n",
" <default class=\"ankle\">\n",
" <joint range=\"-50 50\"/>\n",
" <default class=\"ankle_y\">\n",
" <joint pos=\"0 0 .08\" axis=\"0 1 0\" stiffness=\"6\"/>\n",
" </default>\n",
" <default class=\"ankle_x\">\n",
" <joint pos=\"0 0 .04\" stiffness=\"3\"/>\n",
" </default>\n",
" </default>\n",
" <default class=\"shoulder\">\n",
" <joint range=\"-85 60\"/>\n",
" </default>\n",
" <default class=\"elbow\">\n",
" <joint range=\"-100 50\" stiffness=\"0\"/>\n",
" </default>\n",
" </default>\n",
" </default>\n",
"\n",
" <worldbody>\n",
" <body name=\"target\" pos=\".2 -.2 1\" mocap=\"true\">\n",
" <site name=\"target\" size=\".05\" rgba=\"1 0 1 .4\"/>\n",
" </body>\n",
" <geom name=\"floor\" size=\"0 0 .05\" type=\"plane\" material=\"grid\" condim=\"3\"/>\n",
" <light name=\"spotlight\" mode=\"targetbodycom\" target=\"torso\" diffuse=\".8 .8 .8\" specular=\"0.3 0.3 0.3\" pos=\"0 -6 4\" cutoff=\"30\"/>\n",
" <body name=\"torso\" pos=\"0 0 1.282\" childclass=\"body\">\n",
" <light name=\"top\" pos=\"0 0 2\" mode=\"trackcom\"/>\n",
" <camera name=\"back\" pos=\"-3 0 1\" xyaxes=\"0 -1 0 1 0 2\" mode=\"trackcom\"/>\n",
" <camera name=\"side\" pos=\"0 -3 1\" xyaxes=\"1 0 0 0 1 2\" mode=\"trackcom\"/>\n",
" <freejoint name=\"root\"/>\n",
" <geom name=\"torso\" fromto=\"0 -.07 0 0 .07 0\" size=\".07\"/>\n",
" <geom name=\"waist_upper\" fromto=\"-.01 -.06 -.12 -.01 .06 -.12\" size=\".06\"/>\n",
" <body name=\"head\" pos=\"0 0 .19\">\n",
" <geom name=\"head\" type=\"sphere\" size=\".09\"/>\n",
" <camera name=\"egocentric\" pos=\".09 0 0\" xyaxes=\"0 -1 0 .1 0 1\" fovy=\"80\"/>\n",
" </body>\n",
" <body name=\"waist_lower\" pos=\"-.01 0 -.26\">\n",
" <geom name=\"waist_lower\" fromto=\"0 -.06 0 0 .06 0\" size=\".06\"/>\n",
" <joint name=\"abdomen_z\" pos=\"0 0 .065\" axis=\"0 0 1\" range=\"-45 45\" class=\"joint_big_stiff\"/>\n",
" <joint name=\"abdomen_y\" pos=\"0 0 .065\" axis=\"0 1 0\" range=\"-75 30\" class=\"joint_big\"/>\n",
" <body name=\"pelvis\" pos=\"0 0 -.165\">\n",
" <joint name=\"abdomen_x\" pos=\"0 0 .1\" axis=\"1 0 0\" range=\"-35 35\" class=\"joint_big\"/>\n",
" <geom name=\"butt\" fromto=\"-.02 -.07 0 -.02 .07 0\" size=\".09\"/>\n",
" <body name=\"thigh_right\" pos=\"0 -.1 -.04\">\n",
" <joint name=\"hip_x_right\" axis=\"1 0 0\" class=\"hip_x\"/>\n",
" <joint name=\"hip_z_right\" axis=\"0 0 1\" class=\"hip_z\"/>\n",
" <joint name=\"hip_y_right\" class=\"hip_y\"/>\n",
" <geom name=\"thigh_right\" fromto=\"0 0 0 0 .01 -.34\" class=\"thigh\"/>\n",
" <body name=\"shin_right\" pos=\"0 .01 -.4\">\n",
" <joint name=\"knee_right\" class=\"knee\"/>\n",
" <geom name=\"shin_right\" class=\"shin\"/>\n",
" <body name=\"foot_right\" pos=\"0 0 -.39\">\n",
" <joint name=\"ankle_y_right\" class=\"ankle_y\"/>\n",
" <joint name=\"ankle_x_right\" class=\"ankle_x\" axis=\"1 0 .5\"/>\n",
" <geom name=\"foot1_right\" class=\"foot1\"/>\n",
" <geom name=\"foot2_right\" class=\"foot2\"/>\n",
" </body>\n",
" </body>\n",
" </body>\n",
" <body name=\"thigh_left\" pos=\"0 .1 -.04\">\n",
" <joint name=\"hip_x_left\" axis=\"-1 0 0\" class=\"hip_x\"/>\n",
" <joint name=\"hip_z_left\" axis=\"0 0 -1\" class=\"hip_z\"/>\n",
" <joint name=\"hip_y_left\" class=\"hip_y\"/>\n",
" <geom name=\"thigh_left\" fromto=\"0 0 0 0 -.01 -.34\" class=\"thigh\"/>\n",
" <body name=\"shin_left\" pos=\"0 -.01 -.4\">\n",
" <joint name=\"knee_left\" class=\"knee\"/>\n",
" <geom name=\"shin_left\" fromto=\"0 0 0 0 0 -.3\" class=\"shin\"/>\n",
" <body name=\"foot_left\" pos=\"0 0 -.39\">\n",
" <joint name=\"ankle_y_left\" class=\"ankle_y\"/>\n",
" <joint name=\"ankle_x_left\" class=\"ankle_x\" axis=\"-1 0 -.5\"/>\n",
" <geom name=\"foot1_left\" class=\"foot1\"/>\n",
" <geom name=\"foot2_left\" class=\"foot2\"/>\n",
" </body>\n",
" </body>\n",
" </body>\n",
" </body>\n",
" </body>\n",
" <body name=\"upper_arm_right\" pos=\"0 -.17 .06\">\n",
" <joint name=\"shoulder1_right\" axis=\"2 1 1\" class=\"shoulder\"/>\n",
" <joint name=\"shoulder2_right\" axis=\"0 -1 1\" class=\"shoulder\"/>\n",
" <geom name=\"upper_arm_right\" fromto=\"0 0 0 .16 -.16 -.16\" class=\"arm_upper\"/>\n",
" <body name=\"lower_arm_right\" pos=\".18 -.18 -.18\">\n",
" <joint name=\"elbow_right\" axis=\"0 -1 1\" class=\"elbow\"/>\n",
" <geom name=\"lower_arm_right\" fromto=\".01 .01 .01 .17 .17 .17\" class=\"arm_lower\"/>\n",
" <body name=\"hand_right\" pos=\".18 .18 .18\">\n",
" <geom name=\"hand_right\" zaxis=\"1 1 1\" class=\"hand\" rgba=\"1 0 1 1\"/>\n",
" </body>\n",
" </body>\n",
" </body>\n",
" <body name=\"upper_arm_left\" pos=\"0 .17 .06\">\n",
" <joint name=\"shoulder1_left\" axis=\"-2 1 -1\" class=\"shoulder\"/>\n",
" <joint name=\"shoulder2_left\" axis=\"0 -1 -1\" class=\"shoulder\"/>\n",
" <geom name=\"upper_arm_left\" fromto=\"0 0 0 .16 .16 -.16\" class=\"arm_upper\"/>\n",
" <body name=\"lower_arm_left\" pos=\".18 .18 -.18\">\n",
" <joint name=\"elbow_left\" axis=\"0 -1 -1\" class=\"elbow\"/>\n",
" <geom name=\"lower_arm_left\" fromto=\".01 -.01 .01 .17 -.17 .17\" class=\"arm_lower\"/>\n",
" <body name=\"hand_left\" pos=\".18 -.18 .18\">\n",
" <geom name=\"hand_left\" zaxis=\"1 -1 1\" class=\"hand\"/>\n",
" </body>\n",
" </body>\n",
" </body>\n",
" </body>\n",
" </worldbody>\n",
"\n",
" <contact>\n",
" <exclude body1=\"waist_lower\" body2=\"thigh_right\"/>\n",
" <exclude body1=\"waist_lower\" body2=\"thigh_left\"/>\n",
" </contact>\n",
"\n",
" <tendon>\n",
" <fixed name=\"hamstring_right\" limited=\"true\" range=\"-0.3 2\">\n",
" <joint joint=\"hip_y_right\" coef=\".5\"/>\n",
" <joint joint=\"knee_right\" coef=\"-.5\"/>\n",
" </fixed>\n",
" <fixed name=\"hamstring_left\" limited=\"true\" range=\"-0.3 2\">\n",
" <joint joint=\"hip_y_left\" coef=\".5\"/>\n",
" <joint joint=\"knee_left\" coef=\"-.5\"/>\n",
" </fixed>\n",
" </tendon>\n",
"\n",
" <actuator>\n",
" <position name=\"abdomen_z\" kp=\"40\" joint=\"abdomen_z\"/>\n",
" <position name=\"abdomen_y\" kp=\"40\" joint=\"abdomen_y\"/>\n",
" <position name=\"abdomen_x\" kp=\"40\" joint=\"abdomen_x\"/>\n",
" <position name=\"hip_x_right\" kp=\"40\" joint=\"hip_x_right\"/>\n",
" <position name=\"hip_z_right\" kp=\"40\" joint=\"hip_z_right\"/>\n",
" <position name=\"hip_y_right\" kp=\"120\" joint=\"hip_y_right\"/>\n",
" <position name=\"knee_right\" kp=\"80\" joint=\"knee_right\"/>\n",
" <position name=\"ankle_y_right\" kp=\"20\" joint=\"ankle_y_right\"/>\n",
" <position name=\"ankle_x_right\" kp=\"20\" joint=\"ankle_x_right\"/>\n",
" <position name=\"hip_x_left\" kp=\"40\" joint=\"hip_x_left\"/>\n",
" <position name=\"hip_z_left\" kp=\"40\" joint=\"hip_z_left\"/>\n",
" <position name=\"hip_y_left\" kp=\"120\" joint=\"hip_y_left\"/>\n",
" <position name=\"knee_left\" kp=\"80\" joint=\"knee_left\"/>\n",
" <position name=\"ankle_y_left\" kp=\"20\" joint=\"ankle_y_left\"/>\n",
" <position name=\"ankle_x_left\" kp=\"20\" joint=\"ankle_x_left\"/>\n",
" <position name=\"shoulder1_right\" kp=\"20\" joint=\"shoulder1_right\"/>\n",
" <position name=\"shoulder2_right\" kp=\"20\" joint=\"shoulder2_right\"/>\n",
" <position name=\"elbow_right\" kp=\"40\" joint=\"elbow_right\"/>\n",
" <position name=\"shoulder1_left\" kp=\"20\" joint=\"shoulder1_left\"/>\n",
" <position name=\"shoulder2_left\" kp=\"20\" joint=\"shoulder2_left\"/>\n",
" <position name=\"elbow_left\" kp=\"40\" joint=\"elbow_left\"/>\n",
" </actuator>\n",
"\n",
" <sensor>\n",
" <framepos objtype=\"geom\" objname=\"hand_right\" reftype=\"xbody\" refname=\"target\"/>\n",
" <actuatorfrc actuator=\"abdomen_z\"/>\n",
" <actuatorfrc actuator=\"abdomen_y\"/>\n",
" <actuatorfrc actuator=\"abdomen_x\"/>\n",
" <actuatorfrc actuator=\"hip_x_right\"/>\n",
" <actuatorfrc actuator=\"hip_z_right\"/>\n",
" <actuatorfrc actuator=\"hip_y_right\"/>\n",
" <actuatorfrc actuator=\"knee_right\"/>\n",
" <actuatorfrc actuator=\"ankle_y_right\"/>\n",
" <actuatorfrc actuator=\"ankle_x_right\"/>\n",
" <actuatorfrc actuator=\"hip_x_left\"/>\n",
" <actuatorfrc actuator=\"hip_z_left\"/>\n",
" <actuatorfrc actuator=\"hip_y_left\"/>\n",
" <actuatorfrc actuator=\"knee_left\"/>\n",
" <actuatorfrc actuator=\"ankle_y_left\"/>\n",
" <actuatorfrc actuator=\"ankle_x_left\"/>\n",
" <actuatorfrc actuator=\"shoulder1_right\"/>\n",
" <actuatorfrc actuator=\"shoulder2_right\"/>\n",
" <actuatorfrc actuator=\"elbow_right\"/>\n",
" <actuatorfrc actuator=\"shoulder1_left\"/>\n",
" <actuatorfrc actuator=\"shoulder2_left\"/>\n",
" <actuatorfrc actuator=\"elbow_left\"/>\n",
" </sensor>\n",
"\n",
" <keyframe>\n",
" <!--\n",
" The values below are split into rows for readibility:\n",
" torso position\n",
" torso orientation\n",
" spinal\n",
" right leg\n",
" left leg\n",
" arms\n",
" -->\n",
" <key name=\"squat\"\n",
" qpos=\"0 0 0.596\n",
" 0.988015 0 0.154359 0\n",
" 0 0.4 0\n",
" -0.25 -0.5 -2.5 -2.65 -0.8 0.56\n",
" -0.25 -0.5 -2.5 -2.65 -0.8 0.56\n",
" 0 0 0 0 0 0\"/>\n",
" <key name=\"stand_on_left_leg\"\n",
" qpos=\"0 0 1.21948\n",
" 0.971588 -0.179973 0.135318 -0.0729076\n",
" -0.0516 -0.202 0.23\n",
" -0.24 -0.007 -0.34 -1.76 -0.466 -0.0415\n",
" -0.08 -0.01 -0.37 -0.685 -0.35 -0.09\n",
" 0.109 -0.067 -0.7 -0.05 0.12 0.16\"/>\n",
" <key name=\"prone\"\n",
" qpos=\"0.4 0 0.0757706\n",
" 0.7325 0 0.680767 0\n",
" 0 0.0729 0\n",
" 0.0077 0.0019 -0.026 -0.351 -0.27 0\n",
" 0.0077 0.0019 -0.026 -0.351 -0.27 0\n",
" 0.56 -0.62 -1.752\n",
" 0.56 -0.62 -1.752\"/>\n",
" <key name=\"supine\"\n",
" qpos=\"-0.4 0 0.08122\n",
" 0.722788 0 -0.69107 0\n",
" 0 -0.25 0\n",
" 0.0182 0.0142 0.3 0.042 -0.44 -0.02\n",
" 0.0182 0.0142 0.3 0.042 -0.44 -0.02\n",
" 0.186 -0.73 -1.73\n",
" 0.186 -0.73 -1.73\"/>\n",
" </keyframe>\n",
"</mujoco>\n",
"\"\"\"\n",
"\n",
"# Load model, make data, make list of data for multithreading\n",
"model = mujoco.MjModel.from_xml_string(xml)\n",
"data = mujoco.MjData(model)\n",
"data_list = [mujoco.MjData(model) for _ in range(nthread)]\n",
"\n",
"# Set the state to the \"squat\" keyframe, call mj_forward.\n",
"key = model.key('squat').id\n",
"mujoco.mj_resetDataKeyframe(model, data, key)\n",
"mujoco.mj_forward(model, data)\n",
"\n",
"# If a renderer exists, close it.\n",
"if 'renderer' in locals():\n",
" renderer.close()\n",
"\n",
"# Make a Renderer and a camera.\n",
"renderer = mujoco.Renderer(model)\n",
"camera = mujoco.MjvCamera()\n",
"mujoco.mjv_defaultFreeCamera(model, camera)\n",
"camera.distance = 3\n",
"camera.elevation = -10\n",
"\n",
"# Point the camera at the humanoid, render.\n",
"# camera.lookat = data.body('torso').subtree_com\n",
"# renderer.update_scene(data, camera)\n",
"# media.show_image(renderer.render())\n",
"\n",
"def reach(ctrl0T, target, T, torque_scale, traj=None, multithread=False):\n",
" \"\"\"Residual for target-reaching task.\n",
"\n",
" Args:\n",
" ctrl0T: contatenation of the first and last control vectors.\n",
" target: target to which the right hand should reach.\n",
" T: final time for the rollout.\n",
" torque_scale: coefficient by which to scale the torques.\n",
" traj: optional list of positions to be recorded.\n",
"\n",
" Returns:\n",
" The residual of the target-reaching task.\n",
" \"\"\"\n",
" # Extract the initial and final ctrl vectors, transpose to row vectors\n",
" ctrl0 = ctrl0T[:model.nu, :].T\n",
" ctrlT = ctrl0T[model.nu:, :].T\n",
"\n",
" # Move the mocap body to the target\n",
" mocapid = model.body('target').mocapid\n",
" data.mocap_pos[mocapid] = target\n",
"\n",
" # Append the mocap targets to the controls\n",
" nroll = ctrl0.shape[0]\n",
" mocap = np.tile(data.mocap_pos[mocapid], (nroll, 1))\n",
" ctrl0 = np.hstack((ctrl0, mocap))\n",
" ctrlT = np.hstack((ctrlT, mocap))\n",
"\n",
" # Define control spec (ctrl + mocap_pos)\n",
" mjtState = mujoco.mjtState\n",
" control_spec = mjtState.mjSTATE_CTRL | mjtState.mjSTATE_MOCAP_POS\n",
"\n",
" # Interpolate and stack the control sequences\n",
" nstep = int(np.round(T / model.opt.timestep))\n",
" control = np.stack(np.linspace(ctrl0, ctrlT, nstep), axis=1)\n",
"\n",
" if not multithread:\n",
" datas = [data]\n",
" else:\n",
" datas = data_list\n",
"\n",
" # Reset to the \"squat\" keyframe, get the initial state\n",
" for d in datas:\n",
" key = model.key('squat').id\n",
" mujoco.mj_resetDataKeyframe(model, d, key)\n",
" spec = mjtState.mjSTATE_FULLPHYSICS\n",
" nstate = mujoco.mj_stateSize(model, spec)\n",
" state = np.empty(nstate)\n",
" mujoco.mj_getState(model, d, state, spec)\n",
"\n",
" # Perform rollouts (sensors.shape == nroll, nstep, nsensordata)\n",
" states, sensors = rollout.rollout(model, datas, state, control,\n",
" control_spec=control_spec)\n",
"\n",
" # If requested, extract qpos into traj\n",
" if traj is not None:\n",
" assert states.shape[0] == 1\n",
" # Skip the first element in state (mjData.time)\n",
" traj.extend(np.split(states[0, :, 1:model.nq+1], nstep))\n",
"\n",
" # Scale torque sensors\n",
" sensors[:, :, 3:] *= torque_scale\n",
"\n",
" # Reshape to stack the sensor values, transpose to column vectors\n",
" sensors = sensors.reshape((sensors.shape[0], -1)).T\n",
"\n",
" # The normalizer keeps objective values similar when changing T or timestep.\n",
" normalizer = 100 * model.opt.timestep / T\n",
" return normalizer * sensors\n",
"\n",
"def render_solution(x, target):\n",
" # Ask reach to save positions to traj.\n",
" traj = []\n",
" reach(x, target, T, torque_scale, traj=traj);\n",
"\n",
" frames = []\n",
" counter = 0\n",
" print('Rendering frames:', flush=True, end='')\n",
" for qpos in traj:\n",
" # Set positions, call mj_forward to update kinematics.\n",
" data.qpos = qpos\n",
" mujoco.mj_forward(model, data)\n",
"\n",
" # Render and save frames.\n",
" camera.lookat = data.body('torso').subtree_com\n",
" renderer.update_scene(data, camera)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
" counter += 1\n",
" if counter % 10 == 0:\n",
" print(f' {counter}', flush=True, end='')\n",
" return frames\n",
"\n",
"# Settings for the optimization\n",
"T = 0.7 # Rollout length (seconds)\n",
"torque_scale = 0.003 # Scaling for the torques\n",
"\n",
"# Bounds are the stacked control bounds.\n",
"lower = np.atleast_2d(model.actuator_ctrlrange[:,0]).T\n",
"upper = np.atleast_2d(model.actuator_ctrlrange[:,1]).T\n",
"bounds = [np.vstack((lower, lower)), np.vstack((upper, upper))]\n",
"\n",
"# Initial guess is midpoint of the bounds\n",
"x0 = 0.5 * (bounds[1] + bounds[0])\n",
"target = (.4, -.3, 1.2)\n",
"\n",
"# Use default target.\n",
"target = data.mocap_pos[model.body('target').mocapid]\n",
"\n",
"# Visualize the initial guess.\n",
"media.show_video(render_solution(x0, target))"
]
},
{
"cell_type": "markdown",
"id": "39dab9e5-1d67-45bb-bcec-af471036d592",
"metadata": {
"id": "39dab9e5-1d67-45bb-bcec-af471036d592"
},
"source": [
"Next, let's run the optimization in a single threaded and multi-threaded modes and render the resulting solutions."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "248b5bc1-d48d-4875-8c4f-505ab50cc7ca",
"metadata": {
"id": "248b5bc1-d48d-4875-8c4f-505ab50cc7ca",
"outputId": "8150938a-0b7b-4d01-e26e-9c59014d9ce5"
},
"outputs": [],
"source": [
"reach_target = lambda x: reach(x, target, T, torque_scale, traj=None, multithread=False)\n",
"reach_target_multithread = lambda x: reach(x, target, T, torque_scale, traj=None, multithread=True)\n",
"\n",
"print('Using 1 thread')\n",
"x_single, _ = minimize.least_squares(x0, reach_target, bounds, verbose=minimize.Verbosity.FINAL)\n",
"print()\n",
"\n",
"print(f'Using {nthread} threads')\n",
"x_multi, _ = minimize.least_squares(x0, reach_target_multithread, bounds, verbose=minimize.Verbosity.FINAL)\n",
"\n",
"# Render the solution to verify the results are the same\n",
"print()\n",
"frames_single = render_solution(x_single, target)\n",
"frames_multi = render_solution(x_multi, target)\n",
"media.show_video(np.concatenate((frames_single, frames_multi), axis=2))"
]
},
{
"cell_type": "markdown",
"id": "6f1e003c-96a4-4428-9c56-d9b833a71d16",
"metadata": {
"id": "6f1e003c-96a4-4428-9c56-d9b833a71d16"
},
"source": [
"By the using multithreaded `rollout` the minimization completed ~4x faster on a 5800X3D and the results are the same."
]
},
{
"cell_type": "markdown",
"id": "7944637f",
"metadata": {},
"source": [
"# Benchmarking `rollout`\n",
"\n",
"The `rollout.rollout` function in the `mujoco` Python library runs batches of simulations for a fixed number steps. It can run in single or multi-threaded modes. The speedup over pure Python is significant because `rollout` can be easily configured to use multithreading.\n",
"\n",
"To show the speedup, we will run benchmarks with the \"tippe top\", \"humanoid\", and \"humanoid100\" models."
]
},
{
"cell_type": "markdown",
"id": "7ef2bd0b",
"metadata": {},
"source": [
"## Python rollouts versus `rollout`\n",
"\n",
"The benchmark runs the three models with varying batch and step counts.\n",
"\n",
"The Python code for nbatch rollouts of nstep steps is:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "cb6355dd",
"metadata": {},
"outputs": [],
"source": [
"def python_rollout(model, init_model, nbatch, nstep):\n",
" for i in range(nbatch):\n",
" data = init_model(model)\n",
" for i in range(nstep):\n",
" mujoco.mj_step(model, data)"
]
},
{
"cell_type": "markdown",
"id": "6fe4a78b",
"metadata": {},
"source": [
"To run nbatch rollouts with `rollout`, we need to make an array of nbatch initial states to start the rollouts from.\n",
"\n",
"Additionally, to use `rollout`'s parallelism, we must pass one MjData per thread.\n",
"\n",
"The resulting `rollout` call parameterized by nbatch, nstep, and nthread is:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "74f143e2",
"metadata": {},
"outputs": [],
"source": [
"def nthread_rollout(model, init_model, nbatch, nstep, nthread):\n",
" # Initialize the MjData for the given model using the provided initializer\n",
" data = init_model(model)\n",
" rollout.rollout(model,\n",
" [copy.copy(data) for _ in range(nthread)], # Create one MjData per thread\n",
" np.tile(get_state(model, data), (nbatch, 1)), # Tile the initial condition nbatch times\n",
" nstep=nstep)"
]
},
{
"cell_type": "markdown",
"id": "b75dc44c",
"metadata": {},
"source": [
"Next, we benchmark the Python loop and `rollout` in both single threaded and multithreaded modes. The three benchmarks take about 2.5 minutes in total to run in total on an AMD 5800X3D."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0301e3ee",
"metadata": {},
"outputs": [],
"source": [
"#@title Benchmarking and plotting code\n",
"\n",
"def benchmark_rollout(model, data, init_model, nbatch, nstep, nominal_nbatch, nominal_nstep, ntiming=1):\n",
" print('Benchmarking pure python', end='\\r')\n",
" start = time.time()\n",
" t_python_nbatch = benchmark(lambda x: python_rollout(model, init_model, x, nominal_nstep), nbatch, ntiming)\n",
" t_python_nstep = benchmark(lambda x: python_rollout(model, init_model, nominal_nbatch, x), nstep, ntiming)\n",
" end = time.time()\n",
" print(f'Benchmarking pure python took {end-start:0.1f} seconds')\n",
"\n",
" print('Benchmarking single threaded rollout', end='\\r')\n",
" start = time.time()\n",
" t_rollout_single_nbatch = benchmark(lambda x: nthread_rollout(model, init_model, x, nominal_nstep, nthread=1), nbatch, ntiming)\n",
" t_rollout_single_nstep = benchmark(lambda x: nthread_rollout(model, init_model, nominal_nbatch, x, nthread=1), nstep, ntiming)\n",
" end = time.time()\n",
" print(f'Benchmarking single threaded rollout took {end-start:0.1f} seconds')\n",
"\n",
" print(f'Benchmarking multithreaded rollout using {nthread} threads', end='\\r')\n",
" start = time.time()\n",
" t_rollout_multi_nbatch = benchmark(lambda x: nthread_rollout(model, init_model, x, nominal_nstep, nthread), nbatch, ntiming)\n",
" t_rollout_multi_nstep = benchmark(lambda x: nthread_rollout(model, init_model, nominal_nbatch, x, nthread), nstep, ntiming)\n",
" end = time.time()\n",
" print(f'Benchmarking multithreaded rollout using {nthread} threads took {end-start:0.1f} seconds')\n",
"\n",
" return (t_python_nbatch, t_rollout_single_nbatch, t_rollout_multi_nbatch,\n",
" t_python_nstep, t_rollout_single_nstep, t_rollout_multi_nstep)\n",
"\n",
"def plot_benchmark(results, nbatch, nstep, nominal_nbatch, nominal_nstep):\n",
" (t_python_nbatch, t_rollout_single_nbatch, t_rollout_multi_nbatch,\n",
" t_python_nstep, t_rollout_single_nstep, t_rollout_multi_nstep) = results\n",
"\n",
" width = 0.25\n",
" x = np.array([i for i in range(len(nbatch))])\n",
"\n",
" ticker = matplotlib.ticker.EngFormatter(unit='')\n",
"\n",
" fig, (ax1, ax2) = plt.subplots(1, 2, sharey=True)\n",
" steps_per_t = np.array(nbatch) * nominal_nstep\n",
" steps_per_t_python = steps_per_t / t_python_nbatch\n",
" steps_per_t_single = steps_per_t / t_rollout_single_nbatch\n",
" steps_per_t_multi = steps_per_t / t_rollout_multi_nbatch\n",
" ax1.bar(x + 0*width, steps_per_t_python, width=width, label='python')\n",
" ax1.bar(x + 1*width, steps_per_t_single, width=width, label='rollout single threaded')\n",
" ax1.bar(x + 2*width, steps_per_t_multi, width=width, label='rollout multithreaded')\n",
" ax1.set_xticks(x + width, nbatch)\n",
" ax1.yaxis.set_major_formatter(ticker)\n",
" ax1.grid()\n",
" ax1.set_xlabel('nbatch')\n",
" ax1.set_ylabel('steps per second')\n",
" ax1.set_title(f'nbatch varied, nstep = {nominal_nstep}')\n",
"\n",
" x = np.array([i for i in range(len(nstep))])\n",
" steps_per_t = np.array(nstep) * nominal_nbatch\n",
" steps_per_t_python = steps_per_t / t_python_nstep\n",
" steps_per_t_single = steps_per_t / t_rollout_single_nstep\n",
" steps_per_t_multi = steps_per_t / t_rollout_multi_nstep\n",
" ax2.bar(x + 0*width, steps_per_t_python, width=width, label='python')\n",
" ax2.bar(x + 1*width, steps_per_t_single, width=width, label='rollout single threaded')\n",
" ax2.bar(x + 2*width, steps_per_t_multi, width=width, label='rollout multithreaded')\n",
" ax2.set_xticks(x + width, nstep)\n",
" ax2.yaxis.set_major_formatter(ticker)\n",
" ax2.grid()\n",
" ax2.set_xlabel('nstep')\n",
" ax2.set_title(f'nstep varied, nbatch = {nominal_nbatch}')\n",
"\n",
" ax2.legend(loc=(1.04, 0.0))\n",
" fig.set_size_inches(10, 4)\n",
" plt.tight_layout()"
]
},
{
"cell_type": "markdown",
"id": "08fb0c12",
"metadata": {},
"source": [
"### Tippe Top Benchmark"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f7e54830",
"metadata": {},
"outputs": [],
"source": [
"nominal_nbatch = 100 # Batch size to use when testing different nstep\n",
"nominal_nstep = 1000 # Step count to use when testing different nbatch\n",
"nbatch = [1, 10, 100, 500, 1000] # Batch sizes to benchmark\n",
"nstep = sorted([1, 10, 100, 1000, 2000, 4000]) # Step counts to benchmark\n",
"\n",
"top_benchmark_results = benchmark_rollout(top_model, top_data, init_top, nbatch, nstep, nominal_nbatch, nominal_nstep)\n",
"plot_benchmark(top_benchmark_results, nbatch, nstep, nominal_nbatch, nominal_nstep)"
]
},
{
"cell_type": "markdown",
"id": "edefb26e",
"metadata": {},
"source": [
"### Humanoid Benchmark"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c9e58c6c",
"metadata": {},
"outputs": [],
"source": [
"nominal_nbatch = 200 # Batch size to use when testing different nstep\n",
"nominal_nstep = 500 # Step count to use when testing different nbatch\n",
"nbatch = [1, 10, 100, 200, 400] # Batch sizes to benchmark\n",
"nstep = sorted([1, 10, 100, 500, 1000]) # Step counts to benchmark\n",
"\n",
"humanoid_benchmark_results = benchmark_rollout(humanoid_model, humanoid_data, init_humanoid, nbatch, nstep, nominal_nbatch, nominal_nstep)\n",
"plot_benchmark(humanoid_benchmark_results, nbatch, nstep, nominal_nbatch, nominal_nstep)"
]
},
{
"cell_type": "markdown",
"id": "468903bb",
"metadata": {},
"source": [
"### Humanoid100 Benchmark"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "83d775d4",
"metadata": {},
"outputs": [],
"source": [
"nominal_nbatch = 100 # Batch size to use when testing different nstep\n",
"nominal_nstep = 200 # Step count to use when testing different nbatch\n",
"nbatch = [1, 10, 50, 100, 200] # Batch sizes to benchmark\n",
"nstep = sorted([1, 10, 100, 200, 400]) # Step counts to benchmark\n",
"\n",
"humanoid100_benchmark_results = benchmark_rollout(humanoid100_model, humanoid100_data, init_humanoid100, nbatch, nstep, nominal_nbatch, nominal_nstep)\n",
"plot_benchmark(humanoid100_benchmark_results, nbatch, nstep, nominal_nbatch, nominal_nstep)"
]
},
{
"cell_type": "markdown",
"id": "d1133084",
"metadata": {},
"source": [
"## MJX versus `rollout`"
]
},
{
"cell_type": "markdown",
"id": "c1638f2d",
"metadata": {},
"source": [
"Next we will benchmark `rollout` and MJX using the tippe top and humanoid models (humanoid100 is not supported by MJX).\n",
"\n",
"The benchmark below takes about 5.5 minutes on an AMD 5800X3D and an NVIDIA 4090. Almost half the time is spent compiling the JIT functions. The JIT functions are cached so that subsequent runs of the benchmark run much faster.\n",
"\n",
"**Note:** MJX is most useful when coupled with something else that runs best on a GPU, like a neural network. Without any such additional workload, CPU based simulation will sometimes be faster, especially when using less than state-of-the-art GPUs. In the results below, the tippe top model runs faster on the 4090 with batch sizes in the 1000's, however the humanoid model always runs slower than the CPU."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7c86d157",
"metadata": {},
"outputs": [],
"source": [
"#@title MJX helper functions\n",
"def init_mjx_batch(model, init_model, nbatch, skip_jit=False):\n",
" data = init_model(model)\n",
"\n",
" # Make MJX versions of model and data\n",
" mjx_model = mjx.put_model(model)\n",
" mjx_data = mjx.put_data(model, data)\n",
"\n",
" jit_step = jax.jit(jax.vmap(mjx.step, in_axes=(None, 0)))\n",
" batch = jax.vmap(lambda x: mjx_data)(jp.array(list(range(nbatch))))\n",
"\n",
" # Trigger JIT for model/batch so as not to include JIT time in benchmarking information\n",
" if not skip_jit:\n",
" batch = jit_step(mjx_model, batch)\n",
"\n",
" return mjx_model, mjx_data, jit_step, batch\n",
"\n",
"def mjx_rollout(model, init_model, nbatch, nstep, jit_step=None):\n",
" # Iniitalize model, skip JIT of stepping function if possible\n",
" if jit_step is None:\n",
" mjx_model, _, jit_step, batch = init_mjx_batch(model, init_model, nbatch)\n",
" else:\n",
" mjx_model, _, _, batch = init_mjx_batch(model, init_model, nbatch, skip_jit=True)\n",
"\n",
" for _ in range(nstep):\n",
" batch = jit_step(mjx_model, batch)\n",
"\n",
"def benchmark_mjx(model, init_model, nbatch, nstep, nominal_nbatch, nominal_nstep, ntiming=1, jit_steps=None):\n",
" print(f'Benchmarking multithreaded rollout using {nthread} threads', end=\"\\r\")\n",
" start = time.time()\n",
" t_rollout_multi_nbatch = benchmark(lambda x: nthread_rollout(model, init_model, x, nominal_nstep, nthread), nbatch, ntiming)\n",
" t_rollout_multi_nstep = benchmark(lambda x: nthread_rollout(model, init_model, nominal_nbatch, x, nthread), nstep, ntiming)\n",
" end = time.time()\n",
" print(f'Benchmarking multithreaded rollout using {nthread} threads took {end-start:0.1f} seconds')\n",
"\n",
" print('Running JIT for MJX', end='\\r')\n",
" start = time.time()\n",
" if jit_steps is None: jit_steps = {}\n",
" for n in nbatch + [nominal_nbatch,]:\n",
" if n not in jit_steps:\n",
" _, _, jit_steps[n], _ = init_mjx_batch(model, init_model, n)\n",
" end = time.time()\n",
" print(f'Running JIT for MJX took {end-start:0.1f} seconds')\n",
"\n",
" print('Benchmarking MJX', end='\\r')\n",
" start = time.time()\n",
" t_mjx_nbatch = benchmark(lambda x: mjx_rollout(model, init_model, x, nominal_nstep, jit_steps[x]), nbatch, ntiming)\n",
" t_mjx_nstep = benchmark(lambda x: mjx_rollout(model, init_model, nominal_nbatch, x, jit_steps[nominal_nbatch]), nstep, ntiming)\n",
" end = time.time()\n",
" print(f'Benchmarking MJX took {end-start:0.1f} seconds')\n",
"\n",
" return t_rollout_multi_nbatch, t_rollout_multi_nstep, t_mjx_nbatch, t_mjx_nstep\n",
"\n",
"def plot_mjx_benchmark(results, nbatch, nstep, nominal_nbatch, nominal_nstep):\n",
" t_rollout_multi_nbatch, t_rollout_multi_nstep, t_mjx_nbatch, t_mjx_nstep = results\n",
"\n",
" width = 0.333\n",
" x = np.array([i for i in range(len(nbatch))])\n",
"\n",
" ticker = matplotlib.ticker.EngFormatter(unit='')\n",
"\n",
" fig, (ax1, ax2) = plt.subplots(1, 2, sharey=True)\n",
" steps_per_t = np.array(nbatch) * nominal_nstep\n",
" steps_per_t_mjx = steps_per_t / t_mjx_nbatch\n",
" steps_per_t_multi = steps_per_t / t_rollout_multi_nbatch\n",
" ax1.bar(x + 0*width, steps_per_t_mjx, width=width, label='mjx')\n",
" ax1.bar(x + 1*width, steps_per_t_multi, width=width, label='rollout multithreaded')\n",
" ax1.set_xticks(x + width / 2, nbatch)\n",
" ax1.yaxis.set_major_formatter(ticker)\n",
" ax1.grid()\n",
" ax1.set_xlabel('nbatch')\n",
" ax1.set_ylabel('steps per second')\n",
" ax1.set_title(f'nbatch varied, nstep = {nominal_nstep}')\n",
"\n",
" x = np.array([i for i in range(len(nstep))])\n",
" steps_per_t = np.array(nstep) * nominal_nbatch\n",
" steps_per_t_mjx = steps_per_t / t_mjx_nstep\n",
" steps_per_t_multi = steps_per_t / t_rollout_multi_nstep\n",
" ax2.bar(x + 0*width, steps_per_t_mjx, width=width, label='mjx')\n",
" ax2.bar(x + 1*width, steps_per_t_multi, width=width, label='rollout multithreaded')\n",
" ax2.set_xticks(x + width / 2, nstep)\n",
" ax2.yaxis.set_major_formatter(ticker)\n",
" ax2.grid()\n",
" ax2.set_xlabel('nstep')\n",
" ax2.set_title(f'nstep varied, nbatch = {nominal_nbatch}')\n",
"\n",
" ax2.legend(loc=(1.04, 0.0))\n",
" fig.set_size_inches(10, 4)\n",
" plt.tight_layout()\n",
"\n",
"# Caches for jit_step functions, they take a long time to compile\n",
"top_jit_steps = {}\n",
"humanoid_jit_steps = {}\n",
"humanoid100_jit_steps = {}"
]
},
{
"cell_type": "markdown",
"id": "a2dafd2e",
"metadata": {},
"source": [
"### MJX Tippe Top Benchmark"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "98c580b0",
"metadata": {},
"outputs": [],
"source": [
"nominal_nbatch = 50000 # Batch size to use when testing different nstep\n",
"nominal_nstep = 200 # Step count to use when testing different nbatch\n",
"nbatch = [100, 1000, 10000, 50000, 100000] # Batch sizes to benchmark\n",
"nstep = [1, 10, 100, 200] # Step counts to benchmark\n",
"\n",
"mjx_top_results = benchmark_mjx(top_model, init_top, nbatch, nstep, nominal_nbatch, nominal_nstep, jit_steps=top_jit_steps)\n",
"plot_mjx_benchmark(mjx_top_results, nbatch, nstep, nominal_nbatch, nominal_nstep)"
]
},
{
"cell_type": "markdown",
"id": "205da5da",
"metadata": {},
"source": [
"### MJX Humanoid Benchmark"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "53166ae1",
"metadata": {},
"outputs": [],
"source": [
"nominal_nbatch = 10000 # Batch size to use when testing different nstep\n",
"nominal_nstep = 200 # Step count to use when testing different nbatch\n",
"nbatch = [100, 1000, 10000, 30000] # Batch sizes to benchmark\n",
"nstep = [1, 10, 100, 200, 400] # Step counts to benchmark\n",
"\n",
"mjx_humanoid_results = benchmark_mjx(humanoid_model, init_humanoid, nbatch, nstep, nominal_nbatch, nominal_nstep, jit_steps=humanoid_jit_steps)\n",
"plot_mjx_benchmark(mjx_humanoid_results, nbatch, nstep, nominal_nbatch, nominal_nstep)"
]
},
{
"cell_type": "markdown",
"id": "fb2caa72",
"metadata": {},
"source": [
"### MJX Multiple Humanoids in One Model\n",
"\n",
"The MJX [documentation](https://mujoco.readthedocs.io/en/stable/mjx.html#mjx-the-sharp-bits) contains a chart comparing the speed of native MuJoCo vs MJX on a variety of devices.\n",
"\n",
"Here we will produce a similar plot to compare MJX and with `rollout`. On a 5800X3D and 4090 devices the benchmark takes about 6.5 minutes to run.\n",
"\n",
"**Note:** These results are not directly comparable since with the plot in the documentation was run on different devices and in particular an A100. Additionally, to run on a 4090 the batch size was redued from 8192 to 4096."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3d6be608",
"metadata": {},
"outputs": [],
"source": [
"max_humanoids = 10\n",
"nbatch = 8192 // 2 # The original benchmark ran with a batch size of 8192, but on a 4090 we can only fit about 4096 humanoids\n",
"nstep = 200\n",
"\n",
"jit_step = jax.jit(jax.vmap(mjx.step, in_axes=(None, 0)))\n",
"t_rollout = []\n",
"t_mjx = []\n",
"for i in range(1, max_humanoids+1):\n",
" print(f'Running benchmark on {i} humanoids')\n",
" model = mujoco.MjModel.from_xml_path(f'mujoco/mjx/mujoco/mjx/test_data/humanoid/{i:02d}_humanoids.xml')\n",
" data = mujoco.MjData(model)\n",
"\n",
" mjx_model = mjx.put_model(model)\n",
" mjx_data = mjx.put_data(model, data)\n",
" batch = jax.vmap(lambda x: mjx_data)(jp.array(list(range(nbatch))))\n",
"\n",
" start = time.perf_counter()\n",
" rollout.rollout(model, [copy.copy(data) for _ in range(nthread)], initial_state=get_state(model, data, nbatch), nstep=humanoid_nstep)\n",
" end = time.perf_counter()\n",
" t_rollout.append(end-start)\n",
"\n",
" # Trigger JIT for model/batch so as not to include JIT time in benchmarking information\n",
" batch = jit_step(mjx_model, batch)\n",
"\n",
" start = time.perf_counter()\n",
" for _ in range(nstep):\n",
" batch = jit_step(mjx_model, batch)\n",
" end = time.perf_counter()\n",
" t_mjx.append(end-start)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b6c5fc2e",
"metadata": {},
"outputs": [],
"source": [
"#@title Plot MJX nhumanoid benchmark\n",
"\n",
"def plot_mjx_nhumanoid_benchmark(t_rollout, t_mjx, nbatch, nstep, max_humanoids):\n",
" nhumanoids = [i for i in range(1, max_humanoids+1)]\n",
"\n",
" width = 0.333\n",
" x = np.array([i for i in range(len(nhumanoids))])\n",
"\n",
" ticker = matplotlib.ticker.EngFormatter(unit='')\n",
"\n",
" fig, ax1 = plt.subplots(1, 1, sharey=True)\n",
" steps_per_t = nbatch * nstep\n",
" steps_per_t_mjx = steps_per_t / np.array(t_mjx)\n",
" steps_per_t_multi = steps_per_t / np.array(t_rollout)\n",
" ax1.bar(x + 0*width, steps_per_t_mjx, width=width, label='mjx')\n",
" ax1.bar(x + 1*width, steps_per_t_multi, width=width, label='rollout multithreaded')\n",
" ax1.set_xticks(x + width / 2, nhumanoids)\n",
" ax1.yaxis.set_major_formatter(ticker)\n",
" ax1.set_yscale('log')\n",
" ax1.grid()\n",
" ax1.set_xlabel('number of humanoids')\n",
" ax1.set_ylabel('steps per second')\n",
" ax1.set_title(f'nhumanoids varied, nbatch = {nbatch}, nstep = {nstep}')\n",
"\n",
" ax1.legend(loc=(1.04, 0.0))\n",
" fig.set_size_inches(8, 4)\n",
" plt.tight_layout()\n",
"\n",
"plot_mjx_nhumanoid_benchmark(t_rollout, t_mjx, nbatch, nstep, max_humanoids)"
]
},
{
"cell_type": "markdown",
"id": "0961c3ec-a691-4875-9a55-227a3d29c472",
"metadata": {
"id": "0961c3ec-a691-4875-9a55-227a3d29c472"
},
"source": [
"# Advanced Usage\n",
"## skip_checks=True\n",
"\n",
"By default rollout performs many checks on the dimensions of its arguments. This it allows it to infer dimensions such as `nbatch` and `nstep`, tile arguments that were not fully specified, and allocate the returned `state` and `sensordata` arrays.\n",
"\n",
"However, these check take time, particularly if `state` and `sensordata` are large or if there are many models and `nstep` is low. So advanced users may want to use the `skip_checks=True` argument in order to acheive additional performance.\n",
"\n",
"If used, certain arguments become non-optional, and all signals must be fully defined (no implicit tiling). In particular:\n",
"* `model` must be a list of length `nbatch`\n",
"* `data` must be a list of length `nthread`\n",
"* `nstep` must be specified\n",
"* `initial_state` must be an array of shape `nbatch x nstate`\n",
"* `control` is optional, but if passed must be an array of shape `nbatch x nstep x ncontrol`\n",
"* `state` is optional, but must be passed if state is to be returned and must be of shape `nbatch x nstep x nstate`\n",
"* `sensordata` is optional, but must be passed if sensor data is to be returned and must be of shape `nbatch x nstep x nsensordata`\n",
"\n",
"As an extreme example, we pass 10,000 humanoid models to `rollout` and simulate 1 step each with and without checks."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d02cc8e8-63cd-4852-ab3c-364a18025a95",
"metadata": {
"id": "d02cc8e8-63cd-4852-ab3c-364a18025a95",
"outputId": "e9038ef1-7d5c-4e81-ebcc-b3dc88597481"
},
"outputs": [],
"source": [
"nbatch = 1000\n",
"nstep = [1, 10, 100, 500]\n",
"ntiming = 5\n",
"\n",
"top_data = init_top(top_model)\n",
"top_datas = [copy.copy(top_data) for _ in range(nthread)]\n",
"initial_state = get_state(top_model, top_data)\n",
"initial_state_tiled = np.tile(initial_state, (nbatch, 1))\n",
"\n",
"# Note: state, sensordata array automatically allocated and return\n",
"def rollout_with_checks(nstep):\n",
" state, sensordata = rollout.rollout([top_model]*nbatch, top_datas, initial_state, nstep=nstep)\n",
"\n",
"# Note: state, sensordata arrays have to be preallocated\n",
"state = None\n",
"sensordata = None\n",
"def rollout_skip_checks(nstep):\n",
" # Note initial state must be tiled\n",
" rollout.rollout([top_model]*nbatch, top_datas, initial_state_tiled, nstep=nstep,\n",
" state=state, sensordata=sensordata, skip_checks=True)\n",
"\n",
"t_with_checks = benchmark(lambda x: rollout_with_checks(x), nstep, ntiming=ntiming)\n",
"t_skip_checks = benchmark(lambda x: rollout_skip_checks(x), nstep, ntiming=ntiming)\n",
"\n",
"steps_per_second = (nbatch * np.array(nstep)) / np.array(t_with_checks)\n",
"steps_per_second_skip_checks = (nbatch * np.array(nstep)) / np.array(t_skip_checks)\n",
"\n",
"plt.loglog(nstep, steps_per_second, label='with checks')\n",
"plt.loglog(nstep, steps_per_second_skip_checks, label='skip checks')\n",
"plt.ylabel('steps per second')\n",
"plt.xlabel('nstep')\n",
"plt.legend()\n",
"plt.grid()"
]
},
{
"cell_type": "markdown",
"id": "92627030-4726-4689-be8b-f1ba75905104",
"metadata": {
"id": "92627030-4726-4689-be8b-f1ba75905104"
},
"source": [
"As expected, as `nstep` increases, the benefits of using skip checks fades quickly. However, at low nstep and high batch sizes, it can make a significant difference.\n",
"\n",
"Notice that the version with checks can use the non-tiled `initial_state`, however the skip checks version must used the tiled version, `initial_state_tiled`."
]
},
{
"cell_type": "markdown",
"id": "84c746e0-5d2e-47dc-ac76-f2b6f790b7c7",
"metadata": {
"id": "84c746e0-5d2e-47dc-ac76-f2b6f790b7c7"
},
"source": [
"## Warmstarting"
]
},
{
"cell_type": "markdown",
"id": "a113380b-5cde-4eff-a235-c3a333910047",
"metadata": {
"id": "a113380b-5cde-4eff-a235-c3a333910047"
},
"source": [
"The `initial_warmstart` parameter can be used to warmstart the constraint solver as described in the [computation chapter](https://mujoco.readthedocs.io/en/stable/computation/index.html#warmstart-acceleration) of the documentation. This can be useful when rolling out models in chunks of steps. Without warmstarting, chaotic systems involving multi-body contact may diverge.\n",
"\n",
"Below we demonstrate this with the tippe top model where the contact solver was changed to CG. This makes the contact force calculation a less repeatable than if the default, Newton's method, were used and allows demonstrating the benefits of warmstarting.\n",
"\n",
"The simulation is run three times. Once with a 6000 step rollout, once with 100 chunks of 60 steps with warmstarting, and once more in 100 chunks of 60 steps without warmstarting."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d4d9f660-f83c-432e-a579-124a7ecab4fb",
"metadata": {
"id": "d4d9f660-f83c-432e-a579-124a7ecab4fb",
"outputId": "139bfc94-4597-4f1a-b73f-85d8c00615b1"
},
"outputs": [],
"source": [
"model = copy.copy(top_model)\n",
"model.opt.solver = mujoco.mjtSolver.mjSOL_CG\n",
"data = init_top(model)\n",
"\n",
"chunks = 100\n",
"steps_per_chunk = 60\n",
"nstep = steps_per_chunk*chunks\n",
"\n",
"initial_state = get_state(model, data)\n",
"\n",
"start = time.time()\n",
"state_all, _ = rollout.rollout(model, data, initial_state, nstep=nstep)\n",
"\n",
"state_chunks = []\n",
"state_chunk, _ = rollout.rollout(model, data, initial_state, nstep=steps_per_chunk)\n",
"state_chunks.append(state_chunk)\n",
"for _ in range(chunks-1):\n",
" state_chunk, _ = rollout.rollout(model, data, state_chunks[-1][0, -1, :], nstep=steps_per_chunk, initial_warmstart=data.qacc_warmstart)\n",
" state_chunks.append(state_chunk)\n",
"state_all_chunked_warmstart = np.concatenate(state_chunks, axis=1)\n",
"\n",
"state_chunks = []\n",
"state_chunk, _ = rollout.rollout(model, data, initial_state, nstep=steps_per_chunk)\n",
"state_chunks.append(state_chunk)\n",
"first_warmstart = None\n",
"for i in range(chunks-1):\n",
" state_chunk, _ = rollout.rollout(model, data, state_chunks[-1][0, -1, :], nstep=steps_per_chunk)\n",
" state_chunks.append(state_chunk)\n",
"state_all_chunked = np.concatenate(state_chunks, axis=1)\n",
"end = time.time()\n",
"\n",
"start_render = time.time()\n",
"framerate = 60\n",
"state_render = np.concatenate((state_all, state_all_chunked, state_all_chunked_warmstart), axis=0)\n",
"camera = 'distant'\n",
"frames1 = render_many(model, data, state_all, framerate, shape=(240, 320), transparent=False, camera=camera)\n",
"frames2 = render_many(model, data, state_all_chunked_warmstart, framerate, shape=(240, 320), transparent=False, camera=camera)\n",
"frames3 = render_many(model, data, state_all_chunked, framerate, shape=(240, 320), transparent=False, camera=camera)\n",
"media.show_video(np.concatenate((frames1, frames2, frames3), axis=2))\n",
"end_render = time.time()\n",
"\n",
"print(f'Rollout took {end-start:.1f} seconds')\n",
"print(f'Rendering took {end_render-start_render:.1f} seconds')"
]
},
{
"cell_type": "markdown",
"id": "7c2cf4fa",
"metadata": {
"id": "7c2cf4fa"
},
"source": [
"As expected, the middle animation (with warmstarting) matches the continuous rollout on the left. However, the model that did not use warmstarting diverged."
]
},
{
"cell_type": "markdown",
"id": "78c1f864-5238-4e27-a7ec-d03c45484d9a",
"metadata": {
"id": "78c1f864-5238-4e27-a7ec-d03c45484d9a"
},
"source": [
"## chunk_size"
]
},
{
"cell_type": "markdown",
"id": "2573fff5",
"metadata": {
"id": "2573fff5"
},
"source": [
"To minimize communication overhead, `rollout` distributes rollouts to threads in groups of rollouts called chunks. By default, `max(1, 0.1 * (nbatch / nthread))` rollouts are assigned to each chunk. While this chunking rule works well for most workloads it is not always optimal, especially when doing short rollouts with small models.\n",
"\n",
"Below we plot the steps per second versus chunk_size when running 1000 hoppers for 1 step each. In his case, the default chunk_size turns out to be quite a bit slower than using an increased chunk size."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a1be8f93",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 615
},
"id": "a1be8f93",
"outputId": "1f832bb7-4a37-4c49-d136-f97b74c61817"
},
"outputs": [],
"source": [
"nbatch = 100\n",
"nstep = 1\n",
"ntiming = 20\n",
"\n",
"#print('Getting Hopper XML description from GitHub:')\n",
"!git clone https://github.com/google-deepmind/dm_control\n",
"hopper_model = mujoco.MjModel.from_xml_path('dm_control/dm_control/suite/hopper.xml')\n",
"hopper_data = mujoco.MjData(hopper_model)\n",
"\n",
"initial_state = get_state(hopper_model, hopper_data)\n",
"initial_states = np.tile(initial_state, (nbatch, 1))\n",
"\n",
"hopper_datas = [copy.copy(hopper_data) for _ in range(nthread)]\n",
"\n",
"def rollout_chunk_size(chunk_size=None):\n",
" rollout.rollout(hopper_model, hopper_datas, initial_states, nstep=nstep, chunk_size=chunk_size)\n",
"\n",
"default_chunk_size = int(max(1.0, 0.1 * nbatch / nthread))\n",
"chunk_sizes = sorted([1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, default_chunk_size])\n",
"t_chunk_size = benchmark(lambda x: rollout_chunk_size(x), chunk_sizes, ntiming=ntiming)\n",
"\n",
"steps_per_second = nbatch * nstep / t_chunk_size\n",
"default_index = [i for i, c in enumerate(chunk_sizes) if c == default_chunk_size][0]\n",
"optimal_index = np.argmax(steps_per_second)\n",
"plt.loglog(chunk_sizes, steps_per_second, color='b')\n",
"plt.plot(chunk_sizes[default_index], steps_per_second[default_index], marker='o', color='r', label='default chunk size')\n",
"plt.plot(chunk_sizes[optimal_index], steps_per_second[optimal_index], marker='o', color='g', label='optimal chunk size')\n",
"plt.ylabel('steps per second')\n",
"plt.xlabel('chunk size')\n",
"plt.legend()\n",
"plt.grid()\n",
"\n",
"print(f'default chunk size: {default_chunk_size} \\t steps per second: {steps_per_second[default_index]:0.1f}')\n",
"print(f'optimal chunk size: {chunk_sizes[optimal_index]} \\t steps per second: {steps_per_second[optimal_index]:0.1f}')"
]
},
{
"cell_type": "markdown",
"id": "d32a77b5-24bd-4d17-80ac-15cc4d03731c",
"metadata": {
"id": "d32a77b5-24bd-4d17-80ac-15cc4d03731c"
},
"source": [
"## Reusing threadpools with the class `Rollout`\n",
"\n",
"The `rollout` module provided the class `Rollout` in addition to the method `rollout`. The class `Rollout` is designed allow safe reuse of the internally managed thread pool.\n",
"\n",
"Reuse can speed things up considerably when rollouts are short. Let's find out how the speedup changes for the tippe top model by rolling it out with increasing numbers of steps."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49",
"metadata": {
"id": "dd05bbdf-f389-4e4e-b389-d47fe976cb49",
"outputId": "c2bef530-41f3-4bf3-a149-4803c3faafd9"
},
"outputs": [],
"source": [
"nbatch = 100\n",
"nsteps = [2**i for i in [2, 3, 4, 5, 6, 7]]\n",
"ntiming = 5\n",
"\n",
"top_data = init_top(top_model)\n",
"\n",
"initial_state = get_state(top_model, top_data)\n",
"initial_states = np.tile(initial_state, (nbatch, 1))\n",
"\n",
"top_datas = [copy.copy(top_data) for _ in range(nthread)]\n",
"\n",
"def rollout_method(nstep):\n",
" for i in range(20):\n",
" rollout.rollout(top_model, top_datas, initial_states, nstep=nstep)\n",
"\n",
"def rollout_class(nstep):\n",
" with rollout.Rollout(nthread=nthread) as rollout_:\n",
" for i in range(20):\n",
" rollout_.rollout(top_model, top_datas, initial_states, nstep=nstep)\n",
"\n",
"t_method = benchmark(lambda x: rollout_method(x), nsteps, ntiming)\n",
"t_class = benchmark(lambda x: rollout_class(x), nsteps, ntiming)\n",
"\n",
"plt.loglog(nsteps, nbatch * np.array(nsteps) / t_method, label='recreating threadpools')\n",
"plt.loglog(nsteps, nbatch * np.array(nsteps) / t_class, label='reusing threadpool')\n",
"plt.xlabel('nstep')\n",
"plt.ylabel('steps per second')\n",
"plt.legend()\n",
"plt.grid()"
]
},
{
"cell_type": "markdown",
"id": "9b3e14a1-71f3-430d-a3c2-1aadcf6c2671",
"metadata": {
"id": "9b3e14a1-71f3-430d-a3c2-1aadcf6c2671"
},
"source": [
"## Reusing threadpools with the method `rollout`"
]
},
{
"cell_type": "markdown",
"id": "0aba4dd6",
"metadata": {
"id": "0aba4dd6"
},
"source": [
"`rollout` will create and reuse a persistent threadpool by passing `persistent_pool=True`. However there are some caveats.\n",
"\n",
"First, because `rollout` is a function and does not know when the user is done calling it, the threadpool pool needs to be shutdown manually like this:"
]
},
{
"cell_type": "code",
"execution_count": 30,
"id": "b6aa6801",
"metadata": {
"id": "b6aa6801"
},
"outputs": [],
"source": [
"nbatch = 1000\n",
"nstep = 1\n",
"\n",
"top_data = init_top(top_model)\n",
"top_datas = [copy.copy(top_data) for _ in range(nthread)]\n",
"\n",
"initial_state = get_state(top_model, top_data)\n",
"initial_states = np.tile(initial_state, (nbatch, 1))\n",
"\n",
"rollout.rollout(model, top_datas, initial_states, nstep=nstep, persistent_pool=True) # Creates a pool\n",
"rollout.rollout(model, top_datas, initial_states, nstep=nstep, persistent_pool=True) # Reuses the previously created pool\n",
"rollout.shutdown_persistent_pool() # Shutdown the pool manually when finished"
]
},
{
"cell_type": "markdown",
"id": "144378d3",
"metadata": {
"id": "144378d3"
},
"source": [
"Second, if `rollout` reuses the same threadpool between calls, it is no longer safe to call `rollout` from multiple threads. For example the following is not allowed (the offending lines are commented out to avoid crashing the interpreter):"
]
},
{
"cell_type": "code",
"execution_count": 31,
"id": "7f46a6d8",
"metadata": {
"id": "7f46a6d8"
},
"outputs": [],
"source": [
"thread1 = threading.Thread(target=lambda: rollout.rollout(top_model, top_datas, initial_states, nstep=nstep, persistent_pool=True))\n",
"thread2 = threading.Thread(target=lambda: rollout.rollout(top_model, top_datas, initial_states, nstep=nstep, persistent_pool=True))\n",
"\n",
"thread1.start()\n",
"#thread2.start() # Do not do this! rollout will be using the same persistent threadpool from two threads and may crash the interpreter\n",
"thread1.join()\n",
"#thread2.join()\n",
"rollout.shutdown_persistent_pool()"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuType": "T4",
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Reference in New Issue View Git Blame Copy Permalink