diff --git a/python/mjspec.ipynb b/python/mjspec.ipynb
index 925da6bc..78b0f6f1 100644
--- a/python/mjspec.ipynb
+++ b/python/mjspec.ipynb
@@ -124,12 +124,12 @@
" highlighted = pygments.highlight(xml_string, lexer, formatter)\n",
" display(HTML(f\"{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",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " \n",
+ " \n",
+ " \n",
+ "\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",