PR2_L_POSTURES = dict(
    untucked = [0.4,  1.0,   0.0,  -2.05,  0.0,  -0.1,  0.0],
    tucked = [0.06, 1.25, 1.79, -1.68, -1.73, -0.10, -0.09],
    up = [ 0.33, -0.35,  2.59, -0.15,  0.59, -1.41, -0.27],
    side = [  1.832,  -0.332,   1.011,  -1.437,   1.1  ,  -2.106,  3.074]
)
    GRIPPER_L_FINGER_OPEN_CLOSE_THRESH = mult * settings.GRIPPER_OPEN_CLOSE_THRESH
class RopeState(object):
class RopeParams(object):
class SceneState(object):
    def __init__(self, init_rope_nodes, rope_params, tfs=None):
    def __init__(self):
    def __eq__(self, other):
    def __ne__(self, other):
    def __init__(self, cloud, rope_nodes, rope_state, id=None, color=None):
    def get_unique_id():
def make_table_xml(translation, extents):
def make_box_xml(name, translation, extents):
def make_cylinder_xml(name, translation, radius, height):
def reset_arms_to_side(sim_env):
def arm_moved(joint_traj):    
def split_trajectory_by_gripper(seg_info):
def split_trajectory_by_lr_gripper(seg_info, lr):
def get_opening_closing_inds(finger_traj):
def gripper_joint2gripper_l_finger_joint_values(gripper_joint_vals):
def binarize_gripper(angle):
def get_binary_gripper_angle(open):
def set_gripper_maybesim(sim_env, lr, is_open, prev_is_open, animate=False):
def mirror_arm_joints(x):
def unwrap_arm_traj_in_place(traj):
def unwrap_in_place(t, dof_inds=None):
def dof_inds_from_name(robot, name):
def sim_traj_maybesim(sim_env, lr2traj, animate=False, interactive=False, max_cart_vel_trans_traj=.05):
def sim_full_traj_maybesim(sim_env, full_traj, animate=False, interactive=False, max_cart_vel_trans_traj=.05):
    def sim_callback(i):
    def is_rope_pulled_too_tight(i_step):
