Beispiel #1
0
def interp_hmats(newtimes, oldtimes, oldhmats):
    oldposes = openravepy.poseFromMatrices(oldhmats)
    newposes = np.empty((len(newtimes), 7))
    newposes[:,4:7] = mu.interp2d(newtimes, oldtimes, oldposes[:,4:7])
    newposes[:,0:4] = interp_quats(newtimes, oldtimes, oldposes[:,0:4])
    newhmats = openravepy.matrixFromPoses(newposes)
    return newhmats
def make_joint_traj(xyzs, quats, manip, ref_frame, targ_frame, filter_options=0):
    "do ik and then fill in the points where ik failed"

    n = len(xyzs)
    assert len(quats) == n

    robot = manip.GetRobot()
    joint_inds = manip.GetArmIndices()
    robot.SetActiveDOFs(joint_inds)
    orig_joint = robot.GetActiveDOFValues()

    joints = []
    inds = []

    for i in xrange(0, n):
        mat4 = conv.trans_rot_to_hmat(xyzs[i], quats[i])
        joint = PR2.cart_to_joint(manip, mat4, ref_frame, targ_frame, filter_options)
        if joint is not None:
            joints.append(joint)
            inds.append(i)
            robot.SetActiveDOFValues(joint)

    robot.SetActiveDOFValues(orig_joint)

    rospy.loginfo("found ik soln for %i of %i points", len(inds), n)
    if len(inds) > 2:
        joints2 = mu.interp2d(np.arange(n), inds, joints)
        return joints2, inds
    else:
        return np.zeros((n, len(joints))), []
Beispiel #3
0
def make_joint_traj(xyzs, quats, manip, ref_frame, targ_frame, filter_options = 0):
    "do ik and then fill in the points where ik failed"

    n = len(xyzs)
    assert len(quats) == n
    
    robot = manip.GetRobot()
    joint_inds = manip.GetArmIndices()
    robot.SetActiveDOFs(joint_inds)
    orig_joint = robot.GetActiveDOFValues()
    
    joints = []
    inds = []

    for i in xrange(0,n):
        mat4 = conv.trans_rot_to_hmat(xyzs[i], quats[i])
        joint = PR2.cart_to_joint(manip, mat4, ref_frame, targ_frame, filter_options)
        if joint is not None: 
            joints.append(joint)
            inds.append(i)
            robot.SetActiveDOFValues(joint)


    robot.SetActiveDOFValues(orig_joint)
    
    
    LOG.info("found ik soln for %i of %i points",len(inds), n)
    if len(inds) > 2:
        joints2 = mu.interp2d(np.arange(n), inds, joints)
        return joints2, inds
    else:
        return np.zeros((n, len(joints))), []
Beispiel #4
0
def adaptive_resample(x, tol, max_change=None, min_steps=3):
    """
    resample original signal with a small number of waypoints so that the the sparsely sampled function, 
    when linearly interpolated, deviates from the original function by less than tol at every time

    input:
    x: 2D array in R^(t x k)  where t is the number of timesteps
    tol: tolerance. either a single scalar or a vector of length k
    max_change: max change in the sparsely sampled signal at each timestep
    min_steps: minimum number of timesteps in the new trajectory. (usually irrelevant)

    output:
    new_times, new_x

    assuming that the old signal has times 0,1,2,...,len(x)-1
    this gives the new times, and the new signal
    """
    x = np.asarray(x)
    assert x.ndim == 2

    if np.isscalar(tol): 
        tol = np.ones(x.shape[1])*tol
    else:
        tol = np.asarray(tol)
        assert tol.ndim == 1 and tol.shape[0] == x.shape[1]

    times = np.arange(x.shape[0])

    if max_change is None: 
        max_change = np.ones(x.shape[1]) * np.inf
    elif np.isscalar(max_change): 
        max_change = np.ones(x.shape[1]) * max_change
    else:
        max_change = np.asarray(max_change)
        assert max_change.ndim == 1 and max_change.shape[0] == x.shape[1]

    dl = mu.norms(x[1:] - x[:-1],1)
    l = np.cumsum(np.r_[0,dl])

    def bad_inds(x1, t1):
        ibad = np.flatnonzero( (np.abs(mu.interp2d(l, l1, x1) - x) > tol).any(axis=1) )
        jbad1 = np.flatnonzero((np.abs(x1[1:] - x1[:-1]) > max_change[None,:]).any(axis=1))
        if len(ibad) == 0 and len(jbad1) == 0: return []
        else:
            lbad = l[ibad]
            jbad = np.unique(np.searchsorted(l1, lbad)) - 1
            jbad = np.union1d(jbad, jbad1)
            return jbad

    l1 = np.linspace(0,l[-1],min_steps)
    for _ in xrange(20):
        x1 = mu.interp2d(l1, l, x)
        bi = bad_inds(x1, l1)
        if len(bi) == 0:
            return np.interp(l1, l, times), x1
        else:
            l1 = np.union1d(l1, (l1[bi] + l1[bi+1]) / 2 )


    raise Exception("couldn't subdivide enough. something funny is going on. check your input data")
