Loading files
This file is pinned to the selected release. Download the original to open it in your design tool.
8e39d5049This file is pinned to the selected release. Download the original to open it in your design tool.
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "klCa_-bdiWv5"
},
"source": [
"\n",
"\n",
"# <h1><center>Model Editing <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/mjspec.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/></a></center></h1>\n",
"\n",
"This notebook provides an 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\n",
"demonstrated 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",
"<!-- Copyright 2024 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",
"-->\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 successful.\n",
"try:\n",
" print('Checking that the installation succeeded:')\n",
" import mujoco as mj\n",
" mj.MjModel.from_xml_string('<mujoco/>')\n",
"except Exception as e:\n",
" raise e from RuntimeError(\n",
" 'Something went wrong during installation. Check the shell output above '\n",
" 'for more information.\\n'\n",
" 'If using a hosted Colab runtime, make sure you enable GPU acceleration '\n",
" 'by going to the Runtime menu and selecting \"Choose runtime type\".')\n",
"\n",
"print('Installation successful.')\n",
"\n",
"# Other imports and helper functions\n",
"import numpy as np\n",
"import random\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\"<style>{formatter.get_style_defs()}</style>{highlighted}\"))\n",
"\n",
"def render(model, data=None, height=300, camera=-1):\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, camera)\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": {
"cellView": "form",
"id": "oummB7I7EfSq"
},
"outputs": [],
"source": [
"#@title Parse, compile, modify, compile: {vertical-output: true}\n",
"\n",
"static_model = \"\"\"\n",
"<mujoco>\n",
" <worldbody>\n",
" <light name=\"top\" pos=\"0 0 1\"/>\n",
" <geom name=\"red_box\" type=\"box\" size=\".2 .2 .2\" rgba=\"1 0 0 1\"/>\n",
" <geom name=\"green_sphere\" pos=\".2 .2 .2\" size=\".1\" rgba=\"0 1 0 1\"/>\n",
" </worldbody>\n",
"</mujoco>\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",
"<mujoco>\n",
" <visual>\n",
" <headlight diffuse=\".5 .5 .5\" specular=\"1 1 1\"/>\n",
" <global elevation=\"-10\" offwidth=\"2048\" offheight=\"1536\"/>\n",
" <quality shadowsize=\"8192\"/>\n",
" </visual>\n",
"\n",
" <asset>\n",
" <texture type=\"skybox\" builtin=\"gradient\" rgb1=\".5 .5 .5\" rgb2=\"0 0 0\" width=\"10\" height=\"10\"/>\n",
" <texture type=\"2d\" name=\"groundplane\" builtin=\"checker\" mark=\"edge\" rgb1=\"1 1 1\" rgb2=\"1 1 1\" markrgb=\"0 0 0\" width=\"300\" height=\"300\"/>\n",
" <material name=\"groundplane\" texture=\"groundplane\" texuniform=\"true\" texrepeat=\"5 5\" reflectance=\"0.3\"/>\n",
" </asset>\n",
"\n",
" <worldbody>\n",
" <geom name=\"floor\" size=\"5 5 0.01\" type=\"plane\" material=\"groundplane\"/>\n",
" <light pos=\"0 0 3\" diffuse=\"1 1 1\" specular=\"1 1 1\"/>\n",
" </worldbody>\n",
"</mujoco>\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 multiple 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 height field with contours\"\"\"\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",
"<mujoco>\n",
" <visual>\n",
" <headlight diffuse=\".5 .5 .5\" specular=\"1 1 1\"/>\n",
" <global offwidth=\"2048\" offheight=\"1536\"/>\n",
" <quality shadowsize=\"8192\"/>\n",
" </visual>\n",
"\n",
" <asset>\n",
" <texture type=\"skybox\" builtin=\"gradient\" rgb1=\"1 1 1\" rgb2=\"1 1 1\" width=\"10\" height=\"10\"/>\n",
" <texture type=\"2d\" name=\"groundplane\" builtin=\"checker\" mark=\"edge\" rgb1=\"1 1 1\" rgb2=\"1 1 1\" markrgb=\"0 0 0\" width=\"400\" height=\"400\"/>\n",
" <material name=\"groundplane\" texture=\"groundplane\" texrepeat=\"45 45\" reflectance=\"0\"/>\n",
" </asset>\n",
"\n",
" <worldbody>\n",
" <geom name=\"floor\" size=\"150 150 0.1\" type=\"plane\" material=\"groundplane\"/>\n",
" </worldbody>\n",
"</mujoco>\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": "bZ-mpAKhSBHw"
},
"source": [
"## Terrain Generation\n",
"Here we will create a terrain out of different tiles. We start by creating each sinlge tile. Then we move on to putting tiles side by side to create a complete terrain."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"cellView": "form",
"id": "d9oHHP6ZSZeK"
},
"outputs": [],
"source": [
"#@title Utilities\n",
"def render_tile(tile_func, direction=None, cam_distance=6, cam_elevation=-30):\n",
" arena_xml = \"\"\"\n",
" <mujoco>\n",
" <visual>\n",
" <headlight diffuse=\".5 .5 .5\" specular=\"1 1 1\"/>\n",
" <global elevation=\"-10\" offwidth=\"2048\" offheight=\"1536\"/>\n",
" <quality shadowsize=\"8192\"/>\n",
" </visual>\n",
"\n",
" <asset>\n",
" <texture type=\"skybox\" builtin=\"gradient\" rgb1=\".5 .5 .5\" rgb2=\"0 0 0\" width=\"10\" height=\"10\"/>\n",
" <texture type=\"2d\" name=\"groundplane\" builtin=\"checker\" mark=\"edge\" rgb1=\"1 1 1\" rgb2=\"1 1 1\" markrgb=\"0 0 0\" width=\"300\" height=\"300\"/>\n",
" <material name=\"groundplane\" texture=\"groundplane\" texuniform=\"true\" texrepeat=\"5 5\" reflectance=\"0.3\"/>\n",
" </asset>\n",
"\n",
" <worldbody>\n",
" <light pos=\"0 0 5\" diffuse=\"1 1 1\" specular=\"1 1 1\"/>\n",
" </worldbody>\n",
" </mujoco>\n",
" \"\"\"\n",
"\n",
" spec = mj.MjSpec.from_string(arena_xml)\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" name = 'base_tile'\n",
"\n",
" spec.worldbody.add_body(pos=[-3, 0, 0], name=name)\n",
" if direction:\n",
" tile_func(spec, direction=direction)\n",
" else:\n",
" tile_func(spec)\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 = cam_distance\n",
" cam.elevation = cam_elevation\n",
"\n",
" height = 300\n",
"\n",
" with mj.Renderer(model, 480, 640) as renderer:\n",
" mj.mj_forward(model, data)\n",
" renderer.update_scene(data,cam)\n",
" media.show_image(renderer.render(), height=height)\n",
"\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 multiple 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",
"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"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "dtFWSSwWSgeD"
},
"outputs": [],
"source": [
"# @title Stairs\n",
"def stairs(spec=None, grid_loc=[0, 0] , num_stairs=4, direction=1, name='stair'):\n",
" SQUARE_LENGTH = 2\n",
" V_SIZE = 0.076\n",
" H_SIZE = 0.12\n",
" H_STEP = H_SIZE * 2\n",
" V_STEP = V_SIZE * 2\n",
" BROWN = [0.460, 0.362, 0.216, 1.0]\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" # Offset\n",
" x_beginning, y_end = [-SQUARE_LENGTH + H_SIZE] * 2\n",
" x_end, y_beginning = [SQUARE_LENGTH - H_SIZE] * 2\n",
" # Dimension\n",
" size_one = [H_SIZE, SQUARE_LENGTH, V_SIZE]\n",
" size_two = [SQUARE_LENGTH, H_SIZE, V_SIZE]\n",
" # Geoms positions\n",
" x_pos_l = [x_beginning, 0, direction * V_SIZE]\n",
" x_pos_r = [x_end, 0, direction * V_SIZE]\n",
" y_pos_up = [0, y_beginning, direction * V_SIZE]\n",
" y_pos_down = [0, y_end, direction * V_SIZE]\n",
"\n",
" for i in range(num_stairs):\n",
" size_one[1] = SQUARE_LENGTH - H_STEP * i\n",
" size_two[0] = SQUARE_LENGTH - H_STEP * i\n",
"\n",
" x_pos_l[2], x_pos_r[2], y_pos_up[2], y_pos_down[2] = [\n",
" direction * ( V_SIZE + V_STEP * i)] * 4\n",
"\n",
" # Left side\n",
" x_pos_l[0] = x_beginning + H_STEP * i\n",
" body.add_geom(pos=x_pos_l, size=size_one, rgba=BROWN)\n",
" # Right side\n",
" x_pos_r[0] = x_end - H_STEP * i\n",
" body.add_geom(pos=x_pos_r, size=size_one, rgba=BROWN)\n",
" # Top\n",
" y_pos_up[1] = y_beginning - H_STEP * i\n",
" body.add_geom(pos=y_pos_up, size=size_two, rgba=BROWN)\n",
" # Bottom\n",
" y_pos_down[1] = y_end + H_STEP * i\n",
" body.add_geom(pos=y_pos_down, size=size_two, rgba=BROWN)\n",
"\n",
" # Closing\n",
" size = [SQUARE_LENGTH - H_STEP * num_stairs,\n",
" SQUARE_LENGTH - H_STEP * num_stairs,\n",
" V_SIZE]\n",
" pos = [0, 0,\n",
" direction * (V_SIZE + V_STEP * num_stairs)]\n",
" body.add_geom(pos=pos, size=size, rgba=BROWN)\n",
"\n",
"render_tile(stairs, direction=random.choice([-1, 1]))"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "YwX50AtzSmnT"
},
"outputs": [],
"source": [
"# @title Debris (Geoms)\n",
"def debris_with_simple_geoms(spec=None, grid_loc=[0, 0], name='plane'):\n",
" SQUARE_LENGTH = 2\n",
" THICKNESS = 0.05\n",
" BROWN = [0.460, 0.362, 0.216, 1.0]\n",
" RED = [0.6, 0.12, 0.15, 1.0]\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" # Create tile\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" body.add_geom(size=[SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba=BROWN )\n",
"\n",
" # Simple Geoms\n",
" x_beginning, y_end = [-SQUARE_LENGTH + THICKNESS] * 2\n",
" x_end, y_beginning = [SQUARE_LENGTH - THICKNESS] * 2\n",
"\n",
" x_grid = np.linspace(x_beginning, x_end, 10)\n",
" y_grid = np.linspace(y_beginning, y_end, 10)\n",
"\n",
" for i in range(10):\n",
" x = np.random.choice(x_grid)\n",
" y = np.random.choice(y_grid)\n",
"\n",
" pos=[grid_loc[0] + x, grid_loc[1] + y, 0.2]\n",
"\n",
" g_type = None\n",
" size = None\n",
" if random.randint(0, 1):\n",
" g_type = mj.mjtGeom.mjGEOM_BOX\n",
" size = [0.1, 0.1, 0.02]\n",
" else:\n",
" g_type = mj.mjtGeom.mjGEOM_CYLINDER\n",
" size = [0.1, 0.02, 0]\n",
"\n",
" body = spec.worldbody.add_body(pos=pos, name=f'g{i}_{name}', mass=1)\n",
" body.add_geom(type=g_type, size=size, rgba=RED)\n",
" body.add_freejoint()\n",
"\n",
"render_tile(debris_with_simple_geoms)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "Wm_0DHK8S2kQ"
},
"outputs": [],
"source": [
"# @title Debris (Mesh)\n",
"def debris(spec=None, grid_loc=[0, 0] , name='debris'):\n",
" SQUARE_LENGTH = 2\n",
" THICKNESS = 0.05\n",
" STEP = THICKNESS * 8\n",
" SCALE = 0.1\n",
" BROWN = [0.460, 0.362, 0.216, 1.0]\n",
" RED = [0.6, 0.12, 0.15, 1.0]\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
" main.mesh.scale = np.array([SCALE]*3, dtype=np.float64)\n",
"\n",
" x_beginning = -SQUARE_LENGTH + THICKNESS\n",
" y_beginning = SQUARE_LENGTH - THICKNESS\n",
"\n",
" # Create tile\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" body.add_geom(size=[SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba=BROWN)\n",
"\n",
" # Place debris on the tile\n",
" for i in range(10):\n",
" for j in range(10):\n",
" # draw on xy plane\n",
" drawing = np.random.normal(size=(4, 2))\n",
" drawing /= np.linalg.norm(drawing, axis=1, keepdims=True)\n",
" z = np.zeros((drawing.shape[0], 1))\n",
" # Add z value to drawing\n",
" base = np.concatenate((drawing, z), axis=1)\n",
" # Extrude drawing\n",
" z_extrusion = np.full((drawing.shape[0], 1), THICKNESS * 4)\n",
" top = np.concatenate((drawing, z_extrusion), axis=1)\n",
" # Combine to get a mesh\n",
" mesh = np.vstack((base, top))\n",
"\n",
" # Create body and add the mesh to the geom of the body\n",
" spec.add_mesh(name=f'd{i}_{j}_{name}', uservert=mesh.flatten())\n",
" pos=[grid_loc[0] + x_beginning + i * STEP,\n",
" grid_loc[1] + y_beginning - j * STEP,\n",
" 0.2]\n",
"\n",
" body = spec.worldbody.add_body(pos=pos, name=f'd{i}_{j}_{name}', mass=1)\n",
" body.add_geom(type=mj.mjtGeom.mjGEOM_MESH, meshname=f'd{i}_{j}_{name}',\n",
" rgba=RED)\n",
" body.add_freejoint()\n",
"\n",
"render_tile(debris)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "BqTG_K8uS3v_"
},
"outputs": [],
"source": [
"# @title Boxy Terrain\n",
"def boxy_terrain(spec=None, grid_loc=[0, 0], name='boxy_terrain'):\n",
" SQUARE_LENGTH = 2\n",
" CUBE_LENGTH = 0.05\n",
" GRID_SIZE = int(SQUARE_LENGTH / CUBE_LENGTH)\n",
" STEP = CUBE_LENGTH * 2\n",
" BROWN = [0.460, 0.362, 0.216, 1.0]\n",
"\n",
" if spec == None:\n",
" spec=mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" # Create tile\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
"\n",
" x_beginning = -SQUARE_LENGTH + CUBE_LENGTH\n",
" y_beginning = SQUARE_LENGTH - CUBE_LENGTH\n",
" for i in range(GRID_SIZE):\n",
" for j in range(GRID_SIZE):\n",
" body.add_geom(\n",
" pos=[x_beginning + i * STEP ,\n",
" y_beginning - j * STEP ,\n",
" random.randint(-1, 1) * CUBE_LENGTH\n",
" ],\n",
" size=[CUBE_LENGTH] * 3,\n",
" rgba=BROWN\n",
" )\n",
"\n",
"render_tile(boxy_terrain)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "gfQ1CafnS7hs"
},
"outputs": [],
"source": [
"# @title Box (Extrusion | Cut)\n",
"def box_extrusions(spec=None, grid_loc=[0, 0], complex=False,\n",
" name='box_extrusions'):\n",
" # Warning! complex sometimes leads to creation of holes\n",
" SQUARE_LENGTH = 2\n",
" CUBE_LENGTH = 0.05\n",
" GRID_SIZE = int(SQUARE_LENGTH / CUBE_LENGTH)\n",
" STEP = CUBE_LENGTH * 2\n",
" BROWN = [0.460, 0.362, 0.216, 1.0]\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" # Create tile\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
"\n",
" x_beginning = -SQUARE_LENGTH + CUBE_LENGTH\n",
" y_beginning = SQUARE_LENGTH - CUBE_LENGTH\n",
"\n",
" # Create initial grid and store geoms ref\n",
" grid = [[ 0 for _ in range(GRID_SIZE)] for _ in range(GRID_SIZE)]\n",
" for i in range(GRID_SIZE):\n",
" for j in range(GRID_SIZE):\n",
" ref = body.add_geom(\n",
" pos=[x_beginning + i * STEP, y_beginning - j * STEP, 0],\n",
" size=[CUBE_LENGTH] * 3,\n",
" rgba = BROWN\n",
" )\n",
" grid[i][j] = ref\n",
"\n",
" # Extrude or Cut operation using the boxes\n",
" for _ in range(random.randint(4, 50)):\n",
" box = None\n",
" while box == None:\n",
" # Create a box\n",
" start = (random.randint(0, GRID_SIZE - 2), random.randint(0, GRID_SIZE - 2))\n",
" dim = (random.randint(0, GRID_SIZE - 2), random.randint(0, GRID_SIZE-2))\n",
" # Make suer box is valid\n",
" if start[0] + dim [0] < len(grid) and start[1] + dim [1] < len(grid):\n",
" box = {\"start\":start, \"dim\":dim}\n",
"\n",
" # Use the box to Cut or Extrude\n",
" operation = random.choice([1, -1])\n",
" start = box[\"start\"]\n",
" dim = box[\"dim\"]\n",
" for i in range(start[0], dim[0]):\n",
" for j in range(start[1], dim[1]):\n",
" tile = grid[i][j]\n",
" if complex:\n",
" tile.pos[2] += operation * CUBE_LENGTH\n",
" else:\n",
" tile.pos[2] = operation * CUBE_LENGTH\n",
"\n",
"render_tile(box_extrusions)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "e1PEGQ37S_ee"
},
"outputs": [],
"source": [
"# @title Heightfield\n",
"def h_field(spec=None, grid_loc=[0, 0], name='h_field'):\n",
" SQUARE_LENGTH = 2\n",
" HEIGHT = 0.1\n",
" BROWN_RGBA = [0.460, 0.362, 0.216, 1.0]\n",
"\n",
" if spec is None:\n",
" spec = mj.MjSpec()\n",
"\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=name,\n",
" size=[SQUARE_LENGTH, SQUARE_LENGTH,\n",
" HEIGHT, HEIGHT/10],\n",
" nrow=noise.shape[0],\n",
" ncol=noise.shape[1],\n",
" userdata=noise.flatten())\n",
"\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" body.add_geom(type=mj.mjtGeom.mjGEOM_HFIELD, hfieldname=name,\n",
" rgba=BROWN_RGBA)\n",
"\n",
"render_tile(h_field)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "UwDVd6aUTRKF"
},
"outputs": [],
"source": [
"# @title Floating platform\n",
"def floating_platform(spec=None, gird_loc=[0, 0, 0], name='platform'):\n",
" PLATFORM_LENGTH = 0.5\n",
" WIDTH = 0.12\n",
" INWARD_OFFSET = 0.008\n",
" THICKNESS = 0.005\n",
" SIZE = [PLATFORM_LENGTH, WIDTH, THICKNESS]\n",
" TENDON_LENGTH = 0.5\n",
" Z_OFFSET = 0.1\n",
"\n",
" GOLD = [0.850, 0.838, 0.119, 1]\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" # Platform with sites\n",
" gird_loc[2] += Z_OFFSET\n",
" platform = spec.worldbody.add_body(pos=gird_loc, name=name)\n",
" platform.add_geom(size=SIZE, rgba=GOLD)\n",
" platform.add_freejoint()\n",
"\n",
" for x_dir in [-1, 1]:\n",
" for y_dir in [-1, 1]:\n",
" # Add site to world\n",
" vector = np.array([x_dir * PLATFORM_LENGTH,\n",
" y_dir * (WIDTH - INWARD_OFFSET)])\n",
" x_w = gird_loc[0] + vector[0]\n",
" y_w = gird_loc[1] + vector[1]\n",
" z_w = gird_loc[2] + TENDON_LENGTH\n",
" # Rotate sites by theta\n",
" spec.worldbody.add_site(name=f'{name}_hook_{x_dir}_{y_dir}',\n",
" pos=[ x_w, y_w, z_w],\n",
" size=[0.01, 0, 0])\n",
" # Add site to platform\n",
" x_p = x_dir * PLATFORM_LENGTH\n",
" y_p = y_dir * (WIDTH - INWARD_OFFSET)\n",
" platform.add_site(name=f'{name}_anchor_{x_dir}_{y_dir}',\n",
" pos=[ x_p, y_p, THICKNESS * 2],\n",
" size=[0.01, 0, 0])\n",
"\n",
" # Connect tendon to sites\n",
" thread = spec.add_tendon(name=f'{name}_thread_{x_dir}_{y_dir}',\n",
" limited=True,\n",
" range=[0, TENDON_LENGTH], width=0.01 )\n",
" thread.wrap_site(f'{name}_hook_{x_dir}_{y_dir}')\n",
" thread.wrap_site(f'{name}_anchor_{x_dir}_{y_dir}')\n",
"\n",
"render_tile(floating_platform, cam_distance=2, cam_elevation=-20)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "Wb4GmZewTSVI"
},
"outputs": [],
"source": [
"# @title Simple stairs\n",
"def simple_suspended_stair(spec=None, grid_loc=[0, 0], num_stair=20,\n",
" name=\"simple_suspended_stair\"):\n",
" BROWN= [0.460, 0.362, 0.216, 1.0]\n",
" SQUARE_LENGTH = 2\n",
" THICKNESS = 0.05\n",
" OFFSET_Y = -4/5 * SQUARE_LENGTH\n",
"\n",
" V_STEP = 0.076\n",
" H_STEP = 0.12\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" # Create tile\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" body.add_geom(size=[SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba=BROWN)\n",
"\n",
" # Create Stairs\n",
" for i in range(num_stair):\n",
" floating_platform(spec,[grid_loc[0],\n",
" OFFSET_Y + grid_loc[1] + i * 2 * H_STEP,\n",
" i * V_STEP],\n",
" name =f'{name}_p_{i}')\n",
"\n",
"render_tile(simple_suspended_stair,cam_distance=7, cam_elevation=-30)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "AaoLjrKeTXu-"
},
"outputs": [],
"source": [
"# @title Sinusoidal stairs\n",
"def sin_suspended_stair(spec, grid_loc=[0, 0], num_stair=40,\n",
" name=\"sin_suspended_stair\"):\n",
" BROWN = [0.460, 0.362, 0.216, 1.0]\n",
" SQUARE_LENGTH = 2\n",
" THICKNESS = 0.05\n",
" OFFSET_Y = -4/5 * SQUARE_LENGTH\n",
"\n",
" V_STEP = 0.076\n",
" H_STEP = 0.12\n",
" AMPLITUDE = 0.2\n",
" FREQUENCY = 0.5\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" # Plane\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" body.add_geom(size=[SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba=BROWN)\n",
"\n",
" for i in range(num_stair):\n",
" x_step = AMPLITUDE * np.sin(2 * np.pi * FREQUENCY * (i * H_STEP))\n",
" floating_platform(spec, [grid_loc[0] + x_step,\n",
" OFFSET_Y + grid_loc[1] + i * 2 * H_STEP,\n",
" i * V_STEP],\n",
" name=f'{name}_p_{i}')\n",
"\n",
"render_tile(sin_suspended_stair,cam_distance=7, cam_elevation=-30)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "2_Ux8h-PTqbW"
},
"outputs": [],
"source": [
"# @title Floating platform for circular stair\n",
"def floating_platform_for_circular_stair(spec=None, gird_loc=[0, 0, 0] ,theta=0,\n",
" name='platform'):\n",
" PLATFORM_LENGTH = 0.5\n",
" TENDON_LENGTH = 0.5\n",
" WIDTH = 0.12/4 # Platform (body) is made of 4 separate geoms\n",
" THICKNESS = 0.005\n",
" SIZE = [PLATFORM_LENGTH, WIDTH, THICKNESS]\n",
" Z_OFFSET = 0.1\n",
"\n",
" GOLD = [0.850, 0.838, 0.119, 1]\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
" spec.compiler.degree = False\n",
"\n",
" # Platform with sites\n",
" gird_loc[2] += Z_OFFSET\n",
" platform = spec.worldbody.add_body(pos=gird_loc, name=name, euler=[0, 0, theta])\n",
" platform.add_geom(pos=[0, 0, 0] , size=SIZE, euler=[0, 0, 0],rgba=GOLD)\n",
" platform.add_geom(pos=[0, 0.02, 0], size=SIZE, euler=[0, 0, 0.05],rgba=GOLD)\n",
" platform.add_geom(pos=[0, 0.05, 0], size=SIZE, euler=[0, 0, 0.1],rgba=GOLD)\n",
" platform.add_geom(pos=[0, 0.08, 0], size=SIZE, euler=[0, 0, 0.15],rgba=GOLD)\n",
" platform.add_freejoint()\n",
"\n",
" for i, x_dir in enumerate([-1, 1]):\n",
" for j, y_dir in enumerate([-1, 1]):\n",
" # Rotate sites by theta\n",
" rotation_matrix = np.array([[np.cos(-theta), -np.sin(-theta)],\n",
" [np.sin(-theta), np.cos(-theta)]])\n",
" vector = np.array([x_dir * PLATFORM_LENGTH, y_dir * WIDTH ])\n",
" if i + j == 2:\n",
" vector = np.array([x_dir * PLATFORM_LENGTH, y_dir * 6 * WIDTH ])\n",
" vector = np.dot(vector , rotation_matrix)\n",
" x_w = gird_loc[0] + vector[0]\n",
" y_w = gird_loc[1] + vector[1]\n",
" z_w = gird_loc[2] + TENDON_LENGTH\n",
"\n",
" # Add site to world\n",
" spec.worldbody.add_site(name=f'{name}_hook_{x_dir}_{y_dir}',\n",
" pos=[ x_w, y_w, z_w],\n",
" size=[0.01, 0, 0])\n",
" # Add site to platform\n",
" x_p = x_dir * PLATFORM_LENGTH\n",
" y_p = y_dir * WIDTH\n",
" if i + j == 2:\n",
" y_p = y_dir * 6 * WIDTH\n",
" platform.add_site(name=f'{name}_anchor_{x_dir}_{y_dir}',\n",
" pos=[x_p, y_p, THICKNESS * 2],\n",
" size=[0.01, 0, 0])\n",
"\n",
" # Connect tendon to sites\n",
" thread = spec.add_tendon(name=f'{name}_thread_{x_dir}_{y_dir}', limited=True,\n",
" range=[0, TENDON_LENGTH], width=0.01 )\n",
" thread.wrap_site(f'{name}_hook_{x_dir}_{y_dir}')\n",
" thread.wrap_site(f'{name}_anchor_{x_dir}_{y_dir}')\n",
"\n",
"render_tile(floating_platform_for_circular_stair,cam_distance=2, cam_elevation=-40)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "UXSDdVF8TwmQ"
},
"outputs": [],
"source": [
"# @title Circular stairs\n",
"def circular_stairs(spec, grid_loc=[0, 0], num_stair=60, name=\"circular_stairs\"):\n",
" BROWN_RGBA = [0.460, 0.362, 0.216, 1.0]\n",
" SQUARE_LENGTH = 2\n",
" THICKNESS = 0.05\n",
"\n",
" RADIUS = 1.5\n",
" V_STEP = 0.076\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
" spec.compiler.degree = False\n",
"\n",
" # Plane\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" body.add_geom(size = [SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba = BROWN_RGBA )\n",
"\n",
" theta_step = 2 * np.pi / num_stair\n",
" for i in range(num_stair):\n",
" theta = i * theta_step\n",
" x = grid_loc[0] + RADIUS * np.cos(theta)\n",
" y = grid_loc[1] + RADIUS * np.sin(theta)\n",
" z = i * V_STEP\n",
"\n",
" floating_platform_for_circular_stair(spec, [x, y, z], theta=theta, name=f'{name}_p_{i}')\n",
"\n",
"render_tile(circular_stairs,cam_distance=12, cam_elevation=-30)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "lsDky1KjT30W"
},
"outputs": [],
"source": [
"# @title Tile Generator\n",
"def add_tile(spec=None, grid_loc=[0, 0]):\n",
" if spec is None:\n",
" spec = mj.MjSpec()\n",
"\n",
" tile_type = random.randint(0, 9)\n",
"\n",
" if tile_type == 0:\n",
" debris_with_simple_geoms(spec, grid_loc, name=f\"plane_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 1:\n",
" stairs(spec, grid_loc, name=f\"stairs_up_{grid_loc[0]}_{grid_loc[1]}\",direction=1)\n",
" elif tile_type == 2:\n",
" stairs(spec, grid_loc, name=f\"stairs_down_{grid_loc[0]}_{grid_loc[1]}\",direction=-1)\n",
" elif tile_type == 3:\n",
" debris(spec, grid_loc, name=f\"debris_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 4:\n",
" box_extrusions(spec, grid_loc, name=f\"box_extrusions_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 5:\n",
" boxy_terrain(spec, grid_loc, name=f\"boxy_terrain_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 6:\n",
" h_field(spec, grid_loc, name=f\"h_field_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 7:\n",
" simple_suspended_stair(spec, grid_loc, name=f\"sss_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 8:\n",
" sin_suspended_stair(spec, grid_loc, name=f\"sinss_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 9:\n",
" circular_stairs(spec, grid_loc, name=f\"circular_s_{grid_loc[0]}_{grid_loc[1]}\")\n",
" return spec"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "P2F0ObC2T8w8"
},
"outputs": [],
"source": [
"# @title Generate Terrain\n",
"arena_xml = \"\"\"\n",
"<mujoco>\n",
" <visual>\n",
" <headlight diffuse=\".5 .5 .5\" specular=\"1 1 1\"/>\n",
" <global elevation=\"-10\" offwidth=\"2048\" offheight=\"1536\"/>\n",
" <quality shadowsize=\"8192\"/>\n",
" </visual>\n",
"\n",
" <asset>\n",
" <texture type=\"skybox\" builtin=\"gradient\" rgb1=\".5 .5 .5\" rgb2=\"0 0 0\" width=\"10\" height=\"10\"/>\n",
" <texture type=\"2d\" name=\"groundplane\" builtin=\"checker\" mark=\"edge\" rgb1=\"1 1 1\" rgb2=\"1 1 1\" markrgb=\"0 0 0\" width=\"300\" height=\"300\"/>\n",
" <material name=\"groundplane\" texture=\"groundplane\" texuniform=\"true\" texrepeat=\"5 5\" reflectance=\"0.3\"/>\n",
" </asset>\n",
"\n",
" <worldbody>\n",
" </worldbody>\n",
"</mujoco>\n",
"\"\"\"\n",
"\n",
"spec = mj.MjSpec.from_string(arena_xml)\n",
"\n",
"spec.option.enableflags |= mj.mjtEnableBit.mjENBL_OVERRIDE\n",
"spec.option.enableflags |= mj.mjtEnableBit.mjENBL_MULTICCD\n",
"spec.option.timestep = 0.0001\n",
"spec.compiler.degree = False\n",
"\n",
"main = spec.default\n",
"main.geom.solref = [0.001, 1]\n",
"\n",
"# Add lights\n",
"for x in [-1, 1]:\n",
" for y in [-1, 1]:\n",
" spec.worldbody.add_light(pos=[x, y, 40], dir=[-x, -y, -15])\n",
"\n",
"SQUARE_LENGTH = 2\n",
"for i in range(-2, 2):\n",
" for j in range(-2, 2):\n",
" add_tile(spec=spec, grid_loc=[i * 2 * SQUARE_LENGTH, j * 2 * SQUARE_LENGTH])\n",
"\n",
"model = spec.compile()\n",
"data = mj.MjData(model)\n",
"\n",
"cam = mj.MjvCamera()\n",
"mj.mjv_defaultCamera(cam)\n",
"cam.lookat = [-2, 0, -2]\n",
"cam.distance = 18\n",
"cam.elevation = -30\n",
"\n",
"with mj.Renderer(model, 720, 1280) as renderer:\n",
" mj.mj_forward(model, data)\n",
" renderer.update_scene(data,cam)\n",
" media.show_image(renderer.render())"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "IGd0uD64LdEJ"
},
"source": [
"# Model editing"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"cellView": "form",
"id": "223KzKAzLdEJ"
},
"outputs": [],
"source": [
"# @title Get resources\n",
"\n",
"# Get Models\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",
"print('Getting MuJoCo Fly and Franka XML description from GitHub:')\n",
"!git clone https://github.com/google-deepmind/mujoco_menagerie\n",
"fly_file = 'mujoco_menagerie/flybody/fruitfly.xml'\n",
"franka_file = 'mujoco_menagerie/franka_fr3/fr3.xml'\n",
"\n",
"# Camera options\n",
"cam = mj.MjvCamera()\n",
"mj.mjv_defaultCamera(cam)\n",
"cam.elevation = -10\n",
"cam.lookat = [0, 0, 1]\n",
"cam.distance = 4\n",
"cam.azimuth = 135\n",
"\n",
"# Arena\n",
"arena_xml = \"\"\"\n",
"<mujoco>\n",
" <visual>\n",
" <headlight diffuse=\".5 .5 .5\" specular=\"1 1 1\"/>\n",
" <global offwidth=\"2048\" offheight=\"1536\"/>\n",
" <quality shadowsize=\"8192\"/>\n",
" </visual>\n",
"\n",
" <asset>\n",
" <texture type=\"skybox\" builtin=\"gradient\" rgb1=\"1 1 1\" rgb2=\"1 1 1\" width=\"10\" height=\"10\"/>\n",
" <texture type=\"2d\" name=\"groundplane\" builtin=\"checker\" mark=\"edge\" rgb1=\"1 1 1\" rgb2=\"1 1 1\" markrgb=\"0 0 0\" width=\"400\" height=\"400\"/>\n",
" <material name=\"groundplane\" texture=\"groundplane\" texrepeat=\"45 45\" reflectance=\"0\"/>\n",
" </asset>\n",
"\n",
" <worldbody>\n",
" <geom name=\"floor\" size=\"10 10 0.1\" type=\"plane\" material=\"groundplane\"/>\n",
" </worldbody>\n",
"</mujoco>\n",
"\"\"\"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "eGgXNjQ8LdEK"
},
"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": "Len0o_idLdEK"
},
"outputs": [],
"source": [
"#@title Traversing the spec {vertical-output: true}\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": "GeiFFBYxLdEK"
},
"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": "eiRXgh9OLdEK"
},
"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.delete(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": "kuTWD415LdEK"
},
"source": [
"Let us load the humanoid model and inspect it."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "5d1wmQM2LdEK"
},
"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": "38PXB1rWLdEK"
},
"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": "0eaNq0Q7LdEK"
},
"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.delete(arm_left)\n",
"spec.delete(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": "HfhxL2EqLdEK"
},
"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": "uS4LGbI7LdEK"
},
"outputs": [],
"source": [
"#@title Humanoid with Franka arm {vertical-output: true}\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.delete(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": "CWXYy_1uLdEK"
},
"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": "UwWDD-NHLdEK"
},
"outputs": [],
"source": [
"#@title Imported actuators {vertical-output: true}\n",
"\n",
"for actuator in spec.actuators:\n",
" print(actuator.name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "hDvt3vcxLdEK"
},
"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": "oPPFbWawLdEK"
},
"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.delete(keys[-1])\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)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "PML2pxYgLdEK"
},
"source": [
"## Model scaling"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "pcUNLmQBLdEK"
},
"outputs": [],
"source": [
"#@title Uniformly scale humanoid {vertical-output: true}\n",
"\n",
"def scale_spec(spec, scale):\n",
" scaled_spec = spec.copy()\n",
" # Traverse the kinematic tree, scaling all geoms\n",
" def scale_bodies(parent, scale=1.0):\n",
" body = parent.first_body()\n",
" while body:\n",
" if body.pos is not None:\n",
" body.pos = body.pos * scale\n",
" for geom in body.geoms:\n",
" geom.fromto = geom.fromto * scale\n",
" geom.size = geom.size * scale\n",
" if geom.pos is not None:\n",
" geom.pos = geom.pos * scale\n",
" scale_bodies(body, scale)\n",
" body = parent.next_body(body)\n",
"\n",
" scale_bodies(scaled_spec.body('world'), scale)\n",
" return scaled_spec\n",
"\n",
"spec = mj.MjSpec.from_string(arena_xml)\n",
"humanoid = mj.MjSpec.from_file(humanoid_file)\n",
"small_humanoid = scale_spec(humanoid, 0.75)\n",
"large_humanoid = scale_spec(humanoid, 1.25)\n",
"\n",
"# Create a line-up of humanoids\n",
"frame = spec.worldbody.add_frame(pos=[-1, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(humanoid.body('torso'), str(0))\n",
"\n",
"frame = spec.worldbody.add_frame(pos=[0, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(small_humanoid.body('torso'), str(1))\n",
"\n",
"frame = spec.worldbody.add_frame(pos=[1, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2] )\n",
"frame.attach_body(large_humanoid.body('torso'), str(2))\n",
"\n",
"\n",
"spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n",
" targetbody='1torso', 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, camera=cam)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "RYbaTPNmLdEK"
},
"source": [
"We can scale the size of a model by traversing the kinematic tree and applying the scale to the relevant geoms. Above we can see humanoids of three different sizes."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "u-ejx8lKLdEK"
},
"outputs": [],
"source": [
"# @title Scaling actuator forces {vertical-output: true}\n",
"\n",
"def scale_spec(spec, scale, scale_actuators=False):\n",
" scaled_spec = spec.copy()\n",
" # Traverse the kinematic tree, scaling all geoms\n",
" def scale_bodies(parent, scale=1.0):\n",
" body = parent.first_body()\n",
" while body:\n",
" if body.pos is not None:\n",
" body.pos = body.pos * scale\n",
" for geom in body.geoms:\n",
" geom.fromto = geom.fromto * scale\n",
" geom.size = geom.size * scale\n",
" if geom.pos is not None:\n",
" geom.pos = geom.pos * scale\n",
" scale_bodies(body, scale)\n",
" body = parent.next_body(body)\n",
"\n",
" if scale_actuators:\n",
" # scale gear\n",
" for actuator in scaled_spec.actuators:\n",
" # scale the actuator gear by (scale ** 2),\n",
" # this is because muscle force-generating capacity\n",
" # scales with the cross-sectional area of the muscle\n",
" actuator.gear = actuator.gear * scale * scale\n",
"\n",
" # scale the z-position of the humanoid for all keypoints\n",
" for keypoint in scaled_spec.keys:\n",
" qpos = keypoint.qpos\n",
" qpos[2] = qpos[2] * scale\n",
" keypoint.qpos = qpos\n",
" keypoint.qpos[2] = keypoint.qpos[2] * scale\n",
"\n",
" scale_bodies(scaled_spec.body('world'), scale)\n",
" return scaled_spec\n",
"\n",
"# Create specs\n",
"scale = 0.6\n",
"spec = mj.MjSpec.from_string(arena_xml)\n",
"humanoid = mj.MjSpec.from_file(humanoid_file)\n",
"small_humanoid = scale_spec(humanoid, scale)\n",
"small_humanoid_actuators_scaled = scale_spec(humanoid, scale, True)\n",
"\n",
"# Create a line-up of humanoids\n",
"squat_qpos = []\n",
"\n",
"# Add unscaled humanoid\n",
"frame = spec.worldbody.add_frame(pos=[-1, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(humanoid.body('torso'), str(0))\n",
"# Record squat pose\n",
"humanoid_squat = humanoid.key('squat').qpos\n",
"humanoid_squat[:2] = frame.pos[:2]\n",
"humanoid_squat[3:7] = frame.quat\n",
"squat_qpos.append(humanoid_squat)\n",
"\n",
"# Add small humanoid\n",
"frame = spec.worldbody.add_frame(pos=[0, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(small_humanoid.body('torso'), str(1))\n",
"# Record squat pose\n",
"humanoid_squat = small_humanoid.key('squat').qpos\n",
"humanoid_squat[:2] = frame.pos[:2]\n",
"humanoid_squat[3:7] = frame.quat\n",
"squat_qpos.append(humanoid_squat)\n",
"\n",
"# Add small humanoid with scaled actuators\n",
"frame = spec.worldbody.add_frame(pos=[1, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2] )\n",
"frame.attach_body(small_humanoid_actuators_scaled.body('torso'), str(2))\n",
"# Record squat pose\n",
"humanoid_squat = small_humanoid_actuators_scaled.key('squat').qpos\n",
"humanoid_squat[:2] = frame.pos[:2]\n",
"humanoid_squat[3:7] = frame.quat\n",
"squat_qpos.append(humanoid_squat)\n",
"squat_qpos = np.concatenate(squat_qpos)\n",
"\n",
"spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n",
" targetbody='1torso', diffuse=[.8, .8, .8],\n",
" specular=[0.3, 0.3, 0.3], pos=[0, -6, 4], cutoff=30)\n",
"model = spec.compile()\n",
"\n",
"# Initialize to squat position\n",
"data = mj.MjData(model)\n",
"data.qpos = squat_qpos\n",
"\n",
"# jumping motion\n",
"u_t = lambda t: 10.0 * t / duration\n",
"\n",
"# Simulate and display video.\n",
"duration = 2 # (seconds)\n",
"framerate = 30 # (Hz)\n",
"frames = []\n",
"\n",
"with mj.Renderer(model, 480, 640) as renderer:\n",
" while data.time < duration:\n",
" data.ctrl = u_t(data.time)\n",
" mj.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data, camera=cam)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"\n",
"media.show_video(frames, fps=framerate, height=400)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "uKhDI_IfLdEK"
},
"source": [
"We can also apply scaling to the actuators. In the humanoid case, scaling the geoms without scaling the `gear` parameter for the actuators results in a humanoid that can jump higher proportional to its size."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "ovQIAxn7LdEK"
},
"outputs": [],
"source": [
"# @title Long-limbed humanoid {vertical-output: true}\n",
"\n",
"def scale_spec(spec, scale):\n",
" scaled_spec = spec.copy()\n",
" # Traverse the kinematic tree, scaling all geoms\n",
" def scale_bodies(parent, scale=1.0):\n",
" if parent is not None:\n",
" for geom in parent.geoms:\n",
" # Only scale fromto, not size to scale length of capsules\n",
" geom.fromto = geom.fromto * scale\n",
" if geom.pos is not None:\n",
" geom.pos = geom.pos * scale\n",
" body = parent.first_body()\n",
" while body:\n",
" if body.pos is not None:\n",
" body.pos = body.pos * scale\n",
" scale_bodies(body, scale)\n",
" body = parent.next_body(body)\n",
"\n",
" # Scale all the limbs\n",
" scale_bodies(scaled_spec.body('upper_arm_right'), scale)\n",
" scale_bodies(scaled_spec.body('upper_arm_left'), scale)\n",
" scale_bodies(scaled_spec.body('thigh_right'), scale)\n",
" scale_bodies(scaled_spec.body('thigh_left'), scale)\n",
" return scaled_spec\n",
"\n",
"spec = mj.MjSpec.from_string(arena_xml)\n",
"humanoid = mj.MjSpec.from_file(humanoid_file)\n",
"small_humanoid = scale_spec(humanoid, 1.25)\n",
"large_humanoid = scale_spec(humanoid, 2)\n",
"\n",
"# Create a line-up of humanoids by attaching\n",
"frame = spec.worldbody.add_frame(pos=[-1, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(humanoid.body('torso'), str(0), str(0))\n",
"\n",
"frame = spec.worldbody.add_frame(pos=[0, 0, 0.2],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(small_humanoid.body('torso'), str(0), str(1))\n",
"\n",
"frame = spec.worldbody.add_frame(pos=[1, 0, 0.8],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2] )\n",
"frame.attach_body(large_humanoid.body('torso'), str(0), str(2))\n",
"\n",
"\n",
"spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n",
" targetbody='0torso1', diffuse=[.8, .8, .8],\n",
" specular=[0.3, 0.3, 0.3], pos=[0, -6, 4], cutoff=30)\n",
"model = spec.compile()\n",
"\n",
"# camera options\n",
"render(model, height=400, camera=cam)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8o1daXIOLdEK"
},
"source": [
"We can also apply scaling to the model non-uniformly. In this instance we scale the humanoid to have long limbs, by only applying the scale to the length of the capsule geoms for the arms, legs and feet."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "UBSo2nfQLdEK"
},
"outputs": [],
"source": [
"# @title Meshes {vertical-output: true}\n",
"\n",
"def scale_spec(spec, scale):\n",
" scaled_spec = spec.copy()\n",
" # scale all meshes\n",
" for mesh in scaled_spec.meshes:\n",
" if mesh.scale is None:\n",
" mesh.scale = np.ones(3)\n",
" mesh.scale = mesh.scale * scale\n",
"\n",
" # Traverse the kinematic tree\n",
" def scale_bodies(parent, scale=1.0):\n",
" if parent is not None:\n",
" for geom in parent.geoms:\n",
" if geom.pos is not None:\n",
" geom.pos = geom.pos * scale\n",
" body = parent.first_body()\n",
" while body:\n",
" if body.pos is not None:\n",
" body.pos = body.pos * scale\n",
" scale_bodies(body, scale)\n",
" body = parent.next_body(body)\n",
"\n",
" # Scale all the limbs\n",
" scale_bodies(scaled_spec.body('world'), scale)\n",
"\n",
" return scaled_spec\n",
"\n",
"spec = mj.MjSpec.from_string(arena_xml)\n",
"fly = mj.MjSpec.from_file(fly_file)\n",
"# Remove lights from fly so they are not duplicated in line-up\n",
"for light in fly.lights:\n",
" fly.delete(light)\n",
"\n",
"small_fly = scale_spec(fly, 1.25)\n",
"large_fly = scale_spec(fly, 2)\n",
"\n",
"# Create a line-up of flys by attaching\n",
"frame = spec.worldbody.add_frame(pos=[-1, 0, 0.25],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(fly.body('thorax'), str(0), str(0))\n",
"\n",
"frame = spec.worldbody.add_frame(pos=[0, 0, 0.25],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(small_fly.body('thorax'), str(0), str(1))\n",
"\n",
"frame = spec.worldbody.add_frame(pos=[1, 0, 0.25],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2] )\n",
"frame.attach_body(large_fly.body('thorax'), str(0), str(2))\n",
"\n",
"spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n",
" targetbody='0thorax1', 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, camera=cam)\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "qhXwxLe3LdEK"
},
"source": [
"# dm_control example"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "enuJ_YIqLdEK"
},
"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": "4p3P_dP8LdEK"
},
"outputs": [],
"source": [
"leg_model = \"\"\"\n",
"<mujoco>\n",
" <compiler angle=\"radian\"/>\n",
"\n",
" <default>\n",
" <joint damping=\"2\" type=\"hinge\"/>\n",
" <geom type=\"capsule\"/>\n",
" </default>\n",
"\n",
" <worldbody>\n",
" <body name=\"thigh\">\n",
" <joint name=\"hip\" axis=\"0 0 1\"/>\n",
" <body name=\"shin\">\n",
" <joint name=\"knee\" axis=\"0 1 0\"/>\n",
" </body>\n",
" </body>\n",
" </worldbody>\n",
"\n",
" <actuator>\n",
" <position joint=\"hip\" kp=\"10\" name=\"hip\"/>\n",
" <position joint=\"knee\" kp=\"10\" name=\"knee\"/>\n",
" </actuator>\n",
"</mujoco>\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": "Mqr8rXLILdEK"
},
"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": "1z2NBpAPLdEK"
},
"outputs": [],
"source": [
"BODY_RADIUS = 0.1\n",
"random_state = np.random.RandomState(42)\n",
"creature_model = \"\"\"\n",
"<mujoco>\n",
" <compiler angle=\"radian\"/>\n",
"\n",
" <worldbody>\n",
" <geom name=\"torso\" type=\"ellipsoid\" size=\"{} {} {}\"/>\n",
" </worldbody>\n",
"</mujoco>\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": "865FGuntLdEL"
},
"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": "2fPaSkgfLdEL"
},
"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": "tq5mKlc_LdEL"
},
"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": "i37FpwCeLdEL"
},
"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": "uFrvaih4LdEL"
},
"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": "FMW4l-fSLdEL"
},
"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."
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [
"sJFuNetilv4m"
],
"gpuClass": "premium",
"private_outputs": true,
"toc_visible": true
},
"gpuClass": "premium",
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}