Example #1
0
def plot_scatter(wpass, wfail, ppass, pfail, position=True, x_max=0.025):
    """
    Does a scatter plot of the warping costs (y-axis) and position/angle error (x-axis)
    """
    plot.clf()
    p1 = plot.scatter(ppass, wpass, color='black', marker='D', s=20)
    p2 = plot.scatter(pfail,
                      wfail,
                      color='black',
                      marker='o',
                      s=20,
                      facecolors='none')
    plot.legend([p1, p2], ["success", "failure"])

    pname = "position" if position else "orientation"

    if x_max > 0:
        plot.axis((0 if position else 0.3, x_max, 0, 1e-5))

    plot.ticklabel_format(style='sci', axis='x', scilimits=(0, 0))
    plot.ticklabel_format(style='sci', axis='y', scilimits=(0, 0))
    plot.axhline(y=8.5e-7)
    plot.xlabel('max %s error' % pname, fontsize=18)
    plot.ylabel('warping cost', fontsize=18)

    plot_fname = osp.join(save_dir, 'warp-%s.pdf' % pname)
    plot.savefig(plot_fname)
    print colorize("saved plot: %s" % plot_fname, "green", True)
Example #2
0
def plot_prob(pass_dat,
              fail_dat,
              cost_name='',
              label_order=-6,
              is_cost=True,
              nbins=20):
    """
    pass_dat, fail_dat : np.arrays of costs of successes and failures respectively
    bins the data into n-bins and plots a probability distribution:
    """
    cmin, cmax = min(np.min(pass_dat),
                     np.min(fail_dat)), max(np.max(pass_dat), np.max(fail_dat))
    crange = (cmin, cmax)
    sc, sbins = np.histogram(pass_dat, nbins, range=crange)
    fc, fbins = np.histogram(fail_dat, nbins, range=crange)

    assert np.allclose(sbins, fbins)
    tc = sc + fc
    sc = sc / (tc + EPS)

    ## trim the tail of zeros:
    nz, mz = np.max(np.nonzero(sc)), np.min(np.nonzero(sc))
    lim = min(nz + 1, len(sc) - 1)
    sc = sc[mz:lim + 1]
    tc = tc[mz:lim + 1]

    sd = np.sqrt(sc * (1 - sc) / (tc + EPS))

    sbins = sbins[mz:lim + 2]
    wwidth = sbins[1] - sbins[0]

    plot.clf()
    plot.bar(left=sbins[0:-1],
             height=sc,
             width=wwidth,
             yerr=sd,
             color=(0.8, 0.8, 0.8),
             ecolor=(0, 0, 0))

    xname = '%s%s %s' % ('' if is_cost else 'max ', cost_name,
                         'cost' if is_cost else 'error')
    plot.xlabel(xname, fontsize=18)

    plot.gca().xaxis.set_major_formatter(FixedOrderFormatter(label_order))
    ca1, ca2, ca3, ca4 = plot.axis()
    ca1 = max(0, sbins[0] - wwidth)
    plot.axis((ca1, ca2, ca3, 1.))
    plot.xticks(sbins[0:-1:4])
    plot.ylabel('P(success | %s)' % xname, fontsize=18)

    plot_fname = osp.join(save_dir, 'prob-%s.pdf' % cost_name)
    plot.savefig(plot_fname)
    print colorize("saved plot: %s" % plot_fname, "green", True)
Example #3
0
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
Example #4
0
    def save_results(self):       
        summary_dict = {}

        for run_id in self.results.keys():
            cost_fname = osp.join(self.run_dir, 'run%d-costs.txt' % run_id)
            run_info   = cPickle.load(open(cost_fname, 'r'))
            run_dict = {}
            run_dict['perturb'] = run_info['perturbations']
            run_dict['result']  = self.results[run_id]
            summary_dict[run_id] = run_dict

        res_fname = osp.join(self.out_dir, 'run-%s-results.cpickle'%self.run_names[self.rangenum])
        print colorize('Saving results to : '+ res_fname, 'blue', bold=True)
        cPickle.dump(summary_dict, open(res_fname, 'w'))
Example #5
0
def openloop():
    
    demofile = h5py.File(args.h5file, 'r')
    rospy.init_node("exec_task",disable_signals=True)
    
    if args.execution:
        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]
    
    while True:
        seg_name = select_segment(demofile)
        seg_info = demofile[seg_name]
     
        miniseg_starts, miniseg_ends = split_trajectory_by_gripper(seg_info)
        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 segment %s, part %i"%(seg_name, i_miniseg))

            bodypart2traj = {}
            arms_used = ""
            
            # Not sure about frequency - look into this
            for lr in 'lr':
                manip_name = {"l":"leftarm", "r":"rightarm"}[lr]
                _, joint_traj = resampling.adaptive_resample(asarray(seg_info[manip_name][i_start:i_end+1]), tol=0.01, max_change=.1)
                print "Shape of %s traj: "%lr, joint_traj.shape
                if arm_moved(joint_traj):
                    part_name = {"l":"larm", "r":"rarm"}[lr]
                    bodypart2traj[part_name] = joint_traj
                    arms_used += lr
        
            redprint("Executing open-loop trajectory for segment %s, part %i using arms '%s'"%(seg_name, i_miniseg, arms_used))
            
            for lr in 'lr':
                set_gripper_maybesim(lr, binarize_gripper(seg_info["%s_gripper_joint"%lr][i_start]))
               
            print len(bodypart2traj) 
            if len(bodypart2traj) > 0:
                exec_traj_maybesim(bodypart2traj, speed_factor=0.1)
Example #6
0
    def save_results(self):
        summary_dict = {}

        for run_id in self.results.keys():
            cost_fname = osp.join(self.run_dir, 'run%d-costs.txt' % run_id)
            run_info = cPickle.load(open(cost_fname, 'r'))
            run_dict = {}
            run_dict['perturb'] = run_info['perturbations']
            run_dict['result'] = self.results[run_id]
            summary_dict[run_id] = run_dict

        res_fname = osp.join(
            self.out_dir,
            'run-%s-results.cpickle' % self.run_names[self.rangenum])
        print colorize('Saving results to : ' + res_fname, 'blue', bold=True)
        cPickle.dump(summary_dict, open(res_fname, 'w'))
Example #7
0
def plot_prob(pass_dat, fail_dat, cost_name='', label_order=-6, is_cost=True, nbins=20):
    """
    pass_dat, fail_dat : np.arrays of costs of successes and failures respectively
    bins the data into n-bins and plots a probability distribution:
    """
    cmin, cmax = min(np.min(pass_dat), np.min(fail_dat)), max(np.max(pass_dat), np.max(fail_dat))
    crange = (cmin, cmax)
    sc, sbins = np.histogram(pass_dat, nbins, range=crange)
    fc, fbins = np.histogram(fail_dat, nbins, range=crange)

    assert np.allclose(sbins, fbins)
    tc = sc+fc
    sc = sc/(tc + EPS)

    ## trim the tail of zeros:
    nz, mz = np.max(np.nonzero(sc)), np.min(np.nonzero(sc))
    lim = min(nz+1, len(sc)-1)
    sc = sc[mz:lim+1]
    tc = tc[mz:lim+1]
    
    sd = np.sqrt(sc*(1-sc)/(tc+EPS))

    sbins  = sbins[mz:lim+2]
    wwidth = sbins[1]-sbins[0]

    plot.clf()
    plot.bar(left=sbins[0:-1], height=sc, width=wwidth, yerr=sd, color=(0.8,0.8,0.8), ecolor=(0,0,0))

    xname = '%s%s %s' % ('' if is_cost else 'max ', cost_name, 'cost' if is_cost else 'error')
    plot.xlabel(xname, fontsize=18)

    plot.gca().xaxis.set_major_formatter(FixedOrderFormatter(label_order))
    ca1,ca2,ca3,ca4 =  plot.axis()
    ca1 = max(0,sbins[0]-wwidth)
    plot.axis((ca1,ca2,ca3,1.))
    plot.xticks(sbins[0:-1:4])
    plot.ylabel('P(success | %s)'%xname, fontsize=18)

    plot_fname = osp.join(save_dir, 'prob-%s.pdf'%cost_name)
    plot.savefig(plot_fname)
    print colorize("saved plot: %s"%plot_fname, "green", True)
def run_tests_on_cloud(cloud_params, do_local=False):
    """
    make sure that task_fname and action_fname are available on the volume in the cloud.
    
    do_local : if true, the 'map' is done locally.
    """
    
    ntests     = len(cloud_params.cmd_params)
    batch_size = int(math.ceil(ntests/(cloud_params.num_batches+0.0)))

    batch_edges = batch_size*np.array(xrange(cloud_params.num_batches))[cloud_params.start_batch_num : cloud_params.end_batch_num]

    all_succ = []
    for i in xrange(len(batch_edges)):
        if i==len(batch_edges)-1:
            cmds = cloud_params.cmd_params[batch_edges[i]:]
        else:
            cmds = cloud_params.cmd_params[batch_edges[i]:min(batch_edges[i+1], ntests)]
        print colorize("calling on cloud : batch [%d/%d] "%(i, len(batch_edges)), "yellow", True)
        try:
            if not do_local:
                jids = cloud.map(run_example, cmds, _vol=cloud_params.vol, _env=cloud_params.env, _type=cloud_params.core_type)
                succ = cloud.result(jids)
                print colorize("\t got results for batch %d/%d "%(i, len(batch_edges)), "green", True)
            else:
                succ = map(run_example, cmds)
            all_succ += succ
        except Exception as e:
            print "Found exception %s. Not saving data for this demo."%e

    with open(cloud_params.results_fname, 'w') as f:
        print colorize("\t\t saved results in : %s"%cloud_params.results_fname, "green")
        cp.dump(all_succ, f)
def test_bootrun(bootrun_name='boot_1', do_nn=False, tree_sizes=[30,60,90,120], test_fname="eval_set.h5"):
    """
    @ res_dir       : the directory where the results from the test runs will be saved.
                      the saved results will be like: 
                         res_dir/<bootrun_name>_<tree_size>_res.cp
    @ bootrun_name  : the directory holding the actions from the bootstrap runs.
    @ do_nn         : the action file here is just the nearest neighbor file [use this for baseline]
    @ tree_sizes    : the sizes of bootstrap trees to run tests on.
    """
    cloud_bootstrapping_dir = "/home/picloud/sandbox/bootstrapping"
    data_dir           = osp.join(os.getenv("BOOTSTRAPPING_DIR"), "data")
    local_task_fname   = osp.join(data_dir, test_fname)
    task_fname         = osp.join(cloud_bootstrapping_dir, "data/%s"%test_fname) 

    if not do_nn:
        result_fnames      = [osp.join(data_dir, "test_results", "%s_%d_result.cp"%(bootrun_name, s)) for s in tree_sizes]
        test_action_fnames = [osp.join(cloud_bootstrapping_dir, "data", bootrun_name, 'test_bootstrapping_%d.h5'%s) for s in tree_sizes]
    else:
        test_basename      = osp.splitext(osp.basename(test_fname))[0]
        result_fnames      = [osp.join(data_dir, "test_results", "%s_%s_nn_result.cp"%(bootrun_name, test_basename))]
        test_action_fnames = [osp.join(cloud_bootstrapping_dir, "data", bootrun_name, 'test_bootstrapping_orig.h5')]


    for i in xrange(len(test_action_fnames)):
        cmd_params      = create_test_params(local_task_fname, task_fname, test_action_fnames[i])
        print colorize(" SUBMITTING %d jobs to run on the cloud"%len(cmd_params), "red", True)
        cloud_params    =  CloudParams()
        cloud_params.cmd_params      = cmd_params
        cloud_params.num_batches     = 1
        cloud_params.start_batch_num = 0
        cloud_params.end_batch_num   = 1 ## exclusive
        cloud_params.results_fname   = result_fnames[i]
        cloud_params.env             = 'RSS3'
        cloud_params.vol             = 'iros_dat'
        cloud_params.core_type       = 'f2'

        print colorize("running tests for file: %s"%test_action_fnames[i], "magenta", True)
        run_tests_on_cloud(cloud_params, False)
Example #10
0
def plot_scatter(wpass, wfail, ppass, pfail, position=True, x_max=0.025):
    """
    Does a scatter plot of the warping costs (y-axis) and position/angle error (x-axis)
    """
    plot.clf()
    p1 = plot.scatter(ppass, wpass, color='black', marker='D',s=20)
    p2 = plot.scatter(pfail, wfail, color='black', marker='o',s=20, facecolors='none')
    plot.legend([p1, p2], ["success", "failure"])

    pname = "position" if position else "orientation"

    if x_max > 0:
        plot.axis((0 if position else 0.3, x_max, 0, 1e-5))

    plot.ticklabel_format(style='sci', axis='x', scilimits=(0,0))
    plot.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
    plot.axhline(y=8.5e-7)
    plot.xlabel('max %s error'%pname, fontsize=18)
    plot.ylabel('warping cost', fontsize=18)

    plot_fname = osp.join(save_dir, 'warp-%s.pdf'%pname)
    plot.savefig(plot_fname)
    print colorize("saved plot: %s"%plot_fname, "green", True)
Example #11
0
def test_all(stop=False):
    nPass,nFail = 0,0
    for (name,func) in TEST_FUNCS.items():
        print colorize("function: %s"%name,"green")
        try:
            t_start = time()
            func()
            t_elapsed = time() - t_start
            print colorize("PASSED (%.3f sec)"%t_elapsed,"blue")
            nPass += 1
        except Exception:    
            traceback.print_exc(file=sys.stdout)
            if stop: raise
            print colorize("FAILED","red")
            nFail += 1
            
            
    print "%i passed, %i failed"%(nPass,nFail)
def yellowprint(msg):
    """
    Print the message to the console in red, bold font.
    """
    print colorize.colorize(msg, "yellow", bold=True)
