class Demonstration(object):
class SceneState(object):
class GroundTruthRopeSceneState(SceneState):
class AugmentedTrajectory(object):
    def __init__(self, name, scene_state, aug_traj):
    def __repr__(self):
    def __init__(self, full_cloud, id=None, full_color=None, downsample_size=0):
    def get_unique_id():
    def __repr__(self):
    def __init__(self, rope_nodes, radius, upsample=0, upsample_rad=1, downsample_size=0):
    def __init__(self, lr2arm_traj=None, lr2finger_traj=None, lr2ee_traj=None, lr2open_finger_traj=None, lr2close_finger_traj=None):
    def __eq__(self, other):
    def __ne__(self, other):
    def create_from_full_traj(robot, full_traj, lr2open_finger_traj=None, lr2close_finger_traj=None):
    def get_full_traj(self, robot):
    def get_resampled_traj(self, timesteps_rs):
    def __repr__(self): def redprint(msg):
def yellowprint(msg):
    def parse_args(self, *args, **kw):
    def __init__(self, adict):
    def __init__(self):
    def __enter__(self):
    def __exit__(self, *_):
class ArgumentParser(argparse.ArgumentParser):
class Bunch(object):
class suppress_stdout(object): def interp2d(x,xp,yp):
def interp_mat(x, xp):
def normalize(x):
def normr(x):
def normc(x):
def norms(x,ax):
def intround(x):
def deriv(x):
def linspace2d(start,end,n):
def remove_duplicate_rows(mat):
def invertHmat(hmat):
    T = len(x)
    R = hmat[:3,:3] def lerp(a, b, fracs):
