From 9c31494fa17de739eae6b31982223eb72a6547ab Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Thu, 27 Feb 2025 01:56:57 -0800 Subject: [PATCH] Add MjSpec example for procedural meshes. PiperOrigin-RevId: 731635686 Change-Id: I310b802966755f627c4ec2cbdcd2ce7b14921c57 --- python/mjspec.ipynb | 154 ++++++++++++++++++++++++++++++++++---------- 1 file changed, 120 insertions(+), 34 deletions(-) diff --git a/python/mjspec.ipynb b/python/mjspec.ipynb index 2eb89566..925da6bc 100644 --- a/python/mjspec.ipynb +++ b/python/mjspec.ipynb @@ -270,37 +270,6 @@ "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, @@ -403,7 +372,7 @@ " spec.option.density = 1.294\n", "\n", " # Defaults for joint and geom\n", - " main = spec.default()\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", @@ -482,6 +451,27 @@ "outputs": [], "source": [ "#@title Make video\n", + "arena_xml = \"\"\"\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + " \n", + "\n", + "\"\"\"\n", + "\n", "spec = procedural_tree(spec=mj.MjSpec.from_string(arena_xml))\n", "model = spec.compile()\n", "data = mj.MjData(model)\n", @@ -489,7 +479,7 @@ "duration = 3 # (seconds)\n", "framerate = 60 # (Hz)\n", "frames = []\n", - "with mj.Renderer(model, width=1920 // 2, height=1080 // 2) as renderer:\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", @@ -755,7 +745,103 @@ " while data.time < duration:\n", " mj.mj_step(model, data)\n", " if len(frames) < data.time * framerate:\n", - " cam.azimuth = 20 + 30 * (1 - np.cos(np.pi*data.time / duration))\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",