def quad_features_noregcostsq(self, state, action):
feat = np.zeros(1 + 2*self.num_actions)
if action == 'done':
return feat
s = registration_cost_cheap(clouds.downsample(state[1], DS_SIZE_2), # Added downsampling for multi-scale MMQL
clouds.downsample(self.get_ds_cloud(action), DS_SIZE_2)) # Added downsampling for multi-scale MMQL
feat[0] = s
feat[1+self.action_to_ind[action]] = s
feat[1+self.num_actions+self.action_to_ind[action]] = 1
return feat
def simulate_demo_traj(sim_env, new_xyz, seg_info, full_trajs, ignore_infeasibility=True, animate=False, interactive=False):
sim_util.reset_arms_to_side(sim_env)
old_xyz = np.squeeze(seg_info["cloud_xyz"])
old_xyz = clouds.downsample(old_xyz, DS_SIZE)
new_xyz = clouds.downsample(new_xyz, DS_SIZE)
handles = []
if animate:
handles.append(sim_env.env.plot3(old_xyz,5, (1,0,0)))
handles.append(sim_env.env.plot3(new_xyz,5, (0,0,1)))
miniseg_starts, miniseg_ends, _ = sim_util.split_trajectory_by_gripper(seg_info, thresh = GRIPPER_ANGLE_THRESHOLD)
success = True
feasible = True
misgrasp = False
print colorize.colorize("mini segments:", "red"), miniseg_starts, miniseg_ends
for (i_miniseg, (i_start, i_end)) in enumerate(zip(miniseg_starts, miniseg_ends)):
if i_miniseg >= len(full_trajs): break
full_traj = full_trajs[i_miniseg]
for lr in 'lr':
gripper_open = sim_util.binarize_gripper(seg_info["%s_gripper_joint"%lr][i_start], GRIPPER_ANGLE_THRESHOLD)
prev_gripper_open = sim_util.binarize_gripper(seg_info["%s_gripper_joint"%lr][i_start-1], GRIPPER_ANGLE_THRESHOLD) if i_start != 0 else False
if not sim_util.set_gripper_maybesim(sim_env, lr, gripper_open, prev_gripper_open):
redprint("Grab %s failed" % lr)
misgrasp = True
success = False
if not success: break
if len(full_traj[0]) > 0:
if not eval_util.traj_is_safe(sim_env, full_traj, COLLISION_DIST_THRESHOLD):
redprint("Trajectory not feasible")
feasible = False
if feasible or ignore_infeasibility:
success &= sim_util.sim_full_traj_maybesim(sim_env, full_traj, animate=animate, interactive=interactive)
else:
success = False
if not success: break
sim_env.sim.settle(animate=animate)
sim_env.sim.release_rope('l')
sim_env.sim.release_rope('r')
sim_util.reset_arms_to_side(sim_env)
if animate:
sim_env.viewer.Step()
return success, feasible, misgrasp, full_trajs
def get_downsampled_clouds(values, DS_SIZE):
cloud_list = []
shapes = []
for seg in values:
cloud = seg["cloud_xyz"]
cloud = cloud[...].copy()
if cloud.shape[0] > 200:
down_cloud = clouds.downsample(cloud, DS_SIZE)
else:
down_cloud = clouds.downsample(cloud, DS_SIZE)
cloud_list.append(down_cloud)
shapes.append(down_cloud.shape)
return cloud_list, shapes
def landmark_features(self, state, action):
if state[0] in ActionSet.caches['landmarks']:
return ActionSet.caches['landmarks'][state[0]]
feat = np.empty(len(self.landmarks))
if ActionSet.args and 'parallel' in ActionSet.args and ActionSet.args.parallel:
landmarks = [self.landmarks[str(i)]['cloud_xyz'][()] for i in range(len(self.landmarks))]
costs = Parallel(n_jobs=-1,verbose=0)(delayed(registration_cost_cheap)( \
clouds.downsample(state[1], DS_SIZE_2), # Added downsampling for multi-scale MMQL
clouds.downsample(l, DS_SIZE_2)) # Added downsampling for multi-scale MMQL
for l in landmarks)
feat = np.asarray(costs)
else:
for i in range(len(self.landmarks)):
landmark = self.landmarks[str(i)]
# Added downsampling for multi-state MMQL
downsampled_state = clouds.downsample(state[1], DS_SIZE_2)
downsampled_landmark = clouds.downsample(landmark['cloud_xyz'][()], DS_SIZE_2)
feat[i] = registration_cost_cheap(downsampled_state,
downsampled_landmark)
ActionSet.caches['landmarks'][state[0]] = feat
return feat
def get_downsampled_clouds(values):
#TODO: actually fix this. It is probably better to upsample the derived segments
cloud_list = []
shapes = []
for seg in values:
cloud = seg["cloud_xyz"]
#This eliminates the z dimension
cloud = cloud[...].copy()
#cloud[:, 2] = np.mean(cloud[:, 2])
if cloud.shape[0] > 200:
down_cloud = clouds.downsample(cloud, DS_SIZE)
else:
down_cloud = clouds.downsample(cloud, 0.01 * DS_SIZE)
#down_cloud[20:, 2] = 0.66
#down_cloud[:, 2] = np.mean(down_cloud[:, 2])
cloud_list.append(down_cloud)
shapes.append(down_cloud.shape)
#print "cloud_shape = ", down_cloud.shape
#return [clouds.downsample(seg["cloud_xyz"], DS_SIZE) for seg in demofile.values()]
return cloud_list, shapes
def grabcut(rgb, depth, T_w_k):
xyz_k = clouds.depth_to_xyz(depth, berkeley_pr2.f)
xyz_w = xyz_k.dot(T_w_k[:3, :3].T) + T_w_k[:3, 3][None, None, :]
valid_mask = depth > 0
import interactive_roi as ir
xys = ir.get_polyline(rgb, "rgb")
xy_corner1 = np.clip(np.array(xys).min(axis=0), [0, 0], [639, 479])
xy_corner2 = np.clip(np.array(xys).max(axis=0), [0, 0], [639, 479])
polymask = ir.mask_from_poly(xys)
#cv2.imshow("mask",mask)
xy_tl = np.array([xy_corner1, xy_corner2]).min(axis=0)
xy_br = np.array([xy_corner1, xy_corner2]).max(axis=0)
xl, yl = xy_tl
w, h = xy_br - xy_tl
mask = np.zeros((h, w), dtype='uint8')
mask[polymask[yl:yl + h, xl:xl + w] > 0] = cv2.GC_PR_FGD
print mask.shape
#mask[h//4:3*h//4, w//4:3*w//4] = cv2.GC_PR_FGD
tmp1 = np.zeros((1, 13 * 5))
tmp2 = np.zeros((1, 13 * 5))
cv2.grabCut(rgb[yl:yl + h, xl:xl + w, :],
mask, (0, 0, 0, 0),
tmp1,
tmp2,
10,
mode=cv2.GC_INIT_WITH_MASK)
mask = mask % 2
#mask = ndi.binary_erosion(mask, utils_images.disk(args.erode)).astype('uint8')
contours = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)[0]
cv2.drawContours(rgb[yl:yl + h, xl:xl + w, :],
contours,
-1, (0, 255, 0),
thickness=2)
cv2.imshow('rgb', rgb)
print "press enter to continue"
cv2.waitKey()
zsel = xyz_w[yl:yl + h, xl:xl + w, 2]
mask = (mask % 2 == 1) & np.isfinite(zsel) # & (zsel - table_height > -1)
mask &= valid_mask[yl:yl + h, xl:xl + w]
xyz_sel = xyz_w[yl:yl + h, xl:xl + w, :][mask.astype('bool')]
return clouds.downsample(xyz_sel, .01)
def _warp_hmats(self, state, action):
hit, value = self.check_cache(state, action)
if hit:
return value
else:
if ActionSet.args and ActionSet.args.gripper_weighting:
interest_pts = get_closing_pts(self.actionfile[action])
else:
interest_pts = None
[warped_trajs, rc, warped_rope_xyz] = warp_hmats(self.get_ds_cloud(action),
clouds.downsample(state[1], DS_SIZE_2), # Added downsampling for multi-scale MMQL
[(lr, self.actionfile[action][ln]['hmat']) for lr, ln in zip('lr', self.link_names)],
interest_pts)
self.store_cache(state, action, [warped_trajs, rc, warped_rope_xyz])
return [warped_trajs, rc, warped_rope_xyz]
def test_features(args, feature_type):
"""
Accepted values for feature_type: "sc" and "rope_dist"
"""
if feature_type == "sc":
feature_fn, num_features, act_file = get_sc_feature_fn(args.actionfile)
elif feature_type == "rope_dist":
fns = [get_bias_feature_fn, get_sc_feature_fn, get_rope_dist_feat_fn]
feature_fn, num_features, act_file = concatenate_fns(fns, args.actionfile)
else:
print "Argument to test_features is not one of the accepted types: sc, rope_dist"
return
for name, seg_info in act_file.iteritems():
print feature_fn([name, clouds.downsample(seg_info['cloud_xyz'], DS_SIZE)], name), "\n"
def extract_red(rgb, depth, T_w_k):
"""
extract red points and downsample
"""
hsv = cv2.cvtColor(rgb, cv2.COLOR_BGR2HSV)
h = hsv[:,:,0]
s = hsv[:,:,1]
v = hsv[:,:,2]
h_mask = (h<15) | (h>145)
s_mask = (s > 30 )
v_mask = (v > 100)
red_mask = h_mask & s_mask & v_mask
valid_mask = depth > 0
xyz_k = clouds.depth_to_xyz(depth, berkeley_pr2.f)
xyz_w = xyz_k.dot(T_w_k[:3,:3].T) + T_w_k[:3,3][None,None,:]
z = xyz_w[:,:,2]
z0 = xyz_k[:,:,2]
height_mask = xyz_w[:,:,2] > .7 # TODO pass in parameter
good_mask = red_mask & height_mask & valid_mask
good_mask = skim.remove_small_objects(good_mask,min_size=64)
if DEBUG_PLOTS:
cv2.imshow("z0",z0/z0.max())
cv2.imshow("z",z/z.max())
cv2.imshow("hue", h_mask.astype('uint8')*255)
cv2.imshow("sat", s_mask.astype('uint8')*255)
cv2.imshow("val", v_mask.astype('uint8')*255)
cv2.imshow("final",good_mask.astype('uint8')*255)
cv2.imshow("rgb", rgb)
cv2.waitKey()
good_xyz = xyz_w[good_mask]
return clouds.downsample(good_xyz, .025)
def extract_red(rgb, depth, T_w_k):
"""
extract red points and downsample
"""
hsv = cv2.cvtColor(rgb, cv2.COLOR_BGR2HSV)
h = hsv[:,:,0]
s = hsv[:,:,1]
v = hsv[:,:,2]
red_mask = ((h<10) | (h>150)) & (s > 100) & (v > 100)
valid_mask = depth > 0
xyz_k = clouds.depth_to_xyz(depth, berkeley_pr2.f)
xyz_w = xyz_k.dot(T_w_k[:3,:3].T) + T_w_k[:3,3][None,None,:]
z = xyz_w[:,:,2]
z0 = xyz_k[:,:,2]
if DEBUG_PLOTS:
cv2.imshow("z0",z0/z0.max())
cv2.imshow("z",z/z.max())
cv2.imshow("rgb", rgb)
cv2.waitKey()
height_mask = xyz_w[:,:,2] > .7 # TODO pass in parameter
good_mask = red_mask & height_mask & valid_mask
if DEBUG_PLOTS:
cv2.imshow("red",red_mask.astype('uint8')*255)
cv2.imshow("above_table", height_mask.astype('uint8')*255)
cv2.imshow("red and above table", good_mask.astype('uint8')*255)
print "press enter to continue"
cv2.waitKey()
good_xyz = xyz_w[good_mask]
return clouds.downsample(good_xyz, .025)
def auto_choose(demofile, new_xyz, only_original_segments):
"""
@param demofile:
@param new_xyz:
@param only_original_segments: if true, then only the original_segments will be registered with
@return:
"""
import pprint
"""Return the segment with the lowest warping cost. Takes about 2 seconds."""
parallel = True
if parallel:
from joblib import Parallel, delayed
items = demofile.items()
if only_original_segments:
#remove all derived segments from items
print("Only registering with the original segments")
items = [item for item in items if not "derived" in item[1].keys()]
unzipped_items = zip(*items)
keys = unzipped_items[0]
values = unzipped_items[1]
ds_clouds, shapes = get_downsampled_clouds(values)
ds_new = clouds.downsample(new_xyz, 0.01 * DS_SIZE)
#print 'ds_new_len shape', ds_new.shape
if parallel:
before = time.time()
#TODO: change back n_jobs=12 ?
costs = Parallel(n_jobs=8, verbose=0)(delayed(registration_cost)(ds_cloud, ds_new) for ds_cloud in ds_clouds)
after = time.time()
print "Parallel registration time in seconds =", after - before
else:
costs = []
for (i, ds_cloud) in enumerate(ds_clouds):
costs.append(registration_cost(ds_cloud, ds_new))
print(("completed %i/%i" % (i + 1, len(ds_clouds))))
#print(("costs\n", costs))
ibest = np.argmin(costs)
print "ibest = ", ibest
#pprint.pprint(zip(keys, costs, shapes))
#print keys
print "best key = ", keys[ibest]
print "best cost = ", costs[ibest]
return keys[ibest]
def grabcut(rgb, depth, T_w_k):
xyz_k = clouds.depth_to_xyz(depth, berkeley_pr2.f)
xyz_w = xyz_k.dot(T_w_k[:3,:3].T) + T_w_k[:3,3][None,None,:]
valid_mask = depth > 0
import interactive_roi as ir
xys = ir.get_polyline(rgb, "rgb")
xy_corner1 = np.clip(np.array(xys).min(axis=0), [0,0], [639,479])
xy_corner2 = np.clip(np.array(xys).max(axis=0), [0,0], [639,479])
polymask = ir.mask_from_poly(xys)
#cv2.imshow("mask",mask)
xy_tl = np.array([xy_corner1, xy_corner2]).min(axis=0)
xy_br = np.array([xy_corner1, xy_corner2]).max(axis=0)
xl, yl = xy_tl
w, h = xy_br - xy_tl
mask = np.zeros((h,w),dtype='uint8')
mask[polymask[yl:yl+h, xl:xl+w] > 0] = cv2.GC_PR_FGD
print mask.shape
#mask[h//4:3*h//4, w//4:3*w//4] = cv2.GC_PR_FGD
tmp1 = np.zeros((1, 13 * 5))
tmp2 = np.zeros((1, 13 * 5))
cv2.grabCut(rgb[yl:yl+h, xl:xl+w, :],mask,(0,0,0,0),tmp1, tmp2,10,mode=cv2.GC_INIT_WITH_MASK)
mask = mask % 2
#mask = ndi.binary_erosion(mask, utils_images.disk(args.erode)).astype('uint8')
contours = cv2.findContours(mask,cv2.RETR_LIST,cv2.CHAIN_APPROX_NONE)[0]
cv2.drawContours(rgb[yl:yl+h, xl:xl+w, :],contours,-1,(0,255,0),thickness=2)
cv2.imshow('rgb', rgb)
print "press enter to continue"
cv2.waitKey()
zsel = xyz_w[yl:yl+h, xl:xl+w, 2]
mask = (mask%2==1) & np.isfinite(zsel)# & (zsel - table_height > -1)
mask &= valid_mask[yl:yl+h, xl:xl+w]
xyz_sel = xyz_w[yl:yl+h, xl:xl+w,:][mask.astype('bool')]
return clouds.downsample(xyz_sel, .01)
def extract_red(rgb, depth, T_w_k):
"""
extract red points and downsample
"""
hsv = cv2.cvtColor(rgb, cv2.COLOR_BGR2HSV)
h = hsv[:, :, 0]
s = hsv[:, :, 1]
v = hsv[:, :, 2]
red_mask = ((h < 10) | (h > 150)) & (s > 100) & (v > 100)
valid_mask = depth > 0
xyz_k = clouds.depth_to_xyz(depth, berkeley_pr2.f)
xyz_w = xyz_k.dot(T_w_k[:3, :3].T) + T_w_k[:3, 3][None, None, :]
z = xyz_w[:, :, 2]
z0 = xyz_k[:, :, 2]
if DEBUG_PLOTS:
cv2.imshow("z0", z0 / z0.max())
cv2.imshow("z", z / z.max())
cv2.imshow("rgb", rgb)
cv2.waitKey()
height_mask = xyz_w[:, :, 2] > .7 # TODO pass in parameter
good_mask = red_mask & height_mask & valid_mask
if DEBUG_PLOTS:
cv2.imshow("red", red_mask.astype('uint8') * 255)
cv2.imshow("above_table", height_mask.astype('uint8') * 255)
cv2.imshow("red and above table", good_mask.astype('uint8') * 255)
print "press enter to continue"
cv2.waitKey()
good_xyz = xyz_w[good_mask]
return clouds.downsample(good_xyz, .025)
def compute_trans_traj(sim_env,
new_xyz,
seg_info,
ignore_infeasibility=True,
animate=False,
interactive=False):
sim_util.reset_arms_to_side(sim_env)
redprint("Generating end-effector trajectory")
old_xyz = np.squeeze(seg_info["cloud_xyz"])
old_xyz = clouds.downsample(old_xyz, DS_SIZE)
new_xyz = clouds.downsample(new_xyz, DS_SIZE)
link_names = ["%s_gripper_tool_frame" % lr for lr in ('lr')]
hmat_list = [(lr, seg_info[ln]['hmat'])
for lr, ln in zip('lr', link_names)]
if GlobalVars.gripper_weighting:
interest_pts = get_closing_pts(seg_info)
else:
interest_pts = None
lr2eetraj, _, old_xyz_warped = warp_hmats_tfm(old_xyz, new_xyz, hmat_list,
interest_pts)
handles = []
if animate:
handles.append(sim_env.env.plot3(old_xyz, 5, (1, 0, 0)))
handles.append(sim_env.env.plot3(new_xyz, 5, (0, 0, 1)))
handles.append(sim_env.env.plot3(old_xyz_warped, 5, (0, 1, 0)))
miniseg_starts, miniseg_ends, _ = sim_util.split_trajectory_by_gripper(
seg_info, thresh=GRIPPER_ANGLE_THRESHOLD)
success = True
feasible = True
misgrasp = False
print colorize.colorize("mini segments:",
"red"), miniseg_starts, miniseg_ends
full_trajs = []
prev_vals = {lr: None for lr in 'lr'}
for (i_miniseg, (i_start,
i_end)) in enumerate(zip(miniseg_starts, miniseg_ends)):
################################
redprint("Generating joint trajectory for part %i" % (i_miniseg))
# figure out how we're gonna resample stuff
# Use inverse kinematics to get trajectory for initializing TrajOpt,
# since demonstrations library does not contain joint angle data for
# left and right arms
ee_hmats = {}
for lr in 'lr':
ee_link_name = "%s_gripper_tool_frame" % lr
# TODO: Change # of timesteps for resampling?
ee_hmats[ee_link_name] = resampling.interp_hmats(
np.arange(i_end + 1 - i_start), np.arange(i_end + 1 - i_start),
lr2eetraj[lr][i_start:i_end + 1])
lr2oldtraj = get_old_joint_traj_ik(sim_env, ee_hmats, prev_vals,
i_start, i_end)
#lr2oldtraj = {}
#for lr in 'lr':
# manip_name = {"l":"leftarm", "r":"rightarm"}[lr]
# old_joint_traj = asarray(seg_info[manip_name][i_start:i_end+1])
# #print (old_joint_traj[1:] - old_joint_traj[:-1]).ptp(axis=0), i_start, i_end
# if sim_util.arm_moved(old_joint_traj):
# lr2oldtraj[lr] = old_joint_traj
if len(lr2oldtraj) > 0:
old_total_traj = np.concatenate(lr2oldtraj.values(), 1)
JOINT_LENGTH_PER_STEP = .1
_, timesteps_rs = sim_util.unif_resample(old_total_traj,
JOINT_LENGTH_PER_STEP)
####
### Generate fullbody traj
bodypart2traj = {}
for (lr, old_joint_traj) in lr2oldtraj.items():
manip_name = {"l": "leftarm", "r": "rightarm"}[lr]
old_joint_traj_rs = mu.interp2d(timesteps_rs,
np.arange(len(old_joint_traj)),
old_joint_traj)
ee_link_name = "%s_gripper_tool_frame" % lr
new_ee_traj = lr2eetraj[lr][i_start:i_end + 1]
new_ee_traj_rs = resampling.interp_hmats(
timesteps_rs, np.arange(len(new_ee_traj)), new_ee_traj)
print "planning trajectory following"
new_joint_traj, pose_errs = planning.plan_follow_traj(
sim_env.robot, manip_name, sim_env.robot.GetLink(ee_link_name),
new_ee_traj_rs, old_joint_traj_rs)
prev_vals[lr] = new_joint_traj[-1]
part_name = {"l": "larm", "r": "rarm"}[lr]
bodypart2traj[part_name] = new_joint_traj
################################
redprint("Executing joint trajectory for part %i using arms '%s'" %
(i_miniseg, bodypart2traj.keys()))
full_traj = sim_util.getFullTraj(sim_env, bodypart2traj)
full_trajs.append(full_traj)
for lr in 'lr':
gripper_open = sim_util.binarize_gripper(
seg_info["%s_gripper_joint" % lr][i_start],
GRIPPER_ANGLE_THRESHOLD)
prev_gripper_open = sim_util.binarize_gripper(
seg_info["%s_gripper_joint" % lr][i_start - 1],
GRIPPER_ANGLE_THRESHOLD) if i_start != 0 else False
if not sim_util.set_gripper_maybesim(sim_env, lr, gripper_open,
prev_gripper_open):
redprint("Grab %s failed" % lr)
misgrasp = True
success = False
if not success: break
if len(full_traj[0]) > 0:
if not eval_util.traj_is_safe(sim_env, full_traj,
COLLISION_DIST_THRESHOLD):
redprint("Trajectory not feasible")
feasible = False
if feasible or ignore_infeasibility:
success &= sim_util.sim_full_traj_maybesim(
sim_env,
full_traj,
animate=animate,
interactive=interactive)
else:
success = False
if not success: break
sim_env.sim.settle(animate=animate)
sim_env.sim.release_rope('l')
sim_env.sim.release_rope('r')
sim_util.reset_arms_to_side(sim_env)
if animate:
sim_env.viewer.Step()
return success, feasible, misgrasp, full_trajs
def get_downsampled_clouds(demofile):
return [clouds.downsample(seg["cloud_xyz"], DS_SIZE) for seg in demofile.values()]
def get_downsampled_clouds(demofile):
return [
clouds.downsample(seg["cloud_xyz"], DS_SIZE)
for seg in demofile.values()
]
def main():
import argparse, h5py, os
import matplotlib.pyplot as plt
from rapprentice import clouds, plotting_plt
import registration
import time
parser = argparse.ArgumentParser()
parser.add_argument("input_file", type=str)
parser.add_argument("--output_folder", type=str, default="")
parser.add_argument("--i_start", type=int, default=0)
parser.add_argument("--i_end", type=int, default=-1)
regtype_choices = ['rpm', 'rpm-cheap', 'rpm-bij', 'rpm-bij-cheap']
parser.add_argument("--regtypes", type=str, nargs='*', choices=regtype_choices, default=regtype_choices)
parser.add_argument("--plot_color", type=int, default=1)
parser.add_argument("--proj", type=int, default=1, help="project 3d visualization into 2d")
parser.add_argument("--visual_prior", type=int, default=1)
parser.add_argument("--plotting", type=int, default=1)
args = parser.parse_args()
def plot_cb_gen(output_prefix, args, x_color, y_color):
def plot_cb(x_nd, y_md, corr_nm, f, iteration):
if args.plot_color:
plotting_plt.plot_tps_registration(x_nd, y_md, f, x_color = x_color, y_color = y_color, proj_2d=args.proj)
else:
plotting_plt.plot_tps_registration(x_nd, y_md, f, proj_2d=args.proj)
# save plot to file
if output_prefix is not None:
plt.savefig(output_prefix + "_iter" + str(iteration) + '.png')
return plot_cb
def plot_cb_bij_gen(output_prefix, args, x_color, y_color):
def plot_cb_bij(x_nd, y_md, xtarg_nd, corr_nm, wt_n, f):
if args.plot_color:
plotting_plt.plot_tps_registration(x_nd, y_md, f, res = (.3, .3, .12), x_color = x_color, y_color = y_color, proj_2d=args.proj)
else:
plotting_plt.plot_tps_registration(x_nd, y_md, f, res = (.4, .3, .12), proj_2d=args.proj)
# save plot to file
if output_prefix is not None:
plt.savefig(output_prefix + "_iter" + str(iteration) + '.png')
return plot_cb_bij
# preprocess and downsample clouds
DS_SIZE = 0.025
infile = h5py.File(args.input_file)
source_clouds = {}
target_clouds = {}
for i in range(args.i_start, len(infile) if args.i_end==-1 else args.i_end):
source_cloud = clouds.downsample(infile[str(i)]['source_cloud'][()], DS_SIZE)
source_clouds[i] = source_cloud
target_clouds[i] = []
for (cloud_key, target_cloud) in infile[str(i)]['target_clouds'].iteritems():
target_cloud = clouds.downsample(target_cloud[()], DS_SIZE)
target_clouds[i].append(target_cloud)
infile.close()
tps_costs = []
tps_reg_costs = []
for regtype in args.regtypes:
start_time = time.time()
costs = []
reg_costs = []
for i in range(args.i_start, len(source_clouds) if args.i_end==-1 else args.i_end):
source_cloud = source_clouds[i]
for target_cloud in target_clouds[i]:
if args.visual_prior:
vis_cost_xy = ab_cost(source_cloud, target_cloud)
else:
vis_cost_xy = None
if regtype == 'rpm':
f, corr_nm = tps_rpm(source_cloud[:,:-3], target_cloud[:,:-3],
vis_cost_xy = vis_cost_xy,
plotting=args.plotting, plot_cb = plot_cb_gen(os.path.join(args.output_folder, str(i) + "_" + cloud_key + "_rpm") if args.output_folder else None,
args,
source_cloud[:,-3:],
target_cloud[:,-3:]))
elif regtype == 'rpm-cheap':
f, corr_nm = tps_rpm(source_cloud[:,:-3], target_cloud[:,:-3],
vis_cost_xy = vis_cost_xy, n_iter = N_ITER_CHEAP, em_iter = EM_ITER_CHEAP,
plotting=args.plotting, plot_cb = plot_cb_gen(os.path.join(args.output_folder, str(i) + "_" + cloud_key + "_rpm_cheap") if args.output_folder else None,
args,
source_cloud[:,-3:],
target_cloud[:,-3:]))
elif regtype == 'rpm-bij':
x_nd = source_cloud[:,:3]
y_md = target_cloud[:,:3]
scaled_x_nd, _ = registration.unit_boxify(x_nd)
scaled_y_md, _ = registration.unit_boxify(y_md)
f,g = registration.tps_rpm_bij(scaled_x_nd, scaled_y_md, rot_reg=np.r_[1e-4, 1e-4, 1e-1], n_iter=50, reg_init=10, reg_final=.1, outlierfrac=1e-2, vis_cost_xy=vis_cost_xy,
plotting=args.plotting, plot_cb=plot_cb_bij_gen(os.path.join(args.output_folder, str(i) + "_" + cloud_key + "_rpm_bij") if args.output_folder else None,
args,
source_cloud[:,-3:],
target_cloud[:,-3:]))
elif regtype == 'rpm-bij-cheap':
x_nd = source_cloud[:,:3]
y_md = target_cloud[:,:3]
scaled_x_nd, _ = registration.unit_boxify(x_nd)
scaled_y_md, _ = registration.unit_boxify(y_md)
f,g = registration.tps_rpm_bij(scaled_x_nd, scaled_y_md, rot_reg=np.r_[1e-4, 1e-4, 1e-1], n_iter=10, outlierfrac=1e-2, vis_cost_xy=vis_cost_xy, # Note registration_cost_cheap in rope_qlearn has a different rot_reg and outlierfrac
plotting=args.plotting, plot_cb=plot_cb_bij_gen(os.path.join(args.output_folder, str(i) + "_" + cloud_key + "_rpm_bij_cheap") if args.output_folder else None,
args,
source_cloud[:,-3:],
target_cloud[:,-3:]))
costs.append(f._cost)
reg_costs.append(registration.tps_reg_cost(f))
tps_costs.append(costs)
tps_reg_costs.append(reg_costs)
print regtype, "time elapsed", time.time() - start_time
np.set_printoptions(suppress=True)
print ""
print "tps_costs"
print args.regtypes
print np.array(tps_costs).T
print ""
print "tps_reg_costs"
print args.regtypes
print np.array(tps_reg_costs).T
def get_downsampled_clouds(demofile):
if not use_rapprentice:
return get_clouds(demofile)
return [clouds.downsample(seg["cloud_xyz"], DS_SIZE) for seg in demofile.values()]
def downsample(cloud, size):
print "cloud_size", cloud.size
return clouds.downsample(cloud, size)
def get_ds_cloud(self, action):
return clouds.downsample(self.actionfile[action]['cloud_xyz'], DS_SIZE)
def simulate_demo(new_xyz, seg_info, animate=False):
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())
redprint("Generating end-effector trajectory")
handles = []
old_xyz = np.squeeze(seg_info["cloud_xyz"])
handles.append(Globals.env.plot3(old_xyz,5, (1,0,0)))
handles.append(Globals.env.plot3(new_xyz,5, (0,0,1)))
old_xyz = clouds.downsample(old_xyz, DS_SIZE)
new_xyz = clouds.downsample(new_xyz, DS_SIZE)
link_names = ["%s_gripper_tool_frame"%lr for lr in ('lr')]
hmat_list = [(lr, seg_info[ln]['hmat']) for lr, ln in zip('lr', link_names)]
lr2eetraj = warp_hmats(old_xyz, new_xyz, hmat_list)[0]
miniseg_starts, miniseg_ends = split_trajectory_by_gripper(seg_info)
success = True
print colorize.colorize("mini segments:", "red"), miniseg_starts, miniseg_ends
for (i_miniseg, (i_start, i_end)) in enumerate(zip(miniseg_starts, miniseg_ends)):
################################
redprint("Generating joint trajectory for part %i"%(i_miniseg))
# figure out how we're gonna resample stuff
lr2oldtraj = {}
for lr in 'lr':
manip_name = {"l":"leftarm", "r":"rightarm"}[lr]
old_joint_traj = asarray(seg_info[manip_name][i_start:i_end+1])
#print (old_joint_traj[1:] - old_joint_traj[:-1]).ptp(axis=0), i_start, i_end
if arm_moved(old_joint_traj):
lr2oldtraj[lr] = old_joint_traj
if len(lr2oldtraj) > 0:
old_total_traj = np.concatenate(lr2oldtraj.values(), 1)
JOINT_LENGTH_PER_STEP = .1
_, timesteps_rs = unif_resample(old_total_traj, JOINT_LENGTH_PER_STEP)
####
### Generate fullbody traj
bodypart2traj = {}
for (lr,old_joint_traj) in lr2oldtraj.items():
manip_name = {"l":"leftarm", "r":"rightarm"}[lr]
old_joint_traj_rs = mu.interp2d(timesteps_rs, np.arange(len(old_joint_traj)), old_joint_traj)
ee_link_name = "%s_gripper_tool_frame"%lr
new_ee_traj = lr2eetraj[lr][i_start:i_end+1]
new_ee_traj_rs = resampling.interp_hmats(timesteps_rs, np.arange(len(new_ee_traj)), new_ee_traj)
print "planning trajectory following"
with util.suppress_stdout():
new_joint_traj = planning.plan_follow_traj(Globals.robot, manip_name,
Globals.robot.GetLink(ee_link_name), new_ee_traj_rs,old_joint_traj_rs)[0]
part_name = {"l":"larm", "r":"rarm"}[lr]
bodypart2traj[part_name] = new_joint_traj
################################
redprint("Executing joint trajectory for part %i using arms '%s'"%(i_miniseg, bodypart2traj.keys()))
for lr in 'lr':
gripper_open = binarize_gripper(seg_info["%s_gripper_joint"%lr][i_start])
prev_gripper_open = binarize_gripper(seg_info["%s_gripper_joint"%lr][i_start-1]) if i_start != 0 else False
if not set_gripper_maybesim(lr, gripper_open, prev_gripper_open):
redprint("Grab %s failed" % lr)
success = False
if not success: break
if len(bodypart2traj) > 0:
success &= sim_traj_maybesim(bodypart2traj, animate=True)
if not success: break
Globals.sim.settle(animate=animate)
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())
Globals.sim.release_rope('l')
Globals.sim.release_rope('r')
return success
def downsample_cloud(cloud_xyz):
return clouds.downsample(cloud_xyz, DS_SIZE)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('actionfile', type=str)
parser.add_argument('cloudfile', type=str, help="result file containing the encountered clouds")
parser.add_argument('type', type=str, choices=['explore', 'compare'], help="")
parser.add_argument("--tasks", type=int, nargs='*', metavar="i_task")
parser.add_argument("--taskfile", type=str)
parser.add_argument("--i_start", type=int, default=-1, metavar="i_task")
parser.add_argument("--i_end", type=int, default=-1, metavar="i_task")
parser.add_argument("--steps", type=int, nargs='*', metavar="i_step")
parser.add_argument("--draw_rows", type=int, default=8, help="plot only the draw_rows registrations with the smallest cost")
parser.add_argument("--ds_size", type=float, default=0.025, metavar="meters")
args = parser.parse_args()
actionfile = h5py.File(args.actionfile, 'r')
cloudfile = h5py.File(args.cloudfile, 'r')
cloud_items = eval_util.get_holdout_items(cloudfile, args.tasks, args.taskfile, args.i_start, args.i_end)
if args.type == 'explore':
fs = []
reg_costs = []
old_clouds_ds = []
new_clouds_ds = []
texts = []
for i_task, _ in cloud_items:
n_steps = len(cloudfile[i_task]) - (1 if 'init' in cloudfile[i_task] else 0)
steps = [i_step for i_step in args.steps if i_step in range(n_steps)] if args.steps else range(n_steps)
for i_step in steps:
if args.steps and i_step not in args.steps:
continue
print "task %s step %i" % (i_task, i_step)
old_cloud = cloudfile[i_task][str(i_step)]['demo_cloud'][()]
old_cloud_ds = clouds.downsample(old_cloud, args.ds_size)
new_cloud = cloudfile[i_task][str(i_step)]['cloud'][()]
new_cloud_ds = clouds.downsample(new_cloud, args.ds_size)
f, corr = tps_registration.tps_rpm(old_cloud_ds[:,:3], new_cloud_ds[:,:3], n_iter=N_ITER_CHEAP, em_iter=EM_ITER_CHEAP, vis_cost_xy = tps_registration.ab_cost(old_cloud_ds, new_cloud_ds), user_data={'old_cloud':old_cloud_ds, 'new_cloud':new_cloud_ds})
reg_cost = registration.tps_reg_cost(f)
fs.append(f)
reg_costs.append(reg_cost)
old_clouds_ds.append(old_cloud_ds)
new_clouds_ds.append(new_cloud_ds)
texts.append("task %s step %i" % (i_task, i_step))
if args.draw_rows == len(fs) or (i_task == cloud_items[-1][0] and i_step == steps[-1]):
# sort based on reg_costs
fs = [f for (cost, f) in sorted(zip(reg_costs, fs))]
old_clouds_ds = [cloud for (cost, cloud) in sorted(zip(reg_costs, old_clouds_ds))]
new_clouds_ds = [cloud for (cost, cloud) in sorted(zip(reg_costs, new_clouds_ds))]
texts = [text for (cost, text) in sorted(zip(reg_costs, texts))]
print "plotting"
plot_tps_registrations_proj_2d(old_clouds_ds, new_clouds_ds, fs, (.1, .1, .04), texts)
ipy.embed()
fs[:] = []
reg_costs[:] = []
old_clouds_ds[:] = []
new_clouds_ds[:] = []
texts[:] = []
elif args.type == 'compare':
for i_task, _ in cloud_items:
for i_step in range(len(cloudfile[i_task]) - (1 if 'init' in cloudfile[i_task] else 0)):
if args.steps and i_step not in args.steps:
continue
print "task %s step %i" % (i_task, i_step)
new_cloud = cloudfile[i_task][str(i_step)]['cloud'][()]
new_cloud_ds = clouds.downsample(new_cloud, args.ds_size)
fs = []
reg_costs = []
old_clouds_ds = []
new_clouds_ds = []
texts = []
print "computing costs"
for action in actionfile.keys():
old_cloud = actionfile[action]['cloud_xyz']
old_cloud_ds = clouds.downsample(old_cloud, args.ds_size)
f, corr = tps_registration.tps_rpm(old_cloud_ds[:,:3], new_cloud_ds[:,:3], n_iter=N_ITER_CHEAP, em_iter=EM_ITER_CHEAP, vis_cost_xy = tps_registration.ab_cost(old_cloud_ds, new_cloud_ds), user_data={'old_cloud':old_cloud_ds, 'new_cloud':new_cloud_ds})
reg_cost = registration.tps_reg_cost(f)
fs.append(f)
reg_costs.append(reg_cost)
old_clouds_ds.append(old_cloud_ds)
new_clouds_ds.append(new_cloud_ds)
texts.append("task %s step %i action %s" % (i_task, i_step, action))
# sort based on reg_costs
fs = [f for (cost, f) in sorted(zip(reg_costs, fs))]
old_clouds_ds = [cloud for (cost, cloud) in sorted(zip(reg_costs, old_clouds_ds))]
if args.draw_rows != -1:
fs = fs[:args.draw_rows]
old_clouds_ds = old_clouds_ds[:args.draw_rows]
print "plotting"
plot_tps_registrations_proj_2d(old_clouds_ds, new_clouds_ds, fs, (.1, .1, .04), texts)
ipy.embed()
actionfile.close()
cloudfile.close()
def compute_trans_traj(sim_env, new_xyz, seg_info, ignore_infeasibility=True, animate=False, interactive=False):
sim_util.reset_arms_to_side(sim_env)
redprint("Generating end-effector trajectory")
old_xyz = np.squeeze(seg_info["cloud_xyz"])
old_xyz = clouds.downsample(old_xyz, DS_SIZE)
new_xyz = clouds.downsample(new_xyz, DS_SIZE)
link_names = ["%s_gripper_tool_frame"%lr for lr in ('lr')]
hmat_list = [(lr, seg_info[ln]['hmat']) for lr, ln in zip('lr', link_names)]
if GlobalVars.gripper_weighting:
interest_pts = get_closing_pts(seg_info)
else:
interest_pts = None
lr2eetraj, _, old_xyz_warped = warp_hmats_tfm(old_xyz, new_xyz, hmat_list, interest_pts)
handles = []
if animate:
handles.append(sim_env.env.plot3(old_xyz,5, (1,0,0)))
handles.append(sim_env.env.plot3(new_xyz,5, (0,0,1)))
handles.append(sim_env.env.plot3(old_xyz_warped,5, (0,1,0)))
miniseg_starts, miniseg_ends, _ = sim_util.split_trajectory_by_gripper(seg_info, thresh = GRIPPER_ANGLE_THRESHOLD)
success = True
feasible = True
misgrasp = False
print colorize.colorize("mini segments:", "red"), miniseg_starts, miniseg_ends
full_trajs = []
prev_vals = {lr:None for lr in 'lr'}
for (i_miniseg, (i_start, i_end)) in enumerate(zip(miniseg_starts, miniseg_ends)):
################################
redprint("Generating joint trajectory for part %i"%(i_miniseg))
# figure out how we're gonna resample stuff
# Use inverse kinematics to get trajectory for initializing TrajOpt,
# since demonstrations library does not contain joint angle data for
# left and right arms
ee_hmats = {}
for lr in 'lr':
ee_link_name = "%s_gripper_tool_frame"%lr
# TODO: Change # of timesteps for resampling?
ee_hmats[ee_link_name] = resampling.interp_hmats(np.arange(i_end+1-i_start), np.arange(i_end+1-i_start), lr2eetraj[lr][i_start:i_end+1])
lr2oldtraj = get_old_joint_traj_ik(sim_env, ee_hmats, prev_vals, i_start, i_end)
#lr2oldtraj = {}
#for lr in 'lr':
# manip_name = {"l":"leftarm", "r":"rightarm"}[lr]
# old_joint_traj = asarray(seg_info[manip_name][i_start:i_end+1])
# #print (old_joint_traj[1:] - old_joint_traj[:-1]).ptp(axis=0), i_start, i_end
# if sim_util.arm_moved(old_joint_traj):
# lr2oldtraj[lr] = old_joint_traj
if len(lr2oldtraj) > 0:
old_total_traj = np.concatenate(lr2oldtraj.values(), 1)
JOINT_LENGTH_PER_STEP = .1
_, timesteps_rs = sim_util.unif_resample(old_total_traj, JOINT_LENGTH_PER_STEP)
####
### Generate fullbody traj
bodypart2traj = {}
for (lr,old_joint_traj) in lr2oldtraj.items():
manip_name = {"l":"leftarm", "r":"rightarm"}[lr]
old_joint_traj_rs = mu.interp2d(timesteps_rs, np.arange(len(old_joint_traj)), old_joint_traj)
ee_link_name = "%s_gripper_tool_frame"%lr
new_ee_traj = lr2eetraj[lr][i_start:i_end+1]
new_ee_traj_rs = resampling.interp_hmats(timesteps_rs, np.arange(len(new_ee_traj)), new_ee_traj)
print "planning trajectory following"
new_joint_traj, pose_errs = planning.plan_follow_traj(sim_env.robot, manip_name,
sim_env.robot.GetLink(ee_link_name), new_ee_traj_rs,old_joint_traj_rs)
prev_vals[lr] = new_joint_traj[-1]
part_name = {"l":"larm", "r":"rarm"}[lr]
bodypart2traj[part_name] = new_joint_traj
################################
redprint("Executing joint trajectory for part %i using arms '%s'"%(i_miniseg, bodypart2traj.keys()))
full_traj = sim_util.getFullTraj(sim_env, bodypart2traj)
full_trajs.append(full_traj)
for lr in 'lr':
gripper_open = sim_util.binarize_gripper(seg_info["%s_gripper_joint"%lr][i_start], GRIPPER_ANGLE_THRESHOLD)
prev_gripper_open = sim_util.binarize_gripper(seg_info["%s_gripper_joint"%lr][i_start-1], GRIPPER_ANGLE_THRESHOLD) if i_start != 0 else False
if not sim_util.set_gripper_maybesim(sim_env, lr, gripper_open, prev_gripper_open):
redprint("Grab %s failed" % lr)
misgrasp = True
success = False
if not success: break
if len(full_traj[0]) > 0:
if not eval_util.traj_is_safe(sim_env, full_traj, COLLISION_DIST_THRESHOLD):
redprint("Trajectory not feasible")
feasible = False
if feasible or ignore_infeasibility:
success &= sim_util.sim_full_traj_maybesim(sim_env, full_traj, animate=animate, interactive=interactive)
else:
success = False
if not success: break
sim_env.sim.settle(animate=animate)
sim_env.sim.release_rope('l')
sim_env.sim.release_rope('r')
sim_util.reset_arms_to_side(sim_env)
if animate:
sim_env.viewer.Step()
return success, feasible, misgrasp, full_trajs
def simulate_demo_traj(sim_env,
new_xyz,
seg_info,
full_trajs,
ignore_infeasibility=True,
animate=False,
interactive=False):
sim_util.reset_arms_to_side(sim_env)
old_xyz = np.squeeze(seg_info["cloud_xyz"])
old_xyz = clouds.downsample(old_xyz, DS_SIZE)
new_xyz = clouds.downsample(new_xyz, DS_SIZE)
handles = []
if animate:
handles.append(sim_env.env.plot3(old_xyz, 5, (1, 0, 0)))
handles.append(sim_env.env.plot3(new_xyz, 5, (0, 0, 1)))
miniseg_starts, miniseg_ends, _ = sim_util.split_trajectory_by_gripper(
seg_info, thresh=GRIPPER_ANGLE_THRESHOLD)
success = True
feasible = True
misgrasp = False
print colorize.colorize("mini segments:",
"red"), miniseg_starts, miniseg_ends
for (i_miniseg, (i_start,
i_end)) in enumerate(zip(miniseg_starts, miniseg_ends)):
if i_miniseg >= len(full_trajs): break
full_traj = full_trajs[i_miniseg]
for lr in 'lr':
gripper_open = sim_util.binarize_gripper(
seg_info["%s_gripper_joint" % lr][i_start],
GRIPPER_ANGLE_THRESHOLD)
prev_gripper_open = sim_util.binarize_gripper(
seg_info["%s_gripper_joint" % lr][i_start - 1],
GRIPPER_ANGLE_THRESHOLD) if i_start != 0 else False
if not sim_util.set_gripper_maybesim(sim_env, lr, gripper_open,
prev_gripper_open):
redprint("Grab %s failed" % lr)
misgrasp = True
success = False
if not success: break
if len(full_traj[0]) > 0:
if not eval_util.traj_is_safe(sim_env, full_traj,
COLLISION_DIST_THRESHOLD):
redprint("Trajectory not feasible")
feasible = False
if feasible or ignore_infeasibility:
success &= sim_util.sim_full_traj_maybesim(
sim_env,
full_traj,
animate=animate,
interactive=interactive)
else:
success = False
if not success: break
sim_env.sim.settle(animate=animate)
sim_env.sim.release_rope('l')
sim_env.sim.release_rope('r')
sim_util.reset_arms_to_side(sim_env)
if animate:
sim_env.viewer.Step()
return success, feasible, misgrasp, full_trajs
def downsample_cloud(cloud):
d = DIM
cloud_xyz = cloud[:,:d]
cloud_xyz_downsamp = clouds.downsample(cloud_xyz, DS_SIZE)
return np.array(cloud_xyz_downsamp, dtype=np.float32)
def get_ds_cloud(sim_env, action):
return clouds.downsample(GlobalVars.actions[action]['cloud_xyz'], DS_SIZE)
def test_sc_features(args):
feature_fn, num_features, act_file = get_sc_feature_fn(args.actionfile)
for name, seg_info in act_file.iteritems():
print feature_fn([name, clouds.downsample(seg_info['cloud_xyz'], DS_SIZE)], name)
def get_ds_cloud(sim_env, action):
return clouds.downsample(GlobalVars.actions[action]['cloud_xyz'], DS_SIZE)
def compute_trans_traj(sim_env, new_xyz, seg_info, ignore_infeasibility=True, animate=False, interactive=False):
redprint("Generating end-effector trajectory")
old_xyz = np.squeeze(seg_info["cloud_xyz"])
old_xyz = clouds.downsample(old_xyz, DS_SIZE)
l1 = len(old_xyz)
new_xyz = clouds.downsample(new_xyz, DS_SIZE)
l2 = len(new_xyz)
print l1, l2
link_names = ["%s_gripper_tool_frame"%lr for lr in ('lr')]
hmat_list = [(lr, seg_info[ln]['hmat']) for lr, ln in zip('lr', link_names)]
if GlobalVars.gripper_weighting:
interest_pts = get_closing_pts(seg_info)
else:
interest_pts = None
lr2eetraj, _, old_xyz_warped, f = warp_hmats_tfm(old_xyz, new_xyz, hmat_list, interest_pts)
handles = []
if animate:
handles.extend(plotting_openrave.draw_grid(sim_env.env, f.transform_points, old_xyz.min(axis=0)-np.r_[0,0,.1], old_xyz.max(axis=0)+np.r_[0,0,.1], xres = .1, yres = .1, zres = .04))
handles.append(sim_env.env.plot3(old_xyz,5, (1,0,0))) # red: demonstration point cloud
handles.append(sim_env.env.plot3(new_xyz,5, (0,0,1))) # blue: rope nodes
handles.append(sim_env.env.plot3(old_xyz_warped,5, (0,1,0))) # green: warped point cloud from demonstration
mapped_pts = []
for i in range(len(old_xyz)):
mapped_pts.append(old_xyz[i])
mapped_pts.append(old_xyz_warped[i])
handles.append(sim_env.env.drawlinelist(np.array(mapped_pts), 1, [0.1,0.1,1]))
for lr in 'lr':
handles.append(sim_env.env.drawlinestrip(lr2eetraj[lr][:,:3,3], 2, (0,1,0,1)))
for k, hmats in hmat_list:
hmats_tfm = np.asarray([GlobalVars.init_tfm.dot(h) for h in hmats])
handles.append(sim_env.env.drawlinestrip(hmats_tfm[:,:3,3], 2, (1,0,0,1)))
miniseg_starts, miniseg_ends, lr_open = sim_util.split_trajectory_by_gripper(seg_info, thresh=GRIPPER_ANGLE_THRESHOLD)
success = True
feasible = True
misgrasp = False
print colorize.colorize("mini segments:", "red"), miniseg_starts, miniseg_ends
miniseg_trajs = []
prev_vals = {lr:None for lr in 'lr'}
for (i_miniseg, (i_start, i_end)) in enumerate(zip(miniseg_starts, miniseg_ends)):
################################
redprint("Generating joint trajectory for part %i"%(i_miniseg))
### adaptive resampling based on xyz in end_effector
end_trans_trajs = np.zeros([i_end+1-i_start, 6])
for lr in 'lr':
for i in xrange(i_start,i_end+1):
if lr == 'l':
end_trans_trajs[i-i_start, :3] = lr2eetraj[lr][i][:3,3]
else:
end_trans_trajs[i-i_start, 3:] = lr2eetraj[lr][i][:3,3]
if True:
adaptive_times, end_trans_trajs = resampling.adaptive_resample2(end_trans_trajs, 0.005)
else:
adaptive_times = range(len(end_trans_trajs))
miniseg_traj = {}
for lr in 'lr':
#ee_hmats = resampling.interp_hmats(np.arange(i_end+1-i_start), np.arange(i_end+1-i_start), lr2eetraj[lr][i_start:i_end+1])
ee_hmats = resampling.interp_hmats(adaptive_times, np.arange(i_end+1-i_start), lr2eetraj[lr][i_start:i_end+1])
# the interpolation above will then the velocity of the trajectory (since there are fewer waypoints). Resampling again to make sure
# the trajectory has the same number of waypoints as before.
ee_hmats = resampling.interp_hmats(np.arange(i_end+1-i_start), adaptive_times, ee_hmats)
# if arm_moved(ee_hmats, floating=True):
if True:
miniseg_traj[lr] = ee_hmats
safe_drop = {'l': True, 'r': True}
for lr in 'lr':
next_gripper_open = lr_open[lr][i_miniseg+1] if i_miniseg < len(miniseg_starts) - 1 else False
gripper_open = lr_open[lr][i_miniseg]
if next_gripper_open and not gripper_open:
tfm = miniseg_traj[lr][-1]
if tfm[2,3] > GlobalVars.table_height + 0.01:
safe_drop[lr] = False
#safe_drop = {'l': True, 'r': True}
if not (safe_drop['l'] and safe_drop['r']):
for lr in 'lr':
if not safe_drop[lr]:
tfm = miniseg_traj[lr][-1]
for i in range(1, 8):
safe_drop_tfm = tfm
safe_drop_tfm[2,3] = tfm[2,3] - i / 10. * (tfm[2,3] - GlobalVars.table_height - 0.01)
miniseg_traj[lr].append(safe_drop_tfm)
else:
for i in range(1, 8):
miniseg_traj[lr].append(miniseg_traj[lr][-1])
miniseg_trajs.append(miniseg_traj)
for lr in 'lr':
hmats = np.asarray(miniseg_traj[lr])
handles.append(sim_env.env.drawlinestrip(hmats[:,:3,3], 2, (0,0,1,1)))
redprint("Executing joint trajectory for part %i using arms '%s'"%(i_miniseg, miniseg_traj.keys()))
for lr in 'lr':
gripper_open = lr_open[lr][i_miniseg]
prev_gripper_open = lr_open[lr][i_miniseg-1] if i_miniseg != 0 else False
if not sim_util.set_gripper_maybesim(sim_env, lr, gripper_open, prev_gripper_open, floating=True):
redprint("Grab %s failed"%lr)
success = False
if not success: break
if len(miniseg_traj) > 0:
success &= sim_util.exec_traj_sim(sim_env, miniseg_traj, animate=animate)
if not success: break
sim_env.sim.settle(animate=animate)
sim_env.sim.release_rope('l')
sim_env.sim.release_rope('r')
if animate:
sim_env.viewer.Step()
return success, feasible, misgrasp, miniseg_trajs
