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

Merge pull request #1769 from fanminshi:renderer_context

Browse Source 
PiperOrigin-RevId: 648508835
Change-Id: Ice87cd1221723ed44ba0ddca97b6aac245dd8975
This commit is contained in:
Copybara-Service
2024-07-01 15:48:41 -07:00
parent 19d505d847 521922f133
commit 0676ad499f
1 changed files with 197 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tutorial.ipynb
+197 -176
View File
@@ -420,8 +420,8 @@
"data = mujoco.MjData(model)\n",
"\n",
"# Make renderer, render and show the pixels\n",
"renderer = mujoco.Renderer(model)\n",
"media.show_image(renderer.render())"
"with mujoco.Renderer(model) as renderer:\n",
" media.show_image(renderer.render())"
]
},
{
@@ -445,10 +445,11 @@
},
"outputs": [],
"source": [
"mujoco.mj_forward(model, data)\n",
"renderer.update_scene(data)\n",
"with mujoco.Renderer(model) as renderer:\n",
" mujoco.mj_forward(model, data)\n",
" renderer.update_scene(data)\n",
"\n",
"media.show_image(renderer.render())"
" media.show_image(renderer.render())"
]
},
{
@@ -480,11 +481,11 @@
"model = mujoco.MjModel.from_xml_string(xml)\n",
"data = mujoco.MjData(model)\n",
"\n",
"renderer = mujoco.Renderer(model)\n",
"mujoco.mj_forward(model, data)\n",
"renderer.update_scene(data)\n",
"with mujoco.Renderer(model) as renderer:\n",
" mujoco.mj_forward(model, data)\n",
" renderer.update_scene(data)\n",
"\n",
"media.show_image(renderer.render())"
" media.show_image(renderer.render())"
]
},
{
@@ -508,9 +509,10 @@
"source": [
"# Run this cell multiple times for different colors\n",
"model.geom('red_box').rgba[:3] = np.random.rand(3)\n",
"renderer.update_scene(data)\n",
"with mujoco.Renderer(model) as renderer:\n",
" renderer.update_scene(data)\n",
"\n",
"media.show_image(renderer.render())"
" media.show_image(renderer.render())"
]
},
{
@@ -540,12 +542,13 @@
"# Simulate and display video.\n",
"frames = []\n",
"mujoco.mj_resetData(model, data) # Reset state and time.\n",
"while data.time < duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"with mujoco.Renderer(model) as renderer:\n",
" while data.time < duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"\n",
"media.show_video(frames, fps=framerate)"
]
@@ -583,7 +586,6 @@
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(xml)\n",
"data = mujoco.MjData(model)\n",
"renderer = mujoco.Renderer(model)\n",
"\n",
"# enable joint visualization option:\n",
"scene_option = mujoco.MjvOption()\n",
@@ -595,12 +597,13 @@
"# Simulate and display video.\n",
"frames = []\n",
"mujoco.mj_resetData(model, data)\n",
"while data.time < duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data, scene_option=scene_option)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"with mujoco.Renderer(model) as renderer:\n",
" while data.time < duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data, scene_option=scene_option)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"\n",
"media.show_video(frames, fps=framerate)"
]
@@ -651,12 +654,13 @@
"# Simulate and display video.\n",
"frames = []\n",
"mujoco.mj_resetData(model, data)\n",
"while data.time < duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data, scene_option=scene_option)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"with mujoco.Renderer(model) as renderer:\n",
" while data.time < duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data, scene_option=scene_option)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"\n",
"media.show_video(frames, fps=60)"
]
@@ -760,12 +764,12 @@
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(tippe_top)\n",
"data = mujoco.MjData(model)\n",
"renderer = mujoco.Renderer(model)\n",
"\n",
"mujoco.mj_forward(model, data)\n",
"renderer.update_scene(data, camera=\"closeup\")\n",
"with mujoco.Renderer(model) as renderer:\n",
" renderer.update_scene(data, camera=\"closeup\")\n",
"\n",
"media.show_image(renderer.render())"
" media.show_image(renderer.render())"
]
},
{
@@ -821,12 +825,13 @@
"# Simulate and display video.\n",
"frames = []\n",
"mujoco.mj_resetDataKeyframe(model, data, 0) # Reset the state to keyframe 0\n",
"while data.time < duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data, \"closeup\")\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"with mujoco.Renderer(model) as renderer:\n",
" while data.time < duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data, \"closeup\")\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"\n",
"media.show_video(frames, fps=framerate)"
]
@@ -939,12 +944,14 @@
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(chaotic_pendulum)\n",
"data = mujoco.MjData(model)\n",
"renderer = mujoco.Renderer(model, 480, 640)\n",
"height = 480\n",
"width = 640\n",
"\n",
"mujoco.mj_forward(model, data)\n",
"renderer.update_scene(data, camera=\"fixed\")\n",
"with mujoco.Renderer(model, height, width) as renderer:\n",
" mujoco.mj_forward(model, data)\n",
" renderer.update_scene(data, camera=\"fixed\")\n",
"\n",
"media.show_image(renderer.render())"
" media.show_image(renderer.render())"
]
},
{
@@ -970,7 +977,8 @@
"framerate = 30 # Hz\n",
"n_frames = int(n_seconds * framerate)\n",
"frames = []\n",
"renderer = mujoco.Renderer(model, 240, 320)\n",
"height = 240\n",
"width = 320\n",
"\n",
"# set initial state\n",
"mujoco.mj_resetData(model, data)\n",
@@ -981,17 +989,18 @@
"sim_time = 0\n",
"render_time = 0\n",
"n_steps = 0\n",
"for i in range(n_frames):\n",
" while data.time * framerate < i:\n",
"with mujoco.Renderer(model, height, width) as renderer:\n",
" for i in range(n_frames):\n",
" while data.time * framerate < i:\n",
" tic = time.time()\n",
" mujoco.mj_step(model, data)\n",
" sim_time += time.time() - tic\n",
" n_steps += 1\n",
" tic = time.time()\n",
" mujoco.mj_step(model, data)\n",
" sim_time += time.time() - tic\n",
" n_steps += 1\n",
" tic = time.time()\n",
" renderer.update_scene(data, \"fixed\")\n",
" frame = renderer.render()\n",
" render_time += time.time() - tic\n",
" frames.append(frame)\n",
" renderer.update_scene(data, \"fixed\")\n",
" frame = renderer.render()\n",
" render_time += time.time() - tic\n",
" frames.append(frame)\n",
"\n",
"# print timing and play video\n",
"step_time = 1e6*sim_time/n_steps\n",
@@ -1233,12 +1242,14 @@
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(free_body_MJCF)\n",
"data = mujoco.MjData(model)\n",
"renderer = mujoco.Renderer(model, 400, 600)\n",
"height = 400\n",
"width = 600\n",
"\n",
"mujoco.mj_forward(model, data)\n",
"renderer.update_scene(data, \"fixed\")\n",
"with mujoco.Renderer(model, height, width) as renderer:\n",
" mujoco.mj_forward(model, data)\n",
" renderer.update_scene(data, \"fixed\")\n",
"\n",
"media.show_image(renderer.render())"
" media.show_image(renderer.render())"
]
},
{
@@ -1262,7 +1273,6 @@
"height = 240\n",
"width = 320\n",
"frames = []\n",
"renderer = mujoco.Renderer(model, height, width)\n",
"\n",
"# visualize contact frames and forces, make body transparent\n",
"options = mujoco.MjvOption()\n",
@@ -1282,12 +1292,13 @@
"data.qvel[3:6] = 5*np.random.randn(3)\n",
"\n",
"# Simulate and display video.\n",
"for i in range(n_frames):\n",
" while data.time < i/120.0: #1/4x real time\n",
" mujoco.mj_step(model, data)\n",
" renderer.update_scene(data, \"track\", options)\n",
" frame = renderer.render()\n",
" frames.append(frame)\n",
"with mujoco.Renderer(model, height, width) as renderer:\n",
" for i in range(n_frames):\n",
" while data.time < i/120.0: #1/4x real time\n",
" mujoco.mj_step(model, data)\n",
" renderer.update_scene(data, \"track\", options)\n",
" frame = renderer.render()\n",
" frames.append(frame)\n",
"\n",
"media.show_video(frames, fps=30)"
]
@@ -1440,16 +1451,16 @@
"# load\n",
"model = mujoco.MjModel.from_xml_string(MJCF)\n",
"data = mujoco.MjData(model)\n",
"renderer = mujoco.Renderer(model, height, width)\n",
"\n",
"# Simulate and display video.\n",
"mujoco.mj_resetData(model, data)\n",
"for i in range(n_frames):\n",
" while data.time < i/30.0:\n",
" mujoco.mj_step(model, data)\n",
" renderer.update_scene(data, \"y\")\n",
" frame = renderer.render()\n",
" frames.append(frame)\n",
"with mujoco.Renderer(model, height, width) as renderer:\n",
" mujoco.mj_resetData(model, data)\n",
" for i in range(n_frames):\n",
" while data.time < i/30.0:\n",
" mujoco.mj_step(model, data)\n",
" renderer.update_scene(data, \"y\")\n",
" frame = renderer.render()\n",
" frames.append(frame)\n",
"\n",
"media.show_video(frames, fps=30)"
]
@@ -1520,12 +1531,14 @@
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(MJCF)\n",
"data = mujoco.MjData(model)\n",
"renderer = mujoco.Renderer(model, 480, 480)\n",
"height = 480\n",
"width = 480\n",
"\n",
"mujoco.mj_forward(model, data)\n",
"renderer.update_scene(data, \"fixed\")\n",
"with mujoco.Renderer(model, height, width) as renderer:\n",
" mujoco.mj_forward(model, data)\n",
" renderer.update_scene(data, \"fixed\")\n",
"\n",
"media.show_image(renderer.render())\n"
" media.show_image(renderer.render())"
]
},
{
@@ -1553,21 +1566,20 @@
"times = []\n",
"sensordata = []\n",
"\n",
"renderer = mujoco.Renderer(model, height, width)\n",
"\n",
"# constant actuator signal\n",
"mujoco.mj_resetData(model, data)\n",
"data.ctrl = 20\n",
"\n",
"# Simulate and display video.\n",
"for i in range(n_frames):\n",
" while data.time < i/fps:\n",
" mujoco.mj_step(model, data)\n",
" times.append(data.time)\n",
" sensordata.append(data.sensor('accelerometer').data.copy())\n",
" renderer.update_scene(data, \"fixed\")\n",
" frame = renderer.render()\n",
" frames.append(frame)\n",
"with mujoco.Renderer(model, height, width) as renderer:\n",
" for i in range(n_frames):\n",
" while data.time < i/fps:\n",
" mujoco.mj_step(model, data)\n",
" times.append(data.time)\n",
" sensordata.append(data.sensor('accelerometer').data.copy())\n",
" renderer.update_scene(data, \"fixed\")\n",
" frame = renderer.render()\n",
" frames.append(frame)\n",
"\n",
"media.show_video(frames, fps=fps)"
]
@@ -1644,12 +1656,12 @@
"</mujoco>\n",
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(xml)\n",
"renderer = mujoco.Renderer(model)\n",
"data = mujoco.MjData(model)\n",
"\n",
"mujoco.mj_forward(model, data)\n",
"renderer.update_scene(data)\n",
"media.show_image(renderer.render())"
"with mujoco.Renderer(model) as renderer:\n",
" mujoco.mj_forward(model, data)\n",
" renderer.update_scene(data)\n",
" media.show_image(renderer.render())"
]
},
{
@@ -1664,9 +1676,10 @@
"\n",
"scene_option.frame = mujoco.mjtFrame.mjFRAME_GEOM\n",
"scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = True\n",
"renderer.update_scene(data, scene_option=scene_option)\n",
"frame = renderer.render()\n",
"media.show_image(frame)"
"with mujoco.Renderer(model) as renderer:\n",
" renderer.update_scene(data, scene_option=scene_option)\n",
" frame = renderer.render()\n",
" media.show_image(frame)"
]
},
{
@@ -1679,25 +1692,24 @@
"source": [
"#@title Depth rendering {vertical-output: true}\n",
"\n",
"# update renderer to render depth\n",
"renderer.enable_depth_rendering()\n",
"with mujoco.Renderer(model) as renderer:\n",
" # update renderer to render depth\n",
" renderer.enable_depth_rendering()\n",
"\n",
"# reset the scene\n",
"renderer.update_scene(data)\n",
" # reset the scene\n",
" renderer.update_scene(data)\n",
"\n",
"# depth is a float array, in meters.\n",
"depth = renderer.render()\n",
" # depth is a float array, in meters.\n",
" depth = renderer.render()\n",
"\n",
"# Shift nearest values to the origin.\n",
"depth -= depth.min()\n",
"# Scale by 2 mean distances of near rays.\n",
"depth /= 2*depth[depth <= 1].mean()\n",
"# Scale to [0, 255]\n",
"pixels = 255*np.clip(depth, 0, 1)\n",
" # Shift nearest values to the origin.\n",
" depth -= depth.min()\n",
" # Scale by 2 mean distances of near rays.\n",
" depth /= 2*depth[depth <= 1].mean()\n",
" # Scale to [0, 255]\n",
" pixels = 255*np.clip(depth, 0, 1)\n",
"\n",
"media.show_image(pixels.astype(np.uint8))\n",
"\n",
"renderer.disable_depth_rendering()"
" media.show_image(pixels.astype(np.uint8))"
]
},
{
@@ -1710,25 +1722,26 @@
"source": [
"#@title Segmentation rendering {vertical-output: true}\n",
"\n",
"# update renderer to render segmentation\n",
"renderer.enable_segmentation_rendering()\n",
"with mujoco.Renderer(model) as renderer:\n",
" renderer.disable_depth_rendering()\n",
"\n",
"# reset the scene\n",
"renderer.update_scene(data)\n",
" # update renderer to render segmentation\n",
" renderer.enable_segmentation_rendering()\n",
"\n",
"seg = renderer.render()\n",
" # reset the scene\n",
" renderer.update_scene(data)\n",
"\n",
"# Display the contents of the first channel, which contains object\n",
"# IDs. The second channel, seg[:, :, 1], contains object types.\n",
"geom_ids = seg[:, :, 0]\n",
"# Infinity is mapped to -1\n",
"geom_ids = geom_ids.astype(np.float64) + 1\n",
"# Scale to [0, 1]\n",
"geom_ids = geom_ids / geom_ids.max()\n",
"pixels = 255*geom_ids\n",
"media.show_image(pixels.astype(np.uint8))\n",
" seg = renderer.render()\n",
"\n",
"renderer.disable_segmentation_rendering()"
" # Display the contents of the first channel, which contains object\n",
" # IDs. The second channel, seg[:, :, 1], contains object types.\n",
" geom_ids = seg[:, :, 0]\n",
" # Infinity is mapped to -1\n",
" geom_ids = geom_ids.astype(np.float64) + 1\n",
" # Scale to [0, 1]\n",
" geom_ids = geom_ids / geom_ids.max()\n",
" pixels = 255*geom_ids\n",
" media.show_image(pixels.astype(np.uint8))"
]
},
{
@@ -1792,37 +1805,39 @@
"source": [
"#@title Project from world to camera coordinates {vertical-output: true}\n",
"\n",
"# reset the scene\n",
"renderer.update_scene(data)\n",
"with mujoco.Renderer(model) as renderer:\n",
" renderer.disable_segmentation_rendering()\n",
" # reset the scene\n",
" renderer.update_scene(data)\n",
"\n",
"# Get the world coordinates of the box corners\n",
"box_pos = data.geom_xpos[model.geom('red_box').id]\n",
"box_mat = data.geom_xmat[model.geom('red_box').id].reshape(3, 3)\n",
"box_size = model.geom_size[model.geom('red_box').id]\n",
"offsets = np.array([-1, 1]) * box_size[:, None]\n",
"xyz_local = np.stack(list(itertools.product(*offsets))).T\n",
"xyz_global = box_pos[:, None] + box_mat @ xyz_local\n",
" # Get the world coordinates of the box corners\n",
" box_pos = data.geom_xpos[model.geom('red_box').id]\n",
" box_mat = data.geom_xmat[model.geom('red_box').id].reshape(3, 3)\n",
" box_size = model.geom_size[model.geom('red_box').id]\n",
" offsets = np.array([-1, 1]) * box_size[:, None]\n",
" xyz_local = np.stack(list(itertools.product(*offsets))).T\n",
" xyz_global = box_pos[:, None] + box_mat @ xyz_local\n",
"\n",
"# Camera matrices multiply homogenous [x, y, z, 1] vectors.\n",
"corners_homogeneous = np.ones((4, xyz_global.shape[1]), dtype=float)\n",
"corners_homogeneous[:3, :] = xyz_global\n",
" # Camera matrices multiply homogenous [x, y, z, 1] vectors.\n",
" corners_homogeneous = np.ones((4, xyz_global.shape[1]), dtype=float)\n",
" corners_homogeneous[:3, :] = xyz_global\n",
"\n",
"# Get the camera matrix.\n",
"m = compute_camera_matrix(renderer, data)\n",
" # Get the camera matrix.\n",
" m = compute_camera_matrix(renderer, data)\n",
"\n",
"# Project world coordinates into pixel space. See:\n",
"# https://en.wikipedia.org/wiki/3D_projection#Mathematical_formula\n",
"xs, ys, s = m @ corners_homogeneous\n",
"# x and y are in the pixel coordinate system.\n",
"x = xs / s\n",
"y = ys / s\n",
" # Project world coordinates into pixel space. See:\n",
" # https://en.wikipedia.org/wiki/3D_projection#Mathematical_formula\n",
" xs, ys, s = m @ corners_homogeneous\n",
" # x and y are in the pixel coordinate system.\n",
" x = xs / s\n",
" y = ys / s\n",
"\n",
"# Render the camera view and overlay the projected corner coordinates.\n",
"pixels = renderer.render()\n",
"fig, ax = plt.subplots(1, 1)\n",
"ax.imshow(pixels)\n",
"ax.plot(x, y, '+', c='w')\n",
"ax.set_axis_off()"
" # Render the camera view and overlay the projected corner coordinates.\n",
" pixels = renderer.render()\n",
" fig, ax = plt.subplots(1, 1)\n",
" ax.imshow(pixels)\n",
" ax.plot(x, y, '+', c='w')\n",
" ax.set_axis_off()"
]
},
{
@@ -1894,17 +1909,18 @@
"mujoco.mj_resetData(model, data)\n",
"mujoco.mj_forward(model, data)\n",
"\n",
"while data.time < duration:\n",
" # append data to the traces\n",
" positions.append(data.geom_xpos[data.geom(\"green_sphere\").id].copy())\n",
" times.append(data.time)\n",
" speeds.append(get_geom_speed(model, data, \"green_sphere\"))\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data)\n",
" modify_scene(renderer.scene)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"with mujoco.Renderer(model) as renderer:\n",
" while data.time < duration:\n",
" # append data to the traces\n",
" positions.append(data.geom_xpos[data.geom(\"green_sphere\").id].copy())\n",
" times.append(data.time)\n",
" speeds.append(get_geom_speed(model, data, \"green_sphere\"))\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data)\n",
" modify_scene(renderer.scene)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"\n",
"media.show_video(frames, fps=framerate)"
]
@@ -2018,8 +2034,7 @@
"</mujoco>\n",
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(dominos_xml)\n",
"data = mujoco.MjData(model)\n",
"renderer = mujoco.Renderer(model, height=1024, width=1440)"
"data = mujoco.MjData(model)\n"
]
},
{
@@ -2035,16 +2050,19 @@
"\n",
"duration = 2.5 # (seconds)\n",
"framerate = 60 # (Hz)\n",
"height = 1024\n",
"width = 1440\n",
"\n",
"# Simulate and display video.\n",
"frames = []\n",
"mujoco.mj_resetData(model, data) # Reset state and time.\n",
"while data.time < duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data, camera='top')\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"with mujoco.Renderer(model, height, width) as renderer:\n",
" while data.time < duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate:\n",
" renderer.update_scene(data, camera='top')\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"\n",
"media.show_video(frames, fps=framerate)"
]
@@ -2061,6 +2079,8 @@
"#@title Render from moving camera\n",
"\n",
"duration = 3 # (seconds)\n",
"height = 1024\n",
"width = 1440\n",
"\n",
"# find time when box is thrown (speed > 2cm/s)\n",
"throw_time = 0.0\n",
@@ -2115,13 +2135,14 @@
"slowdown = 4 # 4x slow-down\n",
"mujoco.mj_resetData(model, data)\n",
"frames = []\n",
"while data.time < duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate * slowdown:\n",
" cam.distance, cam.azimuth, cam.elevation, cam.lookat = cam_motion()\n",
" renderer.update_scene(data, cam)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"with mujoco.Renderer(model, height, width) as renderer:\n",
" while data.time < duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) < data.time * framerate * slowdown:\n",
" cam.distance, cam.azimuth, cam.elevation, cam.lookat = cam_motion()\n",
" renderer.update_scene(data, cam)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"\n",
"media.show_video(frames, fps=framerate)"
]