def adaptive_resample(x, t = None, max_err = np.inf, max_dx = np.inf, max_dt = np.inf, normfunc = None):
def get_velocities(positions, times, tol):
def smooth_positions(positions, tol):
def unif_resample(x,n,weights,tol=.001,deg=3):    
def test_resample():
def test_resample_big():
def interp_quats(newtimes, oldtimes, oldquats):
def interp_hmats(newtimes, oldtimes, oldhmats): 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): 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 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 DynamicSimulationRobotWorld(DynamicSimulation, RobotWorld):
    def __init__(self, env=None, T_w_k=None, range_k=2.):
        """
        T_w_k: world transform of the depth camera
        range_k: length of the rays. 2 meters by default
        """
        self.constraints = {"l": [], "r": []}
        self.constraints_links = {"l": [], "r": []}
        super(DynamicSimulationRobotWorld, self).__init__(env=env)
        self.T_w_k = T_w_k
        self.range_k = range_k
    
    def observe_cloud(self):
        if self.T_w_k is None:
            if self.robot is None:
                raise RuntimeError("Can't observe cloud when there is no robot")
            else:
                from lfd.rapprentice import berkeley_pr2
                self.T_w_k = berkeley_pr2.get_kinect_transform(self.robot)
        
        # camera's parameters
        cx = 320.-.5
        cy = 240.-.5
        f = 525.  # focal length
        w = 640.
        h = 480.
        
        pixel_ij = np.array(np.meshgrid(np.arange(w), np.arange(h))).T.reshape((-1, 2))  # all pixel positions
        rays_to = self.range_k * np.c_[(pixel_ij - np.array([cx, cy])) / f, np.ones(pixel_ij.shape[0])]
        rays_from = np.zeros_like(rays_to)
        # transform the rays from the camera frame to the world frame
        rays_to = rays_to.dot(self.T_w_k[:3,:3].T) + self.T_w_k[:3,3]
        rays_from = rays_from.dot(self.T_w_k[:3,:3].T) + self.T_w_k[:3,3]

        cloud = []
        for sim_obj in self.dyn_sim_objs:
            for bt_obj in sim_obj.get_bullet_objects():
                ray_collisions = self.bt_env.RayTest(rays_from, rays_to, bt_obj)
                
                pts = np.empty((len(ray_collisions), 3))
                for i, ray_collision in enumerate(ray_collisions):
                    pts[i, :] = ray_collision.pt
                cloud.append(pts)
        cloud = np.concatenate(cloud)

        # hack to filter out point below the top of the table. TODO: fix this hack
        table_sim_objs = [sim_obj for sim_obj in self.sim_objs if "table" in sim_obj.names]
        assert len(table_sim_objs) == 1
        table_sim_obj = table_sim_objs[0]
        table_height = table_sim_obj.translation[2] + table_sim_obj.extents[2]
        cloud = cloud[cloud[:, 2] > table_height, :]
        return cloud

    def open_gripper(self, lr, target_val=None, step_viewer=1, max_vel=.02):
        self._remove_constraints(lr)
         
        # generate gripper finger trajectory
        joint_ind = self.robot.GetJoint("%s_gripper_l_finger_joint" % lr).GetDOFIndex()
        start_val = self.robot.GetDOFValues([joint_ind])[0]
        if target_val is None:
            target_val = sim_util.get_binary_gripper_angle(True)
        joint_traj = np.linspace(start_val, target_val, np.ceil(abs(target_val - start_val) / max_vel))

        # execute gripper finger trajectory
        for val in joint_traj:
            self.robot.SetDOFValues([val], [joint_ind])
            self.step()
        if self.viewer and step_viewer:
            self.viewer.Step()
    
    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()
        handles = []
        # add constraints at the points where a ray hits a dynamic link within a distance of grab_dist_thresh
        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)
            
            ray_collisions = [rc for rcs in [from_to_ray_collisions, to_from_ray_collisions] for rc in rcs]

            for rc in ray_collisions:
                if np.linalg.norm(rc.pt - rc.rayFrom) < grab_dist_thresh:
                    link_tf = rc.link.GetTransform()
                    link_tf[:3, 3] = rc.pt
                    self._add_constraints(lr, rc.link, link_tf)
        if self.viewer and step_viewer:
            self.viewer.Step()
    
    def execute_trajectory(self, full_traj, step_viewer=1, interactive=False,
                           max_cart_vel_trans_traj=.05, sim_callback=None):
        # TODO: incorporate other parts of sim_full_traj_maybesim
        if sim_callback is None:
            sim_callback = lambda i: self.step()
        
        traj, dof_inds = full_traj
        
    #     # clip finger joint angles to the binary gripper angles if necessary
    #     for lr in 'lr':
    #         joint_ind = self.robot.GetJoint("%s_gripper_l_finger_joint"%lr).GetDOFIndex()
    #         if joint_ind in dof_inds:
    #             ind = dof_inds.index(joint_ind)
    #             traj[:,ind] = np.minimum(traj[:,ind], get_binary_gripper_angle(True))
    #             traj[:,ind] = np.maximum(traj[:,ind], get_binary_gripper_angle(False))
        
        # in simulation mode, we must make sure to gradually move to the new starting position
        self.robot.SetActiveDOFs(dof_inds)
        curr_vals = self.robot.GetActiveDOFValues()
        transition_traj = np.r_[[curr_vals], [traj[0]]]
        sim_util.unwrap_in_place(transition_traj, dof_inds=dof_inds)
        transition_traj = ropesim.retime_traj(self.robot, dof_inds, transition_traj,
                                              max_cart_vel=max_cart_vel_trans_traj)
        animate_traj.animate_traj(transition_traj, self.robot, restore=False, pause=interactive,
                                  callback=sim_callback, step_viewer=step_viewer if self.viewer else 0)
        
        traj[0] = transition_traj[-1]
        sim_util.unwrap_in_place(traj, dof_inds=dof_inds)
        traj = ropesim.retime_traj(self.robot, dof_inds, traj)  # make the trajectory slow enough for the simulation
    
        animate_traj.animate_traj(traj, self.robot, restore=False, pause=interactive,
                                  callback=sim_callback, step_viewer=step_viewer if self.viewer else 0)
        if self.viewer and step_viewer:
            self.viewer.Step()
        return True
    
    def _create_bullet(self):
        for lr in 'lr':
            if self.constraints[lr] or self.constraints_links[lr]:
                raise RuntimeError("Bullet environment can't be removed while the robot is grasping an object")
        super(DynamicSimulationRobotWorld, self)._create_bullet()

    def _remove_bullet(self):
        for lr in 'lr':
            if self.constraints[lr] or self.constraints_links[lr]:
                raise RuntimeError("Bullet environment can't be removed while the robot is grasping an object")
        super(DynamicSimulationRobotWorld, self)._remove_bullet()
        
    def _get_finger_pts_grid(self, lr, min_sample_dist=0.005):
        sample_grid = None
        flr2finger_pts_grid = {}
        for finger_lr in 'lr':
            world_from_finger = self.robot.GetLink("%s_gripper_%s_finger_tip_link" % (lr, finger_lr)).GetTransform()
            finger_pts = world_from_finger[:3, 3] \
                + sim_util.get_finger_rel_pts(finger_lr).dot(world_from_finger[:3, :3].T)
            pt0 = finger_pts[0 if finger_lr == 'l' else 3][None, :]
            pt1 = finger_pts[1 if finger_lr == 'l' else 2][None, :]
            pt3 = finger_pts[3 if finger_lr == 'l' else 0][None, :]
            if sample_grid is None:
                num_sample_01 = np.round(np.linalg.norm(pt1 - pt0)/min_sample_dist)
                num_sample_03 = np.round(np.linalg.norm(pt3 - pt0)/min_sample_dist)
                sample_grid = np.array(np.meshgrid(np.linspace(0, 1, num_sample_01),
                                                   np.linspace(0, 1, num_sample_03))).T.reshape((-1, 2))
            flr2finger_pts_grid[finger_lr] = pt0 + sample_grid[:, 0][:, None].dot(pt1 - pt0) \
                + sample_grid[:, 1][:, None].dot(pt3 - pt0)
        return flr2finger_pts_grid
    
    def _remove_constraints(self, lr, grab_link=None):
        """
        If grab_link is None, remove all constraints that attaches the lr gripper,
        else remove all constraints that attaches between the lr gripper and grab_link
        """
        num_links_removed = 0
        for (cnt, link) in zip(self.constraints[lr], self.constraints_links[lr]):
            if grab_link is None or link == grab_link:
                self.bt_env.RemoveConstraint(cnt)
                num_links_removed += 1
        # TODO: provide option to color the contrained links and save color before overriding it
        for link in self.constraints_links[lr]:
            if grab_link is None or link == grab_link:
                for geom in link.GetGeometries():
                    geom.SetDiffuseColor([1., 1., 1.])
        if grab_link is None:
            self.constraints[lr] = []
            self.constraints_links[lr] = []
        else:
            if grab_link in self.constraints_links[lr]:
                constraints_links_pairs = zip(*[(cnt, link) for (cnt, link)
                                                in zip(self.constraints[lr], self.constraints_links[lr])
                                                if link != grab_link])
                if constraints_links_pairs:
                    constraints, constraints_links = constraints_links_pairs
                    self.constraints[lr] = list(constraints)
                    self.constraints_links[lr] = list(constraints_links)
                else:
                    self.constraints[lr] = []
                    self.constraints_links[lr] = []
    
    def _add_constraints(self, lr, grab_link, grab_tf=None):
        if grab_tf is None:
            grab_tf = grab_link.GetTransform()
        # TODO: provide option to color the contrained links and save color before overriding it
        for geom in grab_link.GetGeometries():
            geom.SetDiffuseColor([1.,0.,0.])
        for flr in 'lr':
            robot_link = self.robot.GetLink("%s_gripper_%s_finger_tip_link" % (lr, flr))
            cnt = self.bt_env.AddConstraint({
                "type": "generic6dof",
                "params": {
                    "link_a": robot_link,
                    "link_b": grab_link,
                    "frame_in_a": np.linalg.inv(robot_link.GetTransform()).dot(grab_tf),
                    "frame_in_b": np.linalg.inv(grab_link.GetTransform()).dot(grab_tf),
                    "use_linear_reference_frame_a": False,
                    "stop_erp": .8,
                    "stop_cfm": .1,
                    "disable_collision_between_linked_bodies": True,
                }
            })
            self.constraints[lr].append(cnt)
            self.constraints_links[lr].append(grab_link) 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 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 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 def plan_follow_traj(robot, manip_name, ee_link, new_hmats, old_traj, 
                     no_collision_cost_first=False, use_collision_cost=True, start_fixed=False, joint_vel_limits=None,
                     beta_pos=settings.BETA_POS, beta_rot=settings.BETA_ROT, gamma=settings.GAMMA):