Example #13
0
def main():

    demofile = h5py.File(args.h5file, 'r')
    
    trajoptpy.SetInteractive(args.interactive)

    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:
            new_xyz = np.squeeze(demofile[args.fake_data_segment]["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,:]
            
        else:
            
            Globals.pr2.rarm.goto_posture('side')
            Globals.pr2.larm.goto_posture('side')            
            Globals.pr2.join_all()
            
            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)

    
        ################################    
        redprint("Finding closest demonstration")
        if args.fake_data_segment:
            seg_name = args.fake_data_segment
        else:
            seg_name = find_closest(demofile, new_xyz)
        
        seg_info = demofile[seg_name]
        # redprint("using demo %s, description: %s"%(seg_name, seg_info["description"]))
    
        ################################

        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,1)))
        handles.append(Globals.env.plot3(new_xyz,5, (0,0,1,1)))


        f = registration.tps_rpm(old_xyz, new_xyz, plot_cb = tpsrpm_plot_cb,plotting=1)            
        
        #f = registration.ThinPlateSpline() XXX XXX
        
        handles.extend(plotting_openrave.draw_grid(Globals.env, f.transform_points, old_xyz.min(axis=0), old_xyz.max(axis=0), 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)))
            
            
            
        # TODO plot
        # plot_warping_and_trajectories(f, old_xyz, new_xyz, old_ee_traj, new_ee_traj)
    
        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))

            bodypart2traj = {}
            
            arms_used = ""
        
            for lr in 'lr':
                manip_name = {"l":"leftarm", "r":"rightarm"}[lr]
                old_joint_traj = asarray(seg_info[manip_name][i_start:i_end+1])
                if arm_moved(old_joint_traj):          
                    ee_link_name = "%s_gripper_tool_frame"%lr
                    new_ee_traj = link2eetraj[ee_link_name]
                    if args.execution: Globals.pr2.update_rave()                    
                    new_joint_traj = plan_follow_traj(Globals.robot, manip_name,
                     Globals.robot.GetLink(ee_link_name), new_ee_traj[i_start:i_end+1], 
                     old_joint_traj)
                    # (robot, manip_name, ee_link, new_hmats, old_traj):
                    part_name = {"l":"larm", "r":"rarm"}[lr]
                    bodypart2traj[part_name] = new_joint_traj
                    arms_used += lr

        

            ################################    
            redprint("Executing joint trajectory for segment %s, part %i using arms '%s'"%(seg_name, i_miniseg, arms_used))

            for lr in 'lr':
                set_gripper_maybesim(lr, binarize_gripper(seg_info["%s_gripper_joint"%lr][i_start]))
            #trajoptpy.GetViewer(Globals.env).Idle()
        
            if len(bodypart2traj) > 0:
                exec_traj_maybesim(bodypart2traj)
        
            # TODO measure failure condtions

            if not success:
                break
            
        redprint("Segment %s result: %s"%(seg_name, success))
    
    
        if args.fake_data_segment: break
Example #14
0
def main():

    demofile = h5py.File(args.h5file, 'r')
    
    trajoptpy.SetInteractive(args.interactive)
    rospy.init_node("exec_task",disable_signals=True)
    
    global filename
    if args.logging:
        name = raw_input ('Enter identifier of this experiment (without spaces): ')
        name = name + '.yaml'
        import os.path as osp
        filename = osp.join('/home/sibi/sandbox/rapprentice/experiments/', name)
        if not osp.exists(filename):
            with open(filename,'w') as fh:
                fh.write("experiment: %s\ninfo: \n"%name)
    
    #thc.start()
    
    if args.execution:
        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]
        
    import trajoptpy.make_kinbodies as mk
    #Globals.env.Load("/home/sibi/sandbox/rapprentice/objects/table.xml")
    addBoxToRave(Globals.env, "table")
    Globals.sponge = addBoxToRave(Globals.env, "sponge") # Not sure about this
    Globals.needle_tip = mk.create_spheres(Globals.env, [(0,0,0)], radii=0.02, name="needle_tip")
    Globals.demo_env = Globals.env.CloneSelf(1)
    Globals.demo_env.StopSimulation()
    Globals.demo_robot = Globals.demo_env.GetRobot("pr2")
    Globals.demo_needle_tip = Globals.demo_env.GetKinBody("needle_tip")
    
    if not args.fake_data_segment:
        grabber = cloudprocpy.CloudGrabber()
        grabber.startRGBD()
        thc.grabber = grabber

    #Globals.viewer = trajoptpy.GetViewer(Globals.env)
    #####################
    
    #Globals.env.SetViewer('qtcoin')
    #threadClass().start()

    while True:
        redprint("Acquire point cloud")
         
        ################################
        redprint("Finding closest demonstration")
        if args.fake_data_segment:
            seg_name = args.fake_data_segment
        else:
            seg_name = find_closest(demofile)
        
        seg_info = demofile[seg_name]
        # redprint("using demo %s, description: %s"%(seg_name, seg_info["description"]))
    
        ################################
        
        redprint("Generating end-effector trajectory")    

        handles = []
        
        if seg_info.get('ar_marker_poses') is None:
            use_markers = False
        else:
            use_markers = True

        if use_markers:
            old_marker_poses_str = seg_info['ar_marker_poses']
            old_marker_poses = {}
            for i in old_marker_poses_str:
                old_marker_poses[int(i)] = old_marker_poses_str[i]
        
        
        # TODO: Have a check for only transformations -> rigid transformations
        if args.fake_data_segment:
            new_keypoints = demofile[args.fake_data_segment]["key_points"]
            if use_markers:
                new_marker_poses_str = demofile[args.fake_data_segment]["ar_marker_poses"]
                new_marker_poses = {}
                for i in new_marker_poses_str:
                    new_marker_poses[int(i)] = new_marker_poses_str[i]
            if seg_info['key_points'].keys().sort() != new_keypoints.keys().sort():
                print "Keypoints don't match."
                exit(1)
        else:
            if args.execution: Globals.pr2.update_rave()
            T_w_k = berkeley_pr2.get_kinect_transform(Globals.robot)
            new_keypoints, new_marker_poses = fk.get_keypoints_execution(grabber, seg_info['key_points'].keys(), T_w_k, use_markers)


        print "Warping the points for the new situation."

        if "l_grab" in seg_info['extra_information'] or "r_grab" in seg_info['extra_information']:
            use_ntt_kp = False
        else:
            use_ntt_kp = True

        # Points from keypoints
        old_xyz, rigid, kp_mapping = fk.key_points_to_points(seg_info['key_points'], use_ntt_kp)
        new_xyz, _, _ = fk.key_points_to_points(new_keypoints, use_ntt_kp)

        # Points from markers
        if use_markers:
            old_common_poses, new_common_poses = find_common_marker_poses(old_marker_poses, new_marker_poses, new_keypoints.keys())
            old_m_xyz, rigid_m, _ = fk.key_points_to_points(old_common_poses, False)
            new_m_xyz, _, _ = fk.key_points_to_points(new_common_poses, False)

            if len(old_m_xyz) > 0 and rigid: rigid = False
            elif rigid_m and not old_xyz.any(): rigid = True

            # concatenate points
            if old_xyz.any() and old_m_xyz.any():
                old_xyz = np.r_[old_xyz, old_m_xyz]
            elif old_m_xyz.any():
                old_xyz = old_m_xyz
        
            if new_xyz.any() and new_m_xyz.any():
                new_xyz = np.r_[new_xyz, new_m_xyz]
            elif new_m_xyz.any():
                new_xyz = new_m_xyz
                
        if new_xyz.any() and new_xyz.shape != (4,4):
            #handles.append(Globals.env.plot3(old_xyz,5, (1,0,0,1)))
            #handles.append(Globals.env.plot3(old_xyz,5, np.array([(1,0,0) for _ in xrange(old_xyz.shape[0])])))
            handles.append(Globals.env.plot3(new_xyz,5, np.array([(0,0,1) for _ in xrange(old_xyz.shape[0])])))

        print 'Old points:', old_xyz
        print 'New points:', new_xyz
        
        #if not yes_or_no.yes_or_no("Use identity?"):
        if rigid:
            f = registration.ThinPlateSpline()
            rel_tfm = new_xyz.dot(np.linalg.inv(old_xyz))
            f.init_rigid_tfm(rel_tfm)
            bend_c = 0
            rot_c = 0
            scale_c = 0
            max_error = None
        elif len(new_xyz) > 0:
            #f.fit(demopoints_m3, newpoints_m3, 10,10)
            # TODO - check if this regularization on bending is correct
            #if "right_hole_normal" in new_keypoints or "left_hole_normal" in new_keypoints:
            #    bend_c = 0.1
            #    rot_c = [1e-5,1e-5,0.1]
            #wt = 5
            #wt_n = np.ones(len(old_xyz))
            #if kp_mapping.get("right_hole_normal"):
            #    wt_n[kp_mapping["right_hole_normal"][0]] = wt
            #if kp_mapping.get("left_hole_normal"):
            #    wt_n[kp_mapping["left_hole_normal"][0]] = wt
            #else:
            #    bend_c = 0.1
            #    rot_c = 1e-5#[0,0,1e-5]
            #    wt_n = None
            # f = registration.fit_ThinPlateSpline(old_xyz, new_xyz, bend_coef = bend_c,rot_coef = rot_c)
            bend_c = 0.05
            rot_c = [1e-3, 1e-3, 1e-3]
            scale_c = 0.1
            f = registration.fit_ThinPlateSpline_RotReg(old_xyz, new_xyz, bend_c, rot_c, scale_c)
            np.set_printoptions(precision=3)
            max_error = np.max(np.abs(f.transform_points(old_xyz) - new_xyz))
            print "nonlinear part", f.w_ng
            print "affine part", f.lin_ag
            print "translation part", f.trans_g
            print "residual", f.transform_points(old_xyz) - new_xyz
            print "max error ", max_error
        else:
            f = registration.ThinPlateSpline()
            bend_c = 0
            rot_c = 0
            scale_c = 0
            max_error = None
                        
#         else:
#             f = registration.ThinPlateSpline()            
#             bend_c = 0
#             rot_c = 0
#         
        if old_xyz.any() and old_xyz.shape != (4,4):
            tfm_xyz = f.transform_points(old_xyz)
            handles.append(Globals.env.plot3(tfm_xyz,5, np.array([(0,1,0) for _ in xrange(tfm_xyz.shape[0])])))
        
        #f = registration.ThinPlateSpline() XXX XXX
        
        if new_xyz.any() and new_xyz.shape != (4,4):
            handles.extend(plotting_openrave.draw_grid(Globals.env, f.transform_points, old_xyz.min(axis=0), old_xyz.max(axis=0), xres = .1, yres = .1, zres = .04))

        if args.ask:
#             if new_xyz.any() and new_xyz.shape != (4,4):
#                 import visualize
#                 visualize.plot_mesh_points(f.transform_points, old_xyz, new_xyz)
                #lines = plotting_openrave.gen_grid(f.transform_points, np.array([0,-1,0]), np.array([1,1,1]))
                #plotting_openrave.plot_lines(lines)
            if not yes_or_no.yes_or_no("Continue?"):
                continue
        
        miniseg_starts, miniseg_ends = split_trajectory_by_gripper(seg_info)
        success = True
        print colorize.colorize("mini segments:", "red"), miniseg_starts, miniseg_ends

        # Assuming only lgrab or rgrab
        use_needle = None
        for lr in 'lr':      
            if "%s_grab"%lr in seg_info['extra_information']:
                if "needle_tip_transform" in seg_info['key_points']:
                    demo_ntfm = np.array(seg_info['key_points']['needle_tip_transform'])
                    ntfm = np.array(new_keypoints['needle_tip_transform'])
                    use_needle = lr
                elif Globals.rel_ntfm is not None:
                    T_w_g = Globals.robot.GetLink("%s_gripper_tool_frame"%lr).GetTransform()
                    T_g_n = Globals.rel_ntfm
                    ntfm = T_w_g.dot(T_g_n)
                    T_w_g_demo = Globals.demo_robot.GetLink("%s_gripper_tool_frame"%lr).GetTransform()
                    T_g_n_demo = Globals.demo_rel_ntfm
                    demo_ntfm = T_w_g_demo.dot(T_g_n_demo)
                    use_needle = lr
 
        if use_needle is not None:
            setup_needle_grabs(use_needle, seg_info["joint_states"]["name"], seg_info["joint_states"]["look_position"], demo_ntfm, ntfm)
        

        if args.logging:
            with open(filename,'a') as fh:
                fh.write("- seg_name: %s\n"%seg_name)
                fh.write("  tps_info: \n")
                if max_error is not None:
                    res_cost, bend_cost, tot_cost = f.fitting_cost(new_xyz, bend_c)
                    fh.write("    res_cost: %f\n    bend_cost: %f\n    tot_cost: %f\n"%(res_cost, bend_cost, tot_cost))
                    fh.write("    max_tps_error: %f\n"%max_error)
                else:
                    fh.write("    res_cost: -1\n    bend_cost: -1\n    tot_cost: -1\n    max_tps_error: -1\n")
                
                fh.write("  trajopt_info: \n")


        for (i_miniseg, (i_start, i_end)) in enumerate(zip(miniseg_starts, miniseg_ends)):
            
            if args.execution: Globals.pr2.update_rave()

            # Changing the trajectory according to the end-effector
            full_traj = np.c_[seg_info["leftarm"][i_start:i_end+1], seg_info["rightarm"][i_start:i_end+1]]
            full_traj = mu.remove_duplicate_rows(full_traj)
            _, ds_traj =  resampling.adaptive_resample(full_traj, tol=.01, max_change=.1) # about 2.5 degrees, 10 degrees

            Globals.robot.SetActiveDOFs(np.r_[Globals.robot.GetManipulator("leftarm").GetArmIndices(), Globals.robot.GetManipulator("rightarm").GetArmIndices()])
            Globals.demo_robot.SetActiveDOFs(np.r_[Globals.robot.GetManipulator("leftarm").GetArmIndices(), Globals.robot.GetManipulator("rightarm").GetArmIndices()])        
            Globals.demo_robot.SetDOFValues(Globals.robot.GetDOFValues())
        

            ################################    
            redprint("Generating joint trajectory for segment %s, part %i"%(seg_name, i_miniseg))

            bodypart2traj = {}
            arms_used = ""
            
            if args.logging:
                with open(filename, 'a') as fh:
                    fh.write("    mini_seg%i: \n"%i_miniseg)
        
            for lr in 'lr':
                
