import distutils.util import subprocess if subprocess.run('nvidia-smi').returncode: raise RuntimeError( 'Cannot communicate with GPU. ' 'Make sure you are using a GPU Colab runtime. ' 'Go to the Runtime menu and select Choose runtime type.') # Configure MuJoCo to use the EGL rendering backend (requires GPU) # use WGL for Windows print('Setting environment variable to use GPU rendering:') %env MUJOCO_GL=egl
try: print('Checking that the installation succeeded:') import mujoco mujoco.MjModel.from_xml_string('<mujoco/>') except Exception as e: raise e from RuntimeError( 'Something went wrong during installation. Check the shell output above ' 'for more information.\n' 'If using a hosted Colab runtime, make sure you enable GPU acceleration ' 'by going to the Runtime menu and selecting "Choose runtime type".')
print('Installation successful.')
导入相关图形和绘图包
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
#@title Import packages for plotting and creating graphics import time import itertools import numpy as np from typing importCallable, NamedTuple, Optional, Union, List
# Graphics and plotting. print('Installing mediapy:') !command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg) !pip install -q mediapy import mediapy as media import matplotlib.pyplot as plt
# More legible printing from numpy. np.set_printoptions(precision=3, suppress=True, linewidth=100)
id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'green_sphere') model.geom_rgba[id, :]
array([0., 1., 0., 1.], dtype=float32)
类似地, 只读的 id 和name 属性可用于从 id 转换到 name 以及反之:
1 2 3
print('id of "green_sphere": ', model.geom('green_sphere').id) print('name of geom 1: ', model.geom(1).name) print('name of body 0: ', model.body(0).name)
id of "green_sphere": 1 name of geom 1: green_sphere name of body 0: world
注意, mjModel 实例中的所有值都是可写入的. 虽然通常不建议这样做, 而是建议更改 XML 中的值, 因为很容易创建无效的模型, 但有些值是可以安全写入的, 例如颜色:
1 2 3 4
# Run this cell multiple times for different colors model.geom('red_box').rgba[:3] = np.random.rand(3) renderer.update_scene(data) media.show_image(renderer.render())
# Simulate and save data mujoco.mj_resetDataKeyframe(model, data, 0) while data.time < duration: mujoco.mj_step(model, data) timevals.append(data.time) angular_velocity.append(data.qvel[3:6].copy()) stem_height.append(data.geom_xpos[2,2]);
# visualize contact frames and forces, make body transparent options = mujoco.MjvOption() mujoco.mjv_defaultOption(options) options.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True options.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = True options.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = True
# tweak scales of contact visualization elements model.vis.scale.contactwidth = 0.1 model.vis.scale.contactheight = 0.03 model.vis.scale.forcewidth = 0.05 model.vis.map.force = 0.3
# random initial rotational velocity: mujoco.mj_resetData(model, data) data.qvel[3:6] = 5*np.random.randn(3)
# simulate and render for i inrange(n_frames): while data.time < i/120.0: #1/4x real time mujoco.mj_step(model, data) renderer.update_scene(data, "track", options) frame = renderer.render() frames.append(frame)
# random initial rotational velocity: mujoco.mj_resetData(model, data) data.qvel[3:6] = 2*np.random.randn(3)
# simulate and save data for i inrange(n_steps): mujoco.mj_step(model, data) sim_time[i] = data.time ncon[i] = data.ncon velocity[i] = data.qvel[:] acceleration[i] = data.qacc[:] # iterate over active contacts, save force and distance for j,c inenumerate(data.contact): mujoco.mj_contactForce(model, data, j, forcetorque) force[i] += forcetorque[0:3] penetration[i] = min(penetration[i], c.dist) # we could also do # force[i] += data.qfrc_constraint[0:3] # do you see why?
# update renderer to render depth renderer.enable_depth_rendering()
# reset the scene renderer.update_scene(data)
# depth is a float array, in meters. depth = renderer.render()
# Shift nearest values to the origin. depth -= depth.min() # Scale by 2 mean distances of near rays. depth /= 2*depth[depth <= 1].mean() # Scale to [0, 255] pixels = 255*np.clip(depth, 0, 1)
defcompute_camera_matrix(renderer, data): """Returns the 3x4 camera matrix.""" # If the camera is a 'free' camera, we get its position and orientation # from the scene data structure. It is a stereo camera, so we average over # the left and right channels. Note: we call `self.update()` in order to # ensure that the contents of `scene.camera` are correct. renderer.update_scene(data) pos = np.mean([camera.pos for camera in renderer.scene.camera], axis=0) z = -np.mean([camera.forward for camera in renderer.scene.camera], axis=0) y = np.mean([camera.up for camera in renderer.scene.camera], axis=0) rot = np.vstack((np.cross(y, z), y, z)) fov = model.vis.global_.fovy
# Get the world coordinates of the box corners box_pos = data.geom_xpos[model.geom('red_box').id] box_mat = data.geom_xmat[model.geom('red_box').id].reshape(3, 3) box_size = model.geom_size[model.geom('red_box').id] offsets = np.array([-1, 1]) * box_size[:, None] xyz_local = np.stack(list(itertools.product(*offsets))).T xyz_global = box_pos[:, None] + box_mat @ xyz_local
# Camera matrices multiply homogenous [x, y, z, 1] vectors. corners_homogeneous = np.ones((4, xyz_global.shape[1]), dtype=float) corners_homogeneous[:3, :] = xyz_global
# Get the camera matrix. m = compute_camera_matrix(renderer, data)
# Project world coordinates into pixel space. See: # https://en.wikipedia.org/wiki/3D_projection#Mathematical_formula xs, ys, s = m @ corners_homogeneous # x and y are in the pixel coordinate system. x = xs / s y = ys / s
# Render the camera view and overlay the projected corner coordinates. pixels = renderer.render() fig, ax = plt.subplots(1, 1) ax.imshow(pixels) ax.plot(x, y, '+', c='w') ax.set_axis_off()
defget_geom_speed(model, data, geom_name): """Returns the speed of a geom.""" geom_vel = np.zeros(6) geom_type = mujoco.mjtObj.mjOBJ_GEOM geom_id = data.geom(geom_name).id mujoco.mj_objectVelocity(model, data, geom_type, geom_id, geom_vel, 0) return np.linalg.norm(geom_vel)
defadd_visual_capsule(scene, point1, point2, radius, rgba): """Adds one capsule to an mjvScene.""" if scene.ngeom >= scene.maxgeom: return scene.ngeom += 1# increment ngeom # initialise a new capsule, add it to the scene using mjv_makeConnector mujoco.mjv_initGeom(scene.geoms[scene.ngeom-1], mujoco.mjtGeom.mjGEOM_CAPSULE, np.zeros(3), np.zeros(3), np.zeros(9), rgba.astype(np.float32)) mujoco.mjv_makeConnector(scene.geoms[scene.ngeom-1], mujoco.mjtGeom.mjGEOM_CAPSULE, radius, point1[0], point1[1], point1[2], point2[0], point2[1], point2[2])
# traces of time, position and speed times = [] positions = [] speeds = [] offset = model.jnt_axis[0]/8# offset along the joint axis
defmodify_scene(scn): """Draw position trace, speed modifies width and colors.""" iflen(positions) > 1: for i inrange(len(positions)-1): rgba=np.array((np.clip(speeds[i]/10, 0, 1), np.clip(1-speeds[i]/10, 0, 1), .5, 1.)) radius=.003*(1+speeds[i]) point1 = positions[i] + offset*times[i] point2 = positions[i+1] + offset*times[i+1] add_visual_capsule(scn, point1, point2, radius, rgba)
duration = 6# (seconds) framerate = 30# (Hz)
# Simulate and display video. frames = []
# Reset state and time. mujoco.mj_resetData(model, data) mujoco.mj_forward(model, data)
while data.time < duration: # append data to the traces positions.append(data.geom_xpos[data.geom("green_sphere").id].copy()) times.append(data.time) speeds.append(get_geom_speed(model, data, "green_sphere")) mujoco.mj_step(model, data) iflen(frames) < data.time * framerate: renderer.update_scene(data) modify_scene(renderer.scene) pixels = renderer.render() frames.append(pixels) media.show_video(frames, fps=framerate)