def plan_follow_trajs(robot, manip_name, ee_link_names, ee_trajs, old_traj, 
                     no_collision_cost_first=False, use_collision_cost=True, start_fixed=False, joint_vel_limits=None,
                     beta_pos=settings.BETA_POS, beta_rot=settings.BETA_ROT, gamma=settings.GAMMA):
def plan_follow_finger_pts_traj(robot, manip_name, flr2finger_link, flr2finger_rel_pts, flr2finger_pts_traj, old_traj, 
                                no_collision_cost_first=False, use_collision_cost=True, start_fixed=False, joint_vel_limits=None,
                                beta_pos=settings.BETA_POS, gamma=settings.GAMMA):
def plan_follow_finger_pts_trajs(robot, manip_name, flr2finger_link_names, flr2finger_rel_pts, flr2finger_pts_trajs, old_traj, 
                                no_collision_cost_first=False, use_collision_cost=True, start_fixed=False, joint_vel_limits=None,
                                beta_pos=settings.BETA_POS, gamma=settings.GAMMA):
def joint_fit_tps_follow_finger_pts_traj(robot, manip_name, flr2finger_link, flr2finger_rel_pts, flr2finger_pts_traj, old_traj, 
                                              f, closing_pts=None, 
                                              no_collision_cost_first=False, use_collision_cost=True, start_fixed=False, joint_vel_limits=None, 
                                              alpha=settings.ALPHA, beta_pos=settings.BETA_POS, gamma=settings.GAMMA):
