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)))