def get_full_traj(robot, lr2arm_traj, lr2finger_traj = {}):
def merge_full_trajs(full_trajs):
def get_ee_traj(robot, lr, arm_traj_or_full_traj, ee_link_name_fmt="%s_gripper_tool_frame"):
def get_finger_rel_pts(finger_lr):
def get_finger_pts_traj(robot, lr, full_traj_or_ee_finger_traj):
def grippers_exceed_rope_length(sim_env, full_traj, thresh):
def remove_tight_rope_pull(sim_env, full_traj):
def load_random_start_segment(demofile):
def load_fake_data_segment(demofile, fake_data_segment, fake_data_transform):
def unif_resample(traj, max_diff, wt = None):        
def get_rope_transforms(sim_env):
def replace_rope(new_rope, sim_env, rope_params, restore=False):
def set_rope_transforms(tfs, sim_env):
def get_rope_params(params_id):
def tpsrpm_plot_cb(sim_env, x_nd, y_md, targ_Nd, corr_nm, wt_n, f):
def draw_grid(sim_env, old_xyz, f, color = (1,1,0,1)):
def draw_axis(sim_env, hmat, arrow_length=.1, arrow_width=0.005):
def draw_finger_pts_traj(sim_env, flr2finger_pts_traj, color): GRIPPER_OPEN_CLOSE_THRESH = 0.04 # .06 for fig8 (thick rope), 0.04 for thin rope (overhand)
ROPE_RADIUS        = .005
ROPE_ANG_STIFFNESS = .1
ROPE_ANG_DAMPING   = 1
ROPE_LIN_DAMPING   = .75
ROPE_ANG_LIMIT     = .4
ROPE_LIN_STOP_ERP  = .2
ROPE_MASS          = 1.0
ROPE_RADIUS_THICK  = .008
WINDOW_PROP = [0,0,1500,1500]
CAMERA_MATRIX = [[    0, 1,   0, 0],
				 [   -1, 0, 0.5, 0],
				 [  0.5, 0,   1, 0],
				 [ 2.25, 0, 4.5, 1]] class SceneState(object):
    ids = set()
    def __init__(self, full_cloud, id=None, full_color=None, downsample_size=0):
        """Inits SceneState
        
        Args:
            full_cloud: full (i.e. not downsampled) cloud
            id: unique id for this SceneState
            full_color: colors for the respective points of full_cloud. Should have the same dimensions as full_cloud
            downsample_size: if downsample_size is positive, the full cloud and color are downsampled to a voxel size of downsample_size, else they are not downsampled
        """
        self.full_cloud = full_cloud
        self.full_color = full_color
        if downsample_size > 0:
            global clouds
            from lfd.rapprentice import clouds
            if full_color is not None:
                cloud_color = clouds.downsample(np.c_[full_cloud, full_color], downsample_size)
                self.cloud = cloud_color[:,:3]
                self.color = cloud_color[:,3:]
            else:
                self.cloud = clouds.downsample(full_cloud, downsample_size)
                self.color = None
        else:
            self.cloud = full_cloud
            self.color = full_color
        if id is None:
            self.id = SceneState.get_unique_id()
        else:
            self.id = id

    @staticmethod
    def get_unique_id():
        id = len(SceneState.ids)
        assert id not in SceneState.ids
        SceneState.ids.add(id)
        return id
    
    def __repr__(self):
        return "%s(..., id=%i)" % (self.__class__.__name__, self.id) class GroundTruthRopeSceneState(SceneState):
    def __init__(self, rope_nodes, radius, upsample=0, upsample_rad=1, downsample_size=0):
        full_cloud = ropesim.observe_cloud(rope_nodes, radius, upsample=upsample, upsample_rad=upsample_rad)
        super(GroundTruthRopeSceneState, self).__init__(full_cloud, downsample_size=downsample_size)
        self.rope_nodes = rope_nodes
        self.crossing_info = None #TODO: optionally compute/load cached crossing_info class AugmentedTrajectory(object):
    def __init__(self, lr2arm_traj=None, lr2finger_traj=None, lr2ee_traj=None, lr2open_finger_traj=None, lr2close_finger_traj=None):
        """Inits AugmentedTrajectory
        
        Args:
            lr2arm_traj: dict that maps from 'l' and/or 'r' to the left arm's and/or right arm's joint angle trajectory
            lr2finger_traj: dict that maps to the left gripper's and/or right gripper's finger joint angle trajectory
            lr2ee_traj: dict that maps to the left gripper's and/or right gripper's end-effector trajectory (i.e. a numpy.array of homogeneous matrices)
            lr2open_finger_traj: dict that maps to a boolean vector indicating whether there is a opening action for the left gripper and/or right gripper for every time step. By default, there is no action for all time steps.
            lr2close_finger_traj: same as lr2open_finger_traj but for closing action
        """
        # make sure all trajs have the same number of steps
        self.n_steps = None
        for lr2traj in [lr2arm_traj, lr2finger_traj, lr2ee_traj, lr2open_finger_traj, lr2close_finger_traj]:
            if lr2traj is None:
                continue
            for lr in lr2traj.keys():
                if self.n_steps is None:
                    self.n_steps = lr2traj[lr].shape[0]
                else:
                    assert lr2traj[lr].shape[0] == self.n_steps

        if lr2arm_traj is None:
            self.lr2arm_traj = {}
        else:
            self.lr2arm_traj = lr2arm_traj
        
        if lr2finger_traj is None:
            self.lr2finger_traj = {}
        else:
            self.lr2finger_traj = lr2finger_traj
        
        if lr2ee_traj is None:
            self.lr2ee_traj = {}
        else:
            self.lr2ee_traj = lr2ee_traj
        
        if lr2open_finger_traj is None:
            self.lr2open_finger_traj = {}
            for lr in 'lr':
                self.lr2open_finger_traj[lr] = np.zeros(self.n_steps, dtype=bool)
        else:
            self.lr2open_finger_traj = lr2open_finger_traj
        if lr2close_finger_traj is None:
            self.lr2close_finger_traj = {}
            for lr in 'lr':
                self.lr2close_finger_traj[lr] = np.zeros(self.n_steps, dtype=bool)
        else:
            self.lr2close_finger_traj = lr2close_finger_traj
    
    def __eq__(self, other):
        if isinstance(other, self.__class__):
            for (lr2traj, other_lr2traj) in [(self.lr2arm_traj, other.lr2arm_traj), (self.lr2finger_traj, other.lr2finger_traj),
                                             (self.lr2ee_traj, other.lr2ee_traj),
                                             (self.lr2open_finger_traj, other.lr2open_finger_traj), (self.lr2close_finger_traj, other.lr2close_finger_traj)]:
                if lr2traj is None:
                    if other_lr2traj is None:
                        continue
                    else:
                        return False
                if set(lr2traj.keys()) != set(other_lr2traj.keys()):
                    return False
                for lr in lr2traj.keys():
                    if np.any(lr2traj[lr] != other_lr2traj[lr]):
                        return False
            return True
        else:
            return False
    
    def __ne__(self, other):
        return not self.__eq__(other)
    
    @staticmethod
    def create_from_full_traj(robot, full_traj, lr2open_finger_traj=None, lr2close_finger_traj=None):
        traj, dof_inds = full_traj
        lr2arm_traj = {}
        for lr in 'lr':
            manip_name = {"l":"leftarm", "r":"rightarm"}[lr]
            arm_inds = robot.GetManipulator(manip_name).GetArmIndices()
            if set(arm_inds).intersection(set(dof_inds)):
                if not set(arm_inds).issubset(set(dof_inds)):
                    raise RuntimeError("Cannot create AugmentedTrajectory from incomplete full_traj")
                arm_traj = np.zeros((traj.shape[0], len(arm_inds)))
                for (j,arm_ind) in enumerate(arm_inds):
                    arm_traj[:,j] = traj[:,dof_inds.index(arm_ind)]
                lr2arm_traj[lr] = arm_traj
        lr2finger_traj = {}
        for lr in 'lr':
            finger_ind = robot.GetJointIndex("%s_gripper_l_finger_joint"%lr)
            if finger_ind in dof_inds:
                lr2finger_traj[lr] = traj[:,dof_inds.index(finger_ind)][:,None]
        lr2ee_traj = {}    
        for lr in lr2arm_traj.keys():
            lr2ee_traj[lr] = sim_util.get_ee_traj(robot, lr, lr2arm_traj[lr])
        return AugmentedTrajectory(lr2arm_traj=lr2arm_traj, lr2finger_traj=lr2finger_traj, lr2ee_traj=lr2ee_traj, lr2open_finger_traj=lr2open_finger_traj, lr2close_finger_traj=lr2close_finger_traj)
    
    def get_full_traj(self, robot):
        """
        TODO: remove sim_util.get_full_traj
        """
        trajs = []
        dof_inds = []
        for lr in 'lr':
            if lr in self.lr2arm_traj:
                arm_traj = self.lr2arm_traj[lr]
                manip_name = {"l":"leftarm", "r":"rightarm"}[lr]
                trajs.append(arm_traj)
                dof_inds.extend(robot.GetManipulator(manip_name).GetArmIndices())
        for lr in 'lr':
            if lr in self.lr2finger_traj:
                finger_traj = self.lr2finger_traj[lr]
                trajs.append(finger_traj)
                dof_inds.append(robot.GetJointIndex("%s_gripper_l_finger_joint"%lr))
        if len(trajs) > 0:
            full_traj = (np.concatenate(trajs, axis=1), dof_inds)
        else:
            full_traj = (np.zeros((0,0)), [])
        return full_traj
    
    def get_resampled_traj(self, timesteps_rs):
        lr2arm_traj_rs = None if self.lr2arm_traj is None else {}
        lr2finger_traj_rs = None if self.lr2finger_traj is None else {}
        for (lr2traj_rs, self_lr2traj) in [(lr2arm_traj_rs, self.lr2arm_traj), (lr2finger_traj_rs, self.lr2finger_traj)]:
            if self_lr2traj is None:
                continue
            for lr in self_lr2traj.keys():
                lr2traj_rs[lr] = mu.interp2d(timesteps_rs, np.arange(len(self_lr2traj[lr])), self_lr2traj[lr])

        if self.lr2ee_traj is None:
            lr2ee_traj_rs = None
        else:
            lr2ee_traj_rs = {}
            for lr in self.lr2ee_traj.keys():
                lr2ee_traj_rs[lr] = np.asarray(resampling.interp_hmats(timesteps_rs, np.arange(len(self.lr2ee_traj[lr])), self.lr2ee_traj[lr]))
        
        lr2open_finger_traj_rs = None if self.lr2open_finger_traj is None else {}
        lr2close_finger_traj_rs = None if self.lr2close_finger_traj is None else {}
        for (lr2oc_finger_traj_rs, self_lr2oc_finger_traj) in [(lr2open_finger_traj_rs, self.lr2open_finger_traj), (lr2close_finger_traj_rs, self.lr2close_finger_traj)]:
            if self_lr2oc_finger_traj is None:
                continue
            for lr in self_lr2oc_finger_traj.keys():
                self_oc_finger_traj = self_lr2oc_finger_traj[lr]
                self_oc_inds = np.where(self_oc_finger_traj)[0]
                oc_inds_rs = np.abs(timesteps_rs[:,None] - self_oc_inds[None,:]).argmin(axis=0)
                oc_finger_traj_rs = np.zeros(len(timesteps_rs), dtype=bool)
                oc_finger_traj_rs[oc_inds_rs] = True
                lr2oc_finger_traj_rs[lr] = oc_finger_traj_rs
        
        return AugmentedTrajectory(lr2arm_traj=lr2arm_traj_rs, lr2finger_traj=lr2finger_traj_rs, lr2ee_traj=lr2ee_traj_rs, lr2open_finger_traj=lr2open_finger_traj_rs, lr2close_finger_traj=lr2close_finger_traj_rs)
    
    def __repr__(self):
        return "%s(...)" % (self.__class__.__name__) class Demonstration(object):
    def __init__(self, name, scene_state, aug_traj):
        """Inits Demonstration
        
        Args:
            name: demonstration name, which is the same as the ones used for indexing the demonstrations
            scene_state: demonstration SceneState
            aug_traj: demonstration AugmentedTrajectory. Only lr2ee_traj and lr2finger_traj are required for the demonstration AugmentedTrajectory. If lr2arm_traj is specified, it might be used for initializing the trajectory optimization.
        """
        self.name = name
        self.scene_state = scene_state
        if aug_traj is not None:
            assert aug_traj.lr2ee_traj is not None
            assert aug_traj.lr2finger_traj is not None
        self.aug_traj = aug_traj
    
    def __repr__(self):
        return "%s(%s, %s, %s)" % (self.__class__.__name__, self.name, self.scene_state.__repr__(), self.aug_traj.__repr__()) class DynamicRopeSimulationRobotWorld(DynamicSimulationRobotWorld):

    def in_grasp_region(self,robot, lr, pt):
        tol = .00

        manip_name = {"l": "leftarm", "r": "rightarm"}[lr]
        manip = robot.GetManipulator(manip_name)
        l_finger = robot.GetLink("%s_gripper_l_finger_tip_link"%lr)
        r_finger = robot.GetLink("%s_gripper_r_finger_tip_link"%lr)

        def transform(hmat, p):
            return hmat[:3,:3].dot(p) + hmat[:3,3]
        def on_inner_side(pt, finger_lr):
            finger = l_finger
            closing_dir = np.cross(manip.GetLocalToolDirection(), [-1, 0, 0])
            local_inner_pt = np.array([0.234402, -0.299, 0])/20.
            if finger_lr == "r":
                finger = r_finger
                closing_dir *= -1
                local_inner_pt[1] *= -1
            inner_pt = transform(finger.GetTransform(), local_inner_pt)
            return manip.GetTransform()[:3,:3].dot(closing_dir).dot(pt - inner_pt) > 0

        # check that pt is behind the gripper tip
        pt_local = transform(np.linalg.inv(manip.GetTransform()), pt)
        if pt_local[2] > .03 + tol:
            return False

        # check that pt is within the finger width
        if abs(pt_local[0]) > .01 + tol:
            return False

        # check that pt is between the fingers
        if not on_inner_side(pt, "l") or not on_inner_side(pt, "r"):
            return False

        return True


 
    def close_gripper(self, lr, step_viewer=1, max_vel=.02, close_dist_thresh=0.004, grab_dist_thresh=0.005):
        print 'CLOSING GRIPPER'
        # generate gripper finger trajectory
        joint_ind = self.robot.GetJoint("%s_gripper_l_finger_joint"%lr).GetDOFIndex()
        start_val = self.robot.GetDOFValues([joint_ind])[0]
        print 'start_val: ', start_val
        # execute gripper finger trajectory
        dyn_bt_objs = [bt_obj for sim_obj in self.dyn_sim_objs for bt_obj in sim_obj.get_bullet_objects()]
        next_val = start_val
        while next_val:
            flr2finger_pts_grid = self._get_finger_pts_grid(lr)
            ray_froms, ray_tos = flr2finger_pts_grid['l'], flr2finger_pts_grid['r']

            # stop closing if any ray hits a dynamic object within a distance of close_dist_thresh from both sides
            next_vel = max_vel
            for bt_obj in dyn_bt_objs:
                from_to_ray_collisions = self.bt_env.RayTest(ray_froms, ray_tos, bt_obj)
                to_from_ray_collisions = self.bt_env.RayTest(ray_tos, ray_froms, bt_obj)
                rays_dists = np.inf * np.ones((len(ray_froms), 2))
                for rc in from_to_ray_collisions:
                    ray_id = np.argmin(np.apply_along_axis(np.linalg.norm, 1, ray_froms - rc.rayFrom))
                    rays_dists[ray_id,0] = np.linalg.norm(rc.pt - rc.rayFrom)
                for rc in to_from_ray_collisions:
                    ray_id = np.argmin(np.apply_along_axis(np.linalg.norm, 1, ray_tos - rc.rayFrom))
                    rays_dists[ray_id,1] = np.linalg.norm(rc.pt - rc.rayFrom)
                colliding_rays_inds = np.logical_and(rays_dists[:,0] != np.inf, rays_dists[:,1] != np.inf)
                if np.any(colliding_rays_inds):
                    rays_dists = rays_dists[colliding_rays_inds,:]
                    if np.any(np.logical_and(rays_dists[:,0] < close_dist_thresh, rays_dists[:,1] < close_dist_thresh)):
                        next_vel = 0
                    else:
                        next_vel = np.minimum(next_vel, np.min(rays_dists.sum(axis=1)))
            if next_vel == 0:
                break
            next_val = np.maximum(next_val - next_vel, 0)

            self.robot.SetDOFValues([next_val], [joint_ind])
            self.step()
            if self.viewer and step_viewer:
                self.viewer.Step()

 
        rope = [bt_obj for sim_obj in self.dyn_sim_objs for bt_obj in sim_obj.get_bullet_objects()][0]
        nodes, ctl_pts = rope.GetNodes(), rope.GetControlPoints()

        graspable_nodes = np.array([self.in_grasp_region(self.robot, lr, n) for n in nodes])
        graspable_ctl_pts = np.array([self.in_grasp_region(self.robot, lr, n) for n in ctl_pts])
        graspable_inds = np.flatnonzero(np.logical_or(graspable_nodes, np.logical_or(graspable_ctl_pts[:-1], graspable_ctl_pts[1:])))
        print 'graspable inds for %s: %s' % (lr, str(graspable_inds))
        if len(graspable_inds) == 0:
            return False

        robot_link = self.robot.GetLink("%s_gripper_l_finger_tip_link"%lr)
        rope_links = rope.GetKinBody().GetLinks()
        for i_node in graspable_inds:
            i_cnt = i_node
            for geom in rope_links[i_cnt].GetGeometries():
                geom.SetDiffuseColor([1.,0.,0.])
            link = rope_links[i_cnt]
            self._add_constraints(lr, link, link.GetTransform())
        if step_viewer and self.viewer:
            self.viewer.Step()

        return True

    def release_rope(self, lr):