def joint_fit_tps_follow_finger_pts_trajs(robot, manip_name, flr2finger_link_names, flr2finger_rel_pts, flr2old_finger_pts_trajs, old_traj, 
                                         f, closing_pts=None,
                                         no_collision_cost_first=False, use_collision_cost=True, start_fixed=False, joint_vel_limits=None,
                                          alpha=settings.ALPHA, beta_pos=settings.BETA_POS, gamma=settings.GAMMA):
        N = f.N class TrajectoryTransferer(object):
    def __init__(self, sim, 
                 beta_pos=settings.BETA_POS, 
                 gamma=settings.GAMMA, 
                 use_collision_cost=settings.USE_COLLISION_COST, 
                 init_trajectory_transferer=None):
        """Inits TrajectoryTransferer
        
        Args:
            sim: StaticSimulation
            beta_pos: penalty coefficient for gripper positions
            gamma: penalty coefficient for joint velocities
            use_collision_cost: if False, collisions are ignored
            init_trajectory_transferer: TrajectoryTransferer used to get a trajectory for initializing the optimization
        """
        self.sim = sim
        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, reg, demo, plotting=False):
        """Transfers demonstration trajectory using the given registration
        
        Args:
            reg: Registration of the demonstration scene onto the test scene
            demo: Demonstration that has the trajectory to transfer
            plotting: if True, visualization is plotted. The color convention is as follows:
                Red:    demonstration
                Yellow: resampled demonstration
                Green:  transformed resampled demonstration
                Blue:   test
        
        Returns:
            The transferred AugmentedTrajectory
        """
        raise NotImplementedError
import sys
import os.path
import h5py
import copy
import openravepy
import numpy as np
import matplotlib.pyplot as plt
import scipy
import trajoptpy
import IPython as ipy
        import time
from scipy import optimize
from pykalman import KalmanFilter
from pykalman.standard import _last_dims
from lfd.demonstration import demonstration
from lfd.util import util
from lfd.rapprentice import math_utils, resampling
from lfd.environment import sim_util
from lfd.demonstration.demonstration import SceneState, AugmentedTrajectory, Demonstration
from lfd.environment.simulation import DynamicSimulationRobotWorld
from lfd.environment.simulation_object import XmlSimulationObject, BoxSimulationObject
from lfd.environment.environment import LfdEnvironment
from lfd.registration.registration import TpsRpmBijRegistrationFactory, TpsRpmRegistrationFactory, TpsSegmentRegistrationFactory, BatchGpuTpsRpmBijRegistrationFactory, BatchGpuTpsRpmRegistrationFactory
from lfd.transfer import planning
from lfd.transfer.transfer import TrajectoryTransferer
                break
        
        lr2Cts[lr] = Cts
        lr2cts[lr] = cts
    return lr2Cts, lr2cts

def processFuncCt(cliptype, pM, pm, empsig, empmu, pD, pd, Ct, ct, Lt):
    if cliptype == "LBFGS":
        def func_minimize(p, empsig, pD):
            L = np.zeros(pD.shape)
            tidx = np.triu_indices(pD.shape[0])
            L[tidx] = p
            Ct = L.T.dot(L)

            prior_m_x = 1e-1
            prior_m_u = 1e-2
             
            prior_wt = 1
            prior_mat = np.diag(np.r_[prior_m_x*np.ones(12), prior_m_u*np.ones(6)])
             
            val = ((empsig + prior_wt*prior_mat)*Ct).sum() \
                - prior_wt*np.log(max(1e-100,np.linalg.det(Ct))) \
                - np.log(max(1e-100,np.linalg.det(Ct+pD)))
            grad = 2*L.dot(-empsig - prior_wt*prior_mat + prior_wt*np.linalg.inv(Ct) + np.linalg.inv(Ct + pD))
            grad = -grad[tidx]
            return val, grad
        
        if not np.any(Ct):
            Ct = pM - pD
        ct = pm - pd;
        
        if not np.any(Lt):
            nval = 1e-4
            mval = 1e3
            
            val, vec = np.linalg.eig((Ct+Ct.T)/2.)
            val = np.real(val)
            vec = np.real(vec)
            val = np.clip(val, nval, mval)
            val = np.diag(val)
            Ct = vec.dot(val).dot(vec.T)
            Lt = scipy.linalg.cholesky((Ct+Ct.T)/2.)
        
        tidx = np.triu_indices(Ct.shape[0])
        theta, _, _ = optimize.fmin_l_bfgs_b(func_minimize, Lt[tidx], fprime=None, args=(empsig, pD), maxfun=20, maxiter=20)
        
        Lt[tidx] = theta
        Ct = Lt.T.dot(Lt)

        ctm = np.linalg.solve(Lt, (np.linalg.solve(Lt.T,((Ct + pD).dot(empmu + np.linalg.solve(Ct + pD, pd))))))
        ct = -Ct.dot(ctm)

        return Ct, ct, Lt
    else:
        raise NotImplementedError

