install miniconda homebrew brew install python brew install glfw egl ffmpeg sudo chown -R vector ~/.conda conda create -n agents conda activate agents conda install pip pip install dm_control pip install mediapy pip install matplotlib ``` from dm_control import mjcf from dm_control import suite, viewer import numpy as np #import matplotlib as plt from PIL import Image import mediapy class Leg(object): def __init__(self, length, rgba): self.model = mjcf.RootElement() self.model.default.joint.damping = 2 self.model.default.joint.type = 'hinge' self.model.default.geom.type = 'capsule' self.model.default.geom.rgba = rgba self.thigh = self.model.worldbody.add('body') self.hip = self.thigh.add('joint', axis=[0, 0, 1]) self.thigh.add('geom', fromto=[0, 0, 0, length, 0, 0], size=[length/4]) self.shin = self.thigh.add('body', pos=[length, 0, 0]) self.knee = self.shin.add('joint', axis=[0, 1, 0]) self.shin.add('geom', fromto=[0, 0, 0, 0, 0, -length], size=[length/5]) self.model.actuator.add('position', joint=self.hip, kp=10) self.model.actuator.add('position', joint=self.knee, kp=10) BODY_RADIUS = 0.1 BODY_SIZE = (BODY_RADIUS, BODY_RADIUS, BODY_RADIUS / 2) def make_creature(num_legs): rgba = np.random.uniform([0,0,0,1], [1,1,1,1]) model = mjcf.RootElement() model.compiler.angle = 'radian' torso = model.worldbody.add('geom', name='torso', type='ellipsoid', size=BODY_SIZE, rgba=rgba) for i in range(num_legs): theta = 2 * i * np.pi / num_legs hip_pos = BODY_RADIUS * np.array([np.cos(theta), np.sin(theta), 0]) hip_site = model.worldbody.add('site', pos=hip_pos, euler=[0, 0, theta]) leg = Leg(length=BODY_RADIUS, rgba=rgba) hip_site.attach(leg.model) return model # Setup the arena / floor arena = mjcf.RootElement() checker = arena.asset.add('texture', type='2d', builtin='checker', width=300, height=300, rgb1=[.2, .3, .4], rgb2=[.3, .4, .5]) grid = arena.asset.add('material', name='grid', texture=checker, texrepeat=[5,5], reflectance=0.2) arena.worldbody.add('geom', type='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 the creatures creatures = [make_creature(num_legs=num_legs) for num_legs in (3,4,5,6,7,8)] height = 0.15 grid = 5 * BODY_RADIUS xpos, ypos, zpos = np.meshgrid([-grid, 0, grid], [0, grid], [height]) for i, model in enumerate(creatures): spawn_pos = (xpos.flat[i], ypos.flat[i], zpos.flat[i]) spawn_site = arena.worldbody.add('site', pos=spawn_pos, group=3) # ??? spawn_site.attach(model).add('freejoint') physics = mjcf.Physics.from_mjcf_model(arena) pixels = physics.render() img = Image.fromarray(pixels.astype('uint8'), 'RGB') img.save('/tmp/dm_control_tutorial.png') duration = 10 # seconds framerate = 30 # Hz video = []; pos_x = []; pos_y = [] torsos = [] actuators = [] for creature in creatures: torsos.append(creature.find('geom', 'torso')) actuators.extend(creature.find_all('actuator')) freq = 8 phase = 2 * np.pi * np.random.rand(len(actuators)) amp = 0.9 physics.reset() while physics.data.time < duration: # inject control signals physics.bind(actuators).ctrl = amp * np.sin(freq*physics.data.time + phase) physics.step() pos_x.append(physics.bind(torsos).pos[:, 0].copy()) pos_y.append(physics.bind(torsos).pos[:, 1].copy()) if len(video) < physics.data.time * framerate: pixels = physics.render(height=480, width=640) video.append(pixels.copy()) if pixels.dtype != np.uint8 or np.min(pixels) < 0 or np.max(pixels) > 255: pixels = pixels.astype(np.uint8) # ensure uint8 type pixels = np.clip(pixels, 0, 255) # clamp to 0-255 mediapy.write_video('/tmp/test_video.mp4', video, fps=10, codec='h264') # creature_colors = physics.bind(torsos).rgba[:, :3] # fig, ax = plt.subplots(figsize=(8,8)) # ax.set_prop_cycle(color=creature_colors) # ax.plot(pos_x, pos_y, linewidth=4) ```