class SimManager(object): """Object to generate sequence of images with their transforms""" def __init__(self, filepath, random_params={}, gpu_render=False, gui=False, display_data=False): self.model = load_model_from_path(filepath) self.sim = MjSim(self.model) self.filepath = filepath self.gui = gui self.display_data = display_data # Take the default random params and update anything we need self.RANDOM_PARAMS = {} self.RANDOM_PARAMS.update(random_params) if gpu_render: self.viewer = MjViewer(self.sim) else: self.viewer = None # Get start state of params to slightly jitter later self.START_GEOM_POS = self.model.geom_pos.copy() self.START_GEOM_SIZE = self.model.geom_size.copy() self.START_GEOM_QUAT = self.model.geom_quat.copy() self.START_BODY_POS = self.model.body_pos.copy() self.START_BODY_QUAT = self.model.body_quat.copy() self.START_MATID = self.model.geom_matid.copy() #self.FLOOR_OFFSET = self.model.body_pos[self.model.body_name2id('floor')] self.tex_modder = TextureModder(self.sim) self.tex_modder.whiten_materials( ) # ensures materials won't impact colors self.cam_modder = CameraModder(self.sim) self.light_modder = LightModder(self.sim) self.start_obj_pose = self.sim.data.get_joint_qpos( 'object:joint').copy() def get_data(self, num_images=10): """ Returns camera intrinsics, and a sequence of images, pose transforms, and camera transforms """ # randomize the scene self._rand_textures() self._rand_lights() self._rand_object() self._rand_walls() self._rand_distract() sequence = defaultdict(list) context = {} # object pose obj_pose, robot_pose = self._get_ground_truth() context["obj_world_pose"] = obj_pose context["robot_world_pose"] = robot_pose self._cam_step = 0 self._cam_choices = np.array([[-1.75, 0, 2], [-1.75, 0, 1.62], [-1.3, 1.7, 1.62]]) self._curr_cam_pos = self._cam_choices[0] for i in range(num_images): self._next_camera() self._forward() img = self._get_cam_frame() sequence["img"].append(img) cam_pos = self.cam_modder.get_pos("camera1") cam_quat = self.cam_modder.get_quat("camera1") cam_pose = np.concatenate([cam_pos, cam_quat]).astype(np.float32) sequence["cam_pose"].append(cam_pose) cam_id = self.cam_modder.get_camid("camera1") fovy = self.sim.model.cam_fovy[cam_id] width, height = 640, 480 f = 0.5 * height / math.tan(fovy * math.pi / 360) camera_intrinsics = np.array( ((f, 0, width / 2), (0, f, height / 2), (0, 0, 1))) context['cam_matrix'] = camera_intrinsics return context, sequence def _forward(self): """Advances simulator a step (NECESSARY TO MAKE CAMERA AND LIGHT MODDING WORK) And add some visualization""" self.sim.forward() if self.viewer and self.gui: # Get angle of camera and display it quat = np.quaternion(*self.model.cam_quat[0]) ypr = quaternion.as_euler_angles(quat) * 180 / np.pi cam_pos = self.model.cam_pos[0] cam_fovy = self.model.cam_fovy[0] #self.viewer.add_marker(pos=cam_pos, label="CAM: {}{}".format(cam_pos, ypr)) #self.viewer.add_marker(pos=cam_pos, label="CAM: {}".format(ypr)) #self.viewer.add_marker(pos=cam_pos, label="CAM: {}".format(cam_pos)) self.viewer.add_marker(pos=cam_pos, label="FOVY: {}, CAM: {}".format( cam_fovy, cam_pos)) self.viewer.render() def _get_ground_truth(self): """ Return the position to the robot, and quaternion to the robot quaternion 7 dim total """ robot_gid = self.sim.model.geom_name2id('base_link') obj_gid = self.sim.model.geom_name2id('object') # only x and y pos needed obj_world_pos = self.sim.data.geom_xpos[obj_gid] robot_world_pos = self.sim.data.geom_xpos[robot_gid] # obj_pos_in_robot_frame = (self.sim.data.geom_xpos[obj_gid] - self.sim.data.geom_xpos[robot_gid])[:2] # robot_quat = quaternion.as_quat_array(self.model.geom_quat[robot_gid].copy()) obj_world_quat = self.model.geom_quat[obj_gid].copy() robot_world_quat = self.model.geom_quat[robot_gid].copy() # # want quat of obj relative to robot frame # # obj_q = robot_q * localrot # # robot_q.inv * obj_q = localrot # rel_quat = quaternion.as_float_array(robot_quat.inverse() * obj_quat) # pose = np.concatenate([obj_pos_in_robot_frame, rel_quat]).astype(np.float32) obj_pose = self.sim.data.get_joint_qpos('object:joint').copy() robot_pose = np.concatenate([robot_world_pos, robot_world_quat]).astype(np.float32) return obj_pose, robot_pose def _get_cam_frame(self, ground_truth=None): """Grab an image from the camera (224, 244, 3) to feed into CNN""" #IMAGE_NOISE_RVARIANCE = Range(0.0, 0.0001) cam_img = self.sim.render( 640, 480, camera_name='camera1' )[::-1, :, :] # Rendered images are upside-down. # make camera crop be more like kinect #cam_img = self.sim.render(854, 480, camera_name='camera1')[::-1, 107:-107, :] # Rendered images are upside-down. #image_noise_variance = sample(IMAGE_NOISE_RVARIANCE) #cam_img = (skimage.util.random_noise(cam_img, mode='gaussian', var=image_noise_variance) * 255).astype(np.uint8) if self.display_data: print(ground_truth) #label = str(ground_truth[3:6]) display_image(cam_img, mode='preproc') #, label) # cam_img = preproc_image(cam_img) return cam_img def _randomize(self): self._rand_textures() self._rand_camera() self._rand_lights() #self._rand_robot() self._rand_object() self._rand_walls() self._rand_distract() def _rand_textures(self): """Randomize all the textures in the scene, including the skybox""" bright = np.random.binomial(1, 0.5) for name in self.sim.model.geom_names + ('skybox', ): self.tex_modder.rand_all(name) if bright: self.tex_modder.brighten(name, np.random.randint(0, 150)) def _rand_camera(self): """Randomize pos, orientation, and fov of camera real camera pos is -1.75, 0, 1.62 FOVY: Kinect2 is 53.8 ASUS is 45 https://www.asus.com/us/3D-Sensor/Xtion_PRO_LIVE/specifications/ http://smeenk.com/kinect-field-of-view-comparison/ """ # Params FOVY_R = Range(40, 50) #X = Range(-3, -1) #Y = Range(-1, 3) #Z = Range(1, 2) #C_R3D = Range3D(X, Y, Z) #cam_pos = sample_xyz(C_R3D) #L_R3D = rto3d([-0.1, 0.1]) C_R3D = Range3D([-0.07, 0.07], [-0.07, 0.07], [-0.07, 0.07]) ANG3 = Range3D([-3, 3], [-3, 3], [-3, 3]) # Look approximately at the robot, but then randomize the orientation around that cam_choices = np.array([[-1.75, 0, 1.62], [-1.3, 1.7, 1.62], [-1.75, 0, 2]]) cam_pos = cam_choices[np.random.choice(len(cam_choices))] # cam_pos = get_real_cam_pos(FLAGS.real_data_path) target_id = self.model.body_name2id(FLAGS.look_at) cam_off = 0 #sample_xyz(L_R3D) target_off = 0 #sample_xyz(L_R3D) quat = look_at(cam_pos + cam_off, self.sim.data.body_xpos[target_id] + target_off) quat = jitter_angle(quat, ANG3) #quat = jitter_quat(quat, 0.01) cam_pos += sample_xyz(C_R3D) self.cam_modder.set_quat('camera1', quat) self.cam_modder.set_pos('camera1', cam_pos) self.cam_modder.set_fovy('camera1', 60) # hard code to wide fovy def _next_camera(self): """Randomize pos, orientation, and fov of camera real camera pos is -1.75, 0, 1.62 FOVY: Kinect2 is 53.8 ASUS is 45 https://www.asus.com/us/3D-Sensor/Xtion_PRO_LIVE/specifications/ http://smeenk.com/kinect-field-of-view-comparison/ """ # Params FOVY_R = Range(40, 50) #X = Range(-3, -1) #Y = Range(-1, 3) #Z = Range(1, 2) #C_R3D = Range3D(X, Y, Z) #cam_pos = sample_xyz(C_R3D) #L_R3D = rto3d([-0.1, 0.1]) C_R3D = Range3D([-0.07, 0.07], [-0.07, 0.07], [-0.07, 0.07]) ANG3 = Range3D([-3, 3], [-3, 3], [-3, 3]) # Look approximately at the robot, but then randomize the orientation around that # linearly interpolate to the next camera every K steps K = 5 goal_cam_pos = self._cam_choices[(self._cam_step // K) + 1] offset = goal_cam_pos - self._curr_cam_pos offset *= (self._cam_step % K) / K self._curr_cam_pos += offset cam_pos = self._curr_cam_pos self._cam_step += 1 # cam_pos = cam_choices[np.random.choice(len(cam_choices))] # cam_pos = get_real_cam_pos(FLAGS.real_data_path) target_id = self.model.body_name2id(FLAGS.look_at) cam_off = 0 #sample_xyz(L_R3D) target_off = 0 #sample_xyz(L_R3D) quat = look_at(cam_pos + cam_off, self.sim.data.body_xpos[target_id] + target_off) quat = jitter_angle(quat, ANG3) #quat = jitter_quat(quat, 0.01) cam_pos += sample_xyz(C_R3D) self.cam_modder.set_quat('camera1', quat) self.cam_modder.set_pos('camera1', cam_pos) self.cam_modder.set_fovy('camera1', 60) # hard code to wide fovy def _rand_lights(self): """Randomize pos, direction, and lights""" # light stuff #X = Range(-1.5, 1.5) #Y = Range(-1.2, 1.2) #Z = Range(0, 2.8) X = Range(-1.5, -0.5) Y = Range(-0.6, 0.6) Z = Range(1.0, 1.5) LIGHT_R3D = Range3D(X, Y, Z) LIGHT_UNIF = Range3D(Range(0, 1), Range(0, 1), Range(0, 1)) # TODO: also try not altering the light dirs and just keeping them at like -1, or [0, -0.15, -1.0] for i, name in enumerate(self.model.light_names): lid = self.model.light_name2id(name) # random sample 80% of any given light being on if lid != 0: self.light_modder.set_active(name, sample([0, 1]) < 0.8) self.light_modder.set_dir(name, sample_light_dir()) self.light_modder.set_pos(name, sample_xyz(LIGHT_R3D)) #self.light_modder.set_dir(name, sample_xyz(rto3d([-1,1]))) #self.light_modder.set_specular(name, sample_xyz(LIGHT_UNIF)) #self.light_modder.set_diffuse(name, sample_xyz(LIGHT_UNIF)) #self.light_modder.set_ambient(name, sample_xyz(LIGHT_UNIF)) spec = np.array([sample(Range(0.5, 1))] * 3) diffuse = np.array([sample(Range(0.5, 1))] * 3) ambient = np.array([sample(Range(0.5, 1))] * 3) self.light_modder.set_specular(name, spec) self.light_modder.set_diffuse(name, diffuse) self.light_modder.set_ambient(name, ambient) #self.model.light_directional[lid] = sample([0,1]) < 0.2 self.model.light_castshadow[lid] = sample([0, 1]) < 0.5 def _rand_robot(self): """Randomize joint angles and jitter orientation""" jnt_shape = self.sim.data.qpos.shape self.sim.data.qpos[:] = sample_joints(self.model.jnt_range, jnt_shape) robot_gid = self.model.geom_name2id('robot_table_link') self.model.geom_quat[robot_gid] = jitter_quat( self.START_GEOM_QUAT[robot_gid], 0.01) def _rand_object(self): obj_gid = self.sim.model.geom_name2id('object') obj_bid = self.sim.model.geom_name2id('object') table_gid = self.model.geom_name2id('object_table') table_bid = self.model.body_name2id('object_table') obj_pose = self.start_obj_pose.copy() xval = self.model.geom_size[table_gid][ 0] #- self.model.geom_size[obj_gid][0] yval = self.model.geom_size[table_gid][ 1] #- self.model.geom_size[obj_gid][1] O_X = Range(-xval, xval) O_Y = Range(-yval, yval) O_Z = Range(0, 0) O_R3D = Range3D(O_X, O_Y, O_Z) newpos = obj_pose[:3] + sample_xyz(O_R3D) newquat = jitter_quat(obj_pose[3:], 0.1) obj_pose[:3] = newpos obj_pose[3:] = newquat self.sim.data.set_joint_qpos('object:joint', obj_pose) #T_X = Range(-0.1, 0.1) #T_Y = Range(-0.1, 0.1) #T_Z = Range(-0.1, 0.1) #T_R3D = Range3D(T_X, T_Y, T_Z) #self.model.body_pos[table_bid] = self.START_BODY_POS[table_bid] + sample_xyz(T_R3D) ## randomize orientation a wee bit #self.model.geom_quat[table_gid] = jitter_quat(self.START_GEOM_QUAT[table_gid], 0.01) def _rand_walls(self): wall_bids = { name: self.model.body_name2id(name) for name in ['wall_' + dir for dir in 'nesw'] } window_gid = self.model.geom_name2id('west_window') #floor_gid = self.model.geom_name2id('floor') WA_X = Range(-0.2, 0.2) WA_Y = Range(-0.2, 0.2) WA_Z = Range(-0.1, 0.1) WA_R3D = Range3D(WA_X, WA_Y, WA_Z) WI_X = Range(-0.1, 0.1) WI_Y = Range(0, 0) WI_Z = Range(-0.5, 0.5) WI_R3D = Range3D(WI_X, WI_Y, WI_Z) R = Range(0, 0) P = Range(-10, 10) Y = Range(0, 0) RPY_R = Range3D(R, P, Y) #self.model.geom_quat[floor_gid] = jitter_quat(self.START_GEOM_QUAT[floor_gid], 0.01) #self.model.geom_pos[floor_gid] = self.START_GEOM_POS[floor_gid] + [0,0,sample(-0.1,0.1) self.model.geom_quat[window_gid] = sample_quat(RPY_R) #self.model.geom_quat[window_gid] = jitter_quat(self.START_GEOM_QUAT[window_gid], 0.01) self.model.geom_pos[ window_gid] = self.START_GEOM_POS[window_gid] + sample_xyz(WI_R3D) for name in wall_bids: gid = wall_bids[name] self.model.body_quat[gid] = jitter_quat(self.START_BODY_QUAT[gid], 0.01) self.model.body_pos[gid] = self.START_BODY_POS[gid] + sample_xyz( WA_R3D) def _rand_distract(self): PREFIX = 'distract' geom_names = [ name for name in self.model.geom_names if name.startswith(PREFIX) ] # Size range SX = Range(0.01, 0.5) SY = Range(0.01, 0.9) SZ = Range(0.01, 0.5) S3D = Range3D(SX, SY, SZ) # Back range B_PX = Range(-0.5, 2) B_PY = Range(-1.5, 2) B_PZ = Range(0, 3) B_P3D = Range3D(B_PX, B_PY, B_PZ) # Front range F_PX = Range(-2, -0.5) F_PY = Range(-2, 1) F_PZ = Range(0, 0.5) F_P3D = Range3D(F_PX, F_PY, F_PZ) for name in geom_names: gid = self.model.geom_name2id(name) range = B_P3D if np.random.binomial(1, 0.5) else F_P3D self.model.geom_pos[gid] = sample_xyz(range) self.model.geom_quat[gid] = random_quat() self.model.geom_size[gid] = sample_xyz(S3D, mode='logspace') self.model.geom_type[gid] = sample_geom_type() self.model.geom_rgba[gid][-1] = np.random.binomial(1, 0.5) def _set_visible(self, prefix, range_top, visible): """Helper function to set visibility of several objects""" if not visible: if range_top == 0: name = prefix gid = self.model.geom_name2id(name) self.model.geom_rgba[gid][-1] = 0.0 for i in range(range_top): name = "{}{}".format(prefix, i) gid = self.model.geom_name2id(name) self.model.geom_rgba[gid][-1] = 0.0 else: if range_top == 0: name = prefix gid = self.model.geom_name2id(name) self.model.geom_rgba[gid][-1] = 1.0 for i in range(range_top): name = "{}{}".format(prefix, i) gid = self.model.geom_name2id(name) self.model.geom_rgba[gid][-1] = 1.0
<?xml version="1.0" ?> <mujoco> <worldbody> <body name="box" pos="0 0 0.2"> <geom size="0.15 0.15 0.15" type="sphere"/> <joint axis="1 0 0" name="box:x" type="slide"/> <joint axis="0 1 0" name="box:y" type="slide"/> </body> <body name="floor" pos="0 0 0.025"> <geom size="1.0 1.0 0.02" rgba="0 1 0 1" type="box"/> </body> </worldbody> </mujoco> """ model = load_model_from_xml(MODEL_XML) sim = MjSim(model) viewer = MjViewer(sim) step = 0 while True: t = time.time() x, y = math.cos(t), math.sin(t) viewer.add_marker(pos=np.array([x, y, 1]), label=str(t), type=const.GEOM_SPHERE) viewer.render() step += 1 if step > 100 and os.getenv('TESTING') is not None: break
class GatherEnv: @autoassign(exclude=('model_path')) def __init__(self, model_path, n_apples=8, n_bombs=8, apple_reward=10, bomb_cost=1, activity_range=6.0, catch_range=10, n_bins=10, robot_object_spacing=2.0, sensor_range=6.0, sensor_span=pi): self.viewer = None self.apples = [] self.bombs = [] self.model = self.build_model(model_path) self.sim = MjSim(self.model) self._max_episode_steps = None self._step_num = 0 def build_model(self, agent_xml_path): sim_size = self.activity_range + 1 model_tree = ET.parse(agent_xml_path) worldbody = model_tree.find('.//worldbody') floor_attrs = dict(name='floor', type='plane', material='MatPlane', pos='0 0 0', size=f'{sim_size} {sim_size} {sim_size}', conaffinity='1', rgba='0.8 0.9 0.8 1', condim='3') wall_attrs = dict(type='box', conaffinity='1', rgba='0.8 0.9 0.8 1', condim='3') wall_poses = [f'0 {-sim_size} 0', f'0 {sim_size} 0', \ f'{-sim_size} 0 0', f'{sim_size} 0 0'] wall_sizes = [f'{sim_size + 0.1} 0.1 1', f'{sim_size + 0.1} 0.1 1', \ f'0.1 {sim_size + 0.1} 1', f'0.1 {sim_size + 0.1} 1'] # Add a floor to the model ET.SubElement(worldbody, 'geom', **floor_attrs) # Add walls to the model for i, pos, size in zip(range(4), wall_poses, wall_sizes): ET.SubElement( worldbody, 'geom', dict(**wall_attrs, name=f'wall{i}', pos=pos, size=size)) model_xml = ET.tostring(model_tree.getroot(), encoding='unicode', method='xml') model = load_model_from_xml(model_xml) return model def reset(self): self.sim.reset() self.apples = [] self.bombs = [] obj_coords = [] self._step_num = 0 agent_x_pos, agent_y_pos = self.sim.data.qpos[0], self.sim.data.qpos[1] rand_coord = lambda: np.random.randint(-self.activity_range, self. activity_range) while (len(self.apples) < self.n_apples): x, y = rand_coord(), rand_coord() # Change this later to make (0, 0) the position of the agent if in_range(x, y, agent_x_pos, agent_y_pos, self.robot_object_spacing): continue if (x, y) in obj_coords: continue self.apples.append(Object(APPLE, x, y)) obj_coords.append((x, y)) while (len(self.bombs) < self.n_bombs): x, y = rand_coord(), rand_coord() if in_range(x, y, 0, 0, self.robot_object_spacing): continue if (x, y) in obj_coords: continue self.bombs.append(Object(BOMB, x, y)) obj_coords.append((x, y)) return self._get_obs() def step(self, action): raise NotImplementedError def _do_simulation(self, action): raise NotImplementedError def _update_objects(self): n_apples = 0 n_bombs = 0 agent_x_pos, agent_y_pos = self.sim.data.qpos[0], self.sim.data.qpos[1] for apple in self.apples: if in_range(apple.x, apple.y, agent_x_pos, agent_y_pos, self.catch_range): n_apples += 1 self.apples.remove(apple) for bomb in self.bombs: if in_range(bomb.x, bomb.y, agent_x_pos, agent_y_pos, self.catch_range): n_bombs += 1 self.bombs.remove(bomb) return n_apples, n_bombs def _get_self_obs(self): raise NotImplementedError def _get_sensor_obs(self): apple_readings = np.zeros(self.n_bins) bomb_readings = np.zeros(self.n_bins) idx = self.model.body_names.index('torso') com = self.sim.data.subtree_com[idx].flat agent_x, agent_y = com[:2] sort_key = lambda obj: -euclidian_dist(obj.x, obj.y, agent_x, agent_y) sorted_objs = sorted(self.apples + self.bombs, key=sort_key) bin_res = self.sensor_span / self.n_bins half_span = self.sensor_span / 2 orientation = std_radians(self._get_orientation()) for obj_type, obj_x, obj_y in sorted_objs: dist = euclidian_dist(obj_x, obj_y, agent_x, agent_y) if dist > self.sensor_range: continue # Get the components of the vector from the agent to the object x_delta = obj_x - agent_x y_delta = obj_y - agent_y # Get the angle of the vector relative to the agent's sensor range angle = np.arctan2(y_delta, x_delta) angle = std_radians(angle - orientation + half_span) if angle > self.sensor_span: continue bin_number = int(angle / bin_res) if bin_number == int(angle / bin_res): bin_number -= 1 intensity = 1.0 - dist / self.sensor_range if obj_type == APPLE: apple_readings[bin_number] = intensity else: bomb_readings[bin_number] = intensity sensor_obs = np.concatenate([apple_readings, bomb_readings]) return sensor_obs # dist = euclidian_dist(obj_x, obj_y, agent_x, agent_y) # # if dist > self.sensor_range: # continue # # # Get the components of the vector from the agent to the object # x_comp = obj_x - agent_x # y_comp = obj_y - agent_y # # # Get the angle of the vector relative to the agent's sensor range # angle = np.arctan2(y_comp, x_comp) # angle = standardize_radians(angle - orientation) # # # Skip if object is outside sensor span # if angle > half_span and angle < 2 * pi - half_span: # continue # # # Standardize angle to range [0, pi] to more easily find bin number # angle = standardize_radians(angle + half_span) # bin_number = int(angle / bin_res) # # # The object is exactly on the upper bound of the sensor span # if bin_number == self.n_bins: # bin_number -= 1 # # intensity = 1.0 - dist / self.sensor_range # # if obj_type == APPLE: # apple_readings[bin_number] = intensity # else: # bomb_readings[bin_number] = intensity # # sensor_obs = np.concatenate([apple_readings, bomb_readings]) # # return sensor_obs def _get_obs(self): return np.concatenate([self._get_self_obs(), self._get_sensor_obs()]) def _get_orientation(self): raise NotImplementedError def _is_done(self): raise NotImplementedError def _unhealthy_cost(self): raise NotImplementedError def render(self): if not self.viewer: self.viewer = MjViewer(self.sim) objects = self.apples + self.bombs for object in objects: x, y = object.x, object.y rgba = APPLE_RGBA if object.type is APPLE else BOMB_RGBA self.viewer.add_marker(type=GEOM_SPHERE, pos=np.asarray([x, y, 0.5]), rgba=rgba, size=np.asarray([0.5] * 3)) self.viewer.render()
MODEL_XML = """ <?xml version="1.0" ?> <mujoco> <worldbody> <body name="box" pos="0 0 0.2"> <geom size="0.15 0.15 0.15" type="box"/> <joint axis="1 0 0" name="box:x" type="slide"/> <joint axis="0 1 0" name="box:y" type="slide"/> </body> <body name="floor" pos="0 0 0.025"> <geom size="1.0 1.0 0.02" rgba="0 1 0 1" type="box"/> </body> </worldbody> </mujoco> """ model = load_model_from_xml(MODEL_XML) sim = MjSim(model) viewer = MjViewer(sim) step = 0 while True: t = time.time() x, y = math.cos(t), math.sin(t) viewer.add_marker(pos=np.array([x, y, 1]), label="t: " + str(t)) viewer.render() step += 1 if step > 100 and os.getenv('TESTING') is not None: break
class PointEnvMDP(MDP): def __init__(self, init_mean=(-0.2, -0.2), control_cost=False, dense_reward=False, render=False): xml = os.path.join( os.path.expanduser("~"), "git-repos/dm_control/dm_control/suite/point_mass.xml") model = load_model_from_path(xml) self.sim = MjSim(model) self.render = render self.init_mean = init_mean self.control_cost = control_cost self.dense_reward = dense_reward if self.render: self.viewer = MjViewer(self.sim) # Config self.env_name = "Point-Mass-Environment" self.target_position = np.array([0., 0.]) self.target_tolerance = 0.02 self.init_noise = 0.05 self._initialize_mujoco_state() self.init_state = self.get_state() print("Loaded {} with dense_reward={}".format(self.env_name, self.dense_reward)) MDP.__init__(self, [0, 1], self._transition_func, self._reward_func, self.init_state) def _reward_func(self, state, action): self.next_state = self._step(action) if self.render: self.viewer.render() if self.dense_reward: reward = -np.linalg.norm(self.next_state.position - self.target_position) else: reward = +0. if self.next_state.is_terminal() else -1. control_cost = 0.1 * np.linalg.norm(action) if self.control_cost: reward = reward - control_cost return reward def _transition_func(self, state, action): return self.next_state def execute_agent_action(self, action, option_idx=None): if self.render and option_idx is not None: self.viewer.add_marker(pos=np.array([0, 0, 0.1]), size=0.001 * np.ones(3), label="Option {}".format(option_idx)) reward, next_state = super(PointEnvMDP, self).execute_agent_action(action) return reward, next_state def is_goal_state(self, state): position = state.features()[:2] if isinstance( state, PointEnvState) else state[:2] return self.is_goal_position(position) def is_goal_position(self, position): distance = np.linalg.norm(position - self.target_position) return distance <= self.target_tolerance def get_state(self): # Individually indexing to prevent Mujoco from altering state history x, y = self.sim.data.qpos[0], self.sim.data.qpos[1] x_dot, y_dot = self.sim.data.qvel[0], self.sim.data.qvel[1] # Create State object from simulator state position = np.array([x, y]) velocity = np.array([x_dot, y_dot]) # State is terminal when it is the goal state done = self.is_goal_position(position) state = PointEnvState(position, velocity, done) return state def _step(self, action): self.sim.data.ctrl[:] = action self.sim.step() return self.get_state() def _initialize_mujoco_state(self): init_position = np.array(self.init_mean) + np.random.uniform( 0., self.init_noise, 2) init_state = MjSimState(time=0., qpos=init_position, qvel=np.array([0., 0.]), act=None, udd_state={}) self.sim.set_state(init_state) @staticmethod def state_space_size(): return 4 @staticmethod def action_space_size(): return 2 @staticmethod def is_primitive_action(action): return -1. <= action.all() <= 1. def reset(self): self._initialize_mujoco_state() self.init_state = self.get_state() super(PointEnvMDP, self).reset() def __str__(self): return self.env_name
class SimManager(): """""" def __init__(self, filepath, blender_path=None, visualize=False, num_sims=1): self.filepath = filepath self.blender_path = blender_path self.visualize = visualize self.num_sims = num_sims self.base_model = load_model_from_path(filepath) if self.num_sims > 1: self.pool = MjRenderPool(self.base_model, n_workers=num_sims, modder=ArenaModder) else: self.sim = MjSim(self.base_model) self.arena_modder = ArenaModder(self.sim) if self.visualize: self.arena_modder.visualize = True self.viewer = MjViewer(self.sim) self.arena_modder.viewer = self.viewer else: self.viewer = None def forward(self): """ Advances simulator a step (NECESSARY TO MAKE CAMERA AND LIGHT MODDING WORK) """ if self.num_sims <= 1: self.sim.forward() if self.visualize: # Get angle of camera and display it quat = np.quaternion(*self.base_model.cam_quat[0]) ypr = quaternion.as_euler_angles(quat) * 180 / np.pi cam_pos = self.base_model.cam_pos[0] #self.viewer.add_marker(pos=cam_pos, label="CAM: {}{}".format(cam_pos, ypr)) self.viewer.add_marker(pos=cam_pos, label="CAM: {}".format(ypr)) self.viewer.render() def _get_pool_data(self): IMAGE_NOISE_RVARIANCE = Range(0.0, 0.0001) cam_imgs, ground_truths = self.pool.render(640, 360, camera_name='camera1', randomize=True) ground_truths = list(ground_truths) cam_imgs = list( cam_imgs[:, ::-1, :, :]) # Rendered images are upside-down. for i in range(len(cam_imgs)): image_noise_variance = sample(IMAGE_NOISE_RVARIANCE) cam_imgs[i] = (skimage.util.random_noise( cam_imgs[i], mode='gaussian', var=image_noise_variance) * 255).astype(np.uint8) cam_imgs[i] = preproc_image(cam_imgs[i]) return cam_imgs, ground_truths def _get_cam_frame(self, display=False, ground_truth=None): """Grab an image from the camera (224, 244, 3) to feed into CNN""" IMAGE_NOISE_RVARIANCE = Range(0.0, 0.0001) cam_img = self.sim.render( 640, 360, camera_name='camera1' )[::-1, :, :] # Rendered images are upside-down. image_noise_variance = sample(IMAGE_NOISE_RVARIANCE) cam_img = (skimage.util.random_noise( cam_img, mode='gaussian', var=image_noise_variance) * 255).astype( np.uint8) cam_img = preproc_image(cam_img) if display: label = str(ground_truth[3:6]) display_image(cam_img, label) return cam_img def _get_ground_truth(self): return self.arena_modder.get_ground_truth() def get_data(self): if self.num_sims > 1: return self._get_pool_data() else: self.arena_modder.randomize() gt = self._get_ground_truth() self.sim.forward() cam = self._get_cam_frame() return cam, gt def randrocks(self): """Generate a new set of 3 random rock meshes using a Blender script""" if self.blender_path is None: raise Exception( 'You must install Blender and include the path to its exectubale in the constructor to use this method' ) import subprocess subprocess.call( [self.blender_path, "--background", "--python", "randrock.py"])
MODEL_XML = """ <?xml version="1.0" ?> <mujoco> <worldbody> <body name="box" pos="0 0 0.2"> <geom size="0.15 0.15 0.15" type="box"/> <joint axis="1 0 0" name="box:x" type="slide"/> <joint axis="0 1 0" name="box:y" type="slide"/> </body> <body name="floor" pos="0 0 0.025"> <geom size="1.0 1.0 0.02" rgba="0 1 0 1" type="box"/> </body> </worldbody> </mujoco> """ model = load_model_from_xml(MODEL_XML) sim = MjSim(model) viewer = MjViewer(sim) step = 0 while True: t = time.time() x, y = math.cos(t), math.sin(t) viewer.add_marker(pos=np.array([x, y, 1]), label=str(t)) viewer.render() step += 1 if step > 100 and os.getenv('TESTING') is not None: break