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