#                 if "%s_grab"%lr in seg_info['extra_information']:
#                     if "needle_tip_transform" in seg_info['key_points']:
#                         demo_ntfm = np.array(seg_info['key_points']['needle_tip_transform'])
#                         ntfm = np.array(new_keypoints['needle_tip_transform'])
#                         use_needle = True
#                     elif Globals.rel_ntfm is not None:
#                         T_w_g = Globals.robot.GetLink("%s_gripper_tool_frame"%lr).GetTransform()
#                         T_g_n = Globals.rel_ntfm
#                         ntfm = T_w_g.dot(T_g_n)
#                         T_w_g_demo = Globals.demo_robot.GetLink("%s_gripper_tool_frame"%lr).GetTransform()
#                         T_g_n_demo = Globals.demo_rel_ntfm
#                         demo_ntfm = T_w_g_demo.dot(T_g_n_demo)
#                         use_needle = True
#                     else:
#                         use_needle = False
#                 else:
#                     use_needle = False
                
                if use_needle == lr:
                    Globals.sponge.Enable(False)
                    arm_traj = {'l':ds_traj[:,:7], 'r':ds_traj[:,7:]}[lr]
                    old_ee_traj = setup_needle_traj (lr, arm_traj)
                    link = Globals.needle_tip.GetLinks()[0]  
                    redprint("Using needle for trajectory execution...")

                else:
                    arm = {"l":"leftarm", "r":"rightarm"}[lr]
                    Globals.demo_robot.SetActiveDOFs(Globals.robot.GetManipulator(arm).GetArmIndices())
                    if seg_name == "8_knot_tie_next0_seg00":
                        Globals.sponge.Enable(False)
                    else:
                        Globals.sponge.Enable(True)
                    ee_link_name = "%s_gripper_tool_frame"%lr
                    link = Globals.robot.GetLink(ee_link_name)
                    #old_ee_traj = np.asarray(seg_info[ee_link_name]["hmat"])
                    demo_link = Globals.demo_robot.GetLink(ee_link_name)
                    old_ee_traj = []
                    arm_traj = {'l':ds_traj[:,:7], 'r':ds_traj[:,7:]}[lr]
                    for row in arm_traj:
                        Globals.demo_robot.SetActiveDOFValues(row)
                        old_ee_traj.append(demo_link.GetTransform())
                        

############ CAN YOU PLOT THIS OLD_EE_TRAJ?        
                old_ee_traj = np.asarray(old_ee_traj)
                old_joint_traj = {'l':ds_traj[:,:7], 'r':ds_traj[:,7:]}[lr]

                if lr == 'l':#arm_moved(old_joint_traj):
                    
                    manip_name = {"l":"leftarm", "r":"rightarm"}[lr]
                    new_ee_traj = f.transform_hmats(old_ee_traj)
                    #new_ee_traj = old_ee_traj
                    
#################### CAN YOU PLOT THIS 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)))
                    #old_poses = [conv.hmat_to_pose(hmat) for hmat in old_ee_traj]
                    #print new_ee_traj
                    #new_poses = [conv.hmat_to_pose(hmat) for hmat in new_ee_traj]
                    update_markers(new_ee_traj, old_ee_traj)
                    #thc.run()
                    

                    if args.execution: Globals.pr2.update_rave()
                    new_joint_traj, costs, cnts, pos_err = plan_follow_traj(Globals.robot, manip_name,
                                                      link, new_ee_traj, old_joint_traj, True)
                    new_joint_traj = unwrap_joint_traj (lr, new_joint_traj)
                    # (robot, manip_name, ee_link, new_hmats, old_traj):
                    part_name = {"l":"larm", "r":"rarm"}[lr]
                    bodypart2traj[part_name] = new_joint_traj
                    
                    arms_used += lr
                    
                    if args.logging:
                        with open(filename, 'a') as fh:
                            fh.write("    - %s: \n"%lr)
                            fh.write("        costs: \n")
                            tot_cost = 0
                            for c_name, c_val in costs:
                                fh.write("        - %s: %f\n"%(c_name, c_val))
                                tot_cost += c_val
                            fh.write("        - tot_cost: %f\n"%tot_cost)
                            fh.write("        constraints: \n")
                            for cnt in cnts:
                                fh.write("        - %s\n"%str(cnt))
                            fh.write("        max_pos_error: %f\n"%pos_err)
                        

            ################################    
            redprint("Executing joint trajectory for segment %s, part %i using arms '%s'"%(seg_name, i_miniseg, arms_used))

            for lr in 'lr':
                set_gripper_maybesim(lr, binarize_gripper(seg_info["%s_gripper_joint"%lr][i_start]))
                #trajoptpy.GetViewer(Globals.env).Idle()

            if len(bodypart2traj) > 0:
                exec_traj_maybesim(bodypart2traj, speed_factor=1)

            # TODO measure failure condtions
            if not success:
                break
            

        Globals.robot.ReleaseAllGrabbed()
        Globals.demo_robot.ReleaseAllGrabbed()
    
        redprint("Segment %s result: %s"%(seg_name, success))
    
        if args.fake_data_segment: break
def redprint(msg):
    print colorize.colorize(msg, "red", bold=True)
def main():
    import IPython
    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]))
            #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(2)
        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(2)

            
            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)
            print "got new xyz"
            redprint(new_xyz)
            #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) 
        ################################



        ### Old end effector forces at r_gripper_tool_frame (eliminating the torques for now)
        old_eefs = seg_info['end_effector_forces'][:,0:3,:]


        redprint("Generating end-effector trajectory")    

        handles = []
        old_xyz = np.squeeze(seg_info["cloud_xyz"])
       


        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)
        
        old_xyz_transformed = f.transform_points(old_xyz)

        #handles.append(Globals.env.plot3(old_xyz_transformed,5, np.array([(0,0,1,1) for i in old_xyz_transformed])))

        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 'r':
            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
            #print old_ee_traj
            old_ee_pos = old_ee_traj[:,0:3,3]
            #print old_ee_pos

            # End effector forces as oppossed to end effector trajectories
            dfdxs = f.compute_jacobian(old_ee_pos)
            new_eefs = []
            for i in xrange(len(old_eefs)):
                dfdx = np.matrix(dfdxs[i])
                old_eef = np.matrix(old_eefs[i])
                new_eefs.append(dfdx * old_eef)
            old_eefs = asarray(old_eefs)[:,:,0]
            new_eefs = asarray(new_eefs)[:,:,0]

            force_data = {}
            force_data['old_eefs'] = old_eefs
            force_data['new_eefs'] = new_eefs
            force_file = open("trial.pickle", 'wa')
            pickle.dump(force_data, force_file)
            force_file.close()
            new_ee_traj_x = new_ee_traj
            
        
        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 'r':
                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): NOT SURE WHY BUT THIS IS RETURNING FALSE
                lr2oldtraj[lr] = old_joint_traj

            if args.visualize:
                    r2r = ros2rave.RosToRave(Globals.robot, asarray(seg_info["joint_states"]["name"]))
                    r2r.set_values(Globals.robot, asarray(seg_info["joint_states"]["position"][0]))
                    for i in range(0, lr2oldtraj['r'].shape[0], 10):
                        handles = []
                        handles.append(Globals.env.plot3(new_xyz,5,np.array([(0,1,0,1) for x in new_xyz])))
                        handles.append(Globals.env.plot3(old_xyz,5,np.array([(1,0,0,1) for x in old_xyz])))
                        handles.append(Globals.env.drawlinestrip(old_ee_traj[:,:3,3], 2, (1,0,0,1)))
                        # Setting robot arm to joint trajectory
                        r_vals = lr2oldtraj['r'][i,:]
                        Globals.robot.SetDOFValues(r_vals, Globals.robot.GetManipulator('rightarm').GetArmIndices())
                        # Plotting forces from r_gripper_tool_frame
                        hmats = Globals.robot.GetLinkTransformations()
                        trans, rot = conv.hmat_to_trans_rot(hmats[-3])
                        f_start = np.array([0,0,0]) + trans
                        f_end = np.array(old_eefs[i])/100 + trans
                        handles.append(Globals.env.drawlinestrip(np.array([f_start, f_end]), 10, (1,0,0,1)))
                        
                        redprint(i)
                        Globals.viewer.Step()
            if len(lr2oldtraj) > 0:
                old_total_traj = np.concatenate(lr2oldtraj.values(), 1)
                JOINT_LENGTH_PER_STEP = .1
                # FIRST RESAMPLING HAPPENS HERE: JOINT_LENGTH
                _, timesteps_rs = unif_resample(old_total_traj, JOINT_LENGTH_PER_STEP) # Timesteps only, can use to inter eefs for first time
            ####
            new_eefs_segment = asarray(new_eefs[i_start:i_end+1,:]) # Extract miniseg, and re-sample

            if args.no_ds:
                new_eefs_segment_rs = new_eefs_segment
            else:
                new_eefs_segment_rs = mu.interp2d(timesteps_rs, np.arange(len(new_eefs_segment)), new_eefs_segment)
        
            ### Generate fullbody traj
            bodypart2traj = {}
            for (lr,old_joint_traj) in lr2oldtraj.items():

                manip_name = {"l":"leftarm", "r":"rightarm"}[lr]
                ee_link_name = "%s_gripper_tool_frame"%lr
                new_ee_traj = link2eetraj[ee_link_name][i_start:i_end+1]
                if args.no_ds:
                    old_joint_traj_rs = old_joint_traj
                    new_ee_traj_rs = new_ee_traj
                    ds_inds = np.arange(0, len(new_ee_traj_rs), args.trajopt_ds)
                    new_ee_traj_rs = new_ee_traj_rs[ds_inds]
                    
                    old_joint_traj_rs = old_joint_traj_rs[ds_inds]
                    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)
                    new_joint_traj = mu.interp2d(np.arange(len(old_joint_traj)), np.arange(0, len(new_joint_traj) * args.trajopt_ds, args.trajopt_ds), new_joint_traj)
                else:
                    old_joint_traj_rs = mu.interp2d(timesteps_rs, np.arange(len(old_joint_traj)), old_joint_traj)
                    new_ee_traj_rs = resampling.interp_hmats(timesteps_rs, np.arange(len(new_ee_traj)), new_ee_traj)
                    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)
                if args.execution: Globals.pr2.update_rave()
                part_name = {"l":"larm", "r":"rarm"}[lr]
                bodypart2traj[part_name] = new_joint_traj
                redprint("Joint trajectory has length: " + str(len(bodypart2traj[part_name])))
                
            redprint("Executing joint trajectory for segment %s, part %i using arms '%s'"%(seg_name, i_miniseg, bodypart2traj.keys()))

            if args.pid:

                if not args.fake_data_segment: # If running on PR2, add initial state as waypoint and rough interpolate
                    # Add initial position
                    (arm_position, arm_velocity, arm_effort) = robot_states.call_return_joint_states(robot_states.arm_joint_names)
                    traj_length = bodypart2traj['rarm'].shape[0]
                    complete_joint_traj = np.zeros((traj_length+1, 7)) # To include initial state as a way point
                    complete_joint_traj[1:,:] = bodypart2traj['rarm']
                    complete_joint_traj[0,:] = arm_position
                    bodypart2traj['rarm'] = complete_joint_traj

                    # Add initial eff
                    J = np.matrix(np.resize(np.array(robot_states.call_return_jacobian()), (6, 7))) # Jacobian
                    eff = robot_states.compute_end_effector_force(J, arm_effort).T
                    eff =  np.array(eff)[0]
                    init_force = eff[:3]
                    complete_force_traj = np.zeros((traj_length+1, 3))
                    complete_force_traj[1:,:] = new_eefs_segment_rs
                    complete_force_traj[0,:] = init_force

                else:
                    complete_force_traj = new_eefs_segment_rs
                # SECOND RESAMPLING HAPPENS HERE: JOINT VELOCITIES
                if args.no_ds:
                    traj = bodypart2traj['rarm']
                    stamps = asarray(seg_info['stamps'])
                    times = np.array([stamps[i_end] - stamps[i_start]])
                    force = complete_force_traj
                else:
                    times, times_up, traj = pr2_trajectories.return_timed_traj(Globals.pr2, bodypart2traj) # use times and times_up to interpolate the second time
                    force = mu.interp2d(times_up, times, complete_force_traj)
                
                #Globals.robot.SetDOFValues(asarray(fake_seg["joint_states"]["position"][0]))
                if args.visualize:
                    r2r = ros2rave.RosToRave(Globals.robot, asarray(seg_info["joint_states"]["name"]))
                    r2r.set_values(Globals.robot, asarray(seg_info["joint_states"]["position"][0]))
                    for i in range(0, traj.shape[0], 10):
                        handles = []
                        handles.append(Globals.env.plot3(new_xyz,5,np.array([(0,1,0,1) for x in new_xyz])))
                        handles.append(Globals.env.plot3(old_xyz,5,np.array([(1,0,0,1) for x in old_xyz])))
                        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)))
                        handles.append(Globals.env.drawlinestrip(new_ee_traj_rs[:,:3,3], 2, (0,0,1,1)))
                        # Setting robot arm to joint trajectory
                        r_vals = traj[i,:]
                        Globals.robot.SetDOFValues(r_vals, Globals.robot.GetManipulator('rightarm').GetArmIndices())
                        
                        # Plotting forces from r_gripper_tool_frame
                        hmats = Globals.robot.GetLinkTransformations()
                        trans, rot = conv.hmat_to_trans_rot(hmats[-3])
                        f_start = np.array([0,0,0]) + trans
                        f_end = np.array(force[i])/100 + trans
                        handles.append(Globals.env.drawlinestrip(np.array([f_start, f_end]), 10, (1,0,0,1)))
                        
                        redprint(i)
                        Globals.viewer.Step()


                traj = np.array(np.matrix(traj).T) # 7 x n
                redprint(traj)

                traj = np.resize(traj, (1, traj.shape[1]*7)) #resize the data to 1x7n (n being the number of steps)
                force = np.array(np.matrix(force).T) # 3 x n
                force = np.resize(force, (1, force.shape[1]*3)) #resize the data to 1x3n
                #[traj, force, secs]
                traj_force_secs = np.zeros((1, traj.shape[1] + force.shape[1] + 2))
                traj_force_secs[0,0:traj.shape[1]] = traj
                traj_force_secs[0,traj.shape[1]:traj.shape[1]+force.shape[1]] = force
                traj_force_secs[0,traj.shape[1]+force.shape[1]] = times[len(times)-1]

                traj_force_secs[0,traj.shape[1]+force.shape[1]+1] = args.use_force
                IPython.embed()
                msg = Float64MultiArray()
                msg.data = traj_force_secs[0].tolist()
                pub = rospy.Publisher("/joint_positions_forces_secs", Float64MultiArray)
                redprint("Press enter to start trajectory")
                raw_input()
                time.sleep(1)
                pub.publish(msg)
                time.sleep(1)
            else:
                #if not success: break
                
                if len(bodypart2traj) > 0:
                    exec_traj_maybesim(bodypart2traj)