Beispiel #5
0
def traj_collisions(sim_env, full_traj, collision_dist_threshold, n=100):
    """
    Returns the set of collisions. 
    manip = Manipulator or list of indices
    """
    traj, dof_inds = full_traj
    
    traj_up = mu.interp2d(np.linspace(0,1,n), np.linspace(0,1,len(traj)), traj)
    cc = trajoptpy.GetCollisionChecker(sim_env.env)

    with openravepy.RobotStateSaver(sim_env.robot):
        sim_env.robot.SetActiveDOFs(dof_inds)
    
        col_times = []
        for (i,row) in enumerate(traj_up):
            sim_env.robot.SetActiveDOFValues(row)
            col_now = cc.BodyVsAll(sim_env.robot)
            #with util.suppress_stdout():
            #    col_now2 = cc.PlotCollisionGeometry()
            col_now = [cn for cn in col_now if cn.GetDistance() < collision_dist_threshold]
            if col_now:
                #print [cn.GetDistance() for cn in col_now]
                col_times.append(i)
                #print "trajopt.CollisionChecker: ", len(col_now)
            #print col_now2
        
    return col_times
    def observe_cloud2(self, radius, axis_upsample=0, radius_sample=0, angle_sample=0):
        axis_pts = self.rope.GetControlPoints()
        indices = range(len(axis_pts))
        if axis_upsample != 0:
            lengths = np.r_[0, self.rope.GetHalfHeights() * 2]
            summed_lengths = np.cumsum(lengths)
            assert len(lengths) == len(axis_pts)
            indices = np.interp(np.linspace(0, summed_lengths[-1], axis_upsample*len(lengths)), summed_lengths, indices)
            axis_pts = math_utils.interp2d(np.linspace(0, summed_lengths[-1], axis_upsample*len(lengths)), summed_lengths, axis_pts)
        
        half_heights = self.rope.GetHalfHeights()
        rotates = self.rope.GetRotations()
        translates = self.rope.GetTranslations()
        new_pts = []
        for r in range(1, radius_sample+1):
            d = r * radius / float(radius_sample) 
            for a in range(angle_sample):
                local_rotate = matrixFromAxisAngle([a/float(angle_sample)*2*np.pi,0,0])[:3,:3]
                for index in indices:
                    node_id = np.min([np.floor(index), len(self.rope.GetNodes()) - 1])
                    h = half_heights[node_id]
                    v = np.array([-h + 2*h*(index-node_id),d,0])
                    rotate = rotates[node_id]
                    translate = translates[node_id]
                    v = rotate.dot(local_rotate.dot(v)) + translate
                    new_pts.append(v)

        for pt in axis_pts:
            new_pts.append(pt)
                
        return np.array(new_pts)
Beispiel #7
0
def interp_hmats(newtimes, oldtimes, oldhmats):
    oldposes = openravepy.poseFromMatrices(oldhmats)
    newposes = np.empty((len(newtimes), 7))
    newposes[:,4:7] = mu.interp2d(newtimes, oldtimes, oldposes[:,4:7])
    newposes[:,0:4] = interp_quats(newtimes, oldtimes, oldposes[:,0:4])
    newhmats = openravepy.matrixFromPoses(newposes)
    return newhmats
Beispiel #8
0
def traj_collisions(sim_env, full_traj, collision_dist_threshold, n=100):
    """
    Returns the set of collisions. 
    manip = Manipulator or list of indices
    """
    traj, dof_inds = full_traj

    traj_up = mu.interp2d(np.linspace(0, 1, n), np.linspace(0, 1, len(traj)),
                          traj)
    cc = trajoptpy.GetCollisionChecker(sim_env.env)

    with openravepy.RobotStateSaver(sim_env.robot):
        sim_env.robot.SetActiveDOFs(dof_inds)

        col_times = []
        for (i, row) in enumerate(traj_up):
            sim_env.robot.SetActiveDOFValues(row)
            col_now = cc.BodyVsAll(sim_env.robot)
            #with util.suppress_stdout():
            #    col_now2 = cc.PlotCollisionGeometry()
            col_now = [
                cn for cn in col_now
                if cn.GetDistance() < collision_dist_threshold
            ]
            if col_now:
                #print [cn.GetDistance() for cn in col_now]
                col_times.append(i)
                #print "trajopt.CollisionChecker: ", len(col_now)
            #print col_now2

    return col_times
 def observe_cloud(self, upsample=0):
     pts = self.rope.GetControlPoints()
     if upsample == 0:
         return pts
     lengths = np.r_[0, self.rope.GetHalfHeights() * 2]
     summed_lengths = np.cumsum(lengths)
     assert len(lengths) == len(pts)
     return math_utils.interp2d(np.linspace(0, summed_lengths[-1], upsample), summed_lengths, pts)
Beispiel #10
0
 def bad_inds(x1, t1):
     ibad = np.flatnonzero( (np.abs(mu.interp2d(l, l1, x1) - x) > tol).any(axis=1) )
     jbad1 = np.flatnonzero((np.abs(x1[1:] - x1[:-1]) > max_change[None,:]).any(axis=1))
     if len(ibad) == 0 and len(jbad1) == 0: return []
     else:
         lbad = l[ibad]
         jbad = np.unique(np.searchsorted(l1, lbad)) - 1
         jbad = np.union1d(jbad, jbad1)
         return jbad
Beispiel #11
0
    def follow_joint_trajectory(self, traj):
        traj = np.r_[np.atleast_2d(self.get_joint_positions()), traj]
        for i in [2,4,6]:
            traj[:,i] = np.unwrap(traj[:,i])

        times = retiming.retime_with_vel_limits(traj, self.vel_limits)
        times_up = np.arange(0,times[-1],.1)
        traj_up = mu.interp2d(times_up, times, traj)
        vels = ku.get_velocities(traj_up, times_up, .001)
        self.follow_timed_joint_trajectory(traj_up, vels, times_up)
