{ "cells": [ { "metadata": { "id": "MpkYHwCqk7W-" }, "cell_type": "markdown", "source": [ "# **MuJoCo** tutorial with `dm_control` Python bindings\n", "\n", "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/dm_control/blob/main/dm_control/mujoco/tutorial.ipynb)\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n" ] }, { "metadata": { "id": "_UbO9uhtBSX5" }, "cell_type": "markdown", "source": [ ">

Copyright 2021 The dm_control Authors.

\n", ">

Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0.

\n", ">

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

" ] }, { "metadata": { "id": "aThGKGp0cD76" }, "cell_type": "markdown", "source": [ "This notebook provides an overview tutorial of the **MuJoCo** physics simulator, using the `dm_control` Python bindings. It is similar to the notebook in `dm_control/tutorial.ipynb`, but focuses on teaching MuJoCo itself, rather than the additional features provided by the Python package. \n", "\n", "**A Colab runtime with GPU acceleration is required.** If you're using a CPU-only runtime, you can switch using the menu \"Runtime > Change runtime type\"." ] }, { "metadata": { "id": "YkBQUjm6gbGF" }, "cell_type": "markdown", "source": [ "" ] }, { "metadata": { "id": "YvyGCsgSCxHQ" }, "cell_type": "markdown", "source": [ "### Installing `dm_control` on Colab" ] }, { "metadata": { "id": "IbZxYDxzoz5R" }, "cell_type": "code", "source": [ "#@title Run to install MuJoCo and `dm_control`\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", "print('Installing dm_control...')\n", "!pip install -q dm_control>=1.0.39\n", "\n", "# Configure dm_control to use the EGL rendering backend (requires GPU)\n", "%env MUJOCO_GL=egl\n", "\n", "print('Checking that the dm_control installation succeeded...')\n", "try:\n", " from dm_control import suite\n", " env = suite.load('cartpole', 'swingup')\n", " pixels = env.physics.render()\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", "else:\n", " del suite, pixels\n", "print('dm_control installation succeeded.')" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "wtDN43hIJh2C" }, "cell_type": "markdown", "source": [ "# Imports\n", "\n", "Run both of these cells:" ] }, { "metadata": { "id": "T5f4w3Kq2X14" }, "cell_type": "code", "source": [ "#@title All `dm_control` imports required for this tutorial\n", "\n", "# The basic mujoco wrapper.\n", "from dm_control import mujoco\n", "\n", "# Access to enums and MuJoCo library functions.\n", "from dm_control.mujoco.wrapper.mjbindings import enums\n", "from dm_control.mujoco.wrapper.mjbindings import mjlib\n", "\n", "# Composer high level imports\n", "from dm_control import composer\n", "from dm_control.composer.observation import observable\n", "from dm_control.composer import variation\n", "\n", "# Imports for Composer tutorial example\n", "from dm_control.composer.variation import distributions\n", "from dm_control.composer.variation import noises\n", "from dm_control.locomotion.arenas import floors\n", "\n", "# Control Suite\n", "from dm_control import suite\n", "\n", "# Run through corridor example\n", "from dm_control.locomotion.walkers import cmu_humanoid\n", "from dm_control.locomotion.arenas import corridors as corridor_arenas\n", "from dm_control.locomotion.tasks import corridors as corridor_tasks\n", "\n", "# Soccer\n", "from dm_control.locomotion import soccer\n", "\n", "# Manipulation\n", "from dm_control import manipulation" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "gKc1FNhKiVJX" }, "cell_type": "code", "source": [ "#@title Other imports and helper functions\n", "\n", "# General\n", "import copy\n", "import os\n", "import time\n", "import itertools\n", "from IPython.display import clear_output\n", "import numpy as np\n", "\n", "# Graphics-related\n", "import matplotlib\n", "import matplotlib.animation as animation\n", "import matplotlib.pyplot as plt\n", "from IPython.display import HTML\n", "import PIL.Image\n", "# Internal loading of video libraries.\n", "\n", "# Use svg backend for figure rendering\n", "%config InlineBackend.figure_format = 'svg'\n", "\n", "# Font sizes\n", "SMALL_SIZE = 8\n", "MEDIUM_SIZE = 10\n", "BIGGER_SIZE = 12\n", "plt.rc('font', size=SMALL_SIZE) # controls default text sizes\n", "plt.rc('axes', titlesize=SMALL_SIZE) # fontsize of the axes title\n", "plt.rc('axes', labelsize=MEDIUM_SIZE) # fontsize of the x and y labels\n", "plt.rc('xtick', labelsize=SMALL_SIZE) # fontsize of the tick labels\n", "plt.rc('ytick', labelsize=SMALL_SIZE) # fontsize of the tick labels\n", "plt.rc('legend', fontsize=SMALL_SIZE) # legend fontsize\n", "plt.rc('figure', titlesize=BIGGER_SIZE) # fontsize of the figure title\n", "\n", "# Inline video helper function\n", "if os.environ.get('COLAB_NOTEBOOK_TEST', False):\n", " # We skip video generation during tests, as it is quite expensive.\n", " display_video = lambda *args, **kwargs: None\n", "else:\n", " def display_video(frames, framerate=30):\n", " height, width, _ = frames[0].shape\n", " dpi = 70\n", " orig_backend = matplotlib.get_backend()\n", " matplotlib.use('Agg') # Switch to headless 'Agg' to inhibit figure rendering.\n", " fig, ax = plt.subplots(1, 1, figsize=(width / dpi, height / dpi), dpi=dpi)\n", " matplotlib.use(orig_backend) # Switch back to the original backend.\n", " ax.set_axis_off()\n", " ax.set_aspect('equal')\n", " ax.set_position([0, 0, 1, 1])\n", " im = ax.imshow(frames[0])\n", " def update(frame):\n", " im.set_data(frame)\n", " return [im]\n", " interval = 1000/framerate\n", " anim = animation.FuncAnimation(fig=fig, func=update, frames=frames,\n", " interval=interval, blit=True, repeat=False)\n", " return HTML(anim.to_html5_video())\n", "\n", "# Seed numpy's global RNG so that cell outputs are deterministic. We also try to\n", "# use RandomState instances that are local to a single cell wherever possible.\n", "np.random.seed(42)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "jZXz9rPYGA-Y" }, "cell_type": "markdown", "source": [ "# Model definition, compilation and rendering\n", "\n" ] }, { "metadata": { "id": "MRBaZsf1d7Gb" }, "cell_type": "markdown", "source": [ "We begin by describing some basic concepts of the [MuJoCo](http://mujoco.org/) physics simulation library, but recommend the [official documentation](http://mujoco.org/book/) for details.\n", "\n", "Let's define a simple model with two geoms and a light." ] }, { "metadata": { "id": "ZS2utl59ZTsr" }, "cell_type": "code", "source": [ "#@title A static model {vertical-output: true}\n", "\n", "static_model = \"\"\"\n", "\n", " \n", " \n", " \n", " \n", " \n", "\n", "\"\"\"\n", "physics = mujoco.Physics.from_xml_string(static_model)\n", "pixels = physics.render()\n", "PIL.Image.fromarray(pixels)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "p4vPllljTJh8" }, "cell_type": "markdown", "source": [ "`static_model` is written in MuJoCo's XML-based [MJCF](http://www.mujoco.org/book/modeling.html) modeling language. The `from_xml_string()` method invokes the model compiler, which instantiates the library's internal data structures. These can be accessed via the `physics` object, see below." ] }, { "metadata": { "id": "MdUF2UYmR4TA" }, "cell_type": "markdown", "source": [ "## Adding DOFs and simulating, advanced rendering\n", "This is a perfectly legitimate model, but if we simulate it, nothing will happen except for time advancing. This is because this model has no degrees of freedom (DOFs). We add DOFs by adding **joints** to bodies, specifying how they can move with respect to their parents. Let us add a hinge joint and re-render, visualizing the joint axis." ] }, { "metadata": { "id": "R7zokzd_yeEg" }, "cell_type": "code", "source": [ "#@title A child body with a joint { vertical-output: true }\n", "\n", "swinging_body = \"\"\"\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", "\"\"\"\n", "physics = mujoco.Physics.from_xml_string(swinging_body)\n", "# Visualize the joint axis.\n", "scene_option = mujoco.wrapper.core.MjvOption()\n", "scene_option.flags[enums.mjtVisFlag.mjVIS_JOINT] = True\n", "pixels = physics.render(scene_option=scene_option)\n", "PIL.Image.fromarray(pixels)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "INOGGV0PTQus" }, "cell_type": "markdown", "source": [ "The things that move (and which have inertia) are called *bodies*. The body's child `joint` specifies how that body can move with respect to its parent, in this case `box_and_sphere` with respect to the `worldbody`. \n", "\n", "Note that the body's frame is **rotated** with an `euler` directive, and its children, the geoms and the joint, rotate with it. This is to emphasize the local-to-parent-frame nature of position and orientation directives in MJCF.\n", "\n", "Let's make a video, to get a sense of the dynamics and to see the body swinging under gravity." ] }, { "metadata": { "id": "Z_57VMUDpGrj" }, "cell_type": "code", "source": [ "#@title Making a video {vertical-output: true}\n", "\n", "duration = 2 # (seconds)\n", "framerate = 30 # (Hz)\n", "\n", "# Visualize the joint axis\n", "scene_option = mujoco.wrapper.core.MjvOption()\n", "scene_option.flags[enums.mjtVisFlag.mjVIS_JOINT] = True\n", "\n", "# Simulate and display video.\n", "frames = []\n", "physics.reset() # Reset state and time\n", "while physics.data.time < duration:\n", " physics.step()\n", " if len(frames) < physics.data.time * framerate:\n", " pixels = physics.render(scene_option=scene_option)\n", " frames.append(pixels)\n", "display_video(frames, framerate)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "yYvS1UaciMX_" }, "cell_type": "markdown", "source": [ "Note how we collect the video frames. Because physics simulation timesteps are generally much smaller than framerates (the default timestep is 2ms), we don't render after each step." ] }, { "metadata": { "id": "nQ8XOnRQx7T1" }, "cell_type": "markdown", "source": [ "## Rendering options\n", "\n", "Like joint visualisation, additional rendering options are exposed as parameters to the `render` method." ] }, { "metadata": { "id": "AQITZiIgx7T2" }, "cell_type": "code", "source": [ "#@title Enable transparency and frame visualization {vertical-output: true}\n", "\n", "scene_option = mujoco.wrapper.core.MjvOption()\n", "scene_option.frame = enums.mjtFrame.mjFRAME_GEOM\n", "scene_option.flags[enums.mjtVisFlag.mjVIS_TRANSPARENT] = True\n", "pixels = physics.render(scene_option=scene_option)\n", "PIL.Image.fromarray(pixels)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "PDDgY48vx7T6" }, "cell_type": "code", "source": [ "#@title Depth rendering {vertical-output: true}\n", "\n", "# depth is a float array, in meters.\n", "depth = physics.render(depth=True)\n", "# Shift nearest values to the origin.\n", "depth -= depth.min()\n", "# Scale by 2 mean distances of near rays.\n", "depth /= 2*depth[depth <= 1].mean()\n", "# Scale to [0, 255]\n", "pixels = 255*np.clip(depth, 0, 1)\n", "PIL.Image.fromarray(pixels.astype(np.uint8))" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "PNwiIrgpx7T8" }, "cell_type": "code", "source": [ "#@title Segmentation rendering {vertical-output: true}\n", "\n", "seg = physics.render(segmentation=True)\n", "# Display the contents of the first channel, which contains object\n", "# IDs. The second channel, seg[:, :, 1], contains object types.\n", "geom_ids = seg[:, :, 0]\n", "# Infinity is mapped to -1\n", "geom_ids = geom_ids.astype(np.float64) + 1\n", "# Scale to [0, 1]\n", "geom_ids = geom_ids / geom_ids.max()\n", "pixels = 255*geom_ids\n", "PIL.Image.fromarray(pixels.astype(np.uint8))\n" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "uCJQlv3cQcJQ" }, "cell_type": "code", "source": [ "#@title Projecting from world to camera coordinates {vertical-output: true}\n", "\n", "# Get the world coordinates of the box corners\n", "box_pos = physics.named.data.geom_xpos['red_box']\n", "box_mat = physics.named.data.geom_xmat['red_box'].reshape(3, 3)\n", "box_size = physics.named.model.geom_size['red_box']\n", "offsets = np.array([-1, 1]) * box_size[:, None]\n", "xyz_local = np.stack(list(itertools.product(*offsets))).T\n", "xyz_global = box_pos[:, None] + box_mat @ xyz_local\n", "\n", "# Camera matrices multiply homogenous [x, y, z, 1] vectors.\n", "corners_homogeneous = np.ones((4, xyz_global.shape[1]), dtype=float)\n", "corners_homogeneous[:3, :] = xyz_global\n", "\n", "# Get the camera matrix.\n", "camera = mujoco.Camera(physics)\n", "camera_matrix = camera.matrix\n", "\n", "# Project world coordinates into pixel space. See:\n", "# https://en.wikipedia.org/wiki/3D_projection#Mathematical_formula\n", "xs, ys, s = camera_matrix @ corners_homogeneous\n", "# x and y are in the pixel coordinate system.\n", "x = xs / s\n", "y = ys / s\n", "\n", "# Render the camera view and overlay the projected corner coordinates.\n", "pixels = camera.render()\n", "fig, ax = plt.subplots(1, 1)\n", "ax.imshow(pixels)\n", "ax.plot(x, y, '+', c='w')\n", "ax.set_axis_off()" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "R_jkaXd_9Eiu" }, "cell_type": "code", "source": [ "#@title Adding arbitrary 3D geometry {vertical-output: true}\n", "\n", "def get_geom_speed(physics, geom_name):\n", " \"\"\"Returns the speed of a geom.\"\"\"\n", " geom_vel = np.zeros(6)\n", " geom_type = mujoco.mjtObj.mjOBJ_GEOM\n", " geom_id = mujoco.mj_name2id(physics.model.ptr, geom_type, geom_name)\n", " mujoco.mj_objectVelocity(physics.model.ptr, physics.data.ptr,\n", " geom_type, geom_id, geom_vel, 0)\n", " return np.linalg.norm(geom_vel)\n", "\n", "def add_visual_capsule(scene, point1, point2, radius, rgba):\n", " \"\"\"Adds one capsule to an mjvScene.\"\"\"\n", " if scene.ngeom >= scene.maxgeom:\n", " return\n", " scene.ngeom += 1 # increment ngeom\n", " # initialise a new capsule, add it to the scene using mjv_connector\n", " mujoco.mjv_initGeom(scene.geoms[scene.ngeom-1],\n", " mujoco.mjtGeom.mjGEOM_CAPSULE, np.zeros(3),\n", " np.zeros(3), np.zeros(9), rgba.astype(np.float32))\n", " mujoco.mjv_connector(scene.geoms[scene.ngeom-1],\n", " mujoco.mjtGeom.mjGEOM_CAPSULE, radius,\n", " point1, point2)\n", "\n", " # traces of time, position and speed\n", "times = []\n", "positions = []\n", "speeds = []\n", "offset = physics.model.jnt_axis[0]/8 # offset along the joint axis\n", "\n", "def scene_callback(physics, scn):\n", " \"\"\"Draw position trace, speed modifies width and colours.\"\"\"\n", " if len(positions) > 1:\n", " for i in range(len(positions)-1):\n", " rgba=np.array((np.clip(speeds[i]/10, 0, 1),\n", " np.clip(1-speeds[i]/10, 0, 1),\n", " .5, 1.))\n", " radius=.003*(1+speeds[i])\n", " point1 = positions[i] + offset*times[i]\n", " point2 = positions[i+1] + offset*times[i+1]\n", " add_visual_capsule(scn, point1, point2, radius, rgba)\n", "\n", "duration = 6 # (seconds)\n", "framerate = 30 # (Hz)\n", "\n", "# Simulate and display video.\n", "frames = []\n", "physics.reset() # Reset state and time\n", "while physics.data.time < duration:\n", " # append data to the traces\n", " positions.append(physics.named.data.geom_xpos[\"green_sphere\"].copy())\n", " times.append(physics.data.time)\n", " speeds.append(get_geom_speed(physics, \"green_sphere\"))\n", " physics.step()\n", " if len(frames) < physics.data.time * framerate:\n", " camera = mujoco.Camera(physics, max_geom=10000,\n", " scene_callback=scene_callback)\n", " pixels = camera.render()\n", " frames.append(pixels)\n", "display_video(frames, framerate)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "gf9h_wi9weet" }, "cell_type": "markdown", "source": [ "# MuJoCo basics and named indexing" ] }, { "metadata": { "id": "NCcZxrDDB1Cj" }, "cell_type": "markdown", "source": [ "## `mjModel`\n", "MuJoCo's `mjModel`, encapsulated in `physics.model`, contains the *model description*, including the default initial state and other fixed quantities which are not a function of the state, e.g. the positions of geoms in the frame of their parent body. The (x, y, z) offsets of the `box` and `sphere` geoms, relative their parent body `box_and_sphere` are given by `model.geom_pos`:" ] }, { "metadata": { "id": "wx8NANvOF8g1" }, "cell_type": "code", "source": [ "physics.model.geom_pos" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "v9O1VNdmHn_K" }, "cell_type": "markdown", "source": [ "Docstrings of attributes provide short descriptions." ] }, { "metadata": { "id": "8_0MwaeYHn_N" }, "cell_type": "code", "source": [ "help(type(physics.model).geom_pos)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "Wee5ATLtIQn_" }, "cell_type": "markdown", "source": [ "The `model.opt` structure contains global quantities like" ] }, { "metadata": { "id": "BhzbZIfDIU2-" }, "cell_type": "code", "source": [ "print('timestep', physics.model.opt.timestep)\n", "print('gravity', physics.model.opt.gravity)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "t5hY0fyXFLcf" }, "cell_type": "markdown", "source": [ "## `mjData`\n", "`mjData`, encapsulated in `physics.data`, contains the **state** and quantities that depend on it. The state is made up of time, generalized positions and generalised velocities. These are respectively `data.time`, `data.qpos` and `data.qvel`. \n", "\n", "Let's print the state of the swinging body where we left it:" ] }, { "metadata": { "id": "acwZtDwp9mQU" }, "cell_type": "code", "source": [ "print(physics.data.time, physics.data.qpos, physics.data.qvel)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "7YlmcLcA-WQu" }, "cell_type": "markdown", "source": [ "`physics.data` also contains **functions of the state**, for example the cartesian positions of objects in the world frame. The (x, y, z) positions of our two geoms are in `data.geom_xpos`:" ] }, { "metadata": { "id": "CPwDcAQ0-uUE" }, "cell_type": "code", "source": [ "print(physics.data.geom_xpos)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "Z0UodCxS_v49" }, "cell_type": "markdown", "source": [ "## Named indexing\n", "\n", "The semantics of the above arrays are made clearer using the `named` wrapper, which assigns names to rows and type names to columns." ] }, { "metadata": { "id": "cLARcaK6-xCU" }, "cell_type": "code", "source": [ "print(physics.named.data.geom_xpos)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "wgXOUZNZHIx6" }, "cell_type": "markdown", "source": [ "Note how `model.geom_pos` and `data.geom_xpos` have similar semantics but very different meanings." ] }, { "metadata": { "id": "-cW61ClRHS8a" }, "cell_type": "code", "source": [ "print(physics.named.model.geom_pos)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "-lQ0AChVASMv" }, "cell_type": "markdown", "source": [ "Name strings can be used to index **into** the relevant quantities, making code much more readable and robust." ] }, { "metadata": { "id": "Rj4ad9fQAnFZ" }, "cell_type": "code", "source": [ "physics.named.data.geom_xpos['green_sphere', 'z']" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "axr_p6APAzFn" }, "cell_type": "markdown", "source": [ "Joint names can be used to index into quantities in joint space (beginning with the letter `q`):" ] }, { "metadata": { "id": "hluF9aDG9O1W" }, "cell_type": "code", "source": [ "physics.named.data.qpos['swing']" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "3IhfyD2Q1pjv" }, "cell_type": "markdown", "source": [ "We can mix NumPy slicing operations with named indexing. As an example, we can set the color of the box using its name (`\"red_box\"`) as an index into the rows of the `geom_rgba` array. " ] }, { "metadata": { "id": "f5vVUullUvWH" }, "cell_type": "code", "source": [ "#@title Changing colors using named indexing{vertical-output: true}\n", "\n", "random_rgb = np.random.rand(3)\n", "physics.named.model.geom_rgba['red_box', :3] = random_rgb\n", "pixels = physics.render()\n", "PIL.Image.fromarray(pixels)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "elzPPdq-KhLI" }, "cell_type": "markdown", "source": [ "Note that while `physics.model` quantities will not be changed by the engine, we can change them ourselves between steps." ] }, { "metadata": { "id": "22ENjtVuhwsm" }, "cell_type": "markdown", "source": [ "## Setting the state with `reset_context()`\n", "\n", "In order for `data` quantities that are functions of the state to be in sync with the state, MuJoCo's `mj_step1()` needs to be called. This is facilitated by the `reset_context()` context, please see in-depth discussion in Section 2.1 of the `dm_control` [tech report](https://arxiv.org/abs/2006.12983)." ] }, { "metadata": { "id": "WBPprCtWgXFN" }, "cell_type": "code", "source": [ "physics.named.data.qpos['swing'] = np.pi\n", "print('Without reset_context, spatial positions are not updated:',\n", " physics.named.data.geom_xpos['green_sphere', ['z']])\n", "with physics.reset_context():\n", " physics.named.data.qpos['swing'] = np.pi\n", "print('After reset_context, positions are up-to-date:',\n", " physics.named.data.geom_xpos['green_sphere', ['z']])" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "SHppAOjvSupc" }, "cell_type": "markdown", "source": [ "## Free bodies: the self-inverting \"tippe-top\"\n", "\n", "A free body is a body with a `free` joint, with 6 movement DOFs: 3 translations and 3 rotations. We could give our `box_and_sphere` body a free joint and watch it fall, but let's look at something more interesting. A \"tippe top\" is a spinning toy which flips itself on its head ([Wikipedia](https://en.wikipedia.org/wiki/Tippe_top)). We model it as follows:" ] }, { "metadata": { "id": "xasXQpVMjIwA" }, "cell_type": "code", "source": [ "#@title The \"tippe-top\" model{vertical-output: true}\n", "\n", "tippe_top = \"\"\"\n", "\n", " \n", "\"\"\"\n", "physics = mujoco.Physics.from_xml_string(tippe_top)\n", "PIL.Image.fromarray(physics.render(camera_id='closeup'))" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "bvHlr6maJYIG" }, "cell_type": "markdown", "source": [ "Note several new features of this model definition:\n", "0. The free joint is added with the `` clause, which is similar to ``, but prohibits unphysical attributes like friction or stiffness.\n", "1. We use the `