Example #17
0
def tps_nr_fit(x_na, y_ng, bend_coef, nr_ma, nr_coef, method="newton"):
    N,D = x_na.shape
    lin_ag, trans_g, w_ng = tps_fit2(x_na, y_ng, bend_coef, 1e-3)
    #return lin_ag, trans_g, w_ng

    ##for testing that it takes one step when nonrigidity cost is zero:
    #lin_ag, trans_g, w_ng = tps_fit(x_na, y_ng, bend_coef, 0)
    #res_cost, bend_cost, nr_cost, fval = tps_nr_cost_eval(lin_ag, trans_g, w_ng, x_na, nr_ma, bend_coef, nr_coef, return_tuple=True)
    #print "CORRECT COST, res,bend,nr,total = %.3e, %.3e, %.3e, %.3e"%(res_cost, bend_cost, nr_cost, fval)
    #lin_ag += np.random.randn(*lin_ag.shape)*5
    #res_cost, bend_cost, nr_cost, fval = tps_nr_cost_eval(lin_ag, trans_g, w_ng, x_na, nr_ma, bend_coef, nr_coef, return_tuple=True)
    #print "NOISE ADDED COST, res,bend,nr,total = %.ef, %.3e, %.3e, %.3e"%(res_cost, bend_cost, nr_cost, fval)
    
    _u,_s,_vh = np.linalg.svd(np.c_[x_na, np.ones((N,1))], full_matrices=True)
    N_nq = _u[:,4:] # null of data
    #w_ng = N_nq.dot(N_nq.T.dot(w_ng))
        
    K_nn = tps_kernel_matrix(x_na)
    Q_nn = np.c_[x_na, np.ones((N,1)),K_nn.dot(N_nq)]
    QQ_nn = np.dot(Q_nn.T, Q_nn)
    Bend_nn = np.zeros((N,N))
    Bend_nn[4:, 4:] = - N_nq.T.dot(K_nn.dot(N_nq))
    
    n_iter=60
    for i in xrange(n_iter):
        X_ng = np.r_[lin_ag, trans_g[None,:], N_nq.T.dot(w_ng)]

        res_cost, bend_cost, nr_cost, fval = tps_nr_cost_eval(lin_ag, trans_g, w_ng, x_na, y_ng, nr_ma, bend_coef, nr_coef, return_tuple=True)
        if VERBOSE: print colorize("iteration %i, cost %.3e"%(i, fval), 'red'),
        if VERBOSE: print "= %.3e (res) + %.3e (bend) + %.3e (nr)"%(res_cost, bend_cost, nr_cost)
                
        
        Jl_zcg, Jt_zg, Jw_zng = tps_nr_grad(nr_ma, lin_ag, trans_g, w_ng, x_na, return_tuple=True)
        nrerr_z = tps_nr_err(nr_ma, lin_ag, trans_g, w_ng, x_na)
        
        
        if method == "newton":
            fullstep_ng = np.empty((N,D))
            for g in xrange(D):
                J_zn = np.c_[Jl_zcg[:,g::D], Jt_zg[:,g::D], Jw_zng[:,g::D].dot(N_nq)]
                JJ_nn = np.dot(J_zn.T, J_zn)
                A = nr_coef*JJ_nn + QQ_nn + bend_coef*Bend_nn
                X0 = X_ng[:,g]
                B = nr_coef*np.dot(J_zn.T, np.dot(J_zn, X0) - nrerr_z) + Q_nn.T.dot(y_ng[:,g])
                fullstep_ng[:,g] = np.linalg.solve(A,B) - X_ng[:,g]

        elif method == "gradient":
            # def eval_partial(cand_X_ng):
            #     cand_X_ng = cand_X_ng.reshape(-1,3)
            #     cand_lin_ag, cand_trans_g, cand_w_ng = cand_X_ng[:D], cand_X_ng[D], N_nq.dot(cand_X_ng[D+1:])
            #     fval_cand = tps_nr_cost_eval(cand_lin_ag, cand_trans_g, cand_w_ng, x_na, y_ng, nr_ma, bend_coef, nr_coef)
            #     return fval_cand
            # def eval_partial2(cand_X_ng):
            #     return ((Q_nn.dot(X_ng) - y_ng)**2).sum()
            # def eval_partial3(cand_X_ng):
            #     cand_X_ng = cand_X_ng.reshape(-1,3)
            #     cand_lin_ag, cand_trans_g, cand_w_ng = cand_X_ng[:D], cand_X_ng[D], N_nq.dot(cand_X_ng[D+1:])
            #     return ((y_ng - tps_eval(x_na, cand_lin_ag, cand_trans_g, cand_w_ng, x_na))**2).sum()
            
            
            grad_ng = np.empty((N,D))
            for g in xrange(D-1,-1,-1):
                Jnr_zn = np.c_[Jl_zcg[:,g::D], Jt_zg[:,g::D], Jw_zng[:,g::D].dot(N_nq)]
                grad_ng[:,g] = 2 * nr_coef * nrerr_z.dot(Jnr_zn) \
                    + 2 * Q_nn.T.dot(Q_nn.dot(X_ng[:,g]) - y_ng[:,g]) \
                    + 2 * bend_coef * Bend_nn.dot(X_ng[:,g])

            #assert np.allclose(eval_partial2(X_ng), eval_partial3(X_ng))
            #assert np.allclose(eval_partial(X_ng), eval_partial2(X_ng))
            #grad0_ng = ndt.Gradient(eval_partial)(X_ng.flatten()).reshape(-1,3)
            fullstep_ng = -grad_ng
            #assert np.allclose(grad0_ng, grad_ng)
            
            
            

        cost_improved = False
        for stepsize in 3.**np.arange(0,-10,-1):
            cand_X_ng = X_ng + fullstep_ng*stepsize
            cand_lin_ag, cand_trans_g, cand_w_ng = cand_X_ng[:D], cand_X_ng[D], N_nq.dot(cand_X_ng[D+1:])
            fval_cand = tps_nr_cost_eval(cand_lin_ag, cand_trans_g, cand_w_ng, x_na, y_ng, nr_ma, bend_coef, nr_coef)
            if VERBOSE: print "stepsize: %.1g, fval: %.3e"%(stepsize, fval_cand)
            if fval_cand < fval:
                cost_improved = True
                break
        if not cost_improved:
            if VERBOSE: print "couldn't improve objective"
            break

            
        lin_ag = cand_lin_ag
        trans_g = cand_trans_g
        w_ng = cand_w_ng
    return lin_ag, trans_g, w_ng
Example #18
0
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
Example #19
0
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
Example #20
0
def call_and_print(cmd, color='green'):
    print colorize(cmd, color, bold=True)
    subprocess.check_call(cmd, shell=True)