Beispiel #12
0
    def follow_joint_trajectory(self, traj):
        traj = np.r_[np.atleast_2d(self.get_joint_positions()), traj]
        for i in [2,4,6]:
            traj[:,i] = np.unwrap(traj[:,i])

        times = retiming.retime_with_vel_limits(traj, self.vel_limits)
        times_up = np.arange(0,times[-1],.1)
        traj_up = mu.interp2d(times_up, times, traj)
        vels = resampling.get_velocities(traj_up, times_up, .001)
        self.follow_timed_joint_trajectory(traj_up, vels, times_up)
    def observe_cloud(self, upsample=0):
        pts = self.rope.GetControlPoints()
        if upsample == 0:
            return pts
        lengths = np.r_[0, self.rope.GetHalfHeights() * 2]
        summed_lengths = np.cumsum(lengths)
        assert len(lengths) == len(pts)
        upsampled_pts = math_utils.interp2d(
            np.linspace(0, summed_lengths[-1], upsample * len(lengths)),
            summed_lengths, pts)

        return upsampled_pts
def interp_hmats(newtimes, oldtimes, oldhmats):
    oldposes = openravepy.poseFromMatrices(oldhmats)
    for i in xrange(1, len(oldposes)):
        qp = oldposes[i-1,0:4]
        qc = oldposes[i,0:4]
        if np.linalg.norm(qp + qc) < np.linalg.norm(qp-qc):
            oldposes[i,0:4] = -qc

    newposes = np.empty((len(newtimes), 7))
    newposes[:,4:7] = mu.interp2d(newtimes, oldtimes, oldposes[:,4:7])
    newposes[:,0:4] = interp_quats(newtimes, oldtimes, oldposes[:,0:4])
    newhmats = openravepy.matrixFromPoses(newposes)
    return newhmats
def retime_traj(robot, inds, traj, max_cart_vel=.02, upsample_time=.1):
    """retime a trajectory so that it executes slowly enough for the simulation"""
    cart_traj = np.empty((len(traj), 6))
    leftarm, rightarm = robot.GetManipulator("leftarm"), robot.GetManipulator("rightarm")
    with robot:
        for i in range(len(traj)):
            robot.SetDOFValues(traj[i], inds)
            cart_traj[i,:3] = leftarm.GetTransform()[:3,3]
            cart_traj[i,3:] = rightarm.GetTransform()[:3,3]

    times = retiming.retime_with_vel_limits(cart_traj, np.repeat(max_cart_vel, 6))
    times_up = np.linspace(0, times[-1], times[-1]/upsample_time) if times[-1] > upsample_time else times
    traj_up = math_utils.interp2d(times_up, times, traj)
    return traj_up
    def observe_cloud2(self,
                       radius,
                       axis_upsample=0,
                       radius_sample=0,
                       angle_sample=0):
        axis_pts = self.rope.GetControlPoints()
        indices = range(len(axis_pts))
        if axis_upsample != 0:
            lengths = np.r_[0, self.rope.GetHalfHeights() * 2]
            summed_lengths = np.cumsum(lengths)
            assert len(lengths) == len(axis_pts)
            indices = np.interp(
                np.linspace(0, summed_lengths[-1],
                            axis_upsample * len(lengths)), summed_lengths,
                indices)
            axis_pts = math_utils.interp2d(
                np.linspace(0, summed_lengths[-1],
                            axis_upsample * len(lengths)), summed_lengths,
                axis_pts)

        half_heights = self.rope.GetHalfHeights()
        rotates = self.rope.GetRotations()
        translates = self.rope.GetTranslations()
        new_pts = []
        for r in range(1, radius_sample + 1):
            d = r * radius / float(radius_sample)
            for a in range(angle_sample):
                local_rotate = matrixFromAxisAngle(
                    [a / float(angle_sample) * 2 * np.pi, 0, 0])[:3, :3]
                for index in indices:
                    node_id = np.min(
                        [np.floor(index),
                         len(self.rope.GetNodes()) - 1])
                    h = half_heights[node_id]
                    v = np.array([-h + 2 * h * (index - node_id), d, 0])
                    rotate = rotates[node_id]
                    translate = translates[node_id]
                    v = rotate.dot(local_rotate.dot(v)) + translate
                    new_pts.append(v)

        for pt in axis_pts:
            new_pts.append(pt)

        return np.array(new_pts)
