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

Model Editing

\n", "\n", "This notebook provides an introductory tutorial for model editing in MuJoCo using the `mjSpec` API. This notebook assumes that the reader is already familiar with MuJoCo basic concepts, as demostrated in the [introductory tutorial](https://github.com/google-deepmind/mujoco?tab=readme-ov-file#getting-started). Documentation for this API can be found in the [Model Editing](https://mujoco.readthedocs.io/en/latest/programming/modeledit.html) chapter in the documentation (C API) and in the [Python chapter](https://mujoco.readthedocs.io/en/latest/python.html#model-editing). Here we use the Python API.\n", "\n", "The goal of the API is to allow users to easily interact with and modify MuJoCo\n", "models in Python, similarly to what the JavaScript DOM does for HTML.\n", "\n", "\n" ] }, { "cell_type": "markdown", "metadata": { "id": "sJFuNetilv4m" }, "source": [ "## All imports" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "ArmOx3cBDvFR" }, "outputs": [], "source": [ "!pip install mujoco\n", "\n", "# Set up GPU rendering.\n", "from google.colab import files\n", "import distutils.util\n", "import os\n", "import subprocess\n", "if subprocess.run('nvidia-smi').returncode:\n", " raise RuntimeError(\n", " 'Cannot communicate with GPU. '\n", " 'Make sure you are using a GPU Colab runtime. '\n", " 'Go to the Runtime menu and select Choose runtime type.')\n", "\n", "# Add an ICD config so that glvnd can pick up the Nvidia EGL driver.\n", "# This is usually installed as part of an Nvidia driver package, but the Colab\n", "# kernel doesn't install its driver via APT, and as a result the ICD is missing.\n", "# (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)\n", "NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json'\n", "if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):\n", " with open(NVIDIA_ICD_CONFIG_PATH, 'w') as f:\n", " f.write(\"\"\"{\n", " \"file_format_version\" : \"1.0.0\",\n", " \"ICD\" : {\n", " \"library_path\" : \"libEGL_nvidia.so.0\"\n", " }\n", "}\n", "\"\"\")\n", "\n", "# Configure MuJoCo to use the EGL rendering backend (requires GPU)\n", "print('Setting environment variable to use GPU rendering:')\n", "%env MUJOCO_GL=egl\n", "\n", "# Check if installation was succesful.\n", "try:\n", " print('Checking that the installation succeeded:')\n", " import mujoco as mj\n", " mj.MjModel.from_xml_string('')\n", "except Exception as e:\n", " raise e from RuntimeError(\n", " 'Something went wrong during installation. Check the shell output above '\n", " 'for more information.\\n'\n", " 'If using a hosted Colab runtime, make sure you enable GPU acceleration '\n", " 'by going to the Runtime menu and selecting \"Choose runtime type\".')\n", "\n", "print('Installation successful.')\n", "\n", "# Other imports and helper functions\n", "import numpy as np\n", "from scipy.signal import convolve2d\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", "import matplotlib.colors as mcolors\n", "\n", "# Printing.\n", "np.set_printoptions(precision=3, suppress=True, linewidth=100)\n", "import pygments\n", "from google.colab import output\n", "\n", "from IPython.display import clear_output, HTML, display\n", "clear_output()\n", "\n", "\n", "is_dark = output.eval_js('document.documentElement.matches(\"[theme=dark]\")')\n", "print_style = 'monokai' if is_dark else 'lovelace'\n", "\n", "def print_xml(xml_string):\n", " formatter = pygments.formatters.HtmlFormatter(style=print_style)\n", " lexer = pygments.lexers.XmlLexer()\n", " highlighted = pygments.highlight(xml_string, lexer, formatter)\n", " display(HTML(f\"{highlighted}\"))\n", "\n", "def render(model, data=None, height=300):\n", " if data is None:\n", " data = mj.MjData(model)\n", " with mj.Renderer(model, 480, 640) as renderer:\n", " mj.mj_forward(model, data)\n", " renderer.update_scene(data)\n", " media.show_image(renderer.render(), height=height)" ] }, { "cell_type": "markdown", "metadata": { "id": "iJRNczuyHbuc" }, "source": [ "# Separate parsing and compiling\n", "\n", "Unlike `mj_loadXML` which combines parsing and compiling, when using `mjSpec`, parsing and compiling are separate, allowing for editing steps:\n" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "oummB7I7EfSq" }, "outputs": [], "source": [ "#@title Parse, compile, modify, compile: {vertical-output: true}\n", "\n", "static_model = \"\"\"\n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", "\"\"\"\n", "spec = mj.MjSpec.from_string(static_model)\n", "model = spec.compile()\n", "render(model)\n", "\n", "# Change the mjSpec, re-compile and re-render\n", "spec.modelname = \"edited model\"\n", "geoms = spec.worldbody.find_all(mj.mjtObj.mjOBJ_GEOM)\n", "geoms[0].name = 'blue_box'\n", "geoms[0].rgba = [0, 0, 1, 1]\n", "geoms[1].name = 'yellow_sphere'\n", "geoms[1].rgba = [1, 1, 0, 1]\n", "spec.worldbody.add_geom(name='magenta cylinder',\n", " type=mj.mjtGeom.mjGEOM_CYLINDER,\n", " rgba=[1, 0, 1, 1],\n", " pos=[-.2, 0, .2],\n", " size=[.1, .1, 0])\n", "\n", "model = spec.compile()\n", "render(model)" ] }, { "cell_type": "markdown", "metadata": { "id": "Tw_yUwqxKwCI" }, "source": [ "`mjSpec` can save XML to string, with all modifications:" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "_7HAcWqNGwyw" }, "outputs": [], "source": [ "print_xml(spec.to_xml())" ] }, { "cell_type": "markdown", "metadata": { "id": "NolxAaRn9N9r" }, "source": [ "# Procedural models" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "A6EHNulmFHFI" }, "outputs": [], "source": [ "#@title {vertical-output: true}\n", "\n", "spec = mj.MjSpec()\n", "spec.worldbody.add_light(name=\"top\", pos=[0, 0, 1])\n", "body = spec.worldbody.add_body(name=\"box_and_sphere\",\n", " euler=[0, 0, -30])\n", "body.add_joint(name=\"swing\", type=mj.mjtJoint.mjJNT_HINGE,\n", " axis=[1, -1, 0], pos=[-.2, -.2, -.2])\n", "body.add_geom(name=\"red_box\", type=mj.mjtGeom.mjGEOM_BOX,\n", " size=[.2, .2, .2], rgba=[1, 0, 0, 1])\n", "body.add_geom(name=\"green_sphere\", pos=[.2, .2, .2],\n", " size=[.1, 0, 0], rgba=[0, 1, 0, 1])\n", "model = spec.compile()\n", "\n", "duration = 2 # (seconds)\n", "framerate = 30 # (Hz)\n", "\n", "# enable joint visualization option:\n", "scene_option = mj.MjvOption()\n", "scene_option.flags[mj.mjtVisFlag.mjVIS_JOINT] = True\n", "\n", "# Simulate and display video.\n", "frames = []\n", "data = mj.MjData(model)\n", "mj.mj_resetData(model, data)\n", "with mj.Renderer(model) as renderer:\n", " while data.time < duration:\n", " mj.mj_step(model, data)\n", " if len(frames) < data.time * framerate:\n", " renderer.update_scene(data, scene_option=scene_option)\n", " pixels = renderer.render()\n", " frames.append(pixels)\n", "\n", "media.show_video(frames, fps=framerate)" ] }, { "cell_type": "markdown", "metadata": { "id": "Y4rV2NDh92Ga" }, "source": [ "## Tree\n", "\n", "Let's use procedural model creation to make a simple model of a tree.\n", "\n", "We'll start with an \"arena\" xml, containing only a plane and light, and define some utility functions." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "cellView": "form", "id": "IQ9G54Yu-Cse" }, "outputs": [], "source": [ "#@title Utilities\n", "def branch_frames(num_samples, phi_lower=np.pi / 8, phi_upper=np.pi / 3):\n", " \"\"\"Returns branch direction vectors and normalized attachment heights.\"\"\"\n", " directions = []\n", " theta_slice = (2 * np.pi) / num_samples\n", " phi_slice = (phi_upper - phi_lower) / num_samples\n", " for i in range(num_samples):\n", " theta = np.random.uniform(i * theta_slice, (i + 1) * theta_slice)\n", " phi = phi_lower + np.random.uniform(i * phi_slice, (i + 1) * phi_slice)\n", " x = np.sin(phi) * np.cos(theta)\n", " y = np.sin(phi) * np.sin(theta)\n", " z = np.cos(phi)\n", " directions.append([x, y, z])\n", "\n", " heights = np.linspace(0.6, 1, num_samples)\n", "\n", " return directions, heights\n", "\n", "\n", "def add_arrow(scene, from_, to, radius=0.03, rgba=[0.2, 0.2, 0.6, 1]):\n", " \"\"\"Add an arrow to the scene.\"\"\"\n", " scene.geoms[scene.ngeom].category = mj.mjtCatBit.mjCAT_STATIC\n", " mj.mjv_initGeom(\n", " geom=scene.geoms[scene.ngeom],\n", " type=mj.mjtGeom.mjGEOM_ARROW,\n", " size=np.zeros(3),\n", " pos=np.zeros(3),\n", " mat=np.zeros(9),\n", " rgba=np.asarray(rgba).astype(np.float32),\n", " )\n", " mj.mjv_connector(\n", " geom=scene.geoms[scene.ngeom],\n", " type=mj.mjtGeom.mjGEOM_ARROW,\n", " width=radius,\n", " from_=from_,\n", " to=to,\n", " )\n", " scene.ngeom += 1\n", "\n", "\n", "def unit_bump(x, start, end):\n", " \"\"\"Finite-support unit bump function.\"\"\"\n", " if x <= start or x >= end:\n", " return 0.0\n", " else:\n", " n = (x - start) / (end - start)\n", " n = 2 * n - 1\n", " return np.exp(n * n / (n * n - 1))" ] }, { "cell_type": "markdown", "metadata": { "id": "QAfFHvifGnBx" }, "source": [ "Our tree creation function is called recursively to add branches and leaves." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "Bajncwzn-Lid" }, "outputs": [], "source": [ "# @title Tree creation\n", "def procedural_tree(\n", " num_child_branch=5,\n", " length=0.5,\n", " thickness=0.04,\n", " depth=4,\n", " this_body=None,\n", " spec=None,\n", "):\n", " \"\"\"Recursive function that builds a tree of branches and leaves.\"\"\"\n", " BROWN = np.array([0.4, 0.24, 0.0, 1])\n", " GREEN = np.array([0.0, 0.7, 0.2, 1])\n", " SCALE = 0.6\n", "\n", " # Initialize spec and add tree trunk\n", " if this_body is None:\n", " if spec is None:\n", " spec = mj.MjSpec()\n", "\n", " # Disable constraints\n", " spec.option.disableflags |= mj.mjtDisableBit.mjDSBL_CONSTRAINT\n", "\n", " # Air density\n", " spec.option.density = 1.294\n", "\n", " # Defaults for joint and geom\n", " main = spec.default\n", " main.geom.type = mj.mjtGeom.mjGEOM_CAPSULE\n", " main.joint.type = mj.mjtJoint.mjJNT_BALL\n", " main.joint.springdamper = [0.003, 0.7]\n", "\n", " # Visual\n", " spec.stat.center = [0, 0, length]\n", " spec.stat.extent = 2 * length\n", "\n", " # Add trunk body\n", " this_body = spec.worldbody.add_body(name=\"trunk\")\n", " fromto = [0, 0, 0, 0, 0, length]\n", " size = [thickness, 0, 0]\n", " this_body.add_geom(fromto=fromto, size=size, rgba=BROWN)\n", "\n", " # Sample a random color\n", " rgba = np.random.uniform(size=4)\n", " rgba[3] = 1\n", "\n", " # Add child branches using recursive call\n", " if depth > 0:\n", " # Get branch direction vectors and attachment heights\n", " dirs, heights = branch_frames(num_child_branch)\n", " heights *= length\n", "\n", " # Rescale branches with some randomness\n", " thickness *= SCALE * np.random.uniform(0.9, 1.1)\n", " length *= SCALE * np.random.uniform(0.9, 1.1)\n", "\n", " # Branch creation\n", " for i in range(num_child_branch):\n", " branch = this_body.add_body(pos=[0, 0, heights[i]], zaxis=dirs[i])\n", "\n", " fromto = [0, 0, 0, 0, 0, length]\n", " size = [thickness, 0, 0]\n", " rgba = (rgba + BROWN) / 2\n", " branch.add_geom(fromto=fromto, size=size, rgba=rgba)\n", "\n", " branch.add_joint()\n", "\n", " # Recurse.\n", " procedural_tree(\n", " length=length,\n", " thickness=thickness,\n", " depth=depth - 1,\n", " this_body=branch,\n", " spec=spec,\n", " )\n", "\n", " # Max depth reached, add three leaves at the tip\n", " else:\n", " rgba = (rgba + GREEN) / 2\n", " for i in range(3):\n", " pos = [0, 0, length + thickness]\n", " euler = [0, 0, i * 120]\n", " leaf_frame = this_body.add_frame(pos=pos, euler=euler)\n", "\n", " size = length * np.array([0.5, 0.15, 0.01])\n", " pos = length * np.array([0.45, 0, 0])\n", " ellipsoid = mj.mjtGeom.mjGEOM_ELLIPSOID\n", " euler = [np.random.uniform(-50, 50), 0, 0]\n", " leaf = this_body.add_geom(\n", " type=ellipsoid, size=size, pos=pos, rgba=rgba, euler=euler\n", " )\n", "\n", " leaf.set_frame(leaf_frame)\n", "\n", " return spec" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "xrjW4dK-_RKj" }, "outputs": [], "source": [ "#@title Make video\n", "arena_xml = \"\"\"\n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", "\n", "\"\"\"\n", "\n", "spec = procedural_tree(spec=mj.MjSpec.from_string(arena_xml))\n", "model = spec.compile()\n", "data = mj.MjData(model)\n", "\n", "duration = 3 # (seconds)\n", "framerate = 60 # (Hz)\n", "frames = []\n", "with mj.Renderer(model, width=1920 // 3, height=1080 // 3) as renderer:\n", " while data.time < duration:\n", " # Add rightward wind.\n", " wind = 40 * unit_bump(data.time, .2 * duration, .7 * duration)\n", " model.opt.wind[0] = wind\n", "\n", " # Step and render.\n", " mj.mj_step(model, data)\n", " if len(frames) < data.time * framerate:\n", " renderer.update_scene(data)\n", " if wind > 0:\n", " add_arrow(renderer.scene, [0, 0, 1], [wind/25, 0, 1])\n", " pixels = renderer.render()\n", " frames.append(pixels)\n", "\n", "media.show_video(frames, fps=framerate / 2)" ] }, { "cell_type": "markdown", "metadata": { "id": "0wR6XAnKdB6s" }, "source": [ "## Height field\n", "\n", "Height fields represent uneven terrain. There are many ways to generate procedural terrain. Here, we will use [Perlin Noise](https://www.youtube.com/watch?v=9x6NvGkxXhU)." ] }, { "cell_type": "markdown", "metadata": { "id": "yXY7HGfVsVlo" }, "source": [ "### Utilities" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "JWlgTJiHevOg" }, "outputs": [], "source": [ "#@title Perlin noise generator\n", "\n", "# adapted from https://github.com/pvigier/perlin-numpy\n", "\n", "def interpolant(t):\n", " return t*t*t*(t*(t*6 - 15) + 10)\n", "\n", "def perlin(shape, res, tileable=(False, False), interpolant=interpolant):\n", " \"\"\"Generate a 2D numpy array of perlin noise.\n", "\n", " Args:\n", " shape: The shape of the generated array (tuple of two ints).\n", " This must be a multple of res.\n", " res: The number of periods of noise to generate along each\n", " axis (tuple of two ints). Note shape must be a multiple of\n", " res.\n", " tileable: If the noise should be tileable along each axis\n", " (tuple of two bools). Defaults to (False, False).\n", " interpolant: The interpolation function, defaults to\n", " t*t*t*(t*(t*6 - 15) + 10).\n", "\n", " Returns:\n", " A numpy array of shape shape with the generated noise.\n", "\n", " Raises:\n", " ValueError: If shape is not a multiple of res.\n", " \"\"\"\n", " delta = (res[0] / shape[0], res[1] / shape[1])\n", " d = (shape[0] // res[0], shape[1] // res[1])\n", " grid = np.mgrid[0:res[0]:delta[0], 0:res[1]:delta[1]]\\\n", " .transpose(1, 2, 0) % 1\n", " # Gradients\n", " angles = 2*np.pi*np.random.rand(res[0]+1, res[1]+1)\n", " gradients = np.dstack((np.cos(angles), np.sin(angles)))\n", " if tileable[0]:\n", " gradients[-1,:] = gradients[0,:]\n", " if tileable[1]:\n", " gradients[:,-1] = gradients[:,0]\n", " gradients = gradients.repeat(d[0], 0).repeat(d[1], 1)\n", " g00 = gradients[ :-d[0], :-d[1]]\n", " g10 = gradients[d[0]: , :-d[1]]\n", " g01 = gradients[ :-d[0],d[1]: ]\n", " g11 = gradients[d[0]: ,d[1]: ]\n", " # Ramps\n", " n00 = np.sum(np.dstack((grid[:,:,0] , grid[:,:,1] )) * g00, 2)\n", " n10 = np.sum(np.dstack((grid[:,:,0]-1, grid[:,:,1] )) * g10, 2)\n", " n01 = np.sum(np.dstack((grid[:,:,0] , grid[:,:,1]-1)) * g01, 2)\n", " n11 = np.sum(np.dstack((grid[:,:,0]-1, grid[:,:,1]-1)) * g11, 2)\n", " # Interpolation\n", " t = interpolant(grid)\n", " n0 = n00*(1-t[:,:,0]) + t[:,:,0]*n10\n", " n1 = n01*(1-t[:,:,0]) + t[:,:,0]*n11\n", " return np.sqrt(2)*((1-t[:,:,1])*n0 + t[:,:,1]*n1)\n", "\n", "noise = perlin((256, 256), (8, 8))\n", "plt.imshow(noise, cmap = 'gray', interpolation = 'lanczos')\n", "plt.title('Perlin noise example')\n", "plt.colorbar();" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "WKoagEORmCz-" }, "outputs": [], "source": [ "#@title Soft edge slope\n", "def edge_slope(size, border_width=5, blur_iterations=20):\n", " \"\"\"Creates a grayscale image with a white center and fading black edges using convolution.\"\"\"\n", " img = np.ones((size, size), dtype=np.float32)\n", " img[:border_width, :] = 0\n", " img[-border_width:, :] = 0\n", " img[:, :border_width] = 0\n", " img[:, -border_width:] = 0\n", "\n", " kernel = np.array([[1, 1, 1],\n", " [1, 1, 1],\n", " [1, 1, 1]]) / 9.0\n", "\n", " for _ in range(blur_iterations):\n", " img = convolve2d(img, kernel, mode='same', boundary='symm')\n", "\n", " return img\n", "\n", "image = edge_slope(256)\n", "plt.imshow(image, cmap='gray')\n", "plt.title('Smooth sloped edges')\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": { "id": "MCSks-w3spoJ" }, "source": [ "### Textured height-field generator" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "hzroDjVLfxHs" }, "outputs": [], "source": [ "def add_hfield(spec=None, hsize=10, vsize=4):\n", " \"\"\" Function that adds a heighfield with countours\"\"\"\n", "\n", " # Initialize spec\n", " if spec is None:\n", " spec = mj.MjSpec()\n", "\n", " # Generate Perlin noise\n", " size = 128\n", " noise = perlin((size, size), (8, 8))\n", "\n", " # Remap noise to 0 to 1\n", " noise = (noise + 1)/2\n", " noise -= np.min(noise)\n", " noise /= np.max(noise)\n", "\n", " # Makes the edges slope down to avoid sharp boundary\n", " noise *= edge_slope(size)\n", "\n", " # Create height field\n", " hfield = spec.add_hfield(name ='hfield',\n", " size = [hsize, hsize, vsize, vsize/10],\n", " nrow = noise.shape[0],\n", " ncol = noise.shape[1],\n", " userdata = noise.flatten())\n", "\n", " # Add texture\n", " texture = spec.add_texture(name = \"contours\",\n", " type = mj.mjtTexture.mjTEXTURE_2D,\n", " width = 128, height = 128)\n", "\n", " # Create texture map, assign to texture\n", " h = noise\n", " s = 0.7 * np.ones(h.shape)\n", " v = 0.7 * np.ones(h.shape)\n", " hsv = np.stack([h, s, v], axis=-1)\n", " rgb = mcolors.hsv_to_rgb(hsv)\n", " rgb = np.flipud((rgb * 255).astype(np.uint8))\n", " texture.data = rgb.tobytes()\n", "\n", " # Assign texture to material\n", " grid = spec.add_material( name = 'contours')\n", " grid.textures[mj.mjtTextureRole.mjTEXROLE_RGB] = 'contours'\n", " spec.worldbody.add_geom(type = mj.mjtGeom.mjGEOM_HFIELD,\n", " material = 'contours', hfieldname = 'hfield')\n", "\n", " return spec" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "Q2b5BfoZgV79" }, "outputs": [], "source": [ "#@title Video\n", "\n", "arena_xml = \"\"\"\n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", " \n", " \n", " \n", "\n", "\"\"\"\n", "\n", "spec = add_hfield(mj.MjSpec.from_string(arena_xml))\n", "\n", "# Add lights\n", "for x in [-15, 15]:\n", " for y in [-15, 15]:\n", " spec.worldbody.add_light(pos = [x, y, 10], dir = [-x, -y, -15])\n", "\n", "# Add balls\n", "for x in np.linspace(-8, 8, 8):\n", " for y in np.linspace(-8, 8, 8):\n", " pos = [x, y, 4 + np.random.uniform(0, 10)]\n", " ball = spec.worldbody.add_body(pos=pos)\n", " ball.add_geom(type = mj.mjtGeom.mjGEOM_SPHERE, size = [0.5, 0, 0],\n", " rgba = [np.random.uniform()]*3 + [1])\n", " ball.add_freejoint()\n", "\n", "model = spec.compile()\n", "data = mj.MjData(model)\n", "\n", "cam = mj.MjvCamera()\n", "mj.mjv_defaultCamera(cam)\n", "cam.lookat = [0, 0, 0]\n", "cam.distance = 30\n", "cam.elevation = -30\n", "\n", "duration = 6 # (seconds)\n", "framerate = 60 # (Hz)\n", "frames = []\n", "with mj.Renderer(model, width=1920 // 3, height=1080 // 3) as renderer:\n", " while data.time < duration:\n", " mj.mj_step(model, data)\n", " if len(frames) < data.time * framerate:\n", " cam.azimuth = 20 + 20 * (1 - np.cos(np.pi*data.time / duration))\n", " renderer.update_scene(data, cam)\n", " pixels = renderer.render()\n", " frames.append(pixels)\n", "\n", "media.show_video(frames, fps=framerate )" ] }, { "cell_type": "markdown", "metadata": { "id": "92PjTlVCprxn" }, "source": [ "## Mesh\n", "\n", "Random meshes can be easily created by sampling random vertices on the unit sphere and letting MuJoCo create the corresponding convex hull." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "702z1GBRprxq" }, "outputs": [], "source": [ "#@title Add \"rock\" mesh\n", "def add_rock(spec=None, scale=1, name=\"rock\", pos=[0, 0, 0]):\n", " if spec is None:\n", " spec = mj.MjSpec()\n", "\n", " # Defaults\n", " main = spec.default\n", " main.mesh.scale = np.array([scale]*3 , dtype = np.float64)\n", " main.geom.type = mj.mjtGeom.mjGEOM_MESH\n", "\n", " # Random gray-brown color\n", " gray = np.array([.5, .5, .5, 1])\n", " light_brown = np.array([200, 150, 100, 255]) / 255.0\n", " mix = np.random.uniform()\n", " rgba = light_brown*mix + gray*(1-mix)\n", "\n", " # Create mesh vertices\n", " mesh = np.random.normal(size = (20, 3))\n", " mesh /= np.linalg.norm(mesh, axis=1, keepdims=True)\n", "\n", " # Create Body and add mesh to the Geom of the Body\n", " spec.add_mesh(name=name, uservert=mesh.flatten())\n", " body = spec.worldbody.add_body(pos=pos, name=name, mass=1)\n", " body.add_geom(meshname=name, rgba=rgba)\n", " body.add_freejoint()\n", "\n", " return body" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "8BR4mlFdprxx" }, "outputs": [], "source": [ "#@title Video\n", "spec = add_hfield(mj.MjSpec.from_string(arena_xml))\n", "\n", "# Add lights\n", "for x in [-15, 15]:\n", " for y in [-15, 15]:\n", " spec.worldbody.add_light(pos = [x, y, 10], dir = [-x, -y, -15])\n", "\n", "# Add rocks\n", "for x in np.linspace(-8, 8, 8):\n", " for y in np.linspace(-8, 8, 8):\n", " pos = [x, y, np.random.uniform(4, 14)]\n", " rock = add_rock(spec = spec,\n", " scale = np.random.uniform(.3, 1),\n", " name = f\"rock_{x}_{y}\",\n", " pos = pos)\n", "model = spec.compile()\n", "data = mj.MjData(model)\n", "\n", "cam = mj.MjvCamera()\n", "mj.mjv_defaultCamera(cam)\n", "cam.lookat = [0, 0, 0]\n", "cam.distance = 30\n", "cam.elevation = -30\n", "\n", "duration = 6 # (seconds)\n", "framerate = 60 # (Hz)\n", "frames = []\n", "with mj.Renderer(model, width=1920 // 3, height=1080 // 3) as renderer:\n", " while data.time < duration:\n", " mj.mj_step(model, data)\n", " if len(frames) < data.time * framerate:\n", " cam.azimuth = 20 + 20 * (1 - np.cos(np.pi*data.time / duration))\n", " renderer.update_scene(data, cam)\n", " pixels = renderer.render()\n", " frames.append(pixels)\n", "\n", "media.show_video(frames, fps=framerate )" ] }, { "cell_type": "markdown", "metadata": { "id": "3N4YEIVt75_T" }, "source": [ "# `dm_control` example" ] }, { "cell_type": "markdown", "metadata": { "id": "TcQuv56BwaJf" }, "source": [ "A key feature is the ability to easily attach multiple models into a larger one. Disambiguation of duplicated names from different\n", "models, or multiple instances of the same model is handled via user-defined namespacing.\n", "\n", "One example use case is when we want robots with a variable number of joints, as this is a fundamental change to the kinematic structure. The snippets below follow the lines of the [example in dm_control](https://arxiv.org/abs/2006.12983), an older package with similar capabilities." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "7C-hfbtj8nRV" }, "outputs": [], "source": [ "leg_model = \"\"\"\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", "class Leg(object):\n", " \"\"\"A 2-DoF leg with position actuators.\"\"\"\n", " def __init__(self, length, rgba):\n", " self.spec = mj.MjSpec.from_string(leg_model)\n", "\n", " # Thigh:\n", " thigh = self.spec.body('thigh')\n", " thigh.add_geom(fromto=[0, 0, 0, length, 0, 0], size=[length/4, 0, 0], rgba=rgba)\n", "\n", " # Hip:\n", " shin = self.spec.body('shin')\n", " shin.add_geom(fromto=[0, 0, 0, 0, 0, -length], size=[length/5, 0, 0], rgba=rgba)\n", " shin.pos[0] = length" ] }, { "cell_type": "markdown", "metadata": { "id": "MQGsxnIB_RLO" }, "source": [ "The `Leg` class describes an abstract articulated leg, with two joints and corresponding proportional-derivative actuators.\n", "\n", "Note that:\n", "\n", "- MJCF attributes correspond directly to arguments of the `add_()` methods.\n", "- When referencing elements, e.g when specifying the joint to which an actuator is attached, the name string of the MJCF elements is used." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "kMiuMyZW_XoB" }, "outputs": [], "source": [ "BODY_RADIUS = 0.1\n", "random_state = np.random.RandomState(42)\n", "creature_model = \"\"\"\n", "\n", " \n", "\n", " \n", " \n", " \n", "\n", "\"\"\".format(BODY_RADIUS, BODY_RADIUS, BODY_RADIUS / 2)\n", "\n", "def make_creature(num_legs):\n", " \"\"\"Constructs a creature with `num_legs` legs.\"\"\"\n", " rgba = random_state.uniform([0, 0, 0, 1], [1, 1, 1, 1])\n", " spec = mj.MjSpec.from_string(creature_model)\n", " spec.copy_during_attach = True\n", "\n", " # Attach legs to equidistant sites on the circumference.\n", " spec.worldbody.first_geom().rgba = rgba\n", " leg = Leg(length=BODY_RADIUS, rgba=rgba)\n", " for i in range(num_legs):\n", " theta = 2 * i * np.pi / num_legs\n", " hip_pos = BODY_RADIUS * np.array([np.cos(theta), np.sin(theta), 0])\n", " hip_site = spec.worldbody.add_site(pos=hip_pos, euler=[0, 0, theta])\n", " hip_site.attach_body(leg.spec.body('thigh'), '', '-' + str(i))\n", "\n", " return spec" ] }, { "cell_type": "markdown", "metadata": { "id": "QMQ3jc6-_toj" }, "source": [ "The `make_creature` function uses the `attach()` method to procedurally attach legs to the torso. Note that at this stage both the torso and hip attachment sites are children of the `worldbody`, since their parent body has yet to be instantiated. We'll now make an arena with a chequered floor and two lights, and place our creatures in a grid." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "vt2JwXd__1cT" }, "outputs": [], "source": [ "#@title Six Creatures on a floor.{vertical-output: true}\n", "\n", "arena = mj.MjSpec()\n", "\n", "if hasattr(arena, 'compiler'):\n", " arena.compiler.degree = False # MuJoCo dev (next release).\n", "else:\n", " arena.degree = False # MuJoCo release\n", "\n", "# Make arena with textured floor.\n", "chequered = arena.add_texture(\n", " name=\"chequered\", type=mj.mjtTexture.mjTEXTURE_2D,\n", " builtin=mj.mjtBuiltin.mjBUILTIN_CHECKER,\n", " width=300, height=300, rgb1=[.2, .3, .4], rgb2=[.3, .4, .5])\n", "grid = arena.add_material(\n", " name='grid', texrepeat=[5, 5], reflectance=.2\n", " ).textures[mj.mjtTextureRole.mjTEXROLE_RGB] = 'chequered'\n", "arena.worldbody.add_geom(\n", " type=mj.mjtGeom.mjGEOM_PLANE, size=[2, 2, .1], material='grid')\n", "for x in [-2, 2]:\n", " arena.worldbody.add_light(pos=[x, -1, 3], dir=[-x, 1, -2])\n", "\n", "# Instantiate 6 creatures with 3 to 8 legs.\n", "creatures = [make_creature(num_legs=num_legs) for num_legs in range(3, 9)]\n", "\n", "# Place them on a grid in the arena.\n", "height = .15\n", "grid = 5 * BODY_RADIUS\n", "xpos, ypos, zpos = np.meshgrid([-grid, 0, grid], [0, grid], [height])\n", "for i, spec in enumerate(creatures):\n", " # Place spawn sites on a grid.\n", " spawn_pos = (xpos.flat[i], ypos.flat[i], zpos.flat[i])\n", " spawn_site = arena.worldbody.add_site(pos=spawn_pos, group=3)\n", " # Attach to the arena at the spawn sites, with a free joint.\n", " spawn_body = spawn_site.attach_body(spec.worldbody, '', '-' + str(i))\n", " spawn_body.add_freejoint()\n", "\n", "# Instantiate the physics and render.\n", "model = arena.compile()\n", "render(model)" ] }, { "cell_type": "markdown", "metadata": { "id": "mPUGkrCzAFMg" }, "source": [ "Multi-legged creatures, ready to roam! Let's inject some controls and watch them move. We'll generate a sinusoidal open-loop control signal of fixed frequency and random phase, recording both video frames and the horizontal positions of the torso geoms, in order to plot the movement trajectories." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "7gz9FfNzGxPO" }, "outputs": [], "source": [ "#@title Video of the movement{vertical-output: true}\n", "\n", "data = mj.MjData(model)\n", "duration = 10 # (Seconds)\n", "framerate = 30 # (Hz)\n", "video = []\n", "pos_x = []\n", "pos_y = []\n", "geoms = arena.worldbody.find_all(mj.mjtObj.mjOBJ_GEOM)\n", "torsos_data = [data.bind(geom) for geom in geoms if 'torso' in geom.name]\n", "torsos_model = [model.bind(geom) for geom in geoms if 'torso' in geom.name]\n", "actuators = [data.bind(actuator) for actuator in arena.actuators]\n", "\n", "# Control signal frequency, phase, amplitude.\n", "freq = 5\n", "phase = 2 * np.pi * random_state.rand(len(arena.actuators))\n", "amp = 0.9\n", "\n", "# Simulate, saving video frames and torso locations.\n", "mj.mj_resetData(model, data)\n", "with mj.Renderer(model) as renderer:\n", " while data.time < duration:\n", " # Inject controls and step the physics.\n", " for i, actuator in enumerate(actuators):\n", " actuator.ctrl = amp * np.sin(freq * data.time + phase[i])\n", " mj.mj_step(model, data)\n", "\n", " # Save torso horizontal positions using name indexing.\n", " pos_x.append([torso.xpos[0] for torso in torsos_data])\n", " pos_y.append([torso.xpos[1] for torso in torsos_data])\n", "\n", " # Save video frames.\n", " if len(video) < data.time * framerate:\n", " renderer.update_scene(data)\n", " pixels = renderer.render()\n", " video.append(pixels.copy())\n", "\n", "media.show_video(video, fps=framerate)" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "qt2L52e_Tcgt" }, "outputs": [], "source": [ "#@title Movement trajectories{vertical-output: true}\n", "\n", "creature_colors = [torso.rgba[:3] for torso in torsos_model]\n", "fig, ax = plt.subplots(figsize=(4, 4))\n", "ax.set_prop_cycle(color=creature_colors)\n", "_ = ax.plot(pos_x, pos_y, linewidth=4)" ] }, { "cell_type": "markdown", "metadata": { "id": "kSEUoxifxYJ4" }, "source": [ "The plot above shows the corresponding movement trajectories of creature positions. Note how `mjSpec` attribute `id` were used to access both `xpos` and `rgba` values. This attribute is valid only after a model is compiled." ] }, { "cell_type": "markdown", "metadata": { "id": "QZ8alJZz8cB1" }, "source": [ "# Model editing" ] }, { "cell_type": "markdown", "metadata": { "id": "JN3Z4v0PyXKa" }, "source": [ "`mjSpec` elements can be traversed in two ways:\n", "- For elements inside the kinematic tree, the tree can be traversed using the `first` and `next` functions.\n", "- For all other elements, we provide a list.\n", "\n" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "8IcB7nezblyT" }, "outputs": [], "source": [ "#@title Traversing the spec.{vertical-output: true}\n", "\n", "# Get MuJoCo's humanoid model.\n", "print('Getting MuJoCo humanoid XML description from GitHub:')\n", "!git clone https://github.com/google-deepmind/mujoco\n", "humanoid_file = 'mujoco/model/humanoid/humanoid.xml'\n", "humanoid100_file = 'mujoco/model/humanoid/humanoid100.xml'\n", "\n", "spec = mj.MjSpec.from_file(humanoid_file)\n", "\n", "# Function that recursively prints all body names\n", "def print_bodies(parent, level=0):\n", " body = parent.first_body()\n", " while body:\n", " print(''.join(['-' for i in range(level)]) + body.name)\n", " print_bodies(body, level + 1)\n", " body = parent.next_body(body)\n", "\n", "print(\"The spec has the following actuators:\")\n", "for actuator in spec.actuators:\n", " print(actuator.name)\n", "\n", "print(\"\\nThe spec has the following bodies:\")\n", "print_bodies(spec.worldbody)" ] }, { "cell_type": "markdown", "metadata": { "id": "hcGI4orhyzvc" }, "source": [ "An `mjSpec` can be compiled multiple times. If the state has to be preserved between different compilations, then the function `recompile()` must be used, which returns a new `mjData` that contains the mapped state, possibly having a different dimension from the origin." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "uh_N1Fkqk-Mi" }, "outputs": [], "source": [ "#@title Model re-compilation with state preservation.{vertical-output: true}\n", "\n", "spec = mj.MjSpec.from_file(humanoid100_file)\n", "model = spec.compile()\n", "data = mj.MjData(model)\n", "\n", "# Run for 5 seconds\n", "for i in range(1000):\n", " mj.mj_step(model, data)\n", "\n", "# Show result\n", "render(model, data)\n", "\n", "# Create list of all bodies we want to delete\n", "body = spec.worldbody.first_body()\n", "delete_list = []\n", "while body:\n", " geom_type = body.first_geom().type\n", " if (geom_type == mj.mjtGeom.mjGEOM_BOX or\n", " geom_type == mj.mjtGeom.mjGEOM_ELLIPSOID):\n", " delete_list.append(body)\n", " body = spec.worldbody.next_body(body)\n", "\n", "# Remove all bodies in the list from the spec\n", "for body in delete_list:\n", " spec.detach_body(body)\n", "\n", "# # Add another humanoid\n", "spec_humanoid = mj.MjSpec.from_file(humanoid_file)\n", "attachment_frame = spec.worldbody.add_frame(pos=[0, -1, 2])\n", "attachment_frame.attach_body(spec_humanoid.body('torso'), 'a', 'b')\n", "\n", "# Recompile preserving the state\n", "new_model, new_data = spec.recompile(model, data)\n", "\n", "# Show result\n", "render(new_model, new_data)" ] }, { "cell_type": "markdown", "metadata": { "id": "XmSlXirVzLqt" }, "source": [ "Let us load the humanoid model and inspect it." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "UywMzsp5Hnk2" }, "outputs": [], "source": [ "#@title Humanoid model.{vertical-output: true}\n", "\n", "spec = mj.MjSpec.from_file(humanoid_file)\n", "\n", "model = spec.compile()\n", "render(model)" ] }, { "cell_type": "markdown", "metadata": { "id": "owcmKeuSzQRy" }, "source": [ "We wish to remove the arms and replace them with the legs. This can be done by first storing the arm positions into frames attached to the torso. Then we can detach the arms and self-attach the legs into the frames." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "qZCyv-B0IGiG" }, "outputs": [], "source": [ "#@title Humanoid with arms replaced by legs.{vertical-output: true}\n", "\n", "spec = mj.MjSpec.from_file(humanoid_file)\n", "spec.copy_during_attach = True\n", "\n", "# Get the torso, arm, and leg bodies\n", "arm_left = spec.body('upper_arm_left')\n", "arm_right = spec.body('upper_arm_right')\n", "leg_left = spec.body('thigh_left')\n", "leg_right = spec.body('thigh_right')\n", "torso = spec.body('torso')\n", "\n", "# Attach frames at the arm positions\n", "shoulder_left = torso.add_frame(pos=arm_left.pos)\n", "shoulder_right = torso.add_frame(pos=arm_right.pos)\n", "\n", "# Remove the arms\n", "spec.detach_body(arm_left)\n", "spec.detach_body(arm_right)\n", "\n", "# Add new legs\n", "shoulder_left.attach_body(leg_left, 'shoulder', 'left')\n", "shoulder_right.attach_body(leg_right, 'shoulder', 'right')\n", "\n", "model = spec.compile()\n", "render(model, height=400)" ] }, { "cell_type": "markdown", "metadata": { "id": "LnEwEjW3zdua" }, "source": [ "Similarly, different models can be attach together. Here, the right arm is detached and a robot arm from a different model is attached in its place." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "w-NdFhSIIrLL" }, "outputs": [], "source": [ "#@title Humanoid with Franka arm.{vertical-output: true}\n", "\n", "# Get Franka arm from the MuJoCo Menagerie.\n", "!git clone https://github.com/google-deepmind/mujoco_menagerie\n", "franka_file = 'mujoco_menagerie/franka_fr3/fr3.xml'\n", "\n", "spec = mj.MjSpec.from_file(humanoid_file)\n", "franka = mj.MjSpec.from_file(franka_file)\n", "\n", "if hasattr(spec, 'compiler'):\n", " spec.compiler.degree = False # MuJoCo dev (next release).\n", "else:\n", " spec.degree = False # MuJoCo release\n", "\n", "# Replace right arm with frame\n", "arm_right = spec.body('upper_arm_right')\n", "torso = spec.body('torso')\n", "shoulder_right = torso.add_frame(pos=arm_right.pos, quat=[0, 0.8509035, 0, 0.525322])\n", "spec.detach_body(arm_right)\n", "\n", "# Attach Franka arm to humanoid\n", "franka_arm = franka.body('fr3_link2')\n", "shoulder_right.attach_body(franka_arm, 'franka', '')\n", "\n", "model = spec.compile()\n", "render(model, height=400)" ] }, { "cell_type": "markdown", "metadata": { "id": "e_idaggAznXu" }, "source": [ "When doing this, the actuators and all other objects referenced by the attached sub-tree are imported in the new model. All assets are currently imported, referenced or not." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "50lOgJ7mQ2bV" }, "outputs": [], "source": [ "#@title Imported actuators.{vertical-output: true}\n", "\n", "for actuator in spec.actuators:\n", " print(actuator.name)" ] }, { "cell_type": "markdown", "metadata": { "id": "APDoWK4mz0aJ" }, "source": [ "Domain randomization can be performed by attaching multiple times the same spec, edited each time with a new instance of randomized parameters." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "id": "oHjdgkISNLKy" }, "outputs": [], "source": [ "#@title Humanoid with randomized heads and arm poses.{vertical-output: true}\n", "\n", "humanoid = mj.MjSpec.from_file(humanoid_file)\n", "spec = mj.MjSpec()\n", "spec.copy_during_attach = True\n", "\n", "# Delete all key frames to avoid name conflicts\n", "while humanoid.keys:\n", " humanoid.keys[-1].delete()\n", "\n", "# Create a grid of humanoids by attaching humanoid to spec multiple times\n", "for i in range(4):\n", " for j in range(4):\n", " humanoid.materials[0].rgba = [\n", " np.random.uniform(), np.random.uniform(),\n", " np.random.uniform(), 1] # Randomize color\n", " humanoid.body('head').first_geom().size = [\n", " .18*np.random.uniform(), 0, 0] # Randomize head size\n", " humanoid.body('upper_arm_left').quat = [\n", " np.random.uniform(), np.random.uniform(),\n", " np.random.uniform(), np.random.uniform()] # Randomize left arm orientation\n", " humanoid.body('upper_arm_right').quat = [\n", " np.random.uniform(), np.random.uniform(),\n", " np.random.uniform(), np.random.uniform()] # Randomize right arm orientation\n", "\n", " # attach randomized humanoid to parent spec\n", " frame = spec.worldbody.add_frame(pos=[i, j, 0])\n", " frame.attach_body(humanoid.body('torso'), str(i), str(j))\n", "\n", "spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n", " targetbody='3torso3', diffuse=[.8, .8, .8],\n", " specular=[0.3, 0.3, 0.3], pos=[0, -6, 4], cutoff=30)\n", "model = spec.compile()\n", "render(model, height=400)" ] } ], "metadata": { "accelerator": "GPU", "colab": { "collapsed_sections": [ "sJFuNetilv4m", "yXY7HGfVsVlo" ], "gpuClass": "premium", "private_outputs": true, "toc_visible": true }, "gpuClass": "premium", "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 0 }