def main():
    import IPython
    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.eval.fake_data_segment:
        grabber = cloudprocpy.CloudGrabber()
        grabber.startRGBD()

    Globals.viewer = trajoptpy.GetViewer(Globals.env)

    #####################

    while True:
    
        redprint("Acquire point cloud")
        if args.eval.fake_data_segment:
            fake_seg = demofile[args.eval.fake_data_segment]
            new_xyz = np.squeeze(fake_seg["cloud_xyz"])
            hmat = openravepy.matrixFromAxisAngle(args.eval.fake_data_transform[3:6])
            hmat[:3,3] = args.eval.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]))
            #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(2)
        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(2)

            
            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)
            print "got new xyz"
            redprint(new_xyz)
            #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) 
        ################################



        ### Old end effector forces at r_gripper_tool_frame (including torques)
        old_forces = seg_info['end_effector_forces'][:,0:3,:]
        old_torques = seg_info['end_effector_forces'][:,3:6,:]


        redprint("Generating end-effector trajectory")    

        handles = []
        old_xyz = np.squeeze(seg_info["cloud_xyz"])
       


        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)
        
        old_xyz_transformed = f.transform_points(old_xyz)

        #handles.append(Globals.env.plot3(old_xyz_transformed,5, np.array([(0,0,1,1) for i in old_xyz_transformed])))

        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 'r':
            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
            #print old_ee_traj
            old_ee_pos = old_ee_traj[:,0:3,3]
            #print old_ee_pos

            # End effector forces as oppossed to end effector trajectories
            dfdxs = f.compute_jacobian(old_ee_pos)
            old_eefs = []
            new_eefs = []
            for i in xrange(len(old_forces)):
                dfdx = np.matrix(dfdxs[i])
                old_force = np.matrix(old_forces[i])
                old_torque = np.matrix(old_torques[i])

                new_force = (dfdx * old_force)
                new_torque = (dfdx * old_torque)
                old_eefs.append(np.vstack((old_force, old_torque)))
                new_eefs.append(np.vstack((new_force, new_torque)))

            old_eefs = np.asarray(old_eefs)[:,:,0]
            new_eefs = np.asarray(new_eefs)[:,:,0]

            force_data = {}
            force_data['old_eefs'] = old_eefs
            force_data['new_eefs'] = new_eefs
            force_file = open("trial.pickle", 'wa')
            pickle.dump(force_data, force_file)
            force_file.close()
            new_ee_traj_x = new_ee_traj
            
        
        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 'r':
                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): NOT SURE WHY BUT THIS IS RETURNING FALSE
                lr2oldtraj[lr] = old_joint_traj

            if args.visualize:
                    r2r = ros2rave.RosToRave(Globals.robot, asarray(seg_info["joint_states"]["name"]))
                    r2r.set_values(Globals.robot, asarray(seg_info["joint_states"]["position"][0]))
                    for i in range(0, lr2oldtraj['r'].shape[0], 10):
                        handles = []
                        handles.append(Globals.env.plot3(new_xyz,5,np.array([(0,1,0,1) for x in new_xyz])))
                        handles.append(Globals.env.plot3(old_xyz,5,np.array([(1,0,0,1) for x in old_xyz])))
                        handles.append(Globals.env.drawlinestrip(old_ee_traj[:,:3,3], 2, (1,0,0,1)))
                        # Setting robot arm to joint trajectory
                        r_vals = lr2oldtraj['r'][i,:]
                        Globals.robot.SetDOFValues(r_vals, Globals.robot.GetManipulator('rightarm').GetArmIndices())
                        # Plotting forces from r_gripper_tool_frame
                        hmats = Globals.robot.GetLinkTransformations()
                        trans, rot = conv.hmat_to_trans_rot(hmats[-3])
                        f_start = np.array([0,0,0]) + trans
                        #IPython.embed()
                        f_end = np.array(old_forces[i].T)[0]/100 + trans
                        handles.append(Globals.env.drawlinestrip(np.array([f_start, f_end]), 10, (1,0,0,1)))
                        
                        redprint(i)
                        Globals.viewer.Step()
            if len(lr2oldtraj) > 0:
                old_total_traj = np.concatenate(lr2oldtraj.values(), 1)
                JOINT_LENGTH_PER_STEP = .1
                # FIRST RESAMPLING HAPPENS HERE: JOINT_LENGTH
                _, timesteps_rs = unif_resample(old_total_traj, JOINT_LENGTH_PER_STEP) # Timesteps only, can use to inter eefs for first time
            ####
            new_eefs_segment = asarray(new_eefs[i_start:i_end+1,:]) # Extract miniseg, and re-sample

            if args.no_ds:
                new_eefs_segment_rs = new_eefs_segment
            else:
                new_eefs_segment_rs = mu.interp2d(timesteps_rs, np.arange(len(new_eefs_segment)), new_eefs_segment)
        
            ### Generate fullbody traj
            bodypart2traj = {}
            for (lr,old_joint_traj) in lr2oldtraj.items():

                manip_name = {"l":"leftarm", "r":"rightarm"}[lr]
                ee_link_name = "%s_gripper_tool_frame"%lr
                new_ee_traj = link2eetraj[ee_link_name][i_start:i_end+1]
                if args.no_ds:
                    old_joint_traj_rs = old_joint_traj
                    new_ee_traj_rs = new_ee_traj
                    ds_inds = np.arange(0, len(new_ee_traj_rs), args.trajopt_ds)
                    new_ee_traj_rs = new_ee_traj_rs[ds_inds]
                    
                    old_joint_traj_rs = old_joint_traj_rs[ds_inds]
                    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)
                    new_joint_traj = mu.interp2d(np.arange(len(old_joint_traj)), np.arange(0, len(new_joint_traj) * args.trajopt_ds, args.trajopt_ds), new_joint_traj)
                else:
                    old_joint_traj_rs = mu.interp2d(timesteps_rs, np.arange(len(old_joint_traj)), old_joint_traj)
                    new_ee_traj_rs = resampling.interp_hmats(timesteps_rs, np.arange(len(new_ee_traj)), new_ee_traj)
                    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)
                if args.execution: Globals.pr2.update_rave()
                part_name = {"l":"larm", "r":"rarm"}[lr]
                bodypart2traj[part_name] = new_joint_traj
                redprint("Joint trajectory has length: " + str(len(bodypart2traj[part_name])))
                
            redprint("Executing joint trajectory for segment %s, part %i using arms '%s'"%(seg_name, i_miniseg, bodypart2traj.keys()))


            #for lr in 'r':
            #    set_gripper_maybesim(lr, binarize_gripper(seg_info["%s_gripper_joint"%lr][i_start]))

            if args.pid:
                # Add trajectory to arrive at to initial state
                redprint("Press enter to use current position")
                raw_input()
                (arm_position, arm_velocity, arm_effort) = robot_states.call_return_joint_states(robot_states.r_arm_joint_names)
                diff = np.array(arm_position - bodypart2traj['rarm'][0,:])
                maximum = max(abs(diff))
                speed = (300.0/360.0*2*(np.pi))
                time_needed = maximum / speed 
                initial_pos_traj = mu.interp2d(np.arange(0, time_needed, 0.01), np.array([0,time_needed]), np.array([arm_position, bodypart2traj['rarm'][0,:]]))
                # length of the intial trajectory, use this length for padding PD gains
                initial_traj_length = initial_pos_traj.shape[0]
                initial_force_traj = np.array([np.zeros(6) for i in range(initial_traj_length)])

                temptraj = bodypart2traj['rarm']
                bodypart2traj['rarm'] = np.concatenate((initial_pos_traj, temptraj), axis=0)

                complete_force_traj = np.concatenate((initial_force_traj, new_eefs), axis=0)

                # SECOND RESAMPLING HAPPENS HERE: JOINT VELOCITIES
                if args.no_ds:
                    traj = bodypart2traj['rarm']
                    stamps = asarray(seg_info['stamps'])
                    times = np.array([stamps[i_end] - stamps[i_start]])
                    F = complete_force_traj
                else:
                    times, times_up, traj = pr2_trajectories.return_timed_traj(Globals.pr2, bodypart2traj) # use times and times_up to interpolate the second time
                    F = mu.interp2d(times_up, times, complete_force_traj)
                
                #Globals.robot.SetDOFValues(asarray(fake_seg["joint_states"]["position"][0]))
                if args.visualize:
                    r2r = ros2rave.RosToRave(Globals.robot, asarray(seg_info["joint_states"]["name"]))
                    r2r.set_values(Globals.robot, asarray(seg_info["joint_states"]["position"][0]))
                    for i in range(0, traj.shape[0], 10):
                        handles = []
                        handles.append(Globals.env.plot3(new_xyz,5,np.array([(0,1,0,1) for x in new_xyz])))
                        handles.append(Globals.env.plot3(old_xyz,5,np.array([(1,0,0,1) for x in old_xyz])))
                        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)))
                        handles.append(Globals.env.drawlinestrip(new_ee_traj_rs[:,:3,3], 2, (0,0,1,1)))
                        # Setting robot arm to joint trajectory
                        r_vals = traj[i,:]
                        Globals.robot.SetDOFValues(r_vals, Globals.robot.GetManipulator('rightarm').GetArmIndices())
                        
                        # Plotting forces from r_gripper_tool_frame
                        hmats = Globals.robot.GetLinkTransformations()
                        trans, rot = conv.hmat_to_trans_rot(hmats[-3])
                        f_start = np.array([0,0,0]) + trans
                        f_end = np.array(old_forces[i].T)[0]/100 + trans
                        handles.append(Globals.env.drawlinestrip(np.array([f_start, f_end]), 10, (1,0,0,1)))
                        
                        redprint(i)
                        Globals.viewer.Idle()









                # Controller code in joint space
                traj_length = traj.shape[0]
                pgains = np.asarray([2400.0, 1200.0, 1000.0, 700.0, 300.0, 300.0, 300.0])
                dgains = np.asarray([18.0, 10.0, 6.0, 4.0, 6.0, 4.0, 4.0])
                m = np.array([3.33, 1.16, 0.1, 0.25, 0.133, 0.0727, 0.0727]) # masses in joint space (feed forward)
                vel_factor = 1e-3
                #new costs
                covariances, means = analyze_data.run_analyze_data(args)
                CostsNew = []
                counterCovs = 0
                endTraj = False
                covThiss = []
                for covMat in covariances:
                    covThis = np.zeros((18,18))
                    covThis[0:6,0:6] = covMat[0:6,0:6]
                    covThis[12:18,12:18] = covMat[6:12, 6:12]
                    covThis[0:6, 12:18] = covMat[0:6, 6:12]
                    covThis[12:18, 0:6] = covMat[6:12, 0:6]
                    covThis[6:12, 6:12] = np.eye(6)*vel_factor
                    covThis = np.diag(np.diag(covThis))
                    covThiss.append(covThis)
                    #if len(covThiss) <= 69 and len(covThiss) >= 47:
                    #    covThis[12:18,12:18] = np.diag([0.13, 0.06, 0.07, 0.005, 0.01, 0.004])
                    for j in range(args.eval.downsample_traj):
                        invCov = np.linalg.inv(covThis)
                        CostsNew.append(invCov)
                        counterCovs = counterCovs + 1
                        if counterCovs >= traj_length:
                            endTraj = True
                            break
                    if endTraj:
                        break
                allCs = []
                x = np.ones(6) * 1 
                v = np.ones(6) * 1e-3
                a = np.ones(6) * 1e-6
                Ct = np.diag(np.hstack((x, v, a))) # in end effector space
                Ms = []
                num_steps = i_end - i_start + 1
                
                
                stamps = asarray(seg_info['stamps'][i_start:i_end+1])


                arm = Globals.robot.GetManipulator("rightarm")
                jacobians = []
                for i in range(traj.shape[0]):
                    r_vals = traj[i,:]
                    Globals.robot.SetDOFValues(r_vals, Globals.robot.GetManipulator('rightarm').GetArmIndices())
                    jacobians.append(np.vstack((arm.CalculateJacobian(), arm.CalculateAngularVelocityJacobian())))
                jacobians =  asarray(jacobians)
                #jacobiansx = asarray(seg_info["jacobians"][i_start:i_end+1])
                for t in range(num_steps):
                    M = np.zeros((18, 21))
                    J = jacobians[t]
                    M[0:6,0:7] = J
                    M[6:12,7:14] = J
                    mdiag = np.diag(m) # Convert joint feedforward values into diagonal array
                    mdiag_inv = np.linalg.inv(mdiag)
                    M[12:18,14:21] = np.linalg.inv((J.dot(mdiag_inv).dot(np.transpose(J)))).dot(J).dot(mdiag_inv)
                    Ms.append(M)
                for t in range(num_steps):
                    topad = np.zeros((21,21))
                    topad[0:7,0:7] = np.diag(pgains) * 0.1
                    topad[7:14,7:14] = np.diag(dgains) * 0.1
                    topad[14:21,14:21] = np.eye(7) * 0.1
                    #allCs.append(np.transpose(Ms[t]).dot(Ct).dot(Ms[t]) + topad)
                    allCs.append(np.transpose(Ms[t]).dot(CostsNew[t]).dot(Ms[t]) + topad)
                Kps = []
                Kvs = []
                Ks = []
                Qt = None
                Vs = []
                for t in range(num_steps-1, -1, -1):
                    if Qt is None:
                        Qt = allCs[t]
                    else:
                        Ft = np.zeros((14, 21))
                        for j in range(14):
                            Ft[j][j] = 1.0
                        deltat = abs(stamps[t+1] - stamps[t])
                        #print(deltat)
                        for j in range(7):
                            Ft[j][j+7] = deltat
                        for j in range(7):
                            Ft[j+7][j+14] = deltat/m[j]  
                        for j in range(7):
                            Ft[j][j+14] = ((deltat)**2)/m[j]
                        Qt = allCs[t] + (np.transpose(Ft).dot(Vs[num_steps-2-t]).dot(Ft))
                    Qxx = Qt[0:14, 0:14]
                    Quu = Qt[14:21, 14:21]
                    Qxu = Qt[0:14, 14:21]
                    Qux = Qt[14:21, 0:14]
                    Quuinv = np.linalg.inv(Quu)
                    Vt = Qxx - Qxu.dot(Quuinv).dot(Qux)
                    Vt = 0.5*(Vt + np.transpose(Vt))
                    Kt = -1*(Quuinv.dot(Qux))
                    Ks.append(Kt)
                    Kps.append(Kt[:, 0:7])
                    Kvs.append(Kt[:, 7:14])
                    Vs.append(Vt)

                Ks = Ks[::-1]
                Kps = Kps[::-1]
                Kvs = Kvs[::-1]


                JKpJ = np.asarray(Kps)
                JKvJ = np.asarray(Kvs)




                total_length = num_steps + initial_traj_length

                # Pad initial traj with PD gains
                pgainsdiag = np.diag(np.asarray([-2400.0, -1200.0, -1000.0, -700.0, -300.0, -300.0, -300.0]))
                dgainsdiag = np.diag(np.asarray([-18.0, -10.0, -6.0, -4.0, -6.0, -4.0, -4.0]))
                addkp = np.asarray([pgainsdiag for i in range(initial_traj_length)])
                addkv = np.asarray([dgainsdiag for i in range(initial_traj_length)])
                JKpJ = np.concatenate((addkp, JKpJ), axis=0)
                JKvJ = np.concatenate((addkv, JKvJ), axis=0)
                
                Kps = []
                Kvs = []

                for i in range(num_steps + initial_traj_length):
                    Kps.append(np.zeros((6,6)))
                    Kvs.append(np.zeros((6,6)))

                Kps = np.asarray(Kps)
                Kvs = np.asarray(Kvs)

                # Gains as JKpJ and JKvJ for testing
                if args.pdgains:
                    JKpJ = np.asarray([pgainsdiag for i in range(total_length)])
                    JKvJ = np.asarray([dgainsdiag for i in range(total_length)])
                


                qs = traj
                IPython.embed()


                Kps = np.resize(Kps, (1, 36 * total_length))[0]
                Kvs = np.resize(Kvs, (1, 36 * total_length))[0]
                
                JKvJ = np.resize(JKvJ, (1, 49*total_length))[0]
                JKpJ = np.resize(JKpJ, (1, 49*total_length))[0]
                
                # For use with new controller                

                qs = np.resize(qs, (1, qs.shape[0]*7))[0] #resize the data to 1x7n (n being the number of steps)
                
                F = np.resize(F, (1, F.shape[0]*6))[0] #resize the data to 1x6n
                

                # [traj, Kp, Kv, F, use_force, seconds]
                data = np.zeros((1, len(qs) + len(JKpJ) + len(JKvJ) + len(F) + len(Kps) + len(Kvs) + 2))
                data[0][0:len(qs)] = qs
                data[0][len(qs):len(qs)+len(JKpJ)] = JKpJ
                data[0][len(qs)+len(JKpJ):len(qs)+len(JKpJ)+len(JKvJ)] = JKvJ
                data[0][len(qs)+len(JKpJ)+len(JKvJ):len(qs)+len(JKpJ)+len(JKvJ)+len(F)] = F
                data[0][len(qs)+len(JKpJ)+len(JKvJ)+len(F):len(qs)+len(JKpJ)+len(JKvJ)+len(F)+len(Kps)] = Kps
                data[0][len(qs)+len(JKpJ)+len(JKvJ)+len(F)+len(Kps):len(qs)+len(JKpJ)+len(JKvJ)+len(F)+len(Kps)+len(Kvs)] = Kvs
                data[0][-2] = args.use_force
                data[0][-1] = stamps[i_end] - stamps[i_start] + (initial_traj_length * 0.01)
                
                # For use with old controller
                """
                qs = np.array(np.matrix(traj).T) # 7 x n
                qs = np.resize(qs, (1, qs.shape[1]*7))[0] #resize the data to 1x7n (n being the number of steps)
                F = np.array(np.matrix(F).T) # 6 x n
                F = F[0:3,:]
                F = np.resize(F, (1, F.shape[1]*3))[0] #resize the data to 1x3n
                data = np.zeros((1, len(qs) + len(F) + 2))
                data[0][0:len(qs)] = qs
                data[0][len(qs):len(qs)+len(F)] = F
                data[0][-1] = args.use_force
                data[0][-2] = stamps[i_end] - stamps[i_start]
                """


                msg = Float64MultiArray()
                msg.data = data[0].tolist()
                pub = rospy.Publisher("/controller_data", Float64MultiArray)
                redprint("Press enter to start trajectory")
                raw_input()
                time.sleep(1)
                pub.publish(msg)
                time.sleep(1)
            else:
                #if not success: break
                
                if len(bodypart2traj) > 0:
                    exec_traj_maybesim(bodypart2traj)
Example #22
0
def call_and_print(cmd, color='green', check=True):
    print colorize(cmd, color, bold=True)
    if check: subprocess.check_call(cmd, shell=True)
    else: subprocess.call(cmd, shell=True)
Example #23
0
def call_and_print(cmd,color='green'):
    print colorize(cmd, color, bold=True)
    subprocess.check_call(cmd, shell=True)
def redprint(msg):
    """Print the message to the console in red, bold font."""
    print colorize.colorize(msg, "red", bold=True)
Example #25
0
os.chdir(osp.dirname(args.master_file))

with open(args.master_file, "r") as fh: master_info = yaml.load(fh)
for suffix in itertools.chain("", (str(i) for i in itertools.count())):
    demo_name = args.demo_prefix + suffix
    if not any(bag["demo_name"] == demo_name for bag in master_info["bags"]):
        break
    print demo_name

subprocess.call("killall XnSensorServer", shell=True)

try:

    bag_cmd = "rosbag record /joint_states /joy -O %s"%demo_name
    print colorize(bag_cmd, "green")
    bag_handle = subprocess.Popen(bag_cmd, shell=True)
    started_bag = True
    
    video_cmd = "record_rgbd_video --out=%s --downsample=%i"%(demo_name, args.downsample)
    print colorize(video_cmd, "green")
    video_handle = subprocess.Popen(video_cmd, shell=True)
    started_video = True
    
    time.sleep(9999)    