#         print 'RELEASE: %s (%d constraints)' % (lr, len(self.constraints[lr]))
        for c in self.constraints[lr]:
            self.bt_env.RemoveConstraint(c)
        rope_links = self.rope.GetKinBody().GetLinks()
        for link in self.constraints_links[lr]:
            for geom in link.GetGeometries():
                geom.SetDiffuseColor([1.,1.,1.])
        self.constraints[lr] = []
        self.constraints_links[lr] = [] class XmlSimulationObject(SimulationObject):
    def __init__(self, xml, dynamic=False):
        super(XmlSimulationObject, self).__init__(None, dynamic=dynamic) # TODO: None names
        self.xml = xml

    def add_to_env(self, sim):
        self.sim = sim
        pre_names = [body.GetName() for body in self.sim.env.GetBodies()]
        if '<' in self.xml: # TODO: fix this hack
            self.sim.env.LoadData(self.xml)
        else:
            self.sim.env.Load(self.xml)
        post_names = [body.GetName() for body in self.sim.env.GetBodies()]
        self.names = [name for name in post_names if name not in pre_names]
    
    def remove_from_env(self):
        for bt_obj in self.get_bullet_objects():
            self.sim.env.Remove(bt_obj.GetKinBody())
        self.sim = None
    
    def _get_constructor_info(self):
        args = [self.xml]
        kwargs = {"dynamic":self.dynamic}
        return (type(self).__name__, type(self).__module__), args, kwargs
    
    def __repr__(self):
        return "XmlSimulationObject(%s, dynamic=%r)" % (self.xml, self.dynamic) class BoxSimulationObject(XmlSimulationObject):
    def __init__(self, name, translation, extents, dynamic=True):
        xml = """
        <Environment>
          <KinBody name="%s">
            <Body type="%s" name="%s_link">
              <Translation>%f %f %f</Translation>
              <Geom type="box">
                <extents>%f %f %f</extents>
              </Geom>
            </Body>
          </KinBody>
        </Environment>
        """ % (name, 'dynamic' if dynamic else 'static', name, translation[0], translation[1], translation[2], extents[0], extents[1], extents[2])
        super(BoxSimulationObject, self).__init__(xml, dynamic=dynamic)
        self.name = name
        self.translation = translation
        self.extents = extents
    
    def _get_constructor_info(self):
        args = [self.name, self.translation, self.extents]
        kwargs = {"dynamic":self.dynamic}
        return (type(self).__name__, type(self).__module__), args, kwargs
    
    def __repr__(self):
        return "BoxSimulationObject(%s, %s, %s, dynamic=%r)" % (self.name, self.translation, self.extents, self.dynamic) class RopeSimulationObject(SimulationObject):
    add_after = True
    def __init__(self, name, ctrl_points, rope_params=None, dynamic=True, upsample=0, upsample_rad=1):
        super(RopeSimulationObject, self).__init__([name], dynamic=True)
        self.name = name
        self.init_ctrl_points = np.asarray(ctrl_points)
        if rope_params is None:
            self.rope_params = sim_util.RopeParams()
        else:
            self.rope_params = rope_params
        self.upsample = upsample
        self.upsample_rad = upsample_rad
        self.rope = None

    def add_to_env(self, sim):
        self.sim = sim
        capsule_rope_params = bulletsimpy.CapsuleRopeParams()
        capsule_rope_params.radius       = self.rope_params.radius
        capsule_rope_params.angStiffness = self.rope_params.angStiffness
        capsule_rope_params.angDamping   = self.rope_params.angDamping
        capsule_rope_params.linDamping   = self.rope_params.linDamping
        capsule_rope_params.angLimit     = self.rope_params.angLimit
        capsule_rope_params.linStopErp   = self.rope_params.linStopErp
        capsule_rope_params.mass         = self.rope_params.mass
        self.rope = bulletsimpy.CapsuleRope(self.sim.bt_env, self.name, self.init_ctrl_points, capsule_rope_params)
    
    def remove_from_env(self):
        # remove all capsule-capsule exclude to prevent memory leak
        # TODO: only interate through the capsule pairs that actually are excluded
        cc = trajoptpy.GetCollisionChecker(self.sim.env)
        for rope_link0 in self.rope.GetKinBody().GetLinks():
            for rope_link1 in self.rope.GetKinBody().GetLinks():
                cc.IncludeCollisionPair(rope_link0, rope_link1)
        self.sim.env.Remove(self.rope.GetKinBody())
        self.rope = None
        self.sim = None
    
    def get_bullet_objects(self):
        # method of parent class doesn't work because self.rope casted to BulletObject
        if self.sim is None:
            raise RuntimeError("get_bullet_objects should only be called when the object is in an environment")
        return [self.rope]
    
    def get_state(self):
        trans, rots = self.rope.GetTranslations(), self.rope.GetRotations()
        tfs = np.zeros((len(trans), 4, 4))
        tfs[:,:3,3] = trans
        tfs[:,:3,:3] = rots
        tfs[:,3,3] = np.ones(len(trans))
        return tfs
    
    def set_state(self, tfs):
        self.rope.SetTranslations(tfs[:,:3,3])
        self.rope.SetRotations(tfs[:,:3,:3])
    
    def _get_constructor_info(self):
        args = [self.name, self.init_ctrl_points.tolist(), self.rope_params]
        kwargs = {"dynamic":self.dynamic, "upsample":0, "upsample_rad":1}
        return (type(self).__name__, type(self).__module__), args, kwargs
    
    def __repr__(self):
        return "RopeSimulationObject(%s, numpy.array([[...]]), RopeParams(...), dynamic=%r, upsample=%i, upsample_rad=%i)" % (self.name, self.dynamic, self.upsample, self.upsample_rad) class LfdEnvironment(object):
    def __init__(self, world, sim, downsample_size=0):
        """Inits LfdEnvironment

        Args:
            world: RobotWorld
            sim: StaticSimulation that contains a robot
            downsample_size: if downsample_size is positive, the clouds are downsampled to a voxel size of downsample_size, else they are not downsampled
        """
        self.world = world
        self.sim = sim
        self.downsample_size = downsample_size
    
    def execute_augmented_trajectory(self, aug_traj, step_viewer=1, interactive=False, sim_callback=None, check_feasible=False):
        open_or_close_finger_traj = np.zeros(aug_traj.n_steps, dtype=bool)
        if aug_traj.lr2open_finger_traj is not None:
            for lr in aug_traj.lr2open_finger_traj.keys():
                open_or_close_finger_traj = np.logical_or(open_or_close_finger_traj, aug_traj.lr2open_finger_traj[lr])
        if aug_traj.lr2close_finger_traj is not None:
            for lr in aug_traj.lr2close_finger_traj.keys():
                open_or_close_finger_traj = np.logical_or(open_or_close_finger_traj, aug_traj.lr2close_finger_traj[lr])
        open_or_close_inds = np.where(open_or_close_finger_traj)[0]

        traj, dof_inds = aug_traj.get_full_traj(self.sim.robot)
        feasible = True
        misgraspl = False
        misgraspr = False
        lr2gripper_open = {'l':True, 'r':True}
        for (start_ind, end_ind) in zip(np.r_[0, open_or_close_inds], np.r_[open_or_close_inds+1, aug_traj.n_steps]):
            if aug_traj.lr2open_finger_traj is not None:
                for lr in aug_traj.lr2open_finger_traj.keys():
                    if aug_traj.lr2open_finger_traj[lr][start_ind]:
                        target_val = None
                        joint_ind = self.sim.robot.GetJoint("%s_gripper_l_finger_joint"%lr).GetDOFIndex()
                        if joint_ind in dof_inds:
                            target_val = traj[start_ind, dof_inds.index(joint_ind)]
                        self.world.open_gripper(lr, target_val=target_val, step_viewer=step_viewer)
                        lr2gripper_open[lr] = True
            if aug_traj.lr2close_finger_traj is not None:
                for lr in aug_traj.lr2close_finger_traj.keys():
                    if aug_traj.lr2close_finger_traj[lr][start_ind]:
                        n_cnts = len(self.sim.constraints[lr])
                        self.world.close_gripper(lr, step_viewer=step_viewer)
                        if len(self.sim.constraints[lr]) == n_cnts and lr == 'l':
                            misgraspl = True
                        elif lr == 'l':
                            misgraspl = False
                        elif len(self.sim.constraints[lr]) == n_cnts and lr=='r':
                            misgraspr = True
                        else:
                            misgraspr = False
                        #misgrasp |= len(self.sim.constraints[lr]) == n_cnts
                        lr2gripper_open[lr] = False
            misgrasp = misgraspl or misgraspr
            # don't execute trajectory for finger joint if the corresponding gripper is closed
            active_inds = np.ones(len(dof_inds), dtype=bool)
            for lr in 'lr':
                if not lr2gripper_open[lr]:
                    joint_ind = self.sim.robot.GetJoint("%s_gripper_l_finger_joint"%lr).GetDOFIndex()
                    if joint_ind in dof_inds:
                        active_inds[dof_inds.index(joint_ind)] = False
            miniseg_traj = traj[start_ind:end_ind, active_inds]
            miniseg_dof_inds = list(np.asarray(dof_inds)[active_inds])
            full_traj = (miniseg_traj, miniseg_dof_inds)
            feasible &= eval_util.traj_is_safe(self.sim, full_traj, 0)
            if check_feasible and not feasible:
                break
            self.world.execute_trajectory(full_traj, step_viewer=step_viewer, interactive=interactive, sim_callback=sim_callback)
        return feasible, misgrasp
    
    def observe_scene(self):
        full_cloud = self.world.observe_cloud()
        return demonstration.SceneState(full_cloud, downsample_size=self.downsample_size) class GroundTruthRopeLfdEnvironment(LfdEnvironment):
    def __init__(self, world, sim, upsample=0, upsample_rad=1, downsample_size=0):
        """Inits GroundTruthRopeLfdEnvironment
        
        Args:
            world: RobotWorld
            sim: DynamicSimulation that should containing exactly one rope when observe_scene is called
            downsample_size: if downsample_size is positive, the clouds are downsampled to a voxel size of downsample_size, else they are not downsampled
        """
        super(GroundTruthRopeLfdEnvironment, self).__init__(world, sim, downsample_size=downsample_size)

        self.upsample = upsample
        self.upsample_rad = upsample_rad

    def observe_scene(self):
        rope_sim_objs = [sim_obj for sim_obj in self.sim.sim_objs if isinstance(sim_obj, simulation_object.RopeSimulationObject)]
        assert len(rope_sim_objs) == 1
        rope_sim_obj = rope_sim_objs[0]
        return demonstration.GroundTruthRopeSceneState(rope_sim_obj.rope.GetControlPoints(), 
                                                       rope_sim_obj.rope_params.radius, 
                                                       upsample=self.upsample, 
                                                       upsample_rad=self.upsample_rad, 
                                                       downsample_size=self.downsample_size) class TpsRpmBijRegistrationFactory(RegistrationFactory):
    r"""As in:
    
    J. Schulman, J. Ho, C. Lee, and P. Abbeel, "Learning from Demonstrations through the Use of Non-Rigid Registration," in Proceedings of the 16th International Symposium on Robotics Research (ISRR), 2013.
    
    Tries to solve the optimization problem
    
    .. math::
        :nowrap:

        \begin{align*}
            & \min_{f, M} 
                & \frac{1}{n} \sum_{i=1}^n \sum_{j=1}^m m_{ij} ||y_j - f(x_i)||_2^2
                + \lambda Tr(A_f^\top K A_f)
                + Tr((B_f - I) R (B_f - I)) \\
                && + \frac{2T}{n} \sum_{i=1}^n \sum_{j=1}^m m_{ij} \log m_{ij}
                - \frac{2T}{n} \sum_{i=1}^n \sum_{j=1}^m m_{ij} \\
                && + \frac{1}{m} \sum_{j=1}^m \sum_{i=1}^n m_{ij} ||x_i - g(y_j)||_2^2
                + \lambda Tr(A_g^\top K A_g)
                + Tr((B_g - I) R (B_g - I)) \\
                && + \frac{2T}{m} \sum_{j=1}^m \sum_{i=1}^n m_{ij} \log m_{ij}
                - \frac{2T}{m} \sum_{j=1}^m \sum_{i=1}^n m_{ij} \\
            & \text{subject to} 
                & X^\top A_f = 0, 1^\top A_f = 0 \\
                && Y^\top A_g = 0, 1^\top A_g = 0 \\
                && \sum_{i=1}^{n+1} m_{ij} = 1, \sum_{j=1}^{m+1} m_{ij} = 1, m_{ij} \geq 0 \\
        \end{align*}
    """
    def __init__(self, demos=None, 
                 n_iter=settings.N_ITER, em_iter=settings.EM_ITER, 
                 reg_init=settings.REG[0], reg_final=settings.REG[1], 
                 rad_init=settings.RAD[0], rad_final=settings.RAD[1], 
                 rot_reg=settings.ROT_REG, 
                 outlierprior=settings.OUTLIER_PRIOR, outlierfrac=settings.OURLIER_FRAC, 
                 prior_fn=None, 
                 f_solver_factory=solver.AutoTpsSolverFactory(), 
                 g_solver_factory=solver.AutoTpsSolverFactory(use_cache=False)):
        """Inits TpsRpmBijRegistrationFactory with demonstrations and parameters
        
        Args:
            demos: dict that maps from demonstration name to Demonstration
            n_iter: outer iterations for tps-rpm
            em_iter: inner iterations for tps-rpm
            reg_init/reg_final: regularization on curvature
            rad_init/rad_final: radius (temperature) for correspondence calculation (meters)
            rot_reg: regularization on rotation
            prior_fn: function that takes the demo and test SceneState and returns the prior probability (i.e. NOT cost)
            f_solver_factory: solver factory for forward registration
            g_solver_factory: solver factory for backward registration
        
        Note:
            Pick a T_init that is about 1/10 of the largest square distance of all point pairs.
            You might not want to cache for the target SolverFactory.
        """
        super(TpsRpmBijRegistrationFactory, self).__init__(demos=demos)
        self.n_iter = n_iter
        self.em_iter = em_iter
        self.reg_init = reg_init
        self.reg_final = reg_final
        self.rad_init = rad_init
        self.rad_final = rad_final
        self.rot_reg = rot_reg
        self.outlierprior = outlierprior
        self.outlierfrac = outlierfrac
        self.prior_fn = prior_fn
        self.f_solver_factory = f_solver_factory
        self.g_solver_factory = g_solver_factory
    
    def register(self, demo, test_scene_state, callback=None):
        if self.prior_fn is not None:
            prior_prob_nm = self.prior_fn(demo.scene_state, test_scene_state)
        else:
            prior_prob_nm = None
        x_nd = demo.scene_state.cloud[:,:3]
        y_md = test_scene_state.cloud[:,:3]
        
        f, g, corr = tps.tps_rpm_bij(x_nd, y_md, 
                                     f_solver_factory=self.f_solver_factory, g_solver_factory=self.g_solver_factory, 
                                     n_iter=self.n_iter, em_iter=self.em_iter, 
                                     reg_init=self.reg_init, reg_final=self.reg_final, 
                                     rad_init=self.rad_init, rad_final=self.rad_final, 
                                     rot_reg=self.rot_reg, 
                                     outlierprior=self.outlierprior, outlierfrac=self.outlierfrac, 
                                     prior_prob_nm=prior_prob_nm, callback=callback)
        
        return TpsRpmBijRegistration(demo, test_scene_state, f, g, corr, self.rad_final)
    
    def cost(self, demo, test_scene_state):
        """Gets the costs of the forward and backward thin plate spline 
        objective of the resulting registration
        
        Args:
            demo: Demonstration which has the demonstration scene
            test_scene_state: SceneState of the test scene
        
        Returns:
            A 1-dimensional numpy.array containing the residual, bending and 
            rotation cost of the forward and backward spline, each already 
            premultiplied by the respective coefficients.
        """
        reg = self.register(demo, test_scene_state, callback=None)
        cost = np.r_[reg.f.get_objective(), reg.g.get_objective()]
        return cost class TpsRpmRegistrationFactory(RegistrationFactory):
    r"""As in:
    
    H. Chui and A. Rangarajan, "A new point matching algorithm for non-rigid registration," Computer Vision and Image Understanding, vol. 89, no. 2, pp. 114-141, 2003.
    
    Tries to solve the optimization problem
    
    .. math::
        :nowrap:

        \begin{align*}
            & \min_{f, M} 
                & \frac{1}{n} \sum_{i=1}^n \sum_{j=1}^m m_{ij} ||y_j - f(x_i)||_2^2
                + \lambda Tr(A^\top K A)
                + Tr((B - I) R (B - I)) \\
                && + \frac{2T}{n} \sum_{i=1}^n \sum_{j=1}^m m_{ij} \log m_{ij}
                - \frac{2T}{n} \sum_{i=1}^n \sum_{j=1}^m m_{ij} \\
            & \text{subject to} 
                & X^\top A = 0, 1^\top A = 0 \\
                && \sum_{i=1}^{n+1} m_{ij} = 1, \sum_{j=1}^{m+1} m_{ij} = 1, m_{ij} \geq 0 \\
        \end{align*}
    """
    def __init__(self, demos=None, 
                 n_iter=settings.N_ITER, em_iter=settings.EM_ITER, 
                 reg_init=settings.REG[0], reg_final=settings.REG[1], 
                 rad_init=settings.RAD[0], rad_final=settings.RAD[1], 
                 rot_reg=settings.ROT_REG, 
                 outlierprior=settings.OUTLIER_PRIOR, outlierfrac=settings.OURLIER_FRAC, 
                 prior_fn=None, 
                 f_solver_factory=solver.AutoTpsSolverFactory()):
        """Inits TpsRpmRegistrationFactory with demonstrations and parameters
        
        Args:
            demos: dict that maps from demonstration name to Demonstration
            n_iter: outer iterations for tps-rpm
            em_iter: inner iterations for tps-rpm
            reg_init/reg_final: regularization on curvature
            rad_init/rad_final: radius (temperature) for correspondence calculation (meters)
            rot_reg: regularization on rotation
            prior_fn: function that takes the demo and test SceneState and returns the prior probability (i.e. NOT cost)
            f_solver_factory: solver factory for forward registration
        
        Note:
            Pick a T_init that is about 1/10 of the largest square distance of all point pairs.
        """
        super(TpsRpmRegistrationFactory, self).__init__(demos=demos)
        self.n_iter = n_iter
        self.em_iter = em_iter
        self.reg_init = reg_init
        self.reg_final = reg_final
        self.rad_init = rad_init
        self.rad_final = rad_final
        self.rot_reg = rot_reg
        self.outlierprior = outlierprior
        self.outlierfrac = outlierfrac
        self.prior_fn = prior_fn
        self.f_solver_factory = f_solver_factory
    
    def register(self, demo, test_scene_state, callback=None):
        if self.prior_fn is not None:
            prior_prob_nm = self.prior_fn(demo.scene_state, test_scene_state)
        else:
            prior_prob_nm = None
        x_nd = demo.scene_state.cloud[:,:3]
        y_md = test_scene_state.cloud[:,:3]
        
        f, corr = tps.tps_rpm(x_nd, y_md, 
                              f_solver_factory=self.f_solver_factory, 
                              n_iter=self.n_iter, em_iter=self.em_iter, 
                              reg_init=self.reg_init, reg_final=self.reg_final, 
                              rad_init=self.rad_init, rad_final=self.rad_final, 
                              rot_reg=self.rot_reg, 
                              outlierprior=self.outlierprior, outlierfrac=self.outlierfrac, 
                              prior_prob_nm=prior_prob_nm, callback=callback)
        
        return TpsRpmRegistration(demo, test_scene_state, f, corr, self.rad_final)
    
    def cost(self, demo, test_scene_state):
        """Gets the costs of the thin plate spline objective of the 
        resulting registration
        
        Args:
            demo: Demonstration which has the demonstration scene
            test_scene_state: SceneState of the test scene
        
        Returns:
            A 1-dimensional numpy.array containing the residual, bending and 
            rotation cost, each already premultiplied by the respective 
            coefficients.
        """
        reg = self.register(demo, test_scene_state, callback=None)
        cost = reg.f.get_objective()
        return cost class TpsSegmentRegistrationFactory(RegistrationFactory):
    def __init__(self, demos):
        raise NotImplementedError
    
    def register(self, demo, test_scene_state, callback=None):
        raise NotImplementedError
    
    def batch_register(self, test_scene_state):
        raise NotImplementedError
    
    def cost(self, demo, test_scene_state):
        raise NotImplementedError
    
    def batch_cost(self, test_scene_state):
        raise NotImplementedError class BatchGpuTpsRpmBijRegistrationFactory(TpsRpmBijRegistrationFactory):
    """
    Similar to TpsRpmBijRegistrationFactory but batch_register and batch_cost are computed in batch using the GPU
    """
    def __init__(self, demos, actionfile=None, 
                 n_iter=settings.N_ITER, em_iter=settings.EM_ITER, 
                 reg_init=settings.REG[0], reg_final=settings.REG[1], 
                 rad_init=settings.RAD[0], rad_final=settings.RAD[1], 
                 rot_reg=settings.ROT_REG, 
                 outlierprior=settings.OUTLIER_PRIOR, outlierfrac=settings.OURLIER_FRAC, 
                 prior_fn=None, 
                 f_solver_factory=solver.AutoTpsSolverFactory(), 
                 g_solver_factory=solver.AutoTpsSolverFactory(use_cache=False)):
        if not lfd.registration._has_cuda:
            raise NotImplementedError("CUDA not installed")
        super(BatchGpuTpsRpmBijRegistrationFactory, self).__init__(demos=demos, 
                                                              n_iter=n_iter, em_iter=em_iter, 
                                                              reg_init=reg_init, reg_final=reg_final, 
                                                              rad_init=rad_init, rad_final=rad_final, 
                                                              rot_reg=rot_reg, 
                                                              outlierprior=outlierprior, outlierfrac=outlierfrac, 
                                                              prior_fn=prior_fn, 
                                                              f_solver_factory=f_solver_factory, g_solver_factory=g_solver_factory)

        self.actionfile = actionfile
        if self.actionfile:
            self.bend_coefs = tps.loglinspace(self.reg_init, self.reg_final, self.n_iter)
            self.src_ctx = GPUContext(self.bend_coefs)
            self.src_ctx.read_h5(actionfile)
        self.warn_clip_cloud = True
    
    def _clip_cloud(self, cloud):
        if len(cloud) > settings.MAX_CLD_SIZE:
            cloud = cloud[np.random.choice(range(len(cloud)), size=settings.MAX_CLD_SIZE, replace=False)]
        if self.warn_clip_cloud:
            import warnings
            warnings.warn("The cloud has more points than the maximum for GPU and it is being clipped")
            self.warn_clip_cloud = False
        return cloud
    
    def batch_register(self, test_scene_state):
        raise NotImplementedError
    
    def batch_cost(self, test_scene_state):
        if not(self.actionfile):
            raise ValueError('No actionfile provided for gpu context')
        tgt_ctx = TgtContext(self.src_ctx)
        cloud = test_scene_state.cloud
        cloud = self._clip_cloud(cloud)
        tgt_ctx.set_cld(cloud)
        
        cost_array = batch_tps_rpm_bij(self.src_ctx, tgt_ctx,
                                       T_init=self.rad_init, T_final=self.rad_final, 
                                       outlierfrac=self.outlierfrac, outlierprior=self.outlierprior, 
                                       outliercutoff=settings.OUTLIER_CUTOFF, 
                                       em_iter=self.em_iter, 
                                       component_cost=True)
        costs = dict(zip(self.src_ctx.seg_names, cost_array))
        return costs class BatchGpuTpsRpmRegistrationFactory(TpsRpmRegistrationFactory):
    """
    Similar to TpsRpmRegistrationFactory but batch_register and batch_cost are computed in batch using the GPU
    """
    def __init__(self, demos):
        if not lfd.registration._has_cuda:
            raise NotImplementedError("CUDA not installed")
        raise NotImplementedError
    
    def register(self, demo, test_scene_state, callback=None):
        raise NotImplementedError
    
    def batch_register(self, test_scene_state):
        raise NotImplementedError
    
    def cost(self, demo, test_scene_state):
        raise NotImplementedError
    
    def batch_cost(self, test_scene_state):
        raise NotImplementedError class PoseTrajectoryTransferer(TrajectoryTransferer):
    def __init__(self, sim, 
                 beta_pos=settings.BETA_POS, 
                 beta_rot=settings.BETA_ROT, 
                 gamma=settings.GAMMA, 
                 use_collision_cost=settings.USE_COLLISION_COST, 
                 init_trajectory_transferer=None):
        super(PoseTrajectoryTransferer, self).__init__(sim, beta_pos, gamma, use_collision_cost, init_trajectory_transferer=init_trajectory_transferer)
        self.beta_rot = beta_rot
        
    def transfer(self, reg, demo, plotting=False):
        handles = []
        if plotting:
            demo_cloud = demo.scene_state.cloud
            test_cloud = reg.test_scene_state.cloud
            demo_color = demo.scene_state.color
            test_color = reg.test_scene_state.color
            handles.append(self.sim.env.plot3(demo_cloud[:,:3], 2, demo_color if demo_color is not None else (1,0,0)))
            handles.append(self.sim.env.plot3(test_cloud[:,:3], 2, test_color if test_color is not None else (0,0,1)))
            self.sim.viewer.Step()
        
        active_lr = ""
        for lr in 'lr':
            if lr in demo.aug_traj.lr2arm_traj and sim_util.arm_moved(demo.aug_traj.lr2arm_traj[lr]):
                active_lr += lr
        _, timesteps_rs = sim_util.unif_resample(np.c_[(1./settings.JOINT_LENGTH_PER_STEP) * np.concatenate([demo.aug_traj.lr2arm_traj[lr] for lr in active_lr], axis=1), 
                                                       (1./settings.FINGER_CLOSE_RATE) * np.concatenate([demo.aug_traj.lr2finger_traj[lr] for lr in active_lr], axis=1)], 
                                                 1.)
        demo_aug_traj_rs = demo.aug_traj.get_resampled_traj(timesteps_rs)

        if self.init_trajectory_transferer:
            warm_init_traj = self.init_trajectory_transferer.transfer(reg, demo, plotting=plotting)

        manip_name = ""
        ee_link_names = []
        transformed_ee_trajs_rs = []
        init_traj = np.zeros((len(timesteps_rs),0))
        for lr in active_lr:
            arm_name = {"l":"leftarm", "r":"rightarm"}[lr]
            ee_link_name = "%s_gripper_tool_frame"%lr
            
            if manip_name:
                manip_name += "+"
            manip_name += arm_name
            ee_link_names.append(ee_link_name)
            
            if self.init_trajectory_transferer:
                init_traj = np.c_[init_traj, warm_init_traj.lr2arm_traj[lr]]
            else:
                init_traj = np.c_[init_traj, demo_aug_traj_rs.lr2arm_traj[lr]]
            
            transformed_ee_traj_rs = reg.f.transform_hmats(demo_aug_traj_rs.lr2ee_traj[lr])
            transformed_ee_trajs_rs.append(transformed_ee_traj_rs)
            
            if plotting:
                handles.append(self.sim.env.drawlinestrip(demo.aug_traj.lr2ee_traj[lr][:,:3,3], 2, (1,0,0)))
                handles.append(self.sim.env.drawlinestrip(demo_aug_traj_rs.lr2ee_traj[lr][:,:3,3], 2, (1,1,0)))
                handles.append(self.sim.env.drawlinestrip(transformed_ee_traj_rs[:,:3,3], 2, (0,1,0)))
                self.sim.viewer.Step()
        
        if not self.init_trajectory_transferer:
            # modify the shoulder joint angle of init_traj to be the limit (highest arm) because this usually gives a better local optima (but this might not be the right thing to do)
            dof_inds = sim_util.dof_inds_from_name(self.sim.robot, manip_name)
            joint_ind = self.sim.robot.GetJointIndex("%s_shoulder_lift_joint"%lr)
            init_traj[:,dof_inds.index(joint_ind)] = self.sim.robot.GetDOFLimits([joint_ind])[0][0]

        print "planning pose trajectory following"
        test_traj, obj_value, pose_errs = planning.plan_follow_trajs(self.sim.robot, manip_name, ee_link_names, transformed_ee_trajs_rs, init_traj, 
                                                                       start_fixed=False,
                                                                       use_collision_cost=self.use_collision_cost,
                                                                       beta_pos=self.beta_pos, beta_rot=self.beta_rot)

        # the finger trajectory is the same for the demo and the test trajectory
        for lr in active_lr:
            finger_name = "%s_gripper_l_finger_joint"%lr
            manip_name += "+" + finger_name
            test_traj = np.c_[test_traj, demo_aug_traj_rs.lr2finger_traj[lr]]

        full_traj = (test_traj, sim_util.dof_inds_from_name(self.sim.robot, manip_name))
        test_aug_traj = demonstration.AugmentedTrajectory.create_from_full_traj(self.sim.robot, full_traj, lr2open_finger_traj=demo_aug_traj_rs.lr2open_finger_traj, lr2close_finger_traj=demo_aug_traj_rs.lr2close_finger_traj)
        
        if plotting:
            for lr in active_lr:
                handles.append(self.sim.env.drawlinestrip(test_aug_traj.lr2ee_traj[lr][:,:3,3], 2, (0,0,1)))
            self.sim.viewer.Step()
        
        return test_aug_traj class FingerTrajectoryTransferer(TrajectoryTransferer):
    def __init__(self, sim, 
                 beta_pos=settings.BETA_POS, 
                 gamma=settings.GAMMA, 
                 use_collision_cost=settings.USE_COLLISION_COST, 
                 init_trajectory_transferer=None):
        super(FingerTrajectoryTransferer, self).__init__(sim, beta_pos, gamma, use_collision_cost, init_trajectory_transferer=init_trajectory_transferer)

    def transfer(self, reg, demo, plotting=False):
        handles = []
        if plotting:
            demo_cloud = demo.scene_state.cloud
            test_cloud = reg.test_scene_state.cloud
            demo_color = demo.scene_state.color
            test_color = reg.test_scene_state.color
            handles.append(self.sim.env.plot3(demo_cloud[:,:3], 2, demo_color if demo_color is not None else (1,0,0)))
            handles.append(self.sim.env.plot3(test_cloud[:,:3], 2, test_color if test_color is not None else (0,0,1)))
            self.sim.viewer.Step()
        
        active_lr = ""
        for lr in 'lr':
            if lr in demo.aug_traj.lr2arm_traj and sim_util.arm_moved(demo.aug_traj.lr2arm_traj[lr]):
                active_lr += lr
        _, timesteps_rs = sim_util.unif_resample(np.c_[(1./settings.JOINT_LENGTH_PER_STEP) * np.concatenate([demo.aug_traj.lr2arm_traj[lr] for lr in active_lr], axis=1), 
                                                       (1./settings.FINGER_CLOSE_RATE) * np.concatenate([demo.aug_traj.lr2finger_traj[lr] for lr in active_lr], axis=1)], 
                                                 1.)
        demo_aug_traj_rs = demo.aug_traj.get_resampled_traj(timesteps_rs)

        if self.init_trajectory_transferer:
            warm_init_traj = self.init_trajectory_transferer.transfer(reg, demo, plotting=plotting)
        
        manip_name = ""
        flr2finger_link_names = []
        flr2transformed_finger_pts_trajs_rs = []
        init_traj = np.zeros((len(timesteps_rs),0))
        for lr in active_lr:
            arm_name = {"l":"leftarm", "r":"rightarm"}[lr]
            finger_name = "%s_gripper_l_finger_joint"%lr
            
            if manip_name:
                manip_name += "+"
            manip_name += arm_name + "+" + finger_name
            
            if self.init_trajectory_transferer:
                init_traj = np.c_[init_traj, warm_init_traj.lr2arm_traj[lr], warm_init_traj.lr2finger_traj[lr]]
            else:
                init_traj = np.c_[init_traj, demo_aug_traj_rs.lr2arm_traj[lr], demo_aug_traj_rs.lr2finger_traj[lr]]
            
            if plotting:
                handles.append(self.sim.env.drawlinestrip(demo.aug_traj.lr2ee_traj[lr][:,:3,3], 2, (1,0,0)))
                handles.append(self.sim.env.drawlinestrip(demo_aug_traj_rs.lr2ee_traj[lr][:,:3,3], 2, (1,1,0)))
                transformed_ee_traj_rs = reg.f.transform_hmats(demo_aug_traj_rs.lr2ee_traj[lr])
                handles.append(self.sim.env.drawlinestrip(transformed_ee_traj_rs[:,:3,3], 2, (0,1,0)))
                self.sim.viewer.Step()

            flr2demo_finger_pts_traj_rs = sim_util.get_finger_pts_traj(self.sim.robot, lr, (demo_aug_traj_rs.lr2ee_traj[lr], demo_aug_traj_rs.lr2finger_traj[lr]))
            
            flr2transformed_finger_pts_traj_rs = {}
            flr2finger_link_name = {}
            flr2finger_rel_pts = {}
            for finger_lr in 'lr':
                flr2transformed_finger_pts_traj_rs[finger_lr] = reg.f.transform_points(np.concatenate(flr2demo_finger_pts_traj_rs[finger_lr], axis=0)).reshape((-1,4,3))
                flr2finger_link_name[finger_lr] = "%s_gripper_%s_finger_tip_link"%(lr,finger_lr)
                flr2finger_rel_pts[finger_lr] = sim_util.get_finger_rel_pts(finger_lr)
            flr2finger_link_names.append(flr2finger_link_name)
            flr2transformed_finger_pts_trajs_rs.append(flr2transformed_finger_pts_traj_rs)
            
            if plotting:
                handles.extend(sim_util.draw_finger_pts_traj(self.sim, flr2demo_finger_pts_traj_rs, (1,1,0)))
                handles.extend(sim_util.draw_finger_pts_traj(self.sim, flr2transformed_finger_pts_traj_rs, (0,1,0)))
                self.sim.viewer.Step()
        
        if not self.init_trajectory_transferer:
            # modify the shoulder joint angle of init_traj to be the limit (highest arm) because this usually gives a better local optima (but this might not be the right thing to do)
            dof_inds = sim_util.dof_inds_from_name(self.sim.robot, manip_name)
            joint_ind = self.sim.robot.GetJointIndex("%s_shoulder_lift_joint"%lr)
            init_traj[:,dof_inds.index(joint_ind)] = self.sim.robot.GetDOFLimits([joint_ind])[0][0]
        
        print "planning finger trajectory following"
        test_traj, obj_value, rel_pts_costs = planning.plan_follow_finger_pts_trajs(self.sim.robot, manip_name, 
                                                                              flr2finger_link_names, flr2finger_rel_pts, 
                                                                              flr2transformed_finger_pts_trajs_rs, init_traj, 
                                                                              use_collision_cost=self.use_collision_cost,
                                                                              start_fixed=False,
                                                                              beta_pos=self.beta_pos, gamma=self.gamma)

        full_traj = (test_traj, sim_util.dof_inds_from_name(self.sim.robot, manip_name))
        test_aug_traj = demonstration.AugmentedTrajectory.create_from_full_traj(self.sim.robot, full_traj, lr2open_finger_traj=demo_aug_traj_rs.lr2open_finger_traj, lr2close_finger_traj=demo_aug_traj_rs.lr2close_finger_traj)
        
        if plotting:
            for lr in active_lr:
                handles.append(self.sim.env.drawlinestrip(test_aug_traj.lr2ee_traj[lr][:,:3,3], 2, (0,0,1)))
                flr2test_finger_pts_traj = sim_util.get_finger_pts_traj(self.sim.robot, lr, full_traj)
                handles.extend(sim_util.draw_finger_pts_traj(self.sim, flr2test_finger_pts_traj, (0,0,1)))
            self.sim.viewer.Step()
        
        return test_aug_traj class TwoStepRegistrationAndTrajectoryTransferer(RegistrationAndTrajectoryTransferer):
    def transfer(self, demo, test_scene_state, callback=None, plotting=False):
        reg = self.registration_factory.register(demo, test_scene_state, callback=callback)
        test_aug_traj = self.trajectory_transferer.transfer(reg, demo, plotting=plotting)
        return test_aug_traj class UnifiedRegistrationAndTrajectoryTransferer(RegistrationAndTrajectoryTransferer):
    def __init__(self, registration_factory, trajectory_transferer, 
                 alpha=settings.ALPHA, 
                 beta_pos=settings.BETA_POS, 
                 gamma=settings.GAMMA, 
                 use_collision_cost=settings.USE_COLLISION_COST, 
                 init_trajectory_transferer=None):
        super(UnifiedRegistrationAndTrajectoryTransferer, self).__init__(registration_factory, trajectory_transferer)
        if not isinstance(registration_factory, registration.TpsRpmRegistrationFactory):
            raise NotImplementedError("UnifiedRegistrationAndTrajectoryTransferer only supports TpsRpmRegistrationFactory")
        if not isinstance(trajectory_transferer, transfer.FingerTrajectoryTransferer):
            raise NotImplementedError("UnifiedRegistrationAndTrajectoryTransferer only supports FingerTrajectoryTransferer")
        self.sim = trajectory_transferer.sim
        self.alpha = alpha
        self.beta_pos = beta_pos
        self.gamma = gamma
        self.use_collision_cost = use_collision_cost
        self.init_trajectory_transferer = init_trajectory_transferer

    def transfer(self, demo, test_scene_state, callback=None, plotting=False):
        reg = self.registration_factory.register(demo, test_scene_state, callback=callback)

        handles = []
        if plotting:
            demo_cloud = demo.scene_state.cloud
            test_cloud = reg.test_scene_state.cloud
            demo_color = demo.scene_state.color
            test_color = reg.test_scene_state.color
            handles.append(self.sim.env.plot3(demo_cloud[:,:3], 2, demo_color if demo_color is not None else (1,0,0)))
            handles.append(self.sim.env.plot3(test_cloud[:,:3], 2, test_color if test_color is not None else (0,0,1)))
            self.sim.viewer.Step()
        
        active_lr = ""
        for lr in 'lr':
            if lr in demo.aug_traj.lr2arm_traj and sim_util.arm_moved(demo.aug_traj.lr2arm_traj[lr]):
                active_lr += lr
        _, timesteps_rs = sim_util.unif_resample(np.c_[(1./settings.JOINT_LENGTH_PER_STEP) * np.concatenate([demo.aug_traj.lr2arm_traj[lr] for lr in active_lr], axis=1), 
                                                       (1./settings.FINGER_CLOSE_RATE) * np.concatenate([demo.aug_traj.lr2finger_traj[lr] for lr in active_lr], axis=1)], 
                                                 1.)
        demo_aug_traj_rs = demo.aug_traj.get_resampled_traj(timesteps_rs)

        if self.init_trajectory_transferer:
            warm_init_traj = self.init_trajectory_transferer.transfer(reg, demo, plotting=plotting)
        
        manip_name = ""
        flr2finger_link_names = []
        flr2demo_finger_pts_trajs_rs = []
        init_traj = np.zeros((len(timesteps_rs),0))
        for lr in active_lr:
            arm_name = {"l":"leftarm", "r":"rightarm"}[lr]
            finger_name = "%s_gripper_l_finger_joint"%lr
            
            if manip_name:
                manip_name += "+"
            manip_name += arm_name + "+" + finger_name
            
            if self.init_trajectory_transferer:
                init_traj = np.c_[init_traj, warm_init_traj.lr2arm_traj[lr], warm_init_traj.lr2finger_traj[lr]]
            else:
                init_traj = np.c_[init_traj, demo_aug_traj_rs.lr2arm_traj[lr], demo_aug_traj_rs.lr2finger_traj[lr]]
            
            if plotting:
                handles.append(self.sim.env.drawlinestrip(demo.aug_traj.lr2ee_traj[lr][:,:3,3], 2, (1,0,0)))
                handles.append(self.sim.env.drawlinestrip(demo_aug_traj_rs.lr2ee_traj[lr][:,:3,3], 2, (1,1,0)))
                transformed_ee_traj_rs = reg.f.transform_hmats(demo_aug_traj_rs.lr2ee_traj[lr])
                handles.append(self.sim.env.drawlinestrip(transformed_ee_traj_rs[:,:3,3], 2, (0,1,0)))
                self.sim.viewer.Step()

            flr2demo_finger_pts_traj_rs = sim_util.get_finger_pts_traj(self.sim.robot, lr, (demo_aug_traj_rs.lr2ee_traj[lr], demo_aug_traj_rs.lr2finger_traj[lr]))
            flr2demo_finger_pts_trajs_rs.append(flr2demo_finger_pts_traj_rs)
            
            flr2transformed_finger_pts_traj_rs = {}
            flr2finger_link_name = {}
            flr2finger_rel_pts = {}
            for finger_lr in 'lr':
                flr2transformed_finger_pts_traj_rs[finger_lr] = reg.f.transform_points(np.concatenate(flr2demo_finger_pts_traj_rs[finger_lr], axis=0)).reshape((-1,4,3))
                flr2finger_link_name[finger_lr] = "%s_gripper_%s_finger_tip_link"%(lr,finger_lr)
                flr2finger_rel_pts[finger_lr] = sim_util.get_finger_rel_pts(finger_lr)
            flr2finger_link_names.append(flr2finger_link_name)
            
            if plotting:
                handles.extend(sim_util.draw_finger_pts_traj(self.sim, flr2demo_finger_pts_traj_rs, (1,1,0)))
                handles.extend(sim_util.draw_finger_pts_traj(self.sim, flr2transformed_finger_pts_traj_rs, (0,1,0)))
                self.sim.viewer.Step()
        
        if not self.init_trajectory_transferer:
            # modify the shoulder joint angle of init_traj to be the limit (highest arm) because this usually gives a better local optima (but this might not be the right thing to do)
            dof_inds = sim_util.dof_inds_from_name(self.sim.robot, manip_name)
            joint_ind = self.sim.robot.GetJointIndex("%s_shoulder_lift_joint"%lr)
            init_traj[:,dof_inds.index(joint_ind)] = self.sim.robot.GetDOFLimits([joint_ind])[0][0]

        print "planning joint TPS and finger points trajectory following"
        test_traj, obj_value, tps_rel_pts_costs, tps_cost = planning.joint_fit_tps_follow_finger_pts_trajs(self.sim.robot, manip_name, 
                                                                              flr2finger_link_names, flr2finger_rel_pts, 
                                                                              flr2demo_finger_pts_trajs_rs, init_traj, 
                                                                              reg.f, 
                                                                              use_collision_cost=self.use_collision_cost,
                                                                              start_fixed=False,
                                                                              alpha=self.alpha, beta_pos=self.beta_pos, gamma=self.gamma)

        full_traj = (test_traj, sim_util.dof_inds_from_name(self.sim.robot, manip_name))
        test_aug_traj = demonstration.AugmentedTrajectory.create_from_full_traj(self.sim.robot, full_traj, lr2open_finger_traj=demo_aug_traj_rs.lr2open_finger_traj, lr2close_finger_traj=demo_aug_traj_rs.lr2close_finger_traj)
        
        if plotting:
            for lr in active_lr:
                flr2new_transformed_finger_pts_traj_rs = {}
                for finger_lr in 'lr':
                    flr2new_transformed_finger_pts_traj_rs[finger_lr] = reg.f.transform_points(np.concatenate(flr2demo_finger_pts_traj_rs[finger_lr], axis=0)).reshape((-1,4,3))
                handles.extend(sim_util.draw_finger_pts_traj(self.sim, flr2new_transformed_finger_pts_traj_rs, (0,1,1)))
                handles.append(self.sim.env.drawlinestrip(test_aug_traj.lr2ee_traj[lr][:,:3,3], 2, (0,0,1)))
                flr2test_finger_pts_traj = sim_util.get_finger_pts_traj(self.sim.robot, lr, full_traj)
                handles.extend(sim_util.draw_finger_pts_traj(self.sim, flr2test_finger_pts_traj, (0,0,1)))
            self.sim.viewer.Step()
        
        return test_aug_traj class Transform(object):
