Skip to content

Commit

Permalink
Add simple procedural tree example to model editing tutorial.
Browse files Browse the repository at this point in the history 
https://youtu.be/7VJG61AapEc

PiperOrigin-RevId: 724729966
Change-Id: I48c288de8e01baa3cd8f40700bc7137744de37a0
  • Loading branch information
@yuvaltassa @copybara-github
yuvaltassa authored and copybara-github committed Feb 8, 2025
1 parent 5ee0654 commit 73aba91
Showing 1 changed file with 274 additions and 21 deletions.
295 changes: 274 additions & 21 deletions python/mjspec.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
]
},
{
Expand All @@ -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",
"<mujoco>\n",
Expand Down Expand Up @@ -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,
Expand All @@ -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",
Expand Down Expand Up @@ -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",
"<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=\"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=\"300\" height=\"300\"/>\n",
" <material name=\"groundplane\" texture=\"groundplane\" texuniform=\"true\" texrepeat=\"5 5\" reflectance=\"0\"/>\n",
" </asset>\n",
"\n",
" <worldbody>\n",
" <geom name=\"floor\" size=\"5 5 0.01\" type=\"plane\" material=\"groundplane\"/>\n",
" <light pos=\"0 0 3\"/>\n",
" </worldbody>\n",
"</mujoco>\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"
]
},
{
Expand All @@ -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."
]
},
{
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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": {
Expand Down

0 comments on commit 73aba91

Please sign in to comment.