except KeyboardInterrupt:
    print colorize("got control-c", "green")

finally:
    
Example #26
0
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
Example #27
0
def call_and_print(cmd,color='green',check=True):
    print colorize(cmd, color, bold=True)
    if check: subprocess.check_call(cmd, shell=True)
    else: subprocess.call(cmd, shell=True)
Example #28
0
def blueprint(msg):
    print colorize.colorize(msg, "blue", bold=True)
Example #29
0
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
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
Example #31
0
def tps_nr_fit_enhanced(x_na, y_ng, bend_coef, nr_ma, nr_coef):
    
    N,D = x_na.shape
    M = nr_ma.shape[0]
    E = N + 4*M
    F = E - M
    Q = N + 3*M - 4
    
    s = .1 # tetrahedron sidelength (meters)
    u = 1/(2*np.sqrt(2))
    
    tetra_pts = []
    for pt in nr_ma:
        tetra_pts.append(s*np.r_[-.5, 0, -u]+pt)
        tetra_pts.append(s*np.r_[+.5, 0, -u]+pt)
        tetra_pts.append(s*np.r_[0, -.5, +u]+pt)
        tetra_pts.append(s*np.r_[0, +.5, +u]+pt)
    
    x_ea = np.r_[x_na, tetra_pts]

    badsub_ex = np.c_[x_ea, np.ones((E,1)), np.r_[np.zeros((N,M)), np.repeat(np.eye(M), 4, axis=0)]]    
    lin_ag, trans_g, w_ng = tps_fit2(x_na, y_ng,  bend_coef, 1e-3)
    w_eg = np.r_[w_ng, np.zeros((4*M, D))]

    assert badsub_ex.shape[0] >= badsub_ex.shape[1]
    _u,_s,_vh = np.linalg.svd(badsub_ex, full_matrices=True)
    assert badsub_ex.shape[1] == _s.size
    N_eq = _u[:,badsub_ex.shape[1]:] # null of data
        
    assert N_eq.shape == (E,Q)

    assert E == N + 4*M
    assert F == Q + 4
    # e is number of kernels
    # q is number of nonrigid dofs
    # f is total number of dofs
    K_ee = tps_kernel_matrix(x_ea)
    K_ne = K_ee[:N, :]
    Q_nf = np.c_[x_na, np.ones((N,1)),K_ne.dot(N_eq)]
    QQ_ff = np.dot(Q_nf.T, Q_nf)
    Bend_ff = np.zeros((F,F))
    Bend_ff[4:, 4:] = - N_eq.T.dot(K_ee.dot(N_eq)) # K_qq
    
    assert Q_nf.shape == (N, F)
    assert w_eg.shape == (E, D)
    
    n_iter=40
    for i in xrange(n_iter):
        
        
        # if plotting and i%plotting==0:
            # import lfd.registration as lr
            # lr.Globals.setup()
            # def eval_partial(x_ma):
            #     return tps_eval(x_ma, lin_ag, trans_g, w_eg, x_ea) 
            # lr.plot_orig_and_warped_clouds(eval_partial, x_na, y_ng, res=.008)            
        
        X_fg = np.r_[lin_ag, 
                    trans_g[None,:], 
                    N_eq.T.dot(w_eg)]

        res_cost, bend_cost, nr_cost, fval = tps_nr_cost_eval_general(lin_ag, trans_g, w_eg, x_ea, y_ng, nr_ma, bend_coef, nr_coef, return_tuple=True)
        if VERBOSE: print colorize("iteration %i, cost %.3e"%(i, fval), 'red'),
        if VERBOSE: print "= %.3e (res) + %.3e (bend) + %.3e (nr)"%(res_cost, bend_cost, nr_cost)
                
        
        Jl_zcg, Jt_zg, Jw_zeg = tps_nr_grad(nr_ma, lin_ag, trans_g, w_eg, x_ea, return_tuple=True)
        nrerr_z = tps_nr_err(nr_ma, lin_ag, trans_g, w_eg, x_ea)        
        
        fullstep_fg = np.empty((F,D))
        for g in xrange(D):
            J_zf = np.c_[Jl_zcg[:,g::D], Jt_zg[:,g::D], Jw_zeg[:,g::D].dot(N_eq)]
            JJ_ff = np.dot(J_zf.T, J_zf)
            A_ff = nr_coef*JJ_ff + QQ_ff + bend_coef*Bend_ff
            X0 = X_fg[:,g]
            B_f = nr_coef*np.dot(J_zf.T, np.dot(J_zf, X0) - nrerr_z) + Q_nf.T.dot(y_ng[:,g])
            fullstep_fg[:,g] = np.linalg.solve(A_ff,B_f) - X_fg[:,g]

        cost_improved = False
        for stepsize in 3.**np.arange(0,-10,-1):
            cand_X_fg = X_fg + fullstep_fg*stepsize
            cand_lin_ag, cand_trans_g, cand_w_eg = cand_X_fg[:D], cand_X_fg[D], N_eq.dot(cand_X_fg[D+1:])
            fval_cand = tps_nr_cost_eval_general(cand_lin_ag, cand_trans_g, cand_w_eg, x_ea, y_ng, nr_ma, bend_coef, nr_coef, return_tuple=False)
            if VERBOSE: print "stepsize: %.1g, fval: %.3e"%(stepsize, fval_cand)
            if fval_cand < fval:
                cost_improved = True
                break
        if not cost_improved:
            if VERBOSE: print "couldn't improve objective"
            break

            
        lin_ag = cand_lin_ag
        trans_g = cand_trans_g
        w_eg = cand_w_eg
    return lin_ag, trans_g, w_eg, x_ea
Example #32
0
def tps_nr_fit(x_na, y_ng, bend_coef, nr_ma, nr_coef, method="newton"):
    N, D = x_na.shape
    lin_ag, trans_g, w_ng = tps_fit2(x_na, y_ng, bend_coef, 1e-3)
    #return lin_ag, trans_g, w_ng

    ##for testing that it takes one step when nonrigidity cost is zero:
    #lin_ag, trans_g, w_ng = tps_fit(x_na, y_ng, bend_coef, 0)
    #res_cost, bend_cost, nr_cost, fval = tps_nr_cost_eval(lin_ag, trans_g, w_ng, x_na, nr_ma, bend_coef, nr_coef, return_tuple=True)
    #print "CORRECT COST, res,bend,nr,total = %.3e, %.3e, %.3e, %.3e"%(res_cost, bend_cost, nr_cost, fval)
    #lin_ag += np.random.randn(*lin_ag.shape)*5
    #res_cost, bend_cost, nr_cost, fval = tps_nr_cost_eval(lin_ag, trans_g, w_ng, x_na, nr_ma, bend_coef, nr_coef, return_tuple=True)
    #print "NOISE ADDED COST, res,bend,nr,total = %.ef, %.3e, %.3e, %.3e"%(res_cost, bend_cost, nr_cost, fval)

    _u, _s, _vh = np.linalg.svd(np.c_[x_na, np.ones((N, 1))],
                                full_matrices=True)
    N_nq = _u[:, 4:]  # null of data
    #w_ng = N_nq.dot(N_nq.T.dot(w_ng))

    K_nn = ssd.squareform(ssd.pdist(x_na))
    Q_nn = np.c_[x_na, np.ones((N, 1)), K_nn.dot(N_nq)]
    QQ_nn = np.dot(Q_nn.T, Q_nn)
    Bend_nn = np.zeros((N, N))
    Bend_nn[4:, 4:] = -N_nq.T.dot(K_nn.dot(N_nq))

    n_iter = 60
    for i in xrange(n_iter):
        X_ng = np.r_[lin_ag, trans_g[None, :], N_nq.T.dot(w_ng)]

        res_cost, bend_cost, nr_cost, fval = tps_nr_cost_eval(
            lin_ag,
            trans_g,
            w_ng,
            x_na,
            y_ng,
            nr_ma,
            bend_coef,
            nr_coef,
            return_tuple=True)
        if VERBOSE:
            print colorize("iteration %i, cost %.3e" % (i, fval), 'red'),
        if VERBOSE:
            print "= %.3e (res) + %.3e (bend) + %.3e (nr)" % (
                res_cost, bend_cost, nr_cost)

        Jl_zcg, Jt_zg, Jw_zng = tps_nr_grad(nr_ma,
                                            lin_ag,
                                            trans_g,
                                            w_ng,
                                            x_na,
                                            return_tuple=True)
        nrerr_z = tps_nr_err(nr_ma, lin_ag, trans_g, w_ng, x_na)

        if method == "newton":
            fullstep_ng = np.empty((N, D))
            for g in xrange(D):
                J_zn = np.c_[Jl_zcg[:, g::D], Jt_zg[:, g::D],
                             Jw_zng[:, g::D].dot(N_nq)]
                JJ_nn = np.dot(J_zn.T, J_zn)
                A = nr_coef * JJ_nn + QQ_nn + bend_coef * Bend_nn
                X0 = X_ng[:, g]
                B = nr_coef * np.dot(J_zn.T,
                                     np.dot(J_zn, X0) - nrerr_z) + Q_nn.T.dot(
                                         y_ng[:, g])
                fullstep_ng[:, g] = np.linalg.solve(A, B) - X_ng[:, g]

        elif method == "gradient":
            # def eval_partial(cand_X_ng):
            #     cand_X_ng = cand_X_ng.reshape(-1,3)
            #     cand_lin_ag, cand_trans_g, cand_w_ng = cand_X_ng[:D], cand_X_ng[D], N_nq.dot(cand_X_ng[D+1:])
            #     fval_cand = tps_nr_cost_eval(cand_lin_ag, cand_trans_g, cand_w_ng, x_na, y_ng, nr_ma, bend_coef, nr_coef)
            #     return fval_cand
            # def eval_partial2(cand_X_ng):
            #     return ((Q_nn.dot(X_ng) - y_ng)**2).sum()
            # def eval_partial3(cand_X_ng):
            #     cand_X_ng = cand_X_ng.reshape(-1,3)
            #     cand_lin_ag, cand_trans_g, cand_w_ng = cand_X_ng[:D], cand_X_ng[D], N_nq.dot(cand_X_ng[D+1:])
            #     return ((y_ng - tps_eval(x_na, cand_lin_ag, cand_trans_g, cand_w_ng, x_na))**2).sum()

            grad_ng = np.empty((N, D))
            for g in xrange(D - 1, -1, -1):
                Jnr_zn = np.c_[Jl_zcg[:, g::D], Jt_zg[:, g::D],
                               Jw_zng[:, g::D].dot(N_nq)]
                grad_ng[:,g] = 2 * nr_coef * nrerr_z.dot(Jnr_zn) \
                    + 2 * Q_nn.T.dot(Q_nn.dot(X_ng[:,g]) - y_ng[:,g]) \
                    + 2 * bend_coef * Bend_nn.dot(X_ng[:,g])

            #assert np.allclose(eval_partial2(X_ng), eval_partial3(X_ng))
            #assert np.allclose(eval_partial(X_ng), eval_partial2(X_ng))
            #grad0_ng = ndt.Gradient(eval_partial)(X_ng.flatten()).reshape(-1,3)
            fullstep_ng = -grad_ng
            #assert np.allclose(grad0_ng, grad_ng)

        cost_improved = False
        for stepsize in 3.**np.arange(0, -10, -1):
            cand_X_ng = X_ng + fullstep_ng * stepsize
            cand_lin_ag, cand_trans_g, cand_w_ng = cand_X_ng[:D], cand_X_ng[
                D], N_nq.dot(cand_X_ng[D + 1:])
            fval_cand = tps_nr_cost_eval(cand_lin_ag, cand_trans_g, cand_w_ng,
                                         x_na, y_ng, nr_ma, bend_coef, nr_coef)
            if VERBOSE:
                print "stepsize: %.1g, fval: %.3e" % (stepsize, fval_cand)
            if fval_cand < fval:
                cost_improved = True
                break
        if not cost_improved:
            if VERBOSE: print "couldn't improve objective"
            break

        lin_ag = cand_lin_ag
        trans_g = cand_trans_g
        w_ng = cand_w_ng
    return lin_ag, trans_g, w_ng
