From 73aba9129756e1eb5853f9fc8e601dd9ccf0d325 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Sat, 8 Feb 2025 13:01:10 -0800 Subject: [PATCH] Add simple procedural tree example to model editing tutorial. https://youtu.be/7VJG61AapEc PiperOrigin-RevId: 724729966 Change-Id: I48c288de8e01baa3cd8f40700bc7137744de37a0 --- python/mjspec.ipynb | 295 ++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 274 insertions(+), 21 deletions(-) diff --git a/python/mjspec.ipynb b/python/mjspec.ipynb index 8ac3364055..d813ea9537 100644 --- a/python/mjspec.ipynb +++ b/python/mjspec.ipynb @@ -146,16 +146,9 @@ "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:" + "# 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" ] }, { @@ -166,7 +159,7 @@ }, "outputs": [], "source": [ - "#@title A static model, from string {vertical-output: true}\n", + "#@title Parse, compile, modify, compile: {vertical-output: true}\n", "\n", "static_model = \"\"\"\n", "\n", @@ -218,6 +211,15 @@ "print_xml(spec.to_xml())" ] }, + { + "cell_type": "markdown", + "metadata": { + "id": "NolxAaRn9N9r" + }, + "source": [ + "# Constructing models from scratch" + ] + }, { "cell_type": "code", "execution_count": 0, @@ -226,11 +228,12 @@ }, "outputs": [], "source": [ - "#@title Building an `mjSpec` from scratch {vertical-output: true}\n", + "#@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\", euler=[0, 0, -30])\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", @@ -261,13 +264,262 @@ "media.show_video(frames, fps=framerate)" ] }, + { + "cell_type": "markdown", + "metadata": { + "id": "Y4rV2NDh92Ga" + }, + "source": [ + "## Procedural 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": { + "id": "AgYLwOhs1Msn" + }, + "outputs": [], + "source": [ + "#@title arena model\n", + "arena_xml = \"\"\"\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + " \n", + "\n", + "\"\"\"" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "IQ9G54Yu-Cse" + }, + "outputs": [], + "source": [ + "#@title utility functions\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", + "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 // 2, height=1080 // 2) 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": "3N4YEIVt75_T" }, "source": [ - "# Control example" + "# `dm_control` example" ] }, { @@ -276,10 +528,10 @@ "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", + "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 is a fundamental change to the kinematic structure. The following snippets realise this scenario." + "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." ] }, { @@ -368,7 +620,6 @@ " \"\"\"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", @@ -673,7 +924,6 @@ "#@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", @@ -783,7 +1033,6 @@ "\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", @@ -819,8 +1068,12 @@ "metadata": { "accelerator": "GPU", "colab": { + "collapsed_sections": [ + "sJFuNetilv4m" + ], "gpuClass": "premium", - "private_outputs": true + "private_outputs": true, + "toc_visible": true }, "gpuClass": "premium", "kernelspec": {