class BaseEnv(gym.Env): ''' a basic environment, step, observation and reward not implemented ''' def __init__(self, config_file, model_id=None, mode='headless', action_timestep=1 / 10.0, physics_timestep=1 / 240.0, device_idx=0): """ :param config_file: config_file path :param model_id: override model_id in config file :param mode: headless or gui mode :param action_timestep: environment executes action per action_timestep second :param physics_timestep: physics timestep for pybullet :param device_idx: device_idx: which GPU to run the simulation and rendering on """ self.config = parse_config(config_file) if model_id is not None: self.config['model_id'] = model_id self.mode = mode self.action_timestep = action_timestep self.physics_timestep = physics_timestep self.simulator = Simulator( mode=mode, timestep=physics_timestep, use_fisheye=self.config.get('fisheye', False), image_width=self.config.get('image_width', 128), image_height=self.config.get('image_height', 128), vertical_fov=self.config.get('vertical_fov', 90), device_idx=device_idx, auto_sync=False) self.simulator_loop = int(self.action_timestep / self.simulator.timestep) self.load() def reload(self, config_file): """ Reload another config file, this allows one to change the envrionment on the fly :param config_file: new config file path """ self.config = parse_config(config_file) self.simulator.reload() self.load() def reload_model(self, model_id): """ Reload another model, this allows one to change the envrionment on the fly :param model_id: new model_id """ self.config['model_id'] = model_id self.simulator.reload() self.load() def load(self): """ Load the scene and robot """ if self.config['scene'] == 'empty': scene = EmptyScene() elif self.config['scene'] == 'stadium': scene = StadiumScene() elif self.config['scene'] == 'building': scene = BuildingScene( self.config['model_id'], waypoint_resolution=self.config.get('waypoint_resolution', 0.2), num_waypoints=self.config.get('num_waypoints', 10), build_graph=self.config.get('build_graph', False), trav_map_resolution=self.config.get('trav_map_resolution', 0.1), trav_map_erosion=self.config.get('trav_map_erosion', 2), is_interactive=self.config.get('is_interactive', False), pybullet_load_texture=self.config.get('pybullet_load_texture', False), ) self.simulator.import_scene(scene, load_texture=self.config.get( 'load_texture', True)) if self.config['robot'] == 'Turtlebot': robot = Turtlebot(self.config) elif self.config['robot'] == 'Husky': robot = Husky(self.config) elif self.config['robot'] == 'Ant': robot = Ant(self.config) elif self.config['robot'] == 'Humanoid': robot = Humanoid(self.config) elif self.config['robot'] == 'JR2': robot = JR2(self.config) elif self.config['robot'] == 'JR2_Kinova': robot = JR2_Kinova(self.config) elif self.config['robot'] == 'Freight': robot = Freight(self.config) elif self.config['robot'] == 'Fetch': robot = Fetch(self.config) elif self.config['robot'] == 'Locobot': robot = Locobot(self.config) else: raise Exception('unknown robot type: {}'.format( self.config['robot'])) self.scene = scene self.robots = [robot] for robot in self.robots: self.simulator.import_robot(robot) def clean(self): """ Clean up """ if self.simulator is not None: self.simulator.disconnect() def simulator_step(self): """ Step the simulation, this is different from environment step where one can get observation and reward """ self.simulator.step() def step(self, action): """ Overwritten by subclasses """ return NotImplementedError() def reset(self): """ Overwritten by subclasses """ return NotImplementedError() def set_mode(self, mode): self.simulator.mode = mode
class BaseEnv(gym.Env): ''' a basic environment, step, observation and reward not implemented ''' def __init__(self, config_file, mode='headless', device_idx=0): self.config = parse_config(config_file) self.simulator = Simulator(mode=mode, use_fisheye=self.config.get( 'fisheye', False), resolution=self.config['resolution'], device_idx=device_idx) self.load() def reload(self, config_file): self.config = parse_config(config_file) self.simulator.reload() self.load() def load(self): if self.config['scene'] == 'stadium': scene = StadiumScene() elif self.config['scene'] == 'building': scene = BuildingScene(self.config['model_id']) self.simulator.import_scene(scene) if self.config['robot'] == 'Turtlebot': robot = Turtlebot(self.config) elif self.config['robot'] == 'Husky': robot = Husky(self.config) elif self.config['robot'] == 'Ant': robot = Ant(self.config) elif self.config['robot'] == 'Humanoid': robot = Humanoid(self.config) elif self.config['robot'] == 'JR2': robot = JR2(self.config) elif self.config['robot'] == 'JR2_Kinova': robot = JR2_Kinova(self.config) else: raise Exception('unknown robot type: {}'.format( self.config['robot'])) self.scene = scene self.robots = [robot] for robot in self.robots: self.simulator.import_robot(robot) def clean(self): if not self.simulator is None: self.simulator.disconnect() def simulator_step(self): self.simulator.step() def step(self, action): return NotImplementedError def reset(self): return NotImplementedError def set_mode(self, mode): self.simulator.mode = mode
class BaseEnv(gym.Env): ''' a basic environment, step, observation and reward not implemented ''' def __init__(self, config_file, model_id=None, mode='headless', device_idx=0): """ :param config_file: config_file path :param model_id: override model_id in config file :param mode: headless or gui mode :param device_idx: which GPU to run the simulation and rendering on """ self.config = parse_config(config_file) if model_id is not None: self.config['model_id'] = model_id self.simulator = Simulator( mode=mode, use_fisheye=self.config.get('fisheye', False), resolution=self.config.get('resolution', 64), fov=self.config.get('fov', 90), device_idx=device_idx) self.load() def reload(self, config_file): """ Reload another config file, this allows one to change the envrionment on the fly :param config_file: new config file path """ self.config = parse_config(config_file) self.simulator.reload() self.load() def load(self): """ Load the scene and robot """ if self.config['scene'] == 'stadium': scene = StadiumScene() elif self.config['scene'] == 'building': scene = BuildingScene( self.config['model_id'], build_graph=self.config.get('build_graph', False), trav_map_erosion=self.config.get('trav_map_erosion', 2), should_load_replaced_objects=self.config.get( 'should_load_replaced_objects', False)) # scene: class_id = 0 # robot: class_id = 1 # objects: class_id > 1 self.simulator.import_scene(scene, load_texture=self.config.get( 'load_texture', True), class_id=0) if self.config['robot'] == 'Turtlebot': robot = Turtlebot(self.config) elif self.config['robot'] == 'Husky': robot = Husky(self.config) elif self.config['robot'] == 'Ant': robot = Ant(self.config) elif self.config['robot'] == 'Humanoid': robot = Humanoid(self.config) elif self.config['robot'] == 'JR2': robot = JR2(self.config) elif self.config['robot'] == 'JR2_Kinova': robot = JR2_Kinova(self.config) elif self.config['robot'] == 'Freight': robot = Freight(self.config) elif self.config['robot'] == 'Fetch': robot = Fetch(self.config) else: raise Exception('unknown robot type: {}'.format( self.config['robot'])) self.scene = scene self.robots = [robot] for robot in self.robots: self.simulator.import_robot(robot, class_id=1) def clean(self): """ Clean up """ if self.simulator is not None: self.simulator.disconnect() def simulator_step(self): """ Step the simulation, this is different from environment step where one can get observation and reward """ self.simulator.step() def step(self, action): return NotImplementedError def reset(self): return NotImplementedError def set_mode(self, mode): self.simulator.mode = mode
class ContainerObjectsEnv(object): def __init__(self, show_gui=False): # set up our initial configurations self.image_height = 480 self.image_width = 640 self.camera_height = 0.9 self.camera_x = 0.0 self.camera_y = -0.8 self.objects = {} if show_gui: self.simulator = Simulator(image_width=self.image_width, image_height=self.image_height) else: self.simulator = Simulator(image_width=self.image_width, image_height=self.image_height, mode='headless', timestep=0.001) pb.setAdditionalSearchPath(pybullet_data.getDataPath()) pb.loadURDF('plane.urdf') # set up renderer self.adjust_camera([self.camera_x, self.camera_y, self.camera_height], [0, 0, 0], [-1, 0, 0]) self.simulator.renderer.set_light_pos([0.65, 0.0, 10.0]) self.simulator.renderer.set_fov(53) def reset(self, container): with suppress_stdout_stderr(): self.simulator.reload() self.container = container self.simulator.import_object(container) container.set_position([0, 0, container.size[2] / 2.0]) self.objects = {} def adjust_camera(self, camera_eye_position, camera_target_position, camera_up_vector, randomize=False): if randomize: self.simulator.renderer.set_camera( camera_eye_position + np.random.random(3) * 0.1, camera_target_position + np.random.random(3) * 0.1, camera_up_vector) else: self.simulator.renderer.set_camera(camera_eye_position, camera_target_position, camera_up_vector) def __del__(self): self.simulator.renderer.release() del self.simulator.renderer def get_renderer_rgb_depth(self): rgb, im3d = self.simulator.renderer.render(modes=('rgb', '3d')) depth = im3d[:, :, 2] return rgb, -depth def get_renderer_rgb(self): rgb = self.simulator.renderer.render(modes=('rgb'))[0] return rgb def get_renderer_depth(self): rgb, im3d = self.simulator.renderer.render(modes=('rgb', '3d')) depth = im3d[:, :, 2] return -depth # selected_object should be an object id def get_renderer_segmask(self, selected_object_id, with_occlusion=True): all_instances = self.simulator.renderer.instances.copy() # get unoccluded segmask self.simulator.renderer.instances = [] selected_object = self.get_body(selected_object_id) for instance in all_instances: if selected_object is not None and isinstance( instance, Instance) and instance.pybullet_uuid == selected_object: self.simulator.renderer.instances.append(instance) elif isinstance( instance, Instance) and instance.pybullet_uuid == selected_object_id: self.simulator.renderer.instances.append(instance) target_segmask = np.sum( self.simulator.renderer.render(modes=('seg'))[0][:, :, :3], axis=2) target_segmask, target_depth = self.simulator.renderer.render( modes=('seg', '3d')) target_segmask = np.sum(target_segmask[:, :, :3], axis=2) target_depth = target_depth[:, :, 2] self.simulator.renderer.instances = all_instances if not with_occlusion: return target_segmask all_depth = self.simulator.renderer.render(modes=('3d'))[0][:, :, 2] occluded_segmask = np.logical_and( np.abs(target_depth - all_depth) < 0.01, target_segmask) return occluded_segmask def set_camera_point_at(self, position, randomize=False, dist=0.3): camera_eye_position = np.asarray( [position[0], position[1] - dist, position[2] + 0.1]) camera_target_position = np.asarray( [position[0], position[1], position[2]]) self.adjust_camera(camera_eye_position, camera_target_position, [0, 0, 1], randomize=randomize) self.simulator.sync() def get_observation(self, segmak_object_id=None, visualize=False, save=False, demo=False, randomize_camera=False): rgb_im, depth_im = self.get_renderer_rgb_depth() if segmak_object_id: segmask_im = self.get_renderer_segmask(segmak_object_id) if visualize: plt.imshow(segmask_im) plt.title('segmask') plt.show() plt.close() plt.imshow(rgb_im) plt.title('rgb') plt.show() plt.close() plt.imshow(depth_im) plt.title('depth') plt.show() plt.close() if save: save_dir = './rendered_images/' save_time = str(time.time()) plt.imshow(segmask_im) plt.title('segmask') plt.savefig(save_dir + "segmask_im_" + save_time + '_.png') plt.close() plt.imshow(rgb_im) plt.title('rgb') plt.savefig(save_dir + "rgb_im_" + save_time + '_.png') plt.close() plt.imshow(depth_im) plt.title('depth') plt.savefig(save_dir + "depth_im_" + save_time + '_.png') plt.close() if segmak_object_id: return rgb_im, depth_im, segmask_im return rgb_im, depth_im def get_obj_ids(self): return list(self.objects.keys()) def remove_object(self, obj): with suppress_stdout_stderr(): obj_id = obj.body_id pb.removeBody(bodyUniqueId=obj_id) del self.objects[obj_id] self.simulator.sync() def add_object(self, obj): with suppress_stdout_stderr(): new_obj_id = self.simulator.import_object(obj) self.objects[new_obj_id] = obj return new_obj_id def get_object(self, body_id): return self.get_body(body_id) def get_body(self, body_id): if body_id in self.objects: return self.objects[body_id] return None def gentle_drop(self, body_id, threshold=0.1): start_time = time.time() lvel = pb.getBaseVelocity(bodyUniqueId=body_id)[0] avel = pb.getBaseVelocity(bodyUniqueId=body_id)[1] pos, _ = pb.getBasePositionAndOrientation(body_id) while np.linalg.norm(lvel) < threshold: pb.stepSimulation() lvel = pb.getBaseVelocity(bodyUniqueId=body_id)[0] avel = pb.getBaseVelocity(bodyUniqueId=body_id)[1] for _ in range(100000): if time.time() > start_time + 5: return False pos, _ = pb.getBasePositionAndOrientation(body_id) if pos[2] < 0: return False for _ in range(10): pb.stepSimulation() lvel = pb.getBaseVelocity(bodyUniqueId=body_id)[0] avel = pb.getBaseVelocity(bodyUniqueId=body_id)[1] # return if it's basically stable if np.linalg.norm(lvel) < threshold * 0.5 and np.linalg.norm( avel) < threshold: return True # modify linear velocity if too large, else leave the same if np.linalg.norm(lvel) > threshold: new_lvel = np.array(lvel) * (threshold / np.linalg.norm(lvel)) else: new_lvel = lvel # modify angular velocity if too large, else leave the same if np.linalg.norm(avel) > threshold: new_avel = np.array(avel) * (threshold / np.linalg.norm(avel)) else: new_avel = avel pb.resetBaseVelocity(objectUniqueId=body_id, angularVelocity=list(new_avel), linearVelocity=list(new_lvel)) return True def load_container_setup(self, obj_file, container_id=None): abs_path = os.path.join(os.getcwd() + '/', obj_file) with open(abs_path, 'rb') as file: data = pickle.load(file) self.add_objects(data, container_id) def get_obj_img(self, color_obs, segmask, zoom_factor=1.2, save=False): np_any = np.any(segmask) if not np_any: return None rows = np.any(segmask, axis=0) cols = np.any(segmask, axis=1) rmin, rmax = np.where(rows)[0][[0, -1]] cmin, cmax = np.where(cols)[0][[0, -1]] rmean = (rmin + rmax) / 2 cmean = (cmin + cmax) / 2 rlen = rmax - rmin clen = cmax - cmin maxlen = max([rlen, clen]) * zoom_factor rmin = int((rmean - maxlen / 2)) rmax = int((rmean + maxlen / 2)) cmin = int((cmean - maxlen / 2)) cmax = int((cmean + maxlen / 2)) obj_color = color_obs[cmin:cmax, rmin:rmax] try: if save: save_dir = './sim_images/' plt.imshow(obj_color) save_time = str(time.time()) plt.savefig(save_dir + "im_" + save_time + '_.png') except: pass return obj_color