imitate_demo.py
Go to the documentation of this file.
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     #    data = {'base_pose': pose_base, 
00028     #            'robot_pose': j0_dict,
00029     #            'arm': arm_used,
00030     #            'movement_states': None}
00031     data = ut.load_pickle(data_fname)
00032     rospy.init_node('imitate')
00033     robot = pr2.PR2()
00034     #self.pr2_pub = rospy.Publisher(pr2_control_topic, PoseStamped)
00035     state = 'drive'
00036 
00037     ##Need to be localized!!
00038     ## NOT LEARNED: go into safe state.
00039     
00040     ## drive. learned locations. (might learn path/driving too?)
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     ## Need a refinement step
00048     
00049     ## Move joints to initial state. learned initial state. (maybe coordinate this with sensors?)
00050     #Put robot in the correct state
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         ## For each contact state
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                 #Publish...
00084 
00085 
00086     
00087     if state == 'manipulate':
00088         rospy.loginfo('STATE manipulate')
00089         rospy.loginfo('there are %d states' % len(data['movement_states']))
00090         ## For each contact state
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             ## send trajectory. wait until contact state changes or traj. finished executing.
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     ## rosbag implementation steps in time and also figures out how long to sleep until it needs to publish next message
00117     ## Just play pose stamped back at 10 hz
00118     ## For each contact state
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         #publish posture & cartesian poses
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 


hai_sandbox
Author(s): Hai Nguyen
autogenerated on Wed Nov 27 2013 11:46:56