Beispiel #17
0
def exec_traj_maybesim(bodypart2traj, speed_factor=0.5):
    if args.animation:
        """
        dof_inds = []
        trajs = []
        for (part_name, traj) in bodypart2traj.items():
            manip_name = {"larm":"leftarm","rarm":"rightarm"}[part_name]
            dof_inds.extend(Globals.robot.GetManipulator(manip_name).GetArmIndices())            
            trajs.append(traj)
        full_traj = np.concatenate(trajs, axis=1)
        Globals.robot.SetActiveDOFs(dof_inds)
        animate_traj.animate_traj(full_traj, Globals.robot, restore=False,pause=True)
        """
        name2part = {"lgrip":Globals.pr2.lgrip,
                     "rgrip":Globals.pr2.rgrip,
                     "larm":Globals.pr2.larm,
                     "rarm":Globals.pr2.rarm,
                     "base":Globals.pr2.base}
        dof_inds = []
        trajs = []
        vel_limits = []
        for (part_name, traj) in bodypart2traj.items():
            manip_name = {"larm":"leftarm","rarm":"rightarm"}[part_name]
            vel_limits.extend(name2part[part_name].vel_limits)
            dof_inds.extend(Globals.robot.GetManipulator(manip_name).GetArmIndices())
            if traj.ndim == 1: traj = traj.reshape(-1,1)            
            trajs.append(traj)
        
        trajectories = np.concatenate(trajs, 1)
        print trajectories.shape
        times = retiming.retime_with_vel_limits(trajectories, np.array(vel_limits))
        times_up = np.linspace(0, times[-1], int(np.ceil(times[-1]/.1)))
        full_traj = mu.interp2d(times_up, times, trajectories)
        print full_traj.shape
        
        Globals.robot.SetActiveDOFs(dof_inds)
        animate_traj.animate_traj(full_traj, Globals.robot, restore=False,pause=True)
        #return True
    if args.execution:
        pr2_trajectories.follow_body_traj(Globals.pr2, bodypart2traj, speed_factor=speed_factor)
        return True
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
Beispiel #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
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)
Beispiel #21
0
def follow_body_traj(pr2,
                     bodypart2traj,
                     wait=True,
                     base_frame="/base_footprint",
                     speed_factor=1):

    rospy.loginfo("following trajectory with bodyparts %s",
                  " ".join(bodypart2traj.keys()))

    name2part = {
        "lgrip": pr2.lgrip,
        "rgrip": pr2.rgrip,
        "larm": pr2.larm,
        "rarm": pr2.rarm,
        "base": pr2.base
    }
    for partname in bodypart2traj:
        if partname not in name2part:
            raise Exception("invalid part name %s" % partname)

    #### Go to initial positions #######

    for (name, part) in name2part.items():
        if name in bodypart2traj:
            part_traj = bodypart2traj[name]
            if name == "lgrip" or name == "rgrip":
                part.set_angle(np.squeeze(part_traj)[0])
            elif name == "larm" or name == "rarm":
                part.goto_joint_positions(part_traj[0])
            elif name == "base":
                part.goto_pose(part_traj[0], base_frame)
    pr2.join_all()

    #### Construct total trajectory so we can retime it #######

    n_dof = 0
    trajectories = []
    vel_limits = []
    acc_limits = []
    bodypart2inds = {}
    for (name, part) in name2part.items():
        if name in bodypart2traj:
            traj = bodypart2traj[name]
            if traj.ndim == 1: traj = traj.reshape(-1, 1)
            trajectories.append(traj)
            vel_limits.extend(part.vel_limits)
            acc_limits.extend(part.acc_limits)
            bodypart2inds[name] = range(n_dof, n_dof + part.n_joints)
            n_dof += part.n_joints

    trajectories = np.concatenate(trajectories, 1)

    vel_limits = np.array(vel_limits) * speed_factor

    times = retiming.retime_with_vel_limits(trajectories, vel_limits)
    times_up = np.linspace(0, times[-1], int(np.ceil(times[-1] / .1)))
    traj_up = mu.interp2d(times_up, times, trajectories)

    #### Send all part trajectories ###########

    for (name, part) in name2part.items():
        if name in bodypart2traj:
            part_traj = traj_up[:, bodypart2inds[name]]
            if name == "lgrip" or name == "rgrip":
                part.follow_timed_trajectory(times_up, part_traj.flatten())
            elif name == "larm" or name == "rarm":
                vels = ku.get_velocities(part_traj, times_up, .001)
                part.follow_timed_joint_trajectory(part_traj, vels, times_up)
            elif name == "base":
                part.follow_timed_trajectory(times_up, part_traj, base_frame)

    if wait: pr2.join_all()

    return True
def follow_body_traj(pr2, bodypart2traj, wait=True, base_frame="/base_footprint", speed_factor=1):

    rospy.loginfo("following trajectory with bodyparts %s", " ".join(bodypart2traj.keys()))

    name2part = {"lgrip": pr2.lgrip, "rgrip": pr2.rgrip, "larm": pr2.larm, "rarm": pr2.rarm, "base": pr2.base}
    for partname in bodypart2traj:
        if partname not in name2part:
            raise Exception("invalid part name %s" % partname)

    #### Go to initial positions #######

    for (name, part) in name2part.items():
        if name in bodypart2traj:
            part_traj = bodypart2traj[name]
            if name == "lgrip" or name == "rgrip":
                part.set_angle(np.squeeze(part_traj)[0])
            elif name == "larm" or name == "rarm":
                part.goto_joint_positions(part_traj[0])
            elif name == "base":
                part.goto_pose(part_traj[0], base_frame)
    pr2.join_all()

    #### Construct total trajectory so we can retime it #######

    n_dof = 0
    trajectories = []
    vel_limits = []
    acc_limits = []
    bodypart2inds = {}
    for (name, part) in name2part.items():
        if name in bodypart2traj:
            traj = bodypart2traj[name]
            if traj.ndim == 1:
                traj = traj.reshape(-1, 1)
            trajectories.append(traj)
            vel_limits.extend(part.vel_limits)
            acc_limits.extend(part.acc_limits)
            bodypart2inds[name] = range(n_dof, n_dof + part.n_joints)
            n_dof += part.n_joints

    trajectories = np.concatenate(trajectories, 1)

    vel_limits = np.array(vel_limits) * speed_factor

    times = retiming.retime_with_vel_limits(trajectories, vel_limits)
    times_up = np.linspace(0, times[-1], int(np.ceil(times[-1] / 0.1)))
    traj_up = mu.interp2d(times_up, times, trajectories)

    #### Send all part trajectories ###########

    for (name, part) in name2part.items():
        if name in bodypart2traj:
            part_traj = traj_up[:, bodypart2inds[name]]
            if name == "lgrip" or name == "rgrip":
                part.follow_timed_trajectory(times_up, part_traj.flatten())
            elif name == "larm" or name == "rarm":
                vels = ku.get_velocities(part_traj, times_up, 0.001)
                part.follow_timed_joint_trajectory(part_traj, vels, times_up)
            elif name == "base":
                part.follow_timed_trajectory(times_up, part_traj, base_frame)

    if wait:
        pr2.join_all()

    return True
