Source code for lfd.transfer.registration_transfer

from __future__ import division

import settings
import numpy as np
from lfd.demonstration import demonstration
from lfd.environment import sim_util
from lfd.registration import registration
from lfd.transfer import transfer
from lfd.transfer import planning

[docs]class RegistrationAndTrajectoryTransferer(object): def __init__(self, registration_factory, trajectory_transferer): self.registration_factory = registration_factory self.trajectory_transferer = trajectory_transferer
[docs] def transfer(self, demo, test_scene_state, callback=None, plotting=False): """Registers demonstration scene onto the test scene and uses this registration to transfer the demonstration trajectory Args: demo: Demonstration that has the demonstration scene and the trajectory to transfer test_scene_state: SceneState of the test scene Returns: The transferred Trajectory """ raise NotImplementedError
[docs]class TwoStepRegistrationAndTrajectoryTransferer(RegistrationAndTrajectoryTransferer):
[docs] 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
[docs]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
[docs] 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