def main():
    import IPython
    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]))
            #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(2)
        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(2)

            
            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)
            print "got new xyz"
            redprint(new_xyz)
            #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) 
        ################################



        ### Old end effector forces at r_gripper_tool_frame (including torques)
        old_forces = seg_info['end_effector_forces'][:,0:3,:]
        old_torques = seg_info['end_effector_forces'][:,3:6,:]


        redprint("Generating end-effector trajectory")    

        handles = []
        old_xyz = np.squeeze(seg_info["cloud_xyz"])
       


        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)
        
        old_xyz_transformed = f.transform_points(old_xyz)

        #handles.append(Globals.env.plot3(old_xyz_transformed,5, np.array([(0,0,1,1) for i in old_xyz_transformed])))

        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 'r':
            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
            #print old_ee_traj
            old_ee_pos = old_ee_traj[:,0:3,3]
            #print old_ee_pos

            # End effector forces as oppossed to end effector trajectories
            dfdxs = f.compute_jacobian(old_ee_pos)
            old_eefs = []
            new_eefs = []
            for i in xrange(len(old_forces)):
                dfdx = np.matrix(dfdxs[i])
                old_force = np.matrix(old_forces[i])
                old_torque = np.matrix(old_torques[i])

                new_force = (dfdx * old_force)
                new_torque = (dfdx * old_torque)
                old_eefs.append(np.vstack((old_force, old_torque)))
                new_eefs.append(np.vstack((new_force, new_torque)))

            old_eefs = np.asarray(old_eefs)[:,:,0]
            new_eefs = np.asarray(new_eefs)[:,:,0]

            force_data = {}
            force_data['old_eefs'] = old_eefs
            force_data['new_eefs'] = new_eefs
            force_file = open("trial.pickle", 'wa')
            pickle.dump(force_data, force_file)
            force_file.close()
            new_ee_traj_x = new_ee_traj
            
        
        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 'r':
                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): NOT SURE WHY BUT THIS IS RETURNING FALSE
                lr2oldtraj[lr] = old_joint_traj

            if args.visualize:
                    r2r = ros2rave.RosToRave(Globals.robot, asarray(seg_info["joint_states"]["name"]))
                    r2r.set_values(Globals.robot, asarray(seg_info["joint_states"]["position"][0]))
                    for i in range(0, lr2oldtraj['r'].shape[0], 10):
                        handles = []
                        handles.append(Globals.env.plot3(new_xyz,5,np.array([(0,1,0,1) for x in new_xyz])))
                        handles.append(Globals.env.plot3(old_xyz,5,np.array([(1,0,0,1) for x in old_xyz])))
                        handles.append(Globals.env.drawlinestrip(old_ee_traj[:,:3,3], 2, (1,0,0,1)))
                        # Setting robot arm to joint trajectory
                        r_vals = lr2oldtraj['r'][i,:]
                        Globals.robot.SetDOFValues(r_vals, Globals.robot.GetManipulator('rightarm').GetArmIndices())
                        # Plotting forces from r_gripper_tool_frame
                        hmats = Globals.robot.GetLinkTransformations()
                        trans, rot = conv.hmat_to_trans_rot(hmats[-3])
                        f_start = np.array([0,0,0]) + trans
                        #IPython.embed()
                        f_end = np.array(old_forces[i].T)[0]/100 + trans
                        handles.append(Globals.env.drawlinestrip(np.array([f_start, f_end]), 10, (1,0,0,1)))
                        
                        redprint(i)
                        Globals.viewer.Step()
            if len(lr2oldtraj) > 0:
                old_total_traj = np.concatenate(lr2oldtraj.values(), 1)
                JOINT_LENGTH_PER_STEP = .1
                # FIRST RESAMPLING HAPPENS HERE: JOINT_LENGTH
                _, timesteps_rs = unif_resample(old_total_traj, JOINT_LENGTH_PER_STEP) # Timesteps only, can use to inter eefs for first time
            ####
            new_eefs_segment = asarray(new_eefs[i_start:i_end+1,:]) # Extract miniseg, and re-sample

            if args.no_ds:
                new_eefs_segment_rs = new_eefs_segment
            else:
                new_eefs_segment_rs = mu.interp2d(timesteps_rs, np.arange(len(new_eefs_segment)), new_eefs_segment)
        
            ### Generate fullbody traj
            bodypart2traj = {}
            for (lr,old_joint_traj) in lr2oldtraj.items():

                manip_name = {"l":"leftarm", "r":"rightarm"}[lr]
                ee_link_name = "%s_gripper_tool_frame"%lr
                new_ee_traj = link2eetraj[ee_link_name][i_start:i_end+1]
                if args.no_ds:
                    old_joint_traj_rs = old_joint_traj
                    new_ee_traj_rs = new_ee_traj
                    ds_inds = np.arange(0, len(new_ee_traj_rs), args.trajopt_ds)
                    new_ee_traj_rs = new_ee_traj_rs[ds_inds]
                    
                    old_joint_traj_rs = old_joint_traj_rs[ds_inds]
                    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)
                    new_joint_traj = mu.interp2d(np.arange(len(old_joint_traj)), np.arange(0, len(new_joint_traj) * args.trajopt_ds, args.trajopt_ds), new_joint_traj)
                else:
                    old_joint_traj_rs = mu.interp2d(timesteps_rs, np.arange(len(old_joint_traj)), old_joint_traj)
                    new_ee_traj_rs = resampling.interp_hmats(timesteps_rs, np.arange(len(new_ee_traj)), new_ee_traj)
                    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)
                if args.execution: Globals.pr2.update_rave()
                part_name = {"l":"larm", "r":"rarm"}[lr]
                bodypart2traj[part_name] = new_joint_traj
                redprint("Joint trajectory has length: " + str(len(bodypart2traj[part_name])))
                
            redprint("Executing joint trajectory for segment %s, part %i using arms '%s'"%(seg_name, i_miniseg, bodypart2traj.keys()))


            redprint("Press enter to set gripper")
            raw_input()

            #for lr in 'r':
            #    set_gripper_maybesim(lr, binarize_gripper(seg_info["%s_gripper_joint"%lr][i_start]))

            if args.pid:

                if not args.fake_data_segment: # If running on PR2, add initial state as waypoint and rough interpolate
                    # Add initial position
                    (arm_position, arm_velocity, arm_effort) = robot_states.call_return_joint_states(robot_states.arm_joint_names)
                    traj_length = bodypart2traj['rarm'].shape[0]
                    complete_joint_traj = np.zeros((traj_length+1, 7)) # To include initial state as a way point
                    complete_joint_traj[1:,:] = bodypart2traj['rarm']
                    complete_joint_traj[0,:] = arm_position
                    bodypart2traj['rarm'] = complete_joint_traj

                    # Add initial eff
                    J = np.matrix(np.resize(np.array(robot_states.call_return_jacobian()), (6, 7))) # Jacobian
                    eff = robot_states.compute_end_effector_force(J, arm_effort).T
                    eff =  np.array(eff)[0]
                    init_force = eff[:3]
                    complete_force_traj = np.zeros((traj_length+1, 3))
                    complete_force_traj[1:,:] = new_eefs_segment_rs
                    complete_force_traj[0,:] = init_force

                else:
                    complete_force_traj = new_eefs_segment_rs
                # SECOND RESAMPLING HAPPENS HERE: JOINT VELOCITIES
                if args.no_ds:
                    traj = bodypart2traj['rarm']
                    stamps = asarray(seg_info['stamps'])
                    times = np.array([stamps[i_end] - stamps[i_start]])
                    F = complete_force_traj
                else:
                    times, times_up, traj = pr2_trajectories.return_timed_traj(Globals.pr2, bodypart2traj) # use times and times_up to interpolate the second time
                    F = mu.interp2d(times_up, times, complete_force_traj)
                
                #Globals.robot.SetDOFValues(asarray(fake_seg["joint_states"]["position"][0]))
                if args.visualize:
                    r2r = ros2rave.RosToRave(Globals.robot, asarray(seg_info["joint_states"]["name"]))
                    r2r.set_values(Globals.robot, asarray(seg_info["joint_states"]["position"][0]))
                    for i in range(0, traj.shape[0], 10):
                        handles = []
                        handles.append(Globals.env.plot3(new_xyz,5,np.array([(0,1,0,1) for x in new_xyz])))
                        handles.append(Globals.env.plot3(old_xyz,5,np.array([(1,0,0,1) for x in old_xyz])))
                        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)))
                        handles.append(Globals.env.drawlinestrip(new_ee_traj_rs[:,:3,3], 2, (0,0,1,1)))
                        # Setting robot arm to joint trajectory
                        r_vals = traj[i,:]
                        Globals.robot.SetDOFValues(r_vals, Globals.robot.GetManipulator('rightarm').GetArmIndices())
                        
                        # Plotting forces from r_gripper_tool_frame
                        hmats = Globals.robot.GetLinkTransformations()
                        trans, rot = conv.hmat_to_trans_rot(hmats[-3])
                        f_start = np.array([0,0,0]) + trans
                        f_end = np.array(old_forces[i].T)[0]/100 + trans
                        handles.append(Globals.env.drawlinestrip(np.array([f_start, f_end]), 10, (1,0,0,1)))
                        
                        redprint(i)
                        Globals.viewer.Step()






                qs = np.array(np.matrix(traj).T) # 7 x n
                qs = np.resize(traj, (1, qs.shape[1]*7))[0] #resize the data to 1x7n (n being the number of steps)
                F = np.array(np.matrix(F).T) # 6 x n
                F = np.resize(F, (1, F.shape[1]*6))[0] #resize the data to 1x3n


                # Controller code
                allCs = []
                x = np.ones(6) * 1 
                v = np.ones(6) * 1e-3
                a = np.ones(6) * 1e-6
                Ct = np.diag(np.hstack((x, v, a)))
                num_steps = i_end - i_start + 1
                #need to figure out the stamps correctly
                stamps = asarray(seg_info['stamps'][i_start:i_end+1])
                for t in range(num_steps-1, -1, -1):
                    allCs.append(Ct)
                m = np.ones(6)
                Kps = []
                Kvs = []
                Ks = []
                Qt = None
                Vs = []
                for t in range(num_steps-1, -1, -1):
                    if Qt is None:
                        Qt = allCs[t]
                    else:
                        Ft = np.zeros((12, 18))
                        for j in range(12):
                            Ft[j][j] = 1.0
                        deltat = abs(stamps[t+1] - stamps[t])
                        #print(deltat)
                        for j in range(6):
                            Ft[j][j+6] = deltat
                        for j in range(6):
                            Ft[j+6][j+12] = deltat/m[j]  
                        for j in range(6):
                            Ft[j][j+12] = ((deltat)**2)/m[j]
                        Qt = allCs[t] + (np.transpose(Ft).dot(Vs[num_steps-2-t]).dot(Ft))
                    Qxx = Qt[0:12, 0:12]
                    Quu = Qt[12:18, 12:18]
                    Qxu = Qt[0:12, 12:18]
                    Qux = Qt[12:18, 0:12]
                    Quuinv = np.linalg.inv(Quu)
                    Vt = Qxx - Qxu.dot(Quuinv).dot(Qux)
                    Kt = -1*(Quuinv.dot(Qux))
                    Ks.append(Kt)
                    Kps.append(Kt[:, 0:6])
                    Kvs.append(Kt[:, 6:12])
                    Vs.append(Vt)

                Ks = Ks[::-1]
                Kps = Kps[::-1]
                Kvs = Kvs[::-1]



                JKpJ = []
                JKvJ = []
                toAddJkpJ = np.diag(np.asarray([-2400.0, -1200.0, -1000.0, -700.0, -300.0, -300.0, -300.0]))
                toAddJkvJ = np.diag(np.asarray([-18.0, -10.0, -6.0, -4.0, -6.0, -4.0, -4.0]))
                JacobianOlds = asarray(seg_info["jacobians"][i_start:i_end+1])
                for i in range(i_end- i_start + 1):
                    J = JacobianOlds[i]
                    psuedoJ = np.linalg.inv(np.transpose(J).dot(J)).dot(np.transpose(J))
                    #JKpJ.append(np.transpose(J).dot(Kps[i]).dot(J) + toAddJkpJ*0.001)
                    #JKvJ.append(np.transpose(J).dot(Kvs[i]).dot(J) + toAddJkvJ*0.001)
                    JKpJ.append(psuedoJ.dot(Kps[i]).dot(J) + toAddJkpJ*0.001)
                    JKvJ.append(psuedoJ.dot(Kvs[i]).dot(J) + toAddJkvJ*0.001)
                if args.useJK:
                    Kps = []
                    Kvs = []
                    for i in range(i_end- i_start + 1):
                        Kps.append(np.zeros((6,6)))
                        Kvs.append(np.zeros((6,6)))
                else:
                    JKpJ = []
                    JKvJ = []
                    for i in range(i_end- i_start + 1):
                        JKpJ.append(np.zeros((7,7)))
                        JKvJ.append(np.zeros((7,7)))



                Kps = np.asarray(Kps)
                Kvs = np.asarray(Kvs)
                JKpJ = np.asarray(JKpJ)
                JKvJ = np.asarray(JKvJ)

                IPython.embed()
                
                # # Temp Kps and Kvs for testing
                """
                toAddJkpJ = np.diag(np.asarray([-2400.0, -1200.0, -1000.0, -700.0, -300.0, -300.0, -300.0]))
                toAddJkvJ = np.diag(np.asarray([-18.0, -10.0, -6.0, -4.0, -6.0, -4.0, -4.0]))
                length = complete_force_traj.shape[0]
                JKpJ = np.asarray([toAddJkpJ for i in range(length)])
                JKpJ = np.resize(JKpJ, (1, 49*length))[0]

                JKvJ = np.asarray([toAddJkvJ for i in range(length)])
                JKvJ = np.resize(JKvJ, (1, 49*length))[0]
                """

                # [traj, Kp, Kv, F, use_force, seconds]
                
                Kps = np.resize(Kps, (1, 36 * num_steps))[0]
                Kvs = np.resize(Kvs, (1, 36 * num_steps))[0]

                
                JKpJ = np.resize(JKpJ, (1, 49 * num_steps))[0]
                JKvJ = np.resize(JKvJ, (1, 49 * num_steps))[0]
                
                data = np.zeros((1, len(qs) + len(JKpJ) + len(JKvJ) + len(F) + len(Kps) + len(Kvs) + 2))
                data[0][0:len(qs)] = qs
                data[0][len(qs):len(qs)+len(JKpJ)] = JKpJ
                data[0][len(qs)+len(JKpJ):len(qs)+len(JKpJ)+len(JKvJ)] = JKvJ
                data[0][len(qs)+len(JKpJ)+len(JKvJ):len(qs)+len(JKpJ)+len(JKvJ)+len(F)] = F
                data[0][len(qs)+len(JKpJ)+len(JKvJ)+len(F):len(qs)+len(JKpJ)+len(JKvJ)+len(F)+len(Kps)] = Kps
                data[0][len(qs)+len(JKpJ)+len(JKvJ)+len(F)+len(Kps):len(qs)+len(JKpJ)+len(JKvJ)+len(F)+len(Kps)+len(Kvs)] = Kvs
                data[0][-2] = args.use_force
                data[0][-1] = stamps[i_end] - stamps[i_start]
                msg = Float64MultiArray()
                msg.data = data[0].tolist()
                pub = rospy.Publisher("/controller_data", Float64MultiArray)
                redprint("Press enter to start trajectory")
                raw_input()
                time.sleep(1)
                pub.publish(msg)
                time.sleep(1)
            else:
                #if not success: break
                
                if len(bodypart2traj) > 0:
                    exec_traj_maybesim(bodypart2traj)