Beispiel #23
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)

    #####################
    
    started_bag = False
    started_video = False

    localtime   = time.localtime()
    time_string  = time.strftime("%Y-%m-%d-%H-%M-%S", localtime)

    os.chdir(osp.dirname(args.new_demo))


    if not osp.exists(args.new_demo):
        yn = yes_or_no("master file does not exist. create?")
        basename = raw_input("what is the base name?\n").strip()
        if yn:
            with open(args.new_demo,'w') as fh: fh.write("""
    name: %s
    h5path: %s
    bags:
            """%(basename, basename+".h5"))
        else:
            print "exiting."
            exit(0)
    with open(args.new_demo, "r") as fh: master_info = yaml.load(fh)
    if master_info["bags"] is None: master_info["bags"] = []
    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

    timestampfile = demo_name+"timestamps.txt"
    fht = open(timestampfile,"w")
    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)
        time.sleep(1)
        poll_result = bag_handle.poll() 
        print "poll result", poll_result
        if poll_result is not None:
            raise Exception("problem starting video recording")
        else: 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


       

        
        
            
        
            
            #grab_end(new_xyz)
        fht.write('look:' + str(rospy.get_rostime().secs))
        
        
        ################################    
        redprint("Finding closest demonstration")
       
        seg_name = find_closest_manual(demofile, None)
        
        
        seg_info = demofile[seg_name]
        

        redprint("Generating end-effector trajectory")    

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

        scaled_old_xyz, src_params = registration.unit_boxify(old_xyz)
        

        link2eetraj = {}
        for lr in 'lr':
            link_name = "%s_gripper_tool_frame"%lr
            old_ee_traj = asarray(seg_info[link_name]["hmat"])
           
    
        miniseg_starts, miniseg_ends = split_trajectory_by_gripper(seg_info)    
        success = True
        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
                lr2oldtraj[lr] = old_joint_traj   

        
            ### 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] = old_joint_traj
                
            traj = {}

            for lr in 'lr':
                part_name = {"l":"larm", "r":"rarm"}[lr]
                traj[lr] = bodypart2traj[part_name]

            if i_miniseg ==0:
                
                redprint("Press enter to use current position as starting point")
                raw_input()
                arm_positions = {}
                (arm_position, arm_velocity, arm_effort) = robot_states.call_return_joint_states(robot_states.r_arm_joint_names)
                arm_positions['r'] = arm_position
                diff_r = np.array(arm_position - traj['r'][0,:])
                maximum_r = max(abs(diff_r))
                (arm_position, arm_velocity, arm_effort) = robot_states.call_return_joint_states(robot_states.l_arm_joint_names)
                arm_positions['l'] = arm_position
                diff_l = np.array(arm_position - traj['l'][0,:])
                maximum_l = max(abs(diff_l))
                maximum = max(maximum_l, maximum_r)

                speed = (20.0/360.0*2*(np.pi))
                time_needed = maximum / speed
                
                
                initial_pos_traj = {}
                
                for lr in 'lr':
                    part_name = {"l":"larm", "r":"rarm"}[lr]
                    initial_pos_traj[part_name] = mu.interp2d(np.arange(0, time_needed, 0.01), np.array([0,time_needed]), np.array([arm_positions[lr], traj[lr][0,:]]))
                
                #initial_traj_length = initial_pos_traj.shape[0]

                #traj[lr] = np.concatenate((initial_pos_traj, traj_part[lr]), axis=0)
               
                #=======================send to controller======================
                length_total = initial_pos_traj["larm"].shape[0]
                qs_l = np.resize(initial_pos_traj["larm"], (1, initial_pos_traj["larm"].shape[0]*7))[0] #resize the data to 1x7n (n being the number of steps)
                qs_r = np.resize(initial_pos_traj["rarm"], (1, initial_pos_traj["rarm"].shape[0]*7))[0]
                #F = np.array(np.matrix(F).T) # 6 x n
                F = np.zeros((length_total,6))
                F = np.resize(F, (1, F.shape[0]*6))[0] #resize the data to 1x3n
                gripper = np.zeros((length_total,1))
                gripper = np.resize(gripper, (1, gripper.shape[0]*1))[0]
                # Controller code in joint space

                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])
                # Gains as Kps and Kvs for testing

                Kps = []
                Kvs = []
                for i in range(length_total):
                    Kps.append(np.zeros((6,6)))
                    Kvs.append(np.zeros((6,6)))
                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]))
                #toAddJkvJ = np.diag(np.asarray([0, 0, 0, 0, 0, 0, 0]))
                #length = complete_force_traj.shape[0]
                JKpJ = np.asarray([toAddJkpJ for i in range(length_total)])
                JKpJ = np.resize(JKpJ, (1, 49*length_total))[0]

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

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

                
                #JKpJ = np.resize(JKpJ, (1, 49 * num_steps))[0]
                #JKvJ = np.resize(JKvJ, (1, 49 * num_steps))[0]
                stamps = asarray(seg_info['stamps'])
                data = np.zeros((1, length_total*(7+49+49+6+36+36+7+49+49+6+36+36+1)+2))
                data[0][0:length_total*7] = qs_r
                data[0][length_total*7:length_total*(7+49)] = JKpJ
                data[0][length_total*(7+49):length_total*(7+49+49)] = JKvJ
                data[0][length_total*(7+49+49):length_total*(7+49+49+6)] = F
                data[0][length_total*(7+49+49+6):length_total*(7+49+49+6+36)] = Kps
                data[0][length_total*(7+49+49+6+36):length_total*(7+49+49+6+36+36)] = Kvs
                data[0][length_total*(7+49+49+6+36+36):length_total*(7+49+49+6+36+36+7)] = qs_l
                data[0][length_total*(7+49+49+6+36+36+7):length_total*(7+49+49+6+36+36+7+49)] = JKpJ
                data[0][length_total*(7+49+49+6+36+36+7+49):length_total*(7+49+49+6+36+36+7+49+49)] = JKvJ
                data[0][length_total*(7+49+49+6+36+36+7+49+49):length_total*(7+49+49+6+36+36+7+49+49+6)] = F
                data[0][length_total*(7+49+49+6+36+36+7+49+49+6):length_total*(7+49+49+6+36+36+7+49+49+6+36)] = Kps
                data[0][length_total*(7+49+49+6+36+36+7+49+49+6+36):length_total*(7+49+49+6+36+36+7+49+49+6+36+36)] = Kvs
                data[0][length_total*(7+49+49+6+36+36+7+49+49+6+36+36):length_total*(7+49+49+6+36+36+7+49+49+6+36+36+1)] = gripper
                data[0][-2] = 0
                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)


                #===================end send to controller=======================

                raw_input("CAME TO START POSITION")
                time.sleep(1)
                fht.write('\nstart:' + str(rospy.get_rostime().secs))

            ################################    
            redprint("Executing joint trajectory for segment %s, part %i using arms '%s'"%(seg_name, i_miniseg, bodypart2traj.keys()))
            if not args.useHenry:
                
                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
                print('1')
                fht.write('\nstart:' + str(rospy.get_rostime().secs))
                print('2')
                if len(bodypart2traj) > 0:
                    success &= exec_traj_maybesim(bodypart2traj)
                print('3')
                fht.write('\nstop:' + str(rospy.get_rostime().secs))
                print('4')
            else:
                

                for lr in 'lr':
                    redprint("Press enter to set gripper")
                    raw_input()
                    set_gripper_maybesim(lr, binarize_gripper(seg_info["%s_gripper_joint"%lr][i_start]))
                    # Doesn't actually check if grab occurred, unfortunately

                if bodypart2traj!={}:
                    length = i_end - i_start + 1
                    length_total = length
                    traj_r = traj['r']
                    qs_r = np.resize(traj_r, (1, traj_r.shape[0]*7))[0] #resize the data to 1x7n (n being the number of steps)

                    traj_l = traj['l']
                    qs_l = np.resize(traj_l, (1, traj_l.shape[0]*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.zeros((length_total,6))
                    F = np.resize(F, (1, F.shape[0]*6))[0] #resize the data to 1x3n
                    gripper = np.zeros((length_total,1))
                    gripper = np.resize(gripper, (1, gripper.shape[0]*1))[0]
                    # Controller code in joint space

                    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])
                    # Gains as Kps and Kvs for testing

                    Kps = []
                    Kvs = []
                    for i in range(length_total):
                        Kps.append(np.zeros((6,6)))
                        Kvs.append(np.zeros((6,6)))
                    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]))
                    #toAddJkvJ = np.diag(np.asarray([0, 0, 0, 0, 0, 0, 0]))
                    #length = complete_force_traj.shape[0]
                    JKpJ = np.asarray([toAddJkpJ for i in range(length_total)])
                    JKpJ = np.resize(JKpJ, (1, 49*length_total))[0]

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

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

                    
                    #JKpJ = np.resize(JKpJ, (1, 49 * num_steps))[0]
                    #JKvJ = np.resize(JKvJ, (1, 49 * num_steps))[0]
                    stamps = asarray(seg_info['stamps'])
                    data = np.zeros((1, length_total*(7+49+49+6+36+36+7+49+49+6+36+36+1)+2))
                    data[0][0:length_total*7] = qs_r
                    data[0][length_total*7:length_total*(7+49)] = JKpJ
                    data[0][length_total*(7+49):length_total*(7+49+49)] = JKvJ
                    data[0][length_total*(7+49+49):length_total*(7+49+49+6)] = F
                    data[0][length_total*(7+49+49+6):length_total*(7+49+49+6+36)] = Kps
                    data[0][length_total*(7+49+49+6+36):length_total*(7+49+49+6+36+36)] = Kvs
                    data[0][length_total*(7+49+49+6+36+36):length_total*(7+49+49+6+36+36+7)] = qs_l
                    data[0][length_total*(7+49+49+6+36+36+7):length_total*(7+49+49+6+36+36+7+49)] = JKpJ
                    data[0][length_total*(7+49+49+6+36+36+7+49):length_total*(7+49+49+6+36+36+7+49+49)] = JKvJ
                    data[0][length_total*(7+49+49+6+36+36+7+49+49):length_total*(7+49+49+6+36+36+7+49+49+6)] = F
                    data[0][length_total*(7+49+49+6+36+36+7+49+49+6):length_total*(7+49+49+6+36+36+7+49+49+6+36)] = Kps
                    data[0][length_total*(7+49+49+6+36+36+7+49+49+6+36):length_total*(7+49+49+6+36+36+7+49+49+6+36+36)] = Kvs
                    data[0][length_total*(7+49+49+6+36+36+7+49+49+6+36+36):length_total*(7+49+49+6+36+36+7+49+49+6+36+36+1)] = gripper
                    data[0][-2] = 0
                    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)


            #if not success: break


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

        print("exit loop")
        for i in range(100):
            time.sleep(1)
    except KeyboardInterrupt:
        redprint("=================================DONE================================")
        raw_input()
        raw_input()
        raw_input()
        fht.write('\nstop:' + str(rospy.get_rostime().secs))
        fht.write('\ndone:' + str(rospy.get_rostime().secs))
        fht.close()
        time.sleep(3)
        if started_bag:
            print "stopping bag"
            bag_handle.send_signal(signal.SIGINT)
            #bag_handle.wait()
            started_bag = False
        if started_video:
            print "stopping video"
            video_handle.send_signal(signal.SIGINT)
            #video_handle.wait()
            started_video = False
        bagfilename = demo_name+".bag"
        
        annfilename = demo_name+".ann.yaml"
        
        call_and_print("generate_ann.py %s %s %s"%(bagfilename, annfilename, timestampfile),check=False)
        with open(args.new_demo,"a") as fh1:
            fh1.write("\n"
                "- bag_file: %(bagfilename)s\n"
                "  annotation_file: %(annfilename)s\n"
                "  video_dir: %(videodir)s\n"
                "  demo_name: %(demoname)s"%dict(bagfilename=bagfilename, annfilename=annfilename, videodir=demo_name, demoname=demo_name))
        
        return
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)
Beispiel #25
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 compute_dt_code(ctl_pts, plotting=False):
    """Takes rope control points (Nx3 array), closes the loop, and computes the Dowker-Thistlethwaite code for the knot.
       The z-value for the points are used for determining over/undercrossings.
       Follows procedure outlined here: http://katlas.math.toronto.edu/wiki/DT_(Dowker-Thistlethwaite)_Codes
       """

    # First, close the loop by introducing extra points under the table and toward the robot (by subtracting z and x values)
    # first_pt, last_pt =  ctl_pts[0], ctl_pts[-1]
    # flipped = False
    # if first_pt[1] > last_pt[1]:
    #     first_pt, last_pt = last_pt, first_pt
    #     flipped = True
    # min_z = ctl_pts[:,2].min()
    # extra_first_pt, extra_last_pt = first_pt + [-.1, -.1, min_z-1], last_pt + [-.1, .1, min_z-1]
    # if flipped:
    #     extra_pts = [extra_first_pt, extra_first_pt + [-1, 0, 0], extra_last_pt + [-1, 0, 0], extra_last_pt, last_pt]
    # else:
    #     extra_pts = [extra_last_pt, extra_last_pt + [-1, 0, 0], extra_first_pt + [-1, 0, 0], extra_first_pt, first_pt]
    # ctl_pts = np.append(ctl_pts, extra_pts, axis=0)

    if plotting:
        import trajoptpy, openravepy
        env = openravepy.Environment()
        viewer = trajoptpy.GetViewer(env)
        handles = []
        handles.append(env.plot3(ctl_pts, 5, [0, 0, 1]))
        viewer.Idle()

    # Upsampling loop: upsample until every segment has at most one crossing
    need_upsample_ind = None
    upsample_iters = 0
    max_upsample_iters = 10
    while True:
        counter = 1
        crossings = defaultdict(list)
        # Walk along rope: for each segment, compute intersections with all other segments
        for i in range(len(ctl_pts) - 1):
            curr_seg = ctl_pts[i:i+2,:]
            intersections, ts, us = intersect_segs(ctl_pts[:,:2], curr_seg[:,:2])

            if len(intersections) == 0:
                continue
            if len(intersections) != 1:
                LOG.debug('warning: more than one intersection for segment %d, now upsampling', i)
                need_upsample_ind = i
                break

            # for each intersection, determine and record over/undercrossing
            i_int = intersections[0]
            if plotting:
                handles.append(env.drawlinestrip(ctl_pts[i_int:i_int+2], 5, [1, 0, 0]))
            int_point_rope = ctl_pts[i_int] + ts[i_int]*(ctl_pts[i_int+1] - ctl_pts[i_int])
            int_point_curr_seg = curr_seg[0] + us[i_int]*(curr_seg[1] - curr_seg[0])
            #assert np.allclose(int_point_rope[:2], int_point_curr_seg[:2])
            above = int_point_curr_seg[2] > int_point_rope[2]
            crossings[tuple(sorted((i, i_int)))].append(-counter if counter % 2 == 0 and above else counter)
            counter += 1
        if plotting: viewer.Idle()
        # upsample if necessary
        if need_upsample_ind is not None and upsample_iters < max_upsample_iters:
            spacing = np.linspace(0, 1, len(ctl_pts))
            new_spacing = np.insert(spacing, need_upsample_ind+1, (spacing[need_upsample_ind]+spacing[need_upsample_ind+1])/2.)
            ctl_pts = math_utils.interp2d(new_spacing, spacing, ctl_pts)
            upsample_iters += 1
            need_upsample = None
            continue
        break

    # Extract relevant part of crossing data to produce DT code
    out = []
    for pair in crossings.itervalues():
        assert len(pair) == 2
        odd = [p for p in pair if p % 2 == 1][0]
        even = [p for p in pair if p % 2 == 0][0]
        out.append((odd, even))
    out.sort()
    dt_code = [-o[1] for o in out]
    return dt_code
