00001 import roslib; roslib.load_manifest('hai_sandbox')
00002 import rospy
00003
00004 import hrl_lib.util as ut
00005 import hai_sandbox.pr2 as pr2
00006
00007 import numpy as np
00008 import time
00009 import pdb
00010 import sys
00011
00012 def dict_to_ps(d):
00013 ps = gm.PoseStamped()
00014 ps.pose.position.x = d['pose']['position']['x']
00015 ps.pose.position.y = d['pose']['position']['y']
00016 ps.pose.position.z = d['pose']['position']['z']
00017 ps.pose.orientation.x = d['pose']['orientation']['x']
00018 ps.pose.orientation.y = d['pose']['orientation']['y']
00019 ps.pose.orientation.z = d['pose']['orientation']['z']
00020 ps.pose.orientation.w = d['pose']['orientation']['w']
00021 ps.header.frame_id = d['header']['frame_id']
00022 ps.header.stamp = d['header']['stamp']
00023 return ps
00024
00025
00026 def imitate(data_fname):
00027
00028
00029
00030
00031 data = ut.load_pickle(data_fname)
00032 rospy.init_node('imitate')
00033 robot = pr2.PR2()
00034
00035 state = 'drive'
00036
00037
00038
00039
00040
00041 if state == 'drive':
00042 t, r = data['base_pose']
00043 print t
00044 r = robot.base.set_pose(t, r, '/map', block=True)
00045 rospy.loginfo('result is %s' % str(r))
00046 state = 'init_manipulation'
00047
00048
00049
00050
00051 if state == 'init_manipulation':
00052 rospy.loginfo('STATE init_manipulation')
00053 j0_dict = data['robot_pose']
00054 cpos = robot.pose()
00055 robot.left_arm.set_poses (np.column_stack([cpos['larm'], j0_dict['poses']['larm']]), np.array([0.1, 5.]), block=False)
00056 robot.right_arm.set_poses(np.column_stack([cpos['rarm'], j0_dict['poses']['rarm']]), np.array([0.1, 5.]), block=False)
00057 robot.head.set_poses(np.column_stack([cpos['head_traj'], j0_dict['poses']['head_traj']]), np.array([.01, 5.]))
00058 robot.torso.set_pose(j0_dict['poses']['torso'][0,0], block=True)
00059 state = 'manipulate'
00060
00061 if state == 'manipulate_cart':
00062 rospy.loginfo('STATE manipulate')
00063 rospy.loginfo('there are %d states' % len(data['movement_states']))
00064
00065 for state in range(len(data['movement_states'])):
00066 cur_state = data['movement_states'][state]
00067 rospy.loginfo("starting %s" % cur_state['name'])
00068 left_cart = cur_state['cartesian'][0]
00069 right_cart = cur_state['cartesian'][1]
00070 start_time = cur_state['start_time']
00071
00072 for ldict, rdict in zip(left_cart, right_cart):
00073 lps = dict_to_ps(ldict)
00074 rps = dict_to_ps(rdict)
00075 time_from_start = ((lps.header.stamp - start_time) + (rps.header.stamp - start_time))/2.0
00076 cur_time = rospy.get_rostime()
00077 ntime = cur_time + rospy.Duration(time_from_start)
00078
00079 diff_time = (ntime - rospy.get_rostime()).to_sec()
00080 if diff_time < 0:
00081 rospy.logerror('DIFF time < 0, %f' % diff_time)
00082 time.sleep(diff_time - .005)
00083
00084
00085
00086
00087 if state == 'manipulate':
00088 rospy.loginfo('STATE manipulate')
00089 rospy.loginfo('there are %d states' % len(data['movement_states']))
00090
00091 for state in range(len(data['movement_states'])):
00092 cur_state = data['movement_states'][state]
00093 rospy.loginfo("starting %s" % cur_state['name'])
00094
00095 larm, lvel, ltime, rarm, rvel, rtime = zip(*[[jdict['poses']['larm'], jdict['vels']['larm'], jdict['time'], \
00096 jdict['poses']['rarm'], jdict['vels']['rarm'], jdict['time']] \
00097 for jdict in cur_state['joint_states']])
00098
00099 larm = np.column_stack(larm)
00100 rarm = np.column_stack(rarm)
00101 lvel = np.column_stack(lvel)
00102 rvel = np.column_stack(rvel)
00103 ltime = np.array(ltime) - cur_state['start_time']
00104 rtime = np.array(rtime) - cur_state['start_time']
00105
00106
00107 robot.left_arm.set_poses(larm[:,0], np.array([2.]), block=False)
00108 robot.right_arm.set_poses(rarm[:,0], np.array([2.]), block=True)
00109
00110 robot.left_arm.set_poses(larm, ltime, vel_mat=lvel, block=False)
00111 robot.right_arm.set_poses(rarm, rtime, vel_mat=rvel, block=True)
00112
00113 rospy.loginfo("%s FINISHED" % cur_state['name'])
00114 time.sleep(5)
00115
00116
00117
00118
00119
00120 class ControllerTest:
00121 def __init__(self):
00122 pass
00123
00124 def run(self):
00125 self.robot = pr2.PR2()
00126 rospy.loginfo('switching to cartesian controllers')
00127 self.robot.controller_manager.switch(['l_cart', 'r_cart'], ['l_arm_controller', 'r_arm_controller'])
00128 rospy.on_shutdown(self.shutdown)
00129 r = rospy.Rate(1)
00130
00131 while not rospy.is_shutdown():
00132 self.robot.left_arm.set_posture(self.robot.left_arm.POSTURES['elbowupl'])
00133 self.robot.right_arm.set_posture(self.robot.right_arm.POSTURES['elbowupr'])
00134 r.sleep()
00135
00136 def shutdown(self):
00137 rospy.loginfo('switching back joint controllers')
00138 self.robot.controller_manager.switch(['l_arm_controller', 'r_arm_controller'], ['l_cart', 'r_cart'])
00139
00140
00141 if __name__ == '__main__':
00142 imitate(sys.argv[1])
00143 if False:
00144 c = ControllerTest()
00145 c.run()
00146
00147
00148
00149
00150
00151