Example #34
0
def tps_nr_fit_enhanced(x_na, y_ng, bend_coef, nr_ma, nr_coef, plotting=0):

    N, D = x_na.shape
    M = nr_ma.shape[0]
    E = N + 4 * M
    F = E - M
    Q = N + 3 * M - 4

    s = .1  # tetrahedron sidelength (meters)
    u = 1 / (2 * np.sqrt(2))

    tetra_pts = []
    for pt in nr_ma:
        tetra_pts.append(s * np.r_[-.5, 0, -u] + pt)
        tetra_pts.append(s * np.r_[+.5, 0, -u] + pt)
        tetra_pts.append(s * np.r_[0, -.5, +u] + pt)
        tetra_pts.append(s * np.r_[0, +.5, +u] + pt)

    x_ea = np.r_[x_na, tetra_pts]

    badsub_ex = np.c_[x_ea,
                      np.ones((E, 1)), np.r_[np.zeros((N, M)),
                                             np.repeat(np.eye(M), 4, axis=0)]]
    lin_ag, trans_g, w_ng = tps_fit2(x_na, y_ng, bend_coef, 1e-3)
    w_eg = np.r_[w_ng, np.zeros((4 * M, D))]

    assert badsub_ex.shape[0] >= badsub_ex.shape[1]
    _u, _s, _vh = np.linalg.svd(badsub_ex, full_matrices=True)
    assert badsub_ex.shape[1] == _s.size
    N_eq = _u[:, badsub_ex.shape[1]:]  # null of data

    assert N_eq.shape == (E, Q)

    assert E == N + 4 * M
    assert F == Q + 4
    # e is number of kernels
    # q is number of nonrigid dofs
    # f is total number of dofs
    K_ee = ssd.squareform(ssd.pdist(x_ea))
    K_ne = K_ee[:N, :]
    Q_nf = np.c_[x_na, np.ones((N, 1)), K_ne.dot(N_eq)]
    QQ_ff = np.dot(Q_nf.T, Q_nf)
    Bend_ff = np.zeros((F, F))
    Bend_ff[4:, 4:] = -N_eq.T.dot(K_ee.dot(N_eq))  # K_qq

    assert Q_nf.shape == (N, F)
    assert w_eg.shape == (E, D)

    n_iter = 40
    for i in xrange(n_iter):

        # if plotting and i%plotting==0:
        # import lfd.registration as lr
        # lr.Globals.setup()
        # def eval_partial(x_ma):
        #     return tps_eval(x_ma, lin_ag, trans_g, w_eg, x_ea)
        # lr.plot_orig_and_warped_clouds(eval_partial, x_na, y_ng, res=.008)

        X_fg = np.r_[lin_ag, trans_g[None, :], N_eq.T.dot(w_eg)]

        res_cost, bend_cost, nr_cost, fval = tps_nr_cost_eval_general(
            lin_ag,
            trans_g,
            w_eg,
            x_ea,
            y_ng,
            nr_ma,
            bend_coef,
            nr_coef,
            return_tuple=True)
        if VERBOSE:
            print colorize("iteration %i, cost %.3e" % (i, fval), 'red'),
        if VERBOSE:
            print "= %.3e (res) + %.3e (bend) + %.3e (nr)" % (
                res_cost, bend_cost, nr_cost)

        Jl_zcg, Jt_zg, Jw_zeg = tps_nr_grad(nr_ma,
                                            lin_ag,
                                            trans_g,
                                            w_eg,
                                            x_ea,
                                            return_tuple=True)
        nrerr_z = tps_nr_err(nr_ma, lin_ag, trans_g, w_eg, x_ea)

        fullstep_fg = np.empty((F, D))
        for g in xrange(D):
            J_zf = np.c_[Jl_zcg[:, g::D], Jt_zg[:, g::D],
                         Jw_zeg[:, g::D].dot(N_eq)]
            JJ_ff = np.dot(J_zf.T, J_zf)
            A_ff = nr_coef * JJ_ff + QQ_ff + bend_coef * Bend_ff
            X0 = X_fg[:, g]
            B_f = nr_coef * np.dot(J_zf.T,
                                   np.dot(J_zf, X0) - nrerr_z) + Q_nf.T.dot(
                                       y_ng[:, g])
            fullstep_fg[:, g] = np.linalg.solve(A_ff, B_f) - X_fg[:, g]

        cost_improved = False
        for stepsize in 3.**np.arange(0, -10, -1):
            cand_X_fg = X_fg + fullstep_fg * stepsize
            cand_lin_ag, cand_trans_g, cand_w_eg = cand_X_fg[:D], cand_X_fg[
                D], N_eq.dot(cand_X_fg[D + 1:])
            fval_cand = tps_nr_cost_eval_general(cand_lin_ag,
                                                 cand_trans_g,
                                                 cand_w_eg,
                                                 x_ea,
                                                 y_ng,
                                                 nr_ma,
                                                 bend_coef,
                                                 nr_coef,
                                                 return_tuple=False)
            if VERBOSE:
                print "stepsize: %.1g, fval: %.3e" % (stepsize, fval_cand)
            if fval_cand < fval:
                cost_improved = True
                break
        if not cost_improved:
            if VERBOSE: print "couldn't improve objective"
            break

        lin_ag = cand_lin_ag
        trans_g = cand_trans_g
        w_eg = cand_w_eg
    return lin_ag, trans_g, w_eg, x_ea
Example #35
0
def tps_rpm_regrot_multi(x_clouds, y_clouds,
                   x_aug=None, y_aug=None,
                   n_iter = 100, 
                   n_iter_powell_init=100, n_iter_powell_final=100, 
                   bend_init = 0.05, bend_final = .0001, 
                   rot_init = (0.1,0.1,0.025), rot_final=(0.001,0.001,0.00025),
                   scale_init=1, scale_final=0.001, 
                   rad_init = .5, rad_final = .0005,
                   verbose=True, f_init = None, return_full = False,
                   plotting_cb=None, plotter=None):
    """
    x_aug : dict of matrices of extra coordinates for x_clouds. The key is the index of the cloud.
    y_aug : similar to x_aug for y_clouds
    
    
    Similar to tps_rpm_regrot except that it accepts a 
    LIST of source and target point clouds and registers 
    a cloud in the source to the corresponding one in the target.  
    
    For details on the various parameters check the doc of tps_rpm_regrot.
    """

    assert len(x_clouds)==len(y_clouds), "Different number of point-clouds in source and target."

    if x_aug==None or y_aug==None:
        x_aug = y_aug = {}
 
 
    #flatten the list of point clouds into one big point cloud
    combined_x = np.concatenate(x_clouds) 
    combined_y = np.concatenate(y_clouds)

    # concatenate the clouds into one big cloud
    _,d  = combined_x.shape

    regs     = loglinspace(bend_init, bend_final, n_iter)
    rads     = loglinspace(rad_init, rad_final, n_iter)
    scales   = loglinspace(scale_init, scale_final, n_iter)
    rots     = loglinspace_arr(rot_init, rot_final, n_iter)
    npowells = loglinspace(n_iter_powell_init, n_iter_powell_final, n_iter).astype(int) 

    # initialize the function f.
    if f_init is not None: 
        f = f_init  
    else:
        f         = ThinPlateSpline(d)
        f.trans_g = np.median(combined_y,axis=0) - np.median(combined_x,axis=0)


    # iterate b/w calculating correspondences and fitting the transformation.
    for i in xrange(n_iter):
        target_pts   = []
        good_inds    = []
        wt           = []

        for j in xrange(len(x_clouds)): #process a pair of point-clouds
            x_nd = x_clouds[j]
            y_md = y_clouds[j]
            
            assert x_nd.ndim==y_md.ndim==2, "tps_rpm_reg_rot_multi : Point clouds are not two dimensional arrays"
                        
            xwarped_nd = f.transform_points(x_nd)
            
            # use augmented coordinates.
            if x_aug.has_key(j) and y_aug.has_key(j):
                corr_nm = calc_correspondence_matrix(np.c_[xwarped_nd, x_aug[j]], np.c_[y_md, y_aug[j]], r=rads[i], p=.2)
            else:
                corr_nm = calc_correspondence_matrix(xwarped_nd, y_md, r=rads[i], p=.2)

            wt_n = corr_nm.sum(axis=1) # gives the row-wise sum of the corr_nm matrix
            goodn = wt_n > 0.1
            targ_Nd = np.dot(corr_nm[goodn, :]/wt_n[goodn][:,None], y_md) # calculate the average points based on softmatching

            target_pts.append(targ_Nd)
            good_inds.append(goodn)  
            wt.append(wt_n[goodn])

        target_pts = np.concatenate(target_pts)
        good_inds  = np.concatenate(good_inds)
        source_pts = combined_x[good_inds]
        wt         = np.concatenate(wt)

        assert len(target_pts)==len(source_pts)==len(wt), "Lengths are not equal. Error!"
        f = fit_ThinPlateSpline_RotReg(source_pts, target_pts, wt_n=wt, bend_coef = regs[i], rot_coefs = rots[i], scale_coef=scales[i], niter_powell=npowells[i])
        

        mscore = match_score(f.transform_points(source_pts), target_pts)
        tscore = tps.tps_cost(f.lin_ag, f.trans_g, f.w_ng, source_pts, target_pts, regs[-1])
        print colorize("\ttps-rpm-rotreg : iter : %d | fit distance : "%i, "red") , colorize("%g"%mscore, "green"), colorize(" | tps score: %g"%tscore, "blue")

        if plotting_cb and i%5==0:
            plotting_cb(f)

        # just plots the "target_pts" : the matched up points found by the correspondences.
        if plotter:
            plotter.request(gen_mlab_request(mlab.points3d, target_pts[:,0], target_pts[:,1], target_pts[:,2], color=(1,1,0), scale_factor=0.001))

        # return if source and target match up well
        if mscore < 0.01:
            break

    if return_full:
        info = {}
        info["cost"] = tps.tps_cost(f.lin_ag, f.trans_g, f.w_ng, source_pts, target_pts, regs[-1])
        return f, info
    else:
        return f
Example #36
0
def main():

    demofile = h5py.File(args.h5file, 'r')

    trajoptpy.SetInteractive(args.interactive)

    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:
        subprocess.call("killall XnSensorServer", shell=True)
        grabber = cloudprocpy.CloudGrabber()
        grabber.startRGBD()

    Globals.viewer = trajoptpy.GetViewer(Globals.env)

    #####################

    while True:

        redprint("Acquire point cloud")
        if args.fake_data_segment:
            new_xyz = np.squeeze(demofile[args.fake_data_segment]["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, :]

        else:

            Globals.pr2.rarm.goto_posture('side')
            Globals.pr2.larm.goto_posture('side')
            Globals.pr2.join_all()

            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)

        ################################
        redprint("Finding closest demonstration")
        if args.fake_data_segment:
            seg_name = args.fake_data_segment
        elif args.choice == "costs":
            f, seg_name = choose_segment(demofile, new_xyz)
        else:
            seg_name = choose_segment(demofile, new_xyz)

        seg_info = demofile[seg_name]
        # redprint("using demo %s, description: %s"%(seg_name, seg_info["description"]))

        ################################

        redprint("Generating end-effector trajectory")
        if not args.choice == "costs":
            handles = []
            old_xyz = np.squeeze(seg_info["cloud_xyz"])
            handles.append(Globals.env.plot3(old_xyz, 5, (1, 0, 0, 1)))
            handles.append(Globals.env.plot3(new_xyz, 5, (0, 0, 1, 1)))
            f = registration.tps_rpm(old_xyz,
                                     new_xyz,
                                     plot_cb=tpsrpm_plot_cb,
                                     plotting=1)

            handles.extend(
                plotting_openrave.draw_grid(Globals.env,
                                            f.transform_points,
                                            old_xyz.min(axis=0),
                                            old_xyz.max(axis=0),
                                            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
            if not args.choice == "costs":
                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)))

        # TODO plot
        # plot_warping_and_trajectories(f, old_xyz, new_xyz, old_ee_traj, new_ee_traj)

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

            bodypart2traj = {}

            arms_used = ""

            for lr in 'lr':
                manip_name = {"l": "leftarm", "r": "rightarm"}[lr]
                old_joint_traj = asarray(seg_info[manip_name][i_start:i_end +
                                                              1])
                if arm_moved(old_joint_traj):
                    ee_link_name = "%s_gripper_tool_frame" % lr
                    new_ee_traj = link2eetraj[ee_link_name]
                    if args.execution: Globals.pr2.update_rave()
                    new_joint_traj = plan_follow_traj(
                        Globals.robot, manip_name,
                        Globals.robot.GetLink(ee_link_name),
                        new_ee_traj[i_start:i_end + 1], old_joint_traj)
                    # (robot, manip_name, ee_link, new_hmats, old_traj):
                    part_name = {"l": "larm", "r": "rarm"}[lr]
                    bodypart2traj[part_name] = new_joint_traj
                    arms_used += lr

            ################################
            redprint(
                "Executing joint trajectory for segment %s, part %i using arms '%s'"
                % (seg_name, i_miniseg, arms_used))

            for lr in 'lr':
                set_gripper_maybesim(
                    lr,
                    binarize_gripper(seg_info["%s_gripper_joint" %
                                              lr][i_start]))
            #trajoptpy.GetViewer(Globals.env).Idle()

            if len(bodypart2traj) > 0:
                exec_traj_maybesim(bodypart2traj)

            # TODO measure failure condtions

            if not success:
                break

        redprint("Segment %s result: %s" % (seg_name, success))

        if args.fake_data_segment: break
Example #37
0
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 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
Example #39
0
def redprint(msg):
    print colorize.colorize(msg, "red", bold=True)
Example #40
0
def blueprint(msg):
    print colorize.colorize(msg, "blue", bold=True)
Example #41
0
def yellowprint(msg):
    print colorize.colorize(msg, "yellow", bold=True)
Example #42
0
def yellowprint(msg):
    print colorize.colorize(msg, "yellow", bold=True)