def sample_rope_state(demofile):
success = False
while not success:
new_xyz_camera, demo_id = load_random_start_segment(demofile)
perturb_radius = random.uniform(args.min_rad, args.max_rad)
new_xyz_robot = apply_tfm(new_xyz_camera[()], args.inittfmfile)
rope_nodes = rope_initialization.find_path_through_point_cloud(
new_xyz_robot,
perturb_peak_dist = perturb_radius,
num_perturb_points = args.perturb_points)
replace_rope(rope_nodes)
Globals.sim.settle()
Globals.viewer.Step()
#Globals.viewer.Idle()
if args.human_check:
resp = raw_input("Use this start simulation?[Y/n]")
success = resp not in ('N', 'n')
else:
success = True
new_xyz_robot = Globals.sim.observe_cloud()
#new_xyz_camera = apply_inv_tfm(new_xyz_robot, args.inittfmfile)
#return (new_xyz_camera, demo_id)
return (new_xyz_robot, demo_id)
def setup_and_return_demofile(demofile_name, init_rope_state_segment, animate, perturb_radius=0, perturb_num_points=0):
"""For the simulation, this code runs before the main loop. It also sets the numpy random seed"""
if Globals.random_seed is not None:
np.random.seed(Globals.random_seed)
demofile = h5py.File(demofile_name, 'r+')
Globals.env = openravepy.Environment() # @UndefinedVariable
Globals.env.StopSimulation()
Globals.env.Load("robots/pr2-beta-static.zae")
Globals.robot = Globals.env.GetRobots()[0]
#Set up the simulation with the table (no rope yet)
new_xyz, _ = load_segment(demofile, init_rope_state_segment)
#table_height = new_xyz[:, 2].mean() - .02
table_height = new_xyz[:, 2].mean() - 0.17
table_xml = make_table_xml(translation=[1, 0, table_height], extents=[.85, .55, .01])
if animate:
Globals.viewer = trajoptpy.GetViewer(Globals.env)
Globals.env.LoadData(table_xml)
Globals.sim = ropesim.Simulation(Globals.env, Globals.robot)
# create rope with optional pertubations
move_sim_arms_to_side()
rope_nodes = rope_initialization.find_path_through_point_cloud(
new_xyz,
perturb_peak_dist=perturb_radius,
num_perturb_points=perturb_num_points)
rope_nodes = rotate_about_median(rope_nodes, 2*np.pi/3.0)
ropt_nodes = place_in_feasible_region(rope_nodes)
Globals.sim.create(rope_nodes)
return demofile
def eval_demo_playback(args, sim_env):
actionfile = GlobalVars.actions
num_successes = 0
num_total = 0
actions_names = []
for action_name in actionfile:
actions_names.append(action_name)
i = 0
i_task = 0
i_step = 0
while i < len(actions_names):
(demo_name, seg_name) = actions_names[i].split('-')
if seg_name == 'seg00':
if i > 0:
if is_knot(sim_env.sim.rope.GetControlPoints()):
redprint("KNOT TIED!")
num_successes += 1
num_total += 1
i_task += 1
i_step = 0
print "task %s" % demo_name
sim_util.reset_arms_to_side(sim_env, floating=True)
redprint("Replace rope")
rope_xyz = np.asarray(actionfile[actions_names[i]]['cloud_xyz'])
rope_nodes = rope_initialization.find_path_through_point_cloud(rope_xyz)
# don't call replace_rope and sim.settle() directly. use time machine interface for deterministic results!
time_machine = sim_util.RopeSimTimeMachine(rope_nodes, sim_env, rope_params=sim_util.get_rope_params("thick"), floating=True)
if args.animation:
sim_env.viewer.Step()
print "task %s step %i" % (i_task, i_step)
sim_util.reset_arms_to_side(sim_env, floating=True)
redprint("Observe point cloud")
#new_xyz = sim_env.sim.observe_cloud(GlobalVars.rope_observe_cloud_upsample)
new_xyz = sim_env.sim.observe_cloud2(0.01, GlobalVars.rope_observe_cloud_upsample, 1, 4)
state = ("eval_%i"%get_unique_id(), new_xyz)
action = actions_names[i]
eval_stats = eval_util.EvalStats()
eval_stats.success, eval_stats.feasible, eval_stats.misgrasp, full_trajs = \
compute_trans_traj(sim_env, new_xyz, GlobalVars.actions[action], animate=args.animation, interactive=args.interactive)
i += 1
num_total += 1
redprint('Eval Successes / Total: ' + str(num_successes) + '/' + str(num_total))
def track_video_frames_offline(video_dir, T_w_k, init_tfm, ref_image, rope_params = None):
video_stamps = np.loadtxt(osp.join(video_dir,demo_names.stamps_name))
rgbs = []
depths = []
color_clouds = []
for (index, stamp) in zip(range(len(video_stamps)), video_stamps):
print index, stamp
rgb = cv2.imread(osp.join(video_dir,demo_names.rgb_name%index))
rgb = color_match.match(ref_image, rgb)
assert rgb is not None
depth = cv2.imread(osp.join(video_dir,demo_names.depth_name%index), 2)
assert depth is not None
color_cloud = extract_rope_tracking(rgb, depth, T_w_k)
color_cloud = clouds.downsample_colored(color_cloud, .01)
color_cloud[:,:3] = color_cloud[:,:3].dot(init_tfm[:3,:3].T) + init_tfm[:3,3][None,:]
if index == 0:
rope_xyz = extract_red(rgb, depth, T_w_k)
rope_xyz = clouds.downsample(rope_xyz, .01)
rope_xyz = rope_xyz.dot(init_tfm[:3, :3].T) + init_tfm[:3,3][None,:]
rope_nodes = rope_initialization.find_path_through_point_cloud(rope_xyz)
rope_radius = 0.005
rgbs.append(rgb)
depths.append(depth)
color_clouds.append(color_cloud)
rgbs = np.asarray(rgbs)
depths = np.asarray(depths)
all_tracked_nodes = cbulletracpy2.py_tracking(rope_nodes, rope_radius, T_w_k, color_clouds, rgbs, depths, 5)
for i in range(len(rgbs)):
if i % 30 != 0:
continue
print i
fig.clf()
ax = fig.gca(projection='3d')
ax.set_autoscale_on(False)
tracked_nodes = all_tracked_nodes[i, :, :]
ax.plot(tracked_nodes[:, 0], tracked_nodes[:,1], tracked_nodes[:,2], 'o')
fig.show()
raw_input()
def sample_rope_state(demofile, perturb_points=5, min_rad=0, max_rad=.15, rotation=False):
"""
samples a rope state, by picking a random segment, perturbing, rotating about the median,
then setting a random translation such that the rope is essentially within grasp room
"""
# TODO: pick a random rope initialization
new_xyz= load_random_start_segment(demofile)
perturb_radius = random.uniform(min_rad, max_rad)
rope_nodes = rope_initialization.find_path_through_point_cloud( new_xyz,
perturb_peak_dist=perturb_radius,
num_perturb_points=perturb_points)
if rotation:
rand_theta = rotation*np.random.rand()
rope_nodes = rotate_about_median(rope_nodes, rand_theta)
rope_nodes = place_in_feasible_region(rope_nodes)
return rope_nodes
def sample_rope_state(demofile, human_check=False, perturb_points=5, min_rad=0, max_rad=.15):
success = False
while not success:
# TODO: pick a random rope initialization
new_xyz= load_random_start_segment(demofile)
perturb_radius = random.uniform(min_rad, max_rad)
rope_nodes = rope_initialization.find_path_through_point_cloud( new_xyz,
perturb_peak_dist=perturb_radius,
num_perturb_points=perturb_points)
replace_rope(rope_nodes)
Globals.sim.settle()
Globals.viewer.Step()
raw_input("Press enter to continue")
if human_check:
resp = raw_input("Use this simulation?[Y/n]")
success = resp not in ('N', 'n')
else:
success = True
def sample_rope_state(demofile, perturb_points=5, min_rad=0, max_rad=.15):
"""
samples a rope state, by picking a random segment, perturbing, rotating about the median,
then setting a random translation such that the rope is essentially within grasp room
"""
# TODO: pick a random rope initialization
new_xyz= load_random_start_segment(demofile)
perturb_radius = random.uniform(min_rad, max_rad)
rope_nodes = rope_initialization.find_path_through_point_cloud( new_xyz,
perturb_peak_dist=perturb_radius,
num_perturb_points=perturb_points)
rand_theta = np.pi*(np.random.rand() - 0.5)
# rand_theta = np.random.randn()
rope_nodes = rotate_about_median(rope_nodes, rand_theta)
r_trans = np.r_[np.random.multivariate_normal([0, 0], np.eye(2)), [0]]
rope_nodes = rope_nodes + r_trans
rope_nodes = place_in_feasible_region(rope_nodes)
return rope_nodes
def main():
demofile = h5py.File(args.demos_h5file, 'r')
trajoptpy.SetInteractive(False)
Globals.env = openravepy.Environment()
Globals.env.StopSimulation()
Globals.env.Load("robots/pr2-beta-static.zae")
Globals.robot = Globals.env.GetRobots()[0]
Globals.robot.SetDOFValues(PR2_L_POSTURES["side"], Globals.robot.GetManipulator("leftarm").GetArmIndices())
Globals.robot.SetDOFValues(mirror_arm_joints(PR2_L_POSTURES["side"]), Globals.robot.GetManipulator("rightarm").GetArmIndices())
init_rope_xyz, _ = load_random_start_segment(demofile)
init_rope_xyz_robot = apply_tfm(init_rope_xyz[()], args.inittfmfile)
table_height = init_rope_xyz_robot[:,2].mean() - .03
table_xml = make_table_xml(translation=[1, 0, table_height], extents=[.85, .55, .01])
Globals.env.LoadData(table_xml)
Globals.sim = ropesim.Simulation(Globals.env, Globals.robot)
Globals.viewer = trajoptpy.GetViewer(Globals.env)
Globals.viewer.Idle()
if not args.same_as_training:
for i in range(0, args.dataset_size):
print "State ", i, " of ", args.dataset_size
rope_nodes, demo_id = sample_rope_state(demofile)
save_example_action(rope_nodes, demo_id)
else:
start_keys = [k for k in demofile.keys()
if k.startswith('demo') and k.endswith('00')]
for demo_id in start_keys:
xyz_camera = demofile[demo_id]['cloud_xyz'][()]
xyz_robot = apply_tfm(xyz_camera, args.inittfmfile)
rope_nodes = rope_initialization.find_path_through_point_cloud(
xyz_robot, perturb_peak_dist = None,
num_perturb_points = 0)
replace_rope(rope_nodes)
Globals.sim.settle()
Globals.viewer.Step()
new_xyz_robot = Globals.sim.observe_cloud()
save_example_action(new_xyz_robot, demo_id)
time.sleep(1)
def eval_on_holdout(args, sim_env):
feature_fn, _, num_features, actions = select_feature_fn(args)
weightfile = h5py.File(args.weightfile, 'r')
weights = weightfile['weights'][:]
w0 = weightfile['w0'][()] if 'w0' in weightfile else 0
weightfile.close()
assert weights.shape[
0] == num_features, "Dimensions of weights and features don't match. Make sure the right feature is being used"
holdoutfile = h5py.File(args.holdoutfile, 'r')
tasks = eval_util.get_specified_tasks(args.tasks, args.taskfile,
args.i_start, args.i_end)
holdout_items = eval_util.get_holdout_items(holdoutfile, tasks)
num_successes = 0
num_total = 0
for i_task, demo_id_rope_nodes in holdout_items:
print "task %s" % i_task
sim_util.reset_arms_to_side(sim_env)
redprint("Replace rope")
rope_xyz = demo_id_rope_nodes["rope_nodes"][:]
# Transform rope_nodes from the kinect's frame into the frame of the PR2
if 'frame' not in demo_id_rope_nodes or demo_id_rope_nodes['frame'][
()] != 'r':
redprint("Transforming rope into frame of robot")
rope_xyz = rope_xyz.dot(GlobalVars.init_tfm[:3, :3].T
) + GlobalVars.init_tfm[:3, 3][None, :]
rope_nodes = rope_initialization.find_path_through_point_cloud(
rope_xyz)
# don't call replace_rope and sim.settle() directly. use time machine interface for deterministic results!
time_machine = sim_util.RopeSimTimeMachine(rope_nodes, sim_env)
if args.animation:
sim_env.viewer.Step()
eval_util.save_task_results_init(args.resultfile, sim_env, i_task,
rope_nodes, args.exec_rope_params)
for i_step in range(args.num_steps):
print "task %s step %i" % (i_task, i_step)
sim_util.reset_arms_to_side(sim_env)
redprint("Observe point cloud")
new_xyz = sim_env.sim.observe_cloud()
state = ("eval_%i" % get_unique_id(), new_xyz)
redprint("Choosing an action")
q_values = [
q_value_fn(state, action, feature_fn, weights, w0)
for action in actions
]
q_values_root = q_values
time_machine.set_checkpoint('depth_0_%i' % i_step, sim_env)
assert args.lookahead_width >= 1, 'Lookahead branches set to zero will fail to select any action'
agenda = sorted(zip(q_values, actions),
key=lambda v: -v[0])[:args.lookahead_width]
agenda = [
(v, a, 'depth_0_%i' % i_step, a) for (v, a) in agenda
] # state is (value, most recent action, checkpoint id, root action)
best_root_action = None
for depth in range(args.lookahead_depth):
expansion_results = []
for (branch, (q, a, chkpt, r_a)) in enumerate(agenda):
time_machine.restore_from_checkpoint(
chkpt, sim_env,
sim_util.get_rope_params(args.lookahead_rope_params))
cur_xyz = sim_env.sim.observe_cloud()
success, _, _, full_trajs = \
compute_trans_traj(sim_env, cur_xyz, GlobalVars.actions[a], animate=args.animation, interactive=False)
if args.animation:
sim_env.viewer.Step()
if is_knot(sim_env.sim.rope.GetControlPoints()):
best_root_action = r_a
break
result_cloud = sim_env.sim.observe_cloud()
result_chkpt = 'depth_%i_branch_%i_%i' % (depth + 1,
branch, i_step)
if depth != args.lookahead_depth - 1: # don't save checkpoint at the last depth to save computation time
time_machine.set_checkpoint(result_chkpt, sim_env)
expansion_results.append(
(result_cloud, a, success, result_chkpt, r_a))
if best_root_action is not None:
redprint('Knot Found, stopping search early')
break
agenda = []
for (cld, incoming_a, success, chkpt,
r_a) in expansion_results:
if not success:
agenda.append((-np.inf, actions[0], chkpt,
r_a)) # TODO why first action?
continue
next_state = ("eval_%i" % get_unique_id(), cld)
q_values = [(q_value_fn(next_state, action, feature_fn,
weights, w0), action, chkpt, r_a)
for action in actions]
agenda.extend(q_values)
agenda.sort(key=lambda v: -v[0])
agenda = agenda[:args.lookahead_width]
first_root_action = agenda[0][-1]
if all(r_a == first_root_action for (_, _, _, r_a) in agenda):
best_root_action = first_root_action
redprint(
'All best actions have same root, stopping search early'
)
break
if best_root_action is None:
best_root_action = agenda[0][-1]
time_machine.restore_from_checkpoint(
'depth_0_%i' % i_step, sim_env,
sim_util.get_rope_params(args.exec_rope_params))
eval_stats = eval_util.EvalStats()
eval_stats.success, eval_stats.feasible, eval_stats.misgrasp, full_trajs = \
compute_trans_traj(sim_env, new_xyz, GlobalVars.actions[best_root_action], animate=args.animation, interactive=args.interactive)
print "BEST ACTION:", best_root_action
eval_util.save_task_results_step(args.resultfile, sim_env, i_task,
i_step, eval_stats,
best_root_action, full_trajs,
q_values_root)
if is_knot(sim_env.sim.rope.GetControlPoints()):
redprint("KNOT TIED!")
break
if is_knot(sim_env.sim.rope.GetControlPoints()):
num_successes += 1
num_total += 1
redprint('Eval Successes / Total: ' + str(num_successes) + '/' +
str(num_total))
def track_video_frames_online(video_dir, T_w_k, init_tfm, table_plane, ref_image = None, rope_params = None):
video_stamps = np.loadtxt(osp.join(video_dir,demo_names.stamps_name))
env = openravepy.Environment()
stdev = np.array([])
rgb = cv2.imread(osp.join(video_dir,demo_names.rgb_name%0))
rgb = color_match.match(ref_image, rgb)
assert rgb is not None
depth = cv2.imread(osp.join(video_dir,demo_names.depth_name%0), 2)
assert depth is not None
rope_xyz = extract_red(rgb, depth, T_w_k)
table_dir = np.array([table_plane[0], table_plane[1], table_plane[2]])
table_dir = table_dir / np.linalg.norm(table_dir)
table_dir = init_tfm[:3,:3].dot(table_dir)
if np.dot([0,0,1], table_dir) < 0:
table_dir = -table_dir
table_axis = np.cross(table_dir, [0,0,1])
table_angle = np.arccos(np.dot([0,0,1], table_dir))
table_tfm = openravepy.matrixFromAxisAngle(table_axis, table_angle)
table_center_xyz = np.mean(rope_xyz, axis=0)
table_tfm[:3,3] = - table_tfm[:3,:3].dot(table_center_xyz) + table_center_xyz
init_tfm = table_tfm.dot(init_tfm)
tracked_nodes = None
for (index, stamp) in zip(range(len(video_stamps)), video_stamps):
print index, stamp
rgb = cv2.imread(osp.join(video_dir,demo_names.rgb_name%index))
rgb = color_match.match(ref_image, rgb)
assert rgb is not None
depth = cv2.imread(osp.join(video_dir,demo_names.depth_name%index), 2)
assert depth is not None
color_cloud = extract_rope_tracking(rgb, depth, T_w_k) # color_cloud is at first in camera frame
#print "cloud in camera frame"
#print "==================================="
#print color_cloud[:, :3]
color_cloud = clouds.downsample_colored(color_cloud, .01)
color_cloud[:,:3] = color_cloud[:,:3].dot(init_tfm[:3,:3].T) + init_tfm[:3,3][None,:] # color_cloud now is in global frame
raw_color_cloud = np.array(color_cloud)
##########################################
### remove the shadow points on the desk
##########################################
color_cloud = remove_shadow(color_cloud)
if tracked_nodes is not None:
color_cloud = rope_shape_filter(tracked_nodes, color_cloud)
if index == 0:
rope_xyz = extract_red(rgb, depth, T_w_k)
rope_xyz = clouds.downsample(rope_xyz, .01)
rope_xyz = rope_xyz.dot(init_tfm[:3, :3].T) + init_tfm[:3,3][None,:]
rope_nodes = rope_initialization.find_path_through_point_cloud(rope_xyz) # rope_nodes and rope_xyz are in global frame
# print rope_nodes
table_height = rope_xyz[:,2].min() - .02
table_xml = make_table_xml(translation=[1, 0, table_height], extents=[.85, .55, .01])
env.LoadData(table_xml)
bulletsimpy.sim_params.scale = 10
bulletsimpy.sim_params.maxSubSteps = 200
if rope_params is None:
rope_params = bulletsimpy.CapsuleRopeParams()
rope_params.radius = 0.005
#angStiffness: a rope with a higher angular stiffness seems to have more resistance to bending.
#orig self.rope_params.angStiffness = .1
rope_params.angStiffness = .1
#A higher angular damping causes the ropes joints to change angle slower.
#This can cause the rope to be dragged at an angle by the arm in the air, instead of falling straight.
#orig self.rope_params.angDamping = 1
rope_params.angDamping = 1
#orig self.rope_params.linDamping = .75
#Not sure what linear damping is, but it seems to limit the linear accelertion of centers of masses.
rope_params.linDamping = .75
#Angular limit seems to be the minimum angle at which the rope joints can bend.
#A higher angular limit increases the minimum radius of curvature of the rope.
rope_params.angLimit = .4
#TODO--Find out what the linStopErp is
#This could be the tolerance for error when the joint is at or near the joint limit
rope_params.linStopErp = .2
bt_env = bulletsimpy.BulletEnvironment(env, [])
bt_env.SetGravity([0,0,-0.1])
rope = bulletsimpy.CapsuleRope(bt_env, 'rope', rope_nodes, rope_params)
continue
#==============================================================================
# rope_nodes = rope.GetNodes()
# #print "rope nodes in camera frame"
# R = init_tfm[:3,:3].T
# t = - R.dot(init_tfm[:3,3])
# rope_in_camera_frame = rope_nodes.dot(R.T) + t[None,:]
# #print rope_in_camera_frame
# uvs = XYZ_to_xy(rope_in_camera_frame[:,0], rope_in_camera_frame[:,1], rope_in_camera_frame[:,2], asus_xtion_pro_f)
# uvs = np.vstack(uvs)
# #print uvs
# #print "uvs"
# uvs = uvs.astype(int)
#
# n_rope_nodes = len(rope_nodes)
#
# DEPTH_OCCLUSION_DIST = .03
# occ_dist = DEPTH_OCCLUSION_DIST
#
# vis = np.ones(n_rope_nodes)
#
# rope_depth_in_camera = np.array(rope_in_camera_frame)
# depth_xyz = depth_to_xyz(depth, asus_xtion_pro_f)
#
# for i in range(n_rope_nodes):
# u = uvs[0, i]
# v = uvs[1, i]
#
# neighbor_radius = 10;
# v_range = [max(0, v-neighbor_radius), v+neighbor_radius+1]
# u_range = [max(0, u-neighbor_radius), u+neighbor_radius+1]
#
# xyzs = depth_xyz[v_range[0]:v_range[1], u_range[0]:u_range[1]]
#
# xyzs = np.reshape(xyzs, (xyzs.shape[0]*xyzs.shape[1], xyzs.shape[2]))
# dists_to_origin = np.linalg.norm(xyzs, axis=1)
#
# dists_to_origin = dists_to_origin[np.isfinite(dists_to_origin)]
#
# #print dists_to_origin
# min_dist_to_origin = np.min(dists_to_origin)
#
# print v, u, min_dist_to_origin, np.linalg.norm(rope_in_camera_frame[i])
#
# if min_dist_to_origin + occ_dist > np.linalg.norm(rope_in_camera_frame[i]):
# vis[i] = 1
# else:
# vis[i] = 0
#
# #print "vis result"
# #print vis
#
# rope_depth_in_global = rope_depth_in_camera.dot(init_tfm[:3,:3].T) + init_tfm[:3,3][None,:]
#==============================================================================
depth_cloud = extract_cloud(depth, T_w_k)
depth_cloud[:,:3] = depth_cloud[:,:3].dot(init_tfm[:3,:3].T) + init_tfm[:3,3][None,:] # depth_cloud now is in global frame
[tracked_nodes, new_stdev] = bulletsimpy.py_tracking(rope, bt_env, init_tfm, color_cloud, rgb, depth, 5, stdev)
stdev = new_stdev
#print tracked_nodes
#if index % 10 != 0:
# continue
print index
xx, yy = np.mgrid[-1:3, -1:3]
zz = np.ones(xx.shape) * table_height
table_cloud = [xx, yy, zz]
fig.clf()
ax = fig.gca(projection='3d')
ax.set_autoscale_on(False)
print init_tfm[:,3]
ax.plot(depth_cloud[:,0], depth_cloud[:,1], depth_cloud[:,2], 'go', alpha=0.1)
ax.plot(color_cloud[:,0], color_cloud[:,1], color_cloud[:,2]-0.1, 'go')
ax.plot(raw_color_cloud[:,0], raw_color_cloud[:,1], raw_color_cloud[:,2] -0.2, 'ro')
#ax.plot(rope_depth_in_global[:,0], rope_depth_in_global[:,1], rope_depth_in_global[:,2], 'ro')
ax.plot_surface(table_cloud[0], table_cloud[1], table_cloud[2], color = (0,1,0,0.5))
ax.plot(tracked_nodes[:,0], tracked_nodes[:,1], tracked_nodes[:,2], 'bo')
ax.plot([init_tfm[0,3]], [init_tfm[1,3]], [init_tfm[2,3]], 'ro')
fig.show()
raw_input()
def eval_on_holdout(args, sim_env):
feature_fn, _, num_features, actions = select_feature_fn(args)
weightfile = h5py.File(args.weightfile, 'r')
weights = weightfile['weights'][:]
w0 = weightfile['w0'][()] if 'w0' in weightfile else 0
weightfile.close()
assert weights.shape[0] == num_features, "Dimensions of weights and features don't match. Make sure the right feature is being used"
holdoutfile = h5py.File(args.holdoutfile, 'r')
tasks = eval_util.get_specified_tasks(args.tasks, args.taskfile, args.i_start, args.i_end)
holdout_items = eval_util.get_holdout_items(holdoutfile, tasks)
num_successes = 0
num_total = 0
for i_task, demo_id_rope_nodes in holdout_items:
print "task %s" % i_task
sim_util.reset_arms_to_side(sim_env)
redprint("Replace rope")
rope_xyz = demo_id_rope_nodes["rope_nodes"][:]
# Transform rope_nodes from the kinect's frame into the frame of the PR2
if 'frame' not in demo_id_rope_nodes or demo_id_rope_nodes['frame'][()] != 'r':
redprint("Transforming rope into frame of robot")
rope_xyz = rope_xyz.dot(GlobalVars.init_tfm[:3,:3].T) + GlobalVars.init_tfm[:3,3][None,:]
rope_nodes = rope_initialization.find_path_through_point_cloud(rope_xyz)
# don't call replace_rope and sim.settle() directly. use time machine interface for deterministic results!
time_machine = sim_util.RopeSimTimeMachine(rope_nodes, sim_env)
if args.animation:
sim_env.viewer.Step()
eval_util.save_task_results_init(args.resultfile, sim_env, i_task, rope_nodes, args.exec_rope_params)
for i_step in range(args.num_steps):
print "task %s step %i" % (i_task, i_step)
sim_util.reset_arms_to_side(sim_env)
redprint("Observe point cloud")
new_xyz = sim_env.sim.observe_cloud()
state = ("eval_%i"%get_unique_id(), new_xyz)
redprint("Choosing an action")
q_values = [q_value_fn(state, action, feature_fn, weights, w0) for action in actions]
q_values_root = q_values
time_machine.set_checkpoint('depth_0_%i'%i_step, sim_env)
assert args.lookahead_width>= 1, 'Lookahead branches set to zero will fail to select any action'
agenda = sorted(zip(q_values, actions), key = lambda v: -v[0])[:args.lookahead_width]
agenda = [(v, a, 'depth_0_%i'%i_step, a) for (v, a) in agenda] # state is (value, most recent action, checkpoint id, root action)
best_root_action = None
for depth in range(args.lookahead_depth):
expansion_results = []
for (branch, (q, a, chkpt, r_a)) in enumerate(agenda):
time_machine.restore_from_checkpoint(chkpt, sim_env, sim_util.get_rope_params(args.lookahead_rope_params))
cur_xyz = sim_env.sim.observe_cloud()
success, _, _, full_trajs = \
compute_trans_traj(sim_env, cur_xyz, GlobalVars.actions[a], animate=args.animation, interactive=False)
if args.animation:
sim_env.viewer.Step()
if is_knot(sim_env.sim.rope.GetControlPoints()):
best_root_action = r_a
break
result_cloud = sim_env.sim.observe_cloud()
result_chkpt = 'depth_%i_branch_%i_%i'%(depth+1, branch, i_step)
if depth != args.lookahead_depth-1: # don't save checkpoint at the last depth to save computation time
time_machine.set_checkpoint(result_chkpt, sim_env)
expansion_results.append((result_cloud, a, success, result_chkpt, r_a))
if best_root_action is not None:
redprint('Knot Found, stopping search early')
break
agenda = []
for (cld, incoming_a, success, chkpt, r_a) in expansion_results:
if not success:
agenda.append((-np.inf, actions[0], chkpt, r_a)) # TODO why first action?
continue
next_state = ("eval_%i"%get_unique_id(), cld)
q_values = [(q_value_fn(next_state, action, feature_fn, weights, w0), action, chkpt, r_a) for action in actions]
agenda.extend(q_values)
agenda.sort(key = lambda v: -v[0])
agenda = agenda[:args.lookahead_width]
first_root_action = agenda[0][-1]
if all(r_a == first_root_action for (_, _, _, r_a) in agenda):
best_root_action = first_root_action
redprint('All best actions have same root, stopping search early')
break
if best_root_action is None:
best_root_action = agenda[0][-1]
time_machine.restore_from_checkpoint('depth_0_%i'%i_step, sim_env, sim_util.get_rope_params(args.exec_rope_params))
eval_stats = eval_util.EvalStats()
eval_stats.success, eval_stats.feasible, eval_stats.misgrasp, full_trajs = \
compute_trans_traj(sim_env, new_xyz, GlobalVars.actions[best_root_action], animate=args.animation, interactive=args.interactive)
print "BEST ACTION:", best_root_action
eval_util.save_task_results_step(args.resultfile, sim_env, i_task, i_step, eval_stats, best_root_action, full_trajs, q_values_root)
if is_knot(sim_env.sim.rope.GetControlPoints()):
redprint("KNOT TIED!")
break;
if is_knot(sim_env.sim.rope.GetControlPoints()):
num_successes += 1
num_total += 1
redprint('Eval Successes / Total: ' + str(num_successes) + '/' + str(num_total))
Globals.env = openravepy.Environment()
Globals.env.StopSimulation()
Globals.env.Load("robots/pr2-beta-static.zae")
Globals.robot = Globals.env.GetRobots()[0]
actionfile = h5py.File(args.actionfile, 'r')
outfile = h5py.File(args.outfile, 'a')
init_rope_xyz, _ = load_fake_data_segment(actionfile, args.fake_data_segment, args.fake_data_transform) # this also sets the torso (torso_lift_joint) to the height in the data
table_height = init_rope_xyz[:,2].mean() - .02
table_xml = make_table_xml(translation=[1, 0, table_height], extents=[.85, .55, .01])
Globals.env.LoadData(table_xml)
Globals.sim = ropesim.Simulation(Globals.env, Globals.robot)
# create rope from rope in data
rope_nodes = rope_initialization.find_path_through_point_cloud(init_rope_xyz)
Globals.sim.create(rope_nodes)
# move arms to the side
reset_arms_to_side()
Globals.viewer = trajoptpy.GetViewer(Globals.env)
print "move viewer to viewpoint that isn't stupid"
print "then hit 'p' to continue"
Globals.viewer.Idle()
if use_dagger:
# Load initial state of first complete trajectory
dagger_states = h5py.File(args.dagger_states_file, 'r')
task_indices = sorted(dagger_states.keys())
curr_task_index = 0
curr_step_index = 0
