with bz2.BZ2File(source_file, 'r') as compressed_file:
all_trajectories = pickle.load(compressed_file)
for trajectory in all_trajectories:
poses = trajectory[1]
poses_after_filter = process_poses([p[target_point_tuple] for p in poses])
result.extend(poses_after_filter)
return result
il_transitions = process_dir(os.path.join(imitation_data_path, 'train'), number_of_imitation_files)
# load the openrave view
params_file = os.path.abspath(os.path.expanduser(
os.path.join('~/ModelBasedDDPG/scenario_params', scenario, 'params.pkl')))
openrave_manager.set_params(params_file)
openrave_manager.get_initialized_viewer()
# remove robot from view
for link in openrave_manager.robot.GetLinks():
link.SetVisible(False)
red_color = np.array([1.0, 0.0, 0.0])
yellow_color = np.array([1.0, 1.0, 0.0])
green_color = np.array([0.0, 1.0, 0.0])
def create_sphere(id, radius, openrave_manager):
body = RaveCreateKinBody(openrave_manager.env, '')
body.SetName('sphere{}'.format(id))
body.InitFromSpheres(np.array([[0.0]*3 + [radius]]), True)
def main():
config_path = os.path.join(os.getcwd(), 'config/config.yml')
with open(config_path, 'r') as yml_file:
config = yaml.load(yml_file)
print('------------ Config ------------')
print(yaml.dump(config))
model_dir = os.path.join(trajectories_dir, model_name)
potential_points_location = os.path.join(model_dir, 'potential_points.p')
potential_points = pickle.load(open(potential_points_location))
# search_key = os.path.join(model_dir, '{}_step_{}_{}.p'.format(global_step, message, path_id))
search_key = os.path.join(model_dir, global_step,
'{}_{}.p'.format(message, path_id))
trajectory_files = glob.glob(search_key)
trajectory_file = trajectory_files[0]
pose_goal, trajectory, workspace_id = pickle.load(open(trajectory_file))
trajectory = [[0.0] + list(j) for j in trajectory]
openrave_manager = OpenraveManager(0.001, potential_points)
# get the parameters
if scenario == 'vision':
workspace_params_path = os.path.abspath(
os.path.expanduser(
os.path.join('~/ModelBasedDDPG/scenario_params/vision/',
workspace_id)))
else:
workspace_params_path = os.path.abspath(
os.path.expanduser(
os.path.join('~/ModelBasedDDPG/scenario_params', scenario,
'params.pkl')))
if workspace_params_path is not None:
openrave_manager.set_params(workspace_params_path)
openrave_manager.get_initialized_viewer()
openrave_manager.env.GetViewer().SetSize(1200, 800)
# for link in openrave_manager.robot.GetLinks():
# link.SetVisible(False)
if display_start_goal_end_spheres:
start = trajectory[0]
end = trajectory[-1]
pose_start = openrave_manager.get_target_pose(start)
pose_start = (pose_start[0], 0.0, pose_start[1])
pose_goal = (pose_goal[0], 0.0, pose_goal[1])
pose_end = openrave_manager.get_target_pose(end)
pose_end = (pose_end[0], 0.0, pose_end[1])
start_sphere = create_sphere('start', trajectory_spheres_radius,
openrave_manager)
move_body(start_sphere, pose_start, 0.0)
goal_sphere = create_sphere('goal', goal_radius, openrave_manager)
move_body(goal_sphere, pose_goal, 0.0)
end_sphere = create_sphere('end', trajectory_spheres_radius,
openrave_manager)
start_sphere.GetLinks()[0].GetGeometries()[0].SetDiffuseColor(
np.array([0, 0, 204]))
goal_sphere.GetLinks()[0].GetGeometries()[0].SetDiffuseColor(
np.array([240, 100, 10]))
goal_sphere.GetLinks()[0].GetGeometries()[0].SetTransparency(0.3)
end_sphere.GetLinks()[0].GetGeometries()[0].SetDiffuseColor(
np.array([100, 204, 204]))
move_body(end_sphere, pose_end, 0.0)
print 'len(trajectory) ', len(trajectory)
for i in range(len(trajectory)):
print 'i ', i
openrave_manager.robot.SetDOFValues(trajectory[i])
time.sleep(1 / speed)
time.sleep(0.2)
if message == 'collision':
time.sleep(1.2)