class Pose(Transform):                  # 2.5D transform
class Point:
class LineXY:
class LineSeg(LineXY):
class SymbolGenerator:
    def __init__(self, matrix):
    def inverse(self):
    def __neg__(self):
    def compose(self, trans):
    def __mul__(self, other):
    def pose(self, zthr = 0.01, fail = True):
    def point(self):
    def applyToPoint(self, point):
    def __call__(self, point):
    def __repr__(self):
    def shortStr(self, trim = False):
    def __init__(self, x, y, z, theta):
    def initTrans(self):
    def setX(self, x):
    def setY(self, y):
    def setZ(self, z):
    def setTheta(self, theta):
    def average(self, other, alpha):
    def point(self):
    def pose(self, fail = False):
    def near(self, pose, distEps, angleEps):
    def diff(self, pose):
    def distance(self, pose):
    def totalDist(self, pose, angleScale = 1):
    def inverse(self):
    def xyztTuple(self):
    def corrupt(self, e, eAng = None):
        def corrupt(x, e):
    def corruptGauss(self, mu, sigma, noZ = False):
        def corrupt(x):
    def __repr__(self):
    def shortStr(self, trim = False):
    def __eq__(self, other):
    def __hash__(self):
    def __init__(self, x, y, z, w=1.0):
    def matrix(self):
    def isNear(self, point, distEps):
    def distance(self, point):
    def distanceXY(self, point):
    def distanceSq(self, point):
    def distanceSqXY(self, point):
    def magnitude(self):
    def xyzTuple(self):
    def pose(self, angle = 0.0):
    def point(self):
    def __repr__(self):
    def shortStr(self, trim = False):
    def angleToXY(self, p):
    def add(self, point):
    def __add__(self, point):
    def sub(self, point):
    def __sub__(self, point):
    def scale(self, s):
    def __rmul__(self, s):
    def dot(self, p):
    def __init__(self, p1, p2):
    def pointOnLine(self, p, eps):
    def __repr__(self):
    def shortStr(self, trim = False):
    def __init__(self, p1, p2):
    def closestPoint(self, p):
    def distToPoint(self, p):
    def __repr__(self):
