# Set up GPU rendering. from google.colab import files import distutils.util import os 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.')
# Add an ICD config so that glvnd can pick up the Nvidia EGL driver. # This is usually installed as part of an Nvidia driver package, but the Colab # kernel doesn't install its driver via APT, and as a result the ICD is missing. # (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md) NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json' ifnot os.path.exists(NVIDIA_ICD_CONFIG_PATH): withopen(NVIDIA_ICD_CONFIG_PATH, 'w') as f: f.write("""{ "file_format_version" : "1.0.0", "ICD" : { "library_path" : "libEGL_nvidia.so.0" } } """)
# Configure MuJoCo to use the EGL rendering backend (requires GPU) print('Setting environment variable to use GPU rendering:') %env MUJOCO_GL=egl
# Check if installation was succesful. try: print('Checking that the installation succeeded:') import mujoco as mj mj.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.')
# Other imports and helper functions import time import itertools import numpy as np
# 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)
from IPython.display import clear_output, HTML, display clear_output()
# Get MuJoCo's humanoid model and a Franka arm from the MuJoCo Menagerie. print('Getting MuJoCo humanoid XML description from GitHub:') !git clone https://github.com/google-deepmind/mujoco humanoid_file = 'mujoco/model/humanoid/humanoid.xml' humanoid100_file = 'mujoco/model/humanoid/humanoid100.xml' print('Getting MuJoCo Menagerie Franka XML description from GitHub:') !git clone https://github.com/google-deepmind/mujoco_menagerie franka_file = 'mujoco_menagerie/franka_fr3/fr3.xml'
defrender(model, data=None, height=250): if data isNone: data = mj.MjData(model) with mj.Renderer(model, 480, 640) as renderer: mj.mj_forward(model, data) renderer.update_scene(data) media.show_image(renderer.render(), height=height)
# Make arena with textured floor. chequered = arena.add_texture( name="chequered", type=mj.mjtTexture.mjTEXTURE_2D, builtin=mj.mjtBuiltin.mjBUILTIN_CHECKER, width=300, height=300, rgb1=[.2, .3, .4], rgb2=[.3, .4, .5]) grid = arena.add_material( name='grid', texrepeat=[5, 5], reflectance=.2 ).textures[mj.mjtTextureRole.mjTEXROLE_RGB] = 'chequered' arena.worldbody.add_geom( type=mj.mjtGeom.mjGEOM_PLANE, size=[2, 2, .1], material='grid') for x in [-2, 2]: arena.worldbody.add_light(pos=[x, -1, 3], dir=[-x, 1, -2])
# Instantiate 6 creatures with 3 to 8 legs. creatures = [make_creature(num_legs=num_legs) for num_legs inrange(3, 9)]
# Place them on a grid in the arena. height = .15 grid = 5 * BODY_RADIUS xpos, ypos, zpos = np.meshgrid([-grid, 0, grid], [0, grid], [height]) for i, spec inenumerate(creatures): # Place spawn sites on a grid. spawn_pos = (xpos.flat[i], ypos.flat[i], zpos.flat[i]) spawn_site = arena.worldbody.add_site(pos=spawn_pos, group=3) # Attach to the arena at the spawn sites, with a free joint. spawn_body = spawn_site.attach_body(spec.worldbody, '', '-' + str(i)) spawn_body.add_freejoint()
# Instantiate the physics and render. model = arena.compile() render(model)
# Video of the movement data = mj.MjData(model) duration = 10# (Seconds) framerate = 30# (Hz) video = [] pos_x = [] pos_y = [] geoms = arena.worldbody.find_all(mj.mjtObj.mjOBJ_GEOM) torsos_data = [data.bind(geom) for geom in geoms if'torso'in geom.name] torsos_model = [model.bind(geom) for geom in geoms if'torso'in geom.name] actuators = [data.bind(actuator) for actuator in arena.actuators]
# Control signal frequency, phase, amplitude. freq = 5 phase = 2 * np.pi * random_state.rand(len(arena.actuators)) amp = 0.9
# Simulate, saving video frames and torso locations. mj.mj_resetData(model, data) with mj.Renderer(model) as renderer: while data.time < duration: # Inject controls and step the physics. for i, actuator inenumerate(actuators): actuator.ctrl = amp * np.sin(freq * data.time + phase[i]) mj.mj_step(model, data)
# Save torso horizontal positions using name indexing. pos_x.append([torso.xpos[0] for torso in torsos_data]) pos_y.append([torso.xpos[1] for torso in torsos_data])
# Save video frames. iflen(video) < data.time * framerate: renderer.update_scene(data) pixels = renderer.render() video.append(pixels.copy())
media.show_video(video, fps=framerate)
运动轨迹
1 2 3 4 5
# Movement trajectories creature_colors = [torso.rgba[:3] for torso in torsos_model] fig, ax = plt.subplots(figsize=(4, 4)) ax.set_prop_cycle(color=creature_colors) _ = ax.plot(pos_x, pos_y, linewidth=4)
# Traversing the spec. spec = mj.MjSpec.from_file(humanoid_file)
# Function that recursively prints all body names defprint_bodies(parent, level=0): body = parent.first_body() while body: print(''.join(['-'for i inrange(level)]) + body.name) print_bodies(body, level + 1) body = parent.next_body(body)
print("The spec has the following actuators:") for actuator in spec.actuators: print(actuator.name)
print("\nThe spec has the following bodies:") print_bodies(spec.worldbody)
The spec has the following actuators: abdomen_z abdomen_y abdomen_x hip_x_right hip_z_right hip_y_right knee_right ankle_y_right ankle_x_right hip_x_left hip_z_left hip_y_left knee_left ankle_y_left ankle_x_left shoulder1_right shoulder2_right elbow_right shoulder1_left shoulder2_left elbow_left
The spec has the following bodies: torso -head -waist_lower --pelvis ---thigh_right ----shin_right -----foot_right ---thigh_left ----shin_left -----foot_left -upper_arm_right --lower_arm_right ---hand_right -upper_arm_left --lower_arm_left ---hand_left
# Model re-compilation with state preservation. spec = mj.MjSpec.from_file(humanoid100_file) model = spec.compile() data = mj.MjData(model)
# Run for 5 seconds for i inrange(1000): mj.mj_step(model, data)
# Show result render(model, data)
# Create list of all bodies we want to delete body = spec.worldbody.first_body() delete_list = [] while body: geom_type = body.first_geom().type if (geom_type == mj.mjtGeom.mjGEOM_BOX or geom_type == mj.mjtGeom.mjGEOM_ELLIPSOID): delete_list.append(body) body = spec.worldbody.next_body(body)
# Remove all bodies in the list from the spec for body in delete_list: spec.detach_body(body)
#@title Humanoid with randomized heads and arm poses.{vertical-output: true} humanoid = mj.MjSpec.from_file(humanoid_file) spec = mj.MjSpec()
# Delete all key frames to avoid name conflicts while humanoid.keys: humanoid.keys[-1].delete()
# Create a grid of humanoids by attaching humanoid to spec multiple times for i inrange(4): for j inrange(4): humanoid.materials[0].rgba = [ np.random.uniform(), np.random.uniform(), np.random.uniform(), 1] # Randomize color humanoid.find_body('head').first_geom().size = [ .18*np.random.uniform(), 0, 0] # Randomize head size humanoid.find_body('upper_arm_left').quat = [ np.random.uniform(), np.random.uniform(), np.random.uniform(), np.random.uniform()] # Randomize left arm orientation humanoid.find_body('upper_arm_right').quat = [ np.random.uniform(), np.random.uniform(), np.random.uniform(), np.random.uniform()] # Randomize right arm orientation
