chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add examples of scaling assets into tutorial notebook. Fixes #2476

Browse Source 
PiperOrigin-RevId: 760617809
Change-Id: I86d4fadb473cf8ad9bcfd058567551c431f2bf31
This commit is contained in:
Tom Power
2025-05-19 07:33:54 -07:00
committed by Copybara-Service
parent 84d658a5ad
commit 6c201d3e1c
1 changed files with 397 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mjspec.ipynb
+397 -23
View File
@@ -124,12 +124,12 @@
" 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",
"def render(model, data=None, height=300, camera=-1):\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",
" renderer.update_scene(data, camera)\n",
" media.show_image(renderer.render(), height=height)"
]
},
@@ -148,6 +148,7 @@
"cell_type": "code",
"execution_count": 0,
"metadata": {
"cellView": "form",
"id": "oummB7I7EfSq"
},
"outputs": [],
@@ -987,7 +988,7 @@
},
"outputs": [],
"source": [
"#@title Six Creatures on a floor.{vertical-output: true}\n",
"#@title Six Creatures on a floor {vertical-output: true}\n",
"\n",
"arena = mj.MjSpec()\n",
"\n",
@@ -1046,7 +1047,7 @@
},
"outputs": [],
"source": [
"#@title Video of the movement{vertical-output: true}\n",
"#@title Video of the movement {vertical-output: true}\n",
"\n",
"data = mj.MjData(model)\n",
"duration = 10 # (Seconds)\n",
@@ -1094,7 +1095,7 @@
},
"outputs": [],
"source": [
"#@title Movement trajectories{vertical-output: true}\n",
"#@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",
@@ -1120,6 +1121,57 @@
"# Model editing"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"cellView": "form",
"id": "m4sppBqvf7yd"
},
"outputs": [],
"source": [
"# @title Get resources\n",
"\n",
"# Get Models\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 Fly and Franka XML description from GitHub:')\n",
"!git clone https://github.com/google-deepmind/mujoco_menagerie\n",
"fly_file = 'mujoco_menagerie/flybody/fruitfly.xml'\n",
"franka_file = 'mujoco_menagerie/franka_fr3/fr3.xml'\n",
"\n",
"# Camera options\n",
"cam = mj.MjvCamera()\n",
"mj.mjv_defaultCamera(cam)\n",
"cam.elevation = -10\n",
"cam.lookat = [0, 0, 1]\n",
"cam.distance = 4\n",
"cam.azimuth = 135\n",
"\n",
"# Arena\n",
"arena_xml = \"\"\"\n",
"<mujoco>\n",
" <visual>\n",
" <headlight diffuse=\".5 .5 .5\" specular=\"1 1 1\"/>\n",
" <global 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=\"400\" height=\"400\"/>\n",
" <material name=\"groundplane\" texture=\"groundplane\" texrepeat=\"45 45\" reflectance=\"0\"/>\n",
" </asset>\n",
"\n",
" <worldbody>\n",
" <geom name=\"floor\" size=\"10 10 0.1\" type=\"plane\" material=\"groundplane\"/>\n",
" </worldbody>\n",
"</mujoco>\n",
"\"\"\"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
@@ -1140,13 +1192,7 @@
},
"outputs": [],
"source": [
"#@title Traversing the spec.{vertical-output: true}\n",
"\n",
"# Get MuJoCo's humanoid model.\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",
"#@title Traversing the spec {vertical-output: true}\n",
"\n",
"spec = mj.MjSpec.from_file(humanoid_file)\n",
"\n",
@@ -1183,7 +1229,7 @@
},
"outputs": [],
"source": [
"#@title Model re-compilation with state preservation.{vertical-output: true}\n",
"#@title Model re-compilation with state preservation {vertical-output: true}\n",
"\n",
"spec = mj.MjSpec.from_file(humanoid100_file)\n",
"model = spec.compile()\n",
@@ -1239,7 +1285,7 @@
},
"outputs": [],
"source": [
"#@title Humanoid model.{vertical-output: true}\n",
"#@title Humanoid model {vertical-output: true}\n",
"\n",
"spec = mj.MjSpec.from_file(humanoid_file)\n",
"\n",
@@ -1264,7 +1310,7 @@
},
"outputs": [],
"source": [
"#@title Humanoid with arms replaced by legs.{vertical-output: true}\n",
"#@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",
@@ -1309,11 +1355,7 @@
},
"outputs": [],
"source": [
"#@title Humanoid with Franka arm.{vertical-output: true}\n",
"\n",
"# Get Franka arm from the MuJoCo Menagerie.\n",
"!git clone https://github.com/google-deepmind/mujoco_menagerie\n",
"franka_file = 'mujoco_menagerie/franka_fr3/fr3.xml'\n",
"#@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",
@@ -1350,11 +1392,12 @@
"cell_type": "code",
"execution_count": 0,
"metadata": {
"cellView": "form",
"id": "50lOgJ7mQ2bV"
},
"outputs": [],
"source": [
"#@title Imported actuators.{vertical-output: true}\n",
"#@title Imported actuators {vertical-output: true}\n",
"\n",
"for actuator in spec.actuators:\n",
" print(actuator.name)"
@@ -1377,7 +1420,7 @@
},
"outputs": [],
"source": [
"#@title Humanoid with randomized heads and arm poses.{vertical-output: true}\n",
"#@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",
@@ -1412,13 +1455,344 @@
"model = spec.compile()\n",
"render(model, height=400)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "iXgYCVzEWFTU"
},
"source": [
"## Model scaling"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "-hSJKyH4A2VY"
},
"outputs": [],
"source": [
"#@title Uniformly scale humanoid {vertical-output: true}\n",
"\n",
"def scale_spec(spec, scale):\n",
" scaled_spec = spec.copy()\n",
" # Traverse the kinematic tree, scaling all geoms\n",
" def scale_bodies(parent, scale=1.0):\n",
" body = parent.first_body()\n",
" while body:\n",
" if body.pos is not None:\n",
" body.pos = body.pos * scale\n",
" for geom in body.geoms:\n",
" geom.fromto = geom.fromto * scale\n",
" geom.size = geom.size * scale\n",
" if geom.pos is not None:\n",
" geom.pos = geom.pos * scale\n",
" scale_bodies(body, scale)\n",
" body = parent.next_body(body)\n",
"\n",
" scale_bodies(scaled_spec.body('world'), scale)\n",
" return scaled_spec\n",
"\n",
"spec = mj.MjSpec.from_string(arena_xml)\n",
"humanoid = mj.MjSpec.from_file(humanoid_file)\n",
"small_humanoid = scale_spec(humanoid, 0.75)\n",
"large_humanoid = scale_spec(humanoid, 1.25)\n",
"\n",
"# Create a line-up of humanoids\n",
"frame = spec.worldbody.add_frame(pos=[-1, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(humanoid.body('torso'), str(0))\n",
"\n",
"frame = spec.worldbody.add_frame(pos=[0, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(small_humanoid.body('torso'), str(1))\n",
"\n",
"frame = spec.worldbody.add_frame(pos=[1, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2] )\n",
"frame.attach_body(large_humanoid.body('torso'), str(2))\n",
"\n",
"\n",
"spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n",
" targetbody='1torso', 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, camera=cam)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "tBH5vmqJXleD"
},
"source": [
"We can scale the size of a model by traversing the kinematic tree and applying the the scale to the relevant geoms. Above we can see humanoids of three different sizes."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "cV4tkG6siFQp"
},
"outputs": [],
"source": [
"# @title Scaling actuator forces {vertical-output: true}\n",
"\n",
"def scale_spec(spec, scale, scale_actuators=False):\n",
" scaled_spec = spec.copy()\n",
" # Traverse the kinematic tree, scaling all geoms\n",
" def scale_bodies(parent, scale=1.0):\n",
" body = parent.first_body()\n",
" while body:\n",
" if body.pos is not None:\n",
" body.pos = body.pos * scale\n",
" for geom in body.geoms:\n",
" geom.fromto = geom.fromto * scale\n",
" geom.size = geom.size * scale\n",
" if geom.pos is not None:\n",
" geom.pos = geom.pos * scale\n",
" scale_bodies(body, scale)\n",
" body = parent.next_body(body)\n",
"\n",
" if scale_actuators:\n",
" # scale gear\n",
" for actuator in scaled_spec.actuators:\n",
" # scale the actuator gear by (scale ** 2),\n",
" # this is because muscle force-generating capacity\n",
" # scales with the cross-sectional area of the muscle\n",
" actuator.gear = actuator.gear * scale * scale\n",
"\n",
" # scale the z-position of the humanoid for all keypoints\n",
" for keypoint in scaled_spec.keys:\n",
" qpos = keypoint.qpos\n",
" qpos[2] = qpos[2] * scale\n",
" keypoint.qpos = qpos\n",
" keypoint.qpos[2] = keypoint.qpos[2] * scale\n",
"\n",
" scale_bodies(scaled_spec.body('world'), scale)\n",
" return scaled_spec\n",
"\n",
"# Create specs\n",
"scale = 0.6\n",
"spec = mj.MjSpec.from_string(arena_xml)\n",
"humanoid = mj.MjSpec.from_file(humanoid_file)\n",
"small_humanoid = scale_spec(humanoid, scale)\n",
"small_humanoid_actuators_scaled = scale_spec(humanoid, scale, True)\n",
"\n",
"# Create a line-up of humanoids\n",
"squat_qpos = []\n",
"\n",
"# Add unscaled humanoid\n",
"frame = spec.worldbody.add_frame(pos=[-1, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(humanoid.body('torso'), str(0))\n",
"# Record squat pose\n",
"humanoid_squat = humanoid.key('squat').qpos\n",
"humanoid_squat[:2] = frame.pos[:2]\n",
"humanoid_squat[3:7] = frame.quat\n",
"squat_qpos.append(humanoid_squat)\n",
"\n",
"# Add small humanoid\n",
"frame = spec.worldbody.add_frame(pos=[0, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(small_humanoid.body('torso'), str(1))\n",
"# Record squat pose\n",
"humanoid_squat = small_humanoid.key('squat').qpos\n",
"humanoid_squat[:2] = frame.pos[:2]\n",
"humanoid_squat[3:7] = frame.quat\n",
"squat_qpos.append(humanoid_squat)\n",
"\n",
"# Add small humanoid with scaled actuators\n",
"frame = spec.worldbody.add_frame(pos=[1, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2] )\n",
"frame.attach_body(small_humanoid_actuators_scaled.body('torso'), str(2))\n",
"# Record squat pose\n",
"humanoid_squat = small_humanoid_actuators_scaled.key('squat').qpos\n",
"humanoid_squat[:2] = frame.pos[:2]\n",
"humanoid_squat[3:7] = frame.quat\n",
"squat_qpos.append(humanoid_squat)\n",
"squat_qpos = np.concatenate(squat_qpos)\n",
"\n",
"spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n",
" targetbody='1torso', diffuse=[.8, .8, .8],\n",
" specular=[0.3, 0.3, 0.3], pos=[0, -6, 4], cutoff=30)\n",
"model = spec.compile()\n",
"\n",
"# Initialize to squat position\n",
"data = mj.MjData(model)\n",
"data.qpos = squat_qpos\n",
"\n",
"# jumping motion\n",
"u_t = lambda t: 10.0 * t / duration\n",
"\n",
"# Simulate and display video.\n",
"duration = 2 # (seconds)\n",
"framerate = 30 # (Hz)\n",
"frames = []\n",
"\n",
"with mj.Renderer(model, 480, 640) as renderer:\n",
" while data.time < duration:\n",
" data.ctrl = u_t(data.time)\n",
" mj.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data, camera=cam)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"\n",
"media.show_video(frames, fps=framerate, height=400)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ZSOra3S2YpIB"
},
"source": [
"We can also apply scaling to the actuators. In the humanoid case, scaling the geoms without scaling the `gear` parameter for the actuators results in a humanoid that can jump higher proportional to its size."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "9IuwQQ0F2ddA"
},
"outputs": [],
"source": [
"# @title Long-limbed humanoid {vertical-output: true}\n",
"\n",
"def scale_spec(spec, scale):\n",
" scaled_spec = spec.copy()\n",
" # Traverse the kinematic tree, scaling all geoms\n",
" def scale_bodies(parent, scale=1.0):\n",
" if parent is not None:\n",
" for geom in parent.geoms:\n",
" # Only scale fromto, not size to scale length of capsules\n",
" geom.fromto = geom.fromto * scale\n",
" if geom.pos is not None:\n",
" geom.pos = geom.pos * scale\n",
" body = parent.first_body()\n",
" while body:\n",
" if body.pos is not None:\n",
" body.pos = body.pos * scale\n",
" scale_bodies(body, scale)\n",
" body = parent.next_body(body)\n",
"\n",
" # Scale all the limbs\n",
" scale_bodies(scaled_spec.body('upper_arm_right'), scale)\n",
" scale_bodies(scaled_spec.body('upper_arm_left'), scale)\n",
" scale_bodies(scaled_spec.body('thigh_right'), scale)\n",
" scale_bodies(scaled_spec.body('thigh_left'), scale)\n",
" return scaled_spec\n",
"\n",
"spec = mj.MjSpec.from_string(arena_xml)\n",
"humanoid = mj.MjSpec.from_file(humanoid_file)\n",
"small_humanoid = scale_spec(humanoid, 1.25)\n",
"large_humanoid = scale_spec(humanoid, 2)\n",
"\n",
"# Create a line-up of humanoids by attaching\n",
"frame = spec.worldbody.add_frame(pos=[-1, 0, 0],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(humanoid.body('torso'), str(0), str(0))\n",
"\n",
"frame = spec.worldbody.add_frame(pos=[0, 0, 0.2],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(small_humanoid.body('torso'), str(0), str(1))\n",
"\n",
"frame = spec.worldbody.add_frame(pos=[1, 0, 0.8],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2] )\n",
"frame.attach_body(large_humanoid.body('torso'), str(0), str(2))\n",
"\n",
"\n",
"spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n",
" targetbody='0torso1', diffuse=[.8, .8, .8],\n",
" specular=[0.3, 0.3, 0.3], pos=[0, -6, 4], cutoff=30)\n",
"model = spec.compile()\n",
"\n",
"# camera options\n",
"render(model, height=400, camera=cam)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "D2DrPBvBZjI0"
},
"source": [
"We can also apply scaling to the model non-uniformly. In this instance we scale the humanoid to have long limbs, by only applying the scale to the length of the capsule geoms for the arms, legs and feet."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "1G8VO45v2ddA"
},
"outputs": [],
"source": [
"# @title Meshes {vertical-output: true}\n",
"\n",
"def scale_spec(spec, scale):\n",
" scaled_spec = spec.copy()\n",
" # scale all meshes\n",
" for mesh in scaled_spec.meshes:\n",
" if mesh.scale is None:\n",
" mesh.scale = np.ones(3)\n",
" mesh.scale = mesh.scale * scale\n",
"\n",
" # Traverse the kinematic tree\n",
" def scale_bodies(parent, scale=1.0):\n",
" if parent is not None:\n",
" for geom in parent.geoms:\n",
" if geom.pos is not None:\n",
" geom.pos = geom.pos * scale\n",
" body = parent.first_body()\n",
" while body:\n",
" if body.pos is not None:\n",
" body.pos = body.pos * scale\n",
" scale_bodies(body, scale)\n",
" body = parent.next_body(body)\n",
"\n",
" # Scale all the limbs\n",
" scale_bodies(scaled_spec.body('world'), scale)\n",
"\n",
" return scaled_spec\n",
"\n",
"spec = mj.MjSpec.from_string(arena_xml)\n",
"fly = mj.MjSpec.from_file(fly_file)\n",
"# Remove lights from fly so they are not duplicated in line-up\n",
"for light in fly.lights:\n",
" light.delete()\n",
"\n",
"small_fly = scale_spec(fly, 1.25)\n",
"large_fly = scale_spec(fly, 2)\n",
"\n",
"# Create a line-up of flys by attaching\n",
"frame = spec.worldbody.add_frame(pos=[-1, 0, 0.25],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(fly.body('thorax'), str(0), str(0))\n",
"\n",
"frame = spec.worldbody.add_frame(pos=[0, 0, 0.25],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2])\n",
"frame.attach_body(small_fly.body('thorax'), str(0), str(1))\n",
"\n",
"frame = spec.worldbody.add_frame(pos=[1, 0, 0.25],\n",
" quat=[-np.sqrt(2)/2, 0, 0, np.sqrt(2) / 2] )\n",
"frame.attach_body(large_fly.body('thorax'), str(0), str(2))\n",
"\n",
"spec.worldbody.add_light(mode=mj.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM,\n",
" targetbody='0thorax1', 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, camera=cam)\n"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [
"sJFuNetilv4m",
"yXY7HGfVsVlo"
],
"gpuClass": "premium",