def localToGlobal(pose, point):
def localPoseToGlobalPose(pose1, pose2):
def globalToLocal(pose, point):
def globalPoseToLocalPose(pose1, pose2):
def sum(items):
def smash(lists):
def within(v1, v2, eps):
def nearAngle(a1,a2,eps):
def nearlyEqual(x,y):
def fixAnglePlusMinusPi(a):
def fixAngle02Pi(a):
def reverseCopy(items):
def dotProd(a, b):
def argmax(l, f):
def argmaxWithVal(l, f):
def argmaxIndex(l, f = lambda x: x):
def argmaxIndexWithTies(l, f = lambda x: x):
def randomMultinomial(dist):
def clip(v, vMin, vMax):
def flatten(M):
def sign(x):
def make2DArray(dim1, dim2, initValue):
def make2DArrayFill(dim1, dim2, initFun):
def make3DArray(dim1, dim2, dim3, initValue):
def mapArray3D(array, f):
def makeVector(dim, initValue):
def makeVectorFill(dim, initFun):
def prettyString(struct):
def swapRange(x, y):
def avg(a, b):
def recoverPath(volume, start, end):
    def __init__(self):
    def gensym(self, prefix = 'i'):
def logGaussian(x, mu, sigma):
def gaussian(x, mu, sigma):
def lineIndices((i0, j0), (i1, j1)):
def angleDiff(x, y):
def inRange(v, r):
def rangeOverlap(r1, r2):
def rangeIntersect(r1, r2):
def average(stuff):
def tuplify(x):
def squash(listOfLists):
def angleAverage(th1, th2, alpha):
def floatRange(lo, hi, stepsize):
def euclideanDistance(x, y):
def pop(x):
def isIterable(x):
def tangentSpaceAdd(a, b):
def scalarMult(l, c):
def componentAdd(a, b):
def componentSubtract(a, b): def request_int_in_range(too_high_val):
    def __init__(self, step, name, data, description=""):
    def __init__(self, filename, max_unwritten=10):
    def append(self, entry):
    def __call__(self, *args, **kwargs):
    def flush(self):
    def close(self):