def mass_calculate(dof_val_mus, dof_val_sigmas, aligned_trajs):
    #Regression code for getting masses
    ipy.embed()
    aligned_trajs = aligned_trajs[:,:,:18]
    n_demos, t_steps, n_dof = aligned_trajs.shape
    
    xdot_trajs = aligned_trajs[:,:,6:12] / DT # TODO have 1/DT somewhere else
    xdotdot_trajs = np.diff(xdot_trajs, axis=1) / DT
    xdotdot_trajs = np.asarray([np.r_[xdotdot_traj, xdotdot_traj[-1][None,:]] for xdotdot_traj in xdotdot_trajs])
    u_trajs = aligned_trajs[:,:,12:18]
    
    lambda_val = 0.1
    alpha = 0.01
    ms = np.zeros((t_steps,6))
    for d in range(6):
        for t in range(t_steps):
            F = []
            Xdot = []
            Xdotdot = []
            w = []
            for i in range(n_demos):
                # time window
                t_start = np.clip(t-100, 0, t_steps)
                t_end = np.clip(t+100+1, 0, t_steps)
                F.append(u_trajs[i,t_start:t_end,d:d+1])
                Xdot.append(xdot_trajs[i,t_start:t_end,d:d+1])
                Xdotdot.append(xdotdot_trajs[i,t_start:t_end,d:d+1])
                w_unnormalized = np.exp(-.5*alpha*((np.arange(t_start, t_end)-t)**2))
                w.append(w_unnormalized/w_unnormalized.sum())
            F = np.concatenate(F)
            F = np.c_[F, np.ones_like(F)]
            Xdot = np.concatenate(Xdot)
            Xdotdot = np.concatenate(Xdotdot)
            w = np.concatenate(w)
            WF = w[:,None] * F
            M = np.linalg.solve(F.T.dot(WF) + lambda_val*np.eye(2), WF.T.dot(Xdotdot))
            ms[t,d] = 1./M[0]
            if d == 0 and t == 300:
                fig = plt.figure()
                plt.scatter(Xdot, F[:,0])
                plt.draw()

    fig = plt.figure()
    for d in range(6):
        ax = fig.add_subplot(6,1,d+1)
        plt.plot(ms[:,d])
        ax.set_yscale('log')
    plt.draw()
    
class ForceAugmentedTrajectory(AugmentedTrajectory):
    def __init__(self, lr2force_traj, lr2x_traj, lr2arm_traj=None, lr2finger_traj=None, lr2ee_traj=None, lr2open_finger_traj=None, lr2close_finger_traj=None):
        super(ForceAugmentedTrajectory, self).__init__(lr2arm_traj=lr2arm_traj, lr2finger_traj=lr2finger_traj, lr2ee_traj=lr2ee_traj, 
                                                       lr2open_finger_traj=lr2open_finger_traj, lr2close_finger_traj=lr2close_finger_traj)
        for lr2traj in [lr2force_traj, lr2x_traj]:
            if lr2traj is None:
                continue
            for lr in lr2traj.keys():
                assert lr2traj[lr].shape[0] == self.n_steps
        
        self.lr2force_traj = lr2force_traj
        self.lr2x_traj = lr2x_traj
    
    @staticmethod
    def create_from_full_traj(robot, full_traj, lr2force_traj, lr2x_traj, lr2open_finger_traj=None, lr2close_finger_traj=None):

        aug_traj = AugmentedTrajectory.create_from_full_traj(robot, full_traj, lr2open_finger_traj=lr2open_finger_traj, lr2close_finger_traj=lr2close_finger_traj)