Mujoco_WASM/python/mjspec.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "klCa_-bdiWv5"
   },
   "source": [
    "![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\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 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",
    "<!-- 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 succesful.\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",
    "\n",
    "# Graphics and plotting.\n",
    "print('Installing mediapy:')\n",
    "!command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n",
    "!pip install -q mediapy\n",
    "import mediapy as media\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# 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",
    "# Get MuJoCo's humanoid model and a Franka arm from the MuJoCo Menagerie.\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 Menagerie Franka XML description from GitHub:')\n",
    "!git clone https://github.com/google-deepmind/mujoco_menagerie\n",
    "franka_file = 'mujoco_menagerie/franka_fr3/fr3.xml'\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):\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": [
    "# Parsing XML to `mjSpec` and compiling to `mjModel`\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MYPbnl3mxvDj"
   },
   "source": [
    "Unlike `mj_loadXML` which combines parsing and compiling, when using `mjSpec`, parsing and compiling are separate, allowing for editing steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "id": "oummB7I7EfSq"
   },
   "outputs": [],
   "source": [
    "#@title A static model, from string {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": "code",
   "execution_count": 0,
   "metadata": {
    "id": "A6EHNulmFHFI"
   },
   "outputs": [],
   "source": [
    "#@title Building an `mjSpec` from scratch {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\", 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": "3N4YEIVt75_T"
   },
   "source": [
    "# Control example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TcQuv56BwaJf"
   },
   "source": [
    "A key feature of this library 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 is a fundamental change to the kinematic structure. The following snippets realise this scenario."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "id": "7C-hfbtj8nRV"
   },
   "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.find_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.find_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",
    "<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.find_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",
    "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.find_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.find_body('upper_arm_left')\n",
    "arm_right = spec.find_body('upper_arm_right')\n",
    "leg_left = spec.find_body('thigh_left')\n",
    "leg_right = spec.find_body('thigh_right')\n",
    "torso = spec.find_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",
    "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.find_body('upper_arm_right')\n",
    "torso = spec.find_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.find_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.find_body('head').first_geom().size = [\n",
    "        .18*np.random.uniform(), 0, 0]  # Randomize head size\n",
    "    humanoid.find_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.find_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.find_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": {
   "gpuClass": "premium",
   "private_outputs": true
  },
  "gpuClass": "premium",
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}