class ExecutionLogEntry(object):
class ExecutionLog(object): MIN_SEG_LEN = 3
    S = skeletonize_point_cloud(xyzs)
    C = make_cost_matrix(segs3d)
    PG = make_path_graph(C, [len(path) for path in segs3d])
    S = skeletonize_point_cloud(xyzs)
    G = nx.Graph()
    G = points_to_graph(xyzs, point_conn_dist)
    S = skeletonize_graph(G, cluster_size)
MAX_COST = 50
SKIP_COST = 1000 # per meter
MAX_NEIGHBORS=5
WEIGHTS = np.r_[
    3, # ang_a_disp: angle between end of source and displacement targ-source
    3, # ang_b_disp: angle between end of target and tisplacement source-targ
    30, # ang_a_b: angle between end of source and end of targ
    30,# fwddist_b_a: positive part of displacement `dot` source direction
    50,# backdist_b_a: negative part of ...
    50,# perpdist_b_a: displacement perpendicular to source direction
    30,# fwddist_a_b
    50,# fwddist_a_b
    50] # perpdist_a_b
    N = len(paths)
    M,N = cost_matrix.shape
def find_path_through_point_cloud(xyzs, plotting=False, perturb_peak_dist=None, num_perturb_points=7):
def prune_skeleton(S, segs):
def get_skeleton_points(xyzs):
def plot_paths_2d(paths3d):
def unif_resample(x,n=None,tol=0,deg=None, seg_len = .02):
def points_to_graph(xyzs, max_dist):
def skeletonize_graph(G, resolution):
def largest_connected_component(G):
def skeletonize_point_cloud(xyzs, point_conn_dist = .025, cluster_size = .04):
def calc_distances(G):
def calc_reeb_partitions(G, resolution):
def plot_graph_3d(G):
def get_segments(G):
def get_segments_connected(G):
def get_junctions(G):
def get_segment_from_junction(G, junc, nei):
def calc_path_features(pos_a, dir_a, pos_b, dir_b):
def sqnorm(x):
def ang_between(x,y):
def normalized(x):
def start_pos_dir(path):
def end_pos_dir(path):
def estimate_end_directions(points, tol):
def make_cost_matrix(paths):
def make_path_graph(cost_matrix, lengths):
def longest_path_between(G,start,used,target):  
def longest_path_from(G,start,used):  
def longest_path_through_segment_graph(G):
def remove_duplicate_rows(mat):
def norms(x,ax): def isFig8Knot(rope_nodes):
    (crossings, crossings_links_inds, cross_pairs, rope_closed) = calculateCrossings(rope_nodes)
    # simplify crossings a bit
    crossings, crossings_links_inds, cross_pairs = remove_consecutive_crossings(crossings, crossings_links_inds, cross_pairs)
    crossings, crossings_links_inds, cross_pairs = remove_consecutive_cross_pairs(crossings, crossings_links_inds, cross_pairs)
    s = crossingsToString(crossings)
    
    knot_topologies = ['uouououo', 'ouuuoouuooou', 'ouuouuoouo']
    for top in knot_topologies:
        flipped_top = top.replace('u','t').replace('o','u').replace('t','o')
        if top in s and crossings_match(cross_pairs, top, s):
            return True
        if top[::-1] in s and crossings_match(cross_pairs, top[::-1], s):
            return True
        if flipped_top in s and crossings_match(cross_pairs, flipped_top, s):
            return True
        if flipped_top[::-1] in s and crossings_match(cross_pairs, flipped_top[::-1], s):
            return True

    return False class Transform(object):