Beispiel #27
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 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 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)
Beispiel #30
0
def interp_quats(newtimes, oldtimes, oldquats):
    "should actually do slerp"
    quats_unnormed = mu.interp2d(newtimes, oldtimes, oldquats)
    return mu.normr(quats_unnormed)
Beispiel #31
0
def interp_quats(newtimes, oldtimes, oldquats):
    "should actually do slerp"
    quats_unnormed = mu.interp2d(newtimes, oldtimes, oldquats)
    return mu.normr(quats_unnormed)
def loop_body(demofile, choose_segment, knot, animate, task_params, curr_step=None):
    """Do the body of the main task execution loop (ie. do a segment).
    Arguments:
        curr_step is 0 indexed
        choose_segment is a function that returns the key in the demofile to the segment
        knot is the knot the rope is checked against
        new_xyz is the new pointcloud
        task_params is used for the only_original_segments argument

    return None or {'found_knot': found_knot, 'segment': segment, 'link2eetraj': link2eetraj, 'new_xyz': new_xyz}
    """
    #TODO -- Return the new trajectory and state info to be used for bootstrapping (knot_success, new_xyz, link2eetraj,
    #TODO segment)

    #TODO -- End condition
    #TODO -- max_segments logic
    redprint("Acquire point cloud")

    move_sim_arms_to_side()
    #TODO -- Possibly refactor this section to be before the loop

    new_xyz = Globals.sim.observe_cloud()
    new_xyz_upsampled = Globals.sim.observe_cloud(upsample=120)

    print "loop_body only_original_segments", task_params.only_original_segments
    segment = choose_segment(demofile, new_xyz, task_params.only_original_segments)
    if segment is None:
        return None
    seg_info = demofile[segment]

    handles = []
    old_xyz = np.squeeze(seg_info["cloud_xyz"])
    handles.append(Globals.env.plot3(new_xyz, 5, (0, 0, 1)))

    #TODO: test that old_xyz is now bigger if from a derived segment
    #old_xyz = clouds.downsample(old_xyz, DS_SIZE)
    if not "derived" in seg_info.keys():
        old_xyz = downsample(old_xyz, DS_SIZE)
        print "derived, so downsamping"
        #new_xyz = clouds.downsample(new_xyz, DS_SIZE)
    #new_xyz = downsample_if_big(new_xyz, DS_SIZE)
    handles.append(Globals.env.plot3(old_xyz, 5, (1, 0, 0)))

    print "new_xyz cloud size", new_xyz.shape
    print "old_xyz cloud size", old_xyz.shape
    scaled_old_xyz, src_params = registration.unit_boxify(old_xyz)
    scaled_new_xyz, targ_params = registration.unit_boxify(new_xyz)
    #TODO: get rid of g
    f, g = registration.tps_rpm_bij(scaled_old_xyz, scaled_new_xyz, plot_cb=tpsrpm_plot_cb,
                                    plotting=5 if animate else 0, rot_reg=np.r_[1e-4, 1e-4, 1e-1], n_iter=50,
                                    reg_init=10, reg_final=.01, old_xyz=old_xyz, new_xyz=new_xyz)
    f = registration.unscale_tps(f, src_params, targ_params)
    g = registration.unscale_tps(g, src_params, targ_params)
    #Globals.exec_log(curr_step, "gen_traj.f", f)

    #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))
    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, .02], xres=.01, yres=.01, zres=.04))
    #handles.extend(plotting_openrave.draw_grid(Globals.env, g.transform_points, old_xyz.min(axis=0) - np.r_[0, 0, .1],
    #                                           old_xyz.max(axis=0) + np.r_[0, 0, .02], xres=.01, yres=.01, zres=.04))

    link2eetraj = {}
    #link2eetraj is a hash of gripper fram to new trajectory

    #Transform the gripper trajectory here
    for lr in 'lr':
        link_name = "%s_gripper_tool_frame" % lr
        #old_ee_traj is the old gripper trajectory
        old_ee_traj = asarray(seg_info[link_name]["hmat"])
        #new_ee_traj is the transformed gripper trajectory
        new_ee_traj = f.transform_hmats(old_ee_traj)

        link2eetraj[link_name] = new_ee_traj

        #Draw the old and new gripper trajectories as lines
        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)))
        #Globals.exec_log(curr_step, "gen_traj.link2eetraj", link2eetraj)

    miniseg_starts, miniseg_ends = split_trajectory_by_gripper(seg_info)
    success = True
    #print colorize.colorize("mini segments:", "red"), miniseg_starts, miniseg_ends

    #TODO: modify for no body
    for (i_miniseg, (i_start, i_end)) in enumerate(zip(miniseg_starts, miniseg_ends)):
        ################################
        redprint("Generating joint trajectory for segment %s, part %i" % (segment, 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 is the old trajectory inside the minisegment
            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)
            #Call the planner (eg. trajopt)
            with dhm_u.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)
            part_name = {"l": "larm", "r": "rarm"}[lr]
            bodypart2traj[part_name] = new_joint_traj

        ### Execute the gripper ###
        redprint("Executing joint trajectory for segment %s, part %i using arms '%s'" % (
            segment, 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_sim(lr, gripper_open, prev_gripper_open, animate):
                redprint("Grab %s failed" % lr)
                success = False
        if not success:
            break
            # Execute the robot trajectory
        if len(bodypart2traj) > 0:
            success &= exec_traj_sim(bodypart2traj, animate)

        #Globals.exec_log(curr_step, "execute_traj.miniseg_%d.sim_rope_nodes_after_traj" % i_miniseg, Globals.sim.rope.GetNodes())

        if not success:
            break

    Globals.sim.settle(animate=animate)
    if Globals.exec_log:
        Globals.exec_log(curr_step, "execute_traj.sim_rope_nodes_after_full_traj", Globals.sim.rope.GetNodes())

    from rapprentice import knot_identification

    knot_name = knot_identification.identify_knot(Globals.sim.rope.GetControlPoints())
    found_knot = False
    if knot_name is not None:
        if knot_name == knot or knot == "any":
            redprint("Identified knot: %s. Success!" % knot_name)
            #Globals.exec_log(curr_step, "result", True, description="identified knot %s" % knot_name)
            found_knot = True
        else:
            redprint("Identified knot: %s, but expected %s. Continuing." % (knot_name, knot))
    else:
        redprint("Not a knot. Continuing.")

    redprint("Segment %s result: %s" % (segment, success))
    return {'found_knot': found_knot, 'segment': segment, 'link2eetraj': link2eetraj, 'new_xyz': new_xyz_upsampled}