def add_kinematics_to_group(group, linknames, manipnames, jointnames, robot):
"do forward kinematics on those links"
if robot is None: robot = get_robot()
r2r = ros2rave.RosToRave(robot, group["joint_states"]["name"])
link2hmats = dict([(linkname, []) for linkname in linknames])
links = [robot.GetLink(linkname) for linkname in linknames]
rave_traj = []
rave_inds = r2r.rave_inds
for ros_vals in group["joint_states"]["position"]:
r2r.set_values(robot, ros_vals)
rave_vals = r2r.convert(ros_vals)
robot.SetDOFValues(rave_vals, rave_inds)
rave_traj.append(rave_vals)
for (linkname,link) in zip(linknames, links):
link2hmats[linkname].append(link.GetTransform())
for (linkname, hmats) in link2hmats.items():
group.create_group(linkname)
group[linkname]["hmat"] = np.array(hmats)
rave_traj = np.array(rave_traj)
rave_ind_list = list(rave_inds)
for manipname in manipnames:
arm_inds = robot.GetManipulator(manipname).GetArmIndices()
group[manipname] = rave_traj[:,[rave_ind_list.index(i) for i in arm_inds]]
for jointname in jointnames:
joint_ind = robot.GetJointIndex(jointname)
group[jointname] = rave_traj[:,rave_ind_list.index(joint_ind)]
def add_rgbd_to_hdf(video_dir, annotations, hdfroot, demo_name):
frame_stamps = [seg_info["look"] for seg_info in annotations]
rgb_imgs, depth_imgs = get_video_frames(video_dir, frame_stamps)
for (i_seg, seg_info) in enumerate(annotations):
group = hdfroot[demo_name + "_" + seg_info["name"]]
group["rgb"] = rgb_imgs[i_seg]
group["depth"] = depth_imgs[i_seg]
robot = get_robot()
r2r = ros2rave.RosToRave(robot, group["joint_states"]["name"])
r2r.set_values(robot, group["joint_states"]["position"][0])
T_w_h = robot.GetLink("head_plate_frame").GetTransform()
T_w_k = T_w_h.dot(berkeley_pr2.T_h_k)
group["T_w_k"] = T_w_k
def load_fake_data_segment(sim_env,
demofile,
fake_data_segment,
fake_data_transform,
set_robot_state=True,
floating=False):
fake_seg = demofile[fake_data_segment]
new_xyz = np.squeeze(fake_seg["cloud_xyz"])
hmat = openravepy.matrixFromAxisAngle(fake_data_transform[3:6])
hmat[:3, 3] = fake_data_transform[0:3]
new_xyz = new_xyz.dot(hmat[:3, :3].T) + hmat[:3, 3][None, :]
r2r = None
if not floating:
r2r = ros2rave.RosToRave(sim_env.robot,
asarray(fake_seg["joint_states"]["name"]))
if set_robot_state:
r2r.set_values(sim_env.robot,
asarray(fake_seg["joint_states"]["position"][0]))
return new_xyz, r2r
def main():
demofile = h5py.File(args.h5file, 'r')
trajoptpy.SetInteractive(args.interactive)
if args.log:
LOG_DIR = osp.join(osp.expanduser("~/Data/do_task_logs"),
datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
os.mkdir(LOG_DIR)
LOG_COUNT = 0
if args.execution:
rospy.init_node("exec_task", disable_signals=True)
Globals.pr2 = PR2.PR2()
Globals.env = Globals.pr2.env
Globals.robot = Globals.pr2.robot
else:
Globals.env = openravepy.Environment()
Globals.env.StopSimulation()
Globals.env.Load("robots/pr2-beta-static.zae")
Globals.robot = Globals.env.GetRobots()[0]
if not args.fake_data_segment:
grabber = cloudprocpy.CloudGrabber()
grabber.startRGBD()
Globals.viewer = trajoptpy.GetViewer(Globals.env)
#####################
while True:
redprint("Acquire point cloud")
if args.fake_data_segment:
fake_seg = demofile[args.fake_data_segment]
new_xyz = np.squeeze(fake_seg["cloud_xyz"])
hmat = openravepy.matrixFromAxisAngle(
args.fake_data_transform[3:6])
hmat[:3, 3] = args.fake_data_transform[0:3]
new_xyz = new_xyz.dot(hmat[:3, :3].T) + hmat[:3, 3][None, :]
r2r = ros2rave.RosToRave(Globals.robot,
asarray(fake_seg["joint_states"]["name"]))
r2r.set_values(Globals.robot,
asarray(fake_seg["joint_states"]["position"][0]))
else:
Globals.pr2.head.set_pan_tilt(0, 1.2)
Globals.pr2.rarm.goto_posture('side')
Globals.pr2.larm.goto_posture('side')
Globals.pr2.join_all()
time.sleep(.5)
Globals.pr2.update_rave()
rgb, depth = grabber.getRGBD()
T_w_k = berkeley_pr2.get_kinect_transform(Globals.robot)
new_xyz = cloud_proc_func(rgb, depth, T_w_k)
#grab_end(new_xyz)
if args.log:
LOG_COUNT += 1
import cv2
cv2.imwrite(osp.join(LOG_DIR, "rgb%i.png" % LOG_COUNT), rgb)
cv2.imwrite(osp.join(LOG_DIR, "depth%i.png" % LOG_COUNT), depth)
np.save(osp.join(LOG_DIR, "xyz%i.npy" % LOG_COUNT), new_xyz)
################################
redprint("Finding closest demonstration")
if args.select_manual:
seg_name = find_closest_manual(demofile, new_xyz)
else:
seg_name = find_closest_auto(demofile, new_xyz)
seg_info = demofile[seg_name]
redprint("closest demo: %s" % (seg_name))
if "done" in seg_name:
redprint("DONE!")
break
if args.log:
with open(osp.join(LOG_DIR, "neighbor%i.txt" % LOG_COUNT),
"w") as fh:
fh.write(seg_name)
################################
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)))
scaled_old_xyz, src_params = registration.unit_boxify(old_xyz)
scaled_new_xyz, targ_params = registration.unit_boxify(new_xyz)
f, _ = registration.tps_rpm_bij(scaled_old_xyz,
scaled_new_xyz,
plot_cb=tpsrpm_plot_cb,
plotting=5 if args.animation else 0,
rot_reg=np.r_[1e-4, 1e-4, 1e-1],
n_iter=50,
reg_init=10,
reg_final=.1)
f = registration.unscale_tps(f, src_params, targ_params)
handles.extend(
plotting_openrave.draw_grid(Globals.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))
link2eetraj = {}
for lr in 'lr':
link_name = "%s_gripper_tool_frame" % lr
old_ee_traj = asarray(seg_info[link_name]["hmat"])
new_ee_traj = f.transform_hmats(old_ee_traj)
link2eetraj[link_name] = new_ee_traj
handles.append(
Globals.env.drawlinestrip(old_ee_traj[:, :3, 3], 2,
(1, 0, 0, 1)))
handles.append(
Globals.env.drawlinestrip(new_ee_traj[:, :3, 3], 2,
(0, 1, 0, 1)))
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)):
if args.execution == "real": Globals.pr2.update_rave()
################################
redprint("Generating joint trajectory for segment %s, part %i" %
(seg_name, 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 = link2eetraj[ee_link_name][i_start:i_end + 1]
new_ee_traj_rs = resampling.interp_hmats(
timesteps_rs, np.arange(len(new_ee_traj)), new_ee_traj)
if args.execution: Globals.pr2.update_rave()
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)
part_name = {"l": "larm", "r": "rarm"}[lr]
bodypart2traj[part_name] = new_joint_traj
################################
redprint(
"Executing joint trajectory for segment %s, part %i using arms '%s'"
% (seg_name, i_miniseg, bodypart2traj.keys()))
for lr in 'lr':
success &= set_gripper_maybesim(
lr,
binarize_gripper(seg_info["%s_gripper_joint" %
lr][i_start]))
# Doesn't actually check if grab occurred, unfortunately
if not success: break
if len(bodypart2traj) > 0:
success &= exec_traj_maybesim(bodypart2traj)
if not success: break
redprint("Segment %s result: %s" % (seg_name, success))
if args.fake_data_segment: break
for d in depth_images
]
depth_images = depth_images / 1000.
joint_positions = npzfile["joint_positions"]
joint_names = npzfile["names"] if "names" in npzfile else npzfile["joint_names"]
env = openravepy.Environment()
env.StopSimulation()
env.Load("robots/pr2-beta-static.zae")
robot = env.GetRobots()[0]
viewer = trajoptpy.GetViewer(env)
viewer.Step()
sensor = sensorsimpy.CreateFakeKinect(env)
r2r = ros2rave.RosToRave(robot, joint_names)
rave_inds = r2r.rave_inds
rave_values = np.array(
[r2r.convert(ros_values) for ros_values in joint_positions])
robot.SetActiveDOFs(rave_inds)
robot.SetActiveDOFValues(rave_values[0])
n_imgs = len(depth_images)
assert len(joint_positions) == n_imgs
MAX_DIST = 1.5
def calc_Thk(xyz, rod):
T_h_k = openravepy.matrixFromAxisAngle(rod)
segnames = [args.seg] if args.seg else hdf.keys()
env = openravepy.Environment()
env.StopSimulation()
env.Load("robots/pr2-beta-static.zae")
robot = env.GetRobots()[0]
viewer = trajoptpy.GetViewer(env)
for segname in segnames:
seg_info = hdf[segname]
from rapprentice import berkeley_pr2
r2r = ros2rave.RosToRave(robot, seg_info["joint_states"]["name"])
rave_traj = [
r2r.convert(row)
for row in asarray(seg_info["joint_states"]["position"])
]
robot.SetActiveDOFs(r2r.rave_inds)
robot.SetActiveDOFValues(rave_traj[0])
handles = []
T_w_k = berkeley_pr2.get_kinect_transform(robot)
o = T_w_k[:3, 3]
x = T_w_k[:3, 0]
y = T_w_k[:3, 1]
z = T_w_k[:3, 2]
handles.append(env.drawarrow(o, o + .3 * x, .005, (1, 0, 0, 1)))