class Pose(Transform):                  # 2.5D transform
class Point:
class LineXY:
class LineSeg(LineXY):
class SymbolGenerator:
    def __init__(self, matrix):
    def inverse(self):
    def __neg__(self):
    def compose(self, trans):
    def __mul__(self, other):
    def pose(self, zthr = 0.01, fail = True):
    def point(self):
    def applyToPoint(self, point):
    def __call__(self, point):
    def __repr__(self):
    def shortStr(self, trim = False):
    def __init__(self, x, y, z, theta):
    def initTrans(self):
    def setX(self, x):
    def setY(self, y):
    def setZ(self, z):
    def setTheta(self, theta):
    def average(self, other, alpha):
    def point(self):
    def pose(self, fail = False):
    def near(self, pose, distEps, angleEps):
    def diff(self, pose):
    def distance(self, pose):
    def totalDist(self, pose, angleScale = 1):
    def inverse(self):
    def xyztTuple(self):
    def corrupt(self, e, eAng = None):
        def corrupt(x, e):
    def corruptGauss(self, mu, sigma, noZ = False):
        def corrupt(x):
    def __repr__(self):
    def shortStr(self, trim = False):
    def __eq__(self, other):
    def __hash__(self):
    def __init__(self, x, y, z, w=1.0):
    def matrix(self):
    def isNear(self, point, distEps):
    def distance(self, point):
    def distanceXY(self, point):
    def distanceSq(self, point):
    def distanceSqXY(self, point):
    def magnitude(self):
    def xyzTuple(self):
    def pose(self, angle = 0.0):
    def point(self):
    def __repr__(self):
    def shortStr(self, trim = False):
    def angleToXY(self, p):
    def add(self, point):
    def __add__(self, point):
    def sub(self, point):
    def __sub__(self, point):
    def scale(self, s):
    def __rmul__(self, s):
    def dot(self, p):
    def __init__(self, p1, p2):
    def pointOnLine(self, p, eps):
    def __repr__(self):
    def shortStr(self, trim = False):
    def __init__(self, p1, p2):
    def closestPoint(self, p):
    def distToPoint(self, p):
    def __repr__(self):
def localToGlobal(pose, point):
def localPoseToGlobalPose(pose1, pose2):
def globalToLocal(pose, point):
def globalPoseToLocalPose(pose1, pose2):
def sum(items):
def smash(lists):
def within(v1, v2, eps):
def nearAngle(a1,a2,eps):
def nearlyEqual(x,y):
def fixAnglePlusMinusPi(a):
def fixAngle02Pi(a):
def reverseCopy(items):
def dotProd(a, b):
def argmax(l, f):
def argmaxWithVal(l, f):
def argmaxIndex(l, f = lambda x: x):
def argmaxIndexWithTies(l, f = lambda x: x):
def randomMultinomial(dist):
def clip(v, vMin, vMax):
def flatten(M):
def sign(x):
def make2DArray(dim1, dim2, initValue):
def make2DArrayFill(dim1, dim2, initFun):
def make3DArray(dim1, dim2, dim3, initValue):
def mapArray3D(array, f):
def makeVector(dim, initValue):
def makeVectorFill(dim, initFun):
def prettyString(struct):
def swapRange(x, y):
def avg(a, b):
def recoverPath(volume, start, end):
    def __init__(self):
    def gensym(self, prefix = 'i'):
def logGaussian(x, mu, sigma):
def gaussian(x, mu, sigma):
def lineIndices((i0, j0), (i1, j1)):
def angleDiff(x, y):
def inRange(v, r):
def rangeOverlap(r1, r2):
def rangeIntersect(r1, r2):
def average(stuff):
def tuplify(x):
def squash(listOfLists):
def angleAverage(th1, th2, alpha):
def floatRange(lo, hi, stepsize):
def euclideanDistance(x, y):
def pop(x):
def isIterable(x):
def tangentSpaceAdd(a, b):
def scalarMult(l, c):
def componentAdd(a, b):
def componentSubtract(a, b):
import os.path
import h5py
import atexit
import random
import trajoptpy
import numpy as np
from string import lower
from lfd.environment import sim_util
from lfd.environment import settings
from lfd.demonstration.demonstration import SceneState, GroundTruthRopeSceneState, AugmentedTrajectory, Demonstration
from lfd.environment.simulation import DynamicRopeSimulationRobotWorld
from lfd.environment.simulation_object import XmlSimulationObject, BoxSimulationObject, RopeSimulationObject
from lfd.environment.environment import LfdEnvironment, GroundTruthRopeLfdEnvironment
from lfd.registration.registration import TpsRpmBijRegistrationFactory, TpsRpmRegistrationFactory, TpsSegmentRegistrationFactory, BatchGpuTpsRpmBijRegistrationFactory, BatchGpuTpsRpmRegistrationFactory
from lfd.transfer.transfer import PoseTrajectoryTransferer, FingerTrajectoryTransferer
from lfd.transfer.registration_transfer import TwoStepRegistrationAndTrajectoryTransferer, UnifiedRegistrationAndTrajectoryTransferer
from lfd.rapprentice import util
from lfd.rapprentice import task_execution, rope_initialization
from lfd.rapprentice.knot_classifier import isFig8Knot
from lfd.rapprentice.util import redprint
#!/usr/bin/env python

from __future__ import division








# Usage: python scripts/label.py --animation 2 label bigdata/misc/overhand_actions.h5 data/misc/Sep14_train2.h5 finger bij --gpu

class GlobalVars:
    exec_log = None
    actions = None
    actions_cache = None
    demos = None

# @profile
def label_demos_parallel(args, transferer, lfd_env, sim):
    outfile = h5py.File(args.eval.outfile, 'a')

    rope_params = sim_util.RopeParams()
    rope_params.radius = settings.ROPE_RADIUS_THICK
    if args.eval.rope_param_radius is not None:
        rope_params.radius = args.eval.rope_param_radius
    if args.eval.rope_param_angStiffness is not None:
        rope_params.angStiffness = args.eval.rope_param_angStiffness

    # TODO pass in rope params to sample_rope_state
    (init_rope_nodes, demo_id) = sample_rope_state(
        args.eval, sim, args.animation)
    resample = False



    labeled_data = []


    while True:
        if args.animation:
            sim.viewer.Step()
        sim_state = sim.get_state()
        sim.set_state(sim_state)

        scene_state = lfd_env.observe_scene()

        # plot cloud of the test scene
        handles = []
        if args.plotting:
            handles.append(
                sim.env.plot3(
                    scene_state.cloud[
                        :,
                        :3],
                    2,
                    scene_state.color if scene_state.color is not None else (
                        0,
                        0,
                        1)))
            sim.viewer.Step()
        costs = transferer.registration_factory.batch_cost(scene_state)
        best_keys = sorted(costs, key=costs.get)
        for seg_name in best_keys:
            traj = transferer.transfer(GlobalVars.demos[seg_name],
                                                scene_state,
                                                plotting=args.plotting)
        

            feasible, misgrasp = lfd_env.execute_augmented_trajectory(
                traj, step_viewer=args.animation, interactive=args.interactive)
            sim_util.reset_arms_to_side(sim)
            sim.settle(step_viewer=args.animation)

            if not feasible or misgrasp:
                print 'Feasible: ', feasible
                print 'Misgrasp: ', misgrasp
                if misgrasp:
                    sim.set_state(sim_state)
                    continue
            print "y accepts this action"
            print "n rejects this action"
            print "r resamples rope state"
            print "f to save this as a failure"
            print "C-c safely quits"
            user_input = lower(raw_input("What to do?"))
            success = False
            if user_input == 'y':
                success = True
                labeled_data.append((scene_state,seg_name))

                if isFig8Knot(get_